From ba71f280d4785319fcbd6037c5c1d66aef979ea3 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Tue, 20 Oct 2026 17:19:43 -0700 Subject: [PATCH] feat(proxy): per-provider and per-model rpm/tpm requirements for enforce_project_model_quota --- .../management_endpoints/project_endpoints.py | 105 +++++---- .../project_model_quota.py | 125 +++++++++++ .../test_project_endpoints_prisma.py | 212 +++++++++++++++++- 3 files changed, 395 insertions(+), 47 deletions(-) create mode 100644 enterprise/litellm_enterprise/proxy/management_endpoints/project_model_quota.py diff --git a/enterprise/litellm_enterprise/proxy/management_endpoints/project_endpoints.py b/enterprise/litellm_enterprise/proxy/management_endpoints/project_endpoints.py index db419fde1e6..8bca681ea12 100644 --- a/enterprise/litellm_enterprise/proxy/management_endpoints/project_endpoints.py +++ b/enterprise/litellm_enterprise/proxy/management_endpoints/project_endpoints.py @@ -15,7 +15,18 @@ from collections.abc import Mapping, Sequence from typing import TYPE_CHECKING, Final, cast from fastapi import APIRouter, Depends, HTTPException, Request -from pydantic import TypeAdapter +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 from litellm._logging import verbose_proxy_logger from litellm._uuid import uuid @@ -265,26 +276,6 @@ def _check_team_project_limits( ) -def _project_models_missing_positive_quota( - models: list[str] | None, - rpm_limits: Mapping[str, object] | None, - tpm_limits: Mapping[str, object] | None, -) -> list[str]: - """Return the models that lack a positive `rpm` AND `tpm` quota. - - A valid quota is a positive integer; null, zero, and negative are rejected - because downstream rate limiters treat a non-positive limit as immediately - exhausted (every request blocked). - """ - - def _is_positive(value: object) -> bool: - return isinstance(value, int) and not isinstance(value, bool) and value > 0 - - rpm = rpm_limits or {} - tpm = tpm_limits or {} - return [model for model in (models or []) if not _is_positive(rpm.get(model)) or not _is_positive(tpm.get(model))] - - def _router_access_group_names(llm_router: "Router | None") -> frozenset[str]: return frozenset(llm_router.get_model_access_groups()) if llm_router is not None else frozenset() @@ -311,64 +302,79 @@ def _raise_on_project_models_expanding_at_request_time( raise HTTPException( status_code=400, detail={ - "error": f"models {list(expanding)} expand to multiple models at request time, so a per-model rpm/tpm quota cannot be enforced for them while 'enforce_project_model_quota' is enabled. List concrete model names instead." + "error": f"models {list(expanding)} expand to multiple models at request time, so a per-model rpm/tpm quota cannot be enforced for them while '{ENFORCE_PROJECT_MODEL_QUOTA_SETTING}' is enabled. List concrete model names instead." }, ) +def _project_model_quota_policy(general_settings: Mapping[str, object]) -> ProjectModelQuotaPolicy | None: + try: + return parse_project_model_quota_setting(general_settings.get(ENFORCE_PROJECT_MODEL_QUOTA_SETTING)) + except ValidationError as e: + raise HTTPException( + status_code=500, + detail={"error": f"general_settings.{ENFORCE_PROJECT_MODEL_QUOTA_SETTING} is invalid: {e}"}, + ) from e + + def _raise_on_missing_project_model_quota( - data: NewProjectRequest | UpdateProjectRequest, access_group_names: frozenset[str] = frozenset() + data: NewProjectRequest | UpdateProjectRequest, + access_group_names: frozenset[str] = frozenset(), + policy: ProjectModelQuotaPolicy = REQUIRE_BOTH_LIMITS, + providers_of: ProvidersOf = no_providers, ) -> None: - """Require a positive `rpm`/`tpm` quota for every model on project CREATE. + """Require the positive `rpm`/`tpm` quotas the policy asks of each model on project CREATE. `model_rpm_limit`/`model_tpm_limit` are relocated into `metadata` by the request model's `set_model_info` validator, so they are read from there. - - Only invoked when `general_settings.enforce_project_model_quota` is enabled - (default off), so it is opt-in and does not change behavior for existing users. """ _raise_on_project_models_expanding_at_request_time(data.models, access_group_names) metadata = data.metadata or {} - missing = _project_models_missing_positive_quota( - data.models, metadata.get("model_rpm_limit"), metadata.get("model_tpm_limit") + missing: Final = missing_project_model_quota( + data.models or (), metadata.get("model_rpm_limit"), metadata.get("model_tpm_limit"), policy, providers_of ) if not missing: return raise HTTPException( status_code=400, detail={ - "error": f"models {missing} added to project without a positive rpm/tpm quota. Set a positive model_rpm_limit and model_tpm_limit for each model." + "error": f"models added to project without a positive rpm/tpm quota: {describe_missing_quota(missing)}. Set a positive value for each listed limit." }, ) def _raise_on_missing_project_model_quota_on_update( - data: UpdateProjectRequest, existing_project: object, access_group_names: frozenset[str] = frozenset() + data: UpdateProjectRequest, + existing_project: object, + access_group_names: frozenset[str] = frozenset(), + policy: ProjectModelQuotaPolicy = REQUIRE_BOTH_LIMITS, + providers_of: ProvidersOf = no_providers, ) -> None: - """Require a positive `rpm`/`tpm` quota over the RESULTING state on project UPDATE. + """Require the positive `rpm`/`tpm` quotas the policy asks of each model over the RESULTING state on project UPDATE. `/project/update` replaces `models` and `metadata` when they are provided, so the check runs on what the project WILL look like: a partial update that doesn't touch models/quota keeps the existing values, while one that adds a model or clears a - model's quota must leave every resulting model with a positive limit. - - Only invoked when `general_settings.enforce_project_model_quota` is enabled - (default off), so it is opt-in and does not change behavior for existing users. + model's quota must leave every resulting model with the limits it needs. """ resulting_models = data.models if data.models is not None else (getattr(existing_project, "models", None) or []) resulting_metadata = ( data.metadata if data.metadata is not None else (getattr(existing_project, "metadata", None) or {}) ) _raise_on_project_models_expanding_at_request_time(resulting_models, access_group_names) - missing = _project_models_missing_positive_quota( - resulting_models, resulting_metadata.get("model_rpm_limit"), resulting_metadata.get("model_tpm_limit") + missing: Final = missing_project_model_quota( + resulting_models, + resulting_metadata.get("model_rpm_limit"), + resulting_metadata.get("model_tpm_limit"), + policy, + providers_of, ) if not missing: return raise HTTPException( status_code=400, detail={ - "error": f"models {missing} would be left on the project without a positive rpm/tpm quota. Set a positive model_rpm_limit and model_tpm_limit for each model." + "error": f"models would be left on the project without a positive rpm/tpm quota: {describe_missing_quota(missing)}. Set a positive value for each listed limit." }, ) @@ -567,9 +573,14 @@ async def new_project( data=data, ) - # Opt-in (default off): require rpm/tpm for every model added to the project. - if general_settings.get("enforce_project_model_quota", False): - _raise_on_missing_project_model_quota(data, _router_access_group_names(llm_router)) + quota_policy: Final = _project_model_quota_policy(general_settings) + if quota_policy is not None: + _raise_on_missing_project_model_quota( + data, + _router_access_group_names(llm_router), + quota_policy, + router_deployment_providers(llm_router, data.team_id), + ) # Check if user has permission to create projects for this team # only team admins can create projects for their team @@ -817,10 +828,14 @@ async def update_project( data=data, ) - # Opt-in (default off): require rpm/tpm for every model the update would leave on the project. - if general_settings.get("enforce_project_model_quota", False): + quota_policy: Final = _project_model_quota_policy(general_settings) + if quota_policy is not None: _raise_on_missing_project_model_quota_on_update( - data, existing_project, _router_access_group_names(llm_router) + data, + existing_project, + _router_access_group_names(llm_router), + quota_policy, + router_deployment_providers(llm_router, target_team_id), ) # Prepare update data diff --git a/enterprise/litellm_enterprise/proxy/management_endpoints/project_model_quota.py b/enterprise/litellm_enterprise/proxy/management_endpoints/project_model_quota.py new file mode 100644 index 00000000000..5cdfe158fd1 --- /dev/null +++ b/enterprise/litellm_enterprise/proxy/management_endpoints/project_model_quota.py @@ -0,0 +1,125 @@ +import re +from collections.abc import Callable, Mapping, Sequence +from dataclasses import dataclass +from itertools import chain +from types import MappingProxyType +from typing import TYPE_CHECKING, Final, Literal, TypeAlias + +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]] + +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) + + +@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(chain.from_iterable(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 missing_project_model_quota( + models: Sequence[str], + rpm_limits: Mapping[str, object] | None, + tpm_limits: Mapping[str, object] | None, + policy: ProjectModelQuotaPolicy, + 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} + ) + gaps: Final = tuple( + MissingModelQuota(model, _unset_limits(model, required_limits(policy, model, providers_of(model)), set_limits)) + for model in models + ) + return tuple(gap for gap in gaps if gap.limits) + + +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/tests/unit/enterprise/proxy/management_endpoints/test_project_endpoints_prisma.py b/tests/unit/enterprise/proxy/management_endpoints/test_project_endpoints_prisma.py index 45697723af1..6a33959e634 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 @@ -1205,7 +1205,9 @@ def test_update_quota_existing_quotaless_model_rejected(): _raise_on_missing_project_model_quota_on_update(data, existing) -def _enforced_new_project_mocks(monkeypatch, team_models: list[str], llm_router: mock.MagicMock | None) -> None: +def _enforced_new_project_mocks( + monkeypatch, team_models: list[str], llm_router: object, quota_setting: object = True +) -> None: from litellm.proxy._types import LiteLLM_TeamTable from litellm_enterprise.proxy.management_endpoints import project_endpoints as pe @@ -1213,7 +1215,7 @@ def _enforced_new_project_mocks(monkeypatch, team_models: list[str], llm_router: monkeypatch.setattr(litellm.proxy.proxy_server, "prisma_client", mock.MagicMock()) monkeypatch.setattr(litellm.proxy.proxy_server, "premium_user", True) monkeypatch.setattr(litellm.proxy.proxy_server, "llm_router", llm_router) - monkeypatch.setattr(litellm.proxy.proxy_server, "general_settings", {"enforce_project_model_quota": True}) + monkeypatch.setattr(litellm.proxy.proxy_server, "general_settings", {"enforce_project_model_quota": quota_setting}) monkeypatch.setattr(pe, "_validate_team_exists", mock.AsyncMock(return_value=team)) monkeypatch.setattr(pe, "_check_user_permission_for_project", mock.AsyncMock(return_value=True)) @@ -1604,6 +1606,212 @@ async def test_new_project_flag_on_access_group_model_returns_400(monkeypatch): assert "expand to multiple models at request time" in str(exc_info.value) +_CUSTOMER_QUOTA_POLICY: Final = { + "providers": {"vertex_ai": ["rpm", "tpm"], "runwayml": []}, + "models": {"veo-*": ["rpm"], "jev-latest": []}, +} + +_CUSTOMER_MODEL_PROVIDERS: Final = { + "gemini-3.8-flash": frozenset({"vertex_ai"}), + "veo-3.1-lite-generate-001": frozenset({"vertex_ai"}), + "runway-gen4-turbo": frozenset({"runwayml"}), +} + + +def _customer_providers_of(model: str) -> frozenset[str]: + return _CUSTOMER_MODEL_PROVIDERS.get(model, frozenset()) + + +def _quota_error( + models: list[str], + rpm: dict[str, int] | None = None, + tpm: dict[str, int] | None = None, + setting: object = _CUSTOMER_QUOTA_POLICY, + providers_of=_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, + ) + + policy = parse_project_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) + try: + _raise_on_missing_project_model_quota(data, frozenset(), policy, providers_of) + except HTTPException as e: + rejection = e + else: + return None + assert rejection.status_code == 400 + return str(rejection.detail) + + +def test_quota_policy_exempts_provider_and_passthrough_name_but_keeps_vertex_strict(): + assert _quota_error(["runway-gen4-turbo", "jev-latest"]) is None + error = _quota_error(["gemini-3.8-flash"], rpm={"gemini-3.8-flash": 10}) + assert error is not None + assert "gemini-3.8-flash (model_tpm_limit)" in error + + +def test_quota_policy_model_pattern_overrides_its_provider_per_limit(): + assert _quota_error(["veo-3.1-lite-generate-001"], rpm={"veo-3.1-lite-generate-001": 5}) is None + error = _quota_error(["veo-3.1-lite-generate-001"], tpm={"veo-3.1-lite-generate-001": 5000}) + assert error is not None + assert "veo-3.1-lite-generate-001 (model_rpm_limit)" in error + assert "model_tpm_limit" not 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 + assert "gemini-3.8-flash (model_rpm_limit and model_tpm_limit)" in error + assert "veo-3.1-lite-generate-001 (model_rpm_limit)" in error + assert "runway-gen4-turbo" not in error + + +@pytest.mark.parametrize( + "setting, model, expected_missing", + [ + ({"default": ["tpm"]}, "unrouted-name", "unrouted-name (model_tpm_limit)"), + ({"default": [], "providers": {"runwayml": ["rpm"]}}, "runway-gen4-turbo", "(model_rpm_limit)"), + ( + {"providers": {"runwayml": []}, "models": {"runway-gen4-turbo": ["tpm"]}}, + "runway-gen4-turbo", + "(model_tpm_limit)", + ), + ({"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)"), + ], +) +def test_quota_policy_precedence_is_model_then_pattern_then_provider_then_default(setting, model, expected_missing): + error = _quota_error([model], setting=setting) + assert error is not None + assert expected_missing in error + + +def test_quota_policy_model_entry_matches_names_literally_apart_from_star(): + setting = {"models": {"claude-opus-5-5[1m]": [], "gpt.5": []}} + assert _quota_error(["claude-opus-5-5[1m]"], setting=setting) is None + assert _quota_error(["claude-opus-5-51"], setting=setting) is not None + assert _quota_error(["gptx5"], setting=setting) is not None + + +def test_quota_policy_model_on_several_providers_needs_every_required_limit(): + def mixed_providers(_model: str) -> frozenset[str]: + return frozenset({"runwayml", "vertex_ai"}) + + error = _quota_error(["video-pool"], rpm={"video-pool": 10}, providers_of=mixed_providers) + assert error is not None + 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) + 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, + ) + + assert parse_project_model_quota_setting(setting) is None + + +def test_update_quota_policy_resolves_providers_for_the_resulting_models(): + import types + + 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, + ) + + policy = parse_project_model_quota_setting(_CUSTOMER_QUOTA_POLICY) + assert policy is not None + existing = types.SimpleNamespace(models=["runway-gen4-turbo"], metadata={}) + assert ( + _raise_on_missing_project_model_quota_on_update( + UpdateProjectRequest(project_id="p", description="unrelated change"), + existing, + frozenset(), + policy, + _customer_providers_of, + ) + is None + ) + with pytest.raises(HTTPException) as exc_info: + _raise_on_missing_project_model_quota_on_update( + UpdateProjectRequest(project_id="p", models=["runway-gen4-turbo", "gemini-3.8-flash"]), + existing, + frozenset(), + policy, + _customer_providers_of, + ) + assert "gemini-3.8-flash (model_rpm_limit and model_tpm_limit)" in str(exc_info.value.detail) + + +def _customer_router() -> "litellm.Router": + return litellm.Router( + model_list=[ + { + "model_name": "gemini-3.8-flash", + "litellm_params": {"model": "vertex_ai/gemini-3.8-flash", "vertex_project": "qa-project"}, + }, + { + "model_name": "veo-3.1-lite-generate-001", + "litellm_params": {"model": "vertex_ai/veo-3.1-lite-generate-001", "vertex_project": "qa-project"}, + }, + { + "model_name": "runway-gen4-turbo", + "litellm_params": {"model": "runwayml/gen4_turbo", "api_key": "qa-runway-key"}, + }, + ] + ) + + +@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"] + _enforced_new_project_mocks( + monkeypatch, team_models=team_models, llm_router=_customer_router(), quota_setting=_CUSTOMER_QUOTA_POLICY + ) + + with pytest.raises(ProxyException) as exc_info: + await _run_new_project(NewProjectRequest(team_id="test-team", models=team_models)) + + assert exc_info.value.code == "400" + assert ( + "gemini-3.8-flash (model_rpm_limit and model_tpm_limit), veo-3.1-lite-generate-001 (model_rpm_limit)." + in str(exc_info.value.message) + ) + assert "runway-gen4-turbo" not in str(exc_info.value.message) + assert "jev-latest" not in str(exc_info.value.message) + + +@pytest.mark.asyncio +async def test_new_project_rejects_malformed_quota_setting_instead_of_skipping_it(monkeypatch): + _enforced_new_project_mocks( + monkeypatch, + team_models=["runway-gen4-turbo"], + llm_router=_customer_router(), + quota_setting={"provider": {"runwayml": []}}, + ) + + with pytest.raises(ProxyException) as exc_info: + await _run_new_project(NewProjectRequest(team_id="test-team", models=["runway-gen4-turbo"])) + + assert exc_info.value.code == "500" + assert "general_settings.enforce_project_model_quota is invalid" in str(exc_info.value.message) + + _ACCESS_NOW: Final = datetime.now(timezone.utc)