From 60063202517cc9748cfb2530e93313133a3e5753 Mon Sep 17 00:00:00 2001 From: Krrish Dholakia Date: Tue, 30 Sep 2025 12:50:21 -0700 Subject: [PATCH 01/96] fix(proxy/utils.py): run guardrails before running other logging hooks on "async_post_call_success_hook" Closes LIT-1152 --- litellm/proxy/utils.py | 58 ++++++++++++++++++++++-------------------- 1 file changed, 31 insertions(+), 27 deletions(-) diff --git a/litellm/proxy/utils.py b/litellm/proxy/utils.py index 5b11c25b2b..d3d2972abf 100644 --- a/litellm/proxy/utils.py +++ b/litellm/proxy/utils.py @@ -1395,9 +1395,12 @@ class ProxyLogging: 3. /image/generation 4. /files """ + from litellm.types.guardrails import GuardrailEventHooks - for callback in litellm.callbacks: - try: + guardrail_callbacks: List[CustomGuardrail] = [] + other_callbacks: List[CustomLogger] = [] + try: + for callback in litellm.callbacks: _callback: Optional[CustomLogger] = None if isinstance(callback, str): _callback = litellm.litellm_core_utils.litellm_logging.get_custom_logger_compatible_class( @@ -1407,36 +1410,37 @@ class ProxyLogging: _callback = callback # type: ignore if _callback is not None: + if isinstance(_callback, CustomGuardrail): + guardrail_callbacks.append(_callback) + else: + other_callbacks.append(_callback) ############## Handle Guardrails ######################################## ############################################################################# - if isinstance(callback, CustomGuardrail): - # Main - V2 Guardrails implementation - from litellm.types.guardrails import GuardrailEventHooks - if ( - callback.should_run_guardrail( - data=data, event_type=GuardrailEventHooks.post_call - ) - is not True - ): - continue + for callback in guardrail_callbacks: + # Main - V2 Guardrails implementation + if ( + callback.should_run_guardrail( + data=data, event_type=GuardrailEventHooks.post_call + ) + is not True + ): + continue - await callback.async_post_call_success_hook( - user_api_key_dict=user_api_key_dict, - data=data, - response=response, - ) + await callback.async_post_call_success_hook( + user_api_key_dict=user_api_key_dict, + data=data, + response=response, + ) - ############ Handle CustomLogger ############################### - ################################################################# - elif isinstance(_callback, CustomLogger): - await _callback.async_post_call_success_hook( - user_api_key_dict=user_api_key_dict, - data=data, - response=response, - ) - except Exception as e: - raise e + ############ Handle CustomLogger ############################### + ################################################################# + for callback in other_callbacks: + await callback.async_post_call_success_hook( + user_api_key_dict=user_api_key_dict, data=data, response=response + ) + except Exception as e: + raise e return response async def async_post_call_streaming_hook( From 6ca7752381fe8ac6aae08985879d59efdedae6fc Mon Sep 17 00:00:00 2001 From: Krrish Dholakia Date: Tue, 30 Sep 2025 15:46:30 -0700 Subject: [PATCH 02/96] fix(prometheus.py): don't require metadata labels to be set for all requests add a default value if metadata label not set --- .../integrations/prometheus.py | 23 ++++++++++++++++--- 1 file changed, 20 insertions(+), 3 deletions(-) diff --git a/enterprise/litellm_enterprise/integrations/prometheus.py b/enterprise/litellm_enterprise/integrations/prometheus.py index d3b0aefb86..a42f9b642d 100644 --- a/enterprise/litellm_enterprise/integrations/prometheus.py +++ b/enterprise/litellm_enterprise/integrations/prometheus.py @@ -1649,9 +1649,22 @@ class PrometheusLogger(CustomLogger): api_base: Optional[str], api_provider: str, ): - self.litellm_deployment_state.labels( - litellm_model_name, model_id, api_base, api_provider - ).set(state) + """ + Set the deployment state. + """ + ### get labels + _labels = prometheus_label_factory( + supported_enum_labels=self.get_labels_for_metric( + metric_name="litellm_deployment_state" + ), + enum_values=UserAPIKeyLabelValues( + litellm_model_name=litellm_model_name, + model_id=model_id, + api_base=api_base, + api_provider=api_provider, + ), + ) + self.litellm_deployment_state.labels(**_labels).set(state) def set_deployment_healthy( self, @@ -2230,6 +2243,10 @@ def prometheus_label_factory( for key, value in enum_values.custom_metadata_labels.items(): if key in supported_enum_labels: filtered_labels[key] = value + else: + filtered_labels[key] = ( + "None" # this happens for dynamically added metadata labels + ) # Add custom tags if configured if enum_values.tags is not None: From d6800ee706194aeaff40bbacc653b74033a33586 Mon Sep 17 00:00:00 2001 From: Krrish Dholakia Date: Tue, 30 Sep 2025 17:02:40 -0700 Subject: [PATCH 03/96] feat(prometheus.py): initial working commit of passing team/key metadata as prometheus metrics Closes LIT-1006 --- .../integrations/prometheus.py | 12 +++-- litellm/litellm_core_utils/litellm_logging.py | 2 + litellm/proxy/_new_secret_config.yaml | 33 +++--------- litellm/proxy/_types.py | 1 + litellm/proxy/litellm_pre_call_utils.py | 52 ++++++++++++++++++- litellm/types/utils.py | 31 ++++++++--- 6 files changed, 93 insertions(+), 38 deletions(-) diff --git a/enterprise/litellm_enterprise/integrations/prometheus.py b/enterprise/litellm_enterprise/integrations/prometheus.py index a42f9b642d..b472ccbb35 100644 --- a/enterprise/litellm_enterprise/integrations/prometheus.py +++ b/enterprise/litellm_enterprise/integrations/prometheus.py @@ -794,9 +794,16 @@ class PrometheusLogger(CustomLogger): output_tokens = standard_logging_payload["completion_tokens"] tokens_used = standard_logging_payload["total_tokens"] response_cost = standard_logging_payload["response_cost"] - _requester_metadata = standard_logging_payload["metadata"].get( + _requester_metadata: Optional[dict] = standard_logging_payload["metadata"].get( "requester_metadata" ) + user_api_key_auth_metadata: Optional[dict] = standard_logging_payload[ + "metadata" + ].get("user_api_key_auth_metadata") + combined_metadata: Dict[str, Any] = { + **(_requester_metadata if _requester_metadata else {}), + **(user_api_key_auth_metadata if user_api_key_auth_metadata else {}), + } if standard_logging_payload is not None and isinstance( standard_logging_payload, dict ): @@ -828,8 +835,7 @@ class PrometheusLogger(CustomLogger): exception_status=None, exception_class=None, custom_metadata_labels=get_custom_labels_from_metadata( - metadata=standard_logging_payload["metadata"].get("requester_metadata") - or {} + metadata=combined_metadata ), route=standard_logging_payload["metadata"].get( "user_api_key_request_route" diff --git a/litellm/litellm_core_utils/litellm_logging.py b/litellm/litellm_core_utils/litellm_logging.py index 24449e1bd0..46e363c865 100644 --- a/litellm/litellm_core_utils/litellm_logging.py +++ b/litellm/litellm_core_utils/litellm_logging.py @@ -4019,6 +4019,7 @@ class StandardLoggingPayloadSetup: usage_object=usage_object, requester_custom_headers=None, cold_storage_object_key=None, + user_api_key_auth_metadata=None, ) if isinstance(metadata, dict): # Filter the metadata dictionary to include only the specified keys @@ -4685,6 +4686,7 @@ def get_standard_logging_metadata( requester_custom_headers=None, user_api_key_request_route=None, cold_storage_object_key=None, + user_api_key_auth_metadata=None, ) if isinstance(metadata, dict): # Update the clean_metadata with values from input metadata that match StandardLoggingMetadata fields diff --git a/litellm/proxy/_new_secret_config.yaml b/litellm/proxy/_new_secret_config.yaml index 804cf2cf2c..96d9ea0cc3 100644 --- a/litellm/proxy/_new_secret_config.yaml +++ b/litellm/proxy/_new_secret_config.yaml @@ -1,30 +1,9 @@ model_list: - - model_name: byok-fixed-gpt-4o-mini + - model_name: openai/gpt-4o litellm_params: - model: openai/gpt-4o-mini - api_base: "https://webhook.site/2f385e05-00aa-402b-86d1-efc9261471a5" - api_key: dummy - - model_name: "byok-wildcard/*" - litellm_params: - model: openai/* - - model_name: xai-grok-3 - litellm_params: - model: xai/grok-3 - - model_name: hosted_vllm/whisper-v3 - litellm_params: - model: hosted_vllm/whisper-v3 - api_base: "https://webhook.site/2f385e05-00aa-402b-86d1-efc9261471a5" - api_key: dummy - -mcp_servers: - github_mcp: - url: "https://api.githubcopilot.com/mcp" - auth_type: oauth2 - authorization_url: https://github.com/login/oauth/authorize - token_url: https://github.com/login/oauth/access_token - client_id: os.environ/GITHUB_OAUTH_CLIENT_ID - client_secret: os.environ/GITHUB_OAUTH_CLIENT_SECRET - scopes: ["public_repo", "user:email"] - allowed_tools: ["list_tools"] - # disallowed_tools: ["repo_delete"] + model: openai/gpt-4o + api_key: os.environ/OPENAI_API_KEY +litellm_settings: + callbacks: ["prometheus"] + custom_prometheus_metadata_labels: ["metadata.initiative"] \ No newline at end of file diff --git a/litellm/proxy/_types.py b/litellm/proxy/_types.py index c5370eb7d7..00ffae718e 100644 --- a/litellm/proxy/_types.py +++ b/litellm/proxy/_types.py @@ -3066,6 +3066,7 @@ LiteLLM_ManagementEndpoint_MetadataFields_Premium = [ "tags", "team_member_key_duration", "prompts", + "logging", ] diff --git a/litellm/proxy/litellm_pre_call_utils.py b/litellm/proxy/litellm_pre_call_utils.py index e077d0ee92..73052d1495 100644 --- a/litellm/proxy/litellm_pre_call_utils.py +++ b/litellm/proxy/litellm_pre_call_utils.py @@ -579,7 +579,12 @@ class LiteLLMProxyRequestSetup: user_api_key_end_user_id=user_api_key_dict.end_user_id, user_api_key_user_email=user_api_key_dict.user_email, user_api_key_request_route=user_api_key_dict.request_route, - user_api_key_budget_reset_at=user_api_key_dict.budget_reset_at.isoformat() if user_api_key_dict.budget_reset_at else None, + user_api_key_budget_reset_at=( + user_api_key_dict.budget_reset_at.isoformat() + if user_api_key_dict.budget_reset_at + else None + ), + user_api_key_auth_metadata=None, ) return user_api_key_logged_metadata @@ -607,6 +612,35 @@ class LiteLLMProxyRequestSetup: ) return data + @staticmethod + def add_management_endpoint_metadata_to_request_metadata( + data: dict, + management_endpoint_metadata: dict, + _metadata_variable_name: str, + ) -> dict: + """ + Adds the `UserAPIKeyAuth` metadata to the request metadata. + + ignore any sensitive fields like logging, api_key, etc. + """ + from litellm.proxy._types import ( + LiteLLM_ManagementEndpoint_MetadataFields, + LiteLLM_ManagementEndpoint_MetadataFields_Premium, + ) + + # ignore any special fields + added_metadata = {} + for k, v in management_endpoint_metadata.items(): + if k not in ( + LiteLLM_ManagementEndpoint_MetadataFields_Premium + + LiteLLM_ManagementEndpoint_MetadataFields + ): + added_metadata[k] = v + data[_metadata_variable_name].setdefault( + "user_api_key_auth_metadata", {} + ).update(added_metadata) + return data + @staticmethod def add_key_level_controls( key_metadata: Optional[dict], data: dict, _metadata_variable_name: str @@ -651,6 +685,13 @@ class LiteLLMProxyRequestSetup: key_metadata["disable_fallbacks"], bool ): data["disable_fallbacks"] = key_metadata["disable_fallbacks"] + + ## KEY-LEVEL METADATA + data = LiteLLMProxyRequestSetup.add_management_endpoint_metadata_to_request_metadata( + data=data, + management_endpoint_metadata=key_metadata, + _metadata_variable_name=_metadata_variable_name, + ) return data @staticmethod @@ -889,6 +930,15 @@ async def add_litellm_data_to_request( # noqa: PLR0915 "spend_logs_metadata" ] + ## TEAM-LEVEL METADATA + data = ( + LiteLLMProxyRequestSetup.add_management_endpoint_metadata_to_request_metadata( + data=data, + management_endpoint_metadata=team_metadata, + _metadata_variable_name=_metadata_variable_name, + ) + ) + # Team spend, budget - used by prometheus.py data[_metadata_variable_name][ "user_api_key_team_max_budget" diff --git a/litellm/types/utils.py b/litellm/types/utils.py index e5786e50a5..c8de97bba2 100644 --- a/litellm/types/utils.py +++ b/litellm/types/utils.py @@ -123,12 +123,18 @@ class ModelInfoBase(ProviderSpecificModelInfo, total=False): max_output_tokens: Required[Optional[int]] input_cost_per_token: Required[float] input_cost_per_token_flex: Optional[float] # OpenAI flex service tier pricing - input_cost_per_token_priority: Optional[float] # OpenAI priority service tier pricing + input_cost_per_token_priority: Optional[ + float + ] # OpenAI priority service tier pricing cache_creation_input_token_cost: Optional[float] cache_creation_input_token_cost_above_1hr: Optional[float] cache_read_input_token_cost: Optional[float] - cache_read_input_token_cost_flex: Optional[float] # OpenAI flex service tier pricing - cache_read_input_token_cost_priority: Optional[float] # OpenAI priority service tier pricing + cache_read_input_token_cost_flex: Optional[ + float + ] # OpenAI flex service tier pricing + cache_read_input_token_cost_priority: Optional[ + float + ] # OpenAI priority service tier pricing input_cost_per_character: Optional[float] # only for vertex ai models input_cost_per_audio_token: Optional[float] input_cost_per_token_above_128k_tokens: Optional[float] # only for vertex ai models @@ -147,7 +153,9 @@ class ModelInfoBase(ProviderSpecificModelInfo, total=False): output_cost_per_token_batches: Optional[float] output_cost_per_token: Required[float] output_cost_per_token_flex: Optional[float] # OpenAI flex service tier pricing - output_cost_per_token_priority: Optional[float] # OpenAI priority service tier pricing + output_cost_per_token_priority: Optional[ + float + ] # OpenAI priority service tier pricing output_cost_per_character: Optional[float] # only for vertex ai models output_cost_per_audio_token: Optional[float] output_cost_per_token_above_128k_tokens: Optional[ @@ -1856,6 +1864,7 @@ class StandardLoggingUserAPIKeyMetadata(TypedDict): user_api_key_team_alias: Optional[str] user_api_key_end_user_id: Optional[str] user_api_key_request_route: Optional[str] + user_api_key_auth_metadata: Optional[Dict[str, str]] class StandardLoggingMCPToolCall(TypedDict, total=False): @@ -2059,10 +2068,12 @@ class StandardLoggingGuardrailInformation(TypedDict, total=False): StandardLoggingPayloadStatus = Literal["success", "failure"] + class CachingDetails(TypedDict): """ Track all caching related metrics, fields for a given request """ + cache_hit: Optional[bool] """ Whether the request hit the cache @@ -2072,12 +2083,16 @@ class CachingDetails(TypedDict): Duration for reading from cache """ + class CostBreakdown(TypedDict): """ Detailed cost breakdown for a request """ + input_cost: float # Cost of input/prompt tokens - output_cost: float # Cost of output/completion tokens (includes reasoning if applicable) + output_cost: ( + float # Cost of output/completion tokens (includes reasoning if applicable) + ) total_cost: float # Total cost (input + output + tool usage) tool_usage_cost: float # Cost of usage of built-in tools @@ -2616,6 +2631,7 @@ class SpecialEnums(Enum): class ServiceTier(Enum): """Enum for service tier types used in cost calculations.""" + FLEX = "flex" PRIORITY = "priority" @@ -2662,13 +2678,14 @@ CostResponseTypes = Union[ class PriorityReservationSettings(BaseModel): """ Settings for priority-based rate limiting reservation. - + Defines what priority to assign to keys without explicit priority metadata. The priority_reservation mapping is configured separately via litellm.priority_reservation. """ + default_priority: float = Field( default=0.5, - description="Priority level to assign to API keys without explicit priority metadata. Should match a key in litellm.priority_reservation." + description="Priority level to assign to API keys without explicit priority metadata. Should match a key in litellm.priority_reservation.", ) model_config = ConfigDict(protected_namespaces=()) From a1a0e99638ebca998db597649b679a4f1d869a81 Mon Sep 17 00:00:00 2001 From: Krrish Dholakia Date: Tue, 30 Sep 2025 21:23:25 -0700 Subject: [PATCH 04/96] fix(prometheus.py): working e2e calls w/ userapikeymetadata --- .../litellm_enterprise/integrations/prometheus.py | 11 +++++------ litellm/proxy/_new_secret_config.yaml | 2 +- litellm/proxy/litellm_pre_call_utils.py | 8 +++++--- litellm/types/integrations/prometheus.py | 8 ++++---- 4 files changed, 15 insertions(+), 14 deletions(-) diff --git a/enterprise/litellm_enterprise/integrations/prometheus.py b/enterprise/litellm_enterprise/integrations/prometheus.py index b472ccbb35..3b37e14b89 100644 --- a/enterprise/litellm_enterprise/integrations/prometheus.py +++ b/enterprise/litellm_enterprise/integrations/prometheus.py @@ -21,6 +21,7 @@ from litellm._logging import print_verbose, verbose_logger from litellm.integrations.custom_logger import CustomLogger from litellm.proxy._types import LiteLLM_TeamTable, UserAPIKeyAuth from litellm.types.integrations.prometheus import * +from litellm.types.integrations.prometheus import _sanitize_prometheus_label_name from litellm.types.utils import StandardLoggingPayload from litellm.utils import get_end_user_id_for_cost_tracking @@ -2247,12 +2248,10 @@ def prometheus_label_factory( if enum_values.custom_metadata_labels is not None: for key, value in enum_values.custom_metadata_labels.items(): - if key in supported_enum_labels: - filtered_labels[key] = value - else: - filtered_labels[key] = ( - "None" # this happens for dynamically added metadata labels - ) + # check sanitized key + sanitized_key = _sanitize_prometheus_label_name(key) + if sanitized_key in supported_enum_labels: + filtered_labels[sanitized_key] = value # Add custom tags if configured if enum_values.tags is not None: diff --git a/litellm/proxy/_new_secret_config.yaml b/litellm/proxy/_new_secret_config.yaml index 96d9ea0cc3..5d8052493a 100644 --- a/litellm/proxy/_new_secret_config.yaml +++ b/litellm/proxy/_new_secret_config.yaml @@ -6,4 +6,4 @@ model_list: litellm_settings: callbacks: ["prometheus"] - custom_prometheus_metadata_labels: ["metadata.initiative"] \ No newline at end of file + custom_prometheus_metadata_labels: ["metadata.initiative", "metadata.business-unit"] \ No newline at end of file diff --git a/litellm/proxy/litellm_pre_call_utils.py b/litellm/proxy/litellm_pre_call_utils.py index 73052d1495..44e26313f9 100644 --- a/litellm/proxy/litellm_pre_call_utils.py +++ b/litellm/proxy/litellm_pre_call_utils.py @@ -636,9 +636,11 @@ class LiteLLMProxyRequestSetup: + LiteLLM_ManagementEndpoint_MetadataFields ): added_metadata[k] = v - data[_metadata_variable_name].setdefault( - "user_api_key_auth_metadata", {} - ).update(added_metadata) + if data[_metadata_variable_name].get("user_api_key_auth_metadata") is None: + data[_metadata_variable_name]["user_api_key_auth_metadata"] = {} + data[_metadata_variable_name]["user_api_key_auth_metadata"].update( + added_metadata + ) return data @staticmethod diff --git a/litellm/types/integrations/prometheus.py b/litellm/types/integrations/prometheus.py index 9c1a14a830..a3dd4dcb1c 100644 --- a/litellm/types/integrations/prometheus.py +++ b/litellm/types/integrations/prometheus.py @@ -426,13 +426,13 @@ class PrometheusMetricLabels: # Buffer monitoring metrics - these typically don't need additional labels litellm_pod_lock_manager_size: List[str] = [] - + litellm_in_memory_daily_spend_update_queue_size: List[str] = [] - + litellm_redis_daily_spend_update_queue_size: List[str] = [] - + litellm_in_memory_spend_update_queue_size: List[str] = [] - + litellm_redis_spend_update_queue_size: List[str] = [] @staticmethod From 005aec69c7879734b97f8f0c1c6d42c287528b22 Mon Sep 17 00:00:00 2001 From: Krrish Dholakia Date: Wed, 1 Oct 2025 16:57:07 -0700 Subject: [PATCH 05/96] feat(key_management_endpoints.py): allow specifying rate limit type when creating tpm/rpm limits on keys prevents overallocating tpm/rpm limits --- litellm/proxy/_types.py | 8 ++ .../key_management_endpoints.py | 100 +++++++++++++++--- 2 files changed, 91 insertions(+), 17 deletions(-) diff --git a/litellm/proxy/_types.py b/litellm/proxy/_types.py index c5370eb7d7..d39f12f05e 100644 --- a/litellm/proxy/_types.py +++ b/litellm/proxy/_types.py @@ -731,6 +731,12 @@ class GenerateRequestBase(LiteLLMPydanticObjectBase): metadata: Optional[dict] = {} tpm_limit: Optional[int] = None rpm_limit: Optional[int] = None + 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 budget_duration: Optional[str] = None allowed_cache_controls: Optional[list] = [] config: Optional[dict] = {} @@ -3054,6 +3060,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 007c0164be..3b39f609e1 100644 --- a/litellm/proxy/management_endpoints/key_management_endpoints.py +++ b/litellm/proxy/management_endpoints/key_management_endpoints.py @@ -90,10 +90,10 @@ def _get_user_in_team( def _calculate_key_rotation_time(rotation_interval: str) -> datetime: """ Helper function to calculate the next rotation time for a key based on the rotation interval. - + Args: rotation_interval: String representing the rotation interval (e.g., '30d', '90d', '1h') - + Returns: datetime: The calculated next rotation time in UTC """ @@ -102,21 +102,25 @@ def _calculate_key_rotation_time(rotation_interval: str) -> datetime: return now + timedelta(seconds=interval_seconds) -def _set_key_rotation_fields(data: dict, auto_rotate: bool, rotation_interval: Optional[str]) -> None: +def _set_key_rotation_fields( + data: dict, auto_rotate: bool, rotation_interval: Optional[str] +) -> None: """ Helper function to set rotation fields in key data if auto_rotate is enabled. - + Args: data: Dictionary to update with rotation fields auto_rotate: Whether auto rotation is enabled rotation_interval: The rotation interval string (required if auto_rotate is True) """ if auto_rotate and rotation_interval: - data.update({ - "auto_rotate": auto_rotate, - "rotation_interval": rotation_interval, - "key_rotation_at": _calculate_key_rotation_time(rotation_interval) - }) + data.update( + { + "auto_rotate": auto_rotate, + "rotation_interval": rotation_interval, + "key_rotation_at": _calculate_key_rotation_time(rotation_interval), + } + ) def _is_allowed_to_make_key_request( @@ -542,6 +546,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) @@ -620,6 +633,46 @@ async def _common_key_generation_helper( # noqa: PLR0915 return response +async def _check_team_key_limits( + team_table: LiteLLM_TeamTableCachedObj, + data: GenerateKeyRequest, + prisma_client: PrismaClient, +) -> None: + """ + Check if the team key is allocating guaranteed throughput limits. If so, raise an error if we're overallocating. + """ + # 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}, + ) + 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}", + ) + + @router.post( "/key/generate", tags=["key management"], @@ -696,12 +749,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: @@ -729,7 +789,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, @@ -738,12 +797,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( @@ -1198,9 +1266,9 @@ async def update_key_fn( # Handle rotation fields if auto_rotate is being enabled _set_key_rotation_fields( - non_default_values, - non_default_values.get("auto_rotate", False), - non_default_values.get("rotation_interval") + non_default_values, + non_default_values.get("auto_rotate", False), + non_default_values.get("rotation_interval"), ) _data = {**non_default_values, "token": key} @@ -1602,8 +1670,6 @@ def _check_model_access_group( return True - - async def generate_key_helper_fn( # noqa: PLR0915 request_type: Literal[ "user", "key" @@ -1766,12 +1832,12 @@ async def generate_key_helper_fn( # noqa: PLR0915 "allowed_routes": allowed_routes or [], "object_permission_id": object_permission_id, } - + # Add rotation fields if auto_rotate is enabled _set_key_rotation_fields( data=key_data, auto_rotate=auto_rotate or False, - rotation_interval=rotation_interval + rotation_interval=rotation_interval, ) if ( From a83238a2db788c81637581f168b5ff88bbb3b834 Mon Sep 17 00:00:00 2001 From: Krrish Dholakia Date: Wed, 1 Oct 2025 17:02:05 -0700 Subject: [PATCH 06/96] test: add unit tests --- .../test_key_management_endpoints.py | 520 +++++++++++++++--- 1 file changed, 448 insertions(+), 72 deletions(-) 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..35ecdd8e0a 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,12 +15,14 @@ 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, generate_key_helper_fn, @@ -1040,7 +1042,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 +1056,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 +1179,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 +1213,400 @@ 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, + ) + + # 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_key2 = MagicMock() + existing_key2.tpm_limit = 3000 + existing_key2.rpm_limit = 200 + + 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, + ) + + # 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_key2 = MagicMock() + existing_key2.tpm_limit = 2000 + existing_key2.rpm_limit = 300 + + 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, + ) + + # 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, + ) From 3ce074564e6202b1dd65710f6d8f56c7662aa932 Mon Sep 17 00:00:00 2001 From: Krrish Dholakia Date: Wed, 1 Oct 2025 17:48:23 -0700 Subject: [PATCH 07/96] feat(key_management_endpoints.py): add guaranteed throughput support for model specific tpm / rpm limits prevents admin from granting keys more tpm/rpm than created for a key --- .../key_management_endpoints.py | 118 ++++++++++++++++-- 1 file changed, 109 insertions(+), 9 deletions(-) diff --git a/litellm/proxy/management_endpoints/key_management_endpoints.py b/litellm/proxy/management_endpoints/key_management_endpoints.py index 3b39f609e1..b0f9c1573d 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, @@ -633,20 +634,93 @@ async def _common_key_generation_helper( # noqa: PLR0915 return response -async def _check_team_key_limits( +def check_team_key_model_specific_limits( + keys: List[LiteLLM_VerificationToken], team_table: LiteLLM_TeamTableCachedObj, data: GenerateKeyRequest, - prisma_client: PrismaClient, ) -> None: """ - Check if the team key is allocating guaranteed throughput limits. If so, raise an error if we're overallocating. + 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 = {} + model_specific_tpm_limit = {} + + 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: GenerateKeyRequest, +) -> None: + """ + Check if the team key is allocating rpm/tpm limits. If so, raise an error if we're overallocating. """ - # 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}, - ) 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) @@ -673,6 +747,32 @@ async def _check_team_key_limits( ) +async def _check_team_key_limits( + team_table: LiteLLM_TeamTableCachedObj, + data: GenerateKeyRequest, + prisma_client: PrismaClient, +) -> None: + """ + Check if the team key is allocating guaranteed throughput limits. If so, raise an error if we're overallocating. + """ + # 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"], From 757436f30235423074be0563e331f78d5af0752c Mon Sep 17 00:00:00 2001 From: Krrish Dholakia Date: Wed, 1 Oct 2025 17:51:10 -0700 Subject: [PATCH 08/96] test: add unit test --- .../test_key_management_endpoints.py | 113 ++++++++++++++++++ 1 file changed, 113 insertions(+) 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 35ecdd8e0a..0b659acd2b 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 @@ -25,6 +25,7 @@ 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, @@ -1610,3 +1611,115 @@ async def test_check_team_key_limits_exact_boundary(): 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) + ) From d8e3a62fbe29af91ad43fdc0d1550fb5e824f69c Mon Sep 17 00:00:00 2001 From: Krrish Dholakia Date: Wed, 1 Oct 2025 17:58:31 -0700 Subject: [PATCH 09/96] feat(key_management_endpoints.py): add support for guaranteed throughput on key update and service account key creation --- .../key_management_endpoints.py | 48 +++++++++++++++---- 1 file changed, 39 insertions(+), 9 deletions(-) diff --git a/litellm/proxy/management_endpoints/key_management_endpoints.py b/litellm/proxy/management_endpoints/key_management_endpoints.py index b0f9c1573d..22f3f5c530 100644 --- a/litellm/proxy/management_endpoints/key_management_endpoints.py +++ b/litellm/proxy/management_endpoints/key_management_endpoints.py @@ -637,7 +637,7 @@ async def _common_key_generation_helper( # noqa: PLR0915 def check_team_key_model_specific_limits( keys: List[LiteLLM_VerificationToken], team_table: LiteLLM_TeamTableCachedObj, - data: GenerateKeyRequest, + data: Union[GenerateKeyRequest, UpdateKeyRequest], ) -> None: """ Check if the team key is allocating model specific limits. If so, raise an error if we're overallocating. @@ -716,7 +716,7 @@ def check_team_key_model_specific_limits( def check_team_key_rpm_tpm_limits( keys: List[LiteLLM_VerificationToken], team_table: LiteLLM_TeamTableCachedObj, - data: GenerateKeyRequest, + data: Union[GenerateKeyRequest, UpdateKeyRequest], ) -> None: """ Check if the team key is allocating rpm/tpm limits. If so, raise an error if we're overallocating. @@ -749,15 +749,23 @@ def check_team_key_rpm_tpm_limits( async def _check_team_key_limits( team_table: LiteLLM_TeamTableCachedObj, - data: GenerateKeyRequest, + 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}, ) @@ -993,12 +1001,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, @@ -1031,6 +1046,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, @@ -1330,14 +1352,22 @@ async def update_key_fn( user_api_key_cache=user_api_key_cache, ) + team_obj = await get_team_object( + team_id=cast(str, 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): - team_obj = await get_team_object( - team_id=cast(str, data.team_id), - prisma_client=prisma_client, - user_api_key_cache=user_api_key_cache, - check_db_only=True, - ) if llm_router is None: raise HTTPException( status_code=400, From ff866aae7bf5964a61fc4ca5c81c9524ba83ea69 Mon Sep 17 00:00:00 2001 From: Krrish Dholakia Date: Wed, 1 Oct 2025 18:30:34 -0700 Subject: [PATCH 10/96] feat(create_key_button.tsx): add tpm/rpm rate limit type options to UI allows user to set the type of tpm/rpm limit they're trying to set --- .../organisms/create_key_button.tsx | 80 +++++++++++++++++++ 1 file changed, 80 insertions(+) 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..0ab8b25c1f 100644 --- a/ui/litellm-dashboard/src/components/organisms/create_key_button.tsx +++ b/ui/litellm-dashboard/src/components/organisms/create_key_button.tsx @@ -810,6 +810,46 @@ const CreateKey: React.FC = ({ > + + TPM Rate Limit Type {' '} + + + + + } + name="tpm_limit_type" + initialValue="default" + className="mt-4" + > + + = ({ > + + RPM Rate Limit Type {' '} + + + + + } + name="rpm_limit_type" + initialValue="default" + className="mt-4" + > + + From 20b6f011f7eeae2131c15667037a60e6a50b26e1 Mon Sep 17 00:00:00 2001 From: Krrish Dholakia Date: Wed, 1 Oct 2025 18:34:29 -0700 Subject: [PATCH 11/96] fix(create_key_button.tsx): working ui controls to set guaranteed throughput/best effort throughput on key --- .../src/components/organisms/create_key_button.tsx | 8 ++++---- 1 file changed, 4 insertions(+), 4 deletions(-) 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 0ab8b25c1f..e367565bae 100644 --- a/ui/litellm-dashboard/src/components/organisms/create_key_button.tsx +++ b/ui/litellm-dashboard/src/components/organisms/create_key_button.tsx @@ -820,7 +820,7 @@ const CreateKey: React.FC = ({ } name="tpm_limit_type" - initialValue="default" + initialValue={null} className="mt-4" > = ({ form.setFieldValue('rpm_limit_type', value); }} > -