« prev ^ index » next coverage.py v7.15.4, created at 2026-08-25 07:19 +0800
1"""Budget tracker — TPM (tokens-per-minute) + sliding-window cost enforcement.
3Extends the existing RPM-based governance with:
4* TPM enforcement using a sliding-window counter.
5* Per-model and per-tenant cost tracking.
6* Budget alert events at configurable thresholds (80%, 90%, 100%) emitted
7 via ``EventBusProtocol``.
8* ``Result``-based API — policy violations return ``Err``, infrastructure
9 errors raise.
10"""
12from __future__ import annotations
14import asyncio
15from collections import deque
16from dataclasses import dataclass
17import time
18from typing import TYPE_CHECKING, Any
20from lexigram.identity import ambient as identity
21from lexigram.logging import (
22 get_logger,
23)
24from lexigram.result import Err, Ok, Result
26if TYPE_CHECKING:
27 from lexigram.contracts.events import EventBusProtocol
29logger = get_logger(__name__)
31# ---------------------------------------------------------------------------
32# Domain types
33# ---------------------------------------------------------------------------
35_ALERT_THRESHOLDS = (0.80, 0.90, 1.00)
38@dataclass
39class BudgetApproval:
40 """Returned when a budget check passes.
42 Attributes:
43 remaining_tokens: Remaining TPM capacity for this window.
44 remaining_cost: Remaining cost budget for the current period.
45 """
47 remaining_tokens: int | None
48 remaining_cost: float | None
51@dataclass
52class BudgetExceeded:
53 """Domain error returned when a request would exceed a budget limit.
55 Attributes:
56 limit_type: ``"tpm"`` or ``"cost"``.
57 current: Current value in the windows.
58 limit: Configured limit.
59 model: Model identifier.
60 tenant_id: Tenant identifier (if applicable).
61 """
63 limit_type: str
64 current: float
65 limit: float
66 model: str
67 tenant_id: str | None = None
70@dataclass
71class BudgetAlertEvent:
72 """Domain event emitted when a budget threshold is crossed.
74 Attributes:
75 threshold: The fraction threshold that was crossed (0.8, 0.9, 1.0).
76 limit_type: ``"tpm"`` or ``"cost"``.
77 current: Current value.
78 limit: Configured limit.
79 model: Model identifier.
80 tenant_id: Tenant identifier.
81 """
83 threshold: float
84 limit_type: str
85 current: float
86 limit: float
87 model: str
88 tenant_id: str | None = None
91# ---------------------------------------------------------------------------
92# Sliding-window counter
93# ---------------------------------------------------------------------------
96class SlidingWindowCounter:
97 """Thread-safe sliding window counter for rate / budget enforcement.
99 Uses a deque of ``(timestamp, value)`` tuples. On each access,
100 entries older than *window_seconds* are dropped. Reservation
101 amounts (``reserve``/``release_reservation``) are held separately
102 from used entries and counted by ``total()`` so concurrent
103 reservations cannot oversubscribe a limit.
104 """
106 def __init__(self, window_seconds: float = 60.0) -> None:
107 self._window = window_seconds
108 self._entries: deque[tuple[float, float]] = deque()
109 self._reserved: dict[str, float] = {}
110 self._lock = asyncio.Lock()
112 async def add(self, value: float) -> float:
113 """Record *value* and return the updated used-window total.
115 Args:
116 value: Value to add (token count, cost in USD, etc.).
118 Returns:
119 Sum of used values in the current window after adding *value*.
120 """
121 now = time.monotonic()
122 async with self._lock:
123 self._entries.append((now, value))
124 self._prune(now)
125 return sum(v for _, v in self._entries)
127 async def total(self) -> float:
128 """Return the window total including reservations."""
129 now = time.monotonic()
130 async with self._lock:
131 self._prune(now)
132 used = sum(v for _, v in self._entries)
133 return used + sum(self._reserved.values())
135 async def reserve(self, reservation_id: str, value: float) -> None:
136 """Reserve *value* under *reservation_id* against the limit.
138 Args:
139 reservation_id: Identifier used to release the reservation.
140 value: Amount to hold (token count, cost in USD, etc.).
141 """
142 async with self._lock:
143 self._reserved[reservation_id] = (
144 self._reserved.get(reservation_id, 0.0) + value
145 )
147 async def release_reservation(self, reservation_id: str) -> None:
148 """Release a reservation; harmless when unknown."""
149 async with self._lock:
150 self._reserved.pop(reservation_id, None)
152 def _prune(self, now: float) -> None:
153 cutoff = now - self._window
154 while self._entries and self._entries[0][0] < cutoff:
155 self._entries.popleft()
158# ---------------------------------------------------------------------------
159# BudgetTracker
160# ---------------------------------------------------------------------------
163class BudgetTracker:
164 """Track and enforce TPM (tokens-per-minute) + cost budgets.
166 Maintains per-model and per-tenant sliding windows. When a request
167 would push a window over its limit, returns ``Err(BudgetExceeded)``.
168 Emits :class:`BudgetAlertEvent` domain events at configurable thresholds
169 (80%, 90%, 100%) via an optional ``EventBusProtocol``.
171 Example::
173 tracker = BudgetTracker(
174 tpm_limit=100_000,
175 cost_limit_hourly=10.0,
176 )
177 result = await tracker.check_budget("gpt-4", estimated_tokens=500)
178 if result.is_err():
179 raise PermissionError(str(result.unwrap_err()))
181 # After the LLM call succeeds:
182 await tracker.record_usage("gpt-4", tokens_used=480, cost=0.024)
183 """
185 def __init__(
186 self,
187 *,
188 tpm_limit: int | None = None,
189 cost_limit_hourly: float | None = None,
190 window_seconds: float = 60.0,
191 event_bus: EventBusProtocol | None = None,
192 ) -> None:
193 """Initialize the budget tracker.
195 Args:
196 tpm_limit: Maximum tokens per minute (across all models). When
197 ``None``, TPM is not enforced.
198 cost_limit_hourly: Maximum USD cost per hour. When ``None``,
199 cost is tracked but not enforced.
200 window_seconds: Sliding window size in seconds (default 60 s for
201 TPM; cost window is 3600 s internally).
202 event_bus: Optional event bus for publishing
203 :class:`BudgetAlertEvent` domain events.
204 """
205 self._tpm_limit = tpm_limit
206 self._cost_limit = cost_limit_hourly
207 self._event_bus = event_bus
209 # Counters keyed by ``"<model>:<tenant_id>"`` (or ``"<model>:global"``)
210 self._token_counters: dict[str, SlidingWindowCounter] = {}
211 self._cost_counters: dict[str, SlidingWindowCounter] = {}
213 self._tpm_window = window_seconds
214 self._cost_window = 3600.0 # 1-hour sliding window for cost
216 self._alerted_thresholds: dict[str, set[float]] = {}
218 def _counter_key(self, model: str, tenant_id: str | None) -> str:
219 return f"{model}:{tenant_id or 'global'}"
221 def _get_token_counter(self, key: str) -> SlidingWindowCounter:
222 if key not in self._token_counters:
223 self._token_counters[key] = SlidingWindowCounter(self._tpm_window)
224 return self._token_counters[key]
226 def _get_cost_counter(self, key: str) -> SlidingWindowCounter:
227 if key not in self._cost_counters:
228 self._cost_counters[key] = SlidingWindowCounter(self._cost_window)
229 return self._cost_counters[key]
231 async def check_budget(
232 self,
233 model: str,
234 estimated_tokens: int,
235 *,
236 estimated_cost: float = 0.0,
237 tenant_id: str | None = None,
238 ) -> Result[BudgetApproval, BudgetExceeded]:
239 """Check whether a request fits within budget limits.
241 This method is read-only — it does NOT record usage. Call
242 :meth:`record_usage` after a successful LLM call.
244 Args:
245 model: Model identifier.
246 estimated_tokens: Estimated token count for the request.
247 estimated_cost: Estimated USD cost for the request.
248 tenant_id: Optional tenant identifier for per-tenant limits.
250 Returns:
251 ``Ok(BudgetApproval)`` if within budget,
252 ``Err(BudgetExceeded)`` if the limit would be exceeded.
253 """
254 key = self._counter_key(model, tenant_id)
256 if self._tpm_limit is not None:
257 current_tpm = await self._get_token_counter(key).total()
258 if current_tpm + estimated_tokens > self._tpm_limit:
259 logger.warning(
260 "budget_tracker_tpm_exceeded",
261 model=model,
262 current_tpm=current_tpm,
263 tpm_limit=self._tpm_limit,
264 tenant_id=tenant_id,
265 )
266 return Err(
267 BudgetExceeded(
268 limit_type="tpm",
269 current=current_tpm,
270 limit=float(self._tpm_limit),
271 model=model,
272 tenant_id=tenant_id,
273 )
274 )
276 if self._cost_limit is not None and estimated_cost > 0:
277 current_cost = await self._get_cost_counter(key).total()
278 if current_cost + estimated_cost > self._cost_limit:
279 logger.warning(
280 "budget_tracker_cost_exceeded",
281 model=model,
282 current_cost=current_cost,
283 cost_limit=self._cost_limit,
284 tenant_id=tenant_id,
285 )
286 return Err(
287 BudgetExceeded(
288 limit_type="cost",
289 current=current_cost,
290 limit=self._cost_limit,
291 model=model,
292 tenant_id=tenant_id,
293 )
294 )
296 remaining_tokens = (
297 (self._tpm_limit - int(await self._get_token_counter(key).total()))
298 if self._tpm_limit is not None
299 else None
300 )
301 remaining_cost = (
302 (self._cost_limit - await self._get_cost_counter(key).total())
303 if self._cost_limit is not None
304 else None
305 )
306 return Ok(
307 BudgetApproval(
308 remaining_tokens=remaining_tokens, remaining_cost=remaining_cost
309 )
310 )
312 async def record_usage(
313 self,
314 model: str,
315 tokens_used: int,
316 cost: float,
317 *,
318 tenant_id: str | None = None,
319 ) -> None:
320 """Record actual token usage and cost after a completed LLM call.
322 Args:
323 model: Model identifier.
324 tokens_used: Actual tokens consumed.
325 cost: Actual cost in USD.
326 tenant_id: Optional tenant identifier.
327 """
328 key = self._counter_key(model, tenant_id)
329 new_tpm = await self._get_token_counter(key).add(float(tokens_used))
330 new_cost = await self._get_cost_counter(key).add(cost)
332 # Emit alert events at threshold crossings
333 await self._check_alerts(
334 key, model, tenant_id, "tpm", new_tpm, float(self._tpm_limit or 0)
335 )
336 await self._check_alerts(
337 key, model, tenant_id, "cost", new_cost, self._cost_limit or 0.0
338 )
340 async def _check_alerts(
341 self,
342 key: str,
343 model: str,
344 tenant_id: str | None,
345 limit_type: str,
346 current: float,
347 limit: float,
348 ) -> None:
349 """Emit alert events when thresholds are crossed."""
350 if limit <= 0 or self._event_bus is None:
351 return
353 alerted = self._alerted_thresholds.setdefault(f"{key}:{limit_type}", set())
354 fraction = current / limit
356 for threshold in _ALERT_THRESHOLDS:
357 if fraction >= threshold and threshold not in alerted:
358 alerted.add(threshold)
359 event = BudgetAlertEvent(
360 threshold=threshold,
361 limit_type=limit_type,
362 current=current,
363 limit=limit,
364 model=model,
365 tenant_id=tenant_id,
366 )
367 result = await self._event_bus.publish(event)
368 if result.is_err():
369 logger.warning(
370 "budget_alert_publish_failed",
371 threshold=threshold,
372 limit_type=limit_type,
373 error=str(result.unwrap_err()),
374 )
376 async def get_usage(
377 self, model: str, *, tenant_id: str | None = None
378 ) -> dict[str, Any]:
379 """Return current usage totals for a model/tenant combination.
381 Args:
382 model: Model identifier.
383 tenant_id: Optional tenant identifier.
385 Returns:
386 Dict with ``tokens_per_minute`` and ``cost_per_hour`` keys.
387 """
388 key = self._counter_key(model, tenant_id)
389 return {
390 "tokens_per_minute": await self._get_token_counter(key).total(),
391 "cost_per_hour": await self._get_cost_counter(key).total(),
392 "tpm_limit": self._tpm_limit,
393 "cost_limit_hourly": self._cost_limit,
394 }
396 def reset_alerts(self, model: str, *, tenant_id: str | None = None) -> None:
397 """Reset threshold-alert state for a model/tenant (e.g. on period rollover).
399 Args:
400 model: Model identifier.
401 tenant_id: Optional tenant identifier.
402 """
403 key = self._counter_key(model, tenant_id)
404 for limit_type in ("tpm", "cost"):
405 self._alerted_thresholds.pop(f"{key}:{limit_type}", None)
407 async def reserve(
408 self,
409 model: str,
410 estimated_tokens: int,
411 *,
412 estimated_cost: float = 0.0,
413 tenant_id: str | None = None,
414 ) -> Result[str, BudgetExceeded]:
415 """Atomically check limits and reserve capacity for a request.
417 Reserved amounts count against the sliding window so concurrent
418 requests cannot oversubscribe a limit. Call
419 :meth:`release_reservation` after the request settles.
421 Args:
422 model: Model identifier.
423 estimated_tokens: Estimated token count for the request.
424 estimated_cost: Estimated USD cost for the request.
425 tenant_id: Optional tenant identifier.
427 Returns:
428 ``Ok(reservation_id)`` on success, ``Err(BudgetExceeded)``
429 when a limit would be exceeded.
430 """
431 key = self._counter_key(model, tenant_id)
433 if self._tpm_limit is not None:
434 current_tpm = await self._get_token_counter(key).total()
435 if current_tpm + estimated_tokens > self._tpm_limit:
436 logger.warning(
437 "budget_tracker_reserve_tpm_exceeded",
438 model=model,
439 current_tpm=current_tpm,
440 tpm_limit=self._tpm_limit,
441 tenant_id=tenant_id,
442 )
443 return Err(
444 BudgetExceeded(
445 limit_type="tpm",
446 current=current_tpm,
447 limit=float(self._tpm_limit),
448 model=model,
449 tenant_id=tenant_id,
450 )
451 )
453 if self._cost_limit is not None and estimated_cost > 0:
454 current_cost = await self._get_cost_counter(key).total()
455 if current_cost + estimated_cost > self._cost_limit:
456 logger.warning(
457 "budget_tracker_reserve_cost_exceeded",
458 model=model,
459 current_cost=current_cost,
460 cost_limit=self._cost_limit,
461 tenant_id=tenant_id,
462 )
463 return Err(
464 BudgetExceeded(
465 limit_type="cost",
466 current=current_cost,
467 limit=self._cost_limit,
468 model=model,
469 tenant_id=tenant_id,
470 )
471 )
473 reservation_id = identity.new_uuid()
474 await self._get_token_counter(key).reserve(
475 reservation_id, float(estimated_tokens)
476 )
477 await self._get_cost_counter(key).reserve(reservation_id, estimated_cost)
478 return Ok(reservation_id)
480 async def release_reservation(
481 self,
482 model: str,
483 reservation_id: str,
484 *,
485 tenant_id: str | None = None,
486 ) -> None:
487 """Release a previously reserved amount; harmless when unknown.
489 Args:
490 model: Model identifier the reservation was made for.
491 reservation_id: Reservation identifier from :meth:`reserve`.
492 tenant_id: Optional tenant identifier.
493 """
494 key = self._counter_key(model, tenant_id)
495 await self._get_token_counter(key).release_reservation(reservation_id)
496 await self._get_cost_counter(key).release_reservation(reservation_id)
499__all__ = [
500 "BudgetAlertEvent",
501 "BudgetApproval",
502 "BudgetExceeded",
503 "BudgetTracker",
504 "SlidingWindowCounter",
505]