From 710f6b2dcd10e0ab3b496d11596b76a92db100d1 Mon Sep 17 00:00:00 2001 From: Cursor Agent Date: Thu, 2 Jul 2026 16:32:50 +0000 Subject: [PATCH] fix: preserve Anthropic blocked stream usage --- .../base_llm/guardrail_translation/utils.py | 64 ++++++++++++++++++- .../test_anthropic_streaming_block.py | 55 ++++++++++++++++ 2 files changed, 118 insertions(+), 1 deletion(-) diff --git a/litellm/llms/base_llm/guardrail_translation/utils.py b/litellm/llms/base_llm/guardrail_translation/utils.py index 02dd7867d69..8a06dd4ea52 100644 --- a/litellm/llms/base_llm/guardrail_translation/utils.py +++ b/litellm/llms/base_llm/guardrail_translation/utils.py @@ -1,11 +1,69 @@ from __future__ import annotations +import json from typing import Any, List, Optional from litellm.types.llms.anthropic_messages.anthropic_response import AnthropicUsage from litellm.types.llms.openai import AllMessageValues +def _anthropic_stream_chunk_events(item: Any) -> list[dict]: + if isinstance(item, dict): + return [item] + if isinstance(item, bytes): + chunk = item.decode("utf-8", errors="replace") + elif isinstance(item, str): + chunk = item + else: + return [] + + events: list[dict] = [] + for block in chunk.split("\n\n"): + for line in block.splitlines(): + stripped = line.strip() + if not stripped.startswith("data:"): + continue + payload = stripped[len("data:") :].strip() + if not payload or payload == "[DONE]": + continue + try: + parsed = json.loads(payload) + except json.JSONDecodeError: + continue + if isinstance(parsed, dict): + events.append(parsed) + return events + + +def _usage_from_anthropic_stream_chunks(original_response: list[Any]) -> Optional[AnthropicUsage]: + input_tokens = 0 + output_tokens = 0 + found_usage = False + + for item in original_response: + for event in _anthropic_stream_chunk_events(item): + event_type = event.get("type") + if event_type == "message_start": + message = event.get("message") or {} + usage_obj = message.get("usage") or {} + elif event_type == "message_delta": + usage_obj = event.get("usage") or {} + else: + usage_obj = {} + if not isinstance(usage_obj, dict): + continue + if usage_obj.get("input_tokens") is not None: + input_tokens = int(usage_obj.get("input_tokens") or 0) + found_usage = True + if usage_obj.get("output_tokens") is not None: + output_tokens = int(usage_obj.get("output_tokens") or 0) + found_usage = True + + if not found_usage: + return None + return AnthropicUsage(input_tokens=input_tokens, output_tokens=output_tokens) + + def blocked_response_usage(original_response: Optional[Any]) -> AnthropicUsage: """ Token usage for a synthetic guardrail-blocked response. @@ -17,7 +75,11 @@ def blocked_response_usage(original_response: Optional[Any]) -> AnthropicUsage: so usage is zero. """ usage_obj: Any = None - if isinstance(original_response, dict): + if isinstance(original_response, list): + stream_usage = _usage_from_anthropic_stream_chunks(original_response) + if stream_usage is not None: + return stream_usage + elif isinstance(original_response, dict): usage_obj = original_response.get("usage") elif original_response is not None: usage_obj = getattr(original_response, "usage", None) diff --git a/tests/test_litellm/proxy/guardrails/guardrail_hooks/unified_guardrails/test_anthropic_streaming_block.py b/tests/test_litellm/proxy/guardrails/guardrail_hooks/unified_guardrails/test_anthropic_streaming_block.py index 48109a632eb..865b4164e5e 100644 --- a/tests/test_litellm/proxy/guardrails/guardrail_hooks/unified_guardrails/test_anthropic_streaming_block.py +++ b/tests/test_litellm/proxy/guardrails/guardrail_hooks/unified_guardrails/test_anthropic_streaming_block.py @@ -217,6 +217,61 @@ async def test_end_of_stream_only_block_does_not_append_after_message_stop(): assert message_delta_usages[-1] == 5 +def test_blocked_stream_reports_usage_from_original_chunks(): + from litellm.integrations.custom_guardrail import ModifyResponseException + from litellm.llms.anthropic.chat.guardrail_translation.handler import ( + AnthropicMessagesHandler, + ) + from litellm.llms.base_llm.guardrail_translation.utils import ( + blocked_response_usage, + ) + + original_chunks: List[Any] = [ + _sse_event( + "message_start", + { + "type": "message_start", + "message": { + "id": "msg_orig", + "type": "message", + "role": "assistant", + "model": "claude-3-5-sonnet", + "content": [], + "stop_reason": None, + "usage": {"input_tokens": 12, "output_tokens": 0}, + }, + }, + ), + {"type": "message_delta", "delta": {"stop_reason": "end_turn"}, "usage": {"output_tokens": 5}}, + ] + seen_chunks = [ + _sse_event( + "content_block_start", + {"type": "content_block_start", "index": 0, "content_block": {"type": "text", "text": ""}}, + ) + ] + exc = ModifyResponseException( + message=BLOCK_MESSAGE, + model="claude-3-5-sonnet", + request_data={}, + guardrail_name="g", + original_response=original_chunks, + ) + + usage = blocked_response_usage(original_chunks) + raw = b"".join( + AnthropicMessagesHandler().build_block_sse_chunks(exc, stream_started=True, responses_so_far=seen_chunks) + ).decode() + message_delta_usages = [ + payload.get("usage", {}).get("output_tokens") + for payload in _parse_sse_payloads(raw) + if payload.get("type") == "message_delta" + ] + + assert usage == {"input_tokens": 12, "output_tokens": 5} + assert message_delta_usages[-1] == 5 + + class TestContentBlockState: """`_content_block_state` must reflect the true open/last block index across the two chunk formats the stream can carry (multi-event bytes, parsed dict),