feat(storage): incremental HNSW updates on embedding writes

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 {