From 7f04cbe781fccbbb285b23052ee2e2bca5094e57 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Tue, 20 Oct 2026 19:09:26 -0700 Subject: [PATCH] fix(proxy): apply the quota policy to model add, follow router wildcard rules, keep old setting values --- .../management_endpoints/project_endpoints.py | 22 ++-- .../project_model_quota.py | 116 +++--------------- .../model_management_endpoints.py | 50 +++++--- .../management_helpers/model_quota_policy.py | 110 +++++++++++++++++ .../test_project_endpoints_prisma.py | 83 +++++++++---- .../test_model_management_endpoints.py | 96 +++++++++++++-- 6 files changed, 326 insertions(+), 151 deletions(-) create mode 100644 litellm/proxy/management_helpers/model_quota_policy.py diff --git a/enterprise/litellm_enterprise/proxy/management_endpoints/project_endpoints.py b/enterprise/litellm_enterprise/proxy/management_endpoints/project_endpoints.py index 8ad022a857a..883e7b8ae8e 100644 --- a/enterprise/litellm_enterprise/proxy/management_endpoints/project_endpoints.py +++ b/enterprise/litellm_enterprise/proxy/management_endpoints/project_endpoints.py @@ -17,14 +17,8 @@ from typing import TYPE_CHECKING, Final, cast from fastapi import APIRouter, Depends, HTTPException, Request from litellm_enterprise.proxy.management_endpoints.project_model_quota import ( ENFORCE_PROJECT_MODEL_QUOTA_SETTING, - REQUIRE_BOTH_LIMITS, - ProjectModelQuotaPolicy, - ProvidersOf, describe_missing_quota, missing_project_model_quota, - no_providers, - parse_project_model_quota_setting, - router_deployment_providers, ) from pydantic import TypeAdapter, ValidationError @@ -38,6 +32,14 @@ from litellm.proxy.management.teams.dependencies import get_team_access from litellm.proxy.management.users.service import org_admin_org_ids from litellm.proxy.management_endpoints.common_utils import _set_object_metadata_field from litellm.proxy.management_endpoints.team_admin_field_permissions import team_admin_may_manage_projects +from litellm.proxy.management_helpers.model_quota_policy import ( + REQUIRE_BOTH_LIMITS, + ModelQuotaPolicy, + ProvidersOf, + no_providers, + parse_model_quota_setting, + router_deployment_providers, +) from litellm.proxy.management_helpers.utils import ( management_endpoint_wrapper, ) @@ -307,9 +309,9 @@ def _raise_on_project_models_expanding_at_request_time( ) -def _project_model_quota_policy(general_settings: Mapping[str, object]) -> ProjectModelQuotaPolicy | None: +def _project_model_quota_policy(general_settings: Mapping[str, object]) -> ModelQuotaPolicy | None: try: - return parse_project_model_quota_setting(general_settings.get(ENFORCE_PROJECT_MODEL_QUOTA_SETTING)) + return parse_model_quota_setting(general_settings.get(ENFORCE_PROJECT_MODEL_QUOTA_SETTING)) except ValidationError as e: raise HTTPException( status_code=500, @@ -320,7 +322,7 @@ def _project_model_quota_policy(general_settings: Mapping[str, object]) -> Proje def _raise_on_missing_project_model_quota( data: NewProjectRequest | UpdateProjectRequest, access_group_names: frozenset[str] = frozenset(), - policy: ProjectModelQuotaPolicy = REQUIRE_BOTH_LIMITS, + policy: ModelQuotaPolicy = REQUIRE_BOTH_LIMITS, providers_of: ProvidersOf = no_providers, ) -> None: """Require the positive `rpm`/`tpm` quotas the policy asks of each model on project CREATE. @@ -347,7 +349,7 @@ def _raise_on_missing_project_model_quota_on_update( data: UpdateProjectRequest, existing_project: object, access_group_names: frozenset[str] = frozenset(), - policy: ProjectModelQuotaPolicy = REQUIRE_BOTH_LIMITS, + policy: ModelQuotaPolicy = REQUIRE_BOTH_LIMITS, providers_of: ProvidersOf = no_providers, ) -> None: """Require the positive `rpm`/`tpm` quotas the policy asks of each model over the RESULTING state on project UPDATE. diff --git a/enterprise/litellm_enterprise/proxy/management_endpoints/project_model_quota.py b/enterprise/litellm_enterprise/proxy/management_endpoints/project_model_quota.py index f0ddd3eed30..e427e37f10b 100644 --- a/enterprise/litellm_enterprise/proxy/management_endpoints/project_model_quota.py +++ b/enterprise/litellm_enterprise/proxy/management_endpoints/project_model_quota.py @@ -1,103 +1,46 @@ -import re -from collections.abc import Callable, Mapping, Sequence +from collections.abc import Mapping, Sequence from dataclasses import dataclass from types import MappingProxyType -from typing import TYPE_CHECKING, Final, Literal, TypeAlias +from typing import Final -from pydantic import BaseModel, ConfigDict, Field, TypeAdapter - -from litellm.router_utils.common_utils import provider_for_generic_call -from litellm.router_utils.pattern_match_deployments import PatternUtils - -if TYPE_CHECKING: - from litellm import Router - -QuotaLimit: TypeAlias = Literal["rpm", "tpm"] -ProvidersOf: TypeAlias = Callable[[str], frozenset[str]] +from litellm.proxy.management_helpers.model_quota_policy import ( + ModelQuotaPolicy, + ProvidersOf, + QuotaLimit, + required_limits, + unmet_limits, +) ENFORCE_PROJECT_MODEL_QUOTA_SETTING: Final = "enforce_project_model_quota" -_BOTH_LIMITS: Final[frozenset[QuotaLimit]] = frozenset(("rpm", "tpm")) _NO_LIMITS_SET: Final[Mapping[str, object]] = MappingProxyType({}) -class ProjectModelQuotaPolicy(BaseModel): - model_config = ConfigDict(frozen=True, extra="forbid") - - default: frozenset[QuotaLimit] = _BOTH_LIMITS - providers: Mapping[str, frozenset[QuotaLimit]] = Field(default_factory=dict) - models: Mapping[str, frozenset[QuotaLimit]] = Field(default_factory=dict) - - -REQUIRE_BOTH_LIMITS: Final = ProjectModelQuotaPolicy() - -_QUOTA_SETTING: Final[TypeAdapter[bool | ProjectModelQuotaPolicy | None]] = TypeAdapter( - bool | ProjectModelQuotaPolicy | None -) - - @dataclass(frozen=True, slots=True) class MissingModelQuota: model: str limits: tuple[QuotaLimit, ...] -def parse_project_model_quota_setting(setting: object) -> ProjectModelQuotaPolicy | None: - parsed: Final = _QUOTA_SETTING.validate_python(setting) - match parsed: - case ProjectModelQuotaPolicy(): - return parsed - case True: - return REQUIRE_BOTH_LIMITS - case _: - return None - - -def _wildcard_matches(pattern: str, model: str) -> bool: - return re.fullmatch(re.escape(pattern).replace(r"\*", ".*"), model) is not None - - -def _most_specific_model_pattern(model: str, entries: Mapping[str, frozenset[QuotaLimit]]) -> str | None: - matching: Final = tuple(pattern for pattern in entries if "*" in pattern and _wildcard_matches(pattern, model)) - return max(matching, key=PatternUtils.calculate_pattern_specificity, default=None) - - -def required_limits(policy: ProjectModelQuotaPolicy, model: str, providers: frozenset[str]) -> frozenset[QuotaLimit]: - """A model entry (exact name, then the most specific `*` pattern) wins over the providers of the - model's deployments, which win over the default. A model served by several providers needs every - limit any of them requires.""" - exact: Final = policy.models.get(model) - if exact is not None: - return exact - pattern: Final = _most_specific_model_pattern(model, policy.models) - if pattern is not None: - return policy.models[pattern] - if not providers: - return policy.default - return frozenset[QuotaLimit]().union(*(policy.providers.get(provider, policy.default) for provider in providers)) - - -def _is_positive(value: object) -> bool: - return isinstance(value, int) and not isinstance(value, bool) and value > 0 - - -def _unset_limits( - model: str, required: frozenset[QuotaLimit], set_limits: Mapping[QuotaLimit, Mapping[str, object]] -) -> tuple[QuotaLimit, ...]: - return tuple(sorted(limit for limit in required if not _is_positive(set_limits[limit].get(model)))) +def _limits_set_for( + model: str, rpm_set: Mapping[str, object], tpm_set: Mapping[str, object] +) -> Mapping[QuotaLimit, object]: + return MappingProxyType({"rpm": rpm_set.get(model), "tpm": tpm_set.get(model)}) def missing_project_model_quota( models: Sequence[str], rpm_limits: Mapping[str, object] | None, tpm_limits: Mapping[str, object] | None, - policy: ProjectModelQuotaPolicy, + policy: ModelQuotaPolicy, providers_of: ProvidersOf, ) -> tuple[MissingModelQuota, ...]: - set_limits: Final[Mapping[QuotaLimit, Mapping[str, object]]] = MappingProxyType( - {"rpm": rpm_limits or _NO_LIMITS_SET, "tpm": tpm_limits or _NO_LIMITS_SET} - ) + rpm_set: Final = rpm_limits or _NO_LIMITS_SET + tpm_set: Final = tpm_limits or _NO_LIMITS_SET gaps: Final = tuple( - MissingModelQuota(model, _unset_limits(model, required_limits(policy, model, providers_of(model)), set_limits)) + MissingModelQuota( + model, + unmet_limits(required_limits(policy, model, providers_of(model)), _limits_set_for(model, rpm_set, tpm_set)), + ) for model in models ) return tuple(gap for gap in gaps if gap.limits) @@ -105,22 +48,3 @@ def missing_project_model_quota( def describe_missing_quota(missing: Sequence[MissingModelQuota]) -> str: return ", ".join(f"{gap.model} ({' and '.join(f'model_{limit}_limit' for limit in gap.limits)})" for gap in missing) - - -def no_providers(_model: str) -> frozenset[str]: - return frozenset() - - -def router_deployment_providers(llm_router: "Router | None", team_id: str | None) -> ProvidersOf: - if llm_router is None: - return no_providers - - def providers_of(model: str) -> frozenset[str]: - deployments: Final = llm_router.get_model_list(model_name=model, team_id=team_id) or () - return frozenset( - provider - for deployment in deployments - if (provider := provider_for_generic_call(deployment["litellm_params"])) is not None - ) - - return providers_of diff --git a/litellm/proxy/management_endpoints/model_management_endpoints.py b/litellm/proxy/management_endpoints/model_management_endpoints.py index fc0868a5fc9..504cc9afcbc 100644 --- a/litellm/proxy/management_endpoints/model_management_endpoints.py +++ b/litellm/proxy/management_endpoints/model_management_endpoints.py @@ -102,6 +102,13 @@ from litellm.proxy.management_helpers.auto_router_permissions import ( authorize_member_auto_router_write, ) from litellm.proxy.management_helpers.model_allowlist_rename_sync import sync_model_allowlists_for_renamed_model +from litellm.proxy.management_helpers.model_quota_policy import ( + ModelQuotaPolicy, + deployment_provider, + parse_model_quota_setting, + required_limits, + unmet_limits, +) from litellm.proxy.spend_tracking.ptu_feature_flag import ( PTU_COST_ATTRIBUTION_ENV_VAR, is_ptu_cost_attribution_enabled, @@ -151,6 +158,7 @@ from litellm.types.router import ( SPECIAL_MODEL_INFO_PARAMS, Deployment, GenericLiteLLMParams, + LiteLLM_Params, ModelInfo, updateDeployment, updateLiteLLMParams, @@ -760,26 +768,39 @@ async def _auto_router_capability_slot( ENFORCE_RPM_TPM_ON_MODEL_ADD_SETTING: Final = "enforce_rpm_tpm_on_model_add" -def _raise_if_rate_limits_required_but_missing(*, litellm_params: GenericLiteLLMParams, enforced: bool) -> None: - """Require both rpm and tpm (each a positive value) when the operator opts in via config.yaml. +def _model_add_quota_policy(general_settings: Mapping[str, object]) -> ModelQuotaPolicy | None: + try: + return parse_model_quota_setting(general_settings.get(ENFORCE_RPM_TPM_ON_MODEL_ADD_SETTING)) + except ValidationError as e: + raise ProxyException( + message=f"general_settings.{ENFORCE_RPM_TPM_ON_MODEL_ADD_SETTING} is invalid: {e}", + type=ProxyErrorTypes.internal_server_error.value, + code=status.HTTP_500_INTERNAL_SERVER_ERROR, + param=None, + ) from e - Off by default, so deployments keep adding models without limits. When - ``enforce_rpm_tpm_on_model_add: true`` is set under general_settings, a model added - without both rpm and tpm set to a positive value is rejected rather than stored - unbounded (or effectively excluded from routing by a zero/negative limit). - """ - if not enforced: + +def _raise_if_rate_limits_required_but_missing( + *, model_name: str, litellm_params: LiteLLM_Params, policy: ModelQuotaPolicy | None +) -> None: + """Off by default, so deployments keep adding models without limits. Once the operator opts in under + general_settings, the policy decides per model name and provider which of rpm and tpm a new model needs, + and a model missing one, or carrying a zero or negative one, is rejected rather than stored unbounded + (or effectively excluded from routing).""" + if policy is None: return - missing: Final = tuple( - field - for field, value in (("rpm", litellm_params.rpm), ("tpm", litellm_params.tpm)) - if value is None or value <= 0 + provider: Final = deployment_provider( + MappingProxyType({"model": litellm_params.model, "custom_llm_provider": litellm_params.custom_llm_provider}) + ) + missing: Final = unmet_limits( + required_limits(policy, model_name, frozenset() if provider is None else frozenset((provider,))), + MappingProxyType({"rpm": litellm_params.rpm, "tpm": litellm_params.tpm}), ) if not missing: return raise ProxyException( message=( - f"{' and '.join(missing)} must be set to a positive value when " + f"{' and '.join(missing)} must be set to a positive value for '{model_name}' when " f"'{ENFORCE_RPM_TPM_ON_MODEL_ADD_SETTING}' is enabled in general_settings" ), type=ProxyErrorTypes.validation_error.value, @@ -2717,8 +2738,9 @@ async def add_new_model( ) _raise_if_rate_limits_required_but_missing( + model_name=model_params.model_name, litellm_params=model_params.litellm_params, - enforced=bool(general_settings.get(ENFORCE_RPM_TPM_ON_MODEL_ADD_SETTING, False)), + policy=_model_add_quota_policy(general_settings), ) clean_model_info: Final = ModelInfo.model_validate( diff --git a/litellm/proxy/management_helpers/model_quota_policy.py b/litellm/proxy/management_helpers/model_quota_policy.py new file mode 100644 index 00000000000..e2467e32997 --- /dev/null +++ b/litellm/proxy/management_helpers/model_quota_policy.py @@ -0,0 +1,110 @@ +import re +from collections.abc import Callable, Mapping, Sequence +from typing import TYPE_CHECKING, Final, Literal, TypeAlias + +from pydantic import BaseModel, ConfigDict, Field + +from litellm.llms.azure.ptu_capacity import is_str_keyed_mapping +from litellm.router_utils.common_utils import provider_for_generic_call +from litellm.router_utils.pattern_match_deployments import PatternMatchRouter, PatternUtils + +if TYPE_CHECKING: + from litellm import Router + +QuotaLimit: TypeAlias = Literal["rpm", "tpm"] +ProvidersOf: TypeAlias = Callable[[str], frozenset[str]] + +_LIMITS: Final[tuple[QuotaLimit, ...]] = ("rpm", "tpm") +_ROUTER_PATTERNS: Final = PatternMatchRouter() + + +class ModelQuotaPolicy(BaseModel): + model_config = ConfigDict(frozen=True, extra="forbid") + + default: frozenset[QuotaLimit] = frozenset(_LIMITS) + providers: Mapping[str, frozenset[QuotaLimit]] = Field(default_factory=dict) + models: Mapping[str, frozenset[QuotaLimit]] = Field(default_factory=dict) + + +REQUIRE_BOTH_LIMITS: Final = ModelQuotaPolicy() + + +def parse_model_quota_setting(setting: object) -> ModelQuotaPolicy | None: + """Any falsy value is off and any other non-mapping value requires both limits, as the boolean + setting always behaved, so only a mapping is read as a per-provider and per-model policy.""" + if not setting: + return None + if is_str_keyed_mapping(setting): + return ModelQuotaPolicy.model_validate(setting) + return REQUIRE_BOTH_LIMITS + + +def _wildcard_matches(pattern: str, model: str) -> bool: + return re.fullmatch(_ROUTER_PATTERNS.pattern_to_regex(pattern), model) is not None + + +def _router_specificity(pattern: str) -> tuple[int, int]: + return PatternUtils.calculate_pattern_specificity(_ROUTER_PATTERNS.pattern_to_regex(pattern)) + + +def _most_specific_model_pattern(model: str, entries: Mapping[str, frozenset[QuotaLimit]]) -> str | None: + matching: Final = tuple(pattern for pattern in entries if "*" in pattern and _wildcard_matches(pattern, model)) + return max(matching, key=_router_specificity, default=None) + + +def required_limits(policy: ModelQuotaPolicy, model: str, providers: frozenset[str]) -> frozenset[QuotaLimit]: + """A model entry (exact name, then the most specific `*` pattern) wins over the providers of the + model's deployments, which win over the default. A model served by several providers needs every + limit any of them requires.""" + exact: Final = policy.models.get(model) + if exact is not None: + return exact + pattern: Final = _most_specific_model_pattern(model, policy.models) + if pattern is not None: + return policy.models[pattern] + if not providers: + return policy.default + return frozenset[QuotaLimit]().union(*(policy.providers.get(provider, policy.default) for provider in providers)) + + +def _is_positive(value: object) -> bool: + return isinstance(value, int) and not isinstance(value, bool) and value > 0 + + +def unmet_limits(required: frozenset[QuotaLimit], values: Mapping[QuotaLimit, object]) -> tuple[QuotaLimit, ...]: + return tuple( + limit + for limit in _LIMITS + if (limit in required or values.get(limit) is not None) and not _is_positive(values.get(limit)) + ) + + +def no_providers(_model: str) -> frozenset[str]: + return frozenset() + + +def deployment_provider(litellm_params: object) -> str | None: + return provider_for_generic_call(litellm_params) if is_str_keyed_mapping(litellm_params) else None + + +def _routed_deployments(llm_router: "Router", model: str, team_id: str | None) -> Sequence[Mapping[str, object]]: + named: Final[tuple[Mapping[str, object], ...]] = ( + *llm_router.get_all_deployments(model_name=model, team_id=team_id), + *llm_router.get_model_list_from_model_alias(model_name=model), + *llm_router.get_model_list_from_routing_groups(model_name=model), + ) + return named or llm_router.wildcard_route_deployments(model, team_id) + + +def router_deployment_providers(llm_router: "Router | None", team_id: str | None) -> ProvidersOf: + if llm_router is None: + return no_providers + + def providers_of(model: str) -> frozenset[str]: + return frozenset( + provider + for deployment in _routed_deployments(llm_router, model, team_id) + if (provider := deployment_provider(deployment.get("litellm_params"))) is not None + ) + + return providers_of diff --git a/tests/unit/enterprise/proxy/management_endpoints/test_project_endpoints_prisma.py b/tests/unit/enterprise/proxy/management_endpoints/test_project_endpoints_prisma.py index 6a33959e634..c1e3c035b57 100644 --- a/tests/unit/enterprise/proxy/management_endpoints/test_project_endpoints_prisma.py +++ b/tests/unit/enterprise/proxy/management_endpoints/test_project_endpoints_prisma.py @@ -1,7 +1,7 @@ import asyncio import os import traceback -from collections.abc import Mapping +from collections.abc import Callable, Mapping from dataclasses import dataclass from datetime import datetime, timezone from typing import Final @@ -1627,22 +1627,20 @@ def _quota_error( rpm: dict[str, int] | None = None, tpm: dict[str, int] | None = None, setting: object = _CUSTOMER_QUOTA_POLICY, - providers_of=_customer_providers_of, + providers_of: Callable[[str], frozenset[str]] = _customer_providers_of, ) -> str | None: from litellm_enterprise.proxy.management_endpoints.project_endpoints import ( _raise_on_missing_project_model_quota, ) - from litellm_enterprise.proxy.management_endpoints.project_model_quota import ( - parse_project_model_quota_setting, - ) + from litellm.proxy.management_helpers.model_quota_policy import parse_model_quota_setting - policy = parse_project_model_quota_setting(setting) + policy: Final = parse_model_quota_setting(setting) assert policy is not None - data = NewProjectRequest(team_id="test-team", models=models, model_rpm_limit=rpm, model_tpm_limit=tpm) + data: Final = NewProjectRequest(team_id="test-team", models=models, model_rpm_limit=rpm, model_tpm_limit=tpm) try: _raise_on_missing_project_model_quota(data, frozenset(), policy, providers_of) except HTTPException as e: - rejection = e + rejection: Final = e else: return None assert rejection.status_code == 400 @@ -1664,6 +1662,27 @@ def test_quota_policy_model_pattern_overrides_its_provider_per_limit(): assert "model_tpm_limit" not in error +@pytest.mark.parametrize( + ("model", "rpm", "tpm", "expected_missing"), + [ + ("runway-gen4-turbo", {"runway-gen4-turbo": 0}, None, "runway-gen4-turbo (model_rpm_limit)"), + ("jev-latest", None, {"jev-latest": -5}, "jev-latest (model_tpm_limit)"), + ( + "veo-3.1-lite-generate-001", + {"veo-3.1-lite-generate-001": 5}, + {"veo-3.1-lite-generate-001": 0}, + "veo-3.1-lite-generate-001 (model_tpm_limit)", + ), + ], +) +def test_quota_policy_rejects_a_non_positive_limit_set_on_an_exempt_model( + model: str, rpm: dict[str, int] | None, tpm: dict[str, int] | None, expected_missing: str +): + error: Final = _quota_error([model], rpm=rpm, tpm=tpm) + assert error is not None + assert expected_missing in error + + def test_quota_policy_error_names_each_model_with_the_limits_it_misses(): error = _quota_error(["gemini-3.8-flash", "veo-3.1-lite-generate-001", "runway-gen4-turbo"]) assert error is not None @@ -1684,10 +1703,13 @@ def test_quota_policy_error_names_each_model_with_the_limits_it_misses(): ), ({"models": {"runway-*": [], "runway-gen4-turbo": ["rpm"]}}, "runway-gen4-turbo", "(model_rpm_limit)"), ({"models": {"veo-*": [], "veo-3.1-lite-*": ["rpm"]}}, "veo-3.1-lite-generate-001", "(model_rpm_limit)"), + ({"models": {"gpt-5-mini*": [], "gpt-*-*-*": ["rpm"]}}, "gpt-5-mini-fast", "(model_rpm_limit)"), ], ) -def test_quota_policy_precedence_is_model_then_pattern_then_provider_then_default(setting, model, expected_missing): - error = _quota_error([model], setting=setting) +def test_quota_policy_precedence_is_model_then_pattern_then_provider_then_default( + setting: object, model: str, expected_missing: str +): + error: Final = _quota_error([model], setting=setting) assert error is not None assert expected_missing in error @@ -1708,20 +1730,18 @@ def test_quota_policy_model_on_several_providers_needs_every_required_limit(): assert "video-pool (model_tpm_limit)" in error -@pytest.mark.parametrize("setting", [True, "true", {}]) -def test_quota_setting_true_or_empty_mapping_still_requires_both_limits(setting): - error = _quota_error(["runway-gen4-turbo"], setting=setting) +@pytest.mark.parametrize("setting", [True, "true", "false", 1, ["rpm"]]) +def test_quota_setting_truthy_non_mapping_still_requires_both_limits(setting: object): + error: Final = _quota_error(["runway-gen4-turbo"], setting=setting) assert error is not None assert "runway-gen4-turbo (model_rpm_limit and model_tpm_limit)" in error -@pytest.mark.parametrize("setting", [False, None]) -def test_quota_setting_off_enforces_nothing(setting): - from litellm_enterprise.proxy.management_endpoints.project_model_quota import ( - parse_project_model_quota_setting, - ) +@pytest.mark.parametrize("setting", [False, None, {}, "", 0, []]) +def test_quota_setting_off_enforces_nothing(setting: object): + from litellm.proxy.management_helpers.model_quota_policy import parse_model_quota_setting - assert parse_project_model_quota_setting(setting) is None + assert parse_model_quota_setting(setting) is None def test_update_quota_policy_resolves_providers_for_the_resulting_models(): @@ -1730,11 +1750,9 @@ def test_update_quota_policy_resolves_providers_for_the_resulting_models(): from litellm_enterprise.proxy.management_endpoints.project_endpoints import ( _raise_on_missing_project_model_quota_on_update, ) - from litellm_enterprise.proxy.management_endpoints.project_model_quota import ( - parse_project_model_quota_setting, - ) + from litellm.proxy.management_helpers.model_quota_policy import parse_model_quota_setting - policy = parse_project_model_quota_setting(_CUSTOMER_QUOTA_POLICY) + policy = parse_model_quota_setting(_CUSTOMER_QUOTA_POLICY) assert policy is not None existing = types.SimpleNamespace(models=["runway-gen4-turbo"], metadata={}) assert ( @@ -1777,6 +1795,25 @@ def _customer_router() -> "litellm.Router": ) +def test_quota_providers_come_from_the_one_wildcard_route_requests_use(): + from litellm.proxy.management_helpers.model_quota_policy import router_deployment_providers + + router: Final = litellm.Router( + model_list=[ + {"model_name": "gen4*", "litellm_params": {"model": "runwayml/gen4*", "api_key": "qa-runway-key"}}, + { + "model_name": "media-team-gen", + "litellm_params": {"model": "vertex_ai/*", "vertex_project": "qa-project"}, + "model_info": {"team_id": "media-team", "team_public_model_name": "gen*"}, + }, + ] + ) + providers_of: Final = router_deployment_providers(router, "media-team") + + assert providers_of("gen4_turbo") == frozenset({"runwayml"}) + assert providers_of("gen3a_turbo") == frozenset({"vertex_ai"}) + + @pytest.mark.asyncio async def test_new_project_quota_policy_reads_providers_from_router_deployments(monkeypatch): team_models = ["gemini-3.8-flash", "veo-3.1-lite-generate-001", "runway-gen4-turbo", "jev-latest"] diff --git a/tests/unit/proxy/management_endpoints/test_model_management_endpoints.py b/tests/unit/proxy/management_endpoints/test_model_management_endpoints.py index 303b8e060ed..5870e6fad9a 100644 --- a/tests/unit/proxy/management_endpoints/test_model_management_endpoints.py +++ b/tests/unit/proxy/management_endpoints/test_model_management_endpoints.py @@ -38,6 +38,7 @@ from litellm.proxy.management_endpoints.model_management_endpoints import ( patch_model, update_model, ) +from litellm.proxy.management_helpers.model_quota_policy import REQUIRE_BOTH_LIMITS, ModelQuotaPolicy from litellm.proxy.utils import PrismaClient from litellm.router import Router from litellm.types.router import Deployment, LiteLLM_Params, ModelInfo, updateDeployment, updateLiteLLMParams @@ -7352,12 +7353,19 @@ class TestAutoRouterClassifierDefaultPrompt: assert response.system_prompt == classification_system_prompt(5) +_CUSTOMER_MODEL_ADD_POLICY: Final = ModelQuotaPolicy.model_validate( + { + "providers": {"vertex_ai": ["rpm", "tpm"], "runwayml": []}, + "models": {"veo-*": ["rpm"], "jev-latest": []}, + } +) + + class TestEnforceRpmTpmOnModelAdd: def test_passes_when_disabled_even_without_limits(self): assert ( _raise_if_rate_limits_required_but_missing( - litellm_params=LiteLLM_Params(model="azure/gpt-5.2"), - enforced=False, + model_name="gpt-5.2", litellm_params=LiteLLM_Params(model="azure/gpt-5.2"), policy=None ) is None ) @@ -7365,8 +7373,9 @@ class TestEnforceRpmTpmOnModelAdd: def test_passes_when_enabled_and_both_set(self): assert ( _raise_if_rate_limits_required_but_missing( + model_name="gpt-5.2", litellm_params=LiteLLM_Params(model="azure/gpt-5.2", rpm=10, tpm=1000), - enforced=True, + policy=REQUIRE_BOTH_LIMITS, ) is None ) @@ -7381,14 +7390,85 @@ class TestEnforceRpmTpmOnModelAdd: (LiteLLM_Params(model="azure/gpt-5.2", rpm=10, tpm=-1), "tpm"), ], ) - def test_raises_when_enabled_and_missing(self, params, expected_missing): - from litellm.proxy._types import ProxyException - + def test_raises_when_enabled_and_missing(self, params: LiteLLM_Params, expected_missing: str): with pytest.raises(ProxyException) as exc_info: - _raise_if_rate_limits_required_but_missing(litellm_params=params, enforced=True) - assert expected_missing in str(exc_info.value.message) + _raise_if_rate_limits_required_but_missing( + model_name="gpt-5.2", litellm_params=params, policy=REQUIRE_BOTH_LIMITS + ) + assert exc_info.value.message.startswith(f"{expected_missing} must be set") assert exc_info.value.code == "400" + @pytest.mark.parametrize( + ("model_name", "params"), + [ + ("runway-gen4", LiteLLM_Params(model="runwayml/gen4_turbo")), + ("my-runway", LiteLLM_Params(model="gen4_turbo", custom_llm_provider="runwayml")), + ("jev-latest", LiteLLM_Params(model="openai/jev-latest")), + ("veo-3.1-generate-001", LiteLLM_Params(model="vertex_ai/veo-3.1-generate-001", rpm=10)), + ("gemini-3-pro", LiteLLM_Params(model="vertex_ai/gemini-3-pro", rpm=10, tpm=1000)), + ], + ) + def test_policy_asks_each_model_only_for_the_limits_its_entry_requires( + self, model_name: str, params: LiteLLM_Params + ): + assert ( + _raise_if_rate_limits_required_but_missing( + model_name=model_name, litellm_params=params, policy=_CUSTOMER_MODEL_ADD_POLICY + ) + is None + ) + + @pytest.mark.parametrize( + ("model_name", "params", "expected_missing"), + [ + ("veo-3.1-generate-001", LiteLLM_Params(model="vertex_ai/veo-3.1-generate-001", tpm=1000), "rpm"), + ("gemini-3-pro", LiteLLM_Params(model="vertex_ai/gemini-3-pro", rpm=10), "tpm"), + ("openai-gpt", LiteLLM_Params(model="openai/gpt-5.2"), "rpm and tpm"), + ("runway-gen4", LiteLLM_Params(model="runwayml/gen4_turbo", rpm=0), "rpm"), + ], + ) + def test_policy_rejects_a_model_missing_a_required_or_carrying_a_non_positive_limit( + self, model_name: str, params: LiteLLM_Params, expected_missing: str + ): + with pytest.raises(ProxyException) as exc_info: + _raise_if_rate_limits_required_but_missing( + model_name=model_name, litellm_params=params, policy=_CUSTOMER_MODEL_ADD_POLICY + ) + assert exc_info.value.message.startswith(f"{expected_missing} must be set to a positive value for '{model_name}'") + assert exc_info.value.code == "400" + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + ("setting", "expected_code", "expected_message"), + [ + ({"default": [], "models": {"veo-*": ["rpm"]}}, "400", "rpm must be set to a positive value for 'veo-3.1'"), + ("yes", "400", "rpm and tpm must be set to a positive value for 'veo-3.1'"), + ( + {"providers": {"vertex_ai": ["qps"]}}, + "500", + "general_settings.enforce_rpm_tpm_on_model_add is invalid", + ), + ], +) +async def test_model_add_resolves_the_quota_policy_for_the_public_model_name( + monkeypatch: pytest.MonkeyPatch, setting: object, expected_code: str, expected_message: str +) -> None: + from litellm.proxy import proxy_server + from litellm.proxy.management_endpoints.model_management_endpoints import add_new_model + + prisma: Final = MagicMock() + monkeypatch.setattr(proxy_server, "prisma_client", prisma) + monkeypatch.setattr(proxy_server, "general_settings", {"enforce_rpm_tpm_on_model_add": setting}) + with pytest.raises(ProxyException) as error: + await add_new_model( + Deployment(model_name="veo-3.1", litellm_params={"model": "vertex_ai/veo-3.1"}), + UserAPIKeyAuth(user_id="admin", user_role=LitellmUserRoles.PROXY_ADMIN), + ) + assert error.value.code == expected_code + assert error.value.message.startswith(expected_message) + prisma.db.litellm_proxymodeltable.create.assert_not_called() + class TestBlockModelResponseSerialization: @pytest.mark.parametrize(