fix: preserve Anthropic blocked stream usage

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)

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),