diff --git a/src/cli/groups/ask.ts b/src/cli/groups/ask.ts index 17cae7e7..e0ba7756 100644 --- a/src/cli/groups/ask.ts +++ b/src/cli/groups/ask.ts @@ -128,7 +128,9 @@ export function addAskCommands(program: Command): void { "warn", `stopped at ${limit} filing(s); this answer sees only what is indexed. ` + "Run `sec index` for the full build (`--limit` bounds one run), raise " + - "`--index-limit`, or pass `--no-index` to answer from the index as it stands." + "`--index-limit`, or pass `--no-index` to answer from the index as it stands. " + + "Retrieval scores every chunk in the index, so a larger index is a slower " + + "question — scope the build with `--company` or `--form` if you can." ) ); } diff --git a/src/kb/PagedChunkVectorStorage.ts b/src/kb/PagedChunkVectorStorage.ts new file mode 100644 index 00000000..d1587547 --- /dev/null +++ b/src/kb/PagedChunkVectorStorage.ts @@ -0,0 +1,117 @@ +/** + * @license + * Copyright 2026 Steven Roussey + * SPDX-License-Identifier: Apache-2.0 + */ + +import type { + ChunkVectorPrimaryKey, + ChunkVectorStorageSchema, + TypedArray, + VectorSearchOptions, +} from "workglow"; +import { + assertVectorShape, + cosineSimilarity, + emitSimilaritySearch, + matchesFilter, + SqliteVectorStorage, +} from "workglow"; + +/** + * Rows read per page. + * + * A chunk carries its text and a JSON-encoded vector, so a page is the working + * set: large enough that the scan is a few hundred statements over a corpus of + * hundreds of thousands of chunks, small enough to stay a fixed cost. + */ +const SCAN_PAGE = 512; + +/** A scored row, as {@link SqliteVectorStorage.similaritySearch} returns them. */ +interface Scored { + readonly score: number; +} + +/** + * Inserts into a list held at `topK`, descending by score. + * + * Linear rather than a heap: `topK` is single digits by default and the + * comparison runs once per candidate that beats the current floor, which on a + * long scan is a vanishing fraction of the rows. + */ +function keepBest(kept: T[], row: T, topK: number): void { + let low = 0; + let high = kept.length; + while (low < high) { + const mid = (low + high) >> 1; + if (kept[mid]!.score >= row.score) low = mid + 1; + else high = mid; + } + if (low >= topK) return; + kept.splice(low, 0, row); + if (kept.length > topK) kept.pop(); +} + +/** + * The chunk store `sec ask` searches, scored a page at a time. + * + * The inherited search is `SELECT * FROM kb_chunk` with no bound followed by a + * cosine per row, which puts the whole index in the heap to answer one + * question. `sec index` is a build measured in hours to days and `sec ask` + * tells the user to run it, so the index this has to read is exactly the one + * that does not fit. + * + * Bounded memory, not bounded time: there is no approximate-nearest-neighbour + * index here, so every question still scores every chunk and latency grows with + * the corpus. What this removes is the heap ceiling that made a large index + * unqueryable rather than slow. + */ +export class PagedChunkVectorStorage extends SqliteVectorStorage< + ChunkVectorStorageSchema, + ChunkVectorPrimaryKey +> { + override async similaritySearch( + query: TypedArray, + options: VectorSearchOptions> = {} + ) { + assertVectorShape(query, this.getVectorDimensions(), "query"); + const { topK = 10, filter, scoreThreshold = 0 } = options; + + type Row = NonNullable>>[number]; + const kept: (Row & Scored)[] = []; + if (topK <= 0) return emitSimilaritySearch(this.events, query, kept); + + // Ordered by the primary key so the pages partition the table: LIMIT with + // OFFSET and no ORDER BY is free to hand back a row twice and skip another. + for (let offset = 0; ; offset += SCAN_PAGE) { + const page = + (await this.getAll({ + orderBy: [{ column: "chunk_id", direction: "ASC" }], + limit: SCAN_PAGE, + offset, + })) ?? []; + for (const entity of page) { + const metadata = (entity.metadata ?? {}) as Record; + if (filter && !matchesFilter(metadata, filter)) continue; + const score = cosineSimilarity(query, toVector(entity.vector)); + if (score < scoreThreshold) continue; + keepBest(kept, { ...entity, score }, topK); + } + if (page.length < SCAN_PAGE) break; + } + + return emitSimilaritySearch(this.events, query, kept); + } +} + +/** + * The stored vector as a TypedArray, from whichever form the read produced. + * + * SQLite holds it as a JSON string and the tabular read usually decodes it + * already; a row written by an older release, or handed back raw, does not. + */ +function toVector(stored: unknown): TypedArray { + if (typeof stored === "string") return new Float32Array(JSON.parse(stored) as number[]); + if (Array.isArray(stored)) return new Float32Array(stored); + return stored as TypedArray; +} diff --git a/src/kb/secKnowledgeBase.ts b/src/kb/secKnowledgeBase.ts index d467f6ec..4fc4ef03 100644 --- a/src/kb/secKnowledgeBase.ts +++ b/src/kb/secKnowledgeBase.ts @@ -16,7 +16,6 @@ import { KnowledgeBase, registerKnowledgeBase, SqliteTabularStorage, - SqliteVectorStorage, unregisterKnowledgeBase, } from "workglow"; import { isDryRun } from "../cli/isDryRun"; @@ -24,6 +23,7 @@ import { SecCliConfigurationError } from "../config/EnvToDI"; import { secEmbeddingDimensions, secEmbeddingModel } from "../config/models"; import { SEC_DB_TYPE } from "../config/tokens"; import { getDb } from "../util/db"; +import { PagedChunkVectorStorage } from "./PagedChunkVectorStorage"; import { KB_CHUNK_TABLE, KB_DOCUMENT_TABLE, @@ -163,7 +163,9 @@ export async function getSecKnowledgeBase(): Promise { DocumentStorageSchema, DocumentStorageKey ); - const chunks = new SqliteVectorStorage( + // Paged rather than the base class, whose search reads the whole table to + // score one query — see {@link PagedChunkVectorStorage}. + const chunks = new PagedChunkVectorStorage( db, KB_CHUNK_TABLE, ChunkVectorStorageSchema, diff --git a/src/kb/secKnowledgeBaseSearch.test.ts b/src/kb/secKnowledgeBaseSearch.test.ts new file mode 100644 index 00000000..13b4019e --- /dev/null +++ b/src/kb/secKnowledgeBaseSearch.test.ts @@ -0,0 +1,101 @@ +/** + * @license + * Copyright 2026 Steven Roussey + * SPDX-License-Identifier: Apache-2.0 + */ + +import { afterEach, beforeEach, describe, expect, it, vi } from "vitest"; +import { ChunkVectorPrimaryKey, ChunkVectorStorageSchema, SqliteVectorStorage } from "workglow"; +import { withSqliteDb } from "../config/testing/withSqliteDb"; +import { getDb } from "../util/db"; +import { KB_CHUNK_TABLE } from "./secKbTables"; +import { getSecKnowledgeBase, resetSecKnowledgeBaseForTesting } from "./secKnowledgeBase"; + +/** Narrow enough to seed a few thousand rows without a real embedding model. */ +const DIMENSIONS = 4; +const CHUNKS = 1200; + +/** + * Similarity to `[1, 0, 0, 0]` is `1 / sqrt(1 + spread²)`, so the best chunks + * are the ones with the smallest spread — and the seeding puts those LAST in + * `chunk_id` order, where a scan that stops early would never reach them. + */ +function seedVector(index: number): Float32Array { + return new Float32Array([1, CHUNKS - 1 - index, 0, 0]); +} + +function chunkId(index: number): string { + return `c${String(index).padStart(5, "0")}`; +} + +/** + * `sec index` is documented as a build that "takes hours to days", and `sec + * ask` tells the user to run it. The query side has to survive what that + * builds: the inherited `SqliteVectorStorage.similaritySearch` is + * `SELECT * FROM kb_chunk` with no bound, followed by a JS cosine per row, so + * every question hydrated the entire index — text and a JSON-encoded vector per + * row — before scoring one query. + */ +describe("the SEC knowledge base's chunk search", () => { + withSqliteDb("kb_search", []); + + beforeEach(async () => { + await resetSecKnowledgeBaseForTesting(); + process.env.SEC_EMBEDDING_MODEL = "onnx:test/fixture-encoder"; + process.env.SEC_EMBEDDING_DIMENSIONS = String(DIMENSIONS); + }); + + afterEach(async () => { + await resetSecKnowledgeBaseForTesting(); + delete process.env.SEC_EMBEDDING_MODEL; + delete process.env.SEC_EMBEDDING_DIMENSIONS; + }); + + it("reads kb_chunk in bounded pages and still ranks the whole index", async () => { + // Opening the base creates the three tables; the seeding then writes + // through a plain vector storage onto the very same table. + const kb = await getSecKnowledgeBase(); + const seeder = new SqliteVectorStorage( + getDb(), + KB_CHUNK_TABLE, + ChunkVectorStorageSchema, + ChunkVectorPrimaryKey, + [], + DIMENSIONS + ); + await seeder.putBulk( + Array.from({ length: CHUNKS }, (_unused, index) => ({ + chunk_id: chunkId(index), + doc_id: "0000320193-26-000001:primary.htm", + vector: seedVector(index), + metadata: { text: `chunk ${index}` }, + })) as never + ); + + const db = getDb(); + const prepared: string[] = []; + const realPrepare = db.prepare.bind(db); + vi.spyOn(db, "prepare").mockImplementation(((sql: string) => { + prepared.push(sql); + return realPrepare(sql); + }) as never); + + const hits = await kb.similaritySearch(new Float32Array([1, 0, 0, 0]), { topK: 3 }); + + // The whole index is still ranked — the three best chunks are the last + // three rows, which only a scan that reaches the end can find. + expect(hits.map((hit) => (hit as { chunk_id: string }).chunk_id)).toEqual([ + chunkId(CHUNKS - 1), + chunkId(CHUNKS - 2), + chunkId(CHUNKS - 3), + ]); + + const chunkReads = prepared.filter((sql) => + new RegExp(`FROM \`?${KB_CHUNK_TABLE}\``).test(sql) + ); + expect(chunkReads.length).toBeGreaterThan(1); + // Not one of them may be the unbounded read: a corpus of any size is then + // in the heap at once, every row hydrated before a single score is taken. + for (const sql of chunkReads) expect(sql).toMatch(/LIMIT/); + }); +}); diff --git a/src/task/kb/IndexFilingSectionsTask.test.ts b/src/task/kb/IndexFilingSectionsTask.test.ts index dd6254b3..823996f1 100644 --- a/src/task/kb/IndexFilingSectionsTask.test.ts +++ b/src/task/kb/IndexFilingSectionsTask.test.ts @@ -4,10 +4,11 @@ * SPDX-License-Identifier: Apache-2.0 */ -import { afterEach, describe, expect, it, vi } from "vitest"; +import { afterEach, beforeEach, describe, expect, it, vi } from "vitest"; import type { DocumentNode } from "workglow"; import { Document, globalServiceRegistry, NodeKind } from "workglow"; import { withSqliteDb } from "../../config/testing/withSqliteDb"; +import { SEC_DRY_RUN } from "../../config/tokens"; import { getSecKnowledgeBase, resetSecKnowledgeBaseForTesting } from "../../kb/secKnowledgeBase"; import { FILING_DOCUMENT_REPOSITORY_TOKEN, @@ -124,3 +125,67 @@ describe("IndexFilingSectionsTask selection", () => { expect(out).toMatchObject({ indexed: 0, sections: 0, skipped: 0, truncated: false }); }); }); + +/** + * `--dry-run` promises to change nothing, and `runCommand` prints that promise + * before the task runs. The knowledge base's three storages are built directly + * against `getDb()` rather than through `createStorage`, so no + * `ReadOnlyTabularStorage` wrapper stands between this task and a real write — + * and the guard that recognised that covered only the DDL and the `kb_index` + * row, not the documents and chunk vectors the ingest lands. + */ +describe("IndexFilingSectionsTask under --dry-run", () => { + withSqliteDb("kb_index_dry", [FILING_DOCUMENT_REPOSITORY_TOKEN, FILING_SECTION_REPOSITORY_TOKEN]); + + beforeEach(async () => { + await resetSecKnowledgeBaseForTesting(); + const documents = globalServiceRegistry.get(FILING_DOCUMENT_REPOSITORY_TOKEN); + await documents.put(header(0, "2026-01-02")); + const sections = globalServiceRegistry.get(FILING_SECTION_REPOSITORY_TOKEN); + await sections.put({ + cik: 320193, + accession_number: "0000320193-26-000000", + doc_file: "primary.htm", + ordinal: 0, + slug: "risk-factors", + title: "Risk Factors", + depth: 1, + char_count: 24, + markdown: "# Risk Factors\n\nProse.", + }); + }); + + afterEach(async () => { + globalServiceRegistry.registerInstance(SEC_DRY_RUN, false); + await resetSecKnowledgeBaseForTesting(); + }); + + it("embeds and persists nothing", async () => { + // Opened first, so the tables exist: `getSecKnowledgeBase`'s own dry-run + // refusal covers the index that does not exist yet, and an already-indexed + // database is the case it lets through. + const kb = await getSecKnowledgeBase(); + const upsert = vi + .spyOn(kb, "upsert") + .mockResolvedValue({ doc_id: "0000320193-26-000000:primary.htm" } as never); + + globalServiceRegistry.registerInstance(SEC_DRY_RUN, true); + const out = await new IndexFilingSectionsTask().run({}); + + expect(upsert).not.toHaveBeenCalled(); + expect(out).toMatchObject({ success: true, indexed: 0, sections: 0 }); + }); + + it("still indexes when this is not a dry run", async () => { + const kb = await getSecKnowledgeBase(); + const upsert = vi + .spyOn(kb, "upsert") + .mockResolvedValue({ doc_id: "0000320193-26-000000:primary.htm" } as never); + + globalServiceRegistry.registerInstance(SEC_DRY_RUN, false); + const out = await new IndexFilingSectionsTask().run({}); + + expect(upsert).toHaveBeenCalledTimes(1); + expect(out).toMatchObject({ success: true, indexed: 1, sections: 1 }); + }); +}); diff --git a/src/task/kb/IndexFilingSectionsTask.ts b/src/task/kb/IndexFilingSectionsTask.ts index e31f4051..519136fb 100644 --- a/src/task/kb/IndexFilingSectionsTask.ts +++ b/src/task/kb/IndexFilingSectionsTask.ts @@ -14,6 +14,7 @@ import { Task, TaskAbortedError, } from "workglow"; +import { isDryRun } from "../../cli/isDryRun"; import { getSecKnowledgeBase } from "../../kb/secKnowledgeBase"; import type { FilingDocument } from "../../storage/document/FilingDocumentSchema"; import { @@ -177,6 +178,15 @@ export class IndexFilingSectionsTask extends Task< const skipped = input.force === true ? 0 : await countAlreadyIndexed(scope); const denominator = Math.max(1, work.length); + // The knowledge base's storages are built against `getDb()` rather than + // through `createStorage`, so no `ReadOnlyTabularStorage` wrapper stands + // between an upsert here and a committed row. The selection above is a + // read, so it still reports what the run would do. + if (isDryRun()) { + console.log(`Would index ${work.length} filing(s) into the knowledge base.`); + return { success: true, indexed: 0, sections: 0, skipped, truncated }; + } + let indexed = 0; let sectionTotal = 0; for (const header of work) {