stonks-oracle/services/intelligence_pipeline_v3/active_learning/exporter.py

"""Active learning data exporter.

Selects training examples from low-confidence, conflicting, adjudicated,
and reviewer-corrected cases. Applies content policy filters and exports
in a versioned format with source spans, labels, relations, decisions,
and provenance.
"""

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 SelectionCriteria(str, enum.Enum):
    """Why a case was selected for active learning."""

    LOW_CONFIDENCE = "low_confidence"
    CONFLICTING = "conflicting"
    ADJUDICATED = "adjudicated"
    REVIEWER_CORRECTED = "reviewer_corrected"
    HIGH_DISAGREEMENT = "high_disagreement"
    NOVEL_PATTERN = "novel_pattern"


class ContentPolicy(str, enum.Enum):
    """Content policy levels for export filtering."""

    ALLOW = "allow"
    REDACT_PII = "redact_pii"
    EXCLUDE = "exclude"


@dataclass(frozen=True)
class ExportRecord:
    """A single active-learning export record.

    Contains source text spans, schema labels, relations, adjudicator
    decisions, reviewer corrections, and full provenance.
    """

    record_id: UUID
    document_id: str
    selection_criteria: SelectionCriteria
    export_version: str
    timestamp: datetime

    # Source content
    source_spans: list[dict[str, Any]]  # text, start, end, chunk_id
    document_type: str = ""

    # Labels and annotations
    entity_labels: list[dict[str, Any]] = field(default_factory=list)
    relation_labels: list[dict[str, Any]] = field(default_factory=list)
    event_labels: list[dict[str, Any]] = field(default_factory=list)
    fact_labels: list[dict[str, Any]] = field(default_factory=list)

    # Decisions and corrections
    adjudicator_decisions: list[dict[str, Any]] = field(default_factory=list)
    reviewer_corrections: list[dict[str, Any]] = field(default_factory=list)

    # Provenance
    pipeline_run_id: UUID | None = None
    model_versions: dict[str, str] = field(default_factory=dict)
    confidence_scores: dict[str, float] = field(default_factory=dict)


@dataclass
class ExportConfig:
    """Configuration for active learning export."""

    export_version: str = "1.0"
    min_confidence_threshold: float = 0.5  # Select cases below this
    max_export_count: int = 1000
    include_adjudicated: bool = True
    include_corrections: bool = True
    include_low_confidence: bool = True
    include_conflicting: bool = True
    content_policy: ContentPolicy = ContentPolicy.REDACT_PII
    excluded_fields: set[str] = field(default_factory=set)
    sensitive_patterns: list[str] = field(default_factory=list)


@dataclass
class ActiveLearningExporter:
    """Exports selected cases for specialist model training.

    Applies selection criteria, content policy filtering, and
    produces versioned export datasets with full provenance.
    """

    config: ExportConfig
    _records: list[ExportRecord] = field(default_factory=list)
    _excluded_count: int = 0

    def select_record(
        self,
        document_id: str,
        criteria: SelectionCriteria,
        source_spans: list[dict[str, Any]],
        document_type: str = "",
        entity_labels: list[dict[str, Any]] | None = None,
        relation_labels: list[dict[str, Any]] | None = None,
        event_labels: list[dict[str, Any]] | None = None,
        fact_labels: list[dict[str, Any]] | None = None,
        adjudicator_decisions: list[dict[str, Any]] | None = None,
        reviewer_corrections: list[dict[str, Any]] | None = None,
        pipeline_run_id: UUID | None = None,
        model_versions: dict[str, str] | None = None,
        confidence_scores: dict[str, float] | None = None,
    ) -> ExportRecord | None:
        """Select a case for export, applying content policy.

        Returns None if the case is excluded by policy.
        """
        if len(self._records) >= self.config.max_export_count:
            return None

        # Apply content policy
        filtered_spans = self._apply_content_policy(source_spans)
        if not filtered_spans:
            self._excluded_count += 1
            return None

        record = ExportRecord(
            record_id=uuid4(),
            document_id=document_id,
            selection_criteria=criteria,
            export_version=self.config.export_version,
            timestamp=datetime.now(timezone.utc),
            source_spans=filtered_spans,
            document_type=document_type,
            entity_labels=entity_labels or [],
            relation_labels=relation_labels or [],
            event_labels=event_labels or [],
            fact_labels=fact_labels or [],
            adjudicator_decisions=adjudicator_decisions or [],
            reviewer_corrections=reviewer_corrections or [],
            pipeline_run_id=pipeline_run_id,
            model_versions=model_versions or {},
            confidence_scores=confidence_scores or {},
        )
        self._records.append(record)
        return record

    def _apply_content_policy(
        self, spans: list[dict[str, Any]]
    ) -> list[dict[str, Any]]:
        """Apply content policy filtering to source spans."""
        if self.config.content_policy == ContentPolicy.EXCLUDE:
            # Check for sensitive content
            for span in spans:
                text = span.get("text", "")
                if self._contains_sensitive(text):
                    return []  # Exclude entire record

        if self.config.content_policy == ContentPolicy.REDACT_PII:
            return [self._redact_span(span) for span in spans]

        return spans

    def _contains_sensitive(self, text: str) -> bool:
        """Check if text contains sensitive content per policy."""
        for pattern in self.config.sensitive_patterns:
            if pattern.lower() in text.lower():
                return True
        return False

    def _redact_span(self, span: dict[str, Any]) -> dict[str, Any]:
        """Redact PII from a span while preserving structure."""
        # In production, this would use NER-based PII detection
        # For now, preserve the span but mark it as redacted if needed
        return {**span, "content_policy_applied": "redact_pii"}

    @property
    def records(self) -> list[ExportRecord]:
        return list(self._records)

    @property
    def total_exported(self) -> int:
        return len(self._records)

    @property
    def total_excluded(self) -> int:
        return self._excluded_count

    def export_manifest(self) -> dict[str, Any]:
        """Generate export manifest with metadata."""
        criteria_counts: dict[str, int] = {}
        for r in self._records:
            key = r.selection_criteria.value
            criteria_counts[key] = criteria_counts.get(key, 0) + 1

        return {
            "export_version": self.config.export_version,
            "exported_at": datetime.now(timezone.utc).isoformat(),
            "total_records": len(self._records),
            "excluded_count": self._excluded_count,
            "content_policy": self.config.content_policy.value,
            "selection_criteria_distribution": criteria_counts,
        }