diff --git a/pyproject.toml b/pyproject.toml index 452ac3125a..4125292125 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -9,7 +9,7 @@ authors = [ ] dependencies = [ - "httpx>=0.23.0, <1", + "httpx>=0.25.0, <1", "pydantic>=1.9.0, <3", "typing-extensions>=4.11, <5", "typing-extensions>=4.14, <5", "anyio>=3.5.0, <5", diff --git a/src/openai/_base_client.py b/src/openai/_base_client.py index 216b36aabd..96fe1b82cf 100644 --- a/src/openai/_base_client.py +++ b/src/openai/_base_client.py @@ -5,6 +5,7 @@ import time import uuid import email +import socket import asyncio import inspect import logging @@ -831,11 +832,36 @@ def _idempotency_key(self) -> str: return f"stainless-python-retry-{uuid.uuid4()}" +def _build_keepalive_socket_options() -> list[tuple[int, int, int]]: + # Enable TCP keepalive to prevent silent connection drops by NAT gateways and + # load balancers during long-running inference calls (generation can hold a + # TCP connection idle for 300–600 s, exceeding the ~350 s AWS NAT Gateway + # idle timeout). Mirrors the pattern used by the Anthropic Python SDK. + options: list[tuple[int, int, int]] = [(socket.SOL_SOCKET, socket.SO_KEEPALIVE, 1)] + # TCP_KEEPIDLE (Linux) / TCP_KEEPALIVE (macOS): seconds idle before first probe + 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)) + # TCP_KEEPINTVL: seconds between subsequent probes + if hasattr(socket, "TCP_KEEPINTVL"): + options.append((socket.IPPROTO_TCP, socket.TCP_KEEPINTVL, 60)) + # TCP_KEEPCNT: drop the connection after this many unanswered probes + if hasattr(socket, "TCP_KEEPCNT"): + options.append((socket.IPPROTO_TCP, socket.TCP_KEEPCNT, 5)) + return options + + 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 "transport" not in kwargs: + kwargs["transport"] = httpx.HTTPTransport( + limits=kwargs.get("limits", DEFAULT_CONNECTION_LIMITS), + socket_options=_build_keepalive_socket_options(), + ) super().__init__(**kwargs) @@ -1423,6 +1449,11 @@ def __init__(self, **kwargs: Any) -> None: kwargs.setdefault("timeout", DEFAULT_TIMEOUT) kwargs.setdefault("limits", DEFAULT_CONNECTION_LIMITS) kwargs.setdefault("follow_redirects", True) + if "transport" not in kwargs: + kwargs["transport"] = httpx.AsyncHTTPTransport( + limits=kwargs.get("limits", DEFAULT_CONNECTION_LIMITS), + socket_options=_build_keepalive_socket_options(), + ) super().__init__(**kwargs) diff --git a/tests/test_client.py b/tests/test_client.py index 2d8955a58e..185d1206eb 100644 --- a/tests/test_client.py +++ b/tests/test_client.py @@ -1304,6 +1304,22 @@ 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: + import socket as socket_module + + client = OpenAI(base_url=base_url, api_key=api_key) + transport = client._client._transport + assert isinstance(transport, httpx.HTTPTransport) + socket_options = transport._pool._socket_options + assert any( + opt == (socket_module.SOL_SOCKET, socket_module.SO_KEEPALIVE, 1) for opt in socket_options + ), "Default sync transport should have SO_KEEPALIVE enabled to survive NAT idle timeouts" + + def test_custom_http_client_transport_is_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, "A caller-supplied http_client must not be replaced" + @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 +2580,22 @@ 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: + import socket as socket_module + + client = AsyncOpenAI(base_url=base_url, api_key=api_key) + transport = client._client._transport + assert isinstance(transport, httpx.AsyncHTTPTransport) + socket_options = transport._pool._socket_options + assert any( + opt == (socket_module.SOL_SOCKET, socket_module.SO_KEEPALIVE, 1) for opt in socket_options + ), "Default async transport should have SO_KEEPALIVE enabled to survive NAT idle timeouts" + + async def test_custom_async_http_client_transport_is_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, "A caller-supplied http_client must not be replaced" + @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