Add workspace client as optional arg to dspy.DatabricksRM (#174)

Author: kenctrl

integrations/dspy/src/databricks_dspy/retrievers/databricks_rm.py

@@ -5,6 +5,7 @@
 from typing import Any
 
 import dspy
+from databricks.sdk import WorkspaceClient
 from dspy.primitives.prediction import Prediction
 
 logger = logging.getLogger(__name__)
@@ -38,6 +39,10 @@ class DatabricksRM(dspy.Retrieve):
 
         ```python
         from databricks.vector_search.client import VectorSearchClient
+        from databricks.sdk import WorkspaceClient
+
+        # Create a Databricks workspace client
+        w = WorkspaceClient()
 
         # Create a Databricks Vector Search Endpoint
         client = VectorSearchClient()
@@ -65,6 +70,7 @@ class DatabricksRM(dspy.Retrieve):
             docs_id_column_name="id",
             text_column_name="field2",
             k=3,
+            workspace_client=w,
         )
         ```
 
@@ -90,6 +96,7 @@ def __init__(
         docs_uri_column_name: str | None = None,
         text_column_name: str = "text",
         use_with_databricks_agent_framework: bool = False,
+        workspace_client: WorkspaceClient | None = None,
     ):
         """
         Args:
@@ -122,6 +129,8 @@ def __init__(
                 containing document text to retrieve.
             use_with_databricks_agent_framework (bool): Whether to use the `DatabricksRM` in a way
                 that is compatible with the Databricks Mosaic Agent Framework.
+            workspace_client (Optional[WorkspaceClient]): The workspace client to use. If not
+                provided, a new one will be created with default credentials from the environment.
         """
         super().__init__(k=k)
         self.databricks_token = databricks_token or os.environ.get("DATABRICKS_TOKEN")
@@ -154,6 +163,32 @@ def __init__(
                     "library. Please install mlflow via `pip install mlflow`."
                 ) from None
 
+        # Use provided workspace client or create one based on credentials
+        if workspace_client:
+            self.workspace_client = workspace_client
+        elif databricks_client_secret and databricks_client_id:
+            # Use client ID and secret for authentication if they are provided
+            self.workspace_client = WorkspaceClient(
+                client_id=databricks_client_id,
+                client_secret=databricks_client_secret,
+            )
+            logger.info(
+                "Creating Databricks workspace client using service principal authentication."
+            )
+        elif databricks_token and databricks_endpoint:
+            # token-based authentication
+            self.workspace_client = WorkspaceClient(
+                host=databricks_endpoint,
+                token=databricks_token,
+            )
+            logger.info("Creating Databricks workspace client using token authentication.")
+        else:
+            # fallback to default authentication, i.e., using `~/.databrickscfg` file.
+            self.workspace_client = WorkspaceClient()
+            logger.info(
+                "Creating Databricks workspace client using credentials from `~/.databrickscfg` file."
+            )
+
     def _extract_doc_ids(self, item: dict[str, Any]) -> str:
         """Extracts the document id from a search result
 
@@ -237,10 +272,6 @@ def forward(
             query_type=query_type,
             query_text=query_text,
             query_vector=query_vector,
-            databricks_token=self.databricks_token,
-            databricks_endpoint=self.databricks_endpoint,
-            databricks_client_id=self.databricks_client_id,
-            databricks_client_secret=self.databricks_client_secret,
             filters_json=filters_json or self.filters_json,
         )
 
@@ -305,15 +336,10 @@ def _query_vector_search_index(
         query_type: str,
         query_text: str | None,
         query_vector: list[float] | None,
-        databricks_token: str | None,
-        databricks_endpoint: str | None,
-        databricks_client_id: str | None,
-        databricks_client_secret: str | None,
         filters_json: str | None,
     ) -> dict[str, Any]:
         """
         Query a Databricks Vector Search Index via the Databricks SDK.
-        Assumes that the databricks-sdk Python library is installed.
 
         Args:
             index_name (str): Name of the Databricks vector search index to query
@@ -324,49 +350,14 @@ def _query_vector_search_index(
             query_vector (Optional[list[float]]): Numeric query vector for which to find relevant
                 documents. Exactly one of query_text or query_vector must be specified.
             filters_json (Optional[str]): JSON string representing additional query filters.
-            databricks_token (str): Databricks authentication token. If not specified,
-                the token is resolved from the current environment.
-            databricks_endpoint (str): Databricks index endpoint url. If not specified,
-                the endpoint is resolved from the current environment.
-            databricks_client_id (str): Databricks service principal id. If not specified,
-                the token is resolved from the current environment (DATABRICKS_CLIENT_ID).
-            databricks_client_secret (str): Databricks service principal secret. If not specified,
-                the endpoint is resolved from the current environment (DATABRICKS_CLIENT_SECRET).
 
         Returns:
             dict[str, Any]: Parsed JSON response from the Databricks Vector Search Index query.
         """
-
-        from databricks.sdk import WorkspaceClient
-
         if (query_text, query_vector).count(None) != 1:
             raise ValueError("Exactly one of query_text or query_vector must be specified.")
 
-        if databricks_client_secret and databricks_client_id:
-            # Use client ID and secret for authentication if they are provided
-            databricks_client = WorkspaceClient(
-                client_id=databricks_client_id,
-                client_secret=databricks_client_secret,
-            )
-            logger.info(
-                "Creating Databricks workspace client using service principal authentication."
-            )
-
-        elif databricks_token and databricks_endpoint:
-            # token-based authentication
-            databricks_client = WorkspaceClient(
-                host=databricks_endpoint,
-                token=databricks_token,
-            )
-            logger.info("Creating Databricks workspace client using token authentication.")
-        else:
-            # fallback to default authentication, i.e., using `~/.databrickscfg` file.
-            databricks_client = WorkspaceClient()
-            logger.info(
-                "Creating Databricks workspace client using credentials from `~/.databrickscfg` file."
-            )
-
-        return databricks_client.vector_search_indexes.query_index(
+        return self.workspace_client.vector_search_indexes.query_index(
             index_name=index_name,
             query_type=query_type,
             query_text=query_text,

integrations/dspy/tests/unit_tests/retrievers/test_databricks_rm.py

@@ -49,7 +49,7 @@ def mock_vector_search_response_with_uri():
     }
 
 
-@patch("databricks.sdk.WorkspaceClient")
+@patch("databricks_dspy.retrievers.databricks_rm.WorkspaceClient")
 def test_databricks_rm_forward_string_query(mock_workspace_client, mock_vector_search_response):
     """Test forward method with string query and ANN search."""
     mock_client = MagicMock()
@@ -84,7 +84,7 @@ def test_databricks_rm_forward_string_query(mock_workspace_client, mock_vector_s
     assert result.doc_ids[0] == "doc1"
 
 
-@patch("databricks.sdk.WorkspaceClient")
+@patch("databricks_dspy.retrievers.databricks_rm.WorkspaceClient")
 def test_databricks_rm_forward_vector_query(mock_workspace_client, mock_vector_search_response):
     """Test forward method with vector query and HYBRID search."""
     mock_client = MagicMock()
@@ -107,7 +107,7 @@ def test_databricks_rm_forward_vector_query(mock_workspace_client, mock_vector_s
     assert set(call_args["columns"]) == {"id", "text"}
 
 
-@patch("databricks.sdk.WorkspaceClient")
+@patch("databricks_dspy.retrievers.databricks_rm.WorkspaceClient")
 def test_databricks_rm_agent_framework_format(
     mock_workspace_client, mock_vector_search_response_with_uri
 ):
@@ -138,8 +138,12 @@ def test_databricks_rm_agent_framework_format(
     assert doc["type"] == "Document"
 
 
-def test_databricks_rm_initialization():
+@patch("databricks_dspy.retrievers.databricks_rm.WorkspaceClient")
+def test_databricks_rm_initialization(mock_workspace_client):
     """Test initialization with token authentication."""
+    mock_client = MagicMock()
+    mock_workspace_client.return_value = mock_client
+
     rm = DatabricksRM(
         databricks_index_name="test_index",
         databricks_endpoint="https://test.databricks.com",
@@ -155,8 +159,75 @@ def test_databricks_rm_initialization():
     assert rm.text_column_name == "text"
     assert not rm.use_with_databricks_agent_framework
 
+    # Verify WorkspaceClient was created with token auth
+    mock_workspace_client.assert_called_once_with(
+        host="https://test.databricks.com",
+        token="test_token",
+    )
+    # Workspace client should be set
+    assert rm.workspace_client == mock_client
+
+
+def test_databricks_rm_initialization_with_custom_workspace_client():
+    """Test initialization with custom workspace_client."""
+    mock_workspace_client = MagicMock()
+
+    rm = DatabricksRM(
+        databricks_index_name="test_index",
+        workspace_client=mock_workspace_client,
+        k=5,
+    )
+
+    assert rm.databricks_index_name == "test_index"
+    assert rm.workspace_client == mock_workspace_client
+    assert rm.k == 5
+    assert rm.docs_id_column_name == "id"
+    assert rm.text_column_name == "text"
+    assert not rm.use_with_databricks_agent_framework
+
+
+def test_databricks_rm_query_with_custom_workspace_client():
+    """Test that custom workspace_client is used for queries."""
+    mock_workspace_client = MagicMock()
+
+    mock_response = {
+        "result": {
+            "data_array": [
+                ["doc1", "This is document 1", 0.95],
+            ]
+        },
+        "manifest": {
+            "columns": [
+                {"name": "id"},
+                {"name": "text"},
+                {"name": "score"},
+            ]
+        },
+    }
+    mock_workspace_client.vector_search_indexes.query_index.return_value.as_dict.return_value = (
+        mock_response
+    )
 
-@patch("databricks.sdk.WorkspaceClient")
+    rm = DatabricksRM(
+        databricks_index_name="test_index",
+        workspace_client=mock_workspace_client,
+    )
+
+    result = rm("test query")
+
+    # Verify that the custom workspace_client was used for the query
+    mock_workspace_client.vector_search_indexes.query_index.assert_called_once()
+    call_args = mock_workspace_client.vector_search_indexes.query_index.call_args[1]
+    assert call_args["index_name"] == "test_index"
+    assert call_args["query_text"] == "test query"
+
+    # Verify results
+    assert len(result.docs) == 1
+    assert result.docs[0] == "This is document 1"
+    assert result.doc_ids[0] == "doc1"
+
+
+@patch("databricks_dspy.retrievers.databricks_rm.WorkspaceClient")
 def test_databricks_rm_service_principal_auth(mock_workspace_client, mock_vector_search_response):
     """Test querying with service principal authentication."""
     mock_client = MagicMock()
@@ -180,15 +251,19 @@ def test_databricks_rm_service_principal_auth(mock_workspace_client, mock_vector
     )
 
 
-def test_databricks_rm_invalid_query_type():
+@patch("databricks_dspy.retrievers.databricks_rm.WorkspaceClient")
+def test_databricks_rm_invalid_query_type(mock_workspace_client):
     """Test forward method with invalid query type."""
+    mock_client = MagicMock()
+    mock_workspace_client.return_value = mock_client
+
     rm = DatabricksRM(databricks_index_name="test_index")
 
     with pytest.raises(ValueError, match="Invalid query_type: INVALID"):
         rm("test query", query_type="INVALID")
 
 
-@patch("databricks.sdk.WorkspaceClient")
+@patch("databricks_dspy.retrievers.databricks_rm.WorkspaceClient")
 def test_databricks_rm_missing_column_error(mock_workspace_client):
     """Test error when ID column is missing from index."""
     mock_client = MagicMock()
@@ -210,7 +285,7 @@ def test_databricks_rm_missing_column_error(mock_workspace_client):
         rm("test query")
 
 
-@patch("databricks.sdk.WorkspaceClient")
+@patch("databricks_dspy.retrievers.databricks_rm.WorkspaceClient")
 def test_databricks_rm_result_sorting(mock_workspace_client):
     """Test that results are sorted by score in descending order."""
     mock_client = MagicMock()
@@ -236,3 +311,37 @@ def test_databricks_rm_result_sorting(mock_workspace_client):
     # Should be sorted by score (highest first)
     assert result.doc_ids == ["doc1", "doc3", "doc2"]
     assert result.docs == ["Document 1", "Document 3", "Document 2"]
+
+
+@patch("databricks_dspy.retrievers.databricks_rm.WorkspaceClient")
+def test_databricks_rm_query_with_invalid_credentials(mock_workspace_client):
+    """Test that authentication errors are raised during query execution."""
+    mock_client = MagicMock()
+    mock_workspace_client.return_value = mock_client
+    # Simulate authentication error when querying the index
+    mock_client.vector_search_indexes.query_index.side_effect = Exception("Authentication failed")
+
+    rm = DatabricksRM(
+        databricks_index_name="test_index",
+        databricks_token="invalid_token",
+        databricks_endpoint="https://test.databricks.com",
+    )
+
+    # Error occurs when trying to query, not during initialization
+    with pytest.raises(Exception, match="Authentication failed"):
+        rm("test query")
+
+
+@patch("databricks_dspy.retrievers.databricks_rm.WorkspaceClient")
+def test_databricks_rm_fallback_to_default_auth(mock_workspace_client):
+    """Test fallback to default authentication when no credentials provided."""
+    mock_client = MagicMock()
+    mock_workspace_client.return_value = mock_client
+
+    rm = DatabricksRM(databricks_index_name="test_index")
+
+    assert rm.databricks_index_name == "test_index"
+
+    # Verify WorkspaceClient was created with no auth params (default auth)
+    mock_workspace_client.assert_called_once_with()
+    assert rm.workspace_client == mock_client