From 3ea501430b5445d9a2afe17e572fa8cf6825afef Mon Sep 17 00:00:00 2001 From: ryan-crabbe-berri Date: Thu, 27 Aug 2026 14:08:51 -0700 Subject: [PATCH] fix(router): authorize config-level fallback targets against the calling key Router fallbacks configured in router_settings were attempted without re-checking whether the calling key could use the fallback model, so a key limited to one access group was served by any model listed as a fallback for something it could call. Auth only validated the requested model and fallbacks sent in the request body. Add a fallback_access_check predicate to Router, consulted before every cross-model-group fallback attempt; rejected targets are skipped and the primary's own error is raised when none remain. The proxy injects a check that runs the same key, team and project model access checks the requested model goes through. --- litellm/proxy/auth/fallback_model_access.py | 64 ++ litellm/proxy/proxy_server.py | 3 + litellm/router.py | 4 + .../router_utils/fallback_event_handlers.py | 21 + litellm/types/router.py | 16 +- .../proxy/auth/test_fallback_model_access.py | 57 ++ tests/test_litellm/proxy/test_proxy_server.py | 15 + .../test_fallback_event_handlers.py | 91 +- tests/test_litellm/test_router.py | 911 ++++++------------ 9 files changed, 570 insertions(+), 612 deletions(-) create mode 100644 litellm/proxy/auth/fallback_model_access.py create mode 100644 tests/test_litellm/proxy/auth/test_fallback_model_access.py diff --git a/litellm/proxy/auth/fallback_model_access.py b/litellm/proxy/auth/fallback_model_access.py new file mode 100644 index 00000000000..762e3694547 --- /dev/null +++ b/litellm/proxy/auth/fallback_model_access.py @@ -0,0 +1,64 @@ +""" +Authorize router fallback targets against the caller's key, team and project model access. + +`_enforce_key_and_fallback_model_access` only sees fallbacks the client sends in the request body. +Fallbacks configured on the router (`router_settings.fallbacks` and friends) are chosen after auth, +inside the router, so this predicate is injected into the router to re-run the same model access +checks for each fallback target before it is attempted. +""" + +from collections.abc import Mapping +from typing import Final + +from pydantic import BaseModel, ValidationError + +from litellm.proxy._types import ProxyException, UserAPIKeyAuth +from litellm.proxy.auth.auth_checks import can_key_call_resolved_model +from litellm.router import Router + + +class _RequestMetadata(BaseModel): + user_api_key_auth: UserAPIKeyAuth | None = None + + +async def is_model_authorized_for_token(*, model: str, valid_token: UserAPIKeyAuth, llm_router: Router) -> bool: + try: + await can_key_call_resolved_model( + model=model, + llm_model_list=None, + valid_token=valid_token, + llm_router=llm_router, + ) + except ProxyException: + return False + return True + + +def _token_in_metadata(metadata: object) -> UserAPIKeyAuth | None: + try: + return _RequestMetadata.model_validate(metadata).user_api_key_auth + except ValidationError: + return None + + +def _user_api_key_auth_from_request(request_kwargs: Mapping[str, object]) -> UserAPIKeyAuth | None: + return next( + ( + token + for field in ("metadata", "litellm_metadata") + if (token := _token_in_metadata(request_kwargs.get(field))) is not None + ), + None, + ) + + +async def router_fallback_access_check(*, model: str, request_kwargs: Mapping[str, object], llm_router: Router) -> bool: + """ + `FallbackAccessCheck` for the proxy's router: a fallback target is attempted only when the + key behind the request could have requested it directly. Requests that carry no key (for + example internal health checks) are not restricted. + """ + valid_token: Final = _user_api_key_auth_from_request(request_kwargs) + if valid_token is None: + return True + return await is_model_authorized_for_token(model=model, valid_token=valid_token, llm_router=llm_router) diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index 1d33c9ce393..81c52f3bd1b 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -290,6 +290,7 @@ from litellm.proxy.auth.auth_utils import ( is_request_body_safe, warn_once_if_custom_auth_skips_common_checks, ) +from litellm.proxy.auth.fallback_model_access import router_fallback_access_check from litellm.proxy.auth.handle_jwt import JWTHandler from litellm.proxy.auth.litellm_license import LicenseCheck from litellm.proxy.auth.model_checks import ( @@ -5580,6 +5581,7 @@ class ProxyConfig: async_only_mode=True # only init async clients ), ignore_invalid_deployments=True, # don't raise an error if a deployment is invalid + fallback_access_check=router_fallback_access_check, ) if redis_usage_cache is not None and router.cache.redis_cache is None: @@ -6039,6 +6041,7 @@ class ProxyConfig: ), search_tools=search_tools, ignore_invalid_deployments=True, + fallback_access_check=router_fallback_access_check, ) verbose_proxy_logger.debug("updated llm_router: %s", llm_router) else: diff --git a/litellm/router.py b/litellm/router.py index 021dafa9791..97a47f5384a 100644 --- a/litellm/router.py +++ b/litellm/router.py @@ -196,6 +196,7 @@ from litellm.types.router import ( CustomRoutingStrategyBase, Deployment, DeploymentTypedDict, + FallbackAccessCheck, GuardrailTypedDict, LiteLLM_Params, MockRouterTestingParams, @@ -604,6 +605,7 @@ class Router: health_check_ignore_transient_errors: bool = False, background_health_check_model_groups: Sequence[str] | None = None, enable_weighted_failover: bool = False, + fallback_access_check: FallbackAccessCheck | None = None, ) -> None: """ Initialize the Router class with the given parameters for caching, reliability, and routing strategy. @@ -640,6 +642,7 @@ class Router: deployment_affinity_ttl_seconds (int): TTL for user-key -> deployment affinity mapping. Defaults to 3600. ignore_invalid_deployments (bool): Ignores invalid deployments, and continues with other deployments. Default is to raise an error. enable_weighted_failover (bool): When True and the routing strategy is "simple-shuffle", a retryable failure on one deployment causes the request to re-pick (weighted) across the other deployments in the same model group before any cross-group fallback runs. Bounded by `max_fallbacks`. Async-only: currently honored by `router.acompletion()` and other async entrypoints. The sync `router.completion()` path falls back to the regular fallback flow. Defaults to False. + fallback_access_check (Optional[FallbackAccessCheck]): Awaited before each cross-model-group fallback attempt on the async path; a fallback target it rejects is skipped. Defaults to None (every configured fallback is attempted). Returns: Router: An instance of the litellm.Router class. @@ -679,6 +682,7 @@ class Router: self.set_verbose = set_verbose self.ignore_invalid_deployments = ignore_invalid_deployments + self.fallback_access_check: Final = fallback_access_check self.debug_level = debug_level self.enable_pre_call_checks = enable_pre_call_checks self.enable_tag_filtering = enable_tag_filtering diff --git a/litellm/router_utils/fallback_event_handlers.py b/litellm/router_utils/fallback_event_handlers.py index acdc7df5bd1..5ebbb1db0d5 100644 --- a/litellm/router_utils/fallback_event_handlers.py +++ b/litellm/router_utils/fallback_event_handlers.py @@ -263,6 +263,25 @@ def _get_fallback_target_model_group(fallback_entry: str | Mapping[str, object]) return target if isinstance(target, str) else None +async def _is_fallback_target_authorized( + litellm_router: LitellmRouter, + fallback_entry: str | Mapping[str, object], + original_model_group: str, + kwargs: Mapping[str, object], +) -> bool: + access_check: Final = litellm_router.fallback_access_check + target: Final = _get_fallback_target_model_group(fallback_entry) + if access_check is None or target is None or target == original_model_group: + return True + if await access_check(model=target, request_kwargs=kwargs, llm_router=litellm_router): + return True + verbose_router_logger.info( + "Skipping fallback to model_group = %s: caller is not authorized to call it", + mask_sensitive_structure(fallback_entry), + ) + return False + + def references_provider_scoped_resource(kwargs: Mapping[str, object]) -> bool: """ True when the request names a file that only exists under one provider's credentials. @@ -357,6 +376,8 @@ async def run_async_fallback( original_model_group, ) continue + if not await _is_fallback_target_authorized(litellm_router, mg, original_model_group, kwargs): + continue attempt_key = fallback_attempt_key(mg) if attempt_key is not None: if attempt_key in attempted: diff --git a/litellm/types/router.py b/litellm/types/router.py index a3335be2b2b..97bd93f3f47 100644 --- a/litellm/types/router.py +++ b/litellm/types/router.py @@ -6,7 +6,7 @@ import datetime import enum from collections.abc import Mapping from dataclasses import dataclass -from typing import Any, ClassVar, Final, Generic, Literal, TypeVar, get_type_hints +from typing import TYPE_CHECKING, Any, ClassVar, Final, Generic, Literal, TypeVar, get_type_hints import httpx from pydantic import BaseModel, ConfigDict, Field, field_validator, model_validator @@ -14,6 +14,9 @@ from typing_extensions import Protocol, ReadOnly, Required, TypedDict, runtime_c from litellm._uuid import uuid +if TYPE_CHECKING: + from litellm.router import Router + from .completion import CompletionRequest from .embedding import EmbeddingRequest from .llms.openai import OpenAIFileObject @@ -845,6 +848,17 @@ class GenericBudgetWindowDetails(BaseModel): ttl_seconds: int +class FallbackAccessCheck(Protocol): + """ + Decides whether the caller behind `request_kwargs` may be served by fallback `model`. + + The router runs it before every cross-model-group fallback attempt and skips targets it + rejects, so a fallback can never reach a model the caller could not have requested directly. + """ + + async def __call__(self, *, model: str, request_kwargs: Mapping[str, object], llm_router: "Router") -> bool: ... + + OptionalPreCallChecks = list[ Literal[ "prompt_caching", diff --git a/tests/test_litellm/proxy/auth/test_fallback_model_access.py b/tests/test_litellm/proxy/auth/test_fallback_model_access.py new file mode 100644 index 00000000000..ed819eed9ee --- /dev/null +++ b/tests/test_litellm/proxy/auth/test_fallback_model_access.py @@ -0,0 +1,57 @@ +import pytest + +from litellm import Router +from litellm.proxy._types import UserAPIKeyAuth +from litellm.proxy.auth.fallback_model_access import ( + is_model_authorized_for_token, + router_fallback_access_check, +) + + +def _router() -> Router: + return Router( + model_list=[ + { + "model_name": "open-model", + "litellm_params": {"model": "openai/open", "api_key": "k"}, + "model_info": {"access_groups": ["open-group"]}, + }, + { + "model_name": "secret-model", + "litellm_params": {"model": "openai/secret", "api_key": "k"}, + "model_info": {"access_groups": ["secret-group"]}, + }, + ] + ) + + +def _key_limited_to(access_group: str) -> UserAPIKeyAuth: + return UserAPIKeyAuth(api_key="hashed", models=[access_group]) + + +@pytest.mark.asyncio +async def test_is_model_authorized_for_token_follows_the_key_access_groups(): + router = _router() + token = _key_limited_to("open-group") + + assert await is_model_authorized_for_token(model="open-model", valid_token=token, llm_router=router) is True + assert await is_model_authorized_for_token(model="secret-model", valid_token=token, llm_router=router) is False + + +@pytest.mark.asyncio +@pytest.mark.parametrize("metadata_field", ["metadata", "litellm_metadata"]) +async def test_router_fallback_access_check_authorizes_the_key_carried_in_request_metadata(metadata_field: str): + router = _router() + request_kwargs = {metadata_field: {"user_api_key_auth": _key_limited_to("open-group")}} + + assert await router_fallback_access_check(model="open-model", request_kwargs=request_kwargs, llm_router=router) + assert not await router_fallback_access_check( + model="secret-model", request_kwargs=request_kwargs, llm_router=router + ) + + +@pytest.mark.asyncio +async def test_router_fallback_access_check_does_not_restrict_requests_without_a_key(): + assert await router_fallback_access_check( + model="secret-model", request_kwargs={"metadata": {}}, llm_router=_router() + ) diff --git a/tests/test_litellm/proxy/test_proxy_server.py b/tests/test_litellm/proxy/test_proxy_server.py index f51648faf80..965b40188ae 100644 --- a/tests/test_litellm/proxy/test_proxy_server.py +++ b/tests/test_litellm/proxy/test_proxy_server.py @@ -11712,3 +11712,18 @@ async def test_authoritative_floor_spend_keeps_a_reset_marker_written_during_the assert real_spend_counter_cache.in_memory_cache.get_cache(key=marker_key) == 0.0, ( "the in-flight DB read clobbered the post-reset floor marker with the stale pre-reset value" ) + + +@pytest.mark.asyncio +async def test_load_config_router_authorizes_fallback_targets_against_the_calling_key(tmp_path): + from litellm.proxy.auth.fallback_model_access import router_fallback_access_check + from litellm.proxy.proxy_server import ProxyConfig + + config_file = tmp_path / "config.yaml" + config_file.write_text( + yaml.dump({"model_list": [{"model_name": "m", "litellm_params": {"model": "openai/m", "api_key": "k"}}]}) + ) + + router, _, _ = await ProxyConfig().load_config(router=None, config_file_path=str(config_file)) + + assert router.fallback_access_check is router_fallback_access_check diff --git a/tests/test_litellm/router_utils/test_fallback_event_handlers.py b/tests/test_litellm/router_utils/test_fallback_event_handlers.py index 3e02838b88f..e55d2cd796e 100644 --- a/tests/test_litellm/router_utils/test_fallback_event_handlers.py +++ b/tests/test_litellm/router_utils/test_fallback_event_handlers.py @@ -22,6 +22,8 @@ class StreamingWrapper: class FakeRouter: + fallback_access_check = None + def log_retry(self, kwargs, e): return kwargs @@ -30,6 +32,8 @@ class FakeRouter: class AlwaysFailRouter: + fallback_access_check = None + def log_retry(self, kwargs, e): return kwargs @@ -92,6 +96,8 @@ async def test_run_async_fallback_raises_when_all_fallbacks_fail(): class RecordingRouter: + fallback_access_check = None + def __init__(self): self.received_kwargs = None @@ -151,6 +157,8 @@ async def test_run_async_fallback_skips_original_model_group(): class AttemptRecordingRouter: + fallback_access_check = None + def __init__(self): self.attempted_model_groups = [] self.received_kwargs = None @@ -339,7 +347,84 @@ async def test_run_async_fallback_records_batch_model_group_outside_provider_met assert router.received_kwargs["litellm_metadata"]["model_group"] == "openai-group" +class AccessCheckedRouter(AttemptRecordingRouter): + def __init__(self, allowed_models: frozenset[str]): + super().__init__() + self.allowed_models = allowed_models + self.access_checks = [] + + async def fallback_access_check(self, *, model, request_kwargs, llm_router): + self.access_checks.append((model, request_kwargs["metadata"]["user_api_key"], llm_router is self)) + return model in self.allowed_models + + +@pytest.mark.asyncio +async def test_run_async_fallback_skips_targets_the_access_check_rejects(): + router = AccessCheckedRouter(allowed_models=frozenset({"allowed-model"})) + + await run_async_fallback( + litellm_router=router, + fallback_model_group=[ + {"model": "secret-model", "messages": [{"role": "user", "content": "hi"}]}, + "allowed-model", + ], + original_model_group="primary-model", + original_exception=RuntimeError("primary failed"), + max_fallbacks=3, + fallback_depth=0, + model="primary-model", + metadata={"user_api_key": "hashed"}, + ) + + assert router.attempted_model_groups == ["allowed-model"] + assert router.access_checks == [ + ("secret-model", "hashed", True), + ("allowed-model", "hashed", True), + ] + + +@pytest.mark.asyncio +async def test_run_async_fallback_raises_original_error_when_no_target_is_authorized(): + router = AccessCheckedRouter(allowed_models=frozenset()) + + with pytest.raises(RuntimeError, match="primary failed"): + await run_async_fallback( + litellm_router=router, + fallback_model_group=["secret-model", "other-secret-model"], + original_model_group="primary-model", + original_exception=RuntimeError("primary failed"), + max_fallbacks=3, + fallback_depth=0, + model="primary-model", + metadata={"user_api_key": "hashed"}, + ) + + assert router.attempted_model_groups == [] + assert [model for model, _, _ in router.access_checks] == ["secret-model", "other-secret-model"] + + +@pytest.mark.asyncio +async def test_run_async_fallback_does_not_consult_access_check_for_same_model_group_retries(): + router = AccessCheckedRouter(allowed_models=frozenset()) + + await run_async_fallback( + litellm_router=router, + fallback_model_group=[{"model": "primary-model", "_target_order": 2}], + original_model_group="primary-model", + original_exception=RuntimeError("first order level failed"), + max_fallbacks=3, + fallback_depth=0, + model="primary-model", + metadata={"user_api_key": "hashed"}, + ) + + assert router.attempted_model_groups == ["primary-model"] + assert router.access_checks == [] + + class RecordingFailRouter: + fallback_access_check = None + def __init__(self): self.attempted_models = [] @@ -488,9 +573,7 @@ async def test_run_async_fallback_keeps_a_request_override_distinct_from_the_bar with pytest.raises(RuntimeError, match="fallback model also failed"): await run_async_fallback( litellm_router=router, - fallback_model_group=[ - {"model": "already-attempted", "messages": [{"role": "user", "content": "shorter"}]} - ], + fallback_model_group=[{"model": "already-attempted", "messages": [{"role": "user", "content": "shorter"}]}], original_model_group="primary-model", original_exception=RuntimeError("original failed"), max_fallbacks=3, @@ -774,6 +857,8 @@ class TestTriggerCooldownForFailedDeployment: class TestRunAsyncFallbackTriggersCooldown: class RouterWithLoggingKwarg: + fallback_access_check = None + def __init__(self): self.cooldown_time = 60.0 diff --git a/tests/test_litellm/test_router.py b/tests/test_litellm/test_router.py index 8716e6d6b25..90b2f593180 100644 --- a/tests/test_litellm/test_router.py +++ b/tests/test_litellm/test_router.py @@ -10,7 +10,6 @@ import httpx import pytest - import litellm from litellm import Router from litellm.exceptions import MidStreamFallbackError @@ -127,31 +126,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 @@ -182,13 +168,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( @@ -229,10 +211,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, @@ -252,9 +231,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 @@ -301,18 +278,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, @@ -323,9 +292,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}, @@ -878,9 +845,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", @@ -901,9 +866,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=[ { @@ -964,7 +927,7 @@ async def test_arouter_filter_team_based_models(): assert result is not None # FAILS - with pytest.raises(Exception, match='No deployments available for selected model, Try again in') as e: + with pytest.raises(Exception, match="No deployments available for selected model, Try again in") as e: result = await router.acompletion( model="gpt-3.5-turbo", messages=[{"role": "user", "content": "Hello, world!"}], @@ -1048,9 +1011,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( @@ -1058,9 +1019,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( @@ -1076,30 +1037,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 @@ -1108,9 +1057,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(): @@ -1160,9 +1107,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: @@ -1236,7 +1181,7 @@ def test_add_invalid_provider_to_router(): ], ) - with pytest.raises(Exception, match='Unsupported provider - vertex_ai_eu') as e: + with pytest.raises(Exception, match="Unsupported provider - vertex_ai_eu") as e: router.add_deployment( Deployment( model_name="vertex_ai/*", @@ -1288,15 +1233,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, @@ -1331,7 +1270,7 @@ async def test_router_ageneric_api_call_with_fallbacks_helper(): with patch.object(router, "async_get_available_deployment") as mock_get_deployment: mock_get_deployment.side_effect = Exception("No deployment available") - with pytest.raises(Exception, match='No deployment available') as exc_info: + with pytest.raises(Exception, match="No deployment available") as exc_info: await router._ageneric_api_call_with_fallbacks_helper( model="gpt-3.5-turbo", original_generic_function=mock_generic_function, @@ -1359,15 +1298,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, @@ -1396,16 +1329,10 @@ 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 pytest.raises(Exception, match='Mock failure') as exc_info: + 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, match="Mock failure") as exc_info: await router._ageneric_api_call_with_fallbacks_helper( model="gpt-3.5-turbo", original_generic_function=mock_failing_function, @@ -1473,9 +1400,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(): @@ -1496,14 +1423,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"] @@ -1598,9 +1521,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 @@ -1628,9 +1549,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 @@ -1700,9 +1619,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", @@ -1748,9 +1665,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", @@ -2010,6 +1925,7 @@ async def test_acompletion_streaming_iterator(): # Collect streamed chunks — the first chunk succeeds, then the error re-raises collected_chunks = [] + async def _drain(): async for chunk in result: collected_chunks.append(chunk) @@ -2703,11 +2619,7 @@ def _make_responses_iterator( BaseResponsesAPIStreamingIterator, ) - base = ( - LiteLLMCompletionStreamingIterator - if bridge - else BaseResponsesAPIStreamingIterator - ) + base = LiteLLMCompletionStreamingIterator if bridge else BaseResponsesAPIStreamingIterator class _Iter(base): def __init__(self): @@ -2787,9 +2699,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( @@ -2872,9 +2782,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 @@ -2992,9 +2902,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, @@ -3003,9 +2911,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, @@ -3306,9 +3212,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"}}, @@ -3336,14 +3240,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"}}, @@ -3370,14 +3270,10 @@ def test_pre_call_checks_uses_precounted_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 1 - ) + monkeypatch.setattr(litellm, "token_counter", lambda *a, **k: calls.append(1) or 1) deployments = [ {"litellm_params": {"model": "gpt-3.5-turbo"}, "model_info": {"id": "d1"}}, @@ -3404,9 +3300,7 @@ async def test_async_get_healthy_deployments_counts_tokens_off_the_event_loop(mo ], enable_pre_call_checks=True, ) - monkeypatch.setattr( - router, "get_router_model_info", lambda **kwargs: {"max_input_tokens": 1_000_000} - ) + monkeypatch.setattr(router, "get_router_model_info", lambda **kwargs: {"max_input_tokens": 1_000_000}) counting_threads = [] monkeypatch.setattr( @@ -3502,14 +3396,10 @@ def test_pre_call_checks_does_not_recount_inline_after_an_off_loop_failure(monke ], 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"}}, @@ -3537,9 +3427,7 @@ async def test_async_get_healthy_deployments_never_recounts_on_the_loop(monkeypa ], 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}) counting_threads = [] @@ -3574,9 +3462,7 @@ async def test_acount_pre_call_check_tokens_leaves_the_event_loop_free(monkeypat ], 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}) deployments = [ {"litellm_params": {"model": "gpt-3.5-turbo"}, "model_info": {"id": "d1"}}, @@ -3612,9 +3498,7 @@ async def test_acount_pre_call_check_tokens_skips_without_max_input_tokens(monke 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) count = await router._acount_pre_call_check_tokens( model="m", @@ -3642,9 +3526,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"}}, @@ -3669,9 +3551,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"}}, @@ -3713,9 +3593,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", @@ -3748,7 +3626,7 @@ def test_count_pre_call_check_tokens_across_api_surfaces(): assert string_input_tokens > 0 assert list_input_tokens > 0 - with pytest.raises(ValueError, match='Either messages or input must be provided to count tokens'): + with pytest.raises(ValueError, match="Either messages or input must be provided to count tokens"): router._count_pre_call_check_tokens(messages=None, input=None) @@ -3763,9 +3641,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 @@ -3855,9 +3731,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: { @@ -3865,15 +3739,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 @@ -3884,26 +3754,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 = { @@ -3922,9 +3784,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 @@ -3944,9 +3804,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 @@ -3979,9 +3837,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 @@ -3990,12 +3846,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 @@ -4018,9 +3870,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 @@ -4056,15 +3906,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 @@ -4111,18 +3957,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 @@ -4132,29 +3974,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) @@ -4191,10 +4021,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 = { @@ -4206,10 +4035,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 = { @@ -4221,9 +4049,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 = { @@ -4234,9 +4062,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(): @@ -4288,14 +4116,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 @@ -4387,15 +4211,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" @@ -4422,9 +4241,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" @@ -4464,9 +4281,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" @@ -4502,9 +4317,7 @@ def test_get_deployment_credentials_with_provider_bedrock_batch_fields(): ], ) - credentials = router.get_deployment_credentials_with_provider( - model_id="bedrock-batch-model" - ) + credentials = router.get_deployment_credentials_with_provider(model_id="bedrock-batch-model") assert credentials is not None assert credentials["custom_llm_provider"] == "bedrock" @@ -4547,9 +4360,7 @@ def test_get_deployment_credentials_with_provider_preserves_aws_auth_params(): ], ) - credentials = router.get_deployment_credentials_with_provider( - model_id="bedrock-batch-model" - ) + credentials = router.get_deployment_credentials_with_provider(model_id="bedrock-batch-model") assert credentials is not None for key, value in aws_auth_params.items(): @@ -4584,15 +4395,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" @@ -4633,15 +4440,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" @@ -4668,16 +4471,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(): @@ -4717,9 +4512,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" @@ -4729,12 +4522,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(): @@ -4766,9 +4554,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" @@ -4780,21 +4566,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(): @@ -4833,22 +4609,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(): @@ -5037,9 +4804,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", @@ -5118,12 +4883,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(): @@ -5148,9 +4909,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 @@ -5179,9 +4938,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 @@ -5208,9 +4965,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 @@ -5248,9 +5003,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 @@ -5281,9 +5034,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"}] @@ -5312,9 +5063,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) @@ -5416,12 +5165,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"] @@ -5453,12 +5198,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"] @@ -5571,6 +5312,7 @@ async def test_acompletion_streaming_iterator_does_not_log_success_on_terminal_f initial_kwargs=dict(initial_kwargs), ) collected = [] + async def _drain(): async for chunk in result: collected.append(chunk) @@ -5597,6 +5339,7 @@ async def test_acompletion_streaming_iterator_does_not_log_success_on_terminal_f initial_kwargs=dict(initial_kwargs), ) collected = [] + async def _drain(): async for chunk in result: collected.append(chunk) @@ -5813,23 +5556,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 @@ -5874,12 +5611,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( @@ -6004,9 +5737,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 @@ -6051,9 +5782,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( @@ -6128,9 +5857,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( @@ -6197,9 +5924,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( @@ -6254,9 +5979,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 ( @@ -6278,10 +6001,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 @@ -6307,10 +6028,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 @@ -6318,10 +6037,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): @@ -6369,10 +6085,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} ) @@ -6443,9 +6156,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 @@ -6454,9 +6165,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 @@ -6473,9 +6182,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 @@ -6517,9 +6224,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" @@ -6528,9 +6233,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(): @@ -6554,9 +6257,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_pass_through_deployment_api_key_resolves_via_get_credentials(): @@ -6567,12 +6268,7 @@ def test_pass_through_deployment_api_key_resolves_via_get_credentials(): router = _router_with_two_pass_through_deployments([False, False]) passthrough_router = PassthroughEndpointRouter(llm_router_getter=lambda: router) assert len(router.get_model_list()) == 2 - assert ( - passthrough_router.get_credentials( - custom_llm_provider="openai", region_name=None - ) - == "sk-fake-for-tests" - ) + assert passthrough_router.get_credentials(custom_llm_provider="openai", region_name=None) == "sk-fake-for-tests" def test_get_deployment_credentials_returns_none_for_blocked_deployment(): @@ -6606,16 +6302,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 ) @@ -6655,9 +6344,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 @@ -6675,22 +6362,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 @@ -6706,22 +6387,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 # --------------------------------------------------------------------------- @@ -7030,9 +6702,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 @@ -7041,12 +6711,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): @@ -7061,12 +6726,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) @@ -7123,13 +6783,13 @@ def test_stream_chunks_have_generated_content_detects_text_and_non_text(): audio_chunk = _chunk(audio_delta) assert _stream_chunks_have_generated_content([audio_chunk]) is True - images_delta = Delta(images=[{"image_url": {"url": "https://example.com/img.png"}, "index": 0, "type": "image_url"}]) + images_delta = Delta( + images=[{"image_url": {"url": "https://example.com/img.png"}, "index": 0, "type": "image_url"}] + ) images_chunk = _chunk(images_delta) assert _stream_chunks_have_generated_content([images_chunk]) is True - annotations_delta = Delta( - annotations=[{"type": "url_citation", "url_citation": {"url": "https://example.com"}}] - ) + annotations_delta = Delta(annotations=[{"type": "url_citation", "url_citation": {"url": "https://example.com"}}]) annotations_chunk = _chunk(annotations_delta) assert _stream_chunks_have_generated_content([annotations_chunk]) is True @@ -7173,12 +6833,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(): @@ -7425,13 +7081,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", @@ -7470,9 +7129,7 @@ class TestPreRoutingStrategyRegistryLifecycle: def _complexity_router_params(default_model: str, tags=None) -> dict: return { "model": "auto_router/complexity_router", - "complexity_router_config": { - "tiers": {"SIMPLE": "gpt-4o-mini", "MEDIUM": "gpt-4o", "COMPLEX": "gpt-4o"} - }, + "complexity_router_config": {"tiers": {"SIMPLE": "gpt-4o-mini", "MEDIUM": "gpt-4o", "COMPLEX": "gpt-4o"}}, "complexity_router_default_model": default_model, **({"tags": tags} if tags else {}), } @@ -7777,9 +7434,7 @@ class TestPreRoutingStrategyRegistryLifecycle: deployment=Deployment( model_name="hybrid-router", litellm_params=LiteLLM_Params( - **self._hybrid_router_params( - {"SIMPLE": "gpt-4o-mini", "MEDIUM": "gpt-4o", "COMPLEX": "gpt-4o"} - ) + **self._hybrid_router_params({"SIMPLE": "gpt-4o-mini", "MEDIUM": "gpt-4o", "COMPLEX": "gpt-4o"}) ), model_info=ModelInfo(id="router-1", db_model=True), ) @@ -7888,9 +7543,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"] @@ -8077,22 +7730,16 @@ class TestUpsertDeploymentRollback: router.delete_deployment(id="prod-1") assert router.has_model_id("prod-1") is False - router._restore_deployment_after_failed_upsert( - previous_deployment=previous, model_id="prod-1" - ) + router._restore_deployment_after_failed_upsert(previous_deployment=previous, model_id="prod-1") restored = router.get_deployment(model_id="prod-1") assert restored is not None assert restored.litellm_params.model == "gpt-4o" - router._restore_deployment_after_failed_upsert( - previous_deployment=previous, model_id="prod-1" - ) + router._restore_deployment_after_failed_upsert(previous_deployment=previous, model_id="prod-1") assert len(router.model_list) == 1 - router._restore_deployment_after_failed_upsert( - previous_deployment=None, model_id="prod-1" - ) + router._restore_deployment_after_failed_upsert(previous_deployment=None, model_id="prod-1") assert len(router.model_list) == 1 @@ -8401,18 +8048,14 @@ class TestAutoRouterSharedModelNameConnectionParams: return httpx.Response( status_code=200, json={ - "candidates": [ - {"content": {"parts": [{"text": "Paris"}], "role": "model"}, "finishReason": "STOP"} - ], + "candidates": [{"content": {"parts": [{"text": "Paris"}], "role": "model"}, "finishReason": "STOP"}], "usageMetadata": {"promptTokenCount": 5, "candidatesTokenCount": 1, "totalTokenCount": 6}, "modelVersion": "gemini-3.6-flash", }, request=httpx.Request("POST", "https://generativelanguage.googleapis.com"), ) - @pytest.mark.parametrize( - "plain_entry_first", [True, False], ids=["plain_entry_first", "marker_entry_first"] - ) + @pytest.mark.parametrize("plain_entry_first", [True, False], ids=["plain_entry_first", "marker_entry_first"]) async def test_routed_tier_call_goes_out_on_its_own_endpoint_and_credentials(self, plain_entry_first): """The outbound provider request for the routed tier hits the tier's own Gemini host with the tier's own key, never the plain sibling's api_base or api_key.""" @@ -8533,9 +8176,7 @@ async def _drive_cyclic_fallback(router, capture, recorder=None, **request_kwarg litellm.callbacks.append(recorder) try: with pytest.raises(litellm.InternalServerError): - await router.acompletion( - model="group-a", messages=[{"role": "user", "content": "hi"}], **request_kwargs - ) + await router.acompletion(model="group-a", messages=[{"role": "user", "content": "hi"}], **request_kwargs) finally: router_logger.removeHandler(capture) router_logger.setLevel(previous_level) @@ -8572,9 +8213,9 @@ async def test_retry_breadcrumbs_do_not_carry_the_walk_state(): await _drive_cyclic_fallback(router, capture) assert router.previous_models, "no retry breadcrumbs were recorded" - assert any( - "fallback_depth" in breadcrumb for breadcrumb in router.previous_models - ), "no breadcrumb carried router walk state, so this test cannot see the leak" + assert any("fallback_depth" in breadcrumb for breadcrumb in router.previous_models), ( + "no breadcrumb carried router walk state, so this test cannot see the leak" + ) for breadcrumb in router.previous_models: assert "attempted_targets" not in breadcrumb @@ -8618,7 +8259,9 @@ async def test_retry_breadcrumbs_never_carry_a_forwarded_credential(container_ke assert router.previous_models, "no retry breadcrumbs were recorded" dumped = json.dumps(router.previous_models, default=str) - assert container_key in dumped, "the credential-bearing kwarg never reached the breadcrumb, so this test cannot see the leak" + assert container_key in dumped, ( + "the credential-bearing kwarg never reached the breadcrumb, so this test cannot see the leak" + ) assert _BREADCRUMB_CREDENTIAL_CANARY not in dumped @@ -8646,9 +8289,7 @@ async def test_fallback_failure_detail_from_upstream_is_bounded(): await _drive_cyclic_fallback( _cyclic_fallback_router(), capture, - mock_response=litellm.InternalServerError( - message=huge_message, llm_provider="openai", model="group-a" - ), + mock_response=litellm.InternalServerError(message=huge_message, llm_provider="openai", model="group-a"), ) assert capture.messages, "the fallback failure path did not log at ERROR" @@ -8714,9 +8355,7 @@ def test_ensure_deployment_affinity_callback_is_idempotent(): try: router._ensure_deployment_affinity_callback() router._ensure_deployment_affinity_callback() - affinity_callbacks = [ - cb for cb in router.optional_callbacks or [] if isinstance(cb, DeploymentAffinityCheck) - ] + affinity_callbacks = [cb for cb in router.optional_callbacks or [] if isinstance(cb, DeploymentAffinityCheck)] assert len(affinity_callbacks) == 1 finally: for cb in router.optional_callbacks or []: @@ -8858,9 +8497,7 @@ class TestModelGroupAliasReachesPreRoutingStrategies: router = self._router("auto_routers") metadata: dict = {} - response = await router.acompletion( - model="smart-alias", messages=self._messages(), metadata=metadata - ) + response = await router.acompletion(model="smart-alias", messages=self._messages(), metadata=metadata) assert response.choices[0].message.content == "routed by the tier" assert metadata["model_group"] == "smart-alias" @@ -8901,17 +8538,14 @@ class TestAzureBaseModelFallbackLogging: def test_map_known_deployment_name_resolves_without_error_log(self): router = self._router_with_azure_deployment("azure/gpt-4o") - with patch( - "litellm.router.verbose_router_logger.error" - ) as mock_error: + with patch("litellm.router.verbose_router_logger.error") as mock_error: model_info = router.get_router_model_info( deployment=None, received_model_name="my-group", id="azure-base-model-test-id" ) - assert not any( - "Could not identify azure model" in str(call) - for call in mock_error.call_args_list - ), f"unexpected error log: {mock_error.call_args_list}" + assert not any("Could not identify azure model" in str(call) for call in mock_error.call_args_list), ( + f"unexpected error log: {mock_error.call_args_list}" + ) # the fallback resolution must actually surface the map values assert model_info["max_input_tokens"] == litellm.model_cost["azure/gpt-4o"]["max_input_tokens"] assert model_info["input_cost_per_token"] == litellm.model_cost["azure/gpt-4o"]["input_cost_per_token"] @@ -8919,17 +8553,14 @@ class TestAzureBaseModelFallbackLogging: def test_unmappable_deployment_name_still_logs_error(self): router = self._router_with_azure_deployment("azure/my-custom-deployment-name") - with patch( - "litellm.router.verbose_router_logger.error" - ) as mock_error: + with patch("litellm.router.verbose_router_logger.error") as mock_error: model_info = router.get_router_model_info( deployment=None, received_model_name="my-group", id="azure-base-model-test-id" ) - assert any( - "Could not identify azure model" in str(call) - for call in mock_error.call_args_list - ), "expected the error log for an unmappable azure deployment name" + assert any("Could not identify azure model" in str(call) for call in mock_error.call_args_list), ( + "expected the error log for an unmappable azure deployment name" + ) # unmappable names resolve to a zeroed stub — unchanged behavior assert model_info.get("max_input_tokens") is None @@ -8956,6 +8587,7 @@ class TestAzureBaseModelFallbackLogging: ) assert model_info["max_input_tokens"] == litellm.model_cost["azure/gpt-4o-mini"]["max_input_tokens"] + def test_model_group_info_intersects_supported_reasoning_efforts(): router = litellm.Router( model_list=[ @@ -9259,6 +8891,7 @@ class TestAddDeploymentApiBaseProviderResolution: assert deployment is not None assert deployment.litellm_params.custom_llm_provider == "openai" + # ===================================================================== # anthropic_messages mid-stream-fallback helpers, added for #24004 # (mid-stream fallback not supported for anthropic_messages route type). @@ -9369,10 +9002,7 @@ class _AnthropicMessagesFallbackByteStream: def _anthropic_messages_overloaded_error_chunk() -> bytes: - return ( - b"event: error\n" - b'data: {"type": "error", "error": {"type": "overloaded_error", "message": "Overloaded"}}\n\n' - ) + return b'event: error\ndata: {"type": "error", "error": {"type": "overloaded_error", "message": "Overloaded"}}\n\n' def _anthropic_messages_invalid_request_error_chunk() -> bytes: @@ -9417,9 +9047,7 @@ async def test_anthropic_messages_streaming_iterator_passthrough(): [_anthropic_messages_content_chunk("hi"), _anthropic_messages_content_chunk(" there")] ) - wrapped = await router._aanthropic_messages_streaming_iterator( - response=source, initial_kwargs={"model": "primary"} - ) + wrapped = await router._aanthropic_messages_streaming_iterator(response=source, initial_kwargs={"model": "primary"}) collected = [chunk async for chunk in wrapped] assert collected == [_anthropic_messages_content_chunk("hi"), _anthropic_messages_content_chunk(" there")] @@ -9438,12 +9066,14 @@ async def test_anthropic_messages_streaming_iterator_flushes_buffered_lifecycle_ [_anthropic_messages_message_start_chunk(), _anthropic_messages_content_chunk("hi"), message_stop] ) - wrapped = await router._aanthropic_messages_streaming_iterator( - response=source, initial_kwargs={"model": "primary"} - ) + wrapped = await router._aanthropic_messages_streaming_iterator(response=source, initial_kwargs={"model": "primary"}) collected = [chunk async for chunk in wrapped] - assert collected == [_anthropic_messages_message_start_chunk(), _anthropic_messages_content_chunk("hi"), message_stop] + assert collected == [ + _anthropic_messages_message_start_chunk(), + _anthropic_messages_content_chunk("hi"), + message_stop, + ] @pytest.mark.asyncio @@ -9455,9 +9085,7 @@ async def test_anthropic_messages_streaming_iterator_flushes_buffered_frames_on_ message_stop = b'event: message_stop\ndata: {"type": "message_stop"}\n\n' source = _AnthropicMessagesFakeByteStream([_anthropic_messages_message_start_chunk(), message_stop]) - wrapped = await router._aanthropic_messages_streaming_iterator( - response=source, initial_kwargs={"model": "primary"} - ) + wrapped = await router._aanthropic_messages_streaming_iterator(response=source, initial_kwargs={"model": "primary"}) collected = [chunk async for chunk in wrapped] assert collected == [_anthropic_messages_message_start_chunk(), message_stop] @@ -9528,7 +9156,9 @@ async def test_anthropic_messages_leading_ping_keepalive_is_forwarded_live(): yield _anthropic_messages_message_start_chunk() yield _anthropic_messages_content_chunk("hi") - wrapped = await router._aanthropic_messages_streaming_iterator(response=source(), initial_kwargs={"model": "primary"}) + wrapped = await router._aanthropic_messages_streaming_iterator( + response=source(), initial_kwargs={"model": "primary"} + ) assert await asyncio.wait_for(wrapped.__anext__(), timeout=1) == _anthropic_messages_ping_chunk() content_released.set() @@ -10669,3 +10299,68 @@ def test_permission_denied_error_is_retried_when_other_deployments_exist(): ) is True ) + + +class _AllowlistFallbackAccessCheck: + def __init__(self, allowed_models: frozenset[str]): + self.allowed_models = allowed_models + self.checked_models = [] + + async def __call__(self, *, model, request_kwargs, llm_router): + self.checked_models.append(model) + return model in self.allowed_models + + +def _router_with_failing_primary(fallback_access_check) -> Router: + return Router( + model_list=[ + { + "model_name": "primary", + "litellm_params": { + "model": "openai/primary", + "api_key": "k", + "mock_response": Exception("primary is down"), + }, + }, + { + "model_name": "secret-fallback", + "litellm_params": { + "model": "openai/secret", + "api_key": "k", + "mock_response": "served by secret-fallback", + }, + }, + ], + fallbacks=[{"primary": ["secret-fallback"]}], + num_retries=0, + fallback_access_check=fallback_access_check, + ) + + +@pytest.mark.asyncio +async def test_fallback_access_check_blocks_config_fallback_the_caller_cannot_use(): + access_check = _AllowlistFallbackAccessCheck(allowed_models=frozenset()) + router = _router_with_failing_primary(access_check) + + with pytest.raises(Exception, match="primary is down"): + await router.acompletion(model="primary", messages=[{"role": "user", "content": "hi"}]) + + assert access_check.checked_models == ["secret-fallback"] + + +@pytest.mark.asyncio +async def test_fallback_access_check_lets_an_authorized_config_fallback_through(): + router = _router_with_failing_primary(_AllowlistFallbackAccessCheck(allowed_models=frozenset({"secret-fallback"}))) + + response = await router.acompletion(model="primary", messages=[{"role": "user", "content": "hi"}]) + + assert response.choices[0].message.content == "served by secret-fallback" + + +@pytest.mark.asyncio +async def test_router_without_fallback_access_check_attempts_every_config_fallback(): + router = _router_with_failing_primary(None) + + response = await router.acompletion(model="primary", messages=[{"role": "user", "content": "hi"}]) + + assert response.choices[0].message.content == "served by secret-fallback"