Update based on comments

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

src/databricks_ai_bridge/__init__.py

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

src/databricks_ai_bridge/external_tool_request.py → 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.

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

tests/databricks_ai_bridge/test_external_tool_request.py → tests/databricks_ai_bridge/test_external_tools.py

@@ -1,15 +1,17 @@
 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()
     mock_workspace_config.host = "https://mock-host"
-    mock_workspace_config._header_factory.return_value = {"Authorization": "Bearer mock-token"}
+    mock_workspace_config._header_factory.return_value = {
+        "Authorization": "Bearer mock-token"
+    }
     mock_workspace_client.return_value.config = mock_workspace_config
 
     # Mock the POST request
@@ -49,13 +51,15 @@ 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()
     mock_workspace_config.host = "https://mock-host"
-    mock_workspace_config._header_factory.return_value = {"Authorization": "Bearer mock-token"}
+    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