From ffe4e704d4826ec69e4d2af59385f22ca528810e Mon Sep 17 00:00:00 2001 From: Tin Chi Lo Date: Mon, 3 Aug 2026 15:25:32 -0700 Subject: [PATCH] feat(complexity_router)!: require every tier to name at least one model BREAKING CHANGE: a complexity_router_config.tiers map that omits a tier, or names one with an empty pool or an empty pin, is now rejected at config load A tier with no models of its own did not route nowhere; it routed somewhere else. Resolution fell through to another tier's pool or to default_model, so traffic an operator had classified as COMPLEX was answered by whatever the chain reached, and the only evidence was in a spend log. That substitution is what this removes: the model named for a tier is the model that tier gets The check runs on ComplexityRouterConfig, so it covers config.yaml, /model/new, the SDK and the UI alike, and it names every missing tier at once rather than one per restart. Omitting `tiers` entirely is unaffected and still fills in all four defaults, so this only reaches configs that opted into naming tiers and left some out Three behaviors become unreachable and their tests now assert the rejection instead: default_model standing in for an unconfigured tier, the empty-pool error raised when a tier was first selected rather than at load, and escalation skipping an intermediate tier that had no models. The fallback code paths stay as defense for configs built outside the validator; removing them is a separate change The repo's own suites carried 28 partial tiers maps, all of them incidental setup for tests about something else, which is a fair signal that real configs do the same. They are completed here without changing what each test exercises --- .../complexity_router/README.md | 3 + .../complexity_router/config.py | 28 + .../test_model_management_endpoints.py | 8 +- .../proxy/proxy_server/test_proxy_config.py | 4 +- .../router_strategy/test_complexity_router.py | 427 ++++++----- tests/test_litellm/test_router.py | 725 +++++------------- 6 files changed, 501 insertions(+), 694 deletions(-) diff --git a/litellm/router_strategy/complexity_router/README.md b/litellm/router_strategy/complexity_router/README.md index a6267453bf7..271066db3ee 100644 --- a/litellm/router_strategy/complexity_router/README.md +++ b/litellm/router_strategy/complexity_router/README.md @@ -38,6 +38,9 @@ The weighted sum is mapped to tiers using configurable boundaries: ### Basic Configuration +Every tier needs at least one model of its own. A tier left empty does not route nowhere; whatever the fallback chain reaches answers in its place, which silently swaps the model you configured for another tier's. Configs missing a tier are rejected at load, naming every tier that has none. Omitting `tiers` entirely still works and takes the defaults for all four. + + ```yaml model_list: - model_name: smart-router diff --git a/litellm/router_strategy/complexity_router/config.py b/litellm/router_strategy/complexity_router/config.py index 970d8de8575..f65be20597c 100644 --- a/litellm/router_strategy/complexity_router/config.py +++ b/litellm/router_strategy/complexity_router/config.py @@ -5,6 +5,7 @@ Contains default keyword lists, weights, tier boundaries, and configuration clas All values are configurable via proxy config.yaml. """ +from collections.abc import Sequence from enum import Enum from typing import Literal @@ -239,6 +240,13 @@ DEFAULT_TIER_MODELS: dict[str, str] = { } +def _models_named_by(configured: str | Sequence[str] | None) -> tuple[str, ...]: + """The models a tier entry names. Absent, an empty pool and an empty pin all mean none.""" + if isinstance(configured, str): + return (configured,) if configured else () + return tuple(configured or ()) + + class ClassifierLLMConfig(BaseModel): """Configuration for the LLM-based complexity classifier.""" @@ -470,6 +478,26 @@ class ComplexityRouterConfig(BaseModel): return None return [stripped for keyword in value if (stripped := keyword.strip())] + @model_validator(mode="after") + def _require_a_model_for_every_tier(self) -> "ComplexityRouterConfig": + """Every tier must name at least one model of its own. + + A tier left empty does not route nowhere; it routes somewhere else, to whichever + model the fallback chain reaches. That is a silent substitution of the model an + operator asked for, discoverable only by reading a spend log, so it is rejected at + config load instead. Omitting `tiers` entirely still works and fills in all four. + """ + missing = tuple(tier.value for tier in TIER_SEVERITY_ORDER if not _models_named_by(self.tiers.get(tier.value))) + if not missing: + return self + raise ValueError( + f"complexity_router_config.tiers must name at least one model for every tier; " + f"no model for {', '.join(missing)}. A tier with no models of its own has its " + f"traffic served by another tier's model, so configure all of " + f"{', '.join(tier.value for tier in TIER_SEVERITY_ORDER)}, or omit `tiers` to take " + f"the defaults" + ) + @model_validator(mode="after") def _validate_llm_classifier_config(self) -> "ComplexityRouterConfig": if self.classifier_type == "llm" and self.classifier_llm_config is None: diff --git a/tests/test_litellm/proxy/management_endpoints/test_model_management_endpoints.py b/tests/test_litellm/proxy/management_endpoints/test_model_management_endpoints.py index 8dbf38555b3..e3cd75ada5b 100644 --- a/tests/test_litellm/proxy/management_endpoints/test_model_management_endpoints.py +++ b/tests/test_litellm/proxy/management_endpoints/test_model_management_endpoints.py @@ -642,7 +642,7 @@ class TestDeleteModelClearsRouterRegistry: "model_name": "smart-router", "litellm_params": { "model": "auto_router/complexity_router", - "complexity_router_config": {"tiers": {"SIMPLE": "gpt-4o-mini", "MEDIUM": "gpt-4o"}}, + "complexity_router_config": {"tiers": {"SIMPLE": "gpt-4o-mini", "MEDIUM": "gpt-4o", "COMPLEX": 'gpt-4o-mini', "REASONING": 'gpt-4o-mini'}}, "complexity_router_default_model": "gpt-4o", **({"tags": tags} if tags else {}), }, @@ -3408,7 +3408,7 @@ class TestStrategyRouterWriteValidation: def _stored_complexity_params(self) -> LiteLLM_Params: return LiteLLM_Params( model="auto_router/complexity_router", - complexity_router_config={"tiers": {"SIMPLE": "gpt-4o-mini"}}, + complexity_router_config={"tiers": {"SIMPLE": "gpt-4o-mini", "MEDIUM": 'gpt-4o-mini', "COMPLEX": 'gpt-4o-mini', "REASONING": 'gpt-4o-mini'}}, ) def _db_complexity_router(self, model_id: str) -> Deployment: @@ -3467,7 +3467,7 @@ class TestStrategyRouterWriteValidation: corrupted = LiteLLM_Params( model="auto_router/auto_router/complexity_router", - complexity_router_config={"tiers": {"SIMPLE": "gpt-4o-mini"}}, + complexity_router_config={"tiers": {"SIMPLE": "gpt-4o-mini", "MEDIUM": 'gpt-4o-mini', "COMPLEX": 'gpt-4o-mini', "REASONING": 'gpt-4o-mini'}}, ) assert ( _strategy_router_write_violation( @@ -3580,7 +3580,7 @@ class TestStrategyRouterWriteValidation: "model_name": "my-auto-router", "litellm_params": { "model": "auto_router/complexity_router", - "complexity_router_config": {"tiers": {"SIMPLE": "gpt-4o-mini"}}, + "complexity_router_config": {"tiers": {"SIMPLE": "gpt-4o-mini", "MEDIUM": 'gpt-4o-mini', "COMPLEX": 'gpt-4o-mini', "REASONING": 'gpt-4o-mini'}}, }, "model_info": {"id": model_id}, } diff --git a/tests/test_litellm/proxy/proxy_server/test_proxy_config.py b/tests/test_litellm/proxy/proxy_server/test_proxy_config.py index 28d4d87e26f..c4be547b00f 100644 --- a/tests/test_litellm/proxy/proxy_server/test_proxy_config.py +++ b/tests/test_litellm/proxy/proxy_server/test_proxy_config.py @@ -121,11 +121,11 @@ def test__scrub_db_overlay_remote_module_loads_invalid_non_dict_returns_input(): def test_resolve_complexity_router_plugins_no_plugins_key_is_a_noop(): - config: Dict[str, Any] = {"tiers": {"SIMPLE": "gpt-4o-mini"}} + config: Dict[str, Any] = {"tiers": {"SIMPLE": "gpt-4o-mini", "MEDIUM": 'gpt-4o-mini', "COMPLEX": 'gpt-4o-mini', "REASONING": 'gpt-4o-mini'}} resolve_complexity_router_plugins( model_name="smart-router", complexity_router_config=config, config_file_path=None ) - assert config == {"tiers": {"SIMPLE": "gpt-4o-mini"}} + assert config == {"tiers": {"SIMPLE": "gpt-4o-mini", "MEDIUM": 'gpt-4o-mini', "COMPLEX": 'gpt-4o-mini', "REASONING": 'gpt-4o-mini'}} def test_resolve_complexity_router_plugins_resolves_dotted_path_to_live_instance(tmp_path): diff --git a/tests/test_litellm/router_strategy/test_complexity_router.py b/tests/test_litellm/router_strategy/test_complexity_router.py index cc73273450a..106a44891b9 100644 --- a/tests/test_litellm/router_strategy/test_complexity_router.py +++ b/tests/test_litellm/router_strategy/test_complexity_router.py @@ -323,26 +323,28 @@ class TestModelSelection: model = complexity_router.get_model_for_tier(ComplexityTier.REASONING) assert model == "o1-preview" - def test_get_model_fallback_to_default(self, mock_router_instance): - """Should fallback to default_model if tier not configured.""" - config = { - "tiers": {}, # Empty tiers - "default_model": "fallback-model", - } - router = ComplexityRouter( - model_name="test-router", - litellm_router_instance=mock_router_instance, - complexity_router_config=config, - ) - model = router.get_model_for_tier(ComplexityTier.SIMPLE) - assert model == "fallback-model" + def test_a_tiers_map_missing_tiers_is_rejected_even_with_a_default_model(self, mock_router_instance): + """default_model no longer stands in for an unconfigured tier. It used to serve + that tier's traffic, which is exactly the silent substitution this requirement + removes: the model an operator names for a tier is the model that tier gets.""" + with pytest.raises(ValidationError, match="no model for SIMPLE, MEDIUM, COMPLEX, REASONING"): + ComplexityRouter( + model_name="test-router", + litellm_router_instance=mock_router_instance, + complexity_router_config={"tiers": {}, "default_model": "fallback-model"}, + ) def test_get_model_for_tier_list_random_choice(self, mock_router_instance): router = ComplexityRouter( model_name="test-router", litellm_router_instance=mock_router_instance, complexity_router_config={ - "tiers": {"SIMPLE": ["cheap", "premium"], "MEDIUM": "mid"}, + "tiers": { + "SIMPLE": ["cheap", "premium"], + "MEDIUM": "mid", + "COMPLEX": "complex-unused", + "REASONING": "reasoning-unused", + }, "default_model": "mid", }, ) @@ -355,119 +357,20 @@ class TestModelSelection: choice.assert_called_once_with(pool) assert router.get_model_for_tier(ComplexityTier.MEDIUM) == "mid" - def test_get_model_for_tier_empty_pool_raises(self, mock_router_instance): - router = ComplexityRouter( - model_name="test-router", - litellm_router_instance=mock_router_instance, - complexity_router_config={ - "tiers": {"SIMPLE": []}, - "default_model": "mid", - }, - ) - with pytest.raises(ValueError, match="Empty model pool for tier SIMPLE"): - router.get_model_for_tier(ComplexityTier.SIMPLE) - - -class TestPreRoutingHook: - """Test the async_pre_routing_hook method.""" - - @pytest.mark.asyncio - async def test_pre_routing_hook_simple_message(self, complexity_router): - """Test pre-routing hook with a simple message.""" - messages = [{"role": "user", "content": "Hello!"}] - result = await complexity_router.async_pre_routing_hook( - model="test-model", - request_kwargs={}, - messages=messages, - ) - assert result is not None - assert result.model == "gpt-4o-mini" # SIMPLE tier model - assert result.messages == messages - - @pytest.mark.asyncio - async def test_pre_routing_hook_complex_message(self, complexity_router): - """Test pre-routing hook with a message containing technical content.""" - messages = [ - { - "role": "user", - "content": ( - "Design a distributed microservice architecture with Kubernetes " - "orchestration, implementing proper authentication, encryption, " - "and database optimization for high throughput. Think step by step " - "about the performance implications and scalability requirements." - ), - } - ] - result = await complexity_router.async_pre_routing_hook( - model="test-model", - request_kwargs={}, - messages=messages, - ) - assert result is not None - # Should return a valid model from the configured tiers - assert result.model in [ - "gpt-4o-mini", - "gpt-4o", - "claude-sonnet-4-20250514", - "o1-preview", - ] - - @pytest.mark.asyncio - async def test_pre_routing_hook_no_messages(self, complexity_router): - """Test pre-routing hook returns None when no messages.""" - result = await complexity_router.async_pre_routing_hook( - model="test-model", - request_kwargs={}, - messages=None, - ) - assert result is None - - @pytest.mark.asyncio - async def test_pre_routing_hook_empty_messages(self, complexity_router): - """Test pre-routing hook returns None when messages empty.""" - result = await complexity_router.async_pre_routing_hook( - model="test-model", - request_kwargs={}, - messages=[], - ) - assert result is None - - @pytest.mark.asyncio - async def test_pre_routing_hook_with_system_prompt(self, complexity_router): - """Test pre-routing hook considers system prompt.""" - messages = [ - {"role": "system", "content": "You are a helpful assistant."}, - {"role": "user", "content": "Hello!"}, - ] - result = await complexity_router.async_pre_routing_hook( - model="test-model", - request_kwargs={}, - messages=messages, - ) - assert result is not None - # Should still be SIMPLE - assert result.model == "gpt-4o-mini" - - @pytest.mark.asyncio - async def test_pre_routing_hook_reasoning_message(self, complexity_router): - """Test pre-routing hook with reasoning markers.""" - messages = [ - { - "role": "user", - "content": "Let's think step by step and reason through this problem carefully.", - } - ] - result = await complexity_router.async_pre_routing_hook( - model="test-model", - request_kwargs={}, - messages=messages, - ) - assert result is not None - assert result.model == "o1-preview" # REASONING tier model - - -class TestConfigOverrides: - """Test configuration override functionality.""" + @pytest.mark.parametrize("no_models", [[], ""], ids=["empty_pool", "empty_pin"]) + def test_a_tier_with_no_models_is_rejected_at_load_not_at_selection(self, mock_router_instance, no_models): + """An empty pool used to be caught when the tier was first selected, so a config + carrying one started fine and failed on a request. Naming a tier and giving it + nothing says the same as not naming it, and both are refused up front.""" + with pytest.raises(ValidationError, match="no model for SIMPLE"): + ComplexityRouter( + model_name="test-router", + litellm_router_instance=mock_router_instance, + complexity_router_config={ + "tiers": {"SIMPLE": no_models, "MEDIUM": "m", "COMPLEX": "c", "REASONING": "r"}, + "default_model": "mid", + }, + ) def test_custom_tier_boundaries(self, mock_router_instance): """Test custom tier boundaries work correctly.""" @@ -562,7 +465,14 @@ class TestCustomTechnicalKeywords: router_absent = ComplexityRouter( model_name="test-router", litellm_router_instance=mock_router_instance, - complexity_router_config={"tiers": {"MEDIUM": "gpt-4o"}}, + complexity_router_config={ + "tiers": { + "MEDIUM": "gpt-4o", + "SIMPLE": "simple-unused", + "COMPLEX": "complex-unused", + "REASONING": "reasoning-unused", + } + }, ) router_none = ComplexityRouter( model_name="test-router", @@ -746,7 +656,7 @@ class TestSingletonMutation: def test_default_config_not_mutated(self, mock_router_instance): """Test that creating routers without config doesn't mutate defaults.""" from litellm.router_strategy.complexity_router.config import ( - DEFAULT_CLASSIFIER_CONTEXT_WINDOW_SIZE, + DEFAULT_CLASSIFIER_CONTEXT_WINDOW_SIZE, ComplexityRouterConfig, ) @@ -952,6 +862,8 @@ class TestRouterComplexityDeploymentMethods: "tiers": { "SIMPLE": ["cheap"], "MEDIUM": ["cheap", "premium"], + "COMPLEX": "premium", + "REASONING": "premium", }, }, }, @@ -1662,6 +1574,8 @@ class TestRouterPreRoutingAliasOverrides: "tiers": { "SIMPLE": "gpt-4o-mini", "MEDIUM": "gpt-4o", + "COMPLEX": "complex-unused", + "REASONING": "reasoning-unused", } }, "complexity_router_default_model": "gpt-4o", @@ -1719,6 +1633,8 @@ class TestRouterPreRoutingAliasOverrides: "tiers": { "SIMPLE": "gpt-4o-mini", "MEDIUM": "gpt-4o", + "COMPLEX": "complex-unused", + "REASONING": "reasoning-unused", } } assert request_kwargs["complexity_router_default_model"] == "gpt-4o" @@ -1823,7 +1739,7 @@ class TestAdaptiveSoftFloors: def test_adaptive_defaults_use_cost_weighted_cold_policy(self): config = ComplexityRouterConfig( adaptive=True, - tiers={"SIMPLE": ["cheap"]}, + tiers={"SIMPLE": ["cheap"], "MEDIUM": "cheap", "COMPLEX": "cheap", "REASONING": "cheap"}, ) assert config.adaptive_weights.quality == pytest.approx(0.3) assert config.adaptive_weights.cost == pytest.approx(0.7) @@ -1870,7 +1786,9 @@ class TestAdaptiveSoftFloors: def test_adaptive_config_requires_non_empty_pools(self): with pytest.raises(ValidationError): - ComplexityRouterConfig(adaptive=True, tiers={"SIMPLE": []}) + ComplexityRouterConfig( + adaptive=True, tiers={"SIMPLE": [], "MEDIUM": None, "COMPLEX": None, "REASONING": None} + ) def test_cold_start_randomly_samples_unobserved_classified_tier_models(self, adaptive_router_instance): cr = ComplexityRouter( @@ -1881,6 +1799,8 @@ class TestAdaptiveSoftFloors: "tiers": { "SIMPLE": ["cheap", "premium"], "MEDIUM": ["premium"], + "COMPLEX": "complex-unused", + "REASONING": "reasoning-unused", }, }, ) @@ -1907,7 +1827,12 @@ class TestAdaptiveSoftFloors: litellm_router_instance=mock_router_instance, complexity_router_config={ "adaptive": False, - "tiers": {"SIMPLE": ["cheap", "premium"], "MEDIUM": "mid"}, + "tiers": { + "SIMPLE": ["cheap", "premium"], + "MEDIUM": "mid", + "COMPLEX": "complex-unused", + "REASONING": "reasoning-unused", + }, "default_model": "mid", }, ) @@ -1976,6 +1901,7 @@ class TestAdaptiveSoftFloors: "SIMPLE": ["cheap"], "MEDIUM": ["cheap", "premium"], "COMPLEX": ["premium"], + "REASONING": "premium", }, }, ) @@ -3164,7 +3090,12 @@ class TestRoutingPlugins: model_name="test-complexity-router", litellm_router_instance=mock_router_instance, complexity_router_config={ - "tiers": {"SIMPLE": ["gpt-4o-mini", "gpt-4o-nano"]}, + "tiers": { + "SIMPLE": ["gpt-4o-mini", "gpt-4o-nano"], + "MEDIUM": "medium-unused", + "COMPLEX": "complex-unused", + "REASONING": "reasoning-unused", + }, "plugins": [ExcludeGpt4oMini()], }, ) @@ -3192,7 +3123,12 @@ class TestRoutingPlugins: model_name="test-complexity-router", litellm_router_instance=mock_router_instance, complexity_router_config={ - "tiers": {"SIMPLE": "gpt-4o-mini"}, + "tiers": { + "SIMPLE": "gpt-4o-mini", + "MEDIUM": "medium-unused", + "COMPLEX": "complex-unused", + "REASONING": "reasoning-unused", + }, "default_model": "gpt-4o-fallback", "plugins": [BlockEverything()], }, @@ -3215,7 +3151,12 @@ class TestRoutingPlugins: model_name="test-complexity-router", litellm_router_instance=mock_router_instance, complexity_router_config={ - "tiers": {"SIMPLE": "gpt-4o-mini"}, + "tiers": { + "SIMPLE": "gpt-4o-mini", + "MEDIUM": "medium-unused", + "COMPLEX": "complex-unused", + "REASONING": "reasoning-unused", + }, "plugins": [BlockEverything()], }, ) @@ -3239,7 +3180,12 @@ class TestRoutingPlugins: model_name="test-complexity-router", litellm_router_instance=mock_router_instance, complexity_router_config={ - "tiers": {"SIMPLE": "gpt-4o-mini"}, + "tiers": { + "SIMPLE": "gpt-4o-mini", + "MEDIUM": "medium-unused", + "COMPLEX": "complex-unused", + "REASONING": "reasoning-unused", + }, "plugins": [CaptureMetadata()], }, ) @@ -3263,7 +3209,12 @@ class TestRoutingPlugins: model_name="test-complexity-router", litellm_router_instance=mock_router_instance, complexity_router_config={ - "tiers": {"SIMPLE": ["gpt-4o-mini", "gpt-4o-nano"]}, + "tiers": { + "SIMPLE": ["gpt-4o-mini", "gpt-4o-nano"], + "MEDIUM": "medium-unused", + "COMPLEX": "complex-unused", + "REASONING": "reasoning-unused", + }, "keyword_tier_rules": [{"keywords": ["hello"], "tier": "SIMPLE"}], "plugins": [ExcludeGpt4oMini()], }, @@ -3292,7 +3243,12 @@ class TestRoutingPlugins: model_name="test-complexity-router", litellm_router_instance=mock_router_instance, complexity_router_config={ - "tiers": {"MEDIUM": ["gpt-4o-default", "gpt-4o-nano"]}, + "tiers": { + "MEDIUM": ["gpt-4o-default", "gpt-4o-nano"], + "SIMPLE": "simple-unused", + "COMPLEX": "complex-unused", + "REASONING": "reasoning-unused", + }, "default_model": "gpt-4o-default", "plugins": [ExcludeDefaultModel()], }, @@ -3319,7 +3275,12 @@ class TestRoutingPlugins: model_name="test-complexity-router", litellm_router_instance=mock_router_instance, complexity_router_config={ - "tiers": {"MEDIUM": ["gpt-4o-medium-tier"]}, + "tiers": { + "MEDIUM": ["gpt-4o-medium-tier"], + "SIMPLE": "simple-unused", + "COMPLEX": "complex-unused", + "REASONING": "reasoning-unused", + }, "default_model": "gpt-4o-configured-default", }, ) @@ -3337,7 +3298,12 @@ class TestRoutingPlugins: def test_plugins_and_adaptive_together_raises(self): with pytest.raises(ValidationError, match="plugins and adaptive=True cannot both be set"): ComplexityRouterConfig( - tiers={"SIMPLE": ["gpt-4o-mini"]}, + tiers={ + "SIMPLE": ["gpt-4o-mini"], + "MEDIUM": "gpt-4o-mini", + "COMPLEX": "gpt-4o-mini", + "REASONING": "gpt-4o-mini", + }, adaptive=True, plugins=[_DummyPlugin()], ) @@ -3369,7 +3335,12 @@ class TestRoutingPlugins: model_name="test-router", litellm_router_instance=mock_router_instance, complexity_router_config={ - "tiers": {"SIMPLE": ["gpt-4o-mini"]}, + "tiers": { + "SIMPLE": ["gpt-4o-mini"], + "MEDIUM": "premium", + "COMPLEX": "premium", + "REASONING": "premium", + }, "session_affinity": True, "plugins": [AllowAll()], }, @@ -3412,19 +3383,32 @@ class TestEscalationKeywords: def test_escalate_tier_caps_at_highest_configured(self, complexity_router): assert complexity_router._escalate_tier(ComplexityTier.REASONING) == ComplexityTier.REASONING - def test_escalate_tier_skips_unconfigured_intermediate(self, mock_router_instance): + def test_escalate_tier_steps_one_tier_at_a_time(self, mock_router_instance): + """Escalation used to skip intermediate tiers that had no models, which only + happened because a tiers map could omit them. Every tier now has models, so a bump + is always exactly one step.""" router = ComplexityRouter( model_name="test-router", litellm_router_instance=mock_router_instance, - complexity_router_config={"tiers": {"SIMPLE": "gpt-4o-mini", "REASONING": "o1-preview"}}, + complexity_router_config={ + "tiers": { + "SIMPLE": "gpt-4o-mini", + "MEDIUM": "gpt-4o", + "COMPLEX": "claude-sonnet-4", + "REASONING": "o1-preview", + } + }, ) - assert router._escalate_tier(ComplexityTier.SIMPLE) == ComplexityTier.REASONING + assert router._escalate_tier(ComplexityTier.SIMPLE) == ComplexityTier.MEDIUM + assert router._escalate_tier(ComplexityTier.COMPLEX) == ComplexityTier.REASONING def test_tier_for_model_returns_most_severe(self, mock_router_instance): router = ComplexityRouter( model_name="test-router", litellm_router_instance=mock_router_instance, - complexity_router_config={"tiers": {"SIMPLE": "shared", "COMPLEX": "shared", "REASONING": "top"}}, + complexity_router_config={ + "tiers": {"SIMPLE": "shared", "COMPLEX": "shared", "REASONING": "top", "MEDIUM": "medium-unused"} + }, ) assert router._tier_for_model("shared") == ComplexityTier.COMPLEX assert router._tier_for_model("top") == ComplexityTier.REASONING @@ -3671,13 +3655,18 @@ class TestEscalationKeywords: every request; surrounding whitespace on real phrases is trimmed.""" assert ( ComplexityRouterConfig( - tiers={"SIMPLE": "gpt-4o-mini", "MEDIUM": "gpt-4o"}, + tiers={ + "SIMPLE": "gpt-4o-mini", + "MEDIUM": "gpt-4o", + "COMPLEX": "gpt-4o-mini", + "REASONING": "gpt-4o-mini", + }, escalation_keywords=["", " "], ).escalation_keywords == [] ) assert ComplexityRouterConfig( - tiers={"SIMPLE": "gpt-4o-mini", "MEDIUM": "gpt-4o"}, + tiers={"SIMPLE": "gpt-4o-mini", "MEDIUM": "gpt-4o", "COMPLEX": "gpt-4o-mini", "REASONING": "gpt-4o-mini"}, escalation_keywords=[" LITELLM ESCALATE ", ""], ).escalation_keywords == ["LITELLM ESCALATE"] @@ -3702,7 +3691,14 @@ class TestEscalationKeywords: router = ComplexityRouter( model_name="test-router", litellm_router_instance=mock_router_instance, - complexity_router_config={"tiers": {"SIMPLE": "gpt-4o-mini", "REASONING": ["o1-a", "o1-b", "o1-c"]}}, + complexity_router_config={ + "tiers": { + "SIMPLE": "gpt-4o-mini", + "REASONING": ["o1-a", "o1-b", "o1-c"], + "MEDIUM": "medium-unused", + "COMPLEX": "complex-unused", + } + }, ) for pinned in ("o1-a", "o1-b", "o1-c"): assert router._escalated_pin(pinned) == pinned @@ -3714,7 +3710,12 @@ class TestEscalationKeywords: model_name="test-router", litellm_router_instance=mock_router_instance, complexity_router_config={ - "tiers": {"SIMPLE": "gpt-4o-mini", "REASONING": ["o1-a", "o1-b", "o1-c"]}, + "tiers": { + "SIMPLE": "gpt-4o-mini", + "REASONING": ["o1-a", "o1-b", "o1-c"], + "MEDIUM": "medium-unused", + "COMPLEX": "complex-unused", + }, "session_affinity": True, }, ) @@ -3947,7 +3948,12 @@ class TestRoutingDecisionSurvivesToSpendLogOnEveryMetadataShape: "litellm_params": { "model": "auto_router/complexity_router", "complexity_router_config": { - "tiers": {"SIMPLE": ["gpt-4o-mini"], "MEDIUM": ["gpt-4o"]}, + "tiers": { + "SIMPLE": ["gpt-4o-mini"], + "MEDIUM": ["gpt-4o"], + "COMPLEX": "complex-unused", + "REASONING": "reasoning-unused", + }, "session_affinity": False, }, }, @@ -4027,7 +4033,12 @@ class TestRoutingDecisionIsPerAttempt: "litellm_params": { "model": "auto_router/complexity_router", "complexity_router_config": { - "tiers": {"SIMPLE": ["gpt-4o-mini"], "MEDIUM": ["gpt-4o"]}, + "tiers": { + "SIMPLE": ["gpt-4o-mini"], + "MEDIUM": ["gpt-4o"], + "COMPLEX": "complex-unused", + "REASONING": "reasoning-unused", + }, "session_affinity": False, }, }, @@ -4036,18 +4047,14 @@ class TestRoutingDecisionIsPerAttempt: {"model_name": "gpt-4o", "litellm_params": {"model": "openai/gpt-4o"}}, ] - @pytest.mark.parametrize( - "seed, bucket", [({}, "metadata"), ({"litellm_metadata": {}}, "litellm_metadata")] - ) + @pytest.mark.parametrize("seed, bucket", [({}, "metadata"), ({"litellm_metadata": {}}, "litellm_metadata")]) @pytest.mark.asyncio async def test_fallback_to_plain_model_group_clears_the_earlier_decision(self, seed, bucket): router = Router(model_list=self.MODEL_LIST) request_kwargs: Dict = dict(seed) messages = [{"role": "user", "content": "Hello!"}] - await router.async_pre_routing_hook( - model="smart-router", request_kwargs=request_kwargs, messages=messages - ) + await router.async_pre_routing_hook(model="smart-router", request_kwargs=request_kwargs, messages=messages) assert "routing_decision" in request_kwargs[bucket] # The fallback attempt reuses the same kwargs and selects no strategy. @@ -4107,7 +4114,12 @@ class TestEscalationIsRecordedConsistently: never happened is the opposite error; both must be avoided identically everywhere.""" CEILING_CONFIG = { - "tiers": {"SIMPLE": ["gpt-4o-mini"], "REASONING": ["o1-preview"]}, + "tiers": { + "SIMPLE": ["gpt-4o-mini"], + "REASONING": ["o1-preview"], + "MEDIUM": "medium-unused", + "COMPLEX": "complex-unused", + }, "session_affinity": False, } @@ -4219,7 +4231,12 @@ class TestRedactedLoggingDropsPromptText: "litellm_params": { "model": "auto_router/complexity_router", "complexity_router_config": { - "tiers": {"SIMPLE": ["gpt-4o-mini"], "REASONING": ["gpt-4o"]}, + "tiers": { + "SIMPLE": ["gpt-4o-mini"], + "REASONING": ["gpt-4o"], + "MEDIUM": "medium-unused", + "COMPLEX": "complex-unused", + }, "session_affinity": False, "keyword_tier_rules": [{"keywords": ["deploy to k8s"], "tier": "REASONING"}], }, @@ -4362,7 +4379,12 @@ class TestContextAwareClassifier: id="multiple-reminders-stripped", ), pytest.param( - [{"role": "user", "content": [{"type": "text", "text": _REMINDER}, {"type": "text", "text": "and now?"}]}], + [ + { + "role": "user", + "content": [{"type": "text", "text": _REMINDER}, {"type": "text", "text": "and now?"}], + } + ], "and now?", id="reminder-in-its-own-content-part", ), @@ -4729,9 +4751,7 @@ class TestContextAwareClassifier: assert reported > 100 @pytest.mark.asyncio - async def test_no_trajectory_signal_when_request_had_no_messages( - self, llm_complexity_router, mock_router_instance - ): + async def test_no_trajectory_signal_when_request_had_no_messages(self, llm_complexity_router, mock_router_instance): """On the prompt-only path there is no conversation to measure, so the depth line is omitted rather than asserting a false "~0 tokens" to the classifier.""" mock_router_instance.acompletion = AsyncMock(return_value=_llm_response('{"tier": "SIMPLE"}')) @@ -4743,9 +4763,7 @@ class TestContextAwareClassifier: assert "what is 2+2" in user_payload @pytest.mark.asyncio - async def test_single_turn_request_sends_no_conversation_context( - self, llm_complexity_router, mock_router_instance - ): + async def test_single_turn_request_sends_no_conversation_context(self, llm_complexity_router, mock_router_instance): """A single-turn request carries no conversation, so the classifier sees only the ask. Found in QA: the depth line gated on `messages` being non-empty, so single-turn requests got a @@ -4772,7 +4790,12 @@ class TestContextAwareClassifier: model_name="test-router", litellm_router_instance=mock_router_instance, complexity_router_config={ - "tiers": {"SIMPLE": "gpt-4o-mini", "COMPLEX": "claude-sonnet-4-20250514"}, + "tiers": { + "SIMPLE": "gpt-4o-mini", + "COMPLEX": "claude-sonnet-4-20250514", + "MEDIUM": "medium-unused", + "REASONING": "reasoning-unused", + }, "classifier_type": "llm", "classifier_llm_config": {"model": "haiku-classifier"}, "classifier_context_window_size": 0, @@ -4795,7 +4818,6 @@ class TestContextAwareClassifier: assert "sharding strategy" not in user_payload assert user_payload.strip() == "Classify this message:\nwhat is 2+2" - @pytest.mark.asyncio @pytest.mark.parametrize("include_assistant,plan_is_quoted", [(True, True), (False, False)]) async def test_assistant_turn_carrying_the_difficulty_reaches_the_classifier( @@ -4837,7 +4859,6 @@ class TestContextAwareClassifier: assert (f"[1] {ask}" in user_payload) is not plan_is_quoted assert user_payload.endswith("Classify this message:\nyes.") - @pytest.mark.asyncio @pytest.mark.parametrize("include_assistant", [True, False]) async def test_depth_signal_agrees_with_what_the_window_quoted( @@ -4940,7 +4961,12 @@ class TestClassifierTrustBoundary: model_name="test-router", litellm_router_instance=mock_router_instance, complexity_router_config={ - "tiers": {"SIMPLE": "gpt-4o-mini", "REASONING": "o1-preview"}, + "tiers": { + "SIMPLE": "gpt-4o-mini", + "REASONING": "o1-preview", + "MEDIUM": "medium-unused", + "COMPLEX": "complex-unused", + }, "classifier_type": "llm", "classifier_llm_config": {"model": "haiku-classifier"}, }, @@ -4959,9 +4985,6 @@ class TestClassifierTrustBoundary: assert hostile not in system_message["content"] assert hostile in user_message["content"] - - - @pytest.mark.parametrize( "window_size,conversation_is_quoted", [ @@ -4988,7 +5011,6 @@ class TestClassifierTrustBoundary: assert ('short reply such as "yes" or "continue"' in system_prompt) is conversation_is_quoted assert ("Classify only the current message" in system_prompt) is not conversation_is_quoted - @pytest.mark.asyncio @pytest.mark.parametrize("include_assistant", [True, False]) async def test_context_framing_does_not_depend_on_which_roles_the_window_holds( @@ -5045,3 +5067,60 @@ class TestClassifierTrustBoundary: assert "Classify only the current message" not in system_prompt assert "using the earlier turns quoted above it as context" in system_prompt assert "rate the work it approves rather than the reply itself" in system_prompt + + +class TestEveryTierNeedsAModel: + """A tier with no models of its own does not route nowhere, it routes somewhere else. + Whatever the fallback chain reaches answers in its place, so the model an operator + configured is silently swapped for another tier's, visible only in a spend log.""" + + ALL_TIERS = {"SIMPLE": "s", "MEDIUM": "m", "COMPLEX": "c", "REASONING": "r"} + + def test_a_complete_map_is_accepted(self): + assert ComplexityRouterConfig(tiers=dict(self.ALL_TIERS)).tiers == self.ALL_TIERS + + def test_omitting_tiers_entirely_still_works(self): + """The defaults name all four, so the requirement costs nothing to anyone who has + not opted into configuring tiers at all.""" + assert set(ComplexityRouterConfig().tiers) == set(self.ALL_TIERS) + + @pytest.mark.parametrize("absent", list(ALL_TIERS)) + def test_each_missing_tier_is_named(self, absent): + with pytest.raises(ValidationError, match=f"no model for {absent}"): + ComplexityRouterConfig(tiers={k: v for k, v in self.ALL_TIERS.items() if k != absent}) + + def test_every_missing_tier_is_named_at_once(self): + """One startup, one edit: an operator should not fix a tier, restart, and learn + about the next one.""" + with pytest.raises(ValidationError, match="no model for COMPLEX, REASONING"): + ComplexityRouterConfig(tiers={"SIMPLE": "s", "MEDIUM": "m"}) + + @pytest.mark.parametrize("no_models", [[], ""], ids=["empty_pool", "empty_pin"]) + def test_a_named_tier_with_nothing_in_it_counts_as_missing(self, no_models): + with pytest.raises(ValidationError, match="no model for MEDIUM"): + ComplexityRouterConfig(tiers={**self.ALL_TIERS, "MEDIUM": no_models}) + + def test_a_pool_counts_as_models(self): + assert ComplexityRouterConfig(tiers={**self.ALL_TIERS, "SIMPLE": ["a", "b"]}).tiers["SIMPLE"] == ["a", "b"] + + def test_extra_tiers_beyond_the_four_are_left_alone(self): + """The requirement is a floor, not a whitelist; a config naming its own tier keys + is not this validator's business.""" + config = ComplexityRouterConfig(tiers={**self.ALL_TIERS, "CUSTOM": "x"}) + assert config.tiers["CUSTOM"] == "x" + + def test_default_model_does_not_satisfy_the_requirement(self): + """default_model is a last resort for a tier that cannot serve, not a substitute + for configuring one.""" + with pytest.raises(ValidationError, match="no model for REASONING"): + ComplexityRouterConfig( + tiers={k: v for k, v in self.ALL_TIERS.items() if k != "REASONING"}, + default_model="fallback-model", + ) + + def test_the_error_tells_an_operator_what_to_do(self): + with pytest.raises(ValidationError) as caught: + ComplexityRouterConfig(tiers={"SIMPLE": "s"}) + message = str(caught.value) + assert "configure all of SIMPLE, MEDIUM, COMPLEX, REASONING" in message + assert "omit `tiers` to take the defaults" in message diff --git a/tests/test_litellm/test_router.py b/tests/test_litellm/test_router.py index 46b5ce65c3f..578c795a63e 100644 --- a/tests/test_litellm/test_router.py +++ b/tests/test_litellm/test_router.py @@ -7,9 +7,7 @@ from unittest.mock import AsyncMock, MagicMock, patch import pytest -sys.path.insert( - 0, os.path.abspath("../../..") -) # Adds the parent directory to the system path +sys.path.insert(0, os.path.abspath("../../..")) # Adds the parent directory to the system path import litellm @@ -113,31 +111,18 @@ def test_router_model_group_encrypted_content_affinity_callback_registration(): num_retries=0, ) callbacks = router.optional_callbacks or [] - encrypted_content_callbacks = [ - cb for cb in callbacks if isinstance(cb, EncryptedContentAffinityCheck) - ] - deployment_callback = next( - cb for cb in callbacks if isinstance(cb, DeploymentAffinityCheck) - ) + encrypted_content_callbacks = [cb for cb in callbacks if isinstance(cb, EncryptedContentAffinityCheck)] + deployment_callback = next(cb for cb in callbacks if isinstance(cb, DeploymentAffinityCheck)) assert len(encrypted_content_callbacks) == 1 assert encrypted_content_callbacks[0].enable_global_affinity is False - assert ( - encrypted_content_callbacks[0].model_group_affinity_config - == model_group_affinity_config - ) - assert callbacks.index(encrypted_content_callbacks[0]) < callbacks.index( - deployment_callback - ) - assert litellm.callbacks.index(encrypted_content_callbacks[0]) < ( - litellm.callbacks.index(deployment_callback) - ) + assert encrypted_content_callbacks[0].model_group_affinity_config == model_group_affinity_config + assert callbacks.index(encrypted_content_callbacks[0]) < callbacks.index(deployment_callback) + assert litellm.callbacks.index(encrypted_content_callbacks[0]) < (litellm.callbacks.index(deployment_callback)) router._add_encrypted_content_affinity_check(enable_global_affinity=True) callbacks = router.optional_callbacks or [] - encrypted_content_callbacks = [ - cb for cb in callbacks if isinstance(cb, EncryptedContentAffinityCheck) - ] + encrypted_content_callbacks = [cb for cb in callbacks if isinstance(cb, EncryptedContentAffinityCheck)] assert len(encrypted_content_callbacks) == 1 assert encrypted_content_callbacks[0].enable_global_affinity is True assert encrypted_content_callbacks[0].router is router @@ -168,13 +153,9 @@ async def test_encrypted_content_affinity_model_group_config_is_additive(): }, target_deployment, ] - encoded_id = ResponsesAPIRequestUtils._build_encrypted_item_id( - "deployment-b", "rs_test" - ) + encoded_id = ResponsesAPIRequestUtils._build_encrypted_item_id("deployment-b", "rs_test") - assert EncryptedContentAffinityCheck.has_model_group_affinity_enabled( - {model_group: ["encrypted_content_affinity"]} - ) + assert EncryptedContentAffinityCheck.has_model_group_affinity_enabled({model_group: ["encrypted_content_affinity"]}) assert not EncryptedContentAffinityCheck.has_model_group_affinity_enabled(None) per_group_check = EncryptedContentAffinityCheck( @@ -215,10 +196,7 @@ async def test_encrypted_content_affinity_model_group_config_is_additive(): ) assert unfiltered == healthy_deployments - assert ( - "encrypted_content_affinity_enabled" - not in disabled_request_kwargs["litellm_metadata"] - ) + assert "encrypted_content_affinity_enabled" not in disabled_request_kwargs["litellm_metadata"] global_check = EncryptedContentAffinityCheck( enable_global_affinity=True, @@ -238,9 +216,7 @@ async def test_encrypted_content_affinity_model_group_config_is_additive(): ) assert globally_filtered == [target_deployment] - assert global_request_kwargs["litellm_metadata"][ - "encrypted_content_affinity_enabled" - ] + assert global_request_kwargs["litellm_metadata"]["encrypted_content_affinity_enabled"] @pytest.mark.asyncio @@ -287,18 +263,10 @@ async def test_encrypted_content_affinity_takes_priority_over_user_key_affinity( num_retries=0, ) callbacks = router.optional_callbacks or [] - deployment_callback = next( - cb for cb in callbacks if isinstance(cb, DeploymentAffinityCheck) - ) - encrypted_content_callback = next( - cb for cb in callbacks if isinstance(cb, EncryptedContentAffinityCheck) - ) - assert callbacks.index(encrypted_content_callback) < callbacks.index( - deployment_callback - ) - assert litellm.callbacks.index(encrypted_content_callback) < ( - litellm.callbacks.index(deployment_callback) - ) + deployment_callback = next(cb for cb in callbacks if isinstance(cb, DeploymentAffinityCheck)) + encrypted_content_callback = next(cb for cb in callbacks if isinstance(cb, EncryptedContentAffinityCheck)) + assert callbacks.index(encrypted_content_callback) < callbacks.index(deployment_callback) + assert litellm.callbacks.index(encrypted_content_callback) < (litellm.callbacks.index(deployment_callback)) cache_key = DeploymentAffinityCheck.get_affinity_cache_key( model_group=model_group, @@ -309,9 +277,7 @@ async def test_encrypted_content_affinity_takes_priority_over_user_key_affinity( value={"model_id": "deployment-a"}, ttl=60, ) - encoded_id = ResponsesAPIRequestUtils._build_encrypted_item_id( - "deployment-b", "rs_test" - ) + encoded_id = ResponsesAPIRequestUtils._build_encrypted_item_id("deployment-b", "rs_test") request_kwargs = { "input": [{"type": "reasoning", "id": encoded_id}], "litellm_metadata": {"user_api_key_hash": user_api_key_hash}, @@ -696,9 +662,7 @@ async def test_arouter_aretrieve_batch(): ], ) - with patch.object( - litellm, "aretrieve_batch", return_value=AsyncMock() - ) as mock_aretrieve_batch: + with patch.object(litellm, "aretrieve_batch", return_value=AsyncMock()) as mock_aretrieve_batch: try: response = await router.aretrieve_batch( model="gpt-3.5-turbo", @@ -719,9 +683,7 @@ async def test_arouter_aretrieve_file_content(): Test that router.acreate_file with JSONL file returns the correct response """ - with patch.object( - litellm, "afile_content", return_value=AsyncMock() - ) as mock_afile_content: + with patch.object(litellm, "afile_content", return_value=AsyncMock()) as mock_afile_content: router = litellm.Router( model_list=[ { @@ -866,9 +828,7 @@ def test_arouter_should_include_deployment(): model=deployment_with_team_and_public_name, team_id="test-team", ) - assert ( - result is True - ), "Should return True when team_id and team_public_model_name match" + assert result is True, "Should return True when team_id and team_public_model_name match" # Test Case 2: Team-specific deployment - team_id matches but model_name doesn't match team_public_model_name result = router.should_include_deployment( @@ -876,9 +836,9 @@ def test_arouter_should_include_deployment(): model=deployment_with_team_and_public_name, team_id="test-team", ) - assert ( - result is False - ), "Should return False when team_id matches but model_name doesn't match team_public_model_name" + assert result is False, ( + "Should return False when team_id matches but model_name doesn't match team_public_model_name" + ) # Test Case 3: Team-specific deployment - team_id doesn't match result = router.should_include_deployment( @@ -894,30 +854,18 @@ def test_arouter_should_include_deployment(): model=deployment_with_team_no_public_name, team_id="test-team", ) - assert ( - result is True - ), "Should return True when team deployment has no team_public_model_name to match" + assert result is True, "Should return True when team deployment has no team_public_model_name to match" # Test Case 5: Non-team deployment - model_name matches and no team_id - result = router.should_include_deployment( - model_name="gpt-4", model=deployment_without_team, team_id=None - ) - assert ( - result is True - ), "Should return True when model_name matches and deployment has no team_id" + result = router.should_include_deployment(model_name="gpt-4", model=deployment_without_team, team_id=None) + assert result is True, "Should return True when model_name matches and deployment has no team_id" # Test Case 6: Non-team deployment - model_name matches but team_id provided (should still work) - result = router.should_include_deployment( - model_name="gpt-4", model=deployment_without_team, team_id="any-team" - ) - assert ( - result is True - ), "Should return True when model_name matches non-team deployment, regardless of team_id param" + result = router.should_include_deployment(model_name="gpt-4", model=deployment_without_team, team_id="any-team") + assert result is True, "Should return True when model_name matches non-team deployment, regardless of team_id param" # Test Case 7: Non-team deployment - model_name doesn't match - result = router.should_include_deployment( - model_name="different-model", model=deployment_without_team, team_id=None - ) + result = router.should_include_deployment(model_name="different-model", model=deployment_without_team, team_id=None) assert result is False, "Should return False when model_name doesn't match" # Test Case 8: Team deployment accessed without matching team_id @@ -926,9 +874,7 @@ def test_arouter_should_include_deployment(): model=deployment_with_team_and_public_name, team_id=None, ) - assert ( - result is True - ), "Should return True when matching model with exact model_name" + assert result is True, "Should return True when matching model with exact model_name" def test_arouter_responses_api_bridge(): @@ -978,9 +924,7 @@ def test_arouter_responses_api_bridge(): "status": "completed", "output": [], } - mock_response.text = ( - '{"id": "resp_test", "object": "response", "status": "completed", "output": []}' - ) + mock_response.text = '{"id": "resp_test", "object": "response", "status": "completed", "output": []}' with patch.object(client, "post", return_value=mock_response) as mock_post: try: @@ -1106,15 +1050,9 @@ async def test_router_ageneric_api_call_with_fallbacks_helper(): }, } - with patch.object( - router, "_update_kwargs_with_deployment" - ) as mock_update_kwargs: - with patch.object( - router, "async_routing_strategy_pre_call_checks" - ) as mock_pre_call_checks: - with patch.object( - router, "_get_client", return_value=None - ) as mock_get_client: + with patch.object(router, "_update_kwargs_with_deployment") as mock_update_kwargs: + with patch.object(router, "async_routing_strategy_pre_call_checks") as mock_pre_call_checks: + with patch.object(router, "_get_client", return_value=None) as mock_get_client: result = await router._ageneric_api_call_with_fallbacks_helper( model="gpt-3.5-turbo", original_generic_function=mock_generic_function, @@ -1177,15 +1115,9 @@ async def test_router_ageneric_api_call_with_fallbacks_helper(): mock_semaphore = asyncio.Semaphore(1) - with patch.object( - router, "_update_kwargs_with_deployment" - ) as mock_update_kwargs: - with patch.object( - router, "_get_client", return_value=mock_semaphore - ) as mock_get_client: - with patch.object( - router, "async_routing_strategy_pre_call_checks" - ) as mock_pre_call_checks: + with patch.object(router, "_update_kwargs_with_deployment") as mock_update_kwargs: + with patch.object(router, "_get_client", return_value=mock_semaphore) as mock_get_client: + with patch.object(router, "async_routing_strategy_pre_call_checks") as mock_pre_call_checks: result = await router._ageneric_api_call_with_fallbacks_helper( model="gpt-3.5-turbo", original_generic_function=mock_semaphore_function, @@ -1214,15 +1146,9 @@ async def test_router_ageneric_api_call_with_fallbacks_helper(): }, } - with patch.object( - router, "_update_kwargs_with_deployment" - ) as mock_update_kwargs: - with patch.object( - router, "_get_client", return_value=None - ) as mock_get_client: - with patch.object( - router, "async_routing_strategy_pre_call_checks" - ) as mock_pre_call_checks: + with patch.object(router, "_update_kwargs_with_deployment") as mock_update_kwargs: + with patch.object(router, "_get_client", return_value=None) as mock_get_client: + with patch.object(router, "async_routing_strategy_pre_call_checks") as mock_pre_call_checks: with pytest.raises(Exception) as exc_info: await router._ageneric_api_call_with_fallbacks_helper( model="gpt-3.5-turbo", @@ -1291,9 +1217,9 @@ async def test_ageneric_api_call_deployment_model_overrides_alias(): original_generic_function=capture_model, ) - assert ( - captured["model"] == "vertex_ai/gemini-2.5-flash" - ), f"Expected deployment model 'vertex_ai/gemini-2.5-flash', got '{captured['model']}'" + assert captured["model"] == "vertex_ai/gemini-2.5-flash", ( + f"Expected deployment model 'vertex_ai/gemini-2.5-flash', got '{captured['model']}'" + ) def test_router_get_model_access_groups_team_only_models(): @@ -1314,14 +1240,10 @@ def test_router_get_model_access_groups_team_only_models(): ] ) - access_groups = router.get_model_access_groups( - model_name="gpt-3.5-turbo", team_id=None - ) + access_groups = router.get_model_access_groups(model_name="gpt-3.5-turbo", team_id=None) assert len(access_groups) == 0 - access_groups = router.get_model_access_groups( - model_name="gpt-3.5-turbo", team_id="team_1" - ) + access_groups = router.get_model_access_groups(model_name="gpt-3.5-turbo", team_id="team_1") assert list(access_groups.keys()) == ["default-models"] @@ -1416,9 +1338,7 @@ def test_model_group_info_cost_from_db_model_info(): ] ) - with patch.object( - router, "get_deployment_model_info", side_effect=Exception("not found") - ): + with patch.object(router, "get_deployment_model_info", side_effect=Exception("not found")): result = router._cached_get_model_group_info("my-custom-model") assert result is not None assert result.input_cost_per_token == 0.0001 @@ -1446,9 +1366,7 @@ def test_model_group_info_cost_none_when_db_model_info_has_no_cost(): ] ) - with patch.object( - router, "get_deployment_model_info", side_effect=Exception("not found") - ): + with patch.object(router, "get_deployment_model_info", side_effect=Exception("not found")): result = router._cached_get_model_group_info("my-custom-model-no-cost") assert result is not None assert result.input_cost_per_token is None @@ -1518,9 +1436,7 @@ def test_model_group_info_with_stringified_cost_values(): } return None - with patch.object( - router, "get_deployment_model_info", side_effect=_model_info_with_str_costs - ): + with patch.object(router, "get_deployment_model_info", side_effect=_model_info_with_str_costs): result = router._set_model_group_info( model_group="my-custom-model", user_facing_model_group_name="my-custom-model", @@ -1566,9 +1482,7 @@ def test_model_group_info_db_fallback_with_stringified_cost_values(): ] ) - with patch.object( - router, "get_deployment_model_info", side_effect=Exception("not found") - ): + with patch.object(router, "get_deployment_model_info", side_effect=Exception("not found")): result = router._set_model_group_info( model_group="my-custom-model", user_facing_model_group_name="my-custom-model", @@ -1812,9 +1726,7 @@ async def test_acompletion_streaming_iterator(): self.index += 1 return item - mock_error_response = AsyncIteratorWithError( - mock_chunks, 1 - ) # Error after first chunk + mock_error_response = AsyncIteratorWithError(mock_chunks, 1) # Error after first chunk setattr(mock_error_response, "model", "gpt-4") setattr(mock_error_response, "custom_llm_provider", "openai") @@ -2208,11 +2120,7 @@ def _make_responses_iterator( BaseResponsesAPIStreamingIterator, ) - base = ( - LiteLLMCompletionStreamingIterator - if bridge - else BaseResponsesAPIStreamingIterator - ) + base = LiteLLMCompletionStreamingIterator if bridge else BaseResponsesAPIStreamingIterator class _Iter(base): def __init__(self): @@ -2292,9 +2200,7 @@ async def test_aresponses_streaming_iterator_fallback(): BaseResponsesAPIStreamingIterator, ) - router = _make_router_with_fallback( - "anthropic/claude-sonnet-4-6", "vertex_ai/claude-sonnet-4-6" - ) + router = _make_router_with_fallback("anthropic/claude-sonnet-4-6", "vertex_ai/claude-sonnet-4-6") src = _make_responses_iterator( chunks=[MagicMock(type="response.created")], error=MidStreamFallbackError( @@ -2377,9 +2283,9 @@ async def test_aresponses_streaming_iterator_writes_litellm_metadata_on_fallback fbk = mock_fallback_utils.call_args.kwargs["kwargs"] assert "litellm_metadata" in fbk, "wrong metadata_variable_name" assert fbk["litellm_metadata"]["model_group"] == "gpt-4" - assert "model_group" not in fbk.get( - "metadata", {} - ), "model_group leaked into 'metadata' instead of 'litellm_metadata'" + assert "model_group" not in fbk.get("metadata", {}), ( + "model_group leaked into 'metadata' instead of 'litellm_metadata'" + ) @pytest.mark.asyncio @@ -2497,9 +2403,7 @@ async def test_aresponses_streaming_iterator_combines_partial_usage(): fallback_response_object = ResponsesAPIResponse( id="resp_test", created_at=0, model="gpt-4", object="response", output=[] ) - fallback_response_object.usage = ResponseAPIUsage( - input_tokens=20, output_tokens=15, total_tokens=35 - ) + fallback_response_object.usage = ResponseAPIUsage(input_tokens=20, output_tokens=15, total_tokens=35) fallback_event = ResponseCompletedEvent( type=ResponsesAPIStreamEvents.RESPONSE_COMPLETED, response=fallback_response_object, @@ -2508,9 +2412,7 @@ async def test_aresponses_streaming_iterator_combines_partial_usage(): with ( patch( "litellm.main.stream_chunk_builder", - return_value=SimpleNamespace( - usage=SimpleNamespace(prompt_tokens=10, completion_tokens=4) - ), + return_value=SimpleNamespace(usage=SimpleNamespace(prompt_tokens=10, completion_tokens=4)), ), patch.object( router, @@ -2811,9 +2713,7 @@ def test_pre_call_checks_skips_token_count_without_max_input_tokens(monkeypatch) monkeypatch.setattr(router, "get_router_model_info", lambda **kwargs: {}) calls = [] - monkeypatch.setattr( - litellm, "token_counter", lambda *a, **k: calls.append(1) or 1000 - ) + monkeypatch.setattr(litellm, "token_counter", lambda *a, **k: calls.append(1) or 1000) deployments = [ {"litellm_params": {"model": "gpt-3.5-turbo"}, "model_info": {"id": "d1"}}, @@ -2841,14 +2741,10 @@ def test_pre_call_checks_counts_once_and_filters_on_max_input_tokens(monkeypatch ], enable_pre_call_checks=True, ) - monkeypatch.setattr( - router, "get_router_model_info", lambda **kwargs: {"max_input_tokens": 5} - ) + monkeypatch.setattr(router, "get_router_model_info", lambda **kwargs: {"max_input_tokens": 5}) calls = [] - monkeypatch.setattr( - litellm, "token_counter", lambda *a, **k: calls.append(1) or 1000 - ) + monkeypatch.setattr(litellm, "token_counter", lambda *a, **k: calls.append(1) or 1000) deployments = [ {"litellm_params": {"model": "gpt-3.5-turbo"}, "model_info": {"id": "d1"}}, @@ -2876,9 +2772,7 @@ def test_pre_call_checks_counts_tokens_from_responses_input_string(monkeypatch): ], enable_pre_call_checks=True, ) - monkeypatch.setattr( - router, "get_router_model_info", lambda **kwargs: {"max_input_tokens": 1} - ) + monkeypatch.setattr(router, "get_router_model_info", lambda **kwargs: {"max_input_tokens": 1}) deployments = [ {"litellm_params": {"model": "gpt-3.5-turbo"}, "model_info": {"id": "d1"}}, @@ -2903,9 +2797,7 @@ def test_pre_call_checks_counts_tokens_from_responses_input_list(monkeypatch): ], enable_pre_call_checks=True, ) - monkeypatch.setattr( - router, "get_router_model_info", lambda **kwargs: {"max_input_tokens": 1} - ) + monkeypatch.setattr(router, "get_router_model_info", lambda **kwargs: {"max_input_tokens": 1}) deployments = [ {"litellm_params": {"model": "gpt-3.5-turbo"}, "model_info": {"id": "d1"}}, @@ -2947,9 +2839,7 @@ def test_pre_call_checks_counts_responses_instructions_tokens(monkeypatch): ) assert with_instructions_tokens > input_only_tokens - monkeypatch.setattr( - router, "get_router_model_info", lambda **kwargs: {"max_input_tokens": input_only_tokens} - ) + monkeypatch.setattr(router, "get_router_model_info", lambda **kwargs: {"max_input_tokens": input_only_tokens}) with pytest.raises(litellm.ContextWindowExceededError): router._pre_call_checks( model="m", @@ -2997,9 +2887,7 @@ def test_pre_call_checks_no_messages_or_input_does_not_crash(monkeypatch): ], enable_pre_call_checks=True, ) - monkeypatch.setattr( - router, "get_router_model_info", lambda **kwargs: {"max_input_tokens": 5} - ) + monkeypatch.setattr(router, "get_router_model_info", lambda **kwargs: {"max_input_tokens": 5}) counted: list[dict] = [] original = router._count_pre_call_check_tokens @@ -3089,9 +2977,7 @@ def test_get_deployment_model_info_base_model_flow(): } # Test Case 1: Base model flow with custom model info that has base_model - with patch.object( - litellm, "model_cost", {"test-custom-model": mock_custom_model_info} - ): + with patch.object(litellm, "model_cost", {"test-custom-model": mock_custom_model_info}): with patch.object(litellm, "get_model_info") as mock_get_model_info: # Configure mock returns mock_get_model_info.side_effect = lambda model: { @@ -3099,15 +2985,11 @@ def test_get_deployment_model_info_base_model_flow(): "test-model": mock_litellm_model_name_info, }.get(model) - result = router.get_deployment_model_info( - model_id="test-custom-model", model_name="test-model" - ) + result = router.get_deployment_model_info(model_id="test-custom-model", model_name="test-model") # Verify that get_model_info was called for both base model and model name assert mock_get_model_info.call_count == 2 - mock_get_model_info.assert_any_call( - model="gpt-3.5-turbo" - ) # base model call + mock_get_model_info.assert_any_call(model="gpt-3.5-turbo") # base model call mock_get_model_info.assert_any_call(model="test-model") # model name call # Verify the result contains merged information @@ -3118,26 +3000,18 @@ def test_get_deployment_model_info_base_model_flow(): # 2. The result of step 1 gets merged into litellm_model_name_info (custom+base override litellm) # Fields from custom model (should override base model values) - assert ( - result["input_cost_per_token"] == 0.001 - ) # From custom model (overrides base 0.0015) - assert ( - result["output_cost_per_token"] == 0.002 - ) # From custom model (same as base) + assert result["input_cost_per_token"] == 0.001 # From custom model (overrides base 0.0015) + assert result["output_cost_per_token"] == 0.002 # From custom model (same as base) assert result["custom_field"] == "custom_value" # From custom model # Fields from base model that weren't overridden by custom assert result["max_tokens"] == 4096 # From base model assert result["litellm_provider"] == "openai" # From base model - assert ( - result["mode"] == "chat" - ) # From base model (overrides litellm "completion") + assert result["mode"] == "chat" # From base model (overrides litellm "completion") # The key field comes from base model since both base and litellm have it # and base model info overrides litellm model name info in final merge - assert ( - result["key"] == "gpt-3.5-turbo" - ) # From base model (overrides litellm key) + assert result["key"] == "gpt-3.5-turbo" # From base model (overrides litellm key) # Test Case 2: Custom model info without base_model mock_custom_model_info_no_base = { @@ -3156,9 +3030,7 @@ def test_get_deployment_model_info_base_model_flow(): "test-model": mock_litellm_model_name_info, }.get(model) - result = router.get_deployment_model_info( - model_id="test-custom-model-no-base", model_name="test-model" - ) + result = router.get_deployment_model_info(model_id="test-custom-model-no-base", model_name="test-model") # Should only call get_model_info once for model name (no base model) assert mock_get_model_info.call_count == 1 @@ -3178,9 +3050,7 @@ def test_get_deployment_model_info_base_model_flow(): "test-model": mock_litellm_model_name_info, }.get(model) - result = router.get_deployment_model_info( - model_id="non-existent-model", model_name="test-model" - ) + result = router.get_deployment_model_info(model_id="non-existent-model", model_name="test-model") # Should only call get_model_info once for model name assert mock_get_model_info.call_count == 1 @@ -3213,9 +3083,7 @@ def test_get_deployment_model_info_base_model_flow(): mock_get_model_info.side_effect = mock_get_model_info_side_effect - result = router.get_deployment_model_info( - model_id="test-custom-model-invalid", model_name="test-model" - ) + result = router.get_deployment_model_info(model_id="test-custom-model-invalid", model_name="test-model") # Should handle exception gracefully and still return merged result assert result is not None @@ -3224,12 +3092,8 @@ def test_get_deployment_model_info_base_model_flow(): # Test Case 5: Both model_cost.get() and get_model_info() return None with patch.object(litellm, "model_cost", {}): - with patch.object( - litellm, "get_model_info", side_effect=Exception("Not found") - ): - result = router.get_deployment_model_info( - model_id="non-existent", model_name="non-existent" - ) + with patch.object(litellm, "get_model_info", side_effect=Exception("Not found")): + result = router.get_deployment_model_info(model_id="non-existent", model_name="non-existent") # Should return None when no model info is found assert result is None @@ -3252,9 +3116,7 @@ def test_get_deployment_model_info_base_model_flow(): # Model NOT in built-in cost map — raise exception mock_get_model_info.side_effect = Exception("Model not in cost map") - result = router.get_deployment_model_info( - model_id="custom-model-id", model_name="unknown-model" - ) + result = router.get_deployment_model_info(model_id="custom-model-id", model_name="unknown-model") # Should return custom_model_info even when litellm_model_name_model_info is None assert result is not None @@ -3290,15 +3152,11 @@ def test_get_deployment_model_info_base_model_flow(): mock_get_model_info.side_effect = get_info_side_effect - result = router.get_deployment_model_info( - model_id="custom-with-base", model_name="unknown-model" - ) + result = router.get_deployment_model_info(model_id="custom-with-base", model_name="unknown-model") # Should return custom_model_info merged with base model info assert result is not None - assert ( - result["input_cost_per_token"] == 0.01 - ) # From custom (overrides base) + assert result["input_cost_per_token"] == 0.01 # From custom (overrides base) assert result["max_tokens"] == 8192 # From base model assert result["litellm_provider"] == "openai" # From base model @@ -3345,18 +3203,14 @@ def test_get_deployment_model_info_base_model_merge_priority(): "litellm_only_field": "litellm_value", } - with patch.object( - litellm, "model_cost", {"custom-model-id": mock_custom_model_info} - ): + with patch.object(litellm, "model_cost", {"custom-model-id": mock_custom_model_info}): with patch.object(litellm, "get_model_info") as mock_get_model_info: mock_get_model_info.side_effect = lambda model: { "gpt-4": mock_base_model_info, "test-model": mock_litellm_model_name_info, }.get(model) - result = router.get_deployment_model_info( - model_id="custom-model-id", model_name="test-model" - ) + result = router.get_deployment_model_info(model_id="custom-model-id", model_name="test-model") assert result is not None @@ -3366,29 +3220,17 @@ def test_get_deployment_model_info_base_model_merge_priority(): # 3. Result from steps 1-2 overrides litellm_model_name_info # Fields that should come from custom model info (highest priority) - assert ( - result["input_cost_per_token"] == 0.01 - ) # From custom model (overrides base 0.03) - assert ( - result["max_tokens"] == 8000 - ) # From custom model (overrides base 4096) + assert result["input_cost_per_token"] == 0.01 # From custom model (overrides base 0.03) + assert result["max_tokens"] == 8000 # From custom model (overrides base 4096) assert result["custom_only_field"] == "custom_value" # From custom model # Fields that should come from base model (not overridden by custom) - assert ( - result["output_cost_per_token"] == 0.06 - ) # From base model (not in custom) - assert ( - result["litellm_provider"] == "openai" - ) # From base model (not in custom) - assert ( - result["base_only_field"] == "base_value" - ) # From base model (not in custom) + assert result["output_cost_per_token"] == 0.06 # From base model (not in custom) + assert result["litellm_provider"] == "openai" # From base model (not in custom) + assert result["base_only_field"] == "base_value" # From base model (not in custom) # Fields that should come from litellm model name info (not overridden by custom+base) - assert ( - result["mode"] == "completion" - ) # From litellm model name info (not in custom or base) + assert result["mode"] == "completion" # From litellm model name info (not in custom or base) assert ( result["litellm_only_field"] == "litellm_value" ) # From litellm model name info (not in custom or base) @@ -3425,10 +3267,9 @@ def test_add_deployment_model_to_endpoint_for_llm_passthrough_route(): model="special-bedrock-model", model_name="bedrock/us.anthropic.claude-haiku-4-5-20251001-v1:0", ) - assert ( - result["endpoint"] - == "/model/us.anthropic.claude-haiku-4-5-20251001-v1:0/invoke" - ), f"Expected '/model/us.anthropic.claude-haiku-4-5-20251001-v1:0/invoke', got '{result['endpoint']}'" + assert result["endpoint"] == "/model/us.anthropic.claude-haiku-4-5-20251001-v1:0/invoke", ( + f"Expected '/model/us.anthropic.claude-haiku-4-5-20251001-v1:0/invoke', got '{result['endpoint']}'" + ) # Test Case 2: Bedrock invoke-with-response-stream endpoint kwargs = { @@ -3440,10 +3281,9 @@ def test_add_deployment_model_to_endpoint_for_llm_passthrough_route(): model="special-bedrock-model", model_name="bedrock/us.anthropic.claude-haiku-4-5-20251001-v1:0", ) - assert ( - result["endpoint"] - == "/model/us.anthropic.claude-haiku-4-5-20251001-v1:0/invoke-with-response-stream" - ), f"Expected streaming endpoint with stripped prefix, got '{result['endpoint']}'" + assert result["endpoint"] == "/model/us.anthropic.claude-haiku-4-5-20251001-v1:0/invoke-with-response-stream", ( + f"Expected streaming endpoint with stripped prefix, got '{result['endpoint']}'" + ) # Test Case 3: Bedrock converse endpoint kwargs = { @@ -3455,9 +3295,9 @@ def test_add_deployment_model_to_endpoint_for_llm_passthrough_route(): model="bedrock-model", model_name="bedrock/us.meta.llama3-8b-instruct-v1:0", ) - assert ( - result["endpoint"] == "/model/us.meta.llama3-8b-instruct-v1:0/converse" - ), f"Expected '/model/us.meta.llama3-8b-instruct-v1:0/converse', got '{result['endpoint']}'" + assert result["endpoint"] == "/model/us.meta.llama3-8b-instruct-v1:0/converse", ( + f"Expected '/model/us.meta.llama3-8b-instruct-v1:0/converse', got '{result['endpoint']}'" + ) # Test Case 4: Bedrock provider prefix auto-detected from model_name kwargs = { @@ -3468,9 +3308,9 @@ def test_add_deployment_model_to_endpoint_for_llm_passthrough_route(): model="router-model", model_name="bedrock/us.meta.llama3-8b-instruct-v1:0", ) - assert ( - result["endpoint"] == "/model/us.meta.llama3-8b-instruct-v1:0/invoke" - ), f"Expected '/model/us.meta.llama3-8b-instruct-v1:0/invoke', got '{result['endpoint']}'" + assert result["endpoint"] == "/model/us.meta.llama3-8b-instruct-v1:0/invoke", ( + f"Expected '/model/us.meta.llama3-8b-instruct-v1:0/invoke', got '{result['endpoint']}'" + ) def test_update_kwargs_with_deployment_uses_pass_through_request_timeout(): @@ -3522,14 +3362,10 @@ async def test_router_acompletion_with_unknown_model_and_default_fallback(): # Initialize the router with a default fallback router = litellm.Router(model_list=model_list, default_fallbacks=["gpt-4o"]) - messages = [ - {"role": "user", "content": "This call should succeed by falling back."} - ] + messages = [{"role": "user", "content": "This call should succeed by falling back."}] # Call completion with a model name that is NOT in the model_list - response = await router.acompletion( - model="completely-unknown-model", messages=messages - ) + response = await router.acompletion(model="completely-unknown-model", messages=messages) # Check that the call did not fail and we received a valid response object. assert response is not None @@ -3621,15 +3457,10 @@ def test_get_deployment_credentials_with_provider_aws_bedrock_runtime_endpoint() ], ) - credentials = router.get_deployment_credentials_with_provider( - model_id="bedrock-claude-model" - ) + credentials = router.get_deployment_credentials_with_provider(model_id="bedrock-claude-model") assert credentials is not None - assert ( - credentials["aws_bedrock_runtime_endpoint"] - == "https://bedrock-runtime.us-east-1.amazonaws.com" - ) + assert credentials["aws_bedrock_runtime_endpoint"] == "https://bedrock-runtime.us-east-1.amazonaws.com" assert credentials["aws_access_key_id"] == "test-access-key" assert credentials["aws_secret_access_key"] == "test-secret-key" assert credentials["aws_region_name"] == "us-east-1" @@ -3656,9 +3487,7 @@ def test_get_deployment_credentials_with_provider_includes_bucket_name(): ], ) - credentials = router.get_deployment_credentials_with_provider( - model_id="vertex-gemini" - ) + credentials = router.get_deployment_credentials_with_provider(model_id="vertex-gemini") assert credentials is not None assert credentials["gcs_bucket_name"] == "my-batch-bucket" @@ -3698,9 +3527,7 @@ def test_get_deployment_credentials_with_provider_resolves_credential_name(): ], ) - credentials = router.get_deployment_credentials_with_provider( - model_id="azure-gpt-4" - ) + credentials = router.get_deployment_credentials_with_provider(model_id="azure-gpt-4") assert credentials is not None assert credentials["api_key"] == "resolved-api-key" @@ -3742,15 +3569,11 @@ def test_get_deployment_credentials_with_provider_team_wildcard_priority(): ], ) - team_credentials = router.get_deployment_credentials_with_provider( - model_id="openai/gpt-5.2", team_id="team-1" - ) + team_credentials = router.get_deployment_credentials_with_provider(model_id="openai/gpt-5.2", team_id="team-1") assert team_credentials is not None assert team_credentials["api_key"] == "team-key" - global_credentials = router.get_deployment_credentials_with_provider( - model_id="openai/gpt-5.2" - ) + global_credentials = router.get_deployment_credentials_with_provider(model_id="openai/gpt-5.2") assert global_credentials is not None assert global_credentials["api_key"] == "global-key" @@ -3791,15 +3614,11 @@ def test_get_deployment_credentials_with_provider_skips_other_team_deployment(): assert other_team_credentials is not None assert other_team_credentials["vertex_project"] == "shared-project" - unscoped_credentials = router.get_deployment_credentials_with_provider( - model_id="gemini-2.5-pro" - ) + unscoped_credentials = router.get_deployment_credentials_with_provider(model_id="gemini-2.5-pro") assert unscoped_credentials is not None assert unscoped_credentials["vertex_project"] == "shared-project" - owner_credentials = router.get_deployment_credentials_with_provider( - model_id="gemini-2.5-pro", team_id="team-b" - ) + owner_credentials = router.get_deployment_credentials_with_provider(model_id="gemini-2.5-pro", team_id="team-b") assert owner_credentials is not None assert owner_credentials["vertex_project"] == "team-b-project" @@ -3826,16 +3645,8 @@ def test_get_deployment_credentials_with_provider_no_fallback_to_other_team_only ], ) - assert ( - router.get_deployment_credentials_with_provider( - model_id="gemini-2.5-pro", team_id="team-a" - ) - is None - ) - assert ( - router.get_deployment_credentials_with_provider(model_id="gemini-2.5-pro") - is None - ) + assert router.get_deployment_credentials_with_provider(model_id="gemini-2.5-pro", team_id="team-a") is None + assert router.get_deployment_credentials_with_provider(model_id="gemini-2.5-pro") is None def test_deployment_usable_by_team_helpers(): @@ -3875,9 +3686,7 @@ def test_deployment_usable_by_team_helpers(): assert router._deployment_usable_by_team(shared, "team-a") is True assert router._deployment_usable_by_team(shared, None) is True - picked = router._get_model_group_deployment_usable_by_team( - model_group_name="gemini-2.5-pro", team_id="team-a" - ) + picked = router._get_model_group_deployment_usable_by_team(model_group_name="gemini-2.5-pro", team_id="team-a") assert picked is not None assert picked.litellm_params.vertex_project == "shared-project" @@ -3887,12 +3696,7 @@ def test_deployment_usable_by_team_helpers(): assert owner_picked is not None assert owner_picked.litellm_params.vertex_project == "team-b-project" - assert ( - router._get_model_group_deployment_usable_by_team( - model_group_name="unknown-model", team_id="team-a" - ) - is None - ) + assert router._get_model_group_deployment_usable_by_team(model_group_name="unknown-model", team_id="team-a") is None def test_get_deployment_credentials_with_provider_skips_other_team_wildcard(): @@ -3924,9 +3728,7 @@ def test_get_deployment_credentials_with_provider_skips_other_team_wildcard(): assert other_team_credentials is not None assert other_team_credentials["api_key"] == "global-key" - owner_credentials = router.get_deployment_credentials_with_provider( - model_id="openai/gpt-5.2", team_id="team-b" - ) + owner_credentials = router.get_deployment_credentials_with_provider(model_id="openai/gpt-5.2", team_id="team-b") assert owner_credentials is not None assert owner_credentials["api_key"] == "team-b-key" @@ -3938,21 +3740,11 @@ def test_team_wildcard_credentials_not_usable_after_delete_deployment(): """ router = litellm.Router(model_list=[_team_wildcard_model(api_key="old-key")]) - assert ( - router.get_deployment_credentials_with_provider( - model_id="openai/gpt-5.2", team_id="team-1" - ) - is not None - ) + assert router.get_deployment_credentials_with_provider(model_id="openai/gpt-5.2", team_id="team-1") is not None router.delete_deployment(id="team-wildcard-id") - assert ( - router.get_deployment_credentials_with_provider( - model_id="openai/gpt-5.2", team_id="team-1" - ) - is None - ) + assert router.get_deployment_credentials_with_provider(model_id="openai/gpt-5.2", team_id="team-1") is None def test_pattern_match_router_remove_deployment(): @@ -3991,22 +3783,13 @@ def test_team_wildcard_credentials_refreshed_on_upsert_and_set_model_list(): router = litellm.Router(model_list=[_team_wildcard_model(api_key="old-key")]) - router.upsert_deployment( - deployment=Deployment(**_team_wildcard_model(api_key="new-key")) - ) - credentials = router.get_deployment_credentials_with_provider( - model_id="openai/gpt-5.2", team_id="team-1" - ) + router.upsert_deployment(deployment=Deployment(**_team_wildcard_model(api_key="new-key"))) + credentials = router.get_deployment_credentials_with_provider(model_id="openai/gpt-5.2", team_id="team-1") assert credentials is not None assert credentials["api_key"] == "new-key" router.set_model_list(model_list=[]) - assert ( - router.get_deployment_credentials_with_provider( - model_id="openai/gpt-5.2", team_id="team-1" - ) - is None - ) + assert router.get_deployment_credentials_with_provider(model_id="openai/gpt-5.2", team_id="team-1") is None def test_get_available_guardrail_single_deployment(): @@ -4195,9 +3978,7 @@ async def test_anthropic_messages_call_type_is_cached(): startTime=1234567890.0, endTime=1234567891.0, completionStartTime=1234567890.5, - model_map_information=StandardLoggingModelInformation( - model_map_key="gpt-3.5-turbo", model_map_value=None - ), + model_map_information=StandardLoggingModelInformation(model_map_key="gpt-3.5-turbo", model_map_value=None), model="gpt-3.5-turbo", model_id="model-123", model_group="openai-gpt", @@ -4276,12 +4057,8 @@ async def test_anthropic_messages_call_type_is_cached(): ) # This assertion will FAIL if anthropic_messages is filtered out - assert ( - cached_result is not None - ), "Model ID should be cached for anthropic_messages call type" - assert ( - cached_result["model_id"] == test_model_id - ), f"Expected {test_model_id}, got {cached_result['model_id']}" + assert cached_result is not None, "Model ID should be cached for anthropic_messages call type" + assert cached_result["model_id"] == test_model_id, f"Expected {test_model_id}, got {cached_result['model_id']}" def test_update_kwargs_with_deployment_propagates_model_tags(): @@ -4306,9 +4083,7 @@ def test_update_kwargs_with_deployment_propagates_model_tags(): ) kwargs: dict = {"metadata": {}} - deployment = router.get_deployment_by_model_group_name( - model_group_name="gpt-4o-mini" - ) + deployment = router.get_deployment_by_model_group_name(model_group_name="gpt-4o-mini") router._update_kwargs_with_deployment(deployment=deployment, kwargs=kwargs) # Deployment tags should be propagated to kwargs metadata @@ -4337,9 +4112,7 @@ def test_update_kwargs_with_deployment_merges_tags_without_duplicates(): # Simulate request that already has tags (from request body or key/team level) kwargs: dict = {"metadata": {"tags": ["user-tag", "shared-tag"]}} - deployment = router.get_deployment_by_model_group_name( - model_group_name="gpt-4o-mini" - ) + deployment = router.get_deployment_by_model_group_name(model_group_name="gpt-4o-mini") router._update_kwargs_with_deployment(deployment=deployment, kwargs=kwargs) # Both sources should be merged, no duplicates @@ -4366,9 +4139,7 @@ def test_update_kwargs_with_deployment_no_tags(): ) kwargs: dict = {"metadata": {}} - deployment = router.get_deployment_by_model_group_name( - model_group_name="gpt-4o-mini" - ) + deployment = router.get_deployment_by_model_group_name(model_group_name="gpt-4o-mini") router._update_kwargs_with_deployment(deployment=deployment, kwargs=kwargs) # No tags key should be added if deployment has no tags @@ -4406,9 +4177,7 @@ def test_update_kwargs_with_deployment_merges_tools(): }, ], } - deployment = router.get_deployment_by_model_group_name( - model_group_name="o3-deep-research" - ) + deployment = router.get_deployment_by_model_group_name(model_group_name="o3-deep-research") router._update_kwargs_with_deployment(deployment=deployment, kwargs=kwargs) # Tools should be merged: deployment first, then request @@ -4439,9 +4208,7 @@ def test_update_kwargs_with_deployment_merge_tools_deployment_only(): ) kwargs: dict = {"metadata": {}} - deployment = router.get_deployment_by_model_group_name( - model_group_name="o3-deep-research" - ) + deployment = router.get_deployment_by_model_group_name(model_group_name="o3-deep-research") router._update_kwargs_with_deployment(deployment=deployment, kwargs=kwargs) assert kwargs["tools"] == [{"type": "web_search"}] @@ -4470,9 +4237,7 @@ def test_update_kwargs_with_deployment_merge_tools_request_overrides_tool_choice "metadata": {}, "tool_choice": "none", } - deployment = router.get_deployment_by_model_group_name( - model_group_name="o3-deep-research" - ) + deployment = router.get_deployment_by_model_group_name(model_group_name="o3-deep-research") router._update_kwargs_with_deployment(deployment=deployment, kwargs=kwargs) # Request tool_choice should be preserved (merged tools still applied) @@ -4574,12 +4339,8 @@ def test_update_kwargs_with_deployment_model_info_in_litellm_metadata(): ) kwargs: dict = {} - deployment = router.get_deployment_by_model_group_name( - model_group_name="claude-sonnet-4" - ) - router._update_kwargs_with_deployment( - deployment=deployment, kwargs=kwargs, function_name="generic_api_call" - ) + deployment = router.get_deployment_by_model_group_name(model_group_name="claude-sonnet-4") + router._update_kwargs_with_deployment(deployment=deployment, kwargs=kwargs, function_name="generic_api_call") assert "litellm_metadata" in kwargs model_info = kwargs["litellm_metadata"]["model_info"] @@ -4611,12 +4372,8 @@ def test_update_kwargs_with_deployment_model_info_in_metadata(): ) kwargs: dict = {} - deployment = router.get_deployment_by_model_group_name( - model_group_name="claude-sonnet-4" - ) - router._update_kwargs_with_deployment( - deployment=deployment, kwargs=kwargs, function_name=None - ) + deployment = router.get_deployment_by_model_group_name(model_group_name="claude-sonnet-4") + router._update_kwargs_with_deployment(deployment=deployment, kwargs=kwargs, function_name=None) assert "metadata" in kwargs model_info = kwargs["metadata"]["model_info"] @@ -4683,9 +4440,7 @@ async def test_acompletion_streaming_iterator_does_not_log_success_on_terminal_f StreamingChoices( finish_reason=None, index=0, - delta=Delta( - content="The Roman Empire began when", role="assistant" - ), + delta=Delta(content="The Roman Empire began when", role="assistant"), ) ], usage=Usage(prompt_tokens=17, completion_tokens=9, total_tokens=26), @@ -4995,23 +4750,17 @@ def test_multiregion_team_deployments_unique_model_names(): assert len(deployments) == 0 # With team_id: O(n) scan finds BOTH regional deployments - deployments = router._get_all_deployments( - model_name="claude-sonnet", team_id="metis-team" - ) + deployments = router._get_all_deployments(model_name="claude-sonnet", team_id="metis-team") assert len(deployments) == 2 deployment_names = {d["model_name"] for d in deployments} assert deployment_names == {"metis-claude-us-east-1", "metis-claude-us-west-2"} # Each deployment has a unique ID (critical for cooldown/retry to work) deployment_ids = {d["model_info"]["id"] for d in deployments} - assert ( - len(deployment_ids) == 2 - ), "Each deployment must have a unique ID for cooldown tracking" + assert len(deployment_ids) == 2, "Each deployment must have a unique ID for cooldown tracking" # Wrong team: returns nothing - deployments = router._get_all_deployments( - model_name="claude-sonnet", team_id="other-team" - ) + deployments = router._get_all_deployments(model_name="claude-sonnet", team_id="other-team") assert len(deployments) == 0 @@ -5056,12 +4805,8 @@ async def test_multiregion_team_failover_between_regions(): ) # Verify the router finds both deployments for the team - deployments = router._get_all_deployments( - model_name="claude-sonnet", team_id="metis-team" - ) - assert ( - len(deployments) == 2 - ), "Router must find both regional deployments by team_public_model_name" + deployments = router._get_all_deployments(model_name="claude-sonnet", team_id="metis-team") + assert len(deployments) == 2, "Router must find both regional deployments by team_public_model_name" # Make a normal request — should succeed from one of the regions response = await router.acompletion( @@ -5186,9 +4931,7 @@ def test_explicit_model_access_does_not_force_access_group_filtering(): }, ) - deployment_groups = [ - d.get("model_info", {}).get("access_groups") for d in deployments - ] + deployment_groups = [d.get("model_info", {}).get("access_groups") for d in deployments] assert ["AG1"] in deployment_groups assert ["AG2"] in deployment_groups @@ -5233,9 +4976,7 @@ def test_access_group_filter_empty_does_not_bypass_via_litellm_model_fallback( orig_groups = router.get_model_access_groups - def fake_get_model_access_groups( - model_name=None, model_access_group=None, team_id=None - ): + def fake_get_model_access_groups(model_name=None, model_access_group=None, team_id=None): if model_name == "gpt-5" and model_access_group is None: return {"AG1": ["gpt-5"], "AG2": ["gpt-5"]} return orig_groups( @@ -5310,9 +5051,7 @@ def test_access_group_block_does_not_silently_use_default_fallback_model( orig_groups = router.get_model_access_groups - def fake_get_model_access_groups( - model_name=None, model_access_group=None, team_id=None - ): + def fake_get_model_access_groups(model_name=None, model_access_group=None, team_id=None): if model_name == "gpt-5" and model_access_group is None: return {"AG1": ["gpt-5"], "AG2": ["gpt-5"]} return orig_groups( @@ -5379,9 +5118,7 @@ def test_access_group_block_via_litellm_model_branch_does_not_use_default_fallba orig_groups = router.get_model_access_groups - def fake_get_model_access_groups( - model_name=None, model_access_group=None, team_id=None - ): + def fake_get_model_access_groups(model_name=None, model_access_group=None, team_id=None): if model_name == "gpt-5" and model_access_group is None: return {"AG1": ["gpt-5"], "AG2": ["gpt-5"]} return orig_groups( @@ -5436,9 +5173,7 @@ def test_try_early_resolve_deployments_for_model_not_in_names(): ) assert ( - router_in_names._try_early_resolve_deployments_for_model_not_in_names( - model="gpt-5", request_team_id=None - ) + router_in_names._try_early_resolve_deployments_for_model_not_in_names(model="gpt-5", request_team_id=None) is None ) assert ( @@ -5460,10 +5195,8 @@ def test_try_early_resolve_deployments_for_model_not_in_names(): ] ) - pattern_result = ( - pattern_router._try_early_resolve_deployments_for_model_not_in_names( - model="openai/gpt-4o-mini", request_team_id=None - ) + pattern_result = pattern_router._try_early_resolve_deployments_for_model_not_in_names( + model="openai/gpt-4o-mini", request_team_id=None ) assert pattern_result is not None resolved_model, pattern_deployments = pattern_result @@ -5489,10 +5222,8 @@ def test_try_early_resolve_deployments_for_model_not_in_names(): }, } - default_result = ( - default_router._try_early_resolve_deployments_for_model_not_in_names( - model="brand-new-model", request_team_id=None - ) + default_result = default_router._try_early_resolve_deployments_for_model_not_in_names( + model="brand-new-model", request_team_id=None ) assert default_result is not None resolved_model, default_deployment = default_result @@ -5500,10 +5231,7 @@ def test_try_early_resolve_deployments_for_model_not_in_names(): assert isinstance(default_deployment, dict) assert default_deployment["litellm_params"]["model"] == "brand-new-model" # The original default_deployment must not be mutated. - assert ( - default_router.default_deployment["litellm_params"]["model"] - == "openai/will-be-overridden" - ) + assert default_router.default_deployment["litellm_params"]["model"] == "openai/will-be-overridden" def _router_with_two_deployments(blocked_flags): @@ -5551,10 +5279,7 @@ def _seed_unhealthy_states(router, unhealthy_ids, timestamp=None): ts = timestamp if timestamp is not None else time.time() router.health_state_cache.set_deployment_health_states( - { - uid: {"is_healthy": False, "timestamp": ts, "reason": "test_unhealthy"} - for uid in unhealthy_ids - } + {uid: {"is_healthy": False, "timestamp": ts, "reason": "test_unhealthy"} for uid in unhealthy_ids} ) @@ -5625,9 +5350,7 @@ async def test_async_get_fully_unhealthy_model_names_noop_with_allowed_fails_pol @pytest.mark.asyncio async def test_async_get_healthy_deployments_skips_blocked_deployment(): router = _router_with_two_deployments([True, False]) - healthy, all_dep = await router._async_get_healthy_deployments( - model="gpt-4o", parent_otel_span=None - ) + healthy, all_dep = await router._async_get_healthy_deployments(model="gpt-4o", parent_otel_span=None) healthy_ids = [d["model_info"]["id"] for d in healthy] assert "dep-0" not in healthy_ids assert "dep-1" in healthy_ids @@ -5636,9 +5359,7 @@ async def test_async_get_healthy_deployments_skips_blocked_deployment(): def test_get_healthy_deployments_sync_skips_blocked_deployment(): router = _router_with_two_deployments([False, True]) - healthy, all_dep = router._get_healthy_deployments( - model="gpt-4o", parent_otel_span=None - ) + healthy, all_dep = router._get_healthy_deployments(model="gpt-4o", parent_otel_span=None) healthy_ids = [d["model_info"]["id"] for d in healthy] assert "dep-0" in healthy_ids assert "dep-1" not in healthy_ids @@ -5655,9 +5376,7 @@ def test_filter_blocked_deployments_drops_blocked_keeps_unblocked(): @pytest.mark.asyncio async def test_public_async_get_healthy_deployments_skips_blocked_on_primary_path(): router = _router_with_two_deployments([True, False]) - deployments = await router.async_get_healthy_deployments( - model="gpt-4o", request_kwargs={} - ) + deployments = await router.async_get_healthy_deployments(model="gpt-4o", request_kwargs={}) assert isinstance(deployments, list) ids = [d["model_info"]["id"] for d in deployments] assert "dep-0" not in ids @@ -5699,9 +5418,7 @@ def _router_with_two_pass_through_deployments(blocked_flags): def test_get_available_deployment_for_pass_through_skips_blocked(): router = _router_with_two_pass_through_deployments([True, False]) - deployment = router.get_available_deployment_for_pass_through( - model="gpt-4o", request_kwargs={} - ) + deployment = router.get_available_deployment_for_pass_through(model="gpt-4o", request_kwargs={}) assert deployment["model_info"]["id"] == "pt-1" @@ -5710,9 +5427,7 @@ def test_get_available_deployment_for_pass_through_raises_when_dict_blocked(): router = _router_with_two_pass_through_deployments([True, True]) with pytest.raises(litellm.ServiceUnavailableError): - router.get_available_deployment_for_pass_through( - model="pt-0", request_kwargs={} - ) + router.get_available_deployment_for_pass_through(model="pt-0", request_kwargs={}) def test_initialize_deployment_for_pass_through_keeps_bedrock_iam_deployment(): @@ -5736,9 +5451,7 @@ def test_initialize_deployment_for_pass_through_keeps_bedrock_iam_deployment(): } ] ) - assert [m["model_info"]["id"] for m in router.get_model_list()] == [ - "bedrock-iam-pt" - ] + assert [m["model_info"]["id"] for m in router.get_model_list()] == ["bedrock-iam-pt"] def test_initialize_deployment_for_pass_through_sets_credentials_with_api_key(): @@ -5750,9 +5463,7 @@ def test_initialize_deployment_for_pass_through_sets_credentials_with_api_key(): router = _router_with_two_pass_through_deployments([False, False]) assert len(router.get_model_list()) == 2 assert ( - passthrough_endpoint_router.get_credentials( - custom_llm_provider="openai", region_name=None - ) + passthrough_endpoint_router.get_credentials(custom_llm_provider="openai", region_name=None) == "sk-fake-for-tests" ) @@ -5788,16 +5499,9 @@ def test_is_deployment_blocked_static_helper_reflects_blocked_flag(): # No model_info on deployment object → treated as not blocked assert litellm.Router._is_deployment_blocked(object()) is False missing_blocked = types.SimpleNamespace() + assert litellm.Router._is_deployment_blocked(types.SimpleNamespace(model_info=missing_blocked)) is False assert ( - litellm.Router._is_deployment_blocked( - types.SimpleNamespace(model_info=missing_blocked) - ) - is False - ) - assert ( - litellm.Router._is_deployment_blocked( - types.SimpleNamespace(model_info=types.SimpleNamespace(blocked=True)) - ) + litellm.Router._is_deployment_blocked(types.SimpleNamespace(model_info=types.SimpleNamespace(blocked=True))) is True ) @@ -5837,9 +5541,7 @@ class TestRouterRequestTimeoutPropagation: litellm.request_timeout = original_value litellm.request_timeout_explicitly_set = original_flag - def test_request_timeout_stored_independently_when_both_set( - self, explicit_request_timeout - ): + def test_request_timeout_stored_independently_when_both_set(self, explicit_request_timeout): router = self._make_router(timeout=330) assert router.timeout == 330 assert router.request_timeout == 300 @@ -5857,22 +5559,16 @@ class TestRouterRequestTimeoutPropagation: litellm.request_timeout = original_value litellm.request_timeout_explicitly_set = original_flag - def test_non_stream_prefers_request_timeout_over_router_timeout( - self, explicit_request_timeout - ): + def test_non_stream_prefers_request_timeout_over_router_timeout(self, explicit_request_timeout): router = self._make_router(timeout=330) assert router._get_non_stream_timeout(kwargs={}, data={}) == 300 - def test_stream_prefers_request_timeout_over_router_timeout( - self, explicit_request_timeout - ): + def test_stream_prefers_request_timeout_over_router_timeout(self, explicit_request_timeout): router = self._make_router(timeout=330) # stream=True resolves through _get_stream_timeout; request_timeout must win. assert router._get_timeout(kwargs={"stream": True}, data={}) == 300 - def test_explicit_stream_timeout_still_wins_over_request_timeout( - self, explicit_request_timeout - ): + def test_explicit_stream_timeout_still_wins_over_request_timeout(self, explicit_request_timeout): router = self._make_router(timeout=330, stream_timeout=45) assert router._get_stream_timeout(kwargs={}, data={}) == 45 @@ -5888,22 +5584,13 @@ class TestRouterRequestTimeoutPropagation: litellm.request_timeout = original_value litellm.request_timeout_explicitly_set = original_flag - def test_per_deployment_timeout_overrides_request_timeout( - self, explicit_request_timeout - ): + def test_per_deployment_timeout_overrides_request_timeout(self, explicit_request_timeout): router = self._make_router(timeout=330) assert router._get_non_stream_timeout(kwargs={}, data={"timeout": 120}) == 120 - def test_per_request_timeout_overrides_request_timeout( - self, explicit_request_timeout - ): + def test_per_request_timeout_overrides_request_timeout(self, explicit_request_timeout): router = self._make_router(timeout=330) - assert ( - router._get_non_stream_timeout( - kwargs={"timeout": 60}, data={"timeout": 120} - ) - == 60 - ) + assert router._get_non_stream_timeout(kwargs={"timeout": 60}, data={"timeout": 120}) == 60 class TestAdvisorSubCallCooldown: @@ -5936,9 +5623,7 @@ class TestAdvisorSubCallCooldown: ) def _cooled_down_ids(self, router): - active = router.cooldown_cache.get_active_cooldowns( - model_ids=["dep-1"], parent_otel_span=None - ) + active = router.cooldown_cache.get_active_cooldowns(model_ids=["dep-1"], parent_otel_span=None) return [entry[0] for entry in active] @pytest.mark.asyncio @@ -5947,12 +5632,7 @@ class TestAdvisorSubCallCooldown: router = self._router() now = datetime.now() - assert ( - router.deployment_callback_on_failure( - self._kwargs(self._auth_error()), None, now, now - ) - is True - ) + assert router.deployment_callback_on_failure(self._kwargs(self._auth_error()), None, now, now) is True assert "dep-1" in self._cooled_down_ids(router) def test_advisor_orchestration_failure_does_not_cool_down_deployment(self): @@ -5967,12 +5647,7 @@ class TestAdvisorSubCallCooldown: mark_advisor_orchestration_failure(exception) now = datetime.now() - assert ( - router.deployment_callback_on_failure( - self._kwargs(exception), None, now, now - ) - is False - ) + assert router.deployment_callback_on_failure(self._kwargs(exception), None, now, now) is False assert "dep-1" not in self._cooled_down_ids(router) @@ -6015,12 +5690,8 @@ def test_get_configured_token_limits_skips_wildcard_pattern_matching(): ] ) - with patch.object( - router.pattern_router, "route", side_effect=AssertionError("pattern route called") - ): - assert router.get_configured_token_limits( - "bedrock/anthropic.claude-3-5-sonnet-20240620-v1:0" - ) == (None, None) + with patch.object(router.pattern_router, "route", side_effect=AssertionError("pattern route called")): + assert router.get_configured_token_limits("bedrock/anthropic.claude-3-5-sonnet-20240620-v1:0") == (None, None) def test_get_configured_token_limits_treats_malformed_values_as_absent(): @@ -6088,13 +5759,16 @@ async def test_acreate_batch_request_bedrock_tags_override_deployment_tags(): mock_client = MagicMock() mock_client.post = AsyncMock(side_effect=lambda *args, **kwargs: fake_response()) - with patch.object( - CommonBatchFilesUtils, - "sign_aws_request", - return_value=({"Authorization": "signed"}, b"{}"), - ) as mock_sign, patch( - "litellm.llms.custom_httpx.llm_http_handler.get_async_httpx_client", - return_value=mock_client, + with ( + patch.object( + CommonBatchFilesUtils, + "sign_aws_request", + return_value=({"Authorization": "signed"}, b"{}"), + ) as mock_sign, + patch( + "litellm.llms.custom_httpx.llm_http_handler.get_async_httpx_client", + return_value=mock_client, + ), ): await router.acreate_batch( model="bedrock-batch-model", @@ -6134,7 +5808,7 @@ class TestPreRoutingStrategyRegistryLifecycle: return { "model": "auto_router/complexity_router", "complexity_router_config": { - "tiers": {"SIMPLE": "gpt-4o-mini", "MEDIUM": "gpt-4o", "COMPLEX": "gpt-4o"} + "tiers": {"SIMPLE": "gpt-4o-mini", "MEDIUM": "gpt-4o", "COMPLEX": "gpt-4o", "REASONING": "gpt-4o-mini"} }, "complexity_router_default_model": default_model, **({"tags": tags} if tags else {}), @@ -6331,7 +6005,12 @@ class TestPreRoutingStrategyRegistryLifecycle: params = { "model": "auto_router/complexity_router", "complexity_router_config": { - "tiers": {"SIMPLE": "gpt-4o-mini", "MEDIUM": "gpt-4o"}, + "tiers": { + "SIMPLE": "gpt-4o-mini", + "MEDIUM": "gpt-4o", + "COMPLEX": "gpt-4o-mini", + "REASONING": "gpt-4o-mini", + }, "adaptive": True, }, "complexity_router_default_model": "gpt-4o", @@ -6416,7 +6095,9 @@ class TestPreRoutingStrategyRegistryLifecycle: {"model_name": "gpt-4o-mini", "litellm_params": {"model": "openai/gpt-4o-mini"}}, { "model_name": "hybrid-router", - "litellm_params": cls._hybrid_router_params({"SIMPLE": "gpt-4o-mini", "MEDIUM": "gpt-4o"}), + "litellm_params": cls._hybrid_router_params( + {"SIMPLE": "gpt-4o-mini", "MEDIUM": "gpt-4o", "COMPLEX": "gpt-4o", "REASONING": "gpt-4o"} + ), "model_info": {"id": "router-1", "db_model": True}, }, ], @@ -6441,7 +6122,7 @@ class TestPreRoutingStrategyRegistryLifecycle: model_name="hybrid-router", litellm_params=LiteLLM_Params( **self._hybrid_router_params( - {"SIMPLE": "gpt-4o-mini", "MEDIUM": "gpt-4o", "COMPLEX": "gpt-4o"} + {"SIMPLE": "gpt-4o-mini", "MEDIUM": "gpt-4o", "COMPLEX": "gpt-4o", "REASONING": "gpt-4o"} ) ), model_info=ModelInfo(id="router-1", db_model=True), @@ -6538,12 +6219,30 @@ class TestPreRoutingStrategyRegistryLifecycle: cases = [ ({"model": "auto_router/adaptive_router", "adaptive_router_config": {}}, True), - (self._hybrid_router_params({"SIMPLE": "gpt-4o-mini"}), True), + ( + self._hybrid_router_params( + { + "SIMPLE": "gpt-4o-mini", + "MEDIUM": "gpt-4o-mini", + "COMPLEX": "gpt-4o-mini", + "REASONING": "gpt-4o-mini", + } + ), + True, + ), (self._complexity_router_params("gpt-4o"), False), ( { "model": "auto_router/complexity_router", - "complexity_router_config": {"tiers": {"SIMPLE": "gpt-4o-mini"}, "adaptive": False}, + "complexity_router_config": { + "tiers": { + "SIMPLE": "gpt-4o-mini", + "MEDIUM": "gpt-4o-mini", + "COMPLEX": "gpt-4o-mini", + "REASONING": "gpt-4o-mini", + }, + "adaptive": False, + }, "complexity_router_default_model": "gpt-4o", }, False, @@ -6551,9 +6250,7 @@ class TestPreRoutingStrategyRegistryLifecycle: ({"model": "openai/gpt-4o"}, False), ] for params, expected in cases: - actual = router._deployment_participates_in_adaptive_routing( - litellm_params=LiteLLM_Params(**params) - ) + actual = router._deployment_participates_in_adaptive_routing(litellm_params=LiteLLM_Params(**params)) assert actual is expected, params["model"]