fix: bound document extraction
This commit is contained in:
lofyer
@@ -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()
|
||||
})
|
||||
})
|
||||
|
||||
@@ -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)
|
||||
)
|
||||
|
||||
@@ -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 字符，已截断'
|
||||
)
|
||||
})
|
||||
})
|
||||
|
||||
@@ -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(
|
||||
|
||||
@@ -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()
|
||||
})
|
||||
})
|
||||
|
||||
+38
-5
@@ -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<Buffer> {
|
||||
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,
|
||||
|
||||
@@ -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<typeof vi.fn>
|
||||
}) => 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()
|
||||
})
|
||||
})
|
||||
|
||||
@@ -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', () => {
|
||||
|
||||
@@ -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<number>
|
||||
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<number>
|
||||
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<PdfTextPage[]> {
|
||||
buffer: Buffer,
|
||||
options: PdfTextExtractionOptions = {}
|
||||
): Promise<PdfTextExtraction> {
|
||||
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<ParsedDocument> {
|
||||
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<ParsedDocument> {
|
||||
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,
|
||||
|
||||
@@ -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<ReturnType<typeof open>>,
|
||||
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<ReturnType<typeof open>>,
|
||||
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 })
|
||||
}
|
||||
})
|
||||
})
|
||||
@@ -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<Buffer> {
|
||||
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()
|
||||
}
|
||||
|
||||
@@ -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<void> {
|
||||
|
||||
async function recognizeImage(
|
||||
data: ArrayBuffer,
|
||||
locator: string
|
||||
locator: string,
|
||||
pageNumber?: number
|
||||
): Promise<DocumentOcrResult['sections'][number] | undefined> {
|
||||
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 {
|
||||
|
||||
@@ -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({
|
||||
|
||||
Reference in New Issue
Block a user