scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-09-21 00:21:49 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 73

fix(proxy): only discard parked stream logging for errors the failure path logs

Browse source 
Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>
This commit is contained in:
yassin 2026-09-15 20:10:18 +00:00
parent ec799686a4
commit 7095373dd5
4 changed files with 112 additions and 61 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

2
litellm/litellm_core_utils/litellm_logging.py
View file

@ -639,6 +639,8 @@ class Logging(LiteLLMLoggingBaseClass):
self._defer_async_logging: bool = False
self._enqueue_deferred_logging: Callable[[], None] | None = None
self._on_detached_stream_failure: Callable[[Exception], Awaitable[None]] | None = None
self._on_deferred_stream_complete: Callable[..., Awaitable[None]] | None = None
self._deferred_stream_complete_args: tuple[object, ...] | None = None
def set_response_timing_metrics(self, timing_metrics: Mapping[str, float]) -> None:
"""Keep ``_response_ms`` / ``litellm_overhead_time_ms`` for a result that has no ``_hidden_params``."""

51
litellm/proxy/utils.py
View file

@ -905,6 +905,9 @@ def _call_type_for_route(route: str | None) -> str | None:
return call_types[0].value if len(operations) == 1 else None
_PROXY_ONLY_LLM_API_ERRORS: Final = (HTTPException, ProxyException, GuardrailRaisedException)
def _failure_fields_to_lift(request_data: Mapping[str, object]) -> Mapping[str, object]:
"""Failure-path callbacks run after ``litellm_logging_obj`` is popped from
request_data (it is not serialisable), so the caller merges these fields
@ -3010,9 +3013,7 @@ class ProxyLogging:
if not (RouteChecks.is_llm_api_route(route) or RouteChecks.is_info_route(route)):
return False
return isinstance(original_exception, (HTTPException, ProxyException, GuardrailRaisedException)) or (
error_type == ProxyErrorTypes.auth_error
)
return isinstance(original_exception, _PROXY_ONLY_LLM_API_ERRORS) or (error_type == ProxyErrorTypes.auth_error)
async def _handle_logging_proxy_only_error(
self,
@ -3568,8 +3569,9 @@ class ProxyLogging:
yield chunk
except (GeneratorExit, asyncio.CancelledError):
raise
except Exception:
ProxyLogging._discard_deferred_stream_logging_for_failure(request_data)
except Exception as e:
if not ProxyLogging._discard_deferred_stream_logging_for_failure(request_data, e):
ProxyLogging._fire_deferred_stream_logging(request_data)
raise
ProxyLogging._fire_deferred_stream_logging(request_data)
return
@ -3643,8 +3645,9 @@ class ProxyLogging:
yield chunk
except (GeneratorExit, asyncio.CancelledError):
raise
except Exception:
ProxyLogging._discard_deferred_stream_logging_for_failure(request_data)
except Exception as e:
if not ProxyLogging._discard_deferred_stream_logging_for_failure(request_data, e):
ProxyLogging._fire_deferred_stream_logging(request_data)
raise
# Fire deferred logging AFTER all guardrail end-of-stream blocks
@ -3741,35 +3744,29 @@ class ProxyLogging:
asyncio.create_task(_deferred_cb(*_args))
@staticmethod
def _discard_deferred_stream_logging_for_failure(request_data: dict) -> None:
"""Drop the parked success dispatch when the stream ends in an exception.
The CSW shape parks ``(assembled ModelResponse, cache_hit)``: its usage is
carried onto the logging object so the failure row bills what the stream
consumed. The native /v1/messages and responses shapes park a logging
coroutine with no recoverable usage, so they keep firing as before.
def _discard_deferred_stream_logging_for_failure(request_data: Mapping[str, object], error: Exception) -> bool:
"""Drop the parked success dispatch when the stream ends in an error the proxy logs
as a failure (``_PROXY_ONLY_LLM_API_ERRORS``, e.g. a post_call guardrail block) and
the CSW parked an assembled ``ModelResponse``, carrying its usage onto the logging
object so the failure row bills what the stream consumed. Returns False, leaving the
parked dispatch for the caller to flush, for any other error and for the native
/v1/messages and responses shapes that park a logging coroutine with no usage.
"""
logging_obj: Final = request_data.get("litellm_logging_obj")
if logging_obj is None:
return
_deferred_cb: Final[Callable[..., Coroutine[object, object, object]] | None] = getattr(
logging_obj, "_on_deferred_stream_complete", None
)
if not isinstance(logging_obj, Logging):
return False
_args: Final[tuple[object, ...] | None] = getattr(logging_obj, "_deferred_stream_complete_args", None)
if _deferred_cb is None or _args is None:
return
assembled: Final = _args[0]
if not isinstance(assembled, ModelResponse):
ProxyLogging._fire_deferred_stream_logging(request_data)
return
assembled: Final = _args[0] if _args else None
if not isinstance(error, _PROXY_ONLY_LLM_API_ERRORS) or not isinstance(assembled, ModelResponse):
return False
logging_obj._on_deferred_stream_complete = None
logging_obj._deferred_stream_complete_args = None
usage: Final[Usage | None] = getattr(assembled, "usage", None)
if isinstance(usage, Usage):
logging_obj.record_partial_usage_for_failure(
usage,
logging_obj._response_cost_calculator(result=assembled) or 0.0,
usage, logging_obj._response_cost_calculator(result=assembled) or 0.0
)
return True
async def _arelease_max_parallel_requests_on_disconnect(
self,

14
tests/test_litellm/proxy/utils/proxy_logging/test_post_call_failure_hook.py
View file

@ -4,6 +4,7 @@ and ``_handle_logging_proxy_only_error``."""
from __future__ import annotations
import asyncio
from datetime import datetime
from typing import Any
from unittest.mock import AsyncMock, MagicMock
@ -334,15 +335,16 @@ async def test_post_call_failure_hook_guardrail_block_fires_failure_callback(
object's ``async_failure_handler`` so custom loggers see a ``failure``
status - without this, guardrail blocks produce only
``post_call_failure_hook`` and no failure logging event."""
from datetime import datetime
from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj
recorded: dict[str, Any] = {}
recorded: list[object] = []
class _StatusRecorder(CustomLogger):
async def async_log_failure_event(self, kwargs, response_obj, start_time, end_time):
recorded["status"] = (kwargs.get("standard_logging_object") or {}).get("status")
async def async_log_failure_event(
self, kwargs: dict[str, object], response_obj: object, start_time: datetime, end_time: datetime
) -> None:
standard_logging_object = kwargs.get("standard_logging_object")
recorded.append(standard_logging_object.get("status") if isinstance(standard_logging_object, dict) else None)
monkeypatch.setattr(litellm, "_async_failure_callback", [_StatusRecorder()])
logging_obj = LiteLLMLoggingObj(
@ -369,4 +371,4 @@ async def test_post_call_failure_hook_guardrail_block_fires_failure_callback(
)
await asyncio.sleep(0)
await asyncio.sleep(0)
assert recorded["status"] == "failure"
assert recorded == ["failure"]

106
tests/test_litellm/proxy/utils/proxy_logging/test_streaming_hooks.py
View file

@ -10,6 +10,7 @@ Covers ``_wrap_streaming_iterator_with_enrichment``,
from __future__ import annotations
import asyncio
from collections.abc import AsyncGenerator, AsyncIterator
from datetime import datetime
from typing import Any, Dict, List
from unittest.mock import AsyncMock, MagicMock
@ -18,12 +19,15 @@ import pytest
from fastapi import HTTPException
import litellm
from litellm.exceptions import GuardrailRaisedException
from litellm.integrations.custom_logger import CustomLogger
from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj
from litellm.llms.anthropic.experimental_pass_through.messages.streaming_iterator import (
BaseAnthropicMessagesStreamingIterator,
)
from litellm.proxy._types import UserAPIKeyAuth
from litellm.proxy.utils import ProxyLogging
from litellm.types.utils import Usage
@pytest.fixture(autouse=True)
@ -479,37 +483,25 @@ async def test_native_messages_stream_logging_fires_when_guardrail_blocks_after_
assert logging_obj._deferred_stream_complete_args is None
@pytest.mark.asyncio
async def test_chat_stream_guardrail_block_after_stream_end_logs_failure_not_success(
proxy_logging, make_user_api_key_auth, monkeypatch
):
"""
On /chat/completions streams the CSW shape parks
``(assembled ModelResponse, cache_hit)`` as the deferred args. A guardrail
that raises ``GuardrailRaisedException`` at end of stream must NOT dispatch
that deferred success logging - the request is logged via the failure path
instead, with the consumed usage carried over so the failure row bills
correctly.
"""
from litellm.exceptions import GuardrailRaisedException
from litellm.types.utils import Usage
events: List[Any] = []
request_data: Dict[str, Any] = {"metadata": {}}
def _armed_chat_stream(
test_name: str, request_data: dict[str, object], events: list[str]
) -> tuple[LiteLLMLoggingObj, AsyncIterator[dict[str, object]]]:
"""A /chat/completions stream whose CSW shape parks ``(assembled ModelResponse, cache_hit)``
at upstream exhaustion, with the deferred dispatch recording into ``events``."""
logging_obj = LiteLLMLoggingObj(
model="gpt-4o-mini",
messages=[{"role": "user", "content": "hi"}],
stream=True,
call_type="acompletion",
start_time=datetime.now(),
litellm_call_id="test_chat_stream_guardrail_block",
function_id="test_chat_stream_guardrail_block",
litellm_call_id=test_name,
function_id=test_name,
)
logging_obj.optional_params = {}
logging_obj.litellm_params = {}
logging_obj.standard_built_in_tools_params = None
async def _dispatch_deferred_logging(*args):
async def _dispatch_deferred_logging(*args: object) -> None:
events.append("success_dispatched")
logging_obj._on_deferred_stream_complete = _dispatch_deferred_logging
@ -521,24 +513,48 @@ async def test_chat_stream_guardrail_block_after_stream_end_logs_failure_not_suc
usage=Usage(prompt_tokens=3, completion_tokens=5, total_tokens=8),
)
async def _upstream():
async def _upstream() -> AsyncIterator[dict[str, object]]:
yield {"id": "c1", "choices": [{"index": 0, "delta": {"content": "BAN"}}]}
yield {"id": "c1", "choices": [{"index": 0, "delta": {"content": "ANA"}}]}
logging_obj._deferred_stream_complete_args = (assembled, False)
class _BlockingGuardrail(CustomLogger):
async def async_post_call_streaming_iterator_hook(self, user_api_key_dict, response, request_data):
return logging_obj, _upstream()
def _raising_at_end_of_stream(error: Exception) -> CustomLogger:
class _EndOfStreamRaiser(CustomLogger):
async def async_post_call_streaming_iterator_hook(
self, user_api_key_dict: UserAPIKeyAuth, response: AsyncIterator[object], request_data: dict[str, object]
) -> AsyncGenerator[object, None]:
async for chunk in response:
yield chunk
raise GuardrailRaisedException(
guardrail_name="g", message="blocked", blocked_content=True
)
raise error
monkeypatch.setattr(litellm, "callbacks", [_BlockingGuardrail()])
return _EndOfStreamRaiser()
@pytest.mark.asyncio
async def test_chat_stream_guardrail_block_after_stream_end_logs_failure_not_success(
proxy_logging, make_user_api_key_auth, monkeypatch
):
"""
A guardrail that raises ``GuardrailRaisedException`` at end of a
/chat/completions stream must NOT dispatch the parked success logging:
the request is logged via the failure path instead, with the consumed
usage carried over so the failure row bills correctly.
"""
events: list[str] = []
request_data: dict[str, object] = {"metadata": {}}
logging_obj, upstream = _armed_chat_stream("test_chat_stream_guardrail_block", request_data, events)
monkeypatch.setattr(
litellm,
"callbacks",
[_raising_at_end_of_stream(GuardrailRaisedException(guardrail_name="g", message="blocked"))],
)
with pytest.raises(GuardrailRaisedException):
async for _ in proxy_logging.async_post_call_streaming_iterator_hook(
response=_upstream(),
response=upstream,
user_api_key_dict=make_user_api_key_auth(),
request_data=request_data,
):
@ -562,6 +578,40 @@ async def test_chat_stream_guardrail_block_after_stream_end_logs_failure_not_suc
}
@pytest.mark.asyncio
async def test_chat_stream_generic_callback_error_after_stream_end_still_flushes_success_logging(
proxy_logging, make_user_api_key_auth, monkeypatch
):
"""
``post_call_failure_hook`` only routes proxy-level errors (HTTPException,
ProxyException, GuardrailRaisedException) through failure logging. A
callback that dies with any other exception after the stream completed
must keep flushing the parked success dispatch, or the request ends with
no terminal log at all.
"""
events: list[str] = []
request_data: dict[str, object] = {"metadata": {}}
logging_obj, upstream = _armed_chat_stream("test_chat_stream_generic_callback_error", request_data, events)
monkeypatch.setattr(litellm, "callbacks", [_raising_at_end_of_stream(RuntimeError("callback crashed"))])
with pytest.raises(RuntimeError):
async for _ in proxy_logging.async_post_call_streaming_iterator_hook(
response=upstream,
user_api_key_dict=make_user_api_key_auth(),
request_data=request_data,
):
pass
await asyncio.sleep(0)
await asyncio.sleep(0)
snapshot = {
"events": events,
"args_cleared": logging_obj._deferred_stream_complete_args is None,
"failure_usage_recorded": "combined_usage_object" in logging_obj.model_call_details,
}
assert snapshot == {"events": ["success_dispatched"], "args_cleared": True, "failure_usage_recorded": False}
# ---------------------------------------------------------------------------
# _fire_deferred_stream_logging
# ---------------------------------------------------------------------------