diff --git a/litellm/proxy/common_request_processing.py b/litellm/proxy/common_request_processing.py index 0695a8ec9d3..bf8cece5bca 100644 --- a/litellm/proxy/common_request_processing.py +++ b/litellm/proxy/common_request_processing.py @@ -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, diff --git a/litellm/router.py b/litellm/router.py index 43cc6585582..1e3485f7eb7 100644 --- a/litellm/router.py +++ b/litellm/router.py @@ -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 diff --git a/tests/test_litellm/proxy/test_common_request_processing.py b/tests/test_litellm/proxy/test_common_request_processing.py index 018b0afed91..9fcc680b98e 100644 --- a/tests/test_litellm/proxy/test_common_request_processing.py +++ b/tests/test_litellm/proxy/test_common_request_processing.py @@ -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 diff --git a/tests/test_litellm/test_router.py b/tests/test_litellm/test_router.py index a8dc161f958..7d93758d4ad 100644 --- a/tests/test_litellm/test_router.py +++ b/tests/test_litellm/test_router.py @@ -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