ai-workspace-services/litellm
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity
541e81de2f
litellm/litellm/proxy/hooks/parallel_request_limiter.py
ishaan-berri b53cfe729a
Litellm ishaan march30 (#24887) (#25151) 
* fix(pricing): add unversioned vertex_ai/claude-haiku-4-5 entry

Missing unversioned entry causes cost tracking to return $0.00 for
all requests using vertex_ai/claude-haiku-4-5. All other Vertex AI
Claude models have both versioned and unversioned entries.

* fix(router): skip misleading tags error when no candidates (e.g. cooldown)

Return early from get_deployments_for_tag when healthy_deployments is empty so
tag-based routing does not raise no_deployments_with_tag_routing after cooldown
filters all deployments. Adds regression test.

Made-with: Cursor

* feat(oci): add embedding support and update model catalog

- Add OCIEmbeddingConfig for OCI GenAI embedding models
- Add 16 new chat models (Cohere, Meta Llama, xAI Grok, Google Gemini)
- Add 8 embedding models (Cohere embed v3.0, v4.0)
- Update documentation with embedding examples
- Update pricing for all new models



* test(oci): add unit tests for OCI embedding support

- 17 unit tests covering OCIEmbeddingConfig
- Tests for URL generation, param mapping, request/response transform
- Tests for model pricing JSON completeness



* style(oci): format with black and ruff

* fix(oci): correct embedding request body format

OCI embedText API expects inputs, truncate, and inputType at the
top level of the request body, not nested under embedTextDetails.
Fixed transformation and updated tests accordingly.

Verified with real OCI API: 3/3 embedding models working.

* docs: clarify tag routing early return and test intent

Made-with: Cursor

* fix(oci): address code review findings from Greptile

- P1: Fix signing URL mismatch with custom api_base by accepting
  api_base parameter in transform_embedding_request
- P2: Remove encoding_format from supported params (OCI does not
  support it, was silently dropped)
- P2: Raise ValueError for token-array inputs instead of silently
  converting to string representation
- Add test for token-list rejection

* fix(mcp): add STS AssumeRole support for MCP SigV4 authentication

MCPSigV4Auth only supported static AWS credentials or the boto3 default
credential chain. Production Kubernetes environments typically authenticate
via IAM role assumption (sts:AssumeRole), which was not possible.

Add aws_role_name and aws_session_name parameters to the MCP SigV4 auth
stack. When aws_role_name is provided, MCPSigV4Auth calls sts:AssumeRole
to obtain temporary credentials before signing requests. Explicit keys,
if also provided, are used as the source identity for the STS call;
otherwise ambient credentials (pod role, instance profile) are used.

* fix: stop logging credential values and add missing redaction patterns

Replaces raw credential values in debug/error log messages with
boolean presence checks or type names. Adds PEM block, GCP token,
JWT, SAS token, and service-account blob patterns to the redaction
filter. Fixes private_key pattern to capture full PEM blocks instead
of stopping at the first whitespace.

Addresses: Vertex AI credential JSON (including RSA private key)
being logged to stderr on health check failures.

* fix: log only field names for UserAPIKeyAuth, not full object

* style: apply black formatting to experimental_mcp_client/client.py

* style: fix black/isort formatting and mypy error in proxy_server.py

- Fix black formatting in experimental_mcp_client/client.py (done in prev commit)
- Fix black/isort formatting in key_management_endpoints.py, proxy_server.py, transformation.py
- Fix mypy: iterate over optional list safely (access_group_ids or []) in proxy_server.py

* fix(test): patch check_migration.verbose_logger directly to fix xdist ordering issue

When test_proxy_cli.py tests run before test_check_migration.py in the same
xdist worker, litellm.proxy.db.check_migration is already in sys.modules.
Patching litellm._logging.verbose_logger has no effect on the already-bound
reference. Patch the correct target (check_migration.verbose_logger) and
import the module before patching so the order doesn't matter.

* fix(mypy): make api_base Optional in PydanticAIProviderConfig to match base class signature

---------

Co-authored-by: Ihsan Soydemir <soydemir.ihsan@gmail.com>
Co-authored-by: Milan <milan@berri.ai>
Co-authored-by: Daniel Gandolfi <danielgandolfi@gmail.com>
Co-authored-by: Claude Sonnet 4.6 <noreply@anthropic.com>
Co-authored-by: michelligabriele <gabriele.michelli@icloud.com>
Co-authored-by: user <70670632+stuxf@users.noreply.github.com>
Co-authored-by: Ishaan Jaffer <ishaanjaffer0324@gmail.com>
2026-04-04 14:44:07 -07:00

892 lines
36 KiB
Python
Raw Blame History

import asyncio
import sys
from datetime import datetime, timedelta
from typing import TYPE_CHECKING, Any, List, Literal, Optional, Tuple, Union
from fastapi import HTTPException
from pydantic import BaseModel
from typing_extensions import TypedDict
import litellm
from litellm import DualCache, ModelResponse
from litellm._logging import verbose_proxy_logger
from litellm.integrations.custom_logger import CustomLogger
from litellm.litellm_core_utils.core_helpers import _get_parent_otel_span_from_kwargs
from litellm.proxy._types import CommonProxyErrors, CurrentItemRateLimit, UserAPIKeyAuth
from litellm.proxy.auth.auth_utils import (
get_key_model_rpm_limit,
get_key_model_tpm_limit,
)
if TYPE_CHECKING:
from opentelemetry.trace import Span as _Span
from litellm.proxy.utils import InternalUsageCache as _InternalUsageCache
Span = Union[_Span, Any]
InternalUsageCache = _InternalUsageCache
else:
Span = Any
InternalUsageCache = Any
class CacheObject(TypedDict):
current_global_requests: Optional[dict]
request_count_api_key: Optional[dict]
request_count_api_key_model: Optional[dict]
request_count_user_id: Optional[dict]
request_count_team_id: Optional[dict]
request_count_end_user_id: Optional[dict]
class _PROXY_MaxParallelRequestsHandler(CustomLogger):
# Class variables or attributes
def __init__(self, internal_usage_cache: InternalUsageCache):
self.internal_usage_cache = internal_usage_cache
def print_verbose(self, print_statement):
try:
verbose_proxy_logger.debug(print_statement)
if litellm.set_verbose:
print(print_statement) # noqa
except Exception:
pass
async def check_key_in_limits(
self,
user_api_key_dict: UserAPIKeyAuth,
cache: DualCache,
data: dict,
call_type: str,
max_parallel_requests: int,
tpm_limit: int,
rpm_limit: int,
current: Optional[dict],
request_count_api_key: str,
rate_limit_type: Literal["key", "model_per_key", "user", "customer", "team"],
values_to_update_in_cache: List[Tuple[Any, Any]],
) -> dict:
verbose_proxy_logger.info(
f"Current Usage of {rate_limit_type} in this minute: {current}"
)
if current is None:
if max_parallel_requests == 0 or tpm_limit == 0 or rpm_limit == 0:
# base case
raise self.raise_rate_limit_error(
additional_details=f"{CommonProxyErrors.max_parallel_request_limit_reached.value}. Hit limit for {rate_limit_type}. Current limits: max_parallel_requests: {max_parallel_requests}, tpm_limit: {tpm_limit}, rpm_limit: {rpm_limit}"
)
new_val = {
"current_requests": 1,
"current_tpm": 0,
"current_rpm": 1,
}
values_to_update_in_cache.append((request_count_api_key, new_val))
elif (
int(current["current_requests"]) < max_parallel_requests
and current["current_tpm"] < tpm_limit
and current["current_rpm"] < rpm_limit
):
# Increase count for this token
new_val = {
"current_requests": current["current_requests"] + 1,
"current_tpm": current["current_tpm"],
"current_rpm": current["current_rpm"] + 1,
}
values_to_update_in_cache.append((request_count_api_key, new_val))
else:
raise HTTPException(
status_code=429,
detail=f"LiteLLM Rate Limit Handler for rate limit type = {rate_limit_type}. {CommonProxyErrors.max_parallel_request_limit_reached.value}. current rpm: {current['current_rpm']}, rpm limit: {rpm_limit}, current tpm: {current['current_tpm']}, tpm limit: {tpm_limit}, current max_parallel_requests: {current['current_requests']}, max_parallel_requests: {max_parallel_requests}",
headers={"retry-after": str(self.time_to_next_minute())},
)
await self.internal_usage_cache.async_batch_set_cache(
cache_list=values_to_update_in_cache,
ttl=60,
litellm_parent_otel_span=user_api_key_dict.parent_otel_span,
local_only=True,
)
return new_val
def time_to_next_minute(self) -> float:
# Get the current time
now = datetime.now()
# Calculate the next minute
next_minute = (now + timedelta(minutes=1)).replace(second=0, microsecond=0)
# Calculate the difference in seconds
seconds_to_next_minute = (next_minute - now).total_seconds()
return seconds_to_next_minute
def raise_rate_limit_error(
self, additional_details: Optional[str] = None
) -> HTTPException:
"""
Raise an HTTPException with a 429 status code and a retry-after header
"""
error_message = "Max parallel request limit reached"
if additional_details is not None:
error_message = error_message + " " + additional_details
raise HTTPException(
status_code=429,
detail=f"Max parallel request limit reached {additional_details}",
headers={"retry-after": str(self.time_to_next_minute())},
)
async def get_all_cache_objects(
self,
current_global_requests: Optional[str],
request_count_api_key: Optional[str],
request_count_api_key_model: Optional[str],
request_count_user_id: Optional[str],
request_count_team_id: Optional[str],
request_count_end_user_id: Optional[str],
parent_otel_span: Optional[Span] = None,
) -> CacheObject:
keys = [
current_global_requests,
request_count_api_key,
request_count_api_key_model,
request_count_user_id,
request_count_team_id,
request_count_end_user_id,
]
results = await self.internal_usage_cache.async_batch_get_cache(
keys=keys,
parent_otel_span=parent_otel_span,
)
if results is None:
return CacheObject(
current_global_requests=None,
request_count_api_key=None,
request_count_api_key_model=None,
request_count_user_id=None,
request_count_team_id=None,
request_count_end_user_id=None,
)
return CacheObject(
current_global_requests=results[0],
request_count_api_key=results[1],
request_count_api_key_model=results[2],
request_count_user_id=results[3],
request_count_team_id=results[4],
request_count_end_user_id=results[5],
)
async def async_pre_call_hook( # noqa: PLR0915
self,
user_api_key_dict: UserAPIKeyAuth,
cache: DualCache,
data: dict,
call_type: str,
):
self.print_verbose("Inside Max Parallel Request Pre-Call Hook")
api_key = user_api_key_dict.api_key
max_parallel_requests = user_api_key_dict.max_parallel_requests
if max_parallel_requests is None:
max_parallel_requests = sys.maxsize
if data is None:
data = {}
global_max_parallel_requests = data.get("metadata", {}).get(
"global_max_parallel_requests", None
)
tpm_limit = getattr(user_api_key_dict, "tpm_limit", sys.maxsize)
if tpm_limit is None:
tpm_limit = sys.maxsize
rpm_limit = getattr(user_api_key_dict, "rpm_limit", sys.maxsize)
if rpm_limit is None:
rpm_limit = sys.maxsize
values_to_update_in_cache: List[
Tuple[Any, Any]
] = (
[]
) # values that need to get updated in cache, will run a batch_set_cache after this function
# ------------
# Setup values
# ------------
new_val: Optional[dict] = None
if global_max_parallel_requests is not None:
# get value from cache
_key = "global_max_parallel_requests"
current_global_requests = await self.internal_usage_cache.async_get_cache(
key=_key,
local_only=True,
litellm_parent_otel_span=user_api_key_dict.parent_otel_span,
)
# check if below limit
if current_global_requests is None:
current_global_requests = 1
# if above -> raise error
if current_global_requests >= global_max_parallel_requests:
return self.raise_rate_limit_error(
additional_details=f"Hit Global Limit: Limit={global_max_parallel_requests}, current: {current_global_requests}"
)
# if below -> increment
else:
await self.internal_usage_cache.async_increment_cache(
key=_key,
value=1,
local_only=True,
litellm_parent_otel_span=user_api_key_dict.parent_otel_span,
)
_model = data.get("model", None)
current_date = datetime.now().strftime("%Y-%m-%d")
current_hour = datetime.now().strftime("%H")
current_minute = datetime.now().strftime("%M")
precise_minute = f"{current_date}-{current_hour}-{current_minute}"
cache_objects: CacheObject = await self.get_all_cache_objects(
current_global_requests=(
"global_max_parallel_requests"
if global_max_parallel_requests is not None
else None
),
request_count_api_key=(
f"{api_key}::{precise_minute}::request_count"
if api_key is not None
else None
),
request_count_api_key_model=(
f"{api_key}::{_model}::{precise_minute}::request_count"
if api_key is not None and _model is not None
else None
),
request_count_user_id=(
f"{user_api_key_dict.user_id}::{precise_minute}::request_count"
if user_api_key_dict.user_id is not None
else None
),
request_count_team_id=(
f"{user_api_key_dict.team_id}::{precise_minute}::request_count"
if user_api_key_dict.team_id is not None
else None
),
request_count_end_user_id=(
f"{user_api_key_dict.end_user_id}::{precise_minute}::request_count"
if user_api_key_dict.end_user_id is not None
else None
),
parent_otel_span=user_api_key_dict.parent_otel_span,
)
if api_key is not None:
request_count_api_key = f"{api_key}::{precise_minute}::request_count"
# CHECK IF REQUEST ALLOWED for key
await self.check_key_in_limits(
user_api_key_dict=user_api_key_dict,
cache=cache,
data=data,
call_type=call_type,
max_parallel_requests=max_parallel_requests,
current=cache_objects["request_count_api_key"],
request_count_api_key=request_count_api_key,
tpm_limit=tpm_limit,
rpm_limit=rpm_limit,
rate_limit_type="key",
values_to_update_in_cache=values_to_update_in_cache,
)
# Check if request under RPM/TPM per model for a given API Key
_model = data.get("model", None)
_tpm_limit_for_key_model = get_key_model_tpm_limit(
user_api_key_dict, model_name=_model
)
_rpm_limit_for_key_model = get_key_model_rpm_limit(
user_api_key_dict, model_name=_model
)
if _tpm_limit_for_key_model is not None or _rpm_limit_for_key_model is not None:
request_count_api_key = (
f"{api_key}::{_model}::{precise_minute}::request_count"
)
tpm_limit_for_model = None
rpm_limit_for_model = None
if _model is not None:
if _tpm_limit_for_key_model:
tpm_limit_for_model = _tpm_limit_for_key_model.get(_model)
if _rpm_limit_for_key_model:
rpm_limit_for_model = _rpm_limit_for_key_model.get(_model)
new_val = await self.check_key_in_limits(
user_api_key_dict=user_api_key_dict,
cache=cache,
data=data,
call_type=call_type,
max_parallel_requests=sys.maxsize, # TODO: Support max parallel requests for a model
current=cache_objects["request_count_api_key_model"],
request_count_api_key=request_count_api_key,
tpm_limit=tpm_limit_for_model or sys.maxsize,
rpm_limit=rpm_limit_for_model or sys.maxsize,
rate_limit_type="model_per_key",
values_to_update_in_cache=values_to_update_in_cache,
)
_remaining_tokens = None
_remaining_requests = None
# Add remaining tokens, requests to metadata
if new_val:
if tpm_limit_for_model is not None:
_remaining_tokens = tpm_limit_for_model - new_val["current_tpm"]
if rpm_limit_for_model is not None:
_remaining_requests = rpm_limit_for_model - new_val["current_rpm"]
_remaining_limits_data = {
f"litellm-key-remaining-tokens-{_model}": _remaining_tokens,
f"litellm-key-remaining-requests-{_model}": _remaining_requests,
}
if "metadata" not in data:
data["metadata"] = {}
data["metadata"].update(_remaining_limits_data)
# check if REQUEST ALLOWED for user_id
user_id = user_api_key_dict.user_id
if user_id is not None:
user_tpm_limit = user_api_key_dict.user_tpm_limit
user_rpm_limit = user_api_key_dict.user_rpm_limit
if user_tpm_limit is None:
user_tpm_limit = sys.maxsize
if user_rpm_limit is None:
user_rpm_limit = sys.maxsize
request_count_api_key = f"{user_id}::{precise_minute}::request_count"
# print(f"Checking if {request_count_api_key} is allowed to make request for minute {precise_minute}")
await self.check_key_in_limits(
user_api_key_dict=user_api_key_dict,
cache=cache,
data=data,
call_type=call_type,
max_parallel_requests=sys.maxsize, # TODO: Support max parallel requests for a user
current=cache_objects["request_count_user_id"],
request_count_api_key=request_count_api_key,
tpm_limit=user_tpm_limit,
rpm_limit=user_rpm_limit,
rate_limit_type="user",
values_to_update_in_cache=values_to_update_in_cache,
)
# TEAM RATE LIMITS
## get team tpm/rpm limits
team_id = user_api_key_dict.team_id
if team_id is not None:
team_tpm_limit = user_api_key_dict.team_tpm_limit
team_rpm_limit = user_api_key_dict.team_rpm_limit
if team_tpm_limit is None:
team_tpm_limit = sys.maxsize
if team_rpm_limit is None:
team_rpm_limit = sys.maxsize
request_count_api_key = f"{team_id}::{precise_minute}::request_count"
# print(f"Checking if {request_count_api_key} is allowed to make request for minute {precise_minute}")
await self.check_key_in_limits(
user_api_key_dict=user_api_key_dict,
cache=cache,
data=data,
call_type=call_type,
max_parallel_requests=sys.maxsize, # TODO: Support max parallel requests for a team
current=cache_objects["request_count_team_id"],
request_count_api_key=request_count_api_key,
tpm_limit=team_tpm_limit,
rpm_limit=team_rpm_limit,
rate_limit_type="team",
values_to_update_in_cache=values_to_update_in_cache,
)
# End-User Rate Limits
# Only enforce if user passed `user` to /chat, /completions, /embeddings
if user_api_key_dict.end_user_id:
end_user_tpm_limit = getattr(
user_api_key_dict, "end_user_tpm_limit", sys.maxsize
)
end_user_rpm_limit = getattr(
user_api_key_dict, "end_user_rpm_limit", sys.maxsize
)
if end_user_tpm_limit is None:
end_user_tpm_limit = sys.maxsize
if end_user_rpm_limit is None:
end_user_rpm_limit = sys.maxsize
# now do the same tpm/rpm checks
request_count_api_key = (
f"{user_api_key_dict.end_user_id}::{precise_minute}::request_count"
)
# print(f"Checking if {request_count_api_key} is allowed to make request for minute {precise_minute}")
await self.check_key_in_limits(
user_api_key_dict=user_api_key_dict,
cache=cache,
data=data,
call_type=call_type,
max_parallel_requests=sys.maxsize, # TODO: Support max parallel requests for an End-User
request_count_api_key=request_count_api_key,
current=cache_objects["request_count_end_user_id"],
tpm_limit=end_user_tpm_limit,
rpm_limit=end_user_rpm_limit,
rate_limit_type="customer",
values_to_update_in_cache=values_to_update_in_cache,
)
asyncio.create_task(
self.internal_usage_cache.async_batch_set_cache(
cache_list=values_to_update_in_cache,
ttl=60,
litellm_parent_otel_span=user_api_key_dict.parent_otel_span,
) # don't block execution for cache updates
)
return
async def async_log_success_event( # noqa: PLR0915
self, kwargs, response_obj, start_time, end_time
):
from litellm.proxy.common_utils.callback_utils import (
get_model_group_from_litellm_kwargs,
)
litellm_parent_otel_span: Union[Span, None] = _get_parent_otel_span_from_kwargs(
kwargs=kwargs
)
try:
self.print_verbose("INSIDE parallel request limiter ASYNC SUCCESS LOGGING")
global_max_parallel_requests = kwargs["litellm_params"]["metadata"].get(
"global_max_parallel_requests", None
)
user_api_key = kwargs["litellm_params"]["metadata"]["user_api_key"]
user_api_key_user_id = kwargs["litellm_params"]["metadata"].get(
"user_api_key_user_id", None
)
user_api_key_team_id = kwargs["litellm_params"]["metadata"].get(
"user_api_key_team_id", None
)
user_api_key_model_max_budget = kwargs["litellm_params"]["metadata"].get(
"user_api_key_model_max_budget", None
)
user_api_key_end_user_id = kwargs.get("user")
user_api_key_metadata = (
kwargs["litellm_params"]["metadata"].get("user_api_key_metadata", {})
or {}
)
user_api_key_team_metadata = kwargs["litellm_params"]["metadata"].get(
"user_api_key_team_metadata", None
)
user_api_key_dict = UserAPIKeyAuth(
api_key=user_api_key,
metadata=user_api_key_metadata,
model_max_budget=user_api_key_model_max_budget,
team_metadata=user_api_key_team_metadata,
)
# ------------
# Setup values
# ------------
if global_max_parallel_requests is not None:
# get value from cache
_key = "global_max_parallel_requests"
# decrement
await self.internal_usage_cache.async_increment_cache(
key=_key,
value=-1,
local_only=True,
litellm_parent_otel_span=litellm_parent_otel_span,
)
current_date = datetime.now().strftime("%Y-%m-%d")
current_hour = datetime.now().strftime("%H")
current_minute = datetime.now().strftime("%M")
precise_minute = f"{current_date}-{current_hour}-{current_minute}"
total_tokens = 0
if isinstance(response_obj, ModelResponse):
total_tokens = response_obj.usage.total_tokens # type: ignore
# ------------
# Update usage - API Key
# ------------
values_to_update_in_cache = []
if user_api_key is not None:
request_count_api_key = (
f"{user_api_key}::{precise_minute}::request_count"
)
current = await self.internal_usage_cache.async_get_cache(
key=request_count_api_key,
litellm_parent_otel_span=litellm_parent_otel_span,
) or {
"current_requests": 1,
"current_tpm": 0,
"current_rpm": 0,
}
new_val = {
"current_requests": max(current["current_requests"] - 1, 0),
"current_tpm": current["current_tpm"] + total_tokens,
"current_rpm": current["current_rpm"],
}
self.print_verbose(
f"updated_value in success call: {new_val}, precise_minute: {precise_minute}"
)
values_to_update_in_cache.append((request_count_api_key, new_val))
# ------------
# Update usage - model group + API Key
# ------------
model_group = get_model_group_from_litellm_kwargs(kwargs)
_success_tpm_limit = (
get_key_model_tpm_limit(user_api_key_dict, model_name=model_group)
if model_group is not None
else None
)
_success_rpm_limit = (
get_key_model_rpm_limit(user_api_key_dict, model_name=model_group)
if model_group is not None
else None
)
if (
user_api_key is not None
and model_group is not None
and (
"model_rpm_limit" in user_api_key_metadata
or "model_tpm_limit" in user_api_key_metadata
or user_api_key_model_max_budget is not None
or _success_tpm_limit is not None
or _success_rpm_limit is not None
)
):
request_count_api_key = (
f"{user_api_key}::{model_group}::{precise_minute}::request_count"
)
current = await self.internal_usage_cache.async_get_cache(
key=request_count_api_key,
litellm_parent_otel_span=litellm_parent_otel_span,
) or {
"current_requests": 1,
"current_tpm": 0,
"current_rpm": 0,
}
new_val = {
"current_requests": max(current["current_requests"] - 1, 0),
"current_tpm": current["current_tpm"] + total_tokens,
"current_rpm": current["current_rpm"],
}
self.print_verbose(
f"updated_value in success call: {new_val}, precise_minute: {precise_minute}"
)
values_to_update_in_cache.append((request_count_api_key, new_val))
# ------------
# Update usage - User
# ------------
if user_api_key_user_id is not None:
total_tokens = 0
if isinstance(response_obj, ModelResponse):
total_tokens = response_obj.usage.total_tokens # type: ignore
request_count_api_key = (
f"{user_api_key_user_id}::{precise_minute}::request_count"
)
current = await self.internal_usage_cache.async_get_cache(
key=request_count_api_key,
litellm_parent_otel_span=litellm_parent_otel_span,
) or {
"current_requests": 1,
"current_tpm": total_tokens,
"current_rpm": 1,
}
new_val = {
"current_requests": max(current["current_requests"] - 1, 0),
"current_tpm": current["current_tpm"] + total_tokens,
"current_rpm": current["current_rpm"],
}
self.print_verbose(
f"updated_value in success call: {new_val}, precise_minute: {precise_minute}"
)
values_to_update_in_cache.append((request_count_api_key, new_val))
# ------------
# Update usage - Team
# ------------
if user_api_key_team_id is not None:
total_tokens = 0
if isinstance(response_obj, ModelResponse):
total_tokens = response_obj.usage.total_tokens # type: ignore
request_count_api_key = (
f"{user_api_key_team_id}::{precise_minute}::request_count"
)
current = await self.internal_usage_cache.async_get_cache(
key=request_count_api_key,
litellm_parent_otel_span=litellm_parent_otel_span,
) or {
"current_requests": 1,
"current_tpm": total_tokens,
"current_rpm": 1,
}
new_val = {
"current_requests": max(current["current_requests"] - 1, 0),
"current_tpm": current["current_tpm"] + total_tokens,
"current_rpm": current["current_rpm"],
}
self.print_verbose(
f"updated_value in success call: {new_val}, precise_minute: {precise_minute}"
)
values_to_update_in_cache.append((request_count_api_key, new_val))
# ------------
# Update usage - End User
# ------------
if user_api_key_end_user_id is not None:
total_tokens = 0
if isinstance(response_obj, ModelResponse):
total_tokens = response_obj.usage.total_tokens # type: ignore
request_count_api_key = (
f"{user_api_key_end_user_id}::{precise_minute}::request_count"
)
current = await self.internal_usage_cache.async_get_cache(
key=request_count_api_key,
litellm_parent_otel_span=litellm_parent_otel_span,
) or {
"current_requests": 1,
"current_tpm": total_tokens,
"current_rpm": 1,
}
new_val = {
"current_requests": max(current["current_requests"] - 1, 0),
"current_tpm": current["current_tpm"] + total_tokens,
"current_rpm": current["current_rpm"],
}
self.print_verbose(
f"updated_value in success call: {new_val}, precise_minute: {precise_minute}"
)
values_to_update_in_cache.append((request_count_api_key, new_val))
await self.internal_usage_cache.async_batch_set_cache(
cache_list=values_to_update_in_cache,
ttl=60,
litellm_parent_otel_span=litellm_parent_otel_span,
)
except Exception as e:
self.print_verbose(e) # noqa
async def async_log_failure_event(self, kwargs, response_obj, start_time, end_time):
try:
self.print_verbose("Inside Max Parallel Request Failure Hook")
litellm_parent_otel_span: Union[
Span, None
] = _get_parent_otel_span_from_kwargs(kwargs=kwargs)
_metadata = kwargs["litellm_params"].get("metadata", {}) or {}
global_max_parallel_requests = _metadata.get(
"global_max_parallel_requests", None
)
user_api_key = _metadata.get("user_api_key", None)
self.print_verbose(f"user_api_key: [set={user_api_key is not None}]")
if user_api_key is None:
return
## decrement call count if call failed
if CommonProxyErrors.max_parallel_request_limit_reached.value in str(
kwargs["exception"]
):
pass # ignore failed calls due to max limit being reached
else:
# ------------
# Setup values
# ------------
if global_max_parallel_requests is not None:
# get value from cache
_key = "global_max_parallel_requests"
(
await self.internal_usage_cache.async_get_cache(
key=_key,
local_only=True,
litellm_parent_otel_span=litellm_parent_otel_span,
)
)
# decrement
await self.internal_usage_cache.async_increment_cache(
key=_key,
value=-1,
local_only=True,
litellm_parent_otel_span=litellm_parent_otel_span,
)
current_date = datetime.now().strftime("%Y-%m-%d")
current_hour = datetime.now().strftime("%H")
current_minute = datetime.now().strftime("%M")
precise_minute = f"{current_date}-{current_hour}-{current_minute}"
request_count_api_key = (
f"{user_api_key}::{precise_minute}::request_count"
)
# ------------
# Update usage
# ------------
current = await self.internal_usage_cache.async_get_cache(
key=request_count_api_key,
litellm_parent_otel_span=litellm_parent_otel_span,
) or {
"current_requests": 1,
"current_tpm": 0,
"current_rpm": 0,
}
new_val = {
"current_requests": max(current["current_requests"] - 1, 0),
"current_tpm": current["current_tpm"],
"current_rpm": current["current_rpm"],
}
self.print_verbose(f"updated_value in failure call: {new_val}")
await self.internal_usage_cache.async_set_cache(
request_count_api_key,
new_val,
ttl=60,
litellm_parent_otel_span=litellm_parent_otel_span,
) # save in cache for up to 1 min.
except Exception as e:
verbose_proxy_logger.exception(
"Inside Parallel Request Limiter: An exception occurred - {}".format(
str(e)
)
)
async def get_internal_user_object(
self,
user_id: str,
user_api_key_dict: UserAPIKeyAuth,
) -> Optional[dict]:
"""
Helper to get the 'Internal User Object'
It uses the `get_user_object` function from `litellm.proxy.auth.auth_checks`
We need this because the UserApiKeyAuth object does not contain the rpm/tpm limits for a User AND there could be a perf impact by additionally reading the UserTable.
"""
from litellm._logging import verbose_proxy_logger
from litellm.proxy.auth.auth_checks import get_user_object
from litellm.proxy.proxy_server import prisma_client
try:
_user_id_rate_limits = await get_user_object(
user_id=user_id,
prisma_client=prisma_client,
user_api_key_cache=self.internal_usage_cache.dual_cache,
user_id_upsert=False,
parent_otel_span=user_api_key_dict.parent_otel_span,
proxy_logging_obj=None,
)
if _user_id_rate_limits is None:
return None
return _user_id_rate_limits.model_dump()
except Exception as e:
verbose_proxy_logger.debug(
"Parallel Request Limiter: Error getting user object", str(e)
)
return None
async def async_post_call_success_hook(
self, data: dict, user_api_key_dict: UserAPIKeyAuth, response
):
"""
Retrieve the key's remaining rate limits.
"""
api_key = user_api_key_dict.api_key
current_date = datetime.now().strftime("%Y-%m-%d")
current_hour = datetime.now().strftime("%H")
current_minute = datetime.now().strftime("%M")
precise_minute = f"{current_date}-{current_hour}-{current_minute}"
request_count_api_key = f"{api_key}::{precise_minute}::request_count"
current: Optional[
CurrentItemRateLimit
] = await self.internal_usage_cache.async_get_cache(
key=request_count_api_key,
litellm_parent_otel_span=user_api_key_dict.parent_otel_span,
)
key_remaining_rpm_limit: Optional[int] = None
key_rpm_limit: Optional[int] = None
key_remaining_tpm_limit: Optional[int] = None
key_tpm_limit: Optional[int] = None
if current is not None:
if user_api_key_dict.rpm_limit is not None:
key_remaining_rpm_limit = (
user_api_key_dict.rpm_limit - current["current_rpm"]
)
key_rpm_limit = user_api_key_dict.rpm_limit
if user_api_key_dict.tpm_limit is not None:
key_remaining_tpm_limit = (
user_api_key_dict.tpm_limit - current["current_tpm"]
)
key_tpm_limit = user_api_key_dict.tpm_limit
if hasattr(response, "_hidden_params"):
_hidden_params = getattr(response, "_hidden_params")
else:
_hidden_params = None
if _hidden_params is not None and (
isinstance(_hidden_params, BaseModel) or isinstance(_hidden_params, dict)
):
if isinstance(_hidden_params, BaseModel):
_hidden_params = _hidden_params.model_dump()
_additional_headers = _hidden_params.get("additional_headers", {}) or {}
if key_remaining_rpm_limit is not None:
_additional_headers[
"x-ratelimit-remaining-requests"
] = key_remaining_rpm_limit
if key_rpm_limit is not None:
_additional_headers["x-ratelimit-limit-requests"] = key_rpm_limit
if key_remaining_tpm_limit is not None:
_additional_headers[
"x-ratelimit-remaining-tokens"
] = key_remaining_tpm_limit
if key_tpm_limit is not None:
_additional_headers["x-ratelimit-limit-tokens"] = key_tpm_limit
setattr(
response,
"_hidden_params",
{**_hidden_params, "additional_headers": _additional_headers},
)
return await super().async_post_call_success_hook(
data, user_api_key_dict, response
)
Reference in New Issue View Git Blame Copy Permalink