From e3b5702767d256e27af078c0755c072171743350 Mon Sep 17 00:00:00 2001 From: lofyer Date: Thu, 13 Aug 2026 04:46:30 +0800 Subject: [PATCH] fix: bound document extraction --- src/main/document-ocr-broker.test.ts | 29 ++ src/main/document-ocr-broker.ts | 17 + src/main/document-parsing-service.test.ts | 113 +++- src/main/document-parsing-service.ts | 150 +++++- src/main/ipc.test.ts | 16 + src/main/ipc.ts | 43 +- .../document-parser-electron-runtime.test.ts | 190 +++++-- src/main/knowledge/document-parser.test.ts | 47 ++ src/main/knowledge/document-parser.ts | 484 ++++++++++++++---- src/main/workspace-file-access.test.ts | 60 +++ src/main/workspace-file-access.ts | 51 +- src/renderer/src/document-ocr-worker.ts | 49 +- src/shared/document-parsing-contracts.ts | 51 +- 13 files changed, 1106 insertions(+), 194 deletions(-) create mode 100644 src/main/workspace-file-access.test.ts diff --git a/src/main/document-ocr-broker.test.ts b/src/main/document-ocr-broker.test.ts index b418878..02a3178 100644 --- a/src/main/document-ocr-broker.test.ts +++ b/src/main/document-ocr-broker.test.ts @@ -139,4 +139,33 @@ describe('DocumentOcrBroker', () => { broker.dispose() await expect(active).rejects.toThrow('OCR 解析已取消') }) + + it('rejects OCR sections outside the requested page set', async () => { + const send = vi.fn() + const broker = new DocumentOcrBroker({ + isDestroyed: vi.fn(() => false), + webContents: { send } + } as never) + const pending = broker.recognize(request()) + const dispatched = send.mock.calls.find( + ([channel]) => channel === ipcChannels.documentParsingOcrRequest + )?.[1] as { requestId: string } + + broker.respond({ + requestId: dispatched.requestId, + sections: [ + { + locator: '第 2 页', + pageNumber: 2, + content: 'wrong page', + confidence: 0.9 + } + ], + pageCount: 2, + warnings: [] + }) + + await expect(pending).rejects.toThrow('OCR 响应页码无效') + broker.dispose() + }) }) diff --git a/src/main/document-ocr-broker.ts b/src/main/document-ocr-broker.ts index f978088..32d9bb5 100644 --- a/src/main/document-ocr-broker.ts +++ b/src/main/document-ocr-broker.ts @@ -99,6 +99,23 @@ export class DocumentOcrBroker { return } if (result.success) { + if ( + pending.request.mimeType === 'application/pdf' && + result.data.sections.some( + (section) => + section.pageNumber === undefined || + section.pageNumber > result.data.pageCount || + ( + pending.request.pageNumbers !== undefined && + !pending.request.pageNumbers.includes(section.pageNumber) + ) + ) + ) { + this.finishRequest(requestId, () => + pending.reject(new Error('OCR 响应页码无效')) + ) + return + } this.finishRequest(requestId, () => pending.resolve(result.data) ) diff --git a/src/main/document-parsing-service.test.ts b/src/main/document-parsing-service.test.ts index 556fb5d..b230100 100644 --- a/src/main/document-parsing-service.test.ts +++ b/src/main/document-parsing-service.test.ts @@ -57,6 +57,7 @@ function createService(overrides?: { requestId: string sections: Array<{ locator: string + pageNumber?: number content: string confidence: number }> @@ -75,6 +76,7 @@ function createService(overrides?: { sections: [ { locator: '第 1 页', + pageNumber: 1, content: '扫描件识别正文', confidence: 0.93 } @@ -136,7 +138,9 @@ describe('DocumentParsingService', () => { locator: '第 1 页', content: '扫描件识别正文', method: 'ocr', - confidence: 0.93 + confidence: 0.93, + pageNumber: 1, + blockKind: 'text' } ]) expect(recognize).toHaveBeenCalledWith( @@ -203,6 +207,41 @@ describe('DocumentParsingService', () => { ]) }) + it('does not silently index a partial mixed PDF when OCR fails', async () => { + const { service } = createService({ + recognize: async () => { + throw new Error('OCR runtime failed') + } + }) + + await expect( + service.parse( + 'mixed.pdf', + createPdfFixture('Native PDF body text', ''), + 'knowledge-index' + ) + ).rejects.toThrow('OCR runtime failed') + }) + + it('does not silently index a mixed PDF when OCR returns no text', async () => { + const { service } = createService({ + recognize: async () => ({ + requestId: crypto.randomUUID(), + sections: [], + pageCount: 2, + warnings: ['第 2 页未识别到文字'] + }) + }) + + await expect( + service.parse( + 'mixed.pdf', + createPdfFixture('Native PDF body text', ''), + 'knowledge-index' + ) + ).rejects.toThrow('第 2 页未识别到可索引文本') + }) + it('limits the number of pages sent to OCR rather than total PDF pages', async () => { const { recognize, service } = createService({ settings: { maximumPages: 1 } @@ -259,4 +298,76 @@ describe('DocumentParsingService', () => { ).rejects.toThrow('请先安装并校验') expect(settingsStore.update).not.toHaveBeenCalled() }) + + it('rejects oversized non-PDF input through the unified service', async () => { + const { service } = createService() + + await expect( + service.parse( + 'large.txt', + Buffer.alloc(20 * 1024 * 1024 + 1), + 'knowledge-index' + ) + ).rejects.toThrow('20MB') + }) + + it('bounds OCR output before returning parsed sections', async () => { + const { service } = createService({ + recognize: async () => ({ + requestId: crypto.randomUUID(), + sections: [ + { + locator: '第 1 页', + pageNumber: 1, + content: 'x'.repeat(1_000_000), + confidence: 0.9 + }, + { + locator: '第 2 页', + pageNumber: 2, + content: 'y'.repeat(1_000_000), + confidence: 0.9 + }, + { + locator: '第 3 页', + pageNumber: 3, + content: 'z'.repeat(1_000_000), + confidence: 0.9 + }, + { + locator: '第 4 页', + pageNumber: 4, + content: 'a'.repeat(1_000_000), + confidence: 0.9 + }, + { + locator: '第 5 页', + pageNumber: 5, + content: 'b'.repeat(1_000_000), + confidence: 0.9 + } + ], + pageCount: 5, + warnings: [] + }), + settings: { + chatWorkflow: 'high-fidelity', + maximumPages: 5 + } + }) + + const parsed = await service.parse( + 'large-ocr.pdf', + createPdfFixture('', '', '', '', ''), + 'chat-attachment' + ) + + expect(parsed.content.length).toBeLessThanOrEqual(5_000_000) + expect( + parsed.sections.map((section) => section.content).join('\n\n') + ).toBe(parsed.content) + expect(parsed.warnings).toContain( + '文档提取文本超过 5,000,000 字符，已截断' + ) + }) }) diff --git a/src/main/document-parsing-service.ts b/src/main/document-parsing-service.ts index e2c18a7..2cccba9 100644 --- a/src/main/document-parsing-service.ts +++ b/src/main/document-parsing-service.ts @@ -3,6 +3,8 @@ import { documentParsingDiagnosticSchema, documentParsingSettingsUpdateSchema, documentParsingSnapshotSchema, + maximumDocumentExtractedCharacters, + maximumDocumentParsingWarnings, type DocumentParsingDiagnostic, type DocumentParsingPurpose, type DocumentParsingSettings, @@ -12,6 +14,7 @@ import type { DocumentOcrBroker } from './document-ocr-broker' import type { DocumentOcrModelManager } from './document-ocr-model-manager' import type { DocumentParsingSettingsStore } from './document-parsing-settings-store' import { + assertDocumentBuffer, DocumentTextUnavailableError, extractPdfTextPages, parseDocument, @@ -81,20 +84,57 @@ function buildPdfDocument( pageCount: number, warnings: string[] = [] ): ParsedDocument { - const content = sections + const truncationWarning = + '文档提取文本超过 5,000,000 字符，已截断' + const boundedWarnings = [ + ...new Set( + warnings.filter((warning) => warning !== truncationWarning) + ) + ] + const limitedSections: ParsedSection[] = [] + let remaining = maximumDocumentExtractedCharacters + let truncated = false + for (const section of sections) { + const separatorLength = limitedSections.length > 0 ? 2 : 0 + if (remaining <= separatorLength) { + truncated = true + break + } + const content = section.content.slice(0, remaining - separatorLength) + if (content) { + limitedSections.push( + content === section.content ? section : { ...section, content } + ) + remaining -= separatorLength + content.length + } + if (content.length < section.content.length) { + truncated = true + break + } + } + if (limitedSections.length < sections.length) { + truncated = true + } + const content = limitedSections .map((section) => section.content) .join('\n\n') - .slice(0, 5_000_000) if (!content) { throw new DocumentTextUnavailableError() } + const documentWarnings = + truncated || warnings.includes(truncationWarning) + ? [ + ...boundedWarnings.slice(0, maximumDocumentParsingWarnings - 1), + truncationWarning + ] + : boundedWarnings.slice(0, maximumDocumentParsingWarnings) return { title: name.replace(/\.[^.]+$/u, ''), sourceFormat: '.pdf', content, - sections, + sections: limitedSections, pageCount, - warnings + warnings: documentWarnings } } @@ -104,7 +144,9 @@ function nativePdfSections(pages: PdfTextPage[]): ParsedSection[] { .map((page) => ({ locator: `第 ${page.pageNumber} 页`, content: page.content, - method: 'native' as const + method: 'native' as const, + pageNumber: page.pageNumber, + blockKind: 'text' as const })) } @@ -158,12 +200,14 @@ export class DocumentParsingService { signal ) => { ensureNotAborted(signal) + assertDocumentBuffer(buffer) if (extname(name).toLowerCase() !== '.pdf') { - return parseDocument(name, buffer) + return parseDocument(name, buffer, signal) } const settings = await this.settingsStore.get() - const pages = await extractPdfTextPages(buffer) + const extracted = await extractPdfTextPages(buffer, { signal }) + const { pages } = extracted ensureNotAborted(signal) const mode = effectiveOcrMode(settings, purpose) const pagesWithoutUsefulText = pages @@ -179,13 +223,19 @@ export class DocumentParsingService { if (mode === 'disabled') { const native = nativePdfSections(pages) if (native.length > 0) { + const warnings = [ + ...(pagesWithoutUsefulText.length > 0 + ? ['部分页面没有有效文本，当前工作流未启用 OCR'] + : []), + ...(extracted.truncated + ? ['文档提取文本超过 5,000,000 字符，已截断'] + : []) + ] return buildPdfDocument( name, native, - pages.length, - pagesWithoutUsefulText.length > 0 - ? ['部分页面没有有效文本，当前工作流未启用 OCR'] - : [] + extracted.pageCount, + warnings ) } throw new DocumentTextUnavailableError( @@ -196,7 +246,10 @@ export class DocumentParsingService { return buildPdfDocument( name, nativePdfSections(pages), - pages.length + extracted.pageCount, + extracted.truncated + ? ['文档提取文本超过 5,000,000 字符，已截断'] + : [] ) } if (ocrPageNumbers.length > settings.maximumPages) { @@ -211,11 +264,20 @@ export class DocumentParsingService { if (!modelStatus.available || !modelStatus.verified) { if ( mode === 'auto' && + purpose !== 'knowledge-index' && native.some((section) => hasUsefulText(section.content)) ) { - return buildPdfDocument(name, native, pages.length, [ - `本地 OCR 不可用，已保留 PDF 文本层内容：${modelStatus.detail}` - ]) + return buildPdfDocument( + name, + native, + extracted.pageCount, + [ + `本地 OCR 不可用，已保留 PDF 文本层内容：${modelStatus.detail}`, + ...(extracted.truncated + ? ['文档提取文本超过 5,000,000 字符，已截断'] + : []) + ] + ) } throw new Error(modelStatus.detail) } @@ -238,23 +300,45 @@ export class DocumentParsingService { ensureNotAborted(signal) if ( mode === 'auto' && + purpose !== 'knowledge-index' && native.some((section) => hasUsefulText(section.content)) ) { const detail = error instanceof Error ? error.message : '本地 OCR 识别失败' - return buildPdfDocument(name, native, pages.length, [ - `本地 OCR 失败，已保留 PDF 文本层内容：${detail}` - ]) + return buildPdfDocument( + name, + native, + extracted.pageCount, + [ + `本地 OCR 失败，已保留 PDF 文本层内容：${detail}`, + ...(extracted.truncated + ? ['文档提取文本超过 5,000,000 字符，已截断'] + : []) + ] + ) } throw error } ensureNotAborted(signal) - const ocrByLocator = new Map( - ocr.sections.map((section) => [section.locator, section]) + const ocrByPageNumber = new Map( + ocr.sections.flatMap((section) => + section.pageNumber === undefined + ? [] + : [[section.pageNumber, section] as const] + ) ) + const missingOcrPage = ocrPageNumbers.find( + (pageNumber) => !ocrByPageNumber.has(pageNumber) + ) + if ( + missingOcrPage !== undefined && + purpose === 'knowledge-index' + ) { + throw new Error(`第 ${missingOcrPage} 页未识别到可索引文本`) + } const merged = pages.flatMap((page): ParsedSection[] => { const locator = `第 ${page.pageNumber} 页` - const recognized = ocrByLocator.get(locator) + const recognized = ocrByPageNumber.get(page.pageNumber) if ( recognized && (mode === 'always' || !hasUsefulText(page.content)) @@ -264,15 +348,33 @@ export class DocumentParsingService { locator, content: recognized.content, method: 'ocr', - confidence: recognized.confidence + confidence: recognized.confidence, + pageNumber: page.pageNumber, + blockKind: 'text' } ] } return page.content - ? [{ locator, content: page.content, method: 'native' }] + ? [{ + locator, + content: page.content, + method: 'native', + pageNumber: page.pageNumber, + blockKind: 'text' + }] : [] }) - return buildPdfDocument(name, merged, pages.length, ocr.warnings) + return buildPdfDocument( + name, + merged, + extracted.pageCount, + [ + ...ocr.warnings, + ...(extracted.truncated + ? ['文档提取文本超过 5,000,000 字符，已截断'] + : []) + ] + ) } async diagnose( diff --git a/src/main/ipc.test.ts b/src/main/ipc.test.ts index 9e0c9aa..045704d 100644 --- a/src/main/ipc.test.ts +++ b/src/main/ipc.test.ts @@ -804,8 +804,13 @@ describe('registerIpcHandlers document parsing', () => { temporaryDirectories.push(directory) const diagnosticPath = join(directory, 'diagnostic.pdf') const artifactPath = join(directory, 'artifact.pdf') + const oversizedArtifactPath = join(directory, 'oversized.pdf') await writeFile(diagnosticPath, 'diagnostic') await writeFile(artifactPath, 'artifact') + await writeFile( + oversizedArtifactPath, + Buffer.alloc(20 * 1024 * 1024 + 1) + ) const diagnostic = { fileName: 'diagnostic.pdf', sourceFormat: 'PDF', @@ -923,6 +928,17 @@ describe('registerIpcHandlers document parsing', () => { 'artifact-import' ) + electronMocks.showOpenDialog.mockResolvedValueOnce({ + canceled: false, + filePaths: [oversizedArtifactPath] + }) + await expect( + electronMocks.handlers.get( + ipcChannels.artifactsImportFiles + )?.(event) + ).rejects.toThrow('超过大小限制') + expect(documentParsingService.parse).toHaveBeenCalledOnce() + await dispose() }) }) diff --git a/src/main/ipc.ts b/src/main/ipc.ts index 138fea7..23bfbc6 100644 --- a/src/main/ipc.ts +++ b/src/main/ipc.ts @@ -5,12 +5,19 @@ import { ipcMain, shell } from 'electron' -import { lstat, mkdir, readFile, realpath, stat } from 'node:fs/promises' +import { + lstat, + mkdir, + readFile, + realpath, + stat +} from 'node:fs/promises' import { randomUUID } from 'node:crypto' import { homedir } from 'node:os' import { basename, extname, isAbsolute, join } from 'node:path' import { z } from 'zod' import { formatShortcutForDisplay } from '../shared/shortcut' +import { readBoundedFile } from './workspace-file-access' import { approvalDecisionSchema, agentQuestionResponseSchema, @@ -473,6 +480,19 @@ function assertTrustedSender( } } +async function readArtifactImportFile( + path: string, + maximumBytes: number, + label: string +): Promise { + return readBoundedFile( + path, + maximumBytes, + `${label}超过大小限制`, + `${label}不是普通文件` + ) +} + function getKnowledgeSnapshot( service: KnowledgeService, selectedLibraryId?: string @@ -3420,14 +3440,15 @@ export function registerIpcHandlers( const artifacts: AssistantArtifact[] = [] for (const filePath of result.filePaths.slice(0, 10)) { const canonicalPath = await realpath(filePath) - const file = await readFile(canonicalPath) const extension = extname(canonicalPath).toLowerCase() const name = basename(canonicalPath) const imageMimeType = imageMimeTypes[extension] if (imageMimeType) { - if (file.byteLength > 3 * 1024 * 1024) { - throw new Error(`图片“${name}”超过 3MB 预览限制`) - } + const file = await readArtifactImportFile( + canonicalPath, + 3 * 1024 * 1024, + `图片“${name}”` + ) artifacts.push( assistantDatabase.createImageArtifact({ projectId, @@ -3439,6 +3460,11 @@ export function registerIpcHandlers( continue } if (extension === '.html' || extension === '.htm') { + const file = await readArtifactImportFile( + canonicalPath, + 5 * 1024 * 1024, + `文件“${name}”` + ) artifacts.push( assistantDatabase.createInlineArtifact({ projectId, @@ -3450,6 +3476,13 @@ export function registerIpcHandlers( ) continue } + const file = await readArtifactImportFile( + canonicalPath, + extension === '.pdf' + ? 20 * 1024 * 1024 + : 5 * 1024 * 1024, + `文件“${name}”` + ) const parsed = documentParsingService ? await documentParsingService.parse( name, diff --git a/src/main/knowledge/document-parser-electron-runtime.test.ts b/src/main/knowledge/document-parser-electron-runtime.test.ts index 064d3fa..4524f22 100644 --- a/src/main/knowledge/document-parser-electron-runtime.test.ts +++ b/src/main/knowledge/document-parser-electron-runtime.test.ts @@ -20,9 +20,14 @@ describe('PDF extraction in Electron main', () => { promise: Promise.resolve({ numPages: 1, getPage: vi.fn(async () => ({ - getTextContent: vi.fn(async () => ({ - items: [{ str: 'PDF body text' }] - })), + streamTextContent: vi.fn(() => + new ReadableStream({ + start(controller) { + controller.enqueue({ items: [{ str: 'PDF body text' }] }) + controller.close() + } + }) + ), cleanup })) }), @@ -31,12 +36,14 @@ describe('PDF extraction in Electron main', () => { await expect( extractPdfTextPages(Buffer.from('synthetic PDF')) - ).resolves.toEqual([ - { + ).resolves.toEqual({ + pageCount: 1, + truncated: false, + pages: [{ pageNumber: 1, content: 'PDF body text' - } - ]) + }] + }) expect(getDocument).toHaveBeenCalledWith({ data: expect.any(Uint8Array), disableFontFace: true, @@ -55,26 +62,33 @@ describe('PDF extraction in Electron main', () => { promise: Promise.resolve({ numPages: 1, getPage: vi.fn(async () => ({ - getTextContent: vi.fn(async () => ({ - items: [ - { - str: 'first', - hasEOL: true, - transform: [1, 0, 0, 1, 10, 100], - height: 10 - }, - { - str: 'second', - transform: [1, 0, 0, 1, 10, 80], - height: 10 - }, - { - str: 'line', - transform: [1, 0, 0, 1, 50, 80], - height: 10 + streamTextContent: vi.fn(() => + new ReadableStream({ + start(controller) { + controller.enqueue({ + items: [ + { + str: 'first', + hasEOL: true, + transform: [1, 0, 0, 1, 10, 100], + height: 10 + }, + { + str: 'second', + transform: [1, 0, 0, 1, 10, 80], + height: 10 + }, + { + str: 'line', + transform: [1, 0, 0, 1, 50, 80], + height: 10 + } + ] + }) + controller.close() } - ] - })), + }) + ), cleanup })) }), @@ -83,11 +97,127 @@ describe('PDF extraction in Electron main', () => { await expect( extractPdfTextPages(Buffer.from('synthetic PDF')) - ).resolves.toEqual([ - { + ).resolves.toEqual({ + pageCount: 1, + truncated: false, + pages: [{ pageNumber: 1, content: 'first\nsecond line' - } - ]) + }] + }) + }) + + it('stops extracting pages at the aggregate character limit', async () => { + const getPage = vi.fn(async (pageNumber: number) => ({ + streamTextContent: vi.fn(() => + new ReadableStream({ + start(controller) { + controller.enqueue({ + items: [{ str: pageNumber === 1 ? 'first' : 'second' }] + }) + controller.close() + } + }) + ), + cleanup: vi.fn() + })) + const destroy = vi.fn(async () => undefined) + getDocument.mockReturnValue({ + promise: Promise.resolve({ numPages: 2, getPage }), + destroy + }) + + await expect( + extractPdfTextPages(Buffer.from('synthetic PDF'), { + maximumCharacters: 5 + }) + ).resolves.toEqual({ + pageCount: 2, + truncated: true, + pages: [{ pageNumber: 1, content: 'first' }] + }) + expect(getPage).toHaveBeenCalledOnce() + }) + + it('cancels the PDF text stream after reaching the character limit', async () => { + let pulls = 0 + const cancel = vi.fn() + const streamTextContent = vi.fn(() => + new ReadableStream({ + pull(controller) { + pulls += 1 + controller.enqueue({ items: [{ str: 'abcde' }] }) + }, + cancel + }) + ) + const destroy = vi.fn(async () => undefined) + getDocument.mockReturnValue({ + promise: Promise.resolve({ + numPages: 1, + getPage: vi.fn(async () => ({ + streamTextContent, + cleanup: vi.fn() + })) + }), + destroy + }) + + await expect( + extractPdfTextPages(Buffer.from('synthetic PDF'), { + maximumCharacters: 5 + }) + ).resolves.toEqual({ + pageCount: 1, + truncated: true, + pages: [{ pageNumber: 1, content: 'abcde' }] + }) + expect(pulls).toBeLessThanOrEqual(2) + expect(cancel).toHaveBeenCalled() + }) + + it('rejects oversized PDFs before reading pages', async () => { + const getPage = vi.fn() + const destroy = vi.fn(async () => undefined) + getDocument.mockReturnValue({ + promise: Promise.resolve({ + numPages: 3, + getPage + }), + destroy + }) + + await expect( + extractPdfTextPages(Buffer.from('synthetic PDF'), { + maximumPages: 2 + }) + ).rejects.toThrow('超过 2 页限制') + expect(getPage).not.toHaveBeenCalled() + expect(destroy).toHaveBeenCalledOnce() + }) + + it('destroys PDF loading when extraction is cancelled', async () => { + let resolveLoading: ((value: { + numPages: number + getPage: ReturnType + }) => void) | undefined + const destroy = vi.fn(async () => undefined) + getDocument.mockReturnValue({ + promise: new Promise((resolve) => { + resolveLoading = resolve + }), + destroy + }) + const controller = new AbortController() + const extraction = extractPdfTextPages( + Buffer.from('synthetic PDF'), + { signal: controller.signal } + ) + + controller.abort(new Error('cancel PDF extraction')) + resolveLoading?.({ numPages: 0, getPage: vi.fn() }) + + await expect(extraction).rejects.toThrow('cancel PDF extraction') + expect(destroy).toHaveBeenCalled() }) }) diff --git a/src/main/knowledge/document-parser.test.ts b/src/main/knowledge/document-parser.test.ts index 1cfef42..b6c5000 100644 --- a/src/main/knowledge/document-parser.test.ts +++ b/src/main/knowledge/document-parser.test.ts @@ -136,6 +136,53 @@ describe('document parser', () => { await expect( parseDocument('expanded.docx', Buffer.from(expandedArchive)) ).rejects.toThrow('损坏') + await expect( + parseDocument('invalid.txt', Buffer.from([0xc3, 0x28])) + ).rejects.toThrow('UTF-8') + }) + + it('keeps extracted sections consistent with the document character limit', async () => { + const parsed = await parseDocument( + 'large.txt', + Buffer.from('x'.repeat(5_000_100)) + ) + + expect(parsed.content).toHaveLength(5_000_000) + expect(parsed.sections).toEqual([ + { + locator: '全文', + content: parsed.content + } + ]) + expect(parsed.warnings).toEqual([ + '文档提取文本超过 5,000,000 字符，已截断' + ]) + }) + + it('rejects chunk output that exceeds the database limit', () => { + expect(() => + chunkDocumentAdvanced( + { + title: 'Too many chunks', + sourceFormat: '.txt', + content: '', + sections: Array.from({ length: 10_001 }, (_, index) => ({ + locator: `section-${index}`, + content: 'content' + })), + warnings: [] + }, + { + version: 1, + mode: 'fixed', + targetCharacters: 400, + overlapCharacters: 0, + parentCharacters: 1_600, + childCharacters: 300, + contextualIndexingEnabled: false + } + ) + ).toThrow('超过 10,000 个分区') }) it('preserves headings and creates recall-only children with parent context', () => { diff --git a/src/main/knowledge/document-parser.ts b/src/main/knowledge/document-parser.ts index 0773f28..3ce3834 100644 --- a/src/main/knowledge/document-parser.ts +++ b/src/main/knowledge/document-parser.ts @@ -5,6 +5,10 @@ import type { KnowledgeChunkingSettings, KnowledgeChunkRole } from '../../shared/knowledge-contracts' +import { + maximumDocumentExtractedCharacters, + maximumPdfPageCount +} from '../../shared/document-parsing-contracts' export type ParsedSection = { locator: string @@ -40,9 +44,10 @@ export type DocumentChunk = { export type DocumentBlockKind = 'text' | 'table' | 'slide' const maximumDocumentBytes = 20 * 1024 * 1024 -const maximumExtractedCharacters = 5_000_000 const maximumHeadingCharacters = 512 const maximumHeadingDepth = 6 +export const maximumDocumentChunks = 10_000 +const maximumDocumentSections = 10_000 export const maximumChunkContextPrefixCharacters = 512 const textExtensions = new Set([ '.c', @@ -112,7 +117,12 @@ function extractXmlText(xml: string): string { } function decodeText(buffer: Buffer): string { - const content = buffer.toString('utf8') + let content: string + try { + content = new TextDecoder('utf-8', { fatal: true }).decode(buffer) + } catch (error) { + throw new Error('文件不是受支持的 UTF-8 文本', { cause: error }) + } const nullCount = [...content.slice(0, 8_192)].filter( (character) => character.charCodeAt(0) === 0 ).length @@ -328,12 +338,18 @@ function parseOfficeArchive( } async function parsePdf( - buffer: Buffer -): Promise<{ sections: ParsedSection[]; pageCount: number }> { - const pages = await extractPdfTextPages(buffer) + buffer: Buffer, + signal?: AbortSignal +): Promise<{ + sections: ParsedSection[] + pageCount: number + truncated: boolean +}> { + const extracted = await extractPdfTextPages(buffer, { signal }) return { - pageCount: pages.length, - sections: pages + pageCount: extracted.pageCount, + truncated: extracted.truncated, + sections: extracted.pages .filter((page) => page.content.length > 0) .map((page) => ({ locator: `第 ${page.pageNumber} 页`, @@ -349,6 +365,18 @@ export type PdfTextPage = { content: string } +export type PdfTextExtraction = { + pages: PdfTextPage[] + pageCount: number + truncated: boolean +} + +export type PdfTextExtractionOptions = { + maximumPages?: number + maximumCharacters?: number + signal?: AbortSignal +} + export class DocumentTextUnavailableError extends Error { constructor(message = '文档中没有可索引的文本内容') { super(message) @@ -356,31 +384,43 @@ export class DocumentTextUnavailableError extends Error { } } -function reconstructPdfText( - items: readonly unknown[] -): string { - const lines: string[] = [] - let line: string[] = [] - let previousY: number | undefined - let previousHeight = 0 - const flush = (): void => { - const value = line.join(' ').replace(/[ \t]+/gu, ' ').trim() - if (value) { - lines.push(value) - } - line = [] - } +type PdfTextItem = { + str?: string + hasEOL?: boolean + transform?: ArrayLike + height?: number +} +type PdfTextReconstructionState = { + lines: string[] + line: string[] + lineCharacterCount: number + characterCount: number + previousY?: number + previousHeight: number +} + +function flushPdfTextLine(state: PdfTextReconstructionState): void { + if (state.lineCharacterCount > 0) { + const value = state.line.join(' ') + state.lines.push(value) + state.characterCount += + (state.lines.length > 1 ? 1 : 0) + value.length + } + state.line = [] + state.lineCharacterCount = 0 +} + +function consumePdfTextItems( + state: PdfTextReconstructionState, + items: readonly unknown[], + maximumCharacters: number +): 'complete' | 'limit' | 'truncated' { for (const candidate of items) { if (typeof candidate !== 'object' || candidate === null) { continue } - const item = candidate as { - str?: string - hasEOL?: boolean - transform?: ArrayLike - height?: number - } + const item = candidate as PdfTextItem const value = typeof item.str === 'string' ? item.str.trim() : '' const y = item.transform && Number.isFinite(item.transform[5]) @@ -391,33 +431,130 @@ function reconstructPdfText( ? Math.abs(item.height) : 0 const coordinateLineBreak = - line.length > 0 && + state.line.length > 0 && y !== undefined && - previousY !== undefined && - Math.abs(y - previousY) > - Math.max(3, previousHeight * 0.8, height * 0.8) + state.previousY !== undefined && + Math.abs(y - state.previousY) > + Math.max(3, state.previousHeight * 0.8, height * 0.8) if (coordinateLineBreak) { - flush() + flushPdfTextLine(state) } if (value) { - line.push(value) + const lineSeparator = state.line.length > 0 ? 1 : 0 + const documentSeparator = + state.line.length === 0 && state.lines.length > 0 ? 1 : 0 + const available = + maximumCharacters - + state.characterCount - + state.lineCharacterCount - + lineSeparator - + documentSeparator + if (available <= 0) { + return 'truncated' + } + const limited = value.slice(0, available) + state.line.push(limited) + state.lineCharacterCount += lineSeparator + limited.length + if (limited.length < value.length) { + return 'truncated' + } } if (item.hasEOL) { - flush() - previousY = undefined - previousHeight = 0 + flushPdfTextLine(state) + state.previousY = undefined + state.previousHeight = 0 } else if (y !== undefined) { - previousY = y - previousHeight = height + state.previousY = y + state.previousHeight = height + } + if ( + state.characterCount + + state.lineCharacterCount + + (state.line.length > 0 && state.lines.length > 0 ? 1 : 0) >= + maximumCharacters + ) { + return 'limit' } } - flush() - return lines.join('\n').replace(/\n{3,}/gu, '\n\n').trim() + return 'complete' +} + +async function reconstructPdfTextStream( + stream: ReadableStream, + maximumCharacters: number, + signal?: AbortSignal +): Promise<{ content: string; truncated: boolean }> { + const reader = stream.getReader() + const state: PdfTextReconstructionState = { + lines: [], + line: [], + lineCharacterCount: 0, + characterCount: 0, + previousHeight: 0 + } + let truncated = false + try { + while (true) { + signal?.throwIfAborted() + const result = await reader.read() + if (result.done) { + break + } + const chunk = result.value as { + items?: readonly unknown[] + } + const consumption = consumePdfTextItems( + state, + chunk.items ?? [], + maximumCharacters + ) + if (consumption === 'truncated') { + truncated = true + break + } + if (consumption === 'limit') { + const lookahead = await reader.read() + truncated = !lookahead.done + break + } + } + } finally { + await reader.cancel().catch(() => undefined) + reader.releaseLock() + } + flushPdfTextLine(state) + return { + content: state.lines.join('\n'), + truncated + } } export async function extractPdfTextPages( - buffer: Buffer -): Promise { + buffer: Buffer, + options: PdfTextExtractionOptions = {} +): Promise { + const maximumPages = options.maximumPages ?? maximumPdfPageCount + const maximumCharacters = + options.maximumCharacters ?? maximumDocumentExtractedCharacters + if ( + !Number.isSafeInteger(maximumPages) || + maximumPages < 1 || + maximumPages > maximumPdfPageCount + ) { + throw new RangeError( + `PDF page limit must be between 1 and ${maximumPdfPageCount}` + ) + } + if ( + !Number.isSafeInteger(maximumCharacters) || + maximumCharacters < 1 || + maximumCharacters > maximumDocumentExtractedCharacters + ) { + throw new RangeError( + `PDF character limit must be between 1 and ${maximumDocumentExtractedCharacters}` + ) + } + options.signal?.throwIfAborted() const pdfjs = await import('pdfjs-dist/legacy/build/pdf.mjs') const loadingTask = pdfjs.getDocument({ data: new Uint8Array(buffer), @@ -429,40 +566,121 @@ export async function extractPdfTextPages( useSystemFonts: false, useWorkerFetch: false }) - const document = await loadingTask.promise + const abortLoading = (): void => { + void loadingTask.destroy() + } + options.signal?.addEventListener('abort', abortLoading, { once: true }) const pages: PdfTextPage[] = [] + let remainingCharacters = maximumCharacters + let truncated = false try { - for (let pageNumber = 1; pageNumber <= document.numPages; pageNumber += 1) { - const page = await document.getPage(pageNumber) - const text = await page.getTextContent() - const content = reconstructPdfText(text.items) - pages.push({ pageNumber, content }) - page.cleanup() + const document = await loadingTask.promise + options.signal?.throwIfAborted() + if (document.numPages > maximumPages) { + throw new Error( + `PDF 有 ${document.numPages} 页，超过 ${maximumPages} 页限制` + ) } + for (let pageNumber = 1; pageNumber <= document.numPages; pageNumber += 1) { + options.signal?.throwIfAborted() + const page = await document.getPage(pageNumber) + try { + const reconstructed = await reconstructPdfTextStream( + page.streamTextContent(), + remainingCharacters, + options.signal + ) + pages.push({ pageNumber, content: reconstructed.content }) + remainingCharacters -= reconstructed.content.length + if (reconstructed.truncated || remainingCharacters === 0) { + truncated = + reconstructed.truncated || pageNumber < document.numPages + break + } + } finally { + page.cleanup() + } + } + return { pages, pageCount: document.numPages, truncated } } finally { + options.signal?.removeEventListener('abort', abortLoading) await loadingTask.destroy() } - return pages } -export async function parseDocument( - name: string, - buffer: Buffer -): Promise { +export function assertDocumentBuffer(buffer: Buffer): void { if (buffer.byteLength === 0) { throw new Error('文档内容为空') } if (buffer.byteLength > maximumDocumentBytes) { throw new Error('单个文档不能超过 20MB') } +} + +function limitParsedSections( + sections: readonly ParsedSection[] +): { + content: string + sections: ParsedSection[] + truncated: boolean +} { + if (sections.length > maximumDocumentSections) { + throw new Error( + `文档包含超过 ${maximumDocumentSections.toLocaleString('en-US')} 个分区` + ) + } + const limited: ParsedSection[] = [] + let remaining = maximumDocumentExtractedCharacters + let truncated = false + for (const section of sections) { + if (!section.content) { + continue + } + const separatorLength = limited.length > 0 ? 2 : 0 + if (remaining <= separatorLength) { + truncated = true + break + } + const maximumSectionLength = remaining - separatorLength + const content = section.content.slice(0, maximumSectionLength) + if (content) { + limited.push( + content === section.content ? section : { ...section, content } + ) + remaining -= separatorLength + content.length + } + if (content.length < section.content.length) { + truncated = true + break + } + } + if (limited.length < sections.filter((section) => section.content).length) { + truncated = true + } + return { + content: limited.map((section) => section.content).join('\n\n'), + sections: limited, + truncated + } +} + +export async function parseDocument( + name: string, + buffer: Buffer, + signal?: AbortSignal +): Promise { + assertDocumentBuffer(buffer) + signal?.throwIfAborted() const extension = extname(name).toLowerCase() let sections: ParsedSection[] let pageCount: number | undefined + let extractionTruncated = false if (extension === '.pdf') { - const parsedPdf = await parsePdf(buffer) + const parsedPdf = await parsePdf(buffer, signal) sections = parsedPdf.sections pageCount = parsedPdf.pageCount + extractionTruncated = parsedPdf.truncated } else if (['.docx', '.xlsx', '.pptx'].includes(extension)) { sections = parseOfficeArchive(buffer, extension) } else if (['.html', '.htm'].includes(extension)) { @@ -481,19 +699,19 @@ export async function parseDocument( throw new Error(`不支持的文档类型：${extension || '未知'}`) } - const content = sections - .map((section) => section.content) - .join('\n\n') - .slice(0, maximumExtractedCharacters) - if (!content) { + signal?.throwIfAborted() + const limited = limitParsedSections(sections) + if (!limited.content) { throw new DocumentTextUnavailableError() } return { title: name.replace(/\.[^.]+$/, ''), sourceFormat: extension || 'unknown', - content, - sections, - warnings: [], + content: limited.content, + sections: limited.sections, + warnings: limited.truncated || extractionTruncated + ? ['文档提取文本超过 5,000,000 字符，已截断'] + : [], ...(pageCount === undefined ? {} : { pageCount }) } } @@ -501,8 +719,14 @@ export async function parseDocument( function splitNatural( content: string, maximumLength: number, - overlap: number + overlap: number, + maximumParts = maximumDocumentChunks ): string[] { + if (maximumParts < 1 && content.trim()) { + throw new Error( + `文档分块超过 ${maximumDocumentChunks.toLocaleString('en-US')} 个，请增大分块长度或缩小文档` + ) + } const chunks: string[] = [] let offset = 0 while (offset < content.length) { @@ -524,6 +748,11 @@ function splitNatural( } const value = content.slice(offset, end).trim() if (value) { + if (chunks.length >= maximumParts) { + throw new Error( + `文档分块超过 ${maximumDocumentChunks.toLocaleString('en-US')} 个，请增大分块长度或缩小文档` + ) + } chunks.push(value) } if (end >= content.length) { @@ -578,9 +807,20 @@ function chunkMetadata( } function structuredSections(document: ParsedDocument): StructuredSection[] { - return document.sections.flatMap((section) => { - const lines = section.content.split(/\r?\n/u) - const result: StructuredSection[] = [] + const result: StructuredSection[] = [] + const append = (section: StructuredSection): void => { + if (!section.content) { + return + } + if (result.length >= maximumDocumentSections) { + throw new Error( + `文档包含超过 ${maximumDocumentSections.toLocaleString('en-US')} 个结构分区` + ) + } + result.push(section) + } + for (const section of document.sections) { + const sectionStart = result.length let heading: string | undefined let headingPath = section.headingPath ? [...section.headingPath].slice(0, maximumHeadingDepth) @@ -590,7 +830,7 @@ function structuredSections(document: ParsedDocument): StructuredSection[] { const flush = (): void => { const content = body.join('\n').trim() if (content) { - result.push({ + append({ locator: heading ? `${section.locator} · ${heading}`.slice(0, 8_192) : section.locator, @@ -602,7 +842,16 @@ function structuredSections(document: ParsedDocument): StructuredSection[] { } body = [] } - for (const line of lines) { + let offset = 0 + while (offset <= section.content.length) { + const lineEnd = section.content.indexOf('\n', offset) + const rawLine = section.content.slice( + offset, + lineEnd < 0 ? section.content.length : lineEnd + ) + const line = rawLine.endsWith('\r') + ? rawLine.slice(0, -1) + : rawLine const match = /^\s{0,3}(#{1,6})\s+(.+?)\s*#*\s*$/u.exec(line) if (match) { flush() @@ -617,18 +866,21 @@ function structuredSections(document: ParsedDocument): StructuredSection[] { } else { body.push(line) } + if (lineEnd < 0) { + break + } + offset = lineEnd + 1 } flush() - return result.length > 0 - ? result - : [ - { - locator: section.locator, - content: section.content.trim(), - ...sectionMetadata(section) - } - ] - }).filter((section) => section.content.length > 0) + if (result.length === sectionStart) { + append({ + locator: section.locator, + content: section.content.trim(), + ...sectionMetadata(section) + }) + } + } + return result } function normalizeContextValue(value: string, maximumLength: number): string { @@ -696,38 +948,55 @@ export function chunkDocumentAdvanced( document: ParsedDocument, settings: KnowledgeChunkingSettings ): DocumentChunk[] { + if (document.sections.length > maximumDocumentSections) { + throw new Error( + `文档包含超过 ${maximumDocumentSections.toLocaleString('en-US')} 个分区` + ) + } if (settings.mode === 'fixed') { - return document.sections.flatMap((section) => - splitNatural( + const result: DocumentChunk[] = [] + for (const section of document.sections) { + const parts = splitNatural( section.content, settings.targetCharacters, - settings.overlapCharacters - ).map((content) => ({ - position: 0, - locator: section.locator, - content, - ...chunkMetadata(section), - role: 'standalone' as const - })) - ).map((chunk, position) => ({ ...chunk, position })) + settings.overlapCharacters, + maximumDocumentChunks - result.length + ) + for (const content of parts) { + result.push({ + position: result.length, + locator: section.locator, + content, + ...chunkMetadata(section), + role: 'standalone' + }) + } + } + return result } const sections = structuredSections(document) if (settings.mode === 'structure') { - return sections.flatMap((section) => - splitNatural( + const result: DocumentChunk[] = [] + for (const section of sections) { + const parts = splitNatural( section.content, settings.targetCharacters, - settings.overlapCharacters - ).map((content) => ({ - position: 0, - locator: section.locator, - heading: section.heading, - content, - ...chunkMetadata(section), - role: 'standalone' as const - })) - ).map((chunk, position) => ({ ...chunk, position })) + settings.overlapCharacters, + maximumDocumentChunks - result.length + ) + for (const content of parts) { + result.push({ + position: result.length, + locator: section.locator, + heading: section.heading, + content, + ...chunkMetadata(section), + role: 'standalone' + }) + } + } + return result } const result: DocumentChunk[] = [] @@ -735,8 +1004,14 @@ export function chunkDocumentAdvanced( for (const parentContent of splitNatural( section.content, settings.parentCharacters, - 0 + 0, + maximumDocumentChunks - result.length )) { + if (result.length >= maximumDocumentChunks) { + throw new Error( + `文档分块超过 ${maximumDocumentChunks.toLocaleString('en-US')} 个，请增大分块长度或缩小文档` + ) + } const parentPosition = result.length result.push({ position: parentPosition, @@ -753,7 +1028,8 @@ export function chunkDocumentAdvanced( for (const childContent of splitNatural( parentContent, settings.childCharacters, - childOverlap + childOverlap, + maximumDocumentChunks - result.length )) { result.push({ position: result.length, diff --git a/src/main/workspace-file-access.test.ts b/src/main/workspace-file-access.test.ts new file mode 100644 index 0000000..aac736a --- /dev/null +++ b/src/main/workspace-file-access.test.ts @@ -0,0 +1,60 @@ +import { + mkdtemp, + open, + rm, + writeFile +} from 'node:fs/promises' +import { tmpdir } from 'node:os' +import { join } from 'node:path' +import { describe, expect, it, vi } from 'vitest' +import { readBoundedFile } from './workspace-file-access' + +type ReadMethod = ( + this: Awaited>, + buffer: Buffer, + offset: number, + length: number, + position: number +) => Promise<{ bytesRead: number; buffer: Buffer }> + +describe('workspace file access', () => { + it('continues reading after a short file-handle read', async () => { + const directory = await mkdtemp(join(tmpdir(), 'goodbuddy-short-read-')) + const path = join(directory, 'short-read.txt') + const content = Buffer.from('hello') + await writeFile(path, content) + const probe = await open(path, 'r') + const prototype = Object.getPrototypeOf(probe) as { + read: ReadMethod + } + const originalRead = prototype.read + await probe.close() + const read = vi + .spyOn(prototype, 'read') + .mockImplementation(function ( + this: Awaited>, + buffer, + offset, + length, + position + ) { + return originalRead.call( + this, + buffer, + offset, + Math.min(length, position === 0 ? 2 : 3), + position + ) + }) + + try { + await expect( + readBoundedFile(path, 5, 'too large') + ).resolves.toEqual(content) + expect(read).toHaveBeenCalledTimes(3) + } finally { + read.mockRestore() + await rm(directory, { recursive: true, force: true }) + } + }) +}) diff --git a/src/main/workspace-file-access.ts b/src/main/workspace-file-access.ts index 1f378f8..b256743 100644 --- a/src/main/workspace-file-access.ts +++ b/src/main/workspace-file-access.ts @@ -52,27 +52,54 @@ export async function readBoundedUtf8File( tooLargeMessage: string, invalidUtf8Message: string ): Promise<{ content: string; size: number }> { + const data = await readBoundedFile( + filePath, + maximumBytes, + tooLargeMessage + ) + try { + return { + content: new TextDecoder('utf-8', { fatal: true }).decode(data), + size: data.byteLength + } + } catch (error) { + throw new Error(invalidUtf8Message, { cause: error }) + } +} + +export async function readBoundedFile( + filePath: string, + maximumBytes: number, + tooLargeMessage: string, + invalidFileMessage = tooLargeMessage +): Promise { const handle = await open(filePath, 'r') try { const metadata = await handle.stat() + if (!metadata.isFile()) { + throw new Error(invalidFileMessage) + } if (metadata.size > maximumBytes) { throw new Error(tooLargeMessage) } const data = Buffer.alloc(metadata.size + 1) - const result = await handle.read(data, 0, data.length, 0) - if (result.bytesRead > maximumBytes) { - throw new Error(tooLargeMessage) - } - try { - return { - content: new TextDecoder('utf-8', { fatal: true }).decode( - data.subarray(0, result.bytesRead) - ), - size: result.bytesRead + let bytesRead = 0 + while (bytesRead < data.length) { + const result = await handle.read( + data, + bytesRead, + data.length - bytesRead, + bytesRead + ) + if (result.bytesRead === 0) { + break + } + bytesRead += result.bytesRead + if (bytesRead > maximumBytes) { + throw new Error(tooLargeMessage) } - } catch (error) { - throw new Error(invalidUtf8Message, { cause: error }) } + return data.subarray(0, bytesRead) } finally { await handle.close() } diff --git a/src/renderer/src/document-ocr-worker.ts b/src/renderer/src/document-ocr-worker.ts index 3b86194..a525f01 100644 --- a/src/renderer/src/document-ocr-worker.ts +++ b/src/renderer/src/document-ocr-worker.ts @@ -10,6 +10,12 @@ import type { DocumentOcrRequest, DocumentOcrResult } from '../../shared/document-parsing-contracts' +import { + maximumDocumentExtractedCharacters, + maximumDocumentOcrSectionCharacters, + maximumDocumentParsingWarnings, + maximumPdfPageCount +} from '../../shared/document-parsing-contracts' import { createWorkerPdfLoadingParameters } from './document-ocr-pdf' type InitializeMessage = { @@ -85,16 +91,21 @@ async function initialize(assets: DocumentOcrAssets): Promise { async function recognizeImage( data: ArrayBuffer, - locator: string + locator: string, + pageNumber?: number ): Promise { if (!service?.isInitialized()) { throw new Error('本地 OCR 模型尚未初始化') } const result = await service.recognize(data) const content = result.text.replace(/\n{3,}/gu, '\n\n').trim() + if (content.length > maximumDocumentOcrSectionCharacters) { + throw new Error('单页 OCR 输出超过字符限制') + } return content ? { locator, + ...(pageNumber === undefined ? {} : { pageNumber }), content, confidence: result.confidence } @@ -152,6 +163,18 @@ async function recognizePdf( createWorkerPdfLoadingParameters(request.data) ) const document = await loadingTask.promise + if (document.numPages > maximumPdfPageCount) { + await loadingTask.destroy() + throw new Error( + `PDF 有 ${document.numPages} 页，超过 ${maximumPdfPageCount} 页限制` + ) + } + if (!request.pageNumbers && document.numPages > request.maximumPages) { + await loadingTask.destroy() + throw new Error( + `PDF 有 ${document.numPages} 页需要 OCR，超过 ${request.maximumPages} 页限制` + ) + } const selectedPages = new Set( request.pageNumbers ?? Array.from( @@ -174,15 +197,11 @@ async function recognizePdf( } const sections: DocumentOcrResult['sections'] = [] const warnings: string[] = [] + let extractedCharacters = 0 try { - for ( - let pageNumber = 1; - pageNumber <= document.numPages; - pageNumber += 1 - ) { - if (!selectedPages.has(pageNumber)) { - continue - } + for (const pageNumber of [...selectedPages].sort( + (left, right) => left - right + )) { worker.postMessage({ type: 'progress', requestId: request.requestId, @@ -192,11 +211,19 @@ async function recognizePdf( try { const section = await recognizeImage( await renderPdfPage(page), - `第 ${pageNumber} 页` + `第 ${pageNumber} 页`, + pageNumber ) if (section) { + extractedCharacters += section.content.length + if ( + extractedCharacters > + maximumDocumentExtractedCharacters + ) { + throw new Error('OCR 输出超过文档字符限制') + } sections.push(section) - } else { + } else if (warnings.length < maximumDocumentParsingWarnings) { warnings.push(`第 ${pageNumber} 页未识别到文字`) } } finally { diff --git a/src/shared/document-parsing-contracts.ts b/src/shared/document-parsing-contracts.ts index f3cf480..35aa4fe 100644 --- a/src/shared/document-parsing-contracts.ts +++ b/src/shared/document-parsing-contracts.ts @@ -1,5 +1,10 @@ import { z } from 'zod' +export const maximumDocumentExtractedCharacters = 5_000_000 +export const maximumDocumentOcrSectionCharacters = 1_000_000 +export const maximumDocumentParsingWarnings = 20 +export const maximumPdfPageCount = 10_000 + export const documentParsingPurposeSchema = z.enum([ 'chat-attachment', 'knowledge-index', @@ -201,13 +206,15 @@ export const documentParsingDiagnosticSchema = z .object({ fileName: z.string().trim().min(1).max(500), sourceFormat: z.string().trim().min(1).max(32), - pageCount: z.number().int().nonnegative().max(10_000), - ocrPageCount: z.number().int().nonnegative().max(10_000), + pageCount: z.number().int().nonnegative().max(maximumPdfPageCount), + ocrPageCount: z.number().int().nonnegative().max(maximumPdfPageCount), characterCount: z.number().int().nonnegative().safe(), method: z.enum(['native', 'ocr', 'mixed']), durationMs: z.number().int().nonnegative().safe(), preview: z.string().max(2_000), - warnings: z.array(z.string().trim().min(1).max(500)).max(20) + warnings: z + .array(z.string().trim().min(1).max(500)) + .max(maximumDocumentParsingWarnings) }) .strict() @@ -239,7 +246,7 @@ export const documentOcrRequestSchema = z ), maximumPages: z.number().int().min(1).max(500), pageNumbers: z - .array(z.number().int().min(1).max(10_000)) + .array(z.number().int().min(1).max(maximumPdfPageCount)) .min(1) .max(500) .optional(), @@ -271,7 +278,17 @@ export const documentOcrRequestSchema = z export const documentOcrSectionSchema = z .object({ locator: z.string().trim().min(1).max(500), - content: z.string().trim().min(1).max(1_000_000), + pageNumber: z + .number() + .int() + .min(1) + .max(maximumPdfPageCount) + .optional(), + content: z + .string() + .trim() + .min(1) + .max(maximumDocumentOcrSectionCharacters), confidence: z.number().min(0).max(1) }) .strict() @@ -280,10 +297,30 @@ export const documentOcrResultSchema = z .object({ requestId: z.string().uuid(), sections: z.array(documentOcrSectionSchema).max(500), - pageCount: z.number().int().nonnegative().max(10_000), - warnings: z.array(z.string().trim().min(1).max(500)).max(20) + pageCount: z + .number() + .int() + .nonnegative() + .max(maximumPdfPageCount), + warnings: z + .array(z.string().trim().min(1).max(500)) + .max(maximumDocumentParsingWarnings) }) .strict() + .superRefine((result, context) => { + let characters = 0 + for (const section of result.sections) { + characters += section.content.length + if (characters > maximumDocumentExtractedCharacters) { + context.addIssue({ + code: 'custom', + path: ['sections'], + message: 'OCR 输出超过文档字符限制' + }) + return + } + } + }) export const documentOcrFailureSchema = z .object({