Fix

Signed-off-by: Sid Murching <[email protected]>

Author: smurching

NEXT_CHANGELOG.md

@@ -4,6 +4,8 @@
 
 ### New Features and Improvements
 
+* Added support for passing additional kwargs to `WorkspaceClient().serving_endpoints.get_open_ai_client()` ([#XXX](https://github.com/databricks/databricks-sdk-py/pull/XXX)). Users can now pass standard OpenAI client parameters like `timeout` and `max_retries` when creating an OpenAI client for Databricks Model Serving.
+
 ### Bug Fixes
 
 ### Documentation

databricks/sdk/mixins/open_ai_client.py

@@ -3,9 +3,7 @@
 
 from requests import Response
 
-from databricks.sdk.service.serving import (ExternalFunctionRequestHttpMethod,
-                                            HttpRequestResponse,
-                                            ServingEndpointsAPI)
+from databricks.sdk.service.serving import ExternalFunctionRequestHttpMethod, HttpRequestResponse, ServingEndpointsAPI
 
 
 class ServingEndpointsExt(ServingEndpointsAPI):
@@ -33,31 +31,31 @@ def auth_flow(self, request: httpx.Request) -> httpx.Request:
 
     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, 
+            ...     timeout=30.0,
             ...     max_retries=5
             ... )
         """
@@ -74,7 +72,7 @@ def get_open_ai_client(self, **kwargs):
             "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)
 

tests/test_open_ai_mixin.py

@@ -25,7 +25,7 @@ def test_open_ai_client_with_custom_params(monkeypatch):
     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)
 
@@ -41,12 +41,10 @@ def test_open_ai_client_with_additional_kwargs(monkeypatch):
     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"}
+        timeout=60.0, max_retries=5, default_headers={"Custom-Header": "test-value"}
     )
 
     assert client.base_url == "https://test_host/serving-endpoints/"