scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-10-29 06:33:28 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 200

feat(proxy): per-provider and per-model rpm/tpm requirements for enforce_project_model_quota (#47693)

Browse source 
* feat(proxy): per-provider and per-model rpm/tpm requirements for enforce_project_model_quota

* test(proxy): cover per-provider and per-model enforce_project_model_quota policies

* fix(proxy): keep enforce_project_model_quota reads typed

* fix(proxy): apply the quota policy to model add, follow router wildcard rules, keep old setting values

* fix(proxy): keep team model names off quota model entries, read the provider behind api_base, match patterns like the router

A team admin picks a team model's public name, so a name matching an exempt model entry no longer skips its provider's limits. A mapping with no policy key keeps requiring both limits as before, the project create permission check runs before the setting is read, and an unprefixed model behind a known OpenAI-compatible api_base resolves to that provider

* fix(proxy): re-check team projects on team model writes and validate quota provider ids

* fix(proxy): check added limits on model updates and match team models by the model they call

* fix(proxy): re-check team projects by every name a team model serves and check renamed names

* fix(proxy): follow the router's order for tied team patterns and annotate quota tests

* fix(proxy): match a team model's called model through JSON provider prefixes too

* fix(proxy): keep model updates that keep their route as on main

A PATCH or POST model update now checks only the limits its new route adds over the stored one, so removing or zeroing a limit on an unchanged route answers 200 again, under true and under a policy
This commit is contained in:
Mateo Wang 2026-10-28 05:24:41 -07:00 • committed by GitHub
parent ca24b396da
commit 0ceca31081
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
6 changed files with 1899 additions and 77 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

111
enterprise/litellm_enterprise/proxy/management_endpoints/project_endpoints.py
View file

@ -15,7 +15,11 @@ 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 (
describe_missing_quota,
missing_project_model_quota,
)
from pydantic import TypeAdapter, ValidationError
from litellm._logging import verbose_proxy_logger
from litellm._uuid import uuid
@ -27,6 +31,15 @@ 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 (
ENFORCE_PROJECT_MODEL_QUOTA_SETTING,
REQUIRE_BOTH_LIMITS,
ModelQuotaPolicy,
RouteOf,
no_route,
parse_model_quota_setting,
router_model_routes,
)
from litellm.proxy.management_helpers.utils import (
management_endpoint_wrapper,
)
@ -276,26 +289,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()
@ -322,64 +315,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]) -> ModelQuotaPolicy | None:
try:
return parse_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: ModelQuotaPolicy = REQUIRE_BOTH_LIMITS,
route_of: RouteOf = no_route,
) -> 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, route_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: ModelQuotaPolicy = REQUIRE_BOTH_LIMITS,
route_of: RouteOf = no_route,
) -> 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,
route_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."
},
)
@ -530,6 +538,7 @@ async def new_project(
"""
from litellm.proxy.proxy_server import (
general_settings,
general_settings_view,
litellm_proxy_admin_name,
llm_router,
premium_user,
@ -578,10 +587,6 @@ 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))
# Check if user has permission to create projects for this team
# only team admins can create projects for their team
has_permission = await _check_user_permission_for_project(
@ -600,6 +605,15 @@ async def new_project(
},
)
quota_policy: Final = _project_model_quota_policy(general_settings_view())
if quota_policy is not None:
_raise_on_missing_project_model_quota(
data,
_router_access_group_names(llm_router),
quota_policy,
router_model_routes(llm_router, data.team_id),
)
# Generate project_id if not provided
if data.project_id is None:
data.project_id = str(uuid.uuid4())
@ -721,6 +735,7 @@ async def update_project(
"""
from litellm.proxy.proxy_server import (
general_settings,
general_settings_view,
litellm_proxy_admin_name,
llm_router,
premium_user,
@ -828,10 +843,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_view())
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_model_routes(llm_router, target_team_id),
)
# Prepare update data

49
enterprise/litellm_enterprise/proxy/management_endpoints/project_model_quota.py Normal file
View file

@ -0,0 +1,49 @@
from collections.abc import Mapping, Sequence
from dataclasses import dataclass
from types import MappingProxyType
from typing import Final
from litellm.proxy.management_helpers.model_quota_policy import (
ModelQuotaPolicy,
QuotaLimit,
RouteOf,
required_limits,
unmet_limits,
)
_NO_LIMITS_SET: Final[Mapping[str, object]] = MappingProxyType({})
@dataclass(frozen=True, slots=True)
class MissingModelQuota:
model: str
limits: tuple[QuotaLimit, ...]
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: ModelQuotaPolicy,
route_of: RouteOf,
) -> tuple[MissingModelQuota, ...]:
rpm_set: Final = rpm_limits or _NO_LIMITS_SET
tpm_set: Final = tpm_limits or _NO_LIMITS_SET
gaps: Final = tuple(
MissingModelQuota(
model,
unmet_limits(required_limits(policy, model, route_of(model)), _limits_set_for(model, rpm_set, tpm_set)),
)
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)

296
litellm/proxy/management_endpoints/model_management_endpoints.py
View file

@ -12,6 +12,7 @@ model/{model_id}/update - PATCH endpoint for model update.
import asyncio
import datetime
import itertools
import json
from collections.abc import AsyncGenerator, Awaitable, Callable, Mapping, Sequence
from contextlib import AbstractAsyncContextManager, asynccontextmanager, suppress
@ -42,6 +43,7 @@ from litellm.litellm_core_utils.ptu_pricing import (
ptu_config_error,
)
from litellm.litellm_core_utils.sensitive_data_masker import redact_credentials_in_payload
from litellm.llms.azure.ptu_capacity import is_str_keyed_mapping
from litellm.proxy._types import (
BlockModelRequest,
CommonProxyErrors,
@ -102,6 +104,18 @@ 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 (
ENFORCE_PROJECT_MODEL_QUOTA_SETTING,
ModelQuotaPolicy,
QuotaLimit,
deployment_route,
exempts_any_route,
limits_added_by_team_write,
parse_model_quota_setting,
required_limits,
team_write_serves,
unmet_limits,
)
from litellm.proxy.spend_tracking.ptu_feature_flag import (
PTU_COST_ATTRIBUTION_ENV_VAR,
is_ptu_cost_attribution_enabled,
@ -111,6 +125,7 @@ from litellm.repositories.chunked_in import find_many_in
from litellm.repositories.credentials_repository import CredentialsRepository
from litellm.repositories.model_repository import ModelRepository
from litellm.repositories.prisma_protocols import TableActions
from litellm.repositories.project_repository import ProjectRepository
from litellm.repositories.table_repositories import ModelTableRepository
from litellm.repositories.team_repository import TeamRepository
from litellm.router import Router
@ -152,6 +167,7 @@ from litellm.types.router import (
SPECIAL_MODEL_INFO_PARAMS,
Deployment,
GenericLiteLLMParams,
LiteLLM_Params,
ModelInfo,
updateDeployment,
updateLiteLLMParams,
@ -762,26 +778,93 @@ 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.
_ROUTING_KEYS: Final = ("model", "custom_llm_provider", "api_base")
_RATE_LIMIT_KEYS: Final = ("rpm", "tpm")
_NOTHING_SENT: Final[Mapping[str, object]] = MappingProxyType({})
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 _merged_litellm_params(
incoming: GenericLiteLLMParams | None, existing: GenericLiteLLMParams | None, keys: tuple[str, ...]
) -> Mapping[str, object]:
sent: Final = _NOTHING_SENT if incoming is None else _dump_without_none(incoming)
stored: Final = _NOTHING_SENT if existing is None else _decrypted_litellm_params(existing)
return MappingProxyType({key: sent.get(key, stored.get(key)) for key in keys})
def _written_routing_params(
incoming: GenericLiteLLMParams | None, existing: GenericLiteLLMParams | None
) -> Mapping[str, object]:
return _merged_litellm_params(incoming, existing, _ROUTING_KEYS)
def _quota_policy(general_settings: Mapping[str, object], setting: str) -> ModelQuotaPolicy | None:
try:
return parse_model_quota_setting(general_settings.get(setting))
except ValidationError as e:
raise ProxyException(
message=f"general_settings.{setting} is invalid: {e}",
type=ProxyErrorTypes.internal_server_error.value,
code=status.HTTP_500_INTERNAL_SERVER_ERROR,
param=None,
) from e
def _raise_if_rate_limits_required_but_missing(
*, model_name: str, litellm_params: LiteLLM_Params, team_model: bool, 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
_raise_on_missing_rate_limits(
model_name,
unmet_limits(
required_limits(
policy, model_name, deployment_route(_written_routing_params(litellm_params, None), team_model)
),
MappingProxyType({"rpm": litellm_params.rpm, "tpm": litellm_params.tpm}),
),
)
def _raise_if_model_update_adds_missing_rate_limits(
*,
existing: Deployment,
written_name: str,
written_team_model: bool,
incoming_params: GenericLiteLLMParams | None,
remaining_params: GenericLiteLLMParams,
policy: ModelQuotaPolicy | None,
) -> None:
if policy is None:
return
before: Final = required_limits(
policy,
existing.model_info.team_public_model_name or existing.model_name,
deployment_route(
_written_routing_params(None, existing.litellm_params), existing.model_info.team_id is not None
),
)
after: Final = required_limits(
policy,
written_name,
deployment_route(_written_routing_params(incoming_params, remaining_params), written_team_model),
)
written_limits: Final = _merged_litellm_params(incoming_params, remaining_params, _RATE_LIMIT_KEYS)
checked: Final = after - before
_raise_on_missing_rate_limits(
written_name, unmet_limits(checked, MappingProxyType({limit: written_limits.get(limit) for limit in checked}))
)
def _raise_on_missing_rate_limits(model_name: str, missing: tuple[QuotaLimit, ...]) -> None:
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,
@ -790,6 +873,113 @@ def _raise_if_rate_limits_required_but_missing(*, litellm_params: GenericLiteLLM
)
def _project_limit(metadata: object, limits_key: str, model: str) -> object:
limits: Final = metadata.get(limits_key) if is_str_keyed_mapping(metadata) else None
return limits.get(model) if is_str_keyed_mapping(limits) else None
def _projects_short_on(
projects: Sequence["prisma_models.LiteLLM_ProjectTable"], model: str, added: frozenset[QuotaLimit]
) -> tuple[str, ...]:
return tuple(
project.project_alias or project.project_id
for project in projects
if model in project.models
and unmet_limits(
added,
MappingProxyType(
{
"rpm": _project_limit(project.metadata, "model_rpm_limit", model),
"tpm": _project_limit(project.metadata, "model_tpm_limit", model),
}
),
)
)
async def _raise_if_team_write_leaves_projects_short(
*,
prisma_client: PrismaClient,
general_settings: Mapping[str, object],
llm_router: Router | None,
team_id: str | None,
public_name: str,
replaced_id: str | None,
written_params: Mapping[str, object] | None,
) -> None:
if team_id is None:
return
policy: Final = _quota_policy(general_settings, ENFORCE_PROJECT_MODEL_QUOTA_SETTING)
if policy is None or not exempts_any_route(policy):
return
projects: Final = await ProjectRepository(prisma_client).table.find_many(where={"team_id": team_id})
listed: Final = frozenset(itertools.chain.from_iterable(project.models for project in projects))
served: Final = sorted(model for model in listed if team_write_serves(llm_router, team_id, public_name, model))
shortfall: Final = next(
(
(model, added, short)
for model in served
if (
added := limits_added_by_team_write(
policy, llm_router, team_id, public_name, replaced_id, written_params, model
)
)
and (short := _projects_short_on(projects, model, added))
),
None,
)
if shortfall is None:
return
short_model, short_limits, short_projects = shortfall
raise ProxyException(
message=(
f"projects {', '.join(short_projects)} list '{short_model}' without a positive "
f"{' and '.join(sorted(short_limits))} "
f"quota, which '{ENFORCE_PROJECT_MODEL_QUOTA_SETTING}' requires once this team model change applies. "
"Set those limits on the projects first"
),
type=ProxyErrorTypes.validation_error.value,
code=status.HTTP_400_BAD_REQUEST,
param=None,
)
async def _raise_if_team_update_leaves_projects_short(
*,
prisma_client: PrismaClient,
general_settings: Mapping[str, object],
llm_router: Router | None,
existing: Deployment,
written_team_id: str | None,
written_name: str,
incoming_params: GenericLiteLLMParams | None,
remaining_params: GenericLiteLLMParams,
) -> None:
if written_team_id is None and existing.model_info.team_id is None:
return
stored_name: Final = existing.model_info.team_public_model_name or existing.model_name
await _raise_if_team_write_leaves_projects_short(
prisma_client=prisma_client,
general_settings=general_settings,
llm_router=llm_router,
team_id=written_team_id,
public_name=written_name,
replaced_id=existing.model_info.id,
written_params=_written_routing_params(incoming_params, remaining_params),
)
if (written_team_id, written_name) == (existing.model_info.team_id, stored_name):
return
await _raise_if_team_write_leaves_projects_short(
prisma_client=prisma_client,
general_settings=general_settings,
llm_router=llm_router,
team_id=existing.model_info.team_id,
public_name=stored_name,
replaced_id=existing.model_info.id,
written_params=None,
)
_PTU_PRICED_PAIR: Final = frozenset({"ptu_count", "cost_per_ptu_per_hour"})
@ -1308,6 +1498,7 @@ async def patch_model(
ProxyException: For various error conditions including authentication and database errors
"""
from litellm.proxy.proxy_server import (
general_settings_view,
litellm_proxy_admin_name,
llm_router,
premium_user,
@ -1434,6 +1625,29 @@ async def patch_model(
existing_params=remaining_params,
)
written_team_id: Final = (
patch_data.model_info.team_id if patch_data.model_info else None
) or db_model.model_info.team_id
written_name: Final = _get_public_model_name(effective_patch, db_model)
_raise_if_model_update_adds_missing_rate_limits(
existing=db_model,
written_name=written_name,
written_team_model=written_team_id is not None,
incoming_params=incoming_params,
remaining_params=remaining_params,
policy=_quota_policy(general_settings_view(), ENFORCE_RPM_TPM_ON_MODEL_ADD_SETTING),
)
await _raise_if_team_update_leaves_projects_short(
prisma_client=prisma_client,
general_settings=general_settings_view(),
llm_router=llm_router,
existing=db_model,
written_team_id=written_team_id,
written_name=written_name,
incoming_params=incoming_params,
remaining_params=remaining_params,
)
effective_params: Final = _effective_complexity_router_params(incoming_params, remaining_params)
requested_model_name: Final = patch_data.model_name
stored_model_name: str | None = None
@ -2509,6 +2723,7 @@ async def delete_model(
from litellm.proxy.proxy_server import (
MODEL_RECONCILE_LOCK,
general_settings_view,
llm_router,
premium_user,
prisma_client,
@ -2542,6 +2757,15 @@ async def delete_model(
incoming_params=None,
allow_missing_team=True,
)
await _raise_if_team_write_leaves_projects_short(
prisma_client=prisma_client,
general_settings=general_settings_view(),
llm_router=llm_router,
team_id=model_params.model_info.team_id,
public_name=model_params.model_info.team_public_model_name or model_params.model_name,
replaced_id=model_params.model_info.id,
written_params=None,
)
# update DB
if store_model_in_db is True:
@ -2708,6 +2932,8 @@ async def add_new_model(
"""
from litellm.proxy.proxy_server import (
general_settings,
general_settings_view,
llm_router,
premium_user,
prisma_client,
proxy_config,
@ -2766,8 +2992,19 @@ 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)),
team_model=model_params.model_info.team_id is not None,
policy=_quota_policy(general_settings_view(), ENFORCE_RPM_TPM_ON_MODEL_ADD_SETTING),
)
await _raise_if_team_write_leaves_projects_short(
prisma_client=prisma_client,
general_settings=general_settings_view(),
llm_router=llm_router,
team_id=model_params.model_info.team_id,
public_name=model_params.model_name,
replaced_id=None,
written_params=_written_routing_params(model_params.litellm_params, None),
)
clean_model_info: Final = ModelInfo.model_validate(
@ -2906,6 +3143,7 @@ async def update_model(
"""
from litellm.proxy.proxy_server import (
LITELLM_PROXY_ADMIN_NAME,
general_settings_view,
llm_router,
premium_user,
prisma_client,
@ -2970,6 +3208,30 @@ async def update_model(
incoming_params=model_params.litellm_params,
existing_params=deployment.litellm_params,
)
renamed_to: Final = (
model_params.model_name
if model_params.model_name not in (None, deployment.model_name) and deployment.model_info.team_id is None
else None
)
written_name: Final = renamed_to or deployment.model_info.team_public_model_name or deployment.model_name
_raise_if_model_update_adds_missing_rate_limits(
existing=deployment,
written_name=written_name,
written_team_model=deployment.model_info.team_id is not None,
incoming_params=model_params.litellm_params,
remaining_params=deployment.litellm_params,
policy=_quota_policy(general_settings_view(), ENFORCE_RPM_TPM_ON_MODEL_ADD_SETTING),
)
await _raise_if_team_update_leaves_projects_short(
prisma_client=prisma_client,
general_settings=general_settings_view(),
llm_router=llm_router,
existing=deployment,
written_team_id=deployment.model_info.team_id,
written_name=written_name,
incoming_params=model_params.litellm_params,
remaining_params=deployment.litellm_params,
)
effective_params: Final = _effective_complexity_router_params(
model_params.litellm_params, deployment.litellm_params
)
@ -3009,12 +3271,6 @@ async def update_model(
if value is not None or _existing_litellm_params_dict.get(key) is not None
}
renamed_to: Final = (
model_params.model_name
if model_params.model_name not in (None, deployment.model_name)
and deployment.model_info.team_id is None
else None
)
base_update: Final[PrismaCompatibleUpdateDBModel] = {
"litellm_params": json.dumps(merged_dictionary),
"updated_by": user_api_key_dict.user_id or LITELLM_PROXY_ADMIN_NAME,

327
litellm/proxy/management_helpers/model_quota_policy.py Normal file
View file

@ -0,0 +1,327 @@
import re
from collections.abc import Callable, Mapping, Sequence
from dataclasses import dataclass
from types import MappingProxyType
from typing import TYPE_CHECKING, Final, Literal, TypeAlias
from pydantic import BaseModel, ConfigDict, Field, field_validator
import litellm
from litellm.exceptions import BadRequestError
from litellm.litellm_core_utils.get_llm_provider_logic import get_llm_provider, is_openai_compatible_endpoint
from litellm.llms.azure.ptu_capacity import is_str_keyed_mapping
from litellm.llms.openai_like.json_loader import JSONProviderRegistry
from litellm.router_utils.common_utils import provider_for_generic_call
from litellm.router_utils.pattern_match_deployments import PatternMatchRouter, PatternUtils
from litellm.secret_managers.main import get_secret_str
if TYPE_CHECKING:
from litellm import Router
QuotaLimit: TypeAlias = Literal["rpm", "tpm"]
ENFORCE_PROJECT_MODEL_QUOTA_SETTING: Final = "enforce_project_model_quota"
_LIMITS: Final[tuple[QuotaLimit, ...]] = ("rpm", "tpm")
_ROUTER_PATTERNS: Final = PatternMatchRouter()
def _is_known_provider(provider: str) -> bool:
return provider in litellm.provider_list or JSONProviderRegistry.exists(provider)
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)
@field_validator("providers")
@classmethod
def _known_providers(cls, providers: Mapping[str, frozenset[QuotaLimit]]) -> Mapping[str, frozenset[QuotaLimit]]:
unknown: Final = sorted(provider for provider in providers if not _is_known_provider(provider))
if unknown:
raise ValueError(f"unknown provider ids {unknown}")
return providers
REQUIRE_BOTH_LIMITS: Final = ModelQuotaPolicy()
_POLICY_KEYS: Final = frozenset(ModelQuotaPolicy.model_fields)
@dataclass(frozen=True, slots=True)
class TeamDeployment:
provider: str | None
called_model: str | None
@dataclass(frozen=True, slots=True)
class ModelRoute:
providers: frozenset[str] = frozenset()
team_deployments: frozenset[TeamDeployment] = frozenset()
RouteOf: TypeAlias = Callable[[str], ModelRoute]
NO_ROUTE: Final = ModelRoute()
def parse_model_quota_setting(setting: object) -> ModelQuotaPolicy | None:
"""Any falsy value is off and any other value requires both limits, as the boolean setting always
behaved, so only a mapping that names a policy key is read as a per-provider and per-model policy."""
if not setting:
return None
if is_str_keyed_mapping(setting) and not _POLICY_KEYS.isdisjoint(setting):
return ModelQuotaPolicy.model_validate(setting)
return REQUIRE_BOTH_LIMITS
def _wildcard_matches(pattern: str, model: str) -> bool:
return re.match(_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 _model_entry(entries: Mapping[str, frozenset[QuotaLimit]], model: str) -> frozenset[QuotaLimit] | None:
exact: Final = entries.get(model)
if exact is not None:
return exact
pattern: Final = _most_specific_model_pattern(model, entries)
return None if pattern is None else entries[pattern]
def _provider_limits(policy: ModelQuotaPolicy, provider: str | None) -> frozenset[QuotaLimit]:
return policy.default if provider is None else policy.providers.get(provider, policy.default)
def _team_deployment_limits(policy: ModelQuotaPolicy, deployment: TeamDeployment) -> frozenset[QuotaLimit]:
entry: Final = None if deployment.called_model is None else _model_entry(policy.models, deployment.called_model)
return _provider_limits(policy, deployment.provider) if entry is None else entry
def required_limits(policy: ModelQuotaPolicy, model: str, route: ModelRoute) -> 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. A team picks its model's public name, so a team deployment matches
model entries by the model it calls instead."""
entry: Final = None if route.team_deployments else _model_entry(policy.models, model)
if entry is not None:
return entry
if not route.providers and not route.team_deployments:
return policy.default
return frozenset[QuotaLimit]().union(
*(_team_deployment_limits(policy, deployment) for deployment in route.team_deployments),
*(_provider_limits(policy, provider) for provider in route.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_route(_model: str) -> ModelRoute:
return NO_ROUTE
def _resolved_api_base(api_base: object) -> object:
if isinstance(api_base, str) and api_base.startswith("os.environ/"):
return get_secret_str(api_base)
return api_base
def _provider_behind_api_base(litellm_params: Mapping[str, object]) -> str | None:
model: Final = litellm_params.get("model")
api_base: Final = _resolved_api_base(litellm_params.get("api_base"))
if litellm_params.get("custom_llm_provider") or not isinstance(model, str) or not isinstance(api_base, str):
return None
if model.split("/", 1)[0] in litellm.provider_list or not is_openai_compatible_endpoint(api_base):
return None
try:
_, provider, _, _ = get_llm_provider(model=model, api_base=api_base)
except BadRequestError:
return None
return provider
def deployment_provider(litellm_params: object) -> str | None:
if not is_str_keyed_mapping(litellm_params):
return None
return _provider_behind_api_base(litellm_params) or provider_for_generic_call(litellm_params)
def _called_model(litellm_params: Mapping[str, object]) -> str | None:
model: Final = litellm_params.get("model")
if not isinstance(model, str):
return None
prefix, _, rest = model.partition("/")
return rest if rest and _is_known_provider(prefix) else model
def team_deployment(litellm_params: object) -> TeamDeployment:
if not is_str_keyed_mapping(litellm_params):
return TeamDeployment(provider=None, called_model=None)
return TeamDeployment(provider=deployment_provider(litellm_params), called_model=_called_model(litellm_params))
def deployment_route(litellm_params: Mapping[str, object], team_model: bool) -> ModelRoute:
if team_model:
return ModelRoute(team_deployments=frozenset((team_deployment(litellm_params),)))
provider: Final = deployment_provider(litellm_params)
return ModelRoute(providers=frozenset() if provider is None else frozenset((provider,)))
def _named_deployments(llm_router: "Router", model: str, team_id: str | None) -> tuple[Mapping[str, object], ...]:
return (
*llm_router.get_all_deployments(model_name=model, team_id=team_id),
*llm_router.get_model_list_from_routing_groups(model_name=model),
)
def _routed_deployments(llm_router: "Router", model: str, team_id: str | None) -> Sequence[Mapping[str, object]]:
target: Final = llm_router.routable_model_group(model)
return _named_deployments(llm_router, target, team_id) or llm_router.wildcard_route_deployments(target, team_id)
def _is_team_deployment(deployment: Mapping[str, object]) -> bool:
model_info: Final = deployment.get("model_info")
return is_str_keyed_mapping(model_info) and model_info.get("team_id") is not None
def router_model_routes(llm_router: "Router | None", team_id: str | None) -> RouteOf:
if llm_router is None:
return no_route
def route_of(model: str) -> ModelRoute:
deployments: Final = _routed_deployments(llm_router, model, team_id)
shared: Final = tuple(deployment for deployment in deployments if not _is_team_deployment(deployment))
return ModelRoute(
providers=frozenset(
provider
for deployment in shared
if (provider := deployment_provider(deployment.get("litellm_params"))) is not None
),
team_deployments=frozenset(
team_deployment(deployment.get("litellm_params"))
for deployment in deployments
if _is_team_deployment(deployment)
),
)
return route_of
def _deployment_id(deployment: Mapping[str, object]) -> object:
model_info: Final = deployment.get("model_info")
return model_info.get("id") if is_str_keyed_mapping(model_info) else None
def _routable_name(llm_router: "Router | None", model: str) -> str:
return model if llm_router is None else llm_router.routable_model_group(model)
def _patterns_tried_before(patterns: Sequence[str], regex: str) -> Sequence[str]:
if regex in patterns:
return patterns[: patterns.index(regex)]
specificity: Final = PatternUtils.calculate_pattern_specificity(regex)
return tuple(rival for rival in patterns if PatternUtils.calculate_pattern_specificity(rival) >= specificity)
def _team_pattern_serves(llm_router: "Router", team_id: str, pattern: str, model: str) -> bool:
regex: Final = _ROUTER_PATTERNS.pattern_to_regex(pattern)
if re.match(regex, model) is None or _named_deployments(llm_router, model, team_id):
return False
team_router: Final = llm_router.team_pattern_routers.get(team_id)
patterns: Final = () if team_router is None else tuple(team_router.patterns)
return not any(re.match(rival, model) is not None for rival in _patterns_tried_before(patterns, regex))
def team_write_serves(llm_router: "Router | None", team_id: str, public_name: str, model: str) -> bool:
"""Whether a request from the team for ``model`` reaches the team route named ``public_name``: the name
itself or an alias of it, or, for a `*` name, a model no named deployment shadows and no team pattern the
router tries first takes."""
target: Final = _routable_name(llm_router, model)
if "*" not in public_name or llm_router is None:
return target == public_name
return _team_pattern_serves(llm_router, team_id, public_name, target)
def _served_params(litellm_params: object, public_name: str, model: str) -> object:
if "*" not in public_name or not is_str_keyed_mapping(litellm_params):
return litellm_params
matched: Final = re.match(_ROUTER_PATTERNS.pattern_to_regex(public_name), model)
called: Final = litellm_params.get("model")
if matched is None or not isinstance(called, str):
return litellm_params
return MappingProxyType(
{
**litellm_params,
"model": PatternMatchRouter.set_deployment_model_name(
matched_pattern=matched, litellm_deployment_litellm_model=called
),
}
)
def team_route_after_write(
llm_router: "Router | None",
team_id: str,
public_name: str,
replaced_id: str | None,
written_params: Mapping[str, object] | None,
model: str,
) -> ModelRoute | None:
deployments: Final = () if llm_router is None else llm_router.get_all_deployments(public_name, team_id=team_id)
kept: Final = tuple(
deployment.get("litellm_params")
for deployment in deployments
if _is_team_deployment(deployment) and _deployment_id(deployment) != replaced_id
)
team_params: Final = kept if written_params is None else (*kept, written_params)
target: Final = _routable_name(llm_router, model)
if team_params:
return ModelRoute(
team_deployments=frozenset(
team_deployment(_served_params(litellm_params, public_name, target)) for litellm_params in team_params
)
)
if llm_router is None or not _routed_deployments(llm_router, model, None):
return None
return router_model_routes(llm_router, None)(model)
def limits_added_by_team_write(
policy: ModelQuotaPolicy,
llm_router: "Router | None",
team_id: str,
public_name: str,
replaced_id: str | None,
written_params: Mapping[str, object] | None,
model: str,
) -> frozenset[QuotaLimit]:
after: Final = team_route_after_write(llm_router, team_id, public_name, replaced_id, written_params, model)
if after is None:
return frozenset()
now: Final = router_model_routes(llm_router, team_id)(model)
return required_limits(policy, model, after) - required_limits(policy, model, now)
def exempts_any_route(policy: ModelQuotaPolicy) -> bool:
return any(
limits != REQUIRE_BOTH_LIMITS.default
for limits in (policy.default, *policy.providers.values(), *policy.models.values())
)

562
tests/unit/enterprise/proxy/management_endpoints/test_project_endpoints_prisma.py
View file

@ -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
@ -30,6 +30,7 @@ from litellm_enterprise.proxy.management_endpoints.project_endpoints import (
delete_project,
project_info,
)
from litellm.proxy.management_helpers.model_quota_policy import ModelRoute, TeamDeployment, router_model_routes
from litellm.proxy.proxy_server import (
LitellmUserRoles,
)
@ -1205,7 +1206,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: pytest.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 +1216,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 +1607,559 @@ 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_route_of(model: str) -> ModelRoute:
return ModelRoute(providers=_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,
route_of: Callable[[str], ModelRoute] = _customer_route_of,
) -> str | None:
from litellm_enterprise.proxy.management_endpoints.project_endpoints import (
_raise_on_missing_project_model_quota,
)
from litellm.proxy.management_helpers.model_quota_policy import parse_model_quota_setting
policy: Final = parse_model_quota_setting(setting)
assert policy is not None
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, route_of)
except HTTPException as e:
rejection: Final = 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: Final = _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: Final = _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
@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: Final = _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)"),
({"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: object, model: str, expected_missing: str
):
error: Final = _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: Final = {"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) -> ModelRoute:
return ModelRoute(providers=frozenset({"runwayml", "vertex_ai"}))
error: Final = _quota_error(["video-pool"], rpm={"video-pool": 10}, route_of=mixed_providers)
assert error is not None
assert "video-pool (model_tpm_limit)" in error
def test_quota_policy_model_pattern_matches_from_the_start_like_the_router():
setting: Final = {"providers": {"vertex_ai": ["rpm", "tpm"]}, "models": {"*-lite": []}}
assert _quota_error(["veo-3.1-lite-generate-001"], setting=setting) is None
assert _quota_error(["gemini-3.8-flash"], setting=setting) is not None
def _team_vertex_route(called_model: str) -> Callable[[str], ModelRoute]:
def route_of(_model: str) -> ModelRoute:
return ModelRoute(team_deployments=frozenset({TeamDeployment(provider="vertex_ai", called_model=called_model)}))
return route_of
def test_quota_policy_matches_a_team_model_by_the_model_it_calls():
borrowed_name: Final = _quota_error(["jev-latest"], route_of=_team_vertex_route("gemini-3.8-flash"))
veo_call: Final = _quota_error(
["video-team"], rpm={"video-team": 5}, route_of=_team_vertex_route("veo-3.1-generate-001")
)
assert borrowed_name is not None and "jev-latest (model_rpm_limit and model_tpm_limit)" in borrowed_name
assert veo_call is None
def test_quota_policy_team_model_beside_a_shared_one_needs_both_routes_limits():
def team_vertex_beside_shared_runway(_model: str) -> ModelRoute:
return ModelRoute(
providers=frozenset({"runwayml"}),
team_deployments=frozenset({TeamDeployment(provider="vertex_ai", called_model="gemini-3.8-flash")}),
)
error: Final = _quota_error(["media-latest"], rpm={"media-latest": 5}, route_of=team_vertex_beside_shared_runway)
assert error is not None and "media-latest (model_tpm_limit)" in error
@pytest.mark.parametrize("setting", [True, "true", "false", 1, ["rpm"], {"enabled": True}, {"provider": {"runwayml": []}}])
def test_quota_setting_truthy_non_policy_value_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, {}, "", 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_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.proxy.management_helpers.model_quota_policy import parse_model_quota_setting
policy: Final = parse_model_quota_setting(_CUSTOMER_QUOTA_POLICY)
assert policy is not None
existing: Final = 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_route_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_route_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"},
},
]
)
def test_quota_providers_come_from_the_one_wildcard_route_requests_use():
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*"},
},
]
)
assert router_model_routes(router, "media-team")("gen4_turbo") == ModelRoute(
team_deployments=frozenset({TeamDeployment(provider="vertex_ai", called_model="4_turbo")})
)
assert router_model_routes(router, None)("gen4_turbo") == ModelRoute(providers=frozenset({"runwayml"}))
def test_quota_providers_read_the_provider_behind_a_known_api_base():
groq_api_base: Final = "https://api.groq.com/openai/v1"
router: Final = litellm.Router(
model_list=[
{
"model_name": "qa-llama",
"litellm_params": {"model": "qa-llama", "api_base": groq_api_base, "api_key": "qa-groq-key"},
},
{
"model_name": "qa-runway",
"litellm_params": {"model": "gen4_turbo", "custom_llm_provider": "runwayml", "api_base": groq_api_base},
},
]
)
route_of: Final = router_model_routes(router, None)
assert route_of("qa-llama").providers == frozenset({"groq"})
assert route_of("qa-runway").providers == frozenset({"runwayml"})
def test_quota_policy_rejects_an_unknown_provider_id():
from pydantic import ValidationError
from litellm.proxy.management_helpers.model_quota_policy import parse_model_quota_setting
with pytest.raises(ValidationError, match=r"unknown provider ids \['runway'\]"):
parse_model_quota_setting({"providers": {"vertex_ai": ["rpm"], "runway": []}})
def test_quota_policy_accepts_a_provider_registered_at_runtime(monkeypatch: pytest.MonkeyPatch):
from litellm.proxy.management_helpers.model_quota_policy import parse_model_quota_setting
monkeypatch.setattr(litellm, "provider_list", [*litellm.provider_list, "qa-custom-llm"])
policy: Final = parse_model_quota_setting({"providers": {"qa-custom-llm": []}})
assert policy is not None
assert policy.providers["qa-custom-llm"] == frozenset()
_TEAM_WRITE_POLICY: Final = {"providers": {"vertex_ai": ["rpm", "tpm"], "runwayml": []}, "models": {"jev-latest": []}}
_VERTEX_PARAMS: Final = {"model": "vertex_ai/gemini-3.8-flash", "vertex_project": "qa-project"}
_RUNWAY_PARAMS: Final = {"model": "gen4_turbo", "custom_llm_provider": "runwayml"}
def _team_deployment(public_name: str, model_id: str, litellm_params: Mapping[str, object]) -> dict[str, object]:
return {
"model_name": f"model_name_media-team_{model_id}",
"litellm_params": dict(litellm_params),
"model_info": {"id": model_id, "team_id": "media-team", "team_public_model_name": public_name},
}
@pytest.mark.parametrize(
("setting", "public_name", "proxy_params", "team_routes", "replaced_id", "written_params", "added"),
[
pytest.param(
_TEAM_WRITE_POLICY,
"jev-latest",
{"model": "openai/jev-latest"},
(),
None,
_VERTEX_PARAMS,
{"rpm", "tpm"},
id="team-add-on-vertex-behind-an-exempt-proxy-name",
),
pytest.param(
_TEAM_WRITE_POLICY,
"jev-latest",
{"model": "openai/jev-latest"},
(),
None,
_RUNWAY_PARAMS,
set(),
id="team-add-on-an-exempt-provider",
),
pytest.param(
_TEAM_WRITE_POLICY,
"jev-latest",
{"model": "openai/jev-latest"},
(),
None,
{"model": "aihubmix/jev-latest"},
set(),
id="team-add-through-a-json-provider-matches-the-model-it-calls",
),
pytest.param(
_TEAM_WRITE_POLICY,
"media-latest",
_VERTEX_PARAMS,
(("t1", _RUNWAY_PARAMS),),
"t1",
None,
{"rpm", "tpm"},
id="team-delete-falls-back-to-the-proxy-vertex-route",
),
pytest.param(
_TEAM_WRITE_POLICY,
"media-latest",
_VERTEX_PARAMS,
(("t1", _RUNWAY_PARAMS), ("t2", _RUNWAY_PARAMS)),
"t1",
None,
set(),
id="team-delete-keeps-a-sibling-team-route",
),
pytest.param(
_TEAM_WRITE_POLICY,
"media-latest",
_VERTEX_PARAMS,
(("t1", _RUNWAY_PARAMS),),
"t1",
_VERTEX_PARAMS,
{"rpm", "tpm"},
id="team-update-moves-the-route-to-vertex",
),
pytest.param(
True,
"jev-latest",
{"model": "openai/jev-latest"},
(),
None,
_VERTEX_PARAMS,
set(),
id="true-already-requires-both-limits",
),
],
)
def test_team_model_write_adds_only_the_limits_its_new_route_requires(
setting: object,
public_name: str,
proxy_params: Mapping[str, object],
team_routes: tuple[tuple[str, Mapping[str, object]], ...],
replaced_id: str | None,
written_params: Mapping[str, object] | None,
added: set[str],
):
from litellm.proxy.management_helpers.model_quota_policy import (
limits_added_by_team_write,
parse_model_quota_setting,
)
policy: Final = parse_model_quota_setting(setting)
router: Final = litellm.Router(
model_list=[
{"model_name": public_name, "litellm_params": dict(proxy_params)},
*(_team_deployment(public_name, model_id, params) for model_id, params in team_routes),
]
)
assert policy is not None
assert limits_added_by_team_write(
policy, router, "media-team", public_name, replaced_id, written_params, public_name
) == frozenset(added)
def test_quota_route_of_an_alias_reads_the_team_deployments_its_target_routes_to():
router: Final = litellm.Router(
model_list=[
{"model_name": "media-latest", "litellm_params": dict(_RUNWAY_PARAMS)},
_team_deployment("media-latest", "t1", _VERTEX_PARAMS),
],
model_group_alias={"media": "media-latest"},
)
route_of: Final = router_model_routes(router, "media-team")
assert route_of("media") == route_of("media-latest")
assert TeamDeployment(provider="vertex_ai", called_model="gemini-3.8-flash") in route_of("media").team_deployments
def _team_pattern_router() -> "litellm.Router":
return litellm.Router(
model_list=[
{"model_name": "media-latest", "litellm_params": dict(_VERTEX_PARAMS)},
{"model_name": "video-named", "litellm_params": dict(_VERTEX_PARAMS)},
_team_deployment("video-*", "t1", {"model": "runwayml/*"}),
_team_deployment("video-veo-*", "t2", {"model": "vertex_ai/veo-*", "vertex_project": "qa-project"}),
_team_deployment("clip-*x", "t3", {"model": "runwayml/*"}),
_team_deployment("clip-x*", "t4", {"model": "vertex_ai/*", "vertex_project": "qa-project"}),
],
model_group_alias={"media": "media-latest"},
)
@pytest.mark.parametrize(
("public_name", "model", "served"),
[
pytest.param("media-latest", "media-latest", True, id="the-name-itself"),
pytest.param("media-latest", "media", True, id="an-alias-of-the-name"),
pytest.param("media-latest", "video-named", False, id="another-name"),
pytest.param("video-*", "video-gen-1", True, id="a-name-the-pattern-matches"),
pytest.param("video-*", "video-veo-3", False, id="a-more-specific-team-pattern-takes-it"),
pytest.param("video-veo-*", "video-veo-3", True, id="the-most-specific-team-pattern"),
pytest.param("video-*", "video-named", False, id="a-named-deployment-shadows-the-pattern"),
pytest.param("video-*", "image-gen-1", False, id="a-name-the-pattern-misses"),
pytest.param("audio-*", "audio-gen-1", True, id="a-pattern-the-write-adds"),
pytest.param("clip-*x", "clip-xx", True, id="the-first-registered-of-two-tied-team-patterns"),
pytest.param("clip-x*", "clip-xx", False, id="a-tied-team-pattern-registered-first-takes-it"),
pytest.param("cli*-xx", "clip-xx", False, id="a-new-pattern-tied-with-a-registered-one"),
],
)
def test_team_write_serves_the_project_names_a_team_request_routes_through_it(
public_name: str, model: str, served: bool
):
from litellm.proxy.management_helpers.model_quota_policy import team_write_serves
assert team_write_serves(_team_pattern_router(), "media-team", public_name, model) is served
@pytest.mark.parametrize(
("model", "added"),
[
pytest.param("video-gen-1", {"rpm", "tpm"}, id="vertex-requires-both"),
pytest.param("video-veo-3", {"rpm"}, id="the-called-veo-model-entry-wins"),
],
)
def test_team_pattern_write_reads_each_served_name_by_the_model_it_calls(model: str, added: set[str]):
from litellm.proxy.management_helpers.model_quota_policy import (
limits_added_by_team_write,
parse_model_quota_setting,
)
policy: Final = parse_model_quota_setting(
{"default": [], "providers": {"vertex_ai": ["rpm", "tpm"]}, "models": {"veo-*": ["rpm"]}}
)
assert policy is not None
assert limits_added_by_team_write(
policy,
litellm.Router(model_list=[]),
"media-team",
"video-*",
None,
{"model": "vertex_ai/*", "vertex_project": "qa-project"},
model,
) == frozenset(added)
@pytest.mark.asyncio
async def test_new_project_quota_policy_reads_providers_from_router_deployments(monkeypatch: pytest.MonkeyPatch):
team_models: Final = ["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: pytest.MonkeyPatch):
_enforced_new_project_mocks(
monkeypatch,
team_models=["runway-gen4-turbo"],
llm_router=_customer_router(),
quota_setting={"providers": {"runwayml": []}, "model": {"jev-latest": []}},
)
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)
@pytest.mark.asyncio
async def test_new_project_checks_permission_before_reading_the_quota_setting(monkeypatch: pytest.MonkeyPatch):
from litellm_enterprise.proxy.management_endpoints import project_endpoints as pe
_enforced_new_project_mocks(
monkeypatch,
team_models=["runway-gen4-turbo"],
llm_router=_customer_router(),
quota_setting={"providers": {"runwayml": []}, "model": {"jev-latest": []}},
)
monkeypatch.setattr(pe, "_check_user_permission_for_project", mock.AsyncMock(return_value=False))
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 == "403"
assert "enforce_project_model_quota" not in str(exc_info.value.message)
@pytest.mark.asyncio
async def test_new_project_quota_policy_ignores_model_entries_for_a_team_model_name(monkeypatch: pytest.MonkeyPatch):
router: Final = litellm.Router(
model_list=[
{
"model_name": "model_name_test-team_qa",
"litellm_params": {"model": "vertex_ai/gemini-3.8-flash", "vertex_project": "qa-project"},
"model_info": {"team_id": "test-team", "team_public_model_name": "jev-latest"},
},
]
)
_enforced_new_project_mocks(
monkeypatch, team_models=["jev-latest"], llm_router=router, quota_setting=_CUSTOMER_QUOTA_POLICY
)
with pytest.raises(ProxyException) as exc_info:
await _run_new_project(NewProjectRequest(team_id="test-team", models=["jev-latest"]))
assert exc_info.value.code == "400"
assert "jev-latest (model_rpm_limit and model_tpm_limit)" in str(exc_info.value.message)
_ACCESS_NOW: Final = datetime.now(timezone.utc)

631
tests/unit/proxy/management_endpoints/test_model_management_endpoints.py
View file

@ -15,6 +15,7 @@ from fastapi.testclient import TestClient
from litellm._uuid import uuid
from litellm.models.credentials import CredentialItem
from litellm.models.project import LiteLLM_ProjectTable
from litellm.proxy._types import (
LiteLLM_ModelTable,
@ -31,6 +32,7 @@ from litellm.proxy.management_endpoints.model_management_endpoints import (
KubernetesPodsResponse,
ModelManagementAuthChecks,
_get_team_deployments,
_raise_if_model_update_adds_missing_rate_limits,
_raise_if_rate_limits_required_but_missing,
clear_cache,
delete_team_models,
@ -38,6 +40,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
@ -7369,12 +7372,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", team_model=False, litellm_params=LiteLLM_Params(model="azure/gpt-5.2"), policy=None
)
is None
)
@ -7382,8 +7392,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", team_model=False,
litellm_params=LiteLLM_Params(model="azure/gpt-5.2", rpm=10, tpm=1000),
enforced=True,
policy=REQUIRE_BOTH_LIMITS,
)
is None
)
@ -7398,14 +7409,618 @@ 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", team_model=False, 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, team_model=False, 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, team_model=False, 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.parametrize("api_base", ["https://api.groq.com/openai/v1", "os.environ/QA_GROQ_API_BASE"])
def test_policy_reads_the_provider_behind_a_known_api_base(self, monkeypatch: pytest.MonkeyPatch, api_base: str):
monkeypatch.setenv("QA_GROQ_API_BASE", "https://api.groq.com/openai/v1")
policy: Final = ModelQuotaPolicy.model_validate({"default": [], "providers": {"groq": ["rpm", "tpm"]}})
with pytest.raises(ProxyException) as exc_info:
_raise_if_rate_limits_required_but_missing(
model_name="qa-llama", team_model=False,
litellm_params=LiteLLM_Params(model="qa-llama", api_base=api_base),
policy=policy,
)
assert exc_info.value.message.startswith("rpm and tpm must be set to a positive value for 'qa-llama'")
@pytest.mark.parametrize(
("model_name", "params", "shared_missing", "team_missing"),
[
pytest.param(
"jev-latest", LiteLLM_Params(model="vertex_ai/gemini-3.8-flash"), None, "rpm and tpm",
id="team-cannot-borrow-an-exempt-public-name",
),
pytest.param(
"video-team", LiteLLM_Params(model="vertex_ai/veo-3.1-generate-001", rpm=5), "tpm", None,
id="team-matches-the-model-it-calls",
),
],
)
def test_policy_matches_a_team_model_by_the_model_it_calls(
self, model_name: str, params: LiteLLM_Params, shared_missing: str | None, team_missing: str | None
):
assert _model_add_error(model_name, params, team_model=False) == _missing_limits_message(shared_missing, model_name)
assert _model_add_error(model_name, params, team_model=True) == _missing_limits_message(team_missing, model_name)
def _missing_limits_message(missing: str | None, model_name: str) -> str | None:
if missing is None:
return None
return (
f"{missing} must be set to a positive value for '{model_name}' when "
"'enforce_rpm_tpm_on_model_add' is enabled in general_settings"
)
def _model_add_error(model_name: str, params: LiteLLM_Params, team_model: bool) -> str | None:
try:
_raise_if_rate_limits_required_but_missing(
model_name=model_name, team_model=team_model, litellm_params=params, policy=_CUSTOMER_MODEL_ADD_POLICY
)
except ProxyException as error:
return error.message
return None
_STORED_RUNWAY: Final = {"model": "runwayml/gen4_turbo"}
def _stored_model(litellm_params: Mapping[str, object], team_id: str | None) -> Deployment:
return Deployment(
model_name="media-latest" if team_id is None else "model_name_media-team_team-row",
litellm_params=dict(litellm_params),
model_info={
"id": "team-row",
"team_id": team_id,
"team_public_model_name": None if team_id is None else "media-latest",
},
)
@pytest.mark.parametrize(
("stored", "team_id", "incoming", "missing"),
[
pytest.param(
_STORED_RUNWAY, "media-team", updateLiteLLMParams(model="vertex_ai/gemini-3.8-flash"), "rpm and tpm",
id="team-model-moves-to-a-provider-that-needs-limits",
),
pytest.param(
_STORED_RUNWAY, None, updateLiteLLMParams(model="vertex_ai/gemini-3.8-flash", rpm=10), "tpm",
id="proxy-model-moves-with-only-one-limit",
),
pytest.param(
{**_STORED_RUNWAY, "tpm": 1000}, None, updateLiteLLMParams(model="vertex_ai/gemini-3.8-flash", rpm=10),
None, id="stored-limit-counts-toward-the-move",
),
pytest.param(
_STORED_RUNWAY, "media-team", updateLiteLLMParams(model="vertex_ai/veo-3.1-generate-001", rpm=5), None,
id="team-model-moves-to-a-model-entry-it-meets",
),
pytest.param(
{"model": "vertex_ai/gemini-3.8-flash"}, None, updateLiteLLMParams(api_key="qa-rotated-key"), None,
id="unchanged-route-adds-no-limits",
),
pytest.param(
{"model": "vertex_ai/gemini-3.8-flash", "rpm": 10, "tpm": 1000}, None, updateLiteLLMParams(rpm=0), None,
id="unchanged-route-can-zero-a-limit-as-on-main",
),
pytest.param(
{"model": "vertex_ai/gemini-3.8-flash"}, None, updateLiteLLMParams(rpm=10), None,
id="legacy-model-without-tpm-stays-editable",
),
],
)
def test_model_update_requires_only_the_limits_its_new_route_adds(
stored: Mapping[str, object], team_id: str | None, incoming: updateLiteLLMParams, missing: str | None
) -> None:
assert _model_update_error(_stored_model(stored, team_id), incoming) == _missing_limits_message(
missing, "media-latest"
)
def _model_update_error(
existing: Deployment,
incoming: updateLiteLLMParams,
policy: ModelQuotaPolicy = _CUSTOMER_MODEL_ADD_POLICY,
remaining: LiteLLM_Params | None = None,
) -> str | None:
try:
_raise_if_model_update_adds_missing_rate_limits(
existing=existing,
written_name="media-latest",
written_team_model=existing.model_info.team_id is not None,
incoming_params=incoming,
remaining_params=existing.litellm_params if remaining is None else remaining,
policy=policy,
)
except ProxyException as error:
return error.message
return None
@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'"),
({"enabled": True}, "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()
@pytest.mark.asyncio
async def test_team_model_add_cannot_borrow_a_model_entry_by_its_public_name(monkeypatch: pytest.MonkeyPatch) -> None:
from litellm.proxy import proxy_server
from litellm.proxy.management_endpoints.model_management_endpoints import add_new_model
prisma: Final = MagicMock()
prisma.db.litellm_teamtable.find_unique = AsyncMock(return_value=LiteLLM_TeamTable(team_id="media-team"))
monkeypatch.setattr(proxy_server, "prisma_client", prisma)
monkeypatch.setattr(proxy_server, "premium_user", True)
monkeypatch.setattr(
proxy_server,
"general_settings",
{"enforce_rpm_tpm_on_model_add": {"providers": {"vertex_ai": ["rpm", "tpm"]}, "models": {"jev-latest": []}}},
)
with pytest.raises(ProxyException) as error:
await add_new_model(
Deployment(
model_name="jev-latest",
litellm_params={"model": "vertex_ai/gemini-3.8-flash"},
model_info={"team_id": "media-team"},
),
UserAPIKeyAuth(user_id="admin", user_role=LitellmUserRoles.PROXY_ADMIN),
)
assert error.value.code == "400"
assert error.value.message.startswith("rpm and tpm must be set to a positive value for 'jev-latest'")
prisma.db.litellm_proxymodeltable.create.assert_not_called()
_PROJECT_QUOTA_SETTINGS: Final = {
"enforce_project_model_quota": {
"providers": {"vertex_ai": ["rpm", "tpm"], "runwayml": []},
"models": {"jev-latest": []},
}
}
_TEAM_VERTEX: Final = {"model": "vertex_ai/gemini-3.8-flash", "vertex_project": "qa-project"}
_TEAM_RUNWAY: Final = {"model": "gen4_turbo", "custom_llm_provider": "runwayml"}
def _media_team_db(public_name: str, project_metadata: Mapping[str, object]) -> MagicMock:
prisma: Final = MagicMock()
prisma.db.litellm_teamtable.find_unique = AsyncMock(return_value=LiteLLM_TeamTable(team_id="media-team"))
prisma.db.litellm_projecttable.find_many = AsyncMock(
return_value=[
LiteLLM_ProjectTable(
project_id="p-1",
project_alias="video-gen",
team_id="media-team",
models=[public_name],
metadata=dict(project_metadata),
),
LiteLLM_ProjectTable(
project_id="p-2", project_alias="chat-only", team_id="media-team", models=["chat-only-model"]
),
]
)
return prisma
def _team_model_row(public_name: str, litellm_params: Mapping[str, object]) -> LiteLLM_ProxyModelTable:
return LiteLLM_ProxyModelTable(
model_id="team-row",
model_name="model_name_media-team_team-row",
litellm_params=dict(litellm_params),
model_info={"id": "team-row", "team_id": "media-team", "team_public_model_name": public_name},
created_by="admin",
updated_by="admin",
)
def _quota_router(public_name: str, proxy_params: Mapping[str, object], *team_params: Mapping[str, object]) -> Router:
return Router(
model_list=[
{"model_name": public_name, "litellm_params": dict(proxy_params)},
*(
{
"model_name": "model_name_media-team_team-row",
"litellm_params": dict(params),
"model_info": {"id": "team-row", "team_id": "media-team", "team_public_model_name": public_name},
}
for params in team_params
),
]
)
def _use_quota_proxy(monkeypatch: pytest.MonkeyPatch, prisma: MagicMock, router: Router) -> None:
from litellm.proxy import proxy_server
monkeypatch.setattr(proxy_server, "prisma_client", prisma)
monkeypatch.setattr(proxy_server, "general_settings", _PROJECT_QUOTA_SETTINGS)
monkeypatch.setattr(proxy_server, "llm_router", router)
monkeypatch.setattr(proxy_server, "store_model_in_db", False)
monkeypatch.setattr(proxy_server, "premium_user", True)
_ADMIN: Final = UserAPIKeyAuth(user_id="admin", user_role=LitellmUserRoles.PROXY_ADMIN)
@pytest.mark.asyncio
async def test_team_model_add_rejects_a_route_its_projects_lack_the_quota_for(monkeypatch: pytest.MonkeyPatch) -> None:
from litellm.proxy.management_endpoints.model_management_endpoints import add_new_model
prisma: Final = _media_team_db("jev-latest", {})
_use_quota_proxy(monkeypatch, prisma, _quota_router("jev-latest", {"model": "openai/jev-latest"}))
with pytest.raises(ProxyException) as error:
await add_new_model(
Deployment(
model_name="jev-latest",
litellm_params={**_TEAM_VERTEX, "rpm": 10, "tpm": 1000},
model_info={"team_id": "media-team"},
),
_ADMIN,
)
assert error.value.code == "400"
assert error.value.message.startswith(
"projects video-gen list 'jev-latest' without a positive rpm and tpm quota, "
"which 'enforce_project_model_quota' requires once this team model change applies"
)
prisma.db.litellm_projecttable.find_many.assert_awaited_once_with(where={"team_id": "media-team"})
@pytest.mark.parametrize(
("litellm_params", "project_metadata"),
[
pytest.param(
_TEAM_VERTEX,
{"model_rpm_limit": {"jev-latest": 5}, "model_tpm_limit": {"jev-latest": 500}},
id="projects-carry-the-added-limits",
),
pytest.param(_TEAM_RUNWAY, {}, id="new-route-is-exempt"),
],
)
@pytest.mark.asyncio
async def test_team_model_add_goes_through_when_its_projects_meet_the_new_route(
monkeypatch: pytest.MonkeyPatch, litellm_params: Mapping[str, object], project_metadata: Mapping[str, object]
) -> None:
from litellm.proxy.management_endpoints.model_management_endpoints import add_new_model
_use_quota_proxy(
monkeypatch, _media_team_db("jev-latest", project_metadata), _quota_router("jev-latest", {"model": "openai/jev-latest"})
)
with pytest.raises(ProxyException) as error:
await add_new_model(
Deployment(model_name="jev-latest", litellm_params=dict(litellm_params), model_info={"team_id": "media-team"}),
_ADMIN,
)
assert "STORE_MODEL_IN_DB" in error.value.message
@pytest.mark.parametrize("endpoint", ["patch", "put"])
@pytest.mark.asyncio
async def test_team_model_update_rejects_moving_its_route_where_its_projects_lack_quota(
monkeypatch: pytest.MonkeyPatch, endpoint: str
) -> None:
from litellm.proxy import proxy_server
runway_by_prefix: Final = {"model": "runwayml/gen4_turbo"}
prisma: Final = _media_team_db("jev-latest", {})
prisma.db.litellm_proxymodeltable.find_unique = AsyncMock(
return_value=_team_model_row("jev-latest", runway_by_prefix)
)
_use_quota_proxy(monkeypatch, prisma, _quota_router("jev-latest", {"model": "openai/jev-latest"}, runway_by_prefix))
monkeypatch.setattr(proxy_server, "store_model_in_db", True)
moved_to_vertex: Final = updateLiteLLMParams(model="vertex_ai/gemini-3.8-flash", vertex_project="qa-project")
write: Final = (
patch_model(
model_id="team-row", patch_data=updateDeployment(litellm_params=moved_to_vertex), user_api_key_dict=_ADMIN
)
if endpoint == "patch"
else update_model(
model_params=updateDeployment(litellm_params=moved_to_vertex, model_info=ModelInfo(id="team-row")),
user_api_key_dict=_ADMIN,
)
)
with pytest.raises(ProxyException) as error:
await write
assert error.value.code == "400"
assert error.value.message.startswith("projects video-gen list 'jev-latest' without a positive rpm and tpm quota")
def _team_route_router(public_name: str, team_params: Mapping[str, object]) -> Router:
return Router(
model_list=[
{"model_name": "jev-latest", "litellm_params": {"model": "openai/jev-latest"}},
{
"model_name": "model_name_media-team_team-row",
"litellm_params": dict(team_params),
"model_info": {"id": "team-row", "team_id": "media-team", "team_public_model_name": public_name},
},
],
model_group_alias={"media": "jev-latest"},
)
@pytest.mark.parametrize(
("public_name", "project_model", "team_params", "moved_to"),
[
pytest.param(
"video-*",
"video-gen-1",
{"model": "runwayml/*"},
updateLiteLLMParams(model="vertex_ai/*", vertex_project="qa-project"),
id="a-name-the-team-wildcard-serves",
),
pytest.param(
"jev-latest",
"media",
{"model": "runwayml/gen4_turbo"},
updateLiteLLMParams(model="vertex_ai/gemini-3.8-flash", vertex_project="qa-project"),
id="an-alias-of-the-team-name",
),
],
)
@pytest.mark.asyncio
async def test_team_model_patch_rechecks_projects_listing_any_name_it_serves(
monkeypatch: pytest.MonkeyPatch,
public_name: str,
project_model: str,
team_params: Mapping[str, object],
moved_to: updateLiteLLMParams,
) -> None:
from litellm.proxy import proxy_server
prisma: Final = _media_team_db(project_model, {})
prisma.db.litellm_proxymodeltable.find_unique = AsyncMock(return_value=_team_model_row(public_name, team_params))
_use_quota_proxy(monkeypatch, prisma, _team_route_router(public_name, team_params))
monkeypatch.setattr(proxy_server, "store_model_in_db", True)
with pytest.raises(ProxyException) as error:
await patch_model(
model_id="team-row", patch_data=updateDeployment(litellm_params=moved_to), user_api_key_dict=_ADMIN
)
assert error.value.code == "400"
assert error.value.message.startswith(
f"projects video-gen list '{project_model}' without a positive rpm and tpm quota"
)
prisma.db.litellm_proxymodeltable.update.assert_not_called()
@pytest.mark.asyncio
async def test_model_update_checks_the_limits_of_the_name_it_renames_to(monkeypatch: pytest.MonkeyPatch) -> None:
from litellm.proxy import proxy_server
prisma: Final = MagicMock()
prisma.db.litellm_proxymodeltable.find_unique = AsyncMock(
return_value=LiteLLM_ProxyModelTable(
model_id="vertex-row",
model_name="jev-latest",
litellm_params={"model": "vertex_ai/gemini-3.8-flash"},
model_info={"id": "vertex-row"},
created_by="admin",
updated_by="admin",
)
)
_use_quota_proxy(monkeypatch, prisma, Router(model_list=[]))
monkeypatch.setattr(
proxy_server,
"general_settings",
{"enforce_rpm_tpm_on_model_add": {"providers": {"vertex_ai": ["rpm", "tpm"]}, "models": {"jev-latest": []}}},
)
monkeypatch.setattr(proxy_server, "store_model_in_db", True)
with pytest.raises(ProxyException) as error:
await update_model(
model_params=updateDeployment(
model_name="gemini-3.8-flash", litellm_params=updateLiteLLMParams(), model_info=ModelInfo(id="vertex-row")
),
user_api_key_dict=_ADMIN,
)
assert error.value.code == "400"
assert error.value.message.startswith("rpm and tpm must be set to a positive value for 'gemini-3.8-flash'")
prisma.db.litellm_proxymodeltable.update.assert_not_called()
@pytest.mark.parametrize(
("incoming", "remaining"),
[
pytest.param(updateLiteLLMParams(rpm=0), {"model": "vertex_ai/gemini-3.8-flash", "tpm": 1000}, id="zeroes-rpm"),
pytest.param(updateLiteLLMParams(), {"model": "vertex_ai/gemini-3.8-flash", "tpm": 1000}, id="removes-rpm"),
pytest.param(updateLiteLLMParams(model="runwayml/gen4_turbo"), {"model": "runwayml/gen4_turbo"}, id="moves"),
],
)
def test_model_update_under_the_boolean_setting_never_rejects_as_on_main(
incoming: updateLiteLLMParams, remaining: Mapping[str, object]
) -> None:
from litellm.proxy.management_helpers.model_quota_policy import parse_model_quota_setting
stored: Final = _stored_model({"model": "vertex_ai/gemini-3.8-flash", "rpm": 10, "tpm": 1000}, None)
policy: Final = parse_model_quota_setting(True)
assert policy is not None
assert _model_update_error(stored, incoming, policy, LiteLLM_Params(**remaining)) is None
@pytest.mark.asyncio
async def test_team_model_patch_rejects_a_removal_that_moves_its_route_where_its_projects_lack_quota(
monkeypatch: pytest.MonkeyPatch,
) -> None:
from litellm.proxy import proxy_server
from litellm.types.proxy.management_endpoints.model_management_endpoints import PatchModelRequest
prisma: Final = _media_team_db("jev-latest", {})
prisma.db.litellm_proxymodeltable.find_unique = AsyncMock(return_value=_team_model_row("jev-latest", _TEAM_RUNWAY))
_use_quota_proxy(monkeypatch, prisma, _quota_router("jev-latest", {"model": "openai/jev-latest"}, _TEAM_RUNWAY))
monkeypatch.setattr(proxy_server, "store_model_in_db", True)
with pytest.raises(ProxyException) as error:
await patch_model(
model_id="team-row",
patch_data=PatchModelRequest(litellm_params_to_remove=("custom_llm_provider",)),
user_api_key_dict=_ADMIN,
)
assert error.value.code == "400"
assert error.value.message.startswith("projects video-gen list 'jev-latest' without a positive rpm and tpm quota")
@pytest.mark.asyncio
async def test_team_model_delete_rejects_falling_back_to_a_route_its_projects_lack_quota_for(
monkeypatch: pytest.MonkeyPatch,
) -> None:
from litellm.proxy.management_endpoints.model_management_endpoints import ModelInfoDelete, delete_model
prisma: Final = _media_team_db("media-latest", {})
prisma.db.litellm_proxymodeltable.find_unique = AsyncMock(return_value=_team_model_row("media-latest", _TEAM_RUNWAY))
_use_quota_proxy(monkeypatch, prisma, _quota_router("media-latest", _TEAM_VERTEX, _TEAM_RUNWAY))
with pytest.raises(ProxyException) as error:
await delete_model(model_info=ModelInfoDelete(id="team-row"), user_api_key_dict=_ADMIN)
assert error.value.code == "400"
assert error.value.message.startswith("projects video-gen list 'media-latest' without a positive rpm and tpm quota")
prisma.db.litellm_proxymodeltable.delete.assert_not_called()
@pytest.mark.asyncio
async def test_team_model_delete_goes_through_when_nothing_else_serves_its_name(
monkeypatch: pytest.MonkeyPatch,
) -> None:
from litellm.proxy.management_endpoints.model_management_endpoints import ModelInfoDelete, delete_model
prisma: Final = _media_team_db("media-latest", {})
prisma.db.litellm_proxymodeltable.find_unique = AsyncMock(return_value=_team_model_row("media-latest", _TEAM_RUNWAY))
only_the_team_model: Final = Router(
model_list=[
{
"model_name": "model_name_media-team_team-row",
"litellm_params": dict(_TEAM_RUNWAY),
"model_info": {"id": "team-row", "team_id": "media-team", "team_public_model_name": "media-latest"},
}
]
)
_use_quota_proxy(monkeypatch, prisma, only_the_team_model)
with pytest.raises(ProxyException) as error:
await delete_model(model_info=ModelInfoDelete(id="team-row"), user_api_key_dict=_ADMIN)
assert "STORE_MODEL_IN_DB" in str(error.value.message)
@pytest.mark.parametrize("endpoint", ["patch", "put"])
@pytest.mark.asyncio
async def test_team_model_update_rejects_moving_to_a_provider_that_needs_limits_it_lacks(
monkeypatch: pytest.MonkeyPatch, endpoint: str
) -> None:
from litellm.proxy import proxy_server
prisma: Final = _media_team_db("media-latest", {})
prisma.db.litellm_proxymodeltable.find_unique = AsyncMock(return_value=_team_model_row("media-latest", _TEAM_RUNWAY))
_use_quota_proxy(monkeypatch, prisma, _quota_router("media-latest", _TEAM_RUNWAY, _TEAM_RUNWAY))
monkeypatch.setattr(proxy_server, "store_model_in_db", True)
monkeypatch.setattr(
proxy_server,
"general_settings",
{"enforce_rpm_tpm_on_model_add": {"providers": {"vertex_ai": ["rpm", "tpm"], "runwayml": []}}},
)
moved_to_vertex: Final = updateLiteLLMParams(model="vertex_ai/gemini-3.8-flash", custom_llm_provider="vertex_ai")
write: Final = (
patch_model(
model_id="team-row", patch_data=updateDeployment(litellm_params=moved_to_vertex), user_api_key_dict=_ADMIN
)
if endpoint == "patch"
else update_model(
model_params=updateDeployment(litellm_params=moved_to_vertex, model_info=ModelInfo(id="team-row")),
user_api_key_dict=_ADMIN,
)
)
with pytest.raises(ProxyException) as error:
await write
assert error.value.code == "400"
assert error.value.message == _missing_limits_message("rpm and tpm", "media-latest")
prisma.db.litellm_proxymodeltable.update.assert_not_called()
class TestBlockModelResponseSerialization:
@pytest.mark.parametrize(