stock-oracle/app/services/attention/wiki_collector.py

"""
Wiki Collector — event-centric Wikipedia pageview collection.

Collects daily pageview counts for the event window: event_date - 20d to event_date + 2d.
Only fetches dates not already in the database.
Uses Wikimedia REST API range endpoint for efficiency.
"""

import logging
from datetime import date, timedelta

import httpx
from sqlalchemy import select
from sqlalchemy.dialects.postgresql import insert as pg_insert
from sqlalchemy.ext.asyncio import AsyncSession

from app.models.attention import CompanyEntityMap, WikiPageviewsDaily

logger = logging.getLogger(__name__)

_WIKI_BASE = "https://wikimedia.org/api/rest_v1/metrics/pageviews/per-article"
_WIKI_HEADERS = {
    "User-Agent": "StockOracleBot/1.0 (attention-subsystem; contact@stockoracle.internal)"
}
_LOOKBACK_DAYS = 20
_LOOKAHEAD_DAYS = 2


def _date_range(start: date, end: date) -> list[date]:
    """Return list of dates from start to end inclusive."""
    dates = []
    current = start
    while current <= end:
        dates.append(current)
        current += timedelta(days=1)
    return dates


async def collect_wiki_pageviews(
    db: AsyncSession,
    ticker: str,
    event_date: date,
) -> int:
    """Collect Wikipedia pageviews for the event window.

    Returns the number of new records inserted.
    Raises ValueError if the entity has no wiki_title.
    """
    ticker = ticker.upper()

    # Look up wiki_title from entity map
    result = await db.execute(
        select(CompanyEntityMap).where(CompanyEntityMap.ticker == ticker)
    )
    entity = result.scalars().first()
    if not entity:
        raise ValueError(f"No entity mapping found for {ticker!r}. Run entity resolution first.")
    if not entity.wiki_title:
        raise ValueError(f"Entity {ticker!r} has no wiki_title — resolution confidence was too low.")

    wiki_title = entity.wiki_title
    start_date = event_date - timedelta(days=_LOOKBACK_DAYS)
    end_date = event_date + timedelta(days=_LOOKAHEAD_DAYS)
    all_dates = _date_range(start_date, end_date)

    # Find which dates are already in DB
    existing_result = await db.execute(
        select(WikiPageviewsDaily.date).where(
            WikiPageviewsDaily.wiki_title == wiki_title,
            WikiPageviewsDaily.date.in_(all_dates),
        )
    )
    existing_dates = set(existing_result.scalars().all())
    missing_dates = [d for d in all_dates if d not in existing_dates]

    if not missing_dates:
        logger.info("All wiki pageviews already collected for %s event=%s", ticker, event_date)
        return 0

    # Fetch range from Wikimedia REST API
    fetch_start = min(missing_dates)
    fetch_end = max(missing_dates)
    start_str = fetch_start.strftime("%Y%m%d")
    end_str = fetch_end.strftime("%Y%m%d")

    encoded_title = wiki_title.replace(" ", "_")
    url = (
        f"{_WIKI_BASE}/en.wikipedia/all-access/all-agents"
        f"/{encoded_title}/daily/{start_str}00/{end_str}00"
    )

    logger.info("Fetching wiki pageviews for %r (%s → %s)", wiki_title, start_str, end_str)

    try:
        async with httpx.AsyncClient(timeout=15.0, headers=_WIKI_HEADERS) as client:
            resp = await client.get(url)
            resp.raise_for_status()
            data = resp.json()
    except httpx.HTTPStatusError as exc:
        if exc.response.status_code == 404:
            logger.warning("Wikipedia page not found: %r", wiki_title)
            return 0
        raise
    except Exception as exc:
        logger.error("Wiki API request failed: %s", exc)
        raise

    items = data.get("items", [])
    if not items:
        logger.info("No pageview data returned for %r", wiki_title)
        return 0

    # Parse and insert missing dates only
    missing_set = set(missing_dates)
    rows_to_insert = []
    for item in items:
        timestamp_str = item.get("timestamp", "")
        if len(timestamp_str) >= 8:
            d = date(
                int(timestamp_str[0:4]),
                int(timestamp_str[4:6]),
                int(timestamp_str[6:8]),
            )
            if d in missing_set:
                rows_to_insert.append({
                    "wiki_title": wiki_title,
                    "date": d,
                    "views": item.get("views", 0),
                })

    if not rows_to_insert:
        return 0

    stmt = (
        pg_insert(WikiPageviewsDaily)
        .values(rows_to_insert)
        .on_conflict_do_nothing(constraint="uq_wiki_pageviews_daily")
    )
    await db.execute(stmt)
    await db.commit()

    logger.info("Inserted %d wiki pageview records for %s", len(rows_to_insert), ticker)
    return len(rows_to_insert)