scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-09-21 00:21:49 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 74

test(integration): move scripted-provider cost suite into cost shard

Browse source 
Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>
This commit is contained in:
kerry 2026-09-19 00:16:12 +00:00
parent d68e20e586
commit 69f9106759
23 changed files with 1586 additions and 652 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

2
.circleci/config.yml
View file

@ -3009,7 +3009,7 @@ workflows:
name: integration-<< matrix.suite >>
matrix:
parameters:
suite: [management, accounting, database, providers, extensions, sdk, browser]
suite: [management, accounting, database, providers, extensions, sdk, cost, browser]
filters:
branches:
only:

33
.circleci/scripts/run_integration.sh
View file

@ -11,6 +11,7 @@ results="test-results/integration-${suite}"
mkdir -p "$results"
integration_identity="$(.venv/bin/python -c 'import uuid; print(uuid.uuid4().hex)')"
upstream_pid=""
scripted_provider_pid=""
proxy_pid=""
peer_pid=""
launched_pid=""
@ -22,9 +23,9 @@ cleanup() {
original_status=$?
trap - EXIT INT TERM
sudo .venv/bin/python .circleci/scripts/stop_integration_processes.py \
"$integration_identity" "$(id -u)" "$proxy_pid" "$peer_pid" "$upstream_pid" \
"$integration_identity" "$(id -u)" "$proxy_pid" "$peer_pid" "$upstream_pid" "$scripted_provider_pid" \
> "$results/process-cleanup.txt" 2>&1 || original_status=1
for owned_pid in "$peer_pid" "$proxy_pid" "$upstream_pid"; do
for owned_pid in "$peer_pid" "$proxy_pid" "$upstream_pid" "$scripted_provider_pid"; do
if [ -n "$owned_pid" ]; then
kill -- "-$owned_pid" 2>/dev/null || true
for _ in {1..50}; do
@ -69,6 +70,7 @@ export STORE_MODEL_IN_DB=True AWS_EC2_METADATA_DISABLED=true DO_NOT_TRACK=1
export INTEGRATION_PROXY_URL=http://127.0.0.1:4000
export INTEGRATION_PEER_URL=""
export INTEGRATION_UPSTREAM_URL=http://127.0.0.1:8190
export INTEGRATION_SCRIPTED_PROVIDER_URL=""
export INTEGRATION_MASTER_KEY="$LITELLM_MASTER_KEY"
export LITELLM_UI_PATH="$PWD/litellm/proxy/_experimental/out"
if [ "$suite" = browser ]; then
@ -108,13 +110,37 @@ awk '$3 == "REJECT" && $1 > 0 { rejected=1 } END { exit !rejected }' "$results/e
setsid env -i PATH="$PATH" HOME="$HOME" PYTHONPATH="$PYTHONPATH" INTEGRATION_RUN_ID="$integration_identity" \
.venv/bin/python -m integration._support.upstream > "$results/upstream.log" 2>&1 &
upstream_pid=$!
if [ "$suite" = cost ]; then
export INTEGRATION_SCRIPTED_PROVIDER_URL=http://127.0.0.1:8191
setsid env -i PATH="$PATH" HOME="$HOME" PYTHONPATH="$PYTHONPATH" INTEGRATION_RUN_ID="$integration_identity" \
.venv/bin/python -m integration._support.scripted_provider --port 8191 \
> "$results/scripted-provider.log" 2>&1 &
scripted_provider_pid=$!
for _ in {1..90}; do
if curl --noproxy '*' -fsS "$INTEGRATION_SCRIPTED_PROVIDER_URL/health" >/dev/null 2>&1; then
break
fi
sleep 1
done
curl --noproxy '*' -fsS "$INTEGRATION_SCRIPTED_PROVIDER_URL/health" >/dev/null
fi
start_proxy() {
local port="$1"
local log_name="$2"
local -a cost_map_env
if [ "$suite" = cost ]; then
cost_map_env=(
"LITELLM_MODEL_COST_MAP_URL=$INTEGRATION_SCRIPTED_PROVIDER_URL/_cost_map"
"MODEL_COST_MAP_MIN_MODEL_COUNT=1"
"MODEL_COST_MAP_MAX_SHRINK_RATIO=0"
)
else
cost_map_env=("LITELLM_LOCAL_MODEL_COST_MAP=True")
fi
setsid env -i PATH="$PATH" HOME="$HOME" PYTHONPATH="$PYTHONPATH" INTEGRATION_RUN_ID="$integration_identity" \
DATABASE_URL="$DATABASE_URL" REDIS_HOST="$REDIS_HOST" REDIS_PORT="$REDIS_PORT" \
LITELLM_MASTER_KEY="$LITELLM_MASTER_KEY" LITELLM_SALT_KEY="$LITELLM_SALT_KEY" LITELLM_UI_PATH="$LITELLM_UI_PATH" PROXY_BASE_URL="http://127.0.0.1:$port" \
LITELLM_MODE=PRODUCTION LITELLM_LOCAL_MODEL_COST_MAP=True STORE_MODEL_IN_DB=True \
LITELLM_MODE=PRODUCTION STORE_MODEL_IN_DB=True "${cost_map_env[@]}" \
AWS_EC2_METADATA_DISABLED=true DO_NOT_TRACK=1 \
.venv/bin/python -m integration._support.proxy --config tests/integration/proxy_config.yaml \
--host 127.0.0.1 --port "$port" --num_workers 1 --telemetry False \
@ -160,6 +186,7 @@ timeout --signal=TERM --kill-after=20s 11m env -i PATH="$PATH" HOME="$HOME" PYTH
DATABASE_URL="$DATABASE_URL" REDIS_HOST="$REDIS_HOST" REDIS_PORT="$REDIS_PORT" \
INTEGRATION_PROXY_URL="$INTEGRATION_PROXY_URL" INTEGRATION_PEER_URL="$INTEGRATION_PEER_URL" \
INTEGRATION_UPSTREAM_URL="$INTEGRATION_UPSTREAM_URL" \
INTEGRATION_SCRIPTED_PROVIDER_URL="$INTEGRATION_SCRIPTED_PROVIDER_URL" \
INTEGRATION_MASTER_KEY="$INTEGRATION_MASTER_KEY" LITELLM_MODE=PRODUCTION \
INTEGRATION_SEED="$INTEGRATION_SEED" \
INTEGRATION_ORDER_SEED="$INTEGRATION_ORDER_SEED" \

5
.circleci/scripts/wait_integration_services.py
View file

@ -9,6 +9,7 @@ from redis import Redis
def main() -> None:
primary: Final = os.environ["INTEGRATION_PROXY_URL"]
peer: Final = os.environ.get("INTEGRATION_PEER_URL")
scripted_provider: Final = os.environ.get("INTEGRATION_SCRIPTED_PROVIDER_URL") or None
proxies: Final = (primary, peer) if peer else (primary,)
deadline: Final = time.monotonic() + 90
headers: Final = {"Authorization": f"Bearer {os.environ['INTEGRATION_MASTER_KEY']}"}
@ -19,6 +20,10 @@ def main() -> None:
try:
ready: Final = (
client.get(f"{os.environ['INTEGRATION_UPSTREAM_URL']}/health").status_code == 200
and (
scripted_provider is None
or client.get(f"{scripted_provider}/health").status_code == 200
)
and all(client.get(f"{url}/health/readiness").status_code == 200 for url in proxies)
)
if ready:

3
tests/e2e/CLAUDE.md
View file

@ -21,7 +21,6 @@ Each subdirectory under `tests/e2e/` is one suite, scoped to an endpoint family
- `load/` - performance-category tests, kept OUT of the main suite: throughput/load SLO tests are a different testing category from functional e2e (variance-driven, historically flaky) and live outside this suite until re-implemented as their own pipeline (LIT-5163); do not add a live load test that runs in the default collection. What lives here: the weekly session-anomaly test (`test_weekly_session_anomaly_e2e.py`, Claude Code-shaped multi-turn sessions against real providers with ceilings on error rate, cache read/write, turn time, and spend; marked `weekly` and deselected unless `E2E_WEEKLY_ANOMALY` is set, driven by `.github/workflows/weekly_load_anomaly.yml`), the Redis chaos test (`test_redis_chaos_e2e.py`, locust load against mock deployments split round robin over `/chat/completions` and `/v1/messages`, one endpoint per simulated user, with `CLIENT PAUSE ALL` on the proxy's Redis mid-run to simulate it being down outright, asserting zero failed requests on every endpoint, budgeting RSS and CPU-per-request as ratios against the same run's healthy phase, and holding p50/p90/p99 latency and log-bytes-per-request to flat ceilings (a ratio cannot bound those two: an open breaker skips Redis instead of waiting on it, so the chaos phase can measure cheaper than baseline while still being far slower than a user should see); needs a proxy booted from `gateway/redis_chaos_ci_config.yml` on the same host with `E2E_PROXY_PID` and `E2E_PROXY_LOG` set, marked `redis_chaos`, deselected unless `E2E_REDIS_CHAOS` is set and excluded from the per-PR selector like the rest of `load/`, driven by `.github/workflows/test-e2e-redis-chaos.yml` and by the Buildkite `e2e-redis-chaos` step in project-releaser, which runs the proxy, Postgres and Valkey co-located with pytest in one pod and sets the opt-in), and markerless harness unit tests for the locust, process-usage, and session-anomaly aggregation logic
- `other/` - the holding-pen suite for the `other.*` registry cluster with no home of its own yet: the master-key auth gate, JWT auth (access tokens issued by a real Keycloak realm, `idp.py` plus `idp_realm.json`, whose JWKS the proxy's `JWT_PUBLIC_KEY_URL` points at; see CONTRIBUTING.md for the start command and config block), and the process-lifecycle health probes (liveness, public readiness, authenticated readiness diagnostics). Promote a cluster out once it is large/stable enough for its own suite
- `gateway/` - proxy configuration only (`litellm-config.yml`); no tests
- `cost_calculation/` - cost accounting against a dedicated proxy whose whole model cost map is the test-owned `tests/e2e/cost_map.json` (loaded via `LITELLM_MODEL_COST_MAP_URL`), with provider calls answered by the scripted-provider sidecar in `scripted_provider.py`; every cost-map entry is a deployment and the cases plus asserted goldens are data in `cases.json` (each `pricing` case owns (model, cost key) pairs via `owns`/`fallback_for` so every rate key present on each map entry has exactly one owning case, and each carries a literal `expected` cell per map key; `transport` cases list `models` and exercise token counting only; `cost_matrix.matrix_data_errors()` runs at collection time so a key absent from the cost map, an unowned or double-owned (model, rate key) pair, an `owns` key absent on all of the case's models, or a `fallback_for` key present on a case model fails the suite's collection loudly), deselected unless `E2E_COST_MAP_STACK` is set, driven by the Buildkite `e2e-cost-calculation` step in project-releaser, which runs a proxy booted from `gateway/cost_calculation_ci_config.yml` with `MODEL_COST_MAP_MIN_MODEL_COUNT=1` and `MODEL_COST_MAP_MAX_SHRINK_RATIO=0` (the 21-entry test map trips the fetched-cost-map integrity check at the defaults), Postgres and the scripted provider co-located with pytest in one pod and sets the opt-in
- `claude_code/` - the Claude Code compatibility matrix: drives the real `claude` CLI (and HTTP probes) against a proxy for each feature x provider cell, reporting tagged-union outcomes via the `compat_result` fixture; ships its own driver/builder/publisher plus `_*_unit_tests/` trees. The HTTP probes ride the shared transport (`ProxyClient.count_tokens` / `ProxyClient.messages`); the CLI-driving path stays bespoke
- `ui/` - the Admin UI browser suite: Playwright in TypeScript, driving the dashboard served by a live proxy on port 4000 (seeded postgres + mock LLM upstream; see its `run_e2e.sh`). It is a self-contained npm package with its own lockfile and does not use the Python harness, pytest markers, or the shared transport; the Python rules in this file (typed models, `Result` unions, basedpyright zero-error gate) do not apply inside it. Its only Python file, `fixtures/mock_llm_server/server.py`, is excluded from the e2e basedpyright gate via the root `pyrightconfig.json`
@ -222,7 +221,7 @@ other.<area>.<case>.<assertion>
```
## Hard Rules
- no unit tests of any kind under `tests/e2e`. a product feature is proven end to end against a live proxy, never with a unit test, and the harness itself is not unit-tested here either. no monkeypatching or mock tests; the one carve-out is a scripted upstream served through a real HTTP sidecar (the cost_calculation suite's scripted provider), allowed because provider-response-shape coverage needs a controlled usage payload and every hop from the proxy's upstream call to the spend row still executes for real. if a contributor asks you to write an end to end test, do NOT stage a unit test with it; if you find a product gap, call it out in the PR description
- no unit tests of any kind under `tests/e2e`. a product feature is proven end to end against a live proxy, never with a unit test, and the harness itself is not unit-tested here either. no monkeypatching or mock tests. if a contributor asks you to write an end to end test, do NOT stage a unit test with it; if you find a product gap, call it out in the PR description
- use model management endpoints to create new models for a test. this could be in a conftest / inline for each test. ask the user what they want.

8
tests/e2e/conftest.py
View file

@ -25,7 +25,6 @@ import requests
from e2e_config import (
CLI_DETERMINISM_OPT_IN_ENV,
CONTROL_PLANE_BASE_URL,
COST_MAP_OPT_IN_ENV,
FIXTURE_DIR,
FIXTURE_MODE_RAW,
MANAGED_FILES_OPT_IN_ENV,
@ -56,7 +55,6 @@ OPT_IN_MARKERS: Final = MappingProxyType(
"managed_files": MANAGED_FILES_OPT_IN_ENV,
"prompt_caching_stack": PROMPT_CACHING_OPT_IN_ENV,
"redis_chaos": REDIS_CHAOS_OPT_IN_ENV,
"cost_map_stack": COST_MAP_OPT_IN_ENV,
"cli_determinism": CLI_DETERMINISM_OPT_IN_ENV,
}
)
@ -134,12 +132,6 @@ def pytest_configure(config: pytest.Config) -> None:
"redis_chaos: load test that pauses the proxy's Redis outright mid-run; needs a proxy booted from "
"gateway/redis_chaos_ci_config.yml on the same host, and is deselected unless E2E_REDIS_CHAOS is set",
)
config.addinivalue_line(
"markers",
"cost_map_stack: needs a proxy whose whole cost map is tests/e2e/cost_map.json "
"(LITELLM_MODEL_COST_MAP_URL) plus a scripted-provider sidecar; deselected unless "
"E2E_COST_MAP_STACK is set",
)
def pytest_sessionstart(session: pytest.Session) -> None:

185
tests/e2e/cost_calculation/conftest.py
View file

@ -1,185 +0,0 @@
"""Cost-calculation suite fixtures.
Runs against a dedicated proxy whose whole model cost map is the test-owned
``tests/e2e/cost_map.json`` (LITELLM_MODEL_COST_MAP_URL); every map entry is a
deployment under test, and the request shapes plus asserted goldens live in
``cases.json``. Provider calls are answered by the
scripted-provider sidecar (``scripted_provider.py``), registered per scenario
over its control API.
The proxy must also run with ``MODEL_COST_MAP_MIN_MODEL_COUNT=1`` and
``MODEL_COST_MAP_MAX_SHRINK_RATIO=0``: the 21-entry test map trips the
fetched-cost-map integrity check (too few models, large shrink versus the
bundled map) at those env vars' defaults.
Deselected unless E2E_COST_MAP_STACK is set (marker `cost_map_stack`).
"""
from __future__ import annotations
import functools
import importlib.util
import json
import sys
from collections.abc import Callable, Mapping
from dataclasses import dataclass
from pathlib import Path
from types import ModuleType
from typing import Final, Protocol, cast
import pytest
from cryptography.hazmat.primitives import serialization
from cryptography.hazmat.primitives.asymmetric import rsa
from cost_matrix import Case, FrontierModel
from e2e_config import COST_MAP_PROXY_URL, SCRIPTED_PROVIDER_PROXY_BASE
from lifecycle import ResourceManager
from models import LiteLLMParamsBody, ModelInfoBody, ModelNewBody
from proxy_client import ProxyClient, build_proxy_client
from scripted_client import ScenarioHandle, delete_scenario, register_scenario
from scripted_provider import Scenario
def _load_cost_rows() -> ModuleType:
"""Load quota_management/spend_tracking/cost_rows.py by path (the e2e tree
has no package layout), the same trick the mcp suite uses for
logging/datadog_reader.py."""
path: Final = (
Path(__file__).resolve().parent.parent
/ "quota_management"
/ "spend_tracking"
/ "cost_rows.py"
)
name: Final = "e2e_spend_tracking_cost_rows"
spec: Final = importlib.util.spec_from_file_location(name, path)
assert spec is not None and spec.loader is not None
module: Final = importlib.util.module_from_spec(spec)
sys.modules[name] = module
spec.loader.exec_module(module)
return module
class SpendCostBreakdown(Protocol):
input_cost: float | None
output_cost: float | None
cache_read_cost: float | None
cache_creation_cost: float | None
reasoning_cost: float | None
tool_usage_cost: float | None
total_cost: float | None
service_tier: str | None
def model_dump(self) -> Mapping[str, object]: ...
class SpendRowMetadata(Protocol):
cost_breakdown: SpendCostBreakdown | None
class SpendCostRow(Protocol):
"""The slice of spend_tracking.cost_rows.CostRow this suite reads."""
spend: float | None
prompt_tokens: int | None
completion_tokens: int | None
metadata: SpendRowMetadata | None
@property
def breakdown(self) -> SpendCostBreakdown: ...
class CostRowsModule(Protocol):
"""cost_rows.py loaded by path has no importable name for basedpyright, so
its surface is declared here and reached through a single cast."""
approx_equal: Callable[[float, float], bool]
assert_total_is_sum_of_components: Callable[[SpendCostRow], None]
poll_cost_row_where: Callable[
[ProxyClient, str, Callable[[SpendCostRow], bool]], SpendCostRow | None
]
cost_rows: Final[CostRowsModule] = cast( # cast-ok: cost_rows.py is loaded by path, so basedpyright has no importable name for it; its surface is declared in CostRowsModule
CostRowsModule, _load_cost_rows()
)
@dataclass(frozen=True, slots=True)
class CostCalcClient:
"""The suite's client: a ProxyClient pointed at the cost-map proxy pod."""
proxy: ProxyClient
@pytest.fixture(scope="session")
def client() -> CostCalcClient:
proxy: Final = build_proxy_client(
base_url=COST_MAP_PROXY_URL,
control_plane_base_url=COST_MAP_PROXY_URL,
replica_urls=(COST_MAP_PROXY_URL,),
)
return CostCalcClient(proxy=proxy)
@functools.cache
def _vertex_private_key_pem() -> str:
return rsa.generate_private_key(public_exponent=65537, key_size=2048).private_bytes(
serialization.Encoding.PEM,
serialization.PrivateFormat.PKCS8,
serialization.NoEncryption(),
).decode()
def _vertex_service_account_json() -> str:
"""A service-account credential JSON whose token_uri is the sidecar's
/_oauth/token route: the proxy's google-auth refresh then gets a scripted
access token without touching Google."""
return json.dumps(
{
"type": "service_account",
"project_id": "cc-scripted-project",
"private_key_id": "scripted",
"private_key": _vertex_private_key_pem(),
"client_email": "scripted@cc-scripted-project.iam.gserviceaccount.com",
"client_id": "0",
"auth_uri": f"{SCRIPTED_PROVIDER_PROXY_BASE}/_oauth/authorize",
"token_uri": f"{SCRIPTED_PROVIDER_PROXY_BASE}/_oauth/token",
}
)
def register_scenario_deployment(
client: CostCalcClient,
resources: ResourceManager,
model: FrontierModel,
case: Case,
marker: str,
) -> tuple[str, ScenarioHandle]:
"""Register the case's scenario on the sidecar plus a deployment pointed at
it; both are torn down by ``resources``. Returns the callable model_name."""
scenario: Final[Scenario] = case.scenario(
scenario_id=f"sc-{marker}", model=model, text=f"scripted answer {marker}"
)
handle: Final = register_scenario(scenario)
resources.defer(lambda: delete_scenario(handle))
model_name: Final = f"{model.model_name}-{marker}"
params: Final = {
"model": model.litellm_model,
"api_key": model.api_key,
"api_base": handle.api_base(),
**model.litellm_params,
**(
{"vertex_credentials": _vertex_service_account_json()}
if model.wire == "vertex_generate"
else {}
),
}
model_id: Final = client.proxy.register_model(
ModelNewBody(
model_name=model_name,
litellm_params=LiteLLMParamsBody.model_validate(params),
model_info=ModelInfoBody(base_model=model.base_model),
)
)
resources.defer(lambda: client.proxy.delete_model(model_id))
return model_name, handle

64
tests/e2e/cost_calculation/scripted_client.py
View file

@ -1,64 +0,0 @@
"""Client side of the scripted-provider sidecar: register scenarios over its
control API through the shared transport helpers and get back a handle whose
``api_base`` is what a /model/new deployment should register for the proxy to
reach the scripted wire."""
from __future__ import annotations
from dataclasses import dataclass
from typing import Final
from e2e_config import SCRIPTED_PROVIDER_CONTROL_URL, SCRIPTED_PROVIDER_PROXY_BASE
from e2e_http import URL, NoBody, unwrap, post
from e2e_http import delete as http_delete
from scripted_provider import (
WIRE_MOUNTS,
Scenario,
ScenarioDeleted,
ScenarioRegistered,
Wire,
)
@dataclass(frozen=True, slots=True)
class ScenarioHandle:
scenario_id: str
wire: Wire
proxy_base: str
def api_base(self) -> str:
return f"{self.proxy_base}/{self.scenario_id}/{self._mount()}"
def _mount(self) -> str:
return WIRE_MOUNTS[self.wire]
def register_scenario(scenario: Scenario) -> ScenarioHandle:
"""POST the scenario to the sidecar's control API and return its handle."""
result: Final = unwrap(
post(
URL(f"{SCRIPTED_PROVIDER_CONTROL_URL}/_scenarios"),
headers=NoBody(),
json=scenario,
response_type=ScenarioRegistered,
)
)
return ScenarioHandle(
scenario_id=result.scenario_id,
wire=scenario.wire,
proxy_base=SCRIPTED_PROVIDER_PROXY_BASE,
)
def delete_scenario(handle: ScenarioHandle) -> None:
unwrap(
http_delete(
URL(f"{SCRIPTED_PROVIDER_CONTROL_URL}/_scenarios/{handle.scenario_id}"),
headers=NoBody(),
json=NoBody(),
response_type=ScenarioDeleted,
)
)
CONTROL_URL: Final = SCRIPTED_PROVIDER_CONTROL_URL

285
tests/e2e/cost_calculation/test_token_pricing_e2e.py
View file

@ -1,285 +0,0 @@
"""Token-pricing e2e: every (map entry, case) cell derived from cost_map.json x
cases.json runs a scripted-usage call through a deployment registered on the
cost-map proxy, and the spend row plus response-cost header must equal the
reviewed golden in the case's ``expected`` cell verbatim -- no rate arithmetic
lives here.
Nothing here touches a real provider or the bundled cost map: the proxy's
upstream is the scripted-provider sidecar and its entire cost map is
tests/e2e/cost_map.json.
"""
from __future__ import annotations
import pytest
from typing import Final
from conftest import CostCalcClient, cost_rows, register_scenario_deployment
from cost_matrix import (
AUDIO_INPUT_DATA_URL,
FRONTIER_MODELS,
IMAGE_INPUT_DATA_URL,
SERVICE_TIER_REQUEST_WIRES,
VIDEO_INPUT_DATA_URL,
Case,
FrontierModel,
cases_for,
matrix_data_errors,
recount_cost,
)
from e2e_config import unique_marker
from lifecycle import ResourceManager
from models import (
CacheControl,
ChatAudio,
ChatBody,
ChatMessage,
ChatStreamOptions,
ChatTool,
ChatToolFunction,
FileContentPart,
FileObject,
FileSearchTool,
GoogleMapsTool,
GoogleSearchTool,
HostedWebSearchTool,
ImageContentPart,
ImageUrl,
InputAudio,
InputAudioContentPart,
TextContentPart,
WebSearchOptions,
)
from scripted_provider import ScriptedUsage, Wire
pytestmark: Final = [pytest.mark.e2e, pytest.mark.cost_map_stack] # mutable-ok: pytest only accepts a list for pytestmark
if _data_errors := matrix_data_errors():
raise ValueError("\n".join(_data_errors))
_MATRIX: Final[tuple[tuple[FrontierModel, Case], ...]] = tuple(
(model, case) for model in FRONTIER_MODELS for case in cases_for(model)
)
def _case_id(param: tuple[FrontierModel, Case]) -> str:
model, case = param
return f"{model.map_key.replace('/', '-')}-{case.name}"
_CACHE_WIRES: Final = frozenset({"anthropic_messages", "bedrock_converse"})
_WEB_SEARCH_OPTION_WIRES: Final = frozenset({"openai_chat", "azure_chat", "openai_responses"})
def _cache_control(usage: ScriptedUsage, wire: Wire) -> CacheControl | None:
if wire not in _CACHE_WIRES:
return None
if not (usage.cache_read_tokens or usage.cache_write_5m_tokens or usage.cache_write_1h_tokens):
return None
return CacheControl(type="ephemeral", ttl="1h" if usage.cache_write_1h_tokens else None)
def _chat_body(model: FrontierModel, case: Case, model_name: str, marker: str) -> ChatBody:
usage: Final = case.usage_for(model.map_key)
user_parts: Final = (
TextContentPart(
text=f"{marker} summarize the attached material in one line and name the city weather",
),
*(
(ImageContentPart(image_url=ImageUrl(url=IMAGE_INPUT_DATA_URL, detail="high")),)
if case.image_input
else ()
),
*(
(
InputAudioContentPart(
input_audio=InputAudio(data=AUDIO_INPUT_DATA_URL.split(",", 1)[1], format="wav")
),
)
if case.audio_input
else ()
),
*(
(FileContentPart(file=FileObject(file_data=VIDEO_INPUT_DATA_URL, format="mp4")),)
if case.video_input
else ()
),
)
tools: Final = (
*(
(
ChatTool(
function=ChatToolFunction(
name="get_weather",
description="Get the current weather and a short forecast for a city.",
parameters={
"type": "object",
"properties": {
"city": {"type": "string", "description": "City name"},
"days": {"type": "integer", "description": "Forecast horizon in days"},
"units": {"type": "string", "enum": ["metric", "imperial"]},
},
"required": ["city"],
},
)
),
)
if case.tool_call
else ()
),
*(
(HostedWebSearchTool(type="web_search_20250305", name="web_search", max_uses=5),)
if case.web_search is not None and model.wire == "anthropic_messages"
else ()
),
*(
(GoogleSearchTool(),)
if case.web_search is not None and model.wire in ("gemini_generate", "vertex_generate")
else ()
),
*((GoogleMapsTool(),) if case.google_maps else ()),
*((FileSearchTool(vector_store_ids=["vs_cost_calc_fixture"]),) if case.file_search else ()),
)
return ChatBody(
model=model_name,
messages=(
ChatMessage(
role="system",
content=[
TextContentPart(
text=(
"You are a deterministic pricing-harness assistant. "
"Keep answers to a single short line."
),
cache_control=_cache_control(usage, model.wire),
)
],
),
ChatMessage(role="user", content=list(user_parts)),
),
stream=case.stream,
stream_options=ChatStreamOptions(include_usage=True) if case.stream else None,
service_tier=(
case.service_tier
if case.service_tier is not None and model.wire in SERVICE_TIER_REQUEST_WIRES
else None
),
reasoning_effort="medium" if case.reasoning else None,
modalities=(
["text"] if case.audio_input else (["text", "audio"] if case.audio_output else None)
),
audio=(
ChatAudio(voice="alloy", format="pcm16") if case.audio_output else None
),
web_search_options=(
WebSearchOptions(search_context_size=case.web_search)
if case.web_search is not None and model.wire in _WEB_SEARCH_OPTION_WIRES
else None
),
tools=tools or None,
tool_choice="auto" if case.tool_call and model.wire != "bedrock_converse" else None,
# The test-owned cost map carries no supports_* flags, so litellm's
# optional-params gate rejects the realistic request fields; allowlist
# exactly the ones this case sends.
allowed_openai_params=[
name
for name, sent in (
("tool_choice", case.tool_call and model.wire != "bedrock_converse"),
("modalities", case.audio_input or case.audio_output),
("audio", case.audio_output),
("web_search_options", case.web_search is not None),
("reasoning_effort", case.reasoning),
)
if sent
],
)
class TestTokenPricing:
@pytest.mark.parametrize("model_case", _MATRIX, ids=_case_id)
@pytest.mark.covers("quota_management.spend_tracking.cost_matrix.logs_cost")
def test_scripted_usage_bills_at_map_rates(
self,
client: CostCalcClient,
resources: ResourceManager,
scoped_key: str,
model_case: tuple[FrontierModel, Case],
) -> None:
model, case = model_case
marker: Final = unique_marker()
model_name, _handle = register_scenario_deployment(client, resources, model, case, marker)
response: Final = client.proxy.transport.send(
"/chat/completions",
headers=client.proxy.transport.bearer(scoped_key),
json=_chat_body(model, case, model_name, marker),
stream=case.stream,
)
assert response.ok, (
f"{model.map_key}/{case.name}: proxy returned {response.status_code}: {response.body[:400]}"
)
assert response.stream_error is None, f"stream carried an error event: {response.stream_error}"
row: Final = cost_rows.poll_cost_row_where(
client.proxy,
scoped_key,
lambda r: r.metadata is not None and r.metadata.cost_breakdown is not None,
)
assert row is not None, f"no spend row with a cost breakdown landed for {model.map_key}/{case.name}"
if not case.exact_spend:
# stream_usage=absent: the provider reported no usage, so the row's
# token counts are the proxy's own recount; assert the recount
# billed both directions at the case's rates.
assert row.prompt_tokens is not None and row.prompt_tokens > 0, (
f"no-usage stream counted no input tokens: {row}"
)
assert row.completion_tokens is not None and row.completion_tokens > 0, (
f"no-usage stream counted no output tokens: {row}"
)
if case.image_input:
assert row.prompt_tokens < 4000, (
f"image data URL looks tokenized as text: prompt_tokens={row.prompt_tokens}"
)
assert row.spend is not None and cost_rows.approx_equal(
row.spend,
recount_cost(model, case, row.prompt_tokens, row.completion_tokens),
), f"no-usage stream spend {row.spend} != recount at map rates: {row}"
cost_rows.assert_total_is_sum_of_components(row)
return
golden: Final = case.expected_for(model)
if not case.stream:
# Streamed responses commit headers before the bill is computed, so
# the x-litellm-response-cost header is asserted only on non-stream
# calls.
assert response.response_cost is not None and cost_rows.approx_equal(
response.response_cost, golden.spend
), (
f"x-litellm-response-cost {response.response_cost} != golden {golden.spend}"
)
assert row.spend is not None and cost_rows.approx_equal(row.spend, golden.spend), (
f"{model.map_key}/{case.name}: spend {row.spend} != golden {golden.spend} "
f"(breakdown {row.breakdown.model_dump()})"
)
breakdown: Final = row.breakdown
assert breakdown.input_cost is not None and cost_rows.approx_equal(
breakdown.input_cost, golden.input_cost
), (
f"{model.map_key}/{case.name}: gross input_cost {breakdown.input_cost} "
f"!= golden {golden.input_cost}; cached/written tokens billed at the input rate"
)
assert breakdown.output_cost is not None and cost_rows.approx_equal(
breakdown.output_cost, golden.output_cost
), (
f"{model.map_key}/{case.name}: output_cost {breakdown.output_cost} "
f"!= golden {golden.output_cost}"
)
assert row.prompt_tokens == golden.prompt_tokens, (
f"prompt_tokens {row.prompt_tokens} != {golden.prompt_tokens}"
)
assert row.completion_tokens == golden.completion_tokens, (
f"completion_tokens {row.completion_tokens} != {golden.completion_tokens}"
)
cost_rows.assert_total_is_sum_of_components(row)

2
tests/e2e/coverage_registry/quota_management.yaml
View file

@ -63,5 +63,3 @@
- {id: quota_management.spend_tracking.key_attribution.health_rows_keep_service_account, module: quota_management, tier: P1, behavior: spend_tracking, variant: key_attribution, assertions: [health_rows_keep_service_account], exercised_on: [chat_completions], source: "proxy/health_check.py", rationale: "A /health probe's spend row stays keyed by the literal litellm-internal-health-check service account rather than a hash of it, so health spend never appears as an unattributed key"}
- {id: quota_management.spend_tracking.key_attribution.retrieve_batch_cost_joins_retrieving_key, module: quota_management, tier: P1, behavior: spend_tracking, variant: key_attribution, assertions: [retrieve_batch_cost_joins_retrieving_key], exercised_on: [batches], source: "proxy/batches_endpoints/endpoints.py", rationale: "The retrieve that first sees a batch in a terminal state prices it inline and writes its {provider_batch_id}_batch_cost row against the retrieving key, so the batch each run creates is one OpenAI fails at validation within seconds and the test retrieves it by its raw provider id with the same key until it is failed; a raw id is never owned by the CheckBatchCost poller, and the row must carry that key's token hash and alias"}
- {id: quota_management.spend_tracking.key_attribution.poller_batch_cost_joins_creating_key, module: quota_management, tier: P1, behavior: spend_tracking, variant: key_attribution, assertions: [poller_batch_cost_joins_creating_key], exercised_on: [batches], source: "enterprise/litellm_enterprise/proxy/common_utils/check_batch_cost.py", rationale: "The CheckBatchCost poller bills a completed, positive-cost batch created through a unified id against the key that created it, a different writer from the inline retrieve. No test claims this cell yet: OpenAI's completion window is 24h and both e2e stacks boot a fresh Postgres per build, so a completed batch is out of one run's reach and the managed list never shows an earlier run's batch; the cell stays visible as a gap until a run can hand a completed batch to the poller"}
- {id: quota_management.spend_tracking.cost_matrix.logs_cost, module: quota_management, tier: P1, behavior: spend_tracking, variant: cost_matrix, assertions: [logs_cost], exercised_on: [chat_completions], source: "litellm_core_utils/llm_cost_calc/utils.py", rationale: "A scripted-usage call through the cost-map proxy bills every reported token kind at the deployment's test-map rate (input, output, cache read, 5m/1h cache write, reasoning, audio, above-threshold tiers, flex/priority service tiers, web search, response-model override) and lands on the row's cost_breakdown, streamed or not"}
- {id: quota_management.spend_tracking.scripted_wire.logs_cost, module: quota_management, tier: P1, behavior: spend_tracking, variant: scripted_wire, assertions: [logs_cost], exercised_on: [chat_completions], source: "litellm_core_utils/llm_cost_calc/utils.py", rationale: "Each provider wire shape (openai chat, responses, anthropic messages, gemini generateContent, together, fireworks) parses usage into the same spend components: the gross input cost is fresh tokens at the input rate plus each cache/audio component at its own rate, streamed anthropic included"}

16
tests/e2e/e2e_config.py
View file

@ -143,22 +143,6 @@ WEEKLY_ANOMALY_OPT_IN_ENV = "E2E_WEEKLY_ANOMALY"
MANAGED_FILES_OPT_IN_ENV = "E2E_MANAGED_FILES_STACK"
PROMPT_CACHING_OPT_IN_ENV = "E2E_PROMPT_CACHING_STACK"
REDIS_CHAOS_OPT_IN_ENV = "E2E_REDIS_CHAOS"
# The cost_calculation suite needs a proxy booted with LITELLM_MODEL_COST_MAP_URL
# pointing at tests/e2e/cost_map.json (its whole map is test-owned rates) plus a
# scripted-provider sidecar; deselected unless the opt-in env var is set.
COST_MAP_OPT_IN_ENV = "E2E_COST_MAP_STACK"
# Base URL of the proxy running the test cost map. Defaults to the shared proxy
# so a local run only has to set the opt-in and boot the proxy accordingly.
COST_MAP_PROXY_URL = os.environ.get("E2E_COST_MAP_PROXY_URL", PROXY_BASE_URL).rstrip("/")
# Where the test runner reaches the scripted-provider sidecar's control API.
SCRIPTED_PROVIDER_CONTROL_URL = os.environ.get(
"E2E_SCRIPTED_PROVIDER_CONTROL_URL", "http://127.0.0.1:9100"
).rstrip("/")
# The api_base root deployments register with: how the proxy (possibly in
# another container) reaches the sidecar's provider wire.
SCRIPTED_PROVIDER_PROXY_BASE = os.environ.get(
"E2E_SCRIPTED_PROVIDER_PROXY_BASE", SCRIPTED_PROVIDER_CONTROL_URL
).rstrip("/")
CLI_DETERMINISM_OPT_IN_ENV = "E2E_CLI_DETERMINISM"
ANOMALY_SESSIONS = int(os.environ.get("E2E_ANOMALY_SESSIONS", "6"))
ANOMALY_TURNS_PER_SESSION = int(os.environ.get("E2E_ANOMALY_TURNS_PER_SESSION", "6"))

7
tests/e2e/gateway/cost_calculation_ci_config.yml
View file

@ -1,7 +0,0 @@
general_settings:
master_key: os.environ/LITELLM_MASTER_KEY
database_url: os.environ/DATABASE_URL
store_model_in_db: true
proxy_batch_write_at: 5
model_list: []

74
tests/e2e/models.py
View file

@ -8,7 +8,7 @@ from __future__ import annotations
from collections.abc import Sequence
from datetime import datetime
from typing import Final, Literal, TypeAlias
from typing import Final, Literal
from e2e_http import PartialBody
from pydantic import (
@ -187,24 +187,12 @@ class ChatMetadata(BaseModel):
class ImageUrl(BaseModel):
url: str
detail: str | None = None
class InputAudio(BaseModel):
data: str
format: str
class FileObject(BaseModel):
file_data: str | None = None
file_id: str | None = None
format: str | None = None
class TextContentPart(BaseModel):
type: str = "text"
text: str
cache_control: CacheControl | None = None
cache_control: "CacheControl | None" = None
class ImageContentPart(BaseModel):
@ -212,17 +200,7 @@ class ImageContentPart(BaseModel):
image_url: ImageUrl
class InputAudioContentPart(BaseModel):
type: str = "input_audio"
input_audio: InputAudio
class FileContentPart(BaseModel):
type: str = "file"
file: FileObject
ContentPart = TextContentPart | ImageContentPart | InputAudioContentPart | FileContentPart
ContentPart = TextContentPart | ImageContentPart
class ChatMessage(BaseModel):
@ -304,38 +282,7 @@ class ChatToolResultTurn(BaseModel):
content: str
ChatTurn: TypeAlias = ChatMessage | ChatAssistantTurn | ChatToolResultTurn
class HostedWebSearchTool(BaseModel):
"""A provider-hosted web-search tool sent inside an OpenAI tools list
(Anthropic's ``web_search_20250305`` shape)."""
type: str
name: str
max_uses: int | None = None
class GoogleSearchTool(BaseModel):
googleSearch: dict[str, object] = {}
class GoogleMapsTool(BaseModel):
googleMaps: dict[str, object] = {}
class FileSearchTool(BaseModel):
type: Literal["file_search"] = "file_search"
vector_store_ids: list[str]
class WebSearchOptions(BaseModel):
search_context_size: Literal["low", "medium", "high"] | None = None
class ChatAudio(BaseModel):
voice: str
format: str
type ChatTurn = ChatMessage | ChatAssistantTurn | ChatToolResultTurn
class ChatStreamOptions(BaseModel):
@ -356,16 +303,10 @@ class ChatBody(BaseModel):
thinking: ThinkingParam | None = None
service_tier: str | None = None
prompt_cache_key: str | None = None
tools: Sequence[
ChatTool | McpChatTool | HostedWebSearchTool | GoogleSearchTool | GoogleMapsTool | FileSearchTool
] | None = None
tools: Sequence[ChatTool | McpChatTool] | None = None
tool_choice: str | None = None
modalities: list[str] | None = None
audio: ChatAudio | None = None
web_search_options: WebSearchOptions | None = None
guardrails: list[str] | None = None
response_format: dict[str, object] | None = None
allowed_openai_params: list[str] | None = None
chat_template_kwargs: dict[str, bool] | None = None
cache: dict[str, bool] | None = {"no-cache": True}
@ -532,7 +473,7 @@ class AnthropicCustomTool(BaseModel):
input_schema: ToolInputSchema
AnthropicTool: TypeAlias = AnthropicToolSearchTool | AnthropicWebSearchTool | AnthropicCustomTool
type AnthropicTool = AnthropicToolSearchTool | AnthropicWebSearchTool | AnthropicCustomTool
class AnthropicContentBlock(BaseModel):
@ -570,7 +511,7 @@ class AnthropicToolResultTurn(BaseModel):
content: list[AnthropicToolResultBlock]
AnthropicMessage: TypeAlias = ChatMessage | AnthropicAssistantTurn | AnthropicToolResultTurn
type AnthropicMessage = ChatMessage | AnthropicAssistantTurn | AnthropicToolResultTurn
class AnthropicToolChoice(BaseModel):
@ -1061,7 +1002,6 @@ class ModelInfoBody(BaseModel):
access_groups: list[str] | None = None
team_id: str | None = None
allowed_fails_policy: dict[str, int] | None = None
base_model: str | None = None
class ModelNewBody(BaseModel):

1
tests/e2e/pytest.ini
View file

@ -12,4 +12,3 @@ markers =
prompt_caching_stack: needs a proxy running with router_settings.optional_pre_call_checks including prompt_caching; deselected unless E2E_PROMPT_CACHING_STACK is set
cli_determinism: drives the real claude CLI for several seconds; deselected unless E2E_CLI_DETERMINISM is set
redis_chaos: load test that pauses the proxy's Redis outright mid-run; needs a proxy booted from gateway/redis_chaos_ci_config.yml on the same host, and is deselected unless E2E_REDIS_CHAOS is set
cost_map_stack: needs a proxy whose whole cost map is tests/e2e/cost_map.json (LITELLM_MODEL_COST_MAP_URL) plus a scripted-provider sidecar; deselected unless E2E_COST_MAP_STACK is set

2
tests/integration/README.md
View file

@ -2,6 +2,8 @@
These tests exercise a running gateway, PostgreSQL and Redis with an owned local upstream. CircleCI owns this suite. Tests are grouped by behavior, with no automatic test retries or fallback to paid provider calls
The `cost` group runs the scripted-provider cost matrix through a dedicated sidecar. The sidecar serves the test-owned cost map over loopback through `LITELLM_MODEL_COST_MAP_URL`; cost goldens are checked into the integration suite and must not be copied into the E2E coverage registry
Use `tests/integration/run.py management`, `accounting`, `database`, `providers`, `extensions` or `sdk` to run a selected group. Set `INTEGRATION_PROXY_URL`, `INTEGRATION_UPSTREAM_URL`, `INTEGRATION_MASTER_KEY` and `DATABASE_URL` to an isolated test deployment. The runner selects the new domain directories explicitly; the legacy OCI and sandbox selections remain separate
Management also requires `INTEGRATION_PEER_URL`, `REDIS_HOST` and `REDIS_PORT`. CircleCI starts two directly addressed proxy processes sharing only that job's stores. The test-only CLI wrapper supplies enterprise route entitlement, following the existing behavior suite's convention. It does not qualify license validation; run it with one worker and no reload

1
tests/integration/_support/manifest.py
View file

@ -20,6 +20,7 @@ OWNED_DIRECTORIES: Final = frozenset(
"observability",
"compatibility",
"sdk",
"cost_calculation",
}
)

57
tests/integration/_support/scripted_client.py Normal file
View file

@ -0,0 +1,57 @@
"""Client for registering scenarios with the integration scripted provider."""
from __future__ import annotations
import os
from dataclasses import dataclass
from typing import Final
import httpx
from integration._support.scripted_provider import (
WIRE_MOUNTS,
Scenario,
ScenarioDeleted,
ScenarioRegistered,
Wire,
)
CONTROL_URL: Final = os.environ.get("INTEGRATION_SCRIPTED_PROVIDER_URL", "http://127.0.0.1:8191").rstrip("/")
@dataclass(frozen=True, slots=True)
class ScenarioHandle:
scenario_id: str
wire: Wire
control_url: str
def api_base(self) -> str:
return f"{self.control_url}/{self.scenario_id}/{self._mount()}"
def _mount(self) -> str:
return WIRE_MOUNTS[self.wire]
def register_scenario(scenario: Scenario) -> ScenarioHandle:
response: Final = httpx.post(
f"{CONTROL_URL}/_scenarios",
json=scenario.model_dump(mode="json"),
trust_env=False,
timeout=15,
)
response.raise_for_status()
result: Final = ScenarioRegistered.model_validate_json(response.content)
return ScenarioHandle(
scenario_id=result.scenario_id,
wire=scenario.wire,
control_url=CONTROL_URL,
)
def delete_scenario(handle: ScenarioHandle) -> None:
response: Final = httpx.delete(
f"{CONTROL_URL}/_scenarios/{handle.scenario_id}",
trust_env=False,
timeout=15,
)
response.raise_for_status()
ScenarioDeleted.model_validate_json(response.content)

21
tests/e2e/cost_calculation/scripted_provider.py → tests/integration/_support/scripted_provider.py
View file

@ -1,6 +1,6 @@
"""Scripted provider sidecar for the cost-calculation e2e suite.
"""Scripted provider sidecar for the cost-calculation integration suite.
A standalone process (``python -m cost_calculation.scripted_provider``) that
A standalone process (``python -m integration._support.scripted_provider``) that
pretends to be an LLM provider for the proxy under test. The suite registers a
Scenario over a small control API; the provider wire routes then answer the
proxy's upstream calls with the scripted usage figures, in the exact wire shape
@ -32,6 +32,7 @@ final stream chunk carries usage or the provider reports none.
from __future__ import annotations
import argparse
import json
import struct
import sys
@ -41,8 +42,9 @@ import zlib
from collections.abc import Mapping
from dataclasses import dataclass
from http.server import BaseHTTPRequestHandler, ThreadingHTTPServer
from pathlib import Path
from types import MappingProxyType
from typing import Final, Literal, TypeAlias
from typing import Final, Literal, TypeAlias, cast
from urllib.parse import unquote, urlsplit
from pydantic import BaseModel, ConfigDict, TypeAdapter, ValidationError, model_validator
@ -1362,6 +1364,12 @@ def handle_request(store: _ScenarioStore, method: str, raw_path: str, body: byte
segments: Final = tuple(segment for segment in path.split("/") if segment)
if method == "GET" and segments == ("health",):
return RenderedResponse(200, "application/json", _json_bytes(_jobj(("status", "ok"))))
if method == "GET" and segments == ("_cost_map",):
return RenderedResponse(
200,
"application/json",
(Path(__file__).resolve().parents[1] / "cost_calculation" / "cost_map.json").read_bytes(),
)
if segments and segments[0] == "_oauth":
if method == "POST" and segments == ("_oauth", "token"):
return RenderedResponse(
@ -1459,7 +1467,7 @@ class _ScriptedHandler(BaseHTTPRequestHandler):
DEFAULT_PORT: Final = 9100
DEFAULT_PORT: Final = 8191
def serve(port: int = DEFAULT_PORT, bind_host: str = "127.0.0.1") -> None:
@ -1469,5 +1477,6 @@ def serve(port: int = DEFAULT_PORT, bind_host: str = "127.0.0.1") -> None:
if __name__ == "__main__":
port_arg: Final = int(sys.argv[1]) if len(sys.argv) > 1 else DEFAULT_PORT
serve(port=port_arg)
parser: Final = argparse.ArgumentParser()
parser.add_argument("--port", type=int, default=8191)
serve(port=cast(int, parser.parse_args().port))

1092
tests/integration/contracts.json
View file

File diff suppressed because it is too large Load diff

0
tests/e2e/cost_calculation/cases.json → tests/integration/cost_calculation/cases.json
View file

147
tests/integration/cost_calculation/conftest.py Normal file
View file

@ -0,0 +1,147 @@
from __future__ import annotations
import functools
import json
import os
from collections.abc import Mapping
from hashlib import sha256
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.scripted_client import delete_scenario, register_scenario
from integration.cost_calculation.cost_matrix import Case, FrontierModel
class CostBreakdown(BaseModel):
model_config = ConfigDict(extra="ignore")
input_cost: float | None = None
output_cost: float | None = None
cache_read_cost: float | None = None
cache_creation_cost: float | None = None
reasoning_cost: float | None = None
tool_usage_cost: float | None = None
total_cost: float | None = None
service_tier: str | None = None
class CostMetadata(BaseModel):
model_config = ConfigDict(extra="ignore")
cost_breakdown: CostBreakdown | None = None
class CostRow(BaseModel):
model_config = ConfigDict(extra="ignore")
spend: float | None = None
prompt_tokens: int | None = None
completion_tokens: int | None = None
metadata: CostMetadata | None = None
@property
def breakdown(self) -> CostBreakdown:
assert self.metadata is not None and self.metadata.cost_breakdown is not None
return self.metadata.cost_breakdown
def approx_equal(actual: float, expected: float) -> bool:
return abs(actual - expected) <= max(1e-9, abs(expected) * 1e-2)
def assert_total_is_sum_of_components(row: CostRow) -> None:
breakdown: Final = row.breakdown
total: Final = sum(
cost or 0.0
for cost in (breakdown.input_cost, breakdown.output_cost, breakdown.tool_usage_cost)
)
assert breakdown.total_cost is not None and approx_equal(breakdown.total_cost, total)
assert row.spend is not None and approx_equal(row.spend, breakdown.total_cost)
def _row(value: Mapping[str, object]) -> CostRow | None:
metadata_value: Final = value.get("metadata")
metadata: Final = json.loads(metadata_value) if isinstance(metadata_value, str) else metadata_value
parsed: Final = CostRow.model_validate({**value, "metadata": metadata})
return parsed if parsed.metadata and parsed.metadata.cost_breakdown else None
def poll_cost_row(key: str) -> CostRow:
digest: Final = sha256(key.encode()).hexdigest()
def read() -> CostRow | None:
rows: Final = read_rows(
'SELECT spend, metadata, prompt_tokens, completion_tokens FROM "LiteLLM_SpendLogs" WHERE api_key=%s',
(digest,),
)
return next((parsed for row in rows if (parsed := _row(row)) is not None), 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(
serialization.Encoding.PEM,
serialization.PrivateFormat.PKCS8,
serialization.NoEncryption(),
).decode()
def _vertex_service_account_json(url: str) -> str:
return json.dumps(
{
"type": "service_account",
"project_id": "cc-scripted-project",
"private_key_id": "scripted",
"private_key": _vertex_private_key_pem(),
"client_email": "scripted@cc-scripted-project.iam.gserviceaccount.com",
"client_id": "0",
"auth_uri": f"{url}/_oauth/authorize",
"token_uri": f"{url}/_oauth/token",
}
)
def register_scenario_deployment(
scenario: Scenario,
model: FrontierModel,
case: Case,
marker: str,
) -> str:
control_url: Final = os.environ["INTEGRATION_SCRIPTED_PROVIDER_URL"].rstrip("/")
sidecar_scenario: Final = case.scenario(
scenario_id=f"sc-{marker}", model=model, text=f"scripted answer {marker}"
)
handle: Final = register_scenario(sidecar_scenario)
scenario.cleanups.callback(delete_scenario, handle)
model_name: Final = f"{model.model_name}-{marker}"
parameters: Final = {
"model": model.litellm_model,
"api_key": model.api_key,
"api_base": handle.api_base(),
**model.litellm_params,
**(
{"vertex_credentials": _vertex_service_account_json(control_url)}
if model.wire == "vertex_generate"
else {}
),
}
created: Final = scenario.gateway.post(
"/model/new",
JSON_OBJECT.validate_python({
"model_name": model_name,
"litellm_params": parameters,
"model_info": {"base_model": model.base_model},
}),
)
identity: Final = string_value(object_value(created["model_info"])["id"])
scenario.cleanups.callback(scenario.delete_model, identity)
return model_name

0
tests/e2e/cost_map.json → tests/integration/cost_calculation/cost_map.json
View file

10
tests/e2e/cost_calculation/cost_matrix.py → tests/integration/cost_calculation/cost_matrix.py
View file

@ -2,9 +2,9 @@
the request/response cases from ``cases.json``, and the loaders both use.
Two data files drive the suite; nothing in Python lists models or cases:
- ``tests/e2e/cost_map.json`` is the proxy's ENTIRE model cost map
- ``tests/integration/cost_calculation/cost_map.json`` is the proxy's ENTIRE model cost map
(LITELLM_MODEL_COST_MAP_URL); every entry becomes a deployment under test.
- ``tests/e2e/cost_calculation/cases.json`` is the case list plus the reviewed
- ``tests/integration/cost_calculation/cases.json`` is the case list plus the reviewed
goldens: each exact-spend case carries an ``expected`` cell per map key it
runs against, each recount case carries its ``models`` list, so matrix
membership and expected values are literal data read side by side.
@ -27,9 +27,9 @@ from types import MappingProxyType
from typing import Final, Literal
from pydantic import BaseModel, ConfigDict, Field, TypeAdapter
from scripted_provider import Scenario, ScriptedOutput, ScriptedToolCall, ScriptedUsage, Wire
from integration._support.scripted_provider import Scenario, ScriptedOutput, ScriptedToolCall, ScriptedUsage, Wire
COST_MAP_PATH: Final = Path(__file__).resolve().parent.parent / "cost_map.json"
COST_MAP_PATH: Final = Path(__file__).resolve().parent / "cost_map.json"
CASES_PATH: Final = Path(__file__).resolve().parent / "cases.json"
class SearchContextCostPerQuery(BaseModel):
@ -506,7 +506,7 @@ VIDEO_INPUT_DATA_URL: Final = video_input_data_url()
def matrix_data_errors() -> tuple[str, ...]:
"""Consistency findings for the data files, as human-readable strings.
Called at collection time by the e2e suite, so a map key named by a case
Called at collection time by the integration suite, so a map key named by a case
but absent from cost_map.json fails the suite's collection loudly.
"""
unknown_deployments: Final = sorted(

223
tests/integration/cost_calculation/test_token_pricing.py Normal file
View file

@ -0,0 +1,223 @@
"""Token pricing coverage for the integration scripted-provider cost shard."""
from __future__ import annotations
import uuid
from typing import Final, cast
import pytest
from pydantic import JsonValue
from integration._support.client import JSON_OBJECT, Gateway
from integration._support.scripted_provider import ScriptedUsage, Wire
from integration.cost_calculation.conftest import (
approx_equal,
assert_total_is_sum_of_components,
poll_cost_row,
register_scenario_deployment,
)
from integration.cost_calculation.cost_matrix import (
AUDIO_INPUT_DATA_URL,
FRONTIER_MODELS,
IMAGE_INPUT_DATA_URL,
SERVICE_TIER_REQUEST_WIRES,
VIDEO_INPUT_DATA_URL,
Case,
FrontierModel,
cases_for,
matrix_data_errors,
recount_cost,
)
if _data_errors := matrix_data_errors():
raise ValueError("\n".join(_data_errors))
def _case_id(param: tuple[FrontierModel, Case]) -> str:
model, case = param
return f"{model.map_key.replace('/', '-')}-{case.name}"
_MATRIX: Final = tuple(
pytest.param(
(model, case),
marks=pytest.mark.covers(
"quota_management.spend_tracking.scripted_wire.logs_cost"
if case.family == "transport"
else "quota_management.spend_tracking.cost_matrix.logs_cost"
),
id=_case_id((model, case)),
)
for model in FRONTIER_MODELS
for case in cases_for(model)
)
_CACHE_WIRES: Final = frozenset({"anthropic_messages", "bedrock_converse"})
_WEB_SEARCH_OPTION_WIRES: Final = frozenset({"openai_chat", "azure_chat", "openai_responses"})
def _cache_control(usage: ScriptedUsage, wire: Wire) -> dict[str, JsonValue] | None:
if wire not in _CACHE_WIRES:
return None
if not (usage.cache_read_tokens or usage.cache_write_5m_tokens or usage.cache_write_1h_tokens):
return None
return {"type": "ephemeral", **({"ttl": "1h"} if usage.cache_write_1h_tokens else {})}
def _chat_body(model: FrontierModel, case: Case, model_name: str, marker: str) -> dict[str, JsonValue]:
usage: Final = case.usage_for(model.map_key)
user_parts: Final = [
{"type": "text", "text": f"{marker} summarize the attached material in one line and name the city weather"},
*(
[{"type": "image_url", "image_url": {"url": IMAGE_INPUT_DATA_URL, "detail": "high"}}]
if case.image_input
else []
),
*(
[{"type": "input_audio", "input_audio": {"data": AUDIO_INPUT_DATA_URL.split(",", 1)[1], "format": "wav"}}]
if case.audio_input
else []
),
*(
[{"type": "file", "file": {"file_data": VIDEO_INPUT_DATA_URL, "format": "mp4"}}]
if case.video_input
else []
),
]
tools: Final[list[JsonValue]] = [
*(
[
{
"type": "function",
"function": {
"name": "get_weather",
"description": "Get the current weather and a short forecast for a city.",
"parameters": {
"type": "object",
"properties": {
"city": {"type": "string", "description": "City name"},
"days": {"type": "integer", "description": "Forecast horizon in days"},
"units": {"type": "string", "enum": ["metric", "imperial"]},
},
"required": ["city"],
},
},
}
]
if case.tool_call
else []
),
*(
[{"type": "web_search_20250305", "name": "web_search", "max_uses": 5}]
if case.web_search is not None and model.wire == "anthropic_messages"
else []
),
*(
[{"googleSearch": {}}]
if case.web_search is not None and model.wire in ("gemini_generate", "vertex_generate")
else []
),
*([{"googleMaps": {}}] if case.google_maps else []),
*([{"type": "file_search", "vector_store_ids": ["vs_cost_calc_fixture"]}] if case.file_search else []),
]
cache_control: Final = _cache_control(usage, model.wire)
message: Final = {
"role": "system",
"content": [
{
"type": "text",
"text": "You are a deterministic pricing-harness assistant. Keep answers to a single short line.",
**({"cache_control": cache_control} if cache_control else {}),
}
],
}
return cast(dict[str, JsonValue], {
"model": model_name,
"messages": [message, {"role": "user", "content": user_parts}],
"stream": case.stream,
**({"stream_options": {"include_usage": True}} if case.stream else {}),
**(
{"service_tier": case.service_tier}
if case.service_tier is not None and model.wire in SERVICE_TIER_REQUEST_WIRES
else {}
),
**({"reasoning_effort": "medium"} if case.reasoning else {}),
**(
{"modalities": ["text", "audio"] if case.audio_output else ["text"]}
if case.audio_input or case.audio_output
else {}
),
**({"audio": {"voice": "alloy", "format": "pcm16"}} if case.audio_output else {}),
**(
{"web_search_options": {"search_context_size": case.web_search}}
if case.web_search is not None and model.wire in _WEB_SEARCH_OPTION_WIRES
else {}
),
**({"tools": tools} if tools else {}),
**({"tool_choice": "auto"} if case.tool_call and model.wire != "bedrock_converse" else {}),
"allowed_openai_params": [
name
for name, sent in (
("tool_choice", case.tool_call and model.wire != "bedrock_converse"),
("modalities", case.audio_input or case.audio_output),
("audio", case.audio_output),
("web_search_options", case.web_search is not None),
("reasoning_effort", case.reasoning),
)
if sent
],
})
def _assert_stream_has_no_error(response_text: str) -> None:
for line in response_text.splitlines():
if not line.startswith("data:"):
continue
payload = line.removeprefix("data:").strip()
if payload == "[DONE]":
continue
parsed = JSON_OBJECT.validate_json(payload)
assert "error" not in parsed, f"stream carried an error event: {parsed}"
@pytest.mark.parametrize("model_case", _MATRIX)
def test_scripted_usage_bills_at_map_rates(
gateway: Gateway,
model_case: tuple[FrontierModel, Case],
) -> None:
model, case = model_case
marker: Final = uuid.uuid4().hex[:12]
with gateway.scenario() as scenario:
key: Final = scenario.key()
model_name: Final = register_scenario_deployment(scenario, model, case, marker)
response: Final = gateway.request(
"POST",
"/v1/chat/completions",
_chat_body(model, case, model_name, marker),
key=key,
)
assert response.is_success, (
f"{model.map_key}/{case.name}: proxy returned {response.status_code}: {response.text[:400]}"
)
if case.stream:
_assert_stream_has_no_error(response.text)
row: Final = poll_cost_row(key)
if not case.exact_spend:
assert row.prompt_tokens is not None and row.prompt_tokens > 0
assert row.completion_tokens is not None and row.completion_tokens > 0
if case.image_input:
assert row.prompt_tokens < 4000
assert row.spend is not None and approx_equal(
row.spend, recount_cost(model, case, row.prompt_tokens, row.completion_tokens)
)
assert_total_is_sum_of_components(row)
return
golden: Final = case.expected_for(model)
if not case.stream:
header: Final = cast(str | None, response.headers.get("x-litellm-response-cost"))
assert header is not None and approx_equal(float(header), golden.spend)
assert row.spend is not None and approx_equal(row.spend, golden.spend)
breakdown: Final = row.breakdown
assert breakdown.input_cost is not None and approx_equal(breakdown.input_cost, golden.input_cost)
assert breakdown.output_cost is not None and approx_equal(breakdown.output_cost, golden.output_cost)
assert row.prompt_tokens == golden.prompt_tokens
assert row.completion_tokens == golden.completion_tokens
assert_total_is_sum_of_components(row)