rony-llm-agent/pkg/rag/backends/sqlitevec/sqlitevec.go

// Package sqlitevec is a zero-dependency rag.Backend backed by SQLite. When
// given a query vector it scores candidates with a brute-force cosine
// similarity scan in Go (no ANN index, so it trades scale — fine up to a few
// tens of thousands of fragments, comfortably covering a single user's saved
// notes/processes — for requiring nothing beyond the pure-Go sqlite driver
// already used elsewhere in Rony). When no query vector is available (no
// embedder configured, or it failed), it falls back to SQLite's built-in
// FTS5 full-text search over the fragment's content, ranked by BM25.
package sqlitevec

import (
	"context"
	"database/sql"
	"encoding/binary"
	"encoding/json"
	"fmt"
	"math"
	"sort"
	"strings"
	"unicode"

	"github.com/VictorVargas/rony-llm-agent/pkg/rag"
	_ "modernc.org/sqlite"
)

// Backend implements rag.Backend on top of a local SQLite file.
type Backend struct {
	db *sql.DB
}

// New opens (or creates) the SQLite-backed vector store at path. An empty
// path opens an in-memory store, useful for tests.
func New(path string) (*Backend, error) {
	dsn := path
	if dsn == "" {
		dsn = "file::memory:?cache=shared"
	}

	db, err := sql.Open("sqlite", dsn)
	if err != nil {
		return nil, fmt.Errorf("open sqlite: %w", err)
	}

	b := &Backend{db: db}
	if err := b.initSchema(); err != nil {
		db.Close()
		return nil, fmt.Errorf("init schema: %w", err)
	}
	return b, nil
}

func (b *Backend) initSchema() error {
	_, err := b.db.Exec(`
	CREATE TABLE IF NOT EXISTS fragments (
		id       TEXT PRIMARY KEY,
		vector   BLOB NOT NULL,
		content  TEXT NOT NULL,
		metadata TEXT NOT NULL
	);
	CREATE VIRTUAL TABLE IF NOT EXISTS fragments_fts USING fts5(
		id UNINDEXED,
		content
	);`)
	return err
}

// Upsert stores or replaces a fragment's vector, text and metadata, keeping
// the FTS5 index in sync for lexical fallback search.
func (b *Backend) Upsert(ctx context.Context, id string, vector []float32, content string, metadata map[string]string) error {
	metaJSON, err := json.Marshal(metadata)
	if err != nil {
		return fmt.Errorf("marshal metadata: %w", err)
	}

	tx, err := b.db.BeginTx(ctx, nil)
	if err != nil {
		return fmt.Errorf("begin tx: %w", err)
	}
	defer tx.Rollback()

	_, err = tx.ExecContext(ctx,
		`INSERT INTO fragments (id, vector, content, metadata) VALUES (?, ?, ?, ?)
		 ON CONFLICT(id) DO UPDATE SET vector = excluded.vector, content = excluded.content, metadata = excluded.metadata`,
		id, encodeVector(vector), content, string(metaJSON),
	)
	if err != nil {
		return fmt.Errorf("upsert fragment: %w", err)
	}

	// FTS5 tables don't support ON CONFLICT, so re-sync via delete+insert.
	if _, err := tx.ExecContext(ctx, `DELETE FROM fragments_fts WHERE id = ?`, id); err != nil {
		return fmt.Errorf("clear fts entry: %w", err)
	}
	if _, err := tx.ExecContext(ctx, `INSERT INTO fragments_fts (id, content) VALUES (?, ?)`, id, content); err != nil {
		return fmt.Errorf("index fts entry: %w", err)
	}

	return tx.Commit()
}

// Search scores fragments against queryVector by cosine similarity when one
// is available; otherwise it falls back to an FTS5 lexical match on query.
func (b *Backend) Search(ctx context.Context, query string, queryVector []float32, topK int) ([]rag.SearchResult, error) {
	if len(queryVector) > 0 {
		return b.searchByVector(ctx, queryVector, topK)
	}
	return b.searchByText(ctx, query, topK)
}

func (b *Backend) searchByVector(ctx context.Context, queryVector []float32, topK int) ([]rag.SearchResult, error) {
	rows, err := b.db.QueryContext(ctx, `SELECT id, vector, content, metadata FROM fragments`)
	if err != nil {
		return nil, fmt.Errorf("query fragments: %w", err)
	}
	defer rows.Close()

	var candidates []rag.SearchResult
	for rows.Next() {
		var id, content, metaJSON string
		var vecBlob []byte
		if err := rows.Scan(&id, &vecBlob, &content, &metaJSON); err != nil {
			return nil, fmt.Errorf("scan fragment: %w", err)
		}

		metadata, err := decodeMetadata(metaJSON)
		if err != nil {
			return nil, err
		}

		candidates = append(candidates, rag.SearchResult{
			ID:       id,
			Content:  content,
			Score:    cosineSimilarity(queryVector, decodeVector(vecBlob)),
			Metadata: metadata,
		})
	}
	if err := rows.Err(); err != nil {
		return nil, fmt.Errorf("iterate fragments: %w", err)
	}

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

func (b *Backend) searchByText(ctx context.Context, query string, topK int) ([]rag.SearchResult, error) {
	ftsQuery := buildFTSQuery(query)
	if ftsQuery == "" {
		return []rag.SearchResult{}, nil
	}

	rows, err := b.db.QueryContext(ctx, `
		SELECT fragments.id, fragments.content, fragments.metadata, bm25(fragments_fts) AS rank
		FROM fragments_fts
		JOIN fragments ON fragments.id = fragments_fts.id
		WHERE fragments_fts MATCH ?
		ORDER BY rank
		LIMIT ?`, ftsQuery, topK)
	if err != nil {
		return nil, fmt.Errorf("fts query: %w", err)
	}
	defer rows.Close()

	var results []rag.SearchResult
	for rows.Next() {
		var id, content, metaJSON string
		var rank float64
		if err := rows.Scan(&id, &content, &metaJSON, &rank); err != nil {
			return nil, fmt.Errorf("scan fts result: %w", err)
		}

		metadata, err := decodeMetadata(metaJSON)
		if err != nil {
			return nil, err
		}

		// bm25() returns lower-is-better (often negative); negate so a
		// higher Score means a better match, matching the vector path.
		results = append(results, rag.SearchResult{
			ID:       id,
			Content:  content,
			Score:    float32(-rank),
			Metadata: metadata,
		})
	}
	if err := rows.Err(); err != nil {
		return nil, fmt.Errorf("iterate fts results: %w", err)
	}
	return results, nil
}

// Forget deletes a single fragment by ID.
func (b *Backend) Forget(ctx context.Context, id string) error {
	tx, err := b.db.BeginTx(ctx, nil)
	if err != nil {
		return fmt.Errorf("begin tx: %w", err)
	}
	defer tx.Rollback()

	if _, err := tx.ExecContext(ctx, `DELETE FROM fragments WHERE id = ?`, id); err != nil {
		return fmt.Errorf("delete fragment: %w", err)
	}
	if _, err := tx.ExecContext(ctx, `DELETE FROM fragments_fts WHERE id = ?`, id); err != nil {
		return fmt.Errorf("delete fts entry: %w", err)
	}
	return tx.Commit()
}

// ForgetAll deletes every stored fragment.
func (b *Backend) ForgetAll(ctx context.Context) error {
	tx, err := b.db.BeginTx(ctx, nil)
	if err != nil {
		return fmt.Errorf("begin tx: %w", err)
	}
	defer tx.Rollback()

	if _, err := tx.ExecContext(ctx, `DELETE FROM fragments`); err != nil {
		return fmt.Errorf("delete all fragments: %w", err)
	}
	if _, err := tx.ExecContext(ctx, `DELETE FROM fragments_fts`); err != nil {
		return fmt.Errorf("delete all fts entries: %w", err)
	}
	return tx.Commit()
}

// Close releases the underlying database connection.
func (b *Backend) Close() error {
	return b.db.Close()
}

func decodeMetadata(metaJSON string) (map[string]string, error) {
	if metaJSON == "" {
		return nil, nil
	}
	var metadata map[string]string
	if err := json.Unmarshal([]byte(metaJSON), &metadata); err != nil {
		return nil, fmt.Errorf("unmarshal metadata: %w", err)
	}
	return metadata, nil
}

// buildFTSQuery turns free-form text into an FTS5 MATCH query that ORs
// together each token as a quoted phrase, so punctuation or FTS5 operator
// characters (-, *, :, "...) in the input can't produce a syntax error, and
// any subset of tokens can match (rather than requiring the exact phrase).
func buildFTSQuery(text string) string {
	tokens := tokenizeForFTS(text)
	if len(tokens) == 0 {
		return ""
	}
	parts := make([]string, len(tokens))
	for i, tok := range tokens {
		parts[i] = `"` + strings.ReplaceAll(tok, `"`, `""`) + `"`
	}
	return strings.Join(parts, " OR ")
}

func tokenizeForFTS(text string) []string {
	return strings.FieldsFunc(strings.ToLower(text), func(r rune) bool {
		return !unicode.IsLetter(r) && !unicode.IsNumber(r)
	})
}

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

func decodeVector(buf []byte) []float32 {
	v := make([]float32, len(buf)/4)
	for i := range v {
		v[i] = math.Float32frombits(binary.LittleEndian.Uint32(buf[i*4:]))
	}
	return v
}

func cosineSimilarity(a, b []float32) float32 {
	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 float32(dot / (math.Sqrt(na) * math.Sqrt(nb)))
}