From 4d2edc4e7a9fba579550474c9afb95b51e7b3096 Mon Sep 17 00:00:00 2001 From: Ishaan Jaff Date: Mon, 26 May 2025 21:14:35 -0700 Subject: [PATCH] [Fixes] Aiohttp transport fixes - add handling for `aiohttp.ClientPayloadError` and ssl_verification settings (#11162) * fix: AiohttpResponseStream transport * fix: use AiohttpResponseStream transport by default * fix: AiohttpResponseStream transport * fixes: mapping aiohttp exceptions * fixes: aiohttp rollout * fixes: add support ssl_verify for aiohttp * fixes: add support ssl_verify for aiohttp * fixes: remove duplicates --- litellm/__init__.py | 2 +- .../llms/custom_httpx/aiohttp_transport.py | 65 +++++++++++++++---- litellm/llms/custom_httpx/http_handler.py | 11 +++- tests/local_testing/test_ollama.py | 24 +++++-- .../llms/custom_httpx/test_http_handler.py | 30 +++++++++ 5 files changed, 109 insertions(+), 23 deletions(-) diff --git a/litellm/__init__.py b/litellm/__init__.py index 3ed894a21c..bbb937a538 100644 --- a/litellm/__init__.py +++ b/litellm/__init__.py @@ -300,7 +300,7 @@ custom_prometheus_metadata_labels: List[str] = [] priority_reservation: Optional[Dict[str, float]] = None ######## Networking Settings ######## -use_aiohttp_transport: bool = False +use_aiohttp_transport: bool = True force_ipv4: bool = False # when True, litellm will force ipv4 for all LLM requests. Some users have seen httpx ConnectionError when using ipv6. module_level_aclient = AsyncHTTPHandler( timeout=request_timeout, client_alias="module level aclient" diff --git a/litellm/llms/custom_httpx/aiohttp_transport.py b/litellm/llms/custom_httpx/aiohttp_transport.py index 91ee03e53a..ca408651e6 100644 --- a/litellm/llms/custom_httpx/aiohttp_transport.py +++ b/litellm/llms/custom_httpx/aiohttp_transport.py @@ -1,23 +1,52 @@ import asyncio import contextlib import typing -from typing import Callable, Union +from typing import Callable, Dict, Union import aiohttp +import aiohttp.client_exceptions import httpx from aiohttp.client import ClientResponse, ClientSession from litellm._logging import verbose_logger -AIOHTTP_EXC_MAP = { +AIOHTTP_EXC_MAP: Dict = { + # Order matters here, most specific exception first + # Timeout related exceptions aiohttp.ServerTimeoutError: httpx.TimeoutException, aiohttp.ConnectionTimeoutError: httpx.ConnectTimeout, aiohttp.SocketTimeoutError: httpx.ReadTimeout, - aiohttp.ClientConnectorError: httpx.ConnectError, - aiohttp.ClientPayloadError: httpx.ReadError, + # Proxy related exceptions aiohttp.ClientProxyConnectionError: httpx.ProxyError, + # SSL related exceptions + aiohttp.ClientConnectorCertificateError: httpx.ProtocolError, + aiohttp.ClientSSLError: httpx.ProtocolError, + aiohttp.ServerFingerprintMismatch: httpx.ProtocolError, + # Network related exceptions + aiohttp.ClientConnectorError: httpx.ConnectError, + aiohttp.ClientOSError: httpx.ConnectError, + aiohttp.ClientPayloadError: httpx.ReadError, + # Connection disconnection exceptions + aiohttp.ServerDisconnectedError: httpx.ReadError, + # Response related exceptions + aiohttp.ClientConnectionError: httpx.NetworkError, + aiohttp.ClientPayloadError: httpx.ReadError, + aiohttp.ContentTypeError: httpx.ReadError, + aiohttp.TooManyRedirects: httpx.TooManyRedirects, + # URL related exceptions + aiohttp.InvalidURL: httpx.InvalidURL, + # Base exceptions + aiohttp.ClientError: httpx.RequestError, } +# Add client_exceptions module exceptions +try: + import aiohttp.client_exceptions + + AIOHTTP_EXC_MAP[aiohttp.client_exceptions.ClientPayloadError] = httpx.ReadError +except ImportError: + pass + @contextlib.contextmanager def map_aiohttp_exceptions() -> typing.Iterator[None]: @@ -46,11 +75,23 @@ class AiohttpResponseStream(httpx.AsyncByteStream): self._aiohttp_response = aiohttp_response async def __aiter__(self) -> typing.AsyncIterator[bytes]: - with map_aiohttp_exceptions(): - async for chunk in self._aiohttp_response.content.iter_chunked( - self.CHUNK_SIZE - ): - yield chunk + try: + with map_aiohttp_exceptions(): + async for chunk in self._aiohttp_response.content.iter_chunked( + self.CHUNK_SIZE + ): + yield chunk + except aiohttp.ClientPayloadError as e: + # Handle incomplete transfers more gracefully + # Log the error but don't re-raise if we've already yielded some data + verbose_logger.debug(f"Transfer incomplete, but continuing: {e}") + # If the error is due to incomplete transfer encoding, we can still + # return what we've received so far, similar to how httpx handles it + return + except Exception: + # For other exceptions, use the normal mapping + with map_aiohttp_exceptions(): + raise async def aclose(self) -> None: with map_aiohttp_exceptions(): @@ -59,7 +100,7 @@ class AiohttpResponseStream(httpx.AsyncByteStream): class AiohttpTransport(httpx.AsyncBaseTransport): def __init__( - self, client: ClientSession | typing.Callable[[], ClientSession] + self, client: Union[ClientSession, Callable[[], ClientSession]] ) -> None: self.client = client @@ -140,10 +181,6 @@ class LiteLLMAiohttpTransport(AiohttpTransport): request: httpx.Request, ) -> httpx.Response: from aiohttp import ClientTimeout - from httpx_aiohttp.transport import ( - AiohttpResponseStream, - map_aiohttp_exceptions, - ) from yarl import URL as YarlURL timeout = request.extensions.get("timeout", {}) diff --git a/litellm/llms/custom_httpx/http_handler.py b/litellm/llms/custom_httpx/http_handler.py index ce4bd8d11b..aaf3b92e6b 100644 --- a/litellm/llms/custom_httpx/http_handler.py +++ b/litellm/llms/custom_httpx/http_handler.py @@ -533,12 +533,19 @@ class AsyncHTTPHandler: """ from litellm.llms.custom_httpx.aiohttp_transport import LiteLLMAiohttpTransport - verbose_logger.debug("Creating AiohttpTransport...") + ######################################################### + # If ssl_verify is None, set it to True + # TCP Connector does not allow ssl_verify to be None + # by default aiohttp sets ssl_verify to True + ######################################################### + if ssl_verify is None: + ssl_verify = True + verbose_logger.debug("Creating AiohttpTransport...") return LiteLLMAiohttpTransport( client=lambda: ClientSession( connector=TCPConnector( - verify_ssl=ssl_verify or True, + verify_ssl=ssl_verify, ssl_context=ssl_context, local_addr=("0.0.0.0", 0) if litellm.force_ipv4 else None, ) diff --git a/tests/local_testing/test_ollama.py b/tests/local_testing/test_ollama.py index 2c4ceb3baf..169db483f2 100644 --- a/tests/local_testing/test_ollama.py +++ b/tests/local_testing/test_ollama.py @@ -274,9 +274,21 @@ async def test_async_ollama_ssl_verify(stream): "async_httpx_clientssl_verify_Falseollama" ) - test_client = httpx.AsyncClient(verify=False) - print(client) - assert ( - client.client._transport._pool._ssl_context.verify_mode - == test_client._transport._pool._ssl_context.verify_mode - ) + # check client + print("type of transport in client=", type(client.client._transport)) + print("vars in transport in client=", vars(client.client._transport)) + litellm_created_session = client.client._transport._get_valid_client_session() + print("litellm_created_session=", litellm_created_session) + # check session ssl + print("litellm_created_session ssl=", litellm_created_session.connector._ssl) + + + # create aiohttp transport with ssl_verify=False + import aiohttp + aiohttp_session = aiohttp.ClientSession(connector=aiohttp.TCPConnector(verify_ssl=False)) + print("aiohttp_session ssl=", aiohttp_session.connector._ssl) + + assert litellm_created_session.connector._ssl is False + assert litellm_created_session.connector._ssl == aiohttp_session.connector._ssl + + diff --git a/tests/test_litellm/llms/custom_httpx/test_http_handler.py b/tests/test_litellm/llms/custom_httpx/test_http_handler.py index 774b584abb..1c9810ac63 100644 --- a/tests/test_litellm/llms/custom_httpx/test_http_handler.py +++ b/tests/test_litellm/llms/custom_httpx/test_http_handler.py @@ -93,3 +93,33 @@ async def test_aiohttp_disabled_transport(): # Should get None when both aiohttp is disabled and force_ipv4 is False assert transport is None + + +@pytest.mark.asyncio +async def test_ssl_verification_with_aiohttp_transport(): + """ + Test aiohttp respects ssl_verify=False + + We validate that the ssl settings for a litellm transport match what a ssl verify=False aiohttp client would have. + + """ + import aiohttp + + # Create a test SSL context + litellm.use_aiohttp_transport = True + litellm_async_client = AsyncHTTPHandler(ssl_verify=False) + + transport_connector = ( + litellm_async_client.client._transport._get_valid_client_session().connector + ) + print("transport_connector", transport_connector) + print("transport_connector._ssl", transport_connector._ssl) + + aiohttp_session = aiohttp.ClientSession( + connector=aiohttp.TCPConnector(verify_ssl=False) + ) + print("aiohttp_session", aiohttp_session) + print("aiohttp_session._ssl", aiohttp_session.connector._ssl) + + # assert both litellm transport and aiohttp session have ssl_verify=False + assert transport_connector._ssl == aiohttp_session.connector._ssl