mirror of
https://github.com/BerriAI/litellm.git
synced 2025-04-27 11:43:54 +00:00
d57be47b0f
* ci(config.yml): add a 'check_code_quality' step Addresses https://github.com/BerriAI/litellm/issues/5991 * ci(config.yml): check why circle ci doesn't pick up this test * ci(config.yml): fix to run 'check_code_quality' tests * fix(__init__.py): fix unprotected import * fix(__init__.py): don't remove unused imports * build(ruff.toml): update ruff.toml to ignore unused imports * fix: fix: ruff + pyright - fix linting + type-checking errors * fix: fix linting errors * fix(lago.py): fix module init error * fix: fix linting errors * ci(config.yml): cd into correct dir for checks * fix(proxy_server.py): fix linting error * fix(utils.py): fix bare except causes ruff linting errors * fix: ruff - fix remaining linting errors * fix(clickhouse.py): use standard logging object * fix(__init__.py): fix unprotected import * fix: ruff - fix linting errors * fix: fix linting errors * ci(config.yml): cleanup code qa step (formatting handled in local_testing) * fix(_health_endpoints.py): fix ruff linting errors * ci(config.yml): just use ruff in check_code_quality pipeline for now * build(custom_guardrail.py): include missing file * style(embedding_handler.py): fix ruff check
74 lines
2.2 KiB
Python
74 lines
2.2 KiB
Python
import os
|
|
import sys
|
|
import traceback
|
|
|
|
# this file is to test litellm/proxy
|
|
|
|
sys.path.insert(
|
|
0, os.path.abspath("../..")
|
|
) # Adds the parent directory to the system path
|
|
|
|
import inspect
|
|
|
|
import litellm
|
|
from litellm.integrations.custom_logger import CustomLogger
|
|
|
|
|
|
# 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
|
|
|
|
|
|
class MyCustomHandler(CustomLogger):
|
|
def __init__(self):
|
|
blue_color_code = "\033[94m"
|
|
reset_color_code = "\033[0m"
|
|
print_verbose(f"{blue_color_code}Initialized LiteLLM custom logger")
|
|
try:
|
|
print_verbose("Logger Initialized with following methods:")
|
|
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 Exception:
|
|
pass
|
|
|
|
def log_pre_api_call(self, model, messages, kwargs):
|
|
print_verbose("Pre-API Call")
|
|
|
|
def log_post_api_call(self, kwargs, response_obj, start_time, end_time):
|
|
print_verbose("Post-API Call")
|
|
|
|
def log_stream_event(self, kwargs, response_obj, start_time, end_time):
|
|
print_verbose("On Stream")
|
|
|
|
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("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):
|
|
try:
|
|
print_verbose("On Async Failure !")
|
|
except Exception as e:
|
|
print_verbose(f"Exception: {e}")
|
|
|
|
|
|
proxy_handler_instance = MyCustomHandler()
|
|
|
|
|
|
# need to set litellm.callbacks = [customHandler] # on the proxy
|
|
|
|
# litellm.success_callback = [async_on_succes_logger]
|