diff --git a/litellm/llms/custom_httpx/aiohttp_transport.py b/litellm/llms/custom_httpx/aiohttp_transport.py index 132191c946c..8e879f77fff 100644 --- a/litellm/llms/custom_httpx/aiohttp_transport.py +++ b/litellm/llms/custom_httpx/aiohttp_transport.py @@ -341,7 +341,7 @@ class LiteLLMAiohttpTransport(AiohttpTransport): return httpx.Response( status_code=response.status, - headers=response.headers, + headers=getattr(response, "raw_headers", response.headers), stream=AiohttpResponseStream(response), request=request, ) diff --git a/tests/test_litellm/test_streaming_connection_cleanup.py b/tests/test_litellm/test_streaming_connection_cleanup.py index 5a81a3ffb17..29331c1984b 100644 --- a/tests/test_litellm/test_streaming_connection_cleanup.py +++ b/tests/test_litellm/test_streaming_connection_cleanup.py @@ -65,6 +65,53 @@ async def test_aiohttp_transport_response_uses_stream_not_content(): assert isinstance(response.stream, AiohttpResponseStream) +@pytest.mark.asyncio +async def test_aiohttp_transport_preserves_non_ascii_response_headers(): + header_value = "视频生成成功" + + class FakeSession: + closed = False + + def __init__(self): + try: + self._loop = asyncio.get_running_loop() + except RuntimeError: + self._loop = None + + def request(self, **kwargs): + class Resp: + status = 200 + headers = {"X-Ark-Message": header_value} + raw_headers = ( + (b"Content-Type", b"text/plain"), + (b"X-Ark-Message", header_value.encode("utf-8")), + ) + + async def __aenter__(self): + return self + + async def __aexit__(self, *args): + pass + + @property + def content(self): + class C: + async def iter_chunked(self, size): + yield b"ok" + + return C() + + return Resp() + + transport = LiteLLMAiohttpTransport(client=lambda: FakeSession()) # type: ignore + response = await transport.handle_async_request( + httpx.Request("GET", "http://example.com") + ) + + assert response.headers["x-ark-message"] == header_value + assert await response.aread() == b"ok" + + @pytest.mark.asyncio async def test_aiohttp_response_stream_aclose_releases_connection(): """AiohttpResponseStream.aclose() must call __aexit__ on the aiohttp response."""