diff --git a/litellm/proxy/hooks/model_max_budget_limiter.py b/litellm/proxy/hooks/model_max_budget_limiter.py index 9713f8986a8..4daef2592f7 100644 --- a/litellm/proxy/hooks/model_max_budget_limiter.py +++ b/litellm/proxy/hooks/model_max_budget_limiter.py @@ -800,6 +800,31 @@ def _as_float(value: object) -> float | None: return None +def _entity_model_budget_entries( + model: str, + entity_budget: tuple[Litellm_EntityType, str | None, object], + model_access_groups: ModelAccessGroups, + team_id: str | None, +) -> tuple[tuple[Litellm_EntityType, str, ResolvedModelBudget], ...]: + entity_type, entity_id, model_max_budget = entity_budget + if entity_id is None or not isinstance(model_max_budget, Mapping) or not model_max_budget: + return () + return tuple( + (entity_type, entity_id, resolved) + for candidate in expand_model_latest_alias(model) + if ( + resolved := resolve_model_budget( + model=candidate, + model_max_budget=model_max_budget, + model_access_groups=model_access_groups, + team_id=team_id, + ) + ) + is not None + and resolved.budget_config.budget_duration is not None + ) + + def _resolve_entity_model_budgets( model: str, entity_budgets: Iterable[tuple[Litellm_EntityType, str | None, object]], @@ -807,21 +832,14 @@ def _resolve_entity_model_budgets( team_id: str | None, ) -> tuple[tuple[Litellm_EntityType, str, ResolvedModelBudget], ...]: """Drop the scopes that do not budget `model`, keeping only what can be incremented.""" + entries: Final = itertools.chain.from_iterable( + _entity_model_budget_entries(model, entity_budget, model_access_groups, team_id) + for entity_budget in entity_budgets + ) return tuple( { (entity_type, entity_id, resolved.budget_model): (entity_type, entity_id, resolved) - for entity_type, entity_id, model_max_budget in entity_budgets - if entity_id is not None and isinstance(model_max_budget, Mapping) and model_max_budget - for candidate in expand_model_latest_alias(model) - for resolved in ( - resolve_model_budget( - model=candidate, - model_max_budget=model_max_budget, - model_access_groups=model_access_groups, - team_id=team_id, - ), - ) - if resolved is not None and resolved.budget_config.budget_duration is not None + for entity_type, entity_id, resolved in entries }.values() ) diff --git a/litellm/proxy/hooks/parallel_request_limiter_v3.py b/litellm/proxy/hooks/parallel_request_limiter_v3.py index aaab4bf70ca..8f1301fc86a 100644 --- a/litellm/proxy/hooks/parallel_request_limiter_v3.py +++ b/litellm/proxy/hooks/parallel_request_limiter_v3.py @@ -5882,14 +5882,6 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger): ) ) - elif tpm_actual > 0: - failed_kwargs: Final = dict(failed_call) - failed_model: Final = get_model_group_from_litellm_kwargs(failed_kwargs) - if failed_model is not None and len(expand_model_latest_alias(failed_model)) > 1: - pipeline_operations.extend( - self._build_success_event_pipeline_operations(failed_kwargs, None, self.get_rate_limit_type()) - ) - # Settle project ITPM/OTPM reservations the same way: at the # recovered partial usage, or a full refund when there is none. itpm_operations: Final = ( diff --git a/litellm/rust_bridge/catalog.py b/litellm/rust_bridge/catalog.py index 9ba8cbc35a7..9adfcb60d74 100644 --- a/litellm/rust_bridge/catalog.py +++ b/litellm/rust_bridge/catalog.py @@ -110,6 +110,7 @@ class Feature(str, Enum): """A model_list that Python fans out across its deployments.""" MODEL_ALIAS_MAP = "model_alias_map" """A model that litellm.model_alias_map rewrites before the call.""" + LATEST_ALIAS = "latest_alias" @dataclass(frozen=True, slots=True) @@ -153,6 +154,7 @@ COMMON_PYTHON_FEATURES: Final = frozenset( Feature.TIMEOUT_VARIANTS, Feature.AGENTIC_LOOP_HOOK, Feature.TLS_POLICY, + Feature.LATEST_ALIAS, } ) diff --git a/litellm/rust_bridge/chat_completions/detectors.py b/litellm/rust_bridge/chat_completions/detectors.py index 02db7750aa1..eb8f8bc49f9 100644 --- a/litellm/rust_bridge/chat_completions/detectors.py +++ b/litellm/rust_bridge/chat_completions/detectors.py @@ -8,7 +8,6 @@ from typing import Final from pydantic import TypeAdapter import litellm -from litellm.litellm_core_utils.model_latest_aliases import is_model_latest_alias from litellm.rust_bridge.detectors import ( agentic_loop_hook, api_base, @@ -40,9 +39,7 @@ def model_list(request: NativeCall) -> bool: def model_alias_map(request: NativeCall) -> bool: """Python rewrites this model through the configured alias map.""" model: Final = str(request.resolved["model"]) - return model in litellm.model_alias_map or is_model_latest_alias( - model, optional_str(request.resolved.get("custom_llm_provider")) - ) + return model in litellm.model_alias_map def chat_retries_fallbacks(request: NativeCall) -> bool: diff --git a/litellm/rust_bridge/detectors.py b/litellm/rust_bridge/detectors.py index 0af3a2e6fa6..f2591bdbcd1 100644 --- a/litellm/rust_bridge/detectors.py +++ b/litellm/rust_bridge/detectors.py @@ -12,6 +12,7 @@ import httpx import litellm from litellm.integrations.custom_logger import CustomLogger from litellm.litellm_core_utils.core_helpers import normalize_drop_params +from litellm.litellm_core_utils.model_latest_aliases import is_model_latest_alias from litellm.rust_bridge.catalog import Feature from litellm.rust_bridge.public_call import NativeCall, optional_mapping, optional_sequence, optional_str @@ -35,6 +36,13 @@ def mock_response(request: NativeCall) -> bool: return present(request, "mock_response") +def latest_alias(request: NativeCall) -> bool: + return is_model_latest_alias( + optional_str(request.resolved.get("model")) or "", + optional_str(request.resolved.get("custom_llm_provider")), + ) + + def caller_client(request: NativeCall) -> bool: """Python sends through the caller's HTTP client or shared session.""" return present(request, "client") or present(request, "shared_session") @@ -179,5 +187,6 @@ COMMON: Final[Detectors] = MappingProxyType( Feature.TIMEOUT_VARIANTS: timeout_variants, Feature.AGENTIC_LOOP_HOOK: agentic_loop_hook, Feature.TLS_POLICY: tls_policy, + Feature.LATEST_ALIAS: latest_alias, } ) diff --git a/tests/unit/proxy/hooks/test_parallel_request_limiter_v3.py b/tests/unit/proxy/hooks/test_parallel_request_limiter_v3.py index d212e91871b..4117072c406 100644 --- a/tests/unit/proxy/hooks/test_parallel_request_limiter_v3.py +++ b/tests/unit/proxy/hooks/test_parallel_request_limiter_v3.py @@ -192,7 +192,8 @@ async def test_latest_alias_collision_charges_and_settles_one_counter_per_metric await handler.async_log_success_event(kwargs, ModelResponse(model=target, usage=usage), None, None) else: await handler.async_log_failure_event(kwargs, None, None, None) - assert await cache.async_get_cache(tpm_key) == (0 if event == "failure" else 5) + (0 if reservation else 1) + actual_tokens: Final = 5 if event == "success" or (event == "partial" and reservation) else 0 + assert await cache.async_get_cache(tpm_key) == actual_tokens + (0 if reservation else 1) alias_key: Final = handler.create_rate_limit_keys(f"model_per_{scope}", f"{entity}:{alias}", "tokens") assert await cache.async_get_cache(alias_key) is None assert await handler.async_pre_call_hook(auth, cache, {"model": alias}, "acompletion") is None @@ -270,9 +271,10 @@ async def test_latest_alias_collision_settles_the_target_team_limit_beside_the_k await handler.async_log_success_event(kwargs, ModelResponse(model=target, usage=usage), None, None) else: await handler.async_log_failure_event(kwargs, None, None, None) + actual_tokens: Final = 5 if reservation or event == "success" else 0 for scope, entity in (("model_per_key", auth.api_key), ("model_per_team", auth.team_id)): counter: Final = handler.create_rate_limit_keys(scope, f"{entity}:{target}", "tokens") - assert await cache.async_get_cache(counter) == 5 + (0 if reservation else 1) + assert await cache.async_get_cache(counter) == actual_tokens + (0 if reservation else 1) @pytest.mark.asyncio @@ -446,8 +448,81 @@ async def test_latest_alias_tpm_settles_both_counters( ) else: await handler.async_log_failure_event(kwargs, None, None, None) + actual_tokens: Final = 5 if event == "success" or (event == "partial" and reservation) else 0 for key in (*keys, handler.create_rate_limit_keys("api_key", auth.api_key, "tokens")): - assert int(await cache.async_get_cache(key) or 0) == (0 if event == "failure" else 5) + (0 if reservation else 1) + assert int(await cache.async_get_cache(key) or 0) == actual_tokens + (0 if reservation else 1) + + +@pytest.mark.asyncio +@pytest.mark.parametrize("scope", ("key", "team", "project", "organization")) +@pytest.mark.parametrize("request_alias", (False, True)) +@pytest.mark.parametrize("reservation_released", (False, True)) +async def test_failed_latest_alias_without_open_reservation_matches_direct_model_accounting( + scope: str, + request_alias: bool, + reservation_released: bool, + monkeypatch: pytest.MonkeyPatch, +) -> None: + monkeypatch.setenv("LITELLM_TPM_TOKEN_RESERVATION_ENABLED", str(reservation_released).lower()) + cache: Final = DualCache() + handler: Final = _PROXY_MaxParallelRequestsHandler( + internal_usage_cache=InternalUsageCache(cache), + time_provider=lambda: datetime(2026, 1, 1, tzinfo=timezone.utc), + ) + alias: Final = "anthropic/claude-opus-latest" + target: Final = f"anthropic/{get_model_latest_aliases()['anthropic']['claude-opus-latest']}" + model: Final = alias if request_alias else target + limits: Final = {"model_tpm_limit": {alias: 10000, target: 10000}} + auth: Final = UserAPIKeyAuth( + api_key="failed-latest-key", + team_id="failed-latest-team", + project_id="failed-latest-project", + org_id="failed-latest-org", + tpm_limit=10000, + metadata=limits if scope == "key" else {}, + team_metadata=limits if scope == "team" else {}, + project_metadata=limits if scope == "project" else {}, + organization_metadata=limits if scope == "organization" else {}, + ) + request_data: Final = {"model": model, "messages": [{"role": "user", "content": "hi"}], "max_tokens": 64} + await handler.async_pre_call_hook(auth, cache, request_data, "acompletion") + if reservation_released: + await handler.async_post_call_failure_hook( + request_data=request_data, + original_exception=Exception("rejected"), + user_api_key_dict=auth, + ) + entity: Final = { + "key": auth.api_key, + "team": auth.team_id, + "project": auth.project_id, + "organization": auth.org_id, + }[scope] + models: Final = (alias, target) if request_alias else (target,) + keys: Final = ( + handler.create_rate_limit_keys("api_key", auth.api_key, "tokens"), + *(handler.create_rate_limit_keys(f"model_per_{scope}", f"{entity}:{name}", "tokens") for name in models), + ) + await handler.async_log_failure_event( + { + "standard_logging_object": { + "model_group": model, + "metadata": { + "user_api_key_hash": auth.api_key, + "user_api_key_team_id": auth.team_id, + "user_api_key_project_id": auth.project_id, + "user_api_key_org_id": auth.org_id, + }, + }, + "litellm_params": {"metadata": {"model_group": model, "user_api_key_metadata": auth.metadata}}, + "combined_usage_object": Usage(prompt_tokens=3, completion_tokens=2, total_tokens=5), + }, + None, + None, + None, + ) + for key in keys: + assert int(await cache.async_get_cache(key) or 0) == (0 if reservation_released else 1) @pytest.mark.asyncio diff --git a/tests/unit/rust_bridge/test_catalog.py b/tests/unit/rust_bridge/test_catalog.py index 704a4f73fe8..1994dd00709 100644 --- a/tests/unit/rust_bridge/test_catalog.py +++ b/tests/unit/rust_bridge/test_catalog.py @@ -5,7 +5,8 @@ from typing import Final import pytest -from litellm.rust_bridge import catalog, configuration +from litellm.litellm_core_utils.model_latest_aliases import get_model_latest_aliases +from litellm.rust_bridge import catalog, configuration, detectors from litellm.rust_bridge.catalog import ( Feature, Route, @@ -15,6 +16,32 @@ from litellm.rust_bridge.catalog import ( Target, ) from litellm.rust_bridge.configuration import Decision, Rollout +from litellm.rust_bridge.public_call import NativeCall + + +@pytest.mark.parametrize( + ("model", "expected"), + ( + ("anthropic/claude-opus-latest", Rollout.PYTHON_ONLY), + (f"anthropic/{get_model_latest_aliases()['anthropic']['claude-opus-latest']}", Rollout.RUST_OPT_IN), + ), +) +def test_latest_alias_messages_falls_back_without_changing_the_pinned_model_rollout( + model: str, + expected: Rollout, +) -> None: + request: Final = NativeCall(args=(), kwargs={"model": model}, base={}) + context: Final = RouteContext( + Route.MESSAGES, + provider="anthropic", + model=model, + features=detectors.detected(detectors.COMMON, request), + ) + rules: Final = ( + RouteRule(Route.MESSAGES, Rollout.PYTHON_ONLY, features=catalog.MESSAGES_PYTHON_FEATURES), + RouteRule(Route.MESSAGES, Rollout.RUST_OPT_IN, providers=frozenset({"anthropic"})), + ) + assert catalog.rollout(context, rules) is expected @pytest.fixture(autouse=True) diff --git a/tests/unit/rust_bridge/test_detectors.py b/tests/unit/rust_bridge/test_detectors.py index 533d9033ff6..3001d23e57e 100644 --- a/tests/unit/rust_bridge/test_detectors.py +++ b/tests/unit/rust_bridge/test_detectors.py @@ -52,10 +52,13 @@ def call(**kwargs: object) -> NativeCall: ("gpt-luna-latest", "openai", True), ("gemini/gemini-pro-latest", None, False), ("gpt-luna-latest", None, False), + ("anthropic/claude-opus-latest", None, True), + ("claude-opus-latest", "anthropic", True), + ("", None, False), ), ) def test_floating_alias_requires_python(model: str, provider: str | None, expected: bool) -> None: - assert chat_detectors.model_alias_map(call(model=model, custom_llm_provider=provider)) is expected + assert detectors.latest_alias(call(model=model, custom_llm_provider=provider)) is expected @pytest.mark.parametrize(