mirror of
https://github.com/BerriAI/litellm.git
synced 2026-10-29 06:33:28 +00:00
fix(proxy): keep team model names off quota model entries, read the provider behind api_base, match patterns like the router
A team admin picks a team model's public name, so a name matching an exempt model entry no longer skips its provider's limits. A mapping with no policy key keeps requiring both limits as before, the project create permission check runs before the setting is read, and an unprefixed model behind a known OpenAI-compatible api_base resolves to that provider
This commit is contained in:
parent
7f04cbe781
commit
acacc486a8
6 changed files with 265 additions and 72 deletions
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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
|
||||
)
|
||||
|
|
|
|||
|
|
@ -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),
|
||||
)
|
||||
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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)
|
||||
|
||||
|
||||
|
|
|
|||
|
|
@ -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)]
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue