diff --git a/litellm/llms/custom_httpx/http_handler.py b/litellm/llms/custom_httpx/http_handler.py index cffeb44f12c..d077e70a4a5 100644 --- a/litellm/llms/custom_httpx/http_handler.py +++ b/litellm/llms/custom_httpx/http_handler.py @@ -665,6 +665,28 @@ class AsyncHTTPHandler: self._client = client self._owns_client = False + @property + def owns_client(self) -> bool: + return self._owns_client + + @classmethod + def from_client( + cls, + client: httpx.AsyncClient, + timeout: float | httpx.Timeout | None = None, + ) -> "AsyncHTTPHandler": + handler: Final = cls.__new__(cls) + handler.timeout = timeout + handler.event_hooks = None + handler.ssl_verify = None + handler.shared_session = None + handler.transport = None + handler.follow_redirects = True + handler._owns_client = False + handler._client = client + handler.client_alias = None + return handler + def create_client( self, timeout: float | httpx.Timeout | None, @@ -1460,6 +1482,10 @@ class HTTPHandler: self._client = client self._owns_client = False + @property + def owns_client(self) -> bool: + return self._owns_client + def close(self): # Close the client when you're done with it if self._owns_client: @@ -1762,7 +1788,16 @@ def get_async_httpx_client( except Exception: pass - _cache_key_name: Final = "async_httpx_client" + _params_key_name + llm_provider + aclient_session: Final = ( + litellm.aclient_session + if isinstance(litellm.aclient_session, httpx.AsyncClient) + and shared_session is None + and llm_provider != "litellm_module_level_client" + and (params is None or ("transport" not in params and "event_hooks" not in params)) + else None + ) + _session_key_name: Final = f"_aclient_session_{id(aclient_session)}" if aclient_session is not None else "" + _cache_key_name: Final = "async_httpx_client" + _params_key_name + llm_provider + _session_key_name # Lazily initialize the global in-memory client cache to avoid relying on # litellm globals being fully populated during import time. @@ -1777,7 +1812,12 @@ def get_async_httpx_client( if _cached_client: return _cached_client - if params is not None: + if aclient_session is not None: + _new_client = AsyncHTTPHandler.from_client( # rebind-ok: branch-specific handler initialization + client=aclient_session, + timeout=_default_cached_client_timeout(), + ) + elif params is not None: # Filter out params that are only used for cache key, not for AsyncHTTPHandler.__init__ handler_params: Final = {k: v for k, v in params.items() if k != "disable_aiohttp_transport"} handler_params["shared_session"] = shared_session @@ -1792,7 +1832,7 @@ def get_async_httpx_client( key=_cache_key_name, value=_new_client, ttl=_DEFAULT_TTL_FOR_HTTPX_CLIENTS, - litellm_owned_client=True, + litellm_owned_client=_new_client.owns_client, ) return _new_client @@ -1812,7 +1852,13 @@ def get_httpx_client(params: dict | None = None) -> HTTPHandler: except Exception: pass - _cache_key_name: Final = "httpx_client" + _params_key_name + client_session: Final = ( + litellm.client_session + if isinstance(litellm.client_session, httpx.Client) and (params is None or "client" not in params) + else None + ) + _session_key_name: Final = f"_client_session_{id(client_session)}" if client_session is not None else "" + _cache_key_name: Final = "httpx_client" + _params_key_name + _session_key_name # Lazily initialize the global in-memory client cache to avoid relying on # litellm globals being fully populated during import time. @@ -1827,7 +1873,12 @@ def get_httpx_client(params: dict | None = None) -> HTTPHandler: if _cached_client: return _cached_client - if params is not None: + if client_session is not None: + _new_client = HTTPHandler( # rebind-ok: branch-specific handler initialization + timeout=_default_cached_client_timeout(), + client=client_session, + ) + elif params is not None: # Filter out params that are only used for cache key, not for HTTPHandler.__init__ handler_params: Final = {k: v for k, v in params.items() if k != "disable_aiohttp_transport"} _new_client = HTTPHandler(**handler_params) @@ -1838,7 +1889,7 @@ def get_httpx_client(params: dict | None = None) -> HTTPHandler: key=_cache_key_name, value=_new_client, ttl=_DEFAULT_TTL_FOR_HTTPX_CLIENTS, - litellm_owned_client=True, + litellm_owned_client=_new_client.owns_client, ) return _new_client diff --git a/tests/unit/llms/custom_httpx/test_http_handler.py b/tests/unit/llms/custom_httpx/test_http_handler.py index bd4ed4c99cc..e44b36929f3 100644 --- a/tests/unit/llms/custom_httpx/test_http_handler.py +++ b/tests/unit/llms/custom_httpx/test_http_handler.py @@ -25,6 +25,7 @@ from litellm.llms.custom_httpx.http_handler import ( AsyncHTTPHandler, HTTPHandler, MaskedHTTPStatusError, + get_async_httpx_client, get_httpx_client, get_shared_realtime_ssl_context, get_ssl_configuration, @@ -2172,3 +2173,101 @@ async def test_async_ollama_ssl_verify_false_builds_a_client_that_skips_certific assert route.call_count == 1 assert text == "Sunny" assert _cached_clients_verify_modes(client_cache, AsyncHTTPHandler) == (ssl.CERT_NONE,) + + +@pytest.mark.asyncio +async def test_aclient_session_is_used_by_http_handler_providers_and_not_closed( + monkeypatch: pytest.MonkeyPatch, +) -> None: + hits: Final[list[str]] = [] # mutable-ok: transport capture buffer + + async def handle(request: httpx.Request) -> httpx.Response: + hits.append(str(request.url)) + return httpx.Response( + 400, + request=request, + json={"error": {"message": "mock", "type": "invalid_request_error"}}, + ) + + client_cache: Final = LLMClientCache() + monkeypatch.setattr(litellm, "in_memory_llm_clients_cache", client_cache) + monkeypatch.setattr(litellm, "num_retries", 0) + + default_handler: Final = get_async_httpx_client(llm_provider="anthropic") + assert default_handler.owns_client is True + + session: Final = httpx.AsyncClient(transport=httpx.MockTransport(handle), trust_env=False) + try: + monkeypatch.setattr(litellm, "aclient_session", session) + session_handler: Final = get_async_httpx_client(llm_provider="anthropic") + assert session_handler.client is session + assert session_handler.owns_client is False + + module_level_handler: Final = get_async_httpx_client(llm_provider="litellm_module_level_client") + assert module_level_handler.client is not session + assert module_level_handler.owns_client is True + + with pytest.raises(litellm.BadRequestError): + await litellm.acompletion( + model="anthropic/claude-sonnet-4-20250514", + messages=[{"role": "user", "content": "hi"}], + api_key="sk-test", + api_base="http://127.0.0.1:9", + timeout=3, + ) + + assert hits == ["http://127.0.0.1:9/v1/messages"] + await session_handler.close() + assert not session.is_closed + finally: + await session.aclose() + await default_handler.close() + + +def test_client_session_is_used_by_sync_http_handler_providers_and_not_closed( + monkeypatch: pytest.MonkeyPatch, +) -> None: + hits: Final[list[str]] = [] # mutable-ok: transport capture buffer + + def handle(request: httpx.Request) -> httpx.Response: + hits.append(str(request.url)) + return httpx.Response( + 400, + request=request, + json={"error": {"message": "mock", "type": "invalid_request_error"}}, + ) + + client_cache: Final = LLMClientCache() + monkeypatch.setattr(litellm, "in_memory_llm_clients_cache", client_cache) + monkeypatch.setattr(litellm, "num_retries", 0) + + default_handler: Final = get_httpx_client() + assert default_handler.owns_client is True + + session: Final = httpx.Client(transport=httpx.MockTransport(handle), trust_env=False) + try: + monkeypatch.setattr(litellm, "client_session", session) + session_handler: Final = get_httpx_client() + assert session_handler.client is session + assert session_handler.owns_client is False + + module_level_sync: Final = HTTPHandler() + assert module_level_sync.client is not session + assert module_level_sync.owns_client is True + module_level_sync.close() + + with pytest.raises(litellm.BadRequestError): + litellm.completion( + model="anthropic/claude-sonnet-4-20250514", + messages=[{"role": "user", "content": "hi"}], + api_key="sk-test", + api_base="http://127.0.0.1:9", + timeout=3, + ) + + assert hits == ["http://127.0.0.1:9/v1/messages"] + session_handler.close() + assert not session.is_closed + finally: + session.close() + default_handler.close()