diff --git a/litellm/proxy/custom_callbacks.py b/litellm/proxy/custom_callbacks.py new file mode 100644 index 00000000000..40fc0d369b8 --- /dev/null +++ b/litellm/proxy/custom_callbacks.py @@ -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 diff --git a/litellm/proxy/proxy_config.yaml b/litellm/proxy/proxy_config.yaml index 5755293e000..6b4b7a8f624 100644 --- a/litellm/proxy/proxy_config.yaml +++ b/litellm/proxy/proxy_config.yaml @@ -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 diff --git a/litellm/utils.py b/litellm/utils.py index 9f88a6791bc..1106d2fec9a 100644 --- a/litellm/utils.py +++ b/litellm/utils.py @@ -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,