mirror of
https://github.com/BerriAI/litellm.git
synced 2025-04-25 18:54:30 +00:00
4f8a3fd4cf
* docs(exception_mapping.md): add missing exception types Fixes https://github.com/Aider-AI/aider/issues/2120#issuecomment-2438971183 * fix(main.py): register custom model pricing with specific key Ensure custom model pricing is registered to the specific model+provider key combination * test: make testing more robust for custom pricing * fix(redis_cache.py): instrument otel logging for sync redis calls ensures complete coverage for all redis cache calls * refactor: pass parent_otel_span for redis caching calls in router allows for more observability into what calls are causing latency issues * test: update tests with new params * refactor: ensure e2e otel tracing for router * refactor(router.py): add more otel tracing acrosss router catch all latency issues for router requests * fix: fix linting error * fix(router.py): fix linting error * fix: fix test * test: fix tests * fix(dual_cache.py): pass ttl to redis cache * fix: fix param
55 lines
1.7 KiB
Python
55 lines
1.7 KiB
Python
import traceback
|
|
|
|
from fastapi import HTTPException
|
|
|
|
import litellm
|
|
from litellm import verbose_logger
|
|
from litellm.caching.caching import DualCache
|
|
from litellm.integrations.custom_logger import CustomLogger
|
|
from litellm.proxy._types import UserAPIKeyAuth
|
|
|
|
|
|
class _PROXY_MaxBudgetLimiter(CustomLogger):
|
|
# 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,
|
|
):
|
|
try:
|
|
self.print_verbose("Inside Max Budget Limiter Pre-Call Hook")
|
|
cache_key = f"{user_api_key_dict.user_id}_user_api_key_user_id"
|
|
user_row = await cache.async_get_cache(
|
|
cache_key, parent_otel_span=user_api_key_dict.parent_otel_span
|
|
)
|
|
if user_row is None: # value not yet cached
|
|
return
|
|
max_budget = user_row["max_budget"]
|
|
curr_spend = user_row["spend"]
|
|
|
|
if max_budget is None:
|
|
return
|
|
|
|
if curr_spend is None:
|
|
return
|
|
|
|
# CHECK IF REQUEST ALLOWED
|
|
if curr_spend >= max_budget:
|
|
raise HTTPException(status_code=429, detail="Max budget limit reached.")
|
|
except HTTPException as e:
|
|
raise e
|
|
except Exception as e:
|
|
verbose_logger.exception(
|
|
"litellm.proxy.hooks.max_budget_limiter.py::async_pre_call_hook(): Exception occured - {}".format(
|
|
str(e)
|
|
)
|
|
)
|