agent/tools/research/classify_memory_corpus.py

#!/usr/bin/env python3
"""Classify an Agent-memory abstract corpus into a research evidence ledger."""

from __future__ import annotations

import argparse
import json
import time
import urllib.error
import urllib.request
from datetime import date
from pathlib import Path


ROOT = Path(__file__).resolve().parents[2]
DEFAULT_OLLAMA_URL = "http://192.168.1.10:11434"

SYSTEM_INSTRUCTIONS = """You are building an evidence ledger for a research review of Agent Memory.
Classify each paper from title and abstract only. Be conservative. Do not infer unreported baselines,
results, limitations, or artifacts. A paper is core only when memory itself is the main research object;
supporting means memory is a substantial mechanism or measured failure; peripheral means memory is
incidental or only an implementation detail.

Return one JSON object with a `papers` array in the same order. Each item must contain:
- arxiv_id: exact input id
- relevance: core | supporting | peripheral
- relevance_reason: <= 18 words
- primary_problem: one of retention_context, retrieval_access, organization_representation,
  update_consolidation_forgetting, execution_state, experience_skill_learning,
  evaluation_measurement, security_trust_privacy, multi_agent_shared_memory,
  systems_efficiency, domain_application, theory_foundations, other
- secondary_problems: up to 3 values from the same enum
- research_role: method | benchmark | diagnostic | theory | survey | security | application | infrastructure
- memory_object: facts | events | trajectories | skills | execution_state | model_internal |
  shared_state | multimodal | mixed | unclear
- temporal_scope: in_episode | cross_episode | both | unclear
- claimed_gap: one short sentence, or empty
- mechanism: one short sentence, or empty
- benchmarks: explicit names only, as an array
- baselines: explicit names or baseline families only, as an array
- reported_results: explicit numeric claims only, as an array of short strings
- evidence_design: none | case_study | benchmark_comparison | controlled_ablation |
  formal_analysis | survey_synthesis
- evidence_strength: none | weak | moderate | strong
- failure_or_boundary: only a boundary explicitly visible in the abstract, or empty

Use compact English. Output JSON only."""


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 ollama_generate(url: str, model: str, prompt: str, timeout: int, retries: int) -> dict:
    payload = {
        "model": model,
        "prompt": prompt,
        "stream": False,
        "format": "json",
        "options": {"temperature": 0, "num_predict": 6000},
    }
    request = urllib.request.Request(
        f"{url.rstrip('/')}/api/generate",
        data=json.dumps(payload).encode("utf-8"),
        headers={"Content-Type": "application/json"},
        method="POST",
    )
    for attempt in range(retries + 1):
        try:
            with urllib.request.urlopen(request, timeout=timeout) as response:
                body = json.loads(response.read().decode("utf-8"))
            return json.loads(body["response"])
        except (urllib.error.HTTPError, urllib.error.URLError, TimeoutError, json.JSONDecodeError, KeyError):
            if attempt >= retries:
                raise
            time.sleep(2 ** attempt)
    raise RuntimeError("unreachable retry state")


def classify_batch(
    records: list[dict], url: str, model: str, timeout: int, retries: int
) -> list[dict]:
    compact = [
        {"arxiv_id": item["arxiv_id"], "title": item["title"], "abstract": item["abstract"]}
        for item in records
    ]
    prompt = f"{SYSTEM_INSTRUCTIONS}\n\nINPUT:\n{json.dumps(compact, ensure_ascii=False)}"
    result = ollama_generate(url, model, prompt, timeout=timeout, retries=retries)
    classified = result.get("papers") if isinstance(result, dict) else None
    if not isinstance(classified, list):
        raise ValueError("model response lacks papers array")
    expected = [item["arxiv_id"] for item in records]
    received = [str(item.get("arxiv_id")) for item in classified]
    if received != expected:
        raise ValueError(f"model returned unexpected ids: expected={expected}, received={received}")
    return classified


def main() -> int:
    parser = argparse.ArgumentParser()
    parser.add_argument(
        "--input", type=Path, default=ROOT / "data" / "research" / "memory-corpus.json"
    )
    parser.add_argument(
        "--output", type=Path, default=ROOT / "data" / "research" / "memory-classified.json"
    )
    parser.add_argument("--ollama-url", default=DEFAULT_OLLAMA_URL)
    parser.add_argument("--model", default="ChatGPT-5.6:Sol")
    parser.add_argument("--batch-size", type=int, default=12)
    parser.add_argument("--timeout", type=int, default=300)
    parser.add_argument("--retries", type=int, default=2)
    parser.add_argument("--limit", type=int)
    parser.add_argument("--refresh", action="store_true")
    args = parser.parse_args()

    args.input = args.input.resolve()
    args.output = args.output.resolve()

    source = json.loads(args.input.read_text(encoding="utf-8"))
    records = source["papers"][: args.limit] if args.limit else source["papers"]
    completed: dict[str, dict] = {}
    if args.output.exists() and not args.refresh:
        old = json.loads(args.output.read_text(encoding="utf-8"))
        completed = {item["arxiv_id"]: item for item in old.get("papers", [])}

    pending = [item for item in records if item["arxiv_id"] not in completed]
    for offset in range(0, len(pending), args.batch_size):
        batch = pending[offset : offset + args.batch_size]
        try:
            classified = classify_batch(
                batch,
                url=args.ollama_url,
                model=args.model,
                timeout=args.timeout,
                retries=args.retries,
            )
        except Exception:
            if len(batch) == 1:
                raise
            classified = []
            for record in batch:
                classified.extend(
                    classify_batch(
                        [record],
                        url=args.ollama_url,
                        model=args.model,
                        timeout=args.timeout,
                        retries=args.retries,
                    )
                )
        for source_record, classification in zip(batch, classified, strict=True):
            completed[source_record["arxiv_id"]] = {**source_record, "analysis": classification}
        ordered = [completed[item["arxiv_id"]] for item in records if item["arxiv_id"] in completed]
        write_json(
            args.output,
            {
                "generated_at": date.today().isoformat(),
                "model": args.model,
                "source": str(args.input.relative_to(ROOT)) if args.input.is_relative_to(ROOT) else str(args.input),
                "total": len(records),
                "completed": len(ordered),
                "papers": ordered,
            },
        )
        print(f"classified {len(ordered)}/{len(records)}", flush=True)
    return 0


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