From 44bf51601a59ab6248f77224dc268b9cc6ebbc61 Mon Sep 17 00:00:00 2001 From: ishaan-jaff Date: Wed, 6 Dec 2023 14:43:45 -0800 Subject: [PATCH] (feat) proxy - custom on failure callback --- litellm/proxy/custom_callbacks.py | 1 - litellm/proxy/proxy_config.yaml | 3 ++- litellm/proxy/proxy_server.py | 12 ++++++++++++ 3 files changed, 14 insertions(+), 2 deletions(-) diff --git a/litellm/proxy/custom_callbacks.py b/litellm/proxy/custom_callbacks.py index 9ad19822f0..0114f40b03 100644 --- a/litellm/proxy/custom_callbacks.py +++ b/litellm/proxy/custom_callbacks.py @@ -88,7 +88,6 @@ async def async_on_succes_logger(kwargs, response_obj, start_time, end_time): async def async_on_fail_logger(kwargs, response_obj, start_time, end_time): print(f"On Async Failure!") - print(kwargs) # Access kwargs passed to litellm.completion() model = kwargs.get("model", None) diff --git a/litellm/proxy/proxy_config.yaml b/litellm/proxy/proxy_config.yaml index ecd0c535f6..4ebc3060eb 100644 --- a/litellm/proxy/proxy_config.yaml +++ b/litellm/proxy/proxy_config.yaml @@ -45,8 +45,9 @@ litellm_settings: # setting callback class # callbacks: custom_callbacks.proxy_handler_instance # sets litellm.callbacks = [proxy_handler_instance] - # setting a callback function + # setting a callback function for success and failure success_callback: [custom_callbacks.async_on_succes_logger] + failure_callback: [custom_callbacks.async_on_fail_logger] general_settings: # otel: True # OpenTelemetry Logger diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index a59b80fd1b..54131ee27b 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -510,6 +510,18 @@ def load_router_config(router: Optional[litellm.Router], config_file_path: str): else: litellm.success_callback.append(callback) print_verbose(f"{blue_color_code} Initialized Success Callbacks - {litellm.success_callback} {reset_color_code}") + elif key == "failure_callback": + litellm.failure_callback = [] + + # intialize success callbacks + for callback in value: + # user passed custom_callbacks.async_on_succes_logger. They need us to import a function + if "." in callback: + litellm.failure_callback.append(get_instance_fn(value=callback)) + # these are litellm callbacks - "langfuse", "sentry", "wandb" + else: + litellm.failure_callback.append(callback) + print_verbose(f"{blue_color_code} Initialized Success Callbacks - {litellm.failure_callback} {reset_color_code}") else: setattr(litellm, key, value)