« prev ^ index » next coverage.py v7.15.4, created at 2026-08-25 07:19 +0800
1"""Database-backed relay billing store.
3Persists reservations and settled usage records through
4:class:`~lexigram.contracts.data.DatabaseProviderProtocol` using only
5generic ``execute``/``execute_query`` SQL so the store works on any
6backend. Idempotency is enforced at the database level: settlement uses
7a unique ``(request_id, attempt_id)`` index combined with insert-or-ignore
8semantics, so retries and concurrent writes return the existing record
9instead of charging twice. Reservation state transitions are guarded by
10compare-and-set predicates (``reserved -> released``,
11``reserved -> settled``, ``reserved -> expired``) so a stale write cannot
12overwrite a newer state.
13"""
15from __future__ import annotations
17from decimal import Decimal
18from typing import TYPE_CHECKING, Any
20from lexigram.contracts.ai.governance import (
21 RelayUsageRecord,
22 RelayUsageScope,
23 RelayUsageStoreProtocol,
24)
25from lexigram.contracts.ai.relay import RelayUsage
26from lexigram.primitives import clock
27from lexigram.serialization import dumps_str, loads_str
29if TYPE_CHECKING:
30 from collections.abc import Mapping, Sequence
32 from lexigram.contracts.ai.governance import RelayUsageReservation
33 from lexigram.contracts.ai.relay import JsonValue
34 from lexigram.contracts.data import DatabaseProviderProtocol
36__all__ = ["DatabaseRelayUsageStore"]
38_CREATE_RESERVATIONS = """
39CREATE TABLE IF NOT EXISTS ai_relay_reservations (
40 reservation_id TEXT NOT NULL PRIMARY KEY,
41 request_id TEXT NOT NULL,
42 status TEXT NOT NULL DEFAULT 'reserved',
43 estimated_tokens INTEGER NOT NULL,
44 estimated_charge TEXT NOT NULL,
45 expires_at TEXT NOT NULL,
46 created_at TEXT NOT NULL
47)
48"""
50_CREATE_RESERVATIONS_INDEX = """
51CREATE INDEX IF NOT EXISTS idx_ai_relay_reservations_request
52 ON ai_relay_reservations (request_id)
53"""
55_CREATE_USAGE = """
56CREATE TABLE IF NOT EXISTS ai_relay_usage (
57 request_id TEXT NOT NULL,
58 attempt_id TEXT NOT NULL,
59 tenant_id TEXT NOT NULL,
60 account_id TEXT,
61 user_id TEXT,
62 model TEXT NOT NULL,
63 provider TEXT NOT NULL DEFAULT '',
64 channel TEXT NOT NULL DEFAULT '',
65 status TEXT NOT NULL,
66 charge TEXT NOT NULL,
67 currency TEXT NOT NULL,
68 converter_id TEXT,
69 metadata TEXT NOT NULL DEFAULT '{}',
70 created_at TEXT NOT NULL,
71 prompt_tokens INTEGER NOT NULL DEFAULT 0,
72 completion_tokens INTEGER NOT NULL DEFAULT 0,
73 cache_read_tokens INTEGER NOT NULL DEFAULT 0,
74 cache_creation_tokens INTEGER NOT NULL DEFAULT 0,
75 reasoning_tokens INTEGER NOT NULL DEFAULT 0,
76 audio_input_tokens INTEGER NOT NULL DEFAULT 0,
77 audio_output_tokens INTEGER NOT NULL DEFAULT 0,
78 image_tokens INTEGER NOT NULL DEFAULT 0,
79 input_tokens INTEGER NOT NULL DEFAULT 0,
80 output_tokens INTEGER NOT NULL DEFAULT 0,
81 total_tokens_override INTEGER,
82 PRIMARY KEY (request_id, attempt_id)
83)
84"""
86_CREATE_USAGE_INDEX = """
87CREATE INDEX IF NOT EXISTS idx_ai_relay_usage_scope
88 ON ai_relay_usage (tenant_id, account_id, user_id, model, created_at)
89"""
91_INSERT_RESERVATION = (
92 "INSERT OR IGNORE INTO ai_relay_reservations "
93 "(reservation_id, request_id, status, estimated_tokens, estimated_charge, "
94 "expires_at, created_at) "
95 "VALUES (?, ?, 'reserved', ?, ?, ?, ?)"
96)
98_EXPIRE_RESERVATIONS = (
99 "UPDATE ai_relay_reservations SET status = 'expired' "
100 "WHERE status = 'reserved' AND expires_at <= ?"
101)
103_RELEASE_RESERVATION = (
104 "UPDATE ai_relay_reservations SET status = 'released' "
105 "WHERE reservation_id = ? AND status = 'reserved'"
106)
108_SETTLE_RESERVATION = (
109 "UPDATE ai_relay_reservations SET status = 'settled' "
110 "WHERE reservation_id = ? AND status = 'reserved'"
111)
113_INSERT_USAGE = (
114 "INSERT OR IGNORE INTO ai_relay_usage "
115 "(request_id, attempt_id, tenant_id, account_id, user_id, model, "
116 "provider, channel, status, charge, currency, converter_id, metadata, "
117 "created_at, prompt_tokens, completion_tokens, cache_read_tokens, "
118 "cache_creation_tokens, reasoning_tokens, audio_input_tokens, "
119 "audio_output_tokens, image_tokens, input_tokens, output_tokens, "
120 "total_tokens_override) "
121 "VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?)"
122)
124_SELECT_USAGE = "SELECT * FROM ai_relay_usage WHERE request_id = ? AND attempt_id = ?"
126_USAGE_ALIASES = {
127 "tenant_id": "tenant_id",
128 "account_id": "account_id",
129 "user_id": "user_id",
130 "model": "model",
131 "provider": "provider",
132 "channel": "channel",
133 "status": "status",
134 "request_id": "request_id",
135 "attempt_id": "attempt_id",
136}
138_USAGE_RANGES = (
139 ("created_at_gte", "created_at", ">="),
140 ("created_at_lte", "created_at", "<="),
141)
144class DatabaseRelayUsageStore(RelayUsageStoreProtocol):
145 """SQL-backed relay billing store.
147 Writes reservations and usage records to ``ai_relay_reservations``
148 and ``ai_relay_usage`` (created lazily on first use). All writes are
149 intentionally idempotent:
151 - :meth:`save_reservation` inserts with ``INSERT OR IGNORE`` keyed by
152 ``reservation_id``; a retry or a duplicate never overwrites an
153 existing reservation.
155 - :meth:`settle_once` uses a unique ``(request_id, attempt_id)``
156 primary key, so retries and concurrent settles return the stored
157 record instead of charging twice. It also transitions the
158 associated reservation with a ``reserved -> settled``
159 compare-and-set update.
161 - :meth:`release` transitions ``reserved -> released`` only when the
162 row is still ``reserved``, and expires stale reservations first with
163 a ``reserved -> expired`` compare-and-set update.
165 Monetary values are stored as exact decimal text so that no currency
166 precision is lost crossing the database boundary. Token dimensions
167 are stored in dedicated integer columns; loss codes and extra
168 metadata are serialised into the ``metadata`` JSON column.
170 Args:
171 db: A connected :class:`~lexigram.contracts.data.DatabaseProviderProtocol`
172 resolved from the DI container.
173 """
175 def __init__(self, db: DatabaseProviderProtocol) -> None:
176 self._db = db
177 self._initialised = False
179 async def _ensure_tables(self) -> None:
180 """Create the storage schema once, on first use."""
181 if not self._initialised:
182 await self._db.execute(_CREATE_RESERVATIONS)
183 await self._db.execute(_CREATE_RESERVATIONS_INDEX)
184 await self._db.execute(_CREATE_USAGE)
185 await self._db.execute(_CREATE_USAGE_INDEX)
186 self._initialised = True
188 async def save_reservation(self, reservation: RelayUsageReservation) -> None:
189 """Persist *reservation* atomically, ignoring duplicates.
191 Args:
192 reservation: The reservation to persist.
193 """
194 await self._ensure_tables()
195 await self._db.execute(
196 _INSERT_RESERVATION,
197 [
198 reservation.reservation_id,
199 reservation.request_id,
200 reservation.estimated_tokens,
201 str(reservation.estimated_charge),
202 reservation.expires_at.isoformat(),
203 clock.now().isoformat(),
204 ],
205 )
207 async def settle_once(self, record: RelayUsageRecord) -> RelayUsageRecord:
208 """Settle *record* exactly once, returning the stored record.
210 The usage write is guarded by the ``(request_id, attempt_id)``
211 primary key: a duplicate insert is ignored and the existing row is
212 returned. The reservation is then marked ``settled`` through a
213 compare-and-set update, and stale reservations are expired first.
215 Args:
216 record: The settled usage record.
218 Returns:
219 The stored record (the existing row on a duplicate).
220 """
221 await self._ensure_tables()
222 now = clock.now()
223 await self._db.execute(_EXPIRE_RESERVATIONS, [now.isoformat()])
224 await self._db.execute(
225 _INSERT_USAGE,
226 [
227 record.request_id,
228 record.attempt_id,
229 record.scope.tenant_id,
230 record.scope.account_id,
231 record.scope.user_id,
232 record.scope.model,
233 record.scope.provider,
234 record.scope.channel,
235 record.status,
236 str(record.charge),
237 record.currency,
238 record.converter_id,
239 dumps_str({"loss_codes": list(record.loss_codes)}),
240 now.isoformat(),
241 record.usage.prompt_tokens,
242 record.usage.completion_tokens,
243 record.usage.cache_read_tokens,
244 record.usage.cache_creation_tokens,
245 record.usage.reasoning_tokens,
246 record.usage.audio_input_tokens,
247 record.usage.audio_output_tokens,
248 record.usage.image_tokens,
249 record.usage.input_tokens,
250 record.usage.output_tokens,
251 record.usage.total_tokens_override,
252 ],
253 )
254 result = await self._db.execute_query(
255 _SELECT_USAGE, [record.request_id, record.attempt_id]
256 )
257 row = result.rows[0]
258 await self._db.execute(_SETTLE_RESERVATION, [record.attempt_id])
259 return _row_to_record(row)
261 async def release(self, reservation_id: str) -> None:
262 """Expire and release a reservation, idempotently.
264 Stale ``reserved`` rows past their expiry are first marked
265 ``expired``; the reservation is then transitioned ``released`` only
266 if it is still ``reserved``. Unknown or already-final
267 reservations are left untouched.
269 Args:
270 reservation_id: Reservation identifier.
271 """
272 await self._ensure_tables()
273 now = clock.now()
274 await self._db.execute(_EXPIRE_RESERVATIONS, [now.isoformat()])
275 await self._db.execute(_RELEASE_RESERVATION, [reservation_id])
277 async def query(
278 self, filters: Mapping[str, JsonValue]
279 ) -> Sequence[RelayUsageRecord]:
280 """Query settled records by scope filters, newest first.
282 Args:
283 filters: Recognised keys are ``tenant_id``, ``account_id``,
284 ``user_id``, ``model``, ``provider``, ``channel``,
285 ``status``, ``request_id``, ``attempt_id``,
286 ``created_at_gte``, and ``created_at_lte``.
288 Returns:
289 Matching usage records ordered by insertion timestamp descending.
290 """
291 await self._ensure_tables()
292 clauses: list[str] = ["1=1"]
293 params: list[JsonValue] = []
294 for key, column in _USAGE_ALIASES.items():
295 value = filters.get(key)
296 if value is not None:
297 clauses.append(f"{column} = ?")
298 params.append(value)
299 for key, column, operator in _USAGE_RANGES:
300 value = filters.get(key)
301 if value is not None:
302 clauses.append(f"{column} {operator} ?")
303 params.append(value)
304 sql = (
305 f"SELECT * FROM ai_relay_usage WHERE {' AND '.join(clauses)} " # noqa: S608 -- columns/operators from constant tuples; values parameterized
306 "ORDER BY created_at DESC"
307 )
308 result = await self._db.execute_query(sql, params)
309 return [_row_to_record(row) for row in result.rows]
312def _row_to_record(row: dict[str, Any]) -> RelayUsageRecord:
313 """Convert a database row into a :class:`RelayUsageRecord`.
315 Args:
316 row: A usage row read from the database.
318 Returns:
319 The reconstructed usage record.
320 """
321 metadata = loads_str(row.get("metadata") or "{}")
322 loss_codes = tuple(metadata.get("loss_codes", ()))
323 return RelayUsageRecord(
324 request_id=row["request_id"],
325 attempt_id=row["attempt_id"],
326 scope=RelayUsageScope(
327 tenant_id=row["tenant_id"],
328 account_id=row["account_id"],
329 user_id=row["user_id"],
330 model=row["model"],
331 provider=row["provider"],
332 channel=row["channel"],
333 ),
334 usage=RelayUsage(
335 prompt_tokens=row["prompt_tokens"],
336 completion_tokens=row["completion_tokens"],
337 cache_read_tokens=row["cache_read_tokens"],
338 cache_creation_tokens=row["cache_creation_tokens"],
339 reasoning_tokens=row["reasoning_tokens"],
340 audio_input_tokens=row["audio_input_tokens"],
341 audio_output_tokens=row["audio_output_tokens"],
342 image_tokens=row["image_tokens"],
343 input_tokens=row["input_tokens"],
344 output_tokens=row["output_tokens"],
345 total_tokens_override=row["total_tokens_override"],
346 ),
347 charge=Decimal(row["charge"]),
348 currency=row["currency"],
349 status=row["status"],
350 converter_id=row["converter_id"],
351 loss_codes=loss_codes,
352 )