From 9bb515c489fc38309d99346d01818d3c119a6a3d Mon Sep 17 00:00:00 2001 From: Bogdan Abaev Date: Fri, 4 Sep 2026 14:14:35 -0700 Subject: [PATCH] maintain a ledger of chunk counts per attachment --- chrome/content/zotero/xpcom/embeddings.js | 223 ++++++++++++---------- test/tests/embeddingsTest.js | 105 +++++++++- 2 files changed, 222 insertions(+), 106 deletions(-) diff --git a/chrome/content/zotero/xpcom/embeddings.js b/chrome/content/zotero/xpcom/embeddings.js index f5aaf12a5c..c92403cdb0 100644 --- a/chrome/content/zotero/xpcom/embeddings.js +++ b/chrome/content/zotero/xpcom/embeddings.js @@ -287,7 +287,7 @@ Zotero.Embeddings = new function () { // Schema version of the attached embeddings database. The tables are only // created when this is bumped (_setUpDB() drops and recreates everything), // so any schema change needs a bump. - const _dbVersion = 3; + const _dbVersion = 4; let _dbInitPromise = null; let _dbHooksRegistered = false; @@ -347,11 +347,8 @@ Zotero.Embeddings = new function () { } async function _setUpDB() { - // Scoring runs on sqlite-vec's vector functions (see - // _scoreExpression()) rather than in JS, because mozStorage hands a - // vector to JS as an array of one number per byte -- a library's worth - // of those blocks the main thread for seconds. Loaded once here; - // DBConnection re-loads it after a reconnect, before this callback. + // Scoring uses sqlite-vec's vector functions: mozStorage hands a + // vector to JS one number per byte, far too slow for a whole library await Zotero.DB.loadExtension('vec'); // Idempotent, since it can run again for a retried initialization or // after a connection reopen @@ -361,87 +358,76 @@ Zotero.Embeddings = new function () { let path = Zotero.DataDirectory.getDatabase('embeddings'); await Zotero.DB.queryAsync("ATTACH DATABASE ? AS embeddings", [path]); } - // The embeddings are keyed by local itemID, which is reassigned - // whenever zotero.sqlite is recreated (e.g., deleted and re-synced - // from the server). Vectors stored against a different database - // instance would map to the wrong items, so they have to be discarded - // rather than reused. Detect that by comparing the localUserKey the - // database was stamped with against the current one. + // itemIDs are reassigned when zotero.sqlite is recreated, so vectors + // stamped with a different localUserKey belong to other items let localUserKey = Zotero.Users.getLocalUserKey(); let version = await Zotero.DB.valueQueryAsync("PRAGMA embeddings.user_version"); let storedUserKey = version >= _dbVersion ? await Zotero.DB.valueQueryAsync( "SELECT value FROM embeddings.itemEmbeddingsMeta WHERE key='localUserKey'") : false; - if (version < _dbVersion || storedUserKey != localUserKey) { - await Zotero.DB.queryAsync("DROP TABLE IF EXISTS embeddings.itemEmbeddings"); - await Zotero.DB.queryAsync("DROP TABLE IF EXISTS embeddings.itemEmbeddingsMeta"); - await Zotero.DB.queryAsync("DROP TABLE IF EXISTS embeddings.modelCalibration"); - // No foreign key on itemID -- references across attached databases - // aren't possible, so item deletions are handled by the indexing - // notifier and eligibility pruning instead. - // An item's text is stored as one or more chunks (see - // Zotero.Embeddings.Chunking); - // every chunk row carries the hash of the item's full source text, - // and scoring takes the item's best chunk. - // An attachment fulltext chunk also records where its text came - // from, rather than the text itself: the top-level block range it - // covers (see Zotero.SDT.getBlockRanges()) with character offsets - // into the first and last block's text -- or, for a chunk of flat - // fallback text with no block structure, NULL blocks and offsets - // into the attachment's plain text. The preview is re-derived - // from those references and verified against textCheck (see - // getMatchingChunks()). sectionPart of sectionParts says which - // piece of a split section the chunk is. All NULL for chunks of - // other item types, which are their own preview and location. - // An attachment that yields no text at all (missing file, - // password-protected, no text layer) gets a single row with a - // NULL embedding: a record that it was processed, so progress - // counts it and later passes skip it (via sourceHash) until the - // file changes. Scoring reads only rows with an embedding. - await Zotero.DB.queryAsync( - "CREATE TABLE embeddings.itemEmbeddings (\n" - + " itemID INTEGER NOT NULL,\n" - + " chunkIndex INTEGER NOT NULL,\n" - + " embedding BLOB,\n" - + " sourceHash TEXT NOT NULL,\n" - + " startBlock INTEGER,\n" - + " endBlock INTEGER,\n" - + " startOffset INTEGER,\n" - + " endOffset INTEGER,\n" - + " textCheck TEXT,\n" - + " sectionPart INTEGER,\n" - + " sectionParts INTEGER,\n" - + " PRIMARY KEY (itemID, chunkIndex)\n" - + ")" - ); - // Database metadata: the localUserKey the vectors were built - // against (above) and the identity of the model that produced them - // (see Indexing._ensureIndexMatchesModel()) - await Zotero.DB.queryAsync( - "CREATE TABLE embeddings.itemEmbeddingsMeta (\n" - + " key TEXT PRIMARY KEY,\n" - + " value NOT NULL\n" - + ")" - ); - // What running the model taught us about it, measured once per - // model version (see Zotero.Embeddings.ensureCalibration()). Keyed - // by version rather than name, so a `revision` bump measures again - // alongside the reindex it already forces. - await Zotero.DB.queryAsync( - "CREATE TABLE embeddings.modelCalibration (\n" - + " modelVersion TEXT PRIMARY KEY,\n" - + " meanVector BLOB NOT NULL,\n" - + " minScore REAL NOT NULL,\n" - + " maxDisplayScore REAL NOT NULL\n" - + ")" - ); - await Zotero.DB.queryAsync( - "REPLACE INTO embeddings.itemEmbeddingsMeta (key, value) VALUES ('localUserKey', ?)", - [localUserKey] - ); - await Zotero.DB.queryAsync("PRAGMA embeddings.user_version = " + _dbVersion); + if (version >= _dbVersion && storedUserKey == localUserKey) { + return; } + await Zotero.DB.queryAsync("DROP TABLE IF EXISTS embeddings.itemEmbeddings"); + await Zotero.DB.queryAsync("DROP TABLE IF EXISTS embeddings.itemEmbeddingsMeta"); + await Zotero.DB.queryAsync("DROP TABLE IF EXISTS embeddings.modelCalibration"); + await Zotero.DB.queryAsync("DROP TABLE IF EXISTS embeddings.itemChunkCounts"); + // One row per chunk of an item's text, each carrying the hash of the + // item's full source. For attachments, the block range and offsets + // locate the chunk's text in the document (see getMatchingChunks()), + // and one with no text gets a single NULL-embedding row so it counts + // as processed. No foreign key: deletions are handled by the notifier + // and eligibility pruning. + await Zotero.DB.queryAsync( + "CREATE TABLE embeddings.itemEmbeddings (\n" + + " itemID INTEGER NOT NULL,\n" + + " chunkIndex INTEGER NOT NULL,\n" + + " embedding BLOB,\n" + + " sourceHash TEXT NOT NULL,\n" + + " startBlock INTEGER,\n" + + " endBlock INTEGER,\n" + + " startOffset INTEGER,\n" + + " endOffset INTEGER,\n" + + " textCheck TEXT,\n" + + " sectionPart INTEGER,\n" + + " sectionParts INTEGER,\n" + + " PRIMARY KEY (itemID, chunkIndex)\n" + + ")" + ); + // The localUserKey the vectors were built against and the model that + // produced them + await Zotero.DB.queryAsync( + "CREATE TABLE embeddings.itemEmbeddingsMeta (\n" + + " key TEXT PRIMARY KEY,\n" + + " value NOT NULL\n" + + ")" + ); + // Per-model measurements (see ensureCalibration()), keyed by version + // so a revision bump measures again + await Zotero.DB.queryAsync( + "CREATE TABLE embeddings.modelCalibration (\n" + + " modelVersion TEXT PRIMARY KEY,\n" + + " meanVector BLOB NOT NULL,\n" + + " minScore REAL NOT NULL,\n" + + " maxDisplayScore REAL NOT NULL\n" + + ")" + ); + // How many chunks each attachment's current source splits into, + // recorded at extraction, so fulltext work is known without reading + // a pack. Other item types count as one apiece. + await Zotero.DB.queryAsync( + "CREATE TABLE embeddings.itemChunkCounts (\n" + + " itemID INTEGER PRIMARY KEY,\n" + + " sourceHash TEXT NOT NULL,\n" + + " chunks INTEGER NOT NULL\n" + + ")" + ); + await Zotero.DB.queryAsync( + "REPLACE INTO embeddings.itemEmbeddingsMeta (key, value) VALUES ('localUserKey', ?)", + [localUserKey] + ); + await Zotero.DB.queryAsync("PRAGMA embeddings.user_version = " + _dbVersion); } async function _rebuildDB() { @@ -2213,28 +2199,34 @@ Zotero.Embeddings.Indexing = new function () { } } let stored = await Zotero.DB.columnQueryAsync( - "SELECT DISTINCT itemID FROM embeddings.itemEmbeddings" + "SELECT DISTINCT itemID FROM embeddings.itemEmbeddings " + + "UNION SELECT itemID FROM embeddings.itemChunkCounts" ); await _deleteEmbeddings(stored.filter(id => !eligible.has(id))); } - // Delete the stored embeddings for the given items, in chunks (avoids the - // SQLite bound-parameter limit) + // Delete the stored embeddings and chunk counts for the given items, in + // chunks (avoids the SQLite bound-parameter limit) async function _deleteEmbeddings(itemIDs) { await Zotero.Embeddings.initDB(); let chunkSize = 500; for (let i = 0; i < itemIDs.length; i += chunkSize) { let chunk = itemIDs.slice(i, i + chunkSize); + let placeholders = chunk.map(() => '?').join(','); await Zotero.DB.queryAsync( - "DELETE FROM embeddings.itemEmbeddings WHERE itemID IN (" - + chunk.map(() => '?').join(',') + ")", + "DELETE FROM embeddings.itemEmbeddings WHERE itemID IN (" + placeholders + ")", + chunk + ); + await Zotero.DB.queryAsync( + "DELETE FROM embeddings.itemChunkCounts WHERE itemID IN (" + placeholders + ")", chunk ); } } - // Delete all stored item embeddings. This removes the computed vectors, - // not the downloaded model files. + // Delete all stored item embeddings and chunk counts (chunks are sized to + // the model's window, so the counts go with the vectors). This removes + // the computed vectors, not the downloaded model files. async function _clearEmbeddings() { await Zotero.Embeddings.initDB(); // Announce the removals, so active semantic views refresh after the @@ -2244,6 +2236,7 @@ Zotero.Embeddings.Indexing = new function () { "SELECT DISTINCT itemID FROM embeddings.itemEmbeddings" ); await Zotero.DB.queryAsync("DELETE FROM embeddings.itemEmbeddings"); + await Zotero.DB.queryAsync("DELETE FROM embeddings.itemChunkCounts"); if (cleared.length) { _notifyIndexed(cleared); } @@ -2486,18 +2479,18 @@ Zotero.Embeddings.Indexing = new function () { })); } - // The stored source hash of each of the given items that has rows, read - // in one query per chunk rather than one per item (every start - // re-enqueues the whole library to find what changed). Every chunk row - // of an item carries the same hash. - async function _getStoredHashes(itemIDs) { + // The stored source hash of each of the given items that has rows in + // the given table, read in one query per chunk rather than one per item + // (every start re-enqueues the whole library to find what changed). + // Every chunk row of an item carries the same hash. + async function _getStoredHashes(itemIDs, table = 'itemEmbeddings') { let storedHashes = new Map(); let chunkSize = 500; for (let i = 0; i < itemIDs.length; i += chunkSize) { let chunk = itemIDs.slice(i, i + chunkSize); let rows = await Zotero.DB.queryAsync( - "SELECT DISTINCT itemID, sourceHash FROM embeddings.itemEmbeddings WHERE itemID IN (" - + chunk.map(() => '?').join(',') + ")", + "SELECT DISTINCT itemID, sourceHash FROM embeddings." + table + + " WHERE itemID IN (" + chunk.map(() => '?').join(',') + ")", chunk ); for (let row of rows) { @@ -2514,7 +2507,7 @@ Zotero.Embeddings.Indexing = new function () { // the ones the embedding pass will read packs for. An up-to-date // attachment won't be re-embedded and a fileless one has nothing to // extract, so neither needs a pack. - async function _staleAttachmentIDs(itemIDs, shouldStop) { + async function _staleAttachments(itemIDs, shouldStop) { let items = await Zotero.Items.getAsync(itemIDs); let storedHashes = await _getStoredHashes(itemIDs); let stale = []; @@ -2524,12 +2517,22 @@ Zotero.Embeddings.Indexing = new function () { } let hash = await _getAttachmentSourceHash(item); if (hash && storedHashes.get(item.id) !== hash) { - stale.push(item.id); + stale.push({ item, hash }); } } return stale; } + // Record how many chunks an attachment's current source splits into (see + // the itemChunkCounts table in _setUpDB()) + async function _storeChunkCount(itemID, hash, chunks) { + await Zotero.DB.queryAsync( + "REPLACE INTO embeddings.itemChunkCounts (itemID, sourceHash, chunks) " + + "VALUES (?, ?, ?)", + [itemID, hash, chunks] + ); + } + // Embed the regular queue -- items, notes and annotations -- a slice at // a time until it's empty or the run stops async function _drainItemQueue(shouldStop, indexOptions) { @@ -2597,24 +2600,34 @@ Zotero.Embeddings.Indexing = new function () { // model would otherwise sit in memory throughout, competing with the // worker for the same cores. async function _extractAttachments(itemIDs, extracted, shouldStop) { - let toExtract = (await _staleAttachmentIDs(itemIDs, shouldStop)) - .filter(itemID => !extracted.has(itemID)); + let toExtract = (await _staleAttachments(itemIDs, shouldStop)) + .filter(({ item }) => !extracted.has(item.id)); if (!toExtract.length) { return; } _setPhase('extracting'); await Zotero.Embeddings.shutdownEngine({ modelChanged: false }); + // Chunk counts are taken while the pack is fresh, through the same + // derivation the embedder uses so the two can't disagree. One already + // recorded for the current source is kept -- chunking again would + // only repeat it. + let counted = await _getStoredHashes( + toExtract.map(({ item }) => item.id), 'itemChunkCounts'); let progress = { done: 0, total: toExtract.length }; _extractionProgress = progress; _emitProgress(); let lastEmit = Date.now(); try { - for (let itemID of toExtract) { + for (let { item, hash } of toExtract) { if (shouldStop()) { return; } - extracted.add(itemID); - await Zotero.SDT.ensure(itemID); + extracted.add(item.id); + await Zotero.SDT.ensure(item.id); + if (counted.get(item.id) !== hash) { + let chunks = await _getAttachmentChunks(item); + await _storeChunkCount(item.id, hash, chunks ? chunks.length : 0); + } progress.done++; if (Date.now() - lastEmit >= EXTRACTION_EMIT_INTERVAL) { lastEmit = Date.now(); @@ -2755,6 +2768,7 @@ Zotero.Embeddings.Indexing = new function () { + "VALUES (?, 0, NULL, ?)", [entry.item.id, entry.hash] ); + await _storeChunkCount(entry.item.id, entry.hash, 0); } }); } @@ -2884,6 +2898,13 @@ Zotero.Embeddings.Indexing = new function () { { debugParams: false } ); } + // Normally recorded at extraction; written here too + // for an attachment that reached embedding without + // passing through it, so the count always matches + // the rows + if (entry.item.isAttachment()) { + await _storeChunkCount(entry.item.id, entry.hash, entry.chunks.length); + } } }); _notifyIndexed(completed.map(entry => entry.item.id)); diff --git a/test/tests/embeddingsTest.js b/test/tests/embeddingsTest.js index f1f1ea5b65..e4648e5e1a 100644 --- a/test/tests/embeddingsTest.js +++ b/test/tests/embeddingsTest.js @@ -1116,6 +1116,11 @@ describe("Zotero.Embeddings", function () { + "(itemID, chunkIndex, embedding, sourceHash) VALUES (?, 0, ?, ?)", [item.id, new Uint8Array([0, 0, 0, 0]), 'hash'] ); + await Zotero.DB.queryAsync( + "REPLACE INTO embeddings.itemChunkCounts (itemID, sourceHash, chunks) " + + "VALUES (?, ?, 1)", + [item.id, 'hash'] + ); // The model switch clears the old vectors and announces the // removals (after the coalescing delay), so active semantic // views refresh @@ -1129,6 +1134,13 @@ describe("Zotero.Embeddings", function () { ), 0 ); + // Chunks are sized to the model, so the counts go too + assert.equal( + await Zotero.DB.valueQueryAsync( + "SELECT COUNT(*) FROM embeddings.itemChunkCounts" + ), + 0 + ); } finally { Zotero.Prefs.set('embeddings.model', ''); @@ -1564,16 +1576,91 @@ describe("Zotero.Embeddings", function () { "SELECT COUNT(*) FROM embeddings.itemEmbeddings WHERE itemID=?", attachment.id )); + assert.ok(await Zotero.DB.valueQueryAsync( + "SELECT COUNT(*) FROM embeddings.itemChunkCounts WHERE itemID=?", + attachment.id + )); // Turning the pref off makes attachments ineligible, and the - // pref observer prunes their stored chunks. The observer runs - // asynchronously, so poll (the test times out on failure). + // pref observer prunes their stored chunks and chunk counts. + // The observer runs asynchronously, so poll (the test times + // out on failure). Zotero.Prefs.set('embeddings.indexFulltext', false); while (await Zotero.DB.valueQueryAsync( "SELECT COUNT(*) FROM embeddings.itemEmbeddings WHERE itemID=?", attachment.id)) { await Zotero.Promise.delay(10); } + assert.equal(await Zotero.DB.valueQueryAsync( + "SELECT COUNT(*) FROM embeddings.itemChunkCounts WHERE itemID=?", + attachment.id + ), 0); + } + finally { + stubs.forEach(stub => stub.restore()); + Zotero.Prefs.clear('embeddings.indexFulltext'); + } + }); + + it("should record an attachment's chunk count at extraction and keep it once embedded", async function () { + this.timeout(60000); + let item = await createDataObject('item', { title: 'Parent of counted attachment' }); + let attachment = await importPDFAttachment(item); + + let vector = new Float32Array(4).fill(0.5); + // Two sections long enough that the chunker keeps them apart + let sections = [ + sdtSection('', 0, ['Owls hunt at night. '.repeat(120)]), + sdtSection('', 1, ['Hawks hunt by day. '.repeat(120)]) + ]; + // Embedding the attachment fails on the first run, so its count + // can only have come from the extraction step + let failAttachment = true; + let stubs = [ + sinon.stub(Zotero.Embeddings, 'embedPassages').callsFake(async (texts) => { + if (failAttachment && texts.some(text => text.includes('hunt'))) { + throw new Error('Embedding failed'); + } + return texts.map(() => vector); + }), + sinon.stub(Zotero.Embeddings, 'isEnabled').returns(true), + sinon.stub(Zotero.Embeddings, 'getModelVersion').returns('test-model/1'), + sinon.stub(Zotero.Embeddings, 'isDownloaded').resolves(true), + sinon.stub(Zotero.Embeddings, 'download').resolves(), + sinon.stub(Zotero.Embeddings, 'ensureCalibration').resolves(), + sinon.stub(Zotero.Embeddings, 'getModelName').returns('bge-small-en-v1.5'), + sinon.stub(Zotero.SDT, 'ensure').resolves(true), + sinon.stub(Zotero.SDT, 'getSections').resolves({ ok: true, sections }) + ]; + try { + Zotero.Prefs.set('embeddings.indexFulltext', true); + await Zotero.Embeddings.Indexing.startIndexing(); + let counted = await Zotero.DB.valueQueryAsync( + "SELECT chunks FROM embeddings.itemChunkCounts WHERE itemID=?", + attachment.id + ); + assert.isAbove(counted, 1); + assert.equal(await Zotero.DB.valueQueryAsync( + "SELECT COUNT(*) FROM embeddings.itemEmbeddings WHERE itemID=?", + attachment.id + ), 0); + // Only attachments are counted + assert.equal(await Zotero.DB.valueQueryAsync( + "SELECT COUNT(*) FROM embeddings.itemChunkCounts WHERE itemID=?", + item.id + ), 0); + + failAttachment = false; + await Zotero.Embeddings.Indexing.startIndexing(); + assert.equal(await Zotero.DB.valueQueryAsync( + "SELECT chunks FROM embeddings.itemChunkCounts WHERE itemID=?", + attachment.id + ), counted); + assert.equal(await Zotero.DB.valueQueryAsync( + "SELECT COUNT(*) FROM embeddings.itemEmbeddings " + + "WHERE itemID=? AND embedding IS NOT NULL", + attachment.id + ), counted); } finally { stubs.forEach(stub => stub.restore()); @@ -1768,7 +1855,9 @@ describe("Zotero.Embeddings", function () { sinon.stub(Zotero.Embeddings, 'getModelName').returns('bge-small-en-v1.5'), sinon.stub(Zotero.SDT, 'ensure').resolves(true), sinon.stub(Zotero.SDT, 'getSections').callsFake(async (itemID) => { - if (attachments.some(a => a.id === itemID)) { + // Sections are read once to count chunks and again to + // embed; the first read is the order of interest + if (attachments.some(a => a.id === itemID) && !extracted.includes(itemID)) { extracted.push(itemID); } return { @@ -2002,7 +2091,13 @@ describe("Zotero.Embeddings", function () { assert.lengthOf(rows, 1); assert.isNull(rows[0].embedding); assert.ok(rows[0].sourceHash); - assert.equal(ourCalls(), 1); + // Sections are read once to count chunks at extraction and + // once to embed + assert.equal(ourCalls(), 2); + assert.equal(await Zotero.DB.valueQueryAsync( + "SELECT chunks FROM embeddings.itemChunkCounts WHERE itemID=?", + attachment.id + ), 0); // A processed-but-empty item can't be scored, and doesn't // break scoring for anything else @@ -2012,7 +2107,7 @@ describe("Zotero.Embeddings", function () { // The record makes later passes skip the attachment without // re-extracting, until the file changes await Zotero.Embeddings.Indexing.startIndexing(); - assert.equal(ourCalls(), 1); + assert.equal(ourCalls(), 2); } finally { stubs.forEach(stub => stub.restore());