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

"""Confidence feature extraction from upstream pipeline stages.

Computes field-level features from extraction, resolution, evidence,
sentiment, and cross-stage agreement to produce a ConfidenceFeatures
vector for calibration or conservative defaults.
"""

from __future__ import annotations

import logging
from dataclasses import dataclass

from services.intelligence_pipeline_v3.confidence.models import ConfidenceFeatures

logger = logging.getLogger(__name__)


@dataclass
class ExtractionStageResult:
    """Subset of extraction results relevant to confidence features.

    This is an adapter interface — callers populate it from
    the full extraction/specialist output.
    """

    entity_scores: list[float]
    """Per-entity confidence scores from specialist extractor."""

    relation_scores: list[float]
    """Per-relation confidence scores."""

    total_facts: int
    """Total facts extracted."""

    valid_numeric_facts: int
    """Facts that passed deterministic parser validation."""

    populated_fields: int
    """Schema fields that have values."""

    expected_fields: int
    """Total expected schema fields for this document type."""


@dataclass
class ResolutionStageResult:
    """Subset of resolution results relevant to confidence features."""

    ambiguity_margins: list[float]
    """Per-mention ambiguity margins (gap between top-2 candidates)."""


@dataclass
class EvidenceStageResult:
    """Subset of evidence verification results relevant to confidence features."""

    total_claims: int
    """Total extracted claims/facts."""

    supported_claims: int
    """Claims backed by valid evidence spans."""


@dataclass
class SentimentStageResult:
    """Subset of sentiment results relevant to confidence features."""

    max_class_probabilities: list[float]
    """Per-company maximum class probability after calibration."""

    calibration_version: str
    """Version of sentiment calibration artifact used."""


@dataclass
class AgreementStageResult:
    """Cross-stage agreement analysis results."""

    agreement_ratio: float
    """Fraction of facts that agree across independent extraction paths."""

    novelty_certainty: float
    """Certainty of the novelty/duplicate classification (0-1)."""

    hard_case_score: float
    """Score indicating presence of known difficult patterns."""


class ConfidenceFeatureExtractor:
    """Extracts confidence features from upstream pipeline stage results.

    Produces a normalized ConfidenceFeatures vector that can be passed
    to the calibrator or used to determine conservative defaults.
    """

    def extract_features(
        self,
        extraction_result: ExtractionStageResult,
        resolution_result: ResolutionStageResult,
        evidence_result: EvidenceStageResult,
        sentiment_result: SentimentStageResult,
        agreement_result: AgreementStageResult | None = None,
        document_type: str = "unknown",
    ) -> ConfidenceFeatures:
        """Compute confidence features from all upstream stage results.

        Parameters
        ----------
        extraction_result
            Entity/relation/fact extraction outputs with scores.
        resolution_result
            Symbol resolution outputs with ambiguity margins.
        evidence_result
            Evidence verification outputs with coverage stats.
        sentiment_result
            Sentiment classification outputs with calibrated probabilities.
        agreement_result
            Optional cross-stage agreement analysis. Defaults used if None.
        document_type
            Document type string for type-specific calibration.

        Returns
        -------
        ConfidenceFeatures
            Normalized feature vector ready for calibration.
        """
        # Entity span score: average of entity scores, or 0 if none
        entity_span_score = (
            sum(extraction_result.entity_scores) / len(extraction_result.entity_scores)
            if extraction_result.entity_scores
            else 0.0
        )

        # Alias resolution margin: average of per-mention margins
        alias_resolution_margin = (
            sum(resolution_result.ambiguity_margins)
            / len(resolution_result.ambiguity_margins)
            if resolution_result.ambiguity_margins
            else 1.0  # No ambiguity if no mentions to resolve
        )

        # Numeric parser validity: fraction of valid numeric facts
        numeric_parser_validity = (
            extraction_result.valid_numeric_facts / extraction_result.total_facts
            if extraction_result.total_facts > 0
            else 1.0  # No numeric facts = no parser failures
        )

        # Evidence coverage: fraction of claims with valid evidence
        evidence_coverage = (
            evidence_result.supported_claims / evidence_result.total_claims
            if evidence_result.total_claims > 0
            else 0.0
        )

        # Relation score: average relation confidence
        relation_score = (
            sum(extraction_result.relation_scores)
            / len(extraction_result.relation_scores)
            if extraction_result.relation_scores
            else 0.0
        )

        # Sentiment calibration confidence: average max class probability
        sentiment_calibration_confidence = (
            sum(sentiment_result.max_class_probabilities)
            / len(sentiment_result.max_class_probabilities)
            if sentiment_result.max_class_probabilities
            else 0.5  # Neutral default when no sentiment data
        )

        # Document completeness: fraction of expected fields populated
        document_completeness = (
            extraction_result.populated_fields / extraction_result.expected_fields
            if extraction_result.expected_fields > 0
            else 0.0
        )

        # Cross-stage agreement features (use defaults if not provided)
        if agreement_result is not None:
            cross_stage_agreement = agreement_result.agreement_ratio
            duplicate_novelty_certainty = agreement_result.novelty_certainty
            known_hard_case_patterns = agreement_result.hard_case_score
        else:
            cross_stage_agreement = 0.5  # Neutral default
            duplicate_novelty_certainty = 0.5
            known_hard_case_patterns = 0.0

        # Validate document type
        valid_types = {
            "news",
            "filing",
            "transcript",
            "press_release",
            "macro_event",
            "unknown",
        }
        if document_type not in valid_types:
            logger.warning(
                "Unknown document_type '%s', defaulting to 'unknown'", document_type
            )
            document_type = "unknown"

        return ConfidenceFeatures(
            entity_span_score=_clamp(entity_span_score),
            alias_resolution_margin=_clamp(alias_resolution_margin),
            numeric_parser_validity=_clamp(numeric_parser_validity),
            evidence_coverage=_clamp(evidence_coverage),
            relation_score=_clamp(relation_score),
            sentiment_calibration_confidence=_clamp(sentiment_calibration_confidence),
            cross_stage_agreement=_clamp(cross_stage_agreement),
            duplicate_novelty_certainty=_clamp(duplicate_novelty_certainty),
            document_completeness=_clamp(document_completeness),
            document_type=document_type,
            known_hard_case_patterns=_clamp(known_hard_case_patterns),
        )


def _clamp(value: float, low: float = 0.0, high: float = 1.0) -> float:
    """Clamp value to [low, high]."""
    return max(low, min(high, value))