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
This commit is contained in:
mateo-berri 2026-10-20 20:29:40 -07:00
parent 7f04cbe781
commit acacc486a8
6 changed files with 265 additions and 72 deletions

View file

@ -35,10 +35,10 @@ from litellm.proxy.management_endpoints.team_admin_field_permissions import team
from litellm.proxy.management_helpers.model_quota_policy import (
REQUIRE_BOTH_LIMITS,
ModelQuotaPolicy,
ProvidersOf,
no_providers,
RouteOf,
no_route,
parse_model_quota_setting,
router_deployment_providers,
router_model_routes,
)
from litellm.proxy.management_helpers.utils import (
management_endpoint_wrapper,
@ -323,7 +323,7 @@ def _raise_on_missing_project_model_quota(
data: NewProjectRequest | UpdateProjectRequest,
access_group_names: frozenset[str] = frozenset(),
policy: ModelQuotaPolicy = REQUIRE_BOTH_LIMITS,
providers_of: ProvidersOf = no_providers,
route_of: RouteOf = no_route,
) -> None:
"""Require the positive `rpm`/`tpm` quotas the policy asks of each model on project CREATE.
@ -333,7 +333,7 @@ def _raise_on_missing_project_model_quota(
_raise_on_project_models_expanding_at_request_time(data.models, access_group_names)
metadata = data.metadata or {}
missing: Final = missing_project_model_quota(
data.models or (), metadata.get("model_rpm_limit"), metadata.get("model_tpm_limit"), policy, providers_of
data.models or (), metadata.get("model_rpm_limit"), metadata.get("model_tpm_limit"), policy, route_of
)
if not missing:
return
@ -350,7 +350,7 @@ def _raise_on_missing_project_model_quota_on_update(
existing_project: object,
access_group_names: frozenset[str] = frozenset(),
policy: ModelQuotaPolicy = REQUIRE_BOTH_LIMITS,
providers_of: ProvidersOf = no_providers,
route_of: RouteOf = no_route,
) -> None:
"""Require the positive `rpm`/`tpm` quotas the policy asks of each model over the RESULTING state on project UPDATE.
@ -369,7 +369,7 @@ def _raise_on_missing_project_model_quota_on_update(
resulting_metadata.get("model_rpm_limit"),
resulting_metadata.get("model_tpm_limit"),
policy,
providers_of,
route_of,
)
if not missing:
return
@ -576,15 +576,6 @@ async def new_project(
data=data,
)
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_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
has_permission = await _check_user_permission_for_project(
@ -603,6 +594,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())
@ -839,7 +839,7 @@ async def update_project(
existing_project,
_router_access_group_names(llm_router),
quota_policy,
router_deployment_providers(llm_router, target_team_id),
router_model_routes(llm_router, target_team_id),
)
# Prepare update data

View file

@ -5,8 +5,8 @@ from typing import Final
from litellm.proxy.management_helpers.model_quota_policy import (
ModelQuotaPolicy,
ProvidersOf,
QuotaLimit,
RouteOf,
required_limits,
unmet_limits,
)
@ -32,14 +32,14 @@ def missing_project_model_quota(
rpm_limits: Mapping[str, object] | None,
tpm_limits: Mapping[str, object] | None,
policy: ModelQuotaPolicy,
providers_of: ProvidersOf,
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, providers_of(model)), _limits_set_for(model, rpm_set, tpm_set)),
unmet_limits(required_limits(policy, model, route_of(model)), _limits_set_for(model, rpm_set, tpm_set)),
)
for model in models
)

View file

@ -104,6 +104,7 @@ from litellm.proxy.management_helpers.auto_router_permissions import (
from litellm.proxy.management_helpers.model_allowlist_rename_sync import sync_model_allowlists_for_renamed_model
from litellm.proxy.management_helpers.model_quota_policy import (
ModelQuotaPolicy,
ModelRoute,
deployment_provider,
parse_model_quota_setting,
required_limits,
@ -781,7 +782,7 @@ def _model_add_quota_policy(general_settings: Mapping[str, object]) -> ModelQuot
def _raise_if_rate_limits_required_but_missing(
*, model_name: str, litellm_params: LiteLLM_Params, policy: ModelQuotaPolicy | None
*, 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,
@ -790,10 +791,20 @@ def _raise_if_rate_limits_required_but_missing(
if policy is None:
return
provider: Final = deployment_provider(
MappingProxyType({"model": litellm_params.model, "custom_llm_provider": litellm_params.custom_llm_provider})
MappingProxyType(
{
"model": litellm_params.model,
"custom_llm_provider": litellm_params.custom_llm_provider,
"api_base": litellm_params.api_base,
}
)
)
missing: Final = unmet_limits(
required_limits(policy, model_name, frozenset() if provider is None else frozenset((provider,))),
required_limits(
policy,
model_name,
ModelRoute(providers=frozenset() if provider is None else frozenset((provider,)), team_named=team_model),
),
MappingProxyType({"rpm": litellm_params.rpm, "tpm": litellm_params.tpm}),
)
if not missing:
@ -2740,6 +2751,7 @@ async def add_new_model(
_raise_if_rate_limits_required_but_missing(
model_name=model_params.model_name,
litellm_params=model_params.litellm_params,
team_model=model_params.model_info.team_id is not None,
policy=_model_add_quota_policy(general_settings),
)

View file

@ -1,9 +1,13 @@
import re
from collections.abc import Callable, Mapping, Sequence
from dataclasses import dataclass
from typing import TYPE_CHECKING, Final, Literal, TypeAlias
from pydantic import BaseModel, ConfigDict, Field
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.router_utils.common_utils import provider_for_generic_call
from litellm.router_utils.pattern_match_deployments import PatternMatchRouter, PatternUtils
@ -12,7 +16,6 @@ if TYPE_CHECKING:
from litellm import Router
QuotaLimit: TypeAlias = Literal["rpm", "tpm"]
ProvidersOf: TypeAlias = Callable[[str], frozenset[str]]
_LIMITS: Final[tuple[QuotaLimit, ...]] = ("rpm", "tpm")
_ROUTER_PATTERNS: Final = PatternMatchRouter()
@ -27,20 +30,31 @@ class ModelQuotaPolicy(BaseModel):
REQUIRE_BOTH_LIMITS: Final = ModelQuotaPolicy()
_POLICY_KEYS: Final = frozenset(ModelQuotaPolicy.model_fields)
@dataclass(frozen=True, slots=True)
class ModelRoute:
providers: frozenset[str] = frozenset()
team_named: bool = False
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 non-mapping value requires both limits, as the boolean
setting always behaved, so only a mapping is read as a per-provider and per-model policy."""
"""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):
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.fullmatch(_ROUTER_PATTERNS.pattern_to_regex(pattern), model) is not None
return re.match(_ROUTER_PATTERNS.pattern_to_regex(pattern), model) is not None
def _router_specificity(pattern: str) -> tuple[int, int]:
@ -52,19 +66,26 @@ def _most_specific_model_pattern(model: str, entries: Mapping[str, frozenset[Quo
return max(matching, key=_router_specificity, default=None)
def required_limits(policy: ModelQuotaPolicy, model: str, providers: frozenset[str]) -> frozenset[QuotaLimit]:
"""A model entry (exact name, then the most specific `*` pattern) wins over the providers of the
model's deployments, which win over the default. A model served by several providers needs every
limit any of them requires."""
exact: Final = policy.models.get(model)
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, policy.models)
if pattern is not None:
return policy.models[pattern]
if not providers:
pattern: Final = _most_specific_model_pattern(model, entries)
return None if pattern is None else entries[pattern]
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 model's name is picked by its team, so model entries never match it."""
entry: Final = None if route.team_named else _model_entry(policy.models, model)
if entry is not None:
return entry
if not route.providers:
return policy.default
return frozenset[QuotaLimit]().union(*(policy.providers.get(provider, policy.default) for provider in providers))
return frozenset[QuotaLimit]().union(
*(policy.providers.get(provider, policy.default) for provider in route.providers)
)
def _is_positive(value: object) -> bool:
@ -79,12 +100,28 @@ def unmet_limits(required: frozenset[QuotaLimit], values: Mapping[QuotaLimit, ob
)
def no_providers(_model: str) -> frozenset[str]:
return frozenset()
def no_route(_model: str) -> ModelRoute:
return NO_ROUTE
def _provider_behind_api_base(litellm_params: Mapping[str, object]) -> str | None:
model: Final = litellm_params.get("model")
api_base: Final = 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:
return provider_for_generic_call(litellm_params) if is_str_keyed_mapping(litellm_params) else 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 _routed_deployments(llm_router: "Router", model: str, team_id: str | None) -> Sequence[Mapping[str, object]]:
@ -96,15 +133,24 @@ def _routed_deployments(llm_router: "Router", model: str, team_id: str | None) -
return named or llm_router.wildcard_route_deployments(model, team_id)
def router_deployment_providers(llm_router: "Router | None", team_id: str | None) -> ProvidersOf:
if llm_router is None:
return no_providers
def _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 providers_of(model: str) -> frozenset[str]:
return frozenset(
provider
for deployment in _routed_deployments(llm_router, model, team_id)
if (provider := deployment_provider(deployment.get("litellm_params"))) is not None
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)
return ModelRoute(
providers=frozenset(
provider
for deployment in deployments
if (provider := deployment_provider(deployment.get("litellm_params"))) is not None
),
team_named=any(_is_team_deployment(deployment) for deployment in deployments),
)
return providers_of
return route_of

View file

@ -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, router_model_routes
from litellm.proxy.proxy_server import (
LitellmUserRoles,
)
@ -1618,8 +1619,8 @@ _CUSTOMER_MODEL_PROVIDERS: Final = {
}
def _customer_providers_of(model: str) -> frozenset[str]:
return _CUSTOMER_MODEL_PROVIDERS.get(model, frozenset())
def _customer_route_of(model: str) -> ModelRoute:
return ModelRoute(providers=_CUSTOMER_MODEL_PROVIDERS.get(model, frozenset()))
def _quota_error(
@ -1627,7 +1628,7 @@ def _quota_error(
rpm: dict[str, int] | None = None,
tpm: dict[str, int] | None = None,
setting: object = _CUSTOMER_QUOTA_POLICY,
providers_of: Callable[[str], frozenset[str]] = _customer_providers_of,
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,
@ -1638,7 +1639,7 @@ def _quota_error(
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, providers_of)
_raise_on_missing_project_model_quota(data, frozenset(), policy, route_of)
except HTTPException as e:
rejection: Final = e
else:
@ -1722,16 +1723,32 @@ def test_quota_policy_model_entry_matches_names_literally_apart_from_star():
def test_quota_policy_model_on_several_providers_needs_every_required_limit():
def mixed_providers(_model: str) -> frozenset[str]:
return frozenset({"runwayml", "vertex_ai"})
def mixed_providers(_model: str) -> ModelRoute:
return ModelRoute(providers=frozenset({"runwayml", "vertex_ai"}))
error = _quota_error(["video-pool"], rpm={"video-pool": 10}, providers_of=mixed_providers)
error = _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
@pytest.mark.parametrize("setting", [True, "true", "false", 1, ["rpm"]])
def test_quota_setting_truthy_non_mapping_still_requires_both_limits(setting: object):
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 test_quota_policy_model_entries_never_match_a_team_model_name():
def team_vertex_route(_model: str) -> ModelRoute:
return ModelRoute(providers=frozenset({"vertex_ai"}), team_named=True)
exact: Final = _quota_error(["jev-latest"], route_of=team_vertex_route)
pattern: Final = _quota_error(["veo-team"], rpm={"veo-team": 5}, route_of=team_vertex_route)
assert exact is not None and "jev-latest (model_rpm_limit and model_tpm_limit)" in exact
assert pattern is not None and "veo-team (model_tpm_limit)" in pattern
@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
@ -1761,7 +1778,7 @@ def test_update_quota_policy_resolves_providers_for_the_resulting_models():
existing,
frozenset(),
policy,
_customer_providers_of,
_customer_route_of,
)
is None
)
@ -1771,7 +1788,7 @@ def test_update_quota_policy_resolves_providers_for_the_resulting_models():
existing,
frozenset(),
policy,
_customer_providers_of,
_customer_route_of,
)
assert "gemini-3.8-flash (model_rpm_limit and model_tpm_limit)" in str(exc_info.value.detail)
@ -1796,8 +1813,6 @@ def _customer_router() -> "litellm.Router":
def test_quota_providers_come_from_the_one_wildcard_route_requests_use():
from litellm.proxy.management_helpers.model_quota_policy import router_deployment_providers
router: Final = litellm.Router(
model_list=[
{"model_name": "gen4*", "litellm_params": {"model": "runwayml/gen4*", "api_key": "qa-runway-key"}},
@ -1808,10 +1823,30 @@ def test_quota_providers_come_from_the_one_wildcard_route_requests_use():
},
]
)
providers_of: Final = router_deployment_providers(router, "media-team")
route_of: Final = router_model_routes(router, "media-team")
assert providers_of("gen4_turbo") == frozenset({"runwayml"})
assert providers_of("gen3a_turbo") == frozenset({"vertex_ai"})
assert route_of("gen4_turbo") == ModelRoute(providers=frozenset({"runwayml"}))
assert route_of("gen3a_turbo") == ModelRoute(providers=frozenset({"vertex_ai"}), team_named=True)
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"})
@pytest.mark.asyncio
@ -1839,7 +1874,7 @@ async def test_new_project_rejects_malformed_quota_setting_instead_of_skipping_i
monkeypatch,
team_models=["runway-gen4-turbo"],
llm_router=_customer_router(),
quota_setting={"provider": {"runwayml": []}},
quota_setting={"providers": {"runwayml": []}, "model": {"jev-latest": []}},
)
with pytest.raises(ProxyException) as exc_info:
@ -1849,6 +1884,47 @@ async def test_new_project_rejects_malformed_quota_setting_instead_of_skipping_i
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):
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):
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)

View file

@ -7365,7 +7365,7 @@ class TestEnforceRpmTpmOnModelAdd:
def test_passes_when_disabled_even_without_limits(self):
assert (
_raise_if_rate_limits_required_but_missing(
model_name="gpt-5.2", litellm_params=LiteLLM_Params(model="azure/gpt-5.2"), policy=None
model_name="gpt-5.2", team_model=False, litellm_params=LiteLLM_Params(model="azure/gpt-5.2"), policy=None
)
is None
)
@ -7373,7 +7373,7 @@ class TestEnforceRpmTpmOnModelAdd:
def test_passes_when_enabled_and_both_set(self):
assert (
_raise_if_rate_limits_required_but_missing(
model_name="gpt-5.2",
model_name="gpt-5.2", team_model=False,
litellm_params=LiteLLM_Params(model="azure/gpt-5.2", rpm=10, tpm=1000),
policy=REQUIRE_BOTH_LIMITS,
)
@ -7393,7 +7393,7 @@ class TestEnforceRpmTpmOnModelAdd:
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(
model_name="gpt-5.2", litellm_params=params, policy=REQUIRE_BOTH_LIMITS
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"
@ -7413,7 +7413,7 @@ class TestEnforceRpmTpmOnModelAdd:
):
assert (
_raise_if_rate_limits_required_but_missing(
model_name=model_name, litellm_params=params, policy=_CUSTOMER_MODEL_ADD_POLICY
model_name=model_name, team_model=False, litellm_params=params, policy=_CUSTOMER_MODEL_ADD_POLICY
)
is None
)
@ -7432,11 +7432,43 @@ class TestEnforceRpmTpmOnModelAdd:
):
with pytest.raises(ProxyException) as exc_info:
_raise_if_rate_limits_required_but_missing(
model_name=model_name, litellm_params=params, policy=_CUSTOMER_MODEL_ADD_POLICY
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"
def test_policy_reads_the_provider_behind_a_known_api_base(self):
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="https://api.groq.com/openai/v1"),
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", "expected_missing"),
[
("jev-latest", LiteLLM_Params(model="vertex_ai/gemini-3.8-flash"), "rpm and tpm"),
("veo-team", LiteLLM_Params(model="vertex_ai/veo-3.1-generate-001", rpm=5), "tpm"),
],
)
def test_policy_model_entries_never_match_a_team_model_name(
self, model_name: str, params: LiteLLM_Params, expected_missing: str
):
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
)
with pytest.raises(ProxyException) as exc_info:
_raise_if_rate_limits_required_but_missing(
model_name=model_name, team_model=True, 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}'")
@pytest.mark.asyncio
@pytest.mark.parametrize(
@ -7444,6 +7476,7 @@ class TestEnforceRpmTpmOnModelAdd:
[
({"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",
@ -7470,6 +7503,32 @@ async def test_model_add_resolves_the_quota_policy_for_the_public_model_name(
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()
monkeypatch.setattr(proxy_server, "prisma_client", prisma)
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()
class TestBlockModelResponseSerialization:
@pytest.mark.parametrize(
("route", "blocked"), [("/model/block", True), ("/model/unblock", False)]