refactor: add black formatting

This commit is contained in:
Krrish Dholakia 2023-12-25 14:10:38 +05:30
parent b87d630b0a
commit 4905929de3
156 changed files with 19723 additions and 10869 deletions

View file

@ -4,17 +4,19 @@ import sys, os, traceback
sys.path.insert(
0, os.path.abspath("../..")
) # Adds the parent directory to the system path
) # Adds the parent directory to the system path
from litellm.integrations.custom_logger import CustomLogger
import litellm
import inspect
# This file includes the custom callbacks for LiteLLM Proxy
# Once defined, these can be passed in proxy_config.yaml
def print_verbose(print_statement):
if litellm.set_verbose:
print(print_statement) # noqa
if litellm.set_verbose:
print(print_statement) # noqa
class MyCustomHandler(CustomLogger):
def __init__(self):
@ -23,36 +25,38 @@ class MyCustomHandler(CustomLogger):
print_verbose(f"{blue_color_code}Initialized LiteLLM custom logger")
try:
print_verbose(f"Logger Initialized with following methods:")
methods = [method for method in dir(self) if inspect.ismethod(getattr(self, method))]
methods = [
method
for method in dir(self)
if inspect.ismethod(getattr(self, method))
]
# Pretty print_verbose the methods
for method in methods:
print_verbose(f" - {method}")
print_verbose(f"{reset_color_code}")
except:
pass
def log_pre_api_call(self, model, messages, kwargs):
def log_pre_api_call(self, model, messages, kwargs):
print_verbose(f"Pre-API Call")
def log_post_api_call(self, kwargs, response_obj, start_time, end_time):
def log_post_api_call(self, kwargs, response_obj, start_time, end_time):
print_verbose(f"Post-API Call")
def log_stream_event(self, kwargs, response_obj, start_time, end_time):
print_verbose(f"On Stream")
def log_success_event(self, kwargs, response_obj, start_time, end_time):
def log_success_event(self, kwargs, response_obj, start_time, end_time):
print_verbose("On Success!")
async def async_log_success_event(self, kwargs, response_obj, start_time, end_time):
print_verbose(f"On Async Success!")
response_cost = litellm.completion_cost(completion_response=response_obj)
assert response_cost > 0.0
return
async def async_log_failure_event(self, kwargs, response_obj, start_time, end_time):
async def async_log_failure_event(self, kwargs, response_obj, start_time, end_time):
try:
print_verbose(f"On Async Failure !")
except Exception as e:
@ -64,4 +68,4 @@ proxy_handler_instance = MyCustomHandler()
# need to set litellm.callbacks = [customHandler] # on the proxy
# litellm.success_callback = [async_on_succes_logger]
# litellm.success_callback = [async_on_succes_logger]