Audit recent Agent memory research
This commit is contained in:
wuyang
@@ -0,0 +1,281 @@
|
||||
#!/usr/bin/env python3
|
||||
"""Fast relevance screening for a broad Agent-memory research corpus."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import argparse
|
||||
import json
|
||||
import re
|
||||
import time
|
||||
import urllib.error
|
||||
import urllib.request
|
||||
from datetime import date
|
||||
from pathlib import Path
|
||||
|
||||
|
||||
ROOT = Path(__file__).resolve().parents[2]
|
||||
|
||||
PROMPT = """Screen papers for a research review about memory in LLM-based agents.
|
||||
Use title and abstract only. Return JSON with a `papers` array in input order.
|
||||
For each paper return:
|
||||
- arxiv_id: exact input id
|
||||
- relevance: core if memory is the main research object; supporting if memory is a substantial
|
||||
mechanism, evaluation dimension, or failure mode; peripheral if memory is incidental, means
|
||||
GPU/KV memory only, refers to human memory only, or the paper is mainly an unrelated application
|
||||
- problem: choose one: access, representation, update_forgetting, execution_state, experience_learning,
|
||||
evaluation, security_privacy, shared_memory, systems_cost, theory, application, unrelated
|
||||
- reason: at most 12 words
|
||||
Be conservative and output JSON only."""
|
||||
|
||||
RELEVANCE_VALUES = {"core", "supporting", "peripheral"}
|
||||
PROBLEM_VALUES = {
|
||||
"access",
|
||||
"representation",
|
||||
"update_forgetting",
|
||||
"execution_state",
|
||||
"experience_learning",
|
||||
"evaluation",
|
||||
"security_privacy",
|
||||
"shared_memory",
|
||||
"systems_cost",
|
||||
"theory",
|
||||
"application",
|
||||
"unrelated",
|
||||
}
|
||||
PROBLEM_ALIASES = {
|
||||
"adaptation": "experience_learning",
|
||||
"continual_learning": "experience_learning",
|
||||
"episodic": "representation",
|
||||
"episodic_memory": "representation",
|
||||
"experience": "experience_learning",
|
||||
"forgetting": "update_forgetting",
|
||||
"implicit_personalization": "application",
|
||||
"memory": "representation",
|
||||
"privacy": "security_privacy",
|
||||
"security": "security_privacy",
|
||||
}
|
||||
|
||||
|
||||
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 generate(url: str, model: str, records: list[dict], timeout: int, retries: int) -> list[dict]:
|
||||
compact = [
|
||||
{"arxiv_id": item["arxiv_id"], "title": item["title"], "abstract": item["abstract"]}
|
||||
for item in records
|
||||
]
|
||||
body = {
|
||||
"model": model,
|
||||
"prompt": f"{PROMPT}\n\nINPUT:\n{json.dumps(compact, ensure_ascii=False)}",
|
||||
"stream": False,
|
||||
"format": "json",
|
||||
"options": {"temperature": 0, "num_predict": 4000},
|
||||
}
|
||||
request = urllib.request.Request(
|
||||
f"{url.rstrip('/')}/api/generate",
|
||||
data=json.dumps(body).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:
|
||||
result = json.loads(response.read().decode("utf-8"))
|
||||
decoded = json.loads(result["response"])
|
||||
if isinstance(decoded, list):
|
||||
parsed = decoded
|
||||
elif isinstance(decoded, dict) and isinstance(decoded.get("papers"), list):
|
||||
parsed = decoded["papers"]
|
||||
elif isinstance(decoded, dict) and isinstance(decoded.get("results"), list):
|
||||
parsed = decoded["results"]
|
||||
elif isinstance(decoded, dict) and decoded.get("arxiv_id"):
|
||||
parsed = [decoded]
|
||||
else:
|
||||
raise ValueError("response lacks a recognized result array")
|
||||
expected = [item["arxiv_id"] for item in records]
|
||||
by_id = {str(item.get("arxiv_id")): item for item in parsed}
|
||||
if any(paper_id not in by_id for paper_id in expected):
|
||||
raise ValueError("response ids do not match input")
|
||||
return [{**by_id[paper_id], "source": "model"} for paper_id in expected]
|
||||
except (
|
||||
urllib.error.HTTPError,
|
||||
urllib.error.URLError,
|
||||
TimeoutError,
|
||||
json.JSONDecodeError,
|
||||
KeyError,
|
||||
ValueError,
|
||||
):
|
||||
if attempt >= retries:
|
||||
raise
|
||||
time.sleep(2 ** attempt)
|
||||
raise RuntimeError("unreachable retry state")
|
||||
|
||||
|
||||
def heuristic_screen(record: dict) -> dict:
|
||||
title = record["title"].lower()
|
||||
text = f"{record['title']} {record['abstract']}".lower()
|
||||
title_memory = bool(re.search(r"memory|memori|remember|forget|stateful", title))
|
||||
agent_memory = bool(
|
||||
re.search(r"agent(?:ic)? memory|memory (?:for|in|of) .*agent|long-term .*agent", text)
|
||||
)
|
||||
hardware_only = bool(re.search(r"gpu memory|memory footprint|memory bandwidth|memory efficient", title))
|
||||
if title_memory and not hardware_only:
|
||||
relevance = "core"
|
||||
elif agent_memory:
|
||||
relevance = "supporting"
|
||||
else:
|
||||
relevance = "peripheral"
|
||||
|
||||
if re.search(r"attack|poison|privacy|security|trust|provenance|leak", text):
|
||||
problem = "security_privacy"
|
||||
elif re.search(r"benchmark|evaluation|diagnos|measure", title):
|
||||
problem = "evaluation"
|
||||
elif re.search(r"update|forget|consolidat|stale|supersed|write", text):
|
||||
problem = "update_forgetting"
|
||||
elif re.search(r"skill|experience|reflect|self-evolv|continual learn", text):
|
||||
problem = "experience_learning"
|
||||
elif re.search(r"execution state|task state|workflow state|long-horizon", text):
|
||||
problem = "execution_state"
|
||||
elif re.search(r"graph|hierarch|represent|structure|database|document", text):
|
||||
problem = "representation"
|
||||
elif re.search(r"retriev|search|recall|access", text):
|
||||
problem = "access"
|
||||
elif re.search(r"multi-agent|shared memory|collective memory", text):
|
||||
problem = "shared_memory"
|
||||
elif re.search(r"cost|latency|token|efficient|compress", text):
|
||||
problem = "systems_cost"
|
||||
else:
|
||||
problem = "application" if relevance != "peripheral" else "unrelated"
|
||||
return {
|
||||
"arxiv_id": record["arxiv_id"],
|
||||
"relevance": relevance,
|
||||
"problem": problem,
|
||||
"reason": "deterministic fallback after model-format failure",
|
||||
"source": "heuristic",
|
||||
}
|
||||
|
||||
|
||||
def normalize_analysis(record: dict, analysis: dict | None) -> dict:
|
||||
"""Constrain model output to the documented screening schema."""
|
||||
fallback = heuristic_screen(record)
|
||||
analysis = analysis if isinstance(analysis, dict) else {}
|
||||
relevance = str(analysis.get("relevance", "")).strip().lower()
|
||||
if relevance not in RELEVANCE_VALUES:
|
||||
relevance = fallback["relevance"]
|
||||
|
||||
problem = str(analysis.get("problem", "")).strip().lower()
|
||||
problem = PROBLEM_ALIASES.get(problem, problem)
|
||||
if problem not in PROBLEM_VALUES:
|
||||
problem = fallback["problem"]
|
||||
|
||||
reason = str(analysis.get("reason", "")).strip() or fallback["reason"]
|
||||
source = str(analysis.get("source", "")).strip().lower()
|
||||
if source not in {"model", "heuristic"}:
|
||||
source = "model" if analysis else "heuristic"
|
||||
return {
|
||||
"arxiv_id": record["arxiv_id"],
|
||||
"relevance": relevance,
|
||||
"problem": problem,
|
||||
"reason": reason,
|
||||
"source": source,
|
||||
}
|
||||
|
||||
|
||||
def main() -> int:
|
||||
parser = argparse.ArgumentParser()
|
||||
parser.add_argument(
|
||||
"--input",
|
||||
type=Path,
|
||||
default=ROOT / "data" / "research" / "agent-memory-expanded.json",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--output",
|
||||
type=Path,
|
||||
default=ROOT / "data" / "research" / "agent-memory-screened.json",
|
||||
)
|
||||
parser.add_argument("--ollama-url", default="http://192.168.1.10:11434")
|
||||
parser.add_argument("--model", default="ChatGPT-5.6:Luna")
|
||||
parser.add_argument("--batch-size", type=int, default=24)
|
||||
parser.add_argument("--timeout", type=int, default=180)
|
||||
parser.add_argument("--retries", type=int, default=2)
|
||||
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"]
|
||||
completed: dict[str, dict] = {}
|
||||
if args.output.exists() and not args.refresh:
|
||||
previous = json.loads(args.output.read_text(encoding="utf-8"))
|
||||
completed = {
|
||||
item["arxiv_id"]: {
|
||||
**item,
|
||||
"screen": normalize_analysis(item, item.get("screen")),
|
||||
}
|
||||
for item in previous.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:
|
||||
screened = generate(
|
||||
args.ollama_url, args.model, batch, timeout=args.timeout, retries=args.retries
|
||||
)
|
||||
except Exception:
|
||||
screened = []
|
||||
for record in batch:
|
||||
try:
|
||||
screened.extend(
|
||||
generate(
|
||||
args.ollama_url,
|
||||
args.model,
|
||||
[record],
|
||||
timeout=args.timeout,
|
||||
retries=args.retries,
|
||||
)
|
||||
)
|
||||
except Exception:
|
||||
screened.append(heuristic_screen(record))
|
||||
for record, analysis in zip(batch, screened, strict=True):
|
||||
completed[record["arxiv_id"]] = {
|
||||
**record,
|
||||
"screen": normalize_analysis(record, analysis),
|
||||
}
|
||||
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)),
|
||||
"total": len(records),
|
||||
"completed": len(ordered),
|
||||
"papers": ordered,
|
||||
},
|
||||
)
|
||||
print(f"screened {len(ordered)}/{len(records)}", flush=True)
|
||||
|
||||
# Rewrites an already complete file after schema normalization as well.
|
||||
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)),
|
||||
"total": len(records),
|
||||
"completed": len(ordered),
|
||||
"papers": ordered,
|
||||
},
|
||||
)
|
||||
return 0
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
raise SystemExit(main())
|
||||
Reference in New Issue
Block a user