mirror of
https://github.com/BerriAI/litellm.git
synced 2026-10-03 02:22:24 +00:00
test(autorouter): consolidate savings regressions
Some checks failed
LiteLLM Rust / rust-lint (push) Has been cancelled
LiteLLM Rust / rust-test (push) Has been cancelled
LiteLLM Rust / rust-wheel (push) Has been cancelled
Terraform Modules / fmt, validate, test (aws) (push) Has been cancelled
Terraform Modules / fmt, validate, test (gcp) (push) Has been cancelled
Terraform Provider / gofmt, vet, build, test (push) Has been cancelled
Terraform Provider / Provider endpoints vs proxy OpenAPI schema (push) Has been cancelled
Some checks failed
LiteLLM Rust / rust-lint (push) Has been cancelled
LiteLLM Rust / rust-test (push) Has been cancelled
LiteLLM Rust / rust-wheel (push) Has been cancelled
Terraform Modules / fmt, validate, test (aws) (push) Has been cancelled
Terraform Modules / fmt, validate, test (gcp) (push) Has been cancelled
Terraform Provider / gofmt, vet, build, test (push) Has been cancelled
Terraform Provider / Provider endpoints vs proxy OpenAPI schema (push) Has been cancelled
This commit is contained in:
parent
e87b4207d4
commit
f996803d0b
13 changed files with 112 additions and 1492 deletions
|
|
@ -1,22 +1,14 @@
|
|||
import json
|
||||
import os
|
||||
import time
|
||||
import uuid
|
||||
from collections.abc import Mapping
|
||||
from datetime import timedelta
|
||||
from typing import Final
|
||||
from urllib.parse import parse_qsl, urlencode, urlsplit, urlunsplit
|
||||
|
||||
import pytest
|
||||
from prisma import Prisma
|
||||
from prisma.errors import RawQueryError
|
||||
|
||||
from litellm.proxy.common_utils.user_api_key_cache import UserApiKeyCache
|
||||
from litellm.proxy.db.autorouter_daily_spend import AUTOROUTER_DAILY_COSTS_SQL, AutoRouterDailyCosts
|
||||
from litellm.proxy.db.autorouter_historical_spend import AUTOROUTER_HISTORICAL_COSTS_SQL, recover_daily_router_costs
|
||||
from litellm.proxy.db.daily_spend_bulk_upsert import DAILY_SPEND_TABLES, build_bulk_upsert, merge_by_conflict_key
|
||||
from litellm.proxy.db.autorouter_historical_spend import AUTOROUTER_HISTORICAL_COSTS_SQL
|
||||
from litellm.proxy.route_llm_request import ROUTE_ENDPOINT_MAPPING
|
||||
from litellm.proxy.utils import PrismaClient, ProxyLogging
|
||||
from litellm.proxy.db.daily_spend_bulk_upsert import DAILY_SPEND_TABLES, build_bulk_upsert, merge_by_conflict_key
|
||||
|
||||
pytestmark = pytest.mark.asyncio(loop_scope="session")
|
||||
|
||||
|
|
@ -91,183 +83,36 @@ async def test_an_old_writer_cannot_certify_partial_daily_costs_as_complete(db:
|
|||
assert costs.baseline_spend(7.0) is None
|
||||
|
||||
|
||||
async def test_retained_historical_costs_reconcile_through_prisma_and_reject_a_missing_free_request(db: Prisma) -> None:
|
||||
@pytest.mark.parametrize("missing", (None, "routed", "free", "classifier"))
|
||||
async def test_historical_costs_require_all_logs_and_count_classifier_charges_once(db: Prisma, missing: str | None) -> None:
|
||||
async with db.tx() as tx:
|
||||
await tx.execute_raw(
|
||||
'CREATE TEMP TABLE "LiteLLM_DailyUserSpend" (LIKE public."LiteLLM_DailyUserSpend" INCLUDING ALL) ON COMMIT DROP'
|
||||
)
|
||||
await tx.execute_raw(
|
||||
'CREATE TEMP TABLE "LiteLLM_SpendLogs" (LIKE public."LiteLLM_SpendLogs" INCLUDING ALL) ON COMMIT DROP'
|
||||
)
|
||||
cases: Final = (
|
||||
("2026-09-21", "a", "owner-a", 90.0, 1.0, 9.0),
|
||||
("2026-09-22", "a", "owner-a", 2.0, 0.2, 3.0),
|
||||
("2026-09-22", "b", "owner-b", 7.0, 0.5, 5.0),
|
||||
)
|
||||
rows: Final = tuple(
|
||||
{
|
||||
"date": day,
|
||||
"api_key": key,
|
||||
"user_id": user,
|
||||
"model": "removed-historical-model",
|
||||
"custom_llm_provider": "removed-provider",
|
||||
"endpoint": "/responses",
|
||||
"spend": llm + classifier,
|
||||
"api_requests": 2,
|
||||
"successful_requests": 2,
|
||||
"prompt_tokens": 25,
|
||||
"completion_tokens": 10,
|
||||
"autorouter_savings_spend": saved,
|
||||
}
|
||||
for day, key, user, llm, classifier, saved in cases
|
||||
)
|
||||
table: Final = DAILY_SPEND_TABLES["user"]
|
||||
sql, values = build_bulk_upsert(table, merge_by_conflict_key(table, rows))
|
||||
await tx.execute_raw(sql, *values)
|
||||
logs: Final = tuple(
|
||||
{
|
||||
"request_id": f"{day}-{key}-{suffix}",
|
||||
"started_at": f"{day}T12:00:00",
|
||||
"api_key": key,
|
||||
"user_id": user,
|
||||
"spend": spend,
|
||||
"prompt_tokens": prompt,
|
||||
"completion_tokens": completion,
|
||||
"metadata": metadata,
|
||||
}
|
||||
for day, key, user, llm, classifier, saved in cases
|
||||
for suffix, spend, prompt, completion, metadata in (
|
||||
(
|
||||
"routed",
|
||||
llm,
|
||||
20,
|
||||
10,
|
||||
{
|
||||
"routing_decision": {"router_model_name": "retired-router", "classifier_cost": classifier},
|
||||
"autorouter_savings": saved,
|
||||
},
|
||||
),
|
||||
("free", 0.0, 0, 0, {}),
|
||||
(
|
||||
"classifier",
|
||||
classifier,
|
||||
5,
|
||||
0,
|
||||
{
|
||||
"internal_call_origin": "autorouter_classifier",
|
||||
"routing_decision": {"router_model_name": "retired-router", "classifier_cost": classifier},
|
||||
},
|
||||
),
|
||||
)
|
||||
)
|
||||
await tx.execute_raw(
|
||||
"""INSERT INTO "LiteLLM_SpendLogs"
|
||||
(request_id, "startTime", "endTime", call_type, api_key, "user", model,
|
||||
custom_llm_provider, spend, prompt_tokens, completion_tokens, status, metadata)
|
||||
SELECT request_id, started_at::timestamp, started_at::timestamp, 'aresponses', api_key, user_id,
|
||||
'removed-historical-model', 'removed-provider', spend, prompt_tokens, completion_tokens,
|
||||
'success', metadata
|
||||
FROM jsonb_to_recordset($1::jsonb) AS rows(
|
||||
request_id text, started_at text, api_key text, user_id text,
|
||||
spend float8, prompt_tokens int, completion_tokens int, metadata jsonb
|
||||
)""",
|
||||
json.dumps(logs),
|
||||
)
|
||||
for key, user, actual, saved, requests in (
|
||||
(None, None, 9.7, 8.0, 2),
|
||||
("a", None, 2.2, 3.0, 1),
|
||||
(None, "owner-a", 2.2, 3.0, 1),
|
||||
("a", "owner-a", 2.2, 3.0, 1),
|
||||
("a", "owner-b", 0.0, 0.0, 0),
|
||||
for table in ("LiteLLM_DailyUserSpend", "LiteLLM_SpendLogs"):
|
||||
await tx.execute_raw(f'CREATE TEMP TABLE "{table}" (LIKE public."{table}" INCLUDING ALL) ON COMMIT DROP')
|
||||
await tx.execute_raw("""INSERT INTO "LiteLLM_DailyUserSpend"
|
||||
(id,date,user_id,api_key,model,custom_llm_provider,mcp_namespaced_tool_name,endpoint,
|
||||
api_requests,successful_requests,prompt_tokens,completion_tokens,spend,autorouter_savings_spend,updated_at)
|
||||
VALUES ('history','2026-09-22','owner','key','model','provider','','/chat/completions',2,2,25,10,2.1,4,NOW())
|
||||
""")
|
||||
decision: Final = {"router_model_name": "router", "classifier_cost": 0.1}
|
||||
for request_id, spend, prompt, completion, metadata in (
|
||||
("routed", 2.0, 20, 10, {"routing_decision": decision, "autorouter_savings": 4}),
|
||||
("free", 0.0, 0, 0, {}),
|
||||
("classifier", 0.1, 5, 0, {"internal_call_origin": "autorouter_classifier"}),
|
||||
):
|
||||
result: Final = await tx.query_raw(
|
||||
AUTOROUTER_HISTORICAL_COSTS_SQL,
|
||||
"2026-09-22",
|
||||
"2026-09-22",
|
||||
key,
|
||||
user,
|
||||
json.dumps(ROUTE_ENDPOINT_MAPPING),
|
||||
)
|
||||
recovered: Final = AutoRouterDailyCosts.model_validate(result[0])
|
||||
assert recovered.complete and recovered.comparison_complete
|
||||
assert recovered.requests == requests
|
||||
assert recovered.recorded_spend == pytest.approx(actual)
|
||||
assert recovered.saved_spend == saved
|
||||
assert recovered.baseline_spend(saved) == pytest.approx(actual + saved)
|
||||
|
||||
await tx.execute_raw('DELETE FROM "LiteLLM_SpendLogs" WHERE request_id = $1', "2026-09-22-a-free")
|
||||
incomplete: Final = await tx.query_raw(
|
||||
AUTOROUTER_HISTORICAL_COSTS_SQL,
|
||||
"2026-09-22",
|
||||
"2026-09-22",
|
||||
"a",
|
||||
None,
|
||||
await tx.execute_raw("""INSERT INTO "LiteLLM_SpendLogs"
|
||||
(request_id,"startTime","endTime","user",api_key,model,custom_llm_provider,
|
||||
call_type,spend,prompt_tokens,completion_tokens,status,metadata)
|
||||
VALUES ($1,'2026-09-22 12:00:00','2026-09-22 12:00:00','owner','key','model','provider',
|
||||
'acompletion',$2::float8,$3::integer,$4::integer,'success',$5::jsonb)
|
||||
""", request_id, spend, prompt, completion, json.dumps(metadata))
|
||||
if missing is not None:
|
||||
await tx.execute_raw('DELETE FROM "LiteLLM_SpendLogs" WHERE request_id=$1', missing)
|
||||
rows: Final = await tx.query_raw(
|
||||
AUTOROUTER_HISTORICAL_COSTS_SQL, "2026-09-22", "2026-09-22", "key", "owner",
|
||||
json.dumps(ROUTE_ENDPOINT_MAPPING),
|
||||
)
|
||||
costs: Final = AutoRouterDailyCosts.model_validate(incomplete[0])
|
||||
assert not costs.complete
|
||||
assert costs.saved_spend == 3.0
|
||||
assert costs.recorded_spend is None
|
||||
assert costs.baseline_spend(3.0) is None
|
||||
|
||||
|
||||
async def test_blocked_historical_recovery_times_out_without_changing_costs_or_read_connection(
|
||||
db: Prisma, monkeypatch: pytest.MonkeyPatch
|
||||
) -> None:
|
||||
identity: Final = uuid.uuid4().hex
|
||||
database_url: Final = os.environ["DATABASE_URL"]
|
||||
parsed: Final = urlsplit(database_url)
|
||||
query: Final = urlencode({**dict(parse_qsl(parsed.query)), "connection_limit": "1"})
|
||||
monkeypatch.setenv("DATABASE_URL_READ_REPLICA", urlunsplit(parsed._replace(query=query)))
|
||||
client: Final = PrismaClient(database_url, ProxyLogging(UserApiKeyCache()))
|
||||
row: Final = {
|
||||
**_daily("2026-09-22", identity, identity, 2.0, 0.1, 7.0),
|
||||
"api_requests": 2,
|
||||
"successful_requests": 2,
|
||||
"spend": 10.0,
|
||||
}
|
||||
table: Final = DAILY_SPEND_TABLES["user"]
|
||||
statement, values = build_bulk_upsert(table, merge_by_conflict_key(table, (row,)))
|
||||
await db.execute_raw(statement, *values)
|
||||
try:
|
||||
await client.db.connect()
|
||||
await client.read_db.execute_raw("SET statement_timeout = 4500")
|
||||
connection_query: Final = (
|
||||
"SELECT pg_backend_pid() AS pid, current_setting('statement_timeout') AS timeout, "
|
||||
"current_setting('transaction_read_only') AS read_only"
|
||||
)
|
||||
connection_before: Final = await client.read_db.query_raw(connection_query)
|
||||
daily_before: Final = await client.read_db.query_raw(
|
||||
AUTOROUTER_DAILY_COSTS_SQL,
|
||||
"2026-09-22",
|
||||
"2026-09-22",
|
||||
identity,
|
||||
identity,
|
||||
)
|
||||
recorded: Final = AutoRouterDailyCosts.model_validate(daily_before[0])
|
||||
assert recorded.recorded_spend == pytest.approx(2.1) and recorded.saved_spend == 7.0
|
||||
assert recorded.coverage == "partial"
|
||||
async with db.tx(timeout=timedelta(seconds=8)) as blocker:
|
||||
await blocker.execute_raw('LOCK TABLE "LiteLLM_SpendLogs" IN ACCESS EXCLUSIVE MODE')
|
||||
started: Final = time.monotonic()
|
||||
with pytest.raises(RawQueryError, match="canceling statement due to statement timeout"):
|
||||
await recover_daily_router_costs(client, "2026-09-22", "2026-09-22", identity, identity)
|
||||
assert time.monotonic() - started < 3.5
|
||||
assert await client.read_db.query_raw(connection_query) == connection_before
|
||||
assert (
|
||||
await client.read_db.query_raw(
|
||||
AUTOROUTER_DAILY_COSTS_SQL,
|
||||
"2026-09-22",
|
||||
"2026-09-22",
|
||||
identity,
|
||||
identity,
|
||||
)
|
||||
== daily_before
|
||||
)
|
||||
recovered: Final = await recover_daily_router_costs(client, "2026-09-22", "2026-09-22", identity, identity)
|
||||
assert recovered is not None
|
||||
assert recovered.recorded_spend == recorded.recorded_spend and recovered.saved_spend == recorded.saved_spend
|
||||
assert await client.read_db.query_raw(connection_query) == connection_before
|
||||
finally:
|
||||
await client.db.disconnect()
|
||||
await db.execute_raw('DELETE FROM "LiteLLM_DailyUserSpend" WHERE api_key=$1', identity)
|
||||
costs: Final = AutoRouterDailyCosts.model_validate(rows[0])
|
||||
assert costs.complete is (missing is None)
|
||||
assert costs.saved_spend == 4
|
||||
assert costs.recorded_spend == (pytest.approx(2.1) if missing is None else None)
|
||||
assert costs.baseline_spend(4) == (pytest.approx(6.1) if missing is None else None)
|
||||
|
|
|
|||
|
|
@ -15,12 +15,10 @@ from litellm.proxy.db.autorouter_session_rollup import AutoRouterTurnTransaction
|
|||
from litellm.proxy.db.baseline_accounting import (
|
||||
BaselineAccountingRecord,
|
||||
BaselineAccountingStore,
|
||||
BaselinePublication,
|
||||
DailyBaselineAttribution,
|
||||
DailyBaselineTarget,
|
||||
)
|
||||
from litellm.proxy.db.create_views import SupportsRawQueries
|
||||
from litellm.proxy.db.daily_spend_bulk_upsert import DAILY_SPEND_TABLES, build_bulk_upsert, merge_by_conflict_key
|
||||
from litellm.proxy.spend_tracking.baseline_accounting import BaselineObservation
|
||||
from litellm.proxy.spend_tracking.savings import BaselineCostSnapshot
|
||||
from litellm.types.utils import Usage
|
||||
|
|
@ -337,155 +335,3 @@ async def test_native_observation_enters_spend_pipeline_once_with_shared_daily_a
|
|||
assert tag_rows[0]["spend"] == tag_rows[0]["api_requests"] == 0
|
||||
finally:
|
||||
await client.db.disconnect()
|
||||
|
||||
|
||||
async def test_old_publisher_transitions_keep_daily_comparison_complete_after_retention(
|
||||
db: Prisma, record: Callable[..., BaselineAccountingRecord]
|
||||
) -> None:
|
||||
event: Final = record("old-publisher")
|
||||
assert event.daily is not None
|
||||
assert await _store(db).append(event) == "recorded"
|
||||
existing_actual: Final = 5.0
|
||||
other_savings: Final = 3.0
|
||||
actual: Final = event.pricing.actual_spend
|
||||
row: Final = {
|
||||
**event.daily.model_dump(exclude={"targets"}),
|
||||
"user_id": event.api_key,
|
||||
"spend": existing_actual + actual,
|
||||
"api_requests": 2,
|
||||
"successful_requests": 2,
|
||||
"autorouter_accounted_requests": 2,
|
||||
"autorouter_requests": 2,
|
||||
"autorouter_llm_spend": existing_actual + actual,
|
||||
"autorouter_classifier_cost_recorded_requests": 2,
|
||||
"autorouter_estimated_requests": 1,
|
||||
"autorouter_estimated_actual_spend": existing_actual,
|
||||
"autorouter_savings_spend": other_savings,
|
||||
}
|
||||
table: Final = DAILY_SPEND_TABLES["user"]
|
||||
statement, values = build_bulk_upsert(table, merge_by_conflict_key(table, (row,)))
|
||||
await db.execute_raw(statement, *values)
|
||||
|
||||
for status, baseline, savings_delta, expected_savings, count, covered_actual in (
|
||||
("estimated", actual + 0.5, 0.5, other_savings + 0.5, 2, existing_actual + actual),
|
||||
("estimated", actual, -0.5, other_savings, 2, existing_actual + actual),
|
||||
("unknown", None, 0.0, other_savings, 1, existing_actual),
|
||||
("estimated", actual + 0.2, 0.2, other_savings + 0.2, 2, existing_actual + actual),
|
||||
("estimated", actual + 0.2, 0.0, other_savings + 0.2, 2, existing_actual + actual),
|
||||
):
|
||||
publication: Final = BaselinePublication(
|
||||
comparison_id=event.scope,
|
||||
comparison_started_at=event.observation.started_at,
|
||||
status=status,
|
||||
reason="legacy-publisher-transition",
|
||||
actual_spend=actual if baseline is not None else None,
|
||||
baseline_spend=baseline,
|
||||
)
|
||||
async with db.tx() as tx:
|
||||
if savings_delta:
|
||||
await tx.execute_raw(
|
||||
'UPDATE "LiteLLM_DailyUserSpend" SET autorouter_savings_spend=autorouter_savings_spend+$1::float8 '
|
||||
"WHERE api_key=$2 AND user_id=$3",
|
||||
savings_delta,
|
||||
event.api_key,
|
||||
event.api_key,
|
||||
)
|
||||
await tx.execute_raw(
|
||||
'UPDATE "LiteLLM_AutoRouterBaselineObservation" SET publication=$1 WHERE request_id=$2',
|
||||
publication.model_dump_json(),
|
||||
event.observation.request_id,
|
||||
)
|
||||
daily: Final = await _daily_user(db, event)
|
||||
assert daily["autorouter_estimated_requests"] == count
|
||||
assert daily["autorouter_estimated_actual_spend"] == pytest.approx(covered_actual)
|
||||
assert daily["autorouter_savings_spend"] == pytest.approx(expected_savings)
|
||||
assert daily["api_requests"] == daily["autorouter_accounted_requests"] == daily["autorouter_requests"] == 2
|
||||
assert daily["spend"] == daily["autorouter_llm_spend"] == pytest.approx(existing_actual + actual)
|
||||
assert daily["autorouter_classifier_cost"] == 0
|
||||
assert daily["autorouter_classifier_cost_recorded_requests"] == 2
|
||||
shadow: Final = await db.query_raw(
|
||||
"SELECT publication::jsonb = daily_costs_publication::jsonb AS accounted "
|
||||
'FROM "LiteLLM_AutoRouterBaselineObservation" WHERE request_id=$1',
|
||||
event.observation.request_id,
|
||||
)
|
||||
assert shadow == [{"accounted": True}]
|
||||
|
||||
retained: Final = await _daily_user(db, event)
|
||||
await db.execute_raw(
|
||||
'DELETE FROM "LiteLLM_AutoRouterBaselineObservation" WHERE request_id=$1', event.observation.request_id
|
||||
)
|
||||
assert await _daily_user(db, event) == retained
|
||||
|
||||
|
||||
@pytest.mark.parametrize("targets", ("missing", "duplicate", "null-and-empty"))
|
||||
async def test_legacy_zero_publication_repairs_shadow_once_for_normalized_user_targets(
|
||||
db: Prisma, record: Callable[..., BaselineAccountingRecord], targets: str
|
||||
) -> None:
|
||||
source: Final = record("legacy-zero")
|
||||
assert source.daily is not None
|
||||
identities: Final = (
|
||||
() if targets == "missing" else ((None, "") if targets == "null-and-empty" else (source.api_key,) * 2)
|
||||
)
|
||||
event: Final = source.model_copy(
|
||||
update={
|
||||
"daily": source.daily.model_copy(
|
||||
update={
|
||||
"targets": tuple(DailyBaselineTarget(entity="user", entity_id=identity) for identity in identities),
|
||||
}
|
||||
)
|
||||
}
|
||||
)
|
||||
publication: Final = BaselinePublication(
|
||||
comparison_id=event.scope,
|
||||
comparison_started_at=event.observation.started_at,
|
||||
status="estimated",
|
||||
reason="legacy-zero",
|
||||
actual_spend=event.pricing.actual_spend,
|
||||
baseline_spend=event.pricing.actual_spend,
|
||||
).model_dump_json()
|
||||
await db.execute_raw(
|
||||
'INSERT INTO "LiteLLM_AutoRouterBaselineObservation" (request_id,scope,started_at,revision,data,publication) '
|
||||
"VALUES ($1,$2,$3::float8,1,$4,$5)",
|
||||
event.observation.request_id,
|
||||
event.scope,
|
||||
event.observation.started_at,
|
||||
event.model_dump_json(),
|
||||
publication,
|
||||
)
|
||||
before: Final = await db.query_raw(
|
||||
'SELECT daily_costs_publication FROM "LiteLLM_AutoRouterBaselineObservation" WHERE request_id=$1',
|
||||
event.observation.request_id,
|
||||
)
|
||||
assert before == [{"daily_costs_publication": None}]
|
||||
for _ in range(2):
|
||||
await db.execute_raw(
|
||||
'UPDATE "LiteLLM_AutoRouterBaselineObservation" SET publication=$1 WHERE request_id=$2',
|
||||
publication,
|
||||
event.observation.request_id,
|
||||
)
|
||||
rows: Final = await db.query_raw(
|
||||
"SELECT user_id,autorouter_estimated_requests,autorouter_estimated_actual_spend,autorouter_savings_spend, "
|
||||
'api_requests,autorouter_requests,spend FROM "LiteLLM_DailyUserSpend" WHERE api_key=$1',
|
||||
event.api_key,
|
||||
)
|
||||
assert rows == (
|
||||
[
|
||||
{
|
||||
"user_id": "" if targets == "null-and-empty" else event.api_key,
|
||||
"autorouter_estimated_requests": 1,
|
||||
"autorouter_estimated_actual_spend": event.pricing.actual_spend,
|
||||
"autorouter_savings_spend": 0,
|
||||
"api_requests": 0,
|
||||
"autorouter_requests": 0,
|
||||
"spend": 0,
|
||||
}
|
||||
]
|
||||
if identities
|
||||
else []
|
||||
)
|
||||
shadow: Final = await db.query_raw(
|
||||
"SELECT publication::jsonb = daily_costs_publication::jsonb AS accounted "
|
||||
'FROM "LiteLLM_AutoRouterBaselineObservation" WHERE request_id=$1',
|
||||
event.observation.request_id,
|
||||
)
|
||||
assert shadow == [{"accounted": True}]
|
||||
|
|
|
|||
|
|
@ -5,7 +5,6 @@ import json
|
|||
import os
|
||||
import time
|
||||
from collections.abc import AsyncGenerator, Iterator
|
||||
from concurrent.futures import ThreadPoolExecutor
|
||||
from contextlib import asynccontextmanager
|
||||
from datetime import timedelta
|
||||
from pathlib import Path
|
||||
|
|
@ -110,312 +109,42 @@ _COVERAGE_MIGRATION: Final = Path(__file__).parents[2] / (
|
|||
)
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def _coverage_postgresql() -> Iterator[psycopg.Connection[tuple[object, ...]]]:
|
||||
base: Final = os.environ["DATABASE_URL"].split("?")[0]
|
||||
name: Final = f"coverage_{uuid4().hex}"
|
||||
with psycopg.connect(base, autocommit=True) as admin:
|
||||
admin.execute(sql.SQL("CREATE DATABASE {}").format(sql.Identifier(name)))
|
||||
try:
|
||||
with psycopg.connect(base, dbname=name) as connection:
|
||||
yield connection
|
||||
finally:
|
||||
admin.execute(sql.SQL("DROP DATABASE {}").format(sql.Identifier(name)))
|
||||
|
||||
|
||||
_COVERAGE_BASE: Final = """
|
||||
CREATE TABLE "LiteLLM_DailyUserSpend" (
|
||||
id TEXT PRIMARY KEY, user_id TEXT, date TEXT NOT NULL, api_key TEXT NOT NULL,
|
||||
model TEXT, custom_llm_provider TEXT, mcp_namespaced_tool_name TEXT, endpoint TEXT,
|
||||
model_group TEXT, spend DOUBLE PRECISION DEFAULT 0, autorouter_savings_spend DOUBLE PRECISION DEFAULT 0,
|
||||
updated_at TIMESTAMP,
|
||||
UNIQUE (user_id, date, api_key, model, custom_llm_provider, mcp_namespaced_tool_name, endpoint)
|
||||
);
|
||||
CREATE TABLE "LiteLLM_AutoRouterBaselineObservation" (
|
||||
request_id TEXT PRIMARY KEY, data TEXT, publication TEXT
|
||||
);
|
||||
INSERT INTO "LiteLLM_DailyUserSpend"
|
||||
(id,user_id,date,api_key,model,custom_llm_provider,mcp_namespaced_tool_name,endpoint,spend,autorouter_savings_spend)
|
||||
VALUES ('recorded','owner','2026-09-23','key','','','','',10,30)
|
||||
"""
|
||||
|
||||
|
||||
def _coverage_tables(conn: psycopg.Connection[tuple[object, ...]]) -> None:
|
||||
conn.execute(_COVERAGE_BASE)
|
||||
conn.execute(
|
||||
'INSERT INTO "LiteLLM_AutoRouterBaselineObservation" (request_id,data) VALUES (%s,%s)',
|
||||
(
|
||||
"request",
|
||||
json.dumps(
|
||||
{
|
||||
"daily": {
|
||||
"date": "2026-09-23",
|
||||
"api_key": "key",
|
||||
"targets": [{"entity": "user", "entity_id": "owner"}],
|
||||
}
|
||||
}
|
||||
),
|
||||
),
|
||||
@pytest.mark.parametrize("repair", (False, True))
|
||||
def test_daily_coverage_install_and_old_publisher_preserve_history(
|
||||
database: tuple[str, psycopg.Connection[tuple[object, ...]]], repair: bool,
|
||||
) -> None:
|
||||
_, connection = database
|
||||
connection.execute("""CREATE TABLE "LiteLLM_DailyUserSpend" (
|
||||
id TEXT PRIMARY KEY, user_id TEXT, date TEXT, api_key TEXT, model TEXT, custom_llm_provider TEXT,
|
||||
mcp_namespaced_tool_name TEXT, endpoint TEXT, model_group TEXT, updated_at TIMESTAMP,
|
||||
spend FLOAT8 DEFAULT 0, autorouter_savings_spend FLOAT8 DEFAULT 0,
|
||||
UNIQUE (user_id, date, api_key, model, custom_llm_provider, mcp_namespaced_tool_name, endpoint));
|
||||
INSERT INTO "LiteLLM_DailyUserSpend"
|
||||
(id,user_id,date,api_key,model,custom_llm_provider,mcp_namespaced_tool_name,endpoint,spend,autorouter_savings_spend)
|
||||
VALUES ('recorded','owner','2026-09-23','key','','','','',10,30)""")
|
||||
connection.execute(
|
||||
'INSERT INTO "LiteLLM_AutoRouterBaselineObservation" (request_id,scope,started_at,revision,data) '
|
||||
"VALUES ('request','scope',0,1,%s)",
|
||||
(json.dumps({"daily": {"date": "2026-09-23", "api_key": "key",
|
||||
"targets": [{"entity": "user", "entity_id": "owner"}]}}),),
|
||||
)
|
||||
conn.commit()
|
||||
|
||||
|
||||
def _coverage_install(conn: psycopg.Connection[tuple[object, ...]]) -> None:
|
||||
conn.execute(_COVERAGE_MIGRATION.read_text())
|
||||
|
||||
|
||||
def _coverage_publish(conn: psycopg.Connection[tuple[object, ...]]) -> tuple[object, ...] | None:
|
||||
conn.execute(
|
||||
'UPDATE "LiteLLM_AutoRouterBaselineObservation" SET publication=%s WHERE request_id=%s',
|
||||
(json.dumps({"status": "estimated", "actual_spend": 2, "baseline_spend": 2}), "request"),
|
||||
)
|
||||
row: Final = conn.execute(
|
||||
"SELECT id,spend,autorouter_savings_spend,autorouter_estimated_requests,autorouter_estimated_actual_spend "
|
||||
'FROM "LiteLLM_DailyUserSpend"'
|
||||
).fetchone()
|
||||
conn.commit()
|
||||
return row
|
||||
|
||||
|
||||
def _coverage_repeat_during_traffic(conn: psycopg.Connection[tuple[object, ...]]) -> None:
|
||||
with (
|
||||
psycopg.connect(conn.info.dsn, password=conn.info.password) as traffic,
|
||||
psycopg.connect(conn.info.dsn, password=conn.info.password, autocommit=True) as installer,
|
||||
connection.execute(_COVERAGE_MIGRATION.read_text())
|
||||
if repair:
|
||||
connection.execute('ALTER TABLE "LiteLLM_AutoRouterBaselineObservation" '
|
||||
'DISABLE TRIGGER litellm_update_daily_autorouter_coverage')
|
||||
connection.execute(_COVERAGE_MIGRATION.read_text())
|
||||
for status, expected_count, expected_actual in (
|
||||
("estimated", 1, 2), ("estimated", 1, 2), ("unknown", 0, 0), ("estimated", 1, 2),
|
||||
):
|
||||
traffic.execute('SELECT * FROM "LiteLLM_DailyUserSpend"')
|
||||
traffic.execute('UPDATE "LiteLLM_AutoRouterBaselineObservation" SET data=data')
|
||||
with ThreadPoolExecutor(max_workers=1) as pool:
|
||||
repeated: Final = pool.submit(_coverage_install, installer)
|
||||
try:
|
||||
repeated.result(timeout=2)
|
||||
finally:
|
||||
installer.cancel()
|
||||
traffic.rollback()
|
||||
|
||||
|
||||
def test_coverage_install_preserves_history_and_repeat_does_not_wait_for_table_traffic(
|
||||
_coverage_postgresql: psycopg.Connection[tuple[object, ...]],
|
||||
) -> None:
|
||||
_coverage_tables(_coverage_postgresql)
|
||||
_coverage_install(_coverage_postgresql)
|
||||
_coverage_repeat_during_traffic(_coverage_postgresql)
|
||||
|
||||
assert _coverage_publish(_coverage_postgresql) == ("recorded", 10, 30, 1, 2)
|
||||
assert _coverage_publish(_coverage_postgresql) == ("recorded", 10, 30, 1, 2)
|
||||
|
||||
|
||||
@pytest.mark.parametrize("damage", ("missing", "disabled", "wrong-function", "wrong-event", "wrong-body"))
|
||||
def test_coverage_installer_repairs_trigger_before_publishing_without_rewriting_history(
|
||||
_coverage_postgresql: psycopg.Connection[tuple[object, ...]],
|
||||
damage: str,
|
||||
) -> None:
|
||||
_coverage_tables(_coverage_postgresql)
|
||||
_coverage_install(_coverage_postgresql)
|
||||
if damage == "disabled":
|
||||
_coverage_postgresql.execute(
|
||||
'ALTER TABLE "LiteLLM_AutoRouterBaselineObservation" '
|
||||
"DISABLE TRIGGER litellm_update_daily_autorouter_coverage"
|
||||
connection.execute(
|
||||
'UPDATE "LiteLLM_AutoRouterBaselineObservation" SET publication=%s WHERE request_id=%s',
|
||||
(json.dumps({"status": status, "actual_spend": 2, "baseline_spend": 2}), "request"),
|
||||
)
|
||||
elif damage == "wrong-body":
|
||||
_coverage_postgresql.execute(
|
||||
"CREATE OR REPLACE FUNCTION litellm_update_daily_autorouter_coverage() "
|
||||
"RETURNS TRIGGER LANGUAGE plpgsql AS 'BEGIN RETURN NEW; END;'"
|
||||
)
|
||||
else:
|
||||
_coverage_postgresql.execute(
|
||||
'DROP TRIGGER litellm_update_daily_autorouter_coverage ON "LiteLLM_AutoRouterBaselineObservation"'
|
||||
)
|
||||
if damage == "wrong-function":
|
||||
_coverage_postgresql.execute(
|
||||
"CREATE FUNCTION coverage_noop() RETURNS TRIGGER LANGUAGE plpgsql AS 'BEGIN RETURN NEW; END;'"
|
||||
)
|
||||
_coverage_postgresql.execute(
|
||||
"CREATE TRIGGER litellm_update_daily_autorouter_coverage BEFORE UPDATE OF publication "
|
||||
'ON "LiteLLM_AutoRouterBaselineObservation" FOR EACH ROW EXECUTE FUNCTION coverage_noop()'
|
||||
)
|
||||
elif damage == "wrong-event":
|
||||
_coverage_postgresql.execute(
|
||||
"CREATE TRIGGER litellm_update_daily_autorouter_coverage BEFORE UPDATE OF data "
|
||||
'ON "LiteLLM_AutoRouterBaselineObservation" FOR EACH ROW '
|
||||
"EXECUTE FUNCTION litellm_update_daily_autorouter_coverage()"
|
||||
)
|
||||
_coverage_postgresql.commit()
|
||||
|
||||
_coverage_install(_coverage_postgresql)
|
||||
|
||||
assert _coverage_publish(_coverage_postgresql) == ("recorded", 10, 30, 1, 2)
|
||||
|
||||
|
||||
@pytest.mark.parametrize("version", (None, 2))
|
||||
def test_coverage_installer_retains_compatible_unversioned_and_newer_functions(
|
||||
_coverage_postgresql: psycopg.Connection[tuple[object, ...]],
|
||||
version: int | None,
|
||||
) -> None:
|
||||
_coverage_tables(_coverage_postgresql)
|
||||
_coverage_install(_coverage_postgresql)
|
||||
before: Final = _coverage_postgresql.execute(
|
||||
"SELECT prosrc FROM pg_proc WHERE oid='litellm_update_daily_autorouter_coverage()'::regprocedure"
|
||||
).fetchone()
|
||||
assert before is not None and isinstance(before[0], str)
|
||||
body: Final = before[0] if version is None else "\n" + before[0]
|
||||
_coverage_postgresql.execute(
|
||||
sql.SQL(
|
||||
"CREATE OR REPLACE FUNCTION litellm_update_daily_autorouter_coverage() "
|
||||
"RETURNS TRIGGER LANGUAGE plpgsql AS {}"
|
||||
).format(sql.Literal(body))
|
||||
)
|
||||
_coverage_postgresql.execute(
|
||||
sql.SQL("COMMENT ON FUNCTION litellm_update_daily_autorouter_coverage() IS {}").format(
|
||||
sql.Literal(None if version is None else f"litellm:autorouter_daily_coverage:{version}")
|
||||
)
|
||||
)
|
||||
_coverage_postgresql.commit()
|
||||
|
||||
_coverage_repeat_during_traffic(_coverage_postgresql)
|
||||
|
||||
after: Final = _coverage_postgresql.execute(
|
||||
"SELECT prosrc,obj_description(oid,'pg_proc') FROM pg_proc "
|
||||
"WHERE oid='litellm_update_daily_autorouter_coverage()'::regprocedure"
|
||||
).fetchone()
|
||||
assert after == (body, None if version is None else f"litellm:autorouter_daily_coverage:{version}")
|
||||
assert _coverage_publish(_coverage_postgresql) == ("recorded", 10, 30, 1, 2)
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"incompatible,reason",
|
||||
(
|
||||
(
|
||||
'ALTER TABLE "LiteLLM_DailyUserSpend" ALTER COLUMN autorouter_estimated_requests TYPE INTEGER',
|
||||
"incompatible column definitions",
|
||||
),
|
||||
(
|
||||
'ALTER TABLE "LiteLLM_DailyUserSpend" ALTER COLUMN autorouter_estimated_requests SET DEFAULT 99',
|
||||
"incompatible column definitions",
|
||||
),
|
||||
(
|
||||
"COMMENT ON FUNCTION litellm_update_daily_autorouter_coverage() IS 'litellm:autorouter_daily_coverage:2';"
|
||||
"ALTER FUNCTION litellm_update_daily_autorouter_coverage() IMMUTABLE",
|
||||
"Newer daily auto-router coverage is incompatible",
|
||||
),
|
||||
(
|
||||
"COMMENT ON FUNCTION litellm_update_daily_autorouter_coverage() IS 'litellm:autorouter_daily_coverage:2';"
|
||||
'ALTER TABLE "LiteLLM_DailyUserSpend" DROP COLUMN autorouter_estimated_actual_spend',
|
||||
"Newer daily auto-router coverage is incompatible",
|
||||
),
|
||||
pytest.param(
|
||||
"COMMENT ON FUNCTION litellm_update_daily_autorouter_coverage() IS 'litellm:autorouter_daily_coverage:2';"
|
||||
'DROP TRIGGER litellm_update_daily_autorouter_coverage ON "LiteLLM_AutoRouterBaselineObservation";'
|
||||
"CREATE TRIGGER litellm_update_daily_autorouter_coverage BEFORE UPDATE OF data "
|
||||
'ON "LiteLLM_AutoRouterBaselineObservation" FOR EACH ROW '
|
||||
"EXECUTE FUNCTION litellm_update_daily_autorouter_coverage()",
|
||||
"Newer daily auto-router coverage is incompatible",
|
||||
id="newer-trigger-event",
|
||||
),
|
||||
pytest.param(
|
||||
"COMMENT ON FUNCTION litellm_update_daily_autorouter_coverage() IS 'litellm:autorouter_daily_coverage:2';"
|
||||
'ALTER TABLE "LiteLLM_AutoRouterBaselineObservation" '
|
||||
"DISABLE TRIGGER litellm_update_daily_autorouter_coverage",
|
||||
"Newer daily auto-router coverage is incompatible",
|
||||
id="newer-trigger-disabled",
|
||||
),
|
||||
pytest.param(
|
||||
"COMMENT ON FUNCTION litellm_update_daily_autorouter_coverage() IS 'litellm:autorouter_daily_coverage:2';"
|
||||
'DROP TRIGGER litellm_update_daily_autorouter_coverage ON "LiteLLM_AutoRouterBaselineObservation"',
|
||||
"Newer daily auto-router coverage is incompatible",
|
||||
id="newer-trigger-missing",
|
||||
),
|
||||
),
|
||||
)
|
||||
def test_coverage_installer_rejects_incompatible_state_without_replacing_it(
|
||||
_coverage_postgresql: psycopg.Connection[tuple[object, ...]],
|
||||
incompatible: str,
|
||||
reason: str,
|
||||
) -> None:
|
||||
_coverage_tables(_coverage_postgresql)
|
||||
_coverage_install(_coverage_postgresql)
|
||||
_coverage_postgresql.execute(incompatible)
|
||||
state_query: Final = (
|
||||
"SELECT p.prosrc,p.provolatile,obj_description(p.oid,'pg_proc'),t.oid,pg_get_triggerdef(t.oid),t.tgenabled "
|
||||
"FROM pg_proc p LEFT JOIN pg_trigger t "
|
||||
"ON t.tgrelid='\"LiteLLM_AutoRouterBaselineObservation\"'::regclass "
|
||||
"AND t.tgname='litellm_update_daily_autorouter_coverage' "
|
||||
"WHERE p.oid='litellm_update_daily_autorouter_coverage()'::regprocedure"
|
||||
)
|
||||
before: Final = _coverage_postgresql.execute(state_query).fetchone()
|
||||
assert before is not None
|
||||
_coverage_postgresql.commit()
|
||||
|
||||
with pytest.raises(psycopg.errors.RaiseException, match=reason):
|
||||
_coverage_install(_coverage_postgresql)
|
||||
|
||||
_coverage_postgresql.rollback()
|
||||
assert _coverage_postgresql.execute(state_query).fetchone() == before
|
||||
assert _coverage_postgresql.execute(
|
||||
'SELECT id,spend,autorouter_savings_spend FROM "LiteLLM_DailyUserSpend"'
|
||||
).fetchall() == [("recorded", 10, 30)]
|
||||
|
||||
|
||||
def test_coverage_installer_does_not_trust_or_modify_public_schema_homonyms(
|
||||
_coverage_postgresql: psycopg.Connection[tuple[object, ...]],
|
||||
) -> None:
|
||||
_coverage_tables(_coverage_postgresql)
|
||||
_coverage_install(_coverage_postgresql)
|
||||
_coverage_postgresql.execute('CREATE SCHEMA "tenant coverage"')
|
||||
_coverage_postgresql.execute('SET search_path TO "tenant coverage", public')
|
||||
_coverage_tables(_coverage_postgresql)
|
||||
|
||||
_coverage_install(_coverage_postgresql)
|
||||
|
||||
assert _coverage_publish(_coverage_postgresql) == ("recorded", 10, 30, 1, 2)
|
||||
assert _coverage_postgresql.execute(
|
||||
'SELECT spend,autorouter_savings_spend,autorouter_estimated_requests FROM public."LiteLLM_DailyUserSpend"'
|
||||
).fetchall() == [(10, 30, 0)]
|
||||
assert _coverage_postgresql.execute(
|
||||
'SELECT daily_costs_publication FROM public."LiteLLM_AutoRouterBaselineObservation"'
|
||||
).fetchall() == [(None,)]
|
||||
|
||||
|
||||
def _coverage_waiting(conn: psycopg.Connection[tuple[object, ...]], pid: int, kind: str) -> None:
|
||||
deadline: Final = time.monotonic() + 3
|
||||
while time.monotonic() < deadline:
|
||||
if conn.execute(
|
||||
"SELECT EXISTS(SELECT 1 FROM pg_locks WHERE pid=%s AND locktype=%s AND NOT granted)", (pid, kind)
|
||||
).fetchone() == (True,):
|
||||
return
|
||||
time.sleep(0.01)
|
||||
pytest.fail(f"Installer did not reach its expected {kind} wait")
|
||||
|
||||
|
||||
def test_waiting_coverage_installer_rechecks_after_peer_finishes_before_taking_table_locks(
|
||||
_coverage_postgresql: psycopg.Connection[tuple[object, ...]],
|
||||
) -> None:
|
||||
_coverage_tables(_coverage_postgresql)
|
||||
_coverage_postgresql.autocommit = True
|
||||
info: Final = _coverage_postgresql.info
|
||||
with (
|
||||
psycopg.connect(info.dsn, password=info.password) as original_reader,
|
||||
psycopg.connect(info.dsn, password=info.password) as next_reader,
|
||||
psycopg.connect(info.dsn, password=info.password, autocommit=True) as first,
|
||||
psycopg.connect(info.dsn, password=info.password, autocommit=True) as second,
|
||||
ThreadPoolExecutor(max_workers=3) as pool,
|
||||
):
|
||||
original_reader.execute('SELECT * FROM "LiteLLM_DailyUserSpend"')
|
||||
installing: Final = pool.submit(_coverage_install, first)
|
||||
try:
|
||||
_coverage_waiting(_coverage_postgresql, first.info.backend_pid, "relation")
|
||||
queued_read: Final = pool.submit(next_reader.execute, 'SELECT * FROM "LiteLLM_DailyUserSpend"')
|
||||
_coverage_waiting(_coverage_postgresql, next_reader.info.backend_pid, "relation")
|
||||
waiting: Final = pool.submit(_coverage_install, second)
|
||||
_coverage_waiting(_coverage_postgresql, second.info.backend_pid, "advisory")
|
||||
original_reader.rollback()
|
||||
installing.result(timeout=2)
|
||||
queued_read.result(timeout=2)
|
||||
waiting.result(timeout=2)
|
||||
finally:
|
||||
first.cancel()
|
||||
second.cancel()
|
||||
next_reader.cancel()
|
||||
original_reader.rollback()
|
||||
next_reader.rollback()
|
||||
|
||||
assert _coverage_publish(_coverage_postgresql) == ("recorded", 10, 30, 1, 2)
|
||||
assert connection.execute(
|
||||
"SELECT id,spend,autorouter_savings_spend,autorouter_estimated_requests,autorouter_estimated_actual_spend "
|
||||
'FROM "LiteLLM_DailyUserSpend"'
|
||||
).fetchall() == [("recorded", 10, 30, expected_count, expected_actual)]
|
||||
connection.execute('DELETE FROM "LiteLLM_AutoRouterBaselineObservation"')
|
||||
assert connection.execute(
|
||||
'SELECT autorouter_estimated_requests,autorouter_estimated_actual_spend FROM "LiteLLM_DailyUserSpend"'
|
||||
).fetchall() == [(1, 2)]
|
||||
|
|
|
|||
|
|
@ -360,47 +360,6 @@ class TestRender:
|
|||
|
||||
|
||||
class TestClaudeCodeMode:
|
||||
@pytest.mark.parametrize("field", ("baseline_spend", "savings_estimated_baseline_spend"))
|
||||
@pytest.mark.parametrize("invalid_baseline", ("1.5", True, float("nan"), float("inf")))
|
||||
def test_invalid_baselines_do_not_display_a_cost_comparison(
|
||||
self, config_dir: Path, field: str, invalid_baseline: object,
|
||||
) -> None:
|
||||
session: Final = statusline_script._session_from_payload({
|
||||
**RECORDED._asdict(), field: invalid_baseline,
|
||||
})
|
||||
assert session is None
|
||||
assert render("m", session, config_dir, use_color=False) == "Routed to: m"
|
||||
|
||||
@pytest.mark.parametrize("estimated_turns, saved_spend, delta", (
|
||||
(0, 30.0, "-75%"), (0, -2.0, "+25%"),
|
||||
(1, 30.0, "-75%"), (1, 0.0, "0%"), (1, -2.0, "+25%"),
|
||||
))
|
||||
def test_historical_and_current_savings_remain_combined_after_cache_refresh(
|
||||
self, tmp_path: Path, transcript: Path, config_dir: Path, estimated_turns: int,
|
||||
saved_spend: float, delta: str,
|
||||
) -> None:
|
||||
session: Final = statusline_script._session_from_payload(
|
||||
{
|
||||
**RECORDED._asdict(),
|
||||
"spend": 10.0,
|
||||
"baseline_spend": 10.0 + saved_spend,
|
||||
"savings_estimated_baseline_spend": 1.5 if estimated_turns else None,
|
||||
"turns": 3,
|
||||
"savings_estimated_turns": estimated_turns,
|
||||
"savings_estimated_actual_spend": 2.0 if estimated_turns else 0.0,
|
||||
}
|
||||
)
|
||||
assert session is not None
|
||||
|
||||
def fetch(credentials: Credentials, session_id: str) -> Fetched:
|
||||
return Fetched(session, True)
|
||||
|
||||
first: Final = _run(_payload(transcript), _env(tmp_path, config_dir), fetch)
|
||||
assert first == _run(_payload(transcript), _env(tmp_path, config_dir), fetch)
|
||||
assert first.splitlines()[0] == f"Routed to: claude-sonnet-5 {delta} vs Claude Opus 5"
|
||||
assert "$10.00" in first and f"${10.0 + saved_spend:.2f}" in first
|
||||
assert "$1.50" not in first and "Savings unavailable" not in first
|
||||
|
||||
@pytest.mark.parametrize("estimated_turns", (0, 1))
|
||||
def test_current_estimates_keep_the_routed_model_and_compare_only_covered_turns(
|
||||
self, tmp_path: Path, transcript: Path, config_dir: Path, estimated_turns: int
|
||||
|
|
|
|||
|
|
@ -103,12 +103,6 @@ if not earlier_calls and os.environ.get("FAKE_PRISMA_HANG_FIRST"):
|
|||
if sys.argv[1:3] == ["db", "execute"]:
|
||||
if os.environ.get("FAKE_PRISMA_FAIL_DB_EXECUTE"):
|
||||
sys.exit(1)
|
||||
if not any(json.loads(call)[:2] == ["db", "execute"] for call in earlier_calls):
|
||||
first_failure = os.environ.get("FAKE_PRISMA_FAIL_FIRST_DB_EXECUTE")
|
||||
if first_failure == "timeout":
|
||||
time.sleep(600)
|
||||
if first_failure:
|
||||
sys.exit(1)
|
||||
sys.exit(0)
|
||||
"""
|
||||
|
||||
|
|
@ -152,7 +146,6 @@ def fake_prisma_cli(tmp_path: Path, monkeypatch: pytest.MonkeyPatch) -> Generato
|
|||
monkeypatch.setenv("LITELLM_PRISMA_COMMAND_TIMEOUT", "1")
|
||||
monkeypatch.delenv("FAKE_PRISMA_HANG_FIRST", raising=False)
|
||||
monkeypatch.delenv("FAKE_PRISMA_FAIL_DB_EXECUTE", raising=False)
|
||||
monkeypatch.delenv("FAKE_PRISMA_FAIL_FIRST_DB_EXECUTE", raising=False)
|
||||
yield cli
|
||||
if cli.grandchild_pidfile.exists():
|
||||
try:
|
||||
|
|
|
|||
|
|
@ -606,22 +606,3 @@ async def test_optional_metric_missing_from_an_older_payload_still_aggregates(
|
|||
assert updates[0][test_key]["autorouter_savings_spend"] == pytest.approx(0.25)
|
||||
assert updates[0][test_key]["total_response_time_ms"] == 900
|
||||
assert updates[0][test_key]["timed_requests"] == 1
|
||||
|
||||
|
||||
@pytest.mark.parametrize("reverse", (False, True))
|
||||
def test_old_daily_payloads_do_not_gain_router_accounting_coverage(reverse: bool) -> None:
|
||||
old = {"spend": 1.0, "prompt_tokens": 2, "completion_tokens": 1,
|
||||
"api_requests": 1, "successful_requests": 1, "failed_requests": 0}
|
||||
current = {**old, "autorouter_accounted_requests": 1, "autorouter_requests": 1,
|
||||
"autorouter_llm_spend": 1.0, "autorouter_classifier_cost": 0.2,
|
||||
"autorouter_classifier_cost_recorded_requests": 1,
|
||||
"autorouter_estimated_requests": 1, "autorouter_estimated_actual_spend": 1.2}
|
||||
updates = [{"key": current}, {"key": old}] if reverse else [{"key": old}, {"key": current}]
|
||||
merged = DailySpendUpdateQueue.get_aggregated_daily_spend_update_transactions(updates)["key"]
|
||||
assert merged["api_requests"] == 2
|
||||
assert merged["autorouter_accounted_requests"] == merged["autorouter_requests"] == 1
|
||||
assert merged["autorouter_llm_spend"] == 1.0
|
||||
assert merged["autorouter_classifier_cost"] == 0.2
|
||||
assert merged["autorouter_classifier_cost_recorded_requests"] == merged["autorouter_estimated_requests"] == 1
|
||||
assert merged["autorouter_estimated_actual_spend"] == 1.2
|
||||
assert "autorouter_accounted_requests" not in old
|
||||
|
|
|
|||
|
|
@ -1,286 +0,0 @@
|
|||
import json
|
||||
import re
|
||||
from collections.abc import Mapping
|
||||
from typing import Final
|
||||
|
||||
import psycopg
|
||||
import pytest
|
||||
from psycopg.rows import dict_row
|
||||
from psycopg.types.json import Jsonb
|
||||
from pytest_postgresql import factories
|
||||
|
||||
from litellm.proxy.db.autorouter_daily_spend import AutoRouterDailyCosts
|
||||
from litellm.proxy.db.autorouter_historical_spend import AUTOROUTER_HISTORICAL_COSTS_SQL
|
||||
from litellm.proxy.route_llm_request import ROUTE_ENDPOINT_MAPPING
|
||||
from litellm.proxy.spend_tracking.savings import classifier_cost_from_decision, known_autorouter_savings
|
||||
|
||||
_historical_postgresql_proc: Final = factories.postgresql_proc()
|
||||
_historical_postgresql: Final = factories.postgresql("_historical_postgresql_proc")
|
||||
_DDL: Final = """
|
||||
CREATE TABLE "LiteLLM_DailyUserSpend" (
|
||||
date TEXT, user_id TEXT, api_key TEXT, model TEXT, custom_llm_provider TEXT,
|
||||
mcp_namespaced_tool_name TEXT, endpoint TEXT,
|
||||
api_requests BIGINT DEFAULT 0, successful_requests BIGINT DEFAULT 0, failed_requests BIGINT DEFAULT 0,
|
||||
prompt_tokens BIGINT DEFAULT 0, completion_tokens BIGINT DEFAULT 0, spend FLOAT8 DEFAULT 0,
|
||||
autorouter_requests BIGINT DEFAULT 0, autorouter_llm_spend FLOAT8 DEFAULT 0,
|
||||
autorouter_classifier_cost FLOAT8 DEFAULT 0, autorouter_classifier_cost_recorded_requests BIGINT DEFAULT 0,
|
||||
autorouter_estimated_requests BIGINT DEFAULT 0, autorouter_estimated_actual_spend FLOAT8 DEFAULT 0,
|
||||
autorouter_savings_spend FLOAT8 DEFAULT 0, autorouter_accounted_requests BIGINT DEFAULT 0
|
||||
);
|
||||
CREATE TABLE "LiteLLM_SpendLogs" (
|
||||
request_id TEXT PRIMARY KEY, "startTime" TIMESTAMP, "user" TEXT, api_key TEXT,
|
||||
model TEXT, custom_llm_provider TEXT, mcp_namespaced_tool_name TEXT, call_type TEXT,
|
||||
spend FLOAT8, prompt_tokens BIGINT, completion_tokens BIGINT, status TEXT, metadata JSONB
|
||||
)
|
||||
"""
|
||||
|
||||
|
||||
def _seed(conn: psycopg.Connection) -> None:
|
||||
conn.execute(_DDL)
|
||||
conn.execute(
|
||||
"""INSERT INTO "LiteLLM_DailyUserSpend"
|
||||
(date, user_id, api_key, model, custom_llm_provider, mcp_namespaced_tool_name, endpoint,
|
||||
api_requests, successful_requests, prompt_tokens, completion_tokens, spend, autorouter_savings_spend)
|
||||
VALUES ('2026-09-22', 'owner', 'key', 'model', NULL, '', '/chat/completions', 2, 2, 55, 10, 5.1, 4)"""
|
||||
)
|
||||
decision: Final = {"router_model_name": "router", "classifier_cost": 0.1}
|
||||
rows: Final = (
|
||||
("routed", 2.0, 20, 10, {"routing_decision": decision, "autorouter_savings": 4}),
|
||||
("ordinary", 3.0, 30, 0, {}),
|
||||
("classifier", 0.1, 5, 0, {"internal_call_origin": "autorouter_classifier", "routing_decision": decision}),
|
||||
)
|
||||
for request_id, spend, prompt, completion, metadata in rows:
|
||||
conn.execute(
|
||||
"""INSERT INTO "LiteLLM_SpendLogs"
|
||||
(request_id, "startTime", "user", api_key, model, custom_llm_provider, mcp_namespaced_tool_name,
|
||||
call_type, spend, prompt_tokens, completion_tokens, status, metadata)
|
||||
VALUES (%s, '2026-09-22 12:00:00', 'owner', 'key', 'model', '', NULL,
|
||||
'acompletion', %s, %s, %s, 'success', %s)""",
|
||||
(request_id, spend, prompt, completion, Jsonb(metadata)),
|
||||
)
|
||||
|
||||
|
||||
def _read(conn: psycopg.Connection, *, key: str | None = None, user: str | None = None) -> AutoRouterDailyCosts:
|
||||
query: Final = re.sub(r"\$(\d+)", r"%(p\1)s", AUTOROUTER_HISTORICAL_COSTS_SQL)
|
||||
with conn.cursor(row_factory=dict_row) as cursor:
|
||||
cursor.execute(
|
||||
query, # pyright: ignore[reportArgumentType] # fixed production SQL uses Prisma parameter syntax
|
||||
{"p1": "2026-09-22", "p2": "2026-09-22", "p3": key, "p4": user, "p5": json.dumps(ROUTE_ENDPOINT_MAPPING)},
|
||||
)
|
||||
return AutoRouterDailyCosts.model_validate(cursor.fetchone())
|
||||
|
||||
|
||||
def test_historical_costs_include_no_session_turns_and_classifier_charge_once(
|
||||
_historical_postgresql: psycopg.Connection,
|
||||
) -> None:
|
||||
_seed(_historical_postgresql)
|
||||
|
||||
costs: Final = _read(_historical_postgresql, key="key", user="owner")
|
||||
|
||||
assert costs == AutoRouterDailyCosts(
|
||||
requests=1,
|
||||
llm_spend=2,
|
||||
classifier_cost=0.1,
|
||||
classifier_requests=1,
|
||||
estimated_requests=1,
|
||||
estimated_actual_spend=2.1,
|
||||
saved_spend=4,
|
||||
)
|
||||
assert costs.recorded_spend == pytest.approx(2.1)
|
||||
assert costs.baseline_spend(4) == pytest.approx(6.1)
|
||||
assert _read(_historical_postgresql, key="other").requests == 0
|
||||
assert _read(_historical_postgresql, user="other").requests == 0
|
||||
|
||||
|
||||
@pytest.mark.parametrize("missing_request", ("ordinary", "classifier", "routed"))
|
||||
def test_missing_logs_cannot_certify_costs_even_when_saved_savings_match(
|
||||
_historical_postgresql: psycopg.Connection, missing_request: str
|
||||
) -> None:
|
||||
_seed(_historical_postgresql)
|
||||
_historical_postgresql.execute('DELETE FROM "LiteLLM_SpendLogs" WHERE request_id = %s', (missing_request,))
|
||||
|
||||
costs: Final = _read(_historical_postgresql)
|
||||
|
||||
assert not costs.complete
|
||||
assert costs.saved_spend == 4
|
||||
assert costs.recorded_spend is None
|
||||
assert costs.baseline_spend(4) is None
|
||||
|
||||
|
||||
def test_missing_zero_cost_request_invalidates_matching_spend_and_token_totals(
|
||||
_historical_postgresql: psycopg.Connection,
|
||||
) -> None:
|
||||
_seed(_historical_postgresql)
|
||||
_historical_postgresql.execute(
|
||||
"""UPDATE "LiteLLM_DailyUserSpend" SET spend = 2.1, prompt_tokens = 25;
|
||||
DELETE FROM "LiteLLM_SpendLogs" WHERE request_id = 'ordinary'"""
|
||||
)
|
||||
|
||||
costs: Final = _read(_historical_postgresql)
|
||||
|
||||
assert not costs.complete
|
||||
assert costs.recorded_spend is None
|
||||
assert costs.saved_spend == 4
|
||||
|
||||
|
||||
@pytest.mark.parametrize("changed", ("model", "custom_llm_provider", "mcp_namespaced_tool_name", "call_type"))
|
||||
def test_equal_totals_in_different_daily_identity_do_not_certify_recovery(
|
||||
_historical_postgresql: psycopg.Connection, changed: str
|
||||
) -> None:
|
||||
_seed(_historical_postgresql)
|
||||
_historical_postgresql.execute(
|
||||
psycopg.sql.SQL('UPDATE "LiteLLM_SpendLogs" SET {} = %s WHERE request_id = %s').format(
|
||||
psycopg.sql.Identifier(changed)
|
||||
),
|
||||
("different", "ordinary"),
|
||||
)
|
||||
|
||||
costs: Final = _read(_historical_postgresql)
|
||||
|
||||
assert not costs.complete
|
||||
assert costs.recorded_spend is None
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"metadata,saved,estimated,comparison",
|
||||
(
|
||||
({"autorouter_savings": 0}, 0, 1, True),
|
||||
({"autorouter_savings": -1}, -1, 1, True),
|
||||
({"autorouter_savings": 0, "autorouter_savings_estimate": {"version": 3, "status": "unknown"}}, 0, 0, True),
|
||||
({"autorouter_savings": 0, "autorouter_savings_estimate": {"version": 3, "status": "estimated"}}, 0, 1, True),
|
||||
({}, 4, 0, False),
|
||||
({"autorouter_savings": True}, 4, 0, False),
|
||||
({"autorouter_savings": "4"}, 4, 0, False),
|
||||
({"autorouter_savings": 4, "autorouter_savings_estimate": {}}, 4, 0, False),
|
||||
(
|
||||
{"autorouter_savings": 4, "autorouter_savings_estimate": {"version": 1.0, "status": "estimated"}},
|
||||
4,
|
||||
0,
|
||||
False,
|
||||
),
|
||||
(
|
||||
{"autorouter_savings": 4, "autorouter_savings_estimate": {"version": True, "status": "estimated"}},
|
||||
4,
|
||||
0,
|
||||
False,
|
||||
),
|
||||
({"autorouter_savings": 4, "autorouter_savings_estimate": {"version": 99, "status": "estimated"}}, 4, 0, False),
|
||||
),
|
||||
)
|
||||
def test_recorded_estimate_provenance_gates_baseline_without_hiding_actual_cost(
|
||||
_historical_postgresql: psycopg.Connection,
|
||||
metadata: Mapping[str, object],
|
||||
saved: float,
|
||||
estimated: int,
|
||||
comparison: bool,
|
||||
) -> None:
|
||||
_seed(_historical_postgresql)
|
||||
_historical_postgresql.execute(
|
||||
'UPDATE "LiteLLM_SpendLogs" SET metadata = %s WHERE request_id = %s',
|
||||
(Jsonb({"routing_decision": {"router_model_name": "router", "classifier_cost": 0.1}, **metadata}), "routed"),
|
||||
)
|
||||
_historical_postgresql.execute('UPDATE "LiteLLM_DailyUserSpend" SET autorouter_savings_spend = %s', (saved,))
|
||||
|
||||
costs: Final = _read(_historical_postgresql)
|
||||
estimate: Final = metadata.get("autorouter_savings_estimate")
|
||||
recorded: Final = known_autorouter_savings(
|
||||
model=None,
|
||||
custom_llm_provider=None,
|
||||
routing_decision=None,
|
||||
usage_object=None,
|
||||
recorded_autorouter_savings=metadata.get("autorouter_savings"),
|
||||
recorded_autorouter_savings_estimate=estimate if isinstance(estimate, Mapping) else None,
|
||||
)
|
||||
|
||||
assert costs.complete and costs.recorded_spend == pytest.approx(2.1)
|
||||
assert costs.estimated_requests == estimated == int(recorded is not None)
|
||||
assert costs.comparison_complete is comparison
|
||||
assert costs.saved_spend == saved
|
||||
assert costs.baseline_spend(saved) == (pytest.approx(2.1 + saved) if estimated and comparison else None)
|
||||
|
||||
|
||||
def test_unknown_classifier_preserves_recorded_subtotal_with_partial_coverage(
|
||||
_historical_postgresql: psycopg.Connection,
|
||||
) -> None:
|
||||
_seed(_historical_postgresql)
|
||||
_historical_postgresql.execute(
|
||||
"""UPDATE "LiteLLM_SpendLogs" SET metadata = metadata #- '{routing_decision,classifier_cost}'
|
||||
WHERE request_id = 'routed'"""
|
||||
)
|
||||
|
||||
costs: Final = _read(_historical_postgresql)
|
||||
|
||||
assert costs.complete and costs.recorded_llm_spend == 2
|
||||
assert costs.recorded_classifier_cost is None
|
||||
assert costs.recorded_spend == 2
|
||||
assert costs.coverage == "partial"
|
||||
assert costs.baseline_spend(4) is None
|
||||
|
||||
|
||||
@pytest.mark.parametrize("classifier,expected", ((-0.1, -0.1), (0, 0), (True, None), ("0.1", None)))
|
||||
def test_classifier_numeric_contract_matches_daily_owner(
|
||||
_historical_postgresql: psycopg.Connection, classifier: object, expected: float | None
|
||||
) -> None:
|
||||
_seed(_historical_postgresql)
|
||||
_historical_postgresql.execute(
|
||||
"""UPDATE "LiteLLM_SpendLogs" SET metadata = jsonb_set(metadata, '{routing_decision,classifier_cost}', %s)
|
||||
WHERE request_id = 'routed'""",
|
||||
(Jsonb(classifier),),
|
||||
)
|
||||
|
||||
costs: Final = _read(_historical_postgresql)
|
||||
|
||||
assert costs.recorded_classifier_cost == expected == classifier_cost_from_decision({"classifier_cost": classifier})
|
||||
assert costs.classifier_requests == int(expected is not None)
|
||||
assert costs.recorded_spend == pytest.approx(2 + (expected or 0))
|
||||
|
||||
|
||||
def test_unknown_routed_request_is_counted_in_actual_but_excluded_from_baseline(
|
||||
_historical_postgresql: psycopg.Connection,
|
||||
) -> None:
|
||||
_seed(_historical_postgresql)
|
||||
_historical_postgresql.execute(
|
||||
'UPDATE "LiteLLM_SpendLogs" SET metadata = %s WHERE request_id = %s',
|
||||
(
|
||||
Jsonb(
|
||||
{
|
||||
"routing_decision": {"router_model_name": "router", "classifier_cost": 0},
|
||||
"autorouter_savings_estimate": {"version": 3, "status": "unknown"},
|
||||
}
|
||||
),
|
||||
"ordinary",
|
||||
),
|
||||
)
|
||||
|
||||
costs: Final = _read(_historical_postgresql)
|
||||
|
||||
assert costs.complete and costs.comparison_complete
|
||||
assert costs.requests == 2 and costs.estimated_requests == 1
|
||||
assert costs.recorded_spend == pytest.approx(5.1)
|
||||
assert costs.baseline_spend(4) == pytest.approx(6.1)
|
||||
|
||||
|
||||
def test_historical_logs_replace_partial_bucket_and_retain_durable_other_days(
|
||||
_historical_postgresql: psycopg.Connection,
|
||||
) -> None:
|
||||
_seed(_historical_postgresql)
|
||||
_historical_postgresql.execute(
|
||||
"""UPDATE "LiteLLM_DailyUserSpend" SET autorouter_accounted_requests = 1,
|
||||
autorouter_requests = 1, autorouter_llm_spend = 1, autorouter_classifier_cost = 0.05,
|
||||
autorouter_classifier_cost_recorded_requests = 1, autorouter_estimated_requests = 1,
|
||||
autorouter_estimated_actual_spend = 1.05;
|
||||
INSERT INTO "LiteLLM_DailyUserSpend"
|
||||
(date, user_id, api_key, model, endpoint, api_requests, successful_requests,
|
||||
autorouter_accounted_requests, autorouter_requests, autorouter_llm_spend,
|
||||
autorouter_classifier_cost, autorouter_classifier_cost_recorded_requests,
|
||||
autorouter_estimated_requests, autorouter_estimated_actual_spend, autorouter_savings_spend)
|
||||
VALUES ('2026-09-22', 'owner', 'key', 'durable-model', '/chat/completions', 1, 1, 1, 1, 3, 0.2, 1, 1, 3.2, 2),
|
||||
('2026-09-21', 'owner', 'key', 'durable-model', '/chat/completions', 1, 1, 1, 1, 90, 1, 1, 1, 91, 9)"""
|
||||
)
|
||||
|
||||
costs: Final = _read(_historical_postgresql)
|
||||
|
||||
assert costs.complete and costs.requests == 2
|
||||
assert costs.recorded_spend == pytest.approx(5.3)
|
||||
assert costs.saved_spend == 6
|
||||
assert costs.baseline_spend(6) == pytest.approx(11.3)
|
||||
|
|
@ -3624,63 +3624,6 @@ async def test_daily_router_costs_account_for_every_external_request_without_a_s
|
|||
assert transaction["autorouter_estimated_actual_spend"] == (0.625 if estimated else 0.0)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize("status", ("success", None, "failure"))
|
||||
async def test_daily_router_costs_use_status_normalized_from_request_metadata(status: str | None) -> None:
|
||||
from litellm.proxy.spend_tracking.spend_tracking_utils import get_logging_payload
|
||||
from litellm.proxy.utils import PrismaClient
|
||||
|
||||
started_at: Final = datetime(2026, 9, 23, tzinfo=timezone.utc)
|
||||
metadata: Final = {
|
||||
"user_api_key": "test-hash", "user_api_key_user_id": "u", "model_group": "auto",
|
||||
"routing_decision": {"router_model_name": "auto", "classifier_cost": 0.125},
|
||||
**({"status": status} if status is not None else {}),
|
||||
}
|
||||
payload: Final = get_logging_payload(
|
||||
kwargs={
|
||||
"model": "synthetic-target", "call_type": "acompletion", "response_cost": 0.5,
|
||||
"litellm_params": {"metadata": metadata, "custom_llm_provider": "openai"},
|
||||
},
|
||||
response_obj={}, start_time=started_at, end_time=started_at,
|
||||
)
|
||||
prisma: Final = PrismaClient.__new__(PrismaClient)
|
||||
transaction: Final = await DBSpendUpdateWriter()._common_add_spend_log_transaction_to_daily_transaction(
|
||||
payload, prisma, "user",
|
||||
)
|
||||
|
||||
assert transaction is not None
|
||||
assert payload["status"] == ("failure" if status == "failure" else "success")
|
||||
assert transaction["autorouter_requests"] == int(status != "failure")
|
||||
assert transaction["autorouter_llm_spend"] == (0.0 if status == "failure" else 0.5)
|
||||
assert transaction["autorouter_classifier_cost"] == (0.0 if status == "failure" else 0.125)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize("classifier_cost, recorded", [(None, 0), (True, 0), (float("inf"), 0), (0.0, 1)])
|
||||
async def test_daily_router_classifier_zero_is_known_only_when_recorded(
|
||||
classifier_cost: object, recorded: int,
|
||||
) -> None:
|
||||
prisma: Final = MagicMock()
|
||||
prisma.get_request_status.return_value = "success"
|
||||
transaction: Final = await DBSpendUpdateWriter()._common_add_spend_log_transaction_to_daily_transaction(
|
||||
payload={
|
||||
"user": "u", "startTime": "2026-09-23T00:00:00", "api_key": "hash", "model": "m",
|
||||
"model_group": "auto", "custom_llm_provider": "openai", "call_type": "acompletion",
|
||||
"prompt_tokens": 1, "completion_tokens": 1, "spend": 1.0, "status": "success",
|
||||
"metadata": json.dumps({
|
||||
"routing_decision": {"router_model_name": "auto", "classifier_cost": classifier_cost},
|
||||
"autorouter_savings": 0.0,
|
||||
}),
|
||||
}, prisma_client=prisma, type="user",
|
||||
)
|
||||
assert transaction is not None
|
||||
assert transaction["autorouter_requests"] == 1
|
||||
assert transaction["autorouter_classifier_cost_recorded_requests"] == recorded
|
||||
assert transaction["autorouter_classifier_cost"] == 0.0
|
||||
assert transaction["autorouter_estimated_requests"] == 1
|
||||
assert transaction["autorouter_estimated_actual_spend"] == 1.0
|
||||
|
||||
|
||||
def _response_time_payload(request_duration_ms: object, metadata: dict | None = None) -> dict:
|
||||
return {
|
||||
"request_id": "req-timed-1",
|
||||
|
|
|
|||
|
|
@ -276,26 +276,6 @@ def test_db_push_cannot_succeed_without_daily_coverage_owner(
|
|||
]
|
||||
|
||||
|
||||
@pytest.mark.parametrize("use_migrate", [False, True])
|
||||
@pytest.mark.parametrize("failure", ["exit", "timeout"])
|
||||
def test_database_setup_retries_transient_daily_coverage_install(
|
||||
fake_prisma_cli: "FakePrismaCli",
|
||||
unset_database_url: None,
|
||||
monkeypatch: pytest.MonkeyPatch,
|
||||
use_migrate: bool,
|
||||
failure: str,
|
||||
) -> None:
|
||||
monkeypatch.setenv("FAKE_PRISMA_FAIL_FIRST_DB_EXECUTE", failure)
|
||||
|
||||
with patch("litellm_proxy_extras.utils.time.sleep"):
|
||||
assert PrismaManager.setup_database(use_migrate=use_migrate) is True
|
||||
|
||||
assert [call[:2] for call in fake_prisma_cli.calls if call != ["--version"]] == [
|
||||
["migrate", "deploy"] if use_migrate else ["db", "push"],
|
||||
["db", "execute"], ["db", "execute"],
|
||||
]
|
||||
|
||||
|
||||
def _entra_jwt(expires_in_seconds: int) -> str:
|
||||
"""A JWT shaped like a real Entra access token, expiring ``expires_in_seconds`` from now."""
|
||||
import base64
|
||||
|
|
|
|||
|
|
@ -2,7 +2,6 @@
|
|||
Unit tests for auto router management endpoints
|
||||
"""
|
||||
|
||||
import json
|
||||
from collections.abc import Mapping, Sequence
|
||||
from functools import partial
|
||||
from pathlib import Path
|
||||
|
|
@ -638,7 +637,6 @@ class TestAutoRouterBenchmarks:
|
|||
rows: Sequence[Mapping[str, object]],
|
||||
model_list: Sequence[object],
|
||||
api_key: str | None = None,
|
||||
costs: AutoRouterDailyCosts | None = None,
|
||||
) -> AutoRouterBenchmarksResponse:
|
||||
from litellm.proxy import proxy_server
|
||||
from litellm.proxy.management_endpoints.auto_router_endpoints import get_auto_router_benchmarks
|
||||
|
|
@ -646,7 +644,7 @@ class TestAutoRouterBenchmarks:
|
|||
class _DB:
|
||||
async def query_raw(self, sql: str, *params: object):
|
||||
if sql in (AUTOROUTER_DAILY_COSTS_SQL, AUTOROUTER_HISTORICAL_COSTS_SQL):
|
||||
return [(costs or AutoRouterDailyCosts(complete=False, saved_spend=30.0)).model_dump()]
|
||||
return [AutoRouterDailyCosts(complete=False, saved_spend=30.0).model_dump()]
|
||||
return rows
|
||||
|
||||
monkeypatch.setattr(proxy_server, "prisma_client", type("P", (), {"db": _DB()})())
|
||||
|
|
@ -751,178 +749,51 @@ class TestAutoRouterBenchmarks:
|
|||
assert group.saved_per_session == saved_spend / row.sessions
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_selected_router_compares_only_estimated_requests(self, monkeypatch: pytest.MonkeyPatch) -> None:
|
||||
row: Final = self.ROW.model_copy(update={
|
||||
"turns": 2, "spend": 100.0, "saved_spend": 1.0,
|
||||
"savings_estimated_turns": 1, "savings_estimated_actual_spend": 1.0,
|
||||
"savings_estimated_saved_spend": 1.0,
|
||||
})
|
||||
response: Final = await self._benchmarks(monkeypatch, rows=[row.model_dump()], model_list=[])
|
||||
group: Final = response.groups[0]
|
||||
assert (group.spend, group.saved_spend, group.baseline_spend, group.saved_pct) == (100.0, 1.0, 2.0, 50.0)
|
||||
assert (group.savings_estimated_turns, group.cost_requests, group.savings_estimated_actual_spend) == (1, 2, 1.0)
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_tracked_session_groups_keep_historical_and_new_estimates_once(
|
||||
self, monkeypatch: pytest.MonkeyPatch
|
||||
) -> None:
|
||||
historical: Final = self.ROW.model_copy(
|
||||
update={
|
||||
"savings_estimated_turns": 0,
|
||||
"savings_estimated_actual_spend": 0.0,
|
||||
"savings_estimated_saved_spend": 0.0,
|
||||
}
|
||||
)
|
||||
current: Final = self.ROW.model_copy(
|
||||
update={
|
||||
"router_name": "current-auto",
|
||||
"sessions": 1,
|
||||
"turns": 4,
|
||||
"spend": 2.0,
|
||||
"saved_spend": -0.5,
|
||||
"savings_estimated_turns": 4,
|
||||
"savings_estimated_actual_spend": 2.0,
|
||||
"savings_estimated_saved_spend": -0.5,
|
||||
}
|
||||
)
|
||||
response: Final = await self._benchmarks(
|
||||
monkeypatch, rows=[historical.model_dump(), current.model_dump()], model_list=[]
|
||||
)
|
||||
assert [group.saved_spend for group in response.groups] == [30.0, -0.5]
|
||||
assert response.groups[1].savings_estimated_turns == 4
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize("api_key,user_id", [(None, None), ("key-hash", None), (None, "selected-user"),
|
||||
("key-hash", "selected-user")])
|
||||
@pytest.mark.parametrize(
|
||||
"daily_saved,has_daily_rows,has_sessions",
|
||||
[(6.5, True, True), (-0.5, True, True), (0.0, True, True), (30.0, True, True),
|
||||
(2.0, True, False), (0.0, False, True), (0.0, False, False)],
|
||||
ids=["historical-and-new-plus-no-session", "signed-loss", "recorded-zero", "matching-is-not-proof",
|
||||
"no-session-history", "out-of-window-session-only", "empty"],
|
||||
)
|
||||
async def test_total_uses_overall_request_date_savings_with_identical_filters(
|
||||
self, monkeypatch: pytest.MonkeyPatch, api_key: str | None, user_id: str | None,
|
||||
daily_saved: float, has_daily_rows: bool, has_sessions: bool,
|
||||
@pytest.mark.parametrize("complete, recovered_saved, estimated, expected_spend, baseline, pct", [
|
||||
(True, 6.5, 2, 5.0, 9.5, 68.4),
|
||||
(True, 6.5, 0, 5.0, None, None),
|
||||
(False, 6.5, 2, 5.0, 9.5, 68.4),
|
||||
(False, 7.5, 2, 5.0, None, None),
|
||||
(False, None, 2, 2.25, None, None),
|
||||
])
|
||||
async def test_daily_savings_survive_recovery_and_only_matching_estimates_are_compared(
|
||||
self, monkeypatch: pytest.MonkeyPatch, complete: bool, recovered_saved: float | None,
|
||||
estimated: int, expected_spend: float, baseline: float | None, pct: float | None,
|
||||
) -> None:
|
||||
from litellm.proxy.management_endpoints import common_daily_activity
|
||||
from litellm.proxy.management_endpoints.auto_router_endpoints import get_auto_router_benchmarks
|
||||
|
||||
daily: Final = AsyncMock(side_effect=AssertionError("Full usage breakdowns must not be loaded"))
|
||||
query: Final = AsyncMock(side_effect=[
|
||||
[self.ROW.model_dump()] if has_sessions else [],
|
||||
[AutoRouterDailyCosts(complete=not has_daily_rows, saved_spend=daily_saved).model_dump()],
|
||||
[AutoRouterDailyCosts(complete=not has_daily_rows, saved_spend=daily_saved).model_dump()],
|
||||
])
|
||||
prisma: Final = SimpleNamespace(db=SimpleNamespace(query_raw=query))
|
||||
monkeypatch.setattr(common_daily_activity, "get_daily_activity_aggregated", daily)
|
||||
monkeypatch.setattr(proxy_server, "prisma_client", prisma)
|
||||
response: Final = await get_auto_router_benchmarks(
|
||||
user_api_key_dict=ADMIN, start_date="2026-07-01", end_date="2026-08-01",
|
||||
api_key=api_key, user_id=user_id,
|
||||
)
|
||||
|
||||
daily.assert_not_awaited()
|
||||
assert query.call_args_list[0].args[1:] == ("2026-07-01T00:00:00", "2026-08-02T00:00:00", api_key, user_id)
|
||||
assert query.call_args_list[1].args[1:] == ("2026-07-01", "2026-08-01", api_key, user_id)
|
||||
assert response.totals.saved_spend == daily_saved
|
||||
assert response.totals.sessions == (4 if has_sessions else 0)
|
||||
assert response.totals.saved_per_session == (7.5 if has_sessions else 0.0)
|
||||
assert [group.saved_spend for group in response.groups] == ([30.0] if has_sessions else [])
|
||||
assert (
|
||||
response.totals.spend, response.totals.classifier_cost, response.totals.baseline_spend,
|
||||
response.totals.saved_pct,
|
||||
) == ((None,) * 4 if has_daily_rows else (0.0,) * 4)
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize(
|
||||
"complete, classifier_requests, estimated, saved, expected_spend, baseline, pct, coverage",
|
||||
[
|
||||
(True, 3, 2, 30.0, 10.0, 32.5, 92.3, "complete"),
|
||||
(True, 3, 0, 0.0, 10.0, None, None, "complete"),
|
||||
(True, 2, 2, 30.0, 10.0, None, None, "partial"),
|
||||
(False, 3, 2, 30.0, 10.0, None, None, "partial"),
|
||||
(True, 3, 2, -0.5, 10.0, 2.0, -25.0, "complete"),
|
||||
],
|
||||
ids=["covered-subset", "unknown-only", "missing-classifier", "partial-history", "loss"],
|
||||
)
|
||||
async def test_daily_costs_preserve_actual_charges_and_compare_only_matching_estimates(
|
||||
self, monkeypatch: pytest.MonkeyPatch, complete: bool, classifier_requests: int,
|
||||
estimated: int, saved: float, expected_spend: float | None,
|
||||
baseline: float | None, pct: float | None, coverage: str,
|
||||
) -> None:
|
||||
costs: Final = AutoRouterDailyCosts(
|
||||
requests=3, llm_spend=9.7, classifier_cost=0.3, classifier_requests=classifier_requests,
|
||||
estimated_requests=estimated, estimated_actual_spend=2.5 if estimated else 0.0,
|
||||
saved_spend=saved, complete=complete,
|
||||
)
|
||||
response: Final = await self._benchmarks(monkeypatch, rows=[], model_list=[], costs=costs)
|
||||
totals: Final = response.totals
|
||||
assert totals.saved_spend == saved
|
||||
assert totals.spend == expected_spend
|
||||
assert totals.llm_spend == 9.7
|
||||
assert totals.classifier_cost == (0.3 if classifier_requests == 3 else None)
|
||||
assert totals.baseline_spend == baseline
|
||||
assert totals.saved_pct == pct
|
||||
assert totals.cost_coverage == coverage
|
||||
assert totals.cost_requests == (3 if complete else None)
|
||||
assert totals.savings_estimated_turns == estimated
|
||||
assert totals.savings_estimated_actual_spend == (2.5 if estimated else 0.0)
|
||||
assert totals.sessions == totals.turns == 0
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize("recovery", ["complete", "partial", "empty", "failure", "invalid", "not-needed", "mismatched"])
|
||||
async def test_historical_cost_recovery_preserves_durable_savings_and_request_scope(
|
||||
self, monkeypatch: pytest.MonkeyPatch, recovery: str,
|
||||
) -> None:
|
||||
from litellm.proxy.route_llm_request import ROUTE_ENDPOINT_MAPPING
|
||||
|
||||
recovered: Final = AutoRouterDailyCosts(
|
||||
requests=3, llm_spend=4.5, classifier_cost=0.5, classifier_requests=3,
|
||||
estimated_requests=2, estimated_actual_spend=3.0, saved_spend=7.5 if recovery == "mismatched" else 6.5,
|
||||
complete=recovery != "partial",
|
||||
estimated_requests=estimated, estimated_actual_spend=3.0 if estimated else 0.0, saved_spend=6.5,
|
||||
)
|
||||
recorded: Final = recovered if recovery == "not-needed" else AutoRouterDailyCosts(
|
||||
requests=1, llm_spend=2.0, classifier_cost=0.25, classifier_requests=1,
|
||||
estimated_requests=1, estimated_actual_spend=2.25, saved_spend=6.5, complete=False,
|
||||
)
|
||||
recovery_result: Final = {
|
||||
"empty": [], "failure": RuntimeError("retained logs unavailable"), "invalid": [{"unexpected": True}],
|
||||
}.get(recovery, [recovered.model_dump()])
|
||||
query: Final = AsyncMock(side_effect=[[self.ROW.model_dump()], [recorded.model_dump()], recovery_result])
|
||||
transaction: Final = SimpleNamespace(query_raw=query, execute_raw=AsyncMock())
|
||||
recorded: Final = costs if complete else costs.model_copy(update={
|
||||
"requests": 1, "llm_spend": 2.0, "classifier_cost": 0.25, "classifier_requests": 1, "complete": False,
|
||||
"estimated_requests": 1, "estimated_actual_spend": 2.25,
|
||||
})
|
||||
recovered: Final = [costs.model_copy(update={"saved_spend": recovered_saved or 6.5}).model_dump()]
|
||||
query: Final = AsyncMock(side_effect=[
|
||||
[self.ROW.model_dump()], [recorded.model_dump()],
|
||||
recovered if recovered_saved is not None else RuntimeError("Recovery unavailable"),
|
||||
])
|
||||
reader: Final = MagicMock()
|
||||
reader.tx.return_value.__aenter__.return_value = transaction
|
||||
reader.tx.return_value.__aenter__.return_value = SimpleNamespace(query_raw=query, execute_raw=AsyncMock())
|
||||
monkeypatch.setattr(proxy_server, "prisma_client", SimpleNamespace(
|
||||
db=SimpleNamespace(query_raw=query), read_db=reader,
|
||||
))
|
||||
breakdown: Final = AsyncMock(side_effect=AssertionError("Unused usage breakdowns must not be loaded"))
|
||||
monkeypatch.setattr(common_daily_activity, "get_daily_activity_aggregated", breakdown)
|
||||
response: Final = await auto_router_endpoints.get_auto_router_benchmarks(
|
||||
user_api_key_dict=ADMIN, start_date="2026-07-01", end_date="2026-08-01",
|
||||
api_key="selected-key", user_id="selected-user",
|
||||
)
|
||||
totals: Final = response.totals
|
||||
assert totals.saved_spend == 6.5
|
||||
assert response.groups[0].saved_spend == self.ROW.saved_spend
|
||||
fallback: Final = recovery in ("empty", "failure", "invalid")
|
||||
assert (totals.spend, totals.llm_spend, totals.classifier_cost) == (
|
||||
(2.25, 2.0, 0.25) if fallback else (5.0, 4.5, 0.5)
|
||||
assert (totals.saved_spend, totals.spend, totals.baseline_spend, totals.saved_pct) == (
|
||||
6.5, expected_spend, baseline, pct,
|
||||
)
|
||||
comparable: Final = recovery in ("complete", "not-needed")
|
||||
assert (totals.baseline_spend, totals.saved_pct) == ((9.5, 68.4) if comparable else (None, None))
|
||||
assert totals.cost_coverage == ("complete" if comparable or recovery == "mismatched" else "partial")
|
||||
assert query.await_count == (2 if recovery == "not-needed" else 3)
|
||||
if recovery != "not-needed":
|
||||
sql, start, end, key, user, endpoints = query.call_args.args
|
||||
assert sql == AUTOROUTER_HISTORICAL_COSTS_SQL
|
||||
assert (start, end, key, user) == ("2026-07-01", "2026-08-01", "selected-key", "selected-user")
|
||||
assert json.loads(endpoints) == ROUTE_ENDPOINT_MAPPING
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_daily_savings_failure_never_falls_back_to_session_savings(self, monkeypatch: pytest.MonkeyPatch) -> None:
|
||||
query: Final = AsyncMock(side_effect=[[self.ROW.model_dump()], RuntimeError("Daily savings unavailable")])
|
||||
monkeypatch.setattr(proxy_server, "prisma_client", SimpleNamespace(db=SimpleNamespace(query_raw=query)))
|
||||
with pytest.raises(RuntimeError, match="Daily savings unavailable"):
|
||||
await auto_router_endpoints.get_auto_router_benchmarks(user_api_key_dict=ADMIN)
|
||||
assert response.groups[0].saved_spend == 30.0
|
||||
assert query.call_args_list[1].args[1:] == ("2026-07-01", "2026-08-01", "selected-key", "selected-user")
|
||||
assert query.await_count == (2 if complete else 3)
|
||||
breakdown.assert_not_awaited()
|
||||
|
||||
def test_an_empty_window_folds_to_zeros(self):
|
||||
from litellm.proxy.management_endpoints.auto_router_endpoints import (
|
||||
|
|
@ -1364,32 +1235,6 @@ class TestAutoRouterSession:
|
|||
assert "$1.00" in rendered and "$2.00" in rendered
|
||||
assert "$100.00" not in rendered
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_unknown_only_session_keeps_savings_unavailable_in_the_cli(
|
||||
self, monkeypatch: pytest.MonkeyPatch, tmp_path: Path,
|
||||
) -> None:
|
||||
from litellm.proxy.client.cli.commands.statusline_script import _session_from_payload, render
|
||||
from litellm.proxy.management_endpoints.auto_router_endpoints import get_auto_router_session
|
||||
|
||||
caller: Final = UserAPIKeyAuth(api_key="sk-caller")
|
||||
self._rig(monkeypatch, [{
|
||||
**self.ROW, "api_key": caller.api_key, "session_id": "unknown-only",
|
||||
"saved_spend": 0.0, "savings_estimated_turns": 0,
|
||||
"savings_estimated_actual_spend": 0.0, "savings_estimated_saved_spend": 0.0,
|
||||
"savings_estimated_baseline_models": {},
|
||||
}])
|
||||
|
||||
response: Final = await get_auto_router_session(user_api_key_dict=caller, session_id="unknown-only")
|
||||
assert response.turns == 3 and response.savings_estimated_turns == 0
|
||||
assert response.baseline_model == "anthropic/claude-opus-5"
|
||||
assert response.baseline_spend is None
|
||||
assert response.savings_estimated_baseline_spend is None
|
||||
session: Final = _session_from_payload(response.model_dump())
|
||||
assert session is not None
|
||||
assert render("claude-sonnet-5", session, tmp_path, use_color=False) == (
|
||||
"Routed to: claude-sonnet-5 · Savings unavailable"
|
||||
)
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_another_keys_session_is_a_404_even_for_an_admin(self, monkeypatch: pytest.MonkeyPatch):
|
||||
# The scope is the caller's own key hash, exactly what the spend writer keyed the row under;
|
||||
|
|
|
|||
|
|
@ -19,30 +19,6 @@ from litellm.types.utils import Usage
|
|||
pytestmark = pytest.mark.usefixtures("local_model_cost_map")
|
||||
|
||||
|
||||
@pytest.mark.parametrize("estimate, saved, expected", [
|
||||
(None, 0.0, 0.0),
|
||||
(None, -0.5, -0.5),
|
||||
(None, 0.5, 0.5),
|
||||
(None, None, None),
|
||||
(None, True, None),
|
||||
(None, float("nan"), None),
|
||||
({"version": 3, "status": "estimated"}, 0.0, 0.0),
|
||||
({"version": 2, "status": "estimated"}, -0.5, -0.5),
|
||||
({"version": 1, "status": "unknown"}, 0.5, None),
|
||||
({"version": 0, "status": "estimated"}, 0.5, None),
|
||||
])
|
||||
def test_known_autorouter_savings_preserves_zero_and_legacy_provenance(
|
||||
estimate: dict[str, object] | None, saved: object, expected: float | None,
|
||||
) -> None:
|
||||
from litellm.proxy.spend_tracking.savings import known_autorouter_savings
|
||||
|
||||
result: Final = known_autorouter_savings(
|
||||
model=None, custom_llm_provider=None, routing_decision=None, usage_object=None,
|
||||
recorded_autorouter_savings=saved, recorded_autorouter_savings_estimate=estimate,
|
||||
)
|
||||
assert result == expected
|
||||
|
||||
|
||||
@pytest.mark.parametrize("model,usage", [
|
||||
(None, {"cache_read_input_tokens": 100}),
|
||||
("claude-sonnet-5", None),
|
||||
|
|
|
|||
|
|
@ -1,6 +1,5 @@
|
|||
import { QueryClient, QueryClientProvider } from "@tanstack/react-query";
|
||||
import { fireEvent, render, screen, within } from "@testing-library/react";
|
||||
import userEvent from "@testing-library/user-event";
|
||||
import React from "react";
|
||||
import { beforeEach, describe, expect, it, vi } from "vitest";
|
||||
|
||||
|
|
@ -184,9 +183,6 @@ describe("AutoRouterBenchmarksTab", () => {
|
|||
expect(screen.getByText("Unavailable")).toBeInTheDocument();
|
||||
expect(screen.getByText("$23.13")).toBeInTheDocument();
|
||||
expect(screen.getByText("Estimated baseline spend")).toBeInTheDocument();
|
||||
expect(screen.queryByText(/turns estimated/)).not.toBeInTheDocument();
|
||||
expect(screen.queryByText(/older estimates.*excluded|excluded.*older estimates/)).not.toBeInTheDocument();
|
||||
expect(screen.queryByText("Actual spend on covered turns")).not.toBeInTheDocument();
|
||||
});
|
||||
|
||||
it("leads with total estimated savings, before the four session-shape metrics", () => {
|
||||
|
|
@ -222,148 +218,6 @@ describe("AutoRouterBenchmarksTab", () => {
|
|||
expect(screen.getByText("5.3M")).toBeInTheDocument();
|
||||
});
|
||||
|
||||
it("keeps daily savings visible when matching historical costs are unavailable", () => {
|
||||
const dailyValues: Partial<Totals> = {
|
||||
spend: null,
|
||||
llm_spend: null,
|
||||
cost_coverage: "unavailable",
|
||||
cost_requests: null,
|
||||
classifier_cost: null,
|
||||
baseline_spend: null,
|
||||
saved_pct: null,
|
||||
};
|
||||
const daily = totals(dailyValues);
|
||||
mockHook({ data: response([group()], daily) });
|
||||
renderTab();
|
||||
|
||||
expect(screen.getByText("$2,174.59")).toBeInTheDocument();
|
||||
expect(screen.getAllByRole("definition").map((node) => node.textContent)).toEqual([
|
||||
"Unavailable",
|
||||
"Unavailable",
|
||||
"Unavailable",
|
||||
"Unavailable",
|
||||
]);
|
||||
expect(screen.getAllByText("Unavailable")).toHaveLength(4);
|
||||
expect(screen.getByText("$23.13")).toBeInTheDocument();
|
||||
expect(screen.queryByText("-86%")).not.toBeInTheDocument();
|
||||
expect(screen.queryByText(/\/ 1K turns/)).not.toBeInTheDocument();
|
||||
expect(screen.getByText("Whole sessions overlapping the selected dates")).toBeInTheDocument();
|
||||
});
|
||||
|
||||
it("shows recorded historical costs without presenting partial costs as the full savings comparison", () => {
|
||||
const recordedCosts: Partial<Totals> = {
|
||||
spend: 12,
|
||||
llm_spend: 10,
|
||||
classifier_cost: null,
|
||||
cost_coverage: "partial",
|
||||
cost_requests: null,
|
||||
baseline_spend: null,
|
||||
saved_pct: null,
|
||||
};
|
||||
mockHook({ data: response([], totals(recordedCosts)) });
|
||||
renderTab();
|
||||
|
||||
expect(screen.getByText("$2,174.59")).toBeInTheDocument();
|
||||
expect(screen.getAllByRole("definition").map((node) => node.textContent)).toEqual([
|
||||
"$12.00",
|
||||
"$10.00",
|
||||
"Unavailable",
|
||||
"Unavailable",
|
||||
]);
|
||||
expect(screen.getByText("Partial")).toBeInTheDocument();
|
||||
expect(screen.getByText("Some request costs are unavailable")).toBeInTheDocument();
|
||||
expect(screen.queryByText("-86%")).not.toBeInTheDocument();
|
||||
expect(screen.queryByText(/\/ 1K turns/)).not.toBeInTheDocument();
|
||||
});
|
||||
|
||||
it("separates actual spend from the estimated-request comparison and uses request-date classification counts", () => {
|
||||
const requestCosts: Partial<Totals> = {
|
||||
spend: 12,
|
||||
llm_spend: 10,
|
||||
classifier_cost: 2,
|
||||
cost_requests: 4,
|
||||
savings_estimated_turns: 2,
|
||||
savings_estimated_actual_spend: 4,
|
||||
saved_spend: 6,
|
||||
baseline_spend: 10,
|
||||
saved_pct: 60,
|
||||
};
|
||||
mockHook({ data: response([], totals(requestCosts)) });
|
||||
renderTab();
|
||||
|
||||
expect(screen.getByText("$6.00")).toBeInTheDocument();
|
||||
expect(screen.getByText("-60%")).toBeInTheDocument();
|
||||
expect(screen.getByText("2 of 4 requests have savings estimates")).toBeInTheDocument();
|
||||
expect(screen.getByText("Actual spend on estimated requests")).toBeInTheDocument();
|
||||
expect(screen.getAllByRole("definition").map((node) => node.textContent)).toEqual([
|
||||
"$12.00",
|
||||
"$10.00",
|
||||
"$2.00",
|
||||
"$4.00",
|
||||
"$10.00",
|
||||
]);
|
||||
expect(screen.getByText("($500.00 / 1K turns)")).toBeInTheDocument();
|
||||
});
|
||||
|
||||
it("labels a selected router as session usage and keeps its recorded cost breakdown", async () => {
|
||||
const user = userEvent.setup();
|
||||
const dailyValues: Partial<Totals> = {
|
||||
saved_spend: 3_000,
|
||||
spend: null,
|
||||
llm_spend: null,
|
||||
cost_coverage: "unavailable",
|
||||
cost_requests: null,
|
||||
classifier_cost: null,
|
||||
baseline_spend: null,
|
||||
saved_pct: null,
|
||||
};
|
||||
const daily = totals(dailyValues);
|
||||
mockHook({ data: response([group()], daily) });
|
||||
renderTab();
|
||||
|
||||
expect(screen.getByText("$3,000.00")).toBeInTheDocument();
|
||||
await user.click(screen.getByRole("combobox"));
|
||||
await user.click(await screen.findByRole("option", { name: "claude-auto" }));
|
||||
|
||||
expect(screen.getByRole("heading", { name: "Auto-router session usage" })).toBeInTheDocument();
|
||||
expect(screen.getByText("Whole sessions overlapping the selected dates")).toBeInTheDocument();
|
||||
expect(screen.getByText("$2,174.59")).toBeInTheDocument();
|
||||
expect(screen.queryByText("$3,000.00")).not.toBeInTheDocument();
|
||||
expect(screen.getByText("-86%")).toBeInTheDocument();
|
||||
expect(screen.getAllByRole("definition").map((node) => node.textContent)).toEqual([
|
||||
"$359.86",
|
||||
"$353.71",
|
||||
"$6.15",
|
||||
"$2,534.45",
|
||||
]);
|
||||
});
|
||||
|
||||
it("keeps selected-router spend visible without a baseline or percentage for unestimated turns", async () => {
|
||||
const user = userEvent.setup();
|
||||
const unknownSavings: Partial<AutoRouterBenchmarkGroup> = {
|
||||
savings_estimated_turns: 0,
|
||||
savings_estimated_actual_spend: 0,
|
||||
saved_spend: 0,
|
||||
baseline_spend: null,
|
||||
saved_pct: null,
|
||||
saved_per_session: 0,
|
||||
};
|
||||
mockHook({ data: response([group(unknownSavings)]) });
|
||||
renderTab();
|
||||
|
||||
await user.click(screen.getByRole("combobox"));
|
||||
await user.click(await screen.findByRole("option", { name: "claude-auto" }));
|
||||
|
||||
expect(screen.getByRole("heading", { name: "Auto-router session usage" })).toBeInTheDocument();
|
||||
expect(screen.getAllByRole("definition").map((node) => node.textContent)).toEqual([
|
||||
"$359.86",
|
||||
"$353.71",
|
||||
"$6.15",
|
||||
"Unavailable",
|
||||
]);
|
||||
expect(screen.queryByText("0%")).not.toBeInTheDocument();
|
||||
});
|
||||
|
||||
it.each([
|
||||
{ spend: 20665.28, classifier_cost: 342.18, turns: 140815, llm: "$20,323.10", cost: "$342.18", rate: "$2.43" },
|
||||
{ spend: 0, classifier_cost: 0, turns: 0, llm: "$0.00", cost: "$0.00", rate: "$0.00" },
|
||||
|
|
|
|||
|
|
@ -29,21 +29,21 @@ const cache = {
|
|||
|
||||
const stats = {
|
||||
sessions: 2,
|
||||
turns: 4,
|
||||
turns: 2,
|
||||
avg_turns_per_session: 2,
|
||||
avg_session_seconds: 30,
|
||||
avg_tokens_per_session: 100,
|
||||
spend: 1.25,
|
||||
llm_spend: 1,
|
||||
spend: 100,
|
||||
llm_spend: 99.75,
|
||||
cost_coverage: "complete" as const,
|
||||
cost_requests: null,
|
||||
savings_estimated_turns: 4,
|
||||
savings_estimated_actual_spend: 1.25,
|
||||
savings_estimated_turns: 1,
|
||||
savings_estimated_actual_spend: 1,
|
||||
classifier_cost: 0.25,
|
||||
saved_spend: 8.75,
|
||||
baseline_spend: 10,
|
||||
saved_pct: 87.5,
|
||||
saved_per_session: 4.375,
|
||||
saved_spend: 1,
|
||||
baseline_spend: 2,
|
||||
saved_pct: 50,
|
||||
saved_per_session: 0.5,
|
||||
cache,
|
||||
};
|
||||
|
||||
|
|
@ -107,8 +107,8 @@ describe("KeyAutoRouterUsageTab", () => {
|
|||
"Unavailable",
|
||||
"Unavailable",
|
||||
]);
|
||||
expect(screen.getByText("$4.38")).toBeInTheDocument();
|
||||
expect(screen.queryByText("-88%")).not.toBeInTheDocument();
|
||||
expect(screen.getByText("$0.5000")).toBeInTheDocument();
|
||||
expect(screen.queryByText("-50%")).not.toBeInTheDocument();
|
||||
expect(screen.getByText("All auto-routers")).toBeInTheDocument();
|
||||
|
||||
const benchmarkUrl = new URL(requestedUrls().find((url) => url.includes("/auto_router/benchmarks")) ?? "");
|
||||
|
|
@ -121,63 +121,18 @@ describe("KeyAutoRouterUsageTab", () => {
|
|||
|
||||
expect(screen.getByRole("heading", { name: "Auto-router session usage" })).toBeInTheDocument();
|
||||
expect(screen.getByText("Whole sessions overlapping the selected dates")).toBeInTheDocument();
|
||||
expect(screen.getByText("$8.75")).toBeInTheDocument();
|
||||
expect(screen.queryByText("$12.75")).not.toBeInTheDocument();
|
||||
expect(screen.getByText("-88%")).toBeInTheDocument();
|
||||
expect(screen.getByText("Actual auto-router spend")).toBeInTheDocument();
|
||||
expect(screen.getByText("$1.25")).toBeInTheDocument();
|
||||
expect(screen.getByText("LLM spend")).toBeInTheDocument();
|
||||
expect(screen.getByText("$1.00")).toBeInTheDocument();
|
||||
expect(screen.getByText("Classification cost")).toBeInTheDocument();
|
||||
expect(screen.getByText("$0.2500")).toBeInTheDocument();
|
||||
expect(screen.getByText("($62.50 / 1K turns)")).toBeInTheDocument();
|
||||
expect(screen.getByText("Estimated baseline spend")).toBeInTheDocument();
|
||||
expect(screen.getByText("$10.00")).toBeInTheDocument();
|
||||
expect(screen.getByText("Auto-router prompt caching")).toBeInTheDocument();
|
||||
expect(screen.getAllByText("50.0%").length).toBeGreaterThan(0);
|
||||
});
|
||||
|
||||
it("shows matching estimated costs and coverage when selecting a mixed-coverage router", async () => {
|
||||
const user = userEvent.setup();
|
||||
const mixed = {
|
||||
...stats,
|
||||
turns: 2,
|
||||
spend: 100,
|
||||
llm_spend: 100,
|
||||
classifier_cost: 0,
|
||||
savings_estimated_turns: 1,
|
||||
savings_estimated_actual_spend: 1,
|
||||
saved_spend: 1,
|
||||
baseline_spend: 2,
|
||||
saved_pct: 50,
|
||||
};
|
||||
vi.stubGlobal("fetch", async (request: Request | string) => {
|
||||
const url = typeof request === "string" ? request : request.url;
|
||||
return jsonResponse(
|
||||
url.includes("/auto_router/benchmarks")
|
||||
? { ...benchmarks, groups: [{ ...benchmarks.groups[0], ...mixed }] }
|
||||
: noDeployments,
|
||||
);
|
||||
});
|
||||
const activity = {
|
||||
dateValue: { from: new Date(2025, 0, 1), to: new Date(2025, 0, 31) },
|
||||
onDateChange: vi.fn(),
|
||||
};
|
||||
renderWithProviders(<KeyAutoRouterUsageTab accessToken="test-token" keyToken="key-hash-1" activity={activity} />);
|
||||
expect(await screen.findByText("$12.75")).toBeInTheDocument();
|
||||
|
||||
await user.click(screen.getByRole("combobox"));
|
||||
await user.click(await screen.findByRole("option", { name: "router-one" }));
|
||||
|
||||
expect(screen.getByText("-50%")).toBeInTheDocument();
|
||||
expect(screen.getByText("1 of 2 requests have savings estimates")).toBeInTheDocument();
|
||||
expect(screen.getByText("Actual spend on estimated requests")).toBeInTheDocument();
|
||||
expect(screen.getAllByRole("definition").map((node) => node.textContent)).toEqual([
|
||||
"$100.00",
|
||||
"$100.00",
|
||||
"$0.00",
|
||||
"$99.75",
|
||||
"$0.2500",
|
||||
"$1.00",
|
||||
"$2.00",
|
||||
]);
|
||||
expect(screen.getByText("Auto-router prompt caching")).toBeInTheDocument();
|
||||
expect(screen.getAllByText("50.0%").length).toBeGreaterThan(0);
|
||||
});
|
||||
});
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue