Fix litellm callback removal error (issue #2513)

Co-Authored-By: Joe Moura <joao@crewai.com>
This commit is contained in:
Devin AI
2025-04-03 09:39:10 +00:00
parent efe27bd570
commit c06eb56cf3
2 changed files with 48 additions and 8 deletions

View File

@@ -963,14 +963,15 @@ class LLM(BaseLLM):
"""
with suppress_warnings():
callback_types = [type(callback) for callback in callbacks]
for callback in litellm.success_callback[:]:
if type(callback) in callback_types:
litellm.success_callback.remove(callback)
for callback in litellm._async_success_callback[:]:
if type(callback) in callback_types:
litellm._async_success_callback.remove(callback)
litellm.success_callback = [
cb for cb in litellm.success_callback if type(cb) not in callback_types
]
litellm._async_success_callback = [
cb for cb in litellm._async_success_callback if type(cb) not in callback_types
]
litellm.callbacks = callbacks
def set_env_callbacks(self):

View File

@@ -0,0 +1,39 @@
import pytest
import litellm
from typing import Any
from crewai.llm import LLM
def test_set_callbacks_handles_removed_callbacks():
"""Test that set_callbacks handles the case where callbacks are removed during iteration."""
class CustomCallback:
pass
original_success_callback = litellm.success_callback
original_async_success_callback = litellm._async_success_callback
try:
litellm.success_callback = []
litellm._async_success_callback = []
llm = LLM(model="test-model")
callback1 = CustomCallback()
callback2 = CustomCallback()
litellm.success_callback.append(callback1)
litellm.success_callback.append(callback2)
new_callback = CustomCallback()
litellm.success_callback.remove(callback1)
llm.set_callbacks([new_callback])
assert litellm.callbacks == [new_callback]
assert len([cb for cb in litellm.success_callback if isinstance(cb, CustomCallback)]) == 0
finally:
litellm.success_callback = original_success_callback
litellm._async_success_callback = original_async_success_callback