From 8f6384d5cc3804b190cb8f894970579c48658a9d Mon Sep 17 00:00:00 2001 From: Ann Zhang Date: Mon, 22 Jun 2026 18:33:37 -0700 Subject: [PATCH 1/2] Enable HTTP redirect following in DatabricksOpenAI clients Passing a custom http_client to the OpenAI SDK bypasses its default of follow_redirects=True, since httpx's Client defaults to False. Restore the intended behavior by following redirects, exposed as a follow_redirects parameter (default True) on both the sync and async clients. Co-authored-by: Isaac --- .../src/databricks_openai/utils/clients.py | 22 ++++++++++++++----- .../openai/tests/unit_tests/test_clients.py | 20 +++++++++++++++++ 2 files changed, 36 insertions(+), 6 deletions(-) diff --git a/integrations/openai/src/databricks_openai/utils/clients.py b/integrations/openai/src/databricks_openai/utils/clients.py index 28138bdc..b8d09edc 100644 --- a/integrations/openai/src/databricks_openai/utils/clients.py +++ b/integrations/openai/src/databricks_openai/utils/clients.py @@ -137,14 +137,18 @@ def _resolve_base_url( return f"{host}/serving-endpoints" -def _get_authorized_http_client(workspace_client: WorkspaceClient) -> Client: +def _get_authorized_http_client( + workspace_client: WorkspaceClient, follow_redirects: bool = True +) -> Client: databricks_token_auth = BearerAuth(workspace_client.config.authenticate) - return Client(auth=databricks_token_auth) + return Client(auth=databricks_token_auth, follow_redirects=follow_redirects) -def _get_authorized_async_http_client(workspace_client: WorkspaceClient) -> AsyncClient: +def _get_authorized_async_http_client( + workspace_client: WorkspaceClient, follow_redirects: bool = True +) -> AsyncClient: databricks_token_auth = BearerAuth(workspace_client.config.authenticate) - return AsyncClient(auth=databricks_token_auth) + return AsyncClient(auth=databricks_token_auth, follow_redirects=follow_redirects) def _validate_oauth_for_apps(workspace_client: WorkspaceClient) -> None: @@ -343,6 +347,8 @@ 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. + follow_redirects: If True, the underlying HTTP client follows redirect responses. + Defaults to True, matching the OpenAI SDK's default behavior. **kwargs: Additional keyword arguments forwarded to the underlying ``openai.OpenAI`` client (e.g. ``timeout``, ``max_retries``, ``default_headers``). @@ -387,6 +393,7 @@ def __init__( base_url: str | None = None, use_ai_gateway_native_api: bool = False, use_ai_gateway: bool = False, + follow_redirects: bool = True, **kwargs: Any, ): if workspace_client is None: @@ -402,7 +409,7 @@ def __init__( super().__init__( base_url=target_base_url, api_key=_get_openai_api_key(), - http_client=_get_authorized_http_client(workspace_client), + http_client=_get_authorized_http_client(workspace_client, follow_redirects), **kwargs, ) @@ -508,6 +515,8 @@ 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. + follow_redirects: If True, the underlying HTTP client follows redirect responses. + Defaults to True, matching the OpenAI SDK's default behavior. **kwargs: Additional keyword arguments forwarded to the underlying ``openai.AsyncOpenAI`` client (e.g. ``timeout``, ``max_retries``, ``default_headers``). @@ -552,6 +561,7 @@ def __init__( base_url: str | None = None, use_ai_gateway_native_api: bool = False, use_ai_gateway: bool = False, + follow_redirects: bool = True, **kwargs: Any, ): if workspace_client is None: @@ -567,7 +577,7 @@ def __init__( super().__init__( base_url=target_base_url, api_key=_get_openai_api_key(), - http_client=_get_authorized_async_http_client(workspace_client), + http_client=_get_authorized_async_http_client(workspace_client, follow_redirects), **kwargs, ) diff --git a/integrations/openai/tests/unit_tests/test_clients.py b/integrations/openai/tests/unit_tests/test_clients.py index bd23211a..aab52864 100644 --- a/integrations/openai/tests/unit_tests/test_clients.py +++ b/integrations/openai/tests/unit_tests/test_clients.py @@ -123,6 +123,16 @@ def test_bearer_auth_flow(self, mock_workspace_client): # Verify authenticate was called mock_workspace_client.config.authenticate.assert_called() + def test_follow_redirects_default_and_override(self, mock_workspace_client): + """follow_redirects defaults to True and can be disabled.""" + assert _get_authorized_http_client(mock_workspace_client).follow_redirects is True + assert ( + _get_authorized_http_client( + mock_workspace_client, follow_redirects=False + ).follow_redirects + is False + ) + class TestAsyncDatabricksOpenAI: """Tests for AsyncDatabricksOpenAI client.""" @@ -167,6 +177,16 @@ def test_bearer_auth_flow(self, mock_workspace_client): # Verify authenticate was called mock_workspace_client.config.authenticate.assert_called() + def test_follow_redirects_default_and_override(self, mock_workspace_client): + """follow_redirects defaults to True and can be disabled for the async client.""" + assert _get_authorized_async_http_client(mock_workspace_client).follow_redirects is True + assert ( + _get_authorized_async_http_client( + mock_workspace_client, follow_redirects=False + ).follow_redirects + is False + ) + class TestDatabricksOpenAIKwargForwarding: """Forwarding of OpenAI kwargs through the wrappers (issue #423).""" From 13b86c5176e76af7b7b623d1870006badb26fab9 Mon Sep 17 00:00:00 2001 From: Ann Zhang Date: Mon, 22 Jun 2026 18:47:18 -0700 Subject: [PATCH 2/2] Simplify follow_redirects docstring Co-authored-by: Isaac --- integrations/openai/src/databricks_openai/utils/clients.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/integrations/openai/src/databricks_openai/utils/clients.py b/integrations/openai/src/databricks_openai/utils/clients.py index b8d09edc..840b6c97 100644 --- a/integrations/openai/src/databricks_openai/utils/clients.py +++ b/integrations/openai/src/databricks_openai/utils/clients.py @@ -348,7 +348,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. follow_redirects: If True, the underlying HTTP client follows redirect responses. - Defaults to True, matching the OpenAI SDK's default behavior. + Defaults to True. **kwargs: Additional keyword arguments forwarded to the underlying ``openai.OpenAI`` client (e.g. ``timeout``, ``max_retries``, ``default_headers``). @@ -516,7 +516,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. follow_redirects: If True, the underlying HTTP client follows redirect responses. - Defaults to True, matching the OpenAI SDK's default behavior. + Defaults to True. **kwargs: Additional keyword arguments forwarded to the underlying ``openai.AsyncOpenAI`` client (e.g. ``timeout``, ``max_retries``, ``default_headers``).