From d861da6d709e3aa8825ef5b652c3d5e7d3e6b9f3 Mon Sep 17 00:00:00 2001 From: Sunish Sheth Date: Mon, 16 Dec 2024 09:51:01 -0800 Subject: [PATCH] Update based on comments Signed-off-by: Sunish Sheth --- src/databricks_ai_bridge/__init__.py | 3 --- .../{external_tool_request.py => external_tools.py} | 2 +- src/databricks_ai_bridge/utils/annotations.py | 2 ++ ...external_tool_request.py => test_external_tools.py} | 10 +++++----- 4 files changed, 8 insertions(+), 9 deletions(-) rename src/databricks_ai_bridge/{external_tool_request.py => external_tools.py} (94%) rename tests/databricks_ai_bridge/{test_external_tool_request.py => test_external_tools.py} (89%) diff --git a/src/databricks_ai_bridge/__init__.py b/src/databricks_ai_bridge/__init__.py index 25165ec..e69de29 100644 --- a/src/databricks_ai_bridge/__init__.py +++ b/src/databricks_ai_bridge/__init__.py @@ -1,3 +0,0 @@ -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_tools.py similarity index 94% rename from src/databricks_ai_bridge/external_tool_request.py rename to src/databricks_ai_bridge/external_tools.py index b5d5144..7813b86 100644 --- a/src/databricks_ai_bridge/external_tool_request.py +++ b/src/databricks_ai_bridge/external_tools.py @@ -18,7 +18,7 @@ def http_request( params: Optional[Dict[str, Any]] = None, ) -> requests.Response: """ - Makes an HTTP request to an external function through the Databricks Workspace. + 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. diff --git a/src/databricks_ai_bridge/utils/annotations.py b/src/databricks_ai_bridge/utils/annotations.py index 4f2b3af..098073f 100644 --- a/src/databricks_ai_bridge/utils/annotations.py +++ b/src/databricks_ai_bridge/utils/annotations.py @@ -1,3 +1,5 @@ +# 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 diff --git a/tests/databricks_ai_bridge/test_external_tool_request.py b/tests/databricks_ai_bridge/test_external_tools.py similarity index 89% rename from tests/databricks_ai_bridge/test_external_tool_request.py rename to tests/databricks_ai_bridge/test_external_tools.py index 94b9e1f..6a1149d 100644 --- a/tests/databricks_ai_bridge/test_external_tool_request.py +++ b/tests/databricks_ai_bridge/test_external_tools.py @@ -1,10 +1,10 @@ from unittest.mock import MagicMock, patch -from databricks_ai_bridge import http_request +from databricks_ai_bridge.external_tools import http_request -@patch("databricks_ai_bridge.external_tool_request.WorkspaceClient") -@patch("databricks_ai_bridge.external_tool_request.requests.post") +@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() @@ -49,8 +49,8 @@ def test_http_request_success(mock_post, mock_workspace_client): ) -@patch("databricks_ai_bridge.external_tool_request.WorkspaceClient") -@patch("databricks_ai_bridge.external_tool_request.requests.post") +@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()