scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-09-11 22:51:28 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 123

fix: qa'ed RPM limiting on agents

Browse source
This commit is contained in:
Krrish Dholakia 2026-03-07 18:21:22 -08:00
parent bb8b8dba43
commit 1bd6c24c06
5 changed files with 791 additions and 181 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

4
litellm/proxy/agent_endpoints/a2a_endpoints.py
View file

@ -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}",

122
litellm/proxy/auth/auth_checks.py
View file

@ -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

148
litellm/proxy/auth/user_api_key_auth.py
View file

@ -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)

177
litellm/proxy/hooks/parallel_request_limiter_v3.py
View file

@ -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:

521
tests/test_litellm/proxy/hooks/test_parallel_request_limiter_v3.py
View file

@ -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.