From a10f41d662afbf5b45055b80c10f79e691c4d2ed Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Tue, 6 Oct 2026 22:54:25 -0700 Subject: [PATCH] fix(router): resolve a deployment's window from base_model for every provider --- litellm/router.py | 8 +- tests/e2e/coverage_registry/reliability.yaml | 3 + tests/e2e/models.py | 2 + ...ty_router_context_window_base_model_e2e.py | 219 ++++++++++++++++++ .../router_strategy/test_complexity_router.py | 46 ++++ tests/unit/test_router/test_router.py | 67 ++++++ 6 files changed, 342 insertions(+), 3 deletions(-) create mode 100644 tests/e2e/router/test_complexity_router_context_window_base_model_e2e.py diff --git a/litellm/router.py b/litellm/router.py index 010e5946cb8..a5db681e0cc 100644 --- a/litellm/router.py +++ b/litellm/router.py @@ -11035,7 +11035,9 @@ class Router: Augment litellm info with additional params set in `model_info`. - For azure models, ignore the `model:`. Only set max tokens, cost values if base_model is set. + A deployment's `base_model` (model_info first, then litellm_params) is the cost-map key for every + provider, the way /model/info resolves it. Without one, azure needs the deployment name to be a + known azure key and every other provider resolves its own `model:`. Returns - ModelInfo - If found -> typed dict with max tokens, input cost, etc. @@ -11077,7 +11079,7 @@ class Router: ) ## SET MODEL TO 'model=' - if base_model is None + not azure - if custom_llm_provider == "azure" and base_model is None: + if base_model is None and custom_llm_provider == "azure": # Router init auto-registers every deployment name into # litellm.model_cost as a zeroed stub, so membership alone can't # tell a resolvable name apart; require usable limits/costs. @@ -11099,7 +11101,7 @@ class Router: "Could not identify azure model '%s'. Set azure 'base_model' for accurate max tokens, cost tracking, etc.- https://docs.litellm.ai/docs/proxy/cost_tracking#spend-tracking-for-azure-openai-models", _model, ) - elif custom_llm_provider != "azure": + elif base_model is None: model = _model if "*" in model: # only call pattern_router for wildcard models diff --git a/tests/e2e/coverage_registry/reliability.yaml b/tests/e2e/coverage_registry/reliability.yaml index 7382cc215f2..896a4618cac 100644 --- a/tests/e2e/coverage_registry/reliability.yaml +++ b/tests/e2e/coverage_registry/reliability.yaml @@ -30,6 +30,9 @@ - {id: reliability.routing.semantic_auto_router.responses_input_routed, module: reliability, tier: P0, behavior: routing, variant: semantic_auto_router, assertions: [responses_input_routed], exercised_on: [responses], source: "litellm/router_strategy/auto_router/auto_router.py:131", fail_before_fix: proven, rationale: "/v1/responses input is resolved into messages for the semantic auto-router pre-routing hook instead of failing 400 Unmapped LLM provider auto_router (GitHub PR #37333)"} - {id: reliability.routing.strategy_alias.custom_pricing_ignored, module: reliability, tier: P1, behavior: routing, variant: strategy_alias, assertions: [custom_pricing_ignored], exercised_on: [chat_completions], source: "litellm/router.py:11489", rationale: "Custom pricing on a strategy-router alias never prices the routed request; spend logs at the routed tier deployment's own rate (GitHub PR #36691)"} - {id: reliability.routing.complexity_heuristic.scores_current_ask_only, module: reliability, tier: P1, behavior: routing, variant: complexity_heuristic, assertions: [scores_current_ask_only], exercised_on: [chat_completions], source: "router_strategy/complexity_router/complexity_router.py:942", rationale: "The heuristic complexity classifier scores the caller's current ask only, so a keyword-heavy agent system prompt cannot inflate the tier (GitHub PR #36721)"} +- {id: reliability.routing.context_window_escalation.escalates_from_base_model_window, module: reliability, tier: P0, behavior: routing, variant: context_window_escalation, assertions: [escalates_from_base_model_window], exercised_on: [chat_completions], source: "litellm/router.py:11081", fail_before_fix: proven, rationale: "An oversized SIMPLE request moves to the next tier that fits when the SIMPLE deployment's window comes only from model_info.base_model (a Bedrock application inference profile ARN), instead of staying put and 400ing at the provider"} +- {id: reliability.routing.context_window_escalation.escalates_from_explicit_window, module: reliability, tier: P1, behavior: routing, variant: context_window_escalation, assertions: [escalates_from_explicit_window], exercised_on: [chat_completions], source: "router_strategy/complexity_router/complexity_router.py:1238", rationale: "An oversized SIMPLE request moves to the next tier that fits when the SIMPLE deployment declares max_input_tokens explicitly"} +- {id: reliability.routing.context_window_escalation.fitting_prompt_stays_on_tier, module: reliability, tier: P1, behavior: routing, variant: context_window_escalation, assertions: [fitting_prompt_stays_on_tier], exercised_on: [chat_completions], source: "router_strategy/complexity_router/complexity_router.py:1238", rationale: "A request that fits the SIMPLE deployment's base_model window is served by that deployment, so resolving the window never over-escalates"} - {id: reliability.cache.exact.returns_cached, module: reliability, tier: P1, behavior: cache, variant: exact, assertions: [returns_cached], exercised_on: [chat_completions, messages, embeddings], source: "litellm/caching/caching.py", rationale: "Response cache returns cached on exact match"} - {id: reliability.cache.prompt_caching_model_select.returns_cached, module: reliability, tier: P1, behavior: cache, variant: prompt_caching_model_select, assertions: [returns_cached], exercised_on: [chat_completions], source: "router_utils/prompt_caching_cache.py", rationale: "Selects model supporting prompt caching for cacheable prefix; runs only on a stack with the prompt_caching pre-call check enabled (E2E_PROMPT_CACHING_STACK)"} - {id: reliability.circuit_breaker.redis.trips_then_recovers, module: reliability, tier: P0, behavior: circuit_breaker, variant: redis, assertions: [trips_then_recovers], exercised_on: [chat_completions, messages, embeddings], source: "litellm/caching/redis_cache.py:99", rationale: "Redis breaker CLOSED->OPEN->HALF_OPEN; guards all cache/rate-limit ops"} diff --git a/tests/e2e/models.py b/tests/e2e/models.py index ccb5cd35280..d3dd4960761 100644 --- a/tests/e2e/models.py +++ b/tests/e2e/models.py @@ -1319,6 +1319,8 @@ class ModelInfoBody(BaseModel): team_id: str | None = None allowed_fails: int | None = None allowed_fails_policy: dict[str, int] | None = None + base_model: str | None = None + max_input_tokens: int | None = None class ModelNewBody(BaseModel): diff --git a/tests/e2e/router/test_complexity_router_context_window_base_model_e2e.py b/tests/e2e/router/test_complexity_router_context_window_base_model_e2e.py new file mode 100644 index 00000000000..863e9b9bcb1 --- /dev/null +++ b/tests/e2e/router/test_complexity_router_context_window_base_model_e2e.py @@ -0,0 +1,219 @@ +"""Live e2e pins for the complexity router's context-window escalation when a tier's +window is declared only through ``model_info.base_model``. + +A Bedrock application inference profile ARN has no cost-map row of its own, so the +SIMPLE tier's window exists only through ``base_model``. ``Router.get_router_model_info`` +read base_model for azure deployments alone and resolved every other provider's +deployment by its own ``model:``, so the ARN came back with no window, the escalation +gate treated an unknown window as able to hold anything, and an oversized request the +classifier scored SIMPLE stayed on the 200k model and 400'd at Bedrock with "prompt is +too long" instead of moving to the 1M MEDIUM tier. + +The explicit-window control declares ``max_input_tokens`` below its model's own row and +sends a request that fits the row but not the declaration, so it escalates only because +a declared window is honored. It lives on its own backend model: the router registers a +deployment's ``model_info`` under the shared cost-map key of its ``model:``, so a declared +window on the ARN would leak to the base_model deployment and hide the bug. A short +request on the base_model router proves the gate only moves what does not fit. + +Every deployment is registered via /model/new (stage has no static config for these) +and the served deployment is read back from the spend log's ``model``, which stores +either the registered alias or the provider-prefixed form. The prompts are plain prose +with a "how many" ask, which keeps the heuristic classifier on SIMPLE, and the router +estimates their size as chars/4. +""" + +import os +from collections.abc import Iterator +from contextlib import ExitStack +from dataclasses import dataclass +from typing import Final + +import pytest + +from e2e_config import unique_marker +from e2e_http import unwrap +from lifecycle import ResourceManager +from models import ( + ChatBody, + ChatMessage, + ChatResponse, + KeyGenerateBody, + LiteLLMParamsBody, + ModelInfoBody, + ModelNewBody, + SpendLogRow, +) +from proxy_client import ProxyClient + +pytestmark = pytest.mark.e2e + +SIMPLE_BACKEND: Final = os.environ.get( + "E2E_BEDROCK_INFERENCE_PROFILE_ARN", + "bedrock/converse/arn:aws:bedrock:us-west-2:888602223428:application-inference-profile/sccejsk6r44v", +) +SIMPLE_BASE_MODEL: Final = "global.anthropic.claude-haiku-4-5-20251001-v1:0" +EXPLICIT_WINDOW_BACKEND: Final = "bedrock/converse/us.anthropic.claude-haiku-4-5-20251001-v1:0" +DECLARED_WINDOW_TOKENS: Final = 60_000 +UPPER_BACKEND: Final = "bedrock/converse/global.amazon.nova-2-lite-v1:0" +MAX_TOKENS: Final = 8 +OVERSIZED_PROMPT_CHARS: Final = 1_000_000 +BEYOND_DECLARED_WINDOW_PROMPT_CHARS: Final = 400_000 +FILLER: Final = ( + "On the morning of day {n}, the garden club met beside the old mill and talked about which flowers " + "to grow along the river walk. Everyone brought warm bread and shared stories about the cold winter " + "and the first robin of spring.\n" +) +ASK: Final = "How many days are written above? Reply with the single word OK." + + +@dataclass(frozen=True, slots=True) +class EscalationRouters: + base_model_router: str + explicit_window_router: str + simple_from_base_model: str + simple_with_explicit_window: str + upper: str + + +def _bedrock_params(model: str) -> LiteLLMParamsBody: + return LiteLLMParamsBody( + model=model, + aws_access_key_id="os.environ/AWS_ACCESS_KEY_ID", + aws_secret_access_key="os.environ/AWS_SECRET_ACCESS_KEY", + aws_region_name="us-west-2", + ) + + +def _router_params(simple: str, upper: str) -> LiteLLMParamsBody: + return LiteLLMParamsBody( + model="auto_router/complexity_router", + complexity_router_config={ + "classifier_type": "heuristic", + "enable_context_window_escalation": True, + "tiers": {"SIMPLE": simple, "MEDIUM": upper, "COMPLEX": upper, "REASONING": upper}, + }, + ) + + +def _prompt_of(chars: int) -> str: + marker: Final = unique_marker() + line_count: Final = -(-chars // len(FILLER.format(n=f"1 ({marker})"))) + lines: Final = tuple(FILLER.format(n=f"{n} ({marker})") for n in range(1, line_count + 1)) + return "".join(lines) + "\n" + ASK + + +def _chat_body(model: str, content: str) -> ChatBody: + return ChatBody(model=model, messages=[ChatMessage(role="user", content=content)], max_tokens=MAX_TOKENS) + + +def _key_for(proxy: ProxyClient, resources: ResourceManager, models: list[str]) -> str: + key: Final = proxy.generate_key(KeyGenerateBody(models=models, user_id="e2e-context-window-base-model")) + resources.defer(lambda: proxy.delete_key(key)) + return key + + +def _served_by(rows: list[SpendLogRow], alias: str, backend: str, context: str) -> None: + served: Final = tuple(row.model for row in rows) + assert served and all(model in {alias, backend} for model in served), ( + f"{context}: expected every request to be served by {alias} ({backend}), spend logs show {served}" + ) + + +def _answered(response: ChatResponse) -> None: + message: Final = response.choices[0].message if response.choices else None + assert message is not None and message.content, f"expected an answer, got {response}" + + +@pytest.fixture(scope="class") +def router_stack() -> Iterator[ExitStack]: + with ExitStack() as stack: + yield stack + + +@pytest.fixture(scope="class") +def routers(proxy: ProxyClient, router_stack: ExitStack) -> EscalationRouters: + marker: Final = unique_marker() + named: Final = EscalationRouters( + base_model_router=f"e2e-cw-router-base-{marker}", + explicit_window_router=f"e2e-cw-router-explicit-{marker}", + simple_from_base_model=f"e2e-cw-simple-base-{marker}", + simple_with_explicit_window=f"e2e-cw-simple-explicit-{marker}", + upper=f"e2e-cw-upper-{marker}", + ) + registrations: Final = ( + ModelNewBody( + model_name=named.simple_from_base_model, + litellm_params=_bedrock_params(SIMPLE_BACKEND), + model_info=ModelInfoBody(base_model=SIMPLE_BASE_MODEL), + ), + ModelNewBody( + model_name=named.simple_with_explicit_window, + litellm_params=_bedrock_params(EXPLICIT_WINDOW_BACKEND), + model_info=ModelInfoBody(max_input_tokens=DECLARED_WINDOW_TOKENS), + ), + ModelNewBody(model_name=named.upper, litellm_params=_bedrock_params(UPPER_BACKEND), model_info=ModelInfoBody()), + ModelNewBody( + model_name=named.base_model_router, + litellm_params=_router_params(named.simple_from_base_model, named.upper), + model_info=ModelInfoBody(), + ), + ModelNewBody( + model_name=named.explicit_window_router, + litellm_params=_router_params(named.simple_with_explicit_window, named.upper), + model_info=ModelInfoBody(), + ), + ) + for body in registrations: + router_stack.callback(proxy.delete_model, proxy.register_model(body)) + return named + + +class TestContextWindowEscalationFromBaseModel: + @pytest.mark.covers("reliability.routing.context_window_escalation.escalates_from_base_model_window") + def test_base_model_window_escalates( + self, proxy: ProxyClient, resources: ResourceManager, routers: EscalationRouters + ) -> None: + key: Final = _key_for( + proxy, resources, [routers.base_model_router, routers.simple_from_base_model, routers.upper] + ) + + response: Final = unwrap( + proxy.chat(key, _chat_body(routers.base_model_router, _prompt_of(OVERSIZED_PROMPT_CHARS))) + ) + + _answered(response) + _served_by(proxy.poll_logs_for_key(key), routers.upper, UPPER_BACKEND, "oversized ask, window from base_model") + + @pytest.mark.covers("reliability.routing.context_window_escalation.escalates_from_explicit_window") + def test_explicit_window_escalates( + self, proxy: ProxyClient, resources: ResourceManager, routers: EscalationRouters + ) -> None: + key: Final = _key_for( + proxy, resources, [routers.explicit_window_router, routers.simple_with_explicit_window, routers.upper] + ) + + response: Final = unwrap( + proxy.chat(key, _chat_body(routers.explicit_window_router, _prompt_of(BEYOND_DECLARED_WINDOW_PROMPT_CHARS))) + ) + + _answered(response) + _served_by(proxy.poll_logs_for_key(key), routers.upper, UPPER_BACKEND, "ask beyond the declared window") + + @pytest.mark.covers("reliability.routing.context_window_escalation.fitting_prompt_stays_on_tier") + def test_fitting_prompt_stays_on_the_simple_tier( + self, proxy: ProxyClient, resources: ResourceManager, routers: EscalationRouters + ) -> None: + key: Final = _key_for( + proxy, resources, [routers.base_model_router, routers.simple_from_base_model, routers.upper] + ) + + response: Final = unwrap(proxy.chat(key, _chat_body(routers.base_model_router, f"say hello {unique_marker()}"))) + + _answered(response) + _served_by( + proxy.poll_logs_for_key(key), + routers.simple_from_base_model, + SIMPLE_BACKEND, + "short ask, window from base_model", + ) diff --git a/tests/unit/router_strategy/test_complexity_router.py b/tests/unit/router_strategy/test_complexity_router.py index 67e0649958b..822720bf86b 100644 --- a/tests/unit/router_strategy/test_complexity_router.py +++ b/tests/unit/router_strategy/test_complexity_router.py @@ -14225,6 +14225,52 @@ class TestContextWindowEscalation: assert result.routing_decision["tier"] == "COMPLEX" assert "context_escalation" in result.routing_decision["signals"] + @pytest.mark.asyncio + @pytest.mark.usefixtures("local_model_cost_map") + async def test_a_window_declared_only_through_base_model_escalates_too(self): + """A Bedrock application inference profile ARN has no cost-map row; its window exists + only through `model_info.base_model`. The gate used to resolve that deployment to no + window at all (base_model was read for azure only), treat it as infinite, and leave an + oversized SIMPLE prompt on it to 400 at the provider.""" + base_model = "meta.llama3-8b-instruct-v1:0" + base_window = litellm.get_model_info(model=f"bedrock/{base_model}")["max_input_tokens"] + assert base_window is not None + litellm_router = Router( + model_list=[ + { + "model_name": "small-model", + "litellm_params": { + "model": "bedrock/converse/arn:aws:bedrock:us-west-2:000000000000:application-inference-profile/abc123xyz", + "mock_response": "ok", + }, + "model_info": {"base_model": base_model}, + }, + { + "model_name": "big-model", + "litellm_params": {"model": "openai/gpt-4o-mini", "mock_response": "ok"}, + "model_info": {"max_input_tokens": base_window * 8}, + }, + ] + ) + router = ComplexityRouter( + model_name="test-router", + litellm_router_instance=litellm_router, + complexity_router_config=_tier_config(), + ) + filler_repeats = (base_window * 4) // len(_CONTEXT_FILLER) + 2 + oversized_turns = [ + {"role": "user", "content": "Here is everything discussed so far. " + _CONTEXT_FILLER * filler_repeats}, + {"role": "assistant", "content": "Noted, I have read all of it."}, + {"role": "user", "content": "ok continue"}, + ] + + result = await router.async_pre_routing_hook(model="test-router", request_kwargs={}, messages=oversized_turns) + + assert result is not None + assert result.model == "big-model" + assert result.routing_decision["context_escalation_original_tier"] == "SIMPLE" + assert "context_escalation" in result.routing_decision["signals"] + @pytest.mark.asyncio async def test_a_prompt_that_fits_routes_exactly_as_before(self): """The gate must be invisible for normal traffic: same model, no escalation facts.""" diff --git a/tests/unit/test_router/test_router.py b/tests/unit/test_router/test_router.py index 290b4e62852..9f8906a7f4a 100644 --- a/tests/unit/test_router/test_router.py +++ b/tests/unit/test_router/test_router.py @@ -13650,6 +13650,73 @@ class TestAzureBaseModelFallbackLogging: ) assert model_info["max_input_tokens"] == litellm.model_cost["azure/gpt-4o-mini"]["max_input_tokens"] + +INFERENCE_PROFILE_ARN = "bedrock/converse/arn:aws:bedrock:us-west-2:000000000000:application-inference-profile/abc123xyz" +INFERENCE_PROFILE_BASE_MODEL = "global.anthropic.claude-haiku-4-5-20251001-v1:0" + + +@pytest.mark.usefixtures("local_model_cost_map") +class TestBaseModelResolvesWindowForEveryProvider: + """A Bedrock application inference profile ARN is not in the cost map, so its window and + prices only exist through `base_model`. get_router_model_info used to read base_model for + azure alone and resolved every other provider's deployment by its own `model:`, so the + ARN came back unmapped and the complexity router's context-window escalation, which + treats an unknown window as infinite, never moved an oversized request off that tier.""" + + def _router(self, model_info: dict[str, object], litellm_params: dict[str, object] | None = None) -> litellm.Router: + return litellm.Router( + model_list=[ + { + "model_name": "haiku", + "litellm_params": {"model": INFERENCE_PROFILE_ARN, **(litellm_params or {})}, + "model_info": {"id": "haiku-profile", **model_info}, + } + ] + ) + + def test_model_info_base_model_resolves_a_non_azure_deployment(self): + router = self._router({"base_model": INFERENCE_PROFILE_BASE_MODEL}) + + model_info = router.get_router_model_info(deployment=None, received_model_name="haiku", id="haiku-profile") + + base = litellm.get_model_info(model=f"bedrock/{INFERENCE_PROFILE_BASE_MODEL}") + assert model_info["max_input_tokens"] == base["max_input_tokens"] + assert model_info["input_cost_per_token"] == base["input_cost_per_token"] + + def test_litellm_params_base_model_resolves_a_non_azure_deployment(self): + router = self._router({}, {"base_model": INFERENCE_PROFILE_BASE_MODEL}) + + model_info = router.get_router_model_info(deployment=None, received_model_name="haiku", id="haiku-profile") + + base = litellm.get_model_info(model=f"bedrock/{INFERENCE_PROFILE_BASE_MODEL}") + assert model_info["max_input_tokens"] == base["max_input_tokens"] + + def test_explicit_max_input_tokens_still_wins_over_base_model(self): + base_window = litellm.get_model_info(model=f"bedrock/{INFERENCE_PROFILE_BASE_MODEL}")["max_input_tokens"] + assert base_window is not None + router = self._router({"base_model": INFERENCE_PROFILE_BASE_MODEL, "max_input_tokens": base_window // 2}) + + model_info = router.get_router_model_info(deployment=None, received_model_name="haiku", id="haiku-profile") + + assert model_info["max_input_tokens"] == base_window // 2 + + def test_a_deployment_without_base_model_still_resolves_its_own_model(self): + router = litellm.Router( + model_list=[ + { + "model_name": "nova", + "litellm_params": {"model": "bedrock/converse/global.amazon.nova-2-lite-v1:0"}, + "model_info": {"id": "nova-plain"}, + } + ] + ) + + model_info = router.get_router_model_info(deployment=None, received_model_name="nova", id="nova-plain") + + own = litellm.get_model_info(model="bedrock/converse/global.amazon.nova-2-lite-v1:0") + assert model_info["max_input_tokens"] == own["max_input_tokens"] + + def test_model_group_info_intersects_supported_reasoning_efforts(): router = litellm.Router( model_list=[