diff --git a/src/config/kbIndexedStamp.sqlite.test.ts b/src/config/kbIndexedStamp.sqlite.test.ts index 819ac810..5d4d1bb2 100644 --- a/src/config/kbIndexedStamp.sqlite.test.ts +++ b/src/config/kbIndexedStamp.sqlite.test.ts @@ -7,6 +7,7 @@ import { afterEach, beforeEach, describe, expect, it } from "vitest"; import type { DocumentNode } from "workglow"; import { Document, globalServiceRegistry, NodeKind } from "workglow"; +import { KB_DOCUMENT_TABLE, SEC_KB_TABLE_NAMES } from "../kb/secKbTables"; import { getSecKnowledgeBase, resetSecKnowledgeBaseForTesting } from "../kb/secKnowledgeBase"; import { FILING_DOCUMENT_REPOSITORY_TOKEN, @@ -17,6 +18,7 @@ import { selectDocumentsToIndex } from "../task/kb/selectDocumentsToIndex"; import { getDb } from "../util/db"; import { syncKbIndexedStamp } from "./kbIndexedStamp"; import { withSqliteDb } from "./testing/withSqliteDb"; +import { SEC_DRY_RUN } from "./tokens"; const accession = (index: number) => `0000320193-26-${String(index).padStart(6, "0")}`; @@ -40,10 +42,11 @@ const doc = (index: number, over: Partial = {}): FilingDocument }); /** - * The stamp is a cache of "is there a `kb_document` row for this", and the three + * The stamp is a cache of "is there a `kb_document` row for this", and the * properties below are what make it safe to read: it never claims more than the - * anti-join, `db setup` never leaves it empty, and the steady state stops - * touching the table. + * anti-join, `db setup` never leaves it empty, a stamp the knowledge base stops + * backing is cleared rather than left to hide the document, and the steady state + * stops touching the table. */ describe("kb_indexed_at stamp (sqlite)", () => { withSqliteDb("kb_indexed_stamp", [FILING_DOCUMENT_REPOSITORY_TOKEN]); @@ -77,6 +80,32 @@ describe("kb_indexed_at stamp (sqlite)", () => { .get(accession(index)) as { s: string | null } | undefined )?.s ?? null; + const selectionPlan = (): string[] => + ( + getDb() + .prepare( + `EXPLAIN QUERY PLAN + SELECT d.* FROM \`filing_document\` d + WHERE d.\`section_count\` > 0 AND d.\`kb_indexed_at\` IS NULL + ORDER BY d.\`filing_date\` DESC, d.\`accession_number\` DESC` + ) + .all() as Array<{ detail?: string }> + ).map((row) => row.detail ?? ""); + + const selectionPlanDetail = (): string => selectionPlan().join(" | "); + + /** A plan step that walks `filing_document` itself rather than an index. */ + const scansTheTable = (): boolean => + selectionPlan().some((step) => /^SCAN\b/.test(step) && !/\bUSING\b.*\bINDEX\b/.test(step)); + + const indexExists = (name: string): boolean => + getDb().prepare("SELECT name FROM sqlite_master WHERE type = 'index' AND name = ?").all(name) + .length > 0; + + const dropKbTables = (): void => { + for (const table of SEC_KB_TABLE_NAMES) getDb().exec(`DROP TABLE \`${table}\``); + }; + it("backfills a database whose documents were indexed before the column existed", async () => { // The state every existing deployment is in: rows in the knowledge base, // no stamp. Left un-backfilled the partial index would cover the whole @@ -108,21 +137,33 @@ describe("kb_indexed_at stamp (sqlite)", () => { for (const i of [1, 2, 3]) await putInKb(i); syncKbIndexedStamp(getDb()); - const plan = getDb() - .prepare( - `EXPLAIN QUERY PLAN - SELECT d.* FROM \`filing_document\` d - WHERE d.\`section_count\` > 0 AND d.\`kb_indexed_at\` IS NULL - ORDER BY d.\`filing_date\` DESC, d.\`accession_number\` DESC` - ) - .all() as Array<{ detail?: string }>; - const detail = plan.map((row) => row.detail ?? "").join(" | "); - - expect(detail).toContain("filing_document_kb_unindexed"); - // The point of the partial index: no table scan, and no temp B-tree to - // sort what a scan would have produced. - expect(detail).not.toContain("SCAN d\n"); - expect(detail).not.toContain("TEMP B-TREE"); + // One access path and it is the partial index — which on a fully indexed + // corpus holds nothing, so the selection reads no rows at all. Spelled as + // the whole plan rather than a substring: SQLite reports a walk of the + // partial index as `SCAN d USING INDEX …` too, so a check for the absence + // of "SCAN" says nothing about which index is being walked. + const plan = selectionPlan(); + expect(plan).toHaveLength(1); + expect(plan[0]).toMatch(/\bUSING\b.*\bINDEX filing_document_kb_unindexed\b/); + // No step reaches the table itself, and none sorts into a temp B-tree. + expect(scansTheTable()).toBe(false); + expect(selectionPlanDetail()).not.toContain("TEMP B-TREE"); + }); + + it("falls back to an index over every row once the partial index is dropped", async () => { + // What gives the assertions above their teeth, and names what they buy. + // Without the partial index the planner still avoids a table scan — the + // schema carries a `filing_date, accession_number` index that serves the + // ORDER BY — so the cost of losing it is not a scan but an index that + // covers the whole corpus instead of only its unindexed tail. + await seed(3); + for (const i of [1, 2, 3]) await putInKb(i); + syncKbIndexedStamp(getDb()); + + getDb().exec("DROP INDEX `filing_document_kb_unindexed`"); + + expect(selectionPlanDetail()).not.toContain("filing_document_kb_unindexed"); + expect(selectionPlanDetail()).toContain("filing_document_filing_date_accession_number"); }); it("never widens the selection: an unstamped document already in the kb is still skipped", async () => { @@ -139,4 +180,95 @@ describe("kb_indexed_at stamp (sqlite)", () => { expect(picked.map((d) => d.accession_number)).toEqual([accession(2)]); }); + + it("clears a stamp the knowledge base no longer backs", async () => { + // The direction that costs a document rather than a traversal: the + // selection reads the stamp as a narrowing of its anti-join, so a stamp + // outliving its `kb_document` row hides a filing that is no longer indexed. + await seed(2); + await putInKb(1); + await putInKb(2); + syncKbIndexedStamp(getDb()); + expect(stampOf(2)).not.toBeNull(); + + getDb() + .prepare(`DELETE FROM \`${KB_DOCUMENT_TABLE}\` WHERE \`doc_id\` = ?`) + .run(kbDocIdFor(accession(2), "primary.htm")); + syncKbIndexedStamp(getDb()); + + expect(stampOf(1)).not.toBeNull(); + expect(stampOf(2)).toBeNull(); + + const picked = await selectDocumentsToIndex({ limit: 10 }); + expect(picked.map((d) => d.accession_number)).toEqual([accession(2)]); + }); + + it("clears every stamp when the knowledge base tables are gone", async () => { + await seed(2); + await putInKb(1); + syncKbIndexedStamp(getDb()); + expect(stampOf(1)).not.toBeNull(); + + await resetSecKnowledgeBaseForTesting(); + dropKbTables(); + syncKbIndexedStamp(getDb()); + + expect(stampOf(1)).toBeNull(); + }); + + it("re-offers every document after the index is dropped and rebuilt empty", async () => { + // The documented way to change the embedding model: drop the three tables + // and re-run `sec index`. Opening the index re-creates `kb_document` empty, + // so the anti-join has nothing to exclude — and the stamps the previous + // index left behind must not exclude anything either, or the rebuild + // indexes nothing and `sec ask` answers from an empty index. + await seed(3); + for (const i of [1, 2, 3]) await putInKb(i); + syncKbIndexedStamp(getDb()); + expect(stampOf(1)).not.toBeNull(); + + await resetSecKnowledgeBaseForTesting(); + dropKbTables(); + await getSecKnowledgeBase(); + + const picked = await selectDocumentsToIndex({ limit: 10 }); + + expect(picked.map((d) => d.accession_number)).toEqual([ + accession(3), + accession(2), + accession(1), + ]); + }); + + it("leaves the stamps alone when the index still holds documents", async () => { + // The other half of the rule above: re-opening a populated index must not + // throw away the stamps, which would put the whole corpus back in front of + // every run. + await seed(2); + await putInKb(1); + syncKbIndexedStamp(getDb()); + const stamp = stampOf(1); + + await resetSecKnowledgeBaseForTesting(); + await getSecKnowledgeBase(); + + expect(stampOf(1)).toBe(stamp); + }); + + it("writes nothing under a dry run", async () => { + // Raw SQL goes around the ReadOnlyTabularStorage wrapper `--dry-run` + // installs, so the bail lives in the function rather than at its callers. + await seed(1); + await putInKb(1); + + globalServiceRegistry.registerInstance(SEC_DRY_RUN, true); + try { + syncKbIndexedStamp(getDb()); + } finally { + globalServiceRegistry.registerInstance(SEC_DRY_RUN, false); + } + + expect(stampOf(1)).toBeNull(); + expect(indexExists("filing_document_kb_unindexed")).toBe(false); + }); }); diff --git a/src/config/kbIndexedStamp.ts b/src/config/kbIndexedStamp.ts index 15a222a6..bc5a6d73 100644 --- a/src/config/kbIndexedStamp.ts +++ b/src/config/kbIndexedStamp.ts @@ -5,16 +5,18 @@ */ import { Sqlite } from "workglow"; +import { isDryRun } from "../cli/isDryRun"; import { KB_DOCUMENT_TABLE } from "../kb/secKbTables"; /** The partial index the steady-state selection walks. */ const PARTIAL_INDEX = "filing_document_kb_unindexed"; /** - * Brings `filing_document.kb_indexed_at` up to date and indexes the nulls. + * Brings `filing_document.kb_indexed_at` into step with the knowledge base and + * indexes the nulls. * - * Two statements that have to travel together, which is why they are one - * function rather than a registry entry. The column is a cache of "is there a + * Statements that have to travel together, which is why they are one function + * rather than a registry entry. The column is a cache of "is there a * `kb_document` row for this", and a cache that is added empty is worse than no * cache: every row reads as unindexed, so the partial index below covers the * whole table and the planner walks it in `filing_date` order, issuing a random @@ -22,8 +24,17 @@ const PARTIAL_INDEX = "filing_document_kb_unindexed"; * it replaced. Backfilling in the same pass that creates the index is what * keeps that state from existing. * - * Both statements are idempotent and cheap on a database already in step: the - * `UPDATE` matches nothing once every indexed document is stamped, and + * The cache moves in both directions, and the two directions cost differently. + * A missing stamp costs a traversal: the selection anti-joins `kb_document` + * anyway, so the document is found and skipped. A stamp the knowledge base no + * longer backs costs the document itself — the selection reads the stamp as a + * narrowing of that anti-join, so a filing whose `kb_document` row has gone is + * never offered again, and a knowledge base rebuilt from empty stays empty. + * So a row the knowledge base has gained is stamped, and one it no longer + * holds is cleared, in the same pass. + * + * Every statement is idempotent and cheap on a database already in step: both + * `UPDATE`s match nothing once the stamps and the knowledge base agree, and * `CREATE INDEX IF NOT EXISTS` is a catalog read. * * Callers must have applied the schema's columns first — the column is added @@ -31,10 +42,16 @@ const PARTIAL_INDEX = "filing_document_kb_unindexed"; * rather than assuming an order it cannot see. */ export function syncKbIndexedStamp(db: Sqlite.Database): void { + // Raw SQL reaches around the repositories' ReadOnlyTabularStorage wrapper, so + // the bail lives here rather than at each call site, where a future caller + // would have to remember it. + if (isDryRun()) return; if (!hasColumn(db, "filing_document", "kb_indexed_at")) return; if (!tableExists(db, KB_DOCUMENT_TABLE)) { - // No knowledge base yet, so nothing is indexed and every row's null is - // already the truth. The index still pays off on the first `ask`. + // No knowledge base, so nothing is indexed: a null is already the truth, + // and any stamp left by one that used to exist is not. The index still + // pays off on the first `ask`. + clearStampsWithNoKbRow(db); createPartialIndex(db); return; } @@ -52,10 +69,60 @@ export function syncKbIndexedStamp(db: Sqlite.Database): void { || ':' || \`filing_document\`.\`doc_file\` )` ).run(); + clearStampsWithNoKbRow(db); createPartialIndex(db); } +/** + * Clears every stamp when the knowledge base holds no documents at all. + * + * For the path that opens the index before filling it, which is where a rebuild + * starts: dropping the knowledge-base tables and re-running `sec index` + * re-creates them empty, and the stamps the previous index left in + * `filing_document` outlive it. Read as a narrowing of the anti-join, those + * stamps exclude the whole corpus from the rebuild, and the run reports nothing + * to do. + * + * Cheap enough to run on every open, which the full reconciliation + * {@link syncKbIndexedStamp} is not: this is one `LIMIT 1` probe wherever the + * index holds anything, which is the steady state, against a probe per stamped + * row. The scan it falls through to happens only where nothing is indexed, so + * it is charged against a run that is about to embed a corpus. + */ +export function clearKbIndexedStampIfIndexEmpty(db: Sqlite.Database): void { + if (isDryRun()) return; + if (!hasColumn(db, "filing_document", "kb_indexed_at")) return; + if (tableExists(db, KB_DOCUMENT_TABLE) && kbHoldsAnyDocument(db)) return; + db.prepare( + "UPDATE `filing_document` SET `kb_indexed_at` = NULL WHERE `kb_indexed_at` IS NOT NULL" + ).run(); +} + +/** + * Drops the stamp from every document the knowledge base does not hold, which + * is all of them when the table itself is gone. + */ +function clearStampsWithNoKbRow(db: Sqlite.Database): void { + const orphaned = tableExists(db, KB_DOCUMENT_TABLE) + ? `AND NOT EXISTS ( + SELECT 1 FROM \`${KB_DOCUMENT_TABLE}\` k + WHERE k.\`doc_id\` = \`filing_document\`.\`accession_number\` + || ':' || \`filing_document\`.\`doc_file\` + )` + : ""; + db.prepare( + `UPDATE \`filing_document\` + SET \`kb_indexed_at\` = NULL + WHERE \`kb_indexed_at\` IS NOT NULL + ${orphaned}` + ).run(); +} + +function kbHoldsAnyDocument(db: Sqlite.Database): boolean { + return db.prepare(`SELECT 1 AS present FROM \`${KB_DOCUMENT_TABLE}\` LIMIT 1`).all().length > 0; +} + /** * Indexed on the order the selection reads — newest first, ties by accession — * so the same walk serves both the filter and the `ORDER BY`. Partial, because diff --git a/src/kb/secKnowledgeBase.ts b/src/kb/secKnowledgeBase.ts index 4fc4ef03..e219bfed 100644 --- a/src/kb/secKnowledgeBase.ts +++ b/src/kb/secKnowledgeBase.ts @@ -20,6 +20,7 @@ import { } from "workglow"; import { isDryRun } from "../cli/isDryRun"; import { SecCliConfigurationError } from "../config/EnvToDI"; +import { clearKbIndexedStampIfIndexEmpty } from "../config/kbIndexedStamp"; import { secEmbeddingDimensions, secEmbeddingModel } from "../config/models"; import { SEC_DB_TYPE } from "../config/tokens"; import { getDb } from "../util/db"; @@ -192,6 +193,12 @@ export async function getSecKnowledgeBase(): Promise { await documents.setupDatabase(); await chunks.setupDatabase(); await index.setupDatabase(); + // An index holding no documents contradicts any `kb_indexed_at` stamp a + // previous one left behind — dropping these tables is how the embedding + // model is changed, and the DDL above re-creates them empty. The selection + // reads that stamp as a narrowing of its anti-join, so left standing it + // would keep every previously indexed filing out of the rebuild. + clearKbIndexedStampIfIndexEmpty(db); } const model = secEmbeddingModel(); diff --git a/src/task/kb/selectDocumentsToIndex.ts b/src/task/kb/selectDocumentsToIndex.ts index 199bff9f..71c479fc 100644 --- a/src/task/kb/selectDocumentsToIndex.ts +++ b/src/task/kb/selectDocumentsToIndex.ts @@ -126,6 +126,12 @@ export async function selectDocumentsToIndex( // Conditional on the column existing because `db setup` backfills it in the // same pass that creates the index, so its presence is what says the values // can be trusted. Where it is absent the query is exactly what it was. + // + // What keeps it a narrowing and not a filter of its own is that the stamp + // is cleared wherever the knowledge base stops backing it — by `db setup` + // per row, and by opening an index that holds nothing. A stamp outliving + // its `kb_document` row would exclude a document the anti-join beside it + // selects, which is the one direction this clause must never take. if (kbIndexedStampAvailable(getDb())) clauses.push("d.`kb_indexed_at` IS NULL"); } // SQLite numbers `?` by position, so the limit binds last because it is