test(spend_logs): cover metadata retention across endpoints, failures, cache hits, batches, runtime updates and auto-router publication

Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>
This commit is contained in:
mrinal 2026-10-05 22:34:16 +00:00
parent d80f61b116
commit f2ce0a8d20
2 changed files with 356 additions and 13 deletions

View file

@ -63,7 +63,7 @@ def _failed_line(index: int) -> str:
)
def _batch_routes(model: str) -> RoutedResponse:
def batch_routes(model: str) -> RoutedResponse:
output_lines: Final = (
_succeeded_line(1, model, **FIRST_LINE),
_succeeded_line(2, model, **SECOND_LINE),
@ -139,7 +139,7 @@ def _batches_total_tokens(endpoints: dict[str, object] | None) -> int | None:
return int(total_tokens) if isinstance(total_tokens, (int, float, str)) else None
def _input_file(model: str) -> bytes:
def batch_input_file(model: str) -> bytes:
return (
"\n".join(
json.dumps(
@ -166,13 +166,13 @@ def test_completed_batch_spend_row_records_reasoning_tokens_and_error_file_failu
with gateway.scenario() as scenario:
key: Final = scenario.key()
scenario_id: Final = f"batch-accounting-{uuid.uuid4().hex[:12]}"
handle: Final = register_scenario(scenario_id, _batch_routes("gpt-4o-mini"))
handle: Final = register_scenario(scenario_id, batch_routes("gpt-4o-mini"))
scenario.cleanups.callback(delete_scenario, handle)
model: Final = scenario.model(api_base=handle.api_base())
file_response: Final = gateway.request_multipart(
"/v1/files",
{"purpose": "batch", "model": model},
{"file": ("in.jsonl", _input_file(model), "application/jsonl")},
{"file": ("in.jsonl", batch_input_file(model), "application/jsonl")},
key=key,
)
assert file_response.status_code == 200, file_response.text
@ -235,7 +235,7 @@ BATCH_SPEND: Final = (BATCH_PROMPT_TOKENS * INPUT_COST_PER_TOKEN + BATCH_COMPLET
def test_completed_batch_spend_lands_under_batches_in_team_endpoint_activity(gateway: Gateway) -> None:
with gateway.scenario() as scenario:
scenario_id: Final = f"batch-endpoint-{uuid.uuid4().hex[:12]}"
handle: Final = register_scenario(scenario_id, _batch_routes("gpt-4o-mini"))
handle: Final = register_scenario(scenario_id, batch_routes("gpt-4o-mini"))
scenario.cleanups.callback(delete_scenario, handle)
model: Final = scenario.model(
api_base=handle.api_base(),
@ -247,7 +247,7 @@ def test_completed_batch_spend_lands_under_batches_in_team_endpoint_activity(gat
file_response: Final = gateway.request_multipart(
"/v1/files",
{"purpose": "batch", "model": model},
{"file": ("in.jsonl", _input_file(model), "application/jsonl")},
{"file": ("in.jsonl", batch_input_file(model), "application/jsonl")},
key=key,
)
assert file_response.status_code == 200, file_response.text

View file

@ -1,33 +1,50 @@
import json
from collections.abc import Mapping
from collections.abc import Callable, Generator, Mapping
from contextlib import contextmanager
from dataclasses import dataclass
from hashlib import sha256
from pathlib import Path
from typing import Final
from uuid import uuid4
from integration._support.client import Gateway, eventually, object_value, string_value
from integration._support.database import read_rows
import httpx
import pytest
from integration._support.client import JSON_OBJECT, Gateway, Scenario, eventually, object_value, string_value
from integration._support.database import read_rows, scratch_database
from integration._support.process import owned_proxy
from integration._support.wire import Reply, Request, wire_server
from integration._support.upstream import delete_scenario, register_scenario
from integration._support.wire import Reply, Request, Wire, wire_server
from integration.spend.test_batch_completion_accounting import batch_input_file, batch_routes
from integration.streaming.test_stream_contracts import text_stream
from pydantic import JsonValue
CACHED_PROMPT_TOKENS: Final = 4
def _config(tmp_path: Path, spend_logs_metadata_fields: Mapping[str, JsonValue]) -> Path:
def _config(
tmp_path: Path,
spend_logs_metadata_fields: Mapping[str, JsonValue] | None,
*,
model_list: tuple[Mapping[str, JsonValue], ...] = (),
litellm_settings: Mapping[str, JsonValue] | None = None,
) -> Path:
retention: Final = (
{} if spend_logs_metadata_fields is None else {"spend_logs_metadata_fields": dict(spend_logs_metadata_fields)}
)
config: Final = tmp_path / f"spend-logs-metadata-fields-{uuid4()}.json"
config.write_text(
json.dumps(
{
"model_list": [],
"model_list": [dict(model) for model in model_list],
"general_settings": {
"master_key": "os.environ/LITELLM_MASTER_KEY",
"database_url": "os.environ/DATABASE_URL",
"store_model_in_db": True,
"proxy_batch_write_at": 1,
"proxy_batch_polling_interval": 1,
"spend_logs_metadata_fields": dict(spend_logs_metadata_fields),
**retention,
},
"litellm_settings": dict(litellm_settings or {}),
}
)
)
@ -107,3 +124,329 @@ def test_included_metadata_fields_are_the_only_ones_stored_besides_always_kept(
row, _ = _stored_row_after_one_chat(gateway, tmp_path, {"include": ["user_api_key_alias"]})
assert set(object_value(row["metadata"])) == {"status", "cold_storage_object_key", "user_api_key_alias"}
EXCLUDED: Final = ("model_map_information", "user_api_key_alias")
@dataclass(frozen=True, slots=True)
class _Isolated:
proxy: Gateway
scenario: Scenario
key: str
database_url: str | None = None
@property
def hashed_key(self) -> str:
return sha256(self.key.encode()).hexdigest()
def rows(self, count: int, where: str = "TRUE", seconds: int = 70) -> tuple[dict[str, JsonValue], ...]:
return tuple(
eventually(
lambda: read_rows(
"SELECT request_id, call_type, status, spend, cache_hit, metadata "
f'FROM "LiteLLM_SpendLogs" WHERE api_key=%s AND {where} ORDER BY "startTime"',
(self.hashed_key,),
database_url=self.database_url,
),
lambda values: len(values) == count,
seconds=seconds,
)
)
@contextmanager
def _isolated(gateway: Gateway, config: Path, tmp_path: Path, database_url: str | None = None) -> Generator[_Isolated]:
database: Final = {} if database_url is None else {"DATABASE_URL": database_url}
replica: Final = () if database_url is None else ("DATABASE_URL_READ_REPLICA",)
with (
owned_proxy(gateway, tmp_path, database, config=config, remove_environment=replica) as proxy,
proxy.scenario() as scenario,
):
key: Final = scenario.key(key_alias=f"metadata-fields-{uuid4()}")
yield _Isolated(proxy, scenario, key, database_url)
def _assert_filtered(row: Mapping[str, JsonValue], *kept: str) -> dict[str, JsonValue]:
metadata: Final = object_value(row["metadata"])
assert not set(EXCLUDED) & set(metadata), metadata
assert {"status", "cold_storage_object_key", *kept} <= set(metadata), metadata
return metadata
def _anthropic_event(event: Mapping[str, JsonValue]) -> bytes:
return f"event: {event['type']}\ndata: {json.dumps(event)}\n\n".encode()
def _anthropic_stream(model: JsonValue) -> tuple[bytes, ...]:
message: Final = {**_anthropic_message(model), "content": [], "stop_reason": None}
events: Final[tuple[Mapping[str, JsonValue], ...]] = (
{"type": "message_start", "message": message},
{"type": "content_block_start", "index": 0, "content_block": {"type": "text", "text": ""}},
{"type": "content_block_delta", "index": 0, "delta": {"type": "text_delta", "text": "ok"}},
{"type": "content_block_stop", "index": 0},
{"type": "message_delta", "delta": {"stop_reason": "end_turn"}, "usage": {"output_tokens": 2}},
{"type": "message_stop"},
)
return tuple(_anthropic_event(event) for event in events)
def _respond_streaming_or_not(request: Request) -> Reply:
if request.method == "GET":
return Reply(body=b'{"object":"list","data":[]}')
if request.target.startswith("/v1/messages"):
return _anthropic(request)
if JSON_OBJECT.validate_json(request.body).get("stream") is True:
return Reply(content_type="text/event-stream", chunks=text_stream(f"chatcmpl-{uuid4()}"))
return _respond(request)
def _stream(proxy: Gateway, key: str, path: str, body: Mapping[str, JsonValue]) -> None:
with proxy.client.stream("POST", path, json=dict(body), headers={"Authorization": f"Bearer {key}"}) as response:
assert response.status_code == 200, response.read().decode()
lines: Final = tuple(response.iter_lines())
assert any(line.startswith("data:") for line in lines), lines
def _call(proxy: Gateway, key: str, path: str, body: Mapping[str, JsonValue]) -> None:
response: Final = proxy.request("POST", path, body, key=key)
assert response.status_code == 200, response.text
def test_every_inference_surface_and_stream_stores_filtered_metadata(gateway: Gateway, tmp_path: Path) -> None:
with (
wire_server(_respond_streaming_or_not) as wire,
_isolated(gateway, _config(tmp_path, {"exclude": list(EXCLUDED)}), tmp_path) as isolated,
):
model: Final = isolated.scenario.model(
model="deepseek/gpt-4o-mini", api_base=wire.url + "/v1", api_key="synthetic-key"
)
anthropic: Final = isolated.scenario.model(
model="anthropic/claude-sonnet-4-6", api_base=wire.url, api_key="synthetic-key"
)
prompt: Final[list[JsonValue]] = [{"role": "user", "content": f"surfaces {uuid4()}"}]
calls: Final[tuple[Callable[[Gateway, str, str, Mapping[str, JsonValue]], None], ...]] = (
_stream,
_call,
_stream,
_call,
_stream,
)
requests: Final[tuple[tuple[str, Mapping[str, JsonValue]], ...]] = (
("/v1/chat/completions", {"model": model, "messages": prompt, "stream": True}),
("/v1/messages", {"model": anthropic, "max_tokens": 16, "messages": prompt}),
("/v1/messages", {"model": anthropic, "max_tokens": 16, "messages": prompt, "stream": True}),
("/v1/responses", {"model": model, "input": f"surfaces {uuid4()}"}),
("/v1/responses", {"model": model, "input": f"surfaces {uuid4()}", "stream": True}),
)
for send, (path, body) in zip(calls, requests, strict=True):
send(isolated.proxy, isolated.key, path, body)
rows: Final = isolated.rows(len(requests))
for row in rows:
assert row["status"] == "success", row
_assert_filtered(row, "usage_object")
def test_failed_request_keeps_its_error_and_status_but_drops_excluded_metadata(
gateway: Gateway, tmp_path: Path
) -> None:
def fail(request: Request) -> Reply:
if request.method == "GET":
return Reply(body=b'{"object":"list","data":[]}')
return Reply(status=500, body=b'{"error":{"message":"scripted upstream failure","type":"server_error"}}')
with (
wire_server(fail) as wire,
_isolated(gateway, _config(tmp_path, {"exclude": list(EXCLUDED)}), tmp_path) as isolated,
):
model: Final = isolated.scenario.model(model="openai/gpt-4o-mini", api_base=wire.url, api_key="synthetic")
response: Final = isolated.proxy.request(
"POST",
"/v1/chat/completions",
{"model": model, "messages": [{"role": "user", "content": "fail"}]},
key=isolated.key,
)
assert response.status_code >= 500, response.text
row: Final = isolated.rows(1)[0]
assert row["status"] == "failure", row
metadata: Final = _assert_filtered(row, "error_information")
assert metadata["status"] == "failure", metadata
assert "scripted upstream failure" in json.dumps(metadata["error_information"]), metadata
def test_response_cache_hit_row_is_filtered_and_charged_nothing(gateway: Gateway, tmp_path: Path) -> None:
config: Final = _config(
tmp_path, {"exclude": list(EXCLUDED)}, litellm_settings={"cache": True, "cache_params": {"type": "local"}}
)
with wire_server(_respond) as wire, _isolated(gateway, config, tmp_path) as isolated:
model: Final = isolated.scenario.model(model="openai/gpt-4o-mini", api_base=wire.url, api_key="synthetic")
text: Final = f"cache {uuid4()}"
first: Final = string_value(isolated.proxy.chat(model, key=isolated.key, text=text)["id"])
isolated.proxy.chat(model, key=isolated.key, text=text)
rows: Final = isolated.rows(2)
assert len([request for request in wire.drain() if request.method == "POST"]) == 1
paid, hit = sorted(rows, key=lambda row: row["cache_hit"] == "True")
assert paid["request_id"] == first and float(str(paid["spend"])) > 0, rows
assert hit["cache_hit"] == "True" and string_value(hit["request_id"]).startswith(first + "_cache_hit"), rows
assert float(str(hit["spend"])) == 0, rows
for row in rows:
_assert_filtered(row, "usage_object")
def test_batch_cost_row_is_filtered_and_charged_once(gateway: Gateway, tmp_path: Path) -> None:
with _isolated(gateway, _config(tmp_path, {"exclude": list(EXCLUDED)}), tmp_path) as isolated:
handle: Final = register_scenario(f"metadata-fields-batch-{uuid4().hex[:12]}", batch_routes("gpt-4o-mini"))
isolated.scenario.cleanups.callback(delete_scenario, handle)
model: Final = isolated.scenario.model(api_base=handle.api_base())
uploaded: Final = isolated.proxy.request_multipart(
"/v1/files",
{"purpose": "batch", "model": model},
{"file": ("in.jsonl", batch_input_file(model), "application/jsonl")},
key=isolated.key,
)
assert uploaded.status_code == 200, uploaded.text
batch: Final = isolated.proxy.post(
"/v1/batches",
{
"input_file_id": string_value(JSON_OBJECT.validate_json(uploaded.content)["id"]),
"endpoint": "/v1/chat/completions",
"completion_window": "24h",
"model": model,
},
key=isolated.key,
)
batch_path: Final = f"/v1/batches/{string_value(batch['id'])}"
retrievals: Final = tuple(isolated.proxy.request("GET", batch_path, key=isolated.key) for _ in range(2))
assert all(r.status_code == 200 and r.json()["status"] == "completed" for r in retrievals), retrievals
row: Final = isolated.rows(1, "call_type='aretrieve_batch'")[0]
assert float(str(row["spend"])) > 0, row
metadata: Final = _assert_filtered(row, "usage_object")
assert (metadata["batch_successful_requests"], metadata["batch_failed_requests"]) == (2, 3), metadata
def _update_retention(proxy: Gateway, value: JsonValue) -> httpx.Response:
return proxy.request(
"POST",
"/config/field/update",
{"field_name": "spend_logs_metadata_fields", "field_value": value, "config_type": "general_settings"},
)
def test_runtime_retention_update_rejects_invalid_values_and_applies_valid_ones_without_restart(
gateway: Gateway, tmp_path: Path
) -> None:
with (
scratch_database() as database_url,
wire_server(_respond) as wire,
_isolated(gateway, _config(tmp_path, None), tmp_path, database_url) as isolated,
):
model: Final = isolated.scenario.model(model="openai/gpt-4o-mini", api_base=wire.url, api_key="synthetic")
invalid: Final[tuple[JsonValue, ...]] = (
{"include": ["usage_object"], "exclude": ["model_map_information"]},
{},
{"exclude": ["not_a_metadata_field"]},
{"exclude": ["status"]},
)
isolated.proxy.chat(model, key=isolated.key)
assert set(EXCLUDED) <= set(object_value(isolated.rows(1)[0]["metadata"]))
for value in invalid:
assert _update_retention(isolated.proxy, value).status_code == 400, value
isolated.proxy.chat(model, key=isolated.key)
assert set(EXCLUDED) <= set(object_value(isolated.rows(2)[-1]["metadata"]))
assert _update_retention(isolated.proxy, {"exclude": list(EXCLUDED)}).status_code == 200
isolated.proxy.chat(model, key=isolated.key)
_assert_filtered(isolated.rows(3)[-1])
for value in invalid:
assert _update_retention(isolated.proxy, value).status_code == 400, value
isolated.proxy.chat(model, key=isolated.key)
_assert_filtered(isolated.rows(4)[-1])
SAVINGS_FIELDS: Final = ("autorouter_savings_estimate", "autorouter_savings")
def _anthropic_message(model: JsonValue) -> dict[str, JsonValue]:
return {
"id": f"msg_{uuid4().hex}",
"type": "message",
"role": "assistant",
"model": model,
"content": [{"type": "text", "text": "ok"}],
"stop_reason": "end_turn",
"stop_sequence": None,
"usage": {
"input_tokens": 12,
"output_tokens": 2,
"cache_creation_input_tokens": 0,
"cache_read_input_tokens": 0,
"cache_creation": {"ephemeral_5m_input_tokens": 0, "ephemeral_1h_input_tokens": 0},
},
}
def _anthropic(request: Request) -> Reply:
if request.target.startswith("/v1/messages/count_tokens"):
return Reply(body=b'{"input_tokens":12}')
body: Final = JSON_OBJECT.validate_json(request.body)
if body.get("stream") is True:
return Reply(content_type="text/event-stream", chunks=_anthropic_stream(body["model"]))
return Reply(body=json.dumps(_anthropic_message(body["model"])).encode())
def _router_models(wire: Wire, router: str) -> tuple[Mapping[str, JsonValue], ...]:
tiers: Final = {"cheap": "anthropic/claude-sonnet-4-6", "frontier": "anthropic/claude-opus-4-8"}
return (
*(
{"model_name": name, "litellm_params": {"model": model, "api_base": wire.url, "api_key": "synthetic"}}
for name, model in tiers.items()
),
{
"model_name": router,
"litellm_params": {
"model": "auto_router/complexity_router",
"complexity_router_config": {
"tiers": {"SIMPLE": "cheap", "MEDIUM": "cheap", "COMPLEX": "frontier", "REASONING": "frontier"},
},
},
},
)
@pytest.mark.parametrize(
("retention", "stored"),
[(None, SAVINGS_FIELDS), ({"exclude": list(SAVINGS_FIELDS)}, ())],
ids=["unset-keeps-savings", "excluded-savings-stay-out"],
)
@pytest.mark.timeout(240) # projection retries a scope at most every 30 seconds after its first attempt
def test_delayed_autorouter_savings_publication_respects_retention(
gateway: Gateway, tmp_path: Path, retention: Mapping[str, JsonValue] | None, stored: tuple[str, ...]
) -> None:
router: Final = f"router-{uuid4().hex[:8]}"
with wire_server(_anthropic) as wire:
config: Final = _config(tmp_path, retention, model_list=_router_models(wire, router))
with _isolated(gateway, config, tmp_path) as isolated:
response: Final = isolated.proxy.request(
"POST",
"/v1/messages",
{"model": router, "max_tokens": 16, "messages": [{"role": "user", "content": "hi"}]},
key=isolated.key,
headers={"x-litellm-session-id": f"session-{uuid4()}"},
)
assert response.status_code == 200, response.text
published: Final = isolated.rows(
1,
"metadata::jsonb ? 'routing_decision' AND NOT metadata::jsonb ? 'autorouter_baseline_observation' "
'AND request_id IN (SELECT request_id FROM "LiteLLM_AutoRouterBaselineObservation" '
"WHERE publication IS NOT NULL)",
seconds=150,
)[0]
metadata: Final = object_value(published["metadata"])
assert {name for name in SAVINGS_FIELDS if name in metadata} == set(stored), metadata