mirror of
https://github.com/BerriAI/litellm.git
synced 2025-04-25 10:44:24 +00:00
9b7ebb6a7d
* build(pyproject.toml): add new dev dependencies - for type checking * build: reformat files to fit black * ci: reformat to fit black * ci(test-litellm.yml): make tests run clear * build(pyproject.toml): add ruff * fix: fix ruff checks * build(mypy/): fix mypy linting errors * fix(hashicorp_secret_manager.py): fix passing cert for tls auth * build(mypy/): resolve all mypy errors * test: update test * fix: fix black formatting * build(pre-commit-config.yaml): use poetry run black * fix(proxy_server.py): fix linting error * fix: fix ruff safe representation error
868 lines
35 KiB
Python
868 lines
35 KiB
Python
import asyncio
|
|
import sys
|
|
from datetime import datetime, timedelta
|
|
from typing import TYPE_CHECKING, Any, List, Literal, Optional, Tuple, TypedDict, Union
|
|
|
|
from fastapi import HTTPException
|
|
from pydantic import BaseModel
|
|
|
|
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
|
|
if (
|
|
get_key_model_tpm_limit(user_api_key_dict) is not None
|
|
or get_key_model_rpm_limit(user_api_key_dict) is not None
|
|
):
|
|
_model = data.get("model", None)
|
|
request_count_api_key = (
|
|
f"{api_key}::{_model}::{precise_minute}::request_count"
|
|
)
|
|
_tpm_limit_for_key_model = get_key_model_tpm_limit(user_api_key_dict)
|
|
_rpm_limit_for_key_model = get_key_model_rpm_limit(user_api_key_dict)
|
|
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 {}
|
|
)
|
|
|
|
# ------------
|
|
# 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)
|
|
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
|
|
)
|
|
):
|
|
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: {user_api_key}")
|
|
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
|
|
)
|