agent-restore-harness/scripts/ingest/push_to_memory.py

#!/usr/bin/env python3
"""
Push a telegram_export.py JSONL file into the memory API.

Usage:
    python3 push_to_memory.py <input.jsonl> \
        --agent-id nasr \
        --api http://172.18.0.20:8000 \
        [--dry-run] [--filter nasr,gmi,clinic]
"""
import argparse
import json
import sys
import time
import urllib.request
import urllib.error

GMI_KEYWORDS = {
    "gmi", "clinic", "cancer", "pathway", "oncology", "breast", "nsclc",
    "colorectal", "clinical", "patient", "diagnosis", "treatment", "smithers",
    "nasr", "salman", "medical", "physician", "physician-scientist",
    "trading", "2dexy", "dexy", "quantum", "crypto"
}

def importance(row: dict) -> float:
    sender = (row.get("metadata", {}).get("sender") or "").lower()
    content = row.get("content", "").lower()
    base = 0.5
    if sender == "nasr salman":
        base = 0.75
    elif sender in ("m2", "mar!0"):
        base = 0.55
    if any(kw in content for kw in GMI_KEYWORDS):
        base = min(base + 0.2, 0.9)
    return round(base, 2)

def entities_from(row: dict) -> list[str]:
    ents = list(row.get("entities", []))
    sender = row.get("metadata", {}).get("sender")
    if sender:
        ents.append(sender.lower().replace(" ", "-"))
    return list(set(ents))[:10]

def store(api_url: str, payload: dict) -> bool:
    data = json.dumps(payload).encode()
    req = urllib.request.Request(
        f"{api_url}/memory/store",
        data=data,
        headers={"Content-Type": "application/json"},
        method="POST"
    )
    try:
        with urllib.request.urlopen(req, timeout=15) as resp:
            return resp.status == 200
    except urllib.error.HTTPError as e:
        print(f"  HTTP {e.code}: {e.read()[:100]}", file=sys.stderr)
        return False
    except Exception as e:
        print(f"  Error: {e}", file=sys.stderr)
        return False

def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("input", help="JSONL file from telegram_export.py")
    ap.add_argument("--agent-id", default="nasr")
    ap.add_argument("--api", default="http://172.18.0.20:8000")
    ap.add_argument("--dry-run", action="store_true")
    ap.add_argument("--filter", default="", help="comma-separated keywords to filter for")
    ap.add_argument("--min-length", type=int, default=20, help="min content length to ingest")
    ap.add_argument("--batch-delay", type=float, default=0.05, help="seconds between requests")
    args = ap.parse_args()

    filter_kws = set(k.strip().lower() for k in args.filter.split(",") if k.strip())

    rows = []
    with open(args.input) as f:
        for line in f:
            line = line.strip()
            if line:
                rows.append(json.loads(line))

    # Filter
    def passes(row: dict) -> bool:
        content = row.get("content", "")
        if not content or len(content) < args.min_length:
            return False
        if row.get("metadata", {}).get("is_service"):
            return False
        if filter_kws:
            return any(kw in content.lower() for kw in filter_kws)
        return True

    filtered = [r for r in rows if passes(r)]
    print(f"Rows: {len(rows)} total → {len(filtered)} after filter (min_len={args.min_length}, filter={filter_kws or 'none'})")

    if args.dry_run:
        print("DRY RUN — first 3 payloads:")
        for r in filtered[:3]:
            payload = {
                "content": r["content"][:300],
                "agent_id": args.agent_id,
                "memory_type": r.get("memory_type", "episodic"),
                "importance": importance(r),
                "entities": entities_from(r),
                "metadata": r.get("metadata", {}),
            }
            print(json.dumps(payload, indent=2)[:400])
        return

    ok = 0
    fail = 0
    for i, row in enumerate(filtered):
        payload = {
            "content": row["content"],
            "agent_id": args.agent_id,
            "memory_type": row.get("memory_type", "episodic"),
            "importance": importance(row),
            "entities": entities_from(row),
            "metadata": row.get("metadata", {}),
        }
        if store(args.api, payload):
            ok += 1
        else:
            fail += 1
        if (i + 1) % 100 == 0:
            print(f"  [{i+1}/{len(filtered)}] ok={ok} fail={fail}")
        if args.batch_delay:
            time.sleep(args.batch_delay)

    print(f"Done: {ok} stored, {fail} failed out of {len(filtered)}")

if __name__ == "__main__":
    main()