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 65

Merge aedc0e3aec into 82ca10cc2d

Browse source
This commit is contained in:
Ivana Dunisijevic 2026-10-17 08:01:46 -07:00 • committed by GitHub
commit ead46fd219
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
4 changed files with 125 additions and 53 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

14
litellm/proxy/common_request_processing.py
View file

@ -4055,14 +4055,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

@ -1159,6 +1159,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: # noqa: BLE001 # background callback failures must not escape
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]:
@ -4420,7 +4435,7 @@ class ProxyLogging:
record_served_output_texts(logging_obj.model_call_details, texts)
@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.
@ -4431,7 +4446,8 @@ class ProxyLogging:
"""
logging_obj: Final = request_data.get("litellm_logging_obj")
if logging_obj is None:
return
return None
_deferred_cb: Final[Callable[..., Coroutine[object, object, object]] | None] = getattr(
logging_obj, "_on_deferred_stream_complete", None
)
@ -4439,7 +4455,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
asyncio.create_task(_deferred_cb(*_args))
return _retain_deferred_stream_callback(_deferred_cb(*_args))
return None
_fire_deferred_stream_logging = fire_deferred_stream_logging

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

@ -36,6 +36,7 @@ from litellm.integrations.custom_logger import CustomLogger
from litellm.proxy._types import UserAPIKeyAuth
from litellm.proxy.common_request_processing import ProxyBaseLLMRequestProcessing
from litellm.proxy.utils import ProxyLogging
from litellm.litellm_core_utils.logging_worker import GLOBAL_LOGGING_WORKER
from litellm.types.guardrails import GuardrailEventHooks
# ---------------------------------------------------------------------------
@ -735,10 +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 asyncio.sleep(0)
await asyncio.sleep(0)
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 (
@ -947,10 +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 asyncio.sleep(0)
await asyncio.sleep(0)
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 (
@ -1319,10 +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 asyncio.sleep(0)
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
@ -1330,6 +1331,53 @@ 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_callback_failure_is_logged_without_escaping(self):
"""A failed background callback must not surface as an unhandled task error."""
callback_error = RuntimeError("deferred callback failed")
async def failing_callback():
raise callback_error
mock_logging_obj = MagicMock()
mock_logging_obj._on_deferred_stream_complete = failing_callback
mock_logging_obj._deferred_stream_complete_args = ()
with patch("litellm.proxy.utils.verbose_proxy_logger.exception") as log_exception:
deferred_task = ProxyLogging._fire_deferred_stream_logging({"litellm_logging_obj": mock_logging_obj})
assert deferred_task is not None
await deferred_task
log_exception.assert_called_once_with("Error in deferred stream callback: %s", callback_error)
@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
@ -1520,8 +1568,9 @@ class TestArmDeferredStreamDispatch:
assembled = object()
logging_obj._deferred_stream_complete_args = (assembled, False)
ProxyLogging._fire_deferred_stream_logging({"litellm_logging_obj": logging_obj})
await asyncio.sleep(0)
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
@ -1554,8 +1603,9 @@ class TestArmDeferredStreamDispatch:
assembled = object()
logging_obj._deferred_stream_complete_args = (assembled, False)
ProxyLogging._fire_deferred_stream_logging({"litellm_logging_obj": logging_obj})
await asyncio.sleep(0)
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
@ -1563,8 +1613,6 @@ class TestArmDeferredStreamDispatch:
@pytest.mark.asyncio
async def test_native_stream_closure_enqueues_single_coroutine(self):
from litellm.litellm_core_utils.logging_worker import GLOBAL_LOGGING_WORKER
logging_obj, recorded = self._dispatch_recording_logging_obj()
async def _agen():
@ -1583,11 +1631,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})
await asyncio.sleep(0)
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()
@ -1600,8 +1651,6 @@ class TestArmDeferredStreamDispatch:
so stores (assembled_response, cache_hit). The closure armed for a raw
generator must accept that shape too, or _fire_deferred_stream_logging
raises TypeError and the request loses its spend log and callbacks."""
from litellm.litellm_core_utils.logging_worker import GLOBAL_LOGGING_WORKER
logging_obj, recorded = self._dispatch_recording_logging_obj()
async def _agen():
@ -1616,11 +1665,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})
await asyncio.sleep(0)
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
@ -1630,8 +1682,6 @@ class TestArmDeferredStreamDispatch:
@pytest.mark.asyncio
@pytest.mark.parametrize("stored_args", [(object(),), (object(), object(), object())])
async def test_raw_generator_stream_with_unknown_arg_shape_logs_and_drops(self, stored_args, caplog):
from litellm.litellm_core_utils.logging_worker import GLOBAL_LOGGING_WORKER
logging_obj, recorded = self._dispatch_recording_logging_obj()
async def _agen():
@ -1651,8 +1701,9 @@ class TestArmDeferredStreamDispatch:
) as mock_enqueue,
caplog.at_level(logging.ERROR, logger="LiteLLM Proxy"),
):
ProxyLogging._fire_deferred_stream_logging({"litellm_logging_obj": logging_obj})
await asyncio.sleep(0)
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 == {}

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

@ -10,7 +10,7 @@ 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
@ -179,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"
@ -563,8 +570,7 @@ async def test_native_messages_stream_logging_fires_after_guardrail_end_of_strea
request_data=request_data,
):
pass
await asyncio.sleep(0)
await asyncio.sleep(0)
await _wait_until(lambda: len(events) == 2)
assert events == ["scan_appended", ("logging_dispatched", "post_call_entry_visible", True)]
@ -599,8 +605,7 @@ async def test_native_messages_stream_logging_fires_when_guardrail_blocks_after_
request_data=request_data,
):
pass
await asyncio.sleep(0)
await asyncio.sleep(0)
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
@ -724,8 +729,7 @@ async def test_chat_stream_generic_callback_error_after_stream_end_still_flushes
request_data=request_data,
):
pass
await asyncio.sleep(0)
await asyncio.sleep(0)
await _wait_until(lambda: len(events) == 1)
snapshot = {
"events": events,
@ -751,8 +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",)
ProxyLogging._fire_deferred_stream_logging(request_data={"litellm_logging_obj": logging_obj})
await asyncio.sleep(0)
deferred_task = ProxyLogging._fire_deferred_stream_logging(request_data={"litellm_logging_obj": logging_obj})
assert deferred_task is not None
assert captured == {}
await deferred_task
snapshot = {
"arg": captured["arg"],
"callback_cleared": logging_obj._on_deferred_stream_complete is None,