diff --git a/litellm/proxy/_types.py b/litellm/proxy/_types.py index 3ccf3ea9952..817a821b162 100644 --- a/litellm/proxy/_types.py +++ b/litellm/proxy/_types.py @@ -2412,6 +2412,10 @@ class ConfigGeneralSettings(LiteLLMPydanticObjectBase): None, description="By default, the user calling /team/new is automatically added to the new team as a team admin. If True, proxy admins are no longer auto-added; members explicitly listed in members_with_roles are unaffected. Default is False.", ) + enforce_team_model_limit_allocation: bool | None = Field( + None, + description="If True, a team key's per-model rpm/tpm limits are rejected at create/update time when they would overallocate the team's per-model limits, regardless of rpm_limit_type/tpm_limit_type. A key's own per-model limit wins over the team's at request time, so this is how the team limit becomes a true cap. Default is False.", + ) maximum_spend_logs_retention_period: str | None = Field( None, description="Maximum retention period for spend logs (e.g., '7d' for 7 days). Logs older than this will be deleted.", diff --git a/litellm/proxy/auth/auth_utils.py b/litellm/proxy/auth/auth_utils.py index fc4948c63cb..86ec56c811a 100644 --- a/litellm/proxy/auth/auth_utils.py +++ b/litellm/proxy/auth/auth_utils.py @@ -902,6 +902,42 @@ def _get_deployment_default_tpm_limit(model_name: str) -> int | None: return _get_deployment_default_limit(model_name, "default_api_key_tpm_limit") +def resolve_own_model_rate_limits( + metadata: object, + model_max_budget: object, + rate_limit_key: Literal["model_rpm_limit", "model_tpm_limit"], +) -> dict[str, int] | None: + """ + Per-model limits an entity sets on itself, from the two places one can carry + them. + + Priority order (returns first found, per metric): + 1. A ``model_rpm_limit`` / ``model_tpm_limit`` map in metadata + 2. The per-model ``rpm_limit`` / ``tpm_limit`` entries of a model_max_budget + + Takes the raw maps rather than a typed row so runtime enforcement + (UserAPIKeyAuth) and grant-time allocation math (verification-token rows, + key request bodies) resolve a key's own limits identically; both sources are + untyped JSON, so every level is shape-checked before use. + """ + if isinstance(metadata, dict): + metadata_limits = metadata.get(rate_limit_key) + if metadata_limits and isinstance(metadata_limits, dict): + return metadata_limits + + if isinstance(model_max_budget, dict): + budget_field = "rpm_limit" if rate_limit_key == "model_rpm_limit" else "tpm_limit" + budget_limits = { + model: budget[budget_field] + for model, budget in model_max_budget.items() + if isinstance(budget, dict) and budget.get(budget_field) is not None + } + if budget_limits: + return budget_limits + + return None + + def get_key_own_model_rate_limits( user_api_key_dict: UserAPIKeyAuth, rate_limit_key: Literal["model_rpm_limit", "model_tpm_limit"], @@ -910,30 +946,15 @@ def get_key_own_model_rate_limits( Get the per-model limits configured on the key itself, ignoring anything it only inherits (team metadata, deployment defaults). - Priority order (returns first found): - 1. Key metadata (model_rpm_limit / model_tpm_limit) - 2. Key model_max_budget (rpm_limit / tpm_limit per model) - Callers that must know whether the key *overrides* an inherited limit use this; callers that want the effective limit use get_key_model_rpm_limit / get_key_model_tpm_limit, which continue the chain past the key. """ - if user_api_key_dict.metadata: - metadata_limits = user_api_key_dict.metadata.get(rate_limit_key) - if metadata_limits: - return metadata_limits - - if user_api_key_dict.model_max_budget: - budget_field = "rpm_limit" if rate_limit_key == "model_rpm_limit" else "tpm_limit" - budget_limits = { - model: budget[budget_field] - for model, budget in user_api_key_dict.model_max_budget.items() - if isinstance(budget, dict) and budget.get(budget_field) is not None - } - if budget_limits: - return budget_limits - - return None + return resolve_own_model_rate_limits( + metadata=user_api_key_dict.metadata, + model_max_budget=user_api_key_dict.model_max_budget, + rate_limit_key=rate_limit_key, + ) def get_key_model_rpm_limit( diff --git a/litellm/proxy/management_endpoints/key_management_endpoints.py b/litellm/proxy/management_endpoints/key_management_endpoints.py index fb9e6fd739e..52a382ce8f8 100644 --- a/litellm/proxy/management_endpoints/key_management_endpoints.py +++ b/litellm/proxy/management_endpoints/key_management_endpoints.py @@ -55,7 +55,10 @@ from litellm.proxy.auth.auth_checks import ( get_project_object, get_team_object, ) -from litellm.proxy.auth.auth_utils import abbreviate_api_key +from litellm.proxy.auth.auth_utils import ( + abbreviate_api_key, + resolve_own_model_rate_limits, +) from litellm.proxy.auth.user_api_key_auth import user_api_key_auth from litellm.proxy.common_utils.callback_utils import ( decrypt_callback_vars, @@ -1158,6 +1161,44 @@ async def _common_key_generation_helper( return response +def _requested_model_specific_limits( + data: Union[GenerateKeyRequest, UpdateKeyRequest], + rate_limit_key: Literal["model_rpm_limit", "model_tpm_limit"], +) -> Mapping[str, int] | None: + """ + Per-model limits this request would grant the key, from any of the shapes a + caller can express them in: the top-level field, the same map under + metadata, or per-model limits inside model_max_budget. + """ + explicit_limits = data.model_rpm_limit if rate_limit_key == "model_rpm_limit" else data.model_tpm_limit + if explicit_limits: + return explicit_limits + + return resolve_own_model_rate_limits( + metadata=data.metadata, + model_max_budget=data.model_max_budget, + rate_limit_key=rate_limit_key, + ) + + +def _allocated_model_specific_limits( + keys: Sequence[LiteLLM_VerificationToken], + rate_limit_key: Literal["model_rpm_limit", "model_tpm_limit"], +) -> Mapping[str, int]: + """Per-model total already handed out across an entity's existing keys.""" + per_key_limits = tuple( + resolve_own_model_rate_limits( + metadata=key.metadata, + model_max_budget=key.model_max_budget, + rate_limit_key=rate_limit_key, + ) + or {} + for key in keys + ) + allocated_models = frozenset(model for limits in per_key_limits for model in limits) + return {model: sum(limits.get(model, 0) for limits in per_key_limits) for model in allocated_models} + + def _check_key_model_specific_limits( keys: list[LiteLLM_VerificationToken], data: Union[GenerateKeyRequest, UpdateKeyRequest], @@ -1171,33 +1212,20 @@ def _check_key_model_specific_limits( Generic function to check if a key is allocating model specific limits. Raises an error if we're overallocating. """ - model_rpm_limit = getattr(data, "model_rpm_limit", None) or ( - data.metadata.get("model_rpm_limit", None) if data.metadata else None - ) - model_tpm_limit = getattr(data, "model_tpm_limit", None) or ( - data.metadata.get("model_tpm_limit", None) if data.metadata else None - ) + model_rpm_limit = _requested_model_specific_limits(data, "model_rpm_limit") + model_tpm_limit = _requested_model_specific_limits(data, "model_tpm_limit") if model_rpm_limit is None and model_tpm_limit is None: return - # get total model specific tpm/rpm limit - model_specific_rpm_limit: dict[str, int] = {} - model_specific_tpm_limit: dict[str, int] = {} - - for key in keys: - if key.metadata.get("model_rpm_limit", None) is not None: - for model, rpm_limit in key.metadata.get("model_rpm_limit", {}).items(): - model_specific_rpm_limit[model] = model_specific_rpm_limit.get(model, 0) + rpm_limit - if key.metadata.get("model_tpm_limit", None) is not None: - for model, tpm_limit in key.metadata.get("model_tpm_limit", {}).items(): - model_specific_tpm_limit[model] = model_specific_tpm_limit.get(model, 0) + tpm_limit + model_specific_rpm_limit = _allocated_model_specific_limits(keys, "model_rpm_limit") + model_specific_tpm_limit = _allocated_model_specific_limits(keys, "model_tpm_limit") if model_rpm_limit is not None: for model, rpm_limit in model_rpm_limit.items(): if entity_rpm_limit is not None and model_specific_rpm_limit.get(model, 0) + rpm_limit > entity_rpm_limit: raise HTTPException( status_code=400, - detail=f"Allocated RPM limit={model_specific_rpm_limit.get(model, 0)} + Key RPM limit={rpm_limit} is greater than {entity_type} RPM limit={entity_rpm_limit}", + detail=f"Allocated RPM limit={model_specific_rpm_limit.get(model, 0)} + Key RPM limit={rpm_limit} is greater than {entity_type} RPM limit={entity_rpm_limit} for model={model}", ) elif entity_model_rpm_limit_dict: entity_model_specific_rpm_limit = entity_model_rpm_limit_dict.get(model) @@ -1207,7 +1235,7 @@ def _check_key_model_specific_limits( ): raise HTTPException( status_code=400, - detail=f"Allocated RPM limit={model_specific_rpm_limit.get(model, 0)} + Key RPM limit={rpm_limit} is greater than {entity_type} RPM limit={entity_model_specific_rpm_limit}", + detail=f"Allocated RPM limit={model_specific_rpm_limit.get(model, 0)} + Key RPM limit={rpm_limit} is greater than {entity_type} RPM limit={entity_model_specific_rpm_limit} for model={model}", ) if model_tpm_limit is not None: @@ -1215,7 +1243,7 @@ def _check_key_model_specific_limits( if entity_tpm_limit is not None and model_specific_tpm_limit.get(model, 0) + tpm_limit > entity_tpm_limit: raise HTTPException( status_code=400, - detail=f"Allocated TPM limit={model_specific_tpm_limit.get(model, 0)} + Key TPM limit={tpm_limit} is greater than {entity_type} TPM limit={entity_tpm_limit}", + detail=f"Allocated TPM limit={model_specific_tpm_limit.get(model, 0)} + Key TPM limit={tpm_limit} is greater than {entity_type} TPM limit={entity_tpm_limit} for model={model}", ) elif entity_model_tpm_limit_dict: entity_model_specific_tpm_limit = entity_model_tpm_limit_dict.get(model) @@ -1225,7 +1253,7 @@ def _check_key_model_specific_limits( ): raise HTTPException( status_code=400, - detail=f"Allocated TPM limit={model_specific_tpm_limit.get(model, 0)} + Key TPM limit={tpm_limit} is greater than {entity_type} TPM limit={entity_model_specific_tpm_limit}", + detail=f"Allocated TPM limit={model_specific_tpm_limit.get(model, 0)} + Key TPM limit={tpm_limit} is greater than {entity_type} TPM limit={entity_model_specific_tpm_limit} for model={model}", ) @@ -1309,6 +1337,21 @@ def check_team_key_rpm_tpm_limits( ) +def _is_team_model_limit_allocation_enforced() -> bool: + """ + Whether every team key must fit inside the team's per-model limits, not just + keys asking for guaranteed throughput. + + Off by default: at runtime a key's own per-model limit deliberately wins over + the team's, so without this an admin can hand one key a ceiling above the + team's. Operators who want the team limit to be a true cap opt in here, and + the rejection happens at grant time rather than silently at request time. + """ + from litellm.proxy.proxy_server import general_settings + + return general_settings.get("enforce_team_model_limit_allocation") is True + + async def _check_team_key_limits( team_table: LiteLLM_TeamTableCachedObj, data: Union[GenerateKeyRequest, UpdateKeyRequest], @@ -1317,9 +1360,16 @@ async def _check_team_key_limits( """ Check if the team key is allocating guaranteed throughput limits. If so, raise an error if we're overallocating. - Only runs check if tpm_limit_type or rpm_limit_type is "guaranteed_throughput" + Only runs check if tpm_limit_type or rpm_limit_type is "guaranteed_throughput". + The model-specific half also runs for every team key when + general_settings.enforce_team_model_limit_allocation is on; the aggregate + half needs no such gate, since the team's own rpm/tpm descriptor still binds + every key at request time regardless of what each key was granted. """ - if data.tpm_limit_type != "guaranteed_throughput" and data.rpm_limit_type != "guaranteed_throughput": + is_guaranteed_throughput = ( + data.tpm_limit_type == "guaranteed_throughput" or data.rpm_limit_type == "guaranteed_throughput" + ) + if not is_guaranteed_throughput and not _is_team_model_limit_allocation_enforced(): return # get all team keys # calculate allocated tpm/rpm limit @@ -1338,6 +1388,8 @@ async def _check_team_key_limits( team_table=team_table, data=data, ) + if not is_guaranteed_throughput: + return check_team_key_rpm_tpm_limits( keys=keys, team_table=team_table, diff --git a/tests/test_litellm/proxy/management_endpoints/test_key_management_endpoints.py b/tests/test_litellm/proxy/management_endpoints/test_key_management_endpoints.py index 867ef759fb3..f79ea1d7606 100644 --- a/tests/test_litellm/proxy/management_endpoints/test_key_management_endpoints.py +++ b/tests/test_litellm/proxy/management_endpoints/test_key_management_endpoints.py @@ -15256,3 +15256,217 @@ async def test_rotate_master_key_rotates_sso_identity_assertions( prisma_client=mock_prisma_client, new_master_key="sk-new-master-key", ) + + +def _team_with_model_limits( + team_id: str, + metadata: dict, + tpm_limit: int | None = None, + rpm_limit: int | None = None, +) -> LiteLLM_TeamTableCachedObj: + return LiteLLM_TeamTableCachedObj( + team_id=team_id, + team_alias="test-team", + tpm_limit=tpm_limit, + rpm_limit=rpm_limit, + max_budget=100.0, + spend=0.0, + models=[], + blocked=False, + members_with_roles=[], + metadata=metadata, + ) + + +@pytest.mark.parametrize( + "team_metadata_key, budget_field, expected_detail", + [ + ( + "model_rpm_limit", + "rpm_limit", + "Allocated RPM limit=0 + Key RPM limit=500 is greater than team RPM limit=100", + ), + ( + "model_tpm_limit", + "tpm_limit", + "Allocated TPM limit=0 + Key TPM limit=500 is greater than team TPM limit=100", + ), + ], +) +def test_check_team_key_model_specific_limits_counts_request_model_max_budget( + team_metadata_key, budget_field, expected_detail +): + """A per-model limit granted through model_max_budget is the same override the + limiter enforces at runtime, so it must count against the team's allocation.""" + team_table = _team_with_model_limits( + team_id="test-team-mmb-request", + metadata={team_metadata_key: {"gpt-4": 100}}, + ) + data = GenerateKeyRequest(model_max_budget={"gpt-4": {budget_field: 500}}) + + with pytest.raises(HTTPException) as exc_info: + check_team_key_model_specific_limits( + keys=[], + team_table=team_table, + data=data, + ) + + assert exc_info.value.status_code == 400 + assert expected_detail in str(exc_info.value.detail) + + +def test_check_team_key_model_specific_limits_counts_existing_key_model_max_budget(): + """Existing keys hold their per-model limits in model_max_budget just as often + as in metadata; both must count toward what the team has already handed out.""" + keys = [ + LiteLLM_VerificationToken( + token="test-token-mmb-1", + team_id="test-team-mmb-existing", + model_max_budget={"gpt-4": {"rpm_limit": 80}}, + ), + ] + team_table = _team_with_model_limits( + team_id="test-team-mmb-existing", + metadata={"model_rpm_limit": {"gpt-4": 100}}, + ) + data = GenerateKeyRequest(model_rpm_limit={"gpt-4": 30}) + + with pytest.raises(HTTPException) as exc_info: + check_team_key_model_specific_limits( + keys=keys, + team_table=team_table, + data=data, + ) + + assert exc_info.value.status_code == 400 + assert ( + "Allocated RPM limit=80 + Key RPM limit=30 is greater than team RPM limit=100" + in str(exc_info.value.detail) + ) + + +def test_check_team_key_model_specific_limits_prefers_metadata_over_model_max_budget(): + """Mirrors runtime precedence: a metadata map wins outright, so a + model_max_budget entry for a model the map omits is not additionally counted.""" + keys = [ + LiteLLM_VerificationToken( + token="test-token-precedence", + team_id="test-team-precedence", + metadata={"model_rpm_limit": {"gpt-4": 40}}, + model_max_budget={"gpt-4": {"rpm_limit": 900}}, + ), + ] + team_table = _team_with_model_limits( + team_id="test-team-precedence", + metadata={"model_rpm_limit": {"gpt-4": 100}}, + ) + + check_team_key_model_specific_limits( + keys=keys, + team_table=team_table, + data=GenerateKeyRequest(model_rpm_limit={"gpt-4": 60}), + ) + + +def _prisma_client_returning(keys: list) -> AsyncMock: + mock_prisma_client = AsyncMock() + mock_prisma_client.db.litellm_verificationtoken.find_many = AsyncMock(return_value=keys) + return mock_prisma_client + + +@pytest.mark.asyncio +async def test_check_team_key_limits_enforces_model_allocation_when_flag_on(monkeypatch): + """With enforce_team_model_limit_allocation on, a plain create (no limit_type) + can no longer be granted a per-model limit above the team's.""" + monkeypatch.setattr( + "litellm.proxy.proxy_server.general_settings", + {"enforce_team_model_limit_allocation": True}, + ) + mock_prisma_client = _prisma_client_returning([]) + team_table = _team_with_model_limits( + team_id="test-team-flag-on", + metadata={"model_rpm_limit": {"gpt-4": 100}}, + ) + + with pytest.raises(HTTPException) as exc_info: + await _check_team_key_limits( + team_table=team_table, + data=GenerateKeyRequest(model_rpm_limit={"gpt-4": 500}), + prisma_client=mock_prisma_client, + ) + + assert exc_info.value.status_code == 400 + assert "is greater than team RPM limit=100" in str(exc_info.value.detail) + + +@pytest.mark.asyncio +async def test_check_team_key_limits_allows_model_overallocation_by_default(monkeypatch): + """Default stays permissive: without the flag and without guaranteed_throughput, + the same create is granted and the team's keys are never even read.""" + monkeypatch.setattr("litellm.proxy.proxy_server.general_settings", {}) + mock_prisma_client = _prisma_client_returning([]) + team_table = _team_with_model_limits( + team_id="test-team-flag-off", + metadata={"model_rpm_limit": {"gpt-4": 100}}, + ) + + await _check_team_key_limits( + team_table=team_table, + data=GenerateKeyRequest(model_rpm_limit={"gpt-4": 500}), + prisma_client=mock_prisma_client, + ) + + mock_prisma_client.db.litellm_verificationtoken.find_many.assert_not_called() + + +@pytest.mark.asyncio +async def test_check_team_key_limits_allows_model_limit_within_remaining_allocation(monkeypatch): + """The flag rejects overallocation, not allocation: a key that fits in what the + team has left is still granted.""" + monkeypatch.setattr( + "litellm.proxy.proxy_server.general_settings", + {"enforce_team_model_limit_allocation": True}, + ) + mock_prisma_client = _prisma_client_returning( + [ + LiteLLM_VerificationToken( + token="test-token-allocated", + team_id="test-team-flag-within", + metadata={"model_rpm_limit": {"gpt-4": 60}}, + ) + ] + ) + team_table = _team_with_model_limits( + team_id="test-team-flag-within", + metadata={"model_rpm_limit": {"gpt-4": 100}}, + ) + + await _check_team_key_limits( + team_table=team_table, + data=GenerateKeyRequest(model_rpm_limit={"gpt-4": 30}), + prisma_client=mock_prisma_client, + ) + + +@pytest.mark.asyncio +async def test_check_team_key_limits_flag_does_not_gate_aggregate_limits(monkeypatch): + """The flag covers per-model allocation only. The team's aggregate rpm/tpm + descriptor still binds every key at request time, so a plain create asking for + more than the team's aggregate stays a runtime concern, not a grant-time one.""" + monkeypatch.setattr( + "litellm.proxy.proxy_server.general_settings", + {"enforce_team_model_limit_allocation": True}, + ) + mock_prisma_client = _prisma_client_returning([]) + team_table = _team_with_model_limits( + team_id="test-team-flag-aggregate", + metadata={}, + tpm_limit=1000, + rpm_limit=100, + ) + + await _check_team_key_limits( + team_table=team_table, + data=GenerateKeyRequest(tpm_limit=5000, rpm_limit=500), + prisma_client=mock_prisma_client, + ) diff --git a/ui/litellm-dashboard/src/lib/http/schema.d.ts b/ui/litellm-dashboard/src/lib/http/schema.d.ts index 2ad979372b3..6ec51f512e2 100644 --- a/ui/litellm-dashboard/src/lib/http/schema.d.ts +++ b/ui/litellm-dashboard/src/lib/http/schema.d.ts @@ -22789,6 +22789,11 @@ export interface components { * @default false */ enable_public_model_hub: boolean; + /** + * Enforce Team Model Limit Allocation + * @description If True, a team key's per-model rpm/tpm limits are rejected at create/update time when they would overallocate the team's per-model limits, regardless of rpm_limit_type/tpm_limit_type. A key's own per-model limit wins over the team's at request time, so this is how the team limit becomes a true cap. Default is False. + */ + enforce_team_model_limit_allocation?: boolean | null; /** * Forward Client Headers To Llm Api * @description If True, forwards client headers (e.g. Authorization) to the LLM API. Required for Claude Code with Max subscription. @@ -23584,7 +23589,7 @@ export interface components { * @description Default role assigned to new users created * @default internal_user_viewer */ - user_role: ("internal_user" | "internal_user_viewer" | "proxy_admin" | "proxy_admin_viewer") | null; + user_role: ("proxy_admin" | "proxy_admin_viewer" | "internal_user" | "internal_user_viewer") | null; }; /** * DefaultTeamSSOParams