scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-09-25 01:02:15 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 131

test(integration): proxy behaviour cost cases

Browse source 
Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>
This commit is contained in:
kerry 2026-09-20 00:19:56 +00:00
parent 98ea6dd405
commit a841750d46
6 changed files with 965 additions and 51 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

12
tests/integration/_support/upstream.py
View file

@ -1,9 +1,10 @@
from __future__ import annotations
import argparse
import asyncio
import base64
from collections import deque
from collections.abc import Mapping
from collections.abc import AsyncIterator, Mapping
import json
from dataclasses import dataclass, field
import os
@ -18,7 +19,7 @@ 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.responses import JSONResponse, Response, StreamingResponse
from starlette.routing import Route
from _fake_openai_endpoint_server import chat_completions, completions, embeddings, health, moderations
@ -223,6 +224,13 @@ class Provider:
media_type=response.content_type,
)
case SseResponse():
if response.frame_delay_ms > 0:
async def stream() -> AsyncIterator[bytes]:
for frame in response.frames:
yield f"{frame.replace('$REQUEST_ID', scenario_id)}\n\n".encode()
await asyncio.sleep(response.frame_delay_ms / 1000)
return StreamingResponse(stream(), media_type=response.content_type)
stream_body: Final = ("\n\n".join(response.frames) + "\n\n").replace(
"$REQUEST_ID", scenario_id
)

42
tests/integration/contracts.json
View file

@ -1539,6 +1539,48 @@
],
"tests/integration/cost_calculation/test_cost_tracking.py::test_case_bills_expected_cost[openai-deployment-pricing-override]": [
"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_400_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_401_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_500_stream_request_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-responses_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[claude-sonnet-5-messages_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-fallback_billed_to_answering_deployment]": [
"quota_management.spend_tracking.routing.fallback_billing"
],
"tests/integration/cost_calculation/test_cost_tracking.py::test_case_bills_expected_cost[gpt-5.6-n_2_choices]": [
"quota_management.spend_tracking.cost_matrix.logs_cost"
],
"tests/integration/cost_calculation/test_cost_tracking.py::test_case_bills_expected_cost[gpt-5.6-finish_reason_length]": [
"quota_management.spend_tracking.cost_matrix.logs_cost"
],
"tests/integration/cost_calculation/test_cost_tracking.py::test_case_bills_expected_cost[gpt-5.6-stream_usage_in_empty_choices_chunk]": [
"quota_management.spend_tracking.scripted_wire.logs_cost"
],
"tests/integration/cost_calculation/test_cost_tracking.py::test_case_bills_expected_cost[gpt-5.6-stream_usage_in_last_delta_chunk]": [
"quota_management.spend_tracking.scripted_wire.logs_cost"
],
"tests/integration/cost_calculation/test_cost_tracking.py::test_case_bills_expected_cost[gpt-5.6-unknown_model_response_model_unknown]": [
"quota_management.spend_tracking.cost_matrix.logs_cost"
],
"tests/integration/cost_calculation/test_cost_tracking.py::test_case_bills_expected_cost[gpt-5.6-unknown_model_response_model_known]": [
"quota_management.spend_tracking.cost_matrix.logs_cost"
],
"tests/integration/cost_calculation/test_cost_tracking.py::test_case_bills_expected_cost[gpt-5.6-chat_request_to_embedding_entry]": [
"quota_management.spend_tracking.cost_matrix.logs_cost"
],
"tests/integration/cost_calculation/test_cost_tracking.py::test_case_bills_expected_cost[gpt-5.6-client_disconnect_mid_stream]": [
"quota_management.spend_tracking.scripted_wire.client_disconnect"
]
},
"browser": {

113
tests/integration/cost_calculation/conftest.py
View file

@ -4,6 +4,7 @@ import functools
import json
import os
from collections.abc import Mapping
from dataclasses import dataclass
from hashlib import sha256
from typing import Final
@ -13,8 +14,8 @@ 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 integration._support.upstream import ScenarioHandle, delete_scenario, register_scenario
from integration.cost_calculation.cost_tracking_case import CostTrackingTestCase, StoredResponse
class CostBreakdown(BaseModel):
@ -43,6 +44,7 @@ class CostRow(BaseModel):
status: str | None = None
prompt_tokens: int | None = None
completion_tokens: int | None = None
model_id: str | None = None
metadata: CostMetadata | None = None
@property
@ -59,6 +61,26 @@ class FailureRow(BaseModel):
completion_tokens: int | None = None
class DailySpend(BaseModel):
model_config = ConfigDict(extra="ignore")
spend: float | None = None
prompt_tokens: int | None = None
completion_tokens: int | None = None
api_requests: int | None = None
class Rollups(BaseModel):
model_config = ConfigDict(frozen=True, extra="forbid")
key_spend: float
team_spend: float
user_spend: float
end_user_spend: float
daily_user: DailySpend
daily_team: DailySpend
def approx_equal(actual: float, expected: float) -> bool:
return abs(actual - expected) <= max(1e-9, abs(expected) * 1e-2)
@ -90,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 '
'SELECT spend, status, metadata, prompt_tokens, completion_tokens, model_id '
'FROM "LiteLLM_SpendLogs" WHERE api_key=%s',
(digest,),
)
@ -101,6 +123,67 @@ def poll_cost_row(key: str) -> CostRow:
return result
def poll_rows(key: str) -> tuple[CostRow, ...]:
digest: Final = sha256(key.encode()).hexdigest()
def read() -> tuple[CostRow, ...]:
rows: Final = read_rows(
'SELECT spend, status, metadata, prompt_tokens, completion_tokens, model_id '
'FROM "LiteLLM_SpendLogs" WHERE api_key=%s ORDER BY "startTime"',
(digest,),
)
return tuple(parsed for row in rows if (parsed := _row(row)) is not None)
result: Final = eventually(read, lambda rows: len(rows) > 0, seconds=20)
return result
def poll_rollups(key: str, team_id: str, user_id: str, end_user_id: str) -> Rollups:
digest: Final = sha256(key.encode()).hexdigest()
def read() -> Rollups | None:
key_rows: Final = read_rows(
'SELECT spend FROM "LiteLLM_VerificationToken" WHERE token=%s',
(digest,),
)
team_rows: Final = read_rows(
'SELECT spend FROM "LiteLLM_TeamTable" WHERE team_id=%s',
(team_id,),
)
user_rows: Final = read_rows(
'SELECT spend FROM "LiteLLM_UserTable" WHERE user_id=%s',
(user_id,),
)
end_user_rows: Final = read_rows(
'SELECT spend FROM "LiteLLM_EndUserTable" WHERE user_id=%s',
(end_user_id,),
)
daily_user_rows: Final = read_rows(
'SELECT spend, prompt_tokens, completion_tokens, api_requests '
'FROM "LiteLLM_DailyUserSpend" WHERE user_id=%s AND api_key=%s AND date=CURRENT_DATE::text',
(user_id, digest),
)
daily_team_rows: Final = read_rows(
'SELECT spend, prompt_tokens, completion_tokens, api_requests '
'FROM "LiteLLM_DailyTeamSpend" WHERE team_id=%s AND api_key=%s AND date=CURRENT_DATE::text',
(team_id, digest),
)
if not all((key_rows, team_rows, user_rows, end_user_rows, daily_user_rows, daily_team_rows)):
return None
return Rollups(
key_spend=float(key_rows[0]["spend"]),
team_spend=float(team_rows[0]["spend"]),
user_spend=float(user_rows[0]["spend"]),
end_user_spend=float(end_user_rows[0]["spend"]),
daily_user=DailySpend.model_validate(daily_user_rows[0]),
daily_team=DailySpend.model_validate(daily_team_rows[0]),
)
result: Final = eventually(read, lambda value: value is not None, seconds=20)
assert result is not None
return result
def poll_failure_row(key: str) -> FailureRow:
digest: Final = sha256(key.encode()).hexdigest()
@ -147,17 +230,31 @@ def _vertex_service_account_json(url: str) -> str:
)
@dataclass(frozen=True, slots=True)
class RegisteredDeployment:
model_name: str
identity: str
handle: ScenarioHandle
def register_scenario_deployment(
scenario: Scenario,
case: CostTrackingTestCase,
marker: str,
key: str,
) -> str:
*,
response: StoredResponse | None = None,
marker_suffix: str = "",
model_name: str | None = None,
) -> RegisteredDeployment:
control_url: Final = os.environ["INTEGRATION_UPSTREAM_URL"].rstrip("/")
run_marker: Final = sha256(key.encode()).hexdigest()[:12]
handle: Final = register_scenario(f"sc-{marker}-{run_marker}", case.response)
handle: Final = register_scenario(
f"sc-{marker}{marker_suffix}-{run_marker}",
case.response if response is None else response,
)
scenario.cleanups.callback(delete_scenario, handle)
model_name: Final = f"cost-{marker}-{run_marker}"
registered_model_name: Final = model_name or f"cost-{marker}{marker_suffix}-{run_marker}"
parameters: Final = {
"model": case.litellm_model,
"api_key": case.api_key,
@ -184,7 +281,7 @@ def register_scenario_deployment(
created: Final = scenario.gateway.post(
"/model/new",
JSON_OBJECT.validate_python({
"model_name": model_name,
"model_name": registered_model_name,
"litellm_params": parameters,
"model_info": (
{"base_model": case.base_model}
@ -195,4 +292,4 @@ def register_scenario_deployment(
)
identity: Final = string_value(object_value(created["model_info"])["id"])
scenario.cleanups.callback(scenario.delete_model, identity)
return model_name
return RegisteredDeployment(model_name=registered_model_name, identity=identity, handle=handle)

34
tests/integration/cost_calculation/cost_tracking_case.py
View file

@ -116,6 +116,7 @@ class SseResponse(BaseModel):
content_type: Literal["text/event-stream"]
frames: tuple[str, ...]
frame_delay_ms: int = Field(default=0, ge=0)
class EventStreamEvent(BaseModel):
@ -160,6 +161,7 @@ class ExactExpected(BaseModel):
tool_usage_cost: float | None = None
breakdown_persisted: bool = True
cost_header: bool = True
rollups: bool = False
class RecountRates(BaseModel):
@ -173,6 +175,8 @@ class RecountExpected(BaseModel):
model_config = ConfigDict(frozen=True, extra="forbid")
recount: RecountRates
prompt_tokens: int | None = None
completion_tokens: int | None = None
class FailureDetails(BaseModel):
@ -217,6 +221,8 @@ class CostTrackingTestCase(BaseModel):
request: dict[str, JsonValue]
response: StoredResponse
expected: Expected
fallback_from: StoredResponse | None = None
disconnect_after_frames: int | None = Field(default=None, ge=1)
@property
def rates(self) -> CostMapEntry:
@ -430,7 +436,31 @@ def data_errors() -> tuple[str, ...]:
and case.rates.mode != "image_generation"
and not case.reports_provider_cost
)
or (not case.expected.cost_header and case.passthrough_provider is None)
or (
not case.expected.cost_header
and case.passthrough_provider is None
and not isinstance(case.response, SseResponse)
and case.expected.spend != 0.0
)
)
)
invalid_fallbacks: Final = sorted(
case.name
for case in CASES
if case.fallback_from is not None
and (
not isinstance(case.fallback_from, JsonResponse)
or not 400 <= case.fallback_from.status <= 599
)
)
invalid_disconnects: Final = sorted(
case.name
for case in CASES
if case.disconnect_after_frames is not None
and (
not isinstance(case.response, SseResponse)
or case.response.frame_delay_ms <= 0
or not isinstance(case.expected, RecountExpected)
)
)
return tuple(
@ -446,6 +476,8 @@ def data_errors() -> tuple[str, ...]:
if failure_response_mismatches
else None,
f"invalid passthrough opt-outs: {invalid_opt_outs}" if invalid_opt_outs else None,
f"invalid fallback responses: {invalid_fallbacks}" if invalid_fallbacks else None,
f"invalid disconnect cases: {invalid_disconnects}" if invalid_disconnects else None,
)
if message is not None
)

624
tests/integration/cost_calculation/cost_tracking_cases.json
View file

@ -617,6 +617,11 @@
"mode": "chat",
"input_cost_per_token": 1.51e-06,
"output_cost_per_token": 7.51e-06
},
"text-embedding-3-large": {
"litellm_provider": "openai",
"mode": "embedding",
"input_cost_per_token": 1.3e-07
}
},
"cases": [
@ -6947,7 +6952,8 @@
"input_cost": 0.00552,
"output_cost": 0.00618,
"prompt_tokens": 1840,
"completion_tokens": 412
"completion_tokens": 412,
"rollups": true
}
},
{
@ -7608,7 +7614,9 @@
"recount": {
"input_cost_per_token": 3e-06,
"output_cost_per_token": 1.5e-05
}
},
"prompt_tokens": 47,
"completion_tokens": 10
}
},
{
@ -7695,7 +7703,9 @@
"recount": {
"input_cost_per_token": 3e-06,
"output_cost_per_token": 1.5e-05
}
},
"prompt_tokens": 49,
"completion_tokens": 111
}
},
{
@ -7752,7 +7762,9 @@
"recount": {
"input_cost_per_token": 3e-06,
"output_cost_per_token": 1.5e-05
}
},
"prompt_tokens": 301,
"completion_tokens": 9
}
},
{
@ -12782,7 +12794,9 @@
"recount": {
"input_cost_per_token": 5.2e-07,
"output_cost_per_token": 3.12e-06
}
},
"prompt_tokens": 48,
"completion_tokens": 12
}
},
{
@ -12862,7 +12876,9 @@
"recount": {
"input_cost_per_token": 5.2e-07,
"output_cost_per_token": 3.12e-06
}
},
"prompt_tokens": 46,
"completion_tokens": 88
}
},
{
@ -12914,7 +12930,9 @@
"recount": {
"input_cost_per_token": 5.2e-07,
"output_cost_per_token": 3.12e-06
}
},
"prompt_tokens": 302,
"completion_tokens": 10
}
},
{
@ -20760,7 +20778,8 @@
"input_cost": 0.00322,
"output_cost": 0.005768,
"prompt_tokens": 1840,
"completion_tokens": 412
"completion_tokens": 412,
"rollups": true
}
},
{
@ -21515,7 +21534,9 @@
"recount": {
"input_cost_per_token": 1.75e-06,
"output_cost_per_token": 1.4e-05
}
},
"prompt_tokens": 49,
"completion_tokens": 12
}
},
{
@ -21601,7 +21622,9 @@
"recount": {
"input_cost_per_token": 1.75e-06,
"output_cost_per_token": 1.4e-05
}
},
"prompt_tokens": 44,
"completion_tokens": 105
}
},
{
@ -21656,7 +21679,9 @@
"recount": {
"input_cost_per_token": 1.75e-06,
"output_cost_per_token": 1.4e-05
}
},
"prompt_tokens": 302,
"completion_tokens": 11
}
},
{
@ -29417,6 +29442,583 @@
"prompt_tokens": 1840,
"completion_tokens": 412
}
},
{
"name": "gpt-5.6-upstream_400_zero_spend",
"covers": "quota_management.spend_tracking.cost_matrix.logs_cost",
"model": "gpt-5.6",
"request": {
"model": "$MODEL",
"messages": [
{
"role": "user",
"content": "proxy behaviour probe"
}
],
"stream": false
},
"response": {
"content_type": "application/json",
"status": 400,
"body": {
"error": {
"message": "scripted upstream failure 400",
"type": "server_error",
"code": "400"
}
}
},
"expected": {
"failure": {
"status": 400
}
}
},
{
"name": "gpt-5.6-upstream_401_zero_spend",
"covers": "quota_management.spend_tracking.cost_matrix.logs_cost",
"model": "gpt-5.6",
"request": {
"model": "$MODEL",
"messages": [
{
"role": "user",
"content": "proxy behaviour probe"
}
],
"stream": false
},
"response": {
"content_type": "application/json",
"status": 401,
"body": {
"error": {
"message": "scripted upstream failure 401",
"type": "server_error",
"code": "401"
}
}
},
"expected": {
"failure": {
"status": 401
}
}
},
{
"name": "gpt-5.6-upstream_500_stream_request_zero_spend",
"covers": "quota_management.spend_tracking.cost_matrix.logs_cost",
"model": "gpt-5.6",
"request": {
"model": "$MODEL",
"messages": [
{
"role": "user",
"content": "proxy behaviour probe"
}
],
"stream": true
},
"response": {
"content_type": "application/json",
"status": 500,
"body": {
"error": {
"message": "scripted upstream failure 500",
"type": "server_error",
"code": "500"
}
}
},
"expected": {
"failure": {
"status": 500
}
}
},
{
"name": "gpt-5.6-responses_upstream_500_zero_spend",
"covers": "quota_management.spend_tracking.cost_matrix.logs_cost",
"model": "gpt-5.6",
"endpoint": "/v1/responses",
"request": {
"model": "$MODEL",
"input": "proxy behaviour probe",
"stream": false
},
"response": {
"content_type": "application/json",
"status": 500,
"body": {
"error": {
"message": "scripted upstream failure 500",
"type": "server_error",
"code": "500"
}
}
},
"expected": {
"failure": {
"status": 500
}
}
},
{
"name": "claude-sonnet-5-messages_upstream_500_zero_spend",
"covers": "quota_management.spend_tracking.cost_matrix.logs_cost",
"model": "gpt-5.6",
"endpoint": "/v1/messages",
"request": {
"model": "$MODEL",
"messages": [
{
"role": "user",
"content": "proxy behaviour probe"
}
]
},
"response": {
"content_type": "application/json",
"status": 500,
"body": {
"error": {
"message": "scripted upstream failure 500",
"type": "server_error",
"code": "500"
}
}
},
"expected": {
"failure": {
"status": 500
}
}
},
{
"name": "gpt-5.6-fallback_billed_to_answering_deployment",
"covers": "quota_management.spend_tracking.routing.fallback_billing",
"model": "gpt-5.6",
"request": {
"model": "$MODEL",
"messages": [
{
"role": "user",
"content": "proxy behaviour probe"
}
],
"stream": false
},
"response": {
"content_type": "application/json",
"body": {
"id": "chatcmpl-$REQUEST_ID",
"object": "chat.completion",
"created": 1789788262,
"model": "gpt-5.6",
"choices": [
{
"index": 0,
"message": {
"role": "assistant",
"content": "scripted answer"
},
"finish_reason": "stop"
}
],
"usage": {
"prompt_tokens": 1840,
"completion_tokens": 412,
"total_tokens": 2252
}
}
},
"fallback_from": {
"content_type": "application/json",
"status": 500,
"body": {
"error": {
"message": "scripted upstream failure 500",
"type": "server_error",
"code": "500"
}
}
},
"expected": {
"spend": 0.008988,
"input_cost": 0.00322,
"output_cost": 0.005768,
"prompt_tokens": 1840,
"completion_tokens": 412
}
},
{
"name": "gpt-5.6-n_2_choices",
"covers": "quota_management.spend_tracking.cost_matrix.logs_cost",
"model": "gpt-5.6",
"request": {
"model": "$MODEL",
"messages": [
{
"role": "user",
"content": "proxy behaviour probe"
}
],
"stream": false
},
"response": {
"content_type": "application/json",
"body": {
"id": "chatcmpl-$REQUEST_ID",
"object": "chat.completion",
"created": 1789788262,
"model": "gpt-5.6",
"choices": [
{
"index": 0,
"message": {
"role": "assistant",
"content": "scripted answer"
},
"finish_reason": "stop"
},
{
"index": 1,
"message": {
"role": "assistant",
"content": "second choice"
},
"finish_reason": "stop"
}
],
"usage": {
"prompt_tokens": 1840,
"completion_tokens": 412,
"total_tokens": 2252
}
}
},
"expected": {
"spend": 0.008988,
"input_cost": 0.00322,
"output_cost": 0.005768,
"prompt_tokens": 1840,
"completion_tokens": 412
}
},
{
"name": "gpt-5.6-finish_reason_length",
"covers": "quota_management.spend_tracking.cost_matrix.logs_cost",
"model": "gpt-5.6",
"request": {
"model": "$MODEL",
"messages": [
{
"role": "user",
"content": "proxy behaviour probe"
}
],
"stream": false
},
"response": {
"content_type": "application/json",
"body": {
"id": "chatcmpl-$REQUEST_ID",
"object": "chat.completion",
"created": 1789788262,
"model": "gpt-5.6",
"choices": [
{
"index": 0,
"message": {
"role": "assistant",
"content": "truncated"
},
"finish_reason": "length"
}
],
"usage": {
"prompt_tokens": 1840,
"completion_tokens": 412,
"total_tokens": 2252
}
}
},
"expected": {
"spend": 0.008988,
"input_cost": 0.00322,
"output_cost": 0.005768,
"prompt_tokens": 1840,
"completion_tokens": 412
}
},
{
"name": "gpt-5.6-stream_usage_in_empty_choices_chunk",
"covers": "quota_management.spend_tracking.scripted_wire.logs_cost",
"model": "gpt-5.6",
"request": {
"model": "$MODEL",
"messages": [
{
"role": "user",
"content": "proxy behaviour probe"
}
],
"stream": true,
"stream_options": {
"include_usage": true
}
},
"response": {
"content_type": "text/event-stream",
"frames": [
"data: {\"id\":\"chatcmpl-$REQUEST_ID\",\"object\":\"chat.completion.chunk\",\"created\":1789788263,\"model\":\"gpt-5.6\",\"choices\":[{\"index\":0,\"delta\":{\"role\":\"assistant\"},\"finish_reason\":null}],\"usage\":null}",
"data: {\"id\":\"chatcmpl-$REQUEST_ID\",\"object\":\"chat.completion.chunk\",\"created\":1789788263,\"model\":\"gpt-5.6\",\"choices\":[{\"index\":0,\"delta\":{\"content\":\"scripted answer\"},\"finish_reason\":null}],\"usage\":null}",
"data: {\"id\":\"chatcmpl-$REQUEST_ID\",\"object\":\"chat.completion.chunk\",\"created\":1789788263,\"model\":\"gpt-5.6\",\"choices\":[],\"usage\":{\"prompt_tokens\":1840,\"completion_tokens\":412,\"total_tokens\":2252}}",
"data: [DONE]"
]
},
"expected": {
"spend": 0.008988,
"input_cost": 0.00322,
"output_cost": 0.005768,
"prompt_tokens": 1840,
"completion_tokens": 412,
"cost_header": false
}
},
{
"name": "gpt-5.6-stream_usage_in_last_delta_chunk",
"covers": "quota_management.spend_tracking.scripted_wire.logs_cost",
"model": "gpt-5.6",
"request": {
"model": "$MODEL",
"messages": [
{
"role": "user",
"content": "proxy behaviour probe"
}
],
"stream": true,
"stream_options": {
"include_usage": true
}
},
"response": {
"content_type": "text/event-stream",
"frames": [
"data: {\"id\":\"chatcmpl-$REQUEST_ID\",\"object\":\"chat.completion.chunk\",\"created\":1789788263,\"model\":\"gpt-5.6\",\"choices\":[{\"index\":0,\"delta\":{\"role\":\"assistant\"},\"finish_reason\":null}],\"usage\":null}",
"data: {\"id\":\"chatcmpl-$REQUEST_ID\",\"object\":\"chat.completion.chunk\",\"created\":1789788263,\"model\":\"gpt-5.6\",\"choices\":[{\"index\":0,\"delta\":{\"content\":\"scripted answer\"},\"finish_reason\":\"stop\"}],\"usage\":{\"prompt_tokens\":1840,\"completion_tokens\":412,\"total_tokens\":2252}}",
"data: [DONE]"
]
},
"expected": {
"spend": 0.008988,
"input_cost": 0.00322,
"output_cost": 0.005768,
"prompt_tokens": 1840,
"completion_tokens": 412,
"cost_header": false
}
},
{
"name": "gpt-5.6-unknown_model_response_model_unknown",
"covers": "quota_management.spend_tracking.cost_matrix.logs_cost",
"model": "gpt-5.6",
"deployment": {
"model": "openai/not-in-any-map-xyz"
},
"request": {
"model": "$MODEL",
"messages": [
{
"role": "user",
"content": "proxy behaviour probe"
}
],
"stream": false
},
"response": {
"content_type": "application/json",
"body": {
"id": "chatcmpl-$REQUEST_ID",
"object": "chat.completion",
"created": 1789788262,
"model": "not-in-any-map-xyz",
"choices": [
{
"index": 0,
"message": {
"role": "assistant",
"content": "scripted answer"
},
"finish_reason": "stop"
}
],
"usage": {
"prompt_tokens": 1840,
"completion_tokens": 412,
"total_tokens": 2252
}
}
},
"expected": {
"spend": 0.0,
"input_cost": 0.0,
"output_cost": 0.0,
"prompt_tokens": 1840,
"completion_tokens": 412,
"cost_header": false
}
},
{
"name": "gpt-5.6-unknown_model_response_model_known",
"covers": "quota_management.spend_tracking.cost_matrix.logs_cost",
"model": "gpt-5.6",
"deployment": {
"model": "openai/not-in-any-map-xyz"
},
"request": {
"model": "$MODEL",
"messages": [
{
"role": "user",
"content": "proxy behaviour probe"
}
],
"stream": false
},
"response": {
"content_type": "application/json",
"body": {
"id": "chatcmpl-$REQUEST_ID",
"object": "chat.completion",
"created": 1789788262,
"model": "gpt-5.6",
"choices": [
{
"index": 0,
"message": {
"role": "assistant",
"content": "scripted answer"
},
"finish_reason": "stop"
}
],
"usage": {
"prompt_tokens": 1840,
"completion_tokens": 412,
"total_tokens": 2252
}
}
},
"expected": {
"spend": 0.008988,
"input_cost": 0.00322,
"output_cost": 0.005768,
"prompt_tokens": 1840,
"completion_tokens": 412
}
},
{
"name": "gpt-5.6-chat_request_to_embedding_entry",
"covers": "quota_management.spend_tracking.cost_matrix.logs_cost",
"model": "text-embedding-3-large",
"deployment": {
"model": "openai/text-embedding-3-large"
},
"request": {
"model": "$MODEL",
"messages": [
{
"role": "user",
"content": "proxy behaviour probe"
}
],
"stream": false
},
"response": {
"content_type": "application/json",
"body": {
"id": "chatcmpl-$REQUEST_ID",
"object": "chat.completion",
"created": 1789788262,
"model": "text-embedding-3-large",
"choices": [
{
"index": 0,
"message": {
"role": "assistant",
"content": "scripted answer"
},
"finish_reason": "stop"
}
],
"usage": {
"prompt_tokens": 1840,
"completion_tokens": 412,
"total_tokens": 2252
}
}
},
"expected": {
"spend": 0.0002392,
"input_cost": 0.0002392,
"output_cost": 0.0,
"prompt_tokens": 1840,
"completion_tokens": 412
}
},
{
"name": "gpt-5.6-client_disconnect_mid_stream",
"covers": "quota_management.spend_tracking.scripted_wire.client_disconnect",
"model": "gpt-5.6",
"request": {
"model": "$MODEL",
"messages": [
{
"role": "user",
"content": "proxy behaviour probe"
}
],
"stream": true,
"stream_options": {
"include_usage": true
}
},
"response": {
"content_type": "text/event-stream",
"frames": [
"data: {\"id\":\"chatcmpl-$REQUEST_ID\",\"object\":\"chat.completion.chunk\",\"created\":1789788263,\"model\":\"gpt-5.6\",\"choices\":[{\"index\":0,\"delta\":{\"content\":\"frame-0\"},\"finish_reason\":null}],\"usage\":null}",
"data: {\"id\":\"chatcmpl-$REQUEST_ID\",\"object\":\"chat.completion.chunk\",\"created\":1789788263,\"model\":\"gpt-5.6\",\"choices\":[{\"index\":0,\"delta\":{\"content\":\"frame-1\"},\"finish_reason\":null}],\"usage\":null}",
"data: {\"id\":\"chatcmpl-$REQUEST_ID\",\"object\":\"chat.completion.chunk\",\"created\":1789788263,\"model\":\"gpt-5.6\",\"choices\":[{\"index\":0,\"delta\":{\"content\":\"frame-2\"},\"finish_reason\":null}],\"usage\":null}",
"data: {\"id\":\"chatcmpl-$REQUEST_ID\",\"object\":\"chat.completion.chunk\",\"created\":1789788263,\"model\":\"gpt-5.6\",\"choices\":[{\"index\":0,\"delta\":{\"content\":\"frame-3\"},\"finish_reason\":null}],\"usage\":null}",
"data: {\"id\":\"chatcmpl-$REQUEST_ID\",\"object\":\"chat.completion.chunk\",\"created\":1789788263,\"model\":\"gpt-5.6\",\"choices\":[{\"index\":0,\"delta\":{\"content\":\"frame-4\"},\"finish_reason\":null}],\"usage\":null}",
"data: {\"id\":\"chatcmpl-$REQUEST_ID\",\"object\":\"chat.completion.chunk\",\"created\":1789788263,\"model\":\"gpt-5.6\",\"choices\":[{\"index\":0,\"delta\":{\"content\":\"frame-5\"},\"finish_reason\":null}],\"usage\":null}",
"data: {\"id\":\"chatcmpl-$REQUEST_ID\",\"object\":\"chat.completion.chunk\",\"created\":1789788263,\"model\":\"gpt-5.6\",\"choices\":[{\"index\":0,\"delta\":{\"content\":\"frame-6\"},\"finish_reason\":null}],\"usage\":null}",
"data: {\"id\":\"chatcmpl-$REQUEST_ID\",\"object\":\"chat.completion.chunk\",\"created\":1789788263,\"model\":\"gpt-5.6\",\"choices\":[{\"index\":0,\"delta\":{\"content\":\"frame-7\"},\"finish_reason\":null}],\"usage\":null}",
"data: {\"id\":\"chatcmpl-$REQUEST_ID\",\"object\":\"chat.completion.chunk\",\"created\":1789788263,\"model\":\"gpt-5.6\",\"choices\":[{\"index\":0,\"delta\":{\"content\":\"frame-8\"},\"finish_reason\":null}],\"usage\":null}",
"data: {\"id\":\"chatcmpl-$REQUEST_ID\",\"object\":\"chat.completion.chunk\",\"created\":1789788263,\"model\":\"gpt-5.6\",\"choices\":[{\"index\":0,\"delta\":{\"content\":\"frame-9\"},\"finish_reason\":null}],\"usage\":null}",
"data: {\"id\":\"chatcmpl-$REQUEST_ID\",\"object\":\"chat.completion.chunk\",\"created\":1789788263,\"model\":\"gpt-5.6\",\"choices\":[{\"index\":0,\"delta\":{\"content\":\"frame-10\"},\"finish_reason\":null}],\"usage\":null}",
"data: {\"id\":\"chatcmpl-$REQUEST_ID\",\"object\":\"chat.completion.chunk\",\"created\":1789788263,\"model\":\"gpt-5.6\",\"choices\":[{\"index\":0,\"delta\":{\"content\":\"frame-11\"},\"finish_reason\":null}],\"usage\":null}",
"data: {\"id\":\"chatcmpl-$REQUEST_ID\",\"object\":\"chat.completion.chunk\",\"created\":1789788263,\"model\":\"gpt-5.6\",\"choices\":[{\"index\":0,\"delta\":{\"content\":\"frame-12\"},\"finish_reason\":null}],\"usage\":null}",
"data: {\"id\":\"chatcmpl-$REQUEST_ID\",\"object\":\"chat.completion.chunk\",\"created\":1789788263,\"model\":\"gpt-5.6\",\"choices\":[{\"index\":0,\"delta\":{\"content\":\"frame-13\"},\"finish_reason\":null}],\"usage\":null}",
"data: {\"id\":\"chatcmpl-$REQUEST_ID\",\"object\":\"chat.completion.chunk\",\"created\":1789788263,\"model\":\"gpt-5.6\",\"choices\":[{\"index\":0,\"delta\":{\"content\":\"frame-14\"},\"finish_reason\":null}],\"usage\":null}",
"data: {\"id\":\"chatcmpl-$REQUEST_ID\",\"object\":\"chat.completion.chunk\",\"created\":1789788263,\"model\":\"gpt-5.6\",\"choices\":[{\"index\":0,\"delta\":{\"content\":\"frame-15\"},\"finish_reason\":null}],\"usage\":null}",
"data: {\"id\":\"chatcmpl-$REQUEST_ID\",\"object\":\"chat.completion.chunk\",\"created\":1789788263,\"model\":\"gpt-5.6\",\"choices\":[{\"index\":0,\"delta\":{\"content\":\"frame-16\"},\"finish_reason\":null}],\"usage\":null}",
"data: {\"id\":\"chatcmpl-$REQUEST_ID\",\"object\":\"chat.completion.chunk\",\"created\":1789788263,\"model\":\"gpt-5.6\",\"choices\":[{\"index\":0,\"delta\":{\"content\":\"frame-17\"},\"finish_reason\":null}],\"usage\":null}",
"data: {\"id\":\"chatcmpl-$REQUEST_ID\",\"object\":\"chat.completion.chunk\",\"created\":1789788263,\"model\":\"gpt-5.6\",\"choices\":[{\"index\":0,\"delta\":{\"content\":\"frame-18\"},\"finish_reason\":null}],\"usage\":null}",
"data: {\"id\":\"chatcmpl-$REQUEST_ID\",\"object\":\"chat.completion.chunk\",\"created\":1789788263,\"model\":\"gpt-5.6\",\"choices\":[{\"index\":0,\"delta\":{\"content\":\"frame-19\"},\"finish_reason\":null}],\"usage\":null}",
"data: [DONE]"
],
"frame_delay_ms": 200
},
"disconnect_after_frames": 3,
"expected": {
"recount": {
"input_cost_per_token": 1.75e-06,
"output_cost_per_token": 1.4e-05
}
}
}
]
}

191
tests/integration/cost_calculation/test_cost_tracking.py
View file

@ -3,10 +3,12 @@
from __future__ import annotations
import io
from itertools import islice
import json
from hashlib import sha256
import struct
from typing import Final, cast
import uuid
import wave
import zlib
@ -18,10 +20,13 @@ from integration._support.client import JSON_OBJECT, Gateway
from integration._support.upstream import delete_scenario, register_scenario
from integration.cost_calculation.conftest import (
CostBreakdown,
CostRow,
approx_equal,
assert_total_is_sum_of_components,
poll_cost_row,
poll_failure_row,
poll_rollups,
poll_rows,
register_scenario_deployment,
)
from integration.cost_calculation.cost_tracking_case import (
@ -179,11 +184,47 @@ def _assert_breakdown(
)
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)
@pytest.mark.parametrize("case", _CASES)
def test_case_bills_expected_cost(gateway: Gateway, case: CostTrackingTestCase) -> None:
marker: Final = sha256(case.name.encode()).hexdigest()[:12]
with gateway.scenario() as scenario:
key: Final = scenario.key()
expected: Final = case.expected
team_id: Final = scenario.team() if isinstance(expected, ExactExpected) and expected.rollups else None
user_id: Final = (
scenario.user(team_id=team_id)
if team_id is not None
else None
)
key: Final = (
scenario.key(team_id=team_id, user_id=user_id)
if team_id is not None and user_id is not None
else scenario.key()
)
passthrough_provider: Final = case.passthrough_provider
scenario_id: Final = f"sc-{marker}-{sha256(key.encode()).hexdigest()[:12]}"
scenario_handle: Final = (
@ -193,21 +234,73 @@ def test_case_bills_expected_cost(gateway: Gateway, case: CostTrackingTestCase)
)
if scenario_handle is not None:
scenario.cleanups.callback(delete_scenario, scenario_handle)
deployment: Final = (
register_scenario_deployment(scenario, case, marker, key)
if passthrough_provider is None
else None
)
fallback_deployment: Final = (
register_scenario_deployment(
scenario,
case,
marker,
key,
response=case.fallback_from,
marker_suffix="-fb",
)
if case.fallback_from is not None
else None
)
if isinstance(expected, ExactExpected) and expected.rollups:
assert deployment is not None
rollup_deployments: Final = tuple(
register_scenario_deployment(
scenario,
case,
marker,
key,
marker_suffix=f"-r{index}",
model_name=deployment.model_name,
)
for index in (2, 3)
)
assert len(rollup_deployments) == 2
model_name: Final = (
case.model
if passthrough_provider in {"gemini", "anthropic"}
else register_scenario_deployment(scenario, case, marker, key)
else deployment.model_name if deployment is not None else None
)
assert model_name is not None
request_model: Final = (
case.model.rsplit("/", 1)[-1]
if passthrough_provider in {"gemini", "anthropic"}
else model_name
else fallback_deployment.model_name if fallback_deployment is not None else model_name
)
request_body: Final = JSON_OBJECT.validate_python(
base_request_values: Final = (
_replace_model(case.request, request_model)
if passthrough_provider is not None
else {**case.request, "model": model_name}
)
end_user_id: Final = (
f"end-user-{uuid.uuid4()}"
if isinstance(expected, ExactExpected) and expected.rollups
else None
)
request_body: Final = JSON_OBJECT.validate_python(
{
**base_request_values,
**(
{"model": fallback_deployment.model_name, "fallbacks": [model_name]}
if fallback_deployment is not None
else {}
),
**(
{"user": end_user_id, "cache": {"no-cache": True}}
if end_user_id is not None
else {}
),
}
)
request_headers: Final = (
{
"x-pass-x-scripted-scenario": scenario_id,
@ -225,12 +318,45 @@ def test_case_bills_expected_cost(gateway: Gateway, case: CostTrackingTestCase)
if passthrough_provider is not None
else case.endpoint
)
response: Final = (
_multipart_request(gateway, case, model_name, key)
if case.upload is not None
else gateway.request("POST", request_path, request_body, key=key, headers=request_headers)
if case.disconnect_after_frames is not None:
with gateway.client.stream(
"POST",
request_path,
json=request_body,
headers={"Authorization": f"Bearer {key}", **request_headers},
) as stream_response:
frames: Final = tuple(
islice(
(line for line in stream_response.iter_lines() if line.startswith("data:")),
case.disconnect_after_frames,
)
)
assert len(frames) == case.disconnect_after_frames
row: Final = poll_cost_row(key)
assert isinstance(expected, RecountExpected)
if expected.prompt_tokens is not None:
assert row.prompt_tokens == expected.prompt_tokens
if expected.completion_tokens is not None:
assert row.completion_tokens == expected.completion_tokens
assert row.prompt_tokens is not None and row.prompt_tokens > 0
assert row.completion_tokens is not None and row.completion_tokens > 0
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)
assert row.breakdown is not None
assert_total_is_sum_of_components(row, row.breakdown, case.name)
return
responses: Final = tuple(
(
_multipart_request(gateway, case, model_name, key)
if case.upload is not None
else gateway.request("POST", request_path, request_body, key=key, headers=request_headers)
)
for _ in range(3 if isinstance(expected, ExactExpected) and expected.rollups else 1)
)
if isinstance(case.expected, FailureExpected):
response: Final = responses[0]
if isinstance(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]}"
@ -245,8 +371,9 @@ def test_case_bills_expected_cost(gateway: Gateway, case: CostTrackingTestCase)
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)
row: Final = poll_cost_row(key)
if isinstance(case.expected, RecountExpected):
rows: Final = poll_rows(key) if len(responses) > 1 else (poll_cost_row(key),)
if isinstance(expected, RecountExpected):
row: Final = rows[0]
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}"
)
@ -263,8 +390,12 @@ def test_case_bills_expected_cost(gateway: Gateway, case: CostTrackingTestCase)
assert breakdown is not None, f"{case.name}: no cost_breakdown persisted"
assert_total_is_sum_of_components(row, breakdown, case.name)
return
expected: Final = case.expected
assert isinstance(expected, ExactExpected)
if fallback_deployment is not None:
assert deployment is not None
assert len(rows) == 1
assert rows[0].status == "success"
assert rows[0].model_id == deployment.identity
if isinstance(case.response, BinaryResponse):
header: Final = response.headers.get("x-litellm-response-cost")
if header is not None:
@ -281,20 +412,22 @@ def test_case_bills_expected_cost(gateway: Gateway, case: CostTrackingTestCase)
assert approx_equal(float(header), expected.spend), (
f"{case.name}: x-litellm-response-cost {header} != expected {expected.spend}"
)
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)
for row in rows:
_assert_exact(case, 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
rollups: Final = poll_rollups(key, team_id, user_id, end_user_id)
target_spend: Final = expected.spend * 3
assert approx_equal(rollups.key_spend, target_spend)
assert approx_equal(rollups.team_spend, target_spend)
assert approx_equal(rollups.user_spend, target_spend)
assert approx_equal(rollups.end_user_spend, target_spend)
assert approx_equal(rollups.daily_user.spend or 0.0, target_spend)
assert approx_equal(rollups.daily_team.spend or 0.0, target_spend)
assert rollups.daily_user.prompt_tokens == expected.prompt_tokens * 3
assert rollups.daily_user.completion_tokens == expected.completion_tokens * 3
assert rollups.daily_user.api_requests == 3
assert rollups.daily_team.prompt_tokens == expected.prompt_tokens * 3
assert rollups.daily_team.completion_tokens == expected.completion_tokens * 3
assert rollups.daily_team.api_requests == 3