Audit recent Agent memory research

2026-07-10 18:52:00 +08:00
parent 36c631e026
commit ab758e774e
17 changed files with 42663 additions and 3 deletions
@@ -0,0 +1,156 @@
+#!/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())