scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-10-23 05:34:41 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 24

fix(proxy): retain deferred stream callbacks

Browse source
This commit is contained in:
Ivana Dunisijevic 2026-10-08 17:08:32 +02:00
parent d92438c4ce
commit 307bb06a05
No known key found for this signature in database
4 changed files with 104 additions and 51 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

14
litellm/proxy/common_request_processing.py
View file

@ -3630,14 +3630,12 @@ class ProxyBaseLLMRequestProcessing:
# callbacks fire regardless of the call-type heuristic in
# _is_sync_litellm_request (which only recognizes a subset of
# async markers stored in litellm_params).
asyncio.create_task(
_as_success_dispatcher(captured_logging_obj).dispatch_success_handlers(
_response,
cache_hit=cache_hit,
start_time=None,
end_time=None,
prefer_async_handlers=True,
)
await _as_success_dispatcher(captured_logging_obj).dispatch_success_handlers(
_response,
cache_hit=cache_hit,
start_time=None,
end_time=None,
prefer_async_handlers=True,
)
except Exception as e:
verbose_proxy_logger.exception(

23
litellm/proxy/utils.py
View file

@ -1103,6 +1103,21 @@ def _call_type_for_route(route: str | None) -> str | None:
_PROXY_ONLY_LLM_API_ERRORS: Final = (HTTPException, ProxyException, GuardrailRaisedException)
_LOG_DB_METRICS_CALL_TYPES: Final = frozenset(("get_data", "insert_data", "update_data", "delete_data"))
_DEFERRED_STREAM_CALLBACK_TASKS: Final[set[asyncio.Task[None]]] = set() # mutable-ok: retains background tasks
async def _run_deferred_stream_callback(callback: Coroutine[object, object, object]) -> None:
try:
await callback
except Exception as e:
verbose_proxy_logger.exception("Error in deferred stream callback: %s", e)
def _retain_deferred_stream_callback(callback: Coroutine[object, object, object]) -> asyncio.Task[None]:
task: Final = asyncio.create_task(_run_deferred_stream_callback(callback))
_DEFERRED_STREAM_CALLBACK_TASKS.add(task)
task.add_done_callback(_DEFERRED_STREAM_CALLBACK_TASKS.discard)
return task
def _failure_fields_to_lift(request_data: Mapping[str, object]) -> Mapping[str, object]:
@ -4172,7 +4187,7 @@ class ProxyLogging:
record_served_output_texts(logging_obj.model_call_details, served_stream_output_texts(served_chunks))
@staticmethod
def fire_deferred_stream_logging(request_data: dict) -> None:
def fire_deferred_stream_logging(request_data: dict) -> asyncio.Task[None] | None:
"""
Fire the deferred streaming logging callback after the full streaming
pipeline (including guardrail end-of-stream blocks) has completed.
@ -4183,8 +4198,7 @@ class ProxyLogging:
"""
logging_obj: Final = request_data.get("litellm_logging_obj")
if logging_obj is None:
return
from litellm.litellm_core_utils.logging_worker import GLOBAL_LOGGING_WORKER
return None
_deferred_cb: Final[Callable[..., Coroutine[object, object, object]] | None] = getattr(
logging_obj, "_on_deferred_stream_complete", None
@ -4193,7 +4207,8 @@ class ProxyLogging:
if _deferred_cb is not None and _args is not None:
logging_obj._on_deferred_stream_complete = None
logging_obj._deferred_stream_complete_args = None
GLOBAL_LOGGING_WORKER.ensure_initialized_and_enqueue(async_coroutine=_deferred_cb(*_args))
return _retain_deferred_stream_callback(_deferred_cb(*_args))
return None
_fire_deferred_stream_logging = fire_deferred_stream_logging

92
tests/unit/proxy/guardrails/test_deferred_guardrail_logging.py
View file

@ -736,8 +736,9 @@ class TestDeferredStreamingClosure:
# CSW stored args; now simulate what ProxyLogging does
request_data = {"litellm_logging_obj": mock_logging_obj}
ProxyLogging._fire_deferred_stream_logging(request_data)
await GLOBAL_LOGGING_WORKER.flush()
deferred_task = ProxyLogging._fire_deferred_stream_logging(request_data)
assert deferred_task is not None
await deferred_task
assert guardrail_called is True, "Guardrail hook should be called"
assert (
@ -946,8 +947,9 @@ class TestDeferredStreamingClosure:
# CSW stored args; now simulate what ProxyLogging does
request_data = {"litellm_logging_obj": mock_logging_obj}
ProxyLogging._fire_deferred_stream_logging(request_data)
await GLOBAL_LOGGING_WORKER.flush()
deferred_task = ProxyLogging._fire_deferred_stream_logging(request_data)
assert deferred_task is not None
await deferred_task
assert hook_called is True, "Production closure must call guardrail hook"
assert (
@ -1316,9 +1318,12 @@ class TestFireDeferredStreamLogging:
mock_logging_obj._deferred_stream_complete_args = ("test_response", True)
request_data = {"litellm_logging_obj": mock_logging_obj}
ProxyLogging._fire_deferred_stream_logging(request_data)
await GLOBAL_LOGGING_WORKER.flush()
with patch.object(GLOBAL_LOGGING_WORKER, "ensure_initialized_and_enqueue") as enqueue:
deferred_task = ProxyLogging._fire_deferred_stream_logging(request_data)
assert deferred_task is not None
await deferred_task
enqueue.assert_not_called()
assert callback_called is True
assert callback_args["response"] == "test_response"
assert callback_args["cache_hit"] is True
@ -1326,6 +1331,34 @@ class TestFireDeferredStreamLogging:
assert mock_logging_obj._on_deferred_stream_complete is None
assert mock_logging_obj._deferred_stream_complete_args is None
@pytest.mark.asyncio
async def test_deferred_guardrails_wait_for_success_logging(self):
logging_started = asyncio.Event()
release_logging = asyncio.Event()
mock_logging_obj = MagicMock()
mock_logging_obj.model_call_details = {"metadata": {}}
async def track_async_success(*args, **kwargs):
logging_started.set()
await release_logging.wait()
_attach_mock_success_dispatch(mock_logging_obj, track_async_success)
with patch("litellm.callbacks", []):
deferred_task = asyncio.create_task(
ProxyBaseLLMRequestProcessing._run_deferred_stream_guardrails(
captured_data={"model": "gpt-4", "metadata": {}},
captured_user_api_key_dict=UserAPIKeyAuth(api_key="test"),
captured_logging_obj=mock_logging_obj,
assembled_response=MagicMock(),
cache_hit=False,
)
)
await logging_started.wait()
assert deferred_task.done() is False
release_logging.set()
await deferred_task
@pytest.mark.asyncio
async def test_noop_when_no_deferred_args(self):
"""_fire_deferred_stream_logging should be a no-op when no deferred
@ -1516,8 +1549,9 @@ class TestArmDeferredStreamDispatch:
assembled = object()
logging_obj._deferred_stream_complete_args = (assembled, False)
ProxyLogging._fire_deferred_stream_logging({"litellm_logging_obj": logging_obj})
await GLOBAL_LOGGING_WORKER.flush()
deferred_task = ProxyLogging._fire_deferred_stream_logging({"litellm_logging_obj": logging_obj})
assert deferred_task is not None
await deferred_task
assert recorded["result"] is assembled
assert recorded["cache_hit"] is False
@ -1550,8 +1584,9 @@ class TestArmDeferredStreamDispatch:
assembled = object()
logging_obj._deferred_stream_complete_args = (assembled, False)
ProxyLogging._fire_deferred_stream_logging({"litellm_logging_obj": logging_obj})
await GLOBAL_LOGGING_WORKER.flush()
deferred_task = ProxyLogging._fire_deferred_stream_logging({"litellm_logging_obj": logging_obj})
assert deferred_task is not None
await deferred_task
assert recorded["result"] is assembled
assert recorded["cache_hit"] is False
@ -1577,13 +1612,14 @@ class TestArmDeferredStreamDispatch:
coro = _logging_coroutine()
logging_obj._deferred_stream_complete_args = (coro,)
with patch.object( # test-quality-ok: GLOBAL_LOGGING_WORKER is a process-global singleton with no injection seam
GLOBAL_LOGGING_WORKER, "ensure_initialized_and_enqueue"
) as mock_enqueue:
ProxyLogging._fire_deferred_stream_logging({"litellm_logging_obj": logging_obj})
deferred = mock_enqueue.call_args.kwargs["async_coroutine"]
mock_enqueue.reset_mock()
await deferred
with (
patch.object( # test-quality-ok: GLOBAL_LOGGING_WORKER is a process-global singleton with no injection seam
GLOBAL_LOGGING_WORKER, "ensure_initialized_and_enqueue"
) as mock_enqueue
):
deferred_task = ProxyLogging._fire_deferred_stream_logging({"litellm_logging_obj": logging_obj})
assert deferred_task is not None
await deferred_task
mock_enqueue.assert_called_once_with(async_coroutine=coro)
assert recorded == {}
coro.close()
@ -1610,13 +1646,14 @@ class TestArmDeferredStreamDispatch:
assembled = object()
logging_obj._deferred_stream_complete_args = (assembled, True)
with patch.object( # test-quality-ok: GLOBAL_LOGGING_WORKER is a process-global singleton with no injection seam
GLOBAL_LOGGING_WORKER, "ensure_initialized_and_enqueue"
) as mock_enqueue:
ProxyLogging._fire_deferred_stream_logging({"litellm_logging_obj": logging_obj})
deferred = mock_enqueue.call_args.kwargs["async_coroutine"]
mock_enqueue.reset_mock()
await deferred
with (
patch.object( # test-quality-ok: GLOBAL_LOGGING_WORKER is a process-global singleton with no injection seam
GLOBAL_LOGGING_WORKER, "ensure_initialized_and_enqueue"
) as mock_enqueue
):
deferred_task = ProxyLogging._fire_deferred_stream_logging({"litellm_logging_obj": logging_obj})
assert deferred_task is not None
await deferred_task
mock_enqueue.assert_not_called()
assert recorded["result"] is assembled
@ -1645,10 +1682,9 @@ class TestArmDeferredStreamDispatch:
) as mock_enqueue,
caplog.at_level(logging.ERROR, logger="LiteLLM Proxy"),
):
ProxyLogging._fire_deferred_stream_logging({"litellm_logging_obj": logging_obj})
deferred = mock_enqueue.call_args.kwargs["async_coroutine"]
mock_enqueue.reset_mock()
await deferred
deferred_task = ProxyLogging._fire_deferred_stream_logging({"litellm_logging_obj": logging_obj})
assert deferred_task is not None
await deferred_task
mock_enqueue.assert_not_called()
assert recorded == {}

26
tests/unit/proxy/utils/proxy_logging/test_streaming_hooks.py
View file

@ -10,10 +10,10 @@ Covers ``_wrap_streaming_iterator_with_enrichment``,
from __future__ import annotations
import asyncio
from collections.abc import AsyncGenerator, AsyncIterator
from collections.abc import AsyncGenerator, AsyncIterator, Callable
from datetime import datetime
from typing import Any, Dict, Final, List
from unittest.mock import AsyncMock, MagicMock, patch
from unittest.mock import AsyncMock, MagicMock
import pytest
from fastapi import HTTPException
@ -23,7 +23,6 @@ from litellm.exceptions import GuardrailRaisedException
from litellm.integrations.custom_guardrail import CustomGuardrail
from litellm.integrations.custom_logger import CustomLogger
from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj
from litellm.litellm_core_utils.logging_worker import GLOBAL_LOGGING_WORKER
from litellm.llms.anthropic.pass_through.messages.streaming_iterator import (
BaseAnthropicMessagesStreamingIterator,
)
@ -180,6 +179,13 @@ async def _one_chunk() -> AsyncGenerator[object, None]:
yield "chunk"
async def _wait_until(condition: Callable[[], bool]) -> None:
for _ in range(200):
if condition():
return
await asyncio.sleep(0.01)
class _AttributeStream:
_hidden_params = {"model_id": "m-1"}
model = "gpt-x"
@ -564,7 +570,7 @@ async def test_native_messages_stream_logging_fires_after_guardrail_end_of_strea
request_data=request_data,
):
pass
await GLOBAL_LOGGING_WORKER.flush()
await _wait_until(lambda: len(events) == 2)
assert events == ["scan_appended", ("logging_dispatched", "post_call_entry_visible", True)]
@ -599,7 +605,7 @@ async def test_native_messages_stream_logging_fires_when_guardrail_blocks_after_
request_data=request_data,
):
pass
await GLOBAL_LOGGING_WORKER.flush()
await _wait_until(lambda: len(events) == 1)
assert [event[0] for event in events] == ["logging_dispatched"]
assert logging_obj._deferred_stream_complete_args is None
@ -723,7 +729,7 @@ async def test_chat_stream_generic_callback_error_after_stream_end_still_flushes
request_data=request_data,
):
pass
await GLOBAL_LOGGING_WORKER.flush()
await _wait_until(lambda: len(events) == 1)
snapshot = {
"events": events,
@ -749,12 +755,10 @@ async def test_fire_deferred_stream_logging_fires_callback():
logging_obj._on_deferred_stream_complete = deferred
logging_obj._deferred_stream_complete_args = ("payload",)
with patch.object(GLOBAL_LOGGING_WORKER, "ensure_initialized_and_enqueue") as enqueue:
ProxyLogging._fire_deferred_stream_logging(request_data={"litellm_logging_obj": logging_obj})
enqueue.assert_called_once()
deferred_task = ProxyLogging._fire_deferred_stream_logging(request_data={"litellm_logging_obj": logging_obj})
assert deferred_task is not None
assert captured == {}
await enqueue.call_args.kwargs["async_coroutine"]
await deferred_task
snapshot = {
"arg": captured["arg"],
"callback_cleared": logging_obj._on_deferred_stream_complete is None,