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

"""Pydantic models for confidence calibration pipeline.

Defines feature vectors, calibration artifact metadata, and
confidence results used throughout the confidence pipeline.
"""

from __future__ import annotations

from datetime import datetime, timezone
from typing import Literal

from pydantic import BaseModel, Field, field_validator


class ConfidenceFeatures(BaseModel):
    """Feature vector for confidence estimation.

    Each feature is a normalized float derived from upstream pipeline
    stages: extraction, resolution, evidence verification, sentiment,
    and cross-stage agreement analysis.
    """

    entity_span_score: float = Field(
        ge=0.0,
        le=1.0,
        description="Best entity span confidence from specialist extractor.",
    )
    alias_resolution_margin: float = Field(
        ge=0.0,
        le=1.0,
        description="Gap between top-2 alias candidates. 1.0 = unambiguous.",
    )
    numeric_parser_validity: float = Field(
        ge=0.0,
        le=1.0,
        description="Fraction of numeric facts that passed parser validation.",
    )
    evidence_coverage: float = Field(
        ge=0.0,
        le=1.0,
        description="Fraction of extracted facts backed by valid evidence spans.",
    )
    relation_score: float = Field(
        ge=0.0,
        le=1.0,
        description="Average confidence of extracted relations.",
    )
    sentiment_calibration_confidence: float = Field(
        ge=0.0,
        le=1.0,
        description="Calibrated sentiment model confidence (max class probability).",
    )
    cross_stage_agreement: float = Field(
        ge=0.0,
        le=1.0,
        description="Agreement ratio between independently derived facts across stages.",
    )
    duplicate_novelty_certainty: float = Field(
        ge=0.0,
        le=1.0,
        description="Certainty of the novelty/duplicate classification.",
    )
    document_completeness: float = Field(
        ge=0.0,
        le=1.0,
        description="Fraction of expected schema fields that were populated.",
    )
    document_type: str = Field(
        description="Document type (news, filing, transcript, press_release, macro_event).",
    )
    known_hard_case_patterns: float = Field(
        ge=0.0,
        le=1.0,
        description="Score indicating presence of known hard patterns (multi-company, contradictions).",
    )

    @field_validator("document_type")
    @classmethod
    def document_type_valid(cls, v: str) -> str:
        valid_types = {
            "news",
            "filing",
            "transcript",
            "press_release",
            "macro_event",
            "unknown",
        }
        if v not in valid_types:
            raise ValueError(f"document_type must be one of {valid_types}, got '{v}'")
        return v

    def to_vector(self) -> list[float]:
        """Convert features to a flat numeric vector for calibration models.

        document_type is encoded as a categorical index.
        """
        type_map = {
            "news": 0.0,
            "filing": 0.2,
            "transcript": 0.4,
            "press_release": 0.6,
            "macro_event": 0.8,
            "unknown": 1.0,
        }
        return [
            self.entity_span_score,
            self.alias_resolution_margin,
            self.numeric_parser_validity,
            self.evidence_coverage,
            self.relation_score,
            self.sentiment_calibration_confidence,
            self.cross_stage_agreement,
            self.duplicate_novelty_certainty,
            self.document_completeness,
            type_map.get(self.document_type, 1.0),
            self.known_hard_case_patterns,
        ]


class CalibrationArtifactMetadata(BaseModel):
    """Metadata for a versioned calibration artifact.

    Stored alongside the serialized calibrator to track provenance,
    training conditions, and quality metrics.
    """

    version: str = Field(description="Artifact version string (e.g., 'v1.0.0').")
    method: Literal["isotonic", "platt"] = Field(
        description="Calibration method used."
    )
    training_count: int = Field(
        ge=0, description="Number of samples used for training."
    )
    training_range: str = Field(
        description="Date range of training data (e.g., '2024-01-01 to 2024-06-30')."
    )
    ece: float = Field(
        ge=0.0,
        le=1.0,
        description="Expected Calibration Error on held-out data.",
    )
    brier_score: float = Field(
        ge=0.0,
        le=1.0,
        description="Brier score on held-out data.",
    )
    created_at: datetime = Field(
        default_factory=lambda: datetime.now(tz=timezone.utc),
        description="When the artifact was created.",
    )


class ConfidenceResult(BaseModel):
    """Final confidence output for an extraction record.

    Contains the calibrated probability, feature breakdown, and
    metadata about whether calibration was applied or defaults used.
    """

    probability: float = Field(
        ge=0.0,
        le=1.0,
        description="Calibrated probability of extraction correctness.",
    )
    features_used: list[str] = Field(
        description="Names of features that contributed to this confidence score.",
    )
    is_calibrated: bool = Field(
        default=True,
        description="Whether a trained calibrator was used (vs conservative default).",
    )
    under_calibrated: bool = Field(
        default=False,
        description="True if class has insufficient calibration data and conservative default was applied.",
    )
    calibration_version: str = Field(
        default="uncalibrated",
        description="Version of the calibration artifact used.",
    )