phoenix/litellm-mirror
0
0
Fork 1
mirror of https://github.com/BerriAI/litellm.git synced 2025-04-25 10:44:24 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions

fix(utils.py): move adding custom logger callback to success event in… (#7905)

Browse source 
* fix(utils.py): move adding custom logger callback to success event into separate function + don't add success callback to failure event

if user is explicitly choosing 'success' callback, don't log failure as well

* test(test_utils.py): add unit test to ensure custom logger callback only adds callback to specific event

* fix(utils.py): remove string from list of callbacks once corresponding callback class is added

prevents floating values - simplifies testing

* fix(utils.py): fix linting error

* test: cleanup args before test

* test: fix test

* test: update test

* test: fix test
This commit is contained in:
Krish Dholakia 2025-01-22 21:49:09 -08:00 • committed by GitHub
parent cefbada875
commit 4911cd80a1
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
5 changed files with 98 additions and 49 deletions
Download patch file Download diff file Expand all files Collapse all files

78
litellm/utils.py
View file

@ -318,9 +318,61 @@ def custom_llm_setup():
litellm._custom_providers.append(custom_llm["provider"])
def _add_custom_logger_callback_to_specific_event(
callback: str, logging_event: Literal["success", "failure"]
) -> None:
"""
Add a custom logger callback to the specific event
"""
from litellm import _custom_logger_compatible_callbacks_literal
from litellm.litellm_core_utils.litellm_logging import (
_init_custom_logger_compatible_class,
)
if callback not in litellm._known_custom_logger_compatible_callbacks:
verbose_logger.debug(
f"Callback {callback} is not a valid custom logger compatible callback. Known list - {litellm._known_custom_logger_compatible_callbacks}"
)
return
callback_class = _init_custom_logger_compatible_class(
cast(_custom_logger_compatible_callbacks_literal, callback),
internal_usage_cache=None,
llm_router=None,
)
# don't double add a callback
if callback_class is not None and not any(
isinstance(cb, type(callback_class)) for cb in litellm.callbacks # type: ignore
):
if logging_event == "success":
litellm.success_callback.append(callback_class)
litellm._async_success_callback.append(callback_class)
if callback in litellm.success_callback:
litellm.success_callback.remove(
callback
) # remove the string from the callback list
if callback in litellm._async_success_callback:
litellm._async_success_callback.remove(
callback
) # remove the string from the callback list
elif logging_event == "failure":
litellm.failure_callback.append(callback_class)
litellm._async_failure_callback.append(callback_class)
if callback in litellm.failure_callback:
litellm.failure_callback.remove(
callback
) # remove the string from the callback list
if callback in litellm._async_failure_callback:
litellm._async_failure_callback.remove(
callback
) # remove the string from the callback list
def function_setup( # noqa: PLR0915
original_function: str, rules_obj, start_time, *args, **kwargs
): # just run once to check if user wants to send their data anywhere - PostHog/Sentry/Slack/etc.
### NOTICES ###
from litellm import Logging as LiteLLMLogging
from litellm.litellm_core_utils.litellm_logging import set_callbacks
@ -401,27 +453,11 @@ def function_setup( # noqa: PLR0915
# we only support async dynamo db logging for acompletion/aembedding since that's used on proxy
litellm._async_success_callback.append(callback)
removed_async_items.append(index)
elif callback in litellm._known_custom_logger_compatible_callbacks:
from litellm.litellm_core_utils.litellm_logging import (
_init_custom_logger_compatible_class,
)
callback_class = _init_custom_logger_compatible_class(
callback, # type: ignore
internal_usage_cache=None,
llm_router=None, # type: ignore
)
# don't double add a callback
if callback_class is not None and not any(
isinstance(cb, type(callback_class)) for cb in litellm.callbacks
):
litellm.callbacks.append(callback_class) # type: ignore
litellm.input_callback.append(callback_class) # type: ignore
litellm.success_callback.append(callback_class) # type: ignore
litellm.failure_callback.append(callback_class) # type: ignore
litellm._async_success_callback.append(callback_class) # type: ignore
litellm._async_failure_callback.append(callback_class) # type: ignore
elif (
callback in litellm._known_custom_logger_compatible_callbacks
and isinstance(callback, str)
):
_add_custom_logger_callback_to_specific_event(callback, "success")
# Pop the async items from success_callback in reverse order to avoid index issues
for index in reversed(removed_async_items):