scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-09-13 23:11:40 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 101

feat(proxy): expose complexity routing headers

Browse source 
Co-Authored-By: Claude Code <noreply@anthropic.com>
This commit is contained in:
Tin 2026-09-11 12:35:24 -07:00
parent 09b694894d
commit c817faec7a
4 changed files with 355 additions and 6 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

24
litellm/router.py
View file

@ -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)

67
litellm/router_utils/add_retry_fallback_headers.py
View file

@ -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

112
tests/test_litellm/router_utils/test_add_retry_fallback_headers.py
View file

@ -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()

158
tests/test_litellm/test_router.py
View file

@ -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