mirror of
https://github.com/zotero/zotero.git
synced 2026-10-08 03:08:19 +00:00
maintain a ledger of chunk counts per attachment
This commit is contained in:
Bogdan Abaev
parent
2a1d4fe339
commit
9bb515c489
2 changed files with 222 additions and 106 deletions
|
|
@ -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));
|
||||
|
|
|
|||
|
|
@ -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());
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue