mirror of
https://github.com/BerriAI/litellm.git
synced 2026-10-29 06:33:28 +00:00
feat(proxy): per-provider and per-model rpm/tpm requirements for enforce_project_model_quota
This commit is contained in:
mateo-berri
parent
6083cf9776
commit
ba71f280d4
3 changed files with 395 additions and 47 deletions
|
|
@ -15,7 +15,18 @@ from collections.abc import Mapping, Sequence
|
|||
from typing import TYPE_CHECKING, Final, cast
|
||||
|
||||
from fastapi import APIRouter, Depends, HTTPException, Request
|
||||
from pydantic import TypeAdapter
|
||||
from litellm_enterprise.proxy.management_endpoints.project_model_quota import (
|
||||
ENFORCE_PROJECT_MODEL_QUOTA_SETTING,
|
||||
REQUIRE_BOTH_LIMITS,
|
||||
ProjectModelQuotaPolicy,
|
||||
ProvidersOf,
|
||||
describe_missing_quota,
|
||||
missing_project_model_quota,
|
||||
no_providers,
|
||||
parse_project_model_quota_setting,
|
||||
router_deployment_providers,
|
||||
)
|
||||
from pydantic import TypeAdapter, ValidationError
|
||||
|
||||
from litellm._logging import verbose_proxy_logger
|
||||
from litellm._uuid import uuid
|
||||
|
|
@ -265,26 +276,6 @@ def _check_team_project_limits(
|
|||
)
|
||||
|
||||
|
||||
def _project_models_missing_positive_quota(
|
||||
models: list[str] | None,
|
||||
rpm_limits: Mapping[str, object] | None,
|
||||
tpm_limits: Mapping[str, object] | None,
|
||||
) -> list[str]:
|
||||
"""Return the models that lack a positive `rpm` AND `tpm` quota.
|
||||
|
||||
A valid quota is a positive integer; null, zero, and negative are rejected
|
||||
because downstream rate limiters treat a non-positive limit as immediately
|
||||
exhausted (every request blocked).
|
||||
"""
|
||||
|
||||
def _is_positive(value: object) -> bool:
|
||||
return isinstance(value, int) and not isinstance(value, bool) and value > 0
|
||||
|
||||
rpm = rpm_limits or {}
|
||||
tpm = tpm_limits or {}
|
||||
return [model for model in (models or []) if not _is_positive(rpm.get(model)) or not _is_positive(tpm.get(model))]
|
||||
|
||||
|
||||
def _router_access_group_names(llm_router: "Router | None") -> frozenset[str]:
|
||||
return frozenset(llm_router.get_model_access_groups()) if llm_router is not None else frozenset()
|
||||
|
||||
|
|
@ -311,64 +302,79 @@ def _raise_on_project_models_expanding_at_request_time(
|
|||
raise HTTPException(
|
||||
status_code=400,
|
||||
detail={
|
||||
"error": f"models {list(expanding)} expand to multiple models at request time, so a per-model rpm/tpm quota cannot be enforced for them while 'enforce_project_model_quota' is enabled. List concrete model names instead."
|
||||
"error": f"models {list(expanding)} expand to multiple models at request time, so a per-model rpm/tpm quota cannot be enforced for them while '{ENFORCE_PROJECT_MODEL_QUOTA_SETTING}' is enabled. List concrete model names instead."
|
||||
},
|
||||
)
|
||||
|
||||
|
||||
def _project_model_quota_policy(general_settings: Mapping[str, object]) -> ProjectModelQuotaPolicy | None:
|
||||
try:
|
||||
return parse_project_model_quota_setting(general_settings.get(ENFORCE_PROJECT_MODEL_QUOTA_SETTING))
|
||||
except ValidationError as e:
|
||||
raise HTTPException(
|
||||
status_code=500,
|
||||
detail={"error": f"general_settings.{ENFORCE_PROJECT_MODEL_QUOTA_SETTING} is invalid: {e}"},
|
||||
) from e
|
||||
|
||||
|
||||
def _raise_on_missing_project_model_quota(
|
||||
data: NewProjectRequest | UpdateProjectRequest, access_group_names: frozenset[str] = frozenset()
|
||||
data: NewProjectRequest | UpdateProjectRequest,
|
||||
access_group_names: frozenset[str] = frozenset(),
|
||||
policy: ProjectModelQuotaPolicy = REQUIRE_BOTH_LIMITS,
|
||||
providers_of: ProvidersOf = no_providers,
|
||||
) -> None:
|
||||
"""Require a positive `rpm`/`tpm` quota for every model on project CREATE.
|
||||
"""Require the positive `rpm`/`tpm` quotas the policy asks of each model on project CREATE.
|
||||
|
||||
`model_rpm_limit`/`model_tpm_limit` are relocated into `metadata` by the request
|
||||
model's `set_model_info` validator, so they are read from there.
|
||||
|
||||
Only invoked when `general_settings.enforce_project_model_quota` is enabled
|
||||
(default off), so it is opt-in and does not change behavior for existing users.
|
||||
"""
|
||||
_raise_on_project_models_expanding_at_request_time(data.models, access_group_names)
|
||||
metadata = data.metadata or {}
|
||||
missing = _project_models_missing_positive_quota(
|
||||
data.models, metadata.get("model_rpm_limit"), metadata.get("model_tpm_limit")
|
||||
missing: Final = missing_project_model_quota(
|
||||
data.models or (), metadata.get("model_rpm_limit"), metadata.get("model_tpm_limit"), policy, providers_of
|
||||
)
|
||||
if not missing:
|
||||
return
|
||||
raise HTTPException(
|
||||
status_code=400,
|
||||
detail={
|
||||
"error": f"models {missing} added to project without a positive rpm/tpm quota. Set a positive model_rpm_limit and model_tpm_limit for each model."
|
||||
"error": f"models added to project without a positive rpm/tpm quota: {describe_missing_quota(missing)}. Set a positive value for each listed limit."
|
||||
},
|
||||
)
|
||||
|
||||
|
||||
def _raise_on_missing_project_model_quota_on_update(
|
||||
data: UpdateProjectRequest, existing_project: object, access_group_names: frozenset[str] = frozenset()
|
||||
data: UpdateProjectRequest,
|
||||
existing_project: object,
|
||||
access_group_names: frozenset[str] = frozenset(),
|
||||
policy: ProjectModelQuotaPolicy = REQUIRE_BOTH_LIMITS,
|
||||
providers_of: ProvidersOf = no_providers,
|
||||
) -> None:
|
||||
"""Require a positive `rpm`/`tpm` quota over the RESULTING state on project UPDATE.
|
||||
"""Require the positive `rpm`/`tpm` quotas the policy asks of each model over the RESULTING state on project UPDATE.
|
||||
|
||||
`/project/update` replaces `models` and `metadata` when they are provided, so the
|
||||
check runs on what the project WILL look like: a partial update that doesn't touch
|
||||
models/quota keeps the existing values, while one that adds a model or clears a
|
||||
model's quota must leave every resulting model with a positive limit.
|
||||
|
||||
Only invoked when `general_settings.enforce_project_model_quota` is enabled
|
||||
(default off), so it is opt-in and does not change behavior for existing users.
|
||||
model's quota must leave every resulting model with the limits it needs.
|
||||
"""
|
||||
resulting_models = data.models if data.models is not None else (getattr(existing_project, "models", None) or [])
|
||||
resulting_metadata = (
|
||||
data.metadata if data.metadata is not None else (getattr(existing_project, "metadata", None) or {})
|
||||
)
|
||||
_raise_on_project_models_expanding_at_request_time(resulting_models, access_group_names)
|
||||
missing = _project_models_missing_positive_quota(
|
||||
resulting_models, resulting_metadata.get("model_rpm_limit"), resulting_metadata.get("model_tpm_limit")
|
||||
missing: Final = missing_project_model_quota(
|
||||
resulting_models,
|
||||
resulting_metadata.get("model_rpm_limit"),
|
||||
resulting_metadata.get("model_tpm_limit"),
|
||||
policy,
|
||||
providers_of,
|
||||
)
|
||||
if not missing:
|
||||
return
|
||||
raise HTTPException(
|
||||
status_code=400,
|
||||
detail={
|
||||
"error": f"models {missing} would be left on the project without a positive rpm/tpm quota. Set a positive model_rpm_limit and model_tpm_limit for each model."
|
||||
"error": f"models would be left on the project without a positive rpm/tpm quota: {describe_missing_quota(missing)}. Set a positive value for each listed limit."
|
||||
},
|
||||
)
|
||||
|
||||
|
|
@ -567,9 +573,14 @@ async def new_project(
|
|||
data=data,
|
||||
)
|
||||
|
||||
# Opt-in (default off): require rpm/tpm for every model added to the project.
|
||||
if general_settings.get("enforce_project_model_quota", False):
|
||||
_raise_on_missing_project_model_quota(data, _router_access_group_names(llm_router))
|
||||
quota_policy: Final = _project_model_quota_policy(general_settings)
|
||||
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
|
||||
|
|
@ -817,10 +828,14 @@ async def update_project(
|
|||
data=data,
|
||||
)
|
||||
|
||||
# Opt-in (default off): require rpm/tpm for every model the update would leave on the project.
|
||||
if general_settings.get("enforce_project_model_quota", False):
|
||||
quota_policy: Final = _project_model_quota_policy(general_settings)
|
||||
if quota_policy is not None:
|
||||
_raise_on_missing_project_model_quota_on_update(
|
||||
data, existing_project, _router_access_group_names(llm_router)
|
||||
data,
|
||||
existing_project,
|
||||
_router_access_group_names(llm_router),
|
||||
quota_policy,
|
||||
router_deployment_providers(llm_router, target_team_id),
|
||||
)
|
||||
|
||||
# Prepare update data
|
||||
|
|
|
|||
|
|
@ -0,0 +1,125 @@
|
|||
import re
|
||||
from collections.abc import Callable, Mapping, Sequence
|
||||
from dataclasses import dataclass
|
||||
from itertools import chain
|
||||
from types import MappingProxyType
|
||||
from typing import TYPE_CHECKING, Final, Literal, TypeAlias
|
||||
|
||||
from pydantic import BaseModel, ConfigDict, Field, TypeAdapter
|
||||
|
||||
from litellm.router_utils.common_utils import provider_for_generic_call
|
||||
from litellm.router_utils.pattern_match_deployments import PatternUtils
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from litellm import Router
|
||||
|
||||
QuotaLimit: TypeAlias = Literal["rpm", "tpm"]
|
||||
ProvidersOf: TypeAlias = Callable[[str], frozenset[str]]
|
||||
|
||||
ENFORCE_PROJECT_MODEL_QUOTA_SETTING: Final = "enforce_project_model_quota"
|
||||
_BOTH_LIMITS: Final[frozenset[QuotaLimit]] = frozenset(("rpm", "tpm"))
|
||||
_NO_LIMITS_SET: Final[Mapping[str, object]] = MappingProxyType({})
|
||||
|
||||
|
||||
class ProjectModelQuotaPolicy(BaseModel):
|
||||
model_config = ConfigDict(frozen=True, extra="forbid")
|
||||
|
||||
default: frozenset[QuotaLimit] = _BOTH_LIMITS
|
||||
providers: Mapping[str, frozenset[QuotaLimit]] = Field(default_factory=dict)
|
||||
models: Mapping[str, frozenset[QuotaLimit]] = Field(default_factory=dict)
|
||||
|
||||
|
||||
REQUIRE_BOTH_LIMITS: Final = ProjectModelQuotaPolicy()
|
||||
|
||||
_QUOTA_SETTING: Final = TypeAdapter(bool | ProjectModelQuotaPolicy | None)
|
||||
|
||||
|
||||
@dataclass(frozen=True, slots=True)
|
||||
class MissingModelQuota:
|
||||
model: str
|
||||
limits: tuple[QuotaLimit, ...]
|
||||
|
||||
|
||||
def parse_project_model_quota_setting(setting: object) -> ProjectModelQuotaPolicy | None:
|
||||
parsed: Final = _QUOTA_SETTING.validate_python(setting)
|
||||
match parsed:
|
||||
case ProjectModelQuotaPolicy():
|
||||
return parsed
|
||||
case True:
|
||||
return REQUIRE_BOTH_LIMITS
|
||||
case _:
|
||||
return None
|
||||
|
||||
|
||||
def _wildcard_matches(pattern: str, model: str) -> bool:
|
||||
return re.fullmatch(re.escape(pattern).replace(r"\*", ".*"), model) is not None
|
||||
|
||||
|
||||
def _most_specific_model_pattern(model: str, entries: Mapping[str, frozenset[QuotaLimit]]) -> str | None:
|
||||
matching: Final = tuple(pattern for pattern in entries if "*" in pattern and _wildcard_matches(pattern, model))
|
||||
return max(matching, key=PatternUtils.calculate_pattern_specificity, default=None)
|
||||
|
||||
|
||||
def required_limits(policy: ProjectModelQuotaPolicy, 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)
|
||||
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:
|
||||
return policy.default
|
||||
return frozenset(chain.from_iterable(policy.providers.get(provider, policy.default) for provider in providers))
|
||||
|
||||
|
||||
def _is_positive(value: object) -> bool:
|
||||
return isinstance(value, int) and not isinstance(value, bool) and value > 0
|
||||
|
||||
|
||||
def _unset_limits(
|
||||
model: str, required: frozenset[QuotaLimit], set_limits: Mapping[QuotaLimit, Mapping[str, object]]
|
||||
) -> tuple[QuotaLimit, ...]:
|
||||
return tuple(sorted(limit for limit in required if not _is_positive(set_limits[limit].get(model))))
|
||||
|
||||
|
||||
def missing_project_model_quota(
|
||||
models: Sequence[str],
|
||||
rpm_limits: Mapping[str, object] | None,
|
||||
tpm_limits: Mapping[str, object] | None,
|
||||
policy: ProjectModelQuotaPolicy,
|
||||
providers_of: ProvidersOf,
|
||||
) -> tuple[MissingModelQuota, ...]:
|
||||
set_limits: Final[Mapping[QuotaLimit, Mapping[str, object]]] = MappingProxyType(
|
||||
{"rpm": rpm_limits or _NO_LIMITS_SET, "tpm": tpm_limits or _NO_LIMITS_SET}
|
||||
)
|
||||
gaps: Final = tuple(
|
||||
MissingModelQuota(model, _unset_limits(model, required_limits(policy, model, providers_of(model)), set_limits))
|
||||
for model in models
|
||||
)
|
||||
return tuple(gap for gap in gaps if gap.limits)
|
||||
|
||||
|
||||
def describe_missing_quota(missing: Sequence[MissingModelQuota]) -> str:
|
||||
return ", ".join(f"{gap.model} ({' and '.join(f'model_{limit}_limit' for limit in gap.limits)})" for gap in missing)
|
||||
|
||||
|
||||
def no_providers(_model: str) -> frozenset[str]:
|
||||
return frozenset()
|
||||
|
||||
|
||||
def router_deployment_providers(llm_router: "Router | None", team_id: str | None) -> ProvidersOf:
|
||||
if llm_router is None:
|
||||
return no_providers
|
||||
|
||||
def providers_of(model: str) -> frozenset[str]:
|
||||
deployments: Final = llm_router.get_model_list(model_name=model, team_id=team_id) or ()
|
||||
return frozenset(
|
||||
provider
|
||||
for deployment in deployments
|
||||
if (provider := provider_for_generic_call(deployment["litellm_params"])) is not None
|
||||
)
|
||||
|
||||
return providers_of
|
||||
|
|
@ -1205,7 +1205,9 @@ def test_update_quota_existing_quotaless_model_rejected():
|
|||
_raise_on_missing_project_model_quota_on_update(data, existing)
|
||||
|
||||
|
||||
def _enforced_new_project_mocks(monkeypatch, team_models: list[str], llm_router: mock.MagicMock | None) -> None:
|
||||
def _enforced_new_project_mocks(
|
||||
monkeypatch, team_models: list[str], llm_router: object, quota_setting: object = True
|
||||
) -> None:
|
||||
from litellm.proxy._types import LiteLLM_TeamTable
|
||||
from litellm_enterprise.proxy.management_endpoints import project_endpoints as pe
|
||||
|
||||
|
|
@ -1213,7 +1215,7 @@ def _enforced_new_project_mocks(monkeypatch, team_models: list[str], llm_router:
|
|||
monkeypatch.setattr(litellm.proxy.proxy_server, "prisma_client", mock.MagicMock())
|
||||
monkeypatch.setattr(litellm.proxy.proxy_server, "premium_user", True)
|
||||
monkeypatch.setattr(litellm.proxy.proxy_server, "llm_router", llm_router)
|
||||
monkeypatch.setattr(litellm.proxy.proxy_server, "general_settings", {"enforce_project_model_quota": True})
|
||||
monkeypatch.setattr(litellm.proxy.proxy_server, "general_settings", {"enforce_project_model_quota": quota_setting})
|
||||
monkeypatch.setattr(pe, "_validate_team_exists", mock.AsyncMock(return_value=team))
|
||||
monkeypatch.setattr(pe, "_check_user_permission_for_project", mock.AsyncMock(return_value=True))
|
||||
|
||||
|
|
@ -1604,6 +1606,212 @@ async def test_new_project_flag_on_access_group_model_returns_400(monkeypatch):
|
|||
assert "expand to multiple models at request time" in str(exc_info.value)
|
||||
|
||||
|
||||
_CUSTOMER_QUOTA_POLICY: Final = {
|
||||
"providers": {"vertex_ai": ["rpm", "tpm"], "runwayml": []},
|
||||
"models": {"veo-*": ["rpm"], "jev-latest": []},
|
||||
}
|
||||
|
||||
_CUSTOMER_MODEL_PROVIDERS: Final = {
|
||||
"gemini-3.8-flash": frozenset({"vertex_ai"}),
|
||||
"veo-3.1-lite-generate-001": frozenset({"vertex_ai"}),
|
||||
"runway-gen4-turbo": frozenset({"runwayml"}),
|
||||
}
|
||||
|
||||
|
||||
def _customer_providers_of(model: str) -> frozenset[str]:
|
||||
return _CUSTOMER_MODEL_PROVIDERS.get(model, frozenset())
|
||||
|
||||
|
||||
def _quota_error(
|
||||
models: list[str],
|
||||
rpm: dict[str, int] | None = None,
|
||||
tpm: dict[str, int] | None = None,
|
||||
setting: object = _CUSTOMER_QUOTA_POLICY,
|
||||
providers_of=_customer_providers_of,
|
||||
) -> str | None:
|
||||
from litellm_enterprise.proxy.management_endpoints.project_endpoints import (
|
||||
_raise_on_missing_project_model_quota,
|
||||
)
|
||||
from litellm_enterprise.proxy.management_endpoints.project_model_quota import (
|
||||
parse_project_model_quota_setting,
|
||||
)
|
||||
|
||||
policy = parse_project_model_quota_setting(setting)
|
||||
assert policy is not None
|
||||
data = 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)
|
||||
except HTTPException as e:
|
||||
rejection = e
|
||||
else:
|
||||
return None
|
||||
assert rejection.status_code == 400
|
||||
return str(rejection.detail)
|
||||
|
||||
|
||||
def test_quota_policy_exempts_provider_and_passthrough_name_but_keeps_vertex_strict():
|
||||
assert _quota_error(["runway-gen4-turbo", "jev-latest"]) is None
|
||||
error = _quota_error(["gemini-3.8-flash"], rpm={"gemini-3.8-flash": 10})
|
||||
assert error is not None
|
||||
assert "gemini-3.8-flash (model_tpm_limit)" in error
|
||||
|
||||
|
||||
def test_quota_policy_model_pattern_overrides_its_provider_per_limit():
|
||||
assert _quota_error(["veo-3.1-lite-generate-001"], rpm={"veo-3.1-lite-generate-001": 5}) is None
|
||||
error = _quota_error(["veo-3.1-lite-generate-001"], tpm={"veo-3.1-lite-generate-001": 5000})
|
||||
assert error is not None
|
||||
assert "veo-3.1-lite-generate-001 (model_rpm_limit)" in error
|
||||
assert "model_tpm_limit" not in error
|
||||
|
||||
|
||||
def test_quota_policy_error_names_each_model_with_the_limits_it_misses():
|
||||
error = _quota_error(["gemini-3.8-flash", "veo-3.1-lite-generate-001", "runway-gen4-turbo"])
|
||||
assert error is not None
|
||||
assert "gemini-3.8-flash (model_rpm_limit and model_tpm_limit)" in error
|
||||
assert "veo-3.1-lite-generate-001 (model_rpm_limit)" in error
|
||||
assert "runway-gen4-turbo" not in error
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"setting, model, expected_missing",
|
||||
[
|
||||
({"default": ["tpm"]}, "unrouted-name", "unrouted-name (model_tpm_limit)"),
|
||||
({"default": [], "providers": {"runwayml": ["rpm"]}}, "runway-gen4-turbo", "(model_rpm_limit)"),
|
||||
(
|
||||
{"providers": {"runwayml": []}, "models": {"runway-gen4-turbo": ["tpm"]}},
|
||||
"runway-gen4-turbo",
|
||||
"(model_tpm_limit)",
|
||||
),
|
||||
({"models": {"runway-*": [], "runway-gen4-turbo": ["rpm"]}}, "runway-gen4-turbo", "(model_rpm_limit)"),
|
||||
({"models": {"veo-*": [], "veo-3.1-lite-*": ["rpm"]}}, "veo-3.1-lite-generate-001", "(model_rpm_limit)"),
|
||||
],
|
||||
)
|
||||
def test_quota_policy_precedence_is_model_then_pattern_then_provider_then_default(setting, model, expected_missing):
|
||||
error = _quota_error([model], setting=setting)
|
||||
assert error is not None
|
||||
assert expected_missing in error
|
||||
|
||||
|
||||
def test_quota_policy_model_entry_matches_names_literally_apart_from_star():
|
||||
setting = {"models": {"claude-opus-5-5[1m]": [], "gpt.5": []}}
|
||||
assert _quota_error(["claude-opus-5-5[1m]"], setting=setting) is None
|
||||
assert _quota_error(["claude-opus-5-51"], setting=setting) is not None
|
||||
assert _quota_error(["gptx5"], setting=setting) is not None
|
||||
|
||||
|
||||
def test_quota_policy_model_on_several_providers_needs_every_required_limit():
|
||||
def mixed_providers(_model: str) -> frozenset[str]:
|
||||
return frozenset({"runwayml", "vertex_ai"})
|
||||
|
||||
error = _quota_error(["video-pool"], rpm={"video-pool": 10}, providers_of=mixed_providers)
|
||||
assert error is not None
|
||||
assert "video-pool (model_tpm_limit)" in error
|
||||
|
||||
|
||||
@pytest.mark.parametrize("setting", [True, "true", {}])
|
||||
def test_quota_setting_true_or_empty_mapping_still_requires_both_limits(setting):
|
||||
error = _quota_error(["runway-gen4-turbo"], setting=setting)
|
||||
assert error is not None
|
||||
assert "runway-gen4-turbo (model_rpm_limit and model_tpm_limit)" in error
|
||||
|
||||
|
||||
@pytest.mark.parametrize("setting", [False, None])
|
||||
def test_quota_setting_off_enforces_nothing(setting):
|
||||
from litellm_enterprise.proxy.management_endpoints.project_model_quota import (
|
||||
parse_project_model_quota_setting,
|
||||
)
|
||||
|
||||
assert parse_project_model_quota_setting(setting) is None
|
||||
|
||||
|
||||
def test_update_quota_policy_resolves_providers_for_the_resulting_models():
|
||||
import types
|
||||
|
||||
from litellm_enterprise.proxy.management_endpoints.project_endpoints import (
|
||||
_raise_on_missing_project_model_quota_on_update,
|
||||
)
|
||||
from litellm_enterprise.proxy.management_endpoints.project_model_quota import (
|
||||
parse_project_model_quota_setting,
|
||||
)
|
||||
|
||||
policy = parse_project_model_quota_setting(_CUSTOMER_QUOTA_POLICY)
|
||||
assert policy is not None
|
||||
existing = types.SimpleNamespace(models=["runway-gen4-turbo"], metadata={})
|
||||
assert (
|
||||
_raise_on_missing_project_model_quota_on_update(
|
||||
UpdateProjectRequest(project_id="p", description="unrelated change"),
|
||||
existing,
|
||||
frozenset(),
|
||||
policy,
|
||||
_customer_providers_of,
|
||||
)
|
||||
is None
|
||||
)
|
||||
with pytest.raises(HTTPException) as exc_info:
|
||||
_raise_on_missing_project_model_quota_on_update(
|
||||
UpdateProjectRequest(project_id="p", models=["runway-gen4-turbo", "gemini-3.8-flash"]),
|
||||
existing,
|
||||
frozenset(),
|
||||
policy,
|
||||
_customer_providers_of,
|
||||
)
|
||||
assert "gemini-3.8-flash (model_rpm_limit and model_tpm_limit)" in str(exc_info.value.detail)
|
||||
|
||||
|
||||
def _customer_router() -> "litellm.Router":
|
||||
return litellm.Router(
|
||||
model_list=[
|
||||
{
|
||||
"model_name": "gemini-3.8-flash",
|
||||
"litellm_params": {"model": "vertex_ai/gemini-3.8-flash", "vertex_project": "qa-project"},
|
||||
},
|
||||
{
|
||||
"model_name": "veo-3.1-lite-generate-001",
|
||||
"litellm_params": {"model": "vertex_ai/veo-3.1-lite-generate-001", "vertex_project": "qa-project"},
|
||||
},
|
||||
{
|
||||
"model_name": "runway-gen4-turbo",
|
||||
"litellm_params": {"model": "runwayml/gen4_turbo", "api_key": "qa-runway-key"},
|
||||
},
|
||||
]
|
||||
)
|
||||
|
||||
|
||||
@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"]
|
||||
_enforced_new_project_mocks(
|
||||
monkeypatch, team_models=team_models, llm_router=_customer_router(), quota_setting=_CUSTOMER_QUOTA_POLICY
|
||||
)
|
||||
|
||||
with pytest.raises(ProxyException) as exc_info:
|
||||
await _run_new_project(NewProjectRequest(team_id="test-team", models=team_models))
|
||||
|
||||
assert exc_info.value.code == "400"
|
||||
assert (
|
||||
"gemini-3.8-flash (model_rpm_limit and model_tpm_limit), veo-3.1-lite-generate-001 (model_rpm_limit)."
|
||||
in str(exc_info.value.message)
|
||||
)
|
||||
assert "runway-gen4-turbo" not in str(exc_info.value.message)
|
||||
assert "jev-latest" not in str(exc_info.value.message)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_new_project_rejects_malformed_quota_setting_instead_of_skipping_it(monkeypatch):
|
||||
_enforced_new_project_mocks(
|
||||
monkeypatch,
|
||||
team_models=["runway-gen4-turbo"],
|
||||
llm_router=_customer_router(),
|
||||
quota_setting={"provider": {"runwayml": []}},
|
||||
)
|
||||
|
||||
with pytest.raises(ProxyException) as exc_info:
|
||||
await _run_new_project(NewProjectRequest(team_id="test-team", models=["runway-gen4-turbo"]))
|
||||
|
||||
assert exc_info.value.code == "500"
|
||||
assert "general_settings.enforce_project_model_quota is invalid" in str(exc_info.value.message)
|
||||
|
||||
|
||||
_ACCESS_NOW: Final = datetime.now(timezone.utc)
|
||||
|
||||
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue