diff --git a/enterprise/litellm_enterprise/proxy/management_endpoints/project_endpoints.py b/enterprise/litellm_enterprise/proxy/management_endpoints/project_endpoints.py index dea8e1b196f..a4a56a634de 100644 --- a/enterprise/litellm_enterprise/proxy/management_endpoints/project_endpoints.py +++ b/enterprise/litellm_enterprise/proxy/management_endpoints/project_endpoints.py @@ -16,7 +16,6 @@ from typing import TYPE_CHECKING, Final, cast from fastapi import APIRouter, Depends, HTTPException, Request from litellm_enterprise.proxy.management_endpoints.project_model_quota import ( - ENFORCE_PROJECT_MODEL_QUOTA_SETTING, describe_missing_quota, missing_project_model_quota, ) @@ -33,6 +32,7 @@ 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, 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 e7078e0bf30..e36ff9b53dc 100644 --- a/enterprise/litellm_enterprise/proxy/management_endpoints/project_model_quota.py +++ b/enterprise/litellm_enterprise/proxy/management_endpoints/project_model_quota.py @@ -11,7 +11,6 @@ from litellm.proxy.management_helpers.model_quota_policy import ( unmet_limits, ) -ENFORCE_PROJECT_MODEL_QUOTA_SETTING: Final = "enforce_project_model_quota" _NO_LIMITS_SET: Final[Mapping[str, object]] = MappingProxyType({}) diff --git a/litellm/proxy/management_endpoints/model_management_endpoints.py b/litellm/proxy/management_endpoints/model_management_endpoints.py index ef00db71eeb..0ce827d214d 100644 --- a/litellm/proxy/management_endpoints/model_management_endpoints.py +++ b/litellm/proxy/management_endpoints/model_management_endpoints.py @@ -42,6 +42,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, @@ -103,9 +104,11 @@ 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 ( + ENFORCE_PROJECT_MODEL_QUOTA_SETTING, ModelQuotaPolicy, ModelRoute, deployment_provider, + limits_added_by_team_write, parse_model_quota_setting, required_limits, unmet_limits, @@ -119,6 +122,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 @@ -769,12 +773,24 @@ async def _auto_router_capability_slot( ENFORCE_RPM_TPM_ON_MODEL_ADD_SETTING: Final = "enforce_rpm_tpm_on_model_add" -def _model_add_quota_policy(general_settings: Mapping[str, object]) -> ModelQuotaPolicy | None: +_ROUTING_KEYS: Final = ("model", "custom_llm_provider", "api_base") +_NOTHING_SENT: Final[Mapping[str, object]] = MappingProxyType({}) + + +def _written_routing_params( + incoming: GenericLiteLLMParams | None, existing: GenericLiteLLMParams | None +) -> 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 _ROUTING_KEYS}) + + +def _quota_policy(general_settings: Mapping[str, object], setting: str) -> ModelQuotaPolicy | None: try: - return parse_model_quota_setting(general_settings.get(ENFORCE_RPM_TPM_ON_MODEL_ADD_SETTING)) + return parse_model_quota_setting(general_settings.get(setting)) except ValidationError as e: raise ProxyException( - message=f"general_settings.{ENFORCE_RPM_TPM_ON_MODEL_ADD_SETTING} is invalid: {e}", + message=f"general_settings.{setting} is invalid: {e}", type=ProxyErrorTypes.internal_server_error.value, code=status.HTTP_500_INTERNAL_SERVER_ERROR, param=None, @@ -790,15 +806,7 @@ def _raise_if_rate_limits_required_but_missing( (or effectively excluded from routing).""" if policy is None: return - provider: Final = deployment_provider( - MappingProxyType( - { - "model": litellm_params.model, - "custom_llm_provider": litellm_params.custom_llm_provider, - "api_base": litellm_params.api_base, - } - ) - ) + provider: Final = deployment_provider(_written_routing_params(litellm_params, None)) missing: Final = unmet_limits( required_limits( policy, @@ -820,6 +828,94 @@ def _raise_if_rate_limits_required_but_missing( ) +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 + + +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: + return + added: Final = limits_added_by_team_write(policy, llm_router, team_id, public_name, replaced_id, written_params) + if not added: + return + projects: Final = await ProjectRepository(prisma_client).table.find_many( + where={"team_id": team_id, "models": {"has": public_name}} + ) + short: Final = tuple( + project.project_alias or project.project_id + for project in projects + if unmet_limits( + added, + MappingProxyType( + { + "rpm": _project_limit(project.metadata, "model_rpm_limit", public_name), + "tpm": _project_limit(project.metadata, "model_tpm_limit", public_name), + } + ), + ) + ) + if not short: + return + raise ProxyException( + message=( + f"projects {', '.join(short)} list '{public_name}' without a positive {' and '.join(sorted(added))} " + 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, +) -> 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, existing.litellm_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 +1404,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, @@ -1417,6 +1514,17 @@ async def patch_model( existing_params=db_model.litellm_params, ) + 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=(patch_data.model_info.team_id if patch_data.model_info else None) + or db_model.model_info.team_id, + written_name=_get_public_model_name(effective_patch, db_model), + incoming_params=patch_data.litellm_params, + ) + effective_params: Final = _effective_complexity_router_params( patch_data.litellm_params, db_model.litellm_params ) @@ -2492,6 +2600,7 @@ async def delete_model( from litellm.proxy.proxy_server import ( MODEL_RECONCILE_LOCK, + general_settings_view, llm_router, premium_user, prisma_client, @@ -2525,6 +2634,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: @@ -2691,6 +2809,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, @@ -2752,7 +2872,16 @@ async def add_new_model( 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), + 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( @@ -2891,6 +3020,7 @@ async def update_model( """ from litellm.proxy.proxy_server import ( LITELLM_PROXY_ADMIN_NAME, + general_settings_view, llm_router, premium_user, prisma_client, @@ -2955,6 +3085,15 @@ async def update_model( incoming_params=model_params.litellm_params, existing_params=deployment.litellm_params, ) + 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=deployment.model_info.team_public_model_name or deployment.model_name, + incoming_params=model_params.litellm_params, + ) effective_params: Final = _effective_complexity_router_params( model_params.litellm_params, deployment.litellm_params ) diff --git a/litellm/proxy/management_helpers/model_quota_policy.py b/litellm/proxy/management_helpers/model_quota_policy.py index 28372b2aa2f..57bb83f2fda 100644 --- a/litellm/proxy/management_helpers/model_quota_policy.py +++ b/litellm/proxy/management_helpers/model_quota_policy.py @@ -3,12 +3,13 @@ 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 +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 @@ -17,10 +18,16 @@ if TYPE_CHECKING: 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") @@ -28,6 +35,14 @@ class ModelQuotaPolicy(BaseModel): 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) @@ -154,3 +169,45 @@ def router_model_routes(llm_router: "Router | None", team_id: str | None) -> Rou ) 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 team_route_after_write( + llm_router: "Router | None", + team_id: str, + public_name: str, + replaced_id: str | None, + written_params: Mapping[str, object] | None, +) -> ModelRoute: + 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) + if not team_params: + return router_model_routes(llm_router, None)(public_name) + return ModelRoute( + providers=frozenset( + provider for litellm_params in team_params if (provider := deployment_provider(litellm_params)) is not None + ), + team_named=True, + ) + + +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, +) -> frozenset[QuotaLimit]: + after: Final = team_route_after_write(llm_router, team_id, public_name, replaced_id, written_params) + now: Final = router_model_routes(llm_router, team_id)(public_name) + return required_limits(policy, public_name, after) - required_limits(policy, public_name, now) 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 d430e1d3df0..97bb36dc6c3 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 @@ -1849,6 +1849,131 @@ def test_quota_providers_read_the_provider_behind_a_known_api_base(): 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): + 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, + "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 + ) == frozenset(added) + + @pytest.mark.asyncio async def test_new_project_quota_policy_reads_providers_from_router_deployments(monkeypatch): team_models = ["gemini-3.8-flash", "veo-3.1-lite-generate-001", "runway-gen4-turbo", "jev-latest"] diff --git a/tests/unit/proxy/management_endpoints/test_model_management_endpoints.py b/tests/unit/proxy/management_endpoints/test_model_management_endpoints.py index aebbaf28446..e78b975f62b 100644 --- a/tests/unit/proxy/management_endpoints/test_model_management_endpoints.py +++ b/tests/unit/proxy/management_endpoints/test_model_management_endpoints.py @@ -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, @@ -7509,7 +7510,9 @@ async def test_team_model_add_cannot_borrow_a_model_entry_by_its_public_name(mon 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", @@ -7529,6 +7532,182 @@ async def test_team_model_add_cannot_borrow_a_model_entry_by_its_public_name(mon 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), + ) + ] + ) + 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", "models": {"has": "jev-latest"}} + ) + + +@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") + + +@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() + + class TestBlockModelResponseSerialization: @pytest.mark.parametrize( ("route", "blocked"), [("/model/block", True), ("/model/unblock", False)]