Add support for passing additional kwargs when instantiating an OpenAI client for Databricks model serving

Signed-off-by: Sid Murching <sid.murching@databricks.com>

Author: smurching

databricks/sdk/mixins/open_ai_client.py

@@ -31,19 +31,54 @@ def auth_flow(self, request: httpx.Request) -> httpx.Request:
         http_client = httpx.Client(auth=databricks_token_auth)
         return http_client
 
-    def get_open_ai_client(self):
+    def get_open_ai_client(self, **kwargs):
+        """Create an OpenAI client configured for Databricks Model Serving.
+        
+        Returns an OpenAI client instance that is pre-configured to send requests to
+        Databricks Model Serving endpoints. The client uses Databricks authentication
+        to query endpoints within the workspace associated with the current WorkspaceClient
+        instance.
+        
+        Args:
+            **kwargs: Additional parameters to pass to the OpenAI client constructor.
+                Common parameters include:
+                - timeout (float): Request timeout in seconds (e.g., 30.0)
+                - max_retries (int): Maximum number of retries for failed requests (e.g., 3)
+
+                Any parameter accepted by the OpenAI client constructor can be passed here.
+                
+        Returns:
+            OpenAI: An OpenAI client instance configured for Databricks Model Serving.
+            
+        Raises:
+            ImportError: If the OpenAI library is not installed.
+            
+        Example:
+            >>> client = workspace_client.serving_endpoints.get_open_ai_client()
+            >>> # With custom timeout and retries
+            >>> client = workspace_client.serving_endpoints.get_open_ai_client(
+            ...     timeout=30.0, 
+            ...     max_retries=5
+            ... )
+        """
         try:
             from openai import OpenAI
         except Exception:
             raise ImportError(
                 "Open AI is not installed. Please install the Databricks SDK with the following command `pip install databricks-sdk[openai]`"
             )
 
-        return OpenAI(
-            base_url=self._api._cfg.host + "/serving-endpoints",
-            api_key="no-token",  # Passing in a placeholder to pass validations, this will not be used
-            http_client=self._get_authorized_http_client(),
-        )
+        # Default parameters that are required for Databricks integration
+        client_params = {
+            "base_url": self._api._cfg.host + "/serving-endpoints",
+            "api_key": "no-token",  # Passing in a placeholder to pass validations, this will not be used
+            "http_client": self._get_authorized_http_client(),
+        }
+        
+        # Update with any additional parameters passed by the user
+        client_params.update(kwargs)
+
+        return OpenAI(**client_params)
 
     def get_langchain_chat_open_ai_client(self, model):
         try:

tests/test_open_ai_mixin.py

@@ -19,6 +19,44 @@ def test_open_ai_client(monkeypatch):
     assert client.api_key == "no-token"
 
 
+def test_open_ai_client_with_custom_params(monkeypatch):
+    from databricks.sdk import WorkspaceClient
+
+    monkeypatch.setenv("DATABRICKS_HOST", "test_host")
+    monkeypatch.setenv("DATABRICKS_TOKEN", "test_token")
+    w = WorkspaceClient(config=Config())
+    
+    # Test with timeout and max_retries parameters
+    client = w.serving_endpoints.get_open_ai_client(timeout=30.0, max_retries=3)
+
+    assert client.base_url == "https://test_host/serving-endpoints/"
+    assert client.api_key == "no-token"
+    assert client.timeout == 30.0
+    assert client.max_retries == 3
+
+
+def test_open_ai_client_with_additional_kwargs(monkeypatch):
+    from databricks.sdk import WorkspaceClient
+
+    monkeypatch.setenv("DATABRICKS_HOST", "test_host")
+    monkeypatch.setenv("DATABRICKS_TOKEN", "test_token")
+    w = WorkspaceClient(config=Config())
+    
+    # Test with additional kwargs that OpenAI client might accept
+    client = w.serving_endpoints.get_open_ai_client(
+        timeout=60.0,
+        max_retries=5,
+        default_headers={"Custom-Header": "test-value"}
+    )
+
+    assert client.base_url == "https://test_host/serving-endpoints/"
+    assert client.api_key == "no-token"
+    assert client.timeout == 60.0
+    assert client.max_retries == 5
+    assert "Custom-Header" in client.default_headers
+    assert client.default_headers["Custom-Header"] == "test-value"
+
+
 @pytest.mark.skipif(sys.version_info < (3, 8), reason="Requires Python > 3.7")
 def test_langchain_open_ai_client(monkeypatch):
     from databricks.sdk import WorkspaceClient