kami
135 lines
3.8 KiB
Go
135 lines
3.8 KiB
Go
package store
|
||
|
||
import (
|
||
"context"
|
||
"encoding/binary"
|
||
"fmt"
|
||
"math"
|
||
"sort"
|
||
"time"
|
||
)
|
||
|
||
// Note — a recall/preference item. No predicate reads it (facts are for that);
|
||
// query-answering ranks notes by embedding cosine. Score is set by QueryNotes.
|
||
type Note struct {
|
||
ID int64
|
||
Ts time.Time
|
||
Text string
|
||
Source string
|
||
Score float64
|
||
}
|
||
|
||
// WriteNote appends a note with its embedding (stored as a little-endian
|
||
// float32 BLOB). Source is provenance (tap:voice, etc.).
|
||
func (s *Store) WriteNote(ctx context.Context, ts time.Time, text string, embedding []float32, source string) (int64, error) {
|
||
res, err := s.db.ExecContext(ctx,
|
||
`INSERT INTO notes (ts, text, embedding, source) VALUES (?,?,?,?)`,
|
||
ts.UnixMilli(), text, floatsToBlob(embedding), source)
|
||
if err != nil {
|
||
return 0, fmt.Errorf("write note: %w", err)
|
||
}
|
||
id, err := res.LastInsertId()
|
||
if err != nil {
|
||
return 0, fmt.Errorf("last insert id: %w", err)
|
||
}
|
||
return id, nil
|
||
}
|
||
|
||
// QueryNotes returns the top-k notes by cosine similarity to embedding, highest
|
||
// first (ties broken newest-first). Fewer than k notes ⇒ returns what exists.
|
||
//
|
||
// ponytail: brute-force O(n) cosine over every note each query. Add sqlite-vec
|
||
// or an ANN index only when note count or latency actually bites — at personal
|
||
// scale (hundreds–thousands) a full scan is sub-millisecond.
|
||
func (s *Store) QueryNotes(ctx context.Context, embedding []float32, k int) ([]Note, error) {
|
||
rows, err := s.db.QueryContext(ctx, `SELECT id, ts, text, embedding, source FROM notes`)
|
||
if err != nil {
|
||
return nil, fmt.Errorf("query notes: %w", err)
|
||
}
|
||
defer rows.Close()
|
||
|
||
var out []Note
|
||
for rows.Next() {
|
||
var n Note
|
||
var tsMilli int64
|
||
var blob []byte
|
||
if err := rows.Scan(&n.ID, &tsMilli, &n.Text, &blob, &n.Source); err != nil {
|
||
return nil, err
|
||
}
|
||
n.Ts = time.UnixMilli(tsMilli).UTC()
|
||
n.Score = cosine(embedding, blobToFloats(blob))
|
||
out = append(out, n)
|
||
}
|
||
if err := rows.Err(); err != nil {
|
||
return nil, err
|
||
}
|
||
|
||
sort.Slice(out, func(i, j int) bool {
|
||
if out[i].Score != out[j].Score {
|
||
return out[i].Score > out[j].Score
|
||
}
|
||
return out[i].Ts.After(out[j].Ts) // newest breaks ties
|
||
})
|
||
if k > 0 && len(out) > k {
|
||
out = out[:k]
|
||
}
|
||
return out, nil
|
||
}
|
||
|
||
// RecentNotes returns the newest n notes, newest first — a browse view (no
|
||
// embedding math; Score stays 0). This is the read surface for /dash: notes
|
||
// captured by voice are otherwise only reachable through semantic query.
|
||
func (s *Store) RecentNotes(ctx context.Context, n int) ([]Note, error) {
|
||
rows, err := s.db.QueryContext(ctx,
|
||
`SELECT id, ts, text, source FROM notes ORDER BY ts DESC LIMIT ?`, n)
|
||
if err != nil {
|
||
return nil, fmt.Errorf("recent notes: %w", err)
|
||
}
|
||
defer rows.Close()
|
||
var out []Note
|
||
for rows.Next() {
|
||
var nt Note
|
||
var tsMilli int64
|
||
if err := rows.Scan(&nt.ID, &tsMilli, &nt.Text, &nt.Source); err != nil {
|
||
return nil, err
|
||
}
|
||
nt.Ts = time.UnixMilli(tsMilli).UTC()
|
||
out = append(out, nt)
|
||
}
|
||
return out, rows.Err()
|
||
}
|
||
|
||
// cosine similarity. Embedder vectors are L2-normalized, so this is just the
|
||
// dot product — but normalize defensively in case a caller passes a raw vector.
|
||
func cosine(a, b []float32) float64 {
|
||
if len(a) != len(b) || len(a) == 0 {
|
||
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))
|
||
}
|
||
|
||
func floatsToBlob(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
|
||
}
|
||
|
||
func blobToFloats(b []byte) []float32 {
|
||
v := make([]float32, len(b)/4)
|
||
for i := range v {
|
||
v[i] = math.Float32frombits(binary.LittleEndian.Uint32(b[4*i:]))
|
||
}
|
||
return v
|
||
}
|