(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

@ -483,15 +483,21 @@ class Router:
self.access_groups = None
## USAGE TRACKING ##
if isinstance(litellm._async_success_callback, list):
litellm._async_success_callback.append(self.deployment_callback_on_success)
litellm.logging_callback_manager.add_litellm_async_success_callback(
self.deployment_callback_on_success
)
else:
litellm._async_success_callback.append(self.deployment_callback_on_success)
litellm.logging_callback_manager.add_litellm_async_success_callback(
self.deployment_callback_on_success
)
if isinstance(litellm.success_callback, list):
litellm.success_callback.append(self.sync_deployment_callback_on_success)
litellm.logging_callback_manager.add_litellm_success_callback(
self.sync_deployment_callback_on_success
)
else:
litellm.success_callback = [self.sync_deployment_callback_on_success]
if isinstance(litellm._async_failure_callback, list):
litellm._async_failure_callback.append(
litellm.logging_callback_manager.add_litellm_async_failure_callback(
self.async_deployment_callback_on_failure
)
else:
@ -500,7 +506,9 @@ class Router:
]
## COOLDOWNS ##
if isinstance(litellm.failure_callback, list):
litellm.failure_callback.append(self.deployment_callback_on_failure)
litellm.logging_callback_manager.add_litellm_failure_callback(
self.deployment_callback_on_failure
)
else:
litellm.failure_callback = [self.deployment_callback_on_failure]
verbose_router_logger.debug(
@ -606,7 +614,7 @@ class Router:
model_list=self.model_list,
)
if _callback is not None:
litellm.callbacks.append(_callback)
litellm.logging_callback_manager.add_litellm_callback(_callback)
def routing_strategy_init(
self, routing_strategy: Union[RoutingStrategy, str], routing_strategy_args: dict
@ -625,7 +633,7 @@ class Router:
else:
litellm.input_callback = [self.leastbusy_logger] # type: ignore
if isinstance(litellm.callbacks, list):
litellm.callbacks.append(self.leastbusy_logger) # type: ignore
litellm.logging_callback_manager.add_litellm_callback(self.leastbusy_logger) # type: ignore
elif (
routing_strategy == RoutingStrategy.USAGE_BASED_ROUTING.value
or routing_strategy == RoutingStrategy.USAGE_BASED_ROUTING
@ -636,7 +644,7 @@ class Router:
routing_args=routing_strategy_args,
)
if isinstance(litellm.callbacks, list):
litellm.callbacks.append(self.lowesttpm_logger) # type: ignore
litellm.logging_callback_manager.add_litellm_callback(self.lowesttpm_logger) # type: ignore
elif (
routing_strategy == RoutingStrategy.USAGE_BASED_ROUTING_V2.value
or routing_strategy == RoutingStrategy.USAGE_BASED_ROUTING_V2
@ -647,7 +655,7 @@ class Router:
routing_args=routing_strategy_args,
)
if isinstance(litellm.callbacks, list):
litellm.callbacks.append(self.lowesttpm_logger_v2) # type: ignore
litellm.logging_callback_manager.add_litellm_callback(self.lowesttpm_logger_v2) # type: ignore
elif (
routing_strategy == RoutingStrategy.LATENCY_BASED.value
or routing_strategy == RoutingStrategy.LATENCY_BASED
@ -658,7 +666,7 @@ class Router:
routing_args=routing_strategy_args,
)
if isinstance(litellm.callbacks, list):
litellm.callbacks.append(self.lowestlatency_logger) # type: ignore
litellm.logging_callback_manager.add_litellm_callback(self.lowestlatency_logger) # type: ignore
elif (
routing_strategy == RoutingStrategy.COST_BASED.value
or routing_strategy == RoutingStrategy.COST_BASED
@ -669,7 +677,7 @@ class Router:
routing_args={},
)
if isinstance(litellm.callbacks, list):
litellm.callbacks.append(self.lowestcost_logger) # type: ignore
litellm.logging_callback_manager.add_litellm_callback(self.lowestcost_logger) # type: ignore
else:
pass
@ -5835,8 +5843,8 @@ class Router:
self.slack_alerting_logger = _slack_alerting_logger
litellm.callbacks.append(_slack_alerting_logger) # type: ignore
litellm.success_callback.append(
litellm.logging_callback_manager.add_litellm_callback(_slack_alerting_logger) # type: ignore
litellm.logging_callback_manager.add_litellm_success_callback(
_slack_alerting_logger.response_taking_too_long_callback
)
verbose_router_logger.info(