Maven/internal/store/memory.go

package store

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

	"github.com/kami/maven/internal/memory"
)

// MemoryStore is the persistent backend for internal/memory's vector Store,
// sharing the main encrypted sqlite database so recall text (note/fact bodies
// carried in the meta blob) inherits at-rest encryption — a plaintext sidecar
// file would undercut store.OpenEncrypted. It survives daemon restarts, which
// the InMemoryStore does not: that was the last gap keeping long-term memory
// from being real.
//
// Search is brute-force cosine over every row loaded into memory — the same
// algorithm as InMemoryStore, just sourced from disk. At the single-user note+
// fact scale (thousands of rows, not millions) a full scan per query is well
// under a millisecond; an ANN index is the swap for later, behind this same
// interface. Vectors are assumed L2-normalized by the embedder, so cosine is a
// dot product.
type MemoryStore struct {
	db *sql.DB
}

// VectorMemory returns a persistent memory.Store backed by this store's db.
// The returned store shares the db handle (single writer — the daemon), so it
// participates in the same encrypted tmpfs working copy and is sealed on Close.
func (s *Store) VectorMemory() *MemoryStore {
	return &MemoryStore{db: s.db}
}

// compile-time check: MemoryStore satisfies the memory.Store interface, and the
// wider Catalog that speaker profiles need (enumerate by prefix, delete by id).
var _ memory.Store = (*MemoryStore)(nil)
var _ memory.Catalog = (*MemoryStore)(nil)

// Insert upserts a vector by id: a repeated id replaces the prior row rather
// than accumulating duplicates (the note/fact ids are stable and unique, so a
// re-index is an update, not a second copy — an improvement on InMemoryStore's
// append-always). meta is stored as a JSON object.
func (m *MemoryStore) Insert(ctx context.Context, id string, vec []float32, meta map[string]string) error {
	metaJSON, err := json.Marshal(meta)
	if err != nil {
		return fmt.Errorf("memory: marshal meta: %w", err)
	}
	_, err = m.db.ExecContext(ctx,
		`INSERT INTO memory_vectors (id, vec, meta, created_ts) VALUES (?,?,?,?)
		 ON CONFLICT(id) DO UPDATE SET vec = excluded.vec, meta = excluded.meta, created_ts = excluded.created_ts`,
		id, encodeVec(vec), string(metaJSON), time.Now().UnixMilli())
	if err != nil {
		return fmt.Errorf("memory: insert %q: %w", id, err)
	}
	return nil
}

// Search returns the topK nearest rows by cosine similarity. A full scan; see
// the type doc for why that's fine at this scale.
//
// Rows under memory.NonRecallPrefix are excluded in SQL. They are speaker
// voiceprints sharing this table, and note recall must not rank them; see that
// constant for why the previous arrangement only appeared to do this.
//
// Every row is still scored, because a full scan is what picks the winners.
// What the scan does NOT do is pay for a row it is about to discard: the score
// is read straight off the stored bytes without materializing a []float32, and
// the meta blob is copied and unmarshalled only for a row that has entered the
// topK. Losers cost one dot product and nothing else. Ranking is unchanged —
// same scores, same order, same ties.
func (m *MemoryStore) Search(ctx context.Context, vec []float32, topK int) ([]memory.Result, error) {
	if topK <= 0 {
		topK = 10
	}
	rows, err := m.db.QueryContext(ctx,
		`SELECT id, vec, meta FROM memory_vectors WHERE id NOT LIKE ? ESCAPE '\'`,
		escapeLike(memory.NonRecallPrefix)+"%")
	if err != nil {
		return nil, fmt.Errorf("memory: scan: %w", err)
	}
	defer rows.Close()

	// sql.RawBytes hands us the driver's own buffer, valid only until the next
	// Next(). Nothing here outlives the row except what topK.offer copies on a
	// survivor, so the three columns cost no allocation per row.
	var id, blob, metaJSON sql.RawBytes
	top := newTopK(topK)
	for rows.Next() {
		if err := rows.Scan(&id, &blob, &metaJSON); err != nil {
			return nil, fmt.Errorf("memory: row: %w", err)
		}
		top.offer(dotBlob(vec, blob), id, metaJSON)
	}
	if err := rows.Err(); err != nil {
		return nil, fmt.Errorf("memory: rows: %w", err)
	}

	survivors := top.sorted()
	out := make([]memory.Result, 0, len(survivors))
	for _, c := range survivors {
		meta := map[string]string{}
		if err := json.Unmarshal(c.meta, &meta); err != nil {
			return nil, fmt.Errorf("memory: unmarshal meta for %q: %w", c.id, err)
		}
		out = append(out, memory.Result{ID: c.id, Score: c.score, Meta: meta})
	}
	return out, nil
}

// candidate is one row that is currently in the topK: its score, its id, and
// its meta blob copied out of the driver's buffer. The copy is the price of
// surviving, and only survivors pay it.
type candidate struct {
	score float64
	id    string
	meta  []byte
}

// topK keeps the k highest-scoring candidates seen so far as a min-heap, so the
// weakest survivor is always heap[0] and one comparison decides whether a new
// row is worth copying. k is 10 in practice, so the heap is tiny and the whole
// structure fits in cache.
//
// It is a plain slice with hand-written sift operations rather than
// container/heap, because that interface boxes every element into an `any` on
// Push and costs an allocation per surviving row.
type topK struct {
	k    int
	heap []candidate
}

func newTopK(k int) *topK {
	return &topK{k: k, heap: make([]candidate, 0, k)}
}

// offer admits a row if it beats the weakest survivor, or if the heap is not
// full yet. id and meta are the driver's buffers and are copied here, never
// retained.
//
// A row that only ties the weakest survivor does not displace it, so among
// equal scores the earliest k rows are kept. The full sort this replaced used
// sort.Slice, which is not stable, so it broke such a tie arbitrarily. That is
// the ONE observable difference between the two, and it is deliberate:
// deterministic beats arbitrary.
//
// It is not academic. Under the real embedder an exact tie means duplicate
// vectors and nothing in the recall eval moved (V-643). Under the hash
// embedder the eval's deterministic floor uses, ties are everywhere — it is
// bag-of-words, so every note sharing no word with the query scores exactly 0
// — and recall@3 on that run moved 74.1% to 81.5% purely because the zeros now
// come out in a fixed order. Neither number measures retrieval. recall@1 and
// false recall, which the eval actually asserts, are unchanged on both runs.
func (t *topK) offer(score float64, id, meta []byte) {
	if t.k == 0 {
		return
	}
	if len(t.heap) < t.k {
		t.heap = append(t.heap, candidate{score: score, id: string(id), meta: append([]byte(nil), meta...)})
		t.up(len(t.heap) - 1)
		return
	}
	if score <= t.heap[0].score {
		return
	}
	t.heap[0] = candidate{score: score, id: string(id), meta: append([]byte(nil), meta...)}
	t.down(0)
}

func (t *topK) up(i int) {
	for i > 0 {
		parent := (i - 1) / 2
		if t.heap[parent].score <= t.heap[i].score {
			return
		}
		t.heap[parent], t.heap[i] = t.heap[i], t.heap[parent]
		i = parent
	}
}

func (t *topK) down(i int) {
	for {
		l, r, small := 2*i+1, 2*i+2, i
		if l < len(t.heap) && t.heap[l].score < t.heap[small].score {
			small = l
		}
		if r < len(t.heap) && t.heap[r].score < t.heap[small].score {
			small = r
		}
		if small == i {
			return
		}
		t.heap[small], t.heap[i] = t.heap[i], t.heap[small]
		i = small
	}
}

// sorted drains the heap into descending score order — what Search returns.
func (t *topK) sorted() []candidate {
	out := t.heap
	sort.Slice(out, func(i, j int) bool { return out[i].score > out[j].score })
	return out
}

// ByPrefix returns every row whose id starts with prefix, vectors included.
//
// This is not a similarity query and deliberately does not score anything:
// listing the enrolled voices is a question about which rows exist, and asking
// it through Search would mean inventing a query vector to rank them by. The
// prefix is matched with LIKE against an escaped pattern, so a profile id
// containing % or _ cannot widen the match.
func (m *MemoryStore) ByPrefix(ctx context.Context, prefix string) ([]memory.Record, error) {
	pattern := escapeLike(prefix) + "%"
	rows, err := m.db.QueryContext(ctx,
		`SELECT id, vec, meta FROM memory_vectors WHERE id LIKE ? ESCAPE '\'`, pattern)
	if err != nil {
		return nil, fmt.Errorf("memory: by prefix %q: %w", prefix, err)
	}
	defer rows.Close()

	var out []memory.Record
	for rows.Next() {
		var id, metaJSON string
		var blob []byte
		if err := rows.Scan(&id, &blob, &metaJSON); err != nil {
			return nil, fmt.Errorf("memory: row: %w", err)
		}
		meta := map[string]string{}
		if err := json.Unmarshal([]byte(metaJSON), &meta); err != nil {
			return nil, fmt.Errorf("memory: unmarshal meta for %q: %w", id, err)
		}
		out = append(out, memory.Record{ID: id, Vec: decodeVec(blob), Meta: meta})
	}
	if err := rows.Err(); err != nil {
		return nil, fmt.Errorf("memory: rows: %w", err)
	}
	return out, nil
}

// Delete removes one vector by id. A row that is not there is not an error —
// "forget this voice" is satisfied either way.
func (m *MemoryStore) Delete(ctx context.Context, id string) error {
	if _, err := m.db.ExecContext(ctx, `DELETE FROM memory_vectors WHERE id = ?`, id); err != nil {
		return fmt.Errorf("memory: delete %q: %w", id, err)
	}
	return nil
}

// DeletePrefix removes every vector whose id starts with prefix and returns
// how many went. Same escaping as ByPrefix, so a key containing % or _ cannot
// widen the delete.
//
// It exists for the repair half of a revert (#470). Voiding a fact row left
// its vector in the index, so recall kept serving the voided fact's utterance
// and the documented repair did not repair.
func (m *MemoryStore) DeletePrefix(ctx context.Context, prefix string) (int64, error) {
	pattern := escapeLike(prefix) + "%"
	res, err := m.db.ExecContext(ctx,
		`DELETE FROM memory_vectors WHERE id LIKE ? ESCAPE '\'`, pattern)
	if err != nil {
		return 0, fmt.Errorf("memory: delete prefix %q: %w", prefix, err)
	}
	n, err := res.RowsAffected()
	if err != nil {
		return 0, fmt.Errorf("memory: delete prefix %q: rows affected: %w", prefix, err)
	}
	return n, nil
}

// memVectorRow is one memory_vectors row with its meta blob decoded — the
// shape both ReembedAll (backfill.go) and RepairFactVectors (factvectors.go)
// read the whole table as, before each decides what to do with a row on its
// own terms (one keys off meta["text"], the other off meta["type"] and the
// id's embedded key/timestamp). Query-then-scan was duplicated across the two
// before this, id-for-id.
type memVectorRow struct {
	ID   string
	Meta map[string]string
}

// queryContexter is the common surface *sql.DB and *sql.Tx share that
// allMemVectorMetas needs. ReembedAll reads inside a transaction so its
// migration is atomic; RepairFactVectors reads directly off the db handle.
type queryContexter interface {
	QueryContext(ctx context.Context, query string, args ...any) (*sql.Rows, error)
}

// allMemVectorMetas reads every memory_vectors row and decodes its meta blob.
func allMemVectorMetas(ctx context.Context, q queryContexter) ([]memVectorRow, error) {
	rows, err := q.QueryContext(ctx, `SELECT id, meta FROM memory_vectors`)
	if err != nil {
		return nil, fmt.Errorf("read memory vectors: %w", err)
	}
	defer rows.Close()
	var out []memVectorRow
	for rows.Next() {
		var id, metaJSON string
		if err := rows.Scan(&id, &metaJSON); err != nil {
			return nil, fmt.Errorf("memory vector row: %w", err)
		}
		meta := map[string]string{}
		if err := json.Unmarshal([]byte(metaJSON), &meta); err != nil {
			return nil, fmt.Errorf("meta for %q: %w", id, err)
		}
		out = append(out, memVectorRow{ID: id, Meta: meta})
	}
	if err := rows.Err(); err != nil {
		return nil, fmt.Errorf("memory vectors: %w", err)
	}
	return out, nil
}

// escapeLike neutralises the LIKE wildcards in a literal prefix.
func escapeLike(s string) string {
	r := strings.NewReplacer(`\`, `\\`, `%`, `\%`, `_`, `\_`)
	return r.Replace(s)
}

// encodeVec serializes a float32 slice as little-endian IEEE-754 bytes (4 bytes
// per element) for the BLOB column.
func encodeVec(v []float32) []byte {
	b := make([]byte, 4*len(v))
	for i, f := range v {
		binary.LittleEndian.PutUint32(b[4*i:], math.Float32bits(f))
	}
	return b
}

// decodeVec reverses encodeVec. A blob whose length isn't a multiple of 4 is
// truncated to the whole-element prefix (defensive — a well-formed row can't
// produce that).
func decodeVec(b []byte) []float32 {
	n := len(b) / 4
	v := make([]float32, n)
	for i := 0; i < n; i++ {
		v[i] = math.Float32frombits(binary.LittleEndian.Uint32(b[4*i:]))
	}
	return v
}

// dotBlob is dot against a vector still in its stored encoding, so scoring a
// row the query is about to discard does not allocate the []float32 that
// decodeVec would build. Same arithmetic, same order of operations, so it
// returns bit-identical scores to dot(a, decodeVec(b)).
//
// A blob whose length isn't a multiple of 4 is truncated to the whole-element
// prefix, matching decodeVec, and a length mismatch is 0, matching dot.
func dotBlob(a []float32, b []byte) float64 {
	n := len(b) / 4
	if len(a) != n || n == 0 {
		return 0
	}
	var sum float64
	for i := 0; i < n; i++ {
		f := math.Float32frombits(binary.LittleEndian.Uint32(b[4*i:]))
		sum += float64(a[i]) * float64(f)
	}
	return sum
}

// dot is the cosine similarity for L2-normalized vectors (mismatched lengths ⇒
// 0, matching internal/memory's cosine).
func dot(a, b []float32) float64 {
	if len(a) != len(b) || len(a) == 0 {
		return 0
	}
	var sum float64
	for i := range a {
		sum += float64(a[i]) * float64(b[i])
	}
	return sum
}