Adding a python http_request wrapper to create external tools

src/databricks_ai_bridge/external_tools.py

@@ -0,0 +1,57 @@
+import json as js
+from typing import Any, Dict, Optional
+
+import requests
+from databricks.sdk import WorkspaceClient
+
+from databricks_ai_bridge.utils.annotations import experimental
+
+
+@experimental
+def http_request(
+    conn: str,
+    method: str,
+    path: str,
+    *,
+    json: Optional[Any] = None,
+    headers: Optional[Dict[str, str]] = None,
+    params: Optional[Dict[str, Any]] = None,
+) -> requests.Response:
+    """
+    Makes an HTTP request to a remote API using authentication from a Unity Catalog HTTP connection.
+
+    Args:
+        conn (str): The connection name to use. This is required to identify the external connection.
+        method (str): The HTTP method to use (e.g., "GET", "POST"). This is required.
+        path (str): The relative path for the API endpoint. This is required.
+        json (Optional[Any]): JSON payload for the request.
+        headers (Optional[Dict[str, str]]): Additional headers for the request.
+                If not provided, only auth headers from connections would be passed.
+        params (Optional[Dict[str, Any]]): Query parameters for the request.
+
+    Returns:
+        requests.Response: The HTTP response from the external function.
+
+    Example Usage:
+        response = http_request(
+            conn="my_connection",
+            method="POST",
+            path="/api/v1/resource",
+            json={"key": "value"},
+            headers={"extra_header_key": "extra_header_value"},
+            params={"query": "example"}
+        )
+    """
+    workspaceConfig = WorkspaceClient().config
+    url = f"{workspaceConfig.host}/external-functions"
+    request_headers = workspaceConfig._header_factory()
+    payload = {
+        "connection_name": conn,
+        "method": method,
+        "path": path,
+        "json": js.dumps(json),
+        "header": headers,
+        "params": params,
+    }
+
+    return requests.post(url, headers=request_headers, json=payload)

src/databricks_ai_bridge/utils/annotations.py

@@ -0,0 +1,67 @@
+# This code is copied from MLflow: https://github.com/mlflow/mlflow/blob/v2.19.0/mlflow/utils/annotations.py#L31
+
+import inspect
+import re
+import types
+from typing import Any, Callable, TypeVar, Union
+
+C = TypeVar("C", bound=Callable[..., Any])
+
+
+def _get_min_indent_of_docstring(docstring_str: str) -> str:
+    """
+    Get the minimum indentation string of a docstring, based on the assumption
+    that the closing triple quote for multiline comments must be on a new line.
+    Note that based on ruff rule D209, the closing triple quote for multiline
+    comments must be on a new line.
+
+    Args:
+        docstring_str: string with docstring
+
+    Returns:
+        Whitespace corresponding to the indent of a docstring.
+    """
+
+    if not docstring_str or "\n" not in docstring_str:
+        return ""
+
+    return re.match(r"^\s*", docstring_str.rsplit("\n", 1)[-1]).group()
+
+
+def experimental(api_or_type: Union[C, str]) -> C:
+    """Decorator / decorator creator for marking APIs experimental in the docstring.
+
+    Args:
+        api_or_type: An API to mark, or an API typestring for which to generate a decorator.
+
+    Returns:
+        Decorated API (if a ``api_or_type`` is an API) or a function that decorates
+        the specified API type (if ``api_or_type`` is a typestring).
+    """
+    if isinstance(api_or_type, str):
+
+        def f(api: C) -> C:
+            return _experimental(api=api, api_type=api_or_type)
+
+        return f
+    elif inspect.isclass(api_or_type):
+        return _experimental(api=api_or_type, api_type="class")
+    elif inspect.isfunction(api_or_type):
+        return _experimental(api=api_or_type, api_type="function")
+    elif isinstance(api_or_type, (property, types.MethodType)):
+        return _experimental(api=api_or_type, api_type="property")
+    else:
+        return _experimental(api=api_or_type, api_type=str(type(api_or_type)))
+
+
+def _experimental(api: C, api_type: str) -> C:
+    indent = _get_min_indent_of_docstring(api.__doc__)
+    notice = (
+        indent + f".. Note:: Experimental: This {api_type} may change or "
+        "be removed in a future release without warning.\n\n"
+    )
+    if api_type == "property":
+        api.__doc__ = api.__doc__ + "\n\n" + notice if api.__doc__ else notice
+    else:
+        api.__doc__ = notice + api.__doc__ if api.__doc__ else notice
+    return api

tests/databricks_ai_bridge/test_external_tools.py

@@ -0,0 +1,89 @@
+from unittest.mock import MagicMock, patch
+
+from databricks_ai_bridge.external_tools import http_request
+
+
+@patch("databricks_ai_bridge.external_tools.WorkspaceClient")
+@patch("databricks_ai_bridge.external_tools.requests.post")
+def test_http_request_success(mock_post, mock_workspace_client):
+    # Mock the WorkspaceClient config
+    mock_workspace_config = MagicMock()
+    mock_workspace_config.host = "https://mock-host"
+    mock_workspace_config._header_factory.return_value = {"Authorization": "Bearer mock-token"}
+    mock_workspace_client.return_value.config = mock_workspace_config
+
+    # Mock the POST request
+    mock_response = MagicMock()
+    mock_response.status_code = 200
+    mock_response.json.return_value = {"success": True}
+    mock_post.return_value = mock_response
+
+    # Call the function
+    response = http_request(
+        conn="mock_connection",
+        method="POST",
+        path="/mock-path",
+        json={"key": "value"},
+        headers={"Custom-Header": "HeaderValue"},
+        params={"query": "test"},
+    )
+
+    # Assertions
+    assert response.status_code == 200
+    assert response.json() == {"success": True}
+    mock_post.assert_called_once_with(
+        "https://mock-host/external-functions",
+        headers={
+            "Authorization": "Bearer mock-token",
+        },
+        json={
+            "connection_name": "mock_connection",
+            "method": "POST",
+            "path": "/mock-path",
+            "json": '{"key": "value"}',
+            "header": {
+                "Custom-Header": "HeaderValue",
+            },
+            "params": {"query": "test"},
+        },
+    )
+
+
+@patch("databricks_ai_bridge.external_tools.WorkspaceClient")
+@patch("databricks_ai_bridge.external_tools.requests.post")
+def test_http_request_error_response(mock_post, mock_workspace_client):
+    # Mock the WorkspaceClient config
+    mock_workspace_config = MagicMock()
+    mock_workspace_config.host = "https://mock-host"
+    mock_workspace_config._header_factory.return_value = {"Authorization": "Bearer mock-token"}
+    mock_workspace_client.return_value.config = mock_workspace_config
+
+    # Mock the POST request to return an error
+    mock_response = MagicMock()
+    mock_response.status_code = 400
+    mock_response.json.return_value = {"error": "Bad Request"}
+    mock_post.return_value = mock_response
+
+    # Call the function
+    response = http_request(
+        conn="mock_connection",
+        method="POST",
+        path="/mock-path",
+        json={"key": "value"},
+    )
+
+    # Assertions
+    assert response.status_code == 400
+    assert response.json() == {"error": "Bad Request"}
+    mock_post.assert_called_once_with(
+        "https://mock-host/external-functions",
+        headers={"Authorization": "Bearer mock-token"},
+        json={
+            "connection_name": "mock_connection",
+            "method": "POST",
+            "path": "/mock-path",
+            "json": '{"key": "value"}',
+            "header": None,
+            "params": None,
+        },
+    )