src/tools.py

"""Core tools: arXiv search, PDF download/parse, LLM structured extraction."""

from __future__ import annotations

import json
import os
import re
import time
import urllib.error
import urllib.parse
import urllib.request
from dataclasses import dataclass, field

from . import utils

USER_AGENT = (
    "research-agent/0.1 (+https://github.com/abhid1234/research-agent) "
    "literature-review bot"
)


@dataclass
class Paper:
    """A single arXiv paper plus anything we extract from it."""

    id: str
    title: str
    authors: list[str]
    year: int
    abstract: str
    citations: list[str] = field(default_factory=list)
    full_text: str = ""
    claims: list[str] = field(default_factory=list)
    methods: list[str] = field(default_factory=list)
    results: list[str] = field(default_factory=list)
    score: float = 0.0
    source: str = "arxiv"          # "arxiv" | "semantic_scholar"
    arxiv_id: str = ""             # set when an arXiv PDF is available
    pdf_url: str = ""              # direct PDF URL for non-arXiv sources
    citation_count: int = 0
    read_from: str = ""            # "pdf" | "abstract" — what was actually read


# --------------------------------------------------------------------------- #
# Task 1.1: SearchTool
# --------------------------------------------------------------------------- #
class SearchTool:
    """Search arXiv and return :class:`Paper` records, sorted by relevance.

    A single rate-limited :class:`arxiv.Client` is shared across all queries so
    the library enforces its minimum delay between requests — issuing fresh
    clients per query trips arXiv's HTTP 429 throttle (notably from cloud IPs).
    """

    def __init__(
        self, max_results: int = 20, delay_seconds: float = 3.0, num_retries: int = 2
    ):
        self.max_results = max_results
        self._delay = delay_seconds
        self._num_retries = num_retries  # low → fail fast instead of long backoff
        self._client = None

    def _get_client(self):
        if self._client is None:
            import arxiv

            self._client = arxiv.Client(
                page_size=min(max(self.max_results, 1), 100),
                delay_seconds=self._delay,
                num_retries=self._num_retries,
            )
            # arXiv deprioritizes the default library User-Agent; a descriptive
            # one with a contact URL is the documented way to avoid throttling.
            try:
                self._client._session.headers.update({"User-Agent": USER_AGENT})
            except Exception:
                pass
        return self._client

    def search(self, query: str) -> list[Paper]:
        import arxiv

        client = self._get_client()
        search = arxiv.Search(
            query=query,
            max_results=self.max_results,
            sort_by=arxiv.SortCriterion.Relevance,
        )

        papers: list[Paper] = []
        seen_titles: set[str] = set()
        for result in client.results(search):
            title = result.title.strip().replace("\n", " ")
            key = title.lower()
            if key in seen_titles:
                continue
            seen_titles.add(key)
            papers.append(
                Paper(
                    id=result.get_short_id(),
                    title=title,
                    authors=[a.name for a in result.authors],
                    year=result.published.year if result.published else 0,
                    abstract=result.summary.strip().replace("\n", " "),
                )
            )
        return papers


# --------------------------------------------------------------------------- #
# Task 1.2: DownloadTool
# --------------------------------------------------------------------------- #
class DownloadTool:
    """Fetch a paper's PDF (cached on disk) and extract its body text."""

    def __init__(self, pdf_dir: str = "papers/"):
        self.pdf_dir = pdf_dir
        os.makedirs(self.pdf_dir, exist_ok=True)

    def _local_path(self, paper: Paper) -> str:
        safe_id = (paper.arxiv_id or paper.id).replace("/", "_")
        return os.path.join(self.pdf_dir, f"{safe_id}.pdf")

    @staticmethod
    def _pdf_url(paper: Paper) -> str | None:
        """Best PDF URL for a paper, or None if only the abstract is available."""
        if paper.arxiv_id:
            return f"https://arxiv.org/pdf/{paper.arxiv_id}.pdf"
        if paper.source == "arxiv":
            return f"https://arxiv.org/pdf/{paper.id}.pdf"
        if paper.pdf_url:
            return paper.pdf_url
        return None

    def _download(self, url: str, path: str) -> None:
        req = urllib.request.Request(url, headers={"User-Agent": USER_AGENT})
        last_err: Exception | None = None
        for attempt in range(4):
            try:
                with urllib.request.urlopen(req, timeout=30) as resp, open(
                    path, "wb"
                ) as fh:
                    fh.write(resp.read())
                return
            except urllib.error.HTTPError as err:
                last_err = err
                if err.code in (429, 503):  # throttled — back off and retry
                    time.sleep(2 * (attempt + 1))
                    continue
                raise
        if last_err:
            raise last_err

    def get_text(self, paper: Paper) -> str:
        """Return body text (skipping the first 2 pages); fall back to abstract.

        Sets ``paper.read_from`` to "pdf" or "abstract" so callers can surface
        when a paper was only read from its abstract.
        """
        url = self._pdf_url(paper)
        if not url:
            paper.read_from = "abstract"
            return paper.abstract
        path = self._local_path(paper)
        try:
            if not os.path.exists(path) or os.path.getsize(path) == 0:
                self._download(url, path)

            import pdfplumber

            with pdfplumber.open(path) as pdf:
                pages = pdf.pages
                body = pages[2:] if len(pages) > 3 else pages  # keep short papers whole
                pages_text = [p.extract_text() or "" for p in body]
            text = "\n\n".join(t for t in pages_text if t).strip()
            if text:
                paper.read_from = "pdf"
                return text
            paper.read_from = "abstract"
            return paper.abstract
        except Exception:
            # Any failure (network, parse, corrupt PDF) -> graceful fallback.
            paper.read_from = "abstract"
            return paper.abstract


# --------------------------------------------------------------------------- #
# Task 1.3: ExtractionTool
# --------------------------------------------------------------------------- #
_EXTRACTION_SYSTEM = (
    "You are a meticulous research assistant. Extract concrete, specific "
    "statements from a paper. Respond with JSON only."
)


class ExtractionTool:
    """Use the LLM to pull structured claims/methods/results from a paper."""

    MAX_CHUNKS = 3
    CAP = 5

    def __init__(
        self,
        provider: str | None = None,
        model: str | None = None,
        meter: dict | None = None,
    ):
        # provider/model are honored via src.config; args kept for the plan's
        # signature and possible future overrides.
        self.provider = provider
        self.model = model
        self.meter = meter

    def _prompt(self, title: str, chunk: str, hint: str = "") -> str:
        focus = f"\nFocus especially on the {hint} section.\n" if hint else ""
        return (
            f'Paper title: "{title}"\n{focus}'
            "From the text below, extract:\n"
            '  - "claims": key assertions or contributions the authors make\n'
            '  - "methods": techniques, models, datasets, or procedures used\n'
            '  - "results": quantitative or qualitative findings\n\n'
            "Return ONLY a JSON object with those three keys, each a list of "
            "short strings (one sentence each). Omit anything not present.\n\n"
            f"TEXT:\n{chunk}"
        )

    @staticmethod
    def _parse(raw: str) -> dict:
        # Strip code fences and grab the first JSON object.
        raw = re.sub(r"^```(?:json)?|```$", "", raw.strip(), flags=re.MULTILINE).strip()
        match = re.search(r"\{.*\}", raw, flags=re.DOTALL)
        if not match:
            return {}
        try:
            return json.loads(match.group(0))
        except json.JSONDecodeError:
            return {}

    def extract(self, paper: Paper) -> Paper:
        text = paper.full_text or paper.abstract
        chunks = utils.chunk_text(text, size=3000)[: self.MAX_CHUNKS]

        claims: list[str] = []
        methods: list[str] = []
        results: list[str] = []

        for chunk in chunks:
            raw = utils.complete(
                self._prompt(paper.title, chunk),
                system=_EXTRACTION_SYSTEM,
                meter=self.meter,
                model=self.model,
            )
            data = self._parse(raw)
            claims += [str(x).strip() for x in data.get("claims", []) if str(x).strip()]
            methods += [str(x).strip() for x in data.get("methods", []) if str(x).strip()]
            results += [str(x).strip() for x in data.get("results", []) if str(x).strip()]

        # One targeted retry if a category came back empty.
        if chunks and not (claims and methods and results):
            for missing, bucket in (("results", results), ("methods", methods), ("claims", claims)):
                if not bucket:
                    raw = utils.complete(
                        self._prompt(paper.title, chunks[0], hint=missing),
                        system=_EXTRACTION_SYSTEM,
                        meter=self.meter,
                        model=self.model,
                    )
                    data = self._parse(raw)
                    bucket += [
                        str(x).strip() for x in data.get(missing, []) if str(x).strip()
                    ]

        paper.claims = _dedupe_cap(claims, self.CAP)
        paper.methods = _dedupe_cap(methods, self.CAP)
        paper.results = _dedupe_cap(results, self.CAP)
        return paper


## --------------------------------------------------------------------------- #
# Second source: Semantic Scholar (resilience + non-arXiv venues + citations)
# --------------------------------------------------------------------------- #
class SemanticScholarTool:
    """Search the Semantic Scholar Graph API as a second source / arXiv fallback."""

    ENDPOINT = "https://api.semanticscholar.org/graph/v1/paper/search"
    FIELDS = "title,abstract,year,authors,externalIds,openAccessPdf,citationCount"

    def __init__(self, max_results: int = 8, api_key: str | None = None, year_min: int = 0):
        self.max_results = max_results
        self.year_min = year_min
        # An API key lifts the shared-pool 429s; without one S2 is heavily throttled.
        self.api_key = api_key if api_key is not None else os.getenv("S2_API_KEY", "")

    def search(self, query: str) -> list[Paper]:
        q = {"query": query, "limit": self.max_results, "fields": self.FIELDS}
        if self.year_min:
            q["year"] = f"{self.year_min}-"
        params = urllib.parse.urlencode(q)
        headers = {"User-Agent": USER_AGENT}
        if self.api_key:
            headers["x-api-key"] = self.api_key
        req = urllib.request.Request(f"{self.ENDPOINT}?{params}", headers=headers)

        last_err: Exception | None = None
        for attempt in range(3):
            try:
                with urllib.request.urlopen(req, timeout=20) as resp:
                    data = json.loads(resp.read().decode("utf-8"))
                break
            except urllib.error.HTTPError as err:
                last_err = err
                if err.code == 429:  # throttled — back off and retry
                    time.sleep(2 * (attempt + 1))
                    continue
                raise
            except Exception as err:
                last_err = err
                raise
        else:
            # Exhausted retries on 429 — surface it so the caller can log it.
            raise RuntimeError(f"Semantic Scholar throttled (429): {last_err}")

        papers: list[Paper] = []
        for item in data.get("data") or []:
            title = (item.get("title") or "").strip()
            abstract = (item.get("abstract") or "").strip()
            oa = item.get("openAccessPdf") or {}
            pdf_url = (oa.get("url") or "").strip()
            arxiv_id = (item.get("externalIds") or {}).get("ArXiv", "") or ""
            if not title or (not abstract and not pdf_url and not arxiv_id):
                continue  # nothing to read
            papers.append(
                Paper(
                    id=arxiv_id or item.get("paperId", title[:40]),
                    title=title.replace("\n", " "),
                    authors=[a.get("name", "") for a in item.get("authors") or []],
                    year=item.get("year") or 0,
                    abstract=abstract,
                    source="semantic_scholar",
                    arxiv_id=arxiv_id,
                    pdf_url=pdf_url,
                    citation_count=item.get("citationCount") or 0,
                )
            )
        return papers


def _dedupe_cap(items: list[str], cap: int) -> list[str]:
    seen: set[str] = set()
    out: list[str] = []
    for item in items:
        key = item.lower()
        if key not in seen:
            seen.add(key)
            out.append(item)
        if len(out) >= cap:
            break
    return out