Coverage for .venv/lib/python3.13/site-packages/litellm/proxy/db/baseline_accounting.py: 53%
272 statements
« prev ^ index » next coverage.py v7.15.2, created at 2026-10-10 12:01 +0000
1from __future__ import annotations
3import asyncio
4import json
5from collections.abc import AsyncIterator, Callable, Sequence
6from datetime import datetime, timedelta
7from functools import reduce
8from itertools import groupby
9from types import MappingProxyType
10from typing import TYPE_CHECKING, Final, Literal, Protocol, cast
12from pydantic import BaseModel, ConfigDict, Field, TypeAdapter, model_validator
13from typing_extensions import Self
15from litellm._logging import verbose_proxy_logger
16from litellm.proxy.db.autorouter_session_rollup import (
17 AutoRouterTurnTransaction,
18 write_autorouter_turn,
19)
20from litellm.proxy.db.create_views import SupportsRawQueries
21from litellm.proxy.db.daily_spend_bulk_upsert import (
22 DAILY_SPEND_TABLES,
23 DailySpendEntity,
24 SpendRow,
25 build_bulk_upsert,
26 merge_by_conflict_key,
27)
28from litellm.proxy.db.routing_prisma_wrapper import writer_wrapper
29from litellm.proxy.spend_tracking.baseline_accounting import (
30 BaselineEstimate,
31 BaselineHistory,
32 BaselineObservation,
33 advance_baseline_history,
34)
35from litellm.proxy.spend_tracking.savings import BaselineCosts, BaselineCostSnapshot, price_baseline_comparison
37if TYPE_CHECKING: 37 ↛ 38line 37 didn't jump to line 38 because the condition on line 37 was never true
38 from litellm.proxy.utils import PrismaClient
41class DailyBaselineTarget(BaseModel):
42 model_config = ConfigDict(extra="forbid", frozen=True, strict=True)
44 entity: DailySpendEntity
45 entity_id: str | None
48class DailyBaselineAttribution(BaseModel):
49 model_config = ConfigDict(extra="forbid", frozen=True, strict=True)
51 date: str
52 api_key: str
53 model: str | None = None
54 custom_llm_provider: str | None = None
55 model_group: str | None = None
56 endpoint: str | None = None
57 mcp_namespaced_tool_name: str | None = None
58 targets: tuple[DailyBaselineTarget, ...] = ()
60 def adjustment(self, target: DailyBaselineTarget, savings_delta: float, request_id: str) -> SpendRow:
61 table: Final = DAILY_SPEND_TABLES[target.entity]
62 return MappingProxyType(
63 {
64 "date": self.date,
65 "api_key": self.api_key,
66 "model": self.model,
67 "custom_llm_provider": self.custom_llm_provider,
68 "model_group": self.model_group,
69 "endpoint": self.endpoint,
70 "mcp_namespaced_tool_name": self.mcp_namespaced_tool_name,
71 table.entity_id_column: target.entity_id,
72 "request_id": request_id,
73 "autorouter_savings_spend": savings_delta,
74 }
75 )
78class BaselineAccountingRecord(BaseModel):
79 model_config = ConfigDict(extra="forbid", frozen=True, strict=True)
81 scope: str = Field(pattern=r"^autorouter-baseline:v3:[a-f0-9]{64}$")
82 api_key: str = Field(min_length=1)
83 session_id: str = Field(min_length=1, max_length=256)
84 router_name: str = Field(min_length=1)
85 baseline_model: str = Field(min_length=1)
86 observation: BaselineObservation
87 pricing: BaselineCostSnapshot
88 turn: AutoRouterTurnTransaction | None
89 daily: DailyBaselineAttribution | None
91 @model_validator(mode="after")
92 def consistent_turn(self) -> Self:
93 turn: Final = self.turn
94 if turn is not None and (
95 (turn.api_key, turn.session_id, turn.router_name, turn.baseline_model)
96 != (self.api_key, self.session_id, self.router_name, self.baseline_model)
97 or turn.spend != self.pricing.actual_spend + self.pricing.classifier_cost
98 or any(
99 (
100 turn.saved_spend,
101 turn.savings_estimated_turns,
102 turn.savings_estimated_actual_spend,
103 turn.savings_estimated_saved_spend,
104 )
105 )
106 ):
107 raise ValueError("Baseline observation must own an unestimated turn with matching scope and actual cost")
108 return self
111class BaselinePublication(BaseModel):
112 model_config = ConfigDict(extra="forbid", frozen=True, strict=True)
114 version: Literal[3] = 3
115 comparison_id: str
116 comparison_started_at: float
117 status: Literal["estimated", "unknown"]
118 reason: str
119 provenance: Literal["observed_identical", "modeled"] | None = None
120 actual_spend: float | None = None
121 baseline_spend: float | None = None
122 input_tokens: int | None = None
123 cache_read_input_tokens: int | None = None
124 cache_creation_5m_input_tokens: int | None = None
125 cache_creation_1h_input_tokens: int | None = None
127 @property
128 def costs(self) -> BaselineCosts | None:
129 if self.status != "estimated" or self.actual_spend is None or self.baseline_spend is None:
130 return None
131 return BaselineCosts(self.actual_spend, self.baseline_spend)
134def baseline_publication(
135 record: BaselineAccountingRecord, estimate: BaselineEstimate, first_at: float
136) -> BaselinePublication:
137 costs: Final = price_baseline_comparison(record.pricing, estimate.usage, estimate.provenance)
138 details: Final = estimate.usage.prompt_tokens_details if estimate.usage is not None else None
139 writes: Final = details.cache_creation_token_details if details is not None else None
140 return BaselinePublication(
141 comparison_id=record.scope,
142 comparison_started_at=first_at,
143 status="estimated" if costs is not None else "unknown",
144 reason=estimate.reason if costs is not None or estimate.usage is None else "pricing_unavailable",
145 provenance=estimate.provenance if costs is not None else None,
146 actual_spend=costs.actual if costs is not None else None,
147 baseline_spend=costs.baseline if costs is not None else None,
148 input_tokens=details.text_tokens if details is not None else None,
149 cache_read_input_tokens=details.cached_tokens if details is not None else None,
150 cache_creation_5m_input_tokens=writes.ephemeral_5m_input_tokens if writes is not None else None,
151 cache_creation_1h_input_tokens=writes.ephemeral_1h_input_tokens if writes is not None else None,
152 )
155class _Comparison(BaseModel):
156 revision: int
157 published_revision: int
158 initial_equivalent: bool
159 retired: bool
160 history: str | None
163class _StoredRecord(BaseModel):
164 data: str
165 publication: str | None
166 conflicted: bool
167 started_at: float
170class _Change(BaseModel):
171 request_id: str
172 publication: BaselinePublication
173 api_key: str
174 user_id: str = ""
175 session_id: str
176 router_name: str
177 baseline_model: str
178 covered_delta: int
179 actual_delta: float
180 savings_delta: float
181 daily: DailyBaselineAttribution | None
184class _TransactionManager(Protocol):
185 async def __aenter__(self) -> SupportsRawQueries: ... 185 ↛ exitline 185 didn't return from function '__aenter__' because
187 async def __aexit__(self, exc_type: object, exc_value: object, traceback: object) -> bool | None: ... 187 ↛ exitline 187 didn't return from function '__aexit__' because
190class _TransactionalDatabase(Protocol):
191 def tx(self, *, timeout: timedelta) -> _TransactionManager: ... 191 ↛ exitline 191 didn't return from function 'tx' because
194_COMPARISONS: Final = TypeAdapter(tuple[_Comparison, ...])
195_RECORDS: Final = TypeAdapter(tuple[_StoredRecord, ...])
196_HISTORY: Final = TypeAdapter(BaselineHistory)
197_PAGE_TIMESTAMPS: Final = 128
198_TRANSACTION_TIMEOUT: Final = timedelta(seconds=10)
200_CREATE_COMPARISON: Final = """
201INSERT INTO "LiteLLM_AutoRouterBaselineComparison"
202 (scope, api_key, session_id, router_name, initial_equivalent)
203VALUES ($1, $2, $3, $4, NOT EXISTS (
204 SELECT 1 FROM "LiteLLM_AutoRouterSession"
205 WHERE api_key = $2 AND session_id = $3 AND router_name = $4
206)) ON CONFLICT (scope) DO NOTHING
207"""
208_LOCK_COMPARISON: Final = """
209SELECT revision, published_revision, initial_equivalent, retired, history
210FROM "LiteLLM_AutoRouterBaselineComparison" WHERE scope = $1 FOR UPDATE
211"""
212_INSERT_RECORD: Final = """
213INSERT INTO "LiteLLM_AutoRouterBaselineObservation"
214 (request_id, scope, started_at, revision, data)
215VALUES ($1, $2, $3::float8, $4::bigint, $5)
216ON CONFLICT (request_id) DO NOTHING
217"""
218_MARK_CONFLICT: Final = """
219UPDATE "LiteLLM_AutoRouterBaselineObservation"
220SET conflicted = TRUE, revision = $4::bigint
221WHERE request_id = $1 AND scope = $2 AND data <> $3 AND NOT conflicted
222"""
223_READ_PAGE: Final = """
224WITH times AS (
225 SELECT DISTINCT started_at FROM "LiteLLM_AutoRouterBaselineObservation"
226 WHERE scope = $1 AND revision > $2::bigint
227 AND ($3::float8 IS NULL OR started_at > $3::float8)
228 AND ($5::float8 IS NULL OR (
229 started_at >= $5::float8 AND publication::jsonb->>'status' = 'estimated'
230 ))
231 ORDER BY started_at LIMIT $4::int
232)
233SELECT data, publication, conflicted, started_at
234FROM "LiteLLM_AutoRouterBaselineObservation"
235WHERE scope = $1 AND revision > $2::bigint
236 AND started_at IN (SELECT started_at FROM times)
237 AND ($5::float8 IS NULL OR publication::jsonb->>'status' = 'estimated')
238ORDER BY started_at, request_id
239"""
240_UPDATE_LOGS: Final = """
241WITH changes AS (
242 SELECT request_id, publication::jsonb AS publication
243 FROM jsonb_to_recordset($1::jsonb) AS x(request_id text, publication jsonb)
244)
245UPDATE "LiteLLM_SpendLogs" AS logs
246SET metadata = (COALESCE(logs.metadata::jsonb, '{}'::jsonb) - 'autorouter_baseline_observation') || jsonb_build_object(
247 'autorouter_savings_estimate', changes.publication,
248 'autorouter_savings', CASE WHEN changes.publication->>'status' = 'estimated' THEN
249 (changes.publication->>'baseline_spend')::float8 - (changes.publication->>'actual_spend')::float8
250 ELSE NULL END
251)
252FROM changes WHERE logs.request_id = changes.request_id
253"""
254_UPDATE_PUBLICATIONS: Final = """
255UPDATE "LiteLLM_AutoRouterBaselineObservation" AS observations
256SET publication = x.publication::text
257FROM jsonb_to_recordset($1::jsonb) AS x(request_id text, publication jsonb)
258WHERE observations.request_id = x.request_id
259"""
262def _session_correction_sql(*, user_scoped: bool) -> str:
263 table_name: Final = "LiteLLM_AutoRouterUserSession" if user_scoped else "LiteLLM_AutoRouterSession"
264 identity_columns: Final = ("user_id, " if user_scoped else "") + "api_key, session_id, router_name"
265 user_filter: Final = "WHERE user_id <> ''" if user_scoped else ""
266 user_match: Final = "session.user_id = totals.user_id AND " if user_scoped else ""
267 return f"""
268WITH changes AS (
269 SELECT * FROM jsonb_to_recordset($1::jsonb) AS x(
270 user_id text, api_key text, session_id text, router_name text, baseline_model text,
271 covered_delta int, actual_delta float8, savings_delta float8
272 )
273 {user_filter}
274), totals AS (
275 SELECT {identity_columns}, SUM(covered_delta)::int AS covered_delta,
276 SUM(actual_delta) AS actual_delta, SUM(savings_delta) AS savings_delta
277 FROM changes GROUP BY {identity_columns}
278), models AS (
279 SELECT {identity_columns}, jsonb_object_agg(baseline_model, delta) AS deltas
280 FROM (
281 SELECT {identity_columns}, baseline_model, SUM(covered_delta)::int AS delta
282 FROM changes GROUP BY {identity_columns}, baseline_model
283 ) grouped GROUP BY {identity_columns}
284)
285UPDATE "{table_name}" AS session
286SET saved_spend = session.saved_spend + totals.savings_delta,
287 savings_estimated_turns = session.savings_estimated_turns + totals.covered_delta,
288 savings_estimated_actual_spend = session.savings_estimated_actual_spend + totals.actual_delta,
289 savings_estimated_saved_spend = session.savings_estimated_saved_spend + totals.savings_delta,
290 savings_estimated_baseline_models = (
291 SELECT COALESCE(jsonb_object_agg(key, value), '{{}}'::jsonb) FROM (
292 SELECT key, SUM(value::int)::int AS value FROM (
293 SELECT * FROM jsonb_each_text(session.savings_estimated_baseline_models)
294 UNION ALL SELECT * FROM jsonb_each_text(models.deltas)
295 ) combined GROUP BY key HAVING SUM(value::int) > 0
296 ) counts
297 )
298FROM totals JOIN models USING ({identity_columns})
299WHERE {user_match}session.api_key = totals.api_key AND session.session_id = totals.session_id
300 AND session.router_name = totals.router_name
301"""
304_UPDATE_SESSIONS: Final = _session_correction_sql(user_scoped=False)
305_UPDATE_USER_SESSIONS: Final = _session_correction_sql(user_scoped=True)
308def _primary_transaction(client: PrismaClient) -> _TransactionManager:
309 primary: Final = cast(_TransactionalDatabase, writer_wrapper(client.db))
310 return primary.tx(timeout=_TRANSACTION_TIMEOUT)
313def _serialized(model: BaseModel) -> str:
314 return json.dumps(model.model_dump(mode="json"), sort_keys=True, separators=(",", ":"))
317def _change(record: BaselineAccountingRecord, old: BaselinePublication | None, new: BaselinePublication) -> _Change:
318 previous: Final = old.costs if old is not None else None
319 current: Final = new.costs
320 return _Change(
321 request_id=record.observation.request_id,
322 publication=new,
323 api_key=record.api_key,
324 user_id=record.turn.user_id if record.turn is not None else "",
325 session_id=record.session_id,
326 router_name=record.router_name,
327 baseline_model=record.baseline_model,
328 covered_delta=int(current is not None) - int(previous is not None),
329 actual_delta=(current.actual if current is not None else 0.0)
330 - (previous.actual if previous is not None else 0.0),
331 savings_delta=(current.savings if current is not None else 0.0)
332 - (previous.savings if previous is not None else 0.0),
333 daily=record.daily,
334 )
337def _project_group(
338 previous: tuple[BaselineHistory, tuple[_Change, ...]], stored: Sequence[_StoredRecord]
339) -> tuple[BaselineHistory, tuple[_Change, ...]]:
340 history, prior_changes = previous
341 records: Final = tuple(BaselineAccountingRecord.model_validate_json(item.data) for item in stored)
342 observations: Final = tuple(
343 record.observation.model_copy(
344 update=MappingProxyType(
345 {"outcome": "uncertain", "baseline_equivalent": False, "reason": "conflicting_observation"}
346 )
347 )
348 if row.conflicted
349 else record.observation
350 for record, row in zip(records, stored)
351 )
352 advanced, estimates = advance_baseline_history(history, observations)
353 publications: Final = tuple(
354 baseline_publication(
355 record, estimate, advanced.first_at if advanced.first_at is not None else observations[0].started_at
356 )
357 for record, estimate in zip(records, estimates)
358 )
359 changes: Final = tuple(
360 _change(record, old, publication)
361 for record, row, publication in zip(records, stored, publications)
362 for old in (BaselinePublication.model_validate_json(row.publication) if row.publication else None,)
363 if publication != old
364 )
365 return advanced, (*prior_changes, *changes)
368async def _publish(db: SupportsRawQueries, changes: Sequence[_Change]) -> None:
369 if not changes:
370 return
371 serialized: Final = json.dumps(tuple(change.model_dump(mode="json") for change in changes), separators=(",", ":"))
372 await db.execute_raw(_UPDATE_LOGS, serialized)
373 await db.execute_raw(_UPDATE_SESSIONS, serialized)
374 if any(change.user_id for change in changes):
375 await db.execute_raw(_UPDATE_USER_SESSIONS, serialized)
376 for entity, table in DAILY_SPEND_TABLES.items():
377 if adjustments := tuple(
378 change.daily.adjustment(target, change.savings_delta, change.request_id)
379 for change in changes
380 if change.daily is not None and change.savings_delta != 0
381 for target in change.daily.targets
382 if target.entity == entity
383 ):
384 statement, values = build_bulk_upsert(table, merge_by_conflict_key(table, adjustments))
385 await db.execute_raw(statement, *values)
386 await db.execute_raw(_UPDATE_PUBLICATIONS, serialized)
389class BaselineAccountingStore:
390 def __init__(self, transaction: Callable[[], _TransactionManager]) -> None:
391 self.transaction: Final = transaction
393 @classmethod
394 def for_client(cls, client: PrismaClient) -> BaselineAccountingStore:
395 def transaction() -> _TransactionManager:
396 return _primary_transaction(client)
398 return cls(transaction)
400 async def append(
401 self, record: BaselineAccountingRecord
402 ) -> Literal["recorded", "retired", "conflict", "unavailable"]:
403 try:
404 async with self.transaction() as db:
405 await db.execute_raw("SET LOCAL statement_timeout = 5000")
406 await db.execute_raw("SET LOCAL lock_timeout = 1000")
407 await db.execute_raw(
408 _CREATE_COMPARISON, record.scope, record.api_key, record.session_id, record.router_name
409 )
410 rows: Final = _COMPARISONS.validate_python(tuple(await db.query_raw(_LOCK_COMPARISON, record.scope)))
411 if not rows:
412 return "unavailable"
413 revision: Final = rows[0].revision + 1
414 data: Final = _serialized(record)
415 inserted: Final = await db.execute_raw(
416 _INSERT_RECORD,
417 record.observation.request_id,
418 record.scope,
419 record.observation.started_at,
420 revision,
421 data,
422 )
423 if inserted and record.turn is not None:
424 await write_autorouter_turn(db, record.turn)
425 conflicted: Final = (
426 0
427 if inserted
428 else await db.execute_raw(
429 _MARK_CONFLICT, record.observation.request_id, record.scope, data, revision
430 )
431 )
432 canonical: Final = (
433 _RECORDS.validate_python(
434 tuple(
435 await db.query_raw(
436 'SELECT data, publication, conflicted, started_at FROM "LiteLLM_AutoRouterBaselineObservation" '
437 "WHERE request_id=$1 AND scope=$2",
438 record.observation.request_id,
439 record.scope,
440 )
441 )
442 )
443 if not inserted
444 else ()
445 )
446 if not inserted and not canonical:
447 return "conflict"
448 if rows[0].retired:
449 await _publish(
450 db,
451 (
452 _change(
453 BaselineAccountingRecord.model_validate_json(canonical[0].data)
454 if canonical
455 else record,
456 BaselinePublication.model_validate_json(canonical[0].publication)
457 if canonical and canonical[0].publication is not None
458 else None,
459 BaselinePublication(
460 comparison_id=record.scope,
461 comparison_started_at=canonical[0].started_at
462 if canonical
463 else record.observation.started_at,
464 status="unknown",
465 reason="comparison_retired",
466 ),
467 ),
468 ),
469 )
470 return "retired"
471 if inserted or conflicted:
472 await self._withdraw(
473 db, record.scope, canonical[0].started_at if canonical else record.observation.started_at
474 )
475 await db.execute_raw(
476 'UPDATE "LiteLLM_AutoRouterBaselineComparison" SET revision = $2::bigint, '
477 "updated_at = CURRENT_TIMESTAMP, attempted_at = NULL WHERE scope = $1",
478 record.scope,
479 revision,
480 )
481 return "recorded"
482 except Exception: # noqa: BLE001 # accounting failure must not change inference or actual billing
483 verbose_proxy_logger.warning("Auto-router baseline observation could not be persisted")
484 return "unavailable"
486 async def _pages(
487 self, db: SupportsRawQueries, scope: str, after_revision: int, withdraw_from: float | None = None
488 ) -> AsyncIterator[tuple[_StoredRecord, ...]]:
489 cursor: float | None = None # rebind-ok: keyset pagination advances after each complete timestamp group
490 while page := _RECORDS.validate_python(
491 tuple(await db.query_raw(_READ_PAGE, scope, after_revision, cursor, _PAGE_TIMESTAMPS, withdraw_from))
492 ):
493 yield page
494 cursor = page[-1].started_at
496 async def _withdraw(self, db: SupportsRawQueries, scope: str, started_at: float) -> None:
497 async for page in self._pages(db, scope, 0, withdraw_from=started_at):
498 await _publish(
499 db,
500 tuple(
501 _change(
502 BaselineAccountingRecord.model_validate_json(row.data),
503 previous,
504 BaselinePublication(
505 comparison_id=scope,
506 comparison_started_at=min(previous.comparison_started_at, started_at),
507 status="unknown",
508 reason="pending_projection",
509 ),
510 )
511 for row in page
512 if row.publication is not None
513 for previous in (BaselinePublication.model_validate_json(row.publication),)
514 ),
515 )
517 async def retire_before(self, cutoff: datetime, batch_size: int, timeout_ms: int) -> None:
518 async with self.transaction() as db:
519 await db.execute_raw(f"SET LOCAL statement_timeout = {max(1, timeout_ms)}")
520 await db.execute_raw(f"SET LOCAL lock_timeout = {max(1, timeout_ms)}")
521 await db.execute_raw(
522 'WITH expired AS (SELECT scope FROM "LiteLLM_AutoRouterBaselineComparison" '
523 "WHERE NOT retired AND updated_at < $1::timestamptz ORDER BY updated_at "
524 "LIMIT $2::int FOR UPDATE SKIP LOCKED) "
525 'UPDATE "LiteLLM_AutoRouterBaselineComparison" AS comparison '
526 "SET retired=TRUE, history=NULL FROM expired WHERE comparison.scope=expired.scope",
527 cutoff,
528 batch_size,
529 )
530 await db.execute_raw(
531 'DELETE FROM "LiteLLM_AutoRouterBaselineObservation" WHERE request_id IN ('
532 'SELECT event.request_id FROM "LiteLLM_AutoRouterBaselineObservation" AS event '
533 'JOIN "LiteLLM_AutoRouterBaselineComparison" AS comparison USING (scope) '
534 "WHERE comparison.retired AND comparison.updated_at < $1::timestamptz "
535 "LIMIT $2::int)",
536 cutoff,
537 batch_size,
538 )
540 async def project(self, scope: str) -> Literal["published", "unchanged", "unavailable"]:
541 try:
542 async with self.transaction() as db:
543 await db.execute_raw("SET LOCAL statement_timeout = 5000")
544 await db.execute_raw("SET LOCAL lock_timeout = 1000")
545 rows: Final = _COMPARISONS.validate_python(tuple(await db.query_raw(_LOCK_COMPARISON, scope)))
546 if not rows or rows[0].retired or rows[0].revision == rows[0].published_revision:
547 return "unchanged"
548 missing_log: Final = await db.query_raw(
549 'SELECT 1 FROM "LiteLLM_AutoRouterBaselineObservation" AS observation '
550 'WHERE scope=$1 AND publication IS NULL AND NOT EXISTS (SELECT 1 FROM "LiteLLM_SpendLogs" AS log '
551 "WHERE log.request_id=observation.request_id) LIMIT 1",
552 scope,
553 )
554 if missing_log:
555 return "unavailable"
556 state: Final = rows[0]
557 checkpoint: Final = (
558 _HISTORY.validate_json(state.history)
559 if state.history is not None
560 else BaselineHistory(equivalent=state.initial_equivalent)
561 )
562 changed: Final = await db.query_raw(
563 'SELECT 1 FROM "LiteLLM_AutoRouterBaselineObservation" '
564 "WHERE scope = $1 AND revision > $2::bigint AND started_at <= $3::float8 LIMIT 1",
565 scope,
566 state.published_revision,
567 checkpoint.last_at,
568 )
569 history = BaselineHistory(equivalent=state.initial_equivalent) if changed else checkpoint
570 async for page in self._pages(db, scope, 0 if changed else state.published_revision):
571 history, updates = reduce(
572 _project_group,
573 (tuple(group) for _, group in groupby(page, key=lambda item: item.started_at)),
574 (history, ()),
575 )
576 await _publish(db, updates)
577 await db.execute_raw(
578 'UPDATE "LiteLLM_AutoRouterBaselineComparison" '
579 "SET published_revision = revision, history = $2 WHERE scope = $1",
580 scope,
581 _HISTORY.dump_json(history).decode(),
582 )
583 return "published"
584 except Exception: # noqa: BLE001 # rollback leaves the durable revision dirty for a later flush
585 verbose_proxy_logger.warning("Auto-router baseline projection remains pending")
586 return "unavailable"
589class _Scope(BaseModel):
590 scope: str
593_SCOPES: Final = TypeAdapter(tuple[_Scope, ...])
594_CLAIM_DIRTY: Final = """
595WITH candidates AS (
596 SELECT scope FROM "LiteLLM_AutoRouterBaselineComparison"
597 WHERE NOT retired AND revision <> published_revision
598 AND (attempted_at IS NULL OR attempted_at < CURRENT_TIMESTAMP - INTERVAL '30 seconds')
599 ORDER BY attempted_at NULLS FIRST, updated_at, scope LIMIT 32 FOR UPDATE SKIP LOCKED
600)
601UPDATE "LiteLLM_AutoRouterBaselineComparison" AS comparison
602SET attempted_at = CURRENT_TIMESTAMP FROM candidates
603WHERE comparison.scope = candidates.scope RETURNING comparison.scope
604"""
607async def _flush_records(
608 store: BaselineAccountingStore, records: Sequence[BaselineAccountingRecord]
609) -> tuple[BaselineAccountingRecord, ...]:
610 slots: Final = asyncio.Semaphore(4)
612 async def append(record: BaselineAccountingRecord) -> bool:
613 async with slots:
614 return await store.append(record) == "unavailable"
616 failed: Final = await asyncio.gather(*(append(record) for record in records))
617 return tuple(record for record, retry in zip(records, failed) if retry)
620async def flush_baseline_accounting(client: PrismaClient) -> None:
621 from litellm.proxy.utils import request_spend_log_flush
623 store: Final = BaselineAccountingStore.for_client(client)
624 async with client.baseline_accounting_lock:
625 batch: Final = tuple(client.baseline_accounting_transactions[:32])
626 client.baseline_accounting_transactions = client.baseline_accounting_transactions[32:]
627 more_queued: Final = bool(client.baseline_accounting_transactions)
628 try:
629 remaining: Final = await asyncio.wait_for(_flush_records(store, batch), timeout=5)
630 except (Exception, asyncio.CancelledError) as error:
631 async with client.baseline_accounting_lock:
632 client.baseline_accounting_transactions.extend(batch)
633 if isinstance(error, asyncio.CancelledError):
634 raise
635 return
636 async with client.baseline_accounting_lock:
637 client.baseline_accounting_transactions.extend(remaining)
638 if more_queued and len(remaining) < len(batch): 638 ↛ 639line 638 didn't jump to line 639 because the condition on line 638 was never true
639 request_spend_log_flush(client)
640 try:
641 async with store.transaction() as db:
642 await db.execute_raw("SET LOCAL statement_timeout = 1000")
643 scopes: Final = _SCOPES.validate_python(tuple(await db.query_raw(_CLAIM_DIRTY)))
644 slots: Final = asyncio.Semaphore(4)
646 async def project(item: _Scope) -> str:
647 async with slots:
648 return await store.project(item.scope)
650 outcomes: Final = await asyncio.wait_for(asyncio.gather(*(project(item) for item in scopes)), timeout=5)
651 if len(scopes) == 32 and "published" in outcomes: 651 ↛ 652line 651 didn't jump to line 652 because the condition on line 651 was never true
652 request_spend_log_flush(client)
653 except Exception: # noqa: BLE001 # durable dirty comparisons remain eligible after the retry interval
654 verbose_proxy_logger.warning("Auto-router baseline projection will retry on a later spend flush")