From 9da4cff544916f0079ade046865eb813e54781ec Mon Sep 17 00:00:00 2001 From: Kevin Yan Date: Mon, 11 May 2026 09:51:48 -0700 Subject: [PATCH 1/5] DatabricksOpenAI: forward extra kwargs to OpenAI init (#423) Both DatabricksOpenAI and AsyncDatabricksOpenAI had closed init signatures, so ChatDatabricks(max_retries=..., timeout=...) crashed the moment .client or .async_client was constructed. The async path was broken since 0.15.0 (when async_client was introduced); 0.19.0 migrated the sync path onto the same wrapper and inherited the same bug. Add **kwargs to both wrappers and forward them through to super().__init__, so timeout, max_retries, default_headers, and other openai.OpenAI kwargs work as advertised in get_*_openai_client's docstring. Conflicts with the hardcoded api_key/http_client (which carry Databricks auth) surface as Python's native TypeError on duplicate kwargs. Co-Authored-By: Claude Opus 4.7 (1M context) --- .../tests/unit_tests/test_chat_models.py | 20 +++++++++++++++++++ .../src/databricks_openai/utils/clients.py | 4 ++++ .../openai/tests/unit_tests/test_clients.py | 12 +++++++++++ 3 files changed, 36 insertions(+) diff --git a/integrations/langchain/tests/unit_tests/test_chat_models.py b/integrations/langchain/tests/unit_tests/test_chat_models.py index 5e85ff7e..feb62b66 100644 --- a/integrations/langchain/tests/unit_tests/test_chat_models.py +++ b/integrations/langchain/tests/unit_tests/test_chat_models.py @@ -123,6 +123,26 @@ def test_timeout_and_max_retries_parameters() -> None: assert llm.max_retries == 5 +def test_async_client_forwards_timeout_and_max_retries() -> None: + """Regression test for issue #423: async_client used to crash on these kwargs.""" + from unittest.mock import MagicMock + + mock_workspace_client = MagicMock(spec=sdk.WorkspaceClient) + mock_workspace_client.config.host = "https://test.databricks.com" + mock_workspace_client.config.authenticate.return_value = {"Authorization": "Bearer t"} + + llm = ChatDatabricks( + model="test-model", + workspace_client=mock_workspace_client, + timeout=42.0, + max_retries=8, + ) + + async_client = llm.async_client + assert async_client.timeout == 42.0 + assert async_client.max_retries == 8 + + 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 diff --git a/integrations/openai/src/databricks_openai/utils/clients.py b/integrations/openai/src/databricks_openai/utils/clients.py index 6711b87f..a645bd07 100644 --- a/integrations/openai/src/databricks_openai/utils/clients.py +++ b/integrations/openai/src/databricks_openai/utils/clients.py @@ -385,6 +385,7 @@ def __init__( base_url: str | None = None, use_ai_gateway_native_api: bool = False, use_ai_gateway: bool = False, + **kwargs: Any, ): if workspace_client is None: workspace_client = WorkspaceClient() @@ -400,6 +401,7 @@ def __init__( base_url=target_base_url, api_key=_get_openai_api_key(), http_client=_get_authorized_http_client(workspace_client), + **kwargs, ) @override @@ -546,6 +548,7 @@ def __init__( base_url: str | None = None, use_ai_gateway_native_api: bool = False, use_ai_gateway: bool = False, + **kwargs: Any, ): if workspace_client is None: workspace_client = WorkspaceClient() @@ -561,6 +564,7 @@ def __init__( base_url=target_base_url, api_key=_get_openai_api_key(), http_client=_get_authorized_async_http_client(workspace_client), + **kwargs, ) @property diff --git a/integrations/openai/tests/unit_tests/test_clients.py b/integrations/openai/tests/unit_tests/test_clients.py index bd3fbe53..bad7fcfd 100644 --- a/integrations/openai/tests/unit_tests/test_clients.py +++ b/integrations/openai/tests/unit_tests/test_clients.py @@ -168,6 +168,18 @@ def test_bearer_auth_flow(self, mock_workspace_client): mock_workspace_client.config.authenticate.assert_called() +class TestDatabricksOpenAIKwargForwarding: + """Forwarding of OpenAI kwargs through the wrappers (issue #423).""" + + @pytest.mark.parametrize( + "client_cls", [DatabricksOpenAI, AsyncDatabricksOpenAI], ids=["sync", "async"] + ) + def test_forwards_timeout_and_max_retries(self, client_cls, mock_workspace_client): + client = client_cls(workspace_client=mock_workspace_client, timeout=42.0, max_retries=7) + assert client.timeout == 42.0 + assert client.max_retries == 7 + + class TestStrictFieldStripping: """Tests for strict field stripping helper functions.""" From 599ada09336ccc59bded5261df4181e88e144d63 Mon Sep 17 00:00:00 2001 From: Kevin Yan Date: Mon, 11 May 2026 11:38:42 -0700 Subject: [PATCH 2/5] Drop cross-package regression test that fails under --resolution lowest-direct CI resolves databricks-openai from PyPI at its floor pin (>=0.14.0), which predates the **kwargs fix in this PR. The openai-side test in integrations/openai/tests/unit_tests/test_clients.py already covers the regression at the layer where the bug actually lived. Co-Authored-By: Claude Opus 4.7 (1M context) --- .../tests/unit_tests/test_chat_models.py | 20 ------------------- 1 file changed, 20 deletions(-) diff --git a/integrations/langchain/tests/unit_tests/test_chat_models.py b/integrations/langchain/tests/unit_tests/test_chat_models.py index feb62b66..5e85ff7e 100644 --- a/integrations/langchain/tests/unit_tests/test_chat_models.py +++ b/integrations/langchain/tests/unit_tests/test_chat_models.py @@ -123,26 +123,6 @@ def test_timeout_and_max_retries_parameters() -> None: assert llm.max_retries == 5 -def test_async_client_forwards_timeout_and_max_retries() -> None: - """Regression test for issue #423: async_client used to crash on these kwargs.""" - from unittest.mock import MagicMock - - mock_workspace_client = MagicMock(spec=sdk.WorkspaceClient) - mock_workspace_client.config.host = "https://test.databricks.com" - mock_workspace_client.config.authenticate.return_value = {"Authorization": "Bearer t"} - - llm = ChatDatabricks( - model="test-model", - workspace_client=mock_workspace_client, - timeout=42.0, - max_retries=8, - ) - - async_client = llm.async_client - assert async_client.timeout == 42.0 - assert async_client.max_retries == 8 - - 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 From 7436faa5691dc4fb52daa3094ad0fdca4ef1b28e Mon Sep 17 00:00:00 2001 From: Kevin Yan Date: Mon, 11 May 2026 14:15:59 -0700 Subject: [PATCH 3/5] Document **kwargs in DatabricksOpenAI / AsyncDatabricksOpenAI docstrings Co-Authored-By: Claude Opus 4.7 (1M context) --- .../openai/src/databricks_openai/utils/clients.py | 8 ++++++++ 1 file changed, 8 insertions(+) diff --git a/integrations/openai/src/databricks_openai/utils/clients.py b/integrations/openai/src/databricks_openai/utils/clients.py index a645bd07..b989af7a 100644 --- a/integrations/openai/src/databricks_openai/utils/clients.py +++ b/integrations/openai/src/databricks_openai/utils/clients.py @@ -343,6 +343,10 @@ class DatabricksOpenAI(OpenAI): with ``base_url`` or ``use_ai_gateway``. Defaults to False. use_ai_gateway: If True, auto-detect AI Gateway V2 availability and route requests through it using the MLflow API. Defaults to False. + **kwargs: Additional keyword arguments forwarded to the underlying ``openai.OpenAI`` + client (e.g. ``timeout``, ``max_retries``, ``default_headers``). ``api_key`` and + ``http_client`` are managed internally for Databricks authentication and cannot + be overridden. Example - Query a serving or AI gateway endpoint: >>> client = DatabricksOpenAI() @@ -506,6 +510,10 @@ class AsyncDatabricksOpenAI(AsyncOpenAI): with ``base_url`` or ``use_ai_gateway``. Defaults to False. use_ai_gateway: If True, auto-detect AI Gateway V2 availability and route requests through it using the MLflow API. Defaults to False. + **kwargs: Additional keyword arguments forwarded to the underlying ``openai.AsyncOpenAI`` + client (e.g. ``timeout``, ``max_retries``, ``default_headers``). ``api_key`` and + ``http_client`` are managed internally for Databricks authentication and cannot + be overridden. Example - Query a serving or AI gateway endpoint: >>> client = AsyncDatabricksOpenAI() From 24e8c56a2d39da197cccd967e9c7b0d2ac459d57 Mon Sep 17 00:00:00 2001 From: Kevin Yan Date: Mon, 11 May 2026 14:18:52 -0700 Subject: [PATCH 4/5] Test that api_key / http_client raise TypeError when passed via **kwargs Locks in the safety invariant that the new **kwargs surface cannot be used to override Databricks-managed auth. Python's native double-pass TypeError serves as the guardrail; this test asserts it stays that way. Co-Authored-By: Claude Opus 4.7 (1M context) --- integrations/openai/tests/unit_tests/test_clients.py | 12 ++++++++++++ 1 file changed, 12 insertions(+) diff --git a/integrations/openai/tests/unit_tests/test_clients.py b/integrations/openai/tests/unit_tests/test_clients.py index bad7fcfd..bd23211a 100644 --- a/integrations/openai/tests/unit_tests/test_clients.py +++ b/integrations/openai/tests/unit_tests/test_clients.py @@ -179,6 +179,18 @@ def test_forwards_timeout_and_max_retries(self, client_cls, mock_workspace_clien assert client.timeout == 42.0 assert client.max_retries == 7 + @pytest.mark.parametrize( + "client_cls", [DatabricksOpenAI, AsyncDatabricksOpenAI], ids=["sync", "async"] + ) + @pytest.mark.parametrize("managed_kwarg", ["api_key", "http_client"]) + def test_rejects_databricks_managed_auth_kwargs( + self, client_cls, managed_kwarg, mock_workspace_client + ): + # api_key and http_client are hardcoded with Databricks auth; passing them through + # **kwargs would collide with the super().__init__() call and raise TypeError. + with pytest.raises(TypeError, match="got multiple values for keyword argument"): + client_cls(workspace_client=mock_workspace_client, **{managed_kwarg: "anything"}) + class TestStrictFieldStripping: """Tests for strict field stripping helper functions.""" From b4e1c5df45f1f38ccfb97241eb35dca0a7c33357 Mon Sep 17 00:00:00 2001 From: Kevin Yan Date: Tue, 12 May 2026 17:58:13 -0700 Subject: [PATCH 5/5] Trim api_key/http_client note from **kwargs docstrings Co-Authored-By: Claude Opus 4.7 (1M context) --- .../openai/src/databricks_openai/utils/clients.py | 8 ++------ 1 file changed, 2 insertions(+), 6 deletions(-) diff --git a/integrations/openai/src/databricks_openai/utils/clients.py b/integrations/openai/src/databricks_openai/utils/clients.py index b989af7a..28138bdc 100644 --- a/integrations/openai/src/databricks_openai/utils/clients.py +++ b/integrations/openai/src/databricks_openai/utils/clients.py @@ -344,9 +344,7 @@ class DatabricksOpenAI(OpenAI): use_ai_gateway: If True, auto-detect AI Gateway V2 availability and route requests through it using the MLflow API. Defaults to False. **kwargs: Additional keyword arguments forwarded to the underlying ``openai.OpenAI`` - client (e.g. ``timeout``, ``max_retries``, ``default_headers``). ``api_key`` and - ``http_client`` are managed internally for Databricks authentication and cannot - be overridden. + client (e.g. ``timeout``, ``max_retries``, ``default_headers``). Example - Query a serving or AI gateway endpoint: >>> client = DatabricksOpenAI() @@ -511,9 +509,7 @@ class AsyncDatabricksOpenAI(AsyncOpenAI): use_ai_gateway: If True, auto-detect AI Gateway V2 availability and route requests through it using the MLflow API. Defaults to False. **kwargs: Additional keyword arguments forwarded to the underlying ``openai.AsyncOpenAI`` - client (e.g. ``timeout``, ``max_retries``, ``default_headers``). ``api_key`` and - ``http_client`` are managed internally for Databricks authentication and cannot - be overridden. + client (e.g. ``timeout``, ``max_retries``, ``default_headers``). Example - Query a serving or AI gateway endpoint: >>> client = AsyncDatabricksOpenAI()