« prev ^ index » next coverage.py v7.15.4, created at 2026-08-25 07:19 +0800
1"""Persistence backends for AI governance state.
3Defines the :class:`GovernancePersistence` protocol and two concrete
4implementations:
6- :class:`InMemoryGovernancePersistence` — process-local state, suitable for
7 development, testing, and single-instance deployments.
8- :class:`RedisGovernancePersistence` — distributed state backed by a
9 :class:`~lexigram.contracts.cache.CacheBackendProtocol`, enabling multi-replica
10 consistency for rate limiting and budget enforcement.
12The :class:`~lexigram.ai.governance.manager.AIGovernanceManager` accepts any
13implementation via constructor injection so the storage strategy is swappable
14without changing governance logic.
15"""
17from __future__ import annotations
19import time
20from typing import TYPE_CHECKING, Any, Protocol, cast, runtime_checkable
22from lexigram.ai.governance.exceptions import GovernancePersistenceError
24if TYPE_CHECKING:
25 from lexigram.contracts.data import DatabaseProviderProtocol
26 from lexigram.contracts.infra.cache import CacheBackendProtocol
27 from lexigram.contracts.infra.cache.exceptions import CacheError
28 from lexigram.result import Result
30__all__ = [
31 "DatabaseGovernancePersistence",
32 "GovernancePersistence",
33 "InMemoryGovernancePersistence",
34 "RedisGovernancePersistence",
35]
38@runtime_checkable
39class GovernancePersistence(Protocol):
40 """Storage protocol for governance counters and spend tallies.
42 Implementations must be safe for concurrent async access. All methods
43 are *coroutines* to allow either local (in-process) or remote (Redis,
44 database) storage without changing the call-site.
45 """
47 async def incr_requests(
48 self,
49 key: str,
50 window: float,
51 ) -> int:
52 """Record a new request and return the request count within the window.
54 Args:
55 key: Bucket key (e.g. ``"global"`` or a user/tenant id).
56 window: Rolling window size in seconds.
58 Returns:
59 Number of requests (including the current one) inside the window.
60 """
61 ...
63 async def add_spend(
64 self,
65 key: str,
66 amount: float,
67 ttl: int,
68 ) -> float:
69 """Add ``amount`` to the spend accumulator and return the new total.
71 Args:
72 key: Bucket key (e.g. ``"global:2025-06"``).
73 amount: Cost amount to add.
74 ttl: Time-to-live for the accumulator entry in seconds.
76 Returns:
77 Updated total spend.
78 """
79 ...
81 async def get_spend(self, key: str) -> float:
82 """Return the current accumulated spend for *key*.
84 Args:
85 key: Bucket key.
87 Returns:
88 Current spend; ``0.0`` if no data recorded.
89 """
90 ...
92 # -- Gauge methods -------------------------------------------------------
94 async def read_gauge(self, key: str) -> float:
95 """Read the current gauge value for *key*.
97 Args:
98 key: Gauge key (e.g. ``"tenant:gpt4:remaining"``).
100 Returns:
101 Current gauge value; ``0.0`` if no data recorded.
102 """
103 ...
105 async def write_gauge(self, key: str, value: float, ttl: int) -> None:
106 """Set the gauge *value* for *key* with a TTL.
108 Args:
109 key: Gauge key.
110 value: Float gauge value.
111 ttl: Time-to-live in seconds.
112 """
113 ...
115 async def incr_gauge(self, key: str, delta: float, ttl: int) -> float:
116 """Atomically increment (or decrement) the gauge for *key*.
118 Args:
119 key: Gauge key.
120 delta: Amount to add (can be negative).
121 ttl: Time-to-live in seconds.
123 Returns:
124 The gauge value after applying *delta*.
125 """
126 ...
128 # -- Calendar methods ----------------------------------------------------
130 async def add_calendar_entry(self, key: str, timestamp: float, ttl: int) -> None:
131 """Record a timestamp entry in a calendar-style bucket.
133 Args:
134 key: Calendar bucket key.
135 timestamp: Unix timestamp to record.
136 ttl: Time-to-live in seconds for the bucket.
137 """
138 ...
140 async def query_calendar(self, key: str, start: float, end: float) -> list[float]:
141 """Return all timestamps in *key* that fall within [*start*, *end*].
143 Args:
144 key: Calendar bucket key.
145 start: Unix timestamp, start of range (inclusive).
146 end: Unix timestamp, end of range (inclusive).
148 Returns:
149 List of matching timestamps (chronological order).
150 """
151 ...
153 async def decr_gauge(self, key: str, amount: float, ttl: int) -> float:
154 """Decrement a gauge key by *amount*.
156 Default implementation delegates to :meth:`incr_gauge` with a
157 negative delta.
158 """
159 return await self.incr_gauge(key, -amount, ttl)
161 async def incr_calendar(
162 self, key: str, period: str, amount: float, ttl: int
163 ) -> float:
164 """Accumulate *amount* in a calendar-window bucket.
166 Delegates to :meth:`add_spend` with a period-scoped key (``key`` +
167 ``:`` + ``period``).
168 """
169 return await self.add_spend(f"{key}:{period}", amount, ttl)
171 async def get_calendar(self, key: str, period: str) -> float:
172 """Read the calendar-window accumulator for *key* + *period*.
174 Delegates to :meth:`get_spend` with a period-scoped key.
175 """
176 return await self.get_spend(f"{key}:{period}")
179# ---------------------------------------------------------------------------
180# In-memory implementation
181# ---------------------------------------------------------------------------
184class InMemoryGovernancePersistence:
185 """Process-local governance persistence using plain Python dicts.
187 Request counts are tracked with a sliding-window approach (list of
188 monotonic timestamps). Spend totals are stored as plain floats.
190 This implementation is **not** suitable for multi-process or
191 multi-replica deployments. Use :class:`RedisGovernancePersistence`
192 in production.
193 """
195 def __init__(self) -> None:
196 self._request_buckets: dict[str, list[float]] = {}
197 self._spend_totals: dict[str, float] = {}
198 self._gauges: dict[str, float] = {}
199 self._calendar: dict[str, list[tuple[float, float]]] = {}
201 async def incr_requests(self, key: str, window: float) -> int:
202 now = time.monotonic()
203 bucket = self._request_buckets.setdefault(key, [])
204 # Prune entries outside the window, then record current request
205 self._request_buckets[key] = [t for t in bucket if now - t <= window]
206 self._request_buckets[key].append(now)
207 return len(self._request_buckets[key])
209 async def add_spend(self, key: str, amount: float, ttl: int) -> float:
210 current = self._spend_totals.get(key, 0.0)
211 updated = current + amount
212 self._spend_totals[key] = updated
213 return updated
215 async def get_spend(self, key: str) -> float:
216 return self._spend_totals.get(key, 0.0)
218 async def read_gauge(self, key: str) -> float:
219 return self._gauges.get(key, 0.0)
221 async def write_gauge(self, key: str, value: float, ttl: int) -> None:
222 self._gauges[key] = value
224 async def incr_gauge(self, key: str, delta: float, ttl: int) -> float:
225 current = self._gauges.get(key, 0.0)
226 updated = max(current + delta, 0.0)
227 self._gauges[key] = updated
228 return updated
230 async def add_calendar_entry(self, key: str, timestamp: float, ttl: int) -> None:
231 now = time.monotonic()
232 bucket = self._calendar.setdefault(key, [])
233 bucket.append((now, timestamp))
235 async def query_calendar(self, key: str, start: float, end: float) -> list[float]:
236 results = []
237 bucket = self._calendar.get(key, [])
238 for _, ts in bucket:
239 if start <= ts <= end:
240 results.append(ts)
241 return results
244# ---------------------------------------------------------------------------
245# Redis-backed implementation
246# ---------------------------------------------------------------------------
248_32_DAYS_SECONDS = 32 * 24 * 3600
251def _cache_payload(value: object, *, key: str) -> Any:
252 """Extract the payload of a cache backend result, raising on failure.
254 Protocol-compliant backends return ``Result[Any | None, CacheError]``;
255 plain values are treated as an ``Ok`` payload. An ``Err`` result is an
256 infrastructure failure and raises :class:`GovernancePersistenceError`
257 instead of being silently treated as a missing value (the previous
258 behavior, which let failures fail open).
260 Args:
261 value: Return value of ``CacheBackendProtocol.get`` / ``set``.
262 key: Cache key the call was made for (included in the error).
264 Returns:
265 The payload — ``Any``, matching the cache protocol's payload type —
266 or ``None`` for a successful miss.
268 Raises:
269 GovernancePersistenceError: When the backend reported failure.
270 """
271 if hasattr(value, "is_ok"):
272 result = cast("Result[Any, CacheError]", value)
273 if result.is_ok():
274 return result.unwrap_or(None)
275 error = result.unwrap_err()
276 raise GovernancePersistenceError(
277 f"cache backend failed for key={key}: {error}"
278 ) from error
279 return value
282class RedisGovernancePersistence:
283 """Distributed governance persistence backed by a Lexigram CacheBackendProtocol.
285 Request windows use a sorted-set approach:
286 - Each request is stored as a member with its Unix timestamp as score.
287 - Expired members (score < ``now - window``) are pruned on every read.
289 Spend totals are stored as plain string floats with a configurable TTL
290 so that monthly counters expire automatically.
292 Args:
293 cache: A :class:`~lexigram.contracts.cache.CacheBackendProtocol` that has been
294 connected and is ready to accept commands. The backend is
295 expected to be Redis-compatible.
296 """
298 _REQUEST_KEY_PREFIX = "ai:gov:req:"
299 _SPEND_KEY_PREFIX = "ai:gov:spend:"
300 _GAUGE_KEY_PREFIX = "ai:gov:gauge:"
301 _CALENDAR_KEY_PREFIX = "ai:gov:cal:"
303 def __init__(self, cache: CacheBackendProtocol) -> None:
304 self._cache = cache
306 async def incr_requests(self, key: str, window: float) -> int:
307 """Use a sorted set to implement a sliding window counter.
309 Falls back to an approximate counter if the backend does not support
310 sorted-set operations (e.g. a simple in-memory mock). Infrastructure
311 failures are not masked: the exception propagates so the caller can
312 apply the configured governance decision (fail-closed by default).
314 Args:
315 key: Bucket key (e.g. ``"global"`` or a user/tenant id).
316 window: Rolling window size in seconds.
318 Returns:
319 Number of requests (including the current one) inside the window.
321 Raises:
322 GovernancePersistenceError: When the cache backend reports failure
323 (``Err`` result).
324 OSError, ConnectionError, RuntimeError, ValueError, TypeError:
325 Propagated from the cache backend when it raises.
326 """
327 redis_key = f"{self._REQUEST_KEY_PREFIX}{key}"
328 now = time.time()
329 cutoff = now - window
331 try:
332 # Try native sorted-set operations (Redis / ioredis)
333 backend = self._cache # CacheBackendProtocol may expose raw client
334 raw = getattr(backend, "_client", None) or getattr(backend, "client", None)
335 if raw is not None and hasattr(raw, "zremrangebyscore"):
336 # Remove expired entries
337 await raw.zremrangebyscore(redis_key, "-inf", cutoff)
338 # Add current request
339 await raw.zadd(redis_key, {str(now): now})
340 # Expire the key after the window to avoid unbounded growth
341 await raw.expire(redis_key, int(window) + 1)
342 count: int = await raw.zcard(redis_key)
343 return count
344 except (OSError, ConnectionError, RuntimeError, AttributeError):
345 pass
347 # Fallback: simple increment counter (less precise but always works)
348 counter_key = f"{redis_key}:count"
349 raw_val = _cache_payload(await self._cache.get(counter_key), key=counter_key)
350 current = int(raw_val) + 1 if raw_val is not None else 1
351 _cache_payload(
352 await self._cache.set(counter_key, str(current), ttl=int(window) + 1),
353 key=counter_key,
354 )
355 return current
357 async def add_spend(self, key: str, amount: float, ttl: int) -> float:
358 """Add ``amount`` to the spend accumulator and return the new total.
360 Infrastructure failures propagate (fail-closed by default at the
361 manager); no value is invented on error.
363 Args:
364 key: Bucket key (e.g. ``"global:2025-06"``).
365 amount: Cost amount to add.
366 ttl: Time-to-live for the accumulator entry in seconds.
368 Returns:
369 Updated total spend.
371 Raises:
372 GovernancePersistenceError: When the cache backend reports failure
373 (``Err`` result, read or write).
374 OSError, ConnectionError, RuntimeError, ValueError, TypeError:
375 Propagated from the cache backend when it raises.
376 """
377 redis_key = f"{self._SPEND_KEY_PREFIX}{key}"
378 raw = _cache_payload(await self._cache.get(redis_key), key=redis_key)
379 current = float(raw) if raw is not None else 0.0
380 updated = current + amount
381 _cache_payload(
382 await self._cache.set(redis_key, str(updated), ttl=ttl),
383 key=redis_key,
384 )
385 return updated
387 async def get_spend(self, key: str) -> float:
388 """Return the current accumulated spend for *key*.
390 Args:
391 key: Bucket key.
393 Returns:
394 Current spend; ``0.0`` if no data recorded.
396 Raises:
397 GovernancePersistenceError: When the cache backend reports failure
398 (``Err`` result).
399 OSError, ConnectionError, RuntimeError, ValueError, TypeError:
400 Propagated from the cache backend when it raises.
401 """
402 redis_key = f"{self._SPEND_KEY_PREFIX}{key}"
403 raw = _cache_payload(await self._cache.get(redis_key), key=redis_key)
404 return float(raw) if raw is not None else 0.0
406 async def read_gauge(self, key: str) -> float:
407 """Read the current gauge value for *key*.
409 Args:
410 key: Gauge key (e.g. ``"tenant:gpt4:remaining"``).
412 Returns:
413 Current gauge value; ``0.0`` if no data recorded.
415 Raises:
416 GovernancePersistenceError: When the cache backend reports failure
417 (``Err`` result).
418 OSError, ConnectionError, RuntimeError, ValueError, TypeError:
419 Propagated from the cache backend when it raises.
420 """
421 redis_key = f"{self._GAUGE_KEY_PREFIX}{key}"
422 raw = _cache_payload(await self._cache.get(redis_key), key=redis_key)
423 return float(raw) if raw is not None else 0.0
425 async def write_gauge(self, key: str, value: float, ttl: int) -> None:
426 redis_key = f"{self._GAUGE_KEY_PREFIX}{key}"
427 try:
428 await self._cache.set(redis_key, str(value), ttl=ttl)
429 except (OSError, ConnectionError, RuntimeError):
430 pass
432 async def incr_gauge(self, key: str, delta: float, ttl: int) -> float:
433 """Atomically increment (or decrement) the gauge for *key*.
435 Args:
436 key: Gauge key.
437 delta: Amount to add (can be negative).
438 ttl: Time-to-live in seconds.
440 Returns:
441 The gauge value after applying *delta*.
443 Raises:
444 GovernancePersistenceError: When the cache backend reports failure
445 (``Err`` result, read or write).
446 OSError, ConnectionError, RuntimeError, ValueError, TypeError:
447 Propagated from the cache backend when it raises.
448 """
449 redis_key = f"{self._GAUGE_KEY_PREFIX}{key}"
450 raw = _cache_payload(await self._cache.get(redis_key), key=redis_key)
451 current = float(raw) if raw is not None else 0.0
452 updated = max(current + delta, 0.0)
453 _cache_payload(
454 await self._cache.set(redis_key, str(updated), ttl=ttl),
455 key=redis_key,
456 )
457 return updated
459 async def add_calendar_entry(self, key: str, timestamp: float, ttl: int) -> None:
460 redis_key = f"{self._CALENDAR_KEY_PREFIX}{key}"
461 try:
462 raw_result = await self._cache.get(redis_key)
463 if hasattr(raw_result, "is_ok"):
464 raw = raw_result.unwrap_or(None) if raw_result.is_ok() else None
465 else:
466 raw = raw_result
467 entries: list[float] = []
468 if raw is not None:
469 entries = [float(v) for v in raw.split(",") if v]
470 entries.append(timestamp)
471 await self._cache.set(redis_key, ",".join(str(v) for v in entries), ttl=ttl)
472 except (OSError, ConnectionError, RuntimeError, ValueError, TypeError):
473 pass
475 async def query_calendar(self, key: str, start: float, end: float) -> list[float]:
476 redis_key = f"{self._CALENDAR_KEY_PREFIX}{key}"
477 try:
478 raw_result = await self._cache.get(redis_key)
479 if hasattr(raw_result, "is_ok"):
480 raw = raw_result.unwrap_or(None) if raw_result.is_ok() else None
481 else:
482 raw = raw_result
483 if raw is None:
484 return []
485 entries = [float(v) for v in raw.split(",") if v]
486 return sorted([v for v in entries if start <= v <= end])
487 except (OSError, ConnectionError, RuntimeError, ValueError, TypeError):
488 return []
491# ---------------------------------------------------------------------------
492# Database-backed implementation
493# ---------------------------------------------------------------------------
495_CREATE_REQUESTS_TABLE = """
496CREATE TABLE IF NOT EXISTS ai_governance_requests (
497 id INTEGER PRIMARY KEY AUTOINCREMENT,
498 key TEXT NOT NULL,
499 ts REAL NOT NULL
500)
501"""
503_CREATE_SPEND_TABLE = """
504CREATE TABLE IF NOT EXISTS ai_governance_spend (
505 key TEXT NOT NULL PRIMARY KEY,
506 amount REAL NOT NULL DEFAULT 0,
507 expires_at REAL NOT NULL
508)
509"""
511_INSERT_REQUEST = "INSERT INTO ai_governance_requests (key, ts) VALUES (?, ?)"
512_DELETE_EXPIRED_REQUESTS = "DELETE FROM ai_governance_requests WHERE key = ? AND ts < ?"
513_COUNT_REQUESTS = "SELECT COUNT(*) AS cnt FROM ai_governance_requests WHERE key = ?"
515_UPSERT_SPEND = (
516 "INSERT INTO ai_governance_spend (key, amount, expires_at) VALUES (?, ?, ?) "
517 "ON CONFLICT (key) DO UPDATE SET "
518 "amount = amount + excluded.amount, "
519 "expires_at = excluded.expires_at"
520)
522_GET_SPEND = "SELECT amount FROM ai_governance_spend WHERE key = ? AND expires_at >= ?"
524_GAUGE_TABLE = """
525CREATE TABLE IF NOT EXISTS ai_governance_gauges (
526 key TEXT NOT NULL PRIMARY KEY,
527 value REAL NOT NULL DEFAULT 0.0,
528 expires_at REAL NOT NULL
529)
530"""
531_UPSERT_GAUGE = (
532 "INSERT INTO ai_governance_gauges (key, value, expires_at) VALUES (?, ?, ?) "
533 "ON CONFLICT (key) DO UPDATE SET "
534 "value = excluded.value, "
535 "expires_at = excluded.expires_at"
536)
537_GET_GAUGE = "SELECT value FROM ai_governance_gauges WHERE key = ? AND expires_at >= ?"
539_CALENDAR_TABLE = """
540CREATE TABLE IF NOT EXISTS ai_governance_calendar (
541 id INTEGER PRIMARY KEY AUTOINCREMENT,
542 key TEXT NOT NULL,
543 ts REAL NOT NULL
544)
545"""
546_INSERT_CALENDAR = "INSERT INTO ai_governance_calendar (key, ts) VALUES (?, ?)"
547_QUERY_CALENDAR = (
548 "SELECT ts FROM ai_governance_calendar "
549 "WHERE key = ? AND ts >= ? AND ts <= ? ORDER BY ts ASC"
550)
553class DatabaseGovernancePersistence:
554 """SQL-backed governance persistence using :class:`~lexigram.contracts.data.DatabaseProviderProtocol`.
556 Stores request timestamps in ``ai_governance_requests`` and spend totals
557 in ``ai_governance_spend``. Both tables are created lazily on first use.
559 This backend is suitable for multi-replica deployments where Redis is not
560 available but a shared relational database is. Row-level locking provided
561 by the database engine ensures counter consistency.
563 Args:
564 db: A connected :class:`~lexigram.contracts.data.DatabaseProviderProtocol`
565 resolved from the DI container.
566 """
568 def __init__(self, db: DatabaseProviderProtocol) -> None:
569 self._db = db
570 self._initialised = False
572 async def _ensure_tables(self) -> None:
573 """Create governance tables if they do not yet exist."""
574 if not self._initialised:
575 await self._db.execute(_CREATE_REQUESTS_TABLE)
576 await self._db.execute(_CREATE_SPEND_TABLE)
577 await self._db.execute(_GAUGE_TABLE)
578 await self._db.execute(_CALENDAR_TABLE)
579 self._initialised = True
581 async def incr_requests(self, key: str, window: float) -> int:
582 """Insert current timestamp, prune expired rows, return window count.
584 Args:
585 key: Governance bucket key.
586 window: Rolling window size in seconds.
588 Returns:
589 Number of requests within *window* (including the current one).
590 """
591 await self._ensure_tables()
592 now = time.time()
593 cutoff = now - window
594 await self._db.execute(_INSERT_REQUEST, [key, now])
595 await self._db.execute(_DELETE_EXPIRED_REQUESTS, [key, cutoff])
596 result = await self._db.execute_query(_COUNT_REQUESTS, [key])
597 rows = result.rows
598 return int(rows[0]["cnt"]) if rows else 1
600 async def add_spend(self, key: str, amount: float, ttl: int) -> float:
601 """Accumulate *amount* against *key* and return the running total.
603 The row is upserted with ``expires_at = now + ttl``; callers should
604 use time-period-scoped keys (e.g. ``"global:2025-06"``) so that
605 distinct budget periods never collide.
607 Args:
608 key: Governance bucket key.
609 amount: Cost to add.
610 ttl: Seconds until the entry should expire.
612 Returns:
613 Updated running total for *key*.
614 """
615 await self._ensure_tables()
616 expires_at = time.time() + ttl
617 await self._db.execute(_UPSERT_SPEND, [key, amount, expires_at])
618 result = await self._db.execute_query(_GET_SPEND, [key, time.time()])
619 rows = result.rows
620 return float(rows[0]["amount"]) if rows else amount
622 async def get_spend(self, key: str) -> float:
623 """Return the current accumulated spend for *key*.
625 Args:
626 key: Governance bucket key.
628 Returns:
629 Current spend; ``0.0`` if no data recorded or entry has expired.
630 """
631 await self._ensure_tables()
632 result = await self._db.execute_query(_GET_SPEND, [key, time.time()])
633 rows = result.rows
634 return float(rows[0]["amount"]) if rows else 0.0
636 async def read_gauge(self, key: str) -> float:
637 """Read the current gauge value for *key*."""
638 await self._ensure_tables()
639 result = await self._db.execute_query(_GET_GAUGE, [key, time.time()])
640 rows = result.rows
641 return float(rows[0]["value"]) if rows else 0.0
643 async def write_gauge(self, key: str, value: float, ttl: int) -> None:
644 """Set the gauge *value* for *key* with a TTL."""
645 await self._ensure_tables()
646 expires_at = time.time() + ttl
647 await self._db.execute(_UPSERT_GAUGE, [key, value, expires_at])
649 async def incr_gauge(self, key: str, delta: float, ttl: int) -> float:
650 """Atomically increment (or decrement) the gauge for *key*."""
651 await self._ensure_tables()
652 now = time.time()
653 result = await self._db.execute_query(_GET_GAUGE, [key, now])
654 rows = result.rows
655 current = float(rows[0]["value"]) if rows else 0.0
656 updated = max(current + delta, 0.0)
657 expires_at = now + ttl
658 await self._db.execute(_UPSERT_GAUGE, [key, updated, expires_at])
659 return updated
661 async def add_calendar_entry(self, key: str, timestamp: float, ttl: int) -> None:
662 """Record a timestamp entry in a calendar-style bucket."""
663 await self._ensure_tables()
664 await self._db.execute(_INSERT_CALENDAR, [key, timestamp])
666 async def query_calendar(self, key: str, start: float, end: float) -> list[float]:
667 """Return all timestamps in *key* that fall within range."""
668 await self._ensure_tables()
669 result = await self._db.execute_query(_QUERY_CALENDAR, [key, start, end])
670 rows = result.rows
671 return [float(r["ts"]) for r in rows] if rows else []