Coverage for .venv/lib/python3.13/site-packages/litellm/proxy/db/autorouter_session_rollup.py: 43%
182 statements
« prev ^ index » next coverage.py v7.15.2, created at 2026-10-10 12:01 +0000
1"""
2Per-session auto-router benchmarks rollup.
4At request time the spend writer builds one AutoRouterTurnTransaction per successful
5auto-routed request (a request whose metadata carries a routing_decision) and queues it
6on the prisma client. The spend-log flush job drains the queue into
7key and user session rollups with one atomic statement per turn: each upsert classifies
8the turn (same model, first visit, return to a model the session already used, out of
9order) against the row's own columns, so nothing is read before the write and concurrent
10pods compose. The benchmarks endpoint aggregates these rows and never touches
11LiteLLM_SpendLogs.
12"""
14from __future__ import annotations
16import asyncio
17import dataclasses
18import hashlib
19import random
20from collections.abc import Mapping, Sequence
21from dataclasses import dataclass
22from datetime import datetime, timezone
23from itertools import groupby
24from typing import TYPE_CHECKING, Final, NamedTuple
26from litellm._logging import verbose_proxy_logger
27from litellm.constants import INTERNAL_CALL_ORIGIN_METADATA_KEY
28from litellm.proxy._types import DB_RETRY_SAFE_ERROR_TYPES
29from litellm.proxy.db.create_views import SupportsExecuteRaw
31if TYPE_CHECKING: 31 ↛ 32line 31 didn't jump to line 32 because the condition on line 31 was never true
32 from litellm.proxy._types import SpendLogsPayload
33 from litellm.proxy.utils import PrismaClient
35CACHE_TTL_5M_SECONDS: Final = 300
36CACHE_TTL_1H_SECONDS: Final = 3600
38_SESSION_COLUMNS: Final = """
39 api_key, session_id, router_name, router_type, first_turn_at, last_turn_at,
40 last_model, models, turns, unordered_turns, covered_turns, cache_hits,
41 same_model_turns, same_model_hits, first_visit_turns, first_visit_hits,
42 return_turns, return_hits, return_expired_misses, return_within_ttl_misses,
43 ttl_5m_turns, ttl_1h_turns, total_tokens, spend, saved_spend, classifier_cost, classifier_cost_recorded_turns, tier_turns,
44 baseline_models, savings_estimated_turns, savings_estimated_actual_spend, savings_estimated_saved_spend,
45 savings_estimated_baseline_models
46"""
48AUTOROUTER_BENCHMARKS_SQL: Final = f"""
49WITH windowed AS (
50 SELECT {_SESSION_COLUMNS} FROM "LiteLLM_AutoRouterSession"
51 WHERE $4::text IS NULL
52 AND last_turn_at >= $1::timestamp
53 AND first_turn_at < $2::timestamp
54 AND ($3::text IS NULL OR api_key = $3::text)
55 UNION ALL
56 SELECT {_SESSION_COLUMNS} FROM "LiteLLM_AutoRouterUserSession"
57 WHERE (($4::text IS NOT NULL AND user_id = $4::text) OR ($4::text IS NULL AND api_key = ''))
58 AND last_turn_at >= $1::timestamp
59 AND first_turn_at < $2::timestamp
60 AND ($3::text IS NULL OR api_key = $3::text)
61),
62tier_maps AS (
63 SELECT router_name, router_type, jsonb_object_agg(tier, tier_turns) AS tier_turns
64 FROM (
65 SELECT router_name, router_type, kv.key AS tier, SUM((kv.value)::int)::int AS tier_turns
66 FROM windowed, LATERAL jsonb_each_text(tier_turns) AS kv
67 GROUP BY router_name, router_type, kv.key
68 ) per_tier
69 GROUP BY router_name, router_type
70)
71SELECT
72 agg.*,
73 COALESCE(tier_maps.tier_turns, '{{}}'::jsonb) AS tier_turns
74FROM (
75SELECT
76 router_name,
77 router_type,
78 COUNT(*)::int AS sessions,
79 COALESCE(SUM(turns), 0)::int AS turns,
80 COALESCE(SUM(unordered_turns), 0)::int AS unordered_turns,
81 COALESCE(SUM(covered_turns), 0)::int AS covered_turns,
82 COALESCE(SUM(cache_hits), 0)::int AS cache_hits,
83 COALESCE(SUM(same_model_turns), 0)::int AS same_model_turns,
84 COALESCE(SUM(same_model_hits), 0)::int AS same_model_hits,
85 COALESCE(SUM(first_visit_turns), 0)::int AS first_visit_turns,
86 COALESCE(SUM(first_visit_hits), 0)::int AS first_visit_hits,
87 COALESCE(SUM(return_turns), 0)::int AS return_turns,
88 COALESCE(SUM(return_hits), 0)::int AS return_hits,
89 COALESCE(SUM(return_expired_misses), 0)::int AS return_expired_misses,
90 COALESCE(SUM(return_within_ttl_misses), 0)::int AS return_within_ttl_misses,
91 COALESCE(SUM(ttl_5m_turns), 0)::int AS ttl_5m_turns,
92 COALESCE(SUM(ttl_1h_turns), 0)::int AS ttl_1h_turns,
93 COALESCE(SUM(total_tokens), 0)::bigint AS total_tokens,
94 COALESCE(SUM(spend), 0)::float8 AS spend,
95 COALESCE(SUM(saved_spend), 0)::float8 AS saved_spend,
96 COALESCE(SUM(savings_estimated_turns), 0)::int AS savings_estimated_turns,
97 COALESCE(SUM(savings_estimated_actual_spend), 0)::float8 AS savings_estimated_actual_spend,
98 COALESCE(SUM(savings_estimated_saved_spend), 0)::float8 AS savings_estimated_saved_spend,
99 COALESCE(SUM(classifier_cost), 0)::float8 AS classifier_cost,
100 COALESCE(SUM(classifier_cost_recorded_turns), 0)::int AS classifier_cost_recorded_turns,
101 COALESCE(SUM(EXTRACT(EPOCH FROM (last_turn_at - first_turn_at))), 0)::float8 AS session_seconds
102FROM windowed
103GROUP BY router_name, router_type
104) agg
105LEFT JOIN tier_maps USING (router_name, router_type)
106ORDER BY agg.spend DESC
107"""
110@dataclass(frozen=True, slots=True)
111class AutoRouterTurnTransaction:
112 api_key: str
113 session_id: str
114 router_name: str
115 router_type: str
116 model: str
117 turn_at: datetime
118 total_tokens: int
119 spend: float
120 saved_spend: float
121 classifier_cost: float
122 covered: bool
123 cache_hit: bool
124 cache_ttl_seconds: int | None
125 cache_touched: bool
126 tier: str | None = None
127 baseline_model: str | None = None
128 savings_estimated_turns: int = 0
129 savings_estimated_actual_spend: float = 0.0
130 savings_estimated_saved_spend: float = 0.0
131 user_id: str = ""
134class TurnCacheFacts(NamedTuple):
135 """One statement of a turn's cache interaction, derived from its usage record.
137 ``touched`` is False only when telemetry positively shows the provider neither
138 read from nor wrote to the cache; absent telemetry reads as touched, which is
139 the conservative input for the per-model idle clock.
140 """
142 covered: bool
143 read_tokens: int
144 write_ttl_seconds: int | None
145 touched: bool
148def turn_cache_facts(usage_object: Mapping[str, object] | None) -> TurnCacheFacts:
149 from litellm.proxy.spend_tracking.savings import extract_cache_read_tokens
151 covered: Final = bool(usage_object)
152 read_tokens: Final = extract_cache_read_tokens(usage_object)
153 write_ttl_seconds: Final = _write_ttl_seconds(usage_object)
154 return TurnCacheFacts(
155 covered=covered,
156 read_tokens=read_tokens,
157 write_ttl_seconds=write_ttl_seconds,
158 touched=not covered or read_tokens > 0 or write_ttl_seconds is not None,
159 )
162def _turn_time_utc(start_time_iso: str) -> datetime | None:
163 try:
164 parsed: Final = datetime.fromisoformat(start_time_iso.replace("Z", "+00:00"))
165 except ValueError:
166 return None
167 if parsed.tzinfo is None:
168 return parsed
169 return parsed.astimezone(timezone.utc).replace(tzinfo=None)
172def _write_ttl_seconds(usage_object: Mapping[str, object] | None) -> int | None:
173 """The TTL this turn's cache write used, or None when nothing was written.
175 Providers that report a TTL split do so under prompt_tokens_details; a write with no
176 split is the provider's default five-minute cache.
177 """
178 from litellm.proxy.spend_tracking.savings import extract_cache_creation_tokens
180 if not usage_object:
181 return None
182 details: Final = usage_object.get("prompt_tokens_details")
183 creation: Final = details.get("cache_creation_token_details") if isinstance(details, Mapping) else None
184 if isinstance(creation, Mapping):
185 if creation.get("ephemeral_1h_input_tokens"):
186 return CACHE_TTL_1H_SECONDS
187 if creation.get("ephemeral_5m_input_tokens"):
188 return CACHE_TTL_5M_SECONDS
189 if extract_cache_creation_tokens(usage_object) > 0:
190 return CACHE_TTL_5M_SECONDS
191 return None
194SESSION_ID_MAX_CHARS: Final = 256
197def bounded_session_id(session_id: str) -> str:
198 """The session id as stored, bounded so a caller-chosen identifier cannot exceed
199 Postgres's B-tree index entry limit through the composite primary key. Oversized
200 ids map to a stable digest, so their turns still aggregate into one session."""
201 if len(session_id) <= SESSION_ID_MAX_CHARS: 201 ↛ 203line 201 didn't jump to line 203 because the condition on line 201 was always true
202 return session_id
203 return "sha256:" + hashlib.sha256(session_id.encode("utf-8", errors="surrogatepass")).hexdigest()
206def build_autorouter_turn_transaction(
207 payload: SpendLogsPayload,
208 metadata: Mapping[str, object],
209 saved_spend: float,
210) -> AutoRouterTurnTransaction | None:
211 """One rollup transaction for a successful auto-routed turn, else None.
213 The routing_decision record is what says a request was auto-routed at all, so a
214 request without one (including the auto-router's own classifier sub-calls) never
215 reaches the rollup. Internal sub-calls that DO carry one (a shadow eval's duplicate
216 of a request through the router) are excluded by their internal_call_origin stamp:
217 they are not traffic a user sent, so counting them would manufacture sessions and
218 savings in the adoption metrics. Failed requests served nothing and are excluded.
219 The classifier's charge still lands here exactly once, via the decision's own
220 classifier_cost folded into this turn's spend: the excluded classifier row is how
221 it was billed, the decision is how it is attributed. Cache facts are derived from
222 the payload's own usage record through the savings owner, never handed in beside it.
223 The baseline the turn's saved_spend was priced against travels with the turn, so the
224 row can name the counterfactual for the money it holds even after the router is
225 reconfigured or removed.
226 """
227 if payload.get("status") != "success":
228 return None
229 if metadata.get(INTERNAL_CALL_ORIGIN_METADATA_KEY):
230 return None
231 routing_decision: Final = metadata.get("routing_decision")
232 if not isinstance(routing_decision, Mapping) or not routing_decision:
233 return None
234 router_name: Final = routing_decision.get("router_model_name") or payload.get("model_group")
235 api_key: Final = payload.get("api_key") or ""
236 user_id: Final = payload.get("user") or ""
237 session_id: Final = payload.get("session_id")
238 model: Final = payload.get("model")
239 if not (isinstance(router_name, str) and router_name and (api_key or user_id) and session_id and model):
240 return None
241 turn_at: Final = _turn_time_utc(str(payload.get("startTime") or ""))
242 if turn_at is None:
243 return None
244 from litellm.proxy.spend_tracking.savings import (
245 classifier_cost_from_decision,
246 recorded_estimated_autorouter_savings,
247 )
249 usage_object_raw: Final = metadata.get("usage_object")
250 cache: Final = turn_cache_facts(usage_object_raw if isinstance(usage_object_raw, Mapping) else None)
251 tier_raw: Final = routing_decision.get("tier")
252 baseline_raw: Final = routing_decision.get("savings_baseline_model")
253 classifier_cost: Final = classifier_cost_from_decision(routing_decision)
254 actual_spend: Final = float(payload.get("spend") or 0.0) + (classifier_cost or 0.0)
255 estimated_savings: Final = recorded_estimated_autorouter_savings(metadata)
256 return AutoRouterTurnTransaction(
257 api_key=api_key,
258 user_id=user_id,
259 session_id=bounded_session_id(session_id),
260 router_name=router_name,
261 router_type=str(routing_decision.get("router_type") or "unknown"),
262 tier=tier_raw if isinstance(tier_raw, str) and tier_raw else None,
263 baseline_model=baseline_raw if isinstance(baseline_raw, str) and baseline_raw else None,
264 model=model,
265 turn_at=turn_at,
266 total_tokens=int(payload.get("prompt_tokens") or 0) + int(payload.get("completion_tokens") or 0),
267 spend=actual_spend,
268 saved_spend=saved_spend,
269 classifier_cost=classifier_cost or 0.0,
270 covered=cache.covered,
271 cache_hit=cache.read_tokens > 0,
272 cache_ttl_seconds=cache.write_ttl_seconds,
273 cache_touched=cache.touched,
274 savings_estimated_turns=int(estimated_savings is not None),
275 savings_estimated_actual_spend=actual_spend if estimated_savings is not None else 0.0,
276 savings_estimated_saved_spend=estimated_savings if estimated_savings is not None else 0.0,
277 )
280_UPSERT_PARAM_FIELDS: Final = tuple(field.name for field in dataclasses.fields(AutoRouterTurnTransaction))
283def _p(field_name: str) -> str:
284 """Positional placeholder for a transaction field, numbered by the dataclass's own
285 field order so the SQL and the argument tuple cannot disagree; typos fail at import."""
286 return f"${_UPSERT_PARAM_FIELDS.index(field_name) + 1}"
289_MODEL: Final = _p("model")
290_TURN_AT: Final = _p("turn_at")
291_COVERED: Final = _p("covered")
292_CACHE_HIT: Final = _p("cache_hit")
293_CACHE_TTL: Final = _p("cache_ttl_seconds")
294_TOUCHED: Final = _p("cache_touched")
295_TIER: Final = f"{_p('tier')}::text"
296_TIER_DELTA: Final = f"(CASE WHEN {_TIER} IS NULL THEN '{{}}'::jsonb ELSE jsonb_build_object({_TIER}, 1) END)"
297_BASELINE: Final = f"{_p('baseline_model')}::text"
298_BASELINE_DELTA: Final = (
299 f"(CASE WHEN {_BASELINE} IS NULL THEN '{{}}'::jsonb ELSE jsonb_build_object({_BASELINE}, 1) END)"
300)
301_ESTIMATED_BASELINE: Final = f"{_p('savings_estimated_turns')}::int = 1 AND {_BASELINE} IS NOT NULL"
302_ESTIMATED_BASELINE_DELTA: Final = (
303 f"(CASE WHEN {_ESTIMATED_BASELINE} THEN jsonb_build_object({_BASELINE}, 1) ELSE '{{}}'::jsonb END)"
304)
306_IN_ORDER: Final = f"{_TURN_AT}::timestamp >= t.last_turn_at"
307_SAME: Final = f"{_IN_ORDER} AND t.last_model = {_MODEL}"
308_FIRST: Final = f"{_IN_ORDER} AND NOT t.models ? {_MODEL}"
309_RETURN: Final = f"{_IN_ORDER} AND t.models ? {_MODEL} AND t.last_model <> {_MODEL}"
310_RETURN_MISS: Final = (
311 f"{_RETURN} AND {_COVERED}::int = 1 AND {_CACHE_HIT}::int = 0 AND (t.models -> {_MODEL} ->> 'ttl') IS NOT NULL"
312)
313_IDLE_SECONDS: Final = f"EXTRACT(EPOCH FROM {_TURN_AT}::timestamp) - (t.models -> {_MODEL} ->> 'at')::float8"
314_CACHE_TOUCHED: Final = f"{_TOUCHED}::int = 1"
317def _session_upsert_sql(*, user_scoped: bool) -> str:
318 table_name: Final = "LiteLLM_AutoRouterUserSession" if user_scoped else "LiteLLM_AutoRouterSession"
319 user_column: Final = "user_id, " if user_scoped else ""
320 user_value: Final = f"{_p('user_id')}::text, " if user_scoped else ""
321 required_identity: Final = _p("user_id" if user_scoped else "api_key")
322 return f"""
323INSERT INTO "{table_name}" AS t (
324 {user_column}{_SESSION_COLUMNS}
325)
326SELECT
327 {user_value}{_p("api_key")}, {_p("session_id")}, {_p("router_name")}, {_p("router_type")}, {_TURN_AT}::timestamp, {_TURN_AT}::timestamp,
328 {_MODEL}, jsonb_build_object({_MODEL}, jsonb_build_object('at', EXTRACT(EPOCH FROM {_TURN_AT}::timestamp), 'ttl', {_CACHE_TTL}::int)),
329 1, 0, {_COVERED}::int, {_CACHE_HIT}::int,
330 0, 0, 1, {_CACHE_HIT}::int,
331 0, 0, 0, 0,
332 (CASE WHEN {_CACHE_TTL}::int = {CACHE_TTL_5M_SECONDS} THEN 1 ELSE 0 END),
333 (CASE WHEN {_CACHE_TTL}::int = {CACHE_TTL_1H_SECONDS} THEN 1 ELSE 0 END),
334 {_p("total_tokens")}::bigint, {_p("spend")}::float8, {_p("saved_spend")}::float8,
335 {_p("classifier_cost")}::float8, 1, {_TIER_DELTA}, {_BASELINE_DELTA},
336 {_p("savings_estimated_turns")}::int, {_p("savings_estimated_actual_spend")}::float8,
337 {_p("savings_estimated_saved_spend")}::float8, {_ESTIMATED_BASELINE_DELTA}
338WHERE {required_identity}::text <> ''
339ON CONFLICT ({user_column}api_key, session_id, router_name) DO UPDATE SET
340 turns = t.turns + 1,
341 total_tokens = t.total_tokens + EXCLUDED.total_tokens,
342 spend = t.spend + EXCLUDED.spend,
343 saved_spend = t.saved_spend + EXCLUDED.saved_spend,
344 savings_estimated_turns = t.savings_estimated_turns + EXCLUDED.savings_estimated_turns,
345 savings_estimated_actual_spend = t.savings_estimated_actual_spend + EXCLUDED.savings_estimated_actual_spend,
346 savings_estimated_saved_spend = t.savings_estimated_saved_spend + EXCLUDED.savings_estimated_saved_spend,
347 classifier_cost = t.classifier_cost + EXCLUDED.classifier_cost,
348 classifier_cost_recorded_turns = t.classifier_cost_recorded_turns + 1,
349 covered_turns = t.covered_turns + EXCLUDED.covered_turns,
350 cache_hits = t.cache_hits + EXCLUDED.cache_hits,
351 ttl_5m_turns = t.ttl_5m_turns + EXCLUDED.ttl_5m_turns,
352 ttl_1h_turns = t.ttl_1h_turns + EXCLUDED.ttl_1h_turns,
353 unordered_turns = t.unordered_turns + (CASE WHEN NOT ({_IN_ORDER}) THEN 1 ELSE 0 END),
354 same_model_turns = t.same_model_turns + (CASE WHEN {_SAME} THEN 1 ELSE 0 END),
355 same_model_hits = t.same_model_hits + (CASE WHEN {_SAME} AND {_CACHE_HIT}::int = 1 THEN 1 ELSE 0 END),
356 first_visit_turns = t.first_visit_turns + (CASE WHEN {_FIRST} THEN 1 ELSE 0 END),
357 first_visit_hits = t.first_visit_hits + (CASE WHEN {_FIRST} AND {_CACHE_HIT}::int = 1 THEN 1 ELSE 0 END),
358 return_turns = t.return_turns + (CASE WHEN {_RETURN} THEN 1 ELSE 0 END),
359 return_hits = t.return_hits + (CASE WHEN {_RETURN} AND {_CACHE_HIT}::int = 1 THEN 1 ELSE 0 END),
360 return_expired_misses = t.return_expired_misses
361 + (CASE WHEN {_RETURN_MISS} AND {_IDLE_SECONDS} > (t.models -> {_MODEL} ->> 'ttl')::float8 THEN 1 ELSE 0 END),
362 return_within_ttl_misses = t.return_within_ttl_misses
363 + (CASE WHEN {_RETURN_MISS} AND {_IDLE_SECONDS} <= (t.models -> {_MODEL} ->> 'ttl')::float8 THEN 1 ELSE 0 END),
364 models = t.models || jsonb_build_object({_MODEL}, jsonb_build_object(
365 'at', (CASE WHEN {_CACHE_TOUCHED}
366 THEN GREATEST(COALESCE((t.models -> {_MODEL} ->> 'at')::float8, 0), EXTRACT(EPOCH FROM {_TURN_AT}::timestamp))
367 ELSE COALESCE((t.models -> {_MODEL} ->> 'at')::float8, EXTRACT(EPOCH FROM {_TURN_AT}::timestamp)) END),
368 'ttl', (CASE WHEN {_IN_ORDER}
369 THEN COALESCE({_CACHE_TTL}::int, (t.models -> {_MODEL} ->> 'ttl')::int)
370 ELSE COALESCE((t.models -> {_MODEL} ->> 'ttl')::int, {_CACHE_TTL}::int) END)
371 )),
372 last_model = (CASE WHEN {_IN_ORDER} THEN {_MODEL} ELSE t.last_model END),
373 tier_turns = (CASE WHEN {_TIER} IS NOT NULL AND t.router_type = {_p("router_type")}
374 THEN t.tier_turns || jsonb_build_object({_TIER}, COALESCE((t.tier_turns ->> {_TIER})::int, 0) + 1)
375 ELSE t.tier_turns END),
376 baseline_models = (CASE WHEN {_BASELINE} IS NOT NULL
377 THEN t.baseline_models || jsonb_build_object({_BASELINE}, COALESCE((t.baseline_models ->> {_BASELINE})::int, 0) + 1)
378 ELSE t.baseline_models END),
379 savings_estimated_baseline_models = (CASE WHEN {_ESTIMATED_BASELINE}
380 THEN t.savings_estimated_baseline_models || jsonb_build_object(
381 {_BASELINE}, COALESCE((t.savings_estimated_baseline_models ->> {_BASELINE})::int, 0) + 1)
382 ELSE t.savings_estimated_baseline_models END),
383 first_turn_at = LEAST(t.first_turn_at, EXCLUDED.first_turn_at),
384 last_turn_at = GREATEST(t.last_turn_at, EXCLUDED.last_turn_at)
385"""
388UPSERT_AUTOROUTER_SESSION_SQL: Final = f"""
389WITH key_rollup AS (
390 {_session_upsert_sql(user_scoped=False)}
391 RETURNING 1
392)
393{_session_upsert_sql(user_scoped=True)}
394"""
396UPSERT_AUTOROUTER_USER_SESSION_SQL: Final = _session_upsert_sql(user_scoped=True)
399def _as_sql_param(value: str | float | bool | datetime | None) -> str | float | None:
400 if isinstance(value, bool):
401 return int(value)
402 if isinstance(value, datetime):
403 return value.isoformat()
404 return value
407def _upsert_params(transaction: AutoRouterTurnTransaction) -> tuple[str | float | None, ...]:
408 return tuple(_as_sql_param(getattr(transaction, name)) for name in _UPSERT_PARAM_FIELDS)
411async def write_autorouter_turn(
412 db: SupportsExecuteRaw,
413 transaction: AutoRouterTurnTransaction,
414 statement: str = UPSERT_AUTOROUTER_SESSION_SQL,
415) -> None:
416 await db.execute_raw(statement, *_upsert_params(transaction))
419async def _upsert_turn_with_retry(
420 prisma_client: PrismaClient,
421 transaction: AutoRouterTurnTransaction,
422 n_retry_times: int,
423 statement: str,
424) -> None:
425 for attempt in range(n_retry_times + 1):
426 try:
427 await write_autorouter_turn(prisma_client.db, transaction, statement)
428 except DB_RETRY_SAFE_ERROR_TYPES:
429 if attempt >= n_retry_times:
430 raise
431 await asyncio.sleep(2**attempt + random.uniform(0, 1))
432 else:
433 return
436def _session_partition(transaction: AutoRouterTurnTransaction) -> tuple[str, str, str, str]:
437 identity: Final = ("key", transaction.api_key) if transaction.api_key else ("user", transaction.user_id)
438 return (*identity, transaction.session_id, transaction.router_name)
441async def _drain_session_partition(
442 prisma_client: PrismaClient,
443 transactions: tuple[AutoRouterTurnTransaction, ...],
444 n_retry_times: int,
445 statement: str,
446) -> tuple[AutoRouterTurnTransaction, ...]:
447 for position, transaction in enumerate(transactions):
448 try:
449 await _upsert_turn_with_retry(prisma_client, transaction, n_retry_times, statement)
450 except Exception as flush_err: # noqa: BLE001 # stop dependent turns without retrying an ambiguous write
451 verbose_proxy_logger.error(
452 "Spend tracking - auto-router session rollup flush failed for router %s; "
453 "%s of %s turn writes stopped in this partition: %s",
454 transaction.router_name,
455 len(transactions) - position,
456 len(transactions),
457 flush_err,
458 )
459 return transactions[position:]
460 return ()
463async def _flush_session_partition(
464 prisma_client: PrismaClient,
465 transactions: tuple[AutoRouterTurnTransaction, ...],
466 n_retry_times: int,
467) -> None:
468 failed_suffix: Final = await _drain_session_partition(
469 prisma_client, transactions, n_retry_times, UPSERT_AUTOROUTER_SESSION_SQL
470 )
471 if not failed_suffix or not failed_suffix[0].api_key:
472 return
473 failed_user: Final = failed_suffix[0].user_id
474 other_users: Final = sorted(
475 (
476 transaction
477 for transaction in failed_suffix[1:]
478 if transaction.user_id and transaction.user_id != failed_user
479 ),
480 key=lambda transaction: transaction.user_id,
481 )
482 for _, user_turns in groupby(other_users, key=lambda transaction: transaction.user_id):
483 await _drain_session_partition(
484 prisma_client, tuple(user_turns), n_retry_times, UPSERT_AUTOROUTER_USER_SESSION_SQL
485 )
488async def flush_autorouter_turn_transactions(
489 prisma_client: PrismaClient,
490 transactions: Sequence[AutoRouterTurnTransaction],
491 n_retry_times: int = 3,
492) -> None:
493 """Drain a queue batch into the rollup, one upsert per turn.
495 Statements run sequentially in per-session event order: a turn's classification
496 depends on the turns before it, and Postgres rejects one multi-row INSERT touching
497 the same key twice. Only ConnectError is retried, per statement, because it proves
498 that statement never reached the database. A failed write stops its key and user
499 histories for this batch. Other users sharing that key can still advance their
500 independent user histories, with the key projection disabled and the real key
501 identity preserved. The failed turn is never replayed. Callers must not add their
502 own retry around this function.
503 """
504 if not transactions: 504 ↛ 506line 504 didn't jump to line 506 because the condition on line 504 was always true
505 return
506 ordered: Final = sorted(
507 transactions,
508 key=lambda transaction: (*_session_partition(transaction), transaction.turn_at),
509 )
510 for _, session_group in groupby(
511 ordered,
512 key=_session_partition,
513 ):
514 await _flush_session_partition(prisma_client, tuple(session_group), n_retry_times)