From 9dfd30745c00ce34a101a07febab7adf1ae3d23c Mon Sep 17 00:00:00 2001 From: Aaron Van Tassle Date: Sun, 27 Sep 2026 14:55:52 -0600 Subject: [PATCH] fix(rag): seek ZIM batches instead of rescanning from the start --- admin/app/jobs/embed_file_job.ts | 10 ++- admin/app/services/rag_service.ts | 35 +++++++- admin/app/services/zim_extraction_service.ts | 45 ++++++++-- admin/app/utils/zim_batch_resume.ts | 28 ++++++ admin/tests/unit/zim_batch_resume.spec.ts | 90 ++++++++++++++++++++ admin/types/rag.ts | 2 + admin/types/zim.ts | 4 +- 7 files changed, 201 insertions(+), 13 deletions(-) create mode 100644 admin/app/utils/zim_batch_resume.ts create mode 100644 admin/tests/unit/zim_batch_resume.spec.ts diff --git a/admin/app/jobs/embed_file_job.ts b/admin/app/jobs/embed_file_job.ts index b1ae07ba..d694536f 100644 --- a/admin/app/jobs/embed_file_job.ts +++ b/admin/app/jobs/embed_file_job.ts @@ -24,6 +24,9 @@ export interface EmbedFileJobParams { // batch's chunk count was stored while Qdrant held the full set). chunksSoFar?: number collection?: string + // Dirents already finished. Present from the second batch on, so extraction + // can seek instead of recounting batchOffset. + resumeAtDirent?: number } export class EmbedFileJob { @@ -57,7 +60,8 @@ export class EmbedFileJob { } async handle(job: Job) { - const { filePath, fileName, batchOffset, totalArticles, collection } = job.data as EmbedFileJobParams + const { filePath, fileName, batchOffset, totalArticles, collection, resumeAtDirent } = + job.data as EmbedFileJobParams // Only the direct KB-upload controller passes `collection` on dispatch; the other // six dispatch sites (download auto-index, scan/sync, re-embed, local ZIM upload, @@ -147,7 +151,8 @@ export class EmbedFileJob { allowDeletion, batchOffset, onProgress, - effectiveCollection + effectiveCollection, + resumeAtDirent ) if (!result.success) { @@ -204,6 +209,7 @@ export class EmbedFileJob { // Carry the collection across batches, otherwise only batch 1 of a ZIM // would be tagged and the rest would land uncategorized. ...(effectiveCollection ? { collection: effectiveCollection } : {}), + ...(result.resumeAtDirent !== undefined ? { resumeAtDirent: result.resumeAtDirent } : {}), }) // Calculate progress based on articles processed. diff --git a/admin/app/services/rag_service.ts b/admin/app/services/rag_service.ts index 7bf77f52..b2a6f8ba 100644 --- a/admin/app/services/rag_service.ts +++ b/admin/app/services/rag_service.ts @@ -608,7 +608,8 @@ export class RagService { deleteAfterEmbedding: boolean, batchOffset?: number, onProgress?: (percent: number) => Promise, - collection?: string + collection?: string, + resumeAtDirent?: number ): Promise { const zimExtractionService = new ZIMExtractionService() @@ -616,18 +617,35 @@ export class RagService { const startOffset = batchOffset || 0 logger.info( - `[RAG] Extracting ZIM content (batch: offset=${startOffset}, size=${ZIM_BATCH_SIZE})` + `[RAG] Extracting ZIM content (batch: offset=${startOffset}, size=${ZIM_BATCH_SIZE}${resumeAtDirent !== undefined ? ', seek' : ''})` ) const { chunks: zimChunks, totalArticles, articlesProcessed, + resumeAtDirent: nextResumeAtDirent, } = await zimExtractionService.extractZIMContent(filepath, { startOffset, batchSize: ZIM_BATCH_SIZE, + resumeAtDirent, }) + // A seek that returns nothing while articles remain is a bad cursor, not + // the deliberate one-batch overrun at the real end of the archive. That + // overrun happens after the article offset has already passed articleCount. + if ( + resumeAtDirent !== undefined && + articlesProcessed === 0 && + startOffset < totalArticles + ) { + const message = + `ZIM resume seek at dirent ${resumeAtDirent} found no further articles ` + + `(${startOffset} of ${totalArticles} done).` + logger.error(`[RAG] ${message}`) + return { success: false, message } + } + logger.info( `[RAG] Extracted ${zimChunks.length} chunks from ZIM file with enhanced metadata (file totalArticles=${totalArticles})` ) @@ -717,6 +735,7 @@ export class RagService { // zero and re-run the same window forever. articlesProcessed, totalArticles, + resumeAtDirent: nextResumeAtDirent, } } @@ -865,7 +884,8 @@ export class RagService { deleteAfterEmbedding: boolean = false, batchOffset?: number, onProgress?: (percent: number) => Promise, - collection?: string + collection?: string, + resumeAtDirent?: number ): Promise { try { const fileType = determineFileType(filepath) @@ -884,7 +904,14 @@ export class RagService { // Process based on file type // ZIM files are handled specially since they have their own embedding workflow if (fileType === 'zim') { - return await this.processZIMFile(filepath, deleteAfterEmbedding, batchOffset, onProgress, collection) + return await this.processZIMFile( + filepath, + deleteAfterEmbedding, + batchOffset, + onProgress, + collection, + resumeAtDirent + ) } // Extract text based on file type diff --git a/admin/app/services/zim_extraction_service.ts b/admin/app/services/zim_extraction_service.ts index ec307880..db52cfaf 100644 --- a/admin/app/services/zim_extraction_service.ts +++ b/admin/app/services/zim_extraction_service.ts @@ -7,6 +7,12 @@ import { ExtractZIMChunkingStrategy, ExtractZIMContentOptions, ZIMContentChunk, import { randomUUID } from 'node:crypto' import { access } from 'node:fs/promises' import { isValidZimFile } from '../utils/fs.js' +import { hasResumeCursor, openZimBatchRange } from '../utils/zim_batch_resume.js' + +// How often a legacy from-the-start scan yields so BullMQ can renew its lock. +// The seek path does not need this; the scan is only for jobs queued before a +// dirent cursor existed. +const LEGACY_SCAN_YIELD_EVERY = 2000 export class ZIMExtractionService { @@ -44,7 +50,7 @@ export class ZIMExtractionService { async extractZIMContent( filePath: string, opts: ExtractZIMContentOptions = {} - ): Promise<{ chunks: ZIMContentChunk[]; totalArticles: number; articlesProcessed: number }> { + ): Promise<{ chunks: ZIMContentChunk[]; totalArticles: number; articlesProcessed: number; resumeAtDirent?: number }> { try { logger.info(`[ZIMExtractionService]: Processing ZIM file at path: ${filePath}`) @@ -71,16 +77,35 @@ export class ZIMExtractionService { let articlesProcessed = 0 let articlesSkipped = 0 + let legacyEntriesVisited = 0 + let direntsConsumed = 0 const processedPaths = new Set() const toReturn: ZIMContentChunk[] = [] - // Support batch processing to avoid lock timeouts on large ZIM files + // resumeAtDirent seeks with iterByPath().offset(). startOffset recounts + // from the first dirent and is only the fallback when no cursor was stored. const startOffset = opts.startOffset || 0 const batchSize = opts.batchSize || (opts.maxArticles || Infinity) + const resumeAtDirent = hasResumeCursor(opts.resumeAtDirent) ? opts.resumeAtDirent : undefined - for (const entry of archive.iterByPath()) { - // Skip articles until we reach the start offset - if (articlesSkipped < startOffset) { + if (resumeAtDirent !== undefined) { + logger.info(`[ZIMExtractionService]: Seeking to dirent offset ${resumeAtDirent}`) + } else if (startOffset > 0) { + logger.warn( + `[ZIMExtractionService]: Batch offset ${startOffset} has no dirent cursor; scanning from the start of the archive` + ) + } + + const entries = openZimBatchRange(archive, resumeAtDirent) + for (const entry of entries) { + // Legacy recount. Yield periodically: this loop is synchronous + // native calls, and a blocked event loop cannot renew the BullMQ lock. + if (resumeAtDirent === undefined && articlesSkipped < startOffset) { + direntsConsumed++ + legacyEntriesVisited++ + if (legacyEntriesVisited % LEGACY_SCAN_YIELD_EVERY === 0) { + await new Promise((resolve) => setImmediate(resolve)) + } if (this.isArticleEntry(entry) && !processedPaths.has(entry.path)) { articlesSkipped++ } @@ -91,6 +116,8 @@ export class ZIMExtractionService { break } + direntsConsumed++ + if (!this.isArticleEntry(entry)) { logger.debug(`[ZIMExtractionService]: Skipping non-article entry at path: ${entry.path}`) continue @@ -165,7 +192,13 @@ export class ZIMExtractionService { textPreview: c.text.substring(0, 100) }))) logger.debug("Total structured sections extracted:", toReturn.length) - return { chunks: toReturn, totalArticles: archive.articleCount, articlesProcessed } + const nextDirent = (resumeAtDirent ?? 0) + direntsConsumed + return { + chunks: toReturn, + totalArticles: archive.articleCount, + articlesProcessed, + resumeAtDirent: nextDirent > 0 ? nextDirent : undefined, + } } catch (error) { logger.error('Error processing ZIM file:', error) throw error diff --git a/admin/app/utils/zim_batch_resume.ts b/admin/app/utils/zim_batch_resume.ts new file mode 100644 index 00000000..77b0d391 --- /dev/null +++ b/admin/app/utils/zim_batch_resume.ts @@ -0,0 +1,28 @@ +/** + * Resume cursor for batched ZIM extraction. + * + * The cursor is how many dirents `iterByPath()` has already finished. That + * range does not start at `entry.index` (user entries begin at the 'C' + * namespace), and `findByPath` is a prefix search that usually returns only + * the article just embedded. `offset(cursor, size)` seeks within the range. + */ + +export type SeekableEntryRange = Iterable & { + size: number + offset(start: number, maxResults: number): Iterable +} + +export function hasResumeCursor(cursor: number | undefined): cursor is number { + return typeof cursor === 'number' && Number.isInteger(cursor) && cursor > 0 +} + +export function openZimBatchRange( + archive: { iterByPath(): SeekableEntryRange }, + resumeAtDirent?: number +): Iterable { + const range = archive.iterByPath() + if (!hasResumeCursor(resumeAtDirent)) return range + // `size` reaches the end of the range. libzim takes a signed 32-bit + // maxResults, and a real archive's entry count fits in that. + return range.offset(resumeAtDirent, range.size) +} diff --git a/admin/tests/unit/zim_batch_resume.spec.ts b/admin/tests/unit/zim_batch_resume.spec.ts new file mode 100644 index 00000000..5e9eabdd --- /dev/null +++ b/admin/tests/unit/zim_batch_resume.spec.ts @@ -0,0 +1,90 @@ +import * as assert from 'node:assert/strict' +import { test } from 'node:test' + +import { hasResumeCursor, openZimBatchRange } from '../../app/utils/zim_batch_resume.js' + +type FakeEntry = { path: string; article: boolean } + +const directory: FakeEntry[] = [ + { path: 'A.html', article: true }, + { path: 'A_files/img.png', article: false }, + { path: 'B.html', article: true }, + { path: 'C.html', article: true }, + { path: 'C_files/style.css', article: false }, + { path: 'D.html', article: true }, +] + +function fakeArchive(entries: FakeEntry[]) { + const range = { + size: entries.length, + offset(start: number, maxResults: number) { + return entries.slice(start, start + maxResults) + }, + [Symbol.iterator]() { + return entries[Symbol.iterator]() + }, + } + return { iterByPath: () => range } +} + +function takeArticleBatch(entries: Iterable, batchSize: number) { + const accepted: string[] = [] + let direntsConsumed = 0 + for (const entry of entries) { + if (accepted.length >= batchSize) break + direntsConsumed++ + if (!entry.article) continue + accepted.push(entry.path) + } + return { accepted, direntsConsumed } +} + +test('hasResumeCursor rejects an empty or zero cursor', () => { + assert.equal(hasResumeCursor(undefined), false) + assert.equal(hasResumeCursor(0), false) + assert.equal(hasResumeCursor(4), true) + assert.equal(hasResumeCursor(-1), false) + assert.equal(hasResumeCursor(1.5), false) +}) + +test('openZimBatchRange reads from the start when there is no cursor', () => { + const archive = fakeArchive(directory) + assert.deepEqual( + [...openZimBatchRange(archive)].map((entry) => entry.path), + directory.map((entry) => entry.path) + ) +}) + +test('openZimBatchRange seeks with offset at the dirent cursor', () => { + const archive = fakeArchive(directory) + let walkedFromStart = false + const wrapped = { + iterByPath: () => { + const range = archive.iterByPath() + return { + size: range.size, + offset: (start: number, maxResults: number) => range.offset(start, maxResults), + [Symbol.iterator]() { + walkedFromStart = true + return range[Symbol.iterator]() + }, + } + }, + } + + const paths = [...openZimBatchRange(wrapped, 3)].map((entry) => entry.path) + assert.equal(walkedFromStart, false) + assert.deepEqual(paths, ['C.html', 'C_files/style.css', 'D.html']) +}) + +test('a later batch continues at the next dirent instead of rereading earlier ones', () => { + const archive = fakeArchive(directory) + + const first = takeArticleBatch(openZimBatchRange(archive), 2) + assert.deepEqual(first.accepted, ['A.html', 'B.html']) + assert.equal(first.direntsConsumed, 3) + + const second = takeArticleBatch(openZimBatchRange(archive, first.direntsConsumed), 2) + assert.deepEqual(second.accepted, ['C.html', 'D.html']) + assert.equal(second.direntsConsumed, 3) +}) diff --git a/admin/types/rag.ts b/admin/types/rag.ts index a0fe0fd2..e04898aa 100644 --- a/admin/types/rag.ts +++ b/admin/types/rag.ts @@ -23,6 +23,8 @@ export type ProcessAndEmbedFileResponse = { hasMoreBatches?: boolean articlesProcessed?: number totalArticles?: number + // Dirents already finished in this ZIM. Continuation jobs seek to it. + resumeAtDirent?: number } export type ProcessZIMFileResponse = ProcessAndEmbedFileResponse diff --git a/admin/types/zim.ts b/admin/types/zim.ts index adf816db..ff18a49d 100644 --- a/admin/types/zim.ts +++ b/admin/types/zim.ts @@ -91,8 +91,10 @@ export type ExtractZIMContentOptions = { maxArticles?: number onProgress?: (processedArticles: number, totalArticles: number) => void // Batch processing options to avoid lock timeouts - startOffset?: number // Article index to start from for resuming + startOffset?: number // Article index. Used only when resumeAtDirent is absent. batchSize?: number // Max articles to process in this batch + // Dirents already finished in iterByPath(). The next batch seeks here. + resumeAtDirent?: number } export type ExtractZIMChunkingStrategy = 'structured' | 'simple'