diff --git a/litellm/llms/bedrock/common_utils.py b/litellm/llms/bedrock/common_utils.py index 0cf3e895ad2..5d76188d9df 100644 --- a/litellm/llms/bedrock/common_utils.py +++ b/litellm/llms/bedrock/common_utils.py @@ -981,6 +981,10 @@ def extract_model_name_from_bedrock_arn(model: str) -> str: return model +def is_bedrock_arn(model: str) -> bool: + return re.search(r"\barn:aws[\w-]*:bedrock:", model) is not None + + def is_bedrock_application_inference_profile_arn(model: str) -> bool: """ An application inference profile ARN ends in an opaque id with no provider diff --git a/litellm/router.py b/litellm/router.py index a507518a748..799c208a241 100644 --- a/litellm/router.py +++ b/litellm/router.py @@ -135,7 +135,7 @@ from litellm.llms.base_llm.vector_store.transformation import ( RouterVectorStoreEmbeddingExecutor, vector_store_request_metadata, ) -from litellm.llms.bedrock.common_utils import get_bedrock_base_model, is_bedrock_application_inference_profile_arn +from litellm.llms.bedrock.common_utils import get_bedrock_base_model from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler, get_async_httpx_client from litellm.llms.openai_like.json_loader import JSONProviderRegistry from litellm.llms.openai_like.model_info import ( @@ -1104,7 +1104,7 @@ def as_output_cap(value: object) -> int | None: return cap if cap >= 0 else None -def _exact_base_model_max_input_tokens(base_model: str) -> int | None: +def exact_base_model_max_input_tokens(base_model: str) -> int | None: try: model_info: Final = litellm.get_model_info(model=base_model) except Exception: # noqa: BLE001 # get_model_info raises ModelNotMappedError or a bare Exception for an unmapped model @@ -1133,22 +1133,15 @@ def _declared_base_model(user_model_info: Mapping[str, object]) -> str | None: return None -def _base_model_window( +def _base_model_widened_window( route_model_info: ModelMapInfo, base_model: str | None, custom_llm_provider: str ) -> Mapping[str, int]: - if base_model is None: - return MappingProxyType({}) - if is_bedrock_application_inference_profile_arn(route_model_info["key"]): - profile_max_input_tokens: Final = _exact_base_model_max_input_tokens(base_model) - return MappingProxyType( - {} if profile_max_input_tokens is None else {"max_input_tokens": profile_max_input_tokens} - ) route_max_input_tokens: Final = route_model_info.get("max_input_tokens") - if route_max_input_tokens is None: + if base_model is None or route_max_input_tokens is None: return MappingProxyType({}) if not _names_route_model(base_model, route_model_info["key"], custom_llm_provider): return MappingProxyType({}) - base_max_input_tokens: Final = _exact_base_model_max_input_tokens(base_model) + base_max_input_tokens: Final = exact_base_model_max_input_tokens(base_model) if base_max_input_tokens is None or base_max_input_tokens <= route_max_input_tokens: return MappingProxyType({}) return MappingProxyType({"max_input_tokens": base_max_input_tokens}) @@ -11993,8 +11986,7 @@ 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. - For every other provider, a base_model naming the same model as `model:` can only widen its context window, - and a Bedrock application inference profile ARN, which names no model, takes its base_model's window. + For every other provider, a base_model naming the same model as `model:` can only widen its context window. Returns - ModelInfo - If found -> typed dict with max tokens, input cost, etc. @@ -12101,7 +12093,7 @@ class Router: # values are skipped or Deployment's None pricing defaults would erase the map's merged_model_info: Final[ModelMapInfo] = { **copy.deepcopy(model_info), - **_base_model_window( + **_base_model_widened_window( model_info, _declared_base_model(user_model_info) or _declared_base_model(litellm_params.model_extra or {}), custom_llm_provider, diff --git a/litellm/router_strategy/complexity_router/complexity_router.py b/litellm/router_strategy/complexity_router/complexity_router.py index a475358e1f0..084130b336a 100644 --- a/litellm/router_strategy/complexity_router/complexity_router.py +++ b/litellm/router_strategy/complexity_router/complexity_router.py @@ -1273,6 +1273,28 @@ def _window_can_hold(window: int | None, needed: int, buffer: float) -> bool: return window is None or needed <= int(window * buffer) +def _declared_base_model(deployment: Mapping[str, object]) -> str | None: + match deployment: + case {"model_info": {"base_model": str() as base_model}} | { + "litellm_params": {"base_model": str() as base_model} + }: + return base_model + case _: + return None + + +def _bedrock_arn_base_model_window(deployment: Mapping[str, object]) -> int | None: + from litellm.llms.bedrock.common_utils import is_bedrock_arn + from litellm.router import exact_base_model_max_input_tokens + + base_model: Final = _declared_base_model(deployment) + match deployment: + case {"litellm_params": {"model": str() as model}} if base_model is not None and is_bedrock_arn(model): + return exact_base_model_max_input_tokens(base_model) + case _: + return None + + def _group_provably_fits(facts: tuple[int | None, bool], needed: int, buffer: float) -> bool: window, has_unknown = facts return window is not None and not has_unknown and needed <= int(window * buffer) @@ -3391,7 +3413,8 @@ class ComplexityRouter(CustomLogger): return name if self.config.has_custom_tiers else ComplexityTier(name) def _deployment_window(self, group: str, deployment: Mapping[str, object]) -> int | None: - return self._deployment_limit(group, deployment, "max_input_tokens") + window: Final = self._deployment_limit(group, deployment, "max_input_tokens") + return window if window is not None else _bedrock_arn_base_model_window(deployment) def _deployment_limit( self, group: str, deployment: Mapping[str, object], key: Literal["max_input_tokens", "max_output_tokens"] diff --git a/tests/code_coverage_tests/router_code_coverage.py b/tests/code_coverage_tests/router_code_coverage.py index 75cac8b026c..dc6fa602ddf 100644 --- a/tests/code_coverage_tests/router_code_coverage.py +++ b/tests/code_coverage_tests/router_code_coverage.py @@ -138,8 +138,8 @@ ignored_function_names = [ "_pre_call_check_requested_chat_params", # Tested through the unsupported-param pre-call check tests in test_router.py "_deployment_supports_param", # Tested through the deployment-capability pre-call check tests in test_router.py "_capability_model_and_provider", # Tested through the authenticating-provider pre-call check test in test_router.py - "_exact_base_model_max_input_tokens", # Tested through get_router_model_info in test_router_model_cost_isolation.py - "_base_model_window", # Tested through get_router_model_info in test_router_model_cost_isolation.py + "exact_base_model_max_input_tokens", # Tested through get_router_model_info in test_router_model_cost_isolation.py + "_base_model_widened_window", # Tested through get_router_model_info in test_router_model_cost_isolation.py "_provider_model_id", # Tested through get_router_model_info in test_router_model_cost_isolation.py "_names_route_model", # Tested through get_router_model_info in test_router_model_cost_isolation.py "_declared_base_model", # Tested through get_router_model_info in test_router_model_cost_isolation.py 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 index 52c617645d6..f9cc33fcc68 100644 --- 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 @@ -18,9 +18,9 @@ from dataclasses import dataclass from typing import Final import pytest - from e2e_config import unique_marker from e2e_http import unwrap +from e2e_metadata import Domain, Mode, Provider, Subject, meta from lifecycle import ResourceManager from models import ( ChatBody, @@ -159,6 +159,14 @@ def routers(proxy: ProxyClient, router_stack: ExitStack) -> EscalationRouters: class TestContextWindowEscalationFromBaseModel: @pytest.mark.covers("reliability.routing.context_window_escalation.escalates_from_base_model_window") + @meta( + Subject( + domain=Domain.ROUTING, + providers=(Provider.BEDROCK,), + models=(SIMPLE_BACKEND, UPPER_BACKEND), + mode=Mode.NONSTREAM, + ) + ) def test_base_model_window_escalates( self, proxy: ProxyClient, resources: ResourceManager, routers: EscalationRouters ) -> None: @@ -174,6 +182,14 @@ class TestContextWindowEscalationFromBaseModel: _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") + @meta( + Subject( + domain=Domain.ROUTING, + providers=(Provider.BEDROCK,), + models=(EXPLICIT_WINDOW_BACKEND, UPPER_BACKEND), + mode=Mode.NONSTREAM, + ) + ) def test_explicit_window_escalates( self, proxy: ProxyClient, resources: ResourceManager, routers: EscalationRouters ) -> None: @@ -189,6 +205,14 @@ class TestContextWindowEscalationFromBaseModel: _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") + @meta( + Subject( + domain=Domain.ROUTING, + providers=(Provider.BEDROCK,), + models=(SIMPLE_BACKEND,), + mode=Mode.NONSTREAM, + ) + ) def test_fitting_prompt_stays_on_the_simple_tier( self, proxy: ProxyClient, resources: ResourceManager, routers: EscalationRouters ) -> None: diff --git a/tests/unit/router_strategy/complexity_router/test_context_compaction.py b/tests/unit/router_strategy/complexity_router/test_context_compaction.py index c6242f775c8..07e62bd0ff5 100644 --- a/tests/unit/router_strategy/complexity_router/test_context_compaction.py +++ b/tests/unit/router_strategy/complexity_router/test_context_compaction.py @@ -312,6 +312,36 @@ async def test_undersized_native_compactor_does_not_block_explicit_escalation( assert json.loads(answer.calls[0].request.content)["model"] == "wide-answer" +@pytest.mark.parametrize("escalate", [False, True]) +async def test_a_bedrock_arn_tier_windowed_only_through_base_model_is_never_compacted( + wire: tuple[respx.Route, respx.Route], escalate: bool, +) -> None: + base_model: Final = "meta.llama3-8b-instruct-v1:0" + base_window: Final = litellm.get_model_info(model=base_model)["max_input_tokens"] + assert base_window is not None + router: Final = make_router(settings={ + "enable_context_window_escalation": escalate, + "tiers": {"SIMPLE": "profile", "MEDIUM": "large", "COMPLEX": "large", "REASONING": "large"}, + }) + router.add_deployment(Deployment( + model_name="profile", litellm_params={ + "model": "bedrock/converse/arn:aws:bedrock:us-west-2:000000000000:application-inference-profile/abc123xyz", + "mock_response": "profile answer", + }, model_info={"id": "profile-answer", "base_model": base_model}, + )) + payload: Final = {"messages": [ + {"role": "user", "content": "Project code MAPLE-47. Background detail. " * (base_window // 14)}, + {"role": "assistant", "content": "Recorded"}, + {"role": "user", "content": "Answer with the project code"}, + ]} + output: Final = base_window // 4 + trigger: Final = int(base_window * ContextCompactionConfig().trigger_ratio) + assert trigger - output < router._count_pre_call_check_tokens(payload["messages"], None, payload) < trigger + response: Final = await router.acompletion(model="auto", max_tokens=output, **payload) + assert all(route.call_count == 0 for route in wire) + assert response.choices[0].message.content == "profile answer" + + @pytest.mark.parametrize( ("surface", "failure"), [(surface, failure) for surface in ("chat", "messages", "responses") for failure in ("unsigned", "oversized", "provider")] diff --git a/tests/unit/router_strategy/test_complexity_router.py b/tests/unit/router_strategy/test_complexity_router.py index 1bf89401584..90b4d7994f8 100644 --- a/tests/unit/router_strategy/test_complexity_router.py +++ b/tests/unit/router_strategy/test_complexity_router.py @@ -14354,6 +14354,24 @@ _CJK_TURNS = [ ] +_BEDROCK_PROFILE_ARN: Final = ( + "bedrock/converse/arn:aws:bedrock:us-west-2:000000000000:application-inference-profile/abc123xyz" +) +_LLAMA_BASE_MODEL: Final = "meta.llama3-8b-instruct-v1:0" + + +def _turns_past(window: int) -> list[dict[str, str]]: + return [ + { + "role": "user", + "content": "Here is everything discussed so far. " + + "The meeting notes were saved to the shared folder for later review this week. " * (window // 4), + }, + {"role": "assistant", "content": "Noted, I have read all of it."}, + {"role": "user", "content": "ok continue"}, + ] + + def _tier_config(**overrides: object) -> dict[str, object]: return { "tiers": {"SIMPLE": "small-model", "COMPLEX": "big-model"}, @@ -14393,43 +14411,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_an_inference_profile_window_declared_only_through_base_model_escalates_too(self): - base_model: Final = "meta.llama3-8b-instruct-v1:0" - base_window: Final = litellm.get_model_info(model=base_model)["max_input_tokens"] - assert base_window is not None + @staticmethod + def _base_model_tier_router(small: dict[str, object], base_window: int) -> ComplexityRouter: litellm_router: Final = 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": "small-model", **small}, { "model_name": "big-model", "litellm_params": {"model": "openai/gpt-4o-mini", "mock_response": "ok"}, - "model_info": {"max_input_tokens": base_window * 8}, + "model_info": {"max_input_tokens": base_window * 16}, }, ] ) - router: Final = ComplexityRouter( - model_name="test-router", - litellm_router_instance=litellm_router, - complexity_router_config=_tier_config(), + return ComplexityRouter( + model_name="test-router", litellm_router_instance=litellm_router, complexity_router_config=_tier_config() + ) + + @pytest.mark.asyncio + @pytest.mark.usefixtures("local_model_cost_map") + @pytest.mark.parametrize( + "arn", + [ + _BEDROCK_PROFILE_ARN, + "bedrock/converse/arn:aws:bedrock:us-west-2:000000000000:provisioned-model/abc123xyz", + "bedrock/llama/arn:aws:bedrock:us-west-2:000000000000:imported-model/abc123xyz", + ], + ) + @pytest.mark.parametrize("declared_in", ["model_info", "litellm_params"]) + async def test_a_bedrock_arn_window_declared_only_through_base_model_escalates_too( + self, arn: str, declared_in: str + ): + base_window: Final = litellm.get_model_info(model=_LLAMA_BASE_MODEL)["max_input_tokens"] + assert base_window is not None + router: Final = self._base_model_tier_router( + { + "litellm_params": { + "model": arn, + "mock_response": "ok", + **({"base_model": _LLAMA_BASE_MODEL} if declared_in == "litellm_params" else {}), + }, + "model_info": {"base_model": _LLAMA_BASE_MODEL} if declared_in == "model_info" else {}, + }, + base_window, ) - filler_repeats: Final = (base_window * 4) // len(_CONTEXT_FILLER) + 2 - oversized_turns: Final = [ - {"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: Final = await router.async_pre_routing_hook( - model="test-router", request_kwargs={}, messages=oversized_turns + model="test-router", request_kwargs={}, messages=_turns_past(base_window) ) assert result is not None @@ -14437,6 +14464,67 @@ class TestContextWindowEscalation: assert result.routing_decision["context_escalation_original_tier"] == "SIMPLE" assert "context_escalation" in result.routing_decision["signals"] + @pytest.mark.asyncio + @pytest.mark.usefixtures("local_model_cost_map") + @pytest.mark.parametrize("own_window", ["declared", "registered", "not_a_bedrock_arn"]) + async def test_base_model_never_overrides_an_own_window_or_reaches_a_non_arn_route( + self, own_window: str, monkeypatch: pytest.MonkeyPatch + ): + base_window: Final = litellm.get_model_info(model=_LLAMA_BASE_MODEL)["max_input_tokens"] + assert base_window is not None + if own_window == "registered": + monkeypatch.setitem( + litellm.model_cost, + _BEDROCK_PROFILE_ARN, + {"max_input_tokens": base_window * 8, "litellm_provider": "bedrock", "mode": "chat"}, + ) + router: Final = self._base_model_tier_router( + { + "litellm_params": { + "model": ( + "openai/a-model-the-cost-map-does-not-know" + if own_window == "not_a_bedrock_arn" + else _BEDROCK_PROFILE_ARN + ), + "mock_response": "ok", + }, + "model_info": { + "base_model": _LLAMA_BASE_MODEL, + **({"max_input_tokens": base_window * 8} if own_window == "declared" else {}), + }, + }, + base_window, + ) + + result: Final = await router.async_pre_routing_hook( + model="test-router", request_kwargs={}, messages=_turns_past(base_window) + ) + + assert result is not None + assert result.model == "small-model" + assert "context_escalated" not in result.routing_decision + + @pytest.mark.asyncio + @pytest.mark.usefixtures("local_model_cost_map") + async def test_an_unmapped_base_model_leaves_a_bedrock_arn_tier_in_place(self): + base_window: Final = litellm.get_model_info(model=_LLAMA_BASE_MODEL)["max_input_tokens"] + assert base_window is not None + router: Final = self._base_model_tier_router( + { + "litellm_params": {"model": _BEDROCK_PROFILE_ARN, "mock_response": "ok"}, + "model_info": {"base_model": "anthropic.claude-not-in-the-cost-map-v1:0"}, + }, + base_window, + ) + + result: Final = await router.async_pre_routing_hook( + model="test-router", request_kwargs={}, messages=_turns_past(base_window) + ) + + assert result is not None + assert result.model == "small-model" + assert "context_escalated" not in result.routing_decision + @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 b0e777b7ad2..13ac55fa937 100644 --- a/tests/unit/test_router/test_router.py +++ b/tests/unit/test_router/test_router.py @@ -16263,55 +16263,6 @@ class TestAzureBaseModelFallbackLogging: ) assert model_info["max_input_tokens"] == litellm.model_cost["azure/gpt-4o-mini"]["max_input_tokens"] - -INFERENCE_PROFILE_ARN: Final = ( - "bedrock/converse/arn:aws:bedrock:us-west-2:000000000000:application-inference-profile/abc123xyz" -) -INFERENCE_PROFILE_BASE_MODEL: Final = "global.anthropic.claude-haiku-4-5-20251001-v1:0" - - -@pytest.mark.usefixtures("local_model_cost_map") -class TestApplicationInferenceProfileWindowFromBaseModel: - def _profile_window(self, model_info: dict[str, object], litellm_params: dict[str, object] | None = None) -> object: - router: Final = litellm.Router( - model_list=[ - { - "model_name": "haiku", - "litellm_params": {"model": INFERENCE_PROFILE_ARN, **(litellm_params or {})}, - "model_info": {"id": "haiku-profile", **model_info}, - } - ] - ) - return router.get_router_model_info(deployment=None, received_model_name="haiku", id="haiku-profile")[ - "max_input_tokens" - ] - - def test_model_info_base_model_sets_the_profile_window(self): - base_window: Final = litellm.get_model_info(model=INFERENCE_PROFILE_BASE_MODEL)["max_input_tokens"] - assert base_window is not None - - assert self._profile_window({"base_model": INFERENCE_PROFILE_BASE_MODEL}) == base_window - - def test_litellm_params_base_model_sets_the_profile_window(self): - base_window: Final = litellm.get_model_info(model=INFERENCE_PROFILE_BASE_MODEL)["max_input_tokens"] - assert base_window is not None - - assert self._profile_window({}, {"base_model": INFERENCE_PROFILE_BASE_MODEL}) == base_window - - def test_explicit_max_input_tokens_still_wins_over_base_model(self): - base_window: Final = litellm.get_model_info(model=INFERENCE_PROFILE_BASE_MODEL)["max_input_tokens"] - assert base_window is not None - - declared_window: Final = self._profile_window( - {"base_model": INFERENCE_PROFILE_BASE_MODEL, "max_input_tokens": base_window // 2} - ) - - assert declared_window == base_window // 2 - - def test_an_unmapped_base_model_leaves_the_profile_window_unknown(self): - assert self._profile_window({"base_model": "anthropic.claude-not-in-the-cost-map-v1:0"}) is None - - def test_model_group_info_intersects_supported_reasoning_efforts(): router = litellm.Router( model_list=[