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

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

Browse source
This commit is contained in:
mateo-berri 2026-10-20 23:31:55 -07:00
parent acacc486a8
commit 39cb563098
6 changed files with 515 additions and 16 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

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

@ -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,

1
enterprise/litellm_enterprise/proxy/management_endpoints/project_model_quota.py
View file

@ -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({})

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

@ -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
)

59
litellm/proxy/management_helpers/model_quota_policy.py
View file

@ -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)

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

@ -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"]

179
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,
@ -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)]