Adding a python http_request wrapper to create external tools

Signed-off-by: Sunish Sheth <[email protected]>

src/databricks_ai_bridge/__init__.py

@@ -0,0 +1,3 @@
+from databricks_ai_bridge.external_tool_request import http_request
+
+__all__ = ["http_request"]

src/databricks_ai_bridge/external_tool_request.py

@@ -0,0 +1,54 @@
+import json as js
+from typing import Any, Dict, Optional
+
+import requests
+from databricks.sdk import WorkspaceClient
+
+
+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 an external function through the Databricks Workspace.
+
+    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)

tests/databricks_ai_bridge/test_external_tool_request.py

@@ -0,0 +1,93 @@
+from unittest.mock import MagicMock, patch
+
+from databricks_ai_bridge import http_request
+
+
+@patch("databricks_ai_bridge.external_tool_request.WorkspaceClient")
+@patch("databricks_ai_bridge.external_tool_request.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_tool_request.WorkspaceClient")
+@patch("databricks_ai_bridge.external_tool_request.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,
+        },
+    )