agent/tools/research/fetch_arxiv_topic.py

#!/usr/bin/env python3
"""Build an abstract-level research corpus for one local paper topic."""

from __future__ import annotations

import argparse
import json
import re
import time
import urllib.error
import urllib.parse
import urllib.request
import xml.etree.ElementTree as ET
from datetime import date
from pathlib import Path


ROOT = Path(__file__).resolve().parents[2]
API_URL = "https://export.arxiv.org/api/query"
NS = {"atom": "http://www.w3.org/2005/Atom"}
ARXIV_ID_RE = re.compile(r"arxiv\.org/(?:abs|html|pdf)/([0-9]{4}\.[0-9]{4,5})")


def normalize(value: str) -> str:
    return re.sub(r"\s+", " ", value).strip()


def extract_arxiv_id(url: str) -> str | None:
    match = ARXIV_ID_RE.search(url or "")
    return match.group(1) if match else None


def parse_feed(payload: bytes) -> dict[str, dict]:
    root = ET.fromstring(payload)
    records: dict[str, dict] = {}
    for entry in root.findall("atom:entry", NS):
        raw_id = entry.findtext("atom:id", default="", namespaces=NS)
        paper_id = raw_id.rsplit("/", 1)[-1].split("v", 1)[0]
        records[paper_id] = {
            "arxiv_id": paper_id,
            "title": normalize(entry.findtext("atom:title", default="", namespaces=NS)),
            "abstract": normalize(entry.findtext("atom:summary", default="", namespaces=NS)),
            "published": entry.findtext("atom:published", default="", namespaces=NS)[:10],
            "updated": entry.findtext("atom:updated", default="", namespaces=NS)[:10],
            "authors": [
                normalize(author.findtext("atom:name", default="", namespaces=NS))
                for author in entry.findall("atom:author", NS)
            ],
            "categories": [
                node.attrib["term"]
                for node in entry.findall("atom:category", NS)
                if node.attrib.get("term")
            ],
            "url": f"https://arxiv.org/abs/{paper_id}",
            "pdf_url": f"https://arxiv.org/pdf/{paper_id}",
        }
    return records


def fetch_batch(ids: list[str], retries: int, timeout: int) -> dict[str, dict]:
    query = urllib.parse.urlencode({"id_list": ",".join(ids), "max_results": len(ids)})
    request = urllib.request.Request(
        f"{API_URL}?{query}",
        headers={"User-Agent": "agent-kb-research/0.1"},
    )
    for attempt in range(retries + 1):
        try:
            with urllib.request.urlopen(request, timeout=timeout) as response:
                return parse_feed(response.read())
        except (urllib.error.HTTPError, urllib.error.URLError, TimeoutError):
            if attempt >= retries:
                raise
            time.sleep(2 ** attempt)
    raise RuntimeError("unreachable retry state")


def write_json(path: Path, payload: dict) -> None:
    path.parent.mkdir(parents=True, exist_ok=True)
    tmp = path.with_suffix(path.suffix + ".tmp")
    tmp.write_text(json.dumps(payload, ensure_ascii=False, indent=2) + "\n", encoding="utf-8")
    tmp.replace(path)


def main() -> int:
    parser = argparse.ArgumentParser()
    parser.add_argument("--topic", default="memory")
    parser.add_argument("--index", type=Path, default=ROOT / "data" / "index.json")
    parser.add_argument("--output", type=Path)
    parser.add_argument("--batch-size", type=int, default=20)
    parser.add_argument("--sleep", type=float, default=1.0)
    parser.add_argument("--retries", type=int, default=4)
    parser.add_argument("--timeout", type=int, default=45)
    parser.add_argument("--refresh", action="store_true")
    args = parser.parse_args()

    output = args.output or ROOT / "data" / "research" / f"{args.topic}-corpus.json"
    index = json.loads(args.index.read_text(encoding="utf-8"))
    selected: dict[str, dict] = {}
    missing_ids: list[dict] = []
    for item in index:
        meta = item.get("meta") or {}
        if item.get("collection") != "papers" or args.topic not in (meta.get("topics") or []):
            continue
        paper_id = extract_arxiv_id(str(meta.get("url") or ""))
        if not paper_id:
            missing_ids.append({"path": item.get("path"), "title": item.get("title")})
            continue
        selected[paper_id] = {
            "path": item.get("path"),
            "local_title": item.get("title"),
            "local_status": meta.get("status"),
            "local_topics": meta.get("topics") or [],
            "collection_queries": meta.get("collection_queries"),
            "collection_score": meta.get("collection_score"),
        }

    cached: dict[str, dict] = {}
    if output.exists() and not args.refresh:
        old = json.loads(output.read_text(encoding="utf-8"))
        cached = {record["arxiv_id"]: record for record in old.get("papers", [])}

    fetched: dict[str, dict] = {}
    ids_to_fetch = [paper_id for paper_id in selected if paper_id not in cached]
    for offset in range(0, len(ids_to_fetch), args.batch_size):
        batch = ids_to_fetch[offset : offset + args.batch_size]
        fetched.update(fetch_batch(batch, retries=args.retries, timeout=args.timeout))
        print(f"fetched {min(offset + len(batch), len(ids_to_fetch))}/{len(ids_to_fetch)}")
        if offset + args.batch_size < len(ids_to_fetch):
            time.sleep(args.sleep)

    papers = []
    unresolved = []
    for paper_id, local in selected.items():
        remote = fetched.get(paper_id) or cached.get(paper_id)
        if not remote:
            unresolved.append({"arxiv_id": paper_id, **local})
            continue
        papers.append({**remote, **local})
    papers.sort(key=lambda item: (item.get("published") or "", item["arxiv_id"]))

    payload = {
        "topic": args.topic,
        "generated_at": date.today().isoformat(),
        "selected_from_index": len(selected),
        "resolved": len(papers),
        "missing_arxiv_id": missing_ids,
        "unresolved": unresolved,
        "papers": papers,
    }
    write_json(output, payload)
    print(json.dumps({key: payload[key] for key in ("selected_from_index", "resolved")}, indent=2))
    return 0 if not unresolved else 2


if __name__ == "__main__":
    raise SystemExit(main())