tiennm99/litellm
1
0
Fork 0
mirror of https://github.com/tiennm99/litellm.git synced 2026-08-21 08:26:34 +00:00
Code Issues Packages Projects Releases Wiki Activity

[Fixes] Aiohttp transport fixes - add handling for aiohttp.ClientPayloadError and ssl_verification settings (#11162)

Browse Source 
* 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
This commit is contained in:
Ishaan Jaff
2025-05-26 21:14:35 -07:00
committed by GitHub
parent e606bfe31d
commit 4d2edc4e7a
5 changed files with 109 additions and 23 deletions
Download Patch File Download Diff File Expand all files Collapse all files
litellm/__init__.py
+1 -1
View File
@@ -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"
litellm/llms/custom_httpx/aiohttp_transport.py
+51 -14
View File
@@ -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", {})
litellm/llms/custom_httpx/http_handler.py
+9 -2
View File
@@ -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,
)
tests/local_testing/test_ollama.py
+18 -6
View File
@@ -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
tests/test_litellm/llms/custom_httpx/test_http_handler.py
+30
View File
@@ -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