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 asyncio
import logging
import time
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

# Process-wide concurrency + rate limiting for Wikimedia REST API.
# Prevents burst-induced 429 floods that correlate with container OOM.
_WIKI_SEMAPHORE = asyncio.Semaphore(2)
_WIKI_MIN_INTERVAL = 0.25  # seconds between requests (≈ 4 req/s ceiling)
_WIKI_RATE_LOCK = asyncio.Lock()
_WIKI_LAST_CALL_TS: float = 0.0
_WIKI_429_LAST_LOG_TS: float = 0.0
_WIKI_429_LOG_WINDOW = 30.0  # log at most once per 30 s during a burst


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)

    global _WIKI_LAST_CALL_TS, _WIKI_429_LAST_LOG_TS
    try:
        async with _WIKI_SEMAPHORE:
            async with _WIKI_RATE_LOCK:
                now = time.monotonic()
                wait = _WIKI_MIN_INTERVAL - (now - _WIKI_LAST_CALL_TS)
                if wait > 0:
                    await asyncio.sleep(wait)
                _WIKI_LAST_CALL_TS = time.monotonic()

            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:
        status = exc.response.status_code
        if status == 404:
            logger.warning("Wikipedia page not found: %r", wiki_title)
            return 0
        if status == 429:
            now = time.monotonic()
            if now - _WIKI_429_LAST_LOG_TS > _WIKI_429_LOG_WINDOW:
                logger.warning("Wikimedia rate limit 429 for %r (burst suppressed %ds)", wiki_title, int(_WIKI_429_LOG_WINDOW))
                _WIKI_429_LAST_LOG_TS = now
        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)