scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-11-15 09:12:18 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 311

fix(router): follow a nested fallback to the deployment that served the stream

Browse source 
A fallback the router picks is itself a fallback-aware wrapper, and it only
repoints at its own fallback once it yields, so reading its hidden params at
selection time named a deployment that produced no output. Re-read them when
the first fallback item arrives, which is still before the proxy commits
response headers.

Also addresses review feedback: the streaming header builder reads self.data
instead of taking a coarse request_data parameter, and the new test recorder
local is Final.
This commit is contained in:
Yucheng He 2026-09-07 01:59:23 -07:00
parent 391e2c33f3
commit 73b00ec992
4 changed files with 231 additions and 26 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

16
litellm/proxy/common_request_processing.py
View file

@ -896,7 +896,6 @@ def _sse_stream_headers(headers: Mapping[str, str]) -> Mapping[str, str]:
async def _resolve_stream_headers(
headers: Mapping[str, str], refresh_headers: Callable[[], Awaitable[Mapping[str, str]]] | None
) -> Mapping[str, str]:
"""The response headers to publish once the first chunk has been buffered."""
if refresh_headers is None:
return headers
try:
@ -2168,11 +2167,10 @@ class ProxyBaseLLMRequestProcessing:
model_id = model_info.get("id", "") or ""
return str(model_id) if model_id else ""
@staticmethod
def _stream_response_headers(
self,
*,
hidden_params: Mapping[str, object],
request_data: dict,
user_api_key_dict: UserAPIKeyAuth,
logging_obj: LiteLLMLoggingObj,
version: str | None,
@ -2184,14 +2182,14 @@ class ProxyBaseLLMRequestProcessing:
**ProxyBaseLLMRequestProcessing.get_custom_headers(
user_api_key_dict=user_api_key_dict,
call_id=logging_obj.litellm_call_id,
model_id=ProxyBaseLLMRequestProcessing._get_model_id_from_response(hidden_params, request_data),
model_id=self._get_model_id_from_response(hidden_params, self.data),
cache_key=hidden_params.get("cache_key") or "",
api_base=hidden_params.get("api_base") or "",
version=version,
response_cost=hidden_params.get("response_cost") or "",
model_region=getattr(user_api_key_dict, "allowed_model_region", ""),
fastest_response_batch_completion=hidden_params.get("fastest_response_batch_completion"),
request_data=request_data,
request_data=self.data,
hidden_params=hidden_params,
litellm_logging_obj=logging_obj,
**(hidden_params.get("additional_headers") or MappingProxyType({})),
@ -2476,7 +2474,6 @@ class ProxyBaseLLMRequestProcessing:
)
custom_headers: Final = self._stream_response_headers(
hidden_params=hidden_params,
request_data=self.data,
user_api_key_dict=user_api_key_dict,
logging_obj=logging_obj,
version=version,
@ -2484,16 +2481,11 @@ class ProxyBaseLLMRequestProcessing:
)
async def refresh_stream_headers() -> Mapping[str, str]:
"""`custom_headers` rebuilt for whichever deployment served the stream.
`create_response` buffers the first chunk before it commits headers, so a
fallback that took over by then has already repointed `response`.
"""
"""`custom_headers` rebuilt for whichever deployment served the stream."""
if not getattr(response, "fallback_headers_adopted", False):
return custom_headers
return self._stream_response_headers(
hidden_params=get_hidden_params_dict(response),
request_data=self.data,
user_api_key_dict=user_api_key_dict,
logging_obj=logging_obj,
version=version,

44
litellm/router.py
View file

@ -2604,6 +2604,18 @@ class Router:
return fallback_hidden_params, {}
return fallback_hidden_params, cast("dict[str, object]", fallback_headers)
@staticmethod
def _adopt_fallback_response_headers(
wrapper_ref: "weakref.ref[FallbackAwareStreamWrapper]",
fallback_response: object,
) -> tuple[dict[str, object], dict[str, object]]:
"""Repoint the wrapper at `fallback_response`, returning its prepared hidden params."""
prepared: Final = Router._prepare_fallback_hidden_params(fallback_response)
adopting_wrapper: Final = wrapper_ref()
if adopting_wrapper is not None:
adopting_wrapper.adopt_fallback_response_headers(fallback_response, prepared)
return prepared
@staticmethod
def _apply_fallback_hidden_params_to_item(
fallback_item: object,
@ -2728,13 +2740,17 @@ class Router:
# If fallback returns a streaming response, iterate over it
if hasattr(fallback_response, "__aiter__"):
prepared_fallback_hidden_params = Router._prepare_fallback_hidden_params(fallback_response)
_adopting_wrapper = wrapper_ref()
if _adopting_wrapper is not None:
_adopting_wrapper.adopt_fallback_response_headers(
fallback_response, prepared_fallback_hidden_params
)
prepared_fallback_hidden_params = Router._adopt_fallback_response_headers(
wrapper_ref, fallback_response
)
fallback_headers_are_settled = False
async for fallback_item in fallback_response:
if not fallback_headers_are_settled:
# a fallback that failed over again only repoints itself once it yields
fallback_headers_are_settled = True
prepared_fallback_hidden_params = Router._adopt_fallback_response_headers(
wrapper_ref, fallback_response
)
Router._apply_fallback_hidden_params_to_item(fallback_item, prepared_fallback_hidden_params)
if (
fallback_item
@ -3272,13 +3288,17 @@ class Router:
)
if hasattr(fallback_response, "__iter__"):
prepared_fallback_hidden_params = Router._prepare_fallback_hidden_params(fallback_response)
_adopting_wrapper = wrapper_ref()
if _adopting_wrapper is not None:
_adopting_wrapper.adopt_fallback_response_headers(
fallback_response, prepared_fallback_hidden_params
)
prepared_fallback_hidden_params = Router._adopt_fallback_response_headers(
wrapper_ref, fallback_response
)
fallback_headers_are_settled = False
for fallback_item in fallback_response:
if not fallback_headers_are_settled:
# a fallback that failed over again only repoints itself once it yields
fallback_headers_are_settled = True
prepared_fallback_hidden_params = Router._adopt_fallback_response_headers(
wrapper_ref, fallback_response
)
Router._apply_fallback_hidden_params_to_item(fallback_item, prepared_fallback_hidden_params)
if (
fallback_item

19
tests/test_litellm/proxy/test_common_request_processing.py
View file

@ -2,7 +2,7 @@ import asyncio
import copy
import datetime
import json
from types import SimpleNamespace
from types import MappingProxyType, SimpleNamespace
from typing import AsyncGenerator, Callable, Final, Optional
from unittest.mock import AsyncMock, MagicMock, patch
@ -1816,7 +1816,7 @@ class TestCommonRequestProcessingHelpers:
headers are still uncommitted, so ``refresh_headers`` is consulted after
the first chunk is buffered and its result wins.
"""
refreshed_at = []
refreshed_at: Final[list[str]] = []
async def mock_generator():
yield 'data: {"content": "data"}\n\n'
@ -8245,3 +8245,18 @@ class TestStreamingResponseHeadersFollowFallback:
assert result.headers["llm_provider-x-request-id"] == "req-SERVED"
assert "llm_provider-stale-marker" not in result.headers
assert result.headers["x-callback-header"] == "kept"
class TestPassthroughHeadersAcceptImmutableMappings:
"""LIT-6767: the streaming branch now hands the passthrough helpers an immutable mapping."""
def test_merge_passthrough_streaming_headers_accepts_a_read_only_mapping(self):
merged = ProxyBaseLLMRequestProcessing._merge_passthrough_streaming_headers(
response_headers=httpx.Headers({"content-type": "text/event-stream", "transfer-encoding": "chunked"}),
custom_headers=MappingProxyType({"x-litellm-model-id": "served-deployment"}),
)
assert merged["x-litellm-model-id"] == "served-deployment"
assert merged["content-type"] == "text/event-stream"
# the excluded hop-by-hop header is still dropped
assert "transfer-encoding" not in merged

178
tests/test_litellm/test_router.py
View file

@ -2682,6 +2682,184 @@ async def test_acompletion_streaming_iterator_adopts_fallback_response_headers()
assert result._hidden_params is not fallback_stream._hidden_params
@pytest.mark.asyncio
async def test_acompletion_streaming_iterator_adopts_the_deployment_that_served_a_nested_fallback():
"""LIT-6767: a fallback that itself fails over before its first chunk.
The selected fallback still describes its own failed attempt at selection time, so
the wrapper has to re-read it once a chunk exists or it publishes a deployment that
produced no output.
"""
from unittest.mock import MagicMock, patch
from litellm.exceptions import MidStreamFallbackError
router = litellm.Router(
model_list=[
{
"model_name": "gpt-4",
"litellm_params": {"model": "gpt-4", "api_key": "fake-key"},
}
],
)
failed_error: Final = MidStreamFallbackError(
message="upstream died before the first chunk",
model="gpt-4",
llm_provider="openai",
generated_content="",
is_pre_first_chunk=True,
)
class FailedStream:
def __init__(self):
self.model = "gpt-4"
self.custom_llm_provider = "openai"
self.logging_obj = MagicMock()
self.chunks = []
self._response_headers = {"x-request-id": "req-FAILED"}
self._hidden_params = {
"model_id": "failed-deployment",
"additional_headers": {"llm_provider-x-request-id": "req-FAILED"},
}
def __aiter__(self):
return self
async def __anext__(self):
raise failed_error
class NestedFallbackStream:
"""A fallback that repoints itself at a third deployment as it yields."""
def __init__(self):
self._response_headers = {"x-request-id": "req-MIDDLE"}
self._hidden_params = {
"model_id": "middle-deployment",
"additional_headers": {"llm_provider-x-request-id": "req-MIDDLE"},
}
self._chunks = iter([litellm.ModelResponseStream(choices=[{"index": 0, "delta": {"content": "OK"}}])])
def __aiter__(self):
return self
async def __anext__(self):
try:
chunk = next(self._chunks)
except StopIteration:
raise StopAsyncIteration from None
self._response_headers = {"x-request-id": "req-SERVED"}
self._hidden_params = {
"model_id": "served-deployment",
"additional_headers": {"llm_provider-x-request-id": "req-SERVED"},
}
return chunk
with patch.object(
router,
"async_function_with_fallbacks_common_utils",
return_value=NestedFallbackStream(),
):
result = await router._acompletion_streaming_iterator(
model_response=FailedStream(),
messages=[{"role": "user", "content": "hi"}],
initial_kwargs={"model": "gpt-4", "stream": True},
)
first_chunk: Final = await result.__anext__()
# the proxy commits response headers once this chunk is buffered
assert result._response_headers == {"x-request-id": "req-SERVED"}
assert result._hidden_params["model_id"] == "served-deployment"
assert result._hidden_params["additional_headers"] == {"llm_provider-x-request-id": "req-SERVED"}
# and the chunk itself carries the same deployment
assert first_chunk._hidden_params["model_id"] == "served-deployment"
async for _ in result:
pass
assert result._response_headers == {"x-request-id": "req-SERVED"}
assert result._hidden_params["model_id"] == "served-deployment"
def test_completion_streaming_iterator_adopts_the_deployment_that_served_a_nested_fallback():
"""LIT-6767, sync counterpart of the nested-fallback adoption test."""
from unittest.mock import MagicMock, patch
from litellm.exceptions import MidStreamFallbackError
router = litellm.Router(
model_list=[
{
"model_name": "gpt-4",
"litellm_params": {"model": "gpt-4", "api_key": "fake-key"},
}
],
)
failed_error: Final = MidStreamFallbackError(
message="upstream died before the first chunk",
model="gpt-4",
llm_provider="openai",
generated_content="",
is_pre_first_chunk=True,
)
class FailedStream:
def __init__(self):
self.model = "gpt-4"
self.custom_llm_provider = "openai"
self.logging_obj = MagicMock()
self.chunks = []
self._response_headers = {"x-request-id": "req-FAILED"}
self._hidden_params = {
"model_id": "failed-deployment",
"additional_headers": {"llm_provider-x-request-id": "req-FAILED"},
}
def __iter__(self):
return self
def __next__(self):
raise failed_error
class NestedFallbackStream:
"""A fallback that repoints itself at a third deployment as it yields."""
def __init__(self):
self._response_headers = {"x-request-id": "req-MIDDLE"}
self._hidden_params = {
"model_id": "middle-deployment",
"additional_headers": {"llm_provider-x-request-id": "req-MIDDLE"},
}
self._chunks = iter([litellm.ModelResponseStream(choices=[{"index": 0, "delta": {"content": "OK"}}])])
def __iter__(self):
return self
def __next__(self):
chunk = next(self._chunks)
self._response_headers = {"x-request-id": "req-SERVED"}
self._hidden_params = {
"model_id": "served-deployment",
"additional_headers": {"llm_provider-x-request-id": "req-SERVED"},
}
return chunk
with patch.object(router, "function_with_fallbacks", return_value=NestedFallbackStream()):
result = router._completion_streaming_iterator(
model_response=FailedStream(),
messages=[{"role": "user", "content": "hi"}],
initial_kwargs={"model": "gpt-4", "stream": True},
)
first_chunk: Final = next(result)
assert result._response_headers == {"x-request-id": "req-SERVED"}
assert result._hidden_params["model_id"] == "served-deployment"
assert first_chunk._hidden_params["model_id"] == "served-deployment"
for _ in result:
pass
assert result._response_headers == {"x-request-id": "req-SERVED"}
assert result._hidden_params["model_id"] == "served-deployment"
def test_completion_streaming_iterator_adopts_fallback_response_headers():
"""LIT-6767, sync counterpart of the fallback-adoption test."""
from unittest.mock import MagicMock, patch