diff --git a/storage/hnsw.go b/storage/hnsw.go index 85d4637..92155b3 100644 --- a/storage/hnsw.go +++ b/storage/hnsw.go @@ -296,9 +296,27 @@ func (h *HNSWIndex) Search(query []float32, k, ef int) ([]string, []float32) { } func (h *HNSWIndex) persist(ctx context.Context, s *Store, model string) error { + h.mu.RLock() + ids := make([]string, 0, len(h.graph)) + for id := range h.graph { + ids = append(ids, id) + } + h.mu.RUnlock() + return h.persistNodes(ctx, s, model, ids) +} + +// persistNodes upserts the embeddings_hnsw rows for the given node IDs. +// IDs that are no longer in the graph are skipped. +func (h *HNSWIndex) persistNodes(ctx context.Context, s *Store, model string, ids []string) error { + ctx, cancel := s.withTimeout(ctx) + defer cancel() h.mu.RLock() defer h.mu.RUnlock() - for nodeID, levels := range h.graph { + for _, nodeID := range ids { + levels, ok := h.graph[nodeID] + if !ok { + continue + } vec := h.vectors[nodeID] byLevel := make(map[string][]HNSWNode, len(levels)) for l, links := range levels { @@ -325,6 +343,125 @@ func (h *HNSWIndex) persist(ctx context.Context, s *Store, model string) error { return nil } +func (h *HNSWIndex) deleteRow(ctx context.Context, s *Store, nodeID string) error { + ctx, cancel := s.withTimeout(ctx) + defer cancel() + _, err := s.db.ExecContext(ctx, `DELETE FROM embeddings_hnsw WHERE node_id=?`, nodeID) + return err +} + +// Upsert adds or replaces a single vector in the index and persists only the +// affected rows, avoiding a full Build. Callers must store the embedding in +// the embeddings table first so the DB stays the source of truth. +func (h *HNSWIndex) Upsert(ctx context.Context, s *Store, model, nodeID string, vec []float32) error { + h.mu.Lock() + affected := h.upsertLocked(nodeID, vec) + h.mu.Unlock() + return h.persistNodes(ctx, s, model, setKeys(affected)) +} + +// upsertLocked replaces (or adds) nodeID in the graph. Caller must hold h.mu. +func (h *HNSWIndex) upsertLocked(nodeID string, vec []float32) map[string]bool { + affected := h.removeLocked(nodeID) + h.vectors[nodeID] = vec + h.insert(nodeID, vec) + affected[nodeID] = true + for _, links := range h.graph[nodeID] { + for _, id := range links { + affected[id] = true + } + } + return affected +} + +// Remove detaches nodeID from the graph (link pruning, entry-point repair) +// and deletes its persisted row. Removing an absent node is a no-op except +// for the row deletion, keeping the DB state consistent. +func (h *HNSWIndex) Remove(ctx context.Context, s *Store, model, nodeID string) error { + h.mu.Lock() + affected := h.removeLocked(nodeID) + h.mu.Unlock() + if err := h.persistNodes(ctx, s, model, setKeys(affected)); err != nil { + return err + } + return h.deleteRow(ctx, s, nodeID) +} + +// removeLocked detaches and drops nodeID, returning every node whose +// persisted state changed (the node itself plus all nodes that linked to +// it). Caller must hold h.mu. +func (h *HNSWIndex) removeLocked(nodeID string) map[string]bool { + affected := map[string]bool{nodeID: true} + levels, ok := h.graph[nodeID] + if !ok { + delete(h.vectors, nodeID) + return affected + } + for l, links := range levels { + for _, nid := range links { + nl, ok := h.graph[nid] + if !ok || nl[l] == nil { + continue + } + without := removeString(nl[l], nodeID) + if len(without) != len(nl[l]) { + nl[l] = without + affected[nid] = true + } + } + } + wasEntry := h.entryPoint == nodeID + delete(h.graph, nodeID) + delete(h.vectors, nodeID) + if wasEntry { + entry, maxLevel := h.topNodeLocked() + h.entryPoint = entry + h.maxLevel = maxLevel + if entry != "" { + affected[entry] = true + } + } + return affected +} + +// topNodeLocked returns the remaining node with the highest level, usable as +// a new entry point. Returns ("", 0) when the graph is empty. +// Caller must hold h.mu. +func (h *HNSWIndex) topNodeLocked() (string, int) { + entry, maxLevel := "", 0 + for id, levels := range h.graph { + top := 0 + for l := range levels { + if l > top { + top = l + } + } + if entry == "" || top > maxLevel { + entry, maxLevel = id, top + } + } + return entry, maxLevel +} + +func removeString(list []string, s string) []string { + out := list[:0] + for _, v := range list { + if v != s { + out = append(out, v) + } + } + return out +} + +// setKeys returns the keys of a string set as a slice. +func setKeys(m map[string]bool) []string { + out := make([]string, 0, len(m)) + for k := range m { + out = append(out, k) + } + return out +} + func containsString(list []string, s string) bool { for _, v := range list { if v == s { diff --git a/storage/hnsw_incremental_test.go b/storage/hnsw_incremental_test.go new file mode 100644 index 0000000..63bb6ef --- /dev/null +++ b/storage/hnsw_incremental_test.go @@ -0,0 +1,247 @@ +package storage + +import ( + "context" + "fmt" + "testing" +) + +// seedHNSW creates n nodes with angle-distinct embeddings (dim 0 = 1, +// dim 1 = i*step) and builds the HNSW index for the given model. +func seedHNSW(t *testing.T, s *Store, model string, n int, step float32) { + t.Helper() + ctx := context.Background() + for i := 0; i < n; i++ { + id := fmt.Sprintf("h-%d", i) + if err := s.CreateNode(ctx, &Node{ + ID: id, Type: "convention", Content: id, + ContentHash: "h-" + id, Scope: "project", Project: "test", + }); err != nil { + t.Fatalf("CreateNode %s: %v", id, err) + } + vec := make([]float32, 8) + vec[0] = 1.0 + vec[1] = float32(i) * step + if err := s.SaveEmbedding(ctx, id, model, vec); err != nil { + t.Fatalf("SaveEmbedding %s: %v", id, err) + } + } + if err := s.BuildHNSWIndex(ctx, model); err != nil { + t.Fatalf("BuildHNSWIndex: %v", err) + } +} + +func dirVec(t *testing.T, d float32) []float32 { + t.Helper() + v := make([]float32, 8) + v[0] = 1.0 + v[1] = d + return v +} + +// TestHNSWIncrementalUpsert verifies SaveEmbedding updates a built index +// without requiring an explicit BuildHNSWIndex. +func TestHNSWIncrementalUpsert(t *testing.T) { + s, cleanup := setupStore(t) + defer cleanup() + ctx := context.Background() + model := "inc-model" + seedHNSW(t, s, model, 10, 0.1) // h-0..h-9 cover dir 0..0.9 + + // Add a brand-new node; it must become searchable immediately. + if err := s.CreateNode(ctx, &Node{ + ID: "h-new", Type: "convention", Content: "h-new", + ContentHash: "h-h-new", Scope: "project", Project: "test", + }); err != nil { + t.Fatal(err) + } + // Exact match for direction 2.0, far outside the seeded range. + newVec := make([]float32, 8) + newVec[0] = 1.0 + newVec[1] = 2.0 + if err := s.SaveEmbedding(ctx, "h-new", model, newVec); err != nil { + t.Fatalf("SaveEmbedding upsert: %v", err) + } + ids, _, err := s.SearchHNSW(ctx, model, dirVec(t, 2.0), 1, 50) + if err != nil { + t.Fatalf("SearchHNSW: %v", err) + } + if len(ids) != 1 || ids[0] != "h-new" { + t.Fatalf("expected top result h-new, got %v", ids) + } + + // The affected node must be persisted in embeddings_hnsw without a rebuild. + var count int + if err := s.DB().QueryRowContext(ctx, + `SELECT COUNT(*) FROM embeddings_hnsw WHERE node_id=?`, "h-new").Scan(&count); err != nil { + t.Fatal(err) + } + if count != 1 { + t.Errorf("embeddings_hnsw rows for h-new = %d, want 1", count) + } + + // Replace h-5's vector (was dir 0.5) with dir 3.0: it must now rank + // first for queries in that direction. + rep := make([]float32, 8) + rep[0] = 1.0 + rep[1] = 3.0 + if err := s.SaveEmbedding(ctx, "h-5", model, rep); err != nil { + t.Fatalf("SaveEmbedding replace: %v", err) + } + ids, _, err = s.SearchHNSW(ctx, model, dirVec(t, 3.0), 1, 50) + if err != nil { + t.Fatalf("SearchHNSW after replace: %v", err) + } + if len(ids) != 1 || ids[0] != "h-5" { + t.Fatalf("expected top result h-5 after replace, got %v", ids) + } + // And it must NOT be the top hit for its old direction any more. + ids, _, err = s.SearchHNSW(ctx, model, dirVec(t, 0.5), 1, 50) + if err != nil { + t.Fatal(err) + } + if len(ids) != 1 || ids[0] == "h-5" { + t.Errorf("h-5 still top-ranked for old direction: %v", ids) + } +} + +// TestHNSWIncrementalDelete verifies DeleteEmbedding detaches the node from +// the index and clears its persisted graph row. +func TestHNSWIncrementalDelete(t *testing.T) { + s, cleanup := setupStore(t) + defer cleanup() + ctx := context.Background() + model := "del-model" + seedHNSW(t, s, model, 10, 0.1) + + // h-9 (dir 0.9) is the only node near dir 0.9. + ids, _, err := s.SearchHNSW(ctx, model, dirVec(t, 0.9), 1, 50) + if err != nil { + t.Fatal(err) + } + if len(ids) != 1 || ids[0] != "h-9" { + t.Fatalf("expected h-9 initially, got %v", ids) + } + + if err := s.DeleteEmbedding(ctx, "h-9"); err != nil { + t.Fatalf("DeleteEmbedding: %v", err) + } + + ids, _, err = s.SearchHNSW(ctx, model, dirVec(t, 0.9), 5, 50) + if err != nil { + t.Fatal(err) + } + for _, id := range ids { + if id == "h-9" { + t.Fatalf("h-9 still returned after delete: %v", ids) + } + } + + // Persisted graph row must be gone. + var count int + if err := s.DB().QueryRowContext(ctx, + `SELECT COUNT(*) FROM embeddings_hnsw WHERE node_id=?`, "h-9").Scan(&count); err != nil { + t.Fatal(err) + } + if count != 0 { + t.Errorf("embeddings_hnsw rows for h-9 after delete = %d, want 0", count) + } +} + +// TestHNSWTransactionalInvalidation verifies embedding writes inside WithTx +// drop the cached index so the next search rebuilds from the embeddings table. +func TestHNSWTransactionalInvalidation(t *testing.T) { + s, cleanup := setupStore(t) + defer cleanup() + ctx := context.Background() + model := "tx-model" + seedHNSW(t, s, model, 10, 0.1) // cache built for tx-model + + err := s.WithTx(ctx, func(tx Storage) error { + if err := tx.CreateNode(ctx, &Node{ + ID: "tx-node", Type: "convention", Content: "tx-node", + ContentHash: "h-tx-node", Scope: "project", Project: "test", + }); err != nil { + return err + } + vec := make([]float32, 8) + vec[0] = 1.0 + vec[1] = 5.0 // far outside seeded range + return tx.SaveEmbedding(ctx, "tx-node", model, vec) + }) + if err != nil { + t.Fatalf("WithTx: %v", err) + } + + // The cache must have been invalidated and rebuilt from the table. + ids, _, err := s.SearchHNSW(ctx, model, dirVec(t, 5.0), 1, 50) + if err != nil { + t.Fatalf("SearchHNSW after tx: %v", err) + } + if len(ids) != 1 || ids[0] != "tx-node" { + t.Fatalf("expected tx-node after transactional save, got %v", ids) + } +} + +// TestHNSWIncrementalVsFullBuild checks the incrementally-mutated index +// still agrees with a full rebuild on the same data. +func TestHNSWIncrementalVsFullBuild(t *testing.T) { + s, cleanup := setupStore(t) + defer cleanup() + ctx := context.Background() + model := "parity-model" + seedHNSW(t, s, model, 10, 0.1) + + // Several incremental mutations. + for i := 0; i < 3; i++ { + id := fmt.Sprintf("extra-%d", i) + if err := s.CreateNode(ctx, &Node{ + ID: id, Type: "convention", Content: id, + ContentHash: "h-" + id, Scope: "project", Project: "test", + }); err != nil { + t.Fatal(err) + } + v := make([]float32, 8) + v[0] = 1.0 + v[1] = 1.0 + float32(i)*0.2 + if err := s.SaveEmbedding(ctx, id, model, v); err != nil { + t.Fatal(err) + } + } + for _, id := range []string{"h-0", "h-1", "extra-0"} { + if err := s.DeleteEmbedding(ctx, id); err != nil { + t.Fatalf("DeleteEmbedding %s: %v", id, err) + } + } + + query := dirVec(t, 1.2) + incIDs, _, err := s.SearchHNSW(ctx, model, query, 5, 100) + if err != nil { + t.Fatal(err) + } + + // Full rebuild must return the same membership for the top-5. + if err := s.BuildHNSWIndex(ctx, model); err != nil { + t.Fatal(err) + } + fullIDs, fullScores, err := s.SearchHNSW(ctx, model, query, 5, 100) + if err != nil { + t.Fatal(err) + } + + if len(fullIDs) != len(incIDs) { + t.Fatalf("result counts differ: incremental %d vs full %d", len(incIDs), len(fullIDs)) + } + seen := make(map[string]bool, len(incIDs)) + for _, id := range incIDs { + seen[id] = true + } + for _, id := range fullIDs { + if !seen[id] { + t.Errorf("full build returned %s missing from incremental results %v", id, incIDs) + } + } + if len(fullScores) != 5 { + t.Fatalf("expected 5 scores, got %d", len(fullScores)) + } +} diff --git a/storage/sqlite_tx.go b/storage/sqlite_tx.go index abdf407..0a8bcd9 100644 --- a/storage/sqlite_tx.go +++ b/storage/sqlite_tx.go @@ -28,17 +28,28 @@ func (s *Store) WithTx(ctx context.Context, fn func(Storage) error) error { } defer func() { _ = tx.Rollback() }() - txStore := &txStore{tx: tx} + txStore := &txStore{tx: tx, store: s, dirtyHNSW: make(map[string]struct{})} if err := fn(txStore); err != nil { return err } - return tx.Commit() + if err := tx.Commit(); err != nil { + return err + } + // Embedding writes inside the transaction were not applied to the + // built HNSW indexes incrementally; drop the affected caches so the + // next SearchHNSW rebuilds them from the embeddings table. + for model := range txStore.dirtyHNSW { + s.invalidateHNSWIndex(model) + } + return nil }, 5, 50*time.Millisecond) } // txStore is a Storage implementation backed by a SQL transaction. type txStore struct { - tx *sql.Tx + tx *sql.Tx + store *Store + dirtyHNSW map[string]struct{} // embedding models touched in this tx } // txStore is a thin wrapper that delegates all operations to shared *Q functions. @@ -173,11 +184,29 @@ func (t *txStore) GetVersions(ctx context.Context, nodeID string) ([]*NodeVersio } func (t *txStore) SaveEmbedding(ctx context.Context, nodeID, model string, vector []float32) error { - return saveEmbeddingQ(ctx, t.tx, nodeID, model, vector) + if err := saveEmbeddingQ(ctx, t.tx, nodeID, model, vector); err != nil { + return err + } + t.dirtyHNSW[model] = struct{}{} + return nil } func (t *txStore) DeleteEmbedding(ctx context.Context, nodeID string) error { - return deleteEmbeddingQ(ctx, t.tx, nodeID) + // Capture the embedding's model (if any) in-tx so the corresponding + // HNSW cache gets invalidated after commit. + _, model, _ := getEmbeddingQ(ctx, t.tx, nodeID) + // FK order: embeddings_hnsw.node_id references embeddings.node_id, so + // the graph row goes first. + if _, err := t.tx.ExecContext(ctx, `DELETE FROM embeddings_hnsw WHERE node_id=?`, nodeID); err != nil { + return err + } + if err := deleteEmbeddingQ(ctx, t.tx, nodeID); err != nil { + return err + } + if model != "" { + t.dirtyHNSW[model] = struct{}{} + } + return nil } func (t *txStore) GetEmbedding(ctx context.Context, nodeID string) ([]float32, string, error) { diff --git a/storage/vectors.go b/storage/vectors.go index 072f615..d56b167 100644 --- a/storage/vectors.go +++ b/storage/vectors.go @@ -28,17 +28,32 @@ func DecodeVector(b []byte) []float32 { return v } -// SaveEmbedding stores a vector embedding for a node. +// SaveEmbedding stores a vector embedding for a node and keeps the HNSW +// index up to date incrementally: if an index for the embedding's model is +// built, only the affected node and its neighbors are re-linked and +// persisted (no full rebuild). A failed index update does not fail the save +// — the embeddings table is the source of truth and a later +// BuildHNSWIndex/SearchHNSW will reconcile the index. // // Like every other write path it retries on SQLITE_BUSY: embedding writes run // outside the engine's write lock and can briefly contend with the async // ingestion goroutine for the single SQLite writer. func (s *Store) SaveEmbedding(ctx context.Context, nodeID, model string, vector []float32) error { - return retryOnBusy(func() error { - ctx, cancel := s.withTimeout(ctx) + err := retryOnBusy(func() error { + qctx, cancel := s.withTimeout(ctx) defer cancel() - return saveEmbeddingQ(ctx, s.q(), nodeID, model, vector) + return saveEmbeddingQ(qctx, s.q(), nodeID, model, vector) }, 5, 50*time.Millisecond) + if err != nil { + return err + } + s.hnswMu.Lock() + idx := s.hnswIndexes[model] + s.hnswMu.Unlock() + if idx != nil { + _ = idx.Upsert(ctx, s, model, nodeID, vector) // best-effort; next Build reconciles + } + return nil } func saveEmbeddingQ(ctx context.Context, q queryable, nodeID, model string, vector []float32) error { @@ -66,15 +81,50 @@ func getEmbeddingQ(ctx context.Context, q queryable, nodeID string) ([]float32, return DecodeVector(blob), model, nil } -// DeleteEmbedding removes a vector embedding for a node. +// DeleteEmbedding removes a vector embedding for a node and detaches the +// node from the built HNSW index for its model (link pruning + entry-point +// repair). If no index is built, any stale persisted graph row is still +// cleared so the DB stays consistent. The index is best-effort: a failed +// detach is reconciled by the next BuildHNSWIndex/SearchHNSW. func (s *Store) DeleteEmbedding(ctx context.Context, nodeID string) error { + model, _ := s.embeddingModel(ctx, nodeID) // empty when no embedding is stored + + // embeddings_hnsw.node_id references embeddings.node_id (no cascade in + // the base schema), so the graph row must be removed before the + // embedding row itself. + if model != "" { + s.hnswMu.Lock() + idx := s.hnswIndexes[model] + s.hnswMu.Unlock() + if idx != nil { + _ = idx.Remove(ctx, s, model, nodeID) // best-effort; next Build reconciles + } else { + qctx, cancel := s.withTimeout(ctx) + _, _ = s.q().ExecContext(qctx, `DELETE FROM embeddings_hnsw WHERE node_id=?`, nodeID) + cancel() + } + } + return retryOnBusy(func() error { - ctx, cancel := s.withTimeout(ctx) + qctx, cancel := s.withTimeout(ctx) defer cancel() - return deleteEmbeddingQ(ctx, s.q(), nodeID) + return deleteEmbeddingQ(qctx, s.q(), nodeID) }, 5, 50*time.Millisecond) } +// embeddingModel returns the model that produced a node's embedding, or an +// error (sql.ErrNoRows when none is stored). +func (s *Store) embeddingModel(ctx context.Context, nodeID string) (string, error) { + qctx, cancel := s.withTimeout(ctx) + defer cancel() + var model string + err := s.q().QueryRowContext(qctx, `SELECT model FROM embeddings WHERE node_id=?`, nodeID).Scan(&model) + if err != nil { + return "", err + } + return model, nil +} + func deleteEmbeddingQ(ctx context.Context, q queryable, nodeID string) error { _, err := q.ExecContext(ctx, `DELETE FROM embeddings WHERE node_id=?`, nodeID) return err @@ -160,6 +210,15 @@ func (s *Store) BuildHNSWIndex(ctx context.Context, model string) error { return nil } +// invalidateHNSWIndex drops the built index for a model so that the next +// SearchHNSW rebuilds it from the embeddings table. Used after transactional +// embedding writes, which bypass the incremental Upsert/Remove path. +func (s *Store) invalidateHNSWIndex(model string) { + s.hnswMu.Lock() + delete(s.hnswIndexes, model) + s.hnswMu.Unlock() +} + // SearchHNSW searches the HNSW index for the given model, building it // on demand if it has not been built yet. func (s *Store) SearchHNSW(ctx context.Context, model string, query []float32, k, ef int) ([]string, []float32, error) {