From 0ef07c62e309528eecd32f784ee4ec0a909a3192 Mon Sep 17 00:00:00 2001 From: mateo <277851410+mateo-berri@users.noreply.github.com> Date: Fri, 6 Nov 2026 20:48:25 +0000 Subject: [PATCH] fix: retain async callbacks across generateContent fallback Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- litellm/google_genai/adapters/handler.py | 6 +---- litellm/utils.py | 23 ++++++++++++++----- .../google_genai/test_google_genai_main.py | 19 +++++++++++++-- 3 files changed, 35 insertions(+), 13 deletions(-) diff --git a/litellm/google_genai/adapters/handler.py b/litellm/google_genai/adapters/handler.py index 22f125a7e52..d0425bea7e3 100644 --- a/litellm/google_genai/adapters/handler.py +++ b/litellm/google_genai/adapters/handler.py @@ -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 diff --git a/litellm/utils.py b/litellm/utils.py index 13f22ddc36e..fd6251565ee 100644 --- a/litellm/utils.py +++ b/litellm/utils.py @@ -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, diff --git a/tests/unit/google_genai/test_google_genai_main.py b/tests/unit/google_genai/test_google_genai_main.py index 353ff1b4b59..35ef54afc97 100644 --- a/tests/unit/google_genai/test_google_genai_main.py +++ b/tests/unit/google_genai/test_google_genai_main.py @@ -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