diff --git a/litellm/litellm_core_utils/litellm_logging.py b/litellm/litellm_core_utils/litellm_logging.py index 97c4d038734..f87c24dbdf9 100644 --- a/litellm/litellm_core_utils/litellm_logging.py +++ b/litellm/litellm_core_utils/litellm_logging.py @@ -4719,7 +4719,11 @@ def _init_custom_logger_compatible_class( return newrelic_logger return None except Exception as e: - verbose_logger.exception("[Non-Blocking Error] Error initializing custom logger: %s", e) + verbose_logger.exception( + "Failed to initialize custom logger for integration %r: %s. Callback will not be registered.", + logging_integration, + e, + ) return None return None diff --git a/litellm/proxy/common_request_processing.py b/litellm/proxy/common_request_processing.py index ed66bb65ced..f3b942009c6 100644 --- a/litellm/proxy/common_request_processing.py +++ b/litellm/proxy/common_request_processing.py @@ -4,7 +4,7 @@ import json import logging import math import traceback -from collections.abc import AsyncGenerator, Awaitable, Callable, Coroutine, Mapping, Sequence +from collections.abc import AsyncGenerator, Awaitable, Callable, Mapping, Sequence from datetime import datetime from functools import lru_cache from types import MappingProxyType @@ -3138,10 +3138,31 @@ class ProxyBaseLLMRequestProcessing: from litellm.litellm_core_utils.logging_worker import GLOBAL_LOGGING_WORKER - async def _on_deferred_native_stream_complete( - logging_coroutine: Coroutine[object, object, object], - ) -> None: - GLOBAL_LOGGING_WORKER.ensure_initialized_and_enqueue(async_coroutine=logging_coroutine) + _captured_native_logging_obj: Final = logging_obj + + async def _on_deferred_native_stream_complete(*args: object) -> None: + # 1-arg (logging_coroutine,) is what the native passthrough/responses + # iterators write; 2-arg (assembled_response, cache_hit) is what the + # /v1/messages completion bridge writes via its inner CustomStreamWrapper. + # The proxy sees a bare byte generator in both cases, so arity dispatch + # here is the only place we can tell them apart. + if len(args) == 1 and asyncio.iscoroutine(args[0]): + GLOBAL_LOGGING_WORKER.ensure_initialized_and_enqueue(async_coroutine=args[0]) + return + if len(args) == 2: + assembled_response, cache_hit = args + await _as_success_dispatcher(_captured_native_logging_obj).dispatch_success_handlers( + assembled_response, + cache_hit=cache_hit, + start_time=None, + end_time=None, + prefer_async_handlers=True, + ) + return + verbose_proxy_logger.error( + "Deferred stream logging received unexpected args shape (len=%d); dropping to avoid double-logging", + len(args), + ) logging_obj._on_deferred_stream_complete = _on_deferred_native_stream_complete diff --git a/tests/test_litellm/proxy/guardrails/test_deferred_guardrail_logging.py b/tests/test_litellm/proxy/guardrails/test_deferred_guardrail_logging.py index 8fde4cc9d5e..53a9d32baf2 100644 --- a/tests/test_litellm/proxy/guardrails/test_deferred_guardrail_logging.py +++ b/tests/test_litellm/proxy/guardrails/test_deferred_guardrail_logging.py @@ -1415,6 +1415,34 @@ class TestArmDeferredStreamDispatch: mock_enqueue.assert_called_once_with(async_coroutine=coro) coro.close() + @pytest.mark.asyncio + async def test_native_stream_closure_accepts_bridge_two_arg_shape(self): + """/v1/messages via the completion bridge returns a bare byte generator, + so arming lands in the native branch. But the inner CustomStreamWrapper + writes the 2-arg (assembled_response, cache_hit) shape at end-of-stream, + which the pre-fix 1-arg closure blew up on.""" + logging_obj, recorded = self._dispatch_recording_logging_obj() + + async def _agen(): + yield b"x" + + self._processor()._arm_deferred_stream_dispatch( + response=_agen(), + route_type="anthropic_messages", + user_api_key_dict=MagicMock(), + logging_obj=logging_obj, + ) + closure = logging_obj._on_deferred_stream_complete + assert closure is not None + + assembled = object() + await closure(assembled, False) + await asyncio.sleep(0) + + assert recorded["result"] is assembled + assert recorded["cache_hit"] is False + assert recorded["prefer_async_handlers"] is True + @pytest.mark.asyncio async def test_csw_closure_routes_through_deferred_stream_guardrails(self, monkeypatch): from litellm.litellm_core_utils.streaming_handler import CustomStreamWrapper