feat(memory/vector): 向量检索引擎（SQLite BLOB 内积默认 + PG pgvector 真语义增强，配置后真生效）

2026-09-06 06:04:44 +08:00
1 parent ef79edd96a
commit fb880572a4
@@ -0,0 +1,207 @@
+package vector
+
+import (
+	"context"
+	"database/sql"
+	"encoding/binary"
+	"fmt"
+	"log/slog"
+	"math"
+	"sort"
+	"strings"
+
+	"git.yeij.top/AskaEth/Cyrene_v2/internal/runtime/embedding"
+)
+
+// createVecItemsTable is the idempotent SQLite schema. vector is stored as a
+// BLOB of little-endian float32 values; dim records the dimension so the blob
+// can be decoded without external knowledge. The table is created by
+// NewSQLite; migration compatibility with internal/db is not required because
+// the vector store is self-contained (see 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 stores vectors as BLOBs in a SQLite vec_items table and answers
+// similarity with Go-side cosine similarity over a linear scan (milestone M4
+// semantics: small datasets, no ANN index).
+//
+// Concurrency is delegated to the *sql.DB, which is safe for concurrent use by
+// contract, so SQLiteIndex holds no further lock. Every query is parameterized
+// (?) and columns are listed explicitly.
+type SQLiteIndex struct {
+	db       *sql.DB
+	dim      int
+	embedder embedding.Embedder
+	log      *slog.Logger
+}
+
+// NewSQLite builds a SQLiteIndex on db for vectors of dimension dim, ensuring
+// the vec_items schema exists. db must be a live connection; a nil db returns
+// ErrNilDB, and a non-positive dim returns ErrInvalidDim. ctx controls schema
+// creation. log may be nil.
+//
+// The returned index has no embedder wired in, so calling Embed directly
+// returns ErrNoEmbedder; use New to build a fully wired index, or wire one in
+// before calling 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 wires the vector-producing embedder into the index. It is used
+// by New; tests may also call it directly.
+func (s *SQLiteIndex) withEmbedder(e embedding.Embedder) { s.embedder = e }
+
+// Dim returns the fixed vector dimension.
+func (s *SQLiteIndex) Dim() int { return s.dim }
+
+// Embed turns texts into vectors with the wired embedder. It returns one Vector
+// per input text, aligned by index. With no embedder wired it returns
+// 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 persists the (id, vector) pair, upserting on a repeated id so the
+// latest vector wins. It rejects a blank id (ErrMissingID) and a vector whose
+// length differs from the index dimension (ErrDimMismatch). ctx aborts the
+// write when canceled.
+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
+}
+
+// Search returns up to topK ids most similar to q, ranked by descending cosine
+// similarity. It scans every persisted vector in Go (linear scan, no ANN). A
+// non-positive topK returns an empty result; a query whose length differs from
+// the index dimension returns 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 releases the underlying SQLite database. It is safe to call when the
+// database is already closed.
+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 serializes v to a little-endian float32 BLOB. It cannot fail.
+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 deserializes a BLOB back into a Vector. It rejects empty or
+// non-4-byte-aligned blobs so a corrupted row surfaces as an error rather than
+// a silent misread.
+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 returns the cosine similarity of a and b in [-1, 1]. Zero-length or
+// mismatched vectors yield 0, so an inconsistent stored row never corrupts the
+// ranking.
+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))
+}