(Refactor / QA) - Use LoggingCallbackManager to append callbacks and ensure no duplicate callbacks are added (#8112)

* LoggingCallbackManager

* add logging_callback_manager

* use logging_callback_manager

* add add_litellm_failure_callback

* use add_litellm_callback

* use add_litellm_async_success_callback

* add_litellm_async_failure_callback

* linting fix

* fix logging callback manager

* test_duplicate_multiple_loggers_test

* use _reset_all_callbacks

* fix testing with dup callbacks

* test_basic_image_generation

* reset callbacks for tests

* fix check for _add_custom_logger_to_list

* fix test_amazing_sync_embedding

* fix _get_custom_logger_key

* fix batches testing

* fix _reset_all_callbacks

* fix _check_callback_list_size

* add callback_manager_test

* fix test gemini-2.0-flash-thinking-exp-01-21
This commit is contained in:
Ishaan Jaff 2025-01-30 19:35:50 -08:00 committed by GitHub
parent 11c8d07ed3
commit fa1c42378f
19 changed files with 607 additions and 59 deletions

View file

@ -352,8 +352,12 @@ def _add_custom_logger_callback_to_specific_event(
and _custom_logger_class_exists_in_success_callbacks(callback_class)
is False
):
litellm.success_callback.append(callback_class)
litellm._async_success_callback.append(callback_class)
litellm.logging_callback_manager.add_litellm_success_callback(
callback_class
)
litellm.logging_callback_manager.add_litellm_async_success_callback(
callback_class
)
if callback in litellm.success_callback:
litellm.success_callback.remove(
callback
@ -367,8 +371,12 @@ def _add_custom_logger_callback_to_specific_event(
and _custom_logger_class_exists_in_failure_callbacks(callback_class)
is False
):
litellm.failure_callback.append(callback_class)
litellm._async_failure_callback.append(callback_class)
litellm.logging_callback_manager.add_litellm_failure_callback(
callback_class
)
litellm.logging_callback_manager.add_litellm_async_failure_callback(
callback_class
)
if callback in litellm.failure_callback:
litellm.failure_callback.remove(
callback
@ -447,13 +455,13 @@ def function_setup( # noqa: PLR0915
if callback not in litellm.input_callback:
litellm.input_callback.append(callback) # type: ignore
if callback not in litellm.success_callback:
litellm.success_callback.append(callback) # type: ignore
litellm.logging_callback_manager.add_litellm_success_callback(callback) # type: ignore
if callback not in litellm.failure_callback:
litellm.failure_callback.append(callback) # type: ignore
litellm.logging_callback_manager.add_litellm_failure_callback(callback) # type: ignore
if callback not in litellm._async_success_callback:
litellm._async_success_callback.append(callback) # type: ignore
litellm.logging_callback_manager.add_litellm_async_success_callback(callback) # type: ignore
if callback not in litellm._async_failure_callback:
litellm._async_failure_callback.append(callback) # type: ignore
litellm.logging_callback_manager.add_litellm_async_failure_callback(callback) # type: ignore
print_verbose(
f"Initialized litellm callbacks, Async Success Callbacks: {litellm._async_success_callback}"
)
@ -488,12 +496,16 @@ def function_setup( # noqa: PLR0915
removed_async_items = []
for index, callback in enumerate(litellm.success_callback): # type: ignore
if inspect.iscoroutinefunction(callback):
litellm._async_success_callback.append(callback)
litellm.logging_callback_manager.add_litellm_async_success_callback(
callback
)
removed_async_items.append(index)
elif callback == "dynamodb" or callback == "openmeter":
# dynamo is an async callback, it's used for the proxy and needs to be async
# we only support async dynamo db logging for acompletion/aembedding since that's used on proxy
litellm._async_success_callback.append(callback)
litellm.logging_callback_manager.add_litellm_async_success_callback(
callback
)
removed_async_items.append(index)
elif (
callback in litellm._known_custom_logger_compatible_callbacks
@ -509,7 +521,9 @@ def function_setup( # noqa: PLR0915
removed_async_items = []
for index, callback in enumerate(litellm.failure_callback): # type: ignore
if inspect.iscoroutinefunction(callback):
litellm._async_failure_callback.append(callback)
litellm.logging_callback_manager.add_litellm_async_failure_callback(
callback
)
removed_async_items.append(index)
elif (
callback in litellm._known_custom_logger_compatible_callbacks