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 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))
}