diff --git a/litellm/proxy/_types.py b/litellm/proxy/_types.py index 0d375f363e..ae8e759f3c 100644 --- a/litellm/proxy/_types.py +++ b/litellm/proxy/_types.py @@ -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", diff --git a/litellm/proxy/management_endpoints/key_management_endpoints.py b/litellm/proxy/management_endpoints/key_management_endpoints.py index 0eb2d27463..8d357cef49 100644 --- a/litellm/proxy/management_endpoints/key_management_endpoints.py +++ b/litellm/proxy/management_endpoints/key_management_endpoints.py @@ -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, 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 e3aa7d5887..7620e27e4c 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 @@ -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) + ) diff --git a/ui/litellm-dashboard/src/components/common_components/RateLimitTypeFormItem.tsx b/ui/litellm-dashboard/src/components/common_components/RateLimitTypeFormItem.tsx new file mode 100644 index 0000000000..4c923c63f9 --- /dev/null +++ b/ui/litellm-dashboard/src/components/common_components/RateLimitTypeFormItem.tsx @@ -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 = ({ + 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 ( + + {limitTypeUpper} Rate Limit Type{' '} + + + + + } + name={name} + initialValue={initialValue} + className={className} + > + + + ) +} + +export default RateLimitTypeFormItem diff --git a/ui/litellm-dashboard/src/components/organisms/create_key_button.tsx b/ui/litellm-dashboard/src/components/organisms/create_key_button.tsx index dc523da766..4c26c6cb76 100644 --- a/ui/litellm-dashboard/src/components/organisms/create_key_button.tsx +++ b/ui/litellm-dashboard/src/components/organisms/create_key_button.tsx @@ -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 = ({ > + = ({ > + diff --git a/ui/litellm-dashboard/src/components/templates/key_edit_view.tsx b/ui/litellm-dashboard/src/components/templates/key_edit_view.tsx index 14032a046d..6c49cb8aaf 100644 --- a/ui/litellm-dashboard/src/components/templates/key_edit_view.tsx +++ b/ui/litellm-dashboard/src/components/templates/key_edit_view.tsx @@ -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({ + + + +