mirror of
https://github.com/BerriAI/litellm.git
synced 2025-04-25 10:44:24 +00:00
109 lines
No EOL
4 KiB
Python
109 lines
No EOL
4 KiB
Python
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
|
|
class MyCustomHandler(CustomLogger):
|
|
def __init__(self):
|
|
blue_color_code = "\033[94m"
|
|
reset_color_code = "\033[0m"
|
|
print(f"{blue_color_code}Initialized LiteLLM custom logger")
|
|
try:
|
|
print(f"Logger Initialized with following methods:")
|
|
methods = [method for method in dir(self) if inspect.ismethod(getattr(self, method))]
|
|
|
|
# Pretty print the methods
|
|
for method in methods:
|
|
print(f" - {method}")
|
|
print(f"{reset_color_code}")
|
|
except:
|
|
pass
|
|
|
|
|
|
def log_pre_api_call(self, model, messages, kwargs):
|
|
print(f"Pre-API Call")
|
|
|
|
def log_post_api_call(self, kwargs, response_obj, start_time, end_time):
|
|
print(f"Post-API Call")
|
|
|
|
def log_stream_event(self, kwargs, response_obj, start_time, end_time):
|
|
print(f"On Stream")
|
|
|
|
def log_success_event(self, kwargs, response_obj, start_time, end_time):
|
|
print("On Success!")
|
|
|
|
async def async_log_success_event(self, kwargs, response_obj, start_time, end_time):
|
|
print(f"On Async Success!")
|
|
# 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
|
|
|
|
# Calculate cost using litellm.completion_cost()
|
|
cost = litellm.completion_cost(completion_response=response_obj)
|
|
response = response_obj
|
|
# tokens used in response
|
|
usage = response_obj["usage"]
|
|
|
|
print(
|
|
f"""
|
|
Model: {model},
|
|
Messages: {messages},
|
|
User: {user},
|
|
Usage: {usage},
|
|
Cost: {cost},
|
|
Response: {response}
|
|
Proxy Metadata: {metadata}
|
|
"""
|
|
)
|
|
return
|
|
|
|
async def async_log_failure_event(self, kwargs, response_obj, start_time, end_time):
|
|
try:
|
|
print(f"On Async Failure !")
|
|
print("\nkwargs", kwargs)
|
|
# 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()
|
|
cost = litellm.completion_cost(completion_response=response_obj)
|
|
print("now checking response obj")
|
|
|
|
print(
|
|
f"""
|
|
Model: {model},
|
|
Messages: {messages},
|
|
User: {user},
|
|
Cost: {cost},
|
|
Response: {response_obj}
|
|
Proxy Metadata: {metadata}
|
|
Exception: {exception_event}
|
|
Traceback: {traceback_event}
|
|
"""
|
|
)
|
|
except Exception as e:
|
|
print(f"Exception: {e}")
|
|
|
|
|
|
proxy_handler_instance = MyCustomHandler()
|
|
|
|
|
|
# need to set litellm.callbacks = [customHandler] # on the proxy
|
|
|
|
# litellm.success_callback = [async_on_succes_logger] |