From 807541291d83f0c1c068f1f3e9505ffd5a44e2fe Mon Sep 17 00:00:00 2001 From: kerry Date: Sun, 20 Sep 2026 01:25:28 +0000 Subject: [PATCH 1/3] test(integration): batch and realtime cost cases Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- tests/integration/_support/upstream.py | 80 ++++- tests/integration/contracts.json | 24 ++ .../cost_calculation/assertions.py | 141 +++++++++ .../integration/cost_calculation/conftest.py | 8 +- .../cost_calculation/cost_tracking_case.py | 192 +++++++++++- .../cost_calculation/cost_tracking_cases.json | 218 +++++++++++++ .../test_batch_realtime_cost.py | 288 ++++++++++++++++++ .../cost_calculation/test_cost_tracking.py | 150 +-------- 8 files changed, 937 insertions(+), 164 deletions(-) create mode 100644 tests/integration/cost_calculation/assertions.py create mode 100644 tests/integration/cost_calculation/test_batch_realtime_cost.py diff --git a/tests/integration/_support/upstream.py b/tests/integration/_support/upstream.py index 759df7003e4..e9c50ea7966 100644 --- a/tests/integration/_support/upstream.py +++ b/tests/integration/_support/upstream.py @@ -3,35 +3,38 @@ from __future__ import annotations import argparse import asyncio import base64 -from collections import deque -from collections.abc import AsyncIterator, Mapping import json -from dataclasses import dataclass, field import os -from pathlib import Path -from queue import SimpleQueue import struct -from typing import Final, cast import uuid import zlib +from collections import deque +from collections.abc import AsyncIterator, Mapping +from dataclasses import dataclass, field +from pathlib import Path +from queue import SimpleQueue +from typing import Final, cast import httpx import uvicorn -from pydantic import BaseModel, JsonValue, TypeAdapter, ValidationError -from starlette.applications import Starlette -from starlette.requests import Request -from starlette.responses import JSONResponse, Response, StreamingResponse -from starlette.routing import Route - from _fake_openai_endpoint_server import chat_completions, completions, embeddings, health, moderations from integration.cost_calculation.cost_tracking_case import ( BinaryResponse, EventStreamEvent, EventStreamResponse, JsonResponse, + RealtimeResponse, + RoutedResponse, SseResponse, StoredResponse, + TextResponse, ) +from pydantic import BaseModel, JsonValue, TypeAdapter, ValidationError +from starlette.applications import Starlette +from starlette.requests import Request +from starlette.responses import JSONResponse, Response, StreamingResponse +from starlette.routing import Route, WebSocketRoute +from starlette.websockets import WebSocket JSON_OBJECT: Final = TypeAdapter(dict[str, JsonValue]) CASES_FILE: Final = Path(__file__).resolve().parents[1] / "cost_calculation" / "cost_tracking_cases.json" @@ -211,8 +214,53 @@ class Provider: response: Final = self.scenario_store.get(scenario_id) if response is None: return JSONResponse({"error": "Unknown scenario"}, status_code=404) + if isinstance(response, RoutedResponse): + route_key: Final = f"{request.method} /{'/'.join(segments[1:])}" + route: Final = next( + ( + candidate + for key, candidate in response.routes.items() + if key.replace("$REQUEST_ID", scenario_id) == route_key + ), + None, + ) + if route is None: + return JSONResponse({"error": "Unknown scripted route"}, status_code=404) + return self._response(route, scenario_id) return self._response(response, scenario_id) + async def realtime(self, websocket: WebSocket) -> None: + scenario_id: Final = websocket.headers.get("authorization", "").removeprefix("Bearer ") + response: Final = self.scenario_store.get(scenario_id) + if not isinstance(response, RealtimeResponse): + await websocket.close(code=4404) + return + await websocket.accept() + model: Final = websocket.query_params.get("model", "") + await websocket.send_json( + { + "type": "session.created", + "session": { + "id": f"sess_{scenario_id}", + "model": response.session_model if response.session_model is not None else model, + }, + } + ) + event_index: Final = iter(response.events) + async for message in websocket.iter_json(): + payload: Final = JSON_OBJECT.validate_python(message) + if payload.get("type") != "response.create": + continue + event: Final = next(event_index, None) + if event is None: + continue + rendered: Final = JSON_OBJECT.validate_json( + json.dumps(event, separators=(",", ":")) + .replace("$REQUEST_ID", scenario_id) + .replace("$UNIQUE_ID", f"{scenario_id}-{uuid.uuid4().hex[:8]}") + ) + await websocket.send_json(rendered) + @staticmethod def _response(response: StoredResponse, scenario_id: str) -> Response: unique_id: Final = f"{scenario_id}-{uuid.uuid4().hex[:8]}" @@ -232,6 +280,12 @@ class Provider: content=b"\x00" * response.length, media_type=response.content_type, ) + case TextResponse(): + return Response( + content=response.body.replace("$REQUEST_ID", scenario_id).encode(), + media_type=response.content_type, + status_code=response.status, + ) case SseResponse(): if response.frame_delay_ms > 0: async def stream() -> AsyncIterator[bytes]: @@ -285,6 +339,8 @@ class Provider: Route("/v1/embeddings", embeddings, methods=["POST"]), Route("/v1/moderations", moderations, methods=["POST"]), Route("/{path:path}", self.scripted, methods=["POST"]), + Route("/{path:path}", self.scripted, methods=["GET"]), + WebSocketRoute("/v1/realtime", self.realtime), ] ) diff --git a/tests/integration/contracts.json b/tests/integration/contracts.json index fe7c808790d..195c7ba4ab4 100644 --- a/tests/integration/contracts.json +++ b/tests/integration/contracts.json @@ -226,6 +226,30 @@ "tests/integration/cost_calculation/test_cost_tracking.py::test_case_bills_expected_cost[anthropic.claude-sonnet-5-v1:0-input_text]": [ "quota_management.spend_tracking.cost_matrix.logs_cost" ], + "tests/integration/cost_calculation/test_batch_realtime_cost.py::test_batch_costs[gpt-5.6-batch-halved_rates_when_map_has_no_batch_keys]": [ + "quota_management.spend_tracking.batch_costs.fallback_rates" + ], + "tests/integration/cost_calculation/test_batch_realtime_cost.py::test_batch_costs[gpt-5.6-batch-cached_input_halved]": [ + "quota_management.spend_tracking.batch_costs.cached_input" + ], + "tests/integration/cost_calculation/test_batch_realtime_cost.py::test_batch_costs[gpt-5.4-batch-explicit_batch_rates_bill_cached_at_batch_input_rate]": [ + "quota_management.spend_tracking.batch_costs.explicit_rates" + ], + "tests/integration/cost_calculation/test_batch_realtime_cost.py::test_batch_costs[gpt-5.6-batch-all_requests_failed_zero_spend]": [ + "quota_management.spend_tracking.batch_costs.failed_requests" + ], + "tests/integration/cost_calculation/test_batch_realtime_cost.py::test_realtime_costs[gpt-realtime-mini-2025-12-15-realtime-single_turn_text_audio_cached]": [ + "quota_management.spend_tracking.realtime_costs.single_turn" + ], + "tests/integration/cost_calculation/test_batch_realtime_cost.py::test_realtime_costs[gpt-realtime-mini-2025-12-15-realtime-two_turns_summed_into_one_row]": [ + "quota_management.spend_tracking.realtime_costs.multiple_turns" + ], + "tests/integration/cost_calculation/test_batch_realtime_cost.py::test_realtime_costs[gpt-realtime-mini-2025-12-15-realtime-priced_from_session_created_model]": [ + "quota_management.spend_tracking.realtime_costs.session_model" + ], + "tests/integration/cost_calculation/test_batch_realtime_cost.py::test_realtime_no_turn_probe": [ + "quota_management.spend_tracking.realtime_costs.no_turn_probe" + ], "tests/integration/cost_calculation/test_cost_tracking.py::test_case_bills_expected_cost[anthropic.claude-sonnet-5-v1:0-cache_read]": [ "quota_management.spend_tracking.cost_matrix.logs_cost" ], diff --git a/tests/integration/cost_calculation/assertions.py b/tests/integration/cost_calculation/assertions.py new file mode 100644 index 00000000000..b0b9057dd9d --- /dev/null +++ b/tests/integration/cost_calculation/assertions.py @@ -0,0 +1,141 @@ +from __future__ import annotations + +import httpx +from integration.cost_calculation.conftest import ( + CostBreakdown, + CostRow, + approx_equal, + assert_total_is_sum_of_components, +) +from integration.cost_calculation.cost_tracking_case import ExactExpected, RecountExpected + + +def assert_breakdown( + case_name: str, + response_content_type: str, + expected: ExactExpected, + breakdown: CostBreakdown, + response: httpx.Response, +) -> None: + assert breakdown.input_cost is not None and approx_equal(breakdown.input_cost, expected.input_cost), ( + f"{case_name}: input_cost {breakdown.input_cost} != expected {expected.input_cost}" + ) + assert breakdown.output_cost is not None and approx_equal(breakdown.output_cost, expected.output_cost), ( + f"{case_name}: output_cost {breakdown.output_cost} != expected {expected.output_cost}" + ) + for field, header_name, actual_component, expected_component in ( + ( + "cache_read_cost", + "x-litellm-response-cost-cache-read", + breakdown.cache_read_cost, + expected.cache_read_cost, + ), + ( + "cache_creation_cost", + "x-litellm-response-cost-cache-creation", + breakdown.cache_creation_cost, + expected.cache_creation_cost, + ), + ( + "reasoning_cost", + "x-litellm-response-cost-reasoning", + breakdown.reasoning_cost, + expected.reasoning_cost, + ), + ( + "tool_usage_cost", + "x-litellm-response-cost-tool-usage", + breakdown.tool_usage_cost, + expected.tool_usage_cost, + ), + ): + if expected_component is None: + continue + omitted_component_allowed: bool = expected_component == 0.0 + assert (actual_component is None and omitted_component_allowed) or ( + actual_component is not None and approx_equal(actual_component, expected_component) + ), f"{case_name}: {field} {actual_component} != expected {expected_component}" + if expected.cost_header and response_content_type == "application/json": + header: str | None = response.headers.get(header_name) + assert (header is None and omitted_component_allowed) or ( + header is not None and approx_equal(float(header), expected_component) + ), f"{case_name}: {header_name} {header} != expected {expected_component}" + if expected.cost_header and response_content_type == "application/json" and any( + component is not None + for component in ( + expected.cache_read_cost, + expected.cache_creation_cost, + expected.reasoning_cost, + expected.tool_usage_cost, + ) + ): + input_header: str | None = response.headers.get("x-litellm-response-cost-input") + output_header: str | None = response.headers.get("x-litellm-response-cost-output") + expected_input_header: float = expected.input_cost - ( + expected.cache_read_cost or 0.0 + ) - (expected.cache_creation_cost or 0.0) + assert input_header is not None and approx_equal(float(input_header), expected_input_header), ( + f"{case_name}: x-litellm-response-cost-input {input_header} != expected {expected_input_header}" + ) + assert output_header is not None and approx_equal(float(output_header), expected.output_cost), ( + f"{case_name}: x-litellm-response-cost-output {output_header} != expected {expected.output_cost}" + ) + + +def assert_exact( + case_name: str, + response_content_type: str, + expected: ExactExpected, + row: CostRow, + response: httpx.Response, +) -> None: + assert row.spend is not None and approx_equal(row.spend, expected.spend), ( + f"{case_name}: spend {row.spend} != expected {expected.spend} " + f"(breakdown {row.breakdown.model_dump() if row.breakdown is not None else None})" + ) + breakdown: CostBreakdown | None = row.breakdown + if expected.breakdown_persisted: + assert breakdown is not None, f"{case_name}: no cost_breakdown persisted" + if breakdown is not None: + assert_breakdown(case_name, response_content_type, expected, breakdown, response) + assert row.prompt_tokens == expected.prompt_tokens, ( + f"{case_name}: prompt_tokens {row.prompt_tokens} != expected {expected.prompt_tokens}" + ) + assert row.completion_tokens == expected.completion_tokens, ( + f"{case_name}: completion_tokens {row.completion_tokens} != expected {expected.completion_tokens}" + ) + if breakdown is not None: + assert_total_is_sum_of_components(row, breakdown, case_name) + + +def assert_recount(case_name: str, expected: RecountExpected, row: CostRow) -> None: + assert row.prompt_tokens is not None and row.prompt_tokens > 0, ( + f"{case_name}: recount case counted no input tokens: prompt_tokens={row.prompt_tokens}" + ) + assert row.completion_tokens is not None and row.completion_tokens > 0, ( + f"{case_name}: recount case counted no output tokens: completion_tokens={row.completion_tokens}" + ) + if expected.prompt_tokens is not None: + assert row.prompt_tokens == expected.prompt_tokens, ( + f"{case_name}: prompt_tokens {row.prompt_tokens} != pinned {expected.prompt_tokens}" + ) + if expected.completion_tokens is not None: + assert row.completion_tokens == expected.completion_tokens, ( + f"{case_name}: completion_tokens {row.completion_tokens} != pinned {expected.completion_tokens}" + ) + if expected.min_completion_tokens is not None: + assert row.completion_tokens >= expected.min_completion_tokens, ( + f"{case_name}: completion_tokens {row.completion_tokens} < minimum {expected.min_completion_tokens}" + ) + if expected.max_completion_tokens is not None: + assert row.completion_tokens <= expected.max_completion_tokens, ( + f"{case_name}: completion_tokens {row.completion_tokens} > maximum {expected.max_completion_tokens}" + ) + recount: float = row.prompt_tokens * expected.recount.input_cost_per_token + ( + row.completion_tokens * expected.recount.output_cost_per_token + ) + assert row.spend is not None and approx_equal(row.spend, recount), ( + f"{case_name}: spend {row.spend} != recount {recount} at map rates" + ) + assert row.breakdown is not None, f"{case_name}: no cost_breakdown persisted" + assert_total_is_sum_of_components(row, row.breakdown, case_name) diff --git a/tests/integration/cost_calculation/conftest.py b/tests/integration/cost_calculation/conftest.py index bc68419f01d..9c6ffe3f7d0 100644 --- a/tests/integration/cost_calculation/conftest.py +++ b/tests/integration/cost_calculation/conftest.py @@ -10,12 +10,11 @@ from typing import Final from cryptography.hazmat.primitives import serialization from cryptography.hazmat.primitives.asymmetric import rsa -from pydantic import BaseModel, ConfigDict - from integration._support.client import JSON_OBJECT, Scenario, eventually, object_value, string_value from integration._support.database import read_rows from integration._support.upstream import ScenarioHandle, delete_scenario, register_scenario from integration.cost_calculation.cost_tracking_case import CostTrackingTestCase, StoredResponse +from pydantic import BaseModel, ConfigDict class CostBreakdown(BaseModel): @@ -45,6 +44,7 @@ class CostRow(BaseModel): prompt_tokens: int | None = None completion_tokens: int | None = None model_id: str | None = None + call_type: str | None = None metadata: CostMetadata | None = None @property @@ -112,7 +112,7 @@ def poll_cost_row(key: str) -> CostRow: def read() -> CostRow | None: rows: Final = read_rows( - 'SELECT spend, status, metadata, prompt_tokens, completion_tokens, model_id ' + 'SELECT spend, status, metadata, prompt_tokens, completion_tokens, model_id, call_type ' 'FROM "LiteLLM_SpendLogs" WHERE api_key=%s', (digest,), ) @@ -126,7 +126,7 @@ def poll_cost_row(key: str) -> CostRow: def read_rows_now(key: str) -> tuple[CostRow, ...]: digest: Final = sha256(key.encode()).hexdigest() rows: Final = read_rows( - 'SELECT spend, status, metadata, prompt_tokens, completion_tokens, model_id ' + 'SELECT spend, status, metadata, prompt_tokens, completion_tokens, model_id, call_type ' 'FROM "LiteLLM_SpendLogs" WHERE api_key=%s ORDER BY "startTime"', (digest,), ) diff --git a/tests/integration/cost_calculation/cost_tracking_case.py b/tests/integration/cost_calculation/cost_tracking_case.py index 56a7cefb443..0e75b0c23b1 100644 --- a/tests/integration/cost_calculation/cost_tracking_case.py +++ b/tests/integration/cost_calculation/cost_tracking_case.py @@ -5,7 +5,7 @@ from pathlib import Path from types import MappingProxyType from typing import Annotated, Final, Literal, TypeAlias -from pydantic import BaseModel, ConfigDict, Field, JsonValue +from pydantic import BaseModel, ConfigDict, Field, JsonValue, field_validator, model_validator CASES_PATH: Final = Path(__file__).resolve().parent / "cost_tracking_cases.json" @@ -44,6 +44,8 @@ class CostMapEntry(BaseModel): supports_function_calling: bool | None = None input_cost_per_token: float | None = None output_cost_per_token: float | None = None + input_cost_per_token_batches: float | None = None + output_cost_per_token_batches: float | None = None input_cost_per_token_above_128k_tokens: float | None = None output_cost_per_token_above_128k_tokens: float | None = None cache_read_input_token_cost: float | None = None @@ -54,6 +56,7 @@ class CostMapEntry(BaseModel): cache_creation_input_token_cost_above_200k_tokens: float | None = None input_cost_per_token_above_200k_tokens: float | None = None output_cost_per_token_above_200k_tokens: float | None = None + cache_read_input_audio_token_cost: float | None = None tiered_pricing: tuple[TieredPrice, ...] | None = None output_cost_per_reasoning_token: float | None = None input_cost_per_audio_token: float | None = None @@ -141,8 +144,31 @@ class BinaryResponse(BaseModel): length: int +class TextResponse(BaseModel): + model_config = ConfigDict(frozen=True, extra="forbid") + + content_type: Literal["application/jsonl"] + body: str + status: int = 200 + + +class RoutedResponse(BaseModel): + model_config = ConfigDict(frozen=True, extra="forbid") + + content_type: Literal["application/x-routed"] + routes: dict[str, JsonResponse | TextResponse] + + +class RealtimeResponse(BaseModel): + model_config = ConfigDict(frozen=True, extra="forbid") + + content_type: Literal["application/x-realtime"] + events: tuple[dict[str, JsonValue], ...] + session_model: str | None = None + + StoredResponse: TypeAlias = Annotated[ - JsonResponse | SseResponse | EventStreamResponse | BinaryResponse, + JsonResponse | SseResponse | EventStreamResponse | BinaryResponse | RoutedResponse | RealtimeResponse, Field(discriminator="content_type"), ] @@ -283,11 +309,156 @@ class CostTrackingTestCase(BaseModel): return isinstance(usage, dict) and isinstance(usage.get("cost"), (int, float)) +class BatchOutputLine(BaseModel): + model_config = ConfigDict(frozen=True, extra="forbid") + + status_code: int + prompt_tokens: int | None = None + completion_tokens: int | None = None + cached_tokens: int | None = None + + @field_validator("status_code") + @classmethod + def validate_status_code(cls, value: int) -> int: + if value != 200 and not 400 <= value <= 499: + raise ValueError("status_code must be 200 or a 4xx status") + return value + + def render(self, index: int, model: str, request_id: str) -> dict[str, JsonValue]: + if self.status_code != 200: + return { + "id": f"batch_req_{index}", + "custom_id": f"r{index}", + "response": None, + "error": {"code": "bad_request", "message": "failed"}, + } + assert self.prompt_tokens is not None + assert self.completion_tokens is not None + usage: dict[str, JsonValue] = { + "prompt_tokens": self.prompt_tokens, + "completion_tokens": self.completion_tokens, + "total_tokens": self.prompt_tokens + self.completion_tokens, + } + if self.cached_tokens is not None: + usage["prompt_tokens_details"] = {"cached_tokens": self.cached_tokens} + return { + "id": f"batch_req_{index}", + "custom_id": f"r{index}", + "response": { + "status_code": 200, + "request_id": f"{request_id}-{index}", + "body": { + "id": f"chatcmpl-{request_id}-{index}", + "object": "chat.completion", + "model": model, + "choices": [ + { + "index": 0, + "message": {"role": "assistant", "content": "ok"}, + "finish_reason": "stop", + } + ], + "usage": usage, + }, + }, + "error": None, + } + + +class BatchCostCase(BaseModel): + model_config = ConfigDict(frozen=True, extra="forbid") + + name: str + covers: str + model: str + litellm_model: str + output_lines: tuple[BatchOutputLine, ...] + expected: ExactExpected + + @property + def request_count(self) -> int: + return len(self.output_lines) or 2 + + @property + def completed_count(self) -> int: + return sum(line.status_code == 200 for line in self.output_lines) + + @property + def failed_count(self) -> int: + return self.request_count - self.completed_count + + +class RealtimeTurn(BaseModel): + model_config = ConfigDict(frozen=True, extra="forbid") + + input_tokens: int + output_tokens: int + input_text_tokens: int + input_audio_tokens: int + input_cached_tokens: int + output_text_tokens: int + output_audio_tokens: int + + @model_validator(mode="after") + def validate_token_totals(self) -> RealtimeTurn: + if self.input_text_tokens + self.input_audio_tokens != self.input_tokens: + raise ValueError("input text and audio tokens must equal input_tokens") + if self.output_text_tokens + self.output_audio_tokens != self.output_tokens: + raise ValueError("output text and audio tokens must equal output_tokens") + if self.input_cached_tokens > self.input_text_tokens: + raise ValueError("input_cached_tokens must not exceed input_text_tokens") + return self + + def render(self, index: int, request_id: str) -> dict[str, JsonValue]: + return { + "type": "response.done", + "event_id": f"evt_{request_id}_{index}", + "response": { + "id": f"resp_{request_id}_{index}", + "object": "realtime.response", + "status": "completed", + "output": [], + "usage": { + "total_tokens": self.input_tokens + self.output_tokens, + "input_tokens": self.input_tokens, + "output_tokens": self.output_tokens, + "input_token_details": { + "text_tokens": self.input_text_tokens, + "audio_tokens": self.input_audio_tokens, + "cached_tokens": self.input_cached_tokens, + "cached_tokens_details": { + "text_tokens": self.input_cached_tokens, + "audio_tokens": 0, + }, + }, + "output_token_details": { + "text_tokens": self.output_text_tokens, + "audio_tokens": self.output_audio_tokens, + }, + }, + }, + } + + +class RealtimeCostCase(BaseModel): + model_config = ConfigDict(frozen=True, extra="forbid") + + name: str + covers: str + model: str + litellm_model: str + turns: tuple[RealtimeTurn, ...] = Field(min_length=1) + session_model: str | None = None + expected: ExactExpected + + class _CasesFile(BaseModel): model_config = ConfigDict(frozen=True, extra="forbid") cost_map: dict[str, CostMapEntry] cases: tuple[CostTrackingTestCase, ...] + batch_cases: tuple[BatchCostCase, ...] = () + realtime_cases: tuple[RealtimeCostCase, ...] = () _PROVIDER_PREFIXES: Final[Mapping[str, str]] = MappingProxyType( @@ -357,18 +528,25 @@ _LITELLM_PARAMS: Final[Mapping[str, Mapping[str, str]]] = MappingProxyType( _LOADED: Final = _CasesFile.model_validate_json(CASES_PATH.read_bytes()) COST_MAP: Final[Mapping[str, CostMapEntry]] = MappingProxyType(dict(_LOADED.cost_map)) CASES: Final[tuple[CostTrackingTestCase, ...]] = _LOADED.cases -_LITELLM_MODELS: Final = tuple(case.litellm_model for case in CASES) +BATCH_CASES: Final[tuple[BatchCostCase, ...]] = _LOADED.batch_cases +REALTIME_CASES: Final[tuple[RealtimeCostCase, ...]] = _LOADED.realtime_cases +_ALL_CASES: Final = CASES + BATCH_CASES + REALTIME_CASES +_LITELLM_MODELS: Final = tuple(case.litellm_model for case in _ALL_CASES) def data_errors() -> tuple[str, ...]: - case_models: Final = frozenset(case.model for case in CASES) - unknown_models: Final = sorted(case.model for case in CASES if case.model not in COST_MAP) + case_models: Final = frozenset(case.model for case in _ALL_CASES) | frozenset( + case.session_model for case in REALTIME_CASES if case.session_model is not None + ) + unknown_models: Final = sorted(model for model in case_models if model not in COST_MAP) missing_cases: Final = sorted(model for model in COST_MAP if model not in case_models) duplicate_names: Final = sorted( - name for name in {case.name for case in CASES} if sum(case.name == name for case in CASES) > 1 + name for name in {case.name for case in _ALL_CASES} if sum(case.name == name for case in _ALL_CASES) > 1 ) input_rates: Final = tuple( - (entry.input_cost_per_token, model) for model, entry in COST_MAP.items() + (entry.input_cost_per_token, model) + for model, entry in COST_MAP.items() + if entry.mode != "realtime" ) shared_input_rates: Final = sorted( f"{rate}: {tuple(model for value, model in input_rates if value == rate)}" diff --git a/tests/integration/cost_calculation/cost_tracking_cases.json b/tests/integration/cost_calculation/cost_tracking_cases.json index 78d00f28381..d8f321890f3 100644 --- a/tests/integration/cost_calculation/cost_tracking_cases.json +++ b/tests/integration/cost_calculation/cost_tracking_cases.json @@ -622,6 +622,35 @@ "litellm_provider": "openai", "mode": "embedding", "input_cost_per_token": 1.3e-07 + }, + "gpt-5.4": { + "litellm_provider": "openai", + "mode": "chat", + "input_cost_per_token": 2.5e-06, + "output_cost_per_token": 1.5e-05, + "cache_read_input_token_cost": 2.5e-07, + "input_cost_per_token_batches": 1.25e-06, + "output_cost_per_token_batches": 7.5e-06 + }, + "gpt-realtime-mini-2025-12-15": { + "litellm_provider": "openai", + "mode": "realtime", + "input_cost_per_token": 6.0e-07, + "output_cost_per_token": 2.4e-06, + "input_cost_per_audio_token": 1.0e-05, + "cache_read_input_token_cost": 6.0e-08, + "cache_read_input_audio_token_cost": 3.0e-07, + "output_cost_per_audio_token": 2.0e-05 + }, + "gpt-realtime-2.1": { + "litellm_provider": "openai", + "mode": "realtime", + "input_cost_per_token": 4.0e-06, + "input_cost_per_audio_token": 3.2e-05, + "cache_read_input_token_cost": 4.0e-07, + "cache_read_input_audio_token_cost": 4.0e-07, + "output_cost_per_token": 2.4e-05, + "output_cost_per_audio_token": 6.4e-05 } }, "cases": [ @@ -30020,5 +30049,194 @@ "max_completion_tokens": 30 } } + ], + "batch_cases": [ + { + "name": "gpt-5.6-batch-halved_rates_when_map_has_no_batch_keys", + "covers": "quota_management.spend_tracking.batch_costs.fallback_rates", + "model": "gpt-5.6", + "litellm_model": "openai/gpt-5.6", + "output_lines": [ + { + "status_code": 200, + "prompt_tokens": 100, + "completion_tokens": 50 + }, + { + "status_code": 200, + "prompt_tokens": 120, + "completion_tokens": 30 + }, + { + "status_code": 400 + } + ], + "expected": { + "spend": 0.0007525, + "input_cost": 0.0001925, + "output_cost": 0.00056, + "prompt_tokens": 220, + "completion_tokens": 80, + "cost_header": false + } + }, + { + "name": "gpt-5.6-batch-cached_input_halved", + "covers": "quota_management.spend_tracking.batch_costs.cached_input", + "model": "gpt-5.6", + "litellm_model": "openai/gpt-5.6", + "output_lines": [ + { + "status_code": 200, + "prompt_tokens": 100, + "completion_tokens": 10, + "cached_tokens": 40 + } + ], + "expected": { + "spend": 0.000126, + "input_cost": 0.000056, + "output_cost": 0.00007, + "prompt_tokens": 100, + "completion_tokens": 10, + "cost_header": false + } + }, + { + "name": "gpt-5.4-batch-explicit_batch_rates_bill_cached_at_batch_input_rate", + "covers": "quota_management.spend_tracking.batch_costs.explicit_rates", + "model": "gpt-5.4", + "litellm_model": "openai/gpt-5.4", + "output_lines": [ + { + "status_code": 200, + "prompt_tokens": 100, + "completion_tokens": 50, + "cached_tokens": 40 + }, + { + "status_code": 200, + "prompt_tokens": 120, + "completion_tokens": 30 + } + ], + "expected": { + "spend": 0.000875, + "input_cost": 0.000275, + "output_cost": 0.0006, + "prompt_tokens": 220, + "completion_tokens": 80, + "cost_header": false + } + }, + { + "name": "gpt-5.6-batch-all_requests_failed_zero_spend", + "covers": "quota_management.spend_tracking.batch_costs.failed_requests", + "model": "gpt-5.6", + "litellm_model": "openai/gpt-5.6", + "output_lines": [ + { + "status_code": 400 + }, + { + "status_code": 400 + } + ], + "expected": { + "spend": 0.0, + "input_cost": 0.0, + "output_cost": 0.0, + "prompt_tokens": 0, + "completion_tokens": 0, + "cost_header": false + } + } + ], + "realtime_cases": [ + { + "name": "gpt-realtime-mini-2025-12-15-realtime-single_turn_text_audio_cached", + "covers": "quota_management.spend_tracking.realtime_costs.single_turn", + "model": "gpt-realtime-mini-2025-12-15", + "litellm_model": "openai/gpt-realtime-mini-2025-12-15", + "turns": [ + { + "input_tokens": 150, + "output_tokens": 100, + "input_text_tokens": 70, + "input_audio_tokens": 80, + "input_cached_tokens": 20, + "output_text_tokens": 40, + "output_audio_tokens": 60 + } + ], + "expected": { + "spend": 0.0021272, + "input_cost": 0.0008312, + "output_cost": 0.001296, + "prompt_tokens": 150, + "completion_tokens": 100, + "cost_header": false + } + }, + { + "name": "gpt-realtime-mini-2025-12-15-realtime-two_turns_summed_into_one_row", + "covers": "quota_management.spend_tracking.realtime_costs.multiple_turns", + "model": "gpt-realtime-mini-2025-12-15", + "litellm_model": "openai/gpt-realtime-mini-2025-12-15", + "turns": [ + { + "input_tokens": 150, + "output_tokens": 100, + "input_text_tokens": 70, + "input_audio_tokens": 80, + "input_cached_tokens": 20, + "output_text_tokens": 40, + "output_audio_tokens": 60 + }, + { + "input_tokens": 100, + "output_tokens": 50, + "input_text_tokens": 100, + "input_audio_tokens": 0, + "input_cached_tokens": 0, + "output_text_tokens": 50, + "output_audio_tokens": 0 + } + ], + "expected": { + "spend": 0.0023072, + "input_cost": 0.0008912, + "output_cost": 0.001416, + "prompt_tokens": 250, + "completion_tokens": 150, + "cost_header": false + } + }, + { + "name": "gpt-realtime-mini-2025-12-15-realtime-priced_from_session_created_model", + "covers": "quota_management.spend_tracking.realtime_costs.session_model", + "model": "gpt-realtime-mini-2025-12-15", + "litellm_model": "openai/gpt-realtime-mini-2025-12-15", + "session_model": "gpt-realtime-2.1", + "turns": [ + { + "input_tokens": 150, + "output_tokens": 100, + "input_text_tokens": 70, + "input_audio_tokens": 80, + "input_cached_tokens": 20, + "output_text_tokens": 40, + "output_audio_tokens": 60 + } + ], + "expected": { + "spend": 0.007568, + "input_cost": 0.002768, + "output_cost": 0.0048, + "prompt_tokens": 150, + "completion_tokens": 100, + "cost_header": false + } + } ] } diff --git a/tests/integration/cost_calculation/test_batch_realtime_cost.py b/tests/integration/cost_calculation/test_batch_realtime_cost.py new file mode 100644 index 00000000000..9db25db32e8 --- /dev/null +++ b/tests/integration/cost_calculation/test_batch_realtime_cost.py @@ -0,0 +1,288 @@ +from __future__ import annotations + +import asyncio +import json +import logging +import os +import time +from hashlib import sha256 +from typing import Final + +import httpx +import pytest +import websockets +from integration._support.client import JSON_OBJECT, Gateway, Scenario, object_value, string_value +from integration._support.upstream import delete_scenario, register_scenario +from integration.cost_calculation.assertions import assert_exact +from integration.cost_calculation.conftest import CostRow, poll_rows, read_rows_now +from integration.cost_calculation.cost_tracking_case import ( + BATCH_CASES, + REALTIME_CASES, + BatchCostCase, + JsonResponse, + RealtimeCostCase, + RealtimeResponse, + RoutedResponse, + TextResponse, +) +from pydantic import JsonValue + + +def _register_deployment( + scenario: Scenario, + litellm_model: str, + response: JsonResponse | TextResponse | RealtimeResponse, + marker: str, + *, + realtime: bool, +) -> tuple[str, str]: + scenario_id: Final = f"cost-{marker}-{sha256(os.urandom(16)).hexdigest()[:12]}" + handle: Final = register_scenario(scenario_id, response) + scenario.cleanups.callback(delete_scenario, handle) + control_url: Final = os.environ["INTEGRATION_UPSTREAM_URL"].rstrip("/") + created: Final = scenario.gateway.post( + "/model/new", + JSON_OBJECT.validate_python( + { + "model_name": f"cost-{marker}-{sha256(scenario_id.encode()).hexdigest()[:12]}", + "litellm_params": { + "model": litellm_model, + "api_key": scenario_id if realtime else "sk-scripted-provider", + "api_base": control_url if realtime else handle.api_base(), + }, + } + ), + ) + identity: Final = string_value(object_value(created["model_info"])["id"]) + scenario.cleanups.callback(scenario.delete_model, identity) + return string_value(created["model_name"]), identity + + +def _batch_response(case: BatchCostCase) -> JsonResponse | RoutedResponse: + request_id: Final = "$REQUEST_ID" + lines: Final = tuple( + json.dumps(line.render(index, case.model, request_id), separators=(",", ":")) + for index, line in enumerate(case.output_lines, start=1) + ) + counts: Final = { + "total": case.request_count, + "completed": case.completed_count, + "failed": case.failed_count, + } + completed: Final = len(case.output_lines) > 0 + batch: Final = { + "id": "batch-$REQUEST_ID", + "object": "batch", + "endpoint": "/v1/chat/completions", + "errors": None, + "input_file_id": "file-in-$REQUEST_ID", + "completion_window": "24h", + "status": "completed" if completed else "completed", + "output_file_id": "file-out-$REQUEST_ID" if completed else None, + "error_file_id": None if completed else "file-err-$REQUEST_ID", + "created_at": 1, + "in_progress_at": 1, + "completed_at": 1, + "expires_at": 1, + "request_counts": counts, + "metadata": None, + } + return RoutedResponse( + content_type="application/x-routed", + routes={ + "POST /files": JsonResponse( + content_type="application/json", + body={ + "id": "file-in-$REQUEST_ID", + "object": "file", + "purpose": "batch", + "bytes": 100, + "created_at": 1, + "filename": "in.jsonl", + "status": "processed", + }, + ), + "POST /batches": JsonResponse( + content_type="application/json", + body={ + **batch, + "status": "validating", + "output_file_id": None, + "error_file_id": None, + }, + ), + "GET /batches/batch-$REQUEST_ID": JsonResponse( + content_type="application/json", + body=batch, + ), + "GET /files/file-out-$REQUEST_ID/content": TextResponse( + content_type="application/jsonl", + body="\n".join(lines) + ("\n" if lines else ""), + ), + }, + ) + + +def _batch_input_lines(case: BatchCostCase, model_name: str) -> bytes: + count: Final = case.request_count + return ( + "\n".join( + json.dumps( + { + "custom_id": f"r{index}", + "method": "POST", + "url": "/v1/chat/completions", + "body": { + "model": model_name, + "messages": [{"role": "user", "content": "batch integration"}], + }, + }, + separators=(",", ":"), + ) + for index in range(1, count + 1) + ) + + "\n" + ).encode() + + +@pytest.mark.parametrize( + "case", + tuple(pytest.param(case, marks=pytest.mark.covers(case.covers), id=case.name) for case in BATCH_CASES), +) +def test_batch_costs(gateway: Gateway, case: BatchCostCase) -> None: + with gateway.scenario() as scenario: + key: Final = scenario.key() + model_name, identity = _register_deployment( + scenario, + case.litellm_model, + _batch_response(case), + case.name, + realtime=False, + ) + file_response: Final = gateway.request_multipart( + "/v1/files", + {"purpose": "batch", "model": model_name}, + {"file": ("in.jsonl", _batch_input_lines(case, model_name), "application/jsonl")}, + key=key, + ) + assert file_response.is_success, file_response.text + file_body: Final = JSON_OBJECT.validate_json(file_response.content) + time.sleep(2) + file_rows: Final = read_rows_now(key) + if file_rows: + assert all(row.spend == 0.0 for row in file_rows) + logging.info("file creation rows: %s", file_rows) + batch_response: Final = gateway.request( + "POST", + "/v1/batches", + { + "input_file_id": string_value(file_body["id"]), + "endpoint": "/v1/chat/completions", + "completion_window": "24h", + "model": model_name, + }, + key=key, + ) + assert batch_response.is_success, batch_response.text + batch_body: Final = JSON_OBJECT.validate_json(batch_response.content) + batch_id: Final = string_value(batch_body["id"]) + first_retrieval: Final = gateway.request("GET", f"/v1/batches/{batch_id}", key=key) + second_retrieval: Final = gateway.request("GET", f"/v1/batches/{batch_id}", key=key) + assert first_retrieval.is_success, first_retrieval.text + assert second_retrieval.is_success, second_retrieval.text + rows: tuple[CostRow, ...] + if case.output_lines: + rows = poll_rows(key, 1) + else: + time.sleep(5) + rows = read_rows_now(key) + if not rows: + logging.info("%s: completed failed batch produced no SpendLogs row", case.name) + return + retrieval_rows: Final = tuple(row for row in rows if row.call_type == "aretrieve_batch") + assert len(retrieval_rows) == 1 + row: Final = retrieval_rows[0] + assert row.status == "success" + assert row.call_type == "aretrieve_batch" + assert row.model_id == identity + assert_exact(case.name, "application/json", case.expected, row, second_retrieval) + time.sleep(3) + assert len(tuple(row for row in read_rows_now(key) if row.call_type == "aretrieve_batch")) == 1 + + +def _realtime_response(case: RealtimeCostCase) -> RealtimeResponse: + return RealtimeResponse( + content_type="application/x-realtime", + session_model=case.session_model, + events=tuple(turn.render(index, "$REQUEST_ID") for index, turn in enumerate(case.turns, start=1)), + ) + + +async def _run_realtime(url: str, key: str, model_name: str, turn_count: int) -> dict[str, JsonValue]: + async with websockets.connect( + f"{url.replace('http://', 'ws://').replace('https://', 'wss://')}/v1/realtime?model={model_name}", + additional_headers={"Authorization": f"Bearer {key}"}, + ) as websocket: + session: Final = JSON_OBJECT.validate_json(await websocket.recv()) + for _ in range(turn_count): + await websocket.send(json.dumps({"type": "response.create"})) + while True: + event: Final = JSON_OBJECT.validate_json(await websocket.recv()) + if event.get("type") == "response.done": + break + return session + + +@pytest.mark.parametrize( + "case", + tuple(pytest.param(case, marks=pytest.mark.covers(case.covers), id=case.name) for case in REALTIME_CASES), +) +def test_realtime_costs(gateway: Gateway, case: RealtimeCostCase) -> None: + with gateway.scenario() as scenario: + key: Final = scenario.key() + model_name, identity = _register_deployment( + scenario, + case.litellm_model, + _realtime_response(case), + case.name, + realtime=True, + ) + session: Final = asyncio.run( + _run_realtime( + os.environ["INTEGRATION_PROXY_URL"].rstrip("/"), + key, + model_name, + len(case.turns), + ) + ) + session_model: Final = object_value(session["session"])["model"] + assert session_model == (case.session_model or case.model) + row: Final = poll_rows(key, 1)[0] + assert row.status == "success" + assert row.call_type == "_arealtime" + assert row.model_id == identity + assert_exact(case.name, "application/json", case.expected, row, httpx.Response(200)) + + +@pytest.mark.covers("quota_management.spend_tracking.realtime_costs.no_turn_probe") +def test_realtime_no_turn_probe(gateway: Gateway) -> None: + with gateway.scenario() as scenario: + key: Final = scenario.key() + model_name, _identity = _register_deployment( + scenario, + "openai/gpt-realtime-mini-2025-12-15", + RealtimeResponse(content_type="application/x-realtime", events=()), + "realtime-no-turn", + realtime=True, + ) + asyncio.run( + _run_realtime( + os.environ["INTEGRATION_PROXY_URL"].rstrip("/"), + key, + model_name, + 0, + ) + ) + time.sleep(3) + rows: Final = read_rows_now(key) + logging.info("realtime no-turn probe rows=%s spend=%s", len(rows), rows[0].spend if rows else None) diff --git a/tests/integration/cost_calculation/test_cost_tracking.py b/tests/integration/cost_calculation/test_cost_tracking.py index f2057067c67..3c5d44c79b8 100644 --- a/tests/integration/cost_calculation/test_cost_tracking.py +++ b/tests/integration/cost_calculation/test_cost_tracking.py @@ -3,27 +3,23 @@ from __future__ import annotations import io -from itertools import islice import json -from hashlib import sha256 import struct import time -from typing import Final, cast import uuid import wave import zlib +from hashlib import sha256 +from itertools import islice +from typing import Final, cast import httpx import pytest -from pydantic import JsonValue - from integration._support.client import JSON_OBJECT, Gateway from integration._support.upstream import delete_scenario, register_scenario +from integration.cost_calculation.assertions import assert_exact, assert_recount from integration.cost_calculation.conftest import ( - CostBreakdown, - CostRow, approx_equal, - assert_total_is_sum_of_components, poll_cost_row, poll_failure_row, poll_rollups, @@ -32,14 +28,15 @@ from integration.cost_calculation.conftest import ( register_scenario_deployment, ) from integration.cost_calculation.cost_tracking_case import ( - BinaryResponse, CASES, + BinaryResponse, CostTrackingTestCase, ExactExpected, FailureExpected, RecountExpected, data_errors, ) +from pydantic import JsonValue if _data_errors := data_errors(): raise ValueError("\n".join(_data_errors)) @@ -115,135 +112,6 @@ def _replace_model(value: JsonValue, model_name: str) -> JsonValue: return value -def _assert_breakdown( - case: CostTrackingTestCase, - expected: ExactExpected, - breakdown: CostBreakdown, - response: httpx.Response, -) -> None: - assert breakdown.input_cost is not None and approx_equal(breakdown.input_cost, expected.input_cost), ( - f"{case.name}: input_cost {breakdown.input_cost} != expected {expected.input_cost}" - ) - assert breakdown.output_cost is not None and approx_equal(breakdown.output_cost, expected.output_cost), ( - f"{case.name}: output_cost {breakdown.output_cost} != expected {expected.output_cost}" - ) - for field, header_name, actual_component, expected_component in ( - ( - "cache_read_cost", - "x-litellm-response-cost-cache-read", - breakdown.cache_read_cost, - expected.cache_read_cost, - ), - ( - "cache_creation_cost", - "x-litellm-response-cost-cache-creation", - breakdown.cache_creation_cost, - expected.cache_creation_cost, - ), - ( - "reasoning_cost", - "x-litellm-response-cost-reasoning", - breakdown.reasoning_cost, - expected.reasoning_cost, - ), - ( - "tool_usage_cost", - "x-litellm-response-cost-tool-usage", - breakdown.tool_usage_cost, - expected.tool_usage_cost, - ), - ): - if expected_component is None: - continue - omitted_component_allowed: Final = expected_component == 0.0 - assert (actual_component is None and omitted_component_allowed) or ( - actual_component is not None and approx_equal(actual_component, expected_component) - ), f"{case.name}: {field} {actual_component} != expected {expected_component}" - if expected.cost_header and case.response.content_type == "application/json": - header: Final = response.headers.get(header_name) - assert (header is None and omitted_component_allowed) or ( - header is not None and approx_equal(float(header), expected_component) - ), f"{case.name}: {header_name} {header} != expected {expected_component}" - if expected.cost_header and case.response.content_type == "application/json" and any( - component is not None - for component in ( - expected.cache_read_cost, - expected.cache_creation_cost, - expected.reasoning_cost, - expected.tool_usage_cost, - ) - ): - input_header: Final = response.headers.get("x-litellm-response-cost-input") - output_header: Final = response.headers.get("x-litellm-response-cost-output") - expected_input_header: Final = expected.input_cost - ( - expected.cache_read_cost or 0.0 - ) - (expected.cache_creation_cost or 0.0) - assert input_header is not None and approx_equal(float(input_header), expected_input_header), ( - f"{case.name}: x-litellm-response-cost-input {input_header} != expected {expected_input_header}" - ) - assert output_header is not None and approx_equal(float(output_header), expected.output_cost), ( - f"{case.name}: x-litellm-response-cost-output {output_header} != expected {expected.output_cost}" - ) - - -def _assert_exact( - case: CostTrackingTestCase, - expected: ExactExpected, - row: CostRow, - response: httpx.Response, -) -> None: - assert row.spend is not None and approx_equal(row.spend, expected.spend), ( - f"{case.name}: spend {row.spend} != expected {expected.spend} " - f"(breakdown {row.breakdown.model_dump() if row.breakdown is not None else None})" - ) - breakdown: Final = row.breakdown - if expected.breakdown_persisted: - assert breakdown is not None, f"{case.name}: no cost_breakdown persisted" - if breakdown is not None: - _assert_breakdown(case, expected, breakdown, response) - assert row.prompt_tokens == expected.prompt_tokens, ( - f"{case.name}: prompt_tokens {row.prompt_tokens} != expected {expected.prompt_tokens}" - ) - assert row.completion_tokens == expected.completion_tokens, ( - f"{case.name}: completion_tokens {row.completion_tokens} != expected {expected.completion_tokens}" - ) - if breakdown is not None: - assert_total_is_sum_of_components(row, breakdown, case.name) - - -def _assert_recount(case: CostTrackingTestCase, expected: RecountExpected, row: CostRow) -> None: - assert row.prompt_tokens is not None and row.prompt_tokens > 0, ( - f"{case.name}: recount case counted no input tokens: prompt_tokens={row.prompt_tokens}" - ) - assert row.completion_tokens is not None and row.completion_tokens > 0, ( - f"{case.name}: recount case counted no output tokens: completion_tokens={row.completion_tokens}" - ) - if expected.prompt_tokens is not None: - assert row.prompt_tokens == expected.prompt_tokens, ( - f"{case.name}: prompt_tokens {row.prompt_tokens} != pinned {expected.prompt_tokens}" - ) - if expected.completion_tokens is not None: - assert row.completion_tokens == expected.completion_tokens, ( - f"{case.name}: completion_tokens {row.completion_tokens} != pinned {expected.completion_tokens}" - ) - if expected.min_completion_tokens is not None: - assert row.completion_tokens >= expected.min_completion_tokens, ( - f"{case.name}: completion_tokens {row.completion_tokens} < minimum {expected.min_completion_tokens}" - ) - if expected.max_completion_tokens is not None: - assert row.completion_tokens <= expected.max_completion_tokens, ( - f"{case.name}: completion_tokens {row.completion_tokens} > maximum {expected.max_completion_tokens}" - ) - recount: Final = row.prompt_tokens * expected.recount.input_cost_per_token + ( - row.completion_tokens * expected.recount.output_cost_per_token - ) - assert row.spend is not None and approx_equal(row.spend, recount), ( - f"{case.name}: spend {row.spend} != recount {recount} at map rates" - ) - assert row.breakdown is not None, f"{case.name}: no cost_breakdown persisted" - assert_total_is_sum_of_components(row, row.breakdown, case.name) - - @pytest.mark.parametrize("case", _CASES) def test_case_bills_expected_cost(gateway: Gateway, case: CostTrackingTestCase) -> None: marker: Final = sha256(case.name.encode()).hexdigest()[:12] @@ -356,7 +224,7 @@ def test_case_bills_expected_cost(gateway: Gateway, case: CostTrackingTestCase) row: Final = poll_cost_row(key) assert isinstance(expected, RecountExpected) assert row.status == "success", f"{case.name}: disconnect row status was {row.status}" - _assert_recount(case, expected, row) + assert_recount(case.name, expected, row) return responses: Final = tuple( ( @@ -385,7 +253,7 @@ def test_case_bills_expected_cost(gateway: Gateway, case: CostTrackingTestCase) rows: Final = poll_rows(key, len(responses)) if isinstance(expected, RecountExpected): row: Final = rows[0] - _assert_recount(case, expected, row) + assert_recount(case.name, expected, row) return assert isinstance(expected, ExactExpected) if fallback_deployment is not None: @@ -412,7 +280,7 @@ def test_case_bills_expected_cost(gateway: Gateway, case: CostTrackingTestCase) f"{case.name}: x-litellm-response-cost {header} != expected {expected.spend}" ) for row in rows: - _assert_exact(case, expected, row, response) + assert_exact(case.name, case.response.content_type, expected, row, response) if expected.rollups: assert deployment is not None and team_id is not None and user_id is not None assert end_user_id is not None From 2e16cd76c1bdd6e33f0be1a83e89ed381ea597e3 Mon Sep 17 00:00:00 2001 From: kerry Date: Sun, 20 Sep 2026 01:29:46 +0000 Subject: [PATCH 2/3] test(integration): tighten batch and realtime cost assertions Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- tests/integration/contracts.json | 4 +- .../cost_calculation/assertions.py | 10 ++-- .../cost_calculation/cost_tracking_case.py | 21 +++++--- .../cost_calculation/cost_tracking_cases.json | 16 ++++++ .../test_batch_realtime_cost.py | 52 +++---------------- 5 files changed, 46 insertions(+), 57 deletions(-) diff --git a/tests/integration/contracts.json b/tests/integration/contracts.json index 195c7ba4ab4..f49feff9e93 100644 --- a/tests/integration/contracts.json +++ b/tests/integration/contracts.json @@ -247,8 +247,8 @@ "tests/integration/cost_calculation/test_batch_realtime_cost.py::test_realtime_costs[gpt-realtime-mini-2025-12-15-realtime-priced_from_session_created_model]": [ "quota_management.spend_tracking.realtime_costs.session_model" ], - "tests/integration/cost_calculation/test_batch_realtime_cost.py::test_realtime_no_turn_probe": [ - "quota_management.spend_tracking.realtime_costs.no_turn_probe" + "tests/integration/cost_calculation/test_batch_realtime_cost.py::test_realtime_costs[gpt-realtime-mini-2025-12-15-realtime-session_without_turns_zero_spend]": [ + "quota_management.spend_tracking.realtime_costs.session_without_turns" ], "tests/integration/cost_calculation/test_cost_tracking.py::test_case_bills_expected_cost[anthropic.claude-sonnet-5-v1:0-cache_read]": [ "quota_management.spend_tracking.cost_matrix.logs_cost" diff --git a/tests/integration/cost_calculation/assertions.py b/tests/integration/cost_calculation/assertions.py index b0b9057dd9d..58a0fe99aab 100644 --- a/tests/integration/cost_calculation/assertions.py +++ b/tests/integration/cost_calculation/assertions.py @@ -15,8 +15,10 @@ def assert_breakdown( response_content_type: str, expected: ExactExpected, breakdown: CostBreakdown, - response: httpx.Response, + response: httpx.Response | None, ) -> None: + if response is None: + assert not expected.cost_header, f"{case_name}: cost headers require an HTTP response" assert breakdown.input_cost is not None and approx_equal(breakdown.input_cost, expected.input_cost), ( f"{case_name}: input_cost {breakdown.input_cost} != expected {expected.input_cost}" ) @@ -55,12 +57,12 @@ def assert_breakdown( assert (actual_component is None and omitted_component_allowed) or ( actual_component is not None and approx_equal(actual_component, expected_component) ), f"{case_name}: {field} {actual_component} != expected {expected_component}" - if expected.cost_header and response_content_type == "application/json": + if response is not None and expected.cost_header and response_content_type == "application/json": header: str | None = response.headers.get(header_name) assert (header is None and omitted_component_allowed) or ( header is not None and approx_equal(float(header), expected_component) ), f"{case_name}: {header_name} {header} != expected {expected_component}" - if expected.cost_header and response_content_type == "application/json" and any( + if response is not None and expected.cost_header and response_content_type == "application/json" and any( component is not None for component in ( expected.cache_read_cost, @@ -87,7 +89,7 @@ def assert_exact( response_content_type: str, expected: ExactExpected, row: CostRow, - response: httpx.Response, + response: httpx.Response | None, ) -> None: assert row.spend is not None and approx_equal(row.spend, expected.spend), ( f"{case_name}: spend {row.spend} != expected {expected.spend} " diff --git a/tests/integration/cost_calculation/cost_tracking_case.py b/tests/integration/cost_calculation/cost_tracking_case.py index 0e75b0c23b1..ad3de61d946 100644 --- a/tests/integration/cost_calculation/cost_tracking_case.py +++ b/tests/integration/cost_calculation/cost_tracking_case.py @@ -324,6 +324,12 @@ class BatchOutputLine(BaseModel): raise ValueError("status_code must be 200 or a 4xx status") return value + @model_validator(mode="after") + def validate_success_tokens(self) -> BatchOutputLine: + if self.status_code == 200 and (self.prompt_tokens is None or self.completion_tokens is None): + raise ValueError("successful batch output lines require prompt and completion tokens") + return self + def render(self, index: int, model: str, request_id: str) -> dict[str, JsonValue]: if self.status_code != 200: return { @@ -332,15 +338,18 @@ class BatchOutputLine(BaseModel): "response": None, "error": {"code": "bad_request", "message": "failed"}, } - assert self.prompt_tokens is not None - assert self.completion_tokens is not None - usage: dict[str, JsonValue] = { + if self.prompt_tokens is None or self.completion_tokens is None: + raise ValueError("successful batch output lines require prompt and completion tokens") + usage: Final = { "prompt_tokens": self.prompt_tokens, "completion_tokens": self.completion_tokens, "total_tokens": self.prompt_tokens + self.completion_tokens, + **( + {"prompt_tokens_details": {"cached_tokens": self.cached_tokens}} + if self.cached_tokens is not None + else {} + ), } - if self.cached_tokens is not None: - usage["prompt_tokens_details"] = {"cached_tokens": self.cached_tokens} return { "id": f"batch_req_{index}", "custom_id": f"r{index}", @@ -447,7 +456,7 @@ class RealtimeCostCase(BaseModel): covers: str model: str litellm_model: str - turns: tuple[RealtimeTurn, ...] = Field(min_length=1) + turns: tuple[RealtimeTurn, ...] = Field(min_length=0) session_model: str | None = None expected: ExactExpected diff --git a/tests/integration/cost_calculation/cost_tracking_cases.json b/tests/integration/cost_calculation/cost_tracking_cases.json index d8f321890f3..b86449a3263 100644 --- a/tests/integration/cost_calculation/cost_tracking_cases.json +++ b/tests/integration/cost_calculation/cost_tracking_cases.json @@ -30237,6 +30237,22 @@ "completion_tokens": 100, "cost_header": false } + }, + { + "name": "gpt-realtime-mini-2025-12-15-realtime-session_without_turns_zero_spend", + "covers": "quota_management.spend_tracking.realtime_costs.session_without_turns", + "model": "gpt-realtime-mini-2025-12-15", + "litellm_model": "openai/gpt-realtime-mini-2025-12-15", + "turns": [], + "expected": { + "spend": 0.0, + "input_cost": 0.0, + "output_cost": 0.0, + "prompt_tokens": 0, + "completion_tokens": 0, + "breakdown_persisted": false, + "cost_header": false + } } ] } diff --git a/tests/integration/cost_calculation/test_batch_realtime_cost.py b/tests/integration/cost_calculation/test_batch_realtime_cost.py index 9db25db32e8..3bf807059b4 100644 --- a/tests/integration/cost_calculation/test_batch_realtime_cost.py +++ b/tests/integration/cost_calculation/test_batch_realtime_cost.py @@ -2,19 +2,17 @@ from __future__ import annotations import asyncio import json -import logging import os import time from hashlib import sha256 from typing import Final -import httpx import pytest import websockets from integration._support.client import JSON_OBJECT, Gateway, Scenario, object_value, string_value from integration._support.upstream import delete_scenario, register_scenario from integration.cost_calculation.assertions import assert_exact -from integration.cost_calculation.conftest import CostRow, poll_rows, read_rows_now +from integration.cost_calculation.conftest import poll_rows, read_rows_now from integration.cost_calculation.cost_tracking_case import ( BATCH_CASES, REALTIME_CASES, @@ -69,7 +67,6 @@ def _batch_response(case: BatchCostCase) -> JsonResponse | RoutedResponse: "completed": case.completed_count, "failed": case.failed_count, } - completed: Final = len(case.output_lines) > 0 batch: Final = { "id": "batch-$REQUEST_ID", "object": "batch", @@ -77,9 +74,9 @@ def _batch_response(case: BatchCostCase) -> JsonResponse | RoutedResponse: "errors": None, "input_file_id": "file-in-$REQUEST_ID", "completion_window": "24h", - "status": "completed" if completed else "completed", - "output_file_id": "file-out-$REQUEST_ID" if completed else None, - "error_file_id": None if completed else "file-err-$REQUEST_ID", + "status": "completed", + "output_file_id": "file-out-$REQUEST_ID" if case.output_lines else None, + "error_file_id": None if case.output_lines else "file-err-$REQUEST_ID", "created_at": 1, "in_progress_at": 1, "completed_at": 1, @@ -168,10 +165,6 @@ def test_batch_costs(gateway: Gateway, case: BatchCostCase) -> None: assert file_response.is_success, file_response.text file_body: Final = JSON_OBJECT.validate_json(file_response.content) time.sleep(2) - file_rows: Final = read_rows_now(key) - if file_rows: - assert all(row.spend == 0.0 for row in file_rows) - logging.info("file creation rows: %s", file_rows) batch_response: Final = gateway.request( "POST", "/v1/batches", @@ -190,17 +183,10 @@ def test_batch_costs(gateway: Gateway, case: BatchCostCase) -> None: second_retrieval: Final = gateway.request("GET", f"/v1/batches/{batch_id}", key=key) assert first_retrieval.is_success, first_retrieval.text assert second_retrieval.is_success, second_retrieval.text - rows: tuple[CostRow, ...] - if case.output_lines: - rows = poll_rows(key, 1) - else: - time.sleep(5) - rows = read_rows_now(key) - if not rows: - logging.info("%s: completed failed batch produced no SpendLogs row", case.name) - return + rows: Final = poll_rows(key, 1) retrieval_rows: Final = tuple(row for row in rows if row.call_type == "aretrieve_batch") assert len(retrieval_rows) == 1 + assert all(row.spend == 0.0 for row in rows if row.call_type != "aretrieve_batch") row: Final = retrieval_rows[0] assert row.status == "success" assert row.call_type == "aretrieve_batch" @@ -261,28 +247,4 @@ def test_realtime_costs(gateway: Gateway, case: RealtimeCostCase) -> None: assert row.status == "success" assert row.call_type == "_arealtime" assert row.model_id == identity - assert_exact(case.name, "application/json", case.expected, row, httpx.Response(200)) - - -@pytest.mark.covers("quota_management.spend_tracking.realtime_costs.no_turn_probe") -def test_realtime_no_turn_probe(gateway: Gateway) -> None: - with gateway.scenario() as scenario: - key: Final = scenario.key() - model_name, _identity = _register_deployment( - scenario, - "openai/gpt-realtime-mini-2025-12-15", - RealtimeResponse(content_type="application/x-realtime", events=()), - "realtime-no-turn", - realtime=True, - ) - asyncio.run( - _run_realtime( - os.environ["INTEGRATION_PROXY_URL"].rstrip("/"), - key, - model_name, - 0, - ) - ) - time.sleep(3) - rows: Final = read_rows_now(key) - logging.info("realtime no-turn probe rows=%s spend=%s", len(rows), rows[0].spend if rows else None) + assert_exact(case.name, "application/json", case.expected, row, None) From 8d6326356d94e431f57907340c336846e56195b1 Mon Sep 17 00:00:00 2001 From: kerry Date: Sun, 20 Sep 2026 01:46:23 +0000 Subject: [PATCH 3/3] test(integration): wait for the batch retrieval row and drop the failed-only output file Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- .../integration/cost_calculation/conftest.py | 12 ++- .../test_batch_realtime_cost.py | 82 ++++++++++--------- 2 files changed, 55 insertions(+), 39 deletions(-) diff --git a/tests/integration/cost_calculation/conftest.py b/tests/integration/cost_calculation/conftest.py index 9c6ffe3f7d0..7a75320a70f 100644 --- a/tests/integration/cost_calculation/conftest.py +++ b/tests/integration/cost_calculation/conftest.py @@ -3,7 +3,7 @@ from __future__ import annotations import functools import json import os -from collections.abc import Mapping +from collections.abc import Callable, Mapping from dataclasses import dataclass from hashlib import sha256 from typing import Final @@ -134,8 +134,16 @@ def read_rows_now(key: str) -> tuple[CostRow, ...]: def poll_rows(key: str, count: int) -> tuple[CostRow, ...]: + return poll_rows_where(key, count, lambda _row: True) + + +def poll_rows_where( + key: str, + count: int, + predicate: Callable[[CostRow], bool], +) -> tuple[CostRow, ...]: result: Final = eventually( - lambda: read_rows_now(key), + lambda: tuple(row for row in read_rows_now(key) if predicate(row)), lambda rows: len(rows) >= count, seconds=60, ) diff --git a/tests/integration/cost_calculation/test_batch_realtime_cost.py b/tests/integration/cost_calculation/test_batch_realtime_cost.py index 3bf807059b4..1941e9ad153 100644 --- a/tests/integration/cost_calculation/test_batch_realtime_cost.py +++ b/tests/integration/cost_calculation/test_batch_realtime_cost.py @@ -12,7 +12,7 @@ import websockets from integration._support.client import JSON_OBJECT, Gateway, Scenario, object_value, string_value from integration._support.upstream import delete_scenario, register_scenario from integration.cost_calculation.assertions import assert_exact -from integration.cost_calculation.conftest import poll_rows, read_rows_now +from integration.cost_calculation.conftest import poll_rows, poll_rows_where, read_rows_now from integration.cost_calculation.cost_tracking_case import ( BATCH_CASES, REALTIME_CASES, @@ -67,6 +67,8 @@ def _batch_response(case: BatchCostCase) -> JsonResponse | RoutedResponse: "completed": case.completed_count, "failed": case.failed_count, } + has_output: Final = any(line.status_code == 200 for line in case.output_lines) + has_failed: Final = any(line.status_code != 200 for line in case.output_lines) batch: Final = { "id": "batch-$REQUEST_ID", "object": "batch", @@ -75,8 +77,8 @@ def _batch_response(case: BatchCostCase) -> JsonResponse | RoutedResponse: "input_file_id": "file-in-$REQUEST_ID", "completion_window": "24h", "status": "completed", - "output_file_id": "file-out-$REQUEST_ID" if case.output_lines else None, - "error_file_id": None if case.output_lines else "file-err-$REQUEST_ID", + "output_file_id": "file-out-$REQUEST_ID" if has_output else None, + "error_file_id": "file-err-$REQUEST_ID" if has_failed else None, "created_at": 1, "in_progress_at": 1, "completed_at": 1, @@ -84,39 +86,46 @@ def _batch_response(case: BatchCostCase) -> JsonResponse | RoutedResponse: "request_counts": counts, "metadata": None, } + routes: Final = { + "POST /files": JsonResponse( + content_type="application/json", + body={ + "id": "file-in-$REQUEST_ID", + "object": "file", + "purpose": "batch", + "bytes": 100, + "created_at": 1, + "filename": "in.jsonl", + "status": "processed", + }, + ), + "POST /batches": JsonResponse( + content_type="application/json", + body={ + **batch, + "status": "validating", + "output_file_id": None, + "error_file_id": None, + }, + ), + "GET /batches/batch-$REQUEST_ID": JsonResponse( + content_type="application/json", + body=batch, + ), + **( + { + "GET /files/file-out-$REQUEST_ID/content": TextResponse( + content_type="application/jsonl", + body="\n".join(lines) + ("\n" if lines else ""), + ) + } + if has_output + else {} + ), + } return RoutedResponse( content_type="application/x-routed", - routes={ - "POST /files": JsonResponse( - content_type="application/json", - body={ - "id": "file-in-$REQUEST_ID", - "object": "file", - "purpose": "batch", - "bytes": 100, - "created_at": 1, - "filename": "in.jsonl", - "status": "processed", - }, - ), - "POST /batches": JsonResponse( - content_type="application/json", - body={ - **batch, - "status": "validating", - "output_file_id": None, - "error_file_id": None, - }, - ), - "GET /batches/batch-$REQUEST_ID": JsonResponse( - content_type="application/json", - body=batch, - ), - "GET /files/file-out-$REQUEST_ID/content": TextResponse( - content_type="application/jsonl", - body="\n".join(lines) + ("\n" if lines else ""), - ), - }, + routes=routes, ) @@ -164,7 +173,6 @@ def test_batch_costs(gateway: Gateway, case: BatchCostCase) -> None: ) assert file_response.is_success, file_response.text file_body: Final = JSON_OBJECT.validate_json(file_response.content) - time.sleep(2) batch_response: Final = gateway.request( "POST", "/v1/batches", @@ -183,9 +191,9 @@ def test_batch_costs(gateway: Gateway, case: BatchCostCase) -> None: second_retrieval: Final = gateway.request("GET", f"/v1/batches/{batch_id}", key=key) assert first_retrieval.is_success, first_retrieval.text assert second_retrieval.is_success, second_retrieval.text - rows: Final = poll_rows(key, 1) - retrieval_rows: Final = tuple(row for row in rows if row.call_type == "aretrieve_batch") + retrieval_rows: Final = poll_rows_where(key, 1, lambda row: row.call_type == "aretrieve_batch") assert len(retrieval_rows) == 1 + rows: Final = read_rows_now(key) assert all(row.spend == 0.0 for row in rows if row.call_type != "aretrieve_batch") row: Final = retrieval_rows[0] assert row.status == "success"