"""Repositories for Audit & Observability and the Usage Ledger.""" from __future__ import annotations from datetime import datetime from sqlalchemy import case, func, select from sqlalchemy.ext.asyncio import AsyncSession from app.application.ports.audit import UsageAggregate, UsageBreakdown, UsageRecord from app.domain.ai_control.models import RateUnit from app.domain.audit.models import AuditEvent, ErrorEvent, QualityMetric, Span, Trace from app.infrastructure.db.models import ( AuditEventRow, ErrorEventRow, QualityMetricRow, SpanRow, TraceRow, UsageEventRow, ) def _audit_to_domain(row: AuditEventRow) -> AuditEvent: return AuditEvent( id=row.id, organization_id=row.organization_id, actor_user_id=row.actor_user_id, action=row.action, target_type=row.target_type, target_id=row.target_id, metadata_json=row.metadata_json or {}, ip_address=row.ip_address, user_agent=row.user_agent, created_at=row.created_at, ) def _error_to_domain(row: ErrorEventRow) -> ErrorEvent: return ErrorEvent( id=row.id, organization_id=row.organization_id, run_id=row.run_id, kind=row.kind, message=row.message, stack=row.stack, metadata_json=row.metadata_json or {}, source=row.source, severity=row.severity, error_type=row.error_type, fingerprint=row.fingerprint, status_code=row.status_code, method=row.method, path=row.path, user_id=row.user_id, request_id=row.request_id, ip=row.ip, user_agent=row.user_agent, resolved_at=row.resolved_at, resolved_note=row.resolved_note, created_at=row.created_at, ) class AuditRepository: """Concrete ``AuditPort``. Append-only.""" def __init__(self, session: AsyncSession) -> None: self.session = session async def record(self, *, event: AuditEvent) -> AuditEvent: row = AuditEventRow( id=event.id, organization_id=event.organization_id, actor_user_id=event.actor_user_id, action=event.action, target_type=event.target_type, target_id=event.target_id, metadata_json=event.metadata_json, ip_address=event.ip_address, user_agent=event.user_agent, created_at=event.created_at, ) self.session.add(row) await self.session.flush() return _audit_to_domain(row) async def list_for_organization( self, *, organization_id: str, limit: int = 100, offset: int = 0 ) -> list[AuditEvent]: stmt = ( select(AuditEventRow) .where(AuditEventRow.organization_id == organization_id) .order_by(AuditEventRow.created_at.desc()) .limit(min(limit, 500)) .offset(offset) ) return [_audit_to_domain(r) for r in (await self.session.execute(stmt)).scalars()] async def list_platform(self, *, limit: int = 100, offset: int = 0) -> list[AuditEvent]: stmt = ( select(AuditEventRow) .order_by(AuditEventRow.created_at.desc()) .limit(min(limit, 500)) .offset(offset) ) return [_audit_to_domain(r) for r in (await self.session.execute(stmt)).scalars()] class UsageLedgerRepository: """Concrete ``UsageLedgerPort``. The ledger is the billing source of truth.""" def __init__(self, session: AsyncSession) -> None: self.session = session async def record(self, *, usage: UsageRecord) -> None: total = usage.total_price or (usage.quantity * usage.unit_price) row = UsageEventRow( organization_id=usage.organization_id, project_id=usage.project_id, user_id=usage.user_id, employee_id=usage.employee_id, run_id=usage.run_id, unit=usage.unit.value, quantity=usage.quantity, unit_price=usage.unit_price, internal_cost_estimate=usage.internal_cost_estimate, total_price=total, metadata_json=usage.metadata_json, ) self.session.add(row) await self.session.flush() async def breakdown( self, *, organization_id: str, since: datetime | None = None, until: datetime | None = None, ) -> UsageBreakdown: filters = [UsageEventRow.organization_id == organization_id] if since is not None: filters.append(UsageEventRow.created_at >= since) if until is not None: filters.append(UsageEventRow.created_at <= until) by_unit_stmt = ( select( UsageEventRow.unit, func.coalesce(func.sum(UsageEventRow.quantity), 0.0), func.coalesce(func.sum(UsageEventRow.total_price), 0.0), func.coalesce(func.sum(UsageEventRow.internal_cost_estimate), 0.0), ) .where(*filters) .group_by(UsageEventRow.unit) ) by_unit = [ UsageAggregate( unit=RateUnit(unit), quantity=float(qty), total_price=float(price), internal_cost_estimate=float(cost), ) for unit, qty, price, cost in (await self.session.execute(by_unit_stmt)).all() ] emp_stmt = ( select( UsageEventRow.employee_id, func.coalesce(func.sum(UsageEventRow.total_price), 0.0), ) .where(*filters, UsageEventRow.employee_id.is_not(None)) .group_by(UsageEventRow.employee_id) ) by_employee = { str(emp): float(total) for emp, total in (await self.session.execute(emp_stmt)).all() } proj_stmt = ( select( UsageEventRow.project_id, func.coalesce(func.sum(UsageEventRow.total_price), 0.0), ) .where(*filters, UsageEventRow.project_id.is_not(None)) .group_by(UsageEventRow.project_id) ) by_project = { str(proj): float(total) for proj, total in (await self.session.execute(proj_stmt)).all() } return UsageBreakdown( by_unit=by_unit, by_employee=by_employee, by_project=by_project, total_price=sum(a.total_price for a in by_unit), total_cost=sum(a.internal_cost_estimate for a in by_unit), ) async def list_events( self, *, organization_id: str, limit: int = 100, offset: int = 0 ) -> list[UsageEventRow]: stmt = ( select(UsageEventRow) .where(UsageEventRow.organization_id == organization_id) .order_by(UsageEventRow.created_at.desc()) .limit(min(limit, 500)) .offset(offset) ) return list((await self.session.execute(stmt)).scalars()) async def platform_totals(self) -> list[tuple[str, float, float]]: """``(organization_id, quantity, total_price)`` across every tenant. Platform-owner console only. """ stmt = ( select( UsageEventRow.organization_id, func.coalesce(func.sum(UsageEventRow.quantity), 0.0), func.coalesce(func.sum(UsageEventRow.total_price), 0.0), ) .group_by(UsageEventRow.organization_id) .order_by(func.sum(UsageEventRow.total_price).desc()) ) rows = (await self.session.execute(stmt)).all() return [(str(org), float(qty), float(price)) for org, qty, price in rows] class ObservabilityRepository: def __init__(self, session: AsyncSession) -> None: self.session = session async def create_trace(self, trace: Trace) -> Trace: self.session.add( TraceRow( id=trace.id, organization_id=trace.organization_id, run_id=trace.run_id, name=trace.name, created_at=trace.created_at, ) ) await self.session.flush() return trace async def add_span(self, span: Span) -> Span: self.session.add( SpanRow( id=span.id, trace_id=span.trace_id, parent_span_id=span.parent_span_id, name=span.name, kind=span.kind, attributes_json=span.attributes_json, started_at=span.started_at, ended_at=span.ended_at, ) ) await self.session.flush() return span async def record_error(self, event: ErrorEvent) -> None: self.session.add( ErrorEventRow( id=event.id, organization_id=event.organization_id, run_id=event.run_id, kind=event.kind, message=event.message, stack=event.stack, metadata_json=event.metadata_json, source=event.source, severity=event.severity, error_type=event.error_type, fingerprint=event.fingerprint, status_code=event.status_code, method=event.method, path=event.path, user_id=event.user_id, request_id=event.request_id, ip=event.ip, user_agent=event.user_agent, resolved_at=event.resolved_at, resolved_note=event.resolved_note, created_at=event.created_at, ) ) await self.session.flush() async def record_metric(self, metric: QualityMetric) -> None: self.session.add( QualityMetricRow( id=metric.id, organization_id=metric.organization_id, name=metric.name, value=metric.value, dimensions_json=metric.dimensions_json, created_at=metric.created_at, ) ) await self.session.flush() async def list_spans(self, *, trace_id: str) -> list[Span]: stmt = select(SpanRow).where(SpanRow.trace_id == trace_id).order_by(SpanRow.started_at) return [ Span( id=r.id, trace_id=r.trace_id, parent_span_id=r.parent_span_id, name=r.name, kind=r.kind, attributes_json=r.attributes_json or {}, started_at=r.started_at, ended_at=r.ended_at, ) for r in (await self.session.execute(stmt)).scalars() ] # -- error console -------------------------------------------------------- def _error_filters( self, *, severity: str | None = None, source: str | None = None, status: str | None = None, kind: str | None = None, organization_id: str | None = None, fingerprint: str | None = None, since: datetime | None = None, search: str | None = None, ) -> list: """Build the WHERE clauses shared by the list and count queries.""" clauses: list = [] if severity: clauses.append(ErrorEventRow.severity == severity) if source: clauses.append(ErrorEventRow.source == source) if status == "open": clauses.append(ErrorEventRow.resolved_at.is_(None)) elif status == "resolved": clauses.append(ErrorEventRow.resolved_at.is_not(None)) if kind: clauses.append(ErrorEventRow.kind == kind) if organization_id: clauses.append(ErrorEventRow.organization_id == organization_id) if fingerprint: clauses.append(ErrorEventRow.fingerprint == fingerprint) if since is not None: clauses.append(ErrorEventRow.created_at >= since) if search: clauses.append(ErrorEventRow.message.ilike(f"%{search}%")) return clauses async def list_errors( self, *, limit: int = 50, offset: int = 0, **filters, ) -> list[ErrorEvent]: stmt = ( select(ErrorEventRow) .where(*self._error_filters(**filters)) # Newest first: the failure someone is reacting to is almost always # the most recent one. .order_by(ErrorEventRow.created_at.desc()) .limit(limit) .offset(offset) ) return [_error_to_domain(row) for row in (await self.session.execute(stmt)).scalars()] async def count_errors(self, **filters) -> int: stmt = select(func.count()).select_from(ErrorEventRow).where( *self._error_filters(**filters) ) return int((await self.session.execute(stmt)).scalar_one()) async def get_error(self, error_id: str) -> ErrorEvent | None: row = ( await self.session.execute( select(ErrorEventRow).where(ErrorEventRow.id == error_id) ) ).scalar_one_or_none() return _error_to_domain(row) if row is not None else None async def list_error_groups( self, *, limit: int = 50, since: datetime | None = None, status: str | None = None, severity: str | None = None, source: str | None = None, organization_id: str | None = None, ) -> list[dict]: """One row per distinct bug, newest activity first. This is the view that makes the console usable: a bug that fired 4 000 times is one line with a count, not 4 000 lines. ``open_count`` is what tells you whether anyone has actually dealt with it. """ clauses: list = [ErrorEventRow.fingerprint != ""] if since is not None: clauses.append(ErrorEventRow.created_at >= since) if severity: clauses.append(ErrorEventRow.severity == severity) if source: clauses.append(ErrorEventRow.source == source) if organization_id: clauses.append(ErrorEventRow.organization_id == organization_id) open_count = func.sum( case((ErrorEventRow.resolved_at.is_(None), 1), else_=0) ).label("open_count") stmt = ( select( ErrorEventRow.fingerprint, func.count().label("occurrences"), open_count, func.min(ErrorEventRow.created_at).label("first_seen"), func.max(ErrorEventRow.created_at).label("last_seen"), func.max(ErrorEventRow.kind).label("kind"), func.max(ErrorEventRow.error_type).label("error_type"), func.max(ErrorEventRow.severity).label("severity"), func.max(ErrorEventRow.source).label("source"), func.max(ErrorEventRow.path).label("path"), func.max(ErrorEventRow.message).label("sample_message"), ) .where(*clauses) .group_by(ErrorEventRow.fingerprint) .order_by(func.max(ErrorEventRow.created_at).desc()) .limit(limit) ) rows = (await self.session.execute(stmt)).all() groups = [ { "fingerprint": row.fingerprint, "occurrences": int(row.occurrences), "open_count": int(row.open_count or 0), "first_seen": row.first_seen, "last_seen": row.last_seen, "kind": row.kind or "", "error_type": row.error_type or "", "severity": row.severity or "", "source": row.source or "", "path": row.path or "", "sample_message": row.sample_message or "", } for row in rows ] if status == "open": return [g for g in groups if g["open_count"] > 0] if status == "resolved": return [g for g in groups if g["open_count"] == 0] return groups async def resolve_error( self, error_id: str, *, note: str, resolved_at: datetime ) -> ErrorEvent | None: row = ( await self.session.execute( select(ErrorEventRow).where(ErrorEventRow.id == error_id) ) ).scalar_one_or_none() if row is None: return None row.resolved_at = resolved_at row.resolved_note = note await self.session.flush() return _error_to_domain(row) async def resolve_error_group( self, fingerprint: str, *, note: str, resolved_at: datetime ) -> int: """Resolve every open occurrence of one bug. Returns how many changed.""" stmt = select(ErrorEventRow).where( ErrorEventRow.fingerprint == fingerprint, ErrorEventRow.resolved_at.is_(None), ) rows = list((await self.session.execute(stmt)).scalars()) for row in rows: row.resolved_at = resolved_at row.resolved_note = note await self.session.flush() return len(rows) async def reopen_error_group(self, fingerprint: str) -> int: """Undo a resolution — for when the bug comes back.""" stmt = select(ErrorEventRow).where( ErrorEventRow.fingerprint == fingerprint, ErrorEventRow.resolved_at.is_not(None), ) rows = list((await self.session.execute(stmt)).scalars()) for row in rows: row.resolved_at = None row.resolved_note = "" await self.session.flush() return len(rows) async def count_open_by_severity(self) -> dict[str, int]: """Unresolved counts per severity, for the console's header tiles.""" stmt = ( select(ErrorEventRow.severity, func.count()) .where(ErrorEventRow.resolved_at.is_(None)) .group_by(ErrorEventRow.severity) ) return { str(severity): int(count) for severity, count in (await self.session.execute(stmt)).all() }