From 0e2b9cfbfed672e3fbdcd65f325a9d6581bfa9f4 Mon Sep 17 00:00:00 2001 From: Agnik47 <140933190+Agnik47@users.noreply.github.com> Date: Fri, 14 Aug 2026 21:29:59 +0530 Subject: [PATCH] fix(rag): make BM25 indexing incremental and idempotent MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit `addToBM25Index` rescanned the whole `docLengths` map on every insert to recompute `avgDocLength`, making ingest O(n^2) in the number of chunks. It also incremented `docCount` unconditionally, so re-adding a chunk ID — which happens on a forced or resumed ingest, since chunk IDs are deterministic — pushed `docCount` above `docLengths.size` and corrupted both `avgDocLength` and idf for every term, while leaving stale postings for the replaced content. Track a running `totalLength`, and replace an existing entry (dropping its postings and length) before re-adding it. Ingest is now linear and BM25 scores no longer depend on run history. Fixes #74 Co-Authored-By: Claude Opus 5 Claude-Session: https://claude.ai/code/session_01JqGuPVWVYRJEJQp2RXchf1 --- src/providers/rag/search.test.ts | 78 ++++++++++++++++++++++++++++++++ src/providers/rag/search.ts | 39 +++++++++++++--- 2 files changed, 111 insertions(+), 6 deletions(-) create mode 100644 src/providers/rag/search.test.ts diff --git a/src/providers/rag/search.test.ts b/src/providers/rag/search.test.ts new file mode 100644 index 0000000..e86a581 --- /dev/null +++ b/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) + ) + }) +}) diff --git a/src/providers/rag/search.ts b/src/providers/rag/search.ts index df2a365..743b0bb 100644 --- a/src/providers/rag/search.ts +++ b/src/providers/rag/search.ts @@ -165,6 +165,10 @@ interface BM25Index { invertedIndex: Map> /** Document lengths (in tokens) */ docLengths: Map + /** Distinct terms per document, so re-indexing can drop stale postings */ + docTerms: Map> + /** 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() @@ -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 {