From c817faec7aaf001baec138ff748fb0583a6d0c4c Mon Sep 17 00:00:00 2001 From: Tin Date: Fri, 11 Sep 2026 12:35:24 -0700 Subject: [PATCH] feat(proxy): expose complexity routing headers Co-Authored-By: Claude Code --- litellm/router.py | 24 ++- .../add_retry_fallback_headers.py | 67 +++++++- .../test_add_retry_fallback_headers.py | 112 +++++++++++++ tests/test_litellm/test_router.py | 158 +++++++++++++++++- 4 files changed, 355 insertions(+), 6 deletions(-) diff --git a/litellm/router.py b/litellm/router.py index 597bcfaa20f..8865543badd 100644 --- a/litellm/router.py +++ b/litellm/router.py @@ -124,9 +124,11 @@ from litellm.router_utils.add_retry_fallback_headers import ( add_retry_headers_to_response, apply_quality_router_decision_headers, apply_remaining_usage_headers, + complexity_router_decision_headers, ensure_response_additional_headers, get_hidden_params_dict, prepare_response_for_header_attachment, + replace_complexity_router_headers, response_in_flight_token_count, ) from litellm.router_utils.auto_router_model_naming import ( @@ -555,6 +557,7 @@ class FallbackAwareAnthropicMessagesStream: def __init__(self, async_generator: AsyncGenerator[bytes, None], source_iterator: object) -> None: self._async_generator = async_generator self._source_iterator = source_iterator + self.fallback_headers_adopted = False self._hidden_params = dict( # mutable-ok: mutated in place by merge_fallback_hidden_params getattr(source_iterator, "_hidden_params", None) or {} ) @@ -565,6 +568,7 @@ class FallbackAwareAnthropicMessagesStream: def adopt_fallback_source(self, fallback_response: object) -> None: self._source_iterator = fallback_response + self.fallback_headers_adopted = True def __aiter__(self) -> "FallbackAwareAnthropicMessagesStream": return self @@ -593,7 +597,9 @@ class FallbackAwareAnthropicMessagesStream: self._hidden_params = { # mutable-ok: matches _hidden_params' existing dict[str, object] shape **self._hidden_params, **fallback_hidden_params, - "additional_headers": {**existing_headers, **fallback_headers}, # mutable-ok: same shape + "additional_headers": dict( # mutable-ok: hidden params expect a writable header bag + replace_complexity_router_headers(existing_headers, fallback_headers) + ), } @@ -3128,6 +3134,8 @@ class Router: async generator. """ + fallback_headers_adopted: bool = False + def __init__(self, async_generator: AsyncGenerator): import time from datetime import datetime @@ -3179,6 +3187,12 @@ class Router: # api_base, additional_headers) keep flowing. self._hidden_params = dict(getattr(source_iterator, "_hidden_params", None) or {}) + def adopt_fallback_headers(self, fallback_response: object) -> tuple[dict[str, object], dict[str, object]]: + prepared: Final = Router._prepare_fallback_hidden_params(fallback_response) + self._hidden_params = {**prepared[0], "additional_headers": prepared[1]} # mutable-ok: stream metadata + self.fallback_headers_adopted = True + return prepared + def __aiter__(self): return self @@ -3269,8 +3283,8 @@ class Router: include_fallback_errors=initial_kwargs.get("include_fallback_errors", False) is True, ) + prepared_fallback_hidden_params = wrapper.adopt_fallback_headers(fallback_response) if hasattr(fallback_response, "__aiter__"): - prepared_fallback_hidden_params = Router._prepare_fallback_hidden_params(fallback_response) async for fallback_item in fallback_response: Router._apply_fallback_hidden_params_to_item(fallback_item, prepared_fallback_hidden_params) if partial_usage is not None: @@ -3305,7 +3319,8 @@ class Router: exc, ) - return FallbackResponsesStreamWrapper(stream_with_fallbacks()) + wrapper: Final = FallbackResponsesStreamWrapper(stream_with_fallbacks()) + return wrapper def _completion_streaming_iterator( self, @@ -11108,7 +11123,7 @@ class Router: self, response: object, model_group: str | None = None, - request_kwargs: dict | None = None, + request_kwargs: dict[str, object] | None = None, ) -> Any: """ Add the most accurate rate limit headers for a given model response. @@ -11124,6 +11139,7 @@ class Router: additional_headers: Final = ensure_response_additional_headers(response) additional_headers["x-litellm-model-group"] = model_group apply_quality_router_decision_headers(additional_headers, request_kwargs) + additional_headers.update(complexity_router_decision_headers(request_kwargs)) if model_group is not None: remaining_usage: Final = await self.get_remaining_model_group_usage(model_group) diff --git a/litellm/router_utils/add_retry_fallback_headers.py b/litellm/router_utils/add_retry_fallback_headers.py index 3251ea457cf..bc88feef7d2 100644 --- a/litellm/router_utils/add_retry_fallback_headers.py +++ b/litellm/router_utils/add_retry_fallback_headers.py @@ -1,7 +1,10 @@ import json +import math +from collections.abc import Mapping +from types import MappingProxyType from typing import Any, Final, Protocol, TypedDict, cast -from pydantic import BaseModel +from pydantic import BaseModel, TypeAdapter, ValidationError class FallbackErrorInfo(TypedDict): @@ -15,6 +18,68 @@ class _HiddenParamsHost(Protocol): _hidden_params: dict[str, object] +_EMPTY_OBJECT_MAPPING: Final[Mapping[str, object]] = MappingProxyType({}) +_ROUTING_HEADER_MAPPING: Final = TypeAdapter(Mapping[str, object]) +_COMPLEXITY_ROUTER_HEADER_PREFIX: Final = "x-litellm-complexity-router-" + + +def _routing_header_mapping(value: object) -> Mapping[str, object]: + try: + mapping: Final[Mapping[str, object]] = _ROUTING_HEADER_MAPPING.validate_python(value, strict=True) + return mapping + except ValidationError: + return _EMPTY_OBJECT_MAPPING + + +def _header_string(value: object) -> str | None: + if not isinstance(value, str): + return None + normalized: Final = value.strip() + return normalized if normalized and all(" " <= character <= "~" for character in normalized) else None + + +def complexity_router_decision_headers(request_kwargs: object) -> Mapping[str, str]: + data: Final = _routing_header_mapping(request_kwargs) + metadata_key: Final = "litellm_metadata" if "litellm_metadata" in data else "metadata" + decision: Final = _routing_header_mapping(_routing_header_mapping(data.get(metadata_key)).get("routing_decision")) + if decision.get("router_type") != "complexity": + return MappingProxyType({}) + score: Final = decision.get("score") + values: Final = ( + ("tier", decision.get("tier")), + ("cause", decision.get("cause")), + ( + "score", + str(score) + if isinstance(score, (int, float)) and not isinstance(score, bool) and math.isfinite(score) + else None, + ), + ( + "reasoning-effort", + _routing_header_mapping(decision.get("tier_litellm_params")).get("reasoning_effort"), + ), + ) + return MappingProxyType( + { + f"{_COMPLEXITY_ROUTER_HEADER_PREFIX}{key}": header_value + for key, value in values + if (header_value := _header_string(value)) is not None + } + ) + + +def replace_complexity_router_headers( + existing_headers: Mapping[str, object], new_headers: Mapping[str, object] +) -> Mapping[str, object]: + return MappingProxyType( + { + key: value + for key, value in (*existing_headers.items(), *new_headers.items()) + if key in new_headers or not key.startswith(_COMPLEXITY_ROUTER_HEADER_PREFIX) + } + ) + + class HiddenParamsAsyncIteratorWrapper: """ Wraps a bare async generator/iterator (e.g. a provider's raw SSE diff --git a/tests/test_litellm/router_utils/test_add_retry_fallback_headers.py b/tests/test_litellm/router_utils/test_add_retry_fallback_headers.py index 3a0deeb13d8..a5f78bd0f4b 100644 --- a/tests/test_litellm/router_utils/test_add_retry_fallback_headers.py +++ b/tests/test_litellm/router_utils/test_add_retry_fallback_headers.py @@ -1,12 +1,15 @@ import json +import pytest from pydantic import BaseModel from litellm.router_utils.add_retry_fallback_headers import ( add_fallback_headers_to_response, add_retry_headers_to_response, + complexity_router_decision_headers, get_fallback_errors_from_headers, get_hidden_params_dict, + replace_complexity_router_headers, ) @@ -15,6 +18,115 @@ class StreamingWrapper: self._hidden_params = {"additional_headers": {"x-existing": "keep"}} +@pytest.mark.parametrize("metadata_key", ["metadata", "litellm_metadata"]) +def test_complexity_router_decision_headers_exposes_only_bounded_fields(metadata_key): + headers = complexity_router_decision_headers( + { + metadata_key: { + "routing_decision": { + "router_type": "complexity", + "tier": " REASONING ", + "cause": "heuristic_scorer", + "score": 0.75, + "tier_litellm_params": {"reasoning_effort": "xhigh", "api_key": "secret"}, + "signals": ["private prompt"], + "matched_keyword": "private prompt", + } + } + } + ) + + assert dict(headers) == { + "x-litellm-complexity-router-tier": "REASONING", + "x-litellm-complexity-router-cause": "heuristic_scorer", + "x-litellm-complexity-router-score": "0.75", + "x-litellm-complexity-router-reasoning-effort": "xhigh", + } + + +@pytest.mark.parametrize( + "decision, expected", + [ + ( + {"router_type": "complexity", "tier": "SIMPLE", "cause": "heuristic_scorer", "score": 0}, + { + "x-litellm-complexity-router-tier": "SIMPLE", + "x-litellm-complexity-router-cause": "heuristic_scorer", + "x-litellm-complexity-router-score": "0", + }, + ), + ( + {"router_type": "complexity", "tier": "COMPLEX", "cause": "llm_classifier"}, + { + "x-litellm-complexity-router-tier": "COMPLEX", + "x-litellm-complexity-router-cause": "llm_classifier", + }, + ), + ( + {"router_type": "complexity", "tier": "REASONING", "cause": "literal_keyword_match"}, + { + "x-litellm-complexity-router-tier": "REASONING", + "x-litellm-complexity-router-cause": "literal_keyword_match", + }, + ), + ({"router_type": "quality", "tier": "premium", "cause": "quality_tier"}, {}), + ({"router_type": "complexity", "score": True}, {}), + ({"router_type": "complexity", "score": float("nan")}, {}), + ({"router_type": "complexity", "score": float("inf")}, {}), + ({"router_type": "complexity", "tier": "研究", "cause": "bad\r\nX-Injected: true"}, {}), + ({"router_type": "complexity", "tier_litellm_params": {"reasoning_effort": 1}}, {}), + ({"router_type": "complexity", "tier_litellm_params": "invalid"}, {}), + ([], {}), + (None, {}), + ], +) +def test_complexity_router_decision_headers_omits_absent_or_invalid_fields(decision, expected): + assert dict(complexity_router_decision_headers({"metadata": {"routing_decision": decision}})) == expected + + +@pytest.mark.parametrize( + "litellm_decision, metadata_decision, expected", + [ + ( + {"router_type": "complexity", "tier": "SIMPLE", "cause": "heuristic_scorer"}, + {"router_type": "complexity", "tier": "REASONING", "tier_litellm_params": {"reasoning_effort": "xhigh"}}, + {"x-litellm-complexity-router-tier": "SIMPLE", "x-litellm-complexity-router-cause": "heuristic_scorer"}, + ), + ( + {"router_type": "quality", "tier": "premium"}, + {"router_type": "complexity", "tier": "FORGED", "cause": "heuristic_scorer"}, + {}, + ), + ( + {}, + {"router_type": "complexity", "tier": "FORGED", "cause": "heuristic_scorer"}, + {}, + ), + ], +) +def test_complexity_router_decision_headers_never_falls_back_from_internal_metadata( + litellm_decision, metadata_decision, expected +): + headers = complexity_router_decision_headers( + { + "litellm_metadata": {"routing_decision": litellm_decision}, + "metadata": {"routing_decision": metadata_decision}, + } + ) + assert dict(headers) == expected + + +def test_replace_complexity_router_headers_drops_stale_values(): + assert replace_complexity_router_headers( + { + "x-existing": "keep", + "x-litellm-complexity-router-tier": "REASONING", + "x-litellm-complexity-router-reasoning-effort": "xhigh", + }, + {"x-litellm-complexity-router-tier": "SIMPLE"}, + ) == {"x-existing": "keep", "x-litellm-complexity-router-tier": "SIMPLE"} + + def test_add_fallback_headers_to_streaming_wrapper(): response = StreamingWrapper() diff --git a/tests/test_litellm/test_router.py b/tests/test_litellm/test_router.py index 2a97e92396a..bd1da89ebf0 100644 --- a/tests/test_litellm/test_router.py +++ b/tests/test_litellm/test_router.py @@ -2622,6 +2622,76 @@ def test_adopt_fallback_response_headers_keeps_identity_when_fallback_has_none() assert wrapper.fallback_headers_adopted is True +@pytest.mark.asyncio +@pytest.mark.parametrize("response_kind", ["object", "dict", "async-generator"]) +async def test_set_response_headers_exposes_complexity_decision_on_every_response_shape(response_kind): + class HeaderResponse: + def __init__(self): + self._hidden_params: dict[str, object] = {} + + response: object + if response_kind == "object": + response = HeaderResponse() + elif response_kind == "dict": + response = {} + else: + response = _AsyncList() + + router = Router(model_list=[]) + result = await router.set_response_headers( + response=response, + request_kwargs={ + "metadata": { + "routing_decision": { + "router_type": "complexity", + "tier": "SIMPLE", + "cause": "heuristic_scorer", + "score": 0.25, + "tier_litellm_params": {"reasoning_effort": "low"}, + } + } + }, + ) + hidden_params = result["_hidden_params"] if isinstance(result, dict) else result._hidden_params + additional_headers = hidden_params["additional_headers"] + + assert additional_headers == { + "x-litellm-model-group": None, + "x-litellm-complexity-router-tier": "SIMPLE", + "x-litellm-complexity-router-cause": "heuristic_scorer", + "x-litellm-complexity-router-score": "0.25", + "x-litellm-complexity-router-reasoning-effort": "low", + } + + +@pytest.mark.asyncio +async def test_set_response_headers_is_the_only_complexity_header_source_for_proxy_headers(): + from litellm.proxy._types import UserAPIKeyAuth + from litellm.proxy.common_request_processing import ProxyBaseLLMRequestProcessing + + router = Router(model_list=[]) + response = await router.set_response_headers(response={}, request_kwargs={}) + additional_headers = response["_hidden_params"]["additional_headers"] + proxy_headers = ProxyBaseLLMRequestProcessing.get_custom_headers( + user_api_key_dict=UserAPIKeyAuth(), + request_data={ + "metadata": { + "routing_decision": { + "router_type": "complexity", + "tier": "REASONING", + "cause": "heuristic_scorer", + "tier_litellm_params": {"reasoning_effort": "xhigh"}, + } + } + }, + **additional_headers, + ) + + assert not { + key for key in proxy_headers if key.startswith("x-litellm-complexity-router-") + } + + @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 @@ -3472,9 +3542,10 @@ def _make_responses_iterator( class _AsyncList: """Generic async iterator over a list — used as the fallback response.""" - def __init__(self, items=()): + def __init__(self, items=(), hidden_params=None): self._items = list(items) self._idx = 0 + self._hidden_params = hidden_params or {} def __aiter__(self): return self @@ -3562,6 +3633,54 @@ async def test_aresponses_streaming_iterator_fallback(): assert call_kwargs["disable_fallbacks"] is False +@pytest.mark.asyncio +async def test_aresponses_streaming_iterator_replaces_complexity_headers_before_fallback_output(): + error = MidStreamFallbackError( + message="primary failed before output", + model="gpt-4", + llm_provider="openai", + is_pre_first_chunk=True, + generated_content="", + ) + primary_headers = { + "x-litellm-complexity-router-tier": "REASONING", + "x-litellm-complexity-router-reasoning-effort": "xhigh", + } + source = _make_responses_iterator( + error=error, + hidden_params={"additional_headers": primary_headers}, + ) + fallback = _AsyncList( + [MagicMock(type="response.output_text.delta")], + hidden_params={"model_id": "fallback-deployment", "additional_headers": {"x-fallback-only": "yes"}}, + ) + router = _make_router_with_fallback() + + with patch.object( + router, + "async_function_with_fallbacks_common_utils", + return_value=fallback, + ): + wrapped = await router._aresponses_streaming_iterator( + response=source, + initial_kwargs={ + "model": "gpt-4", + "stream": True, + "input": "Hello", + "original_generic_function": litellm.aresponses, + }, + ) + assert wrapped._hidden_params["additional_headers"] == primary_headers + first_output = await wrapped.__anext__() + + assert first_output.type == "response.output_text.delta" + assert wrapped.fallback_headers_adopted is True + assert wrapped._hidden_params == { + "model_id": "fallback-deployment", + "additional_headers": {"x-fallback-only": "yes"}, + } + + @pytest.mark.asyncio async def test_aresponses_streaming_iterator_writes_litellm_metadata_on_fallback(): """Regression: model_group must land under "litellm_metadata" (the key @@ -12579,6 +12698,43 @@ async def test_anthropic_messages_fallback_merges_fallback_hidden_params(): assert headers["x-fallback-only"] == "yes" +@pytest.mark.asyncio +async def test_anthropic_messages_fallback_removes_primary_complexity_headers_before_output(): + router = _anthropic_messages_make_router() + source = _AnthropicMessagesFallbackByteStream( + [_anthropic_messages_overloaded_error_chunk()], + hidden_params={ + "additional_headers": { + "x-litellm-complexity-router-tier": "REASONING", + "x-litellm-complexity-router-reasoning-effort": "xhigh", + } + }, + ) + fallback_stream = _AnthropicMessagesFallbackByteStream( + [_anthropic_messages_content_chunk("fallback answer")], + hidden_params={"model_id": "fallback-deployment", "additional_headers": {"x-fallback-only": "yes"}}, + ) + + with patch.object( + router, + "async_function_with_fallbacks_common_utils", + new=AsyncMock(return_value=fallback_stream), + ): + wrapped = await router._aanthropic_messages_streaming_iterator( + response=source, + initial_kwargs={"model": "primary"}, + ) + assert wrapped._hidden_params["additional_headers"]["x-litellm-complexity-router-tier"] == "REASONING" + first_output = await wrapped.__anext__() + + assert first_output == _anthropic_messages_content_chunk("fallback answer") + assert wrapped.fallback_headers_adopted is True + assert wrapped._hidden_params == { + "model_id": "fallback-deployment", + "additional_headers": {"x-fallback-only": "yes"}, + } + + @pytest.mark.asyncio async def test_aanthropic_messages_with_streaming_fallbacks_deep_copies_nested_metadata(): """Bugbot regression: a shallow .copy() of kwargs still shares the