stonks-oracle/services/intelligence_pipeline_v3/schemas/annotations.py

"""V3 annotation schema — labels for entities, events, relations, facts, and evidence.

This module defines the complete annotation schema for the Intelligence Pipeline v3
Gold Corpus. Every extracted field is traceable to source evidence via character offsets.

Schema version: 1.0.0
"""

from __future__ import annotations

import uuid
from datetime import date, datetime
from enum import Enum
from typing import Literal

from pydantic import BaseModel, Field, model_validator

# ---------------------------------------------------------------------------
# Enumerations
# ---------------------------------------------------------------------------


class EntityType(str, Enum):
    """Recognized entity types in the v3 pipeline."""

    COMPANY = "company"
    PERSON = "person"
    PRODUCT = "product"
    EVENT = "event"
    FINANCIAL_METRIC = "financial_metric"
    DATE = "date"
    PERCENTAGE = "percentage"
    CURRENCY = "currency"
    RELATIONSHIP = "relationship"


class EventClass(str, Enum):
    """Versioned event taxonomy for market-relevant occurrences."""

    EARNINGS_BEAT = "earnings_beat"
    EARNINGS_MISS = "earnings_miss"
    GUIDANCE_RAISE = "guidance_raise"
    GUIDANCE_CUT = "guidance_cut"
    MA_ANNOUNCEMENT = "ma_announcement"
    LEGAL_REGULATORY = "legal_regulatory"
    PRODUCT_LAUNCH = "product_launch"
    SUPPLY_CHAIN = "supply_chain"
    RATING_CHANGE = "rating_change"
    MANAGEMENT_CHANGE = "management_change"
    MACRO_EVENT = "macro_event"
    DIVIDEND_CHANGE = "dividend_change"
    BUYBACK = "buyback"


class SentimentLabel(str, Enum):
    """Sentiment classification for company-linked evidence groups."""

    POSITIVE = "positive"
    NEGATIVE = "negative"
    NEUTRAL = "neutral"
    MIXED = "mixed"


class RelationType(str, Enum):
    """Relation types between entities or between events and companies."""

    DIRECTLY_AFFECTS = "directly_affects"
    INFERRED_EXPOSURE = "inferred_exposure"
    COMPETES_WITH = "competes_with"
    SUPPLIES = "supplies"


class PeriodType(str, Enum):
    """Financial period type identifiers."""

    FISCAL_QUARTER = "fiscal_quarter"
    FISCAL_YEAR = "fiscal_year"
    CALENDAR_QUARTER = "calendar_quarter"
    CALENDAR_YEAR = "calendar_year"
    TRAILING_TWELVE_MONTHS = "ttm"
    YEAR_TO_DATE = "ytd"
    CUSTOM = "custom"


class AmbiguityType(str, Enum):
    """Ambiguity reasons that trigger adjudication routing."""

    UNRESOLVED_ALIAS = "unresolved_alias"
    MULTIPLE_PRIMARY_COMPANIES = "multiple_primary_companies"
    CONTRADICTORY_NUMERIC_FACTS = "contradictory_numeric_facts"
    CONFLICTING_SENTIMENT = "conflicting_sentiment"
    IMPLIED_CAUSAL_IMPACT = "implied_causal_impact"
    GUIDANCE_VS_CONSENSUS_REQUIRES_REASONING = "guidance_vs_consensus_requires_reasoning"
    MATERIAL_FIELD_MISSING = "material_field_missing"
    EVIDENCE_COVERAGE_BELOW_THRESHOLD = "evidence_coverage_below_threshold"
    CALIBRATED_CONFIDENCE_BELOW_THRESHOLD = "calibrated_confidence_below_threshold"
    LONG_DOCUMENT_CROSS_CHUNK_RELATION = "long_document_cross_chunk_relation"


# ---------------------------------------------------------------------------
# Evidence Span
# ---------------------------------------------------------------------------


class EvidenceSpanAnnotation(BaseModel):
    """Exact source text with stable character offsets.

    Every extracted fact, entity, or relation MUST reference at least one evidence span.
    Offsets are zero-based and refer to the original (pre-chunking) document text.
    """

    id: str = Field(default_factory=lambda: str(uuid.uuid4()))
    chunk_id: str | None = Field(
        default=None, description="Chunk ID if document was segmented."
    )
    start_char: int = Field(ge=0, description="Zero-based start character offset.")
    end_char: int = Field(ge=0, description="Zero-based end character offset (exclusive).")
    text: str = Field(min_length=1, description="Exact source text at this span.")
    checksum: str | None = Field(
        default=None,
        description="SHA-256 hex digest of the span text for integrity verification.",
    )

    @model_validator(mode="after")
    def end_after_start(self) -> "EvidenceSpanAnnotation":
        if self.end_char <= self.start_char:
            raise ValueError(
                f"end_char ({self.end_char}) must be greater than start_char ({self.start_char})"
            )
        return self


# ---------------------------------------------------------------------------
# Entity Annotation
# ---------------------------------------------------------------------------


class EntityAnnotation(BaseModel):
    """A labeled entity mention with optional canonical resolution."""

    id: str = Field(default_factory=lambda: str(uuid.uuid4()))
    entity_type: EntityType
    literal_text: str = Field(min_length=1, description="Exact surface form as it appears.")
    canonical_id: str | None = Field(
        default=None,
        description="UUID of the canonical company/entity from the symbol registry.",
    )
    canonical_name: str | None = Field(
        default=None, description="Resolved canonical name (e.g., ticker or full name)."
    )
    evidence_ids: list[str] = Field(
        min_length=1,
        description="References to EvidenceSpanAnnotation IDs supporting this entity.",
    )
    confidence: float = Field(ge=0.0, le=1.0, description="Annotator certainty [0, 1].")
    derivation: str = Field(
        default="manual",
        description="How this label was derived: manual, deterministic, specialist, adjudicated.",
    )


# ---------------------------------------------------------------------------
# Event Annotation
# ---------------------------------------------------------------------------


class EventAnnotation(BaseModel):
    """A market-relevant event classification anchored to evidence."""

    id: str = Field(default_factory=lambda: str(uuid.uuid4()))
    event_class: EventClass
    description: str = Field(
        default="", description="Brief human-readable description of the event."
    )
    primary_company_ids: list[str] = Field(
        default_factory=list,
        description="Entity annotation IDs of directly affected companies.",
    )
    evidence_ids: list[str] = Field(
        min_length=1, description="Evidence spans supporting this event classification."
    )
    confidence: float = Field(ge=0.0, le=1.0)
    derivation: str = Field(default="manual")


# ---------------------------------------------------------------------------
# Relation Annotation
# ---------------------------------------------------------------------------


class RelationAnnotation(BaseModel):
    """A typed relation between two entities or between an event and an entity."""

    id: str = Field(default_factory=lambda: str(uuid.uuid4()))
    relation_type: RelationType
    source_id: str = Field(description="Entity or event annotation ID (subject).")
    target_id: str = Field(description="Entity annotation ID (object).")
    evidence_ids: list[str] = Field(min_length=1)
    confidence: float = Field(ge=0.0, le=1.0)
    derivation: str = Field(default="manual")


# ---------------------------------------------------------------------------
# Numeric Fact Annotation
# ---------------------------------------------------------------------------


class NumericFactAnnotation(BaseModel):
    """A numeric fact (EPS, revenue, percentage change, etc.) extracted from the document."""

    id: str = Field(default_factory=lambda: str(uuid.uuid4()))
    fact_type: str = Field(
        description="Category: eps, revenue, percentage_change, price_target, guidance, market_cap, etc."
    )
    subject_entity_id: str | None = Field(
        default=None, description="Entity annotation ID this fact belongs to."
    )
    predicate: str = Field(
        description="Semantic predicate: reported, expected, raised_to, cut_to, beat_by, etc."
    )
    literal_value: str = Field(description="Exact textual representation from the source.")
    normalized_value: float | None = Field(
        default=None, description="Numeric value after normalization."
    )
    unit: str | None = Field(
        default=None, description="Unit: USD, %, bps, shares, etc."
    )
    period: "PeriodAnnotation | None" = Field(
        default=None, description="Financial period this fact applies to."
    )
    evidence_ids: list[str] = Field(min_length=1)
    confidence: float = Field(ge=0.0, le=1.0)
    derivation: str = Field(default="manual")


# ---------------------------------------------------------------------------
# Period Annotation
# ---------------------------------------------------------------------------


class PeriodAnnotation(BaseModel):
    """Financial or calendar period reference."""

    period_type: PeriodType
    fiscal_year: int | None = Field(default=None, description="e.g. 2024")
    fiscal_quarter: int | None = Field(default=None, ge=1, le=4)
    start_date: date | None = None
    end_date: date | None = None
    literal_text: str | None = Field(
        default=None, description="Original text describing the period."
    )


# ---------------------------------------------------------------------------
# Sentiment Annotation
# ---------------------------------------------------------------------------


class CompanySentimentAnnotation(BaseModel):
    """Company-specific sentiment with probability distribution.

    Mixed sentiment is derived from disagreement across evidence groups — it is NOT
    an unconstrained fourth softmax label.
    """

    id: str = Field(default_factory=lambda: str(uuid.uuid4()))
    company_entity_id: str = Field(
        description="Entity annotation ID of the company this sentiment applies to."
    )
    label: SentimentLabel
    positive_probability: float = Field(ge=0.0, le=1.0)
    negative_probability: float = Field(ge=0.0, le=1.0)
    neutral_probability: float = Field(ge=0.0, le=1.0)
    evidence_ids: list[str] = Field(min_length=1)
    confidence: float = Field(ge=0.0, le=1.0)
    derivation: str = Field(default="manual")

    @model_validator(mode="after")
    def probabilities_sum_to_one(self) -> "CompanySentimentAnnotation":
        total = (
            self.positive_probability
            + self.negative_probability
            + self.neutral_probability
        )
        if abs(total - 1.0) > 0.01:
            raise ValueError(
                f"Sentiment probabilities must sum to ~1.0, got {total:.4f}"
            )
        return self


# ---------------------------------------------------------------------------
# Direct Effects and Inferred Exposure
# ---------------------------------------------------------------------------


class DirectEffect(BaseModel):
    """An event directly affecting a specific company, backed by explicit evidence."""

    event_id: str = Field(description="Event annotation ID.")
    company_entity_id: str = Field(description="Entity annotation ID of the affected company.")
    evidence_ids: list[str] = Field(min_length=1)
    confidence: float = Field(ge=0.0, le=1.0)


class InferredExposure(BaseModel):
    """An inferred (not explicitly stated) exposure of a company to an event.

    Inferred exposures have separate confidence and do NOT enter primary extraction.
    They flow through the interpolation/propagation architecture with distinct provenance.
    """

    event_id: str = Field(description="Event annotation ID.")
    company_entity_id: str = Field(description="Entity annotation ID.")
    reasoning: str = Field(
        description="Brief explanation of the inference chain (e.g., supply-chain relationship)."
    )
    evidence_ids: list[str] = Field(
        default_factory=list,
        description="Supporting evidence (may be empty for purely inferred relations).",
    )
    confidence: float = Field(ge=0.0, le=1.0)


# ---------------------------------------------------------------------------
# Ambiguity Marker
# ---------------------------------------------------------------------------


class AmbiguityMarker(BaseModel):
    """Flags cases that require adjudication routing.

    These markers determine whether a document is processed via the fast path
    or routed to the 9B model for semantic adjudication.
    """

    ambiguity_type: AmbiguityType
    description: str = Field(
        default="", description="Human-readable description of the ambiguity."
    )
    affected_entity_ids: list[str] = Field(default_factory=list)
    affected_event_ids: list[str] = Field(default_factory=list)
    severity: Literal["low", "medium", "high"] = Field(
        default="medium",
        description="Impact on extraction confidence.",
    )


# ---------------------------------------------------------------------------
# Annotation Metadata
# ---------------------------------------------------------------------------


class AnnotationMetadata(BaseModel):
    """Metadata for a complete document annotation."""

    schema_version: str = Field(default="1.0.0")
    annotator_id: str = Field(description="Identifier of the annotator (human or system).")
    annotation_date: datetime = Field(default_factory=lambda: datetime.now(tz=datetime.now().astimezone().tzinfo))
    review_status: Literal["draft", "reviewed", "adjudicated", "gold"] = Field(
        default="draft"
    )
    reviewer_id: str | None = None
    review_date: datetime | None = None
    notes: str = Field(default="")


# ---------------------------------------------------------------------------
# Top-Level Annotated Document
# ---------------------------------------------------------------------------


class AnnotatedDocument(BaseModel):
    """Complete v3 annotation for a single source document.

    This is the top-level container used in the Gold Corpus. Every field references
    evidence spans for traceability.
    """

    document_id: str = Field(description="UUID of the source document.")
    document_type: str = Field(description="article, filing, transcript, press_release, macro_event")
    source_text: str = Field(description="Full original document text (offsets reference this).")
    metadata: AnnotationMetadata

    # Core annotations
    evidence_spans: list[EvidenceSpanAnnotation] = Field(default_factory=list)
    entities: list[EntityAnnotation] = Field(default_factory=list)
    events: list[EventAnnotation] = Field(default_factory=list)
    relations: list[RelationAnnotation] = Field(default_factory=list)
    numeric_facts: list[NumericFactAnnotation] = Field(default_factory=list)
    sentiments: list[CompanySentimentAnnotation] = Field(default_factory=list)

    # Effects and exposure
    direct_effects: list[DirectEffect] = Field(default_factory=list)
    inferred_exposures: list[InferredExposure] = Field(default_factory=list)

    # Routing
    ambiguity_markers: list[AmbiguityMarker] = Field(default_factory=list)