stonks-oracle/services/intelligence_pipeline_v3/audit/models.py

"""Audit and review data models.

Supports evidence display with offsets, specialist probabilities,
routing reasons, adjudicator decisions, impact-model outputs,
and immutable reviewer correction events.
"""

from __future__ import annotations

import enum
from dataclasses import dataclass, field
from datetime import datetime, timezone
from typing import Any
from uuid import UUID, uuid4


class ReviewStatus(str, enum.Enum):
    """Status of a document's review."""

    PENDING = "pending"
    REVIEWED = "reviewed"
    CORRECTED = "corrected"
    CONFIRMED = "confirmed"


class CorrectionType(str, enum.Enum):
    """Types of reviewer corrections."""

    CORRECT = "correct"
    INCORRECT = "incorrect"
    UNSUPPORTED = "unsupported"
    AMBIGUOUS = "ambiguous"
    VALUE_OVERRIDE = "value_override"


@dataclass(frozen=True)
class CorrectionEvent:
    """Immutable reviewer correction event.

    Corrections are append-only audit events. They feed the active-learning
    dataset only through the approved export process.
    """

    event_id: UUID
    record_id: UUID
    field_name: str
    correction_type: CorrectionType
    original_value: Any
    corrected_value: Any | None
    reviewer_id: str
    timestamp: datetime
    notes: str = ""

    @classmethod
    def create(
        cls,
        record_id: UUID,
        field_name: str,
        correction_type: CorrectionType,
        original_value: Any,
        corrected_value: Any | None = None,
        reviewer_id: str = "",
        notes: str = "",
    ) -> CorrectionEvent:
        return cls(
            event_id=uuid4(),
            record_id=record_id,
            field_name=field_name,
            correction_type=correction_type,
            original_value=original_value,
            corrected_value=corrected_value,
            reviewer_id=reviewer_id,
            timestamp=datetime.now(timezone.utc),
            notes=notes,
        )


@dataclass
class AuditRecord:
    """Complete audit record for a processed document.

    Contains source evidence, specialist outputs, routing reasons,
    adjudicator decisions, and impact predictions displayed separately.
    """

    record_id: UUID
    document_id: str
    run_id: UUID
    timestamp: datetime

    # Source evidence with offsets
    evidence_spans: list[dict[str, Any]] = field(default_factory=list)

    # Specialist stage outputs (probabilities, scores)
    specialist_outputs: dict[str, Any] = field(default_factory=dict)

    # Routing decision and reasons
    routing_reasons: list[str] = field(default_factory=list)
    route_decision: str = ""

    # Adjudicator decision (if applicable)
    adjudicator_decision: dict[str, Any] | None = None

    # Impact model outputs
    impact_outputs: dict[str, Any] = field(default_factory=dict)

    # Model lineage
    lineage: dict[str, Any] = field(default_factory=dict)

    # Review status
    review_status: ReviewStatus = ReviewStatus.PENDING
    corrections: list[CorrectionEvent] = field(default_factory=list)

    @classmethod
    def create(
        cls,
        document_id: str,
        run_id: UUID,
        evidence_spans: list[dict[str, Any]] | None = None,
        specialist_outputs: dict[str, Any] | None = None,
        routing_reasons: list[str] | None = None,
        route_decision: str = "",
        adjudicator_decision: dict[str, Any] | None = None,
        impact_outputs: dict[str, Any] | None = None,
        lineage: dict[str, Any] | None = None,
    ) -> AuditRecord:
        return cls(
            record_id=uuid4(),
            document_id=document_id,
            run_id=run_id,
            timestamp=datetime.now(timezone.utc),
            evidence_spans=evidence_spans or [],
            specialist_outputs=specialist_outputs or {},
            routing_reasons=routing_reasons or [],
            route_decision=route_decision,
            adjudicator_decision=adjudicator_decision,
            impact_outputs=impact_outputs or {},
            lineage=lineage or {},
        )

    def add_correction(self, correction: CorrectionEvent) -> None:
        """Add an immutable correction event."""
        self.corrections.append(correction)
        self.review_status = ReviewStatus.CORRECTED

    def mark_reviewed(self) -> None:
        """Mark the record as reviewed without corrections."""
        if self.review_status == ReviewStatus.PENDING:
            self.review_status = ReviewStatus.REVIEWED

    def mark_confirmed(self) -> None:
        """Mark the record as confirmed correct."""
        self.review_status = ReviewStatus.CONFIRMED


@dataclass
class ReviewFilter:
    """Filter criteria for audit records.

    Supports filtering by confidence, unsupported claims, adjudication
    status, review status, and date ranges.
    """

    min_confidence: float | None = None
    max_confidence: float | None = None
    has_unsupported_claims: bool | None = None
    is_adjudicated: bool | None = None
    review_status: ReviewStatus | None = None
    document_type: str | None = None
    company_id: UUID | None = None
    from_date: datetime | None = None
    to_date: datetime | None = None

    def matches(self, record: AuditRecord) -> bool:
        """Check if a record matches this filter."""
        if self.is_adjudicated is not None:
            has_adj = record.adjudicator_decision is not None
            if has_adj != self.is_adjudicated:
                return False

        if self.review_status is not None:
            if record.review_status != self.review_status:
                return False

        if self.from_date is not None:
            if record.timestamp < self.from_date:
                return False

        if self.to_date is not None:
            if record.timestamp > self.to_date:
                return False

        if self.has_unsupported_claims is not None:
            has_unsupported = any(
                c.correction_type == CorrectionType.UNSUPPORTED
                for c in record.corrections
            )
            if has_unsupported != self.has_unsupported_claims:
                return False

        return True