""" 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