diff --git a/litellm/litellm_core_utils/streaming_handler.py b/litellm/litellm_core_utils/streaming_handler.py index 6bd0e45fbc0..4497f3de9f3 100644 --- a/litellm/litellm_core_utils/streaming_handler.py +++ b/litellm/litellm_core_utils/streaming_handler.py @@ -92,6 +92,17 @@ def _next_sync_or_exhausted(it: Iterator[object]) -> object: return _SYNC_ITER_EXHAUSTED +def _stamp_hidden_response_cost( + usage_chunk: ModelResponse | TextCompletionResponse | ModelResponseStream, + complete_response: ModelResponse | TextCompletionResponse | ModelResponseStream, +) -> None: + """The bill the builder priced for the spend row rides the usage chunk, so the proxy shows the same number.""" + bill: Final = (get_hidden_params(complete_response) or {}).get("response_cost") + target: Final = get_hidden_params(usage_chunk) + if target is not None and isinstance(bill, (int, float)) and not isinstance(bill, bool): + target["response_cost"] = float(bill) + + def _stamp_served_service_tier(response: ModelResponseStream, complete_streaming_response: ModelResponse) -> None: served_tier: Final = complete_streaming_response.model_dump().get("service_tier") if isinstance(served_tier, str) and served_tier: @@ -2122,6 +2133,7 @@ class CustomStreamWrapper: getattr(complete_streaming_response, "usage"), ) _stamp_served_service_tier(response, complete_streaming_response) + _stamp_hidden_response_cost(response, complete_streaming_response) try: _cache_copy = complete_streaming_response.model_copy(deep=True) _log_copy = complete_streaming_response.model_copy(deep=True) @@ -2343,6 +2355,7 @@ class CustomStreamWrapper: getattr(complete_streaming_response, "usage"), ) _stamp_served_service_tier(response, complete_streaming_response) + _stamp_hidden_response_cost(response, complete_streaming_response) try: _copy = complete_streaming_response.model_copy(deep=True) except RuntimeError: @@ -2500,9 +2513,7 @@ class CustomStreamWrapper: return any(get(field) for field in cls._CONTINUATION_DISQUALIFYING_DELTA_FIELDS) def _new_usage_chunk(self) -> ModelResponseStream: - usage_chunk: Final = self.model_response_creator( - hidden_params=_provider_hidden_params(None, self._provider_response_model) - ) + usage_chunk: Final = self.model_response_creator() if self.caller_requested_stream_usage: usage_chunk.choices = [] return usage_chunk diff --git a/litellm/proxy/common_request_processing.py b/litellm/proxy/common_request_processing.py index 26adbb8def5..3b9c17baa02 100644 --- a/litellm/proxy/common_request_processing.py +++ b/litellm/proxy/common_request_processing.py @@ -273,7 +273,6 @@ from litellm.proxy.litellm_pre_call_utils import ( ) from litellm.proxy.policy_engine.response_retrieval import attach_post_call_pipelines_to_retrieval from litellm.types.litellm_params import STREAM_USAGE_INJECTED_FIELD -from litellm.types.llms.openai import ResponsesAPIResponse from litellm.types.utils import ( ModelResponse, ModelResponseStream, @@ -4725,7 +4724,7 @@ class ProxyBaseLLMRequestProcessing: @staticmethod def _completion_cost_or_none( - model_response: ModelResponse | ResponsesAPIResponse, model_name: str, service_tier: str | None + model_response: ModelResponse, model_name: str, service_tier: str | None ) -> float | None: try: return litellm.completion_cost( @@ -4736,7 +4735,7 @@ class ProxyBaseLLMRequestProcessing: @staticmethod def _logging_obj_cost_or_none( - model_response: ModelResponse | ResponsesAPIResponse, litellm_logging_obj: LiteLLMLoggingObj + model_response: ModelResponse, litellm_logging_obj: LiteLLMLoggingObj ) -> float | None: # Pricing a frame stamps cost_breakdown and, on failure, the cost-failure debug key onto # the live logging object. The pass-through handlers never recompute either one, so a @@ -4763,7 +4762,7 @@ class ProxyBaseLLMRequestProcessing: @staticmethod def _streamed_usage_cost( - model_response: ModelResponse | ResponsesAPIResponse, + model_response: ModelResponse, model_name: str, service_tier: str | None, litellm_logging_obj: LiteLLMLoggingObj | None, diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index 5c7388f7861..87d90904ee2 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -922,7 +922,7 @@ from litellm.proxy.ui_crud_endpoints.proxy_setting_endpoints import ( from litellm.proxy.ui_crud_endpoints.user_banner_endpoints import ( router as user_banner_endpoints_router, ) -from litellm.proxy.usage_cost_egress import Pricer, with_client_stream_usage_cost +from litellm.proxy.usage_cost_egress import with_client_stream_usage_cost from litellm.proxy.utils import ( # noqa: F401, RUF100 # legacy module exports PrismaClient, ProxyLogging, @@ -10135,22 +10135,6 @@ def _fast_serialize_simple_model_response_stream( return orjson.dumps(payload) -def _stream_usage_pricer(request_data: Mapping[str, object]) -> Pricer: - """Bills a usage-bearing stream event the way the request's logging object bills the spend row.""" - logging_obj: Final = request_data.get("litellm_logging_obj") - model_name: Final = str(request_data.get("model", "")) - - def price(reply: ModelResponse | ResponsesAPIResponse) -> float | None: - return ProxyBaseLLMRequestProcessing._streamed_usage_cost( # pyright: ignore[reportPrivateUsage] # the proxy's one stream pricer - reply, - model_name, - None, - logging_obj if isinstance(logging_obj, LiteLLMLoggingObj) else None, - ) - - return price - - def _serialize_streaming_chunk(chunk: BaseModel) -> str | bytes: if isinstance(chunk, ModelResponseStream): serialized_chunk: Final = _fast_serialize_simple_model_response_stream(chunk) @@ -10489,7 +10473,6 @@ async def async_data_generator( needs_per_chunk_hook: Final = proxy_logging_obj.needs_per_chunk_streaming_hook() is_raw_sse_stream: Final = bool(request_data.get("_litellm_raw_sse_stream")) strip_stream_usage: Final = bool(request_data.get("_litellm_strip_stream_usage")) - price_stream_usage: Final = _stream_usage_pricer(request_data) raw_sse_buffer = "" # A stream can start on a deployment with keepalive off and fall back @@ -10576,7 +10559,7 @@ async def async_data_generator( error_state.observe_chunk(chunk) raw_passthrough = False if isinstance(chunk, BaseModel): - chunk = _serialize_streaming_chunk(with_client_stream_usage_cost(chunk, price_stream_usage)) + chunk = _serialize_streaming_chunk(with_client_stream_usage_cost(chunk)) elif isinstance(chunk, bytes): chunk = chunk.decode("utf-8", errors="replace") if is_raw_sse_stream: diff --git a/litellm/proxy/usage_cost_egress.py b/litellm/proxy/usage_cost_egress.py index 1dcdd69706f..ac29c1cef09 100644 --- a/litellm/proxy/usage_cost_egress.py +++ b/litellm/proxy/usage_cost_egress.py @@ -1,21 +1,17 @@ -"""What a client reads in a streamed ``usage.cost``: litellm's bill, written once on the way out of the proxy, onto -a copy, so the final usage frame and the spend row agree. Inside the pipeline ``usage.cost`` stays what the -upstream reported.""" +"""What a client reads in a streamed ``usage.cost``: the bill litellm already priced for the spend row, copied once on +the way out of the proxy onto a copy of the event. Nothing is priced here. Inside the pipeline ``usage.cost`` stays +what the upstream reported.""" from __future__ import annotations import math -from collections.abc import Callable -from typing import Final, TypeAlias, TypeVar +from typing import Final, TypeVar from pydantic import BaseModel from litellm.litellm_core_utils.hidden_params import get_hidden_params from litellm.types.llms.openai import ResponseAPIUsage, ResponsesAPIResponse -from litellm.types.utils import ModelResponse, ModelResponseStream, Usage - -Pricer: TypeAlias = Callable[[ModelResponse | ResponsesAPIResponse], float | None] -"""Prices a reply the way the request's logging object bills it (deployment rates, discount, margin).""" +from litellm.types.utils import ModelResponseStream, Usage _RESPONSES_USAGE_EVENTS: Final = frozenset({"response.completed", "response.incomplete"}) _ReplyT: Final = TypeVar("_ReplyT", bound=BaseModel) @@ -41,23 +37,27 @@ def _with_usage_cost(reply: _ReplyT, usage: Usage | ResponseAPIUsage, cost: floa return reply.model_copy(update={"usage": usage.model_copy(update={"cost": cost})}) -def with_client_stream_usage_cost(chunk: BaseModel, price: Pricer) -> BaseModel: - """A streamed event as the client gets it. Only the usage-bearing events are priced: the chat chunk carrying - ``usage`` and the Responses ``response.completed`` / ``response.incomplete`` events.""" +def _hidden_bill(reply: object) -> object: + hidden: Final = get_hidden_params(reply) + return hidden.get("response_cost") if hidden is not None else None + + +def with_client_stream_usage_cost(chunk: BaseModel) -> BaseModel: + """A streamed event as the client gets it. Only the usage-bearing events change: the chat chunk carrying + ``usage`` and the Responses ``response.completed`` / ``response.incomplete`` events, each read against the + hidden ``response_cost`` the pipeline stamped on it.""" if isinstance(chunk, ModelResponseStream): usage: Final = getattr(chunk, "usage", None) if not isinstance(usage, Usage): return chunk - priced: Final = ModelResponse(model=chunk.model, usage=usage, service_tier=getattr(chunk, "service_tier", None)) - priced.hidden_params = dict(get_hidden_params(chunk) or {}) - cost: Final = _client_usage_cost(getattr(usage, "cost", None), price(priced)) + cost: Final = _client_usage_cost(getattr(usage, "cost", None), _hidden_bill(chunk)) return chunk if cost is None else _with_usage_cost(chunk, usage, cost) if getattr(chunk, "type", None) not in _RESPONSES_USAGE_EVENTS: return chunk response: Final = getattr(chunk, "response", None) if not isinstance(response, ResponsesAPIResponse) or not isinstance(response.usage, ResponseAPIUsage): return chunk - response_cost: Final = _client_usage_cost(getattr(response.usage, "cost", None), price(response)) + response_cost: Final = _client_usage_cost(getattr(response.usage, "cost", None), _hidden_bill(response)) if response_cost is None: return chunk return chunk.model_copy(update={"response": _with_usage_cost(response, response.usage, response_cost)}) diff --git a/litellm/responses/litellm_completion_transformation/streaming_iterator.py b/litellm/responses/litellm_completion_transformation/streaming_iterator.py index dc84a955d88..da375b7265c 100644 --- a/litellm/responses/litellm_completion_transformation/streaming_iterator.py +++ b/litellm/responses/litellm_completion_transformation/streaming_iterator.py @@ -1480,7 +1480,7 @@ class LiteLLMCompletionStreamingIterator(ResponsesAPIStreamingIterator): } ) ) - normalize_responses_usage(partial) + normalize_responses_usage(partial, self.litellm_logging_obj) return ResponseIncompleteEvent(type=ResponsesAPIStreamEvents.RESPONSE_INCOMPLETE, response=partial) def _emit_terminal_response_event( diff --git a/litellm/responses/streaming_iterator.py b/litellm/responses/streaming_iterator.py index 9daf350baa7..11c8d5405c0 100644 --- a/litellm/responses/streaming_iterator.py +++ b/litellm/responses/streaming_iterator.py @@ -41,6 +41,7 @@ from litellm.exceptions import MidStreamFallbackError, RateLimitError from litellm.litellm_core_utils.asyncify import run_async_function from litellm.litellm_core_utils.core_helpers import process_response_headers from litellm.litellm_core_utils.exception_mapping_utils import with_deployment_api_base +from litellm.litellm_core_utils.hidden_params import get_hidden_params from litellm.litellm_core_utils.internal_call_metadata import without_budget_reservation from litellm.litellm_core_utils.litellm_logging import BilledRoundsSpend from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj @@ -585,7 +586,7 @@ class BaseResponsesAPIStreamingIterator: else openai_responses_api_chunk.model_copy(update={"response": _billed_response}) ) self.completed_response = _terminal_chunk - normalize_responses_usage(_billed_response) + normalize_responses_usage(_billed_response, self.logging_obj) if _chunk_type == openai_types.ResponsesAPIStreamEvents.RESPONSE_FAILED: self._handle_logging_failed_response() @@ -1069,7 +1070,7 @@ class BaseResponsesAPIStreamingIterator: "output": _partial_output_items(self._generated_content), } ) - normalize_responses_usage(partial) + normalize_responses_usage(partial, self.logging_obj) return ResponseIncompleteEvent(type=ResponsesAPIStreamEvents.RESPONSE_INCOMPLETE, response=partial) def _note_yielded_event(self, event: ResponsesAPIStreamingResponse) -> None: @@ -1741,14 +1742,24 @@ def _estimate_usage_safely( return None -def normalize_responses_usage(response_obj: ResponsesAPIResponse | None) -> None: - """A terminal response built off an untyped dict carries its usage as a dict, logging and the proxy read a model.""" +def normalize_responses_usage(response_obj: ResponsesAPIResponse | None, logging_obj: LiteLLMLoggingObj | None) -> None: + """A terminal response built off an untyped dict carries its usage as a dict, logging and the proxy read a model. + The bill the logging object prices for it rides its hidden ``response_cost``, so the proxy shows the same number.""" if response_obj is None: return usage_obj: Final[ResponseAPIUsage | None] = _usage_as_model(getattr(response_obj, "usage", None)) if usage_obj is None: return response_obj.usage = usage_obj # rebind-ok: the typed usage has to ride on the response the client receives + if logging_obj is None: + return + try: + bill: Final = logging_obj.response_cost_calculator(result=response_obj) + except Exception: # noqa: BLE001 # a pricing failure leaves the client usage as the upstream sent it + return + hidden: Final = get_hidden_params(response_obj) + if hidden is not None and isinstance(bill, (int, float)) and not isinstance(bill, bool): + hidden["response_cost"] = float(bill) def opening_item_payload(payload: Mapping[str, object]) -> Mapping[str, object]: @@ -1782,7 +1793,7 @@ def build_synthetic_response_events( chunk_size: int, ) -> list[ResponsesAPIStreamingResponse]: openai_types: Final = get_openai_response_types() - normalize_responses_usage(transformed) + normalize_responses_usage(transformed, logging_obj) events: Final[list[ResponsesAPIStreamingResponse]] = [ _build_response_status_event(openai_types.ResponsesAPIStreamEvents.RESPONSE_CREATED, transformed), diff --git a/tests/unit/litellm_core_utils/test_streaming_handler.py b/tests/unit/litellm_core_utils/test_streaming_handler.py index 6ccb4b8828c..d24d147b3df 100644 --- a/tests/unit/litellm_core_utils/test_streaming_handler.py +++ b/tests/unit/litellm_core_utils/test_streaming_handler.py @@ -8,15 +8,16 @@ import time from collections.abc import Mapping, Sequence from dataclasses import dataclass from types import CodeType, FrameType, MappingProxyType -from typing import AsyncIterator, Callable, Final, Iterator, Literal, NoReturn, Optional, cast +from typing import AsyncIterator, Callable, Final, Iterator, Literal, NoReturn, Optional, Tuple, cast from unittest.mock import AsyncMock, MagicMock, Mock, patch import httpx import pytest from openai.types.chat import ChatCompletionChunk from pydantic import BaseModel, TypeAdapter + import litellm -from litellm import verbose_logger +from litellm import completion, verbose_logger from litellm._logging import session_id_var, trace_id_var from litellm.exceptions import MidStreamFallbackError from litellm.integrations.custom_logger import CustomLogger @@ -27,7 +28,6 @@ from litellm.litellm_core_utils.streaming_handler import ( _ProviderChunkEarlyReturn, _ProviderChunkParsed, ) -from litellm.llms.custom_httpx.http_handler import HTTPHandler from litellm.responses.mcp.chat_completions_handler import _add_mcp_metadata_to_response from litellm.types.utils import ( CompletionTokensDetailsWrapper, @@ -35,13 +35,11 @@ from litellm.types.utils import ( ModelResponse, ModelResponseStream, PromptTokensDetailsWrapper, - StreamingChoices, StandardLoggingPayload, + StreamingChoices, Usage, ) from litellm.utils import ModelResponseListIterator -from litellm import completion -from typing import Tuple @pytest.mark.parametrize("async_stream", [False, True]) @@ -2436,7 +2434,7 @@ def test_openai_compatible_streaming_cost_is_the_reported_cost(): custom_llm_provider="openai", call_type="completion", optional_params={}, - ) == pytest.approx(0.037) + ) == 0.037 def test_xai_streaming_reported_cost_still_takes_the_margin(monkeypatch): @@ -6059,6 +6057,25 @@ def test_chunk_creator_preserves_hidden_provider_specific_fields_from_parsed_chu assert assembled._hidden_params["provider_specific_fields"] == {"traffic_type": "ON_DEMAND_FLEX"} +def test_the_usage_chunk_carries_the_bill_the_builder_priced_for_the_spend_row() -> None: + from litellm.litellm_core_utils.hidden_params import get_hidden_params + from litellm.litellm_core_utils.streaming_handler import ( + _stamp_hidden_response_cost, # pyright: ignore[reportPrivateUsage] # the helper under test + ) + + usage: Final = Usage(prompt_tokens=100, completion_tokens=20, total_tokens=120) + complete: Final = ModelResponse(model="omni-model", usage=usage) + complete.hidden_params = {"response_cost": 0.00032} + usage_chunk: Final = ModelResponseStream(id="chunk-1", model="omni-model", choices=[], usage=usage) + usage_chunk.hidden_params = {"response_cost": None} + + _stamp_hidden_response_cost(usage_chunk, complete) + + stamped: Final = get_hidden_params(usage_chunk) + assert stamped is not None and stamped["response_cost"] == 0.00032 + assert getattr(usage, "cost", None) is None + + def test_chunk_creator_keeps_provider_model_private_across_stream(): from litellm.router_utils.add_retry_fallback_headers import ( get_hidden_params_dict, @@ -6118,28 +6135,6 @@ def test_chunk_creator_keeps_provider_model_private_across_stream(): ) -def test_the_usage_chunk_carries_the_served_model_so_its_price_matches_the_assembled_response(): - from litellm.router_utils.add_retry_fallback_headers import get_hidden_params_dict - - wrapper = CustomStreamWrapper( - completion_stream=None, - model="requested-route", - logging_obj=MagicMock(), - custom_llm_provider="openai", - ) - wrapper.chunk_creator( - chunk=ModelResponseStream( - id="chunk-1", - model="selected-model", - choices=[StreamingChoices(finish_reason=None, index=0, delta=Delta(content="hello"))], - ) - ) - - usage_chunk = wrapper._new_usage_chunk() - - assert get_hidden_params_dict(usage_chunk)["provider_response_model"] == "selected-model" - - def test_assembled_stream_uses_later_provider_model_for_cost( monkeypatch: pytest.MonkeyPatch, ): diff --git a/tests/unit/proxy/test_usage_cost_egress.py b/tests/unit/proxy/test_usage_cost_egress.py index 070139c5019..62c32cd3ebc 100644 --- a/tests/unit/proxy/test_usage_cost_egress.py +++ b/tests/unit/proxy/test_usage_cost_egress.py @@ -4,6 +4,7 @@ from typing import Final import pytest +from litellm.litellm_core_utils.hidden_params import get_hidden_params from litellm.proxy.usage_cost_egress import with_client_stream_usage_cost from litellm.types.llms.openai import ( ResponseAPIUsage, @@ -12,14 +13,14 @@ from litellm.types.llms.openai import ( ResponsesAPIResponse, ResponsesAPIStreamEvents, ) -from litellm.types.utils import ModelResponseStream, StreamingChoices, Usage +from litellm.types.utils import Delta, ModelResponseStream, StreamingChoices, Usage UPSTREAM_COST: Final = 0.037 OPERATOR_BILL: Final = 100 * 2e-6 + 20 * 6e-6 -def _responses_reply(upstream_cost: float | None) -> ResponsesAPIResponse: - return ResponsesAPIResponse( +def _responses_reply(upstream_cost: float | None, bill: object) -> ResponsesAPIResponse: + reply: Final = ResponsesAPIResponse( id="resp_egress", created_at=1, status="completed", @@ -28,101 +29,71 @@ def _responses_reply(upstream_cost: float | None) -> ResponsesAPIResponse: output=[], usage=ResponseAPIUsage(input_tokens=100, output_tokens=20, total_tokens=120, cost=upstream_cost), ) + hidden: Final = get_hidden_params(reply) + assert hidden is not None + hidden["response_cost"] = bill + return reply -def _usage_chunk(upstream_cost: float | None, service_tier: str | None = None) -> ModelResponseStream: - return ModelResponseStream( +def _usage_chunk(upstream_cost: float | None, bill: object) -> ModelResponseStream: + chunk: Final = ModelResponseStream( id="chatcmpl-egress", model="omni-model", choices=[], - service_tier=service_tier, usage=Usage(prompt_tokens=100, completion_tokens=20, total_tokens=120, cost=upstream_cost), ) - - -def _bill(_reply: object) -> float | None: - return OPERATOR_BILL - - -def _unpriceable(_reply: object) -> float | None: - return None + chunk.hidden_params = {"response_cost": bill} + return chunk def test_the_streamed_usage_chunk_tells_the_client_what_litellm_billed() -> None: - chunk: Final = _usage_chunk(UPSTREAM_COST) + chunk: Final = _usage_chunk(UPSTREAM_COST, OPERATOR_BILL) - billed: Final = with_client_stream_usage_cost(chunk, _bill) + billed: Final = with_client_stream_usage_cost(chunk) assert isinstance(billed, ModelResponseStream) and billed is not chunk - assert billed.usage is not None and billed.usage.cost == pytest.approx(OPERATOR_BILL) + assert billed.usage is not None and billed.usage.cost == OPERATOR_BILL assert chunk.usage is not None and chunk.usage.cost == UPSTREAM_COST def test_a_zero_bill_replaces_the_upstream_cost() -> None: - billed: Final = with_client_stream_usage_cost(_usage_chunk(UPSTREAM_COST), lambda _reply: 0.0) + billed: Final = with_client_stream_usage_cost(_usage_chunk(UPSTREAM_COST, 0.0)) assert isinstance(billed, ModelResponseStream) assert billed.usage is not None and billed.usage.cost == 0.0 -@pytest.mark.parametrize("bill", [None, -0.001, float("nan"), float("inf")]) -def test_an_unpriceable_stream_keeps_the_upstream_cost(bill: float | None) -> None: - chunk: Final = _usage_chunk(UPSTREAM_COST) +@pytest.mark.parametrize("bill", [None, -0.001, float("nan"), float("inf"), True]) +def test_an_unpriced_stream_keeps_the_upstream_cost(bill: object) -> None: + chunk: Final = _usage_chunk(UPSTREAM_COST, bill) - assert with_client_stream_usage_cost(chunk, lambda _reply: bill) is chunk + assert with_client_stream_usage_cost(chunk) is chunk assert chunk.usage is not None and chunk.usage.cost == UPSTREAM_COST def test_a_priced_stream_without_an_upstream_cost_still_tells_the_client_the_bill() -> None: - chunk: Final = _usage_chunk(None) + billed: Final = with_client_stream_usage_cost(_usage_chunk(None, OPERATOR_BILL)) - billed: Final = with_client_stream_usage_cost(chunk, _bill) - - assert isinstance(billed, ModelResponseStream) and billed is not chunk - assert billed.usage is not None and billed.usage.cost == pytest.approx(OPERATOR_BILL) + assert isinstance(billed, ModelResponseStream) + assert billed.usage is not None and billed.usage.cost == OPERATOR_BILL @pytest.mark.parametrize("bill", [None, 0.0]) def test_an_unpriced_stream_without_an_upstream_cost_goes_out_as_sent(bill: float | None) -> None: - chunk: Final = _usage_chunk(None) + chunk: Final = _usage_chunk(None, bill) - assert with_client_stream_usage_cost(chunk, lambda _reply: bill) is chunk + assert with_client_stream_usage_cost(chunk) is chunk -def test_a_content_chunk_is_not_priced() -> None: +def test_a_content_chunk_is_not_touched() -> None: chunk: Final = ModelResponseStream( - id="chatcmpl-egress", model="omni-model", choices=[StreamingChoices(index=0, delta={"content": "hi"})] + id="chatcmpl-egress", + model="omni-model", + choices=[StreamingChoices(index=0, delta=Delta(content="hi"), finish_reason=None)], ) - priced: list[object] = [] + chunk.hidden_params = {"response_cost": OPERATOR_BILL} - def price(reply: object) -> float | None: - priced.append(reply) - return OPERATOR_BILL - - assert with_client_stream_usage_cost(chunk, price) is chunk - assert priced == [] - - -def test_the_stream_pricer_sees_the_served_model_tier_and_upstream_cost_so_the_bill_matches_the_spend_row() -> None: - seen: list[tuple[object, object, object, object]] = [] - - def price(reply: object) -> float | None: - seen.append( - ( - getattr(reply, "model", None), - getattr(reply, "hidden_params", {}).get("provider_response_model"), - getattr(reply, "service_tier", None), - getattr(getattr(reply, "usage", None), "cost", None), - ) - ) - return OPERATOR_BILL - - chunk: Final = _usage_chunk(UPSTREAM_COST, service_tier="default") - chunk.hidden_params = {"provider_response_model": "gpt-4o-mini"} - - with_client_stream_usage_cost(chunk, price) - - assert seen == [("omni-model", "gpt-4o-mini", "default", UPSTREAM_COST)] + assert with_client_stream_usage_cost(chunk) is chunk @pytest.mark.parametrize("upstream_cost", [UPSTREAM_COST, None]) @@ -132,21 +103,24 @@ def test_the_stream_pricer_sees_the_served_model_tier_and_upstream_cost_so_the_b def test_a_terminal_responses_event_carries_the_bill_with_or_without_an_upstream_cost( upstream_cost: float | None, event_type: ResponsesAPIStreamEvents ) -> None: - event_cls: Final = ( - ResponseCompletedEvent if event_type == ResponsesAPIStreamEvents.RESPONSE_COMPLETED else ResponseIncompleteEvent + reply: Final = _responses_reply(upstream_cost, OPERATOR_BILL) + event: Final = ( + ResponseCompletedEvent(type=ResponsesAPIStreamEvents.RESPONSE_COMPLETED, response=reply) + if event_type is ResponsesAPIStreamEvents.RESPONSE_COMPLETED + else ResponseIncompleteEvent(type=ResponsesAPIStreamEvents.RESPONSE_INCOMPLETE, response=reply) ) - event: Final = event_cls(type=event_type, response=_responses_reply(upstream_cost)) - billed: Final = with_client_stream_usage_cost(event, _bill) + billed: Final = with_client_stream_usage_cost(event) assert billed is not event - assert getattr(billed, "response").usage.cost == pytest.approx(OPERATOR_BILL) - assert event.response.usage is not None and event.response.usage.cost == upstream_cost + assert isinstance(billed, (ResponseCompletedEvent, ResponseIncompleteEvent)) + assert billed.response.usage is not None and billed.response.usage.cost == OPERATOR_BILL + assert reply.usage is not None and reply.usage.cost == upstream_cost -def test_an_unpriceable_terminal_responses_event_goes_out_as_sent() -> None: +def test_an_unpriced_terminal_responses_event_goes_out_as_sent() -> None: event: Final = ResponseCompletedEvent( - type=ResponsesAPIStreamEvents.RESPONSE_COMPLETED, response=_responses_reply(None) + type=ResponsesAPIStreamEvents.RESPONSE_COMPLETED, response=_responses_reply(UPSTREAM_COST, None) ) - assert with_client_stream_usage_cost(event, _unpriceable) is event + assert with_client_stream_usage_cost(event) is event diff --git a/tests/unit/responses/litellm_completion_transformation/test_streaming_iterator_transformation.py b/tests/unit/responses/litellm_completion_transformation/test_streaming_iterator_transformation.py index ed90acbd2fa..e4784e92bc4 100644 --- a/tests/unit/responses/litellm_completion_transformation/test_streaming_iterator_transformation.py +++ b/tests/unit/responses/litellm_completion_transformation/test_streaming_iterator_transformation.py @@ -10,6 +10,7 @@ before response.completed, and that every event of a bridged stream carries the spend tracking stores, so a follow-up previous_response_id still finds the conversation. """ +import itertools import json from collections.abc import Callable from typing import Final @@ -23,6 +24,7 @@ from litellm.responses.litellm_completion_transformation.streaming_iterator impo from litellm.responses.utils import ResponsesAPIRequestUtils from litellm.types.llms.openai import ( BaseLiteLLMOpenAIResponseObject, + OutputTextDeltaEvent, ResponsesAPIStreamEvents, ) from litellm.types.responses.main import build_web_search_call @@ -1558,3 +1560,29 @@ async def test_a_compaction_block_without_its_summary_streams_no_compaction_item assert [event["item"]["type"] for event in events if "item" in event] == ["message", "message"] assert [item["type"] for item in events[-1]["response"]["output"]] == ["message"] + + +def test_the_partial_response_billed_on_disconnect_carries_the_logging_objects_bill_as_hidden_response_cost() -> None: + from litellm.litellm_core_utils.hidden_params import get_hidden_params + + wrapper: Final = _FakeStreamWrapper([_chunk("Hello"), _chunk(" world", finish_reason="stop")]) + wrapper.logging_obj.response_cost_calculator.return_value = 0.00032 + iterator: Final = LiteLLMCompletionStreamingIterator( + model="claude-haiku-4-5", + litellm_custom_stream_wrapper=wrapper, # pyright: ignore[reportArgumentType] # the fake stands in for CustomStreamWrapper, as in _build_iterator + request_input="What is the weather in San Francisco?", + responses_api_request={}, + custom_llm_provider="anthropic", + litellm_metadata={}, + ) + relayed: Final = tuple( + itertools.takewhile(lambda event: not isinstance(event, OutputTextDeltaEvent), iterator) + ) + assert relayed, "the stream must relay at least one chunk before the client goes away" + + partial: Final = iterator.partial_response_on_disconnect() + + assert partial is not None + hidden: Final = get_hidden_params(partial.response) + assert hidden is not None and hidden["response_cost"] == 0.00032 + assert wrapper.logging_obj.response_cost_calculator.call_args.kwargs["result"] is partial.response diff --git a/tests/unit/responses/test_streaming_iterator.py b/tests/unit/responses/test_streaming_iterator.py index bd9122a927c..710bcb43e25 100644 --- a/tests/unit/responses/test_streaming_iterator.py +++ b/tests/unit/responses/test_streaming_iterator.py @@ -25,6 +25,7 @@ from pydantic_core import PydanticSerializationError import litellm from litellm.exceptions import MidStreamFallbackError from litellm.integrations.custom_logger import CustomLogger +from litellm.litellm_core_utils.hidden_params import get_hidden_params from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj from litellm.llms.base_llm.responses.transformation import BaseResponsesAPIConfig from litellm.responses import streaming_iterator as streaming_module @@ -864,7 +865,7 @@ def test_normalize_responses_usage_types_a_dict_usage_and_keeps_the_provider_cos } ) - normalize_responses_usage(response) + normalize_responses_usage(response, None) assert isinstance(response.usage, ResponseAPIUsage) assert response.usage.cost == pytest.approx(3e-05) @@ -876,11 +877,26 @@ def test_normalize_responses_usage_never_writes_a_cost_the_provider_did_not_repo response = _responses_api_response_with_usage() - normalize_responses_usage(response) + normalize_responses_usage(response, None) assert getattr(response.usage, "cost", None) is None +def test_normalize_responses_usage_stamps_the_logging_objects_bill_as_hidden_response_cost(): + from litellm.responses.streaming_iterator import normalize_responses_usage + + response = _responses_api_response_with_usage() + logging_obj = MagicMock() + logging_obj.response_cost_calculator.return_value = 0.00032 + + normalize_responses_usage(response, logging_obj) + + hidden: Final = get_hidden_params(response) + assert hidden is not None and hidden["response_cost"] == 0.00032 + assert getattr(response.usage, "cost", None) is None, "the client field is written at the proxy egress, not here" + assert logging_obj.response_cost_calculator.call_args.kwargs["result"] is response + + def _capture_dispatch(logged: list): """Record the object handed to the success handlers. @@ -898,7 +914,7 @@ def _capture_dispatch(logged: list): return _dispatch -def _headers_config(*, transform_hidden_params: dict | None = None) -> Mock: +def _headers_config(*, transform_hidden_params: Mapping[str, object] | None = None) -> Mock: """Config whose completed event carries a real ResponsesAPIResponse, so the logging copy performs a genuine model_dump/model_validate round trip.""" mock_config = Mock(spec=BaseResponsesAPIConfig) @@ -1683,7 +1699,7 @@ def test_process_chunk_wraps_encrypted_content_with_model_id(): @pytest.mark.usefixtures("_vcr_outcome_gate", "setup_and_teardown") -def test_process_chunk_completed_response_updates_id_and_leaves_usage_cost_alone(monkeypatch): +def test_process_chunk_completed_response_updates_id_and_leaves_usage_cost_alone(monkeypatch: pytest.MonkeyPatch): original_include_cost = litellm.include_cost_in_streaming_usage litellm.include_cost_in_streaming_usage = True openai_types = streaming_module._get_openai_response_types()