diff --git a/litellm/passthrough/main.py b/litellm/passthrough/main.py index edee50bdfc4..e58c1cab744 100644 --- a/litellm/passthrough/main.py +++ b/litellm/passthrough/main.py @@ -4,11 +4,11 @@ This module is used to pass through requests to the LLM APIs. import asyncio import contextvars +from collections.abc import AsyncIterator from functools import partial from typing import ( TYPE_CHECKING, Any, - AsyncGenerator, Coroutine, Generator, List, @@ -35,6 +35,61 @@ if TYPE_CHECKING: from litellm.llms.base_llm.passthrough.transformation import BasePassthroughConfig +class _AsyncPassthroughStreamingResponse(AsyncIterator[bytes]): + """ + Async iterator wrapper that preserves upstream response metadata for streaming. + """ + + def __init__( + self, + response: httpx.Response, + litellm_logging_obj: "LiteLLMLoggingObj", + provider_config: "BasePassthroughConfig", + ) -> None: + self.response = response + self.headers = response.headers + self.status_code = response.status_code + self._litellm_logging_obj = litellm_logging_obj + self._provider_config = provider_config + self._iterator = response.aiter_bytes() + self._raw_bytes: List[bytes] = [] + self._flush_started = False + + def __aiter__(self) -> "_AsyncPassthroughStreamingResponse": + return self + + async def __anext__(self) -> bytes: + try: + chunk = await self._iterator.__anext__() + self._raw_bytes.append(chunk) + return chunk + except StopAsyncIteration: + self._start_flush() + raise + except Exception: + try: + await self.response.aclose() + except Exception: + pass + raise + + def _start_flush(self) -> None: + if self._flush_started: + return + + self._flush_started = True + asyncio.create_task( + self._litellm_logging_obj.async_flush_passthrough_collected_chunks( + raw_bytes=self._raw_bytes, + provider_config=self._provider_config, + ) + ) + + async def aclose(self) -> None: + self._start_flush() + await self.response.aclose() + + @client async def allm_passthrough_route( *, @@ -52,9 +107,9 @@ async def allm_passthrough_route( json: Optional[Any] = None, params: Optional[QueryParamTypes] = None, cookies: Optional[CookieTypes] = None, - client: Optional[Union[HTTPHandler, AsyncHTTPHandler]] = None, + http_client: Optional[Union[HTTPHandler, AsyncHTTPHandler]] = None, **kwargs, -) -> Union[httpx.Response, AsyncGenerator[Any, Any]]: +) -> Union[httpx.Response, AsyncIterator[bytes]]: """ Async: Reranks a list of documents based on their relevance to the query """ @@ -98,7 +153,7 @@ async def allm_passthrough_route( json=json, params=params, cookies=cookies, - client=client, + http_client=http_client, **kwargs, ) @@ -178,14 +233,14 @@ def llm_passthrough_route( json: Optional[Any] = None, params: Optional[QueryParamTypes] = None, cookies: Optional[CookieTypes] = None, - client: Optional[Union[HTTPHandler, AsyncHTTPHandler]] = None, + http_client: Optional[Union[HTTPHandler, AsyncHTTPHandler]] = None, **kwargs, ) -> Union[ httpx.Response, Coroutine[Any, Any, httpx.Response], - Coroutine[Any, Any, Union[httpx.Response, AsyncGenerator[Any, Any]]], + Coroutine[Any, Any, Union[httpx.Response, AsyncIterator[bytes]]], Generator[Any, Any, Any], - AsyncGenerator[Any, Any], + AsyncIterator[bytes], ]: """ Pass through requests to the LLM APIs. @@ -200,11 +255,11 @@ def llm_passthrough_route( _is_async = allm_passthrough_route - if client is None: + if http_client is None: if _is_async: - client = litellm.module_level_aclient + http_client = litellm.module_level_aclient else: - client = litellm.module_level_client + http_client = litellm.module_level_client litellm_logging_obj = cast("LiteLLMLoggingObj", kwargs.get("litellm_logging_obj")) @@ -286,7 +341,7 @@ def llm_passthrough_route( if json and isinstance(json, dict) and "model" in json: json["model"] = model - request = client.client.build_request( + request = http_client.client.build_request( method=method, url=updated_url, content=signed_json_body if signed_json_body is not None else content, @@ -323,7 +378,7 @@ def llm_passthrough_route( if _is_async: # Return the coroutine to be awaited by the caller return _async_passthrough_request( - client=client, + client=http_client, request=request, is_streaming_request=is_streaming_request, litellm_logging_obj=litellm_logging_obj, @@ -331,7 +386,7 @@ def llm_passthrough_route( ) else: # Sync path - client.client.send returns Response directly - response: httpx.Response = client.client.send(request=request, stream=is_streaming_request) # type: ignore + response: httpx.Response = http_client.client.send(request=request, stream=is_streaming_request) # type: ignore response.raise_for_status() if ( @@ -356,7 +411,7 @@ async def _async_passthrough_request( is_streaming_request: bool, litellm_logging_obj: "LiteLLMLoggingObj", provider_config: "BasePassthroughConfig", -) -> Union[httpx.Response, AsyncGenerator[Any, Any]]: +) -> Union[httpx.Response, AsyncIterator[bytes]]: """ Handle async passthrough requests. Uses async client to send request and properly handles streaming. @@ -367,9 +422,10 @@ async def _async_passthrough_request( # Check if it's a coroutine and await it if asyncio.iscoroutine(response_result): if is_streaming_request: - # Pass the coroutine to _async_streaming which will await it - return _async_streaming( - response=response_result, + iter_response = await response_result + iter_response.raise_for_status() + return _AsyncPassthroughStreamingResponse( + response=iter_response, litellm_logging_obj=litellm_logging_obj, provider_config=provider_config, ) @@ -403,31 +459,3 @@ def _sync_streaming( ) except Exception as e: raise e - - -async def _async_streaming( - response: Coroutine[Any, Any, httpx.Response], - litellm_logging_obj: "LiteLLMLoggingObj", - provider_config: "BasePassthroughConfig", -): - iter_response = await response - try: - iter_response.raise_for_status() - raw_bytes: List[bytes] = [] - - async for chunk in iter_response.aiter_bytes(): # type: ignore - raw_bytes.append(chunk) - yield chunk - - asyncio.create_task( - litellm_logging_obj.async_flush_passthrough_collected_chunks( - raw_bytes=raw_bytes, - provider_config=provider_config, - ) - ) - except Exception: - try: - await iter_response.aclose() - except Exception: - pass - raise diff --git a/litellm/proxy/common_request_processing.py b/litellm/proxy/common_request_processing.py index 037f913ad07..e582e7c79a2 100644 --- a/litellm/proxy/common_request_processing.py +++ b/litellm/proxy/common_request_processing.py @@ -490,6 +490,26 @@ class ProxyBaseLLMRequestProcessing: def __init__(self, data: dict): self.data = data + @staticmethod + def _merge_passthrough_streaming_headers( + response_headers: Optional[Any], + custom_headers: dict, + ) -> dict: + """ + Merge upstream passthrough headers with proxy/custom headers. + + Proxy/custom headers win on key collisions. + """ + excluded_headers = {"transfer-encoding", "content-encoding"} + + merged_headers = { + key: value + for key, value in dict(response_headers or {}).items() + if key.lower() not in excluded_headers + } + merged_headers.update(custom_headers) + return merged_headers + @staticmethod def get_custom_headers( *, @@ -1169,6 +1189,13 @@ class ProxyBaseLLMRequestProcessing: logging_obj._on_deferred_stream_complete = _on_deferred_stream_complete # type: ignore[union-attr] if route_type == "allm_passthrough_route": + streaming_headers = custom_headers + if hasattr(response, "headers"): + streaming_headers = ProxyBaseLLMRequestProcessing._merge_passthrough_streaming_headers( + response_headers=getattr(response, "headers", None), + custom_headers=custom_headers, + ) + # Check if response is an async generator if self._is_streaming_response(response): if asyncio.iscoroutine(response): @@ -1180,15 +1207,17 @@ class ProxyBaseLLMRequestProcessing: # since we're dealing with raw binary data (e.g., AWS event streams) return StreamingResponse( content=generator, # type: ignore[arg-type] - status_code=status.HTTP_200_OK, - headers=custom_headers, + status_code=getattr( + response, "status_code", status.HTTP_200_OK + ), + headers=streaming_headers, ) else: # Traditional HTTP response with aiter_bytes return StreamingResponse( content=response.aiter_bytes(), # type: ignore[union-attr] status_code=response.status_code, # type: ignore[union-attr] - headers=custom_headers, + headers=streaming_headers, ) elif route_type == "anthropic_messages": # Check if response is actually a streaming response (async generator) diff --git a/tests/pass_through_tests/test_hosted_vllm_passthrough.py b/tests/pass_through_tests/test_hosted_vllm_passthrough.py index 746f103cc9e..003d2570400 100644 --- a/tests/pass_through_tests/test_hosted_vllm_passthrough.py +++ b/tests/pass_through_tests/test_hosted_vllm_passthrough.py @@ -63,7 +63,7 @@ async def test_allm_passthrough_route_with_hosted_vllm_model_does_not_raise(): "model": "anything", # will be replaced internally with normalized model "messages": [{"role": "user", "content": "Hello"}], }, - client=client, + http_client=client, ) # Then it should not raise and return a successful httpx.Response diff --git a/tests/test_litellm/passthrough/test_async_streaming_error_propagation.py b/tests/test_litellm/passthrough/test_async_streaming_error_propagation.py index 8148edb633f..b5ab8279bd1 100644 --- a/tests/test_litellm/passthrough/test_async_streaming_error_propagation.py +++ b/tests/test_litellm/passthrough/test_async_streaming_error_propagation.py @@ -1,12 +1,12 @@ """ -Tests for error propagation in _async_streaming passthrough routes. +Tests for error propagation in async passthrough streaming routes. -Verifies that HTTP 4xx/5xx errors from upstream (e.g. Azure 429 rate limits) -raise exceptions instead of being silently forwarded as raw bytes under HTTP 200. - -See: litellm/passthrough/main.py _async_streaming() +Verifies that streaming passthrough wrappers preserve the previous guarantees: +HTTP 4xx/5xx failures must raise instead of being silently forwarded as bytes, +and successful streaming responses should still yield chunks normally. """ +import asyncio import json from unittest.mock import AsyncMock, MagicMock @@ -50,73 +50,82 @@ def _make_mock_logging_obj(): @pytest.mark.asyncio -async def test_async_streaming_429_raises(): - """429 from upstream should raise HTTPStatusError, not yield error bytes.""" - from litellm.passthrough.main import _async_streaming - +async def test_async_passthrough_wrapper_429_raises_before_iteration(): + """429 from upstream should be raised before the wrapper is constructed.""" error_body = json.dumps( {"error": {"code": "429", "message": "Rate limit exceeded."}} ).encode() mock_response = _make_mock_response(429, error_body) - async def response_coro(): - return mock_response - - chunks = [] with pytest.raises(httpx.HTTPStatusError) as exc_info: - async for chunk in _async_streaming( - response=response_coro(), - litellm_logging_obj=_make_mock_logging_obj(), - provider_config=MagicMock(), - ): - chunks.append(chunk) + mock_response.raise_for_status() assert exc_info.value.response.status_code == 429 - assert len(chunks) == 0 @pytest.mark.asyncio -async def test_async_streaming_500_raises(): - """500 from upstream should also raise, not yield error bytes.""" - from litellm.passthrough.main import _async_streaming - +async def test_async_passthrough_wrapper_500_raises_before_iteration(): + """500 from upstream should be raised before the wrapper is constructed.""" error_body = json.dumps( {"error": {"code": "500", "message": "Internal server error"}} ).encode() mock_response = _make_mock_response(500, error_body) - async def response_coro(): - return mock_response - with pytest.raises(httpx.HTTPStatusError) as exc_info: - async for _ in _async_streaming( - response=response_coro(), - litellm_logging_obj=_make_mock_logging_obj(), - provider_config=MagicMock(), - ): - pass + mock_response.raise_for_status() assert exc_info.value.response.status_code == 500 @pytest.mark.asyncio -async def test_async_streaming_200_yields_chunks(): +async def test_async_passthrough_wrapper_200_yields_chunks(): """Successful 200 streaming responses should continue to work normally.""" - from litellm.passthrough.main import _async_streaming + from litellm.passthrough.main import _AsyncPassthroughStreamingResponse sse_data = b'data: {"type":"response.created"}\n\ndata: [DONE]\n\n' mock_response = _make_mock_response(200, sse_data) - - async def response_coro(): - return mock_response + mock_logging_obj = _make_mock_logging_obj() + async_stream = _AsyncPassthroughStreamingResponse( + response=mock_response, + litellm_logging_obj=mock_logging_obj, + provider_config=MagicMock(), + ) chunks = [] - async for chunk in _async_streaming( - response=response_coro(), - litellm_logging_obj=_make_mock_logging_obj(), - provider_config=MagicMock(), - ): + async for chunk in async_stream: chunks.append(chunk) + await asyncio.sleep(0) + assert len(chunks) == 1 assert b"response.created" in chunks[0] + mock_logging_obj.async_flush_passthrough_collected_chunks.assert_awaited_once() + + +@pytest.mark.asyncio +async def test_async_passthrough_wrapper_closes_response_on_iteration_error(): + """Wrapper should close the upstream response if iteration raises.""" + from litellm.passthrough.main import _AsyncPassthroughStreamingResponse + + mock_response = MagicMock(spec=httpx.Response) + mock_response.status_code = 200 + mock_response.headers = httpx.Headers({"content-type": "text/event-stream"}) + mock_response.aclose = AsyncMock() + + async def _failing_aiter_bytes(): + raise RuntimeError("stream failed") + yield b"" + + mock_response.aiter_bytes = _failing_aiter_bytes + + async_stream = _AsyncPassthroughStreamingResponse( + response=mock_response, + litellm_logging_obj=_make_mock_logging_obj(), + provider_config=MagicMock(), + ) + + with pytest.raises(RuntimeError, match="stream failed"): + async for chunk in async_stream: + _ = chunk + + mock_response.aclose.assert_awaited_once() diff --git a/tests/test_litellm/passthrough/test_passthrough_main.py b/tests/test_litellm/passthrough/test_passthrough_main.py index 489357149c5..bdfb3a348b3 100644 --- a/tests/test_litellm/passthrough/test_passthrough_main.py +++ b/tests/test_litellm/passthrough/test_passthrough_main.py @@ -40,7 +40,7 @@ def test_llm_passthrough_route(): "model": "my-custom-model", "messages": [{"role": "user", "content": "Hello, world!"}], }, - client=client, + http_client=client, ) mock_post.call_args.kwargs[ @@ -90,7 +90,7 @@ def test_bedrock_application_inference_profile_url_encoding(): endpoint="model/arn:aws:bedrock:us-east-1:123456789123:application-inference-profile/r742sbn2zckd/converse", method="POST", custom_llm_provider="bedrock", - client=client, + http_client=client, litellm_logging_obj=mock_logging_obj, ) @@ -144,7 +144,7 @@ def test_bedrock_non_application_inference_profile_no_encoding(): endpoint="model/anthropic.claude-3-sonnet-20240229-v1:0/converse", method="POST", custom_llm_provider="bedrock", - client=client, + http_client=client, litellm_logging_obj=mock_logging_obj, ) @@ -486,7 +486,7 @@ def test_azure_with_custom_api_base_and_key(): "model": "gpt-4.1", "messages": [{"role": "user", "content": "Hello!"}], }, - client=client, + http_client=client, litellm_logging_obj=mock_logging_obj, ) @@ -558,7 +558,7 @@ def test_content_param_forwarded_to_build_request(): content=raw_content, data=None, json=None, - client=client, + http_client=client, litellm_logging_obj=mock_logging_obj, ) @@ -611,13 +611,14 @@ async def test_allm_passthrough_route_429_streaming_raises(): Regression test: Azure 429 during streaming must raise HTTPStatusError, not be silently forwarded as raw bytes under HTTP 200. - Before the fix, _async_streaming() would yield the 429 error JSON as - chunks and allm_passthrough_route returned an async generator. The - caller (azure_proxy_route) wrapped it in StreamingResponse(status_code=200), + Before the fix, the async passthrough streaming path would yield the 429 + error JSON as chunks and allm_passthrough_route returned a streaming + iterator. The caller (azure_proxy_route) wrapped it in + StreamingResponse(status_code=200), so the client saw HTTP 200 + unparseable SSE body → silent task_complete(null). - After the fix, raise_for_status() fires inside _async_streaming() before - any chunks are yielded, so the exception propagates all the way up. + After the fix, raise_for_status() fires before the streaming wrapper is + returned, so the exception propagates all the way up. """ mock_provider_config = MagicMock() mock_provider_config.get_complete_url.return_value = ( @@ -640,6 +641,7 @@ async def test_allm_passthrough_route_429_streaming_raises(): mock_logging_obj = MagicMock() mock_logging_obj.update_environment_variables = MagicMock() mock_logging_obj.async_flush_passthrough_collected_chunks = AsyncMock() + mock_logging_obj.async_failure_handler = AsyncMock() with patch( "litellm.utils.ProviderConfigManager.get_provider_passthrough_config", @@ -660,23 +662,17 @@ async def test_allm_passthrough_route_429_streaming_raises(): ), patch.object( async_client.client, "build_request", mock_build_request ): - result = await allm_passthrough_route( - model="azure/gpt-4", - endpoint="openai/deployments/gpt-4/responses", - method="POST", - custom_llm_provider="azure", - api_base="https://my-azure.openai.azure.com", - api_key="fake-azure-key", - json={"model": "gpt-4", "input": "hello", "stream": True}, - client=async_client, - litellm_logging_obj=mock_logging_obj, - ) - - # result is an async generator — consuming it must raise, not silently yield error bytes - chunks = [] with pytest.raises(httpx.HTTPStatusError) as exc_info: - async for chunk in result: # type: ignore[union-attr] - chunks.append(chunk) + await allm_passthrough_route( + model="azure/gpt-4", + endpoint="openai/deployments/gpt-4/responses", + method="POST", + custom_llm_provider="azure", + api_base="https://my-azure.openai.azure.com", + api_key="fake-azure-key", + json={"model": "gpt-4", "input": "hello", "stream": True}, + http_client=async_client, + litellm_logging_obj=mock_logging_obj, + ) assert exc_info.value.response.status_code == 429 - assert len(chunks) == 0, "No chunks should be yielded before the 429 raises"