scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-09-07 08:26:10 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 36

Merge pull request #2240 from BerriAI/litellm_fix_race_condition

Browse source 
[FIX] Race Condition with Custom Callbacks where Async Streaming got triggered twice
This commit is contained in:
Ishaan Jaff 2024-02-28 18:36:48 -08:00 • committed by GitHub
commit 552717583a
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
3 changed files with 90 additions and 15 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

66
litellm/proxy/custom_callbacks.py Normal file
View file

@ -0,0 +1,66 @@
from litellm.integrations.custom_logger import CustomLogger
import litellm
# This file includes the custom callbacks for LiteLLM Proxy
# Once defined, these can be passed in proxy_config.yaml
class MyCustomHandler(CustomLogger):
def log_pre_api_call(self, model, messages, kwargs):
print(f"Pre-API Call") # noqa
def log_post_api_call(self, kwargs, response_obj, start_time, end_time):
print(f"Post-API Call") # noqa
def log_stream_event(self, kwargs, response_obj, start_time, end_time):
print(f"On Stream") # noqa
def log_success_event(self, kwargs, response_obj, start_time, end_time):
print("On Success") # noqa
def log_failure_event(self, kwargs, response_obj, start_time, end_time):
print(f"On Failure") # noqa
async def async_log_success_event(self, kwargs, response_obj, start_time, end_time):
print(f"ishaan async_log_success_event") # noqa
# log: key, user, model, prompt, response, tokens, cost
# Access kwargs passed to litellm.completion()
model = kwargs.get("model", None)
messages = kwargs.get("messages", None)
user = kwargs.get("user", None)
# Access litellm_params passed to litellm.completion(), example access `metadata`
litellm_params = kwargs.get("litellm_params", {})
metadata = litellm_params.get(
"metadata", {}
) # headers passed to LiteLLM proxy, can be found here
return
async def async_log_failure_event(self, kwargs, response_obj, start_time, end_time):
try:
print(f"On Async Failure !") # noqa
print("\nkwargs", kwargs) # noqa
# Access kwargs passed to litellm.completion()
model = kwargs.get("model", None)
messages = kwargs.get("messages", None)
user = kwargs.get("user", None)
# Access litellm_params passed to litellm.completion(), example access `metadata`
litellm_params = kwargs.get("litellm_params", {})
metadata = litellm_params.get(
"metadata", {}
) # headers passed to LiteLLM proxy, can be found here
# Acess Exceptions & Traceback
exception_event = kwargs.get("exception", None)
traceback_event = kwargs.get("traceback_exception", None)
# Calculate cost using litellm.completion_cost()
except Exception as e:
print(f"Exception: {e}") # noqa
proxy_handler_instance = MyCustomHandler()
# Set litellm.callbacks = [proxy_handler_instance] on the proxy
# need to set litellm.callbacks = [proxy_handler_instance] # on the proxy

4
litellm/proxy/proxy_config.yaml
View file

@ -43,9 +43,9 @@ model_list:
api_key: os.environ/OPENAI_API_KEY
litellm_settings:
fallbacks: [{"openai-gpt-3.5": ["azure-gpt-3.5"]}]
success_callback: ['clickhouse', 'langfuse']
success_callback: ['langfuse']
# setting callback class
# callbacks: custom_callbacks.proxy_handler_instance # sets litellm.callbacks = [proxy_handler_instance]
callbacks: custom_callbacks.proxy_handler_instance # sets litellm.callbacks = [proxy_handler_instance]
general_settings:
master_key: sk-1234

35
litellm/utils.py
View file

@ -1640,7 +1640,7 @@ class Logging:
verbose_logger.debug(
"Async success callbacks: Got a complete streaming response"
)
self.model_call_details["complete_streaming_response"] = (
self.model_call_details["async_complete_streaming_response"] = (
complete_streaming_response
)
try:
@ -1688,28 +1688,31 @@ class Logging:
print_verbose("async success_callback: reaches cache for logging!")
kwargs = self.model_call_details
if self.stream:
if "complete_streaming_response" not in kwargs:
if "async_complete_streaming_response" not in kwargs:
print_verbose(
f"async success_callback: reaches cache for logging, there is no complete_streaming_response. Kwargs={kwargs}\n\n"
f"async success_callback: reaches cache for logging, there is no async_complete_streaming_response. Kwargs={kwargs}\n\n"
)
pass
else:
print_verbose(
"async success_callback: reaches cache for logging, there is a complete_streaming_response. Adding to cache"
"async success_callback: reaches cache for logging, there is a async_complete_streaming_response. Adding to cache"
)
result = kwargs["complete_streaming_response"]
result = kwargs["async_complete_streaming_response"]
# only add to cache once we have a complete streaming response
litellm.cache.add_cache(result, **kwargs)
if isinstance(callback, CustomLogger): # custom logger class
print_verbose(
f"Async success callbacks: {callback}; self.stream: {self.stream}; complete_streaming_response: {self.model_call_details.get('complete_streaming_response', None)}"
f"Running Async success callback: {callback}; self.stream: {self.stream}; async_complete_streaming_response: {self.model_call_details.get('async_complete_streaming_response', None)} result={result}"
)
if self.stream == True:
if "complete_streaming_response" in self.model_call_details:
if (
"async_complete_streaming_response"
in self.model_call_details
):
await callback.async_log_success_event(
kwargs=self.model_call_details,
response_obj=self.model_call_details[
"complete_streaming_response"
"async_complete_streaming_response"
],
start_time=start_time,
end_time=end_time,
@ -1730,15 +1733,18 @@ class Logging:
)
if callable(callback): # custom logger functions
print_verbose(
f"Making async function logging call - {self.model_call_details}"
f"Making async function logging call for {callback}, result={result} - {self.model_call_details}"
)
if self.stream:
if "complete_streaming_response" in self.model_call_details:
if (
"async_complete_streaming_response"
in self.model_call_details
):
await customLogger.async_log_event(
kwargs=self.model_call_details,
response_obj=self.model_call_details[
"complete_streaming_response"
"async_complete_streaming_response"
],
start_time=start_time,
end_time=end_time,
@ -1759,14 +1765,17 @@ class Logging:
if dynamoLogger is None:
dynamoLogger = DyanmoDBLogger()
if self.stream:
if "complete_streaming_response" in self.model_call_details:
if (
"async_complete_streaming_response"
in self.model_call_details
):
print_verbose(
"DynamoDB Logger: Got Stream Event - Completed Stream Response"
)
await dynamoLogger._async_log_event(
kwargs=self.model_call_details,
response_obj=self.model_call_details[
"complete_streaming_response"
"async_complete_streaming_response"
],
start_time=start_time,
end_time=end_time,