scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-10-06 02:48:13 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 103

fix(auto-router): attribute router-day savings without spend logs or a session (#44226)

Browse source 
Per-router auto-router savings were missing on proxies with
disable_spend_logs and for requests without a session id under
missing_session_id: omit. The router-day row is written by the
auto-router turn, whose enqueue sat behind the spend-logs flag and whose
builder dropped sessionless requests, while LiteLLM_DailyUserSpend has
neither gate. That money then showed only as unattributed savings.

Enqueue the turn whether or not spend logs are kept, and keep a
sessionless turn with an empty session id that writes the router-day row
while the session upserts skip it. The day and session rows still commit
in one statement, so late baseline corrections keep their ordering.

Without spend logs, write only the router-day aggregate: the turn drops
its session id, so no per-session row is stored, and baseline capture is
skipped, since a baseline observation can only publish once its
request's spend log exists. The flag keeps its meaning of no per-request
or per-session data, while the daily per-router money matches
LiteLLM_DailyUserSpend.

Co-authored-by: Claude Opus 5.5 <noreply@anthropic.com>
This commit is contained in:
tin-berri 2026-10-02 20:56:45 -07:00 • committed by GitHub
parent f9a32ffcb5
commit db14931401
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
6 changed files with 246 additions and 21 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

9
litellm/proxy/db/autorouter_session_rollup.py
View file

@ -266,7 +266,8 @@ def build_autorouter_turn_transaction(
the payload's own usage record through the savings owner, never handed in beside it.
The baseline the turn's saved_spend was priced against travels with the turn, so the
row can name the counterfactual for the money it holds even after the router is
reconfigured or removed.
reconfigured or removed. A request with no session id still owns its router-day money,
so it becomes a turn with an empty session id that writes the day row and no session row.
"""
if payload.get("status") != "success":
return None
@ -278,9 +279,9 @@ def build_autorouter_turn_transaction(
router_name: Final = routing_decision.get("router_model_name") or payload.get("model_group")
api_key: Final = payload.get("api_key") or ""
user_id: Final = payload.get("user") or ""
session_id: Final = payload.get("session_id")
session_id: Final = payload.get("session_id") or ""
model: Final = payload.get("model")
if not (isinstance(router_name, str) and router_name and (api_key or user_id) and session_id and model):
if not (isinstance(router_name, str) and router_name and (api_key or user_id) and model):
return None
turn_at: Final = _turn_time_utc(str(payload.get("startTime") or ""))
if turn_at is None:
@ -379,7 +380,7 @@ SELECT
{_p("classifier_cost")}::float8, 1, {_TIER_DELTA}, {_BASELINE_DELTA},
{_p("savings_estimated_turns")}::int, {_p("savings_estimated_actual_spend")}::float8,
{_p("savings_estimated_saved_spend")}::float8, {_ESTIMATED_BASELINE_DELTA}
WHERE {required_identity}::text <> ''
WHERE {required_identity}::text <> '' AND {_p("session_id")}::text <> ''
ON CONFLICT ({user_column}api_key, session_id, router_name) DO UPDATE SET
turns = t.turns + 1,
total_tokens = t.total_tokens + EXCLUDED.total_tokens,

24
litellm/proxy/db/db_spend_update_writer.py
View file

@ -7,6 +7,7 @@ Module responsible for
import asyncio
import copy
import dataclasses
import json
import os
import random
@ -550,6 +551,13 @@ class DBSpendUpdateWriter:
):
return False
# The auto-router router-day rollup is an aggregate like the daily spend tables, so it is
# written whether or not per-request spend logs are kept; per-session rows are not.
await self._enqueue_autorouter_turn_transaction(
payload=payload,
prisma_client=prisma_client,
spend_logs_kept=disable_spend_logs is False,
)
if disable_spend_logs is False:
await self._enqueue_tool_usage_transaction(
payload=payload,
@ -557,10 +565,6 @@ class DBSpendUpdateWriter:
prisma_client=prisma_client,
kwargs=kwargs,
)
await self._enqueue_autorouter_turn_transaction(
payload=payload,
prisma_client=prisma_client,
)
else:
verbose_proxy_logger.debug(
"disable_spend_logs=True. Skipping writing spend logs to db. Other spend updates - Key/User/Team table will still occur."
@ -747,6 +751,7 @@ class DBSpendUpdateWriter:
self,
payload: SpendLogsPayload,
prisma_client: "PrismaClient | None",
spend_logs_kept: bool = True,
) -> None:
try:
if prisma_client is None:
@ -787,14 +792,21 @@ class DBSpendUpdateWriter:
saved_spend=savings_spend.autorouter,
)
try:
if await self._enqueue_baseline_accounting(payload, metadata, transaction, prisma_client):
# A baseline observation publishes only once its spend log exists, so without spend logs
# it could never publish; the plain turn still carries this request's recorded savings.
if spend_logs_kept and await self._enqueue_baseline_accounting(
payload, metadata, transaction, prisma_client
):
return
except Exception: # noqa: BLE001 # optional baseline capture must preserve the original actual-spend rollup
verbose_proxy_logger.warning("Auto-router baseline observation was unavailable; actual turn retained")
if transaction is None:
return
# Without spend logs only the router-day aggregate is kept: an empty session id makes the
# session upserts skip the row, so no per-session record is stored.
kept: Final = transaction if spend_logs_kept else dataclasses.replace(transaction, session_id="")
async with prisma_client._autorouter_turn_transactions_lock:
prisma_client.autorouter_turn_transactions.append(transaction)
prisma_client.autorouter_turn_transactions.append(kept)
except Exception as e: # noqa: BLE001 # a metrics enqueue must never fail the spend write
verbose_proxy_logger.debug("_enqueue_autorouter_turn_transaction error (non-blocking): %s", e)

103
tests/proxy_behavior/spend/test_autorouter_session_rollup.py
View file

@ -20,8 +20,10 @@ from typing_extensions import ReadOnly
from litellm.proxy.db.autorouter_session_rollup import (
AUTOROUTER_BENCHMARKS_SQL,
UPSERT_AUTOROUTER_SESSION_SQL,
UPSERT_AUTOROUTER_USER_SESSION_SQL,
AutoRouterTurnTransaction,
flush_autorouter_turn_transactions,
write_autorouter_turn,
)
from litellm.proxy.db.db_transaction_queue.spend_log_cleanup import SpendLogCleanup
@ -684,3 +686,104 @@ async def test_a_router_type_change_mid_session_keeps_session_shape_with_the_ses
assert (rows["complexity"]["sessions"], rows["complexity"]["session_turns"], rows["complexity"]["turns"]) == (1, 2, 1)
assert (rows["quality"]["sessions"], rows["quality"]["session_turns"], rows["quality"]["turns"]) == (0, 0, 1)
assert rows["quality"]["spend"] == 2.0
@pytest.mark.parametrize("statement", [UPSERT_AUTOROUTER_SESSION_SQL, UPSERT_AUTOROUTER_USER_SESSION_SQL])
async def test_a_sessionless_turn_writes_its_router_day_row_and_no_session_row(db, statement: str):
key = f"k-{uuid.uuid4()}"
router = f"auto-{uuid.uuid4()}"
for offset in range(2):
await write_autorouter_turn(
db,
AutoRouterTurnTransaction(
api_key=key,
user_id="u-sessionless",
session_id="",
router_name=router,
router_type="complexity",
model="A",
turn_at=T0 + timedelta(seconds=offset),
total_tokens=10,
spend=1.0,
saved_spend=2.0,
classifier_cost=0.1,
covered=True,
cache_hit=False,
cache_ttl_seconds=None,
cache_touched=True,
savings_estimated_turns=1,
savings_estimated_actual_spend=1.0,
savings_estimated_saved_spend=2.0,
),
statement,
)
(day,) = await _days(db, key, router=router)
assert (day["turns"], day["spend"], day["saved_spend"], day["classifier_cost"]) == (2, 2.0, 4.0, 0.2)
assert (day["sessions"], day["session_turns"]) == (0, 0)
for table in ("LiteLLM_AutoRouterSession", "LiteLLM_AutoRouterUserSession"):
assert await db.query_raw(f'SELECT 1 FROM "{table}" WHERE router_name = $1', router) == []
async def test_router_day_money_reconciles_with_the_overall_daily_total_including_sessionless_requests(db):
from litellm.proxy.db.daily_spend_bulk_upsert import DAILY_SPEND_TABLES, build_bulk_upsert, merge_by_conflict_key
key = f"k-{uuid.uuid4()}"
router = f"auto-{uuid.uuid4()}"
requests = (("session-1", 0.25, 1.5), ("session-1", 0.5, 2.0), ("", 0.1, 0.25))
for offset, (session_id, spend, saved) in enumerate(requests):
await write_autorouter_turn(
db,
AutoRouterTurnTransaction(
api_key=key,
user_id="u1",
session_id=session_id,
router_name=router,
router_type="complexity",
model="A",
turn_at=T0 + timedelta(seconds=offset),
total_tokens=10,
spend=spend,
saved_spend=saved,
classifier_cost=0.0,
covered=True,
cache_hit=False,
cache_ttl_seconds=None,
cache_touched=True,
savings_estimated_turns=1,
savings_estimated_actual_spend=spend,
savings_estimated_saved_spend=saved,
),
)
table = DAILY_SPEND_TABLES["user"]
statement, values = build_bulk_upsert(
table,
merge_by_conflict_key(
table,
tuple(
{
"user_id": "u1",
"date": T0.date().isoformat(),
"api_key": key,
"model": "A",
"custom_llm_provider": "anthropic",
"model_group": router,
"spend": spend,
"api_requests": 1,
"successful_requests": 1,
"autorouter_savings_spend": saved,
}
for _, spend, saved in requests
),
),
)
await db.execute_raw(statement, *values)
(overall,) = await db.query_raw(
'SELECT SUM(autorouter_savings_spend)::float8 AS saved FROM "LiteLLM_DailyUserSpend" WHERE date = $1 AND api_key = $2',
T0.date().isoformat(),
key,
)
(row,) = await _days(db, key, router=router)
assert overall["saved"] == row["saved_spend"] == pytest.approx(3.75)
assert (row["turns"], row["spend"], row["sessions"], row["session_turns"]) == (3, pytest.approx(0.85), 1, 2)

65
tests/proxy_behavior/spend/test_baseline_accounting.py
View file

@ -249,17 +249,10 @@ async def test_retired_history_never_recreates_an_initial_zero(db: Prisma, recor
assert after["savings_estimated_turns"] == 1 and after["savings_estimated_actual_spend"] == 0.17
async def test_native_observation_enters_spend_pipeline_once_with_shared_daily_attribution(
db: Prisma, record: Callable[..., BaselineAccountingRecord], monkeypatch: pytest.MonkeyPatch,
) -> None:
import os
from litellm.proxy.common_utils.user_api_key_cache import UserApiKeyCache
from litellm.proxy.db.db_spend_update_writer import DBSpendUpdateWriter
def _native_observation_payload(event: BaselineAccountingRecord) -> dict[str, object]:
"""The spend payload a captured, sessioned, auto-routed anthropic_messages request produces."""
from litellm.proxy.hooks.autorouter_baseline_cache import CapturedBaselineObservation
from litellm.proxy.utils import PrismaClient, ProxyLogging
event: Final = record("routed", identical=False)
capture: Final = CapturedBaselineObservation(
scope=event.scope, api_key=event.api_key, session_id=event.session_id,
router_name=event.router_name, baseline_model=event.baseline_model,
@ -272,7 +265,7 @@ async def test_native_observation_enters_spend_pipeline_once_with_shared_daily_a
"autorouter_savings": None, "autorouter_savings_estimate": {"version": 3, "status": "unknown", "reason": "pending_projection"},
"autorouter_baseline_observation": capture.model_dump_json(),
}
payload: Final = {
return {
"request_id": event.observation.request_id, "api_key": event.api_key, "session_id": event.session_id,
"startTime": datetime.fromtimestamp(event.observation.started_at, timezone.utc).isoformat(),
"endTime": datetime.fromtimestamp(event.observation.available_at, timezone.utc).isoformat(),
@ -282,6 +275,19 @@ async def test_native_observation_enters_spend_pipeline_once_with_shared_daily_a
"user": None, "team_id": "", "organization_id": "org", "agent_id": None,
"end_user": "", "request_tags": '["tag","tag"]',
}
async def test_native_observation_enters_spend_pipeline_once_with_shared_daily_attribution(
db: Prisma, record: Callable[..., BaselineAccountingRecord], monkeypatch: pytest.MonkeyPatch,
) -> None:
import os
from litellm.proxy.common_utils.user_api_key_cache import UserApiKeyCache
from litellm.proxy.db.db_spend_update_writer import DBSpendUpdateWriter
from litellm.proxy.utils import PrismaClient, ProxyLogging
event: Final = record("routed", identical=False)
payload: Final = _native_observation_payload(event)
monkeypatch.delenv("DATABASE_URL_READ_REPLICA", raising=False)
client: Final = PrismaClient(os.environ["DATABASE_URL"], ProxyLogging(UserApiKeyCache()))
writer: Final = DBSpendUpdateWriter()
@ -319,3 +325,42 @@ 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_without_spend_logs_a_captured_turn_keeps_only_its_router_day_row(
db: Prisma, record: Callable[..., BaselineAccountingRecord], monkeypatch: pytest.MonkeyPatch,
) -> None:
import os
from litellm.proxy.common_utils.user_api_key_cache import UserApiKeyCache
from litellm.proxy.db.autorouter_session_rollup import flush_autorouter_turn_transactions
from litellm.proxy.db.db_spend_update_writer import DBSpendUpdateWriter
from litellm.proxy.utils import PrismaClient, ProxyLogging
event: Final = record("unlogged", identical=False)
monkeypatch.delenv("DATABASE_URL_READ_REPLICA", raising=False)
client: Final = PrismaClient(os.environ["DATABASE_URL"], ProxyLogging(UserApiKeyCache()))
try:
await client.db.connect()
await DBSpendUpdateWriter()._enqueue_autorouter_turn_transaction(
_native_observation_payload(event), client, spend_logs_kept=False
)
assert client.baseline_accounting_transactions == []
(turn,) = client.autorouter_turn_transactions
await flush_autorouter_turn_transactions(client, (turn,), n_retry_times=0)
finally:
client.autorouter_turn_transactions.clear()
await client.db.disconnect()
assert await db.query_raw(
'SELECT 1 FROM "LiteLLM_AutoRouterBaselineObservation" WHERE request_id=$1', event.observation.request_id
) == []
days: Final = await db.query_raw(
'SELECT turns, spend FROM "LiteLLM_AutoRouterDailySpend" WHERE api_key=$1 AND router_name=$2',
event.api_key, event.router_name,
)
assert [(day["turns"], day["spend"]) for day in days] == [(1, 0.17)]
for table in ("LiteLLM_AutoRouterSession", "LiteLLM_AutoRouterUserSession"):
assert await db.query_raw(
f'SELECT 1 FROM "{table}" WHERE api_key=$1 AND router_name=$2', event.api_key, event.router_name
) == []

7
tests/unit/proxy/db/test_autorouter_session_rollup.py
View file

@ -111,7 +111,6 @@ class TestBuildTransaction:
[
{"status": "failure"},
{"api_key": ""},
{"session_id": None},
{"model": ""},
{"startTime": "not-a-time"},
],
@ -119,6 +118,12 @@ class TestBuildTransaction:
def test_incomplete_payloads_are_skipped(self, payload_overrides: dict):
assert _build(payload=_payload(**payload_overrides)) is None
@pytest.mark.parametrize("session_id", [None, ""])
def test_a_request_without_a_session_keeps_its_router_day_money(self, session_id: str | None) -> None:
transaction: Final = _build(payload=_payload(session_id=session_id))
assert transaction is not None
assert (transaction.session_id, transaction.router_name, transaction.spend) == ("", "live-auto", 0.01)
@pytest.mark.parametrize("metadata", [{}, {"routing_decision": None}, {"routing_decision": {}}])
def test_requests_without_a_routing_decision_are_skipped(self, metadata: dict):
assert _build(metadata=metadata) is None

59
tests/unit/proxy/db/test_db_spend_update_writer.py
View file

@ -289,6 +289,65 @@ async def test_update_database_skips_tool_usage_when_spend_logs_disabled():
assert prisma.tool_usage_transactions == []
@pytest.mark.asyncio
@pytest.mark.parametrize("disable_spend_logs", [True, False])
@pytest.mark.parametrize("session_id", ["session-1", None])
async def test_a_routed_request_reaches_the_auto_router_rollup_whether_or_not_spend_logs_are_kept(
disable_spend_logs: bool, session_id: str | None
) -> None:
db_writer = DBSpendUpdateWriter()
db_writer._insert_spend_log_to_db = AsyncMock()
db_writer._batch_database_updates = AsyncMock()
prisma = _tool_usage_prisma()
prisma.autorouter_turn_transactions = []
prisma._autorouter_turn_transactions_lock = asyncio.Lock()
routed_payload: Final = {
**_minimal_spend_payload(),
"status": "success",
"api_key": "hashed-key",
"user": "u1",
"session_id": session_id,
"model": "claude-haiku-4-5",
"model_group": "smart-router",
"spend": 0.25,
"startTime": "2026-07-25T10:00:00+00:00",
"metadata": json.dumps(
{
"routing_decision": {"router_model_name": "smart-router", "router_type": "complexity"},
"autorouter_savings": 1.5,
}
),
}
with (
patch("litellm.proxy.proxy_server.disable_spend_logs", disable_spend_logs), # test-quality-ok: update_database reads this proxy_server module global at call time; no injection seam
patch("litellm.proxy.proxy_server.prisma_client", prisma),
patch("litellm.proxy.proxy_server.litellm_proxy_budget_name", "test-budget"),
patch(
"litellm.proxy.spend_tracking.spend_tracking_utils.get_logging_payload",
return_value=routed_payload,
),
):
await db_writer.update_database(
token="test-token",
user_id="u1",
end_user_id=None,
team_id=None,
org_id=None,
kwargs={"model": "smart-router"},
completion_response=_tool_call_response("get_weather"),
start_time=datetime.now(timezone.utc),
end_time=datetime.now(timezone.utc),
response_cost=0.25,
)
(turn,) = prisma.autorouter_turn_transactions
stored_session: Final = session_id if session_id and not disable_spend_logs else ""
assert (turn.router_name, turn.router_type, turn.session_id) == ("smart-router", "complexity", stored_session)
assert (turn.spend, turn.saved_spend) == (0.25, 1.5)
assert (prisma.tool_usage_transactions == []) is disable_spend_logs
Statement = tuple[str, tuple[object, ...]]