refactor: centralize Messages usage cost stamping and stream wrapping

Co-Authored-By: Ishaan Jaffer <155045088+ishaan-berri@users.noreply.github.com>

litellm/litellm_core_utils/llm_response_utils/response_metadata.py

@@ -10,7 +10,7 @@ from litellm.constants import LITELLM_DETAILED_TIMING
 from litellm.litellm_core_utils.core_helpers import get_litellm_metadata_from_kwargs, process_response_headers
 from litellm.litellm_core_utils.llm_response_utils.get_api_base import get_api_base
 from litellm.litellm_core_utils.logging_utils import LiteLLMLoggingObject
-from litellm.litellm_core_utils.usage_cost import stamp_usage_cost, validated_usage_cost
+from litellm.litellm_core_utils.usage_cost import MessagesUsageCostStamper, stamp_usage_cost, validated_usage_cost
 from litellm.types.utils import (
     EmbeddingResponse,
     HiddenParams,
@@ -18,26 +18,7 @@ from litellm.types.utils import (
     TranscriptionResponse,
 )
 
-_MESSAGES_BODY: Final = TypeAdapter(Mapping[str, object])
-
-
-def _messages_body(result: object) -> Mapping[str, object] | None:
-    return _MESSAGES_BODY.validate_python(result) if isinstance(result, dict) else None
-
-
-def _stamp_messages_cost(result: object, logging_obj: LiteLLMLoggingObject, cache_hit: bool | None) -> None:
-    try:
-        body: Final = _messages_body(result)
-        if body is None or body.get("type") != "message" or not isinstance(body.get("usage"), Mapping):
-            return
-        priced: Final = logging_obj.normalize_endpoint_logging_result(body)
-        cost: Final = validated_usage_cost(logging_obj.response_cost_calculator(result=priced, cache_hit=cache_hit))
-        stamp_usage_cost(priced, cost)
-        stamp_usage_cost(result, cost)
-        if isinstance(priced, ModelResponse):
-            logging_obj.record_bridged_model_response(priced)
-    except Exception:  # noqa: BLE001  # optional billing metadata must not fail a completed response
-        return
+_HIDDEN_PARAMS: Final = TypeAdapter(Mapping[str, object])
 
 
 def _timing_window_start(
@@ -292,7 +273,7 @@ class ResponseMetadata:
         """Apply metadata to the response object"""
         if hasattr(self.result, "_hidden_params"):
             self.result._hidden_params = self._hidden_params
-            hidden: Final = _MESSAGES_BODY.validate_python(getattr(self, "_hidden_params", {}))
+            hidden: Final = _HIDDEN_PARAMS.validate_python(getattr(self, "_hidden_params", {}))
             stamp_usage_cost(self.result, validated_usage_cost(hidden.get("response_cost")))
 
 
@@ -322,7 +303,9 @@ def update_response_metadata(
         logging_obj.set_response_timing_metrics(
             response_timing_metrics(start_time, end_time, logging_obj, include_overhead)
         )
-        _stamp_messages_cost(result, logging_obj, cache_hit)
+        stamper: Final = MessagesUsageCostStamper.for_response(result)
+        if stamper is not None:
+            stamper.stamp_calculated(logging_obj, cache_hit=cache_hit)
         return
 
     metadata: Final = ResponseMetadata(result)

litellm/litellm_core_utils/usage_cost.py

@@ -19,7 +19,7 @@ if TYPE_CHECKING:
     from litellm.litellm_core_utils.litellm_logging import Logging
 
 _CHAT_USAGE: Final[TypeAdapter[Usage | None]] = TypeAdapter(Usage | None)
-_WIRE_BODY: Final = TypeAdapter(Mapping[str, object])
+_USAGE_MAPPING: Final = TypeAdapter(Mapping[str, object])
 _WIRE_TARGET: Final[TypeAdapter[dict[str, object]]] = TypeAdapter(InstanceOf[dict[str, object]])
 _COST: Final[TypeAdapter[float | None]] = TypeAdapter(float | None, config=ConfigDict(strict=True))
 
@@ -91,14 +91,41 @@ class ResponsesUsageCostStamper(UsageCostStamper):
 
 
 class MessagesUsageCostStamper(UsageCostStamper):
-    def __init__(self, response: object) -> None:
-        self.response = _WIRE_TARGET.validate_python(response)
+    def __init__(self, response: dict[str, object]) -> None:  # mutable-ok: wire body is stamped in place [LIT001]
+        self.response = response
+
+    @classmethod
+    def for_response(cls, response: object) -> MessagesUsageCostStamper | None:
+        if not isinstance(response, dict):
+            return None
+        return cls(_WIRE_TARGET.validate_python(response))
 
     def stamp(self, cost: float) -> None:
-        body: Final = _WIRE_BODY.validate_python(self.response)
-        usage: Final = body.get("usage")
+        usage: Final = self.response.get("usage")
         if isinstance(usage, Mapping):
-            self.response["usage"] = {**usage, "cost": cost}
+            self.response["usage"] = {**_USAGE_MAPPING.validate_python(usage), "cost": cost}
+
+    def stamp_calculated(
+        self,
+        logging_obj: Logging,
+        cache_hit: bool | None = None,
+        *,
+        billing_body: Mapping[str, object] | None = None,
+    ) -> bool:
+        try:
+            body: Final = self.response if billing_body is None else billing_body
+            if body.get("type") != "message" or not isinstance(body.get("usage"), Mapping):
+                return False
+            priced: Final = logging_obj.normalize_endpoint_logging_result(body)
+            cost: Final = validated_usage_cost(logging_obj.response_cost_calculator(result=priced, cache_hit=cache_hit))
+            stamp_usage_cost(priced, cost)
+            if cost is not None:
+                self.stamp(cost)
+            if isinstance(priced, ModelResponse):
+                logging_obj.record_bridged_model_response(priced)
+        except Exception:  # noqa: BLE001  # optional billing metadata must not fail a completed response or stream
+            return False
+        return cost is not None
 
 
 def stamp_usage_cost(response: object, cost: float | None) -> None:
@@ -110,11 +137,9 @@ def stamp_usage_cost(response: object, cost: float | None) -> None:
     if isinstance(response, ResponsesAPIResponse):
         ResponsesUsageCostStamper(response).stamp(float(cost))
         return
-    if not isinstance(response, dict):
-        return
-    wire_body: Final = _WIRE_BODY.validate_python(response)
-    if wire_body.get("type") in ("message", "message_delta"):
-        MessagesUsageCostStamper(_WIRE_TARGET.validate_python(response)).stamp(float(cost))
+    stamper: Final = MessagesUsageCostStamper.for_response(response)
+    if stamper is not None and stamper.response.get("type") in ("message", "message_delta"):
+        stamper.stamp(float(cost))
 
 
 def stamp_responses_usage_cost(response_obj: ResponsesAPIResponse | None, logging_obj: Logging | None) -> None:

litellm/llms/anthropic/messages/native/handler.py

@@ -9,8 +9,8 @@ import asyncio
 import contextvars
 from collections.abc import AsyncIterator, Awaitable, Callable, Coroutine, Iterator, Mapping, Sequence
 from dataclasses import dataclass
-from functools import partial
-from typing import Final, cast
+from functools import partial, wraps
+from typing import Final, ParamSpec, TypeAlias, cast
 
 from pydantic import TypeAdapter
 
@@ -65,8 +65,29 @@ __all__ = ("anthropic_messages", "anthropic_messages_handler")
 # Providers that are routed directly to the OpenAI Responses API instead of
 # going through chat/completions.
 _RESPONSES_API_PROVIDERS: Final = frozenset({"openai"})
-_SAFEGUARD_KWARGS: Final = TypeAdapter(Mapping[str, object])
+_REQUEST_KWARGS: Final = TypeAdapter(Mapping[str, object])
 _SAFEGUARD_MESSAGES: Final = TypeAdapter(Sequence[Mapping[str, object]])
+_MessagesParams: Final = ParamSpec("_MessagesParams")
+_MessagesResult: TypeAlias = AnthropicMessagesResponse | Iterator[bytes] | AsyncIterator[object]
+
+
+def _with_usage_cost(
+    function: Callable[_MessagesParams, Awaitable[_MessagesResult]],
+) -> Callable[_MessagesParams, Awaitable[_MessagesResult]]:
+    @wraps(function)
+    async def wrapped(
+        *args: _MessagesParams.args,
+        **kwargs: _MessagesParams.kwargs,  # kwargs-ok: preserve the Messages entry point's public signature
+    ) -> _MessagesResult:
+        response: Final = await function(*args, **kwargs)
+        logging_obj: Final = _REQUEST_KWARGS.validate_python(kwargs).get("litellm_logging_obj")
+        if isinstance(response, (Iterator, AsyncIterator)) and isinstance(logging_obj, LiteLLMLoggingObj):
+            from .streaming_iterator import UsageCostMessagesStream
+
+            return UsageCostMessagesStream(response, logging_obj)
+        return response
+
+    return wrapped
 
 
 def _bridges_to_responses_api(model: str, custom_llm_provider: str, deployment_model_info: object) -> bool:
@@ -426,7 +447,7 @@ async def _try_websearch_short_circuit(
             anthropic_response = await with_safeguard_results(
                 cast(AnthropicMessagesResponse, response),
                 _short_circuit_safeguards_evaluator(
-                    _SAFEGUARD_KWARGS.validate_python(kwargs) if kwargs is not None else None,
+                    _REQUEST_KWARGS.validate_python(kwargs) if kwargs is not None else None,
                     _SAFEGUARD_MESSAGES.validate_python(messages),
                 ),
             )
@@ -450,6 +471,7 @@ def _short_circuit_safeguards_evaluator(
 
 
 @client
+@_with_usage_cost
 async def anthropic_messages(
     max_tokens: int,
     messages: list[dict],
@@ -845,7 +867,7 @@ def anthropic_messages_handler(
         else api_base
     )
     safeguard_kwargs: Final = (
-        _SAFEGUARD_KWARGS.validate_python(kwargs) if is_async is True and "safeguards" in kwargs else None
+        _REQUEST_KWARGS.validate_python(kwargs) if is_async is True and "safeguards" in kwargs else None
     )
     native_response: Final = base_llm_http_handler.anthropic_messages_handler(
         model=model,

litellm/llms/anthropic/messages/native/streaming_iterator.py

@@ -7,7 +7,7 @@ from datetime import datetime
 from typing import TYPE_CHECKING, Any, Final, Protocol, runtime_checkable
 
 import httpx
-from pydantic import ConfigDict, TypeAdapter, ValidationError
+from pydantic import TypeAdapter, ValidationError
 from typing_extensions import TypedDict
 
 from litellm.constants import (
@@ -17,7 +17,7 @@ from litellm.constants import (
 from litellm.litellm_core_utils.core_helpers import process_response_headers
 from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj
 from litellm.litellm_core_utils.logging_worker import GLOBAL_LOGGING_WORKER
-from litellm.litellm_core_utils.usage_cost import stamp_usage_cost, validated_usage_cost
+from litellm.litellm_core_utils.usage_cost import MessagesUsageCostStamper
 from litellm.llms.anthropic.common_utils import ANTHROPIC_ERROR_STATUS_CODE_MAP
 from litellm.llms.anthropic.messages.native.utils import (
     INCOMPLETE_STREAM_ERROR_MESSAGE,
@@ -31,7 +31,7 @@ from litellm.proxy.pass_through_endpoints.success_handler import (
 )
 from litellm.types.llms.anthropic_messages.anthropic_response import AnthropicMessagesResponse
 from litellm.types.passthrough_endpoints.pass_through_endpoints import EndpointType
-from litellm.types.utils import GenericStreamingChunk, ModelResponse, ModelResponseStream
+from litellm.types.utils import GenericStreamingChunk, ModelResponseStream
 
 if TYPE_CHECKING:
     from litellm.llms.anthropic.messages.safeguards import SafeguardsEvaluator
@@ -453,12 +453,6 @@ def _usage_cost_sse_bytes(chunk: object) -> bytes:
     return anthropic_sse_frame(_SSE_EVENT.validate_python(chunk))
 
 
-_MESSAGES_STREAM: Final[TypeAdapter[Iterator[object] | AsyncIterator[object]]] = TypeAdapter(
-    Iterator[object] | AsyncIterator[object], config=ConfigDict(arbitrary_types_allowed=True)
-)
-_MESSAGES_LOGGING: Final = TypeAdapter(LiteLLMLoggingObj, config=ConfigDict(arbitrary_types_allowed=True))
-
-
 class UsageCostMessagesStream:
     def __init__(self, stream: Iterator[object] | AsyncIterator[object], logging_obj: LiteLLMLoggingObj) -> None:
         self.stream = stream
@@ -468,10 +462,6 @@ class UsageCostMessagesStream:
         self.frames = self._iterate()
         self.sync_frames = self._iterate_sync()
 
-    @classmethod
-    def for_request(cls, stream: object, logging_obj: object) -> "UsageCostMessagesStream":
-        return cls(_MESSAGES_STREAM.validate_python(stream), _MESSAGES_LOGGING.validate_python(logging_obj))
-
     def __iter__(self) -> "UsageCostMessagesStream":
         return self
 
@@ -485,7 +475,7 @@ class UsageCostMessagesStream:
             close()
 
     def __getattr__(self, name: str) -> object:
-        return _DECODED_SSE_DATA.validate_python(getattr(self.stream, name))
+        return getattr(self.stream, name)
 
     def __aiter__(self) -> "UsageCostMessagesStream":
         return self
@@ -520,17 +510,9 @@ class UsageCostMessagesStream:
                 **_SSE_EVENT.validate_python(usage),
             },
         }
-        try:
-            priced: Final = self.logging_obj.normalize_endpoint_logging_result(body)
-            cost: Final = validated_usage_cost(self.logging_obj.response_cost_calculator(result=priced))
-        except Exception:  # noqa: BLE001  # optional billing metadata must not interrupt a completed stream
-            return frame
-        stamp_usage_cost(priced, cost)
-        if isinstance(priced, ModelResponse):
-            self.logging_obj.record_bridged_model_response(priced)
         stamped: Final = dict(event)
-        stamp_usage_cost(stamped, cost)
-        return anthropic_sse_frame(stamped)
+        billed: Final = MessagesUsageCostStamper(stamped).stamp_calculated(self.logging_obj, billing_body=body)
+        return anthropic_sse_frame(stamped) if billed else frame
 
     def _frames(self, chunk: object) -> Iterator[bytes]:
         self.pending += _usage_cost_sse_bytes(chunk)

litellm/utils.py

@@ -2527,15 +2527,6 @@ def client(original_function):
 
             streaming_requested: Final = _returns_requested_stream(kwargs, call_type, result)
             if streaming_requested or _is_converted_stream_result(result):
-                if call_type == CallTypes.anthropic_messages.value:
-                    from litellm.llms.anthropic.messages.native.streaming_iterator import UsageCostMessagesStream
-
-                    result = (  # rebind-ok: decorate the returned Messages stream without changing its public contract
-                        UsageCostMessagesStream.for_request(
-                            TypeAdapter(object).validate_python(result),
-                            TypeAdapter(object).validate_python(logging_obj),
-                        )
-                    )
                 logging_obj.stream = True
                 logging_obj.model_call_details["stream"] = True
                 if not streaming_requested:

tests/unit/llms/anthropic/messages/native/test_streaming_iterator.py

@@ -8,7 +8,6 @@ from unittest.mock import patch
 import httpx
 import pytest
 
-
 from litellm.integrations.custom_logger import CustomLogger
 from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj
 from litellm.llms.anthropic.messages.native import streaming_iterator as streaming_iterator_module
@@ -18,6 +17,7 @@ from litellm.llms.anthropic.messages.native.streaming_iterator import (
     AnthropicMessagesStreamingResponse,
     BaseAnthropicMessagesStreamingIterator,
     SafeguardedAnthropicSSEStream,
+    UsageCostMessagesStream,
     _incomplete_stream_error_sse_event,
     anthropic_messages_response_as_sse_events,
     is_anthropic_content_delta_chunk,
@@ -32,6 +32,45 @@ from litellm.llms.openai_like.messages.transformation import OpenAILikeAnthropic
 from litellm.types.utils import Choices, Message, ModelResponse
 
 
+def _usage_cost_logging_obj() -> LiteLLMLoggingObj:
+    logging_obj: Final = LiteLLMLoggingObj(
+        model="claude-sonnet-5-5",
+        messages=[],
+        stream=True,
+        call_type="anthropic_messages",
+        start_time=datetime(2026, 1, 1),
+        litellm_call_id="stream-cost",
+        function_id="stream-cost",
+    )
+    logging_obj.update_environment_variables(
+        model="claude-sonnet-5-5",
+        litellm_params={},
+        optional_params={},
+        custom_llm_provider="anthropic",
+    )
+    return logging_obj
+
+
+@pytest.mark.parametrize("failure", ["normalize_endpoint_logging_result", "response_cost_calculator"])
+def test_usage_cost_preserves_completed_stream_frames_when_billing_fails(failure: str) -> None:
+    frames: Final = (
+        b'data: {"type":"message_start","message":{"id":"msg_cost","type":"message",'
+        b'"model":"claude-sonnet-5-5","role":"assistant","content":[],"stop_reason":null,'
+        b'"stop_sequence":null,"usage":{"input_tokens":3}}}\n\n',
+        b'data: {"type":"message_delta","delta":{"stop_reason":"end_turn"},"usage":{"output_tokens":4}}\n\n',
+        b'data: {"type":"message_stop"}\n\n',
+    )
+    logging_obj: Final = _usage_cost_logging_obj()
+    with patch.object(logging_obj, failure, side_effect=RuntimeError("billing unavailable")) as unavailable:
+        observed: Final = tuple(UsageCostMessagesStream(iter(frames), logging_obj))
+    assert observed == frames
+    assert unavailable.call_count == 1
+
+
+def test_usage_cost_does_not_add_frames_to_an_empty_messages_stream() -> None:
+    assert tuple(UsageCostMessagesStream(iter(()), _usage_cost_logging_obj())) == ()
+
+
 class _RecordingLoggingIterator(BaseAnthropicMessagesStreamingIterator):
     def __init__(self, litellm_logging_obj: LiteLLMLoggingObj, request_body: dict):
         super().__init__(litellm_logging_obj=litellm_logging_obj, request_body=request_body)
@@ -434,7 +473,7 @@ def _output_tokens_from_logged_chunks(chunks: list[bytes]) -> int | None:
         for line in raw.decode().splitlines():
             if not line.startswith("data:"):
                 continue
-            data = json.loads(line[len("data:"):].strip())
+            data = json.loads(line[len("data:") :].strip())
             usage = data.get("usage") if isinstance(data, dict) else None
             if isinstance(usage, dict) and usage.get("output_tokens") is not None:
                 latest = usage["output_tokens"]
@@ -802,7 +841,9 @@ async def test_async_sse_wrapper_bills_partial_when_detached_drain_cap_reached(m
             await asyncio.sleep(0)
 
         assert iterator.logged_chunks, "capped pump never billed"
-        assert len(iterator.logged_chunks) <= 2 + streaming_iterator_module.ANTHROPIC_MESSAGES_STREAM_RELAY_QUEUE_MAXSIZE
+        assert (
+            len(iterator.logged_chunks) <= 2 + streaming_iterator_module.ANTHROPIC_MESSAGES_STREAM_RELAY_QUEUE_MAXSIZE
+        )
         assert len(iterator.logged_chunks) < 100
         assert not any(c.startswith(b"event: message_stop\n") for c in iterator.logged_chunks)
         assert tail_reached is False, "pump kept draining past the cap instead of stopping"
@@ -1036,7 +1077,10 @@ async def test_proxy_disconnect_closing_upstream_prices_recovered_tokens():
         (
             {"type": "message_start", "message": {"id": "msg_1", "usage": {"input_tokens": 29, "output_tokens": 1}}},
             {"type": "content_block_start", "index": 0, "content_block": {"type": "text", "text": ""}},
-            *({"type": "content_block_delta", "index": 0, "delta": {"type": "text_delta", "text": sentence}} for _ in range(6)),
+            *(
+                {"type": "content_block_delta", "index": 0, "delta": {"type": "text_delta", "text": sentence}}
+                for _ in range(6)
+            ),
         )
     )
     enqueued: list = []
@@ -1113,7 +1157,9 @@ async def test_async_sse_wrapper_aborts_upstream_when_detached_drain_cap_reached
             self.aclose_called = True
 
     stream = _AbortableStream()
-    iterator = _RecordingLoggingIterator(litellm_logging_obj=_make_logging_obj("abort_upstream_at_cap"), request_body={})
+    iterator = _RecordingLoggingIterator(
+        litellm_logging_obj=_make_logging_obj("abort_upstream_at_cap"), request_body={}
+    )
     try:
         gen = iterator.async_sse_wrapper(stream)
         await gen.__anext__()
@@ -1347,13 +1393,15 @@ def test_anthropic_messages_streaming_response_reports_withheld_output_of_its_st
 
     hidden_params: AnthropicMessagesStreamHiddenParams = {"additional_headers": {}}
     assert (
-        AnthropicMessagesStreamingResponse(completion_stream=_HoldingBack(), hidden_params=hidden_params)
-        .has_buffered_provider_output
+        AnthropicMessagesStreamingResponse(
+            completion_stream=_HoldingBack(), hidden_params=hidden_params
+        ).has_buffered_provider_output
         is True
     )
     assert (
-        AnthropicMessagesStreamingResponse(completion_stream=_bare_stream(), hidden_params=hidden_params)
-        .has_buffered_provider_output
+        AnthropicMessagesStreamingResponse(
+            completion_stream=_bare_stream(), hidden_params=hidden_params
+        ).has_buffered_provider_output
         is False
     )
 
@@ -1780,7 +1828,10 @@ async def test_native_stream_preserves_nullable_stop_fields_across_chunk_boundar
 
 def _data_events(frames: Sequence[bytes]) -> list[Mapping[str, object]]:
     return [
-        json.loads(line[len(b"data:") :]) for frame in frames for line in frame.splitlines() if line.startswith(b"data:")
+        json.loads(line[len(b"data:") :])
+        for frame in frames
+        for line in frame.splitlines()
+        if line.startswith(b"data:")
     ]
 
 
@@ -1788,9 +1839,13 @@ def _data_events(frames: Sequence[bytes]) -> list[Mapping[str, object]]:
 @pytest.mark.parametrize("separator", [b"\n", b"\r\n"])
 @pytest.mark.parametrize("chunk_size", [1, 7, 4096])
 async def test_safeguarded_sse_stream_stamps_the_final_delta_across_chunk_boundaries(separator: bytes, chunk_size: int):
-    classifier: Final = _OneVerdictClassifier('{"verdicts": {"call_curl": {"flagged": true, "explanation": "fetched"}}}')
+    classifier: Final = _OneVerdictClassifier(
+        '{"verdicts": {"call_curl": {"flagged": true, "explanation": "fetched"}}}'
+    )
     frames: Final = _sse_frames(separator)
-    stream: Final = SafeguardedAnthropicSSEStream(_chunked(b"".join(frames), chunk_size), _safeguards_evaluator(classifier))
+    stream: Final = SafeguardedAnthropicSSEStream(
+        _chunked(b"".join(frames), chunk_size), _safeguards_evaluator(classifier)
+    )
     emitted: Final = await _drained(stream)
     assert len(emitted) == len(frames)
     assert [emitted[i] for i in (0, 1, 2, 3, 5)] == [frames[i] for i in (0, 1, 2, 3, 5)]
@@ -1857,6 +1912,9 @@ async def test_safeguarded_sse_stream_reports_upstream_buffering_and_closes_the_
     assert stream.has_buffered_provider_output is True
     await stream.aclose()
     assert upstream.closed is True
-    assert SafeguardedAnthropicSSEStream(
-        _chunked(b"", 1), _safeguards_evaluator(_OneVerdictClassifier("{}"))
-    ).has_buffered_provider_output is False
+    assert (
+        SafeguardedAnthropicSSEStream(
+            _chunked(b"", 1), _safeguards_evaluator(_OneVerdictClassifier("{}"))
+        ).has_buffered_provider_output
+        is False
+    )