diff --git a/enterprise/litellm_enterprise/proxy/management_endpoints/project_endpoints.py b/enterprise/litellm_enterprise/proxy/management_endpoints/project_endpoints.py index 883e7b8ae8e..dea8e1b196f 100644 --- a/enterprise/litellm_enterprise/proxy/management_endpoints/project_endpoints.py +++ b/enterprise/litellm_enterprise/proxy/management_endpoints/project_endpoints.py @@ -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 diff --git a/enterprise/litellm_enterprise/proxy/management_endpoints/project_model_quota.py b/enterprise/litellm_enterprise/proxy/management_endpoints/project_model_quota.py index e427e37f10b..e7078e0bf30 100644 --- a/enterprise/litellm_enterprise/proxy/management_endpoints/project_model_quota.py +++ b/enterprise/litellm_enterprise/proxy/management_endpoints/project_model_quota.py @@ -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 ) diff --git a/litellm/proxy/management_endpoints/model_management_endpoints.py b/litellm/proxy/management_endpoints/model_management_endpoints.py index 504cc9afcbc..ef00db71eeb 100644 --- a/litellm/proxy/management_endpoints/model_management_endpoints.py +++ b/litellm/proxy/management_endpoints/model_management_endpoints.py @@ -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), ) diff --git a/litellm/proxy/management_helpers/model_quota_policy.py b/litellm/proxy/management_helpers/model_quota_policy.py index e2467e32997..28372b2aa2f 100644 --- a/litellm/proxy/management_helpers/model_quota_policy.py +++ b/litellm/proxy/management_helpers/model_quota_policy.py @@ -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 diff --git a/tests/unit/enterprise/proxy/management_endpoints/test_project_endpoints_prisma.py b/tests/unit/enterprise/proxy/management_endpoints/test_project_endpoints_prisma.py index c1e3c035b57..d430e1d3df0 100644 --- a/tests/unit/enterprise/proxy/management_endpoints/test_project_endpoints_prisma.py +++ b/tests/unit/enterprise/proxy/management_endpoints/test_project_endpoints_prisma.py @@ -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) diff --git a/tests/unit/proxy/management_endpoints/test_model_management_endpoints.py b/tests/unit/proxy/management_endpoints/test_model_management_endpoints.py index 5870e6fad9a..aebbaf28446 100644 --- a/tests/unit/proxy/management_endpoints/test_model_management_endpoints.py +++ b/tests/unit/proxy/management_endpoints/test_model_management_endpoints.py @@ -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)]