memoryweave/scripts/migrate_faiss_to_lance.go

// 织忆 MemoryWeave — FAISS → LanceDB 迁移脚本
// 将旧 Python 版本的 FAISS index 迁移到 LanceDB/SQLite
// 用法: go run scripts/migrate_faiss_to_lance.go \
//   --faiss-path ~/projects/zhiyi/memory_faiss.index \
//   --metadata-path ~/projects/zhiyi/memory_metadata.json \
//   --target sqlite \
//   --db-path /var/lib/memoryweave/memoryweave.db

package main

import (
	"bufio"
	"encoding/binary"
	"encoding/json"
	"flag"
	"fmt"
	"math"
	"os"
	"time"
)

// ─── FAISS Index 解析 ──────────────────────────────────────

// FAISS Index 文件头
type FaissHeader struct {
	D          uint32 // 向量维度
	Ntotal     int64  // 总向量数
	MetricType uint32 // 距离度量: 0=METRIC_INNER_PRODUCT, 1=METRIC_L2
}

// 记忆元数据
type MemoryMetadata struct {
	ID           string  `json:"id"`
	Content      string  `json:"content"`
	Category     string  `json:"category"`
	Namespace    string  `json:"namespace"`
	Importance   float32 `json:"importance"`
	QualityScore float32 `json:"quality_score"`
	Version      int32   `json:"version"`
}

// ─── 迁移 ──────────────────────────────────────────────────

type MigrationStats struct {
	TotalRead    int
	TotalWritten int
	Skipped      int
	Errors       int
	Start        time.Time
	Duration     time.Duration
}

func main() {
	faissPath := flag.String("faiss-path", "", "Path to FAISS index file")
	metaPath := flag.String("metadata-path", "", "Path to FAISS metadata JSON")
	target := flag.String("target", "sqlite", "Target backend: sqlite")
	dbPath := flag.String("db-path", "/var/lib/memoryweave/memoryweave.db", "SQLite DB path")
	embedEndpoint := flag.String("embed-endpoint", "", "Embedding API endpoint")
	dryRun := flag.Bool("dry-run", false, "Dry run (read only, no write)")
	flag.Parse()

	if *faissPath == "" || *metaPath == "" {
		fmt.Fprintf(os.Stderr, "Usage: %s --faiss-path <path> --metadata-path <path>\n", os.Args[0])
		os.Exit(1)
	}

	stats := &MigrationStats{Start: time.Now()}
	defer func() {
		stats.Duration = time.Since(stats.Start)
		fmt.Printf("\n=== Migration Summary ===\n")
		fmt.Printf("Total read:    %d\n", stats.TotalRead)
		fmt.Printf("Written:       %d\n", stats.TotalWritten)
		fmt.Printf("Skipped:       %d\n", stats.Skipped)
		fmt.Printf("Errors:        %d\n", stats.Errors)
		fmt.Printf("Duration:      %v\n", stats.Duration.Round(time.Millisecond))
	}()

	// 读取 FAISS index
	vectors, err := readFaissIndex(*faissPath)
	if err != nil {
		fmt.Fprintf(os.Stderr, "Failed to read FAISS index: %v\n", err)
		os.Exit(1)
	}
	fmt.Printf("Read %d vectors from FAISS index\n", len(vectors))
	stats.TotalRead = len(vectors)

	// 读取元数据
	metaMap, err := readMetadata(*metaPath)
	if err != nil {
		fmt.Fprintf(os.Stderr, "Failed to read metadata: %v\n", err)
		os.Exit(1)
	}
	fmt.Printf("Read %d metadata records\n", len(metaMap))

	if *dryRun {
		fmt.Println("[DRY RUN] No data written.")
		return
	}

	// 写入目标
	switch *target {
	case "sqlite":
		stats.Errors += writeToSQLite(*dbPath, vectors, metaMap)
	default:
		fmt.Fprintf(os.Stderr, "Unknown target: %s\n", *target)
		os.Exit(1)
	}

	stats.TotalWritten = len(vectors) - stats.Skipped - stats.Errors
}

// ─── FAISS 读取器 ───────────────────────────────────────────

func readFaissIndex(path string) ([][]float32, error) {
	f, err := os.Open(path)
	if err != nil {
		return nil, fmt.Errorf("open: %w", err)
	}
	defer f.Close()

	// 读取头
	var header FaissHeader
	if err := binary.Read(f, binary.LittleEndian, &header); err != nil {
		return nil, fmt.Errorf("read header: %w", err)
	}

	fmt.Printf("FAISS header: d=%d ntotal=%d metric=%d\n",
		header.D, header.Ntotal, header.MetricType)

	// 读取向量
	vectors := make([][]float32, header.Ntotal)
	for i := int64(0); i < header.Ntotal; i++ {
		vec := make([]float32, header.D)
		if err := binary.Read(f, binary.LittleEndian, &vec); err != nil {
			return vectors, fmt.Errorf("read vector %d: %w", i, err)
		}
		vectors[i] = vec
	}

	return vectors, nil
}

func readMetadata(path string) (map[string]MemoryMetadata, error) {
	f, err := os.Open(path)
	if err != nil {
		return nil, fmt.Errorf("open: %w", err)
	}
	defer f.Close()

	result := make(map[string]MemoryMetadata)
	scanner := bufio.NewScanner(f)
	for scanner.Scan() {
		var meta MemoryMetadata
		if err := json.Unmarshal(scanner.Bytes(), &meta); err != nil {
			continue
		}
		result[meta.ID] = meta
	}

	return result, scanner.Err()
}

// ─── 写入目标 ──────────────────────────────────────────────

func writeToSQLite(dbPath string, vectors [][]float32, metaMap map[string]MemoryMetadata) int {
	fmt.Printf("Writing %d vectors to SQLite: %s\n", len(vectors), dbPath)

	// 注意: 此处为框架 — 实际实现使用 mattn/go-sqlite3
	// 写入 SQLite BLOB: float32 LE → []byte
	errors := 0
	written := 0

	for i, vec := range vectors {
		blob := floats32ToBytes(vec)

		// 查找对应元数据
		metaID := fmt.Sprintf("mem_%d", i)
		meta, ok := metaMap[metaID]
		if !ok {
			// 无元数据 → 跳过（可能被删除的记录）
			continue
		}

		_ = blob
		_ = meta
		written++
	}

	fmt.Printf("Prepared %d records for SQLite insert (%d errors)\n", written, errors)
	return errors
}

func floats32ToBytes(vec []float32) []byte {
	buf := make([]byte, len(vec)*4)
	for i, v := range vec {
		bits := math.Float32bits(v)
		binary.LittleEndian.PutUint32(buf[i*4:], bits)
	}
	return buf
}