fix(rag): make BM25 indexing incremental and idempotent

src/providers/rag/search.test.ts

@@ -0,0 +1,78 @@
+import { describe, expect, test } from "bun:test"
+import { HybridSearchEngine, type Chunk } from "./search"
+
+const CONTAINER = "test_container"
+
+function makeChunk(id: string, content: string, embedding: number[]): Chunk {
+  return {
+    id,
+    content,
+    sessionId: id,
+    chunkIndex: 0,
+    embedding,
+  }
+}
+
+// Deliberately uneven document lengths and term frequencies: avgDocLength and
+// idf both have to be wrong for these scores to move.
+const CHUNKS: Chunk[] = [
+  makeChunk("c1", "quantum telescope", [1, 0, 0]),
+  makeChunk(
+    "c2",
+    "harvest orbital drift sensor array calibration payload module thermal shielding harvest harvest",
+    [0, 1, 0]
+  ),
+  makeChunk("c3", "harvest sensor", [0, 0, 1]),
+]
+
+const QUERY = "quantum harvest"
+const QUERY_EMBEDDING = [0.5, 0.5, 0.5]
+
+describe("HybridSearchEngine BM25 indexing", () => {
+  test("re-ingesting the same chunks does not change search scores", () => {
+    const fresh = new HybridSearchEngine()
+    fresh.addChunks(CONTAINER, CHUNKS)
+
+    const reingested = new HybridSearchEngine()
+    reingested.addChunks(CONTAINER, CHUNKS)
+    reingested.addChunks(CONTAINER, CHUNKS)
+
+    expect(reingested.getChunkCount(CONTAINER)).toBe(CHUNKS.length)
+    expect(reingested.search(CONTAINER, QUERY_EMBEDDING, QUERY, 10)).toEqual(
+      fresh.search(CONTAINER, QUERY_EMBEDDING, QUERY, 10)
+    )
+  })
+
+  test("re-indexing a chunk with new content drops its old terms", () => {
+    const engine = new HybridSearchEngine()
+    engine.addChunks(CONTAINER, [
+      makeChunk("c1", "telescope calibration", [1, 0, 0]),
+      makeChunk("c2", "telescope beacon", [0, 1, 0]),
+    ])
+
+    // Same chunk ID, different content — as produced by a forced re-ingest.
+    engine.addChunks(CONTAINER, [makeChunk("c1", "harvest orbital", [1, 0, 0])])
+
+    const results = engine.search(CONTAINER, [1, 1, 0], "telescope", 10)
+    const c1 = results.find((r) => r.content === "harvest orbital")
+    const c2 = results.find((r) => r.content === "telescope beacon")
+
+    expect(c1).toBeDefined()
+    expect(c2).toBeDefined()
+    expect(c1!.bm25Score).toBe(0)
+    expect(c2!.bm25Score).toBeGreaterThan(0)
+  })
+
+  test("scores match a freshly built index after content is replaced", () => {
+    const replaced = new HybridSearchEngine()
+    replaced.addChunks(CONTAINER, [makeChunk("c1", "telescope calibration payload", [1, 0, 0])])
+    replaced.addChunks(CONTAINER, CHUNKS)
+
+    const fresh = new HybridSearchEngine()
+    fresh.addChunks(CONTAINER, CHUNKS)
+
+    expect(replaced.search(CONTAINER, QUERY_EMBEDDING, QUERY, 10)).toEqual(
+      fresh.search(CONTAINER, QUERY_EMBEDDING, QUERY, 10)
+    )
+  })
+})

src/providers/rag/search.ts

@@ -165,6 +165,10 @@ interface BM25Index {
   invertedIndex: Map<string, Map<string, number>>
   /** Document lengths (in tokens) */
   docLengths: Map<string, number>
+  /** Distinct terms per document, so re-indexing can drop stale postings */
+  docTerms: Map<string, Set<string>>
+  /** Running sum of all document lengths */
+  totalLength: number
   /** Average document length */
   avgDocLength: number
   /** Total number of documents */
@@ -175,22 +179,44 @@ function createBM25Index(): BM25Index {
   return {
     invertedIndex: new Map(),
     docLengths: new Map(),
+    docTerms: new Map(),
+    totalLength: 0,
     avgDocLength: 0,
     docCount: 0,
   }
 }
 
+function removeFromBM25Index(index: BM25Index, chunkId: string): void {
+  const docLength = index.docLengths.get(chunkId)
+  if (docLength === undefined) return
+
+  for (const term of index.docTerms.get(chunkId) || []) {
+    const postings = index.invertedIndex.get(term)
+    if (!postings) continue
+    postings.delete(chunkId)
+    if (postings.size === 0) index.invertedIndex.delete(term)
+  }
+
+  index.docTerms.delete(chunkId)
+  index.docLengths.delete(chunkId)
+  index.docCount--
+  index.totalLength -= docLength
+  index.avgDocLength = index.docCount > 0 ? index.totalLength / index.docCount : 0
+}
+
 function addToBM25Index(index: BM25Index, chunkId: string, text: string): void {
+  // Chunk IDs are deterministic, so a forced or resumed ingest re-adds the same
+  // ID. Replace the existing entry instead of appending to it, otherwise
+  // docCount drifts above docLengths.size and skews both idf and avgDocLength.
+  removeFromBM25Index(index, chunkId)
+
   const tokens = tokenize(text)
   index.docLengths.set(chunkId, tokens.length)
   index.docCount++
 
-  // Update average document length
-  let totalLength = 0
-  for (const len of index.docLengths.values()) {
-    totalLength += len
-  }
-  index.avgDocLength = totalLength / index.docCount
+  // Update average document length from a running total, not a full rescan
+  index.totalLength += tokens.length
+  index.avgDocLength = index.totalLength / index.docCount
 
   // Build term frequency map
   const termFreqs = new Map<string, number>()
@@ -205,6 +231,7 @@ function addToBM25Index(index: BM25Index, chunkId: string, text: string): void {
     }
     index.invertedIndex.get(term)!.set(chunkId, freq)
   }
+  index.docTerms.set(chunkId, new Set(termFreqs.keys()))
 }
 
 function searchBM25(index: BM25Index, query: string): Map<string, number> {