diff --git a/integrations/openai/src/databricks_openai/utils/clients.py b/integrations/openai/src/databricks_openai/utils/clients.py index 28138bdc..840b6c97 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. **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. **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)."""