stock-oracle/app/schemas/attention.py

"""
Pydantic v2 schemas for the Attention subsystem API.
"""

from datetime import date, datetime
from typing import Any, Dict, List, Optional
from pydantic import BaseModel, ConfigDict, Field


class EntityInfo(BaseModel):
    model_config = ConfigDict(json_schema_extra={
        "example": {
            "ticker": "AAPL",
            "canonical_name": "Apple",
            "wiki_title": "Apple Inc.",
            "gdelt_query": '"Apple" OR "Apple Inc."',
            "aliases": ["Apple Inc."],
            "resolver_confidence": 0.92,
            "is_manual_override": False,
        }
    })

    ticker: str
    canonical_name: str = Field(description="Normalized company name with legal suffixes stripped (e.g. 'Apple')")
    wiki_title: Optional[str] = Field(default=None, description="Matched Wikipedia article title; null if unresolved")
    gdelt_query: Optional[str] = Field(default=None, description="GDELT DOC API query string (quoted OR phrases)")
    aliases: List[str] = Field(default_factory=list, description="Intermediate forms used during name normalization")
    resolver_confidence: float = Field(default=0.0, description="Wikipedia match confidence [0, 1]")
    is_manual_override: bool = Field(default=False, description="If true, automated re-resolution is skipped")


class WikiFeatures(BaseModel):
    model_config = ConfigDict(json_schema_extra={
        "example": {
            "views": 45230,
            "baseline_10d": 12400.0,
            "spike_10d": 3.65,
            "zscore_20d": 4.21,
        }
    })

    views: Optional[int] = Field(default=None, description="Wikipedia pageviews on the event date")
    baseline_10d: Optional[float] = Field(default=None, description="Median pageviews over the prior 10 days")
    spike_10d: Optional[float] = Field(default=None, description="views / baseline_10d; >1 means above-average attention")
    zscore_20d: Optional[float] = Field(default=None, description="Z-score vs prior 20-day mean/stdev; null if stdev=0")


class NewsFeatures(BaseModel):
    model_config = ConfigDict(json_schema_extra={
        "example": {
            "article_count_1d": 18,
            "article_count_3d": 52,
            "unique_domains_3d": 34,
            "us_article_count_3d": 41,
            "gdelt_status": "collected",
        }
    })

    article_count_1d: int = Field(default=0, description="GDELT articles published on the event date")
    article_count_3d: int = Field(default=0, description="GDELT articles in the event_date ± 1 day window")
    unique_domains_3d: int = Field(default=0, description="Distinct publisher domains in the 3-day window")
    us_article_count_3d: int = Field(default=0, description="US-sourced articles in the 3-day window")
    gdelt_status: str = Field(
        default="not_collected",
        description=(
            "GDELT data availability for this event date. "
            "'collected' — scheduler has run; counts are accurate (0 means genuinely no articles). "
            "'not_collected' — scheduler has not run yet; POST /admin/collect/gdelt/{ticker}?event_date=... to populate. "
            "'not_available' — event date is before GDELT V2 coverage start (2017-01-01)."
        ),
    )


class EventAttentionResponse(BaseModel):
    model_config = ConfigDict(json_schema_extra={
        "example": {
            "ticker": "AAPL",
            "event_date": "2024-02-01",
            "entity": {
                "ticker": "AAPL",
                "canonical_name": "Apple",
                "wiki_title": "Apple Inc.",
                "gdelt_query": '"Apple" OR "Apple Inc."',
                "aliases": ["Apple Inc."],
                "resolver_confidence": 0.92,
                "is_manual_override": False,
            },
            "wiki": {
                "views": 45230,
                "baseline_10d": 12400.0,
                "spike_10d": 3.65,
                "zscore_20d": 4.21,
            },
            "news": {
                "article_count_1d": 18,
                "article_count_3d": 52,
                "unique_domains_3d": 34,
                "us_article_count_3d": 41,
                "gdelt_status": "collected",
            },
            "metadata": {
                "wiki_title": "Apple Inc.",
                "resolver_confidence": 0.92,
            },
        }
    })

    ticker: str
    event_date: date
    entity: EntityInfo
    wiki: WikiFeatures
    news: NewsFeatures
    metadata: Dict[str, Any] = Field(default_factory=dict)


class EntityResolveResponse(BaseModel):
    model_config = ConfigDict(json_schema_extra={
        "example": {
            "ticker": "AAPL",
            "entity": {
                "ticker": "AAPL",
                "canonical_name": "Apple",
                "wiki_title": "Apple Inc.",
                "gdelt_query": '"Apple" OR "Apple Inc."',
                "aliases": ["Apple Inc."],
                "resolver_confidence": 0.92,
                "is_manual_override": False,
            },
            "status": "resolved",
            "message": "Entity resolved: wiki_title='Apple Inc.' confidence=0.92",
        }
    })

    ticker: str
    entity: EntityInfo
    status: str  # "resolved", "already_exists", "failed", "manual_override_skipped"
    message: str


class CollectionStatusResponse(BaseModel):
    model_config = ConfigDict(json_schema_extra={
        "example": {
            "ticker": "AAPL",
            "source": "wiki",
            "records_collected": 22,
            "date_range": {"event_date": "2024-02-01"},
            "status": "success",
        }
    })

    ticker: str
    source: str  # "wiki" or "gdelt"
    records_collected: int
    date_range: Dict[str, Any] = Field(default_factory=dict)
    status: str


class EntityOverrideResponse(BaseModel):
    ticker: str
    wiki_title: str
    gdelt_query: Optional[str]
    message: str