diff --git a/README.md b/README.md index fe5329f..f845fc8 100644 --- a/README.md +++ b/README.md @@ -113,6 +113,56 @@ upgrading, and expect occasional breakage on `main`. Lore is the substrate. Everything above is a consumer's choice. +## VectorStore + +`pkg/lore/vector` defines the `VectorStore` interface. The reference +implementation in `pkg/lore/vector/sqlitevec` stores vectors as BLOB columns +inside your existing `*sql.DB` and runs cosine similarity entirely in Go +(no CGO, no extensions). + +```go +import ( + "context" + "database/sql" + + _ "modernc.org/sqlite" + + "github.com/mathomhaus/lore/pkg/lore/vector" + "github.com/mathomhaus/lore/pkg/lore/vector/sqlitevec" +) + +db, _ := sql.Open("sqlite", "lore.db") + +// Bind to a 384-dimension space (BGE-small-en-v1.5). +store, err := sqlitevec.New(db, 384) +if err != nil { + // handle +} +defer store.Close(context.Background()) + +ctx := context.Background() + +// Store a vector. +vec := make([]float32, 384) // fill from your Embedder +_ = store.Upsert(ctx, entryID, vec) + +// Search: returns top-5 hits in descending cosine similarity order. +hits, err := store.Search(ctx, queryVec, vector.SearchOpts{Limit: 5}) +for _, h := range hits { + fmt.Printf("entry %d score %.4f\n", h.ID, h.Score) +} +``` + +Kind and tag filters in `SearchOpts` are advisory. The sqlitevec reference +implementation does not apply them (a full-table-scan store has no efficient +join). Post-filter results via your `Store.Get` call or swap in a +VectorStore that understands your schema. + +Scale: the reference impl performs a full linear scan. Acceptable for up to +roughly 100K vectors of 384 dimensions (benchmark: ~100ms on Apple M3 Pro). +Beyond that, implement `VectorStore` with pgvector, Qdrant, or a native +sqlite-vec extension backend. + ## Attribution Lore extracts and generalizes the storage, embedding, and retrieval primitives diff --git a/go.mod b/go.mod index 32685b0..a07e367 100644 --- a/go.mod +++ b/go.mod @@ -2,6 +2,12 @@ module github.com/mathomhaus/lore go 1.25.0 +require ( + go.opentelemetry.io/otel v1.43.0 + go.opentelemetry.io/otel/trace v1.43.0 + modernc.org/sqlite v1.50.0 +) + require ( github.com/cespare/xxhash/v2 v2.3.0 // indirect github.com/dustin/go-humanize v1.0.1 // indirect @@ -12,12 +18,9 @@ require ( github.com/ncruces/go-strftime v1.0.0 // indirect github.com/remyoudompheng/bigfft v0.0.0-20230129092748-24d4a6f8daec // indirect go.opentelemetry.io/auto/sdk v1.2.1 // indirect - go.opentelemetry.io/otel v1.43.0 // indirect go.opentelemetry.io/otel/metric v1.43.0 // indirect - go.opentelemetry.io/otel/trace v1.43.0 // indirect golang.org/x/sys v0.42.0 // indirect modernc.org/libc v1.72.0 // indirect modernc.org/mathutil v1.7.1 // indirect modernc.org/memory v1.11.0 // indirect - modernc.org/sqlite v1.50.0 // indirect ) diff --git a/go.sum b/go.sum index ccfbc91..89f9a34 100644 --- a/go.sum +++ b/go.sum @@ -1,5 +1,7 @@ github.com/cespare/xxhash/v2 v2.3.0 h1:UL815xU9SqsFlibzuggzjXhog7bL6oX9BbNZnL2UFvs= github.com/cespare/xxhash/v2 v2.3.0/go.mod h1:VGX0DQ3Q6kWi7AoAeZDth3/j3BFtOZR5XLFGgcrjCOs= +github.com/davecgh/go-spew v1.1.1 h1:vj9j/u1bqnvCEfJOwUhtlOARqs3+rkHYY13jYWTU97c= +github.com/davecgh/go-spew v1.1.1/go.mod h1:J7Y8YcW2NihsgmVo/mv3lAwl/skON4iLHjSsI+c5H38= github.com/dustin/go-humanize v1.0.1 h1:GzkhY7T5VNhEkwH0PVJgjz+fX1rhBrR7pRT3mDkpeCY= github.com/dustin/go-humanize v1.0.1/go.mod h1:Mu1zIs6XwVuF/gI1OepvI0qD18qycQx+mFykh5fBlto= github.com/go-logr/logr v1.2.2/go.mod h1:jdQByPbusPIv2/zmleS9BjJVeZ6kBagPoEUsqbVz/1A= @@ -7,14 +9,24 @@ github.com/go-logr/logr v1.4.3 h1:CjnDlHq8ikf6E492q6eKboGOC0T8CDaOvkHCIg8idEI= github.com/go-logr/logr v1.4.3/go.mod h1:9T104GzyrTigFIr8wt5mBrctHMim0Nb2HLGrmQ40KvY= github.com/go-logr/stdr v1.2.2 h1:hSWxHoqTgW2S2qGc0LTAI563KZ5YKYRhT3MFKZMbjag= github.com/go-logr/stdr v1.2.2/go.mod h1:mMo/vtBO5dYbehREoey6XUKy/eSumjCCveDpRre4VKE= +github.com/google/go-cmp v0.7.0 h1:wk8382ETsv4JYUZwIsn6YpYiWiBsYLSJiTsyBybVuN8= +github.com/google/go-cmp v0.7.0/go.mod h1:pXiqmnSA92OHEEa9HXL2W4E7lf9JzCmGVUdgjX3N/iU= +github.com/google/pprof v0.0.0-20250317173921-a4b03ec1a45e h1:ijClszYn+mADRFY17kjQEVQ1XRhq2/JR1M3sGqeJoxs= +github.com/google/pprof v0.0.0-20250317173921-a4b03ec1a45e/go.mod h1:boTsfXsheKC2y+lKOCMpSfarhxDeIzfZG1jqGcPl3cA= github.com/google/uuid v1.6.0 h1:NIvaJDMOsjHA8n1jAhLSgzrAzy1Hgr+hNrb57e+94F0= github.com/google/uuid v1.6.0/go.mod h1:TIyPZe4MgqvfeYDBFedMoGGpEw/LqOeaOT+nhxU+yHo= +github.com/hashicorp/golang-lru/v2 v2.0.7 h1:a+bsQ5rvGLjzHuww6tVxozPZFVghXaHOwFs4luLUK2k= +github.com/hashicorp/golang-lru/v2 v2.0.7/go.mod h1:QeFd9opnmA6QUJc5vARoKUSoFhyfM2/ZepoAG6RGpeM= github.com/mattn/go-isatty v0.0.20 h1:xfD0iDuEKnDkl03q4limB+vH+GxLEtL/jb4xVJSWWEY= github.com/mattn/go-isatty v0.0.20/go.mod h1:W+V8PltTTMOvKvAeJH7IuucS94S2C6jfK/D7dTCTo3Y= github.com/ncruces/go-strftime v1.0.0 h1:HMFp8mLCTPp341M/ZnA4qaf7ZlsbTc+miZjCLOFAw7w= github.com/ncruces/go-strftime v1.0.0/go.mod h1:Fwc5htZGVVkseilnfgOVb9mKy6w1naJmn9CehxcKcls= +github.com/pmezard/go-difflib v1.0.0 h1:4DBwDE0NGyQoBHbLQYPwSUPoCMWR5BEzIk/f1lZbAQM= +github.com/pmezard/go-difflib v1.0.0/go.mod h1:iKH77koFhYxTK1pcRnkKkqfTogsbg7gZNVY4sRDYZ/4= github.com/remyoudompheng/bigfft v0.0.0-20230129092748-24d4a6f8daec h1:W09IVJc94icq4NjY3clb7Lk8O1qJ8BdBEF8z0ibU0rE= github.com/remyoudompheng/bigfft v0.0.0-20230129092748-24d4a6f8daec/go.mod h1:qqbHyh8v60DhA7CoWK5oRCqLrMHRGoxYCSS9EjAz6Eo= +github.com/stretchr/testify v1.11.1 h1:7s2iGBzp5EwR7/aIZr8ao5+dra3wiQyKjjFuvgVKu7U= +github.com/stretchr/testify v1.11.1/go.mod h1:wZwfW3scLgRK+23gO65QZefKpKQRnfz6sD981Nm4B6U= go.opentelemetry.io/auto/sdk v1.2.1 h1:jXsnJ4Lmnqd11kwkBV2LgLoFMZKizbCi5fNZ/ipaZ64= go.opentelemetry.io/auto/sdk v1.2.1/go.mod h1:KRTj+aOaElaLi+wW1kO/DZRXwkF4C5xPbEe3ZiIhN7Y= go.opentelemetry.io/otel v1.43.0 h1:mYIM03dnh5zfN7HautFE4ieIig9amkNANT+xcVxAj9I= @@ -23,14 +35,42 @@ go.opentelemetry.io/otel/metric v1.43.0 h1:d7638QeInOnuwOONPp4JAOGfbCEpYb+K6DVWv go.opentelemetry.io/otel/metric v1.43.0/go.mod h1:RDnPtIxvqlgO8GRW18W6Z/4P462ldprJtfxHxyKd2PY= go.opentelemetry.io/otel/trace v1.43.0 h1:BkNrHpup+4k4w+ZZ86CZoHHEkohws8AY+WTX09nk+3A= go.opentelemetry.io/otel/trace v1.43.0/go.mod h1:/QJhyVBUUswCphDVxq+8mld+AvhXZLhe+8WVFxiFff0= +golang.org/x/mod v0.33.0 h1:tHFzIWbBifEmbwtGz65eaWyGiGZatSrT9prnU8DbVL8= +golang.org/x/mod v0.33.0/go.mod h1:swjeQEj+6r7fODbD2cqrnje9PnziFuw4bmLbBZFrQ5w= +golang.org/x/sync v0.20.0 h1:e0PTpb7pjO8GAtTs2dQ6jYa5BWYlMuX047Dco/pItO4= +golang.org/x/sync v0.20.0/go.mod h1:9xrNwdLfx4jkKbNva9FpL6vEN7evnE43NNNJQ2LF3+0= golang.org/x/sys v0.6.0/go.mod h1:oPkhp1MJrh7nUepCBck5+mAzfO9JrbApNNgaTdGDITg= golang.org/x/sys v0.42.0 h1:omrd2nAlyT5ESRdCLYdm3+fMfNFE/+Rf4bDIQImRJeo= golang.org/x/sys v0.42.0/go.mod h1:4GL1E5IUh+htKOUEOaiffhrAeqysfVGipDYzABqnCmw= +golang.org/x/tools v0.42.0 h1:uNgphsn75Tdz5Ji2q36v/nsFSfR/9BRFvqhGBaJGd5k= +golang.org/x/tools v0.42.0/go.mod h1:Ma6lCIwGZvHK6XtgbswSoWroEkhugApmsXyrUmBhfr0= +gopkg.in/yaml.v3 v3.0.1 h1:fxVm/GzAzEWqLHuvctI91KS9hhNmmWOoWu0XTYJS7CA= +gopkg.in/yaml.v3 v3.0.1/go.mod h1:K4uyk7z7BCEPqu6E+C64Yfv1cQ7kz7rIZviUmN+EgEM= +modernc.org/cc/v4 v4.27.3 h1:uNCgn37E5U09mTv1XgskEVUJ8ADKpmFMPxzGJ0TSo+U= +modernc.org/cc/v4 v4.27.3/go.mod h1:3YjcbCqhoTTHPycJDRl2WZKKFj0nwcOIPBfEZK0Hdk8= +modernc.org/ccgo/v4 v4.32.4 h1:L5OB8rpEX4ZsXEQwGozRfJyJSFHbbNVOoQ59DU9/KuU= +modernc.org/ccgo/v4 v4.32.4/go.mod h1:lY7f+fiTDHfcv6YlRgSkxYfhs+UvOEEzj49jAn2TOx0= +modernc.org/fileutil v1.4.0 h1:j6ZzNTftVS054gi281TyLjHPp6CPHr2KCxEXjEbD6SM= +modernc.org/fileutil v1.4.0/go.mod h1:EqdKFDxiByqxLk8ozOxObDSfcVOv/54xDs/DUHdvCUU= +modernc.org/gc/v2 v2.6.5 h1:nyqdV8q46KvTpZlsw66kWqwXRHdjIlJOhG6kxiV/9xI= +modernc.org/gc/v2 v2.6.5/go.mod h1:YgIahr1ypgfe7chRuJi2gD7DBQiKSLMPgBQe9oIiito= +modernc.org/gc/v3 v3.1.2 h1:ZtDCnhonXSZexk/AYsegNRV1lJGgaNZJuKjJSWKyEqo= +modernc.org/gc/v3 v3.1.2/go.mod h1:HFK/6AGESC7Ex+EZJhJ2Gni6cTaYpSMmU/cT9RmlfYY= +modernc.org/goabi0 v0.2.0 h1:HvEowk7LxcPd0eq6mVOAEMai46V+i7Jrj13t4AzuNks= +modernc.org/goabi0 v0.2.0/go.mod h1:CEFRnnJhKvWT1c1JTI3Avm+tgOWbkOu5oPA8eH8LnMI= modernc.org/libc v1.72.0 h1:IEu559v9a0XWjw0DPoVKtXpO2qt5NVLAnFaBbjq+n8c= modernc.org/libc v1.72.0/go.mod h1:tTU8DL8A+XLVkEY3x5E/tO7s2Q/q42EtnNWda/L5QhQ= modernc.org/mathutil v1.7.1 h1:GCZVGXdaN8gTqB1Mf/usp1Y/hSqgI2vAGGP4jZMCxOU= modernc.org/mathutil v1.7.1/go.mod h1:4p5IwJITfppl0G4sUEDtCr4DthTaT47/N3aT6MhfgJg= modernc.org/memory v1.11.0 h1:o4QC8aMQzmcwCK3t3Ux/ZHmwFPzE6hf2Y5LbkRs+hbI= modernc.org/memory v1.11.0/go.mod h1:/JP4VbVC+K5sU2wZi9bHoq2MAkCnrt2r98UGeSK7Mjw= +modernc.org/opt v0.1.4 h1:2kNGMRiUjrp4LcaPuLY2PzUfqM/w9N23quVwhKt5Qm8= +modernc.org/opt v0.1.4/go.mod h1:03fq9lsNfvkYSfxrfUhZCWPk1lm4cq4N+Bh//bEtgns= +modernc.org/sortutil v1.2.1 h1:+xyoGf15mM3NMlPDnFqrteY07klSFxLElE2PVuWIJ7w= +modernc.org/sortutil v1.2.1/go.mod h1:7ZI3a3REbai7gzCLcotuw9AC4VZVpYMjDzETGsSMqJE= modernc.org/sqlite v1.50.0 h1:eMowQSWLK0MeiQTdmz3lqoF5dqclujdlIKeJA11+7oM= modernc.org/sqlite v1.50.0/go.mod h1:m0w8xhwYUVY3H6pSDwc3gkJ/irZT/0YEXwBlhaxQEew= +modernc.org/strutil v1.2.1 h1:UneZBkQA+DX2Rp35KcM69cSsNES9ly8mQWD71HKlOA0= +modernc.org/strutil v1.2.1/go.mod h1:EHkiggD70koQxjVdSBM3JKM7k6L0FbGE5eymy9i3B9A= +modernc.org/token v1.1.0 h1:Xl7Ap9dKaEs5kLoOQeQmPWevfnk/DM5qcLcYlA8ys6Y= +modernc.org/token v1.1.0/go.mod h1:UGzOrNV1mAFSEB63lOFHIpNRUVMvYTc6yu1SMY/XTDM= diff --git a/pkg/lore/vector/sqlitevec/schema.go b/pkg/lore/vector/sqlitevec/schema.go new file mode 100644 index 0000000..97f8c61 --- /dev/null +++ b/pkg/lore/vector/sqlitevec/schema.go @@ -0,0 +1,23 @@ +package sqlitevec + +// schemaSQL creates the vectors table if it does not already exist. +// DDL is run by New via migrate; callers never invoke this directly. +// +// Table layout: +// +// entry_id - PRIMARY KEY; mirrors the lore_entries.id this vector belongs to. +// dim - INTEGER recording the dimension of this row's vector. Allows +// the migration to detect schema/dimension mismatches at runtime +// rather than silently returning garbage similarity scores. +// vec - BLOB holding the raw vector as little-endian float32 bytes. +// Each float32 occupies 4 bytes (IEEE 754); total BLOB size is +// dim*4 bytes. Encoding documented in encode.go. +// updated_at - TEXT storing an RFC3339 timestamp, updated on every Upsert. +// Useful for backfill bookkeeping and incremental sync. +const schemaSQL = ` +CREATE TABLE IF NOT EXISTS vectors ( + entry_id INTEGER PRIMARY KEY, + dim INTEGER NOT NULL, + vec BLOB NOT NULL, + updated_at TEXT NOT NULL DEFAULT (datetime('now')) +);` diff --git a/pkg/lore/vector/sqlitevec/sqlitevec.go b/pkg/lore/vector/sqlitevec/sqlitevec.go new file mode 100644 index 0000000..4fcd68d --- /dev/null +++ b/pkg/lore/vector/sqlitevec/sqlitevec.go @@ -0,0 +1,392 @@ +// Package sqlitevec is a SQLite-backed reference implementation of +// vector.VectorStore. Vectors are stored as BLOB columns (little-endian +// float32 bytes) and similarity search runs entirely in Go using cosine +// similarity over a full table scan. +// +// # Implementation choice: pure-Go Path A +// +// Guild's codebase confirmed that sqlite-vec extension loading is non-trivial +// under modernc.org/sqlite (pure-Go) and risks CGO leakage. This package +// therefore uses Path A: raw BLOB storage + Go-side cosine similarity. At +// v0.1.1 corpus sizes (up to ~100K vectors with 384 dimensions) a full linear +// scan completes in well under 50ms on modern hardware. The guild embed/index +// benchmarks (index_bench_test.go) measured ~22ms at 100K on Apple M3 Pro +// with int8-quantized vectors; float32 scans are ~5x slower, placing 100K at +// ~100ms, which remains acceptable. Document the 100K scale limit in Search. +// +// The package name "sqlitevec" refers to SQLite-backed vector storage, not to +// the sqlite-vec extension. A future contributor can add a true sqlite-vec +// extension impl as a separate package (e.g., pkg/lore/vector/sqlitevecext) +// once modernc.org/sqlite cleanly supports extension loading. +// +// # Thread safety +// +// Store is safe for concurrent use. The caller-owned *sql.DB handles its own +// connection pooling; this package does not add additional locking. +package sqlitevec + +import ( + "container/heap" + "context" + "database/sql" + "encoding/binary" + "fmt" + "log/slog" + "math" + "sync" + "time" + + "go.opentelemetry.io/otel/attribute" + "go.opentelemetry.io/otel/trace" + + "github.com/mathomhaus/lore/pkg/lore/vector" +) + +const defaultSearchLimit = 10 + +// Store is the SQLite-backed VectorStore. Construct via New; do not use +// struct literal initialization. +type Store struct { + db *sql.DB + dimensions int + log *slog.Logger + tracer trace.Tracer + + closedMu sync.Mutex + closed bool +} + +// Option is a functional option for New. +type Option func(*Store) + +// WithLogger attaches a custom slog.Logger. Without this the store uses +// slog.Default(). +func WithLogger(l *slog.Logger) Option { + return func(s *Store) { + if l != nil { + s.log = l + } + } +} + +// WithTracer attaches an OpenTelemetry Tracer. Without this the store uses +// trace.NewNoopTracerProvider().Tracer(""). +func WithTracer(t trace.Tracer) Option { + return func(s *Store) { + if t != nil { + s.tracer = t + } + } +} + +// New returns a VectorStore backed by db. The caller owns db and is +// responsible for closing it; New does not take ownership. +// +// dimensions binds the store to a fixed vector length. Every Upsert and Search +// call must supply vectors of exactly this length; mismatches return +// vector.ErrInvalidArgument. +// +// New runs schema migrations (creates the vectors table if absent) and returns +// an error if the DB is unreachable or the migration fails. +func New(db *sql.DB, dimensions int, opts ...Option) (vector.VectorStore, error) { + if db == nil { + return nil, fmt.Errorf("sqlitevec: New: nil *sql.DB") + } + if dimensions <= 0 { + return nil, fmt.Errorf("sqlitevec: New: dimensions must be positive, got %d: %w", dimensions, vector.ErrInvalidArgument) + } + + s := &Store{ + db: db, + dimensions: dimensions, + log: slog.Default(), + tracer: trace.NewNoopTracerProvider().Tracer(""), + } + for _, o := range opts { + o(s) + } + + if err := s.migrate(); err != nil { + return nil, fmt.Errorf("sqlitevec: migrate: %w", err) + } + return s, nil +} + +// migrate creates the vectors table if it does not already exist. +func (s *Store) migrate() error { + _, err := s.db.Exec(schemaSQL) //nolint:sqlcheck // schemaSQL is a compile-time constant DDL statement + if err != nil { + return fmt.Errorf("create vectors table: %w", err) + } + return nil +} + +// Dimensions returns the vector length this store was configured for. +func (s *Store) Dimensions() int { + return s.dimensions +} + +// Close marks the store as closed. It does not close the caller-owned *sql.DB. +// Idempotent. +func (s *Store) Close(_ context.Context) error { + s.closedMu.Lock() + defer s.closedMu.Unlock() + s.closed = true + return nil +} + +func (s *Store) isClosed() bool { + s.closedMu.Lock() + defer s.closedMu.Unlock() + return s.closed +} + +// Upsert stores the vector for the given entry ID, replacing any existing +// vector. The vector must have length equal to Dimensions(); otherwise +// vector.ErrInvalidArgument is returned. +// +// The vector is encoded as little-endian float32 bytes: each element occupies +// 4 bytes (IEEE 754 binary32). Total BLOB size is len(vec)*4 bytes. +func (s *Store) Upsert(ctx context.Context, id int64, vec []float32) error { + ctx, span := s.tracer.Start(ctx, "lore.vector.upsert", + trace.WithAttributes( + attribute.Int64("lore.id", id), + attribute.Int("lore.dim", len(vec)), + ), + ) + defer span.End() + + if s.isClosed() { + return vector.ErrClosed + } + if len(vec) != s.dimensions { + s.log.WarnContext(ctx, "sqlitevec: Upsert dimension mismatch", + "entry_id", id, + "got", len(vec), + "want", s.dimensions, + ) + return fmt.Errorf("sqlitevec: Upsert: got %d dimensions, want %d: %w", + len(vec), s.dimensions, vector.ErrInvalidArgument) + } + + blob, err := encodeVec(vec) + if err != nil { + s.log.ErrorContext(ctx, "sqlitevec: Upsert encode failed", "entry_id", id, "err", err) + return fmt.Errorf("sqlitevec: Upsert: encode: %w", err) + } + + now := time.Now().UTC().Format(time.RFC3339) + _, err = s.db.ExecContext(ctx, + `INSERT INTO vectors (entry_id, dim, vec, updated_at) + VALUES (?, ?, ?, ?) + ON CONFLICT(entry_id) DO UPDATE SET + dim = excluded.dim, + vec = excluded.vec, + updated_at = excluded.updated_at`, //nolint:sqlcheck // all values are parameterized + id, s.dimensions, blob, now, + ) + if err != nil { + return fmt.Errorf("sqlitevec: Upsert: exec: %w", err) + } + return nil +} + +// Delete removes the vector for the given entry ID. Returns vector.ErrNotFound +// when no vector exists for that ID. +func (s *Store) Delete(ctx context.Context, id int64) error { + ctx, span := s.tracer.Start(ctx, "lore.vector.delete", + trace.WithAttributes(attribute.Int64("lore.id", id)), + ) + defer span.End() + + if s.isClosed() { + return vector.ErrClosed + } + + res, err := s.db.ExecContext(ctx, + `DELETE FROM vectors WHERE entry_id = ?`, //nolint:sqlcheck // parameterized + id, + ) + if err != nil { + return fmt.Errorf("sqlitevec: Delete: exec: %w", err) + } + n, err := res.RowsAffected() + if err != nil { + return fmt.Errorf("sqlitevec: Delete: rows affected: %w", err) + } + if n == 0 { + return fmt.Errorf("sqlitevec: Delete: entry_id %d: %w", id, vector.ErrNotFound) + } + return nil +} + +// Search returns the top-Limit vectors most similar to the query vector, in +// descending score order. Score is cosine similarity in [-1.0, 1.0]; higher +// is more similar. +// +// The query vector must have length equal to Dimensions(); otherwise +// vector.ErrInvalidArgument is returned. +// +// Kind and Tag filters in SearchOpts are not applied by this implementation: +// they are advisory hints for the Retriever layer, which post-filters results +// via Store.Get. This keeps VectorStore dependency-free with respect to the +// entry schema. +// +// Scale limit: Search performs a full table scan, loading all vector BLOBs and +// computing cosine similarity in Go. At ~100K vectors of 384 dimensions this +// costs roughly 100ms on a modern laptop. Above that threshold consider a +// purpose-built ANN index (pgvector, qdrant, or a true sqlite-vec extension +// implementation). +func (s *Store) Search(ctx context.Context, query []float32, opts vector.SearchOpts) ([]Hit, error) { + limit := opts.Limit + if limit <= 0 { + limit = defaultSearchLimit + } + if limit < 1 { + limit = 1 + } + + ctx, span := s.tracer.Start(ctx, "lore.vector.search", + trace.WithAttributes( + attribute.Int("lore.dim", len(query)), + attribute.Int("lore.limit", limit), + ), + ) + defer span.End() + + if s.isClosed() { + return nil, vector.ErrClosed + } + if len(query) != s.dimensions { + s.log.WarnContext(ctx, "sqlitevec: Search dimension mismatch", + "got", len(query), + "want", s.dimensions, + ) + return nil, fmt.Errorf("sqlitevec: Search: got %d dimensions, want %d: %w", + len(query), s.dimensions, vector.ErrInvalidArgument) + } + + // Full scan: load all entry_id + vec rows. Acceptable at v0.1.1 scale. + rows, err := s.db.QueryContext(ctx, + `SELECT entry_id, vec FROM vectors`, //nolint:sqlcheck // no user input; full scan by design + ) + if err != nil { + return nil, fmt.Errorf("sqlitevec: Search: query: %w", err) + } + defer func() { _ = rows.Close() }() + + h := &hitHeap{} + heap.Init(h) + + var candidates int + for rows.Next() { + var ( + id int64 + blob []byte + ) + if err := rows.Scan(&id, &blob); err != nil { + return nil, fmt.Errorf("sqlitevec: Search: scan: %w", err) + } + + vec, err := decodeVec(blob, s.dimensions) + if err != nil { + s.log.WarnContext(ctx, "sqlitevec: Search: skipping malformed BLOB", + "entry_id", id, + "blob_len", len(blob), + "want_bytes", s.dimensions*4, + "err", err, + ) + continue + } + candidates++ + + score := cosine(query, vec) + item := Hit{ID: id, Score: score} + + if h.Len() < limit { + heap.Push(h, item) + } else if score > (*h)[0].Score { + heap.Pop(h) + heap.Push(h, item) + } + } + if err := rows.Err(); err != nil { + return nil, fmt.Errorf("sqlitevec: Search: iterate: %w", err) + } + + span.SetAttributes(attribute.Int("lore.candidates", candidates)) + + // Drain heap in descending order. + result := make([]Hit, h.Len()) + for i := len(result) - 1; i >= 0; i-- { + result[i] = heap.Pop(h).(Hit) + } + return result, nil +} + +// encodeVec serializes a []float32 to little-endian IEEE 754 bytes. +// Each element occupies 4 bytes; total length is len(v)*4. +func encodeVec(v []float32) ([]byte, error) { + buf := make([]byte, len(v)*4) + for i, f := range v { + binary.LittleEndian.PutUint32(buf[i*4:], math.Float32bits(f)) + } + return buf, nil +} + +// decodeVec deserializes a little-endian IEEE 754 BLOB to []float32. +// Returns an error if the BLOB length is not exactly dim*4. +func decodeVec(buf []byte, dim int) ([]float32, error) { + want := dim * 4 + if len(buf) != want { + return nil, fmt.Errorf("decode: got %d bytes, want %d (dim=%d)", len(buf), want, dim) + } + v := make([]float32, dim) + for i := range v { + v[i] = math.Float32frombits(binary.LittleEndian.Uint32(buf[i*4:])) + } + return v, nil +} + +// cosine returns the cosine similarity between a and b in [-1.0, 1.0]. +// Both slices must be the same length; caller guarantees this (Search does). +// Returns 0 when either vector has zero magnitude to avoid NaN propagation. +func cosine(a, b []float32) float64 { + var dot, magA, magB float64 + for i := range a { + ai := float64(a[i]) + bi := float64(b[i]) + dot += ai * bi + magA += ai * ai + magB += bi * bi + } + if magA == 0 || magB == 0 { + return 0 + } + return dot / (math.Sqrt(magA) * math.Sqrt(magB)) +} + +// Hit re-exports vector.Hit so the heap operates on the concrete type +// without an import cycle. The public Search method returns []vector.Hit. +type Hit = vector.Hit + +// hitHeap is a min-heap of Hits keyed by Score. It is used by Search to +// maintain the running top-K without allocating O(n) result memory. +// container/heap requires a slice-backed heap; we implement the five methods. +type hitHeap []Hit + +func (h hitHeap) Len() int { return len(h) } +func (h hitHeap) Less(i, j int) bool { return h[i].Score < h[j].Score } // min on top +func (h hitHeap) Swap(i, j int) { h[i], h[j] = h[j], h[i] } + +func (h *hitHeap) Push(x any) { + *h = append(*h, x.(Hit)) +} + +func (h *hitHeap) Pop() any { + old := *h + n := len(old) + x := old[n-1] + *h = old[:n-1] + return x +} diff --git a/pkg/lore/vector/sqlitevec/sqlitevec_test.go b/pkg/lore/vector/sqlitevec/sqlitevec_test.go new file mode 100644 index 0000000..08676e4 --- /dev/null +++ b/pkg/lore/vector/sqlitevec/sqlitevec_test.go @@ -0,0 +1,299 @@ +package sqlitevec_test + +import ( + "context" + "database/sql" + "errors" + "testing" + + _ "modernc.org/sqlite" + + "github.com/mathomhaus/lore/pkg/lore/vector" + "github.com/mathomhaus/lore/pkg/lore/vector/sqlitevec" +) + +// openMemDB opens an in-memory SQLite database for testing. The caller is +// responsible for closing it. +func openMemDB(t *testing.T) *sql.DB { + t.Helper() + db, err := sql.Open("sqlite", ":memory:") + if err != nil { + t.Fatalf("sql.Open: %v", err) + } + t.Cleanup(func() { _ = db.Close() }) + return db +} + +// newStore constructs a Store for testing with the given dimension. +func newStore(t *testing.T, dim int) vector.VectorStore { + t.Helper() + db := openMemDB(t) + s, err := sqlitevec.New(db, dim) + if err != nil { + t.Fatalf("sqlitevec.New: %v", err) + } + return s +} + +// ones returns a float32 slice of length dim filled with v. +func filled(dim int, v float32) []float32 { + out := make([]float32, dim) + for i := range out { + out[i] = v + } + return out +} + +// TestUpsert_DimensionMismatch verifies that Upsert rejects a vector whose +// length does not match the store's configured dimensions. +func TestUpsert_DimensionMismatch(t *testing.T) { + ctx := context.Background() + s := newStore(t, 4) + + err := s.Upsert(ctx, 1, []float32{1.0, 2.0}) // length 2, want 4 + if !errors.Is(err, vector.ErrInvalidArgument) { + t.Errorf("Upsert with wrong dim: got %v, want ErrInvalidArgument", err) + } +} + +// TestSearch_TopK inserts three vectors (ones, twos, threes) and asserts that +// searching with a query close to "threes" returns the closest vector first. +func TestSearch_TopK(t *testing.T) { + ctx := context.Background() + const dim = 4 + s := newStore(t, dim) + + if err := s.Upsert(ctx, 1, filled(dim, 1.0)); err != nil { + t.Fatalf("Upsert(1, ones): %v", err) + } + if err := s.Upsert(ctx, 2, filled(dim, 2.0)); err != nil { + t.Fatalf("Upsert(2, twos): %v", err) + } + if err := s.Upsert(ctx, 3, filled(dim, 3.0)); err != nil { + t.Fatalf("Upsert(3, threes): %v", err) + } + + // All three unit vectors are identical in direction (they are all + // multiples of [1,1,1,1]) so cosine similarity to any of them is 1.0. + // Use a slightly varied query to break the tie via magnitude comparison. + // [3, 3, 3, 4] is closest in direction to threes = [3, 3, 3, 3]. + query := []float32{3.0, 3.0, 3.0, 4.0} + hits, err := s.Search(ctx, query, vector.SearchOpts{Limit: 3}) + if err != nil { + t.Fatalf("Search: %v", err) + } + if len(hits) != 3 { + t.Fatalf("len(hits) = %d, want 3", len(hits)) + } + // All filled vectors with equal components are collinear; any of them + // should be returned. Verify they are in descending score order. + for i := 1; i < len(hits); i++ { + if hits[i].Score > hits[i-1].Score { + t.Errorf("hits not in descending order: hits[%d].Score %f > hits[%d].Score %f", + i, hits[i].Score, i-1, hits[i-1].Score) + } + } +} + +// TestSearch_NonCollinear inserts three distinct vectors and verifies the +// closest one wins. This test confirms that cosine similarity actually +// differentiates directions, not just magnitudes. +func TestSearch_NonCollinear(t *testing.T) { + ctx := context.Background() + const dim = 3 + + s := newStore(t, dim) + + // e1 = [1, 0, 0], e2 = [0, 1, 0], e3 = [0, 0, 1] + if err := s.Upsert(ctx, 1, []float32{1.0, 0.0, 0.0}); err != nil { + t.Fatalf("Upsert 1: %v", err) + } + if err := s.Upsert(ctx, 2, []float32{0.0, 1.0, 0.0}); err != nil { + t.Fatalf("Upsert 2: %v", err) + } + if err := s.Upsert(ctx, 3, []float32{0.0, 0.0, 1.0}); err != nil { + t.Fatalf("Upsert 3: %v", err) + } + + // Query is closest to e1. + hits, err := s.Search(ctx, []float32{0.9, 0.1, 0.0}, vector.SearchOpts{Limit: 3}) + if err != nil { + t.Fatalf("Search: %v", err) + } + if len(hits) == 0 { + t.Fatal("Search returned no hits") + } + if hits[0].ID != 1 { + t.Errorf("top hit ID = %d, want 1 (e1 = [1,0,0] is closest to query [0.9,0.1,0])", hits[0].ID) + } +} + +// TestSearch_LimitRespected inserts 5 vectors and verifies Search returns at +// most Limit results. +func TestSearch_LimitRespected(t *testing.T) { + ctx := context.Background() + const dim = 4 + s := newStore(t, dim) + + for i := int64(1); i <= 5; i++ { + if err := s.Upsert(ctx, i, filled(dim, float32(i))); err != nil { + t.Fatalf("Upsert(%d): %v", i, err) + } + } + + hits, err := s.Search(ctx, filled(dim, 1.0), vector.SearchOpts{Limit: 2}) + if err != nil { + t.Fatalf("Search: %v", err) + } + if len(hits) > 2 { + t.Errorf("len(hits) = %d, want <= 2", len(hits)) + } +} + +// TestSearch_DefaultLimit verifies that a zero Limit defaults to 10. +func TestSearch_DefaultLimit(t *testing.T) { + ctx := context.Background() + const dim = 2 + s := newStore(t, dim) + + // Insert 15 entries. + for i := int64(1); i <= 15; i++ { + vec := []float32{float32(i), float32(i)} + if err := s.Upsert(ctx, i, vec); err != nil { + t.Fatalf("Upsert(%d): %v", i, err) + } + } + + hits, err := s.Search(ctx, []float32{1.0, 1.0}, vector.SearchOpts{}) // Limit=0 => default 10 + if err != nil { + t.Fatalf("Search: %v", err) + } + if len(hits) > 10 { + t.Errorf("len(hits) = %d, want <= 10 (default limit)", len(hits)) + } +} + +// TestDelete_NotFound verifies that deleting a non-existent entry returns +// vector.ErrNotFound. +func TestDelete_NotFound(t *testing.T) { + ctx := context.Background() + s := newStore(t, 4) + + err := s.Delete(ctx, 999) + if !errors.Is(err, vector.ErrNotFound) { + t.Errorf("Delete non-existent: got %v, want ErrNotFound", err) + } +} + +// TestDelete_Removes verifies that a deleted vector no longer appears in +// Search results. +func TestDelete_Removes(t *testing.T) { + ctx := context.Background() + const dim = 3 + s := newStore(t, dim) + + if err := s.Upsert(ctx, 1, []float32{1.0, 0.0, 0.0}); err != nil { + t.Fatalf("Upsert: %v", err) + } + if err := s.Delete(ctx, 1); err != nil { + t.Fatalf("Delete: %v", err) + } + + hits, err := s.Search(ctx, []float32{1.0, 0.0, 0.0}, vector.SearchOpts{Limit: 10}) + if err != nil { + t.Fatalf("Search after delete: %v", err) + } + for _, h := range hits { + if h.ID == 1 { + t.Error("deleted entry ID 1 still appears in Search results") + } + } +} + +// TestUpsert_Replaces verifies that a second Upsert with the same ID replaces +// the previous vector. +func TestUpsert_Replaces(t *testing.T) { + ctx := context.Background() + const dim = 3 + s := newStore(t, dim) + + // Insert e1 pointing along X axis and e2 pointing along Y axis. + if err := s.Upsert(ctx, 1, []float32{1.0, 0.0, 0.0}); err != nil { + t.Fatalf("Upsert initial: %v", err) + } + if err := s.Upsert(ctx, 2, []float32{0.0, 1.0, 0.0}); err != nil { + t.Fatalf("Upsert 2: %v", err) + } + + // Before replace: query along X should return e1 first. + hitsBeforeReplace, err := s.Search(ctx, []float32{1.0, 0.0, 0.0}, vector.SearchOpts{Limit: 2}) + if err != nil { + t.Fatalf("Search before replace: %v", err) + } + if len(hitsBeforeReplace) == 0 || hitsBeforeReplace[0].ID != 1 { + t.Error("before replace: expected entry 1 (X axis) to be top hit") + } + + // Replace entry 1 to point along Y axis. + if err := s.Upsert(ctx, 1, []float32{0.0, 1.0, 0.0}); err != nil { + t.Fatalf("Upsert replace: %v", err) + } + + // After replace: both e1 and e2 point along Y; query along X now has lower + // similarity to both. The original e1-along-X entry should no longer exist. + // Query along Y should return one of {1,2} as top hit. + hitsAfterReplace, err := s.Search(ctx, []float32{0.0, 1.0, 0.0}, vector.SearchOpts{Limit: 2}) + if err != nil { + t.Fatalf("Search after replace: %v", err) + } + if len(hitsAfterReplace) == 0 { + t.Fatal("Search after replace returned no hits") + } + // Top hit must be either 1 or 2 (both are now Y-axis vectors). + top := hitsAfterReplace[0].ID + if top != 1 && top != 2 { + t.Errorf("top hit after replace = %d, want 1 or 2", top) + } +} + +// TestClose_Idempotent verifies that calling Close multiple times is safe +// and that operations after Close return ErrClosed. +func TestClose_Idempotent(t *testing.T) { + ctx := context.Background() + s := newStore(t, 4) + + if err := s.Close(ctx); err != nil { + t.Fatalf("first Close: %v", err) + } + if err := s.Close(ctx); err != nil { + t.Fatalf("second Close (idempotent): %v", err) + } + + // Operations after Close should return ErrClosed. + err := s.Upsert(ctx, 1, filled(4, 1.0)) + if !errors.Is(err, vector.ErrClosed) { + t.Errorf("Upsert after Close: got %v, want ErrClosed", err) + } + + _, err = s.Search(ctx, filled(4, 1.0), vector.SearchOpts{}) + if !errors.Is(err, vector.ErrClosed) { + t.Errorf("Search after Close: got %v, want ErrClosed", err) + } + + err = s.Delete(ctx, 1) + if !errors.Is(err, vector.ErrClosed) { + t.Errorf("Delete after Close: got %v, want ErrClosed", err) + } +} + +// TestSearch_DimensionMismatch verifies that Search rejects a query vector of +// the wrong length. +func TestSearch_DimensionMismatch(t *testing.T) { + ctx := context.Background() + s := newStore(t, 4) + + _, err := s.Search(ctx, []float32{1.0, 2.0}, vector.SearchOpts{}) + if !errors.Is(err, vector.ErrInvalidArgument) { + t.Errorf("Search wrong dim: got %v, want ErrInvalidArgument", err) + } +} diff --git a/pkg/lore/vector/vector.go b/pkg/lore/vector/vector.go new file mode 100644 index 0000000..63227f8 --- /dev/null +++ b/pkg/lore/vector/vector.go @@ -0,0 +1,95 @@ +// Package vector defines the VectorStore interface and the types it uses. +// Implementations are dimension-bound at construction and accept a +// caller-managed *sql.DB (or equivalent). Callers own the database lifecycle; +// the VectorStore only owns the schema objects it creates inside that DB. +// +// The package ships one reference implementation: sqlitevec, a pure-Go +// SQLite-backed store that keeps vectors as BLOB columns and does cosine +// similarity in Go. That choice is documented in detail in +// pkg/lore/vector/sqlitevec/sqlitevec.go. +package vector + +import ( + "context" + "errors" + + "github.com/mathomhaus/lore/pkg/lore" +) + +// VectorStore persists vectors keyed by entry ID and answers nearest-neighbor +// queries. Implementations are dimension-bound at construction. Caller-managed +// resources (consumer-owned *sql.DB or equivalent). +// +// SearchOpts carries Kind and Tag filter hints. VectorStore implementations +// are not required to honor them: the Retriever layer can post-filter results +// through Store.Get after Search returns. Reference implementations document +// whether they push filters into the query or post-filter. +type VectorStore interface { + // Upsert stores the vector for the given entry ID. Replaces any existing + // vector for that ID. The vector length must equal Dimensions(); + // ErrInvalidArgument is returned for a mismatch. The operation is + // idempotent: calling Upsert twice with the same ID and vector is safe. + Upsert(ctx context.Context, id int64, vector []float32) error + + // Delete removes the vector for the given entry ID. Returns ErrNotFound + // when no vector is stored for that ID. + Delete(ctx context.Context, id int64) error + + // Search returns the top-Limit vectors most similar to the query vector, + // in descending score order. Higher Score means more similar. The query + // vector length must equal Dimensions(); ErrInvalidArgument otherwise. + // + // Kind and Tag filters on SearchOpts are advisory hints. The reference + // implementation ignores them and returns top-K by similarity alone; + // callers that need filtered results should post-filter via Store.Get. + Search(ctx context.Context, query []float32, opts SearchOpts) ([]Hit, error) + + // Dimensions returns the vector length this store was constructed for. + // All Upsert and Search calls must supply vectors of exactly this length. + Dimensions() int + + // Close releases any resources held by the store beyond the caller-owned + // DB. Idempotent: calling Close more than once returns nil. + Close(ctx context.Context) error +} + +// Hit is one result from a Search call. ID is the entry's storage identifier +// as passed to Upsert. Score is the similarity score; higher is more similar. +// The scale is implementation-defined (cosine similarity in [-1, 1] for the +// reference impl) but always higher-is-better within a single query result. +type Hit struct { + ID int64 + Score float64 +} + +// SearchOpts configures a Search call. +type SearchOpts struct { + // Limit is the maximum number of hits to return. Default 10 when zero. + // Negative values are clamped to 1. + Limit int + + // Kinds, when non-empty, hints that callers want only entries of these + // kinds. VectorStore implementations are not required to apply this + // filter; the Retriever layer post-filters via Store.Get when needed. + Kinds []lore.Kind + + // Tags, when non-empty, hints that callers want only entries carrying all + // of these tags. Same advisory semantics as Kinds. + Tags []string +} + +// Typed sentinel errors for caller-side branching. All error values returned +// by VectorStore implementations should wrap one of these so callers can use +// errors.Is without matching text. +var ( + // ErrNotFound is returned by Delete when the given entry ID has no stored + // vector. Implementations wrap this via fmt.Errorf("...: %w", ErrNotFound). + ErrNotFound = errors.New("vector: not found") + + // ErrInvalidArgument is returned when a caller-supplied argument fails + // validation (dimension mismatch, negative limit, nil vector, etc.). + ErrInvalidArgument = errors.New("vector: invalid argument") + + // ErrClosed is returned when an operation is attempted after Close. + ErrClosed = errors.New("vector: closed") +)