diff --git a/litellm/llms/custom_httpx/http_handler.py b/litellm/llms/custom_httpx/http_handler.py index 0077f247e3e..16cada08878 100644 --- a/litellm/llms/custom_httpx/http_handler.py +++ b/litellm/llms/custom_httpx/http_handler.py @@ -633,7 +633,6 @@ class AsyncHTTPHandler: shared_session: Optional["ClientSession"] = None, transport: httpx.AsyncBaseTransport | None = None, follow_redirects: bool = True, - client: httpx.AsyncClient | None = None, ): self.timeout = timeout self.event_hooks = event_hooks @@ -641,41 +640,17 @@ class AsyncHTTPHandler: self.shared_session = shared_session self.transport = transport self.follow_redirects = follow_redirects - resolved_client: Final = ( - client - if client is not None - else ( - litellm.aclient_session - if isinstance(litellm.aclient_session, httpx.AsyncClient) - and transport is None - and shared_session is None - and event_hooks is None - else None - ) - ) - self._owns_client = resolved_client is None - self._client = ( - self.create_client( - timeout=timeout, - event_hooks=event_hooks, - ssl_verify=ssl_verify, - shared_session=shared_session, - ) - if resolved_client is None - else resolved_client + self._owns_client = True + self._client = self.create_client( + timeout=timeout, + event_hooks=event_hooks, + ssl_verify=ssl_verify, + shared_session=shared_session, ) self.client_alias = client_alias @property def client(self) -> httpx.AsyncClient: - if ( - self._owns_client - and self.transport is None - and self.shared_session is None - and self.event_hooks is None - and isinstance(litellm.aclient_session, httpx.AsyncClient) - ): - return litellm.aclient_session if self._owns_client and self._client.is_closed: self._client = self.create_client( timeout=self.timeout, @@ -690,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, @@ -1444,14 +1441,9 @@ class HTTPHandler: self.ssl_verify = ssl_verify self.disable_default_headers = disable_default_headers self.follow_redirects = follow_redirects - resolved_client: Final = ( - client - if client is not None - else (litellm.client_session if isinstance(litellm.client_session, httpx.Client) else None) - ) - self._owns_client = resolved_client is None + self._owns_client = client is None self._heal_lock = threading.Lock() - self._client = self.create_client() if resolved_client is None else resolved_client + self._client = self.create_client() if client is None else client def create_client(self) -> httpx.Client: # Get unified SSL configuration @@ -1479,8 +1471,6 @@ class HTTPHandler: @property def client(self) -> httpx.Client: - if self._owns_client and isinstance(litellm.client_session, httpx.Client): - return litellm.client_session if self._owns_client and self._client.is_closed: with self._heal_lock: if self._owns_client and self._client.is_closed: @@ -1492,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: @@ -1794,7 +1788,14 @@ def get_async_httpx_client( except Exception: pass - aclient_session: Final = litellm.aclient_session if isinstance(litellm.aclient_session, httpx.AsyncClient) else None + 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 @@ -1811,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( + 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 @@ -1826,7 +1832,7 @@ def get_async_httpx_client( key=_cache_key_name, value=_new_client, ttl=_DEFAULT_TTL_FOR_HTTPX_CLIENTS, - litellm_owned_client=_new_client._owns_client, + litellm_owned_client=_new_client.owns_client, ) return _new_client @@ -1846,7 +1852,11 @@ def get_httpx_client(params: dict | None = None) -> HTTPHandler: except Exception: pass - client_session: Final = litellm.client_session if isinstance(litellm.client_session, httpx.Client) else None + 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 @@ -1863,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( + 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) @@ -1874,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=_new_client._owns_client, + 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 aea56eb012a..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, @@ -2192,18 +2193,19 @@ async def test_aclient_session_is_used_by_http_handler_providers_and_not_closed( monkeypatch.setattr(litellm, "in_memory_llm_clients_cache", client_cache) monkeypatch.setattr(litellm, "num_retries", 0) - from litellm.llms.custom_httpx.http_handler import get_async_httpx_client - - # Pre-warm cache without aclient_session to verify setting it bypasses stale default handlers default_handler: Final = get_async_httpx_client(llm_provider="anthropic") - assert default_handler._owns_client is True + 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 + 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( @@ -2240,14 +2242,19 @@ def test_client_session_is_used_by_sync_http_handler_providers_and_not_closed( monkeypatch.setattr(litellm, "num_retries", 0) default_handler: Final = get_httpx_client() - assert default_handler._owns_client is True + 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 + 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(