From a161780dd3cd4b28bfbda0a480ec6a6acfc1a009 Mon Sep 17 00:00:00 2001 From: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> Date: Mon, 9 Nov 2026 12:53:48 +0000 Subject: [PATCH] refactor: centralize Messages usage cost stamping and stream wrapping Co-Authored-By: Ishaan Jaffer <155045088+ishaan-berri@users.noreply.github.com> --- .../llm_response_utils/response_metadata.py | 29 ++---- litellm/litellm_core_utils/usage_cost.py | 47 +++++++--- .../llms/anthropic/messages/native/handler.py | 32 +++++-- .../messages/native/streaming_iterator.py | 30 ++----- litellm/utils.py | 9 -- .../native/test_streaming_iterator.py | 88 +++++++++++++++---- 6 files changed, 148 insertions(+), 87 deletions(-) diff --git a/litellm/litellm_core_utils/llm_response_utils/response_metadata.py b/litellm/litellm_core_utils/llm_response_utils/response_metadata.py index 3e3cacdf638..9fd4b277a5a 100644 --- a/litellm/litellm_core_utils/llm_response_utils/response_metadata.py +++ b/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) diff --git a/litellm/litellm_core_utils/usage_cost.py b/litellm/litellm_core_utils/usage_cost.py index cfd5991fa8a..96023ca1c0a 100644 --- a/litellm/litellm_core_utils/usage_cost.py +++ b/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: diff --git a/litellm/llms/anthropic/messages/native/handler.py b/litellm/llms/anthropic/messages/native/handler.py index 922de386fe4..6e9ddba3042 100644 --- a/litellm/llms/anthropic/messages/native/handler.py +++ b/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, diff --git a/litellm/llms/anthropic/messages/native/streaming_iterator.py b/litellm/llms/anthropic/messages/native/streaming_iterator.py index 39710756be6..3c00b3cc216 100644 --- a/litellm/llms/anthropic/messages/native/streaming_iterator.py +++ b/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) diff --git a/litellm/utils.py b/litellm/utils.py index fa44dbd713c..66b5b274f97 100644 --- a/litellm/utils.py +++ b/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: diff --git a/tests/unit/llms/anthropic/messages/native/test_streaming_iterator.py b/tests/unit/llms/anthropic/messages/native/test_streaming_iterator.py index dd50f60250a..691879f166c 100644 --- a/tests/unit/llms/anthropic/messages/native/test_streaming_iterator.py +++ b/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 + )