tiennm99/litellm
1
0
Fork 0
mirror of https://github.com/tiennm99/litellm.git synced 2026-07-14 22:21:29 +00:00
Code Issues Packages Projects Releases Wiki Activity

Merge pull request #15120 from BerriAI/litellm_dev_10_01_2025_p2

Browse Source 
(feat) Support 'guaranteed_throughput' when setting limits on keys belonging to a team
This commit is contained in:
Krish Dholakia
2025-10-04 15:22:47 -07:00
committed by GitHub
parent b188f76dbe 691a1fa6fe
commit 75e3f4e10d
6 changed files with 938 additions and 86 deletions
Download Patch File Download Diff File Expand all files Collapse all files
litellm/proxy/_types.py
+9
View File
@@ -731,6 +731,7 @@ class GenerateRequestBase(LiteLLMPydanticObjectBase):
metadata: Optional[dict] = {}
tpm_limit: Optional[int] = None
rpm_limit: Optional[int] = None
budget_duration: Optional[str] = None
allowed_cache_controls: Optional[list] = []
config: Optional[dict] = {}
@@ -755,6 +756,12 @@ class KeyRequestBase(GenerateRequestBase):
tags: Optional[List[str]] = None
enforced_params: Optional[List[str]] = None
allowed_routes: Optional[list] = []
rpm_limit_type: Optional[
Literal["guaranteed_throughput", "best_effort_throughput"]
] = None # raise an error if 'guaranteed_throughput' is set and we're overallocating rpm
tpm_limit_type: Optional[
Literal["guaranteed_throughput", "best_effort_throughput"]
] = None # raise an error if 'guaranteed_throughput' is set and we're overallocating tpm
class LiteLLMKeyType(str, enum.Enum):
@@ -3056,6 +3063,8 @@ class PassThroughEndpointLoggingTypedDict(TypedDict):
LiteLLM_ManagementEndpoint_MetadataFields = [
"model_rpm_limit",
"model_tpm_limit",
"rpm_limit_type",
"tpm_limit_type",
"guardrails",
"tags",
"enforced_params",
litellm/proxy/management_endpoints/key_management_endpoints.py
+215 -4
View File
@@ -27,6 +27,7 @@ from litellm.caching import DualCache
from litellm.constants import LENGTH_OF_LITELLM_GENERATED_KEY, UI_SESSION_TOKEN_TEAM_ID
from litellm.litellm_core_utils.duration_parser import duration_in_seconds
from litellm.proxy._types import *
from litellm.proxy._types import LiteLLM_VerificationToken
from litellm.proxy.auth.auth_checks import (
_cache_key_object,
_delete_cache_key_object,
@@ -549,6 +550,15 @@ async def _common_key_generation_helper( # noqa: PLR0915
value=getattr(data, field),
)
for field in LiteLLM_ManagementEndpoint_MetadataFields:
if getattr(data, field, None) is not None:
_set_object_metadata_field(
object_data=data,
field_name=field,
value=getattr(data, field),
)
delattr(data, field)
data_json = data.model_dump(exclude_unset=True, exclude_none=True) # type: ignore
data_json = handle_key_type(data, data_json)
@@ -627,6 +637,153 @@ async def _common_key_generation_helper( # noqa: PLR0915
return response
def check_team_key_model_specific_limits(
keys: List[LiteLLM_VerificationToken],
team_table: LiteLLM_TeamTableCachedObj,
data: Union[GenerateKeyRequest, UpdateKeyRequest],
) -> None:
"""
Check if the team key is allocating model specific limits. If so, raise an error if we're overallocating.
"""
if data.model_rpm_limit is None and data.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
)
if data.model_rpm_limit is not None:
for model, rpm_limit in data.model_rpm_limit.items():
if (
model_specific_rpm_limit.get(model, 0) + rpm_limit
> team_table.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 team RPM limit={team_table.rpm_limit}",
)
elif team_table.metadata and team_table.metadata.get("model_rpm_limit"):
team_model_specific_rpm_limit_dict = team_table.metadata.get(
"model_rpm_limit", {}
)
team_model_specific_rpm_limit = team_model_specific_rpm_limit_dict.get(
model
)
if (
model_specific_rpm_limit.get(model, 0) + rpm_limit
> team_model_specific_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 team RPM limit={team_model_specific_rpm_limit.get(model, 0)}",
)
if data.model_tpm_limit is not None:
for model, tpm_limit in data.model_tpm_limit.items():
if (
team_table.tpm_limit is not None
and model_specific_tpm_limit.get(model, 0) + tpm_limit
> team_table.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 team TPM limit={team_table.tpm_limit}",
)
elif team_table.metadata and team_table.metadata.get("model_tpm_limit"):
team_model_specific_tpm_limit_dict = team_table.metadata.get(
"model_tpm_limit", {}
)
team_model_specific_tpm_limit = team_model_specific_tpm_limit_dict.get(
model
)
if (
team_model_specific_tpm_limit
and model_specific_tpm_limit.get(model, 0) + tpm_limit
> team_model_specific_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 team TPM limit={team_model_specific_tpm_limit}",
)
def check_team_key_rpm_tpm_limits(
keys: List[LiteLLM_VerificationToken],
team_table: LiteLLM_TeamTableCachedObj,
data: Union[GenerateKeyRequest, UpdateKeyRequest],
) -> None:
"""
Check if the team key is allocating rpm/tpm limits. If so, raise an error if we're overallocating.
"""
if keys is not None and len(keys) > 0:
allocated_tpm = sum(key.tpm_limit for key in keys if key.tpm_limit is not None)
allocated_rpm = sum(key.rpm_limit for key in keys if key.rpm_limit is not None)
else:
allocated_tpm = 0
allocated_rpm = 0
if (
data.tpm_limit is not None
and team_table.tpm_limit is not None
and data.tpm_limit + allocated_tpm > team_table.tpm_limit
):
raise HTTPException(
status_code=400,
detail=f"Allocated TPM limit={allocated_tpm} + Key TPM limit={data.tpm_limit} is greater than team TPM limit={team_table.tpm_limit}",
)
if (
data.rpm_limit is not None
and team_table.rpm_limit is not None
and data.rpm_limit + allocated_rpm > team_table.rpm_limit
):
raise HTTPException(
status_code=400,
detail=f"Allocated RPM limit={allocated_rpm} + Key RPM limit={data.rpm_limit} is greater than team RPM limit={team_table.rpm_limit}",
)
async def _check_team_key_limits(
team_table: LiteLLM_TeamTableCachedObj,
data: Union[GenerateKeyRequest, UpdateKeyRequest],
prisma_client: PrismaClient,
) -> None:
"""
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"
"""
if (
data.tpm_limit_type != "guaranteed_throughput"
and data.rpm_limit_type != "guaranteed_throughput"
):
return
# get all team keys
# calculate allocated tpm/rpm limit
# check if specified tpm/rpm limit is greater than allocated tpm/rpm limit
keys = await prisma_client.db.litellm_verificationtoken.find_many(
where={"team_id": team_table.team_id},
)
check_team_key_model_specific_limits(
keys=keys,
team_table=team_table,
data=data,
)
check_team_key_rpm_tpm_limits(
keys=keys,
team_table=team_table,
data=data,
)
@router.post(
"/key/generate",
tags=["key management"],
@@ -668,6 +825,8 @@ async def generate_key_fn(
- model_max_budget: Optional[Dict[str, BudgetConfig]] - Model-specific budgets {"gpt-4": {"budget_limit": 0.0005, "time_period": "30d"}}}. IF null or {} then no model specific budget.
- model_rpm_limit: Optional[dict] - key-specific model rpm limit. Example - {"text-davinci-002": 1000, "gpt-3.5-turbo": 1000}. IF null or {} then no model specific rpm limit.
- model_tpm_limit: Optional[dict] - key-specific model tpm limit. Example - {"text-davinci-002": 1000, "gpt-3.5-turbo": 1000}. IF null or {} then no model specific tpm limit.
- tpm_limit_type: Optional[str] - Type of tpm limit. Options: "best_effort_throughput" (no error if we're overallocating tpm), "guaranteed_throughput" (raise an error if we're overallocating tpm). Defaults to "best_effort_throughput".
- rpm_limit_type: Optional[str] - Type of rpm limit. Options: "best_effort_throughput" (no error if we're overallocating rpm), "guaranteed_throughput" (raise an error if we're overallocating rpm). Defaults to "best_effort_throughput".
- allowed_cache_controls: Optional[list] - List of allowed cache control values. Example - ["no-cache", "no-store"]. See all values - https://docs.litellm.ai/docs/proxy/caching#turn-on--off-caching-per-request
- blocked: Optional[bool] - Whether the key is blocked.
- rpm_limit: Optional[int] - Specify rpm limit for a given key (Requests per minute)
@@ -703,12 +862,19 @@ async def generate_key_fn(
- user_id: (str) Unique user id - used for tracking spend across multiple keys for same user id.
"""
try:
from litellm.proxy._types import CommonProxyErrors
from litellm.proxy.proxy_server import (
prisma_client,
user_api_key_cache,
user_custom_key_generate,
)
if prisma_client is None:
raise HTTPException(
status_code=status.HTTP_500_INTERNAL_SERVER_ERROR,
detail={"error": CommonProxyErrors.db_not_connected_error.value},
)
verbose_proxy_logger.debug("entered /key/generate")
if user_custom_key_generate is not None:
@@ -736,7 +902,6 @@ async def generate_key_fn(
verbose_proxy_logger.debug(
f"Error getting team object in `/key/generate`: {e}"
)
team_table = None
key_generation_check(
team_table=team_table,
@@ -745,12 +910,21 @@ async def generate_key_fn(
route=KeyManagementRoutes.KEY_GENERATE,
)
if team_table is not None:
await _check_team_key_limits(
team_table=team_table,
data=data,
prisma_client=prisma_client,
)
return await _common_key_generation_helper(
data=data,
user_api_key_dict=user_api_key_dict,
litellm_changed_by=litellm_changed_by,
team_table=team_table,
)
except HTTPException as e:
raise e
except Exception as e:
verbose_proxy_logger.exception(
"litellm.proxy.proxy_server.generate_key_fn(): Exception occured - {}".format(
@@ -804,6 +978,8 @@ async def generate_service_account_key_fn(
- model_max_budget: Optional[Dict[str, BudgetConfig]] - Model-specific budgets {"gpt-4": {"budget_limit": 0.0005, "time_period": "30d"}}}. IF null or {} then no model specific budget.
- model_rpm_limit: Optional[dict] - key-specific model rpm limit. Example - {"text-davinci-002": 1000, "gpt-3.5-turbo": 1000}. IF null or {} then no model specific rpm limit.
- model_tpm_limit: Optional[dict] - key-specific model tpm limit. Example - {"text-davinci-002": 1000, "gpt-3.5-turbo": 1000}. IF null or {} then no model specific tpm limit.
- tpm_limit_type: Optional[str] - TPM rate limit type - "best_effort_throughput" or "guaranteed_throughput"
- rpm_limit_type: Optional[str] - RPM rate limit type - "best_effort_throughput" or "guaranteed_throughput"
- allowed_cache_controls: Optional[list] - List of allowed cache control values. Example - ["no-cache", "no-store"]. See all values - https://docs.litellm.ai/docs/proxy/caching#turn-on--off-caching-per-request
- blocked: Optional[bool] - Whether the key is blocked.
- rpm_limit: Optional[int] - Specify rpm limit for a given key (Requests per minute)
@@ -832,12 +1008,19 @@ async def generate_service_account_key_fn(
- user_id: (str) Unique user id - used for tracking spend across multiple keys for same user id.
"""
from litellm.proxy._types import CommonProxyErrors
from litellm.proxy.proxy_server import (
prisma_client,
user_api_key_cache,
user_custom_key_generate,
)
if prisma_client is None:
raise HTTPException(
status_code=status.HTTP_500_INTERNAL_SERVER_ERROR,
detail={"error": CommonProxyErrors.db_not_connected_error.value},
)
await validate_team_id_used_in_service_account_request(
team_id=data.team_id,
prisma_client=prisma_client,
@@ -870,6 +1053,13 @@ async def generate_service_account_key_fn(
)
team_table = None
if team_table is not None:
await _check_team_key_limits(
team_table=team_table,
data=data,
prisma_client=prisma_client,
)
key_generation_check(
team_table=team_table,
user_api_key_dict=user_api_key_dict,
@@ -1096,6 +1286,8 @@ async def update_key_fn(
- rpm_limit: Optional[int] - Requests per minute limit
- model_rpm_limit: Optional[dict] - Model-specific RPM limits {"gpt-4": 100, "claude-v1": 200}
- model_tpm_limit: Optional[dict] - Model-specific TPM limits {"gpt-4": 100000, "claude-v1": 200000}
- tpm_limit_type: Optional[str] - TPM rate limit type - "best_effort_throughput" or "guaranteed_throughput"
- rpm_limit_type: Optional[str] - RPM rate limit type - "best_effort_throughput" or "guaranteed_throughput"
- allowed_cache_controls: Optional[list] - List of allowed cache control values
- duration: Optional[str] - Key validity duration ("30d", "1h", etc.)
- permissions: Optional[dict] - Key-specific permissions
@@ -1170,14 +1362,25 @@ async def update_key_fn(
user_api_key_cache=user_api_key_cache,
)
# if team change - check if this is possible
if is_different_team(data=data, existing_key_row=existing_key_row):
# Only check team limits if key has a team_id
team_obj: Optional[LiteLLM_TeamTableCachedObj] = None
if data.team_id is not None:
team_obj = await get_team_object(
team_id=cast(str, data.team_id),
team_id=data.team_id,
prisma_client=prisma_client,
user_api_key_cache=user_api_key_cache,
check_db_only=True,
)
if team_obj is not None:
await _check_team_key_limits(
team_table=team_obj,
data=data,
prisma_client=prisma_client,
)
# if team change - check if this is possible
if is_different_team(data=data, existing_key_row=existing_key_row):
if llm_router is None:
raise HTTPException(
status_code=400,
@@ -1185,6 +1388,14 @@ async def update_key_fn(
"error": "LLM router not found. Please set it up by passing in a valid config.yaml or adding models via the UI."
},
)
# team_obj should be set since is_different_team() returns True only when data.team_id is not None
if team_obj is None:
raise HTTPException(
status_code=500,
detail={
"error": "Team object not found for team change validation"
},
)
validate_key_team_change(
key=existing_key_row,
team=team_obj,
tests/test_litellm/proxy/management_endpoints/test_key_management_endpoints.py
+586 -82
View File
@@ -15,14 +15,17 @@ from fastapi import HTTPException
from litellm.proxy._types import (
GenerateKeyRequest,
LiteLLM_TeamTableCachedObj,
LiteLLM_VerificationToken,
LitellmUserRoles,
UpdateKeyRequest,
)
from litellm.proxy.auth.user_api_key_auth import UserAPIKeyAuth
from litellm.proxy.management_endpoints.key_management_endpoints import (
_check_team_key_limits,
_common_key_generation_helper,
_list_key_helper,
check_team_key_model_specific_limits,
generate_key_helper_fn,
prepare_key_update_data,
validate_key_team_change,
@@ -847,17 +850,24 @@ async def test_generate_service_account_key_endpoint_validation():
)
# Test case 1: Missing team_id
with pytest.raises(HTTPException) as exc_info:
await generate_service_account_key_fn(
data=GenerateKeyRequest(team_id=None),
user_api_key_dict=UserAPIKeyAuth(
user_role=LitellmUserRoles.PROXY_ADMIN, api_key="sk-1"
),
litellm_changed_by=None,
)
with patch("litellm.proxy.proxy_server.prisma_client") as mock_prisma:
# Mock prisma_client to be not None so we can reach team_id validation
mock_prisma_instance = AsyncMock()
mock_prisma.return_value = mock_prisma_instance
assert exc_info.value.status_code == 400
assert "team_id is required for service account keys" in str(exc_info.value.detail)
with pytest.raises(HTTPException) as exc_info:
await generate_service_account_key_fn(
data=GenerateKeyRequest(team_id=None),
user_api_key_dict=UserAPIKeyAuth(
user_role=LitellmUserRoles.PROXY_ADMIN, api_key="sk-1"
),
litellm_changed_by=None,
)
assert exc_info.value.status_code == 400
assert "team_id is required for service account keys" in str(
exc_info.value.detail
)
# Test case 2: Team doesn't exist in database
with patch("litellm.proxy.proxy_server.prisma_client") as mock_prisma:
@@ -1040,7 +1050,7 @@ async def test_unblock_key_invalid_key_format(monkeypatch):
def test_validate_key_team_change_with_member_permissions():
"""
Test validate_key_team_change function with team member permissions.
This test covers the new logic that allows team members with specific
permissions to update keys, not just team admins.
"""
@@ -1054,111 +1064,107 @@ def test_validate_key_team_change_with_member_permissions():
mock_key.models = ["gpt-4"]
mock_key.tpm_limit = None
mock_key.rpm_limit = None
mock_team = MagicMock()
mock_team.team_id = "test-team-456"
mock_team.team_id = "test-team-456"
mock_team.members_with_roles = []
mock_team.tpm_limit = None
mock_team.rpm_limit = None
mock_change_initiator = MagicMock()
mock_change_initiator.user_id = "test-user-123"
mock_router = MagicMock()
# Mock the member object returned by _get_user_in_team
mock_member_object = MagicMock()
with patch('litellm.proxy.management_endpoints.key_management_endpoints.can_team_access_model'):
with patch('litellm.proxy.management_endpoints.key_management_endpoints._get_user_in_team') as mock_get_user:
with patch('litellm.proxy.management_endpoints.key_management_endpoints._is_user_team_admin') as mock_is_admin:
with patch('litellm.proxy.management_endpoints.key_management_endpoints.TeamMemberPermissionChecks.does_team_member_have_permissions_for_endpoint') as mock_has_perms:
with patch(
"litellm.proxy.management_endpoints.key_management_endpoints.can_team_access_model"
):
with patch(
"litellm.proxy.management_endpoints.key_management_endpoints._get_user_in_team"
) as mock_get_user:
with patch(
"litellm.proxy.management_endpoints.key_management_endpoints._is_user_team_admin"
) as mock_is_admin:
with patch(
"litellm.proxy.management_endpoints.key_management_endpoints.TeamMemberPermissionChecks.does_team_member_have_permissions_for_endpoint"
) as mock_has_perms:
mock_get_user.return_value = mock_member_object
mock_is_admin.return_value = False
mock_has_perms.return_value = True
# This should not raise an exception due to member permissions
validate_key_team_change(
key=mock_key,
team=mock_team,
change_initiated_by=mock_change_initiator,
llm_router=mock_router
llm_router=mock_router,
)
# Verify the permission check was called with correct parameters
mock_has_perms.assert_called_once_with(
team_member_object=mock_member_object,
team_table=mock_team,
route=KeyManagementRoutes.KEY_UPDATE.value
route=KeyManagementRoutes.KEY_UPDATE.value,
)
def test_key_rotation_fields_helper():
"""
Test the key data update logic for rotation fields.
This test focuses on the core logic that adds rotation fields to key_data
when auto_rotate is enabled, without the complexity of full key generation.
"""
# Test Case 1: With rotation enabled
key_data = {
"models": ["gpt-3.5-turbo"],
"user_id": "test-user"
}
key_data = {"models": ["gpt-3.5-turbo"], "user_id": "test-user"}
auto_rotate = True
rotation_interval = "30d"
# Simulate the rotation logic from generate_key_helper_fn
if auto_rotate and rotation_interval:
key_data.update({
"auto_rotate": auto_rotate,
"rotation_interval": rotation_interval
})
key_data.update(
{"auto_rotate": auto_rotate, "rotation_interval": rotation_interval}
)
# Verify rotation fields are added
assert key_data["auto_rotate"] == True
assert key_data["rotation_interval"] == "30d"
assert key_data["models"] == ["gpt-3.5-turbo"] # Original fields preserved
# Test Case 2: Without rotation enabled
key_data2 = {
"models": ["gpt-4"],
"user_id": "test-user"
}
key_data2 = {"models": ["gpt-4"], "user_id": "test-user"}
auto_rotate2 = False
rotation_interval2 = None
# Simulate the rotation logic
if auto_rotate2 and rotation_interval2:
key_data2.update({
"auto_rotate": auto_rotate2,
"rotation_interval": rotation_interval2
})
key_data2.update(
{"auto_rotate": auto_rotate2, "rotation_interval": rotation_interval2}
)
# Verify rotation fields are NOT added
assert "auto_rotate" not in key_data2
assert "rotation_interval" not in key_data2
assert key_data2["models"] == ["gpt-4"] # Original fields preserved
# Test Case 3: auto_rotate=True but no interval
key_data3 = {
"models": ["claude-3"],
"user_id": "test-user"
}
key_data3 = {"models": ["claude-3"], "user_id": "test-user"}
auto_rotate3 = True
rotation_interval3 = None
# Simulate the rotation logic
if auto_rotate3 and rotation_interval3:
key_data3.update({
"auto_rotate": auto_rotate3,
"rotation_interval": rotation_interval3
})
key_data3.update(
{"auto_rotate": auto_rotate3, "rotation_interval": rotation_interval3}
)
# Verify rotation fields are NOT added (missing interval)
assert "auto_rotate" not in key_data3
assert "rotation_interval" not in key_data3
@@ -1181,27 +1187,24 @@ async def test_update_key_fn_auto_rotate_enable():
team_id=None,
auto_rotate=False,
rotation_interval=None,
metadata={}
metadata={},
)
# Test enabling auto rotation
update_request = UpdateKeyRequest(
key="test-token",
auto_rotate=True,
rotation_interval="30d"
key="test-token", auto_rotate=True, rotation_interval="30d"
)
result = await prepare_key_update_data(
data=update_request,
existing_key_row=existing_key
data=update_request, existing_key_row=existing_key
)
# Verify rotation fields are included
assert result["auto_rotate"] is True
assert result["rotation_interval"] == "30d"
@pytest.mark.asyncio
@pytest.mark.asyncio
async def test_update_key_fn_auto_rotate_disable():
"""Test that update_key_fn properly handles disabling auto rotation."""
from litellm.proxy._types import LiteLLM_VerificationToken, UpdateKeyRequest
@@ -1218,19 +1221,520 @@ async def test_update_key_fn_auto_rotate_disable():
team_id=None,
auto_rotate=True,
rotation_interval="30d",
metadata={}
metadata={},
)
# Test disabling auto rotation
update_request = UpdateKeyRequest(
key="test-token",
auto_rotate=False
)
update_request = UpdateKeyRequest(key="test-token", auto_rotate=False)
result = await prepare_key_update_data(
data=update_request,
existing_key_row=existing_key
data=update_request, existing_key_row=existing_key
)
# Verify auto_rotate is set to False
assert result["auto_rotate"] is False
@pytest.mark.asyncio
async def test_check_team_key_limits_no_existing_keys():
"""
Test _check_team_key_limits when team has no existing keys.
Should allow any TPM/RPM limits within team bounds.
"""
# Mock prisma client
mock_prisma_client = AsyncMock()
mock_prisma_client.db.litellm_verificationtoken.find_many = AsyncMock(
return_value=[]
)
# Create team table with limits
team_table = LiteLLM_TeamTableCachedObj(
team_id="test-team-123",
team_alias="test-team",
tpm_limit=10000,
rpm_limit=1000,
max_budget=100.0,
spend=0.0,
models=[],
blocked=False,
members_with_roles=[],
)
# Create request with limits within team bounds
data = GenerateKeyRequest(
tpm_limit=5000,
rpm_limit=500,
tpm_limit_type="guaranteed_throughput",
rpm_limit_type="guaranteed_throughput",
)
# Should not raise any exception
await _check_team_key_limits(
team_table=team_table,
data=data,
prisma_client=mock_prisma_client,
)
# Verify database was queried
mock_prisma_client.db.litellm_verificationtoken.find_many.assert_called_once_with(
where={"team_id": "test-team-123"}
)
@pytest.mark.asyncio
async def test_check_team_key_limits_with_existing_keys_within_bounds():
"""
Test _check_team_key_limits when team has existing keys but total allocation
is still within team limits.
"""
# Create mock existing keys
existing_key1 = MagicMock()
existing_key1.tpm_limit = 3000
existing_key1.rpm_limit = 200
existing_key2 = MagicMock()
existing_key2.tpm_limit = 2000
existing_key2.rpm_limit = 300
existing_key3 = MagicMock()
existing_key3.tpm_limit = None # Should be ignored in calculation
existing_key3.rpm_limit = None # Should be ignored in calculation
mock_prisma_client = AsyncMock()
mock_prisma_client.db.litellm_verificationtoken.find_many = AsyncMock(
return_value=[existing_key1, existing_key2, existing_key3]
)
# Create team table with limits
team_table = LiteLLM_TeamTableCachedObj(
team_id="test-team-456",
team_alias="test-team",
tpm_limit=10000, # Total: 3000 + 2000 + 4000 (new) = 9000 < 10000 ✓
rpm_limit=1000, # Total: 200 + 300 + 400 (new) = 900 < 1000 ✓
max_budget=100.0,
spend=0.0,
models=[],
blocked=False,
members_with_roles=[],
)
# Create request that would still be within bounds
data = GenerateKeyRequest(
tpm_limit=4000,
rpm_limit=400,
)
# Should not raise any exception
await _check_team_key_limits(
team_table=team_table,
data=data,
prisma_client=mock_prisma_client,
)
@pytest.mark.asyncio
async def test_check_team_key_limits_tpm_overallocation():
"""
Test _check_team_key_limits when new key would cause TPM overallocation.
Should raise HTTPException with appropriate error message.
"""
# Create mock existing keys with high TPM usage
existing_key1 = MagicMock()
existing_key1.tpm_limit = 6000
existing_key1.rpm_limit = 100
existing_key1.metadata = {}
existing_key2 = MagicMock()
existing_key2.tpm_limit = 3000
existing_key2.rpm_limit = 200
existing_key2.metadata = {}
mock_prisma_client = AsyncMock()
mock_prisma_client.db.litellm_verificationtoken.find_many = AsyncMock(
return_value=[existing_key1, existing_key2]
)
# Create team table with limits
team_table = LiteLLM_TeamTableCachedObj(
team_id="test-team-789",
team_alias="test-team",
tpm_limit=10000, # Allocated: 6000 + 3000 = 9000, New: 2000, Total: 11000 > 10000 ✗
rpm_limit=1000,
max_budget=100.0,
spend=0.0,
models=[],
blocked=False,
members_with_roles=[],
)
# Create request that would exceed TPM limits
data = GenerateKeyRequest(
tpm_limit=2000,
rpm_limit=100,
tpm_limit_type="guaranteed_throughput",
)
# Should raise HTTPException for TPM overallocation
with pytest.raises(HTTPException) as exc_info:
await _check_team_key_limits(
team_table=team_table,
data=data,
prisma_client=mock_prisma_client,
)
assert exc_info.value.status_code == 400
assert (
"Allocated TPM limit=9000 + Key TPM limit=2000 is greater than team TPM limit=10000"
in str(exc_info.value.detail)
)
@pytest.mark.asyncio
async def test_check_team_key_limits_rpm_overallocation():
"""
Test _check_team_key_limits when new key would cause RPM overallocation.
Should raise HTTPException with appropriate error message.
"""
# Create mock existing keys with high RPM usage
existing_key1 = MagicMock()
existing_key1.tpm_limit = 1000
existing_key1.rpm_limit = 600
existing_key1.metadata = {}
existing_key2 = MagicMock()
existing_key2.tpm_limit = 2000
existing_key2.rpm_limit = 300
existing_key2.metadata = {}
mock_prisma_client = AsyncMock()
mock_prisma_client.db.litellm_verificationtoken.find_many = AsyncMock(
return_value=[existing_key1, existing_key2]
)
# Create team table with limits
team_table = LiteLLM_TeamTableCachedObj(
team_id="test-team-101",
team_alias="test-team",
tpm_limit=10000,
rpm_limit=1000, # Allocated: 600 + 300 = 900, New: 200, Total: 1100 > 1000 ✗
max_budget=100.0,
spend=0.0,
models=[],
blocked=False,
members_with_roles=[],
)
# Create request that would exceed RPM limits
data = GenerateKeyRequest(
tpm_limit=1000,
rpm_limit=200,
rpm_limit_type="guaranteed_throughput",
)
# Should raise HTTPException for RPM overallocation
with pytest.raises(HTTPException) as exc_info:
await _check_team_key_limits(
team_table=team_table,
data=data,
prisma_client=mock_prisma_client,
)
assert exc_info.value.status_code == 400
assert (
"Allocated RPM limit=900 + Key RPM limit=200 is greater than team RPM limit=1000"
in str(exc_info.value.detail)
)
@pytest.mark.asyncio
async def test_check_team_key_limits_no_team_limits():
"""
Test _check_team_key_limits when team has no TPM/RPM limits set.
Should allow any key limits since there are no team constraints.
"""
# Create mock existing keys
existing_key = MagicMock()
existing_key.tpm_limit = 5000
existing_key.rpm_limit = 500
mock_prisma_client = AsyncMock()
mock_prisma_client.db.litellm_verificationtoken.find_many = AsyncMock(
return_value=[existing_key]
)
# Create team table with no limits
team_table = LiteLLM_TeamTableCachedObj(
team_id="test-team-202",
team_alias="test-team",
tpm_limit=None, # No team limit
rpm_limit=None, # No team limit
max_budget=100.0,
spend=0.0,
models=[],
blocked=False,
members_with_roles=[],
)
# Create request with any limits
data = GenerateKeyRequest(
tpm_limit=10000, # High limit should be allowed
rpm_limit=2000, # High limit should be allowed
)
# Should not raise any exception
await _check_team_key_limits(
team_table=team_table,
data=data,
prisma_client=mock_prisma_client,
)
@pytest.mark.asyncio
async def test_check_team_key_limits_no_key_limits():
"""
Test _check_team_key_limits when new key has no TPM/RPM limits.
Should not raise any exceptions since no limits are being allocated.
"""
# Create mock existing keys
existing_key = MagicMock()
existing_key.tpm_limit = 8000
existing_key.rpm_limit = 800
mock_prisma_client = AsyncMock()
mock_prisma_client.db.litellm_verificationtoken.find_many = AsyncMock(
return_value=[existing_key]
)
# Create team table with limits
team_table = LiteLLM_TeamTableCachedObj(
team_id="test-team-303",
team_alias="test-team",
tpm_limit=10000,
rpm_limit=1000,
max_budget=100.0,
spend=0.0,
models=[],
blocked=False,
members_with_roles=[],
)
# Create request with no limits
data = GenerateKeyRequest(
tpm_limit=None, # No limit being set
rpm_limit=None, # No limit being set
)
# Should not raise any exception
await _check_team_key_limits(
team_table=team_table,
data=data,
prisma_client=mock_prisma_client,
)
@pytest.mark.asyncio
async def test_check_team_key_limits_mixed_scenarios():
"""
Test _check_team_key_limits with mixed scenarios:
- Some existing keys have limits, others don't
- New key has only one type of limit
- Team has only one type of limit
"""
# Create mock existing keys with mixed limits
existing_key1 = MagicMock()
existing_key1.tpm_limit = 3000
existing_key1.rpm_limit = None # No RPM limit
existing_key2 = MagicMock()
existing_key2.tpm_limit = None # No TPM limit
existing_key2.rpm_limit = 400
existing_key3 = MagicMock()
existing_key3.tpm_limit = 2000
existing_key3.rpm_limit = 300
mock_prisma_client = AsyncMock()
mock_prisma_client.db.litellm_verificationtoken.find_many = AsyncMock(
return_value=[existing_key1, existing_key2, existing_key3]
)
# Create team table with only TPM limit
team_table = LiteLLM_TeamTableCachedObj(
team_id="test-team-404",
team_alias="test-team",
tpm_limit=10000, # Allocated: 3000 + 0 + 2000 = 5000, New: 4000, Total: 9000 < 10000 ✓
rpm_limit=None, # No team RPM limit
max_budget=100.0,
spend=0.0,
models=[],
blocked=False,
members_with_roles=[],
)
# Create request with only TPM limit
data = GenerateKeyRequest(
tpm_limit=4000,
rpm_limit=None, # No RPM limit being set
)
# Should not raise any exception
await _check_team_key_limits(
team_table=team_table,
data=data,
prisma_client=mock_prisma_client,
)
@pytest.mark.asyncio
async def test_check_team_key_limits_exact_boundary():
"""
Test _check_team_key_limits when allocation exactly matches team limits.
Should allow the allocation (boundary case).
"""
# Create mock existing keys
existing_key = MagicMock()
existing_key.tpm_limit = 7000
existing_key.rpm_limit = 700
mock_prisma_client = AsyncMock()
mock_prisma_client.db.litellm_verificationtoken.find_many = AsyncMock(
return_value=[existing_key]
)
# Create team table with limits
team_table = LiteLLM_TeamTableCachedObj(
team_id="test-team-505",
team_alias="test-team",
tpm_limit=10000, # Allocated: 7000, New: 3000, Total: 10000 = 10000 ✓
rpm_limit=1000, # Allocated: 700, New: 300, Total: 1000 = 1000 ✓
max_budget=100.0,
spend=0.0,
models=[],
blocked=False,
members_with_roles=[],
)
# Create request that exactly matches remaining capacity
data = GenerateKeyRequest(
tpm_limit=3000,
rpm_limit=300,
)
# Should not raise any exception (exact boundary should be allowed)
await _check_team_key_limits(
team_table=team_table,
data=data,
prisma_client=mock_prisma_client,
)
def test_check_team_key_model_specific_limits_no_limits():
"""
Test check_team_key_model_specific_limits when no model-specific limits are set.
Should return without raising any exceptions.
"""
# Create existing key with no model-specific limits
existing_key = LiteLLM_VerificationToken(
token="test-token-1",
user_id="test-user",
team_id="test-team-123",
metadata={},
)
keys = [existing_key]
# Create team table
team_table = LiteLLM_TeamTableCachedObj(
team_id="test-team-123",
team_alias="test-team",
tpm_limit=10000,
rpm_limit=1000,
max_budget=100.0,
spend=0.0,
models=[],
blocked=False,
members_with_roles=[],
metadata={},
)
# Create request with no model-specific limits
data = GenerateKeyRequest(
model_rpm_limit=None,
model_tpm_limit=None,
)
# Should not raise any exception
check_team_key_model_specific_limits(
keys=keys,
team_table=team_table,
data=data,
)
def test_check_team_key_model_specific_limits_rpm_overallocation():
"""
Test check_team_key_model_specific_limits when model-specific RPM would cause overallocation.
Should raise HTTPException with appropriate error message.
"""
# Create existing keys with model-specific RPM limits
existing_key1 = LiteLLM_VerificationToken(
token="test-token-1",
user_id="test-user-1",
team_id="test-team-456",
metadata={
"model_rpm_limit": {
"gpt-4": 500,
"gpt-3.5-turbo": 300,
}
},
)
existing_key2 = LiteLLM_VerificationToken(
token="test-token-2",
user_id="test-user-2",
team_id="test-team-456",
metadata={
"model_rpm_limit": {
"gpt-4": 300,
}
},
)
keys = [existing_key1, existing_key2]
# Create team table with RPM limit
team_table = LiteLLM_TeamTableCachedObj(
team_id="test-team-456",
team_alias="test-team",
tpm_limit=10000,
rpm_limit=1000, # Total team RPM limit
max_budget=100.0,
spend=0.0,
models=[],
blocked=False,
members_with_roles=[],
metadata={},
)
# Create request that would exceed model-specific RPM limits
# Existing gpt-4: 500 + 300 = 800, New: 300, Total: 1100 > 1000 (team limit)
data = GenerateKeyRequest(
model_rpm_limit={
"gpt-4": 300, # This would cause overallocation
},
model_tpm_limit=None,
)
# Should raise HTTPException for model-specific RPM overallocation
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=800 + Key RPM limit=300 is greater than team RPM limit=1000"
in str(exc_info.value.detail)
)
ui/litellm-dashboard/src/components/common_components/RateLimitTypeFormItem.tsx
+98
View File
@@ -0,0 +1,98 @@
import React from "react"
import { Form, Select, Tooltip } from "antd"
import { InfoCircleOutlined } from "@ant-design/icons"
const { Option } = Select
interface RateLimitTypeFormItemProps {
/** The type of rate limit - either 'tpm' or 'rpm' */
type: 'tpm' | 'rpm'
/** The form field name */
name: string
/** Whether to show detailed descriptions (default: true) */
showDetailedDescriptions?: boolean
/** Additional CSS classes */
className?: string
/** Initial value for the field */
initialValue?: string | null
/** Form instance for setting field values */
form?: any
/** Custom onChange handler */
onChange?: (value: string) => void
}
export const RateLimitTypeFormItem: React.FC<RateLimitTypeFormItemProps> = ({
type,
name,
showDetailedDescriptions = true,
className = "",
initialValue = null,
form,
onChange
}) => {
const limitTypeUpper = type.toUpperCase()
const limitTypeLower = type.toLowerCase()
const handleChange = (value: string) => {
if (form) {
form.setFieldValue(name, value)
}
if (onChange) {
onChange(value)
}
}
const tooltipTitle = `Select 'guaranteed_throughput' to prevent overallocating ${limitTypeUpper} limit when the key belongs to a Team with specific ${limitTypeUpper} limits.`
return (
<Form.Item
label={
<span>
{limitTypeUpper} Rate Limit Type{' '}
<Tooltip title={tooltipTitle}>
<InfoCircleOutlined style={{ marginLeft: '4px' }} />
</Tooltip>
</span>
}
name={name}
initialValue={initialValue}
className={className}
>
<Select
defaultValue={showDetailedDescriptions ? "default" : undefined}
placeholder="Select rate limit type"
style={{ width: "100%" }}
optionLabelProp={showDetailedDescriptions ? "label" : undefined}
onChange={handleChange}
>
{showDetailedDescriptions ? (
<>
<Option value="best_effort_throughput" label="Default">
<div style={{ padding: '4px 0' }}>
<div style={{ fontWeight: 500 }}>Default</div>
<div style={{ fontSize: '11px', color: '#6b7280', marginTop: '2px' }}>
Best effort throughput - no error if we&apos;re overallocating {limitTypeLower} (Team/Key Limits checked at runtime).
</div>
</div>
</Option>
<Option value="guaranteed_throughput" label="Guaranteed throughput">
<div style={{ padding: '4px 0' }}>
<div style={{ fontWeight: 500 }}>Guaranteed throughput</div>
<div style={{ fontSize: '11px', color: '#6b7280', marginTop: '2px' }}>
Guaranteed throughput - raise an error if we&apos;re overallocating {limitTypeLower} (also checks model-specific limits)
</div>
</div>
</Option>
</>
) : (
<>
<Option value="best_effort_throughput">Best effort throughput</Option>
<Option value="guaranteed_throughput">Guaranteed throughput</Option>
</>
)}
</Select>
</Form.Item>
)
}
export default RateLimitTypeFormItem
ui/litellm-dashboard/src/components/organisms/create_key_button.tsx
+17
View File
@@ -35,6 +35,7 @@ import MCPServerSelector from "../mcp_server_management/MCPServerSelector"
import ModelAliasManager from "../common_components/ModelAliasManager"
import NotificationsManager from "../molecules/notifications_manager"
import KeyLifecycleSettings from "../common_components/KeyLifecycleSettings"
import RateLimitTypeFormItem from "../common_components/RateLimitTypeFormItem"
const { Option } = Select;
@@ -810,6 +811,14 @@ const CreateKey: React.FC<CreateKeyProps> = ({
>
<NumericalInput step={1} width={400} />
</Form.Item>
<RateLimitTypeFormItem
type="tpm"
name="tpm_limit_type"
className="mt-4"
initialValue={null}
form={form}
showDetailedDescriptions={true}
/>
<Form.Item
className="mt-4"
label={
@@ -841,6 +850,14 @@ const CreateKey: React.FC<CreateKeyProps> = ({
>
<NumericalInput step={1} width={400} />
</Form.Item>
<RateLimitTypeFormItem
type="rpm"
name="rpm_limit_type"
className="mt-4"
initialValue={null}
form={form}
showDetailedDescriptions={true}
/>
<Form.Item
label={
<span>
ui/litellm-dashboard/src/components/templates/key_edit_view.tsx
+13
View File
@@ -13,6 +13,7 @@ import { fetchMCPAccessGroups } from "../networking"
import { mapInternalToDisplayNames, mapDisplayToInternalNames } from "../callback_info_helpers"
import GuardrailSelector from "@/components/guardrails/GuardrailSelector"
import KeyLifecycleSettings from "../common_components/KeyLifecycleSettings"
import RateLimitTypeFormItem from "../common_components/RateLimitTypeFormItem"
interface KeyEditViewProps {
keyData: KeyResponse
@@ -222,10 +223,22 @@ export function KeyEditView({
<NumericalInput min={0} />
</Form.Item>
<RateLimitTypeFormItem
type="tpm"
name="tpm_limit_type"
showDetailedDescriptions={false}
/>
<Form.Item label="RPM Limit" name="rpm_limit">
<NumericalInput min={0} />
</Form.Item>
<RateLimitTypeFormItem
type="rpm"
name="rpm_limit_type"
showDetailedDescriptions={false}
/>
<Form.Item label="Max Parallel Requests" name="max_parallel_requests">
<NumericalInput min={0} />
</Form.Item>