diff --git a/basedpyright-code-budget.json b/basedpyright-code-budget.json index da788bf1ce3..694df540544 100644 --- a/basedpyright-code-budget.json +++ b/basedpyright-code-budget.json @@ -57,7 +57,7 @@ "limit": 5601 }, "reportMissingTypeArgument": { - "limit": 15306 + "limit": 15304 }, "reportMissingTypeStubs": { "limit": 40 @@ -99,7 +99,7 @@ "limit": 0 }, "reportUnknownArgumentType": { - "limit": 44364 + "limit": 44353 }, "reportUnknownLambdaType": { "limit": 109 @@ -108,10 +108,10 @@ "limit": 38350 }, "reportUnknownParameterType": { - "limit": 19626 + "limit": 19624 }, "reportUnknownVariableType": { - "limit": 29890 + "limit": 29883 }, "reportUnnecessaryCast": { "limit": 111 diff --git a/cookbook/logging_observability/response_header_logger.md b/cookbook/logging_observability/response_header_logger.md new file mode 100644 index 00000000000..013fdaf147a --- /dev/null +++ b/cookbook/logging_observability/response_header_logger.md @@ -0,0 +1,55 @@ +# Store upstream response headers in spend logs + +Copy `response_header_logger.py` next to your proxy configuration as `custom_callbacks.py`, then register the logger: + +```yaml +model_list: + - model_name: openrouter-header-test + litellm_params: + model: openrouter/openai/gpt-5.6-luna + api_key: os.environ/OPENROUTER_API_KEY +litellm_settings: + callbacks: [custom_callbacks.proxy_handler_instance] +general_settings: + master_key: os.environ/LITELLM_MASTER_KEY + database_url: os.environ/DATABASE_URL + proxy_batch_write_at: 1 +``` + +Start the proxy with this repository's code and your configuration: + +```sh +uv run --no-sync litellm --config config.yaml --port 4013 +``` + +The logger captures response headers before stream iteration, enriches successful spend logs through `async_logging_hook`, and recovers failure headers from the exception chain. A failure without recoverable headers retains the earlier capture. When a fallback response is captured, its headers replace the failed attempt’s headers. Request state stays in request metadata, so concurrent calls do not share headers on the logger instance + +Send a request with existing custom metadata: + +```sh +curl --fail-with-body -sS -N http://localhost:4013/v1/chat/completions \ + -H "Authorization: Bearer $LITELLM_MASTER_KEY" \ + -H 'Content-Type: application/json' \ + -d '{"model":"openrouter-header-test","messages":[{"role":"user","content":"Reply exactly OK"}],"stream":true,"max_tokens":32,"metadata":{"spend_logs_metadata":{"verification_case":"stream","existing_field":"keep-me"}}}' +``` + +After the background write, read the stored headers: + +```sh +curl --fail-with-body -sS http://localhost:4013/spend/logs \ + -H "Authorization: Bearer $LITELLM_MASTER_KEY" | + jq '.[] | . as $row | + (.metadata | if type == "string" then fromjson else . end).spend_logs_metadata as $m | + {request_id: $row.request_id, case: $m.verification_case, + existing_field: $m.existing_field, headers: $m.upstream_response_headers}' +``` + +Headers appear in `metadata.spend_logs_metadata.upstream_response_headers`. Provider names receive the `llm_provider-` prefix; existing custom metadata and LiteLLM's internal-header protections are preserved + +Set `stream` to false for the non-streaming control. To reproduce an upstream routing failure, add `"num_retries":0,"provider":{"only":["header-verification-nonexistent-provider"],"allow_fallbacks":false}` to the request. This returned HTTP 404 in verification + +The logger was verified with OpenRouter chat completions, including simultaneous streams, a non-streaming control, and an upstream 404. The unit tests also cover a headerless error after header capture and both metadata keys used by the logging pipeline. Live commands, spend-log/database readbacks, and the exact tested revisions are in [PR #39376](https://github.com/BerriAI/litellm/pull/39376) + +An earlier logger without early capture lost headers on a live streaming 429. That 429 did not recur after the logger update; retention for that error shape is covered by a unit test, not a second live 429 reproduction + +Unmodified v1.98.0 still needs the shared HTTP-handler propagation fix as well as the router fix. The logger alone cannot recover headers dropped before the callback. These are headers returned by OpenRouter, which may differ from those of the provider behind it. Other providers, every retry/fallback attempt, and every failure type have not been live-tested diff --git a/cookbook/logging_observability/response_header_logger.py b/cookbook/logging_observability/response_header_logger.py new file mode 100644 index 00000000000..749e337fbca --- /dev/null +++ b/cookbook/logging_observability/response_header_logger.py @@ -0,0 +1,86 @@ +from collections.abc import Mapping +from typing import Final, cast + +import httpx + +from litellm.integrations.custom_logger import CustomLogger +from litellm.proxy._types import UserAPIKeyAuth + + +def _as_dict(value: object) -> dict[str, object]: + return cast(dict[str, object], value) if isinstance(value, dict) else {} + + +def _child(parent: dict[str, object], key: str) -> dict[str, object]: + child: Final = _as_dict(parent.get(key)) + parent[key] = child # rebind-ok: callback metadata must be updated in place + return child + + +def _headers(value: object) -> dict[str, str]: + if isinstance(value, httpx.Headers): + return dict(value) + return {name: header for name, header in _as_dict(value).items() if isinstance(header, str)} + + +def get_error_headers(error: BaseException | None, seen: tuple[int, ...] = ()) -> dict[str, str]: + if error is None or id(error) in seen: + return {} + response: Final[object] = getattr(error, "response", None) + headers: Final = _headers(getattr(response, "headers", None)) or _headers(getattr(error, "headers", None)) + if headers: + return headers + visited: Final = (*seen, id(error)) + return get_error_headers(error.__cause__, visited) or get_error_headers(error.__context__, visited) + + +def save_headers(metadata: dict[str, object], headers: Mapping[str, object]) -> None: + existing: Final = _as_dict(metadata.get("spend_logs_metadata")) + metadata["spend_logs_metadata"] = { # rebind-ok: callback metadata must be updated in place + **existing, + "upstream_response_headers": { + **_as_dict(existing.get("upstream_response_headers")), + **headers, + }, + } + + +class ResponseHeaderLogger(CustomLogger): + async def async_post_call_response_headers_hook( + self, + data: dict[str, object], + user_api_key_dict: UserAPIKeyAuth, + response: object, + request_headers: dict[str, str] | None = None, + litellm_call_info: dict[str, object] | None = None, + ) -> None: + headers: Final = _headers(getattr(response, "_response_headers", None)) + spend_metadata: Final = _child(_child(data, "metadata"), "spend_logs_metadata") + spend_metadata["upstream_response_headers"] = {f"llm_provider-{name}": value for name, value in headers.items()} + + async def async_logging_hook( + self, kwargs: dict[str, object], result: object, call_type: str + ) -> tuple[dict[str, object], object]: + payload: Final = _as_dict(kwargs.get("standard_logging_object")) + hidden: Final = _as_dict(payload.get("hidden_params")) + headers: Final = _as_dict(hidden.get("additional_headers")) + params: Final = _child(kwargs, "litellm_params") + key: Final = "litellm_metadata" if params.get("litellm_metadata") else "metadata" + save_headers(_child(params, key), headers) + return kwargs, result + + async def async_post_call_failure_hook( + self, + request_data: dict[str, object], + original_exception: Exception, + user_api_key_dict: UserAPIKeyAuth, + traceback_str: str | None = None, + ) -> None: + headers: Final = get_error_headers(original_exception) + save_headers( + _child(request_data, "metadata"), + {f"llm_provider-{name}": value for name, value in headers.items()}, + ) + + +proxy_handler_instance: Final = ResponseHeaderLogger() diff --git a/litellm/proxy/common_request_processing.py b/litellm/proxy/common_request_processing.py index 05ddef822f1..e143e29c638 100644 --- a/litellm/proxy/common_request_processing.py +++ b/litellm/proxy/common_request_processing.py @@ -719,7 +719,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: @@ -857,22 +857,23 @@ def _sse_error_frames(error_obj: Mapping[str, object]) -> tuple[str, str]: 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. """ + # 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", - } + async def streaming_headers() -> Mapping[str, str]: + current_headers: Final = await refresh_headers() if refresh_headers is not None else headers + return MappingProxyType({**current_headers, "Cache-Control": "no-cache", "X-Accel-Buffering": "no"}) + first_chunk_value: str | None = None final_status_code = default_status_code @@ -909,7 +910,7 @@ async def create_response( return JSONResponse( status_code=final_status_code, content={"error": error_dict}, - headers=headers, + headers=await refresh_headers() if refresh_headers is not None else headers, ) except Exception as e: verbose_proxy_logger.debug("Error parsing first chunk value: %s", e) @@ -938,7 +939,7 @@ async def create_response( return StreamingResponse( empty_gen(), media_type=media_type, - headers=streaming_headers, + headers=await streaming_headers(), status_code=default_status_code, ) except Exception as e: @@ -954,7 +955,7 @@ async def create_response( return StreamingResponse( error_gen_message(), media_type=media_type, - headers=streaming_headers, + headers=await streaming_headers(), status_code=error_status, ) @@ -976,7 +977,7 @@ async def create_response( return _UpstreamClosingStreamingResponse( combined_generator(), media_type=media_type, - headers=streaming_headers, + headers=await streaming_headers(), status_code=final_status_code, upstream_generator=generator, ) @@ -2379,28 +2380,46 @@ 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, - ) + + def get_stream_headers() -> dict[str, str]: + current_hidden: Final = get_hidden_params_dict(response) + return 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(current_hidden, self.data), + cache_key=current_hidden.get("cache_key") or "", + api_base=current_hidden.get("api_base") or "", + version=version, + response_cost=current_hidden.get("response_cost") or "", + model_region=getattr(user_api_key_dict, "allowed_model_region", ""), + fastest_response_batch_completion=current_hidden.get("fastest_response_batch_completion"), + request_data=self.data, + hidden_params=current_hidden, + litellm_logging_obj=logging_obj, + **(current_hidden.get("additional_headers") or MappingProxyType({})), + ) + + custom_headers: Final = get_stream_headers() + request_headers: Final = dict(request.headers) + + async def refresh_stream_headers() -> Mapping[str, str]: + if get_hidden_params_dict(response) is hidden_params: + return custom_headers + refreshed: Final = get_stream_headers() + updated_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=request_headers, + ) + return MappingProxyType({**refreshed, **(updated_callback_headers or MappingProxyType({}))}) # Call response headers hook for streaming success callback_headers = 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), + request_headers=request_headers, ) if callback_headers: custom_headers.update(callback_headers) @@ -2508,6 +2527,8 @@ class ProxyBaseLLMRequestProcessing: ) # Non-streaming response - fall through to normal response handling elif select_data_generator: + from litellm.litellm_core_utils.streaming_handler import CustomStreamWrapper + selected_data_generator = select_data_generator( response=response, user_api_key_dict=user_api_key_dict, @@ -2536,6 +2557,7 @@ class ProxyBaseLLMRequestProcessing: media_type="text/event-stream", headers=custom_headers, request=request, + refresh_headers=refresh_stream_headers if isinstance(response, CustomStreamWrapper) else None, ) ### CALL HOOKS ### - modify outgoing data diff --git a/litellm/router.py b/litellm/router.py index 0c3958be0bf..1e101ce714a 100644 --- a/litellm/router.py +++ b/litellm/router.py @@ -2499,6 +2499,10 @@ class Router: if hasattr(model_response, "_hidden_params"): self._hidden_params = model_response._hidden_params.copy() + def update_headers(self, source: object) -> None: + self._response_headers = getattr(source, "_response_headers", None) + self._hidden_params = get_hidden_params_dict(source).copy() + def __aiter__(self): return self @@ -2552,8 +2556,10 @@ 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) + wrapped_response.update_headers(fallback_response) async for fallback_item in fallback_response: + wrapped_response.update_headers(fallback_response) + prepared_fallback_hidden_params = Router._prepare_fallback_hidden_params(fallback_response) Router._apply_fallback_hidden_params_to_item(fallback_item, prepared_fallback_hidden_params) if ( fallback_item @@ -2601,7 +2607,8 @@ class Router: e, ) - return FallbackStreamWrapper(stream_with_fallbacks()) + wrapped_response: Final = FallbackStreamWrapper(stream_with_fallbacks()) + return wrapped_response @staticmethod def _extract_partial_responses_usage( @@ -3043,6 +3050,10 @@ class Router: if hasattr(model_response, "_hidden_params"): self._hidden_params = model_response._hidden_params.copy() + def update_headers(self, source: object) -> None: + self._response_headers = getattr(source, "_response_headers", None) + self._hidden_params = get_hidden_params_dict(source).copy() + def __iter__(self): return self @@ -3093,8 +3104,10 @@ class Router: ) if hasattr(fallback_response, "__iter__"): - prepared_fallback_hidden_params = Router._prepare_fallback_hidden_params(fallback_response) + wrapped_response.update_headers(fallback_response) for fallback_item in fallback_response: + wrapped_response.update_headers(fallback_response) + prepared_fallback_hidden_params = Router._prepare_fallback_hidden_params(fallback_response) Router._apply_fallback_hidden_params_to_item(fallback_item, prepared_fallback_hidden_params) if ( fallback_item @@ -3132,7 +3145,8 @@ class Router: close_err, ) - return SyncFallbackStreamWrapper(stream_with_fallbacks()) + wrapped_response: Final = SyncFallbackStreamWrapper(stream_with_fallbacks()) + return wrapped_response async def _silent_experiment_acompletion(self, silent_model: str, messages: Sequence[Mapping[str, str]], **kwargs): """ diff --git a/tests/local_testing/test_custom_logger.py b/tests/local_testing/test_custom_logger.py index 1b627d56717..ca44a1cfa22 100644 --- a/tests/local_testing/test_custom_logger.py +++ b/tests/local_testing/test_custom_logger.py @@ -565,3 +565,96 @@ def test_redis_cache_completion_stream(): # test_redis_cache_completion_stream() + + +@pytest.mark.asyncio +async def test_response_header_logger_preserves_headers_on_headerless_stream_failure(): + from types import SimpleNamespace + from typing import Final + + from cookbook.logging_observability.response_header_logger import ResponseHeaderLogger + from litellm.proxy._types import UserAPIKeyAuth + + logger: Final = ResponseHeaderLogger() + auth: Final = UserAPIKeyAuth() + data: Final = {"metadata": {"spend_logs_metadata": {"existing_field": "keep"}}} + response: Final = SimpleNamespace( + _response_headers={ + "x-generation-id": "stream-request", + "x-litellm-attempted-fallbacks": "spoofed", + } + ) + await logger.async_post_call_response_headers_hook(data, auth, response) + await logger.async_post_call_failure_hook(data, RuntimeError("stream failed"), auth) + + assert data["metadata"]["spend_logs_metadata"] == { + "existing_field": "keep", + "upstream_response_headers": { + "llm_provider-x-generation-id": "stream-request", + "llm_provider-x-litellm-attempted-fallbacks": "spoofed", + }, + } + + +@pytest.mark.asyncio +@pytest.mark.parametrize("metadata_key", ["metadata", "litellm_metadata"]) +async def test_response_header_logger_success_preserves_capture_and_metadata(metadata_key: str): + from types import SimpleNamespace + from typing import Final + + from cookbook.logging_observability.response_header_logger import ResponseHeaderLogger + from litellm.proxy._types import UserAPIKeyAuth + + logger: Final = ResponseHeaderLogger() + metadata: Final = {"spend_logs_metadata": {"existing_field": "keep"}} + await logger.async_post_call_response_headers_hook( + {"metadata": metadata}, + UserAPIKeyAuth(), + SimpleNamespace(_response_headers={"x-generation-id": "stream-request"}), + ) + kwargs: Final = { + "litellm_params": {metadata_key: metadata}, + "standard_logging_object": { + "hidden_params": { + "additional_headers": { + "llm_provider-x-litellm-response-cost": 0.01, + } + } + }, + } + await logger.async_logging_hook(kwargs, None, "acompletion") + + assert metadata["spend_logs_metadata"] == { + "existing_field": "keep", + "upstream_response_headers": { + "llm_provider-x-generation-id": "stream-request", + "llm_provider-x-litellm-response-cost": 0.01, + }, + } + + +@pytest.mark.asyncio +@pytest.mark.parametrize("fallback_headers", [None, {"x-request-id": "winner"}]) +async def test_response_header_logger_replaces_headers_after_fallback(fallback_headers: dict[str, str] | None): + from types import SimpleNamespace + from typing import Final + + from cookbook.logging_observability.response_header_logger import ResponseHeaderLogger + from litellm.proxy._types import UserAPIKeyAuth + + logger: Final = ResponseHeaderLogger() + auth: Final = UserAPIKeyAuth() + data: Final = {"metadata": {"spend_logs_metadata": {"existing_field": "keep"}}} + await logger.async_post_call_response_headers_hook( + data, + auth, + SimpleNamespace(_response_headers={"x-request-id": "failed", "x-failed-only": "remove"}), + ) + await logger.async_post_call_response_headers_hook(data, auth, SimpleNamespace(_response_headers=fallback_headers)) + await logger.async_post_call_failure_hook(data, RuntimeError("stream failed"), auth) + assert data["metadata"]["spend_logs_metadata"] == { + "existing_field": "keep", + "upstream_response_headers": { + f"llm_provider-{name}": value for name, value in (fallback_headers or {}).items() + }, + } diff --git a/tests/test_litellm/proxy/test_common_request_processing.py b/tests/test_litellm/proxy/test_common_request_processing.py index df14224af5c..6f3d1b9ec77 100644 --- a/tests/test_litellm/proxy/test_common_request_processing.py +++ b/tests/test_litellm/proxy/test_common_request_processing.py @@ -7665,3 +7665,35 @@ def test_log_llm_api_exception_traceback_only_for_unexpected_errors(exc, expect_ records = [r for r in caplog.records if "_handle_llm_api_exception(): Exception occured" in r.getMessage()] assert len(records) == 1 assert (records[0].exc_info is not None) is expect_traceback + + +@pytest.mark.asyncio +@pytest.mark.parametrize("stream_kind", ["success", "empty", "error"]) +async def test_create_response_refreshes_headers_before_committing(stream_kind: str): + from typing import Final + + headers: Final = {"llm_provider-x-request-id": "failed", "llm_provider-x-failed-only": "remove"} + + async def stream(): + if stream_kind == "error": + yield 'data: {"error":{"message":"failed","code":429}}\n\n' + elif stream_kind == "success": + yield 'data: {"id":"winner","choices":[]}\n\n' + + async def refresh_headers(): + return {"llm_provider-x-request-id": "winner", "x-litellm-model-group": "fallback"} + + response: Final = await create_response( + stream(), + "text/event-stream", + headers, + refresh_headers=refresh_headers, + ) + assert response.headers["llm_provider-x-request-id"] == "winner" + assert response.headers["x-litellm-model-group"] == "fallback" + assert "llm_provider-x-failed-only" not in response.headers + if stream_kind == "error": + assert response.status_code == 429 + else: + assert response.headers["cache-control"] == "no-cache" + assert response.headers["x-accel-buffering"] == "no" diff --git a/tests/test_litellm/test_router.py b/tests/test_litellm/test_router.py index af39e3245af..71188e162e8 100644 --- a/tests/test_litellm/test_router.py +++ b/tests/test_litellm/test_router.py @@ -2346,7 +2346,7 @@ async def test_completion_streaming_iterator_preserves_response_headers( _response_headers=response_headers, ) router: Final = Router(model_list=[]) - upstream._hidden_params["additional_headers"]["x-litellm-model-group"] = "real-group" + await router.set_response_headers(response=upstream, model_group="real-group") iterator_kwargs: Final = dict( model_response=upstream, messages=[{"role": "user", "content": "Hello"}], @@ -2373,6 +2373,103 @@ async def test_completion_streaming_iterator_preserves_response_headers( assert StandardLoggingPayloadSetup.get_hidden_params(complete._hidden_params)["additional_headers"] == expected +@pytest.mark.asyncio +@pytest.mark.parametrize("is_async", [False, True]) +@pytest.mark.parametrize("fallback_headers", [None, {"x-request-id": "winner", "x-litellm-model-group": "spoofed"}]) +async def test_completion_streaming_iterator_replaces_failed_provider_headers( + is_async: bool, fallback_headers: dict[str, str] | None +) -> None: + from typing import Final + from unittest.mock import Mock + + from litellm.litellm_core_utils.streaming_handler import CustomStreamWrapper + from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler, HTTPHandler + + client: Final = AsyncMock(spec=AsyncHTTPHandler) if is_async else Mock(spec=HTTPHandler) + client.post.return_value = httpx.Response( + 200, + content=( + b'data: {"id":"winner","object":"chat.completion.chunk","created":1,' + b'"model":"test-model","choices":[{"index":0,"delta":{"content":"Hello"},' + b'"finish_reason":null}]}\n\n' + b"data: [DONE]\n\n" + ), + request=httpx.Request("POST", "https://provider.test/v1/chat/completions"), + ) + request: Final = dict( + model="hosted_vllm/test-model", + messages=[{"role": "user", "content": "Hello"}], + api_base="https://provider.test/v1", + api_key="fake-key", + stream=True, + client=client, + ) + original: Final = await litellm.acompletion(**request) if is_async else litellm.completion(**request) + fallback: Final = CustomStreamWrapper( + completion_stream=original.completion_stream, + model=original.model, + custom_llm_provider=original.custom_llm_provider, + logging_obj=original.logging_obj, + _response_headers=fallback_headers, + ) + + class FallbackRouter(Router): + async def async_function_with_fallbacks_common_utils(self, **kwargs): + return fallback + + def function_with_fallbacks(self, **kwargs): + return fallback + + class FailedStream: + model = "failed-model" + custom_llm_provider = "hosted_vllm" + logging_obj = original.logging_obj + chunks = [] + _response_headers = {"x-request-id": "failed", "x-failed-only": "remove"} + _hidden_params = { + "additional_headers": {"llm_provider-x-request-id": "failed", "llm_provider-x-failed-only": "remove"} + } + + def __iter__(self): + return self + + def __next__(self): + raise MidStreamFallbackError( + message="limited", + model=self.model, + llm_provider=self.custom_llm_provider, + is_pre_first_chunk=True, + ) + + def __aiter__(self): + return self + + async def __anext__(self): + return self.__next__() + + router: Final = FallbackRouter(model_list=[]) + await router.set_response_headers(response=fallback, model_group="fallback") + iterator_kwargs: Final = dict( + model_response=FailedStream(), + messages=[{"role": "user", "content": "Hello"}], + initial_kwargs={"model": "primary", "stream": True}, + ) + wrapped: Final = ( + await router._acompletion_streaming_iterator(**iterator_kwargs) + if is_async + else router._completion_streaming_iterator(**iterator_kwargs) + ) + chunks: Final = [chunk async for chunk in wrapped] if is_async else list(wrapped) + expected: Final = { + **{f"llm_provider-{name}": value for name, value in (fallback_headers or {}).items()}, + "x-litellm-model-group": "fallback", + } + assert wrapped._response_headers == fallback_headers + assert wrapped._hidden_params["additional_headers"] == expected + assert all(chunk._hidden_params["additional_headers"] == expected for chunk in chunks) + assert litellm.stream_chunk_builder(chunks).choices[0].message.content == "Hello" + + def test_completion_streaming_iterator_fallback_on_429(): """Sync streaming: MidStreamFallbackError (429 pre-first-chunk) triggers fallback. diff --git a/type-discipline-budget.json b/type-discipline-budget.json index 52cb9628252..745d8ad0a17 100644 --- a/type-discipline-budget.json +++ b/type-discipline-budget.json @@ -1,9 +1,9 @@ { "LIT001": { - "limit": 22364 + "limit": 22363 }, "LIT002": { - "limit": 26777 + "limit": 26776 }, "LIT003": { "limit": 269 @@ -27,7 +27,7 @@ "limit": 0 }, "LIT010": { - "limit": 16507 + "limit": 16505 }, "LIT011": { "limit": 5535