phoenix/litellm-mirror
0
0
Fork 1
mirror of https://github.com/BerriAI/litellm.git synced 2025-04-26 03:04:13 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions

fix(utils.py): override default success callbacks with dynamic callbacks if set

Browse source
This commit is contained in:
Krrish Dholakia 2024-02-02 06:21:43 -08:00
parent 443225acb4
commit 92058cbcd4
5 changed files with 43 additions and 19 deletions
Download patch file Download diff file Expand all files Collapse all files

2
litellm/proxy/hooks/cache_control_check.py
View file

@ -10,7 +10,7 @@ from fastapi import HTTPException
import json, traceback import json, traceback
class CacheControlCheck(CustomLogger): class _PROXY_CacheControlCheck(CustomLogger):
# Class variables or attributes # Class variables or attributes
def __init__(self): def __init__(self):
pass pass

2
litellm/proxy/hooks/max_budget_limiter.py
View file

@ -7,7 +7,7 @@ from fastapi import HTTPException
import json, traceback import json, traceback
class MaxBudgetLimiter(CustomLogger): class _PROXY_MaxBudgetLimiter(CustomLogger):
# Class variables or attributes # Class variables or attributes
def __init__(self): def __init__(self):
pass pass

2
litellm/proxy/hooks/parallel_request_limiter.py
View file

@ -9,7 +9,7 @@ from litellm import ModelResponse
from datetime import datetime from datetime import datetime
class MaxParallelRequestsHandler(CustomLogger): class _PROXY_MaxParallelRequestsHandler(CustomLogger):
user_api_key_cache = None user_api_key_cache = None
# Class variables or attributes # Class variables or attributes

14
litellm/proxy/utils.py
View file

@ -8,9 +8,11 @@ from litellm.proxy._types import (
LiteLLM_SpendLogs, LiteLLM_SpendLogs,
) )
from litellm.caching import DualCache from litellm.caching import DualCache
from litellm.proxy.hooks.parallel_request_limiter import MaxParallelRequestsHandler from litellm.proxy.hooks.parallel_request_limiter import (
from litellm.proxy.hooks.max_budget_limiter import MaxBudgetLimiter _PROXY_MaxParallelRequestsHandler,
from litellm.proxy.hooks.cache_control_check import CacheControlCheck )
from litellm.proxy.hooks.max_budget_limiter import _PROXY_MaxBudgetLimiter
from litellm.proxy.hooks.cache_control_check import _PROXY_CacheControlCheck
from litellm.integrations.custom_logger import CustomLogger from litellm.integrations.custom_logger import CustomLogger
from litellm.proxy.db.base_client import CustomDB from litellm.proxy.db.base_client import CustomDB
from litellm._logging import verbose_proxy_logger from litellm._logging import verbose_proxy_logger
@ -41,9 +43,9 @@ class ProxyLogging:
## INITIALIZE LITELLM CALLBACKS ## ## INITIALIZE LITELLM CALLBACKS ##
self.call_details: dict = {} self.call_details: dict = {}
self.call_details["user_api_key_cache"] = user_api_key_cache self.call_details["user_api_key_cache"] = user_api_key_cache
self.max_parallel_request_limiter = MaxParallelRequestsHandler() self.max_parallel_request_limiter = _PROXY_MaxParallelRequestsHandler()
self.max_budget_limiter = MaxBudgetLimiter() self.max_budget_limiter = _PROXY_MaxBudgetLimiter()
self.cache_control_check = CacheControlCheck() self.cache_control_check = _PROXY_CacheControlCheck()
self.alerting: Optional[List] = None self.alerting: Optional[List] = None
self.alerting_threshold: float = 300 # default to 5 min. threshold self.alerting_threshold: float = 300 # default to 5 min. threshold
pass pass

42
litellm/utils.py
View file

@ -774,14 +774,14 @@ class Logging:
self.streaming_chunks = [] # for generating complete stream response self.streaming_chunks = [] # for generating complete stream response
self.sync_streaming_chunks = [] # for generating complete stream response self.sync_streaming_chunks = [] # for generating complete stream response
self.model_call_details = {} self.model_call_details = {}
self.dynamic_input_callbacks = [] # callbacks set for just that call self.dynamic_input_callbacks = [] # [TODO] callbacks set for just that call
self.dynamic_failure_callbacks = [] # callbacks set for just that call self.dynamic_failure_callbacks = [] # [TODO] callbacks set for just that call
self.dynamic_success_callbacks = ( self.dynamic_success_callbacks = (
dynamic_success_callbacks or [] dynamic_success_callbacks # callbacks set for just that call
) # callbacks set for just that call )
self.dynamic_async_success_callbacks = ( self.dynamic_async_success_callbacks = (
dynamic_async_success_callbacks or [] dynamic_async_success_callbacks # callbacks set for just that call
) # callbacks set for just that call )
## DYNAMIC LANGFUSE KEYS ## ## DYNAMIC LANGFUSE KEYS ##
self.langfuse_public_key = langfuse_public_key self.langfuse_public_key = langfuse_public_key
self.langfuse_secret = langfuse_secret self.langfuse_secret = langfuse_secret
@ -1145,7 +1145,19 @@ class Logging:
f"Model={self.model} not found in completion cost map." f"Model={self.model} not found in completion cost map."
) )
self.model_call_details["response_cost"] = None self.model_call_details["response_cost"] = None
callbacks = litellm.success_callback + self.dynamic_success_callbacks if self.dynamic_success_callbacks is not None and isinstance(
self.dynamic_success_callbacks, list
):
callbacks = self.dynamic_success_callbacks
## keep the internal functions ##
for callback in litellm.success_callback:
if (
isinstance(callback, CustomLogger)
and "_PROXY_" in callback.__class__.__name__
):
callbacks.append(callback)
else:
callbacks = litellm.success_callback
for callback in callbacks: for callback in callbacks:
try: try:
if callback == "lite_debugger": if callback == "lite_debugger":
@ -1452,9 +1464,19 @@ class Logging:
) )
self.model_call_details["response_cost"] = None self.model_call_details["response_cost"] = None
callbacks = ( if self.dynamic_async_success_callbacks is not None and isinstance(
litellm._async_success_callback + self.dynamic_async_success_callbacks self.dynamic_async_success_callbacks, list
) ):
callbacks = self.dynamic_async_success_callbacks
## keep the internal functions ##
for callback in litellm._async_success_callback:
if (
isinstance(callback, CustomLogger)
and "_PROXY_" in callback.__class__.__name__
):
callbacks.append(callback)
else:
callbacks = litellm._async_success_callback
for callback in callbacks: for callback in callbacks:
try: try:
if callback == "cache" and litellm.cache is not None: if callback == "cache" and litellm.cache is not None: