scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-09-09 22:31:41 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 87

fix: return headers for streaming passthrough and use HTTP client in llm_passthrough_route

Browse source
This commit is contained in:
KnyazSh 2026-04-16 17:03:02 +00:00
parent 52796fb060
commit 211388ac51
5 changed files with 181 additions and 119 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

118
litellm/passthrough/main.py
View file

@ -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

35
litellm/proxy/common_request_processing.py
View file

@ -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)

2
tests/pass_through_tests/test_hosted_vllm_passthrough.py
View file

@ -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

95
tests/test_litellm/passthrough/test_async_streaming_error_propagation.py
View file

@ -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()

50
tests/test_litellm/passthrough/test_passthrough_main.py
View file

@ -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"