Coverage for /home/admin/Documents/AI/applications/lexigram-dev/lexigram/experimental/ai/lexigram-ai-governance/src/lexigram/ai/governance/relay_billing/reports.py: 42%

62 statements  

« prev     ^ index     » next       coverage.py v7.15.4, created at 2026-08-25 07:19 +0800

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 )