mirror of
https://github.com/BerriAI/litellm.git
synced 2026-10-27 06:13:07 +00:00
fix(http): isolate global sessions from module-level URL fetchers and satisfy strict type checks
This commit is contained in:
parent
740be08822
commit
963c86ec57
2 changed files with 75 additions and 53 deletions
|
|
@ -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
|
||||
|
||||
|
|
|
|||
|
|
@ -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(
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue