diff options
Diffstat (limited to '.venv/lib/python3.12/site-packages/litellm/proxy/hooks/parallel_request_limiter.py')
| -rw-r--r-- | .venv/lib/python3.12/site-packages/litellm/proxy/hooks/parallel_request_limiter.py | 866 |
1 files changed, 866 insertions, 0 deletions
diff --git a/.venv/lib/python3.12/site-packages/litellm/proxy/hooks/parallel_request_limiter.py b/.venv/lib/python3.12/site-packages/litellm/proxy/hooks/parallel_request_limiter.py new file mode 100644 index 00000000..06f3b6af --- /dev/null +++ b/.venv/lib/python3.12/site-packages/litellm/proxy/hooks/parallel_request_limiter.py @@ -0,0 +1,866 @@ +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 = _Span + 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 + ) |