mirror of
https://github.com/BerriAI/litellm.git
synced 2026-09-05 08:07:05 +00:00
321345d4c8
* fix(prometheus): bound per-request budget metric emission with a timeout (#31632) * fix(prometheus): bound per-request budget metric emission with a timeout Wrap the per-request budget-metric gather in asyncio.wait_for so a slow Redis or DB lookup cannot consume the whole LoggingWorker watchdog and get the success-logging event cancelled. On timeout the emission is skipped in isolation; budget gauges are still refreshed by the periodic cron. The timeout is configurable via PROMETHEUS_BUDGET_METRICS_PER_REQUEST_TIMEOUT and defaults to 5.0 seconds, falling back to the default on an invalid value instead of raising * fix(prometheus): reject non-finite and non-positive budget-metrics timeout env float() accepts 0, negatives, nan and inf, which bypass the fallback: a value <= 0 makes asyncio.wait_for time out immediately and skip every per-request emission, and inf reintroduces the unbounded wait the timeout was meant to bound. Validate the parsed value is finite and greater than zero before using it, otherwise fall back to the default * fix: report the blocked LLM response's real token usage (#31217) When a guardrail blocks a post-call response, the synthetic violation response reported hard-coded zero usage, discarding the token usage the upstream call had already consumed. Fix the root cause rather than re-counting tokens: - Add an optional `original_response` field to ModifyResponseException. - The unified guardrail's post-call success hook attaches the blocked LLM response to the exception. - The /v1/messages and OpenAI-format (/v1/chat/completions, /v1/completions) block handlers report `original_response.usage` directly. Pre-call blocks never invoked the LLM, so usage is zero. Mock-based tests cover the helper (returns original usage / zero), the success hook attaching original_response, and the endpoint reporting it end-to-end. Co-authored-by: Claude Opus 4.8 (1M context) <noreply@anthropic.com> * feat(guardrails): buffer + cleanly terminate streamed responses on block (#31389) Streaming moderation improvements for the unified guardrail post-call streaming iterator hook: - streaming_buffer_until_moderated: withhold all chunks until end-of-stream moderation passes, then release the original response (clean) or only the block message (blocked) -- the original content is never delivered on a block. Snapshot chunks with a shallow list() copy (end-of-stream builds a separate assembled response; chunks aren't mutated in place). - Clean Anthropic SSE on block: synthesize a well-formed termination sequence instead of a bare data: {"error": ...} blob that truncates the stream. Provider-specific synthesis lives in AnthropicMessagesHandler via build_block_sse_chunks (format-agnostic routing stays in the hook). - Mid-stream blocks continue the in-progress message (close open content block, append block message, terminate) rather than emitting a second message_start, which clients reject. Standalone envelope only when no chunks were sent (buffered path). - ModifyResponseException imported under TYPE_CHECKING + locally at runtime to avoid a module-level cyclic import. Adds regression tests for buffering (content withheld on block) and mid-stream continuation (single message_start). Co-authored-by: Claude Opus 4.8 (1M context) <noreply@anthropic.com> * fix: report real usage on streaming blocks, disable buffered mode for content-rewriting guardrails - _standalone_block_chunks and _block_continuation_chunks now read real token usage from ModifyResponseException.original_response instead of hardcoding zero, matching the non-streaming _blocked_response_usage path. Shared helper moved to guardrail_translation/utils.py. - streaming_buffer_until_moderated is now forced off when the guardrail has mask_response_content=True, since buffered replay releases the withheld original chunks verbatim -- unsafe for a guardrail that rewrites content (e.g. PII masking). - Fix inverted streaming-flag precedence comment. * style: ruff format after greploop fixes * fix: handle Anthropic streaming guardrail blocks * fix(responses): check terminal event type for streaming guardrail end-of-stream detection _check_streaming_has_ended assumed responses_so_far held ModelResponse objects with .choices, but for the Responses API the accumulated chunks are raw SSE event dicts, causing an AttributeError on every call * fix: preserve Anthropic blocked stream usage --------- Co-authored-by: FERNANDO IZAR <fizar@me.com> Co-authored-by: Joseph Barker <156112794+seph-barker@users.noreply.github.com> Co-authored-by: Claude Opus 4.8 (1M context) <noreply@anthropic.com> Co-authored-by: Cursor Agent <cursoragent@cursor.com>
84 lines
3 KiB
Python
84 lines
3 KiB
Python
"""
|
|
Token usage on synthetic guardrail-blocked responses for the OpenAI-format
|
|
proxy endpoints (/v1/chat/completions and /v1/completions).
|
|
|
|
A post-call block replaces the LLM response with the violation message, but the
|
|
upstream call already consumed tokens. `_blocked_response_usage` reports that
|
|
real usage (carried on `ModifyResponseException.original_response`) rather than
|
|
zero; a pre-call block never invoked the LLM, so usage is zero.
|
|
"""
|
|
|
|
import pytest
|
|
|
|
import litellm
|
|
from litellm.proxy.proxy_server import _blocked_response_usage
|
|
|
|
|
|
def test_uses_original_response_usage():
|
|
resp = litellm.ModelResponse()
|
|
resp.usage = litellm.Usage(prompt_tokens=42, completion_tokens=7, total_tokens=49)
|
|
|
|
usage = _blocked_response_usage(resp)
|
|
|
|
assert usage.prompt_tokens == 42
|
|
assert usage.completion_tokens == 7
|
|
assert usage.total_tokens == 49
|
|
|
|
|
|
def test_zero_usage_when_no_original_response():
|
|
usage = _blocked_response_usage(None)
|
|
|
|
assert usage.prompt_tokens == 0
|
|
assert usage.completion_tokens == 0
|
|
assert usage.total_tokens == 0
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_success_hook_attaches_original_response_on_block():
|
|
"""The unified guardrail's post-call success hook must attach the blocked
|
|
LLM response to ModifyResponseException so its real usage isn't discarded."""
|
|
from unittest.mock import AsyncMock, MagicMock, patch
|
|
|
|
import litellm.proxy.guardrails.guardrail_hooks.unified_guardrail.unified_guardrail as ug
|
|
from litellm.integrations.custom_guardrail import ModifyResponseException
|
|
from litellm.proxy._types import UserAPIKeyAuth
|
|
from litellm.types.utils import CallTypes
|
|
|
|
response = litellm.ModelResponse()
|
|
response.usage = litellm.Usage(prompt_tokens=15, completion_tokens=3, total_tokens=18)
|
|
|
|
guardrail = MagicMock()
|
|
guardrail.should_run_guardrail.return_value = True
|
|
guardrail.guardrail_name = "rubrik"
|
|
|
|
# The translation layer raises a block without pre-setting original_response.
|
|
translation = MagicMock()
|
|
translation.process_output_response = AsyncMock(
|
|
side_effect=ModifyResponseException(
|
|
message="blocked",
|
|
model="gpt-4o",
|
|
request_data={},
|
|
guardrail_name="rubrik",
|
|
)
|
|
)
|
|
|
|
unified = ug.UnifiedLLMGuardrails()
|
|
user_api_key_dict = UserAPIKeyAuth(api_key="test", request_route="/chat/completions")
|
|
data = {"guardrail_to_apply": guardrail, "model": "gpt-4o"}
|
|
|
|
# Inject our translation for the inferred call type (the module global is
|
|
# cached across tests, so patch it directly rather than the loader).
|
|
with patch.object(
|
|
ug,
|
|
"endpoint_guardrail_translation_mappings",
|
|
{
|
|
CallTypes.acompletion: lambda: translation,
|
|
CallTypes.completion: lambda: translation,
|
|
},
|
|
):
|
|
with pytest.raises(ModifyResponseException) as excinfo:
|
|
await unified.async_post_call_success_hook(
|
|
data=data, user_api_key_dict=user_api_key_dict, response=response
|
|
)
|
|
|
|
assert excinfo.value.original_response is response
|