From d887b797c6a4ce2b6ec4cd5585f15b0040e5eebc Mon Sep 17 00:00:00 2001 From: Sunish Sheth Date: Wed, 11 Dec 2024 10:57:07 -0800 Subject: [PATCH] Adding a python http_request wrapper to create external tools Signed-off-by: Sunish Sheth --- src/databricks_ai_bridge/__init__.py | 3 + .../external_tool_request.py | 57 ++++++++++++ src/databricks_ai_bridge/utils/annotations.py | 65 ++++++++++++++ .../test_external_tool_request.py | 89 +++++++++++++++++++ 4 files changed, 214 insertions(+) create mode 100644 src/databricks_ai_bridge/external_tool_request.py create mode 100644 src/databricks_ai_bridge/utils/annotations.py create mode 100644 tests/databricks_ai_bridge/test_external_tool_request.py diff --git a/src/databricks_ai_bridge/__init__.py b/src/databricks_ai_bridge/__init__.py index e69de29..25165ec 100644 --- a/src/databricks_ai_bridge/__init__.py +++ b/src/databricks_ai_bridge/__init__.py @@ -0,0 +1,3 @@ +from databricks_ai_bridge.external_tool_request import http_request + +__all__ = ["http_request"] diff --git a/src/databricks_ai_bridge/external_tool_request.py b/src/databricks_ai_bridge/external_tool_request.py new file mode 100644 index 0000000..b5d5144 --- /dev/null +++ b/src/databricks_ai_bridge/external_tool_request.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 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) diff --git a/src/databricks_ai_bridge/utils/annotations.py b/src/databricks_ai_bridge/utils/annotations.py new file mode 100644 index 0000000..4f2b3af --- /dev/null +++ b/src/databricks_ai_bridge/utils/annotations.py @@ -0,0 +1,65 @@ +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 diff --git a/tests/databricks_ai_bridge/test_external_tool_request.py b/tests/databricks_ai_bridge/test_external_tool_request.py new file mode 100644 index 0000000..94b9e1f --- /dev/null +++ b/tests/databricks_ai_bridge/test_external_tool_request.py @@ -0,0 +1,89 @@ +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, + }, + )