scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-11-13 08:52:02 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 174

fix: retain async callbacks across generateContent fallback

Browse source 
Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>
This commit is contained in:
mateo 2026-11-06 20:48:25 +00:00
parent 4d0aadbabd
commit 0ef07c62e3
3 changed files with 35 additions and 13 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

6
litellm/google_genai/adapters/handler.py
View file

@ -72,11 +72,7 @@ class GenerateContentToCompletionHandler:
completion_kwargs["litellm_metadata"] = extra_kwargs["litellm_metadata"]
if isinstance(custom_logging_obj, LiteLLMLoggingObj):
completion_kwargs["success_callback"] = [
*(custom_logging_obj.dynamic_success_callbacks or ()),
*(custom_logging_obj.dynamic_async_success_callbacks or ()),
]
completion_kwargs["failure_callback"] = list(custom_logging_obj.dynamic_failure_callbacks or ())
completion_kwargs["_dynamic_callback_context"] = custom_logging_obj
if stream:
completion_kwargs["stream"] = stream

23
litellm/utils.py
View file

@ -1372,10 +1372,22 @@ def function_setup(
get_set_callbacks: Final[Callable[[], Callable[..., None]]] = litellm_utils.get_set_callbacks
get_set_callbacks()(callback_list=callbacks_to_set_up, function_id=function_id)
### DYNAMIC CALLBACKS ###
dynamic_success_callbacks: list[str | Callable[..., object] | CustomLogger] | None = None
dynamic_async_success_callbacks: list[str | Callable[..., object] | CustomLogger] | None = None
dynamic_failure_callbacks: list[str | Callable[..., object] | CustomLogger] | None = None
dynamic_async_failure_callbacks: Final[list[str | Callable[..., object] | CustomLogger] | None] = None
get_litellm_logging_class: Final[_LoggingClassGetter] = litellm_utils.get_litellm_logging_class
logging_class: Final = get_litellm_logging_class()
callback_context: Final[object] = kwargs.pop("_dynamic_callback_context", None)
_dynamic_callback_context: Final = callback_context if isinstance(callback_context, logging_class) else None
dynamic_success_callbacks: list[str | Callable[..., object] | CustomLogger] | None = (
copy.copy(_dynamic_callback_context.dynamic_success_callbacks) if _dynamic_callback_context else None
)
dynamic_async_success_callbacks: list[str | Callable[..., object] | CustomLogger] | None = (
copy.copy(_dynamic_callback_context.dynamic_async_success_callbacks) if _dynamic_callback_context else None
)
dynamic_failure_callbacks: list[str | Callable[..., object] | CustomLogger] | None = (
copy.copy(_dynamic_callback_context.dynamic_failure_callbacks) if _dynamic_callback_context else None
)
dynamic_async_failure_callbacks: Final[list[str | Callable[..., object] | CustomLogger] | None] = (
copy.copy(_dynamic_callback_context.dynamic_async_failure_callbacks) if _dynamic_callback_context else None
)
if kwargs.get("success_callback", None) is not None and isinstance(kwargs["success_callback"], list):
removed_async_items = []
for index, callback in enumerate(kwargs["success_callback"]):
@ -1573,9 +1585,8 @@ def function_setup(
call_type=call_type,
):
stream = True
get_litellm_logging_class: Final[_LoggingClassGetter] = litellm_utils.get_litellm_logging_class
# Victim for object pool
logging_obj = get_litellm_logging_class()( # rebind-ok: 2nd assignment to logging_obj (see initial None above)
logging_obj = logging_class( # rebind-ok: 2nd assignment to logging_obj (see initial None above)
model=model,
messages=messages,
stream=stream,

19
tests/unit/google_genai/test_google_genai_main.py
View file

@ -529,6 +529,7 @@ class ContentFailureRecorder(CustomLogger):
def __init__(self) -> None:
super().__init__()
self.calls: tuple[str, ...] = ()
self.asynchronous_calls: tuple[str, ...] = ()
self.logged: Final = Event()
def log_failure_event(self, kwargs: dict[str, object], response_obj: object, start_time: object, end_time: object) -> None:
@ -538,13 +539,16 @@ class ContentFailureRecorder(CustomLogger):
async def async_log_failure_event(
self, kwargs: dict[str, object], response_obj: object, start_time: object, end_time: object
) -> None:
self.asynchronous_calls = (*self.asynchronous_calls, str(kwargs["call_type"]))
self.log_failure_event(kwargs, response_obj, start_time, end_time)
@pytest.mark.parametrize("is_async", [False, True])
@pytest.mark.parametrize("use_logging_context", [False, True])
@pytest.mark.asyncio
async def test_custom_content_adapter_delivers_per_request_failure_callbacks_once(
native_content_handler: NativeContentLLM, monkeypatch: pytest.MonkeyPatch, is_async: bool
native_content_handler: NativeContentLLM, monkeypatch: pytest.MonkeyPatch,
is_async: bool, use_logging_context: bool,
) -> None:
monkeypatch.setattr(
litellm, "custom_provider_map", [{"provider": "native_image_test", "custom_handler": FailingCompletionContentLLM()}]
@ -554,9 +558,18 @@ async def test_custom_content_adapter_delivers_per_request_failure_callbacks_onc
monkeypatch.setattr(litellm, "callbacks", [global_recorder])
monkeypatch.setattr(litellm, "failure_callback", [global_recorder])
monkeypatch.setattr(litellm, "_async_failure_callback", [])
callback_args: Final = (
{"litellm_logging_obj": Logging(
model="native_image_test/chat-model", messages=[], stream=False,
call_type="agenerate_content" if is_async else "generate_content",
start_time=datetime(2026, 1, 1), litellm_call_id="failure-test", function_id="failure-test",
dynamic_failure_callbacks=[request_recorder], dynamic_async_failure_callbacks=[request_recorder],
)}
if use_logging_context else {"failure_callback": [request_recorder.log_failure_event]}
)
request: Final = dict(
model="native_image_test/chat-model", contents=[{"parts": [{"text": "hello"}]}],
failure_callback=[request_recorder.log_failure_event],
**callback_args,
)
if is_async:
with pytest.raises(litellm.RateLimitError, match="adapter quota exhausted"):
@ -569,6 +582,8 @@ async def test_custom_content_adapter_delivers_per_request_failure_callbacks_onc
expected: Final = ("acompletion" if is_async else "completion",)
assert request_recorder.calls == expected
assert global_recorder.calls == expected
assert request_recorder.asynchronous_calls == (expected if is_async and use_logging_context else ())
assert global_recorder.asynchronous_calls == (expected if is_async else ())
@pytest.mark.asyncio