diff --git a/tests/integration/cost_calculation/cases/anthropic.py b/tests/integration/cost_calculation/cases/anthropic.py index c214352fdb0..702420e36e8 100644 --- a/tests/integration/cost_calculation/cases/anthropic.py +++ b/tests/integration/cost_calculation/cases/anthropic.py @@ -4002,6 +4002,46 @@ CLAUDE_SONNET_5_PASSTHROUGH_MESSAGES_CACHE_READ: Final = CostTrackingTestCase( ), ) +CLAUDE_SONNET_5_PASSTHROUGH_MESSAGES_STREAM_CACHE_WRITE_ON_MESSAGE_DELTA: Final = CostTrackingTestCase( + name="claude-sonnet-5-passthrough-messages_stream_cache_write_on_message_delta", + covers="quota_management.spend_tracking.scripted_wire.logs_cost", + model="claude-sonnet-5", + endpoint="/anthropic/v1/messages", + request={ + "model": "$MODEL", + "max_tokens": 12, + "messages": [ + { + "role": "user", + "content": [{"type": "text", "text": "cached text", "cache_control": {"type": "ephemeral"}}], + } + ], + "stream": True, + }, + response=SseResponse( + content_type="text/event-stream", + frames=( + 'event: message_start\ndata: {"type":"message_start","message":{"id":"msg_$REQUEST_ID","type":"message","role":"assistant","model":"claude-sonnet-5","content":[],"stop_reason":null,"stop_sequence":null,"usage":{"input_tokens":3,"output_tokens":1}}}', + 'event: content_block_start\ndata: {"type":"content_block_start","index":0,"content_block":{"type":"text","text":""}}', + 'event: content_block_delta\ndata: {"type":"content_block_delta","index":0,"delta":{"type":"text_delta","text":"scripted response"}}', + 'event: content_block_stop\ndata: {"type":"content_block_stop","index":0}', + 'event: message_delta\ndata: {"type":"message_delta","delta":{"stop_reason":"end_turn","stop_sequence":null},"usage":{"input_tokens":3,"cache_read_input_tokens":25490,"cache_creation_input_tokens":10553,"output_tokens":12}}', + 'data: {"type":"message_stop"}', + ), + ), + expected=ExactExpected( + spend=0.04741075, + input_cost=0.04723075, + output_cost=0.00018, + prompt_tokens=36046, + completion_tokens=12, + cache_read_cost=0.007647, + cache_creation_cost=0.03957375, + breakdown_persisted=False, + cost_header=False, + ), +) + CLAUDE_SONNET_5_CACHE_CREATION_1H_ABOVE_200K: Final = CostTrackingTestCase( name="claude-sonnet-5-cache_creation_1h_above_200k", covers="quota_management.spend_tracking.cost_matrix.logs_cost", @@ -6997,6 +7037,7 @@ CASES: Final[tuple[CostTrackingTestCase, ...]] = ( CLAUDE_HAIKU_4_5_MESSAGES_INPUT_TEXT, CLAUDE_SONNET_5_PASSTHROUGH_MESSAGES, CLAUDE_SONNET_5_PASSTHROUGH_MESSAGES_CACHE_READ, + CLAUDE_SONNET_5_PASSTHROUGH_MESSAGES_STREAM_CACHE_WRITE_ON_MESSAGE_DELTA, CLAUDE_SONNET_5_CACHE_CREATION_1H_ABOVE_200K, CLAUDE_OPUS_5_5_INPUT_TEXT, CLAUDE_SONNET_5_5_INPUT_TEXT, diff --git a/tests/integration/cost_calculation/cases/fireworks_ai.py b/tests/integration/cost_calculation/cases/fireworks_ai.py index bd740cd8166..6e6c2c9abbf 100644 --- a/tests/integration/cost_calculation/cases/fireworks_ai.py +++ b/tests/integration/cost_calculation/cases/fireworks_ai.py @@ -1097,6 +1097,138 @@ FIREWORKS_AI_ACCOUNTS_FIREWORKS_MODELS_KIMI_K3_CACHE_READ: Final = CostTrackingT ), ) +FIREWORKS_AI_ACCOUNTS_FIREWORKS_MODELS_KIMI_K3_MESSAGES_STREAM_CACHE_READ: Final = CostTrackingTestCase( + name="fireworks_ai-accounts-fireworks-models-kimi-k3-messages_stream_cache_read", + covers="quota_management.spend_tracking.scripted_wire.logs_cost", + model="fireworks_ai/accounts/fireworks/models/kimi-k3", + endpoint="/v1/messages", + request={ + "model": "$MODEL", + "max_tokens": 380, + "messages": [{"role": "user", "content": "3e7a1c5d9b02 summarize the attached material in one line"}], + "stream": True, + }, + response=SseResponse( + content_type="text/event-stream", + frames=sse_frames( + { + "id": "chatcmpl-$REQUEST_ID", + "object": "chat.completion.chunk", + "created": 1789788231, + "model": "accounts/fireworks/models/kimi-k3", + "choices": [{"index": 0, "delta": {"role": "assistant"}, "finish_reason": None}], + "usage": None, + }, + { + "id": "chatcmpl-$REQUEST_ID", + "object": "chat.completion.chunk", + "created": 1789788231, + "model": "accounts/fireworks/models/kimi-k3", + "choices": [ + { + "index": 0, + "delta": {"role": "assistant", "content": "scripted answer 3e7a1c5d9b02"}, + "finish_reason": None, + } + ], + "usage": None, + }, + { + "id": "chatcmpl-$REQUEST_ID", + "object": "chat.completion.chunk", + "created": 1789788231, + "model": "accounts/fireworks/models/kimi-k3", + "choices": [{"index": 0, "delta": {}, "finish_reason": "stop"}], + "usage": None, + }, + { + "id": "chatcmpl-$REQUEST_ID", + "object": "chat.completion.chunk", + "created": 1789788231, + "model": "accounts/fireworks/models/kimi-k3", + "choices": [], + "usage": { + "prompt_tokens": 12928, + "completion_tokens": 380, + "total_tokens": 13308, + "prompt_tokens_details": {"cached_tokens": 12288}, + }, + }, + done=True, + ), + ), + expected=ExactExpected( + spend=0.00207128, input_cost=0.00112128, output_cost=0.00095, prompt_tokens=12928, completion_tokens=380 + ), +) + +FIREWORKS_AI_ACCOUNTS_FIREWORKS_MODELS_KIMI_K3_MESSAGES_STREAM_CACHE_READ_CHOICE_BEARING_USAGE: Final = ( + CostTrackingTestCase( + name="fireworks_ai-accounts-fireworks-models-kimi-k3-messages_stream_cache_read_choice_bearing_usage", + covers="quota_management.spend_tracking.scripted_wire.logs_cost", + model="fireworks_ai/accounts/fireworks/models/kimi-k3", + endpoint="/v1/messages", + request={ + "model": "$MODEL", + "max_tokens": 380, + "messages": [{"role": "user", "content": "9b4d2f8e1a67 summarize the attached material in one line"}], + "stream": True, + }, + response=SseResponse( + content_type="text/event-stream", + frames=sse_frames( + { + "id": "chatcmpl-$REQUEST_ID", + "object": "chat.completion.chunk", + "created": 1789788231, + "model": "accounts/fireworks/models/kimi-k3", + "choices": [{"index": 0, "delta": {"role": "assistant"}, "finish_reason": None}], + "usage": None, + }, + { + "id": "chatcmpl-$REQUEST_ID", + "object": "chat.completion.chunk", + "created": 1789788231, + "model": "accounts/fireworks/models/kimi-k3", + "choices": [ + { + "index": 0, + "delta": {"role": "assistant", "content": "scripted answer 9b4d2f8e1a67"}, + "finish_reason": None, + } + ], + "usage": None, + }, + { + "id": "chatcmpl-$REQUEST_ID", + "object": "chat.completion.chunk", + "created": 1789788231, + "model": "accounts/fireworks/models/kimi-k3", + "choices": [{"index": 0, "delta": {}, "finish_reason": "stop"}], + "usage": None, + }, + { + "id": "chatcmpl-$REQUEST_ID", + "object": "chat.completion.chunk", + "created": 1789788231, + "model": "accounts/fireworks/models/kimi-k3", + "choices": [{"index": 0, "delta": {}}], + "usage": { + "prompt_tokens": 12928, + "completion_tokens": 380, + "total_tokens": 13308, + "prompt_tokens_details": {"cached_tokens": 12288}, + }, + }, + done=True, + ), + ), + expected=ExactExpected( + spend=0.00207128, input_cost=0.00112128, output_cost=0.00095, prompt_tokens=12928, completion_tokens=380 + ), + ) +) + FIREWORKS_AI_ACCOUNTS_FIREWORKS_MODELS_KIMI_K3_STREAM: Final = CostTrackingTestCase( name="fireworks_ai-accounts-fireworks-models-kimi-k3-stream", covers="quota_management.spend_tracking.scripted_wire.logs_cost", @@ -2971,6 +3103,8 @@ CASES: Final[tuple[CostTrackingTestCase, ...]] = ( FIREWORKS_AI_ACCOUNTS_FIREWORKS_MODELS_DEEPSEEK_V4P1_FLASH_STREAM_FULL_USAGE, FIREWORKS_AI_ACCOUNTS_FIREWORKS_MODELS_KIMI_K3_INPUT_TEXT, FIREWORKS_AI_ACCOUNTS_FIREWORKS_MODELS_KIMI_K3_CACHE_READ, + FIREWORKS_AI_ACCOUNTS_FIREWORKS_MODELS_KIMI_K3_MESSAGES_STREAM_CACHE_READ, + FIREWORKS_AI_ACCOUNTS_FIREWORKS_MODELS_KIMI_K3_MESSAGES_STREAM_CACHE_READ_CHOICE_BEARING_USAGE, FIREWORKS_AI_ACCOUNTS_FIREWORKS_MODELS_KIMI_K3_STREAM, FIREWORKS_AI_ACCOUNTS_FIREWORKS_MODELS_KIMI_K3_STREAM_NO_USAGE, FIREWORKS_AI_ACCOUNTS_FIREWORKS_MODELS_KIMI_K3_STREAM_NO_USAGE_TOOL_CALL, diff --git a/tests/integration/cost_calculation/cases/openai.py b/tests/integration/cost_calculation/cases/openai.py index 6ad0ff95615..fe4b736959c 100644 --- a/tests/integration/cost_calculation/cases/openai.py +++ b/tests/integration/cost_calculation/cases/openai.py @@ -1936,6 +1936,82 @@ GPT_5_4_MINI_CACHE_READ: Final = CostTrackingTestCase( ), ) +GPT_5_4_MINI_STREAM_CACHE_READ_CHOICE_BEARING_USAGE: Final = CostTrackingTestCase( + name="gpt-5.4-mini-stream_cache_read_choice_bearing_usage", + covers="quota_management.spend_tracking.scripted_wire.logs_cost", + model="gpt-5.4-mini", + request={ + "model": "$MODEL", + "messages": [ + { + "role": "system", + "content": [ + { + "type": "text", + "text": "You are a deterministic pricing-harness assistant. Keep answers to a single short line.", + } + ], + }, + { + "role": "user", + "content": [ + { + "type": "text", + "text": "6c1f0b2a9d4e summarize the attached material in one line and name the city weather", + } + ], + }, + ], + "stream": True, + "stream_options": {"include_usage": True}, + "allowed_openai_params": [], + }, + response=SseResponse( + content_type="text/event-stream", + frames=sse_frames( + { + "id": "chatcmpl-$REQUEST_ID", + "object": "chat.completion.chunk", + "created": 1789788260, + "model": "gpt-5.4-mini", + "choices": [ + { + "index": 0, + "delta": {"role": "assistant", "content": "scripted answer 6c1f0b2a9d4e"}, + "finish_reason": None, + } + ], + "usage": None, + }, + { + "id": "chatcmpl-$REQUEST_ID", + "object": "chat.completion.chunk", + "created": 1789788260, + "model": "gpt-5.4-mini", + "choices": [{"index": 0, "delta": {}, "finish_reason": "stop"}], + "usage": None, + }, + { + "id": "chatcmpl-$REQUEST_ID", + "object": "chat.completion.chunk", + "created": 1789788260, + "model": "gpt-5.4-mini", + "choices": [{"index": 0, "delta": {}}], + "usage": { + "prompt_tokens": 12928, + "completion_tokens": 380, + "total_tokens": 13308, + "prompt_tokens_details": {"cached_tokens": 12288}, + }, + }, + done=True, + ), + ), + expected=ExactExpected( + spend=0.00171808, input_cost=0.00065408, output_cost=0.001064, prompt_tokens=12928, completion_tokens=380 + ), +) + GPT_5_4_MINI_AUDIO_INPUT: Final = CostTrackingTestCase( name="gpt-5.4-mini-audio_input", covers="quota_management.spend_tracking.cost_matrix.logs_cost", @@ -11161,6 +11237,7 @@ CASES: Final[tuple[CostTrackingTestCase, ...]] = ( GPT_5_3_CODEX_STREAM_FULL_USAGE, GPT_5_4_MINI_INPUT_TEXT, GPT_5_4_MINI_CACHE_READ, + GPT_5_4_MINI_STREAM_CACHE_READ_CHOICE_BEARING_USAGE, GPT_5_4_MINI_AUDIO_INPUT, GPT_5_4_MINI_AUDIO_OUTPUT, GPT_5_4_MINI_REASONING, diff --git a/tests/integration/cost_calculation/test_bridged_previous_response_id_spend.py b/tests/integration/cost_calculation/test_bridged_previous_response_id_spend.py new file mode 100644 index 00000000000..2a8c63fff3d --- /dev/null +++ b/tests/integration/cost_calculation/test_bridged_previous_response_id_spend.py @@ -0,0 +1,123 @@ +"""A /v1/responses deployment bridged to chat completions (`use_chat_completions_api: true`) answers a +`previous_response_id` follow-up with the prior turn replayed from LiteLLM_SpendLogs, and bills the follow-up +from its own usage (LIT-9604, #46398)""" + +import json +import uuid +from collections.abc import Iterator +from dataclasses import dataclass +from typing import Final + +import pytest +from integration._support.client import Gateway, eventually, gateway_from_environment +from integration._support.database import read_rows +from integration._support.process import OwnedProxy, owned_proxy_process +from integration._support.wire import Reply, Request, Wire, wire_server +from integration._support.workers import cost_map_environment +from integration.cost_calculation.cost_tracking_case import COST_MAP +from pydantic import JsonValue, TypeAdapter + +pytestmark: Final = pytest.mark.timeout(240) + +_JSON_OBJECT: Final = TypeAdapter(dict[str, JsonValue]) +_MESSAGES: Final = TypeAdapter(list[dict[str, JsonValue]]) +_MODEL: Final = "gpt-5.6" +_PROMPT_TOKENS: Final = 1840 +_COMPLETION_TOKENS: Final = 412 +_FIRST_TURN: Final = "My favorite color is yellow. Remember it" +_SECOND_TURN: Final = "What is my favorite color?" +_ANSWER: Final = "scripted answer yellow" + + +@dataclass(frozen=True, slots=True) +class _Rig: + wire: Wire + proxy: OwnedProxy + + @property + def gateway(self) -> Gateway: + return self.proxy.gateway + + +def _respond(request: Request) -> Reply: + assert request.target.endswith("/chat/completions"), request.target + body: Final = _JSON_OBJECT.validate_json(request.body) + assert body["model"] == _MODEL, body + completion: Final = { + "id": f"chatcmpl-{uuid.uuid4().hex}", + "object": "chat.completion", + "created": 1, + "model": _MODEL, + "choices": [{"index": 0, "message": {"role": "assistant", "content": _ANSWER}, "finish_reason": "stop"}], + "usage": { + "prompt_tokens": _PROMPT_TOKENS, + "completion_tokens": _COMPLETION_TOKENS, + "total_tokens": _PROMPT_TOKENS + _COMPLETION_TOKENS, + }, + } + return Reply(body=json.dumps(completion).encode()) + + +@pytest.fixture(scope="module") +def rig(tmp_path_factory: pytest.TempPathFactory) -> Iterator[_Rig]: + directory: Final = tmp_path_factory.mktemp("bridged-previous-response-id-rig") + with gateway_from_environment() as gateway, wire_server(_respond) as wire: + overrides: Final = {**cost_map_environment(gateway.upstream_url), "STORE_PROMPTS_IN_SPEND_LOGS": "true"} + with owned_proxy_process( + gateway, directory, overrides, remove_environment=("LITELLM_LOCAL_MODEL_COST_MAP",) + ) as owned: + yield _Rig(wire, owned) + + +def _spend_rows(model: str) -> list[dict[str, JsonValue]]: + return read_rows( + 'SELECT request_id, spend, status FROM "LiteLLM_SpendLogs" WHERE model_group = %s ORDER BY "startTime"', + (model,), + ) + + +def _upstream_body(request: Request) -> dict[str, JsonValue]: + body: Final = _JSON_OBJECT.validate_json(request.body) + assert "previous_response_id" not in body, body + return body + + +def _upstream_messages(request: Request) -> list[dict[str, JsonValue]]: + return _MESSAGES.validate_python(_upstream_body(request)["messages"]) + + +def test_bridged_previous_response_id_replays_the_prior_turn_and_bills_the_follow_up(rig: _Rig) -> None: + with rig.gateway.scenario() as scenario: + model: Final = scenario.model( + model=f"openai/{_MODEL}", + api_base=rig.wire.url, + api_key="sk-scripted-provider", + use_chat_completions_api=True, + ) + first: Final = rig.gateway.request("POST", "/v1/responses", {"model": model, "input": _FIRST_TURN}) + assert first.status_code == 200, first.text + first_id: Final = _JSON_OBJECT.validate_json(first.content)["id"] + assert isinstance(first_id, str) and first_id.startswith("resp_"), first_id + first_rows: Final = eventually(lambda: _spend_rows(model), lambda rows: len(rows) == 1, seconds=70) + assert first_rows[0]["status"] == "success", first_rows + + second: Final = rig.gateway.request( + "POST", "/v1/responses", {"model": model, "input": _SECOND_TURN, "previous_response_id": first_id} + ) + assert second.status_code == 200, second.text + sent: Final = tuple(request for request in rig.wire.drain() if request.method == "POST") + assert len(sent) == 2, [request.target for request in sent] + assert _upstream_messages(sent[0]) == [{"role": "user", "content": _FIRST_TURN}], sent[0].body + assert _upstream_messages(sent[1]) == [ + {"role": "user", "content": _FIRST_TURN}, + {"role": "assistant", "content": _ANSWER}, + {"role": "user", "content": _SECOND_TURN}, + ], sent[1].body + + rows: Final = eventually(lambda: _spend_rows(model), lambda values: len(values) == 2, seconds=70) + assert [row["status"] for row in rows] == ["success", "success"], rows + rates: Final = COST_MAP[_MODEL] + assert rates.input_cost_per_token is not None and rates.output_cost_per_token is not None, rates + assert float(rows[1]["spend"]) == pytest.approx( + _PROMPT_TOKENS * rates.input_cost_per_token + _COMPLETION_TOKENS * rates.output_cost_per_token + ), rows