Cyrene_v2/internal/runtime/memory/vector/sqlite.go

package vector

import (
	"context"
	"database/sql"
	"encoding/binary"
	"fmt"
	"log/slog"
	"math"
	"sort"
	"strings"

	"git.yeij.top/AskaEth/Cyrene_v2/internal/runtime/embedding"
)

// createVecItemsTable 是幂等的 SQLite schema。vector 以小端
// float32 值的 BLOB 存储；dim 记录维度，因此 blob
// 无需外部信息即可解码。该表由
// NewSQLite 创建；不要求与 internal/db 的迁移兼容，因为
// 向量存储是自包含的（见 internal/runtime/memory/store.go）。
const createVecItemsTable = `
CREATE TABLE IF NOT EXISTS vec_items (
    id     TEXT PRIMARY KEY,
    vector BLOB NOT NULL,
    dim    INTEGER NOT NULL
)`

// SQLiteIndex 把向量以 BLOB 存入 SQLite 的 vec_items 表，并以
// Go 侧线性扫描的余弦相似度作答相似度（M4 里程碑
// 语义：小数据集，无 ANN 索引）。
//
// 并发交由 *sql.DB，按约定它可安全并发使用，
// 因此 SQLiteIndex 不再持锁。每个查询都参数化
// （?）且列清单显式列出。
type SQLiteIndex struct {
	db       *sql.DB
	dim      int
	embedder embedding.Embedder
	log      *slog.Logger
}

// NewSQLite 在 db 上为维度 dim 的向量构建 SQLiteIndex，确保
// vec_items schema 存在。db 必须是存活连接；nil db 返回
// ErrNilDB，dim 非正返回 ErrInvalidDim。ctx 控制 schema
// 创建。log 可为 nil。
//
// 返回的索引没有接入 embedder，因此直接调用 Embed
// 会返回 ErrNoEmbedder；用 New 构建完全接好的索引，或在调用
// Embed 前接入一个。
func NewSQLite(ctx context.Context, db *sql.DB, dim int, log *slog.Logger) (*SQLiteIndex, error) {
	if db == nil {
		return nil, ErrNilDB
	}
	if dim <= 0 {
		return nil, ErrInvalidDim
	}
	log = sourcedLogger(log)
	if _, err := db.ExecContext(ctx, createVecItemsTable); err != nil {
		return nil, fmt.Errorf("memory/vector: create sqlite schema: %w", err)
	}
	log.Info("sqlite vector index ready", "dim", dim)
	return &SQLiteIndex{db: db, dim: dim, log: log}, nil
}

// withEmbedder 把产向量的 embedder 接入索引。它由
// New 使用；测试也可直接调用。
func (s *SQLiteIndex) withEmbedder(e embedding.Embedder) { s.embedder = e }

// Dim 返回固定向量维度。
func (s *SQLiteIndex) Dim() int { return s.dim }

// Embed 用接入的 embedder 把 texts 转成向量。每段输入文本
// 对应一个按索引对齐的 Vector。未接入 embedder 时返回
// ErrNoEmbedder。
func (s *SQLiteIndex) Embed(ctx context.Context, texts []string) ([]Vector, error) {
	if s.embedder == nil {
		return nil, ErrNoEmbedder
	}
	vecs, err := s.embedder.Embed(ctx, texts)
	if err != nil {
		return nil, fmt.Errorf("memory/vector: embed %d texts: %w", len(texts), err)
	}
	out := make([]Vector, len(vecs))
	for i, v := range vecs {
		out[i] = Vector(v)
	}
	return out, nil
}

// Add 持久化 (id, vector) 对，id 重复时 upsert，
// 最新的向量生效。它拒绝空 id（ErrMissingID）以及与
// 索引维度不一致的向量（ErrDimMismatch）。ctx 被取消时
// 中止写入。
func (s *SQLiteIndex) Add(ctx context.Context, id string, v Vector) error {
	if strings.TrimSpace(id) == "" {
		return ErrMissingID
	}
	if len(v) != s.dim {
		return fmt.Errorf("memory/vector: %w (got %d, want %d)", ErrDimMismatch, len(v), s.dim)
	}
	_, err := s.db.ExecContext(ctx, `INSERT INTO vec_items (id, vector, dim)
		VALUES (?, ?, ?)
		ON CONFLICT(id) DO UPDATE SET vector = excluded.vector, dim = excluded.dim`,
		id, encodeVector(v), s.dim)
	if err != nil {
		return fmt.Errorf("memory/vector: add %q: %w", id, err)
	}
	return nil
}

// Vectors 返回与 ids 对齐的已存向量：命中按位置给出，未命中（含维度与
// 本索引不符的行）为 nil。它是「增量向量化」的读口 —— 调用方据此只嵌
// 缺的那些，而不是每轮重算整个分区。它只读，不写任何行。
func (s *SQLiteIndex) Vectors(ctx context.Context, ids []string) ([]Vector, error) {
	if len(ids) == 0 {
		return []Vector{}, nil
	}
	found := make(map[string]Vector, len(ids))
	for start := 0; start < len(ids); start += vectorsQueryChunk {
		end := min(start+vectorsQueryChunk, len(ids))
		if err := s.lookupVectors(ctx, ids[start:end], found); err != nil {
			return nil, err
		}
	}
	return alignVectors(ids, found), nil
}

// lookupVectors 查一批 id 并把命中的向量写进 found。维度与本索引不符的行
// 当作未命中（它不能参与本索引的余弦，重算一次才是对的）。
func (s *SQLiteIndex) lookupVectors(ctx context.Context, ids []string, found map[string]Vector) error {
	query := `SELECT id, vector FROM vec_items WHERE id IN (` + sqlitePlaceholders(len(ids)) + `)`
	rows, err := s.db.QueryContext(ctx, query, idArgs(ids)...)
	if err != nil {
		return fmt.Errorf("memory/vector: lookup vectors: %w", err)
	}
	defer func() { _ = rows.Close() }()

	for rows.Next() {
		var (
			id   string
			blob []byte
		)
		if err := rows.Scan(&id, &blob); err != nil {
			return fmt.Errorf("memory/vector: scan vector: %w", err)
		}
		v, err := decodeVector(blob)
		if err != nil || len(v) != s.dim {
			continue
		}
		found[id] = v
	}
	if err := rows.Err(); err != nil {
		return fmt.Errorf("memory/vector: iterate vectors: %w", err)
	}
	return nil
}

// Search 返回与 q 最相似的最多 topK 个 id，按余弦相似度
// 降序排列。它在 Go 侧扫描每个已持久化向量（线性扫描，无 ANN）。topK
// 非正返回空结果；长度与
// 索引维度不一致的查询返回 ErrDimMismatch。
func (s *SQLiteIndex) Search(ctx context.Context, q Vector, topK int) ([]Hit, error) {
	if topK <= 0 {
		return []Hit{}, nil
	}
	if len(q) != s.dim {
		return nil, fmt.Errorf("memory/vector: %w (query %d, want %d)", ErrDimMismatch, len(q), s.dim)
	}
	rows, err := s.db.QueryContext(ctx, `SELECT id, vector FROM vec_items`)
	if err != nil {
		return nil, fmt.Errorf("memory/vector: search: %w", err)
	}
	defer func() { _ = rows.Close() }()

	hits := make([]Hit, 0)
	for rows.Next() {
		var (
			id   string
			blob []byte
		)
		if err := rows.Scan(&id, &blob); err != nil {
			return nil, fmt.Errorf("memory/vector: scan item: %w", err)
		}
		v, err := decodeVector(blob)
		if err != nil {
			return nil, fmt.Errorf("memory/vector: decode %q: %w", id, err)
		}
		hits = append(hits, Hit{ID: id, Score: cosine(q, v)})
	}
	if err := rows.Err(); err != nil {
		return nil, fmt.Errorf("memory/vector: iterate items: %w", err)
	}

	sort.Slice(hits, func(i, j int) bool { return hits[i].Score > hits[j].Score })
	if len(hits) > topK {
		hits = hits[:topK]
	}
	return hits, nil
}

// Close 释放底层的 SQLite 数据库。数据库
// 已关闭时调用也是安全的。
func (s *SQLiteIndex) Close() error {
	if s.db == nil {
		return nil
	}
	err := s.db.Close()
	if err != nil {
		return fmt.Errorf("memory/vector: close sqlite: %w", err)
	}
	return nil
}

// encodeVector 把 v 序列化为小端 float32 BLOB。它不会失败。
func encodeVector(v Vector) []byte {
	buf := make([]byte, 4*len(v))
	for i, x := range v {
		binary.LittleEndian.PutUint32(buf[i*4:], math.Float32bits(x))
	}
	return buf
}

// decodeVector 把 BLOB 反序列化回 Vector。它拒绝空或
// 非 4 字节对齐的 blob，使损坏的行以错误暴露，而不是
// 被静默误读。
func decodeVector(blob []byte) (Vector, error) {
	if len(blob) == 0 || len(blob)%4 != 0 {
		return nil, fmt.Errorf("memory/vector: invalid vector blob length %d", len(blob))
	}
	v := make(Vector, len(blob)/4)
	for i := range v {
		v[i] = math.Float32frombits(binary.LittleEndian.Uint32(blob[i*4:]))
	}
	return v, nil
}

// cosine 返回 a 与 b 的余弦相似度，取值 [-1, 1]。长度为零或
// 不一致的向量得 0，因此不一致的存储行绝不会破坏
// 排名。
func cosine(a, b Vector) float64 {
	if len(a) == 0 || len(b) == 0 || len(a) != len(b) {
		return 0
	}
	var dot, na, nb float64
	for i := range a {
		dot += float64(a[i]) * float64(b[i])
		na += float64(a[i]) * float64(a[i])
		nb += float64(b[i]) * float64(b[i])
	}
	if na == 0 || nb == 0 {
		return 0
	}
	return dot / (math.Sqrt(na) * math.Sqrt(nb))
}