From e65a3dee804fe832848600eb67b181c1ddc25a07 Mon Sep 17 00:00:00 2001 From: Sid Murching Date: Tue, 12 Aug 2025 13:56:01 -0700 Subject: [PATCH 1/7] [WIP] Add timeout and retry params to ChatDatabricks Signed-off-by: Sid Murching --- .../src/databricks_langchain/chat_models.py | 12 ++- .../src/databricks_langchain/utils.py | 22 +++-- .../integration_tests/test_chat_models.py | 44 ++++++++++ .../tests/unit_tests/test_chat_models.py | 83 ++++++++++++++++++- .../langchain/tests/unit_tests/test_utils.py | 74 +++++++++++++++++ 5 files changed, 225 insertions(+), 10 deletions(-) create mode 100644 integrations/langchain/tests/unit_tests/test_utils.py diff --git a/integrations/langchain/src/databricks_langchain/chat_models.py b/integrations/langchain/src/databricks_langchain/chat_models.py index 8a09c776a..a53bcd105 100644 --- a/integrations/langchain/src/databricks_langchain/chat_models.py +++ b/integrations/langchain/src/databricks_langchain/chat_models.py @@ -72,6 +72,8 @@ class ChatDatabricks(BaseChatModel): model="databricks-claude-3-7-sonnet", temperature=0, max_tokens=500, + timeout=30.0, # Timeout in seconds + max_retries=3, # Maximum number of retries ) # Using a WorkspaceClient instance for custom authentication @@ -248,6 +250,10 @@ class GetPopulation(BaseModel): """Any extra parameters to pass to the endpoint.""" use_responses_api: bool = False """Whether to use the Responses API to format inputs and outputs.""" + timeout: Optional[float] = None + """Timeout in seconds for the HTTP request. If None, uses the default timeout.""" + max_retries: Optional[int] = None + """Maximum number of retries for failed requests. If None, uses the default retry count.""" client: Optional[object] = Field(default=None, exclude=True) #: :meta private: @property @@ -288,7 +294,11 @@ def __init__(self, **kwargs: Any): ) # Always use OpenAI client (supports both chat completions and responses API) - self.client = get_openai_client(workspace_client=self.workspace_client) + self.client = get_openai_client( + workspace_client=self.workspace_client, + timeout=self.timeout, + max_retries=self.max_retries + ) self.use_responses_api = kwargs.get("use_responses_api", False) self.extra_params = self.extra_params or {} diff --git a/integrations/langchain/src/databricks_langchain/utils.py b/integrations/langchain/src/databricks_langchain/utils.py index 1494f1016..5ed24b32b 100644 --- a/integrations/langchain/src/databricks_langchain/utils.py +++ b/integrations/langchain/src/databricks_langchain/utils.py @@ -20,23 +20,33 @@ def get_deployment_client(target_uri: str) -> Any: ) from e -def get_openai_client(workspace_client: Any = None) -> Any: +def get_openai_client(workspace_client: Any = None, timeout: Any = None, max_retries: Any = None) -> Any: """Get an OpenAI client configured for Databricks. Args: workspace_client: Optional WorkspaceClient instance to use for authentication. If not provided, creates a default WorkspaceClient. + timeout: Optional timeout in seconds for HTTP requests. + max_retries: Optional maximum number of retries for failed requests. """ try: from databricks.sdk import WorkspaceClient # If workspace_client is provided, use it directly if workspace_client is not None: - return workspace_client.serving_endpoints.get_open_ai_client() - - # Otherwise, create default workspace client - workspace_client = WorkspaceClient() - return workspace_client.serving_endpoints.get_open_ai_client() + client = workspace_client.serving_endpoints.get_open_ai_client() + else: + # Otherwise, create default workspace client + workspace_client = WorkspaceClient() + client = workspace_client.serving_endpoints.get_open_ai_client() + + # Configure timeout and max_retries if provided + if timeout is not None: + client.timeout = timeout + if max_retries is not None: + client.max_retries = max_retries + + return client except ImportError as e: raise ImportError( diff --git a/integrations/langchain/tests/integration_tests/test_chat_models.py b/integrations/langchain/tests/integration_tests/test_chat_models.py index dbda1daef..143e0f0ac 100644 --- a/integrations/langchain/tests/integration_tests/test_chat_models.py +++ b/integrations/langchain/tests/integration_tests/test_chat_models.py @@ -715,3 +715,47 @@ def test_chat_databricks_utf8_encoding(model): if hasattr(chunk, "content") and chunk.content: full_content += chunk.content assert "blåbær" in full_content.lower() + + +def test_chat_databricks_with_timeout_and_retries(): + """Test that ChatDatabricks can be initialized with timeout and max_retries parameters.""" + from unittest.mock import Mock, patch + + # Mock the OpenAI client + mock_openai_client = Mock() + mock_workspace_client = Mock() + mock_workspace_client.serving_endpoints.get_open_ai_client.return_value = mock_openai_client + + with patch("databricks.sdk.WorkspaceClient", return_value=mock_workspace_client): + # Create ChatDatabricks with timeout and max_retries + chat = ChatDatabricks( + model="databricks-meta-llama-3-3-70b-instruct", + timeout=45.0, + max_retries=3 + ) + + # Verify the parameters are set correctly + assert chat.timeout == 45.0 + assert chat.max_retries == 3 + + # Verify the client was configured with these parameters + assert chat.client == mock_openai_client + + # Test with workspace_client parameter + with patch("databricks_langchain.chat_models.get_openai_client", return_value=mock_openai_client) as mock_get_client: + chat_with_ws = ChatDatabricks( + model="databricks-meta-llama-3-3-70b-instruct", + workspace_client=mock_workspace_client, + timeout=30.0, + max_retries=2 + ) + + # Verify get_openai_client was called with all parameters + mock_get_client.assert_called_once_with( + workspace_client=mock_workspace_client, + timeout=30.0, + max_retries=2 + ) + + assert chat_with_ws.timeout == 30.0 + assert chat_with_ws.max_retries == 2 diff --git a/integrations/langchain/tests/unit_tests/test_chat_models.py b/integrations/langchain/tests/unit_tests/test_chat_models.py index 943cb9b86..eeb7300a4 100644 --- a/integrations/langchain/tests/unit_tests/test_chat_models.py +++ b/integrations/langchain/tests/unit_tests/test_chat_models.py @@ -77,7 +77,12 @@ def test_workspace_client_parameter() -> None: llm = ChatDatabricks(model="test-model", workspace_client=mock_workspace_client) assert llm.client == mock_openai_client - mock_get_client.assert_called_once_with(workspace_client=mock_workspace_client) + # Now expects timeout and max_retries parameters with None defaults + mock_get_client.assert_called_once_with( + workspace_client=mock_workspace_client, + timeout=None, + max_retries=None + ) def test_workspace_client_and_target_uri_conflict() -> None: @@ -91,6 +96,68 @@ def test_workspace_client_and_target_uri_conflict() -> None: ) +def test_timeout_and_max_retries_parameters() -> None: + """Test that timeout and max_retries parameters are properly passed to the OpenAI client.""" + from unittest.mock import Mock, patch + + mock_openai_client = Mock() + mock_openai_client.timeout = None + mock_openai_client.max_retries = None + + with patch( + "databricks_langchain.chat_models.get_openai_client", return_value=mock_openai_client + ) as mock_get_client: + # Test with timeout and max_retries + llm = ChatDatabricks( + model="test-model", + timeout=60.0, + max_retries=5 + ) + + # Verify get_openai_client was called with the correct parameters + mock_get_client.assert_called_once_with( + workspace_client=None, + timeout=60.0, + max_retries=5 + ) + + # Test that client is set + assert llm.client == mock_openai_client + assert llm.timeout == 60.0 + assert llm.max_retries == 5 + + +def test_timeout_and_max_retries_with_workspace_client() -> None: + """Test timeout and max_retries parameters work with workspace_client.""" + from unittest.mock import Mock, patch + + mock_workspace_client = Mock() + mock_openai_client = Mock() + mock_openai_client.timeout = None + mock_openai_client.max_retries = None + + with patch( + "databricks_langchain.chat_models.get_openai_client", return_value=mock_openai_client + ) as mock_get_client: + llm = ChatDatabricks( + model="test-model", + workspace_client=mock_workspace_client, + timeout=30.0, + max_retries=2 + ) + + # Verify get_openai_client was called with all parameters + mock_get_client.assert_called_once_with( + workspace_client=mock_workspace_client, + timeout=30.0, + max_retries=2 + ) + + assert llm.client == mock_openai_client + assert llm.timeout == 30.0 + assert llm.max_retries == 2 + + def test_default_workspace_client() -> None: """Test that default WorkspaceClient is created when none provided.""" from unittest.mock import Mock, patch @@ -107,7 +174,12 @@ def test_default_workspace_client() -> None: llm = ChatDatabricks(model="test-model") assert llm.client == mock_openai_client - mock_get_client.assert_called_once_with(workspace_client=None) + # Now expects timeout and max_retries parameters with None defaults + mock_get_client.assert_called_once_with( + workspace_client=None, + timeout=None, + max_retries=None + ) def test_target_uri_deprecation_warning() -> None: @@ -960,7 +1032,12 @@ def test_chat_databricks_init_sets_client(): llm = ChatDatabricks(model="test-model") - mock_get_client.assert_called_once_with(workspace_client=None) + # Now expects timeout and max_retries parameters with None defaults + mock_get_client.assert_called_once_with( + workspace_client=None, + timeout=None, + max_retries=None + ) assert llm.client == mock_client diff --git a/integrations/langchain/tests/unit_tests/test_utils.py b/integrations/langchain/tests/unit_tests/test_utils.py new file mode 100644 index 000000000..2add7d506 --- /dev/null +++ b/integrations/langchain/tests/unit_tests/test_utils.py @@ -0,0 +1,74 @@ +"""Test utilities module.""" + +from unittest.mock import Mock, patch + +import pytest + +from databricks_langchain.utils import get_openai_client + + +def test_get_openai_client_with_timeout_and_max_retries() -> None: + """Test that get_openai_client properly sets timeout and max_retries on the OpenAI client.""" + + mock_openai_client = Mock() + mock_openai_client.timeout = None + mock_openai_client.max_retries = None + + mock_workspace_client = Mock() + mock_workspace_client.serving_endpoints.get_open_ai_client.return_value = mock_openai_client + + # Test with workspace_client, timeout, and max_retries + client = get_openai_client( + workspace_client=mock_workspace_client, + timeout=45.0, + max_retries=3 + ) + + # Verify the OpenAI client was obtained + mock_workspace_client.serving_endpoints.get_open_ai_client.assert_called_once() + + # Verify timeout and max_retries were set + assert client.timeout == 45.0 + assert client.max_retries == 3 + + +def test_get_openai_client_with_default_workspace_client() -> None: + """Test get_openai_client creates default WorkspaceClient when none provided.""" + + mock_openai_client = Mock() + mock_openai_client.timeout = None + mock_openai_client.max_retries = None + + mock_workspace_client = Mock() + mock_workspace_client.serving_endpoints.get_open_ai_client.return_value = mock_openai_client + + with patch("databricks.sdk.WorkspaceClient", return_value=mock_workspace_client): + client = get_openai_client(timeout=30.0, max_retries=2) + + # Verify default WorkspaceClient was created + mock_workspace_client.serving_endpoints.get_open_ai_client.assert_called_once() + + # Verify timeout and max_retries were set + assert client.timeout == 30.0 + assert client.max_retries == 2 + + +def test_get_openai_client_without_timeout_and_retries() -> None: + """Test get_openai_client doesn't set timeout/max_retries when not provided.""" + + mock_openai_client = Mock() + # Set initial values to check they're not changed + mock_openai_client.timeout = "original_timeout" + mock_openai_client.max_retries = "original_retries" + + mock_workspace_client = Mock() + mock_workspace_client.serving_endpoints.get_open_ai_client.return_value = mock_openai_client + + client = get_openai_client(workspace_client=mock_workspace_client) + + # Verify the OpenAI client was obtained + mock_workspace_client.serving_endpoints.get_open_ai_client.assert_called_once() + + # Verify timeout and max_retries were NOT changed (None values don't override) + assert client.timeout == "original_timeout" + assert client.max_retries == "original_retries" \ No newline at end of file From 8a3d3bb00777087455a2179876073158892b4e1f Mon Sep 17 00:00:00 2001 From: Sid Murching Date: Tue, 12 Aug 2025 16:03:57 -0700 Subject: [PATCH 2/7] Switch to passing kwargs directly, dependent on Databricks Python SDK release Signed-off-by: Sid Murching --- .../src/databricks_langchain/chat_models.py | 10 ++++- .../src/databricks_langchain/utils.py | 18 +++----- .../tests/unit_tests/test_chat_models.py | 18 +++----- .../langchain/tests/unit_tests/test_utils.py | 44 +++++++++---------- 4 files changed, 39 insertions(+), 51 deletions(-) diff --git a/integrations/langchain/src/databricks_langchain/chat_models.py b/integrations/langchain/src/databricks_langchain/chat_models.py index a53bcd105..c0c06ae71 100644 --- a/integrations/langchain/src/databricks_langchain/chat_models.py +++ b/integrations/langchain/src/databricks_langchain/chat_models.py @@ -294,10 +294,16 @@ def __init__(self, **kwargs: Any): ) # Always use OpenAI client (supports both chat completions and responses API) + # Prepare kwargs for the SDK call + openai_kwargs = {} + if self.timeout is not None: + openai_kwargs["timeout"] = self.timeout + if self.max_retries is not None: + openai_kwargs["max_retries"] = self.max_retries + self.client = get_openai_client( workspace_client=self.workspace_client, - timeout=self.timeout, - max_retries=self.max_retries + **openai_kwargs ) self.use_responses_api = kwargs.get("use_responses_api", False) diff --git a/integrations/langchain/src/databricks_langchain/utils.py b/integrations/langchain/src/databricks_langchain/utils.py index 5ed24b32b..fe35891ed 100644 --- a/integrations/langchain/src/databricks_langchain/utils.py +++ b/integrations/langchain/src/databricks_langchain/utils.py @@ -20,33 +20,25 @@ def get_deployment_client(target_uri: str) -> Any: ) from e -def get_openai_client(workspace_client: Any = None, timeout: Any = None, max_retries: Any = None) -> Any: +def get_openai_client(workspace_client: Any = None, **kwargs) -> Any: """Get an OpenAI client configured for Databricks. Args: workspace_client: Optional WorkspaceClient instance to use for authentication. If not provided, creates a default WorkspaceClient. - timeout: Optional timeout in seconds for HTTP requests. - max_retries: Optional maximum number of retries for failed requests. + **kwargs: Additional keyword arguments to pass to get_open_ai_client(), + such as timeout and max_retries. """ try: from databricks.sdk import WorkspaceClient # If workspace_client is provided, use it directly if workspace_client is not None: - client = workspace_client.serving_endpoints.get_open_ai_client() + return workspace_client.serving_endpoints.get_open_ai_client(**kwargs) else: # Otherwise, create default workspace client workspace_client = WorkspaceClient() - client = workspace_client.serving_endpoints.get_open_ai_client() - - # Configure timeout and max_retries if provided - if timeout is not None: - client.timeout = timeout - if max_retries is not None: - client.max_retries = max_retries - - return client + return workspace_client.serving_endpoints.get_open_ai_client(**kwargs) except ImportError as e: raise ImportError( diff --git a/integrations/langchain/tests/unit_tests/test_chat_models.py b/integrations/langchain/tests/unit_tests/test_chat_models.py index eeb7300a4..d2c07eeb5 100644 --- a/integrations/langchain/tests/unit_tests/test_chat_models.py +++ b/integrations/langchain/tests/unit_tests/test_chat_models.py @@ -77,11 +77,9 @@ def test_workspace_client_parameter() -> None: llm = ChatDatabricks(model="test-model", workspace_client=mock_workspace_client) assert llm.client == mock_openai_client - # Now expects timeout and max_retries parameters with None defaults + # Now expects no additional kwargs when timeout/max_retries are None mock_get_client.assert_called_once_with( - workspace_client=mock_workspace_client, - timeout=None, - max_retries=None + workspace_client=mock_workspace_client ) @@ -174,11 +172,9 @@ def test_default_workspace_client() -> None: llm = ChatDatabricks(model="test-model") assert llm.client == mock_openai_client - # Now expects timeout and max_retries parameters with None defaults + # Now expects no additional kwargs when timeout/max_retries are None mock_get_client.assert_called_once_with( - workspace_client=None, - timeout=None, - max_retries=None + workspace_client=None ) @@ -1032,11 +1028,9 @@ def test_chat_databricks_init_sets_client(): llm = ChatDatabricks(model="test-model") - # Now expects timeout and max_retries parameters with None defaults + # Now expects no additional kwargs when timeout/max_retries are None mock_get_client.assert_called_once_with( - workspace_client=None, - timeout=None, - max_retries=None + workspace_client=None ) assert llm.client == mock_client diff --git a/integrations/langchain/tests/unit_tests/test_utils.py b/integrations/langchain/tests/unit_tests/test_utils.py index 2add7d506..aad3e38cb 100644 --- a/integrations/langchain/tests/unit_tests/test_utils.py +++ b/integrations/langchain/tests/unit_tests/test_utils.py @@ -8,11 +8,9 @@ def test_get_openai_client_with_timeout_and_max_retries() -> None: - """Test that get_openai_client properly sets timeout and max_retries on the OpenAI client.""" + """Test that get_openai_client properly passes timeout and max_retries as kwargs to the SDK.""" mock_openai_client = Mock() - mock_openai_client.timeout = None - mock_openai_client.max_retries = None mock_workspace_client = Mock() mock_workspace_client.serving_endpoints.get_open_ai_client.return_value = mock_openai_client @@ -24,20 +22,20 @@ def test_get_openai_client_with_timeout_and_max_retries() -> None: max_retries=3 ) - # Verify the OpenAI client was obtained - mock_workspace_client.serving_endpoints.get_open_ai_client.assert_called_once() + # Verify the OpenAI client was obtained with the correct kwargs + mock_workspace_client.serving_endpoints.get_open_ai_client.assert_called_once_with( + timeout=45.0, + max_retries=3 + ) - # Verify timeout and max_retries were set - assert client.timeout == 45.0 - assert client.max_retries == 3 + # Verify the client is returned + assert client == mock_openai_client def test_get_openai_client_with_default_workspace_client() -> None: """Test get_openai_client creates default WorkspaceClient when none provided.""" mock_openai_client = Mock() - mock_openai_client.timeout = None - mock_openai_client.max_retries = None mock_workspace_client = Mock() mock_workspace_client.serving_endpoints.get_open_ai_client.return_value = mock_openai_client @@ -45,30 +43,28 @@ def test_get_openai_client_with_default_workspace_client() -> None: with patch("databricks.sdk.WorkspaceClient", return_value=mock_workspace_client): client = get_openai_client(timeout=30.0, max_retries=2) - # Verify default WorkspaceClient was created - mock_workspace_client.serving_endpoints.get_open_ai_client.assert_called_once() + # Verify default WorkspaceClient was created and kwargs were passed + mock_workspace_client.serving_endpoints.get_open_ai_client.assert_called_once_with( + timeout=30.0, + max_retries=2 + ) - # Verify timeout and max_retries were set - assert client.timeout == 30.0 - assert client.max_retries == 2 + # Verify the client is returned + assert client == mock_openai_client def test_get_openai_client_without_timeout_and_retries() -> None: - """Test get_openai_client doesn't set timeout/max_retries when not provided.""" + """Test get_openai_client doesn't pass kwargs when not provided.""" mock_openai_client = Mock() - # Set initial values to check they're not changed - mock_openai_client.timeout = "original_timeout" - mock_openai_client.max_retries = "original_retries" mock_workspace_client = Mock() mock_workspace_client.serving_endpoints.get_open_ai_client.return_value = mock_openai_client client = get_openai_client(workspace_client=mock_workspace_client) - # Verify the OpenAI client was obtained - mock_workspace_client.serving_endpoints.get_open_ai_client.assert_called_once() + # Verify the OpenAI client was obtained without kwargs + mock_workspace_client.serving_endpoints.get_open_ai_client.assert_called_once_with() - # Verify timeout and max_retries were NOT changed (None values don't override) - assert client.timeout == "original_timeout" - assert client.max_retries == "original_retries" \ No newline at end of file + # Verify the client is returned + assert client == mock_openai_client \ No newline at end of file From ce97f2b72c2aa7a81fc6e117041d30cdf5dc4398 Mon Sep 17 00:00:00 2001 From: Sid Murching Date: Tue, 12 Aug 2025 16:13:23 -0700 Subject: [PATCH 3/7] Bump min SDK version - need to finalize after SDK release Signed-off-by: Sid Murching --- integrations/langchain/pyproject.toml | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/integrations/langchain/pyproject.toml b/integrations/langchain/pyproject.toml index cf0e665cd..9933515ca 100644 --- a/integrations/langchain/pyproject.toml +++ b/integrations/langchain/pyproject.toml @@ -17,13 +17,14 @@ dependencies = [ "unitycatalog-langchain[databricks]>=0.2.0", "databricks-connect>=16.1.1,<16.4", "openai>=1.97.1", + "databricks-sdk>=0.63.0", ] [project.optional-dependencies] dev = [ "pytest", "typing_extensions", - "databricks-sdk>=0.34.0", + "databricks-sdk>=0.63.0", "ruff==0.6.4", ] From b3a52248c7b9316dc0f6b7a50818eaf851294ccd Mon Sep 17 00:00:00 2001 From: Sid Murching Date: Tue, 2 Sep 2025 20:38:03 -0700 Subject: [PATCH 4/7] fix lint Signed-off-by: Sid Murching --- integrations/langchain/tests/unit_tests/test_utils.py | 2 -- 1 file changed, 2 deletions(-) diff --git a/integrations/langchain/tests/unit_tests/test_utils.py b/integrations/langchain/tests/unit_tests/test_utils.py index aad3e38cb..d2c4c7771 100644 --- a/integrations/langchain/tests/unit_tests/test_utils.py +++ b/integrations/langchain/tests/unit_tests/test_utils.py @@ -2,8 +2,6 @@ from unittest.mock import Mock, patch -import pytest - from databricks_langchain.utils import get_openai_client From e274a951b834f29927a7cbd653a66e57e31d59eb Mon Sep 17 00:00:00 2001 From: Sid Murching Date: Tue, 2 Sep 2025 20:41:43 -0700 Subject: [PATCH 5/7] Fix lint Signed-off-by: Sid Murching --- .../src/databricks_langchain/chat_models.py | 7 +-- .../integration_tests/test_chat_models.py | 28 +++++------ .../tests/unit_tests/test_chat_models.py | 49 ++++++------------- .../langchain/tests/unit_tests/test_utils.py | 44 +++++++---------- 4 files changed, 49 insertions(+), 79 deletions(-) diff --git a/integrations/langchain/src/databricks_langchain/chat_models.py b/integrations/langchain/src/databricks_langchain/chat_models.py index c0c06ae71..896e20dfe 100644 --- a/integrations/langchain/src/databricks_langchain/chat_models.py +++ b/integrations/langchain/src/databricks_langchain/chat_models.py @@ -300,11 +300,8 @@ def __init__(self, **kwargs: Any): openai_kwargs["timeout"] = self.timeout if self.max_retries is not None: openai_kwargs["max_retries"] = self.max_retries - - self.client = get_openai_client( - workspace_client=self.workspace_client, - **openai_kwargs - ) + + self.client = get_openai_client(workspace_client=self.workspace_client, **openai_kwargs) self.use_responses_api = kwargs.get("use_responses_api", False) self.extra_params = self.extra_params or {} diff --git a/integrations/langchain/tests/integration_tests/test_chat_models.py b/integrations/langchain/tests/integration_tests/test_chat_models.py index 143e0f0ac..42233d376 100644 --- a/integrations/langchain/tests/integration_tests/test_chat_models.py +++ b/integrations/langchain/tests/integration_tests/test_chat_models.py @@ -720,42 +720,40 @@ def test_chat_databricks_utf8_encoding(model): def test_chat_databricks_with_timeout_and_retries(): """Test that ChatDatabricks can be initialized with timeout and max_retries parameters.""" from unittest.mock import Mock, patch - + # Mock the OpenAI client mock_openai_client = Mock() mock_workspace_client = Mock() mock_workspace_client.serving_endpoints.get_open_ai_client.return_value = mock_openai_client - + with patch("databricks.sdk.WorkspaceClient", return_value=mock_workspace_client): # Create ChatDatabricks with timeout and max_retries chat = ChatDatabricks( - model="databricks-meta-llama-3-3-70b-instruct", - timeout=45.0, - max_retries=3 + model="databricks-meta-llama-3-3-70b-instruct", timeout=45.0, max_retries=3 ) - + # Verify the parameters are set correctly assert chat.timeout == 45.0 assert chat.max_retries == 3 - + # Verify the client was configured with these parameters assert chat.client == mock_openai_client - + # Test with workspace_client parameter - with patch("databricks_langchain.chat_models.get_openai_client", return_value=mock_openai_client) as mock_get_client: + with patch( + "databricks_langchain.chat_models.get_openai_client", return_value=mock_openai_client + ) as mock_get_client: chat_with_ws = ChatDatabricks( model="databricks-meta-llama-3-3-70b-instruct", workspace_client=mock_workspace_client, timeout=30.0, - max_retries=2 + max_retries=2, ) - + # Verify get_openai_client was called with all parameters mock_get_client.assert_called_once_with( - workspace_client=mock_workspace_client, - timeout=30.0, - max_retries=2 + workspace_client=mock_workspace_client, timeout=30.0, max_retries=2 ) - + assert chat_with_ws.timeout == 30.0 assert chat_with_ws.max_retries == 2 diff --git a/integrations/langchain/tests/unit_tests/test_chat_models.py b/integrations/langchain/tests/unit_tests/test_chat_models.py index d2c07eeb5..e03203884 100644 --- a/integrations/langchain/tests/unit_tests/test_chat_models.py +++ b/integrations/langchain/tests/unit_tests/test_chat_models.py @@ -78,9 +78,7 @@ def test_workspace_client_parameter() -> None: assert llm.client == mock_openai_client # Now expects no additional kwargs when timeout/max_retries are None - mock_get_client.assert_called_once_with( - workspace_client=mock_workspace_client - ) + mock_get_client.assert_called_once_with(workspace_client=mock_workspace_client) def test_workspace_client_and_target_uri_conflict() -> None: @@ -97,28 +95,20 @@ def test_workspace_client_and_target_uri_conflict() -> None: def test_timeout_and_max_retries_parameters() -> None: """Test that timeout and max_retries parameters are properly passed to the OpenAI client.""" from unittest.mock import Mock, patch - + mock_openai_client = Mock() mock_openai_client.timeout = None mock_openai_client.max_retries = None - + with patch( "databricks_langchain.chat_models.get_openai_client", return_value=mock_openai_client ) as mock_get_client: # Test with timeout and max_retries - llm = ChatDatabricks( - model="test-model", - timeout=60.0, - max_retries=5 - ) - + llm = ChatDatabricks(model="test-model", timeout=60.0, max_retries=5) + # Verify get_openai_client was called with the correct parameters - mock_get_client.assert_called_once_with( - workspace_client=None, - timeout=60.0, - max_retries=5 - ) - + mock_get_client.assert_called_once_with(workspace_client=None, timeout=60.0, max_retries=5) + # Test that client is set assert llm.client == mock_openai_client assert llm.timeout == 60.0 @@ -128,29 +118,24 @@ def test_timeout_and_max_retries_parameters() -> None: def test_timeout_and_max_retries_with_workspace_client() -> None: """Test timeout and max_retries parameters work with workspace_client.""" from unittest.mock import Mock, patch - + mock_workspace_client = Mock() mock_openai_client = Mock() mock_openai_client.timeout = None mock_openai_client.max_retries = None - + with patch( "databricks_langchain.chat_models.get_openai_client", return_value=mock_openai_client ) as mock_get_client: llm = ChatDatabricks( - model="test-model", - workspace_client=mock_workspace_client, - timeout=30.0, - max_retries=2 + model="test-model", workspace_client=mock_workspace_client, timeout=30.0, max_retries=2 ) - + # Verify get_openai_client was called with all parameters mock_get_client.assert_called_once_with( - workspace_client=mock_workspace_client, - timeout=30.0, - max_retries=2 + workspace_client=mock_workspace_client, timeout=30.0, max_retries=2 ) - + assert llm.client == mock_openai_client assert llm.timeout == 30.0 assert llm.max_retries == 2 @@ -173,9 +158,7 @@ def test_default_workspace_client() -> None: assert llm.client == mock_openai_client # Now expects no additional kwargs when timeout/max_retries are None - mock_get_client.assert_called_once_with( - workspace_client=None - ) + mock_get_client.assert_called_once_with(workspace_client=None) def test_target_uri_deprecation_warning() -> None: @@ -1029,9 +1012,7 @@ def test_chat_databricks_init_sets_client(): llm = ChatDatabricks(model="test-model") # Now expects no additional kwargs when timeout/max_retries are None - mock_get_client.assert_called_once_with( - workspace_client=None - ) + mock_get_client.assert_called_once_with(workspace_client=None) assert llm.client == mock_client diff --git a/integrations/langchain/tests/unit_tests/test_utils.py b/integrations/langchain/tests/unit_tests/test_utils.py index d2c4c7771..8ee0f58ea 100644 --- a/integrations/langchain/tests/unit_tests/test_utils.py +++ b/integrations/langchain/tests/unit_tests/test_utils.py @@ -7,62 +7,56 @@ def test_get_openai_client_with_timeout_and_max_retries() -> None: """Test that get_openai_client properly passes timeout and max_retries as kwargs to the SDK.""" - + mock_openai_client = Mock() - + mock_workspace_client = Mock() mock_workspace_client.serving_endpoints.get_open_ai_client.return_value = mock_openai_client - + # Test with workspace_client, timeout, and max_retries - client = get_openai_client( - workspace_client=mock_workspace_client, - timeout=45.0, - max_retries=3 - ) - + client = get_openai_client(workspace_client=mock_workspace_client, timeout=45.0, max_retries=3) + # Verify the OpenAI client was obtained with the correct kwargs mock_workspace_client.serving_endpoints.get_open_ai_client.assert_called_once_with( - timeout=45.0, - max_retries=3 + timeout=45.0, max_retries=3 ) - + # Verify the client is returned assert client == mock_openai_client def test_get_openai_client_with_default_workspace_client() -> None: """Test get_openai_client creates default WorkspaceClient when none provided.""" - + mock_openai_client = Mock() - + mock_workspace_client = Mock() mock_workspace_client.serving_endpoints.get_open_ai_client.return_value = mock_openai_client - + with patch("databricks.sdk.WorkspaceClient", return_value=mock_workspace_client): client = get_openai_client(timeout=30.0, max_retries=2) - + # Verify default WorkspaceClient was created and kwargs were passed mock_workspace_client.serving_endpoints.get_open_ai_client.assert_called_once_with( - timeout=30.0, - max_retries=2 + timeout=30.0, max_retries=2 ) - + # Verify the client is returned assert client == mock_openai_client def test_get_openai_client_without_timeout_and_retries() -> None: """Test get_openai_client doesn't pass kwargs when not provided.""" - + mock_openai_client = Mock() - + mock_workspace_client = Mock() mock_workspace_client.serving_endpoints.get_open_ai_client.return_value = mock_openai_client - + client = get_openai_client(workspace_client=mock_workspace_client) - + # Verify the OpenAI client was obtained without kwargs mock_workspace_client.serving_endpoints.get_open_ai_client.assert_called_once_with() - + # Verify the client is returned - assert client == mock_openai_client \ No newline at end of file + assert client == mock_openai_client From 2831044496437a401d46f9a702a804e46c79a9a2 Mon Sep 17 00:00:00 2001 From: Sid Murching Date: Tue, 2 Sep 2025 20:45:46 -0700 Subject: [PATCH 6/7] remove from dev deps Signed-off-by: Sid Murching --- integrations/langchain/pyproject.toml | 1 - 1 file changed, 1 deletion(-) diff --git a/integrations/langchain/pyproject.toml b/integrations/langchain/pyproject.toml index a8a1a1535..ffdf16bae 100644 --- a/integrations/langchain/pyproject.toml +++ b/integrations/langchain/pyproject.toml @@ -23,7 +23,6 @@ dependencies = [ dev = [ "pytest", "typing_extensions", - "databricks-sdk>=0.63.0", "ruff==0.6.4", ] From 919b8c90c2414d6246b6410da46f8c0c6cbbd1b5 Mon Sep 17 00:00:00 2001 From: Sid Murching Date: Tue, 2 Sep 2025 20:47:54 -0700 Subject: [PATCH 7/7] remove comments Signed-off-by: Sid Murching --- integrations/langchain/tests/unit_tests/test_chat_models.py | 3 --- 1 file changed, 3 deletions(-) diff --git a/integrations/langchain/tests/unit_tests/test_chat_models.py b/integrations/langchain/tests/unit_tests/test_chat_models.py index e03203884..2ee6c12e2 100644 --- a/integrations/langchain/tests/unit_tests/test_chat_models.py +++ b/integrations/langchain/tests/unit_tests/test_chat_models.py @@ -77,7 +77,6 @@ def test_workspace_client_parameter() -> None: llm = ChatDatabricks(model="test-model", workspace_client=mock_workspace_client) assert llm.client == mock_openai_client - # Now expects no additional kwargs when timeout/max_retries are None mock_get_client.assert_called_once_with(workspace_client=mock_workspace_client) @@ -157,7 +156,6 @@ def test_default_workspace_client() -> None: llm = ChatDatabricks(model="test-model") assert llm.client == mock_openai_client - # Now expects no additional kwargs when timeout/max_retries are None mock_get_client.assert_called_once_with(workspace_client=None) @@ -1011,7 +1009,6 @@ def test_chat_databricks_init_sets_client(): llm = ChatDatabricks(model="test-model") - # Now expects no additional kwargs when timeout/max_retries are None mock_get_client.assert_called_once_with(workspace_client=None) assert llm.client == mock_client