diff --git a/litellm/proxy/agent_endpoints/a2a_endpoints.py b/litellm/proxy/agent_endpoints/a2a_endpoints.py index 808954cdc96..300d16e9aa9 100644 --- a/litellm/proxy/agent_endpoints/a2a_endpoints.py +++ b/litellm/proxy/agent_endpoints/a2a_endpoints.py @@ -378,6 +378,10 @@ async def invoke_agent_a2a( ) # Set up data dict for litellm processing + if "metadata" not in body: + body["metadata"] = {} + body["metadata"]["agent_id"] = agent.agent_id + body.update( { "model": f"a2a_agent/{agent_name}", diff --git a/litellm/proxy/auth/auth_checks.py b/litellm/proxy/auth/auth_checks.py index 1e46992e858..08d177898f2 100644 --- a/litellm/proxy/auth/auth_checks.py +++ b/litellm/proxy/auth/auth_checks.py @@ -11,7 +11,8 @@ Run checks for: import asyncio import re import time -from typing import TYPE_CHECKING, Any, Dict, List, Literal, Optional, Union, cast +from typing import (TYPE_CHECKING, Any, Dict, List, Literal, Optional, Union, + cast) from fastapi import HTTPException, Request, status from pydantic import BaseModel @@ -20,48 +21,33 @@ import litellm from litellm._logging import verbose_proxy_logger from litellm.caching.caching import DualCache from litellm.caching.dual_cache import LimitedSizeOrderedDict -from litellm.constants import ( - CLI_JWT_EXPIRATION_HOURS, - CLI_JWT_TOKEN_NAME, - DEFAULT_ACCESS_GROUP_CACHE_TTL, - DEFAULT_IN_MEMORY_TTL, - DEFAULT_MANAGEMENT_OBJECT_IN_MEMORY_CACHE_TTL, - DEFAULT_MAX_RECURSE_DEPTH, - EMAIL_BUDGET_ALERT_MAX_SPEND_ALERT_PERCENTAGE, -) +from litellm.constants import (CLI_JWT_EXPIRATION_HOURS, CLI_JWT_TOKEN_NAME, + DEFAULT_ACCESS_GROUP_CACHE_TTL, + DEFAULT_IN_MEMORY_TTL, + DEFAULT_MANAGEMENT_OBJECT_IN_MEMORY_CACHE_TTL, + DEFAULT_MAX_RECURSE_DEPTH, + EMAIL_BUDGET_ALERT_MAX_SPEND_ALERT_PERCENTAGE) from litellm.litellm_core_utils.get_llm_provider_logic import get_llm_provider -from litellm.proxy._types import ( - RBAC_ROLES, - CallInfo, - LiteLLM_AccessGroupTable, - LiteLLM_BudgetTable, - LiteLLM_EndUserTable, - Litellm_EntityType, - LiteLLM_JWTAuth, - LiteLLM_ObjectPermissionTable, - LiteLLM_OrganizationMembershipTable, - LiteLLM_OrganizationTable, - LiteLLM_ProjectTableCachedObj, - LiteLLM_TagTable, - LiteLLM_TeamMembership, - LiteLLM_TeamTable, - LiteLLM_TeamTableCachedObj, - LiteLLM_UserTable, - LiteLLMRoutes, - LitellmUserRoles, - NewTeamRequest, - ProxyErrorTypes, - ProxyException, - RoleBasedPermissions, - SpecialModelNames, - UserAPIKeyAuth, -) +from litellm.proxy._types import (RBAC_ROLES, CallInfo, + LiteLLM_AccessGroupTable, + LiteLLM_BudgetTable, LiteLLM_EndUserTable, + Litellm_EntityType, LiteLLM_JWTAuth, + LiteLLM_ObjectPermissionTable, + LiteLLM_OrganizationMembershipTable, + LiteLLM_OrganizationTable, + LiteLLM_ProjectTableCachedObj, + LiteLLM_TagTable, LiteLLM_TeamMembership, + LiteLLM_TeamTable, + LiteLLM_TeamTableCachedObj, + LiteLLM_UserTable, LiteLLMRoutes, + LitellmUserRoles, NewTeamRequest, + ProxyErrorTypes, ProxyException, + RoleBasedPermissions, SpecialModelNames, + UserAPIKeyAuth) from litellm.proxy.auth.route_checks import RouteChecks from litellm.proxy.db.exception_handler import PrismaDBExceptionHandler from litellm.proxy.guardrails.tool_name_extraction import ( - TOOL_CAPABLE_CALL_TYPES, - extract_request_tool_names, -) + TOOL_CAPABLE_CALL_TYPES, extract_request_tool_names) from litellm.proxy.route_llm_request import route_request from litellm.proxy.utils import PrismaClient, ProxyLogging, log_db_metrics from litellm.router import Router @@ -295,7 +281,8 @@ def _guardrail_modification_check( if not _request_metadata.get("guardrails"): return - from litellm.proxy.guardrails.guardrail_helpers import can_modify_guardrails + from litellm.proxy.guardrails.guardrail_helpers import \ + can_modify_guardrails if not can_modify_guardrails(team_object): raise HTTPException( @@ -317,23 +304,34 @@ async def check_tools_allowlist( effective allowlist is read from valid_token.metadata and valid_token.team_metadata. Raises ProxyException with tool_access_denied if a tool is not allowed. """ - from litellm.litellm_core_utils.api_route_to_call_types import ( - get_call_types_for_route, - ) + from litellm.litellm_core_utils.api_route_to_call_types import \ + get_call_types_for_route if valid_token is None: return call_types = get_call_types_for_route(route) - if not call_types or not any(ct.value in TOOL_CAPABLE_CALL_TYPES for ct in call_types): + if not call_types or not any( + ct.value in TOOL_CAPABLE_CALL_TYPES for ct in call_types + ): return tool_names = extract_request_tool_names(route, request_body) if not tool_names: return - key_meta = (valid_token.metadata or {}) if isinstance(valid_token.metadata, dict) else {} - team_meta = (valid_token.team_metadata or {}) if isinstance(valid_token.team_metadata, dict) else {} + key_meta = ( + (valid_token.metadata or {}) if isinstance(valid_token.metadata, dict) else {} + ) + team_meta = ( + (valid_token.team_metadata or {}) + if isinstance(valid_token.team_metadata, dict) + else {} + ) key_allowed = key_meta.get("allowed_tools") team_allowed = team_meta.get("allowed_tools") - effective = key_allowed if (isinstance(key_allowed, list) and len(key_allowed) > 0) else team_allowed + effective = ( + key_allowed + if (isinstance(key_allowed, list) and len(key_allowed) > 0) + else team_allowed + ) if not isinstance(effective, list) or len(effective) == 0: return allowed_set = {str(t) for t in effective} @@ -410,8 +408,10 @@ async def common_checks( # noqa: PLR0915 # Require trace id for agent keys when agent has require_trace_id_on_calls_by_agent if valid_token is not None and valid_token.agent_id: - from litellm.proxy.agent_endpoints.agent_registry import global_agent_registry - from litellm.proxy.litellm_pre_call_utils import get_chain_id_from_headers + from litellm.proxy.agent_endpoints.agent_registry import \ + global_agent_registry + from litellm.proxy.litellm_pre_call_utils import \ + get_chain_id_from_headers agent = global_agent_registry.get_agent_by_id(agent_id=valid_token.agent_id) if agent is not None: @@ -1962,9 +1962,8 @@ class ExperimentalUIJWTToken: def get_experimental_ui_login_jwt_auth_token(user_info: LiteLLM_UserTable) -> str: from datetime import timedelta - from litellm.proxy.common_utils.encrypt_decrypt_utils import ( - encrypt_value_helper, - ) + from litellm.proxy.common_utils.encrypt_decrypt_utils import \ + encrypt_value_helper if user_info.user_role is None: raise Exception("User role is required for experimental UI login") @@ -2010,9 +2009,8 @@ class ExperimentalUIJWTToken: """ from datetime import timedelta - from litellm.proxy.common_utils.encrypt_decrypt_utils import ( - encrypt_value_helper, - ) + from litellm.proxy.common_utils.encrypt_decrypt_utils import \ + encrypt_value_helper if user_info.user_role is None: raise Exception("User role is required for CLI JWT login") @@ -2051,9 +2049,8 @@ class ExperimentalUIJWTToken: import json from litellm.proxy.auth.user_api_key_auth import UserAPIKeyAuth - from litellm.proxy.common_utils.encrypt_decrypt_utils import ( - decrypt_value_helper, - ) + from litellm.proxy.common_utils.encrypt_decrypt_utils import \ + decrypt_value_helper decrypted_token = decrypt_value_helper( hashed_token, key="ui_hash_key", exception_type="debug" @@ -2369,8 +2366,10 @@ async def _get_resources_from_access_groups( # Lazy import to avoid circular imports if prisma_client is None or user_api_key_cache is None: from litellm.proxy.proxy_server import prisma_client as _prisma_client - from litellm.proxy.proxy_server import proxy_logging_obj as _proxy_logging_obj - from litellm.proxy.proxy_server import user_api_key_cache as _user_api_key_cache + from litellm.proxy.proxy_server import \ + proxy_logging_obj as _proxy_logging_obj + from litellm.proxy.proxy_server import \ + user_api_key_cache as _user_api_key_cache prisma_client = prisma_client or _prisma_client user_api_key_cache = user_api_key_cache or _user_api_key_cache @@ -3326,7 +3325,8 @@ async def _tag_max_budget_check( BudgetExceededError if any tag is over its max budget. Triggers a budget alert if any tag is over its max budget. """ - from litellm.proxy.common_utils.http_parsing_utils import get_tags_from_request_body + from litellm.proxy.common_utils.http_parsing_utils import \ + get_tags_from_request_body if prisma_client is None: return diff --git a/litellm/proxy/auth/user_api_key_auth.py b/litellm/proxy/auth/user_api_key_auth.py index ab162bcfdc3..6c882534e03 100644 --- a/litellm/proxy/auth/user_api_key_auth.py +++ b/litellm/proxy/auth/user_api_key_auth.py @@ -25,53 +25,38 @@ from litellm.litellm_core_utils.dd_tracing import tracer from litellm.litellm_core_utils.dot_notation_indexing import get_nested_value from litellm.proxy._types import * from litellm.proxy.auth.auth_checks import ( - ExperimentalUIJWTToken, - _cache_key_object, - _delete_cache_key_object, - _get_user_role, - _is_user_proxy_admin, - _virtual_key_max_budget_alert_check, - _virtual_key_max_budget_check, - _virtual_key_soft_budget_check, - can_key_call_model, - common_checks, - get_end_user_object, - get_jwt_key_mapping_object, - get_key_object, - get_project_object, - get_team_object, - get_user_object, - is_valid_fallback_model, -) -from litellm.proxy.auth.auth_exception_handler import UserAPIKeyAuthExceptionHandler -from litellm.proxy.auth.auth_utils import ( - abbreviate_api_key, - get_end_user_id_from_request_body, - get_model_from_request, - get_request_route, - normalize_request_route, - pre_db_read_auth_checks, - route_in_additonal_public_routes, -) + ExperimentalUIJWTToken, _cache_key_object, _delete_cache_key_object, + _get_user_role, _is_user_proxy_admin, _virtual_key_max_budget_alert_check, + _virtual_key_max_budget_check, _virtual_key_soft_budget_check, + can_key_call_model, common_checks, get_end_user_object, + get_jwt_key_mapping_object, get_key_object, get_project_object, + get_team_object, get_user_object, is_valid_fallback_model) +from litellm.proxy.auth.auth_exception_handler import \ + UserAPIKeyAuthExceptionHandler +from litellm.proxy.auth.auth_utils import (abbreviate_api_key, + get_end_user_id_from_request_body, + get_model_from_request, + get_request_route, + normalize_request_route, + pre_db_read_auth_checks, + route_in_additonal_public_routes) from litellm.proxy.auth.handle_jwt import JWTAuthManager, JWTHandler from litellm.proxy.auth.oauth2_check import Oauth2Handler from litellm.proxy.auth.oauth2_proxy_hook import handle_oauth2_proxy_request from litellm.proxy.auth.route_checks import RouteChecks -from litellm.proxy.common_utils.cache_coordinator import EventDrivenCacheCoordinator +from litellm.proxy.common_utils.cache_coordinator import \ + EventDrivenCacheCoordinator from litellm.proxy.common_utils.http_parsing_utils import ( - _read_request_body, - _safe_get_request_headers, - populate_request_with_path_params, -) + _read_request_body, _safe_get_request_headers, + populate_request_with_path_params) from litellm.proxy.common_utils.realtime_utils import _realtime_request_body from litellm.proxy.utils import PrismaClient, ProxyLogging from litellm.secret_managers.main import get_secret_bool from litellm.types.services import ServiceTypes try: - from litellm_enterprise.proxy.auth.user_api_key_auth import ( - enterprise_custom_auth as _enterprise_custom_auth, - ) + from litellm_enterprise.proxy.auth.user_api_key_auth import \ + enterprise_custom_auth as _enterprise_custom_auth enterprise_custom_auth: Optional[Callable] = _enterprise_custom_auth except ImportError as e: @@ -351,9 +336,8 @@ def get_api_key( Tuple[Optional[str], Optional[str]]: Tuple of the api_key and the passed_in_key """ from litellm.proxy.auth.route_checks import RouteChecks - from litellm.proxy.common_utils.http_parsing_utils import ( - _safe_get_request_query_params, - ) + from litellm.proxy.common_utils.http_parsing_utils import \ + _safe_get_request_query_params api_key = api_key passed_in_key: Optional[str] = None @@ -519,20 +503,15 @@ async def _user_api_key_auth_builder( # noqa: PLR0915 request_data: dict, custom_litellm_key_header: Optional[str] = None, ) -> UserAPIKeyAuth: - from litellm.proxy.proxy_server import ( - general_settings, - jwt_handler, - litellm_proxy_admin_name, - llm_model_list, - llm_router, - master_key, - model_max_budget_limiter, - open_telemetry_logger, - prisma_client, - proxy_logging_obj, - user_api_key_cache, - user_custom_auth, - ) + from litellm.proxy.proxy_server import (general_settings, jwt_handler, + litellm_proxy_admin_name, + llm_model_list, llm_router, + master_key, + model_max_budget_limiter, + open_telemetry_logger, + prisma_client, proxy_logging_obj, + user_api_key_cache, + user_custom_auth) parent_otel_span: Optional[Span] = None start_time = datetime.now() @@ -729,9 +708,11 @@ async def _user_api_key_auth_builder( # noqa: PLR0915 if team_object is not None else None ), - team_metadata=team_object.metadata - if team_object is not None - else None, + team_metadata=( + team_object.metadata + if team_object is not None + else None + ), org_id=org_id, end_user_id=end_user_id, parent_otel_span=parent_otel_span, @@ -749,9 +730,9 @@ async def _user_api_key_auth_builder( # noqa: PLR0915 team_rpm_limit=( team_object.rpm_limit if team_object is not None else None ), - team_models=team_object.models - if team_object is not None - else [], + team_models=( + team_object.models if team_object is not None else [] + ), user_role=( LitellmUserRoles(user_object.user_role) if user_object is not None @@ -778,16 +759,17 @@ async def _user_api_key_auth_builder( # noqa: PLR0915 if team_membership is not None else None ), - team_metadata=team_object.metadata - if team_object is not None - else None, + team_metadata=( + team_object.metadata if team_object is not None else None + ), ) # Check if model has zero cost - if so, skip all budget checks model = get_model_from_request(request_data, route) skip_budget_checks = False if model is not None and llm_router is not None: - from litellm.proxy.auth.auth_checks import _is_model_cost_zero + from litellm.proxy.auth.auth_checks import \ + _is_model_cost_zero skip_budget_checks = _is_model_cost_zero( model=model, llm_router=llm_router @@ -892,9 +874,9 @@ async def _user_api_key_auth_builder( # noqa: PLR0915 route=route, ) if _end_user_object is not None: - end_user_params[ - "allowed_model_region" - ] = _end_user_object.allowed_model_region + end_user_params["allowed_model_region"] = ( + _end_user_object.allowed_model_region + ) if _end_user_object.litellm_budget_table is not None: _apply_budget_limits_to_end_user_params( end_user_params=end_user_params, @@ -903,9 +885,8 @@ async def _user_api_key_auth_builder( # noqa: PLR0915 ) elif litellm.max_end_user_budget_id is not None: # End user doesn't exist yet, but apply default budget limits if configured - from litellm.proxy.auth.auth_checks import ( - get_default_end_user_budget, - ) + from litellm.proxy.auth.auth_checks import \ + get_default_end_user_budget default_budget = await get_default_end_user_budget( prisma_client=prisma_client, @@ -1462,9 +1443,9 @@ async def _user_api_key_auth_builder( # noqa: PLR0915 if _end_user_object is not None: valid_token_dict.update(end_user_params) - valid_token_dict[ - "end_user_object_permission" - ] = _end_user_object.object_permission + valid_token_dict["end_user_object_permission"] = ( + _end_user_object.object_permission + ) # check if token is from litellm-ui, litellm ui makes keys to allow users to login with sso. These keys can only be used for LiteLLM UI functions # sso/login, ui/login, /key functions and /user functions @@ -1686,7 +1667,8 @@ async def _lookup_end_user_and_apply_budget( valid_token=valid_token, end_user_params=end_user_params ) elif litellm.max_end_user_budget_id is not None: - from litellm.proxy.auth.auth_checks import get_default_end_user_budget + from litellm.proxy.auth.auth_checks import \ + get_default_end_user_budget default_budget = await get_default_end_user_budget( prisma_client=prisma_client, @@ -1717,14 +1699,10 @@ async def _run_post_custom_auth_checks( route: str, parent_otel_span: Optional[Span], ) -> UserAPIKeyAuth: - from litellm.proxy.proxy_server import ( - prisma_client, - user_api_key_cache, - proxy_logging_obj, - general_settings, - llm_router, - model_max_budget_limiter, - ) + from litellm.proxy.proxy_server import (general_settings, llm_router, + model_max_budget_limiter, + prisma_client, proxy_logging_obj, + user_api_key_cache) # 1. Look up end_user object from DB if end_user_id is set end_user_object = None @@ -1755,9 +1733,11 @@ async def _run_post_custom_auth_checks( message=f"Authentication Error - Expired Key. Key Expiry time {expiry_time} and current time {current_time}", type=ProxyErrorTypes.expired_key, code=400, - param=abbreviate_api_key(api_key=valid_token.token) - if valid_token.token - else "", + param=( + abbreviate_api_key(api_key=valid_token.token) + if valid_token.token + else "" + ), ) current_model = request_data.get("model", None) diff --git a/litellm/proxy/hooks/parallel_request_limiter_v3.py b/litellm/proxy/hooks/parallel_request_limiter_v3.py index b5bbb4237c1..856975ea093 100644 --- a/litellm/proxy/hooks/parallel_request_limiter_v3.py +++ b/litellm/proxy/hooks/parallel_request_limiter_v3.py @@ -7,18 +7,8 @@ This is currently in development and not yet ready for production. import binascii import os from datetime import datetime -from typing import ( - TYPE_CHECKING, - Any, - Callable, - Dict, - List, - Literal, - Optional, - TypedDict, - Union, - cast, -) +from typing import (TYPE_CHECKING, Any, Callable, Dict, List, Literal, + Optional, TypedDict, Union, cast) from fastapi import HTTPException @@ -175,9 +165,8 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger): """Get or lazy-load the batch rate limiter.""" if self._batch_rate_limiter is None: try: - from litellm.proxy.hooks.batch_rate_limiter import ( - _PROXY_BatchRateLimiter, - ) + from litellm.proxy.hooks.batch_rate_limiter import \ + _PROXY_BatchRateLimiter self._batch_rate_limiter = _PROXY_BatchRateLimiter( internal_usage_cache=self.internal_usage_cache, @@ -679,10 +668,8 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger): requested_model: The model being requested descriptors: List of rate limit descriptors to append to """ - from litellm.proxy.auth.auth_utils import ( - get_key_model_rpm_limit, - get_key_model_tpm_limit, - ) + from litellm.proxy.auth.auth_utils import (get_key_model_rpm_limit, + get_key_model_tpm_limit) if not requested_model: return @@ -791,6 +778,92 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger): """ return rpm_limit_type == "dynamic" or tpm_limit_type == "dynamic" + def _get_agent_from_registry(self, agent_id: str) -> Optional[Any]: + """Look up an agent from the in-memory registry by ID.""" + from litellm.proxy.agent_endpoints.agent_registry import \ + global_agent_registry + + return global_agent_registry.get_agent_by_id(agent_id=agent_id) + + def _get_resolved_agent_id( + self, user_api_key_dict: UserAPIKeyAuth, data: dict + ) -> Optional[str]: + """ + Resolve the agent_id from either the API key or request metadata. + Key-level agent_id takes precedence over metadata/header-supplied agent_id. + """ + key_agent_id = getattr(user_api_key_dict, "agent_id", None) + if key_agent_id: + return key_agent_id + metadata = data.get("metadata") or {} + return metadata.get("agent_id") + + def _get_session_id_from_data(self, data: dict) -> Optional[str]: + """Extract session_id from request metadata or litellm_session_id.""" + session_id = data.get("litellm_session_id") + if session_id: + return str(session_id) + metadata = data.get("metadata") or {} + session_id = metadata.get("session_id") + if session_id: + return str(session_id) + litellm_metadata = data.get("litellm_metadata") or {} + session_id = litellm_metadata.get("session_id") + if session_id: + return str(session_id) + return None + + def _create_agent_rate_limit_descriptors( + self, + agent_id: str, + data: dict, + ) -> List[RateLimitDescriptor]: + """ + Create rate limit descriptors for agent-level and session-level limits. + + Agent-level: caps total RPM/TPM across all sessions for a given agent. + Session-level: caps RPM/TPM within a single session (identified by session_id). + """ + descriptors: List[RateLimitDescriptor] = [] + + agent = self._get_agent_from_registry(agent_id) + if agent is None: + return descriptors + + agent_rpm = getattr(agent, "rpm_limit", None) + agent_tpm = getattr(agent, "tpm_limit", None) + if agent_rpm is not None or agent_tpm is not None: + descriptors.append( + RateLimitDescriptor( + key="agent", + value=agent_id, + rate_limit={ + "requests_per_unit": agent_rpm, + "tokens_per_unit": agent_tpm, + "window_size": self.window_size, + }, + ) + ) + + session_rpm = getattr(agent, "session_rpm_limit", None) + session_tpm = getattr(agent, "session_tpm_limit", None) + if session_rpm is not None or session_tpm is not None: + session_id = self._get_session_id_from_data(data) + if session_id is not None: + descriptors.append( + RateLimitDescriptor( + key="agent_session", + value=f"{agent_id}:{session_id}", + rate_limit={ + "requests_per_unit": session_rpm, + "tokens_per_unit": session_tpm, + "window_size": self.window_size, + }, + ) + ) + + return descriptors + def _create_rate_limit_descriptors( self, user_api_key_dict: UserAPIKeyAuth, @@ -802,12 +875,11 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger): """ Create all rate limit descriptors for the request. - Returns list of descriptors for API key, user, team, team member, end user, and model-specific limits. + Returns list of descriptors for API key, user, team, team member, end user, + model-specific, agent, and agent-session limits. """ - from litellm.proxy.auth.auth_utils import ( - get_team_model_rpm_limit, - get_team_model_tpm_limit, - ) + from litellm.proxy.auth.auth_utils import (get_team_model_rpm_limit, + get_team_model_tpm_limit) descriptors = [] @@ -956,6 +1028,17 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger): ) ) + # Agent-level and session-level rate limits + resolved_agent_id = self._get_resolved_agent_id(user_api_key_dict, data) + + if resolved_agent_id: + descriptors.extend( + self._create_agent_rate_limit_descriptors( + agent_id=resolved_agent_id, + data=data, + ) + ) + return descriptors async def _check_model_has_recent_failures( @@ -970,9 +1053,8 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger): Returns True if any deployment has failures in the current minute. """ from litellm.proxy.proxy_server import llm_router - from litellm.router_utils.router_callbacks.track_deployment_metrics import ( - get_deployment_failures_for_current_minute, - ) + from litellm.router_utils.router_callbacks.track_deployment_metrics import \ + get_deployment_failures_for_current_minute if llm_router is None: return False @@ -1386,12 +1468,10 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger): """ Update TPM usage on successful API calls by incrementing counters using pipeline """ - from litellm.litellm_core_utils.core_helpers import ( - _get_parent_otel_span_from_kwargs, - ) - from litellm.proxy.common_utils.callback_utils import ( - get_model_group_from_litellm_kwargs, - ) + from litellm.litellm_core_utils.core_helpers import \ + _get_parent_otel_span_from_kwargs + from litellm.proxy.common_utils.callback_utils import \ + get_model_group_from_litellm_kwargs from litellm.types.caching import RedisPipelineIncrementOperation rate_limit_type = self.get_rate_limit_type() @@ -1533,6 +1613,32 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger): ) ) + # Agent TPM + agent_id = standard_logging_metadata.get("agent_id") + if agent_id: + pipeline_operations.extend( + self._create_pipeline_operations( + key="agent", + value=agent_id, + rate_limit_type="tokens", + total_tokens=total_tokens, + ) + ) + + # Agent Session TPM + session_id = standard_logging_metadata.get( + "session_id" + ) or standard_logging_metadata.get("trace_id") + if session_id: + pipeline_operations.extend( + self._create_pipeline_operations( + key="agent_session", + value=f"{agent_id}:{session_id}", + rate_limit_type="tokens", + total_tokens=total_tokens, + ) + ) + # Execute all increments in a single pipeline if pipeline_operations: await self.async_increment_tokens_with_ttl_preservation( @@ -1549,9 +1655,8 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger): """ Decrement max parallel requests counter for the API Key """ - from litellm.litellm_core_utils.core_helpers import ( - _get_parent_otel_span_from_kwargs, - ) + from litellm.litellm_core_utils.core_helpers import \ + _get_parent_otel_span_from_kwargs from litellm.types.caching import RedisPipelineIncrementOperation try: diff --git a/tests/test_litellm/proxy/hooks/test_parallel_request_limiter_v3.py b/tests/test_litellm/proxy/hooks/test_parallel_request_limiter_v3.py index d92e152d89a..3eb481991f7 100644 --- a/tests/test_litellm/proxy/hooks/test_parallel_request_limiter_v3.py +++ b/tests/test_litellm/proxy/hooks/test_parallel_request_limiter_v3.py @@ -1981,6 +1981,527 @@ async def test_execute_token_increment_script_cluster_compatibility(): ), f"Each key should have 2 args, got {len(args)} args for {len(keys)} keys" +@pytest.mark.asyncio +async def test_agent_level_rate_limit_descriptors(): + """ + Test that agent-level rate limit descriptors are created when + an agent has rpm_limit and/or tpm_limit configured. + """ + from unittest.mock import patch + + from litellm.types.agents import AgentResponse + + _api_key = "sk-12345" + _api_key = hash_token(_api_key) + _agent_id = "agent_abc123" + + user_api_key_dict = UserAPIKeyAuth( + api_key=_api_key, + agent_id=_agent_id, + ) + + local_cache = DualCache() + parallel_request_handler = _PROXY_MaxParallelRequestsHandler( + internal_usage_cache=InternalUsageCache(local_cache) + ) + + mock_agent = AgentResponse( + agent_id=_agent_id, + agent_name="test-agent", + agent_card_params={"name": "Test Agent"}, + rpm_limit=50, + tpm_limit=5000, + ) + + captured_descriptors = None + + async def mock_should_rate_limit(descriptors, **kwargs): + nonlocal captured_descriptors + captured_descriptors = descriptors + return {"overall_code": "OK", "statuses": []} + + parallel_request_handler.should_rate_limit = mock_should_rate_limit + + with patch( + "litellm.proxy.agent_endpoints.agent_registry.global_agent_registry.get_agent_by_id", + return_value=mock_agent, + ): + await parallel_request_handler.async_pre_call_hook( + user_api_key_dict=user_api_key_dict, + cache=local_cache, + data={"model": "gpt-4"}, + call_type="", + ) + + assert captured_descriptors is not None + + agent_descriptor = None + for d in captured_descriptors: + if d["key"] == "agent": + agent_descriptor = d + break + + assert agent_descriptor is not None, "Agent descriptor should be present" + assert agent_descriptor["value"] == _agent_id + assert agent_descriptor["rate_limit"]["requests_per_unit"] == 50 + assert agent_descriptor["rate_limit"]["tokens_per_unit"] == 5000 + + +@pytest.mark.asyncio +async def test_agent_session_rate_limit_descriptors(): + """ + Test that session-level rate limit descriptors are created when + an agent has session_rpm_limit/session_tpm_limit and a session_id is present. + """ + from unittest.mock import patch + + from litellm.types.agents import AgentResponse + + _api_key = "sk-12345" + _api_key = hash_token(_api_key) + _agent_id = "agent_abc123" + _session_id = "sess_xyz789" + + user_api_key_dict = UserAPIKeyAuth( + api_key=_api_key, + agent_id=_agent_id, + ) + + local_cache = DualCache() + parallel_request_handler = _PROXY_MaxParallelRequestsHandler( + internal_usage_cache=InternalUsageCache(local_cache) + ) + + mock_agent = AgentResponse( + agent_id=_agent_id, + agent_name="test-agent", + agent_card_params={"name": "Test Agent"}, + session_rpm_limit=10, + session_tpm_limit=1000, + ) + + captured_descriptors = None + + async def mock_should_rate_limit(descriptors, **kwargs): + nonlocal captured_descriptors + captured_descriptors = descriptors + return {"overall_code": "OK", "statuses": []} + + parallel_request_handler.should_rate_limit = mock_should_rate_limit + + with patch( + "litellm.proxy.agent_endpoints.agent_registry.global_agent_registry.get_agent_by_id", + return_value=mock_agent, + ): + await parallel_request_handler.async_pre_call_hook( + user_api_key_dict=user_api_key_dict, + cache=local_cache, + data={ + "model": "gpt-4", + "metadata": {"session_id": _session_id}, + }, + call_type="", + ) + + assert captured_descriptors is not None + + session_descriptor = None + for d in captured_descriptors: + if d["key"] == "agent_session": + session_descriptor = d + break + + assert session_descriptor is not None, "Agent session descriptor should be present" + assert session_descriptor["value"] == f"{_agent_id}:{_session_id}" + assert session_descriptor["rate_limit"]["requests_per_unit"] == 10 + assert session_descriptor["rate_limit"]["tokens_per_unit"] == 1000 + + +@pytest.mark.asyncio +async def test_agent_session_rate_limit_skipped_without_session_id(): + """ + Test that session-level rate limit descriptors are NOT created + when no session_id is available in the request. + """ + from unittest.mock import patch + + from litellm.types.agents import AgentResponse + + _api_key = "sk-12345" + _api_key = hash_token(_api_key) + _agent_id = "agent_abc123" + + user_api_key_dict = UserAPIKeyAuth( + api_key=_api_key, + agent_id=_agent_id, + ) + + local_cache = DualCache() + parallel_request_handler = _PROXY_MaxParallelRequestsHandler( + internal_usage_cache=InternalUsageCache(local_cache) + ) + + mock_agent = AgentResponse( + agent_id=_agent_id, + agent_name="test-agent", + agent_card_params={"name": "Test Agent"}, + session_rpm_limit=10, + session_tpm_limit=1000, + ) + + captured_descriptors = None + + async def mock_should_rate_limit(descriptors, **kwargs): + nonlocal captured_descriptors + captured_descriptors = descriptors + return {"overall_code": "OK", "statuses": []} + + parallel_request_handler.should_rate_limit = mock_should_rate_limit + + with patch( + "litellm.proxy.agent_endpoints.agent_registry.global_agent_registry.get_agent_by_id", + return_value=mock_agent, + ): + await parallel_request_handler.async_pre_call_hook( + user_api_key_dict=user_api_key_dict, + cache=local_cache, + data={"model": "gpt-4"}, + call_type="", + ) + + # should_rate_limit should not have been called (no agent-level limits, only session limits + # but no session_id) + assert captured_descriptors is None, ( + "No descriptors should be created when agent has only session limits " + "but no session_id in request" + ) + + +@pytest.mark.asyncio +async def test_agent_rate_limit_from_metadata_agent_id(): + """ + Test that agent rate limits work when agent_id comes from + request metadata (header) rather than from the API key. + """ + from unittest.mock import patch + + from litellm.types.agents import AgentResponse + + _api_key = "sk-12345" + _api_key = hash_token(_api_key) + _agent_id = "agent_from_header" + + user_api_key_dict = UserAPIKeyAuth( + api_key=_api_key, + ) + + local_cache = DualCache() + parallel_request_handler = _PROXY_MaxParallelRequestsHandler( + internal_usage_cache=InternalUsageCache(local_cache) + ) + + mock_agent = AgentResponse( + agent_id=_agent_id, + agent_name="header-agent", + agent_card_params={"name": "Header Agent"}, + rpm_limit=25, + tpm_limit=2500, + ) + + captured_descriptors = None + + async def mock_should_rate_limit(descriptors, **kwargs): + nonlocal captured_descriptors + captured_descriptors = descriptors + return {"overall_code": "OK", "statuses": []} + + parallel_request_handler.should_rate_limit = mock_should_rate_limit + + with patch( + "litellm.proxy.agent_endpoints.agent_registry.global_agent_registry.get_agent_by_id", + return_value=mock_agent, + ): + await parallel_request_handler.async_pre_call_hook( + user_api_key_dict=user_api_key_dict, + cache=local_cache, + data={ + "model": "gpt-4", + "metadata": {"agent_id": _agent_id}, + }, + call_type="", + ) + + assert captured_descriptors is not None + + agent_descriptor = None + for d in captured_descriptors: + if d["key"] == "agent": + agent_descriptor = d + break + + assert agent_descriptor is not None, "Agent descriptor should be created from metadata agent_id" + assert agent_descriptor["value"] == _agent_id + assert agent_descriptor["rate_limit"]["requests_per_unit"] == 25 + + +@pytest.mark.asyncio +async def test_agent_both_agent_and_session_rate_limits(): + """ + Test that both agent-level and session-level descriptors are created + when both types of limits are configured on the agent. + """ + from unittest.mock import patch + + from litellm.types.agents import AgentResponse + + _api_key = "sk-12345" + _api_key = hash_token(_api_key) + _agent_id = "agent_dual" + _session_id = "sess_dual" + + user_api_key_dict = UserAPIKeyAuth( + api_key=_api_key, + agent_id=_agent_id, + ) + + local_cache = DualCache() + parallel_request_handler = _PROXY_MaxParallelRequestsHandler( + internal_usage_cache=InternalUsageCache(local_cache) + ) + + mock_agent = AgentResponse( + agent_id=_agent_id, + agent_name="dual-agent", + agent_card_params={"name": "Dual Agent"}, + rpm_limit=100, + tpm_limit=10000, + session_rpm_limit=20, + session_tpm_limit=2000, + ) + + captured_descriptors = None + + async def mock_should_rate_limit(descriptors, **kwargs): + nonlocal captured_descriptors + captured_descriptors = descriptors + return {"overall_code": "OK", "statuses": []} + + parallel_request_handler.should_rate_limit = mock_should_rate_limit + + with patch( + "litellm.proxy.agent_endpoints.agent_registry.global_agent_registry.get_agent_by_id", + return_value=mock_agent, + ): + await parallel_request_handler.async_pre_call_hook( + user_api_key_dict=user_api_key_dict, + cache=local_cache, + data={ + "model": "gpt-4", + "metadata": {"session_id": _session_id}, + }, + call_type="", + ) + + assert captured_descriptors is not None + + agent_descriptor = None + session_descriptor = None + for d in captured_descriptors: + if d["key"] == "agent": + agent_descriptor = d + elif d["key"] == "agent_session": + session_descriptor = d + + assert agent_descriptor is not None, "Agent-level descriptor should be present" + assert agent_descriptor["rate_limit"]["requests_per_unit"] == 100 + assert agent_descriptor["rate_limit"]["tokens_per_unit"] == 10000 + + assert session_descriptor is not None, "Session-level descriptor should be present" + assert session_descriptor["value"] == f"{_agent_id}:{_session_id}" + assert session_descriptor["rate_limit"]["requests_per_unit"] == 20 + assert session_descriptor["rate_limit"]["tokens_per_unit"] == 2000 + + +@pytest.mark.asyncio +async def test_agent_rate_limit_tpm_increment_on_success(monkeypatch): + """ + Test that async_log_success_event increments agent and session + TPM counters when agent_id and session_id are in metadata. + """ + _api_key = "sk-12345" + _api_key = hash_token(_api_key) + _agent_id = "agent_tpm_test" + _session_id = "sess_tpm_test" + + local_cache = DualCache() + parallel_request_handler = _PROXY_MaxParallelRequestsHandler( + internal_usage_cache=InternalUsageCache(local_cache) + ) + + def mock_get_rate_limit_type(): + return "total" + + monkeypatch.setattr( + parallel_request_handler, "get_rate_limit_type", mock_get_rate_limit_type + ) + + mock_usage = Usage(prompt_tokens=20, completion_tokens=30, total_tokens=50) + mock_response = ModelResponse( + id="mock-response", + object="chat.completion", + created=int(datetime.now().timestamp()), + model="gpt-4", + usage=mock_usage, + choices=[], + ) + + mock_kwargs = { + "standard_logging_object": { + "metadata": { + "user_api_key_hash": _api_key, + "user_api_key_user_id": None, + "user_api_key_team_id": None, + "user_api_key_end_user_id": None, + "agent_id": _agent_id, + "session_id": _session_id, + } + }, + "model": "gpt-4", + } + + captured_operations = [] + + async def mock_increment_pipeline(increment_list, **kwargs): + captured_operations.extend(increment_list) + return True + + monkeypatch.setattr( + parallel_request_handler.internal_usage_cache.dual_cache, + "async_increment_cache_pipeline", + mock_increment_pipeline, + ) + + await parallel_request_handler.async_log_success_event( + kwargs=mock_kwargs, + response_obj=mock_response, + start_time=datetime.now(), + end_time=datetime.now(), + ) + + agent_tpm_op = None + session_tpm_op = None + for op in captured_operations: + if op["key"] == f"{{agent:{_agent_id}}}:tokens": + agent_tpm_op = op + elif op["key"] == f"{{agent_session:{_agent_id}:{_session_id}}}:tokens": + session_tpm_op = op + + assert agent_tpm_op is not None, "Agent TPM increment should be present" + assert agent_tpm_op["increment_value"] == 50 + + assert session_tpm_op is not None, "Session TPM increment should be present" + assert session_tpm_op["increment_value"] == 50 + + +@pytest.mark.asyncio +async def test_agent_rate_limit_429_on_over_limit(monkeypatch, time_controller): + """ + Test end-to-end that agent rate limiting returns 429 when the agent + RPM limit is exceeded. + """ + from unittest.mock import patch + + from litellm.types.agents import AgentResponse + + monkeypatch.setenv("LITELLM_RATE_LIMIT_WINDOW_SIZE", "2") + _api_key = "sk-12345" + _api_key = hash_token(_api_key) + _agent_id = "agent_429_test" + + user_api_key_dict = UserAPIKeyAuth( + api_key=_api_key, + agent_id=_agent_id, + ) + + local_cache = DualCache() + parallel_request_handler = _PROXY_MaxParallelRequestsHandler( + internal_usage_cache=InternalUsageCache(local_cache), + time_provider=time_controller.now, + ) + + mock_agent = AgentResponse( + agent_id=_agent_id, + agent_name="rate-limited-agent", + agent_card_params={"name": "Rate Limited Agent"}, + rpm_limit=2, + ) + + window_starts: Dict[str, int] = {} + request_counts: Dict[str, int] = {} + + async def mock_batch_rate_limiter(*args, **kwargs): + keys = kwargs.get("keys") if kwargs else args[0] + args_list = kwargs.get("args") if kwargs else args[1] + now = args_list[0] + window_size = args_list[1] + results = [] + for i in range(0, len(keys), 2): + window_key = keys[i] + counter_key = keys[i + 1] + prev_window = window_starts.get(window_key) + prev_counter = request_counts.get(counter_key, 0) + if prev_window is None or (now - prev_window) >= window_size: + window_starts[window_key] = now + new_counter = 1 + request_counts[counter_key] = new_counter + await local_cache.async_set_cache( + key=window_key, value=now, ttl=window_size + ) + await local_cache.async_set_cache( + key=counter_key, value=new_counter, ttl=window_size + ) + else: + new_counter = prev_counter + 1 + request_counts[counter_key] = new_counter + await local_cache.async_set_cache( + key=counter_key, value=new_counter, ttl=window_size + ) + results.append(now) + results.append(new_counter) + return results + + parallel_request_handler.batch_rate_limiter_script = mock_batch_rate_limiter + + with patch( + "litellm.proxy.agent_endpoints.agent_registry.global_agent_registry.get_agent_by_id", + return_value=mock_agent, + ): + await parallel_request_handler.async_pre_call_hook( + user_api_key_dict=user_api_key_dict, + cache=local_cache, + data={"model": "gpt-4"}, + call_type="", + ) + + await parallel_request_handler.async_pre_call_hook( + user_api_key_dict=user_api_key_dict, + cache=local_cache, + data={"model": "gpt-4"}, + call_type="", + ) + + with pytest.raises(HTTPException) as exc_info: + await parallel_request_handler.async_pre_call_hook( + user_api_key_dict=user_api_key_dict, + cache=local_cache, + data={"model": "gpt-4"}, + call_type="", + ) + + assert exc_info.value.status_code == 429 + assert "agent" in exc_info.value.detail + + class TestGetTotalTokensFromUsageCacheExclusion: """ Tests for _get_total_tokens_from_usage cache token exclusion.