mirror of
https://github.com/BerriAI/litellm.git
synced 2025-04-25 02:34:29 +00:00
108 lines
4.1 KiB
Python
108 lines
4.1 KiB
Python
from typing import Optional
|
|
import litellm
|
|
from litellm.caching import DualCache
|
|
from litellm.proxy._types import UserAPIKeyAuth
|
|
from litellm.integrations.custom_logger import CustomLogger
|
|
from fastapi import HTTPException
|
|
|
|
|
|
class MaxParallelRequestsHandler(CustomLogger):
|
|
user_api_key_cache = None
|
|
|
|
# Class variables or attributes
|
|
def __init__(self):
|
|
pass
|
|
|
|
def print_verbose(self, print_statement):
|
|
if litellm.set_verbose is True:
|
|
print(print_statement) # noqa
|
|
|
|
async def async_pre_call_hook(
|
|
self,
|
|
user_api_key_dict: UserAPIKeyAuth,
|
|
cache: DualCache,
|
|
data: dict,
|
|
call_type: str,
|
|
):
|
|
self.print_verbose(f"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 api_key is None:
|
|
return
|
|
|
|
if max_parallel_requests is None:
|
|
return
|
|
|
|
self.user_api_key_cache = cache # save the api key cache for updating the value
|
|
|
|
# CHECK IF REQUEST ALLOWED
|
|
request_count_api_key = f"{api_key}_request_count"
|
|
current = cache.get_cache(key=request_count_api_key)
|
|
self.print_verbose(f"current: {current}")
|
|
if current is None:
|
|
cache.set_cache(request_count_api_key, 1)
|
|
elif int(current) < max_parallel_requests:
|
|
# Increase count for this token
|
|
cache.set_cache(request_count_api_key, int(current) + 1)
|
|
else:
|
|
raise HTTPException(
|
|
status_code=429, detail="Max parallel request limit reached."
|
|
)
|
|
|
|
async def async_log_success_event(self, kwargs, response_obj, start_time, end_time):
|
|
try:
|
|
self.print_verbose(f"INSIDE ASYNC SUCCESS LOGGING")
|
|
user_api_key = kwargs["litellm_params"]["metadata"]["user_api_key"]
|
|
if user_api_key is None:
|
|
return
|
|
|
|
if self.user_api_key_cache is None:
|
|
return
|
|
|
|
request_count_api_key = f"{user_api_key}_request_count"
|
|
# check if it has collected an entire stream response
|
|
self.print_verbose(
|
|
f"'complete_streaming_response' is in kwargs: {'complete_streaming_response' in kwargs}"
|
|
)
|
|
if "complete_streaming_response" in kwargs or kwargs["stream"] != True:
|
|
# Decrease count for this token
|
|
current = (
|
|
self.user_api_key_cache.get_cache(key=request_count_api_key) or 1
|
|
)
|
|
new_val = current - 1
|
|
self.print_verbose(f"updated_value in success call: {new_val}")
|
|
self.user_api_key_cache.set_cache(request_count_api_key, new_val)
|
|
except Exception as e:
|
|
self.print_verbose(e) # noqa
|
|
|
|
async def async_log_failure_call(
|
|
self, user_api_key_dict: UserAPIKeyAuth, original_exception: Exception
|
|
):
|
|
try:
|
|
self.print_verbose(f"Inside Max Parallel Request Failure Hook")
|
|
api_key = user_api_key_dict.api_key
|
|
if api_key is None:
|
|
return
|
|
|
|
if self.user_api_key_cache is None:
|
|
return
|
|
|
|
## decrement call count if call failed
|
|
if (
|
|
hasattr(original_exception, "status_code")
|
|
and original_exception.status_code == 429
|
|
and "Max parallel request limit reached" in str(original_exception)
|
|
):
|
|
pass # ignore failed calls due to max limit being reached
|
|
else:
|
|
request_count_api_key = f"{api_key}_request_count"
|
|
# Decrease count for this token
|
|
current = (
|
|
self.user_api_key_cache.get_cache(key=request_count_api_key) or 1
|
|
)
|
|
new_val = current - 1
|
|
self.print_verbose(f"updated_value in failure call: {new_val}")
|
|
self.user_api_key_cache.set_cache(request_count_api_key, new_val)
|
|
except Exception as e:
|
|
self.print_verbose(f"An exception occurred - {str(e)}") # noqa
|