trends-mcp/trends_mcp/providers/google_trends.py

from __future__ import annotations

from dataclasses import dataclass
from datetime import datetime

from pytrends.request import TrendReq


class GoogleTrendsError(RuntimeError):
    pass


@dataclass(frozen=True)
class TrendSeries:
    keyword: str
    timeframe: str
    series: list[int]
    fetched_at: str


def _normalize_timeframe(timeframe: str) -> str:
    tf = str(timeframe).strip().lower()
    if tf in {"7d", "7day", "7days"}:
        return "today 7-d"
    if tf in {"30d", "30day", "30days"}:
        return "today 1-m"
    if tf in {"90d", "90day", "90days"}:
        return "today 3-m"
    if tf in {"12m", "1y", "365d"}:
        return "today 12-m"
    return timeframe


class GoogleTrendsProvider:
    def __init__(self):
        self._client = TrendReq(hl="en-US", tz=120, retries=2, backoff_factor=0.2)

    def suggestions(self, keyword: str) -> list[dict]:
        try:
            return self._client.suggestions(keyword)
        except Exception as exc:
            raise GoogleTrendsError(f"suggestions failed for {keyword!r}: {exc}") from exc

    def related_queries(self, keyword: str) -> dict:
        try:
            self._client.build_payload([keyword], timeframe="today 12-m")
            return self._client.related_queries() or {}
        except Exception as first_exc:
            try:
                suggestions = self.suggestions(keyword)
                topic = next((s["mid"] for s in suggestions if s.get("mid")), None)
                if not topic:
                    raise first_exc
                self._client.build_payload([topic], timeframe="today 12-m")
                return self._client.related_queries() or {}
            except Exception as exc:
                raise GoogleTrendsError(f"related_queries failed for {keyword!r}: {exc}") from exc

    def related_topics(self, keyword: str) -> dict:
        try:
            self._client.build_payload([keyword], timeframe="today 12-m")
            return self._client.related_topics() or {}
        except Exception as first_exc:
            try:
                suggestions = self.suggestions(keyword)
                topic = next((s["mid"] for s in suggestions if s.get("mid")), None)
                if not topic:
                    raise first_exc
                self._client.build_payload([topic], timeframe="today 12-m")
                return self._client.related_topics() or {}
            except Exception as exc:
                raise GoogleTrendsError(f"related_topics failed for {keyword!r}: {exc}") from exc

    def interest_over_time(self, keyword: str, timeframe: str = "7d") -> TrendSeries:
        timeframe = _normalize_timeframe(timeframe)
        try:
            self._client.build_payload([keyword], timeframe=timeframe)
            frame = self._client.interest_over_time()
        except Exception as first_exc:
            try:
                suggestions = self.suggestions(keyword)
                topic = next((s["mid"] for s in suggestions if s.get("mid")), None)
                if not topic:
                    raise first_exc
                self._client.build_payload([topic], timeframe=timeframe)
                frame = self._client.interest_over_time()
                keyword = topic
            except Exception as exc:
                raise GoogleTrendsError(f"interest_over_time failed for {keyword!r} timeframe={timeframe!r}: {exc}") from exc
        if frame is None or frame.empty:
            series = [0, 0, 0, 0, 0, 0]
        else:
            col = keyword if keyword in frame.columns else frame.columns[0]
            series = [int(v) for v in frame[col].tail(6).tolist()]
            if len(series) < 6:
                series = ([series[0]] * (6 - len(series)) + series) if series else [0] * 6
        return TrendSeries(
            keyword=keyword,
            timeframe=timeframe,
            series=series,
            fetched_at=datetime.utcnow().isoformat() + "Z",
        )