mirror of
https://github.com/BerriAI/litellm.git
synced 2026-11-05 07:37:50 +00:00
refactor(proxy): price a stream once and copy the builder's bill into the streamed usage.cost at egress
The egress no longer re-prices a reconstructed chunk. The stream wrapper copies the hidden response_cost the builder computed for the spend row onto the usage chunk, the Responses iterator stamps the logging object's bill the same way, and the proxy egress only copies that number into the client-visible usage.cost. Deletes the egress pricer, the _streamed_usage_cost widening and the served-model carry Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>
This commit is contained in:
kerry
parent
b55caf052b
commit
b6571348e5
10 changed files with 167 additions and 150 deletions
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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,
|
||||
|
|
|
|||
|
|
@ -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:
|
||||
|
|
|
|||
|
|
@ -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)})
|
||||
|
|
|
|||
|
|
@ -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(
|
||||
|
|
|
|||
|
|
@ -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),
|
||||
|
|
|
|||
|
|
@ -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,
|
||||
):
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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()
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue