« prev ^ index » next coverage.py v7.15.4, created at 2026-08-25 07:19 +0800
1"""AI usage governance — budget limits, rate limits, model restrictions."""
3from __future__ import annotations
5import asyncio
6import fnmatch
7from typing import TYPE_CHECKING, cast
9from lexigram.ai.governance.exceptions import GovernancePersistenceError
10from lexigram.contracts.ai.governance.resource_unit import ResourceExhaustedError
11from lexigram.di.decorators import inject
12from lexigram.logging import (
13 get_logger,
14)
15from lexigram.result import Err, Ok, Result
17if TYPE_CHECKING:
18 from collections.abc import Callable
20 from lexigram.ai.governance.audit import AIAuditStore
21 from lexigram.ai.governance.config import GovernanceConfig
22 from lexigram.ai.governance.exceptions import GovernanceError
23 from lexigram.ai.governance.persistence import GovernancePersistence
24 from lexigram.ai.governance.resource.registry import ( # noqa: F401
25 ResourceUnitRegistry,
26 )
27 from lexigram.ai.governance.resource.tracker import ( # noqa: F401
28 ResourceUnitTracker,
29 )
30 from lexigram.contracts.infra.cache import CacheBackendProtocol
32logger = get_logger(__name__)
34_PERSISTENCE_FAILURE_EXCEPTIONS: tuple[type[Exception], ...] = (
35 GovernancePersistenceError,
36 OSError,
37 ConnectionError,
38 RuntimeError,
39 ValueError,
40 TypeError,
41)
44@inject
45class AIGovernanceManager:
46 """Enforces AI usage policies: budget limits, rate limits, model restrictions.
48 Implements ``AIGovernanceProtocol`` from contracts.
50 Governance state (request counts, spend totals) is delegated to a
51 :class:`~lexigram.ai.governance.persistence.GovernancePersistence`
52 backend so the storage strategy is swappable without changing policy
53 logic. When no explicit *persistence* is given, an
54 :class:`~lexigram.ai.governance.persistence.InMemoryGovernancePersistence`
55 instance is created automatically using the optional *cache* argument to
56 build a :class:`~lexigram.ai.governance.persistence.RedisGovernancePersistence`
57 when a cache backend is available.
59 Args:
60 config: Governance policy configuration.
61 cache: Optional cache backend used to auto-create a Redis persistence
62 backend. Ignored when *persistence* is supplied explicitly.
63 persistence: Explicit persistence backend. Takes precedence over *cache*.
64 on_soft_limit: Optional async callback invoked when the monthly spend
65 crosses the ``soft_limit_pct`` threshold. Signature:
66 ``async def cb(user_id, current_spend, budget) -> None``.
67 audit_store: Optional audit store for recording governance decisions.
68 When provided, every governance check (allowed or denied) and
69 every cost-tracking call is recorded as an audit event.
70 """
72 def __init__(
73 self,
74 config: GovernanceConfig,
75 cache: CacheBackendProtocol | None = None,
76 persistence: GovernancePersistence | None = None,
77 on_soft_limit: Callable[..., object] | None = None,
78 audit_store: AIAuditStore | None = None,
79 ) -> None:
80 self._config = config
81 self._on_soft_limit = on_soft_limit
82 self._audit_store = audit_store
83 self._background_tasks: set[asyncio.Task[object]] = set()
84 self._resource_registry: ResourceUnitRegistry | None = None
85 self._resource_tracker: ResourceUnitTracker | None = None
87 if persistence is not None:
88 self._persistence = persistence
89 elif cache is not None:
90 from lexigram.ai.governance.persistence import (
91 RedisGovernancePersistence,
92 )
94 self._persistence = cast(
95 "GovernancePersistence", RedisGovernancePersistence(cache)
96 )
97 else:
98 from lexigram.ai.governance.persistence import (
99 InMemoryGovernancePersistence,
100 )
102 self._persistence = cast(
103 "GovernancePersistence", InMemoryGovernancePersistence()
104 )
106 # Initialise resource unit tracker
107 if config.resource_units:
108 from lexigram.ai.governance.resource.registry import (
109 ResourceUnitRegistry,
110 )
111 from lexigram.ai.governance.resource.tracker import (
112 ResourceUnitTracker,
113 )
115 self._resource_registry = ResourceUnitRegistry.from_list(
116 config.resource_units
117 )
118 self._resource_tracker = ResourceUnitTracker(
119 registry=self._resource_registry,
120 persistence=self._persistence,
121 )
122 else:
123 self._resource_registry = None
124 self._resource_tracker = None
126 @property
127 def resource_tracker(self) -> ResourceUnitTracker | None:
128 """The :class:`ResourceUnitTracker` instance, or ``None`` when no
129 resource units are configured.
131 Exposed for DI registration so the same tracker is shared across the
132 application (consume/release calls go through here regardless of
133 whether the caller resolves ``AIGovernanceManager`` or
134 ``ResourceUnitTracker`` from the container).
135 """
136 return self._resource_tracker
138 async def check_request(
139 self,
140 model: str,
141 provider: str,
142 user_id: str | None = None,
143 ) -> bool:
144 """Check if a request is allowed under governance policy.
146 Args:
147 model: Model identifier.
148 provider: Provider name.
149 user_id: Optional user identifier for per-user limits.
151 Returns:
152 True if request is allowed, False if blocked by policy.
153 """
154 if self._config.restricted_models and model in self._config.restricted_models:
155 logger.warning(
156 "governance_model_restricted", model=model, provider=provider
157 )
158 self._emit_audit(
159 "model_denied",
160 model=model,
161 provider=provider,
162 user_id=user_id,
163 status="denied",
164 metadata={"reason": "restricted_model"},
165 )
166 return False
168 if not self.check_model_access(user_id, model):
169 self._emit_audit(
170 "model_denied",
171 model=model,
172 provider=provider,
173 user_id=user_id,
174 status="denied",
175 metadata={"reason": "access_policy"},
176 )
177 return False
179 if self._config.rpm_limit:
180 try:
181 count = await self._persistence.incr_requests(
182 user_id or "global", window=60.0
183 )
184 except _PERSISTENCE_FAILURE_EXCEPTIONS as exc:
185 return self._on_persistence_failure(
186 "rpm_check", user_id or "global", exc
187 )
188 if count > self._config.rpm_limit:
189 logger.warning(
190 "governance_rpm_exceeded",
191 user_id=user_id,
192 rpm_limit=self._config.rpm_limit,
193 )
194 self._emit_audit(
195 "rate_limited",
196 model=model,
197 provider=provider,
198 user_id=user_id,
199 status="denied",
200 metadata={
201 "rpm_limit": self._config.rpm_limit,
202 "current_rpm": count,
203 },
204 )
205 return False
207 return True
209 def check_model_access(self, user_id: str | None, model: str) -> bool:
210 """Check if user is allowed to use the given model.
212 Evaluates per-user ``model_allowlist`` and ``model_denylist`` from
213 :class:`~lexigram.ai.config.GovernanceConfig`. Both support glob
214 patterns (e.g. ``"gpt-4*"``, ``"claude-3-*"``).
216 Logic:
217 1. If ``model_allowlist`` has an entry for *user_id*, the model must
218 match at least one pattern in the allowlist.
219 2. If ``model_denylist`` has an entry for *user_id*, the model must
220 not match any pattern in the denylist.
221 3. When no entry exists for *user_id*, access is allowed.
223 Args:
224 user_id: User identifier, or ``None`` for anonymous / global.
225 model: Model name to check.
227 Returns:
228 True if access is permitted, False if denied.
229 """
230 key = user_id or "global"
232 allowlist = self._config.model_allowlist.get(
233 key
234 ) or self._config.model_allowlist.get("*")
235 if allowlist:
236 if not any(fnmatch.fnmatch(model, pattern) for pattern in allowlist):
237 logger.warning(
238 "governance_model_not_in_allowlist",
239 user_id=user_id,
240 model=model,
241 )
242 return False
244 denylist = self._config.model_denylist.get(
245 key
246 ) or self._config.model_denylist.get("*")
247 if denylist:
248 if any(fnmatch.fnmatch(model, pattern) for pattern in denylist):
249 logger.warning(
250 "governance_model_in_denylist",
251 user_id=user_id,
252 model=model,
253 )
254 return False
256 return True
258 async def check_budget(self, cost: float, user_id: str | None = None) -> bool:
259 """Check if a cost would exceed the monthly budget.
261 Emits a structured warning when the spend crosses the configured
262 ``soft_limit_pct`` threshold and invokes the optional
263 ``on_soft_limit`` callback. Returns ``False`` only when the hard
264 limit (``monthly_budget``) would be exceeded.
266 Args:
267 cost: Estimated cost of the request.
268 user_id: Optional user identifier.
270 Returns:
271 True if within hard budget, False if would exceed.
272 """
273 if not self._config.enforce_budget:
274 return True
275 if self._config.monthly_budget is None:
276 return True
278 try:
279 current = await self._get_monthly_spend(user_id)
280 except _PERSISTENCE_FAILURE_EXCEPTIONS as exc:
281 return self._on_persistence_failure(
282 "budget_check", f"{user_id or 'global'}:{_current_month()}", exc
283 )
284 budget = self._config.monthly_budget
286 # Soft-limit warning (does not block)
287 if (
288 self._config.soft_limit_pct is not None
289 and current + cost >= budget * self._config.soft_limit_pct
290 and current < budget * self._config.soft_limit_pct
291 ):
292 logger.warning(
293 "governance_soft_limit_reached",
294 user_id=user_id,
295 current_spend=current,
296 soft_limit_pct=self._config.soft_limit_pct,
297 monthly_budget=budget,
298 )
299 self._emit_audit(
300 "soft_limit_reached",
301 user_id=user_id,
302 cost=cost,
303 metadata={
304 "current_spend": current,
305 "soft_limit_pct": self._config.soft_limit_pct,
306 "monthly_budget": budget,
307 },
308 )
309 if self._on_soft_limit is not None:
310 import asyncio
311 import inspect
313 result = self._on_soft_limit(user_id, current, budget)
314 if inspect.isawaitable(result):
315 task = asyncio.ensure_future(result)
316 self._background_tasks.add(task)
317 task.add_done_callback(self._background_tasks.discard)
319 allowed = current + cost <= budget
320 if not allowed:
321 logger.warning(
322 "governance_budget_exceeded",
323 user_id=user_id,
324 current_spend=current,
325 request_cost=cost,
326 monthly_budget=budget,
327 )
328 self._emit_audit(
329 "budget_exceeded",
330 user_id=user_id,
331 cost=cost,
332 status="denied",
333 metadata={
334 "current_spend": current,
335 "request_cost": cost,
336 "monthly_budget": budget,
337 },
338 )
339 return allowed
341 async def check_request_budget(
342 self,
343 estimated_cost: float,
344 request_id: str | None = None,
345 ) -> Result[None, GovernanceError]:
346 """Check if a single request cost is within the per-request budget.
348 Validates *estimated_cost* against ``max_request_cost`` (per-request
349 cap) first, then against the monthly budget via :meth:`check_budget`.
351 Args:
352 estimated_cost: Estimated cost in USD for this request.
353 request_id: Optional request identifier for logging context.
355 Returns:
356 ``Ok(None)`` if within all budget limits.
357 ``Err(GovernanceError)`` if either per-request or monthly limit
358 is exceeded.
359 """
360 from lexigram.ai.governance.exceptions import GovernanceError
362 if not self._config.enforce_budget:
363 return Ok(None)
365 if (
366 self._config.max_request_cost is not None
367 and estimated_cost > self._config.max_request_cost
368 ):
369 logger.warning(
370 "governance_request_cost_exceeded",
371 estimated_cost=estimated_cost,
372 max_request_cost=self._config.max_request_cost,
373 request_id=request_id,
374 )
375 self._emit_audit(
376 "request_budget_exceeded",
377 cost=estimated_cost,
378 status="denied",
379 metadata={
380 "estimated_cost": estimated_cost,
381 "max_request_cost": self._config.max_request_cost,
382 "request_id": request_id,
383 },
384 )
385 return Err(
386 GovernanceError(
387 f"Request cost ${estimated_cost:.4f} exceeds per-request "
388 f"limit ${self._config.max_request_cost:.4f}"
389 )
390 )
392 allowed = await self.check_budget(estimated_cost)
393 if not allowed:
394 return Err(
395 GovernanceError(
396 f"Request cost ${estimated_cost:.4f} would exceed monthly budget"
397 )
398 )
400 return Ok(None)
402 async def track_cost(
403 self,
404 cost: float,
405 model: str,
406 user_id: str | None = None,
407 ) -> None:
408 """Record AI usage cost.
410 Args:
411 cost: Cost to record.
412 model: Model that generated the cost.
413 user_id: Optional user identifier.
414 """
415 key = user_id or "global"
416 month_key = f"{key}:{_current_month()}"
417 try:
418 await self._persistence.add_spend(month_key, cost, ttl=32 * 24 * 3600)
419 except _PERSISTENCE_FAILURE_EXCEPTIONS as exc:
420 self._on_persistence_failure("cost_track", month_key, exc)
421 return
422 logger.debug("governance_cost_tracked", cost=cost, model=model, user_id=user_id)
424 def reload_config(self, config: GovernanceConfig) -> None:
425 """Hot-reload governance configuration without restart.
427 Atomically swaps the internal config reference so that subsequent policy
428 checks use the new limits. Does **not** touch persistence state — only
429 the thresholds and rules are updated.
431 Args:
432 config: New governance configuration to apply.
433 """
434 self._config = config
435 logger.info(
436 "governance_config_reloaded",
437 monthly_budget=config.monthly_budget,
438 rpm_limit=config.rpm_limit,
439 soft_limit_pct=config.soft_limit_pct,
440 restricted_models=config.restricted_models,
441 )
442 self._emit_audit(
443 "config_reloaded",
444 metadata={
445 "monthly_budget": config.monthly_budget,
446 "rpm_limit": config.rpm_limit,
447 "soft_limit_pct": config.soft_limit_pct,
448 },
449 )
451 # ------------------------------------------------------------------
452 # Audit helpers
453 # ------------------------------------------------------------------
455 def _emit_audit(
456 self,
457 event_type: str,
458 *,
459 model: str | None = None,
460 provider: str | None = None,
461 user_id: str | None = None,
462 status: str = "success",
463 tokens: int | None = None,
464 cost: float | None = None,
465 latency_ms: float | None = None,
466 metadata: dict[str, object] | None = None,
467 ) -> None:
468 """Fire-and-forget audit event recording.
470 Creates an :class:`~lexigram.ai.governance.audit.AIAuditEvent` and
471 schedules ``record()`` on the audit store without blocking the
472 caller. Silently drops the event when no audit store is configured.
473 """
474 if self._audit_store is None:
475 return
477 from lexigram.ai.governance.audit import AIAuditEvent, AuditEventType
479 event = AIAuditEvent(
480 event_type=AuditEventType(event_type),
481 model=model,
482 provider=provider,
483 user_id=user_id,
484 status=status,
485 tokens=tokens,
486 cost=cost,
487 latency_ms=latency_ms,
488 metadata=dict(metadata) if metadata else {},
489 )
491 import asyncio
493 try:
494 loop = asyncio.get_running_loop()
495 except RuntimeError:
496 return
498 task = loop.create_task(self._audit_store.record(event))
499 self._background_tasks.add(task)
500 task.add_done_callback(self._background_tasks.discard)
502 def _on_persistence_failure(
503 self,
504 operation: str,
505 bucket_key: str,
506 exception: BaseException,
507 ) -> bool:
508 """Apply the configured decision when the persistence backend fails.
510 Logs a warning with the bucket key, exception type, and resulting
511 decision, then returns whether the caller should allow the request
512 (fail-open) or deny it (fail-closed). Infrastructure failures are
513 never re-raised here: agent-executor callers treat a raised
514 governance exception as allow-through, so raising would defeat the
515 fail-closed default.
517 Args:
518 operation: Name of the failing governance operation
519 (e.g. ``"rpm_check"``, ``"budget_check"``, ``"cost_track"``).
520 bucket_key: Bucket key the operation was reading or writing.
521 exception: The exception raised by the persistence backend.
523 Returns:
524 True when the request should be allowed (fail-open configured),
525 False when it must be denied (fail-closed default).
526 """
527 fail_open = self._config.fail_open_on_persistence_error
528 logger.warning(
529 "governance_persistence_unavailable",
530 operation=operation,
531 bucket_key=bucket_key,
532 error_type=type(exception).__name__,
533 decision="allowed" if fail_open else "denied",
534 fail_open=fail_open,
535 )
536 return fail_open
538 async def _get_monthly_spend(self, user_id: str | None) -> float:
539 key = user_id or "global"
540 month_key = f"{key}:{_current_month()}"
541 return await self._persistence.get_spend(month_key)
543 # -- Resource unit delegation --------------------------------------------
545 async def consume_resource(
546 self,
547 tenant_id: str,
548 unit_name: str,
549 amount: float,
550 actor_id: str | None = None,
551 ) -> Result:
552 """Consume *amount* of a resource unit for *tenant_id*.
554 Delegates to :class:`ResourceUnitTracker` if configured.
555 """
556 if self._resource_tracker is None:
557 return Err(self._no_tracker_error(tenant_id, unit_name, amount))
558 return await self._resource_tracker.consume(
559 tenant_id, unit_name, amount, actor_id
560 )
562 async def release_resource(
563 self,
564 tenant_id: str,
565 unit_name: str,
566 amount: float,
567 ) -> None:
568 """Release *amount* of a held resource (INSTANTANEOUS units only)."""
569 if self._resource_tracker is None:
570 return
571 await self._resource_tracker.release(tenant_id, unit_name, amount)
573 async def resource_usage(
574 self,
575 tenant_id: str,
576 unit_name: str,
577 ):
578 """Return current usage snapshot for *tenant_id* + *unit_name*."""
579 if self._resource_tracker is None:
580 from lexigram.contracts.ai.governance.resource_unit import (
581 ResourceUsageSnapshot,
582 )
584 return ResourceUsageSnapshot(
585 tenant_id=tenant_id,
586 unit_name=unit_name,
587 current=0.0,
588 limit=0.0,
589 )
590 return await self._resource_tracker.usage(tenant_id, unit_name)
592 def _no_tracker_error(
593 self, tenant_id: str, unit_name: str, amount: float
594 ) -> ResourceExhaustedError:
595 return ResourceExhaustedError(
596 tenant_id=tenant_id,
597 unit_name=unit_name,
598 limit=0,
599 current=0,
600 )
603def _current_month() -> str:
604 """Return current year-month string for cache key scoping."""
605 from datetime import UTC, datetime
607 return datetime.now(UTC).strftime("%Y-%m")
610__all__ = ["AIGovernanceManager"]