diff --git a/src/openai/_base_client.py b/src/openai/_base_client.py index 216b36aabd..b04b5cfeb7 100644 --- a/src/openai/_base_client.py +++ b/src/openai/_base_client.py @@ -1,5 +1,6 @@ from __future__ import annotations +import os import sys import json import time @@ -9,6 +10,7 @@ import inspect import logging import platform +import socket import warnings import email.utils from types import TracebackType @@ -831,11 +833,60 @@ def _idempotency_key(self) -> str: return f"stainless-python-retry-{uuid.uuid4()}" +def _build_keepalive_socket_options() -> list[tuple[int, int, int]]: + options: list[tuple[int, int, int]] = [(socket.SOL_SOCKET, socket.SO_KEEPALIVE, 1)] + if hasattr(socket, "TCP_KEEPIDLE"): + options.append((socket.IPPROTO_TCP, socket.TCP_KEEPIDLE, 60)) + elif hasattr(socket, "TCP_KEEPALIVE"): + options.append((socket.IPPROTO_TCP, socket.TCP_KEEPALIVE, 60)) + if hasattr(socket, "TCP_KEEPINTVL"): + options.append((socket.IPPROTO_TCP, socket.TCP_KEEPINTVL, 60)) + if hasattr(socket, "TCP_KEEPCNT"): + options.append((socket.IPPROTO_TCP, socket.TCP_KEEPCNT, 5)) + return options + + +_HTTPX_TRANSPORT_SUPPORTS_SOCKET_OPTIONS = "socket_options" in inspect.signature( + httpx.HTTPTransport.__init__ +).parameters + +_TRANSPORT_PASSTHROUGH_KEYS = ("verify", "cert", "trust_env", "http1", "http2", "limits") + + +def _has_env_proxy() -> bool: + proxy_vars = ("HTTP_PROXY", "HTTPS_PROXY", "ALL_PROXY", "http_proxy", "https_proxy", "all_proxy") + return any(os.environ.get(var) for var in proxy_vars) + + +def _should_inject_keepalive_transport(kwargs: dict[str, Any]) -> bool: + if "transport" in kwargs: + return False + if "proxy" in kwargs: + return False + if kwargs.get("trust_env", True) and _has_env_proxy(): + return False + return True + + +def _build_keepalive_transport( + transport_cls: type[httpx.HTTPTransport] | type[httpx.AsyncHTTPTransport], + kwargs: dict[str, Any], +) -> httpx.HTTPTransport | httpx.AsyncHTTPTransport: + transport_kwargs: dict[str, Any] = { + key: kwargs[key] for key in _TRANSPORT_PASSTHROUGH_KEYS if key in kwargs + } + if _HTTPX_TRANSPORT_SUPPORTS_SOCKET_OPTIONS: + transport_kwargs["socket_options"] = _build_keepalive_socket_options() + return transport_cls(**transport_kwargs) + + class _DefaultHttpxClient(httpx.Client): def __init__(self, **kwargs: Any) -> None: kwargs.setdefault("timeout", DEFAULT_TIMEOUT) kwargs.setdefault("limits", DEFAULT_CONNECTION_LIMITS) kwargs.setdefault("follow_redirects", True) + if _should_inject_keepalive_transport(kwargs): + kwargs["transport"] = _build_keepalive_transport(httpx.HTTPTransport, kwargs) super().__init__(**kwargs) @@ -1423,6 +1474,8 @@ def __init__(self, **kwargs: Any) -> None: kwargs.setdefault("timeout", DEFAULT_TIMEOUT) kwargs.setdefault("limits", DEFAULT_CONNECTION_LIMITS) kwargs.setdefault("follow_redirects", True) + if _should_inject_keepalive_transport(kwargs): + kwargs["transport"] = _build_keepalive_transport(httpx.AsyncHTTPTransport, kwargs) super().__init__(**kwargs) diff --git a/tests/test_client.py b/tests/test_client.py index 2d8955a58e..18cf576ba7 100644 --- a/tests/test_client.py +++ b/tests/test_client.py @@ -6,6 +6,7 @@ import os import sys import json +import socket import asyncio import inspect import dataclasses @@ -1304,6 +1305,52 @@ def test_default_client_creation(self) -> None: limits=httpx.Limits(max_connections=100, max_keepalive_connections=20), ) + def test_default_transport_has_tcp_keepalive(self) -> None: + proxy_vars = ("HTTP_PROXY", "HTTPS_PROXY", "ALL_PROXY", "http_proxy", "https_proxy", "all_proxy") + with mock.patch.dict(os.environ, {v: "" for v in proxy_vars}, clear=False): + client = OpenAI(base_url=base_url, api_key=api_key) + transport = client._client._transport + assert isinstance(transport, httpx.HTTPTransport) + pool = transport._pool + socket_options = pool._socket_options + assert any( + opt[0] == socket.SOL_SOCKET and opt[1] == socket.SO_KEEPALIVE and opt[2] == 1 + for opt in socket_options + ) + + def test_custom_http_client_transport_not_overridden(self) -> None: + with httpx.Client() as http_client: + client = OpenAI(base_url=base_url, api_key=api_key, http_client=http_client) + assert client._client is http_client + + def test_keepalive_skipped_when_proxy_env_set(self) -> None: + with mock.patch.dict(os.environ, {"HTTP_PROXY": "http://proxy:8080"}, clear=False): + client = OpenAI(base_url=base_url, api_key=api_key) + transport = client._client._transport + assert isinstance(transport, httpx.HTTPTransport) + pool = transport._pool + socket_options = list(pool._socket_options) + assert not any( + opt[0] == socket.SOL_SOCKET and opt[1] == socket.SO_KEEPALIVE + for opt in socket_options + ), "Keepalive should NOT be set when HTTP_PROXY is configured" + + def test_keepalive_not_skipped_when_trust_env_false_and_proxy_set(self) -> None: + with mock.patch.dict(os.environ, {"HTTP_PROXY": "http://proxy:8080"}, clear=False): + client = OpenAI(base_url=base_url, api_key=api_key, http_client=httpx.Client(trust_env=False)) + assert client._client is not None + + def test_keepalive_skipped_when_explicit_proxy_arg(self) -> None: + with mock.patch.dict(os.environ, clear=True): + client = DefaultHttpxClient(proxy="http://proxy:8080") + transport = client._transport + assert isinstance(transport, httpx.HTTPTransport) + so = list(transport._pool._socket_options) + assert not any( + opt[0] == socket.SOL_SOCKET and opt[1] == socket.SO_KEEPALIVE + for opt in so + ), "Keepalive should NOT be set when proxy= is passed" + @pytest.mark.respx(base_url=base_url) def test_follow_redirects(self, respx_mock: MockRouter, client: OpenAI) -> None: # Test that the default follow_redirects=True allows following redirects @@ -2564,6 +2611,47 @@ async def test_default_client_creation(self) -> None: limits=httpx.Limits(max_connections=100, max_keepalive_connections=20), ) + async def test_default_transport_has_tcp_keepalive(self) -> None: + proxy_vars = ("HTTP_PROXY", "HTTPS_PROXY", "ALL_PROXY", "http_proxy", "https_proxy", "all_proxy") + with mock.patch.dict(os.environ, {v: "" for v in proxy_vars}, clear=False): + client = AsyncOpenAI(base_url=base_url, api_key=api_key) + transport = client._client._transport + assert isinstance(transport, httpx.AsyncHTTPTransport) + pool = transport._pool + socket_options = pool._socket_options + assert any( + opt[0] == socket.SOL_SOCKET and opt[1] == socket.SO_KEEPALIVE and opt[2] == 1 + for opt in socket_options + ) + + async def test_custom_async_http_client_transport_not_overridden(self) -> None: + async with httpx.AsyncClient() as http_client: + client = AsyncOpenAI(base_url=base_url, api_key=api_key, http_client=http_client) + assert client._client is http_client + + async def test_async_keepalive_skipped_when_proxy_env_set(self) -> None: + with mock.patch.dict(os.environ, {"HTTPS_PROXY": "http://proxy:8080"}, clear=False): + client = AsyncOpenAI(base_url=base_url, api_key=api_key) + transport = client._client._transport + assert isinstance(transport, httpx.AsyncHTTPTransport) + pool = transport._pool + socket_options = list(pool._socket_options) + assert not any( + opt[0] == socket.SOL_SOCKET and opt[1] == socket.SO_KEEPALIVE + for opt in socket_options + ) + + async def test_async_keepalive_skipped_when_explicit_proxy_arg(self) -> None: + with mock.patch.dict(os.environ, clear=True): + client = DefaultAsyncHttpxClient(proxy="http://proxy:8080") + transport = client._transport + assert isinstance(transport, httpx.AsyncHTTPTransport) + so = list(transport._pool._socket_options) + assert not any( + opt[0] == socket.SOL_SOCKET and opt[1] == socket.SO_KEEPALIVE + for opt in so + ) + @pytest.mark.respx(base_url=base_url) async def test_follow_redirects(self, respx_mock: MockRouter, async_client: AsyncOpenAI) -> None: # Test that the default follow_redirects=True allows following redirects