memoryweave/scripts/migrate_hermes_to_zhiyi.py

#!/usr/bin/env python3
"""迁移 Hermes 自有记忆到织忆 API — 读取 ~/.hermes/memory_db/lancedb/ 旧记忆，commit 到织忆。

用法:
  python3 migrate_hermes_to_zhiyi.py            # 正常迁移
  python3 migrate_hermes_to_zhiyi.py --dry-run  # 只扫描不写
  python3 migrate_hermes_to_zhiyi.py --verify   # 验证召回质量
  python3 migrate_hermes_to_zhiyi.py --cleanup  # 确认后删旧 DB
"""
import lancedb, json, time, sys, os, hashlib, requests
from pathlib import Path

ZHIYI_URL = os.environ.get("ZHIYI_URL", "http://localhost:7821")
API_KEY = os.environ.get("ZHIYI_API_KEY", "zhiyi-dev-key-2026")
BATCH_DELAY = 0.15
STATE_FILE = Path.home() / ".hermes" / "migration_state.json"
HEADERS = {"X-API-Key": API_KEY, "Content-Type": "application/json"}

def load_state():
    if STATE_FILE.exists():
        return json.loads(STATE_FILE.read_text())
    return {"committed": [], "failed": [], "tables_done": []}

def save_state(state):
    STATE_FILE.write_text(json.dumps(state, ensure_ascii=False, indent=2))

def commit(content, category, agent_id, namespace="hermes-main"):
    payload = {"content": content, "category": category or "general",
               "namespace": namespace, "agent_id": agent_id or "hermes"}
    for attempt in range(5):
        try:
            resp = requests.post(f"{ZHIYI_URL}/api/v1/commit", headers=HEADERS, json=payload, timeout=15)
            if resp.status_code in (200, 201):
                return resp.json().get("memory_ids", resp.json().get("episode_id", "ok"))
            elif resp.status_code == 429:
                time.sleep(resp.json().get("retry_after", 2))
                continue
            else:
                return f"HTTP_{resp.status_code}: {resp.text[:100]}"
        except Exception as e:
            if attempt < 4:
                time.sleep(1); continue
            return f"ERROR: {e}"
    return "MAX_RETRY"

def verify(db_path):
    db = lancedb.connect(str(db_path)); arrow = db.open_table("hermes_memory_default").to_arrow()
    contents = arrow.column("content").to_pylist()
    import random
    samples = random.sample([c for c in contents if c and len(str(c)) > 15], min(5, len(contents)))
    hits = 0
    for content in samples:
        resp = requests.post(f"{ZHIYI_URL}/api/v1/recall", headers=HEADERS,
                             json={"query": str(content)[:30], "top_k": 5, "namespace": "hermes-main"}, timeout=10)
        if resp.status_code == 200:
            results = resp.json().get("results", [])
            if any(str(content)[:20] in str(r.get("content", ""))[:50] for r in results):
                hits += 1; print(f"  ✓ {str(content)[:30]}...")
            else:
                print(f"  ✗ {str(content)[:30]}...")
        else:
            print(f"  ✗ HTTP {resp.status_code}")
    print(f"\n命中: {hits}/{len(samples)}")
    return hits == len(samples)

def migrate_table(db_path, tbl_name, state):
    key = tbl_name
    if key in state["tables_done"]:
        print(f"  [skip] {tbl_name}")
        return [], []
    db = lancedb.connect(str(db_path)); arrow = db.open_table(tbl_name).to_arrow()
    total = arrow.num_rows
    if total == 0:
        state["tables_done"].append(key); save_state(state)
        print(f"  [empty] {tbl_name}"); return [], []
    contents = arrow.column("content").to_pylist()
    tags = arrow.column("tag").to_pylist() if "tag" in arrow.schema.names else [None]*total
    agent_ids = arrow.column("agent_id").to_pylist() if "agent_id" in arrow.schema.names else ["hermes"]*total
    committed, failed = [], []; done_ids = set(state["committed"])
    for i, (c, tag, aid) in enumerate(zip(contents, tags, agent_ids)):
        if not c or not c.strip(): continue
        cid = hashlib.sha256(c.encode()).hexdigest()[:16]
        if cid in done_ids: continue
        result = commit(str(c), str(tag) if tag else None, str(aid) if aid else "hermes")
        if isinstance(result, str) and (result.startswith("HTTP_") or result.startswith("ERROR") or result.startswith("MAX")):
            failed.append({"i": i, "content": str(c)[:50], "error": result})
        else:
            committed.append(cid); done_ids.add(cid)
        if (i+1) % 20 == 0: print(f"  [{i+1}/{total}] ({len(committed)} ok, {len(failed)} fail)")
        time.sleep(BATCH_DELAY)
    state["committed"] = list(done_ids); state["tables_done"].append(key); save_state(state)
    print(f"  [{tbl_name}] {len(committed)} ✓, {len(failed)} ✗")
    return committed, failed

if __name__ == "__main__":
    db_path = Path.home() / ".hermes" / "memory_db" / "lancedb"
    tables = ["hermes_memory_default", "hermes_memory_hermes", "hermes_memory_muc"]
    if "--dry-run" in sys.argv:
        db = lancedb.connect(str(db_path))
        for t in tables:
            try:
                a = db.open_table(t).to_arrow(); print(f"{t}: {a.num_rows} rows, schema={a.schema.names}")
            except: print(f"{t}: error")
        sys.exit(0)
    if "--verify" in sys.argv: sys.exit(0 if verify(db_path) else 1)
    if "--cleanup" in sys.argv:
        import shutil; shutil.rmtree(str(db_path)); print(f"deleted {db_path}"); sys.exit(0)
    state = load_state()
    total_ok, total_fail = 0, 0
    for t in tables:
        ok, fail = migrate_table(db_path, t, state); total_ok += len(ok); total_fail += len(fail)
    print(f"\n{'='*20} 迁移完成 {'='*20}\n成功: {total_ok}, 失败: {total_fail}")
    if total_ok > 0 and total_fail == 0: print("运行 --verify 验证，--cleanup 删旧DB")