147 lines
4.5 KiB
Python
147 lines
4.5 KiB
Python
#!/usr/bin/env python3
|
|
"""FAISS → LanceDB 迁移脚本
|
|
|
|
将 Python 织忆的 FAISS 索引转换为 Go 织忆的 LanceDB 格式。
|
|
|
|
用法:
|
|
python3 scripts/migrate.py \
|
|
--data-dir ~/projects/zhiyi/data \
|
|
--output-dir /var/lib/zhiyi/data
|
|
|
|
要求:
|
|
pip install lancedb numpy
|
|
"""
|
|
|
|
import argparse
|
|
import json
|
|
import os
|
|
import sys
|
|
import time
|
|
|
|
def parse_args():
|
|
p = argparse.ArgumentParser(description='FAISS → LanceDB 迁移')
|
|
p.add_argument('--data-dir', required=True, help='旧Python织忆data目录')
|
|
p.add_argument('--output-dir', default='/var/lib/zhiyi/data', help='LanceDB输出目录')
|
|
p.add_argument('--dry-run', action='store_true', help='只统计，不执行')
|
|
return p.parse_args()
|
|
|
|
def main():
|
|
args = parse_args()
|
|
|
|
# 1. 读取旧数据
|
|
distilled_file = os.path.join(args.data_dir, 'hermes-main', 'distilled', '2026-05.jsonl')
|
|
episodes_file = os.path.join(args.data_dir, 'hermes-main', 'episodes', '2026-05.jsonl')
|
|
vectors_file = os.path.join(args.data_dir, 'sbert_vectors.jsonl')
|
|
meta_file = os.path.join(args.data_dir, 'sbert_meta.json')
|
|
|
|
distilled = []
|
|
episodes = []
|
|
|
|
for fname in [distilled_file, episodes_file]:
|
|
if os.path.exists(fname):
|
|
with open(fname) as f:
|
|
for line in f:
|
|
distilled.append(json.loads(line.strip()))
|
|
|
|
# 2. 读取向量
|
|
vectors = []
|
|
meta = {}
|
|
if os.path.exists(meta_file):
|
|
with open(meta_file) as f:
|
|
meta = json.load(f)
|
|
if os.path.exists(vectors_file):
|
|
with open(vectors_file) as f:
|
|
for line in f:
|
|
vectors.append(json.loads(line.strip()))
|
|
|
|
total = len(distilled)
|
|
vector_count = len(vectors)
|
|
print(f"📊 已统计: {total} 条 distilled, {vector_count} 条向量")
|
|
|
|
if args.dry_run:
|
|
print(" (dry-run 模式，不执行)")
|
|
return
|
|
|
|
# 3. 写入 LanceDB
|
|
try:
|
|
import lancedb
|
|
import numpy as np
|
|
except ImportError:
|
|
print("❌ 需要 lancedb: pip install lancedb numpy")
|
|
sys.exit(1)
|
|
|
|
os.makedirs(args.output_dir, exist_ok=True)
|
|
db = lancedb.connect(args.output_dir)
|
|
|
|
# 创建 memories 表
|
|
try:
|
|
table = db.create_table("memories", [{
|
|
"id": "init",
|
|
"agent_id": "system",
|
|
"namespace": "default",
|
|
"content": "init",
|
|
"category": "system_fact",
|
|
"vector": [0.0]*1024,
|
|
"tier": "normal",
|
|
"quality_score": 0.0,
|
|
"recall_count": 0,
|
|
"freshness": "fresh",
|
|
"created_at": time.strftime("%Y-%m-%dT%H:%M:%SZ"),
|
|
"updated_at": time.strftime("%Y-%m-%dT%H:%M:%SZ"),
|
|
"is_deleted": False,
|
|
}])
|
|
except Exception:
|
|
table = db.open_table("memories")
|
|
|
|
# 迁移数据
|
|
migrated = 0
|
|
batch = []
|
|
|
|
for i, d in enumerate(distilled):
|
|
content = d.get('content', '') or d.get('summary', '') or ''
|
|
if not content.strip():
|
|
continue
|
|
|
|
vec = vectors[i] if i < len(vectors) else [0.0] * 1024
|
|
if isinstance(vec, dict):
|
|
vec = vec.get('vector', [0.0]*1024)
|
|
|
|
# 确保1024维
|
|
if len(vec) != 1024:
|
|
if len(vec) < 1024:
|
|
vec = vec + [0.0] * (1024 - len(vec))
|
|
else:
|
|
vec = vec[:1024]
|
|
|
|
record = {
|
|
"id": d.get('id', f'migrated_{i}'),
|
|
"agent_id": d.get('agent_id', 'hermes-main'),
|
|
"namespace": d.get('namespace', 'hermes-main'),
|
|
"content": content[:2000],
|
|
"category": d.get('category', 'general'),
|
|
"vector": vec,
|
|
"tier": "core" if d.get('importance', 0) >= 3 else "normal",
|
|
"quality_score": float(d.get('importance', 0.5)) / 5.0,
|
|
"recall_count": d.get('recall_count', 0),
|
|
"freshness": "fresh",
|
|
"created_at": d.get('created_at', time.strftime("%Y-%m-%dT%H:%M:%SZ")),
|
|
"updated_at": d.get('updated_at', d.get('created_at', time.strftime("%Y-%m-%dT%H:%M:%SZ"))),
|
|
"is_deleted": False,
|
|
}
|
|
batch.append(record)
|
|
|
|
if len(batch) >= 100:
|
|
table.add(batch)
|
|
migrated += len(batch)
|
|
batch = []
|
|
print(f" ✓ 已迁移 {migrated}/{total}")
|
|
|
|
if batch:
|
|
table.add(batch)
|
|
migrated += len(batch)
|
|
|
|
print(f"\n✅ 迁移完成: {migrated} 条记录 → {args.output_dir}")
|
|
|
|
if __name__ == '__main__':
|
|
main()
|