107 lines
3.3 KiB
Python
107 lines
3.3 KiB
Python
"""Hermes LanceDB → 织忆 MemoryWeave 迁移脚本 (v2)
|
||
用单条 commit，走服务端 BGE 编码，确保向量对齐。
|
||
"""
|
||
import json, time, requests
|
||
from concurrent.futures import ThreadPoolExecutor, as_completed
|
||
import lancedb
|
||
import pyarrow as pa
|
||
|
||
HERMES_DB = "/home/muc/.hermes/memory_db/lancedb"
|
||
ZHIYI_URL = "http://localhost:7821"
|
||
ZHIYI_KEY = "zhiyi-dev-key-2026"
|
||
WORKERS = 4
|
||
|
||
def log(msg):
|
||
print(f"[migrate] {msg}", flush=True)
|
||
|
||
def arrow_to_dicts(table: pa.Table) -> list[dict]:
|
||
rows = []
|
||
cols = table.column_names
|
||
for i in range(table.num_rows):
|
||
row = {}
|
||
for c in cols:
|
||
row[c] = table.column(c)[i].as_py()
|
||
rows.append(row)
|
||
return rows
|
||
|
||
def commit_one(r: dict) -> bool:
|
||
"""提交单条记忆到织忆"""
|
||
content = str(r.get("content", r.get("text", "")))
|
||
if not content.strip():
|
||
return False
|
||
try:
|
||
resp = requests.post(
|
||
f"{ZHIYI_URL}/api/v1/commit",
|
||
headers={"X-API-Key": ZHIYI_KEY, "Content-Type": "application/json"},
|
||
json={
|
||
"content": content,
|
||
"category": str(r.get("category", "general")),
|
||
"namespace": "hermes-main",
|
||
"agent_id": "hermes-a06",
|
||
},
|
||
timeout=30,
|
||
)
|
||
return resp.status_code in (200, 201)
|
||
except:
|
||
return False
|
||
|
||
def main():
|
||
log(f"打开 Hermes LanceDB: {HERMES_DB}")
|
||
db = lancedb.connect(HERMES_DB)
|
||
tables = db.table_names()
|
||
log(f"发现 {len(tables)} 个表: {tables}")
|
||
|
||
all_records = []
|
||
for tbl_name in tables:
|
||
tbl = db.open_table(tbl_name)
|
||
rows = arrow_to_dicts(tbl.to_arrow())
|
||
log(f" {tbl_name}: {len(rows)} 条")
|
||
all_records.extend(rows)
|
||
|
||
total = len(all_records)
|
||
log(f"总计 {total} 条，准备迁移到织忆（namespace=hermes-main）")
|
||
|
||
if total == 0:
|
||
log("无可迁移数据，退出")
|
||
return
|
||
|
||
# 并发迁移
|
||
migrated = 0
|
||
failed = 0
|
||
start = time.time()
|
||
|
||
with ThreadPoolExecutor(max_workers=WORKERS) as pool:
|
||
futures = {pool.submit(commit_one, r): i for i, r in enumerate(all_records)}
|
||
for f in as_completed(futures):
|
||
idx = futures[f]
|
||
try:
|
||
if f.result():
|
||
migrated += 1
|
||
else:
|
||
failed += 1
|
||
except:
|
||
failed += 1
|
||
|
||
if (migrated + failed) % 50 == 0:
|
||
elapsed = time.time() - start
|
||
rate = (migrated + failed) / elapsed
|
||
log(f" [{migrated+failed}/{total}] ✅{migrated} ❌{failed} ({rate:.1f}/s)")
|
||
|
||
# 验证
|
||
log(f"\n迁移完成: ✅{migrated} ❌{failed} / {total} ({time.time()-start:.0f}s)")
|
||
resp = requests.get(f"{ZHIYI_URL}/api/v1/stats", headers={"X-API-Key": ZHIYI_KEY})
|
||
stats = resp.json()
|
||
log(f"织忆当前: {stats['total_memories']} 条记忆")
|
||
|
||
# 召回验证
|
||
test = requests.post(
|
||
f"{ZHIYI_URL}/api/v1/recall",
|
||
headers={"X-API-Key": ZHIYI_KEY, "Content-Type": "application/json"},
|
||
json={"query": "牧尘用什么系统", "limit": 3},
|
||
).json()
|
||
log(f"召回测试: {test['count']} 条, top={test['results'][0]['content'][:60] if test['results'] else 'N/A'}")
|
||
log("✅ 迁移完成")
|
||
|
||
if __name__ == "__main__":
|
||
main()
|