mirror of
https://github.com/BerriAI/litellm.git
synced 2026-09-09 22:31:41 +00:00
fix(router): keep provider response headers on streaming chat completions (#40091)
* fix(router): keep provider response headers on streaming chat completions The Router re-wraps a deployment's CustomStreamWrapper in FallbackStreamWrapper (and its sync twin) so a mid-stream failure can fail over. Neither wrapper forwarded `_response_headers`, so every streaming chat completion handed the proxy's callbacks and its response-header builder a wrapper with no provider headers, and a successful mid-stream fallback still published the failed deployment's identity, `x-request-id` and rate limit counters. Forward `_response_headers` into both wrappers, repoint the wrapper at the deployment that served the stream once a fallback takes over, and rebuild the proxy's response headers from that deployment while `create_response` still has the first chunk buffered. * fix(router): follow a nested fallback to the deployment that served the stream 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. * test(router): cover the fallback header adoption helper directly The router_code_coverage gate wants every router.py function named in a router test, and this also pins the weak-reference behavior: a wrapper collected mid-stream must not break the generator still draining it. * refactor(proxy): take a read-only mapping for the model-id lookup _get_model_id_from_response only reads its request payload, so a Mapping says what it needs and the two metadata hops are narrowed instead of assumed to be dicts. * test: drop mutable recorder locals and routine comments from the new tests An AsyncMock await_count and an asyncio.Event say the same thing as a list and a dict that the test mutates. * chore(router): justify the two rebinds in the fallback loops Both are the one-shot re-read that follows a nested fallback, so they get the repo's rebind-ok note like the rest of the file.
This commit is contained in:
yucheng-berri
GitHub
parent
11702ae4e1
commit
e04e5d7113
5 changed files with 985 additions and 47 deletions
|
|
@ -756,7 +756,7 @@ class _UpstreamClosingStreamingResponse(StreamingResponse):
|
|||
content: AsyncGenerator[str, None],
|
||||
*,
|
||||
media_type: str | None = None,
|
||||
headers: dict | None = None,
|
||||
headers: Mapping[str, str] | None = None,
|
||||
status_code: int = status.HTTP_200_OK,
|
||||
upstream_generator: AsyncGenerator[str, None] | None = None,
|
||||
) -> None:
|
||||
|
|
@ -888,25 +888,39 @@ def _sse_error_frames(error_obj: Mapping[str, object]) -> tuple[str, str]:
|
|||
return f"data: {json.dumps({'error': error_obj})}\n\n", "data: [DONE]\n\n"
|
||||
|
||||
|
||||
def _sse_stream_headers(headers: Mapping[str, str]) -> Mapping[str, str]:
|
||||
"""`headers` plus the two that stop reverse proxies from buffering SSE (issue #28384)."""
|
||||
return MappingProxyType({**headers, **_TTFT_KEEPALIVE_HEADERS})
|
||||
|
||||
|
||||
async def _resolve_stream_headers(
|
||||
headers: Mapping[str, str], refresh_headers: Callable[[], Awaitable[Mapping[str, str]]] | None
|
||||
) -> Mapping[str, str]:
|
||||
if refresh_headers is None:
|
||||
return headers
|
||||
try:
|
||||
return await refresh_headers()
|
||||
except Exception as e: # noqa: BLE001 # a stream whose first chunk is already paid for must not fail over its headers
|
||||
verbose_proxy_logger.exception("Error refreshing streaming response headers: %s", e)
|
||||
return headers
|
||||
|
||||
|
||||
async def create_response(
|
||||
generator: AsyncGenerator[str, None],
|
||||
media_type: str,
|
||||
headers: dict,
|
||||
headers: Mapping[str, str],
|
||||
default_status_code: int = status.HTTP_200_OK,
|
||||
request: Request | None = None,
|
||||
refresh_headers: Callable[[], Awaitable[Mapping[str, str]]] | None = None,
|
||||
) -> StreamingResponse | JSONResponse:
|
||||
"""
|
||||
Create streaming response, checking if the first chunk is an error.
|
||||
If the first chunk is an error, return a standard JSON error response.
|
||||
Otherwise, return StreamingResponse and stream all content.
|
||||
|
||||
``refresh_headers`` is consulted once the first chunk has been buffered, for
|
||||
callers whose headers can only be known then.
|
||||
"""
|
||||
# Tell buffering reverse proxies (nginx, ingress-nginx, Envoy) to flush SSE
|
||||
# immediately instead of releasing the whole stream in one batch (issue #28384).
|
||||
streaming_headers: Final = {
|
||||
**headers,
|
||||
"Cache-Control": "no-cache",
|
||||
"X-Accel-Buffering": "no",
|
||||
}
|
||||
first_chunk_value: str | None = None
|
||||
final_status_code = default_status_code
|
||||
|
||||
|
|
@ -917,6 +931,7 @@ async def create_response(
|
|||
|
||||
# Now get the first chunk from the actual generator
|
||||
first_chunk_value = await _buffer_first_chunk_honoring_disconnect(generator, request)
|
||||
resolved_headers: Final = await _resolve_stream_headers(headers, refresh_headers)
|
||||
|
||||
if first_chunk_value is not None:
|
||||
try:
|
||||
|
|
@ -943,7 +958,7 @@ async def create_response(
|
|||
return JSONResponse(
|
||||
status_code=final_status_code,
|
||||
content={"error": error_dict},
|
||||
headers=headers,
|
||||
headers=resolved_headers,
|
||||
)
|
||||
except Exception as e:
|
||||
verbose_proxy_logger.debug("Error parsing first chunk value: %s", e)
|
||||
|
|
@ -972,7 +987,7 @@ async def create_response(
|
|||
return StreamingResponse(
|
||||
empty_gen(),
|
||||
media_type=media_type,
|
||||
headers=streaming_headers,
|
||||
headers=_sse_stream_headers(await _resolve_stream_headers(headers, refresh_headers)),
|
||||
status_code=default_status_code,
|
||||
)
|
||||
except Exception as e:
|
||||
|
|
@ -988,7 +1003,7 @@ async def create_response(
|
|||
return StreamingResponse(
|
||||
error_gen_message(),
|
||||
media_type=media_type,
|
||||
headers=streaming_headers,
|
||||
headers=_sse_stream_headers(await _resolve_stream_headers(headers, refresh_headers)),
|
||||
status_code=error_status,
|
||||
)
|
||||
|
||||
|
|
@ -1010,7 +1025,7 @@ async def create_response(
|
|||
return _UpstreamClosingStreamingResponse(
|
||||
combined_generator(),
|
||||
media_type=media_type,
|
||||
headers=streaming_headers,
|
||||
headers=_sse_stream_headers(resolved_headers),
|
||||
status_code=final_status_code,
|
||||
upstream_generator=generator,
|
||||
)
|
||||
|
|
@ -1535,7 +1550,7 @@ class ProxyBaseLLMRequestProcessing:
|
|||
@staticmethod
|
||||
def _merge_passthrough_streaming_headers(
|
||||
response_headers: httpx.Headers | dict | None,
|
||||
custom_headers: dict,
|
||||
custom_headers: Mapping[str, str],
|
||||
) -> dict:
|
||||
"""
|
||||
Merge upstream passthrough headers with proxy/custom headers.
|
||||
|
|
@ -2143,14 +2158,45 @@ class ProxyBaseLLMRequestProcessing:
|
|||
return fallback_model_group
|
||||
|
||||
@staticmethod
|
||||
def _get_model_id_from_response(hidden_params: dict, data: dict) -> str:
|
||||
def _get_model_id_from_response(hidden_params: Mapping[str, object], data: Mapping[str, object]) -> str:
|
||||
"""Extract model_id from hidden_params with fallback to litellm_metadata."""
|
||||
model_id = hidden_params.get("model_id", None) or ""
|
||||
if not model_id:
|
||||
litellm_metadata: Final = data.get("litellm_metadata", {}) or {}
|
||||
model_info: Final = litellm_metadata.get("model_info", {}) or {}
|
||||
model_id = model_info.get("id", "") or ""
|
||||
return model_id
|
||||
litellm_metadata: Final = data.get("litellm_metadata")
|
||||
model_info: Final = litellm_metadata.get("model_info") if isinstance(litellm_metadata, Mapping) else None
|
||||
model_id = (model_info.get("id") or "") if isinstance(model_info, Mapping) else ""
|
||||
return str(model_id) if model_id else ""
|
||||
|
||||
def _stream_response_headers(
|
||||
self,
|
||||
*,
|
||||
hidden_params: Mapping[str, object],
|
||||
user_api_key_dict: UserAPIKeyAuth,
|
||||
logging_obj: LiteLLMLoggingObj,
|
||||
version: str | None,
|
||||
callback_headers: Mapping[str, str],
|
||||
) -> Mapping[str, str]:
|
||||
"""The streaming response headers describing `hidden_params`' deployment."""
|
||||
return MappingProxyType(
|
||||
{
|
||||
**ProxyBaseLLMRequestProcessing.get_custom_headers(
|
||||
user_api_key_dict=user_api_key_dict,
|
||||
call_id=logging_obj.litellm_call_id,
|
||||
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=self.data,
|
||||
hidden_params=hidden_params,
|
||||
litellm_logging_obj=logging_obj,
|
||||
**(hidden_params.get("additional_headers") or MappingProxyType({})),
|
||||
),
|
||||
**callback_headers,
|
||||
}
|
||||
)
|
||||
|
||||
@staticmethod
|
||||
def _get_deployment_model_name(
|
||||
|
|
@ -2419,31 +2465,32 @@ class ProxyBaseLLMRequestProcessing:
|
|||
if self._is_streaming_request(
|
||||
data=self.data, is_streaming_request=is_streaming_request
|
||||
) or self._is_streaming_response(response): # use generate_responses to stream responses
|
||||
custom_headers: Final = ProxyBaseLLMRequestProcessing.get_custom_headers(
|
||||
user_api_key_dict=user_api_key_dict,
|
||||
call_id=logging_obj.litellm_call_id,
|
||||
model_id=model_id,
|
||||
cache_key=cache_key,
|
||||
api_base=api_base,
|
||||
version=version,
|
||||
response_cost=response_cost,
|
||||
model_region=getattr(user_api_key_dict, "allowed_model_region", ""),
|
||||
fastest_response_batch_completion=fastest_response_batch_completion,
|
||||
request_data=self.data,
|
||||
hidden_params=hidden_params,
|
||||
litellm_logging_obj=logging_obj,
|
||||
**additional_headers,
|
||||
)
|
||||
|
||||
# Call response headers hook for streaming success
|
||||
callback_headers = await proxy_logging_obj.post_call_response_headers_hook(
|
||||
stream_callback_headers: Final = await proxy_logging_obj.post_call_response_headers_hook(
|
||||
data=self.data,
|
||||
user_api_key_dict=user_api_key_dict,
|
||||
response=response,
|
||||
request_headers=dict(request.headers),
|
||||
)
|
||||
if callback_headers:
|
||||
custom_headers.update(callback_headers)
|
||||
custom_headers: Final = self._stream_response_headers(
|
||||
hidden_params=hidden_params,
|
||||
user_api_key_dict=user_api_key_dict,
|
||||
logging_obj=logging_obj,
|
||||
version=version,
|
||||
callback_headers=stream_callback_headers or MappingProxyType({}),
|
||||
)
|
||||
|
||||
async def refresh_stream_headers() -> Mapping[str, str]:
|
||||
"""`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),
|
||||
user_api_key_dict=user_api_key_dict,
|
||||
logging_obj=logging_obj,
|
||||
version=version,
|
||||
callback_headers=stream_callback_headers or MappingProxyType({}),
|
||||
)
|
||||
|
||||
# Preserve the original client-requested model (pre-alias mapping) for downstream
|
||||
# streaming generators. Pre-call processing can rewrite `self.data["model"]` for
|
||||
|
|
@ -2581,6 +2628,7 @@ class ProxyBaseLLMRequestProcessing:
|
|||
media_type="text/event-stream",
|
||||
headers=custom_headers,
|
||||
request=request,
|
||||
refresh_headers=refresh_stream_headers,
|
||||
)
|
||||
|
||||
### CALL HOOKS ### - modify outgoing data
|
||||
|
|
@ -3032,7 +3080,7 @@ class ProxyBaseLLMRequestProcessing:
|
|||
response: Any,
|
||||
proxy_logging_obj: "ProxyLogging",
|
||||
user_api_key_dict: "UserAPIKeyAuth",
|
||||
custom_headers: dict,
|
||||
custom_headers: Mapping[str, str],
|
||||
request_headers: dict[str, str],
|
||||
) -> Response | None:
|
||||
if not self._has_post_call_guardrails_for_passthrough():
|
||||
|
|
|
|||
|
|
@ -322,7 +322,7 @@ class HttpPassThroughEndpointHelpers(BasePassthroughUtils):
|
|||
def get_response_headers(
|
||||
headers: httpx.Headers,
|
||||
litellm_call_id: str | None = None,
|
||||
custom_headers: dict | None = None,
|
||||
custom_headers: Mapping[str, str] | None = None,
|
||||
) -> dict:
|
||||
# Exclude headers that uvicorn writes itself (server, date) and
|
||||
# encoding/length headers that don't survive re-serialization.
|
||||
|
|
|
|||
|
|
@ -616,6 +616,38 @@ RETRY_BREADCRUMB_EXCLUDED_KWARGS: Final = frozenset(
|
|||
RETRY_BREADCRUMB_LIMIT: Final = 4
|
||||
|
||||
|
||||
class FallbackAwareStreamWrapper(CustomStreamWrapper):
|
||||
"""Base for the Router's chat-completion stream wrappers, which are built around the
|
||||
attempt the Router picked first and have to repoint themselves when a fallback takes over."""
|
||||
|
||||
fallback_headers_adopted: bool = False
|
||||
|
||||
def adopt_fallback_response_headers(
|
||||
self,
|
||||
fallback_response: object,
|
||||
prepared_fallback_hidden_params: tuple[dict[str, object], dict[str, object]],
|
||||
) -> None:
|
||||
"""Repoint this wrapper at the deployment that served the stream.
|
||||
|
||||
Replaces rather than merges, so the failed attempt's `x-request-id`, rate limit
|
||||
counters, `model_id` and `api_base` cannot reach the proxy's response headers or
|
||||
its callbacks.
|
||||
"""
|
||||
self._response_headers = getattr(fallback_response, "_response_headers", None)
|
||||
fallback_hidden_params, fallback_headers = prepared_fallback_hidden_params
|
||||
if fallback_hidden_params:
|
||||
self._hidden_params = { # mutable-ok: the rest of litellm writes into _hidden_params
|
||||
**fallback_hidden_params,
|
||||
# dict() because add_retry_fallback_headers mutates additional_headers in place
|
||||
"additional_headers": dict(fallback_headers), # mutable-ok: see above
|
||||
}
|
||||
self._base_hidden_params = { # mutable-ok: CustomStreamWrapper keeps this snapshot as a dict
|
||||
**self._hidden_params,
|
||||
"response_cost": None,
|
||||
}
|
||||
self.fallback_headers_adopted = True
|
||||
|
||||
|
||||
class Router:
|
||||
model_names: set = set()
|
||||
cache_responses: bool | None = False
|
||||
|
|
@ -2576,6 +2608,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,
|
||||
|
|
@ -2615,7 +2659,7 @@ class Router:
|
|||
|
||||
held_slot: Final = deployment_slot if deployment_slot is not None else contextlib.AsyncExitStack()
|
||||
|
||||
class FallbackStreamWrapper(CustomStreamWrapper):
|
||||
class FallbackStreamWrapper(FallbackAwareStreamWrapper):
|
||||
def __init__(self, async_generator: AsyncGenerator):
|
||||
# Copy attributes from the original model_response
|
||||
super().__init__(
|
||||
|
|
@ -2623,6 +2667,7 @@ class Router:
|
|||
model=model_response.model,
|
||||
custom_llm_provider=model_response.custom_llm_provider,
|
||||
logging_obj=model_response.logging_obj,
|
||||
_response_headers=getattr(model_response, "_response_headers", None),
|
||||
)
|
||||
self._async_generator = async_generator
|
||||
inner_chunks: Final[object] = getattr(model_response, "chunks", None)
|
||||
|
|
@ -2699,8 +2744,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)
|
||||
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:
|
||||
fallback_headers_are_settled = True # rebind-ok: one-shot latch
|
||||
# a fallback that failed over again only repoints itself once it yields
|
||||
prepared_fallback_hidden_params = ( # rebind-ok: re-read once the fallback yields
|
||||
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
|
||||
|
|
@ -2742,7 +2796,11 @@ class Router:
|
|||
e,
|
||||
)
|
||||
|
||||
return FallbackStreamWrapper(stream_with_fallbacks())
|
||||
wrapped_response: Final = FallbackStreamWrapper(stream_with_fallbacks())
|
||||
# weak, so the generator closing over it does not keep the wrapper out of
|
||||
# refcount teardown and delay the `finally` that releases the deployment slot
|
||||
wrapper_ref: Final = weakref.ref(wrapped_response)
|
||||
return wrapped_response
|
||||
|
||||
@staticmethod
|
||||
def _extract_partial_responses_usage(
|
||||
|
|
@ -3171,13 +3229,14 @@ class Router:
|
|||
"""
|
||||
from litellm.exceptions import MidStreamFallbackError
|
||||
|
||||
class SyncFallbackStreamWrapper(CustomStreamWrapper):
|
||||
class SyncFallbackStreamWrapper(FallbackAwareStreamWrapper):
|
||||
def __init__(self, sync_generator: Generator):
|
||||
super().__init__(
|
||||
completion_stream=sync_generator,
|
||||
model=model_response.model,
|
||||
custom_llm_provider=model_response.custom_llm_provider,
|
||||
logging_obj=model_response.logging_obj,
|
||||
_response_headers=getattr(model_response, "_response_headers", None),
|
||||
)
|
||||
self._sync_generator = sync_generator
|
||||
if hasattr(model_response, "_hidden_params"):
|
||||
|
|
@ -3233,8 +3292,17 @@ class Router:
|
|||
)
|
||||
|
||||
if hasattr(fallback_response, "__iter__"):
|
||||
prepared_fallback_hidden_params = Router._prepare_fallback_hidden_params(fallback_response)
|
||||
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:
|
||||
fallback_headers_are_settled = True # rebind-ok: one-shot latch
|
||||
# a fallback that failed over again only repoints itself once it yields
|
||||
prepared_fallback_hidden_params = ( # rebind-ok: re-read once the fallback yields
|
||||
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,7 +3340,10 @@ class Router:
|
|||
close_err,
|
||||
)
|
||||
|
||||
return SyncFallbackStreamWrapper(stream_with_fallbacks())
|
||||
wrapped_response: Final = SyncFallbackStreamWrapper(stream_with_fallbacks())
|
||||
# weak, for the same reason as the async twin
|
||||
wrapper_ref: Final = weakref.ref(wrapped_response)
|
||||
return wrapped_response
|
||||
|
||||
async def _silent_experiment_acompletion(self, silent_model: str, messages: Sequence[Mapping[str, str]], **kwargs):
|
||||
"""
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
||||
|
|
@ -1809,6 +1809,146 @@ class TestCommonRequestProcessingHelpers:
|
|||
response = await create_response(mock_generator(), "text/event-stream", custom_headers)
|
||||
assert response.headers["x-custom-header"] == "TestValue"
|
||||
|
||||
async def test_create_streaming_response_refresh_headers_after_first_chunk(self):
|
||||
"""LIT-6767: headers a caller can only resolve once the first chunk exists.
|
||||
|
||||
A pre-first-chunk fallback replaces the deployment while the response
|
||||
headers are still uncommitted, so ``refresh_headers`` is consulted after
|
||||
the first chunk is buffered and its result wins.
|
||||
"""
|
||||
|
||||
async def mock_generator():
|
||||
yield 'data: {"content": "data"}\n\n'
|
||||
yield "data: [DONE]\n\n"
|
||||
|
||||
refresh_headers: Final = AsyncMock(
|
||||
return_value={"x-litellm-model-id": "fallback-deployment", "llm_provider-x-request-id": "req-FALLBACK"}
|
||||
)
|
||||
|
||||
response = await create_response(
|
||||
mock_generator(),
|
||||
"text/event-stream",
|
||||
{"x-litellm-model-id": "failed-deployment", "llm_provider-x-request-id": "req-FAILED"},
|
||||
refresh_headers=refresh_headers,
|
||||
)
|
||||
assert isinstance(response, StreamingResponse)
|
||||
assert refresh_headers.await_count == 1
|
||||
assert response.headers["x-litellm-model-id"] == "fallback-deployment"
|
||||
assert response.headers["llm_provider-x-request-id"] == "req-FALLBACK"
|
||||
# the buffering headers are still applied on top of the refreshed set
|
||||
assert response.headers["x-accel-buffering"] == "no"
|
||||
assert response.headers["cache-control"] == "no-cache"
|
||||
|
||||
async def test_create_streaming_response_refreshes_only_after_the_first_chunk(self):
|
||||
"""LIT-6767: the refresh has to be consulted after the generator produced a chunk.
|
||||
|
||||
A pre-first-chunk fallback only repoints the response while that first chunk is
|
||||
being produced, so a refresh consulted any earlier still describes the attempt
|
||||
that failed and the headers go out wrong.
|
||||
"""
|
||||
first_chunk_produced: Final = asyncio.Event()
|
||||
|
||||
async def mock_generator():
|
||||
first_chunk_produced.set()
|
||||
yield 'data: {"content": "data"}\n\n'
|
||||
yield "data: [DONE]\n\n"
|
||||
|
||||
async def refresh_headers():
|
||||
served = "fallback-deployment" if first_chunk_produced.is_set() else "failed-deployment"
|
||||
return {"x-litellm-model-id": served}
|
||||
|
||||
response = await create_response(
|
||||
mock_generator(),
|
||||
"text/event-stream",
|
||||
{"x-litellm-model-id": "failed-deployment"},
|
||||
refresh_headers=refresh_headers,
|
||||
)
|
||||
assert response.headers["x-litellm-model-id"] == "fallback-deployment"
|
||||
|
||||
async def test_create_streaming_response_empty_stream_uses_refreshed_headers(self):
|
||||
"""LIT-6767: a fallback that served nothing still gets to name itself.
|
||||
|
||||
The empty-generator branch returns its own StreamingResponse, so it needs the
|
||||
refreshed headers too or the client is told the failed deployment answered.
|
||||
"""
|
||||
|
||||
async def mock_generator():
|
||||
return
|
||||
yield # make it an async generator
|
||||
|
||||
async def refresh_headers():
|
||||
return {"x-litellm-model-id": "fallback-deployment"}
|
||||
|
||||
response = await create_response(
|
||||
mock_generator(),
|
||||
"text/event-stream",
|
||||
{"x-litellm-model-id": "failed-deployment"},
|
||||
refresh_headers=refresh_headers,
|
||||
)
|
||||
assert isinstance(response, StreamingResponse)
|
||||
assert response.headers["x-litellm-model-id"] == "fallback-deployment"
|
||||
assert response.headers["x-accel-buffering"] == "no"
|
||||
|
||||
async def test_create_streaming_response_without_refresh_headers_is_unchanged(self):
|
||||
"""LIT-6767: the default keeps the caller-supplied headers verbatim."""
|
||||
|
||||
async def mock_generator():
|
||||
yield 'data: {"content": "data"}\n\n'
|
||||
yield "data: [DONE]\n\n"
|
||||
|
||||
response = await create_response(
|
||||
mock_generator(),
|
||||
"text/event-stream",
|
||||
{"x-litellm-model-id": "failed-deployment"},
|
||||
)
|
||||
assert response.headers["x-litellm-model-id"] == "failed-deployment"
|
||||
|
||||
async def test_create_streaming_response_refresh_headers_failure_keeps_stream(self):
|
||||
"""LIT-6767: the first chunk is already paid for, so a failing refresh
|
||||
falls back to the caller's headers instead of erroring the stream."""
|
||||
|
||||
async def mock_generator():
|
||||
yield 'data: {"content": "data"}\n\n'
|
||||
yield "data: [DONE]\n\n"
|
||||
|
||||
async def refresh_headers():
|
||||
raise RuntimeError("boom")
|
||||
|
||||
response = await create_response(
|
||||
mock_generator(),
|
||||
"text/event-stream",
|
||||
{"x-litellm-model-id": "failed-deployment"},
|
||||
refresh_headers=refresh_headers,
|
||||
)
|
||||
assert isinstance(response, StreamingResponse)
|
||||
assert response.status_code == status.HTTP_200_OK
|
||||
assert response.headers["x-litellm-model-id"] == "failed-deployment"
|
||||
assert await self.consume_stream(response) == [
|
||||
'data: {"content": "data"}\n\n',
|
||||
"data: [DONE]\n\n",
|
||||
]
|
||||
|
||||
async def test_create_response_first_chunk_error_uses_refreshed_headers(self):
|
||||
"""LIT-6767: the JSON error response built from a bad first chunk carries
|
||||
the refreshed headers too, so it cannot describe a deployment that no
|
||||
longer served the request."""
|
||||
|
||||
async def mock_generator():
|
||||
yield 'data: {"error": {"code": 403, "message": "forbidden"}}\n\n'
|
||||
yield "data: [DONE]\n\n"
|
||||
|
||||
async def refresh_headers():
|
||||
return {"x-litellm-model-id": "fallback-deployment"}
|
||||
|
||||
response = await create_response(
|
||||
mock_generator(),
|
||||
"text/event-stream",
|
||||
{"x-litellm-model-id": "failed-deployment"},
|
||||
refresh_headers=refresh_headers,
|
||||
)
|
||||
assert isinstance(response, JSONResponse)
|
||||
assert response.headers["x-litellm-model-id"] == "fallback-deployment"
|
||||
|
||||
async def test_create_streaming_response_disables_proxy_buffering(self):
|
||||
"""Regression for #28384: every StreamingResponse create_response returns
|
||||
must carry the headers that stop nginx/ingress/Envoy from buffering the
|
||||
|
|
@ -8014,3 +8154,108 @@ class TestDetachedStreamFailureHook:
|
|||
await logging_obj._on_detached_stream_failure(failure)
|
||||
|
||||
assert [call["original_exception"] for call in recorder.calls] == [failure]
|
||||
|
||||
|
||||
class TestStreamingResponseHeadersFollowFallback:
|
||||
"""LIT-6767: the streaming branch has to publish the deployment that served the stream."""
|
||||
|
||||
@staticmethod
|
||||
def _fallback_adopting_stream():
|
||||
class _Stream:
|
||||
def __init__(self) -> None:
|
||||
self._hidden_params = {
|
||||
"model_id": "failed-deployment",
|
||||
"api_base": "http://127.0.0.1:20769/v1",
|
||||
"additional_headers": {"llm_provider-stale-marker": "failed-deployment"},
|
||||
}
|
||||
self.fallback_headers_adopted = False
|
||||
|
||||
def adopt(self) -> None:
|
||||
self._hidden_params = {
|
||||
"model_id": "served-deployment",
|
||||
"api_base": "https://api.openai.com",
|
||||
"additional_headers": {"llm_provider-x-request-id": "req-SERVED"},
|
||||
}
|
||||
self.fallback_headers_adopted = True
|
||||
|
||||
return _Stream()
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_streaming_headers_name_the_deployment_that_served(self, monkeypatch):
|
||||
"""A pre-first-chunk fallback repoints the stream while the headers are still
|
||||
uncommitted, so the published headers must describe the fallback, not the attempt
|
||||
the Router picked first."""
|
||||
stream = self._fallback_adopting_stream()
|
||||
|
||||
def select_data_generator(**kwargs):
|
||||
async def generator():
|
||||
stream.adopt()
|
||||
yield 'data: {"choices": [{"delta": {"content": "OK"}}]}\n\n'
|
||||
yield "data: [DONE]\n\n"
|
||||
|
||||
return generator()
|
||||
|
||||
logging_obj = MagicMock()
|
||||
logging_obj.litellm_call_id = "lit-6767-call"
|
||||
logging_obj._defer_async_logging = False
|
||||
logging_obj._on_deferred_stream_complete = None
|
||||
logging_obj.cost_breakdown = None
|
||||
|
||||
processor = ProxyBaseLLMRequestProcessing(
|
||||
data={"model": "oa-midfail", "stream": True, "litellm_logging_obj": logging_obj}
|
||||
)
|
||||
|
||||
proxy_logging_obj = MagicMock(spec=ProxyLogging)
|
||||
proxy_logging_obj.during_call_hook = AsyncMock(return_value=None)
|
||||
proxy_logging_obj.update_request_status = AsyncMock(return_value=None)
|
||||
proxy_logging_obj.post_call_success_hook = AsyncMock(
|
||||
side_effect=lambda data, user_api_key_dict, response: response
|
||||
)
|
||||
proxy_logging_obj.post_call_response_headers_hook = AsyncMock(
|
||||
return_value={"x-callback-header": "kept"}
|
||||
)
|
||||
|
||||
async def fake_route_request(**kwargs):
|
||||
async def call():
|
||||
return stream
|
||||
|
||||
return call()
|
||||
|
||||
monkeypatch.setattr(
|
||||
litellm.proxy.common_request_processing, "route_request", fake_route_request
|
||||
)
|
||||
|
||||
result = await processor.base_process_llm_request(
|
||||
request=Request(scope={"type": "http", "headers": []}),
|
||||
fastapi_response=Response(),
|
||||
user_api_key_dict=ProxyUserAPIKeyAuth(api_key="sk-test"),
|
||||
route_type="acompletion",
|
||||
proxy_logging_obj=proxy_logging_obj,
|
||||
general_settings={},
|
||||
proxy_config=MagicMock(spec=ProxyConfig),
|
||||
select_data_generator=select_data_generator,
|
||||
is_streaming_request=True,
|
||||
skip_pre_call_logic=True,
|
||||
)
|
||||
|
||||
assert isinstance(result, StreamingResponse)
|
||||
assert result.headers["x-litellm-model-id"] == "served-deployment"
|
||||
assert result.headers["x-litellm-model-api-base"] == "https://api.openai.com"
|
||||
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
|
||||
|
|
|
|||
|
|
@ -2387,6 +2387,580 @@ async def test_acompletion_streaming_iterator_preserves_hidden_params():
|
|||
assert result._hidden_params.get("_response_ms") == 500.0
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_acompletion_streaming_iterator_preserves_response_headers():
|
||||
"""LIT-6767: the returned wrapper must carry the provider's raw response headers.
|
||||
|
||||
Proxy callbacks read ``_response_headers`` off the object the router hands
|
||||
back. The wrapper used to be built without it, so every streaming chat
|
||||
completion reported zero raw provider headers while the non-streaming path
|
||||
reported the full set.
|
||||
"""
|
||||
from unittest.mock import MagicMock
|
||||
|
||||
from litellm.litellm_core_utils.streaming_handler import CustomStreamWrapper
|
||||
|
||||
router = litellm.Router(
|
||||
model_list=[
|
||||
{
|
||||
"model_name": "gpt-4",
|
||||
"litellm_params": {"model": "gpt-4", "api_key": "fake-key"},
|
||||
}
|
||||
],
|
||||
)
|
||||
|
||||
async def _empty():
|
||||
return
|
||||
yield # make it an async generator
|
||||
|
||||
provider_headers = {
|
||||
"x-request-id": "req-provider-123",
|
||||
"x-ratelimit-remaining-requests": "42",
|
||||
# a provider must never be able to spoof an internal header
|
||||
"x-litellm-model-id": "spoofed",
|
||||
}
|
||||
source = CustomStreamWrapper(
|
||||
completion_stream=_empty(),
|
||||
model="gpt-4",
|
||||
custom_llm_provider="openai",
|
||||
logging_obj=MagicMock(),
|
||||
_response_headers=provider_headers,
|
||||
)
|
||||
|
||||
result = await router._acompletion_streaming_iterator(
|
||||
model_response=source,
|
||||
messages=[{"role": "user", "content": "hi"}],
|
||||
initial_kwargs={"model": "gpt-4", "stream": True},
|
||||
)
|
||||
|
||||
assert result._response_headers == provider_headers
|
||||
additional_headers = result._hidden_params["additional_headers"]
|
||||
assert additional_headers["llm_provider-x-request-id"] == "req-provider-123"
|
||||
assert additional_headers["llm_provider-x-ratelimit-remaining-requests"] == "42"
|
||||
# internal-header protection survives: the provider value is namespaced, never promoted
|
||||
assert additional_headers["llm_provider-x-litellm-model-id"] == "spoofed"
|
||||
assert "x-litellm-model-id" not in additional_headers
|
||||
|
||||
|
||||
def test_completion_streaming_iterator_preserves_response_headers():
|
||||
"""LIT-6767, sync counterpart of the async header-preservation test."""
|
||||
from unittest.mock import MagicMock
|
||||
|
||||
from litellm.litellm_core_utils.streaming_handler import CustomStreamWrapper
|
||||
|
||||
router = litellm.Router(
|
||||
model_list=[
|
||||
{
|
||||
"model_name": "gpt-4",
|
||||
"litellm_params": {"model": "gpt-4", "api_key": "fake-key"},
|
||||
}
|
||||
],
|
||||
)
|
||||
|
||||
provider_headers = {"x-request-id": "req-provider-sync", "openai-organization": "org-real"}
|
||||
source = CustomStreamWrapper(
|
||||
completion_stream=iter([]),
|
||||
model="gpt-4",
|
||||
custom_llm_provider="openai",
|
||||
logging_obj=MagicMock(),
|
||||
_response_headers=provider_headers,
|
||||
)
|
||||
|
||||
result = router._completion_streaming_iterator(
|
||||
model_response=source,
|
||||
messages=[{"role": "user", "content": "hi"}],
|
||||
initial_kwargs={"model": "gpt-4", "stream": True},
|
||||
)
|
||||
|
||||
assert result._response_headers == provider_headers
|
||||
assert result._hidden_params["additional_headers"]["llm_provider-x-request-id"] == "req-provider-sync"
|
||||
|
||||
|
||||
def test_adopt_fallback_response_headers_replaces_rather_than_merges():
|
||||
"""LIT-6767: direct unit for FallbackAwareStreamWrapper.adopt_fallback_response_headers.
|
||||
|
||||
Values from the failed attempt must not survive, so the wrapper replaces both
|
||||
``_response_headers`` and ``_hidden_params`` instead of merging them.
|
||||
"""
|
||||
from unittest.mock import MagicMock
|
||||
|
||||
from litellm.router import FallbackAwareStreamWrapper, Router
|
||||
|
||||
wrapper = FallbackAwareStreamWrapper(
|
||||
completion_stream=iter([]),
|
||||
model="gpt-4",
|
||||
custom_llm_provider="openai",
|
||||
logging_obj=MagicMock(),
|
||||
_response_headers={"x-request-id": "req-FAILED"},
|
||||
)
|
||||
wrapper._hidden_params = {
|
||||
"model_id": "failed-deployment",
|
||||
"only_on_failed_attempt": "stale",
|
||||
"additional_headers": {"llm_provider-x-request-id": "req-FAILED"},
|
||||
}
|
||||
|
||||
fallback = MagicMock()
|
||||
fallback._response_headers = {"x-request-id": "req-FALLBACK"}
|
||||
fallback._hidden_params = {
|
||||
"model_id": "fallback-deployment",
|
||||
"additional_headers": {"llm_provider-x-request-id": "req-FALLBACK"},
|
||||
}
|
||||
|
||||
wrapper.adopt_fallback_response_headers(
|
||||
fallback, Router._prepare_fallback_hidden_params(fallback)
|
||||
)
|
||||
|
||||
assert wrapper._response_headers == {"x-request-id": "req-FALLBACK"}
|
||||
assert wrapper._hidden_params["model_id"] == "fallback-deployment"
|
||||
assert "only_on_failed_attempt" not in wrapper._hidden_params
|
||||
assert wrapper._hidden_params is not fallback._hidden_params
|
||||
# the snapshot CustomStreamWrapper caches at init has to follow, or a chunk built
|
||||
# from it would still be stamped with the deployment that failed
|
||||
assert wrapper._base_hidden_params["model_id"] == "fallback-deployment"
|
||||
# the nested header dict is copied too, so a later mutation on the fallback
|
||||
# response cannot reach headers the proxy has already published
|
||||
assert wrapper._hidden_params["additional_headers"] is not fallback._hidden_params["additional_headers"]
|
||||
fallback._hidden_params["additional_headers"]["llm_provider-x-request-id"] = "req-MUTATED"
|
||||
assert wrapper._hidden_params["additional_headers"] == {"llm_provider-x-request-id": "req-FALLBACK"}
|
||||
|
||||
|
||||
def test_adopt_fallback_response_headers_survives_a_collected_wrapper():
|
||||
"""LIT-6767: adoption still returns the fallback's params once the wrapper is gone."""
|
||||
import weakref
|
||||
from unittest.mock import MagicMock
|
||||
|
||||
from litellm.router import FallbackAwareStreamWrapper, Router
|
||||
|
||||
fallback: Final = MagicMock()
|
||||
fallback._response_headers = {"x-request-id": "req-FALLBACK"}
|
||||
fallback._hidden_params = {
|
||||
"model_id": "fallback-deployment",
|
||||
"additional_headers": {"llm_provider-x-request-id": "req-FALLBACK"},
|
||||
}
|
||||
|
||||
wrapper = FallbackAwareStreamWrapper(
|
||||
completion_stream=iter([]),
|
||||
model="gpt-4",
|
||||
custom_llm_provider="openai",
|
||||
logging_obj=MagicMock(),
|
||||
)
|
||||
live_ref: Final = weakref.ref(wrapper)
|
||||
prepared: Final = Router._adopt_fallback_response_headers(live_ref, fallback)
|
||||
assert prepared == (fallback._hidden_params, fallback._hidden_params["additional_headers"])
|
||||
assert wrapper.fallback_headers_adopted is True
|
||||
assert wrapper._response_headers == {"x-request-id": "req-FALLBACK"}
|
||||
|
||||
dead_ref: Final = weakref.ref(wrapper)
|
||||
del wrapper
|
||||
assert dead_ref() is None
|
||||
assert Router._adopt_fallback_response_headers(dead_ref, fallback) == prepared
|
||||
|
||||
|
||||
def test_adopt_fallback_response_headers_drops_headers_the_fallback_cannot_replace():
|
||||
"""LIT-6767: a fallback that carries no raw provider headers publishes none.
|
||||
|
||||
Keeping the failed attempt's raw headers would hand the client and the callbacks a
|
||||
provider ``x-request-id`` for a request that deployment never served, which is the
|
||||
leak this fix exists to close.
|
||||
"""
|
||||
from unittest.mock import MagicMock
|
||||
|
||||
from litellm.router import FallbackAwareStreamWrapper, Router
|
||||
|
||||
wrapper = FallbackAwareStreamWrapper(
|
||||
completion_stream=iter([]),
|
||||
model="gpt-4",
|
||||
custom_llm_provider="openai",
|
||||
logging_obj=MagicMock(),
|
||||
_response_headers={"x-request-id": "req-FAILED"},
|
||||
)
|
||||
wrapper._hidden_params = {"model_id": "failed-deployment", "additional_headers": {}}
|
||||
|
||||
fallback = MagicMock()
|
||||
fallback._response_headers = None
|
||||
fallback._hidden_params = {"model_id": "fallback-deployment"}
|
||||
|
||||
wrapper.adopt_fallback_response_headers(
|
||||
fallback, Router._prepare_fallback_hidden_params(fallback)
|
||||
)
|
||||
|
||||
assert wrapper._response_headers is None
|
||||
assert wrapper._hidden_params["model_id"] == "fallback-deployment"
|
||||
assert wrapper.fallback_headers_adopted is True
|
||||
|
||||
|
||||
def test_adopt_fallback_response_headers_keeps_identity_when_fallback_has_none():
|
||||
"""A fallback response carrying no hidden params keeps the identity headers.
|
||||
|
||||
Publishing no ``x-litellm-*`` header at all for a request the fallback served is
|
||||
worse than keeping what is there, so only the raw provider headers are dropped.
|
||||
"""
|
||||
from unittest.mock import MagicMock
|
||||
|
||||
from litellm.router import FallbackAwareStreamWrapper, Router
|
||||
|
||||
wrapper = FallbackAwareStreamWrapper(
|
||||
completion_stream=iter([]),
|
||||
model="gpt-4",
|
||||
custom_llm_provider="openai",
|
||||
logging_obj=MagicMock(),
|
||||
_response_headers={"x-request-id": "req-FAILED"},
|
||||
)
|
||||
hidden_params_before = wrapper._hidden_params
|
||||
|
||||
fallback = object()
|
||||
wrapper.adopt_fallback_response_headers(
|
||||
fallback, Router._prepare_fallback_hidden_params(fallback)
|
||||
)
|
||||
|
||||
assert wrapper._response_headers is None
|
||||
assert wrapper._hidden_params is hidden_params_before
|
||||
assert wrapper.fallback_headers_adopted is True
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_acompletion_streaming_iterator_adopts_fallback_response_headers():
|
||||
"""LIT-6767: after a successful pre-first-chunk fallback, the wrapper must
|
||||
describe the deployment that served the stream, with no value left over
|
||||
from the attempt that failed."""
|
||||
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 = 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",
|
||||
"api_base": "https://failed.example",
|
||||
"additional_headers": {"llm_provider-x-request-id": "req-FAILED"},
|
||||
"only_on_failed_attempt": "stale",
|
||||
}
|
||||
|
||||
def __aiter__(self):
|
||||
return self
|
||||
|
||||
async def __anext__(self):
|
||||
raise failed_error
|
||||
|
||||
class FallbackStream:
|
||||
def __init__(self):
|
||||
self._response_headers = {"x-request-id": "req-FALLBACK"}
|
||||
self._hidden_params = {
|
||||
"model_id": "fallback-deployment",
|
||||
"api_base": "https://fallback.example",
|
||||
"additional_headers": {"llm_provider-x-request-id": "req-FALLBACK"},
|
||||
}
|
||||
self._chunks = iter([litellm.ModelResponseStream(choices=[{"index": 0, "delta": {"content": "OK"}}])])
|
||||
|
||||
def __aiter__(self):
|
||||
return self
|
||||
|
||||
async def __anext__(self):
|
||||
try:
|
||||
return next(self._chunks)
|
||||
except StopIteration:
|
||||
raise StopAsyncIteration from None
|
||||
|
||||
fallback_stream = FallbackStream()
|
||||
with patch.object(
|
||||
router,
|
||||
"async_function_with_fallbacks_common_utils",
|
||||
return_value=fallback_stream,
|
||||
):
|
||||
result = await router._acompletion_streaming_iterator(
|
||||
model_response=FailedStream(),
|
||||
messages=[{"role": "user", "content": "hi"}],
|
||||
initial_kwargs={"model": "gpt-4", "stream": True},
|
||||
)
|
||||
# the failed attempt is what the wrapper is built from
|
||||
assert result._response_headers == {"x-request-id": "req-FAILED"}
|
||||
# the very first chunk the fallback produces must already be published under
|
||||
# the fallback's identity: the proxy commits response headers once that chunk
|
||||
# is buffered, so adopting any later is adopting too late
|
||||
await result.__anext__()
|
||||
assert result._response_headers == {"x-request-id": "req-FALLBACK"}
|
||||
assert result._hidden_params["model_id"] == "fallback-deployment"
|
||||
async for _ in result:
|
||||
pass
|
||||
|
||||
assert result._response_headers == {"x-request-id": "req-FALLBACK"}
|
||||
assert result._hidden_params["model_id"] == "fallback-deployment"
|
||||
assert result._hidden_params["api_base"] == "https://fallback.example"
|
||||
assert result._hidden_params["additional_headers"] == {"llm_provider-x-request-id": "req-FALLBACK"}
|
||||
# stale values are removed, not merged over
|
||||
assert "only_on_failed_attempt" not in result._hidden_params
|
||||
# and the wrapper holds its own copy, so later fallback mutations cannot leak in
|
||||
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
|
||||
|
||||
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 = 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"},
|
||||
"only_on_failed_attempt": "stale",
|
||||
}
|
||||
|
||||
def __iter__(self):
|
||||
return self
|
||||
|
||||
def __next__(self):
|
||||
raise failed_error
|
||||
|
||||
class FallbackStream:
|
||||
def __init__(self):
|
||||
self._response_headers = {"x-request-id": "req-FALLBACK"}
|
||||
self._hidden_params = {
|
||||
"model_id": "fallback-deployment",
|
||||
"additional_headers": {"llm_provider-x-request-id": "req-FALLBACK"},
|
||||
}
|
||||
|
||||
def __iter__(self):
|
||||
return iter([])
|
||||
|
||||
with patch.object(router, "function_with_fallbacks", return_value=FallbackStream()):
|
||||
result = router._completion_streaming_iterator(
|
||||
model_response=FailedStream(),
|
||||
messages=[{"role": "user", "content": "hi"}],
|
||||
initial_kwargs={"model": "gpt-4", "stream": True},
|
||||
)
|
||||
assert result._response_headers == {"x-request-id": "req-FAILED"}
|
||||
for _ in result:
|
||||
pass
|
||||
|
||||
assert result._response_headers == {"x-request-id": "req-FALLBACK"}
|
||||
assert result._hidden_params["model_id"] == "fallback-deployment"
|
||||
assert "only_on_failed_attempt" not in result._hidden_params
|
||||
|
||||
|
||||
def test_completion_streaming_iterator_fallback_on_429():
|
||||
"""Sync streaming: MidStreamFallbackError (429 pre-first-chunk) triggers fallback.
|
||||
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue