1"""Usage and quota reporting for relay billing.
2
3``RelayUsageReportService`` runs bounded, read-only aggregations over
4:class:`~lexigram.contracts.ai.governance.RelayUsageStoreProtocol`:
5scope and status filters narrow the window, token and charge totals are
6computed with integer/Decimal arithmetic, and loss/status counts survive
7aggregation. Reports are always bounded — a window start/end and a
8maximum page size are enforced before any store access.
9"""
10
11from __future__ import annotations
12
13from dataclasses import dataclass, field
14from datetime import datetime
15from decimal import Decimal
16from typing import TYPE_CHECKING
17
18from lexigram.contracts.ai.governance import (
19 RelayUsageRecord,
20 RelayUsageStoreProtocol,
21)
22
23if TYPE_CHECKING:
24 from collections.abc import Sequence
25
26 from lexigram.contracts.ai.relay import JsonValue
27
28__all__ = [
29 "TERMINAL_STATUSES",
30 "RelayUsageReport",
31 "RelayUsageReportService",
32 "RelayUsageTotals",
33]
34
35TERMINAL_STATUSES = ("completed", "failed", "cancelled", "truncated")
36
37
38@dataclass(frozen=True, slots=True)
39class RelayUsageTotals:
40 """Aggregated token, charge, and count totals for a report window.
41
42 Attributes:
43 prompt_tokens: Sum of prompt tokens across matched records.
44 completion_tokens: Sum of completion tokens across matched records.
45 total_tokens: Sum of normalized total tokens.
46 total_charge: Sum of charges as exact ``Decimal`` arithmetic.
47 request_count: Number of matched records.
48 status_counts: Per-terminal-status record counts (non-zero only).
49 loss_counts: Per-loss-code occurrence counts (non-zero only).
50 """
51
52 prompt_tokens: int = 0
53 completion_tokens: int = 0
54 total_tokens: int = 0
55 total_charge: Decimal = Decimal("0")
56 request_count: int = 0
57 status_counts: dict[str, int] = field(default_factory=dict)
58 loss_counts: dict[str, int] = field(default_factory=dict)
59
60
61@dataclass(frozen=True, slots=True)
62class RelayUsageReport:
63 """One page of a usage report plus window totals.
64
65 Attributes:
66 rows: The requested page of matched records (empty past the end).
67 total_rows: Number of records matching the filters in the window.
68 totals: Aggregation across every matched record, not just the page.
69 """
70
71 rows: tuple[RelayUsageRecord, ...] = ()
72 total_rows: int = 0
73 totals: RelayUsageTotals = field(default_factory=RelayUsageTotals)
74
75
76def _aggregate(records: Sequence[RelayUsageRecord]) -> RelayUsageTotals:
77 """Aggregate token/charge/count totals across *records*.
78
79 Args:
80 records: The full matched set within the report window.
81
82 Returns:
83 Totals with zero entries omitted from the count maps.
84 """
85 prompt_tokens = 0
86 completion_tokens = 0
87 total_tokens = 0
88 total_charge = Decimal("0")
89 status_counts: dict[str, int] = {}
90 loss_counts: dict[str, int] = {}
91 for record in records:
92 prompt_tokens += record.usage.prompt_tokens
93 completion_tokens += record.usage.completion_tokens
94 total_tokens += record.usage.total_tokens
95 total_charge += record.charge
96 status_counts[record.status] = status_counts.get(record.status, 0) + 1
97 for code in record.loss_codes:
98 loss_counts[code] = loss_counts.get(code, 0) + 1
99 return RelayUsageTotals(
100 prompt_tokens=prompt_tokens,
101 completion_tokens=completion_tokens,
102 total_tokens=total_tokens,
103 total_charge=total_charge,
104 request_count=len(records),
105 status_counts=status_counts,
106 loss_counts=loss_counts,
107 )
108
109
110class RelayUsageReportService:
111 """Read-only aggregation over settled relay usage records.
112
113 Args:
114 store: The usage store to query (resolved from the DI container).
115 max_page_size: Largest page size accepted by :meth:`report`.
116 """
117
118 def __init__(
119 self,
120 store: RelayUsageStoreProtocol,
121 *,
122 max_page_size: int = 25,
123 ) -> None:
124 self._store = store
125 self._max_page_size = max_page_size
126
127 async def report(
128 self,
129 *,
130 start: datetime,
131 end: datetime,
132 page: int = 1,
133 page_size: int = 25,
134 tenant_id: str | None = None,
135 account_id: str | None = None,
136 user_id: str | None = None,
137 model: str | None = None,
138 provider: str | None = None,
139 channel: str | None = None,
140 status: str | None = None,
141 ) -> RelayUsageReport:
142 """Run a bounded usage report over a UTC window.
143
144 Args:
145 start: Window start (inclusive).
146 end: Window end (inclusive).
147 page: One-based page number.
148 page_size: Number of rows per page, capped by the maximum.
149 tenant_id: Narrow the report to one tenant.
150 account_id: Narrow the report to one account.
151 user_id: Narrow the report to one user.
152 model: Narrow the report to one model alias.
153 provider: Narrow the report to one provider.
154 channel: Narrow the report to one channel.
155 status: Narrow the report to one terminal status.
156
157 Returns:
158 The requested page plus totals covering the whole window.
159
160 Raises:
161 ValueError: For unbounded/inverted windows, non-positive
162 pages or page sizes, page sizes above the configured
163 maximum, or unknown statuses.
164 """
165 if start is None or end is None:
166 raise ValueError("report requires both start and end window bounds")
167 if end <= start:
168 raise ValueError("end must be after start")
169 if page < 1:
170 raise ValueError("page must be at least 1")
171 if page_size < 1:
172 raise ValueError("page_size must be at least 1")
173 if page_size > self._max_page_size:
174 raise ValueError(
175 f"page_size must not exceed the maximum {self._max_page_size}"
176 )
177 if status is not None and status not in TERMINAL_STATUSES:
178 raise ValueError(
179 f"unknown status {status!r}; expected one of {TERMINAL_STATUSES}"
180 )
181
182 filters: dict[str, JsonValue] = {
183 "created_at_gte": start.isoformat(),
184 "created_at_lte": end.isoformat(),
185 }
186 filters.update(
187 {
188 key: value
189 for key, value in (
190 ("tenant_id", tenant_id),
191 ("account_id", account_id),
192 ("user_id", user_id),
193 ("model", model),
194 ("provider", provider),
195 ("channel", channel),
196 )
197 if value is not None
198 }
199 )
200 if status is not None:
201 filters["status"] = status
202
203 records = await self._store.query(filters)
204 offset = (page - 1) * page_size
205 return RelayUsageReport(
206 rows=tuple(records[offset : offset + page_size]),
207 total_rows=len(records),
208 totals=_aggregate(records),
209 )