From c958f9db7e0a024f7fbbb28168c47ac9e65c4853 Mon Sep 17 00:00:00 2001 From: kerry Date: Sat, 19 Sep 2026 18:44:51 +0000 Subject: [PATCH] test: extend cost tracking integration harness Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- tests/integration/_support/upstream.py | 20 +-- tests/integration/contracts.json | 9 ++ .../integration/cost_calculation/conftest.py | 34 ++++- .../cost_calculation/cost_tracking_case.py | 74 +++++++++- .../cost_calculation/cost_tracking_cases.json | 135 +++++++++++++++++- .../cost_calculation/test_cost_tracking.py | 54 ++++++- 6 files changed, 304 insertions(+), 22 deletions(-) diff --git a/tests/integration/_support/upstream.py b/tests/integration/_support/upstream.py index 1ad02b6a3f2..3c4198c8133 100644 --- a/tests/integration/_support/upstream.py +++ b/tests/integration/_support/upstream.py @@ -1,25 +1,19 @@ from __future__ import annotations import argparse +import json +import os +import struct +import zlib from collections import deque from collections.abc import 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 zlib 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 -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 ( EventStreamResponse, @@ -27,6 +21,11 @@ from integration.cost_calculation.cost_tracking_case import ( SseResponse, StoredResponse, ) +from pydantic import BaseModel, JsonValue, TypeAdapter, ValidationError +from starlette.applications import Starlette +from starlette.requests import Request +from starlette.responses import JSONResponse, Response +from starlette.routing import Route JSON_OBJECT: Final = TypeAdapter(dict[str, JsonValue]) CASES_FILE: Final = Path(__file__).resolve().parents[1] / "cost_calculation" / "cost_tracking_cases.json" @@ -210,6 +209,7 @@ class Provider: "$REQUEST_ID", scenario_id ).encode(), media_type=response.content_type, + status_code=response.status, ) case SseResponse(): stream_body: Final = ("\n\n".join(response.frames) + "\n\n").replace( diff --git a/tests/integration/contracts.json b/tests/integration/contracts.json index fe7b6dfe7ac..1b0c9e31d4e 100644 --- a/tests/integration/contracts.json +++ b/tests/integration/contracts.json @@ -382,6 +382,15 @@ "tests/integration/cost_calculation/test_cost_tracking.py::test_case_bills_expected_cost[azure-gpt-5.6-stream_full_usage]": [ "quota_management.spend_tracking.scripted_wire.logs_cost" ], + "tests/integration/cost_calculation/test_cost_tracking.py::test_case_bills_expected_cost[gpt-5.6-responses_native_json]": [ + "quota_management.spend_tracking.cost_matrix.logs_cost" + ], + "tests/integration/cost_calculation/test_cost_tracking.py::test_case_bills_expected_cost[gpt-5.6-upstream_500_zero_spend]": [ + "quota_management.spend_tracking.cost_matrix.logs_cost" + ], + "tests/integration/cost_calculation/test_cost_tracking.py::test_case_bills_expected_cost[gpt-5.6-upstream_429_zero_spend]": [ + "quota_management.spend_tracking.cost_matrix.logs_cost" + ], "tests/integration/cost_calculation/test_cost_tracking.py::test_case_bills_expected_cost[claude-haiku-4-5-input_text]": [ "quota_management.spend_tracking.cost_matrix.logs_cost" ], diff --git a/tests/integration/cost_calculation/conftest.py b/tests/integration/cost_calculation/conftest.py index f1b8901d626..192f2bf5461 100644 --- a/tests/integration/cost_calculation/conftest.py +++ b/tests/integration/cost_calculation/conftest.py @@ -9,12 +9,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 delete_scenario, register_scenario from integration.cost_calculation.cost_tracking_case import CostTrackingTestCase +from pydantic import BaseModel, ConfigDict class CostBreakdown(BaseModel): @@ -50,6 +49,15 @@ class CostRow(BaseModel): return self.metadata.cost_breakdown +class FailureRow(BaseModel): + model_config = ConfigDict(extra="ignore") + + spend: float + status: str + prompt_tokens: int | None = None + completion_tokens: int | None = None + + def approx_equal(actual: float, expected: float) -> bool: return abs(actual - expected) <= max(1e-9, abs(expected) * 1e-2) @@ -92,6 +100,28 @@ def poll_cost_row(key: str) -> CostRow: return result +def poll_failure_row(key: str) -> FailureRow: + digest: Final = sha256(key.encode()).hexdigest() + + def read() -> FailureRow | None: + rows: Final = read_rows( + 'SELECT spend, status, prompt_tokens, completion_tokens FROM "LiteLLM_SpendLogs" WHERE api_key=%s', + (digest,), + ) + return next( + ( + parsed + for row in rows + if (parsed := FailureRow.model_validate(row)).status == "failure" + ), + None, + ) + + result: Final = eventually(read, lambda row: row is not None, seconds=60) + assert result is not None + return result + + @functools.cache def _vertex_private_key_pem() -> str: return rsa.generate_private_key(public_exponent=65537, key_size=2048).private_bytes( diff --git a/tests/integration/cost_calculation/cost_tracking_case.py b/tests/integration/cost_calculation/cost_tracking_case.py index 6af95f995ff..cb408f2c488 100644 --- a/tests/integration/cost_calculation/cost_tracking_case.py +++ b/tests/integration/cost_calculation/cost_tracking_case.py @@ -71,6 +71,7 @@ class JsonResponse(BaseModel): content_type: Literal["application/json"] body: dict[str, JsonValue] + status: int = 200 class SseResponse(BaseModel): @@ -108,6 +109,10 @@ class ExactExpected(BaseModel): output_cost: float prompt_tokens: int completion_tokens: int + cache_read_cost: float | None = None + cache_creation_cost: float | None = None + reasoning_cost: float | None = None + tool_usage_cost: float | None = None class RecountRates(BaseModel): @@ -123,7 +128,19 @@ class RecountExpected(BaseModel): recount: RecountRates -Expected: TypeAlias = ExactExpected | RecountExpected +class FailureDetails(BaseModel): + model_config = ConfigDict(frozen=True, extra="forbid") + + status: int + + +class FailureExpected(BaseModel): + model_config = ConfigDict(frozen=True, extra="forbid") + + failure: FailureDetails + + +Expected: TypeAlias = ExactExpected | RecountExpected | FailureExpected class CostTrackingTestCase(BaseModel): @@ -132,6 +149,15 @@ class CostTrackingTestCase(BaseModel): name: str covers: str model: str + endpoint: Literal[ + "/v1/chat/completions", + "/v1/responses", + "/v1/messages", + "/v1/embeddings", + "/v1/rerank", + "/v1/completions", + "/v1/moderations", + ] = "/v1/chat/completions" deployment: Deployment | None = None request: dict[str, JsonValue] response: StoredResponse @@ -240,6 +266,44 @@ def data_errors() -> tuple[str, ...]: or case.expected.recount.output_cost_per_token != (COST_MAP[case.model].output_cost_per_token or 0.0) ) ) + component_mismatches: Final = sorted( + case.name + for case in CASES + if isinstance(case.expected, ExactExpected) + and any( + component is not None + for component in ( + case.expected.cache_read_cost, + case.expected.cache_creation_cost, + case.expected.reasoning_cost, + case.expected.tool_usage_cost, + ) + ) + and ( + (case.expected.cache_read_cost or 0.0) + (case.expected.cache_creation_cost or 0.0) + > case.expected.input_cost + or (case.expected.reasoning_cost or 0.0) > case.expected.output_cost + or not _approx_equal( + case.expected.input_cost + + case.expected.output_cost + + (case.expected.tool_usage_cost or 0.0), + case.expected.spend, + ) + ) + ) + failure_response_mismatches: Final = sorted( + case.name + for case in CASES + if ( + isinstance(case.expected, FailureExpected) + and (not isinstance(case.response, JsonResponse) or case.response.status < 400) + ) + or ( + not isinstance(case.expected, FailureExpected) + and isinstance(case.response, JsonResponse) + and case.response.status != 200 + ) + ) return tuple( message for message in ( @@ -248,6 +312,14 @@ def data_errors() -> tuple[str, ...]: f"duplicate case names: {duplicate_names}" if duplicate_names else None, f"cost-map entries share input_cost_per_token: {shared_input_rates}" if shared_input_rates else None, f"recount rates differ from cost-map rates: {recount_mismatches}" if recount_mismatches else None, + f"breakdown components are inconsistent: {component_mismatches}" if component_mismatches else None, + f"failure response statuses are inconsistent: {failure_response_mismatches}" + if failure_response_mismatches + else None, ) if message is not None ) + + +def _approx_equal(actual: float, expected: float) -> bool: + return abs(actual - expected) <= max(1e-9, abs(expected) * 1e-2) diff --git a/tests/integration/cost_calculation/cost_tracking_cases.json b/tests/integration/cost_calculation/cost_tracking_cases.json index 3627774816f..379250d3123 100644 --- a/tests/integration/cost_calculation/cost_tracking_cases.json +++ b/tests/integration/cost_calculation/cost_tracking_cases.json @@ -534,7 +534,8 @@ "input_cost": 0.00616704, "output_cost": 0.00627, "prompt_tokens": 12928, - "completion_tokens": 380 + "completion_tokens": 380, + "cache_read_cost": 0.00405504 } }, { @@ -605,7 +606,8 @@ "input_cost": 0.0397056, "output_cost": 0.005775, "prompt_tokens": 9728, - "completion_tokens": 350 + "completion_tokens": 350, + "cache_creation_cost": 0.038016 } }, { @@ -681,7 +683,8 @@ "input_cost": 0.0574464, "output_cost": 0.005775, "prompt_tokens": 9728, - "completion_tokens": 350 + "completion_tokens": 350, + "cache_creation_cost": 0.0557568 } }, { @@ -3417,7 +3420,8 @@ "input_cost": 0.002232, "output_cost": 0.065484, "prompt_tokens": 1240, - "completion_tokens": 4040 + "completion_tokens": 4040, + "reasoning_cost": 0.05742 } }, { @@ -3638,7 +3642,8 @@ "input_cost": 0.003312, "output_cost": 0.0059328, "prompt_tokens": 1840, - "completion_tokens": 412 + "completion_tokens": 412, + "tool_usage_cost": 0.0125 } }, { @@ -4494,7 +4499,8 @@ "input_cost": 0.0018688, "output_cost": 0.0019, "prompt_tokens": 12928, - "completion_tokens": 380 + "completion_tokens": 380, + "cache_read_cost": 0.0012288 } }, { @@ -16955,7 +16961,8 @@ "input_cost": 0.00276, "output_cost": 0.004944, "prompt_tokens": 1840, - "completion_tokens": 412 + "completion_tokens": 412, + "tool_usage_cost": 0.0025 } }, { @@ -21797,6 +21804,120 @@ "completion_tokens": 1592 } }, + { + "name": "gpt-5.6-responses_native_json", + "covers": "quota_management.spend_tracking.cost_matrix.logs_cost", + "model": "gpt-5.6", + "endpoint": "/v1/responses", + "request": { + "model": "$MODEL", + "input": "responses native fixture", + "stream": false + }, + "response": { + "content_type": "application/json", + "body": { + "object": "response", + "status": "completed", + "model": "gpt-5.6", + "output": [ + { + "type": "message", + "id": "msg_$REQUEST_ID", + "status": "completed", + "role": "assistant", + "content": [ + { + "type": "output_text", + "text": "scripted response", + "annotations": [] + } + ] + } + ], + "usage": { + "input_tokens": 11, + "output_tokens": 7, + "total_tokens": 18, + "input_tokens_details": { + "cached_tokens": 0 + }, + "output_tokens_details": { + "reasoning_tokens": 0 + } + } + } + }, + "expected": { + "spend": 0.00011725, + "input_cost": 1.925e-05, + "output_cost": 9.8e-05, + "prompt_tokens": 11, + "completion_tokens": 7 + } + }, + { + "name": "gpt-5.6-upstream_500_zero_spend", + "covers": "quota_management.spend_tracking.cost_matrix.logs_cost", + "model": "gpt-5.6", + "request": { + "model": "$MODEL", + "messages": [ + { + "role": "user", + "content": "scripted upstream failure 500" + } + ], + "stream": false + }, + "response": { + "content_type": "application/json", + "status": 500, + "body": { + "error": { + "message": "scripted upstream failure", + "type": "server_error", + "code": "500" + } + } + }, + "expected": { + "failure": { + "status": 500 + } + } + }, + { + "name": "gpt-5.6-upstream_429_zero_spend", + "covers": "quota_management.spend_tracking.cost_matrix.logs_cost", + "model": "gpt-5.6", + "request": { + "model": "$MODEL", + "messages": [ + { + "role": "user", + "content": "scripted upstream failure 429" + } + ], + "stream": false + }, + "response": { + "content_type": "application/json", + "status": 429, + "body": { + "error": { + "message": "scripted upstream failure", + "type": "rate_limit_error", + "code": "429" + } + } + }, + "expected": { + "failure": { + "status": 429 + } + } + }, { "name": "meta.llama4-maverick-17b-instruct-v1:0-input_text", "covers": "quota_management.spend_tracking.cost_matrix.logs_cost", diff --git a/tests/integration/cost_calculation/test_cost_tracking.py b/tests/integration/cost_calculation/test_cost_tracking.py index a8a56fbfbbd..c9dc7a5ffe3 100644 --- a/tests/integration/cost_calculation/test_cost_tracking.py +++ b/tests/integration/cost_calculation/test_cost_tracking.py @@ -6,18 +6,19 @@ from hashlib import sha256 from typing import Final, cast import pytest - from integration._support.client import JSON_OBJECT, Gateway from integration.cost_calculation.conftest import ( approx_equal, assert_total_is_sum_of_components, poll_cost_row, + poll_failure_row, register_scenario_deployment, ) from integration.cost_calculation.cost_tracking_case import ( CASES, CostTrackingTestCase, ExactExpected, + FailureExpected, RecountExpected, data_errors, ) @@ -51,10 +52,22 @@ def test_case_bills_expected_cost(gateway: Gateway, case: CostTrackingTestCase) model_name: Final = register_scenario_deployment(scenario, case, marker, key) response: Final = gateway.request( "POST", - "/v1/chat/completions", + case.endpoint, {**case.request, "model": model_name}, key=key, ) + if isinstance(case.expected, FailureExpected): + assert response.status_code == case.expected.failure.status, ( + f"{case.name}: proxy returned {response.status_code}, expected {case.expected.failure.status}: " + f"{response.text[:400]}" + ) + response_cost: Final = response.headers.get("x-litellm-response-cost") + assert response_cost is None or approx_equal(float(response_cost), 0.0), ( + f"{case.name}: failure response cost was {response_cost}" + ) + row: Final = poll_failure_row(key) + assert row.spend == 0, f"{case.name}: failure spend was {row.spend}" + return assert response.is_success, f"{case.name}: proxy returned {response.status_code}: {response.text[:400]}" if case.response.content_type == "text/event-stream": _assert_stream_has_no_error(response.text) @@ -92,6 +105,43 @@ def test_case_bills_expected_cost(gateway: Gateway, case: CostTrackingTestCase) 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, expected_component in ( + ("cache_read_cost", "x-litellm-response-cost-cache-read", expected.cache_read_cost), + ("cache_creation_cost", "x-litellm-response-cost-cache-creation", expected.cache_creation_cost), + ("reasoning_cost", "x-litellm-response-cost-reasoning", expected.reasoning_cost), + ("tool_usage_cost", "x-litellm-response-cost-tool-usage", expected.tool_usage_cost), + ): + if expected_component is None: + continue + actual_component: Final = getattr(breakdown, field) + assert actual_component is not None and approx_equal(actual_component, expected_component), ( + f"{case.name}: {field} {actual_component} != expected {expected_component}" + ) + if case.response.content_type == "application/json": + header: Final = response.headers.get(header_name) + assert header is not None and approx_equal(float(header), expected_component), ( + f"{case.name}: {header_name} {header} != expected {expected_component}" + ) + if 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}" + ) assert row.prompt_tokens == expected.prompt_tokens, ( f"{case.name}: prompt_tokens {row.prompt_tokens} != expected {expected.prompt_tokens}" )