scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-10-08 03:08:45 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 183

fix: preserve shared optional callbacks

Browse source 
Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>
This commit is contained in:
mateo 2026-09-02 03:28:04 +00:00
parent 70a4f74a0d
commit 1d3e26fd98
2 changed files with 50 additions and 10 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

19
litellm/router.py
View file

@ -2099,16 +2099,19 @@ class Router:
self.add_optional_pre_call_checks(optional_pre_call_checks)
def _remove_optional_callbacks_of_type(self, callback_cls: type[CustomLogger]) -> None:
if self.optional_callbacks is None:
if self.optional_callbacks is None or not any(type(cb) is callback_cls for cb in self.optional_callbacks):
return
removed: Final = [cb for cb in self.optional_callbacks if isinstance(cb, callback_cls)]
if not removed:
self.optional_callbacks = [cb for cb in self.optional_callbacks if type(cb) is not callback_cls]
if any(
router is not self and any(type(cb) is callback_cls for cb in (router.optional_callbacks or []))
for router in tuple(_live_routers)
):
return
self.optional_callbacks = [cb for cb in self.optional_callbacks if not isinstance(cb, callback_cls)]
for cb in removed:
litellm.logging_callback_manager.remove_callback_from_list_by_object(
litellm.callbacks, cb, require_self=False
)
for cb in tuple(litellm.callbacks):
if type(cb) is callback_cls:
litellm.logging_callback_manager.remove_callback_from_list_by_object(
litellm.callbacks, cb, require_self=False
)
def print_deployment(self, deployment: dict):
"""

41
tests/test_litellm/test_router_retry_policy_update.py
View file

@ -151,8 +151,45 @@ def test_remove_optional_pre_call_check_removes_local_and_global_callbacks():
router.set_optional_pre_call_checks(["prompt_caching"])
router._remove_optional_callbacks_of_type(PromptCachingDeploymentCheck)
assert not any(isinstance(callback, PromptCachingDeploymentCheck) for callback in (router.optional_callbacks or []))
assert not any(isinstance(callback, PromptCachingDeploymentCheck) for callback in litellm.callbacks)
assert not any(type(callback) is PromptCachingDeploymentCheck for callback in (router.optional_callbacks or []))
assert not any(type(callback) is PromptCachingDeploymentCheck for callback in litellm.callbacks)
def test_remove_optional_pre_call_check_keeps_global_callback_for_another_router():
router_a = _build_router()
router_b = _build_router()
router_a.update_settings(optional_pre_call_checks=["prompt_caching"])
router_b.update_settings(optional_pre_call_checks=["prompt_caching"])
router_a.update_settings(optional_pre_call_checks=[])
assert not any(type(callback) is PromptCachingDeploymentCheck for callback in (router_a.optional_callbacks or []))
assert any(type(callback) is PromptCachingDeploymentCheck for callback in (router_b.optional_callbacks or []))
assert any(type(callback) is PromptCachingDeploymentCheck for callback in litellm.callbacks)
router_b.update_settings(optional_pre_call_checks=[])
assert not any(type(callback) is PromptCachingDeploymentCheck for callback in (router_b.optional_callbacks or []))
assert not any(type(callback) is PromptCachingDeploymentCheck for callback in litellm.callbacks)
def test_remove_optional_pre_call_check_keeps_global_callback_when_second_router_clears_first():
router_a = _build_router()
router_b = _build_router()
router_a.update_settings(optional_pre_call_checks=["prompt_caching"])
router_b.update_settings(optional_pre_call_checks=["prompt_caching"])
router_b.update_settings(optional_pre_call_checks=[])
assert any(type(callback) is PromptCachingDeploymentCheck for callback in (router_a.optional_callbacks or []))
assert not any(type(callback) is PromptCachingDeploymentCheck for callback in (router_b.optional_callbacks or []))
assert any(type(callback) is PromptCachingDeploymentCheck for callback in litellm.callbacks)
router_a.update_settings(optional_pre_call_checks=[])
assert not any(type(callback) is PromptCachingDeploymentCheck for callback in litellm.callbacks)
def test_update_settings_replaces_toggleable_pre_call_checks():