feat: add document OCR and offline model archives
This commit is contained in:
lofyer
@@ -24,6 +24,7 @@ import type {
|
||||
} from './agent/runtime'
|
||||
import { encodeBoundedJpeg } from './bounded-jpeg'
|
||||
import { parseDocument } from './knowledge/document-parser'
|
||||
import type { ParsedDocument } from './knowledge/document-parser'
|
||||
|
||||
type StoredTextContext = ContextAttachment & {
|
||||
kind: 'text'
|
||||
@@ -94,7 +95,7 @@ function truncateUtf8(value: string, maximumBytes: number): string {
|
||||
}
|
||||
|
||||
function formatParsedDocument(
|
||||
sections: Awaited<ReturnType<typeof parseDocument>>['sections']
|
||||
sections: ParsedDocument['sections']
|
||||
): string {
|
||||
return sections
|
||||
.map(
|
||||
@@ -121,6 +122,23 @@ function remoteAttachmentName(value: string): string {
|
||||
export class ContextManager {
|
||||
private readonly contexts = new Map<string, StoredContext>()
|
||||
private totalBytes = 0
|
||||
private readonly documentParser: (
|
||||
name: string,
|
||||
buffer: Buffer,
|
||||
purpose: 'chat-attachment'
|
||||
) => Promise<ParsedDocument>
|
||||
|
||||
constructor(options?: {
|
||||
parseDocument?: (
|
||||
name: string,
|
||||
buffer: Buffer,
|
||||
purpose: 'chat-attachment'
|
||||
) => Promise<ParsedDocument>
|
||||
}) {
|
||||
this.documentParser =
|
||||
options?.parseDocument ??
|
||||
((name, buffer) => parseDocument(name, buffer))
|
||||
}
|
||||
|
||||
private toPublic(context: StoredContext): ContextAttachment {
|
||||
return {
|
||||
@@ -245,7 +263,11 @@ export class ContextManager {
|
||||
)
|
||||
}
|
||||
if (supportedDocumentExtensions.has(extension)) {
|
||||
const parsed = await parseDocument(name, data)
|
||||
const parsed = await this.documentParser(
|
||||
name,
|
||||
data,
|
||||
'chat-attachment'
|
||||
)
|
||||
return this.storeText(
|
||||
name,
|
||||
truncateUtf8(
|
||||
@@ -340,9 +362,10 @@ export class ContextManager {
|
||||
) {
|
||||
throw new Error('PDF 或 Office 文档必须小于 20MB 且不能是目录')
|
||||
}
|
||||
const parsed = await parseDocument(
|
||||
const parsed = await this.documentParser(
|
||||
basename(canonicalPath),
|
||||
await handle.readFile()
|
||||
await handle.readFile(),
|
||||
'chat-attachment'
|
||||
)
|
||||
attachments.push(
|
||||
this.storeText(
|
||||
|
||||
@@ -0,0 +1,53 @@
|
||||
import { describe, expect, it, vi } from 'vitest'
|
||||
import { ipcChannels } from '../shared/ipc-channels'
|
||||
import { DocumentOcrBroker } from './document-ocr-broker'
|
||||
|
||||
function request() {
|
||||
return {
|
||||
modelId: 'pp-ocrv6-tiny',
|
||||
fileName: 'scan.pdf',
|
||||
mimeType: 'application/pdf' as const,
|
||||
data: new ArrayBuffer(8),
|
||||
maximumPages: 10,
|
||||
pageNumbers: [1],
|
||||
pageTimeoutSeconds: 60
|
||||
}
|
||||
}
|
||||
|
||||
describe('DocumentOcrBroker', () => {
|
||||
it('forwards an AbortSignal cancellation to the renderer', async () => {
|
||||
const send = vi.fn()
|
||||
const broker = new DocumentOcrBroker({
|
||||
isDestroyed: vi.fn(() => false),
|
||||
webContents: { send }
|
||||
} as never)
|
||||
const controller = new AbortController()
|
||||
const result = broker.recognize(request(), controller.signal)
|
||||
const ocrRequest = send.mock.calls.find(
|
||||
([channel]) => channel === ipcChannels.documentParsingOcrRequest
|
||||
)?.[1] as { requestId: string }
|
||||
|
||||
controller.abort()
|
||||
|
||||
await expect(result).rejects.toThrow('OCR 解析已取消')
|
||||
expect(send).toHaveBeenCalledWith(
|
||||
ipcChannels.documentParsingOcrCancel,
|
||||
ocrRequest.requestId
|
||||
)
|
||||
broker.dispose()
|
||||
})
|
||||
|
||||
it('rejects a request that is already cancelled', () => {
|
||||
const broker = new DocumentOcrBroker({
|
||||
isDestroyed: vi.fn(() => false),
|
||||
webContents: { send: vi.fn() }
|
||||
} as never)
|
||||
const controller = new AbortController()
|
||||
controller.abort()
|
||||
|
||||
expect(() =>
|
||||
broker.recognize(request(), controller.signal)
|
||||
).toThrow('OCR 解析已取消')
|
||||
broker.dispose()
|
||||
})
|
||||
})
|
||||
@@ -0,0 +1,131 @@
|
||||
import type { BrowserWindow } from 'electron'
|
||||
import { ipcChannels } from '../shared/ipc-channels'
|
||||
import {
|
||||
documentOcrFailureSchema,
|
||||
documentOcrRequestSchema,
|
||||
documentOcrResultSchema,
|
||||
type DocumentOcrRequest,
|
||||
type DocumentOcrResult
|
||||
} from '../shared/document-parsing-contracts'
|
||||
|
||||
type PendingRequest = {
|
||||
resolve: (result: DocumentOcrResult) => void
|
||||
reject: (error: Error) => void
|
||||
timer: ReturnType<typeof setTimeout>
|
||||
detachAbort: () => void
|
||||
}
|
||||
|
||||
const maximumPendingRequests = 4
|
||||
const maximumTotalTimeoutMs = 10 * 60 * 1_000
|
||||
|
||||
export class DocumentOcrBroker {
|
||||
private readonly pending = new Map<string, PendingRequest>()
|
||||
private disposed = false
|
||||
|
||||
constructor(private readonly window: BrowserWindow) {}
|
||||
|
||||
recognize(
|
||||
input: Omit<DocumentOcrRequest, 'requestId'>,
|
||||
signal?: AbortSignal
|
||||
): Promise<DocumentOcrResult> {
|
||||
if (this.disposed || this.window.isDestroyed()) {
|
||||
throw new Error('OCR 渲染服务不可用')
|
||||
}
|
||||
if (this.pending.size >= maximumPendingRequests) {
|
||||
throw new Error('OCR 任务过多，请稍后重试')
|
||||
}
|
||||
const request = documentOcrRequestSchema.parse({
|
||||
...input,
|
||||
requestId: crypto.randomUUID()
|
||||
})
|
||||
if (signal?.aborted) {
|
||||
throw new Error('OCR 解析已取消')
|
||||
}
|
||||
const timeoutMs = Math.min(
|
||||
maximumTotalTimeoutMs,
|
||||
Math.max(
|
||||
request.pageTimeoutSeconds * 1_000,
|
||||
request.pageTimeoutSeconds *
|
||||
request.maximumPages *
|
||||
1_000
|
||||
)
|
||||
)
|
||||
return new Promise<DocumentOcrResult>((resolve, reject) => {
|
||||
const cancel = (message: string): void => {
|
||||
const pending = this.pending.get(request.requestId)
|
||||
if (!pending) {
|
||||
return
|
||||
}
|
||||
clearTimeout(pending.timer)
|
||||
pending.detachAbort()
|
||||
this.pending.delete(request.requestId)
|
||||
this.window.webContents.send(
|
||||
ipcChannels.documentParsingOcrCancel,
|
||||
request.requestId
|
||||
)
|
||||
reject(new Error(message))
|
||||
}
|
||||
const timer = setTimeout(() => {
|
||||
cancel('OCR 解析超时')
|
||||
}, timeoutMs)
|
||||
const onAbort = (): void => cancel('OCR 解析已取消')
|
||||
signal?.addEventListener('abort', onAbort, { once: true })
|
||||
this.pending.set(request.requestId, {
|
||||
resolve,
|
||||
reject,
|
||||
timer,
|
||||
detachAbort: () =>
|
||||
signal?.removeEventListener('abort', onAbort)
|
||||
})
|
||||
if (signal?.aborted) {
|
||||
cancel('OCR 解析已取消')
|
||||
return
|
||||
}
|
||||
this.window.webContents.send(
|
||||
ipcChannels.documentParsingOcrRequest,
|
||||
request
|
||||
)
|
||||
})
|
||||
}
|
||||
|
||||
respond(input: unknown): void {
|
||||
const result = documentOcrResultSchema.safeParse(input)
|
||||
const failure = result.success
|
||||
? undefined
|
||||
: documentOcrFailureSchema.safeParse(input)
|
||||
const requestId = result.success
|
||||
? result.data.requestId
|
||||
: failure?.success
|
||||
? failure.data.requestId
|
||||
: undefined
|
||||
if (!requestId) {
|
||||
throw new Error('OCR 响应无效')
|
||||
}
|
||||
const pending = this.pending.get(requestId)
|
||||
if (!pending) {
|
||||
return
|
||||
}
|
||||
clearTimeout(pending.timer)
|
||||
pending.detachAbort()
|
||||
this.pending.delete(requestId)
|
||||
if (result.success) {
|
||||
pending.resolve(result.data)
|
||||
} else {
|
||||
if (!failure?.success) {
|
||||
pending.reject(new Error('OCR 响应无效'))
|
||||
return
|
||||
}
|
||||
pending.reject(new Error(failure.data.error))
|
||||
}
|
||||
}
|
||||
|
||||
dispose(): void {
|
||||
this.disposed = true
|
||||
for (const pending of this.pending.values()) {
|
||||
clearTimeout(pending.timer)
|
||||
pending.detachAbort()
|
||||
pending.reject(new Error('OCR 解析已取消'))
|
||||
}
|
||||
this.pending.clear()
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,204 @@
|
||||
import {
|
||||
documentOcrModelCatalogEntrySchema,
|
||||
type DocumentOcrModelCatalogEntry
|
||||
} from '../shared/document-parsing-contracts'
|
||||
|
||||
const detectionRevision =
|
||||
'7d7f5d128d9309ebf6de4f21f404dd583afdbae3'
|
||||
const recognitionRevision =
|
||||
'afba04b618200c5f4824531c6e42c957c6439d9a'
|
||||
const smallDetectionRevision =
|
||||
'956a0b620a4017cc04056c692be1703b0025d028'
|
||||
const smallRecognitionRevision =
|
||||
'296d43bc0ebced0fd9c605174aa5962e49810ab6'
|
||||
const mediumDetectionRevision =
|
||||
'c317b40325be40bfaaff58c8dcece2a075294f8a'
|
||||
const mediumRecognitionRevision =
|
||||
'db5d610d492a14e3c34dc1fd4e9339bd369f79e6'
|
||||
|
||||
export const DOCUMENT_OCR_MODEL_CATALOG: readonly DocumentOcrModelCatalogEntry[] =
|
||||
documentOcrModelCatalogEntrySchema.array().parse([
|
||||
{
|
||||
id: 'pp-ocrv6-tiny',
|
||||
displayName: 'PP-OCRv6 Tiny',
|
||||
description:
|
||||
'PaddleOCR 官方轻量中文 OCR 模型，适合扫描 PDF 和图片的本地 CPU 识别。',
|
||||
languages: ['中文', '英语'],
|
||||
runtime: 'onnxruntime-web-wasm',
|
||||
quality: 'basic',
|
||||
speed: 'fast',
|
||||
recommended: false,
|
||||
repositoryUrl:
|
||||
'https://modelscope.cn/models/PaddlePaddle/' +
|
||||
'PP-OCRv6_tiny_rec_onnx',
|
||||
license: {
|
||||
name: 'Apache License 2.0',
|
||||
notice:
|
||||
'检测与识别模型由 PaddlePaddle 在 ModelScope 发布，使用前请阅读模型仓库及 PaddleOCR 的许可证说明。',
|
||||
url: 'https://github.com/PaddlePaddle/PaddleOCR/blob/main/LICENSE'
|
||||
},
|
||||
files: [
|
||||
{
|
||||
name: 'detection.onnx',
|
||||
role: 'detection',
|
||||
download: {
|
||||
url:
|
||||
'https://modelscope.cn/models/PaddlePaddle/' +
|
||||
'PP-OCRv6_tiny_det_onnx/resolve/' +
|
||||
`${detectionRevision}/inference.onnx`,
|
||||
size: 1_780_590,
|
||||
sha256:
|
||||
'193bab7a04fca699a6c82e6abb5b81bdb28177f0abd4062552b04908dafb19f8'
|
||||
}
|
||||
},
|
||||
{
|
||||
name: 'recognition.onnx',
|
||||
role: 'recognition',
|
||||
download: {
|
||||
url:
|
||||
'https://modelscope.cn/models/PaddlePaddle/' +
|
||||
'PP-OCRv6_tiny_rec_onnx/resolve/' +
|
||||
`${recognitionRevision}/inference.onnx`,
|
||||
size: 4_462_639,
|
||||
sha256:
|
||||
'9ef676d6ed3c88256a2d92c640c44f25b0c40947e111b14b8be8f594091563e6'
|
||||
}
|
||||
},
|
||||
{
|
||||
name: 'dictionary.yml',
|
||||
role: 'dictionary',
|
||||
download: {
|
||||
url:
|
||||
'https://modelscope.cn/models/PaddlePaddle/' +
|
||||
'PP-OCRv6_tiny_rec_onnx/resolve/' +
|
||||
`${recognitionRevision}/inference.yml`,
|
||||
size: 55_571,
|
||||
sha256:
|
||||
'66170210bad538e83fff3c4a3867e547d6bf20b50d64b20347c4b913f3034ea1'
|
||||
}
|
||||
}
|
||||
]
|
||||
},
|
||||
{
|
||||
id: 'pp-ocrv6-small',
|
||||
displayName: 'PP-OCRv6 Small',
|
||||
description:
|
||||
'PaddleOCR 官方 50 语言 OCR 模型，在识别质量、速度和本地资源占用之间取得平衡。',
|
||||
languages: ['50 种语言'],
|
||||
runtime: 'onnxruntime-web-wasm',
|
||||
quality: 'balanced',
|
||||
speed: 'balanced',
|
||||
recommended: true,
|
||||
repositoryUrl:
|
||||
'https://modelscope.cn/models/PaddlePaddle/' +
|
||||
'PP-OCRv6_small_rec_onnx',
|
||||
license: {
|
||||
name: 'Apache License 2.0',
|
||||
notice:
|
||||
'检测与识别模型由 PaddlePaddle 在 ModelScope 发布，使用前请阅读模型仓库及 PaddleOCR 的许可证说明。',
|
||||
url: 'https://github.com/PaddlePaddle/PaddleOCR/blob/main/LICENSE'
|
||||
},
|
||||
files: [
|
||||
{
|
||||
name: 'detection.onnx',
|
||||
role: 'detection',
|
||||
download: {
|
||||
url:
|
||||
'https://modelscope.cn/models/PaddlePaddle/' +
|
||||
'PP-OCRv6_small_det_onnx/resolve/' +
|
||||
`${smallDetectionRevision}/inference.onnx`,
|
||||
size: 9_880_512,
|
||||
sha256:
|
||||
'd73e0058b7a8086bbd57f3d10b8bcd4ff95363f67e06e2762b5e814fe9c9410e'
|
||||
}
|
||||
},
|
||||
{
|
||||
name: 'recognition.onnx',
|
||||
role: 'recognition',
|
||||
download: {
|
||||
url:
|
||||
'https://modelscope.cn/models/PaddlePaddle/' +
|
||||
'PP-OCRv6_small_rec_onnx/resolve/' +
|
||||
`${smallRecognitionRevision}/inference.onnx`,
|
||||
size: 21_159_378,
|
||||
sha256:
|
||||
'5435fd747c9e0efe15a96d0b378d5bd157e9492ed8fd80edf08f30d02fa24634'
|
||||
}
|
||||
},
|
||||
{
|
||||
name: 'dictionary.yml',
|
||||
role: 'dictionary',
|
||||
download: {
|
||||
url:
|
||||
'https://modelscope.cn/models/PaddlePaddle/' +
|
||||
'PP-OCRv6_small_rec_onnx/resolve/' +
|
||||
`${smallRecognitionRevision}/inference.yml`,
|
||||
size: 150_579,
|
||||
sha256:
|
||||
'ab078671bb49f06228eadccd34f1bb501e157f7a047095ffb943ba81512c77d1'
|
||||
}
|
||||
}
|
||||
]
|
||||
},
|
||||
{
|
||||
id: 'pp-ocrv6-medium',
|
||||
displayName: 'PP-OCRv6 Medium',
|
||||
description:
|
||||
'PaddleOCR 官方 50 语言高质量 OCR 模型，识别较慢，并需要更多内存且具有更高延迟。',
|
||||
languages: ['50 种语言'],
|
||||
runtime: 'onnxruntime-web-wasm',
|
||||
quality: 'high',
|
||||
speed: 'slow',
|
||||
recommended: false,
|
||||
repositoryUrl:
|
||||
'https://modelscope.cn/models/PaddlePaddle/' +
|
||||
'PP-OCRv6_medium_rec_onnx',
|
||||
license: {
|
||||
name: 'Apache License 2.0',
|
||||
notice:
|
||||
'检测与识别模型由 PaddlePaddle 在 ModelScope 发布，使用前请阅读模型仓库及 PaddleOCR 的许可证说明。',
|
||||
url: 'https://github.com/PaddlePaddle/PaddleOCR/blob/main/LICENSE'
|
||||
},
|
||||
files: [
|
||||
{
|
||||
name: 'detection.onnx',
|
||||
role: 'detection',
|
||||
download: {
|
||||
url:
|
||||
'https://modelscope.cn/models/PaddlePaddle/' +
|
||||
'PP-OCRv6_medium_det_onnx/resolve/' +
|
||||
`${mediumDetectionRevision}/inference.onnx`,
|
||||
size: 62_032_837,
|
||||
sha256:
|
||||
'eb13b44b25bb36f89528b68720af8a61d9cf381176107f465db1757b65d086e1'
|
||||
}
|
||||
},
|
||||
{
|
||||
name: 'recognition.onnx',
|
||||
role: 'recognition',
|
||||
download: {
|
||||
url:
|
||||
'https://modelscope.cn/models/PaddlePaddle/' +
|
||||
'PP-OCRv6_medium_rec_onnx/resolve/' +
|
||||
`${mediumRecognitionRevision}/inference.onnx`,
|
||||
size: 76_554_979,
|
||||
sha256:
|
||||
'9c09abf0957f7968c7586464b7397b84ad2387a0497a351af40e9acc71b673ba'
|
||||
}
|
||||
},
|
||||
{
|
||||
name: 'dictionary.yml',
|
||||
role: 'dictionary',
|
||||
download: {
|
||||
url:
|
||||
'https://modelscope.cn/models/PaddlePaddle/' +
|
||||
'PP-OCRv6_medium_rec_onnx/resolve/' +
|
||||
`${mediumRecognitionRevision}/inference.yml`,
|
||||
size: 150_580,
|
||||
sha256:
|
||||
'991b700facf5b50a7de193468207d5f4255b538dde0d312ae3b7c7a9b6873129'
|
||||
}
|
||||
}
|
||||
]
|
||||
}
|
||||
])
|
||||
@@ -0,0 +1,341 @@
|
||||
import { createHash } from 'node:crypto'
|
||||
import {
|
||||
mkdtemp,
|
||||
mkdir,
|
||||
rm,
|
||||
writeFile
|
||||
} from 'node:fs/promises'
|
||||
import { tmpdir } from 'node:os'
|
||||
import { join } from 'node:path'
|
||||
import { afterEach, describe, expect, it, vi } from 'vitest'
|
||||
import type { DocumentOcrModelCatalogEntry } from '../shared/document-parsing-contracts'
|
||||
import { DOCUMENT_OCR_MODEL_CATALOG } from './document-ocr-model-catalog'
|
||||
import {
|
||||
DocumentOcrModelManager,
|
||||
extractPaddleCharacterDictionary
|
||||
} from './document-ocr-model-manager'
|
||||
|
||||
const temporaryDirectories: string[] = []
|
||||
|
||||
function sha256(value: Uint8Array): string {
|
||||
return createHash('sha256').update(value).digest('hex')
|
||||
}
|
||||
|
||||
function dictionaryYaml(): Uint8Array {
|
||||
const characters = [
|
||||
"'!'",
|
||||
"'\"'",
|
||||
"''''",
|
||||
...Array.from({ length: 120 }, (_, index) =>
|
||||
String.fromCodePoint(0x4e00 + index)
|
||||
)
|
||||
]
|
||||
return Buffer.from(
|
||||
`PostProcess:\n name: CTCLabelDecode\n character_dict:\n${characters
|
||||
.map((character) => ` - ${character}`)
|
||||
.join('\n')}\n`,
|
||||
'utf8'
|
||||
)
|
||||
}
|
||||
|
||||
function catalog(
|
||||
detection: Uint8Array,
|
||||
recognition: Uint8Array,
|
||||
dictionary: Uint8Array
|
||||
): readonly DocumentOcrModelCatalogEntry[] {
|
||||
const files = [
|
||||
{
|
||||
name: 'detection.onnx',
|
||||
role: 'detection' as const,
|
||||
bytes: detection
|
||||
},
|
||||
{
|
||||
name: 'recognition.onnx',
|
||||
role: 'recognition' as const,
|
||||
bytes: recognition
|
||||
},
|
||||
{
|
||||
name: 'dictionary.yml',
|
||||
role: 'dictionary' as const,
|
||||
bytes: dictionary
|
||||
}
|
||||
]
|
||||
return [
|
||||
{
|
||||
id: 'pp-ocrv6-tiny',
|
||||
displayName: 'PP-OCRv6 Tiny',
|
||||
description: 'Test OCR model catalog entry.',
|
||||
languages: ['中文', '英语'],
|
||||
runtime: 'onnxruntime-web-wasm',
|
||||
quality: 'balanced',
|
||||
speed: 'fast',
|
||||
recommended: true,
|
||||
repositoryUrl:
|
||||
'https://modelscope.cn/models/PaddlePaddle/PP-OCRv6_tiny_rec_onnx',
|
||||
license: {
|
||||
name: 'Apache License 2.0',
|
||||
notice: 'Test license notice.',
|
||||
url: 'https://example.com/license'
|
||||
},
|
||||
files: files.map((file) => ({
|
||||
name: file.name,
|
||||
role: file.role,
|
||||
download: {
|
||||
url: `https://modelscope.cn/models/example/resolve/revision/${file.name}`,
|
||||
size: file.bytes.byteLength,
|
||||
sha256: sha256(file.bytes)
|
||||
}
|
||||
}))
|
||||
}
|
||||
]
|
||||
}
|
||||
|
||||
async function createManager(
|
||||
bytes?: {
|
||||
detection: Uint8Array
|
||||
recognition: Uint8Array
|
||||
dictionary: Uint8Array
|
||||
}
|
||||
): Promise<{
|
||||
directory: string
|
||||
manager: DocumentOcrModelManager
|
||||
modelBytes: {
|
||||
detection: Uint8Array
|
||||
recognition: Uint8Array
|
||||
dictionary: Uint8Array
|
||||
}
|
||||
}> {
|
||||
const directory = await mkdtemp(
|
||||
join(tmpdir(), 'goodbuddy-document-ocr-model-')
|
||||
)
|
||||
temporaryDirectories.push(directory)
|
||||
const modelBytes = bytes ?? {
|
||||
detection: Buffer.from('detection model'),
|
||||
recognition: Buffer.from('recognition model'),
|
||||
dictionary: dictionaryYaml()
|
||||
}
|
||||
const testCatalog = catalog(
|
||||
modelBytes.detection,
|
||||
modelBytes.recognition,
|
||||
modelBytes.dictionary
|
||||
)
|
||||
const entry = testCatalog[0]
|
||||
if (!entry) {
|
||||
throw new Error('Test OCR catalog is empty')
|
||||
}
|
||||
const files = new Map(
|
||||
entry.files.map((file) => [
|
||||
file.download.url,
|
||||
modelBytes[file.role]
|
||||
])
|
||||
)
|
||||
const transport = vi.fn(async (input: string | URL | Request) => {
|
||||
const url =
|
||||
input instanceof Request ? input.url : input.toString()
|
||||
const body = files.get(url)
|
||||
if (!body) {
|
||||
return new Response(null, { status: 404 })
|
||||
}
|
||||
return new Response(body, {
|
||||
status: 200,
|
||||
headers: {
|
||||
'content-length': String(body.byteLength)
|
||||
}
|
||||
})
|
||||
}) as unknown as typeof fetch
|
||||
return {
|
||||
directory,
|
||||
manager: new DocumentOcrModelManager({
|
||||
userDataDirectory: directory,
|
||||
fetch: transport,
|
||||
catalog: testCatalog
|
||||
}),
|
||||
modelBytes
|
||||
}
|
||||
}
|
||||
|
||||
afterEach(async () => {
|
||||
await Promise.all(
|
||||
temporaryDirectories.splice(0).map((directory) =>
|
||||
rm(directory, { recursive: true, force: true })
|
||||
)
|
||||
)
|
||||
})
|
||||
|
||||
describe('DocumentOcrModelManager', () => {
|
||||
it('uses immutable SHA-256 verified ModelScope catalog files', () => {
|
||||
expect(DOCUMENT_OCR_MODEL_CATALOG).toHaveLength(3)
|
||||
expect(
|
||||
new Set(DOCUMENT_OCR_MODEL_CATALOG.map((entry) => entry.id)).size
|
||||
).toBe(3)
|
||||
expect(
|
||||
DOCUMENT_OCR_MODEL_CATALOG.filter((entry) => entry.recommended).map(
|
||||
(entry) => entry.id
|
||||
)
|
||||
).toEqual(['pp-ocrv6-small'])
|
||||
|
||||
for (const entry of DOCUMENT_OCR_MODEL_CATALOG) {
|
||||
for (const file of entry.files) {
|
||||
expect(file.download.url).toMatch(
|
||||
/^https:\/\/modelscope\.cn\/models\/PaddlePaddle\/[^/]+\/resolve\/[a-f0-9]{40}\/[^/]+$/u
|
||||
)
|
||||
expect(file.download.sha256).toMatch(/^[a-f0-9]{64}$/u)
|
||||
expect(file.download.size).toBeGreaterThan(0)
|
||||
}
|
||||
}
|
||||
|
||||
expect(
|
||||
DOCUMENT_OCR_MODEL_CATALOG.find(
|
||||
(entry) => entry.id === 'pp-ocrv6-small'
|
||||
)
|
||||
).toMatchObject({
|
||||
languages: ['50 种语言'],
|
||||
quality: 'balanced',
|
||||
speed: 'balanced',
|
||||
recommended: true,
|
||||
files: [
|
||||
{
|
||||
role: 'detection',
|
||||
download: {
|
||||
url: 'https://modelscope.cn/models/PaddlePaddle/PP-OCRv6_small_det_onnx/resolve/956a0b620a4017cc04056c692be1703b0025d028/inference.onnx',
|
||||
size: 9_880_512,
|
||||
sha256:
|
||||
'd73e0058b7a8086bbd57f3d10b8bcd4ff95363f67e06e2762b5e814fe9c9410e'
|
||||
}
|
||||
},
|
||||
{
|
||||
role: 'recognition',
|
||||
download: {
|
||||
url: 'https://modelscope.cn/models/PaddlePaddle/PP-OCRv6_small_rec_onnx/resolve/296d43bc0ebced0fd9c605174aa5962e49810ab6/inference.onnx',
|
||||
size: 21_159_378,
|
||||
sha256:
|
||||
'5435fd747c9e0efe15a96d0b378d5bd157e9492ed8fd80edf08f30d02fa24634'
|
||||
}
|
||||
},
|
||||
{
|
||||
role: 'dictionary',
|
||||
download: {
|
||||
url: 'https://modelscope.cn/models/PaddlePaddle/PP-OCRv6_small_rec_onnx/resolve/296d43bc0ebced0fd9c605174aa5962e49810ab6/inference.yml',
|
||||
size: 150_579,
|
||||
sha256:
|
||||
'ab078671bb49f06228eadccd34f1bb501e157f7a047095ffb943ba81512c77d1'
|
||||
}
|
||||
}
|
||||
]
|
||||
})
|
||||
expect(
|
||||
DOCUMENT_OCR_MODEL_CATALOG.find(
|
||||
(entry) => entry.id === 'pp-ocrv6-medium'
|
||||
)
|
||||
).toMatchObject({
|
||||
languages: ['50 种语言'],
|
||||
quality: 'high',
|
||||
speed: 'slow',
|
||||
recommended: false,
|
||||
files: [
|
||||
{
|
||||
role: 'detection',
|
||||
download: {
|
||||
url: 'https://modelscope.cn/models/PaddlePaddle/PP-OCRv6_medium_det_onnx/resolve/c317b40325be40bfaaff58c8dcece2a075294f8a/inference.onnx',
|
||||
size: 62_032_837,
|
||||
sha256:
|
||||
'eb13b44b25bb36f89528b68720af8a61d9cf381176107f465db1757b65d086e1'
|
||||
}
|
||||
},
|
||||
{
|
||||
role: 'recognition',
|
||||
download: {
|
||||
url: 'https://modelscope.cn/models/PaddlePaddle/PP-OCRv6_medium_rec_onnx/resolve/db5d610d492a14e3c34dc1fd4e9339bd369f79e6/inference.onnx',
|
||||
size: 76_554_979,
|
||||
sha256:
|
||||
'9c09abf0957f7968c7586464b7397b84ad2387a0497a351af40e9acc71b673ba'
|
||||
}
|
||||
},
|
||||
{
|
||||
role: 'dictionary',
|
||||
download: {
|
||||
url: 'https://modelscope.cn/models/PaddlePaddle/PP-OCRv6_medium_rec_onnx/resolve/db5d610d492a14e3c34dc1fd4e9339bd369f79e6/inference.yml',
|
||||
size: 150_580,
|
||||
sha256:
|
||||
'991b700facf5b50a7de193468207d5f4255b538dde0d312ae3b7c7a9b6873129'
|
||||
}
|
||||
}
|
||||
]
|
||||
})
|
||||
})
|
||||
|
||||
it('downloads, verifies, and loads OCR assets', async () => {
|
||||
const { manager, modelBytes } = await createManager()
|
||||
|
||||
await expect(manager.install('pp-ocrv6-tiny')).resolves.toMatchObject({
|
||||
id: 'pp-ocrv6-tiny',
|
||||
source: 'download'
|
||||
})
|
||||
await expect(manager.getStatus('pp-ocrv6-tiny')).resolves.toMatchObject({
|
||||
available: true,
|
||||
verified: true
|
||||
})
|
||||
const assets = await manager.getAssets('pp-ocrv6-tiny')
|
||||
expect(new Uint8Array(assets.detection)).toEqual(
|
||||
Uint8Array.from(modelBytes.detection)
|
||||
)
|
||||
expect(new Uint8Array(assets.recognition)).toEqual(
|
||||
Uint8Array.from(modelBytes.recognition)
|
||||
)
|
||||
expect(new TextDecoder().decode(assets.dictionary)).toContain(
|
||||
"!\n\"\n'\n"
|
||||
)
|
||||
})
|
||||
|
||||
it('rejects an imported model whose hash does not match', async () => {
|
||||
const { directory, manager, modelBytes } = await createManager()
|
||||
const source = join(directory, 'manual-model')
|
||||
await mkdir(source)
|
||||
await Promise.all([
|
||||
writeFile(join(source, 'detection.onnx'), modelBytes.detection),
|
||||
writeFile(join(source, 'recognition.onnx'), modelBytes.recognition),
|
||||
writeFile(join(source, 'dictionary.yml'), 'tampered')
|
||||
])
|
||||
|
||||
await expect(
|
||||
manager.registerLocalDirectory('pp-ocrv6-tiny', source)
|
||||
).rejects.toThrow('校验失败')
|
||||
await expect(manager.getSnapshot()).resolves.toMatchObject({
|
||||
installed: [],
|
||||
operations: []
|
||||
})
|
||||
})
|
||||
|
||||
it('round-trips a verified OCR model through an offline ZIP archive', async () => {
|
||||
const { directory, manager, modelBytes } = await createManager()
|
||||
const archive = join(directory, 'ocr-model.zip')
|
||||
|
||||
await manager.install('pp-ocrv6-tiny')
|
||||
await manager.exportArchive('pp-ocrv6-tiny', archive)
|
||||
await manager.remove('pp-ocrv6-tiny')
|
||||
|
||||
await expect(
|
||||
manager.importArchive('pp-ocrv6-tiny', archive)
|
||||
).resolves.toMatchObject({
|
||||
id: 'pp-ocrv6-tiny',
|
||||
source: 'local'
|
||||
})
|
||||
const assets = await manager.getAssets('pp-ocrv6-tiny')
|
||||
expect(new Uint8Array(assets.detection)).toEqual(
|
||||
Uint8Array.from(modelBytes.detection)
|
||||
)
|
||||
expect(new Uint8Array(assets.recognition)).toEqual(
|
||||
Uint8Array.from(modelBytes.recognition)
|
||||
)
|
||||
})
|
||||
})
|
||||
|
||||
describe('extractPaddleCharacterDictionary', () => {
|
||||
it('converts Paddle YAML scalars into the line dictionary used by OCR', () => {
|
||||
const dictionary = extractPaddleCharacterDictionary(
|
||||
new TextDecoder().decode(dictionaryYaml())
|
||||
)
|
||||
expect(dictionary.startsWith("!\n\"\n'\n")).toBe(true)
|
||||
expect(dictionary.split('\n')).toHaveLength(124)
|
||||
})
|
||||
})
|
||||
@@ -0,0 +1,899 @@
|
||||
import { createHash, randomUUID } from 'node:crypto'
|
||||
import {
|
||||
copyFile,
|
||||
lstat,
|
||||
mkdir,
|
||||
open,
|
||||
readFile,
|
||||
readdir,
|
||||
rename,
|
||||
rm,
|
||||
stat,
|
||||
writeFile
|
||||
} from 'node:fs/promises'
|
||||
import { dirname, resolve } from 'node:path'
|
||||
import {
|
||||
documentOcrAssetsSchema,
|
||||
documentOcrModelCatalogEntrySchema,
|
||||
documentOcrModelSnapshotSchema,
|
||||
documentParsingModelStatusSchema,
|
||||
installedDocumentOcrModelSchema,
|
||||
localOcrModelIdSchema,
|
||||
type DocumentOcrAssets,
|
||||
type DocumentOcrModelCatalogEntry,
|
||||
type DocumentOcrModelFile,
|
||||
type DocumentOcrModelOperation,
|
||||
type DocumentOcrModelSnapshot,
|
||||
type InstalledDocumentOcrModel
|
||||
} from '../shared/document-parsing-contracts'
|
||||
import { DOCUMENT_OCR_MODEL_CATALOG } from './document-ocr-model-catalog'
|
||||
import {
|
||||
exportModelArchive,
|
||||
extractModelArchive
|
||||
} from './model-archive'
|
||||
|
||||
const DEFAULT_MAX_FILE_BYTES = 96 * 1024 * 1024
|
||||
const MANIFEST_FILE_NAME = 'manifest.json'
|
||||
const MAX_REDIRECTS = 3
|
||||
const PARTIAL_SUFFIX = '.partial'
|
||||
const MAXIMUM_ARCHIVE_BYTES = 512 * 1024 * 1024
|
||||
const ARCHIVE_OVERHEAD_BYTES = 1024 * 1024
|
||||
const executableExtensionPattern =
|
||||
/\.(?:app|bat|bin|cmd|com|cpl|dll|dmg|exe|hta|inf|ins|iso|jar|js|jse|lnk|msi|msp|mst|pif|ps1|reg|scr|sh|sys|vb|vbe|vbs|ws|wsc|wsf|wsh)$/iu
|
||||
|
||||
type ActiveOperation = {
|
||||
controller: AbortController
|
||||
progress: DocumentOcrModelOperation
|
||||
}
|
||||
|
||||
export type DocumentOcrModelManagerOptions = {
|
||||
userDataDirectory: string
|
||||
fetch: typeof fetch
|
||||
catalog?: readonly DocumentOcrModelCatalogEntry[]
|
||||
maxFileBytes?: number
|
||||
}
|
||||
|
||||
function abortError(): DOMException {
|
||||
return new DOMException('The operation was aborted', 'AbortError')
|
||||
}
|
||||
|
||||
function ensureNotAborted(signal: AbortSignal): void {
|
||||
if (signal.aborted) {
|
||||
throw abortError()
|
||||
}
|
||||
}
|
||||
|
||||
function cloneCatalogEntry(
|
||||
entry: DocumentOcrModelCatalogEntry
|
||||
): DocumentOcrModelCatalogEntry {
|
||||
return documentOcrModelCatalogEntrySchema.parse(entry)
|
||||
}
|
||||
|
||||
function safeChild(parent: string, name: string): string {
|
||||
const child = resolve(parent, name)
|
||||
if (dirname(child) !== resolve(parent)) {
|
||||
throw new Error('OCR 模型路径超出受管目录')
|
||||
}
|
||||
return child
|
||||
}
|
||||
|
||||
function validateDownloadUrl(value: string): URL {
|
||||
const url = new URL(value)
|
||||
if (url.protocol !== 'http:' && url.protocol !== 'https:') {
|
||||
throw new Error('OCR 模型下载地址必须使用 HTTP 或 HTTPS')
|
||||
}
|
||||
return url
|
||||
}
|
||||
|
||||
function toArrayBuffer(buffer: Buffer): ArrayBuffer {
|
||||
return Uint8Array.from(buffer).buffer
|
||||
}
|
||||
|
||||
async function hashFile(
|
||||
path: string,
|
||||
signal?: AbortSignal
|
||||
): Promise<{ size: number; sha256: string }> {
|
||||
const handle = await open(path, 'r')
|
||||
const hash = createHash('sha256')
|
||||
const buffer = Buffer.allocUnsafe(64 * 1024)
|
||||
let size = 0
|
||||
try {
|
||||
while (true) {
|
||||
if (signal) {
|
||||
ensureNotAborted(signal)
|
||||
}
|
||||
const { bytesRead } = await handle.read(buffer, 0, buffer.length)
|
||||
if (bytesRead === 0) {
|
||||
break
|
||||
}
|
||||
hash.update(buffer.subarray(0, bytesRead))
|
||||
size += bytesRead
|
||||
}
|
||||
} finally {
|
||||
await handle.close()
|
||||
}
|
||||
return { size, sha256: hash.digest('hex') }
|
||||
}
|
||||
|
||||
function parseYamlScalar(value: string): string {
|
||||
if (value.startsWith("'") && value.endsWith("'")) {
|
||||
return value.slice(1, -1).replace(/''/gu, "'")
|
||||
}
|
||||
if (value.startsWith('"') && value.endsWith('"')) {
|
||||
return JSON.parse(value) as string
|
||||
}
|
||||
return value
|
||||
}
|
||||
|
||||
export function extractPaddleCharacterDictionary(source: string): string {
|
||||
const characters: string[] = []
|
||||
let readingDictionary = false
|
||||
for (const line of source.replace(/\r/gu, '').split('\n')) {
|
||||
if (line === ' character_dict:') {
|
||||
readingDictionary = true
|
||||
continue
|
||||
}
|
||||
if (!readingDictionary) {
|
||||
continue
|
||||
}
|
||||
const match = /^ {2}- (.*)$/u.exec(line)
|
||||
if (!match) {
|
||||
break
|
||||
}
|
||||
const character = parseYamlScalar(match[1]!)
|
||||
if (!character) {
|
||||
throw new Error('OCR 字符字典包含空条目')
|
||||
}
|
||||
characters.push(character)
|
||||
}
|
||||
if (characters.length < 100) {
|
||||
throw new Error('OCR 字符字典格式无效')
|
||||
}
|
||||
return `${characters.join('\n')}\n`
|
||||
}
|
||||
|
||||
export class DocumentOcrModelManager {
|
||||
readonly rootDirectory: string
|
||||
|
||||
private readonly transport: typeof fetch
|
||||
private readonly catalog: DocumentOcrModelCatalogEntry[]
|
||||
private readonly maxFileBytes: number
|
||||
private readonly operations = new Map<string, ActiveOperation>()
|
||||
private readonly verifiedModels = new Map<string, Promise<void>>()
|
||||
|
||||
constructor(options: DocumentOcrModelManagerOptions) {
|
||||
if (!options.userDataDirectory.trim()) {
|
||||
throw new Error('userDataDirectory is required')
|
||||
}
|
||||
this.rootDirectory = resolve(
|
||||
options.userDataDirectory,
|
||||
'models',
|
||||
'document-ocr'
|
||||
)
|
||||
this.transport = options.fetch
|
||||
this.catalog = (options.catalog ?? DOCUMENT_OCR_MODEL_CATALOG).map(
|
||||
cloneCatalogEntry
|
||||
)
|
||||
if (
|
||||
new Set(this.catalog.map((entry) => entry.id)).size !==
|
||||
this.catalog.length
|
||||
) {
|
||||
throw new Error('OCR 模型目录包含重复 ID')
|
||||
}
|
||||
this.maxFileBytes = options.maxFileBytes ?? DEFAULT_MAX_FILE_BYTES
|
||||
if (
|
||||
!Number.isSafeInteger(this.maxFileBytes) ||
|
||||
this.maxFileBytes <= 0 ||
|
||||
this.maxFileBytes > 512 * 1024 * 1024
|
||||
) {
|
||||
throw new RangeError('maxFileBytes must be a positive safe integer')
|
||||
}
|
||||
}
|
||||
|
||||
async getSnapshot(): Promise<DocumentOcrModelSnapshot> {
|
||||
await this.ensureRoot()
|
||||
return documentOcrModelSnapshotSchema.parse({
|
||||
rootDirectory: this.rootDirectory,
|
||||
catalog: this.catalog.map(cloneCatalogEntry),
|
||||
installed: await this.readInstalled(),
|
||||
operations: [...this.operations.values()].map((operation) => ({
|
||||
...operation.progress
|
||||
}))
|
||||
})
|
||||
}
|
||||
|
||||
async getStatus(
|
||||
modelId: string
|
||||
): Promise<ReturnType<typeof documentParsingModelStatusSchema.parse>> {
|
||||
const entry = this.requireCatalogEntry(modelId)
|
||||
try {
|
||||
await this.getVerifiedStatus(entry)
|
||||
return documentParsingModelStatusSchema.parse({
|
||||
id: entry.id,
|
||||
displayName: entry.displayName,
|
||||
available: true,
|
||||
verified: true,
|
||||
runtime: entry.runtime,
|
||||
detail: '模型已安装并通过 SHA-256 校验，可离线使用'
|
||||
})
|
||||
} catch {
|
||||
return documentParsingModelStatusSchema.parse({
|
||||
id: entry.id,
|
||||
displayName: entry.displayName,
|
||||
available: false,
|
||||
verified: false,
|
||||
runtime: entry.runtime,
|
||||
detail: '模型尚未安装或校验失败，请从 ModelScope 下载'
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
getAssets(modelId: string): Promise<DocumentOcrAssets> {
|
||||
return this.loadVerifiedAssets(this.requireCatalogEntry(modelId))
|
||||
}
|
||||
|
||||
async install(
|
||||
modelId: string,
|
||||
externalSignal?: AbortSignal
|
||||
): Promise<InstalledDocumentOcrModel> {
|
||||
const entry = this.requireCatalogEntry(modelId)
|
||||
const totalBytes = entry.files.reduce(
|
||||
(total, file) => total + file.download.size,
|
||||
0
|
||||
)
|
||||
if (!Number.isSafeInteger(totalBytes)) {
|
||||
throw new RangeError('OCR 模型总大小超出安全范围')
|
||||
}
|
||||
const operation = this.beginOperation(entry.id, 'download', totalBytes)
|
||||
const detachAbort = this.attachExternalSignal(
|
||||
externalSignal,
|
||||
operation.controller
|
||||
)
|
||||
let stagingDirectory: string | undefined
|
||||
try {
|
||||
await this.ensureRoot()
|
||||
await this.assertNotInstalled(entry.id)
|
||||
stagingDirectory = await this.createStagingDirectory(entry.id)
|
||||
for (const file of entry.files) {
|
||||
ensureNotAborted(operation.controller.signal)
|
||||
operation.progress.phase = 'transferring'
|
||||
operation.progress.currentFile = file.name
|
||||
await this.downloadFile(
|
||||
file,
|
||||
safeChild(stagingDirectory, file.name),
|
||||
operation,
|
||||
operation.controller.signal
|
||||
)
|
||||
}
|
||||
operation.progress.phase = 'installing'
|
||||
operation.progress.currentFile = null
|
||||
const installed = await this.createInstalledManifest(
|
||||
entry,
|
||||
'download',
|
||||
stagingDirectory,
|
||||
operation.controller.signal
|
||||
)
|
||||
ensureNotAborted(operation.controller.signal)
|
||||
await rename(stagingDirectory, this.modelDirectory(entry.id))
|
||||
stagingDirectory = undefined
|
||||
this.verifiedModels.delete(entry.id)
|
||||
return installed
|
||||
} finally {
|
||||
detachAbort()
|
||||
this.operations.delete(entry.id)
|
||||
if (stagingDirectory) {
|
||||
await rm(stagingDirectory, { recursive: true, force: true })
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
async registerLocalDirectory(
|
||||
modelId: string,
|
||||
sourceDirectory: string,
|
||||
externalSignal?: AbortSignal
|
||||
): Promise<InstalledDocumentOcrModel> {
|
||||
const entry = this.requireCatalogEntry(modelId)
|
||||
const source = resolve(sourceDirectory)
|
||||
const operation = this.beginOperation(entry.id, 'import', null)
|
||||
const detachAbort = this.attachExternalSignal(
|
||||
externalSignal,
|
||||
operation.controller
|
||||
)
|
||||
let stagingDirectory: string | undefined
|
||||
try {
|
||||
await this.ensureRoot()
|
||||
await this.assertNotInstalled(entry.id)
|
||||
await this.validateLocalDirectory(
|
||||
source,
|
||||
entry,
|
||||
operation.controller.signal
|
||||
)
|
||||
stagingDirectory = await this.createStagingDirectory(entry.id)
|
||||
operation.progress.phase = 'transferring'
|
||||
for (const file of entry.files) {
|
||||
ensureNotAborted(operation.controller.signal)
|
||||
operation.progress.currentFile = file.name
|
||||
const sourceFile = safeChild(source, file.name)
|
||||
const destination = safeChild(stagingDirectory, file.name)
|
||||
await copyFile(sourceFile, destination)
|
||||
operation.progress.completedBytes +=
|
||||
(await stat(destination)).size
|
||||
}
|
||||
operation.progress.totalBytes =
|
||||
operation.progress.completedBytes
|
||||
operation.progress.phase = 'installing'
|
||||
operation.progress.currentFile = null
|
||||
const installed = await this.createInstalledManifest(
|
||||
entry,
|
||||
'local',
|
||||
stagingDirectory,
|
||||
operation.controller.signal
|
||||
)
|
||||
ensureNotAborted(operation.controller.signal)
|
||||
await rename(stagingDirectory, this.modelDirectory(entry.id))
|
||||
stagingDirectory = undefined
|
||||
this.verifiedModels.delete(entry.id)
|
||||
return installed
|
||||
} finally {
|
||||
detachAbort()
|
||||
this.operations.delete(entry.id)
|
||||
if (stagingDirectory) {
|
||||
await rm(stagingDirectory, { recursive: true, force: true })
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
async exportArchive(
|
||||
modelId: string,
|
||||
destinationPath: string
|
||||
): Promise<void> {
|
||||
const entry = this.requireCatalogEntry(modelId)
|
||||
await this.ensureRoot()
|
||||
const installed = (await this.readInstalled()).find(
|
||||
(model) => model.id === entry.id
|
||||
)
|
||||
if (!installed) {
|
||||
throw new Error('只能导出已安装的 OCR 模型')
|
||||
}
|
||||
const directory = this.modelDirectory(entry.id)
|
||||
const files = []
|
||||
for (const expected of entry.files) {
|
||||
const recorded = installed.files.find(
|
||||
(file) =>
|
||||
file.name === expected.name &&
|
||||
file.role === expected.role
|
||||
)
|
||||
if (
|
||||
!recorded ||
|
||||
recorded.size !== expected.download.size ||
|
||||
recorded.sha256 !== expected.download.sha256
|
||||
) {
|
||||
throw new Error(`OCR 模型文件校验失败：${expected.name}`)
|
||||
}
|
||||
files.push({
|
||||
name: expected.name,
|
||||
role: expected.role,
|
||||
size: recorded.size,
|
||||
sha256: recorded.sha256
|
||||
})
|
||||
}
|
||||
await exportModelArchive({
|
||||
destinationPath,
|
||||
sourceDirectory: directory,
|
||||
descriptor: {
|
||||
kind: 'document-ocr',
|
||||
modelId: entry.id,
|
||||
displayName: entry.displayName,
|
||||
files
|
||||
}
|
||||
})
|
||||
}
|
||||
|
||||
async importArchive(
|
||||
modelId: string,
|
||||
archivePath: string
|
||||
): Promise<InstalledDocumentOcrModel> {
|
||||
const entry = this.requireCatalogEntry(modelId)
|
||||
const expectedTotal = entry.files.reduce(
|
||||
(total, file) => total + file.download.size,
|
||||
0
|
||||
)
|
||||
const operation = this.beginOperation(
|
||||
entry.id,
|
||||
'import',
|
||||
expectedTotal
|
||||
)
|
||||
let stagingDirectory: string | undefined
|
||||
try {
|
||||
await this.ensureRoot()
|
||||
await this.assertNotInstalled(entry.id)
|
||||
stagingDirectory = await this.createStagingDirectory(entry.id)
|
||||
operation.progress.phase = 'transferring'
|
||||
const descriptor = await extractModelArchive({
|
||||
archivePath,
|
||||
destinationDirectory: stagingDirectory,
|
||||
expectedKind: 'document-ocr',
|
||||
expectedModelId: entry.id,
|
||||
expectedFiles: entry.files.map((file) => ({
|
||||
name: file.name,
|
||||
role: file.role
|
||||
})),
|
||||
maximumArchiveBytes: Math.min(
|
||||
MAXIMUM_ARCHIVE_BYTES,
|
||||
expectedTotal + ARCHIVE_OVERHEAD_BYTES
|
||||
),
|
||||
maximumFileBytes: this.maxFileBytes,
|
||||
maximumTotalBytes: expectedTotal + ARCHIVE_OVERHEAD_BYTES,
|
||||
signal: operation.controller.signal,
|
||||
onProgress: (completedBytes) => {
|
||||
operation.progress.completedBytes = completedBytes
|
||||
}
|
||||
})
|
||||
for (const expected of entry.files) {
|
||||
const archived = descriptor.files.find(
|
||||
(file) =>
|
||||
file.name === expected.name &&
|
||||
file.role === expected.role
|
||||
)
|
||||
if (
|
||||
!archived ||
|
||||
archived.size !== expected.download.size ||
|
||||
archived.sha256 !== expected.download.sha256
|
||||
) {
|
||||
throw new Error(
|
||||
`OCR 模型 ZIP 与当前模型目录不匹配：${expected.name}`
|
||||
)
|
||||
}
|
||||
}
|
||||
operation.progress.phase = 'installing'
|
||||
operation.progress.currentFile = null
|
||||
const installed = installedDocumentOcrModelSchema.parse({
|
||||
id: entry.id,
|
||||
displayName: entry.displayName,
|
||||
source: 'local',
|
||||
installedAt: new Date().toISOString(),
|
||||
files: descriptor.files
|
||||
})
|
||||
await writeFile(
|
||||
safeChild(stagingDirectory, MANIFEST_FILE_NAME),
|
||||
`${JSON.stringify(installed, null, 2)}\n`,
|
||||
{ encoding: 'utf8', flag: 'wx' }
|
||||
)
|
||||
ensureNotAborted(operation.controller.signal)
|
||||
await rename(stagingDirectory, this.modelDirectory(entry.id))
|
||||
stagingDirectory = undefined
|
||||
this.verifiedModels.delete(entry.id)
|
||||
return installed
|
||||
} finally {
|
||||
this.operations.delete(entry.id)
|
||||
if (stagingDirectory) {
|
||||
await rm(stagingDirectory, { recursive: true, force: true })
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
cancel(modelId: string): boolean {
|
||||
const id = localOcrModelIdSchema.parse(modelId)
|
||||
const operation = this.operations.get(id)
|
||||
if (!operation) {
|
||||
return false
|
||||
}
|
||||
operation.controller.abort()
|
||||
return true
|
||||
}
|
||||
|
||||
async remove(modelId: string): Promise<void> {
|
||||
const id = localOcrModelIdSchema.parse(modelId)
|
||||
this.cancel(id)
|
||||
this.verifiedModels.delete(id)
|
||||
await rm(this.modelDirectory(id), {
|
||||
recursive: true,
|
||||
force: true
|
||||
})
|
||||
}
|
||||
|
||||
dispose(): void {
|
||||
for (const operation of this.operations.values()) {
|
||||
operation.controller.abort()
|
||||
}
|
||||
this.operations.clear()
|
||||
this.verifiedModels.clear()
|
||||
}
|
||||
|
||||
private async ensureRoot(): Promise<void> {
|
||||
await mkdir(this.rootDirectory, { recursive: true })
|
||||
}
|
||||
|
||||
private modelDirectory(modelId: string): string {
|
||||
return safeChild(
|
||||
this.rootDirectory,
|
||||
localOcrModelIdSchema.parse(modelId)
|
||||
)
|
||||
}
|
||||
|
||||
private requireCatalogEntry(
|
||||
modelId: string
|
||||
): DocumentOcrModelCatalogEntry {
|
||||
const id = localOcrModelIdSchema.parse(modelId)
|
||||
const entry = this.catalog.find((candidate) => candidate.id === id)
|
||||
if (!entry) {
|
||||
throw new Error('未知的 OCR 模型')
|
||||
}
|
||||
return entry
|
||||
}
|
||||
|
||||
private beginOperation(
|
||||
modelId: string,
|
||||
kind: DocumentOcrModelOperation['kind'],
|
||||
totalBytes: number | null
|
||||
): ActiveOperation {
|
||||
if (this.operations.has(modelId)) {
|
||||
throw new Error('该 OCR 模型已有进行中的操作')
|
||||
}
|
||||
const operation: ActiveOperation = {
|
||||
controller: new AbortController(),
|
||||
progress: {
|
||||
modelId: localOcrModelIdSchema.parse(modelId),
|
||||
kind,
|
||||
phase: 'preparing',
|
||||
currentFile: null,
|
||||
completedBytes: 0,
|
||||
totalBytes
|
||||
}
|
||||
}
|
||||
this.operations.set(modelId, operation)
|
||||
return operation
|
||||
}
|
||||
|
||||
private attachExternalSignal(
|
||||
signal: AbortSignal | undefined,
|
||||
controller: AbortController
|
||||
): () => void {
|
||||
if (!signal) {
|
||||
return () => undefined
|
||||
}
|
||||
const abort = (): void => controller.abort()
|
||||
if (signal.aborted) {
|
||||
controller.abort()
|
||||
} else {
|
||||
signal.addEventListener('abort', abort, { once: true })
|
||||
}
|
||||
return () => signal.removeEventListener('abort', abort)
|
||||
}
|
||||
|
||||
private async assertNotInstalled(modelId: string): Promise<void> {
|
||||
try {
|
||||
await lstat(this.modelDirectory(modelId))
|
||||
throw new Error('OCR 模型已安装')
|
||||
} catch (error) {
|
||||
if (
|
||||
error instanceof Error &&
|
||||
'code' in error &&
|
||||
error.code === 'ENOENT'
|
||||
) {
|
||||
return
|
||||
}
|
||||
throw error
|
||||
}
|
||||
}
|
||||
|
||||
private async createStagingDirectory(modelId: string): Promise<string> {
|
||||
const directory = safeChild(
|
||||
this.rootDirectory,
|
||||
`.install-${modelId}-${randomUUID()}`
|
||||
)
|
||||
await mkdir(directory, { recursive: false })
|
||||
return directory
|
||||
}
|
||||
|
||||
private async fetchFollowingRedirects(
|
||||
initialUrl: string,
|
||||
signal: AbortSignal
|
||||
): Promise<Response> {
|
||||
let url = validateDownloadUrl(initialUrl)
|
||||
for (let redirectCount = 0; ; redirectCount += 1) {
|
||||
ensureNotAborted(signal)
|
||||
const response = await this.transport(url, {
|
||||
method: 'GET',
|
||||
redirect: 'manual',
|
||||
credentials: 'omit',
|
||||
cache: 'no-store',
|
||||
signal
|
||||
})
|
||||
if ([301, 302, 303, 307, 308].includes(response.status)) {
|
||||
if (redirectCount >= MAX_REDIRECTS) {
|
||||
await response.body?.cancel().catch(() => undefined)
|
||||
throw new Error('OCR 模型下载重定向次数过多')
|
||||
}
|
||||
const location = response.headers.get('location')
|
||||
await response.body?.cancel().catch(() => undefined)
|
||||
if (!location) {
|
||||
throw new Error('OCR 模型下载重定向缺少地址')
|
||||
}
|
||||
url = validateDownloadUrl(new URL(location, url).toString())
|
||||
continue
|
||||
}
|
||||
return response
|
||||
}
|
||||
}
|
||||
|
||||
private async downloadFile(
|
||||
file: DocumentOcrModelFile,
|
||||
destination: string,
|
||||
operation: ActiveOperation,
|
||||
signal: AbortSignal
|
||||
): Promise<void> {
|
||||
if (file.download.size > this.maxFileBytes) {
|
||||
throw new RangeError(`OCR 模型文件过大：${file.name}`)
|
||||
}
|
||||
const response = await this.fetchFollowingRedirects(
|
||||
file.download.url,
|
||||
signal
|
||||
)
|
||||
if (!response.ok) {
|
||||
await response.body?.cancel().catch(() => undefined)
|
||||
throw new Error(`OCR 模型下载失败：HTTP ${response.status}`)
|
||||
}
|
||||
if (!response.body) {
|
||||
throw new Error('OCR 模型下载响应没有内容')
|
||||
}
|
||||
const declaredLength = response.headers.get('content-length')
|
||||
if (
|
||||
declaredLength !== null &&
|
||||
Number(declaredLength) !== file.download.size
|
||||
) {
|
||||
await response.body.cancel().catch(() => undefined)
|
||||
throw new Error(`OCR 模型文件大小不匹配：${file.name}`)
|
||||
}
|
||||
|
||||
const partialPath = `${destination}${PARTIAL_SUFFIX}`
|
||||
const handle = await open(partialPath, 'wx')
|
||||
const reader = response.body.getReader()
|
||||
const hash = createHash('sha256')
|
||||
let written = 0
|
||||
try {
|
||||
while (true) {
|
||||
ensureNotAborted(signal)
|
||||
const result = await reader.read()
|
||||
if (result.done) {
|
||||
break
|
||||
}
|
||||
written += result.value.byteLength
|
||||
if (
|
||||
written > file.download.size ||
|
||||
written > this.maxFileBytes
|
||||
) {
|
||||
await reader.cancel()
|
||||
throw new RangeError(`OCR 模型文件过大：${file.name}`)
|
||||
}
|
||||
await handle.write(result.value)
|
||||
hash.update(result.value)
|
||||
operation.progress.completedBytes += result.value.byteLength
|
||||
}
|
||||
} catch (error) {
|
||||
await reader.cancel().catch(() => undefined)
|
||||
throw error
|
||||
} finally {
|
||||
await handle.close()
|
||||
}
|
||||
if (
|
||||
written !== file.download.size ||
|
||||
hash.digest('hex') !== file.download.sha256
|
||||
) {
|
||||
throw new Error(`OCR 模型文件校验失败：${file.name}`)
|
||||
}
|
||||
await rename(partialPath, destination)
|
||||
}
|
||||
|
||||
private async validateLocalDirectory(
|
||||
sourceDirectory: string,
|
||||
entry: DocumentOcrModelCatalogEntry,
|
||||
signal: AbortSignal
|
||||
): Promise<void> {
|
||||
const sourceInfo = await lstat(sourceDirectory)
|
||||
if (!sourceInfo.isDirectory() || sourceInfo.isSymbolicLink()) {
|
||||
throw new Error('本地 OCR 模型来源必须是普通目录')
|
||||
}
|
||||
const entries = await readdir(sourceDirectory, { withFileTypes: true })
|
||||
for (const localEntry of entries) {
|
||||
ensureNotAborted(signal)
|
||||
if (
|
||||
localEntry.isSymbolicLink() ||
|
||||
executableExtensionPattern.test(localEntry.name)
|
||||
) {
|
||||
throw new Error('本地 OCR 模型目录包含不安全文件')
|
||||
}
|
||||
}
|
||||
for (const file of entry.files) {
|
||||
ensureNotAborted(signal)
|
||||
const path = safeChild(sourceDirectory, file.name)
|
||||
const info = await lstat(path)
|
||||
if (!info.isFile() || info.isSymbolicLink()) {
|
||||
throw new Error(`OCR 模型文件必须是普通文件：${file.name}`)
|
||||
}
|
||||
const actual = await hashFile(path, signal)
|
||||
if (
|
||||
actual.size !== file.download.size ||
|
||||
actual.sha256 !== file.download.sha256
|
||||
) {
|
||||
throw new Error(`本地 OCR 模型文件校验失败：${file.name}`)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
private async createInstalledManifest(
|
||||
entry: DocumentOcrModelCatalogEntry,
|
||||
source: InstalledDocumentOcrModel['source'],
|
||||
stagingDirectory: string,
|
||||
signal: AbortSignal
|
||||
): Promise<InstalledDocumentOcrModel> {
|
||||
const files = []
|
||||
for (const file of entry.files) {
|
||||
ensureNotAborted(signal)
|
||||
files.push({
|
||||
name: file.name,
|
||||
role: file.role,
|
||||
...(await hashFile(
|
||||
safeChild(stagingDirectory, file.name),
|
||||
signal
|
||||
))
|
||||
})
|
||||
}
|
||||
const manifest = installedDocumentOcrModelSchema.parse({
|
||||
id: entry.id,
|
||||
displayName: entry.displayName,
|
||||
source,
|
||||
installedAt: new Date().toISOString(),
|
||||
files
|
||||
})
|
||||
await writeFile(
|
||||
safeChild(stagingDirectory, MANIFEST_FILE_NAME),
|
||||
`${JSON.stringify(manifest, null, 2)}\n`,
|
||||
{ encoding: 'utf8', flag: 'wx' }
|
||||
)
|
||||
return manifest
|
||||
}
|
||||
|
||||
private async readInstalled(): Promise<InstalledDocumentOcrModel[]> {
|
||||
const entries = await readdir(this.rootDirectory, {
|
||||
withFileTypes: true
|
||||
})
|
||||
const installed: InstalledDocumentOcrModel[] = []
|
||||
for (const entry of entries) {
|
||||
if (
|
||||
!entry.isDirectory() ||
|
||||
entry.name.startsWith('.install-') ||
|
||||
!localOcrModelIdSchema.safeParse(entry.name).success
|
||||
) {
|
||||
continue
|
||||
}
|
||||
try {
|
||||
const manifest = installedDocumentOcrModelSchema.parse(
|
||||
JSON.parse(
|
||||
await readFile(
|
||||
safeChild(
|
||||
this.modelDirectory(entry.name),
|
||||
MANIFEST_FILE_NAME
|
||||
),
|
||||
'utf8'
|
||||
)
|
||||
) as unknown
|
||||
)
|
||||
if (manifest.id === entry.name) {
|
||||
installed.push(manifest)
|
||||
}
|
||||
} catch {
|
||||
// Ignore incomplete or externally modified model directories.
|
||||
}
|
||||
}
|
||||
return installed
|
||||
}
|
||||
|
||||
private async readInstalledManifest(
|
||||
entry: DocumentOcrModelCatalogEntry
|
||||
): Promise<InstalledDocumentOcrModel> {
|
||||
const directory = this.modelDirectory(entry.id)
|
||||
const manifest = installedDocumentOcrModelSchema.parse(
|
||||
JSON.parse(
|
||||
await readFile(
|
||||
safeChild(directory, MANIFEST_FILE_NAME),
|
||||
'utf8'
|
||||
)
|
||||
) as unknown
|
||||
)
|
||||
if (manifest.id !== entry.id) {
|
||||
throw new Error('OCR 模型清单 ID 不匹配')
|
||||
}
|
||||
return manifest
|
||||
}
|
||||
|
||||
private async verifyInstalledModel(
|
||||
entry: DocumentOcrModelCatalogEntry
|
||||
): Promise<void> {
|
||||
const directory = this.modelDirectory(entry.id)
|
||||
const manifest = await this.readInstalledManifest(entry)
|
||||
for (const file of entry.files) {
|
||||
const installed = manifest.files.find(
|
||||
(candidate) =>
|
||||
candidate.name === file.name &&
|
||||
candidate.role === file.role
|
||||
)
|
||||
const actual = await hashFile(safeChild(directory, file.name))
|
||||
if (
|
||||
!installed ||
|
||||
actual.size !== file.download.size ||
|
||||
actual.sha256 !== file.download.sha256 ||
|
||||
actual.size !== installed.size ||
|
||||
actual.sha256 !== installed.sha256
|
||||
) {
|
||||
throw new Error(`OCR 模型文件校验失败：${file.name}`)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
private getVerifiedStatus(
|
||||
entry: DocumentOcrModelCatalogEntry
|
||||
): Promise<void> {
|
||||
let verification = this.verifiedModels.get(entry.id)
|
||||
if (!verification) {
|
||||
verification = this.verifyInstalledModel(entry).catch((error) => {
|
||||
this.verifiedModels.delete(entry.id)
|
||||
throw error
|
||||
})
|
||||
this.verifiedModels.set(entry.id, verification)
|
||||
}
|
||||
return verification
|
||||
}
|
||||
|
||||
private async loadVerifiedAssets(
|
||||
entry: DocumentOcrModelCatalogEntry
|
||||
): Promise<DocumentOcrAssets> {
|
||||
const directory = this.modelDirectory(entry.id)
|
||||
const manifest = await this.readInstalledManifest(entry)
|
||||
const loaded = new Map<
|
||||
DocumentOcrModelFile['role'],
|
||||
ArrayBuffer
|
||||
>()
|
||||
for (const file of entry.files) {
|
||||
const installed = manifest.files.find(
|
||||
(candidate) =>
|
||||
candidate.name === file.name &&
|
||||
candidate.role === file.role
|
||||
)
|
||||
const path = safeChild(directory, file.name)
|
||||
const contents = await readFile(path)
|
||||
const actual = {
|
||||
size: contents.byteLength,
|
||||
sha256: createHash('sha256').update(contents).digest('hex')
|
||||
}
|
||||
if (
|
||||
!installed ||
|
||||
actual.size !== file.download.size ||
|
||||
actual.sha256 !== file.download.sha256 ||
|
||||
actual.size !== installed.size ||
|
||||
actual.sha256 !== installed.sha256
|
||||
) {
|
||||
throw new Error(`OCR 模型文件校验失败：${file.name}`)
|
||||
}
|
||||
loaded.set(
|
||||
file.role,
|
||||
file.role === 'dictionary'
|
||||
? toArrayBuffer(
|
||||
Buffer.from(
|
||||
extractPaddleCharacterDictionary(
|
||||
contents.toString('utf8')
|
||||
),
|
||||
'utf8'
|
||||
)
|
||||
)
|
||||
: toArrayBuffer(contents)
|
||||
)
|
||||
}
|
||||
return documentOcrAssetsSchema.parse({
|
||||
modelId: entry.id,
|
||||
detection: loaded.get('detection'),
|
||||
recognition: loaded.get('recognition'),
|
||||
dictionary: loaded.get('dictionary')
|
||||
})
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,130 @@
|
||||
import { describe, expect, it, vi } from 'vitest'
|
||||
import {
|
||||
defaultDocumentParsingSettings
|
||||
} from './document-parsing-settings-store'
|
||||
import { DocumentParsingService } from './document-parsing-service'
|
||||
|
||||
function createPdfFixture(text: string): Buffer {
|
||||
const stream = `BT /F1 18 Tf 50 100 Td (${text}) Tj ET`
|
||||
const objects = [
|
||||
'<< /Type /Catalog /Pages 2 0 R >>',
|
||||
'<< /Type /Pages /Kids [3 0 R] /Count 1 >>',
|
||||
'<< /Type /Page /Parent 2 0 R /MediaBox [0 0 300 200] /Resources << /Font << /F1 4 0 R >> >> /Contents 5 0 R >>',
|
||||
'<< /Type /Font /Subtype /Type1 /BaseFont /Helvetica >>',
|
||||
`<< /Length ${Buffer.byteLength(stream)} >>\nstream\n${stream}\nendstream`
|
||||
]
|
||||
let content = '%PDF-1.4\n'
|
||||
const offsets = [0]
|
||||
for (const [index, object] of objects.entries()) {
|
||||
offsets.push(Buffer.byteLength(content))
|
||||
content += `${index + 1} 0 obj\n${object}\nendobj\n`
|
||||
}
|
||||
const xrefOffset = Buffer.byteLength(content)
|
||||
content += `xref\n0 ${objects.length + 1}\n`
|
||||
content += '0000000000 65535 f \n'
|
||||
content += offsets
|
||||
.slice(1)
|
||||
.map((offset) => `${String(offset).padStart(10, '0')} 00000 n \n`)
|
||||
.join('')
|
||||
content += `trailer\n<< /Size ${objects.length + 1} /Root 1 0 R >>\n`
|
||||
content += `startxref\n${xrefOffset}\n%%EOF\n`
|
||||
return Buffer.from(content)
|
||||
}
|
||||
|
||||
function createService(overrides?: {
|
||||
settings?: Partial<typeof defaultDocumentParsingSettings>
|
||||
}) {
|
||||
const settings = {
|
||||
...defaultDocumentParsingSettings,
|
||||
...overrides?.settings
|
||||
}
|
||||
const recognize = vi.fn(async () => ({
|
||||
requestId: crypto.randomUUID(),
|
||||
sections: [
|
||||
{
|
||||
locator: '第 1 页',
|
||||
content: '扫描件识别正文',
|
||||
confidence: 0.93
|
||||
}
|
||||
],
|
||||
pageCount: 1,
|
||||
warnings: []
|
||||
}))
|
||||
const service = new DocumentParsingService(
|
||||
{
|
||||
get: vi.fn(async () => settings),
|
||||
update: vi.fn(async () => settings)
|
||||
} as never,
|
||||
{
|
||||
getStatus: vi.fn(async () => ({
|
||||
id: 'pp-ocrv6-tiny',
|
||||
displayName: 'PP-OCRv6 Tiny',
|
||||
available: true,
|
||||
verified: true,
|
||||
runtime: 'onnxruntime-web-wasm',
|
||||
detail: '可用'
|
||||
}))
|
||||
} as never,
|
||||
{ recognize } as never
|
||||
)
|
||||
return { recognize, service }
|
||||
}
|
||||
|
||||
describe('DocumentParsingService', () => {
|
||||
it('keeps useful PDF text local without invoking OCR', async () => {
|
||||
const { recognize, service } = createService()
|
||||
|
||||
const parsed = await service.parse(
|
||||
'native.pdf',
|
||||
createPdfFixture('Native PDF body text'),
|
||||
'knowledge-index'
|
||||
)
|
||||
|
||||
expect(parsed.content).toContain('Native PDF body text')
|
||||
expect(parsed.sections[0]?.method).toBe('native')
|
||||
expect(recognize).not.toHaveBeenCalled()
|
||||
})
|
||||
|
||||
it('uses OCR for a PDF without useful text', async () => {
|
||||
const { recognize, service } = createService()
|
||||
|
||||
const parsed = await service.parse(
|
||||
'scan.pdf',
|
||||
createPdfFixture(''),
|
||||
'chat-attachment'
|
||||
)
|
||||
|
||||
expect(parsed.content).toBe('扫描件识别正文')
|
||||
expect(parsed.sections).toEqual([
|
||||
{
|
||||
locator: '第 1 页',
|
||||
content: '扫描件识别正文',
|
||||
method: 'ocr',
|
||||
confidence: 0.93
|
||||
}
|
||||
])
|
||||
expect(recognize).toHaveBeenCalledWith(
|
||||
expect.objectContaining({
|
||||
fileName: 'scan.pdf',
|
||||
modelId: 'pp-ocrv6-tiny',
|
||||
mimeType: 'application/pdf',
|
||||
pageNumbers: [1]
|
||||
})
|
||||
)
|
||||
})
|
||||
|
||||
it('does not use OCR in a fast-text workflow', async () => {
|
||||
const { recognize, service } = createService({
|
||||
settings: { chatWorkflow: 'fast-text' }
|
||||
})
|
||||
|
||||
await expect(
|
||||
service.parse(
|
||||
'scan.pdf',
|
||||
createPdfFixture(''),
|
||||
'chat-attachment'
|
||||
)
|
||||
).rejects.toThrow('未启用 OCR')
|
||||
expect(recognize).not.toHaveBeenCalled()
|
||||
})
|
||||
})
|
||||
@@ -0,0 +1,270 @@
|
||||
import { extname } from 'node:path'
|
||||
import {
|
||||
documentParsingDiagnosticSchema,
|
||||
documentParsingSnapshotSchema,
|
||||
type DocumentParsingDiagnostic,
|
||||
type DocumentParsingPurpose,
|
||||
type DocumentParsingSettings,
|
||||
type DocumentParsingSnapshot
|
||||
} from '../shared/document-parsing-contracts'
|
||||
import type { DocumentOcrBroker } from './document-ocr-broker'
|
||||
import type { DocumentOcrModelManager } from './document-ocr-model-manager'
|
||||
import type { DocumentParsingSettingsStore } from './document-parsing-settings-store'
|
||||
import {
|
||||
DocumentTextUnavailableError,
|
||||
extractPdfTextPages,
|
||||
parseDocument,
|
||||
type ParsedDocument,
|
||||
type ParsedSection,
|
||||
type PdfTextPage
|
||||
} from './knowledge/document-parser'
|
||||
|
||||
const minimumUsefulPdfCharacters = 12
|
||||
const maximumReplacementCharacterRatio = 0.08
|
||||
|
||||
export type ParseDocumentForPurpose = (
|
||||
name: string,
|
||||
buffer: Buffer,
|
||||
purpose: DocumentParsingPurpose,
|
||||
signal?: AbortSignal
|
||||
) => Promise<ParsedDocument>
|
||||
|
||||
function ensureNotAborted(signal?: AbortSignal): void {
|
||||
if (signal?.aborted) {
|
||||
throw signal.reason instanceof Error
|
||||
? signal.reason
|
||||
: new Error('文档解析已取消')
|
||||
}
|
||||
}
|
||||
|
||||
function hasUsefulText(content: string): boolean {
|
||||
const compact = content.replace(/\s+/gu, '')
|
||||
if (compact.length < minimumUsefulPdfCharacters) {
|
||||
return false
|
||||
}
|
||||
const replacementCount = [...compact].filter(
|
||||
(character) => character === '\uFFFD'
|
||||
).length
|
||||
return replacementCount / compact.length <=
|
||||
maximumReplacementCharacterRatio
|
||||
}
|
||||
|
||||
function effectiveOcrMode(
|
||||
settings: DocumentParsingSettings,
|
||||
purpose: DocumentParsingPurpose
|
||||
): DocumentParsingSettings['pdfOcrMode'] {
|
||||
if (
|
||||
(purpose === 'chat-attachment' &&
|
||||
settings.chatWorkflow === 'fast-text') ||
|
||||
(purpose === 'knowledge-index' &&
|
||||
settings.knowledgeWorkflow === 'fast-index')
|
||||
) {
|
||||
return 'disabled'
|
||||
}
|
||||
if (
|
||||
(purpose === 'chat-attachment' &&
|
||||
settings.chatWorkflow === 'high-fidelity') ||
|
||||
(purpose === 'knowledge-index' &&
|
||||
settings.knowledgeWorkflow === 'high-fidelity')
|
||||
) {
|
||||
return 'always'
|
||||
}
|
||||
return settings.pdfOcrMode
|
||||
}
|
||||
|
||||
function buildPdfDocument(
|
||||
name: string,
|
||||
sections: ParsedSection[],
|
||||
pageCount: number,
|
||||
warnings: string[] = []
|
||||
): ParsedDocument {
|
||||
const content = sections
|
||||
.map((section) => section.content)
|
||||
.join('\n\n')
|
||||
.slice(0, 5_000_000)
|
||||
if (!content) {
|
||||
throw new DocumentTextUnavailableError()
|
||||
}
|
||||
return {
|
||||
title: name.replace(/\.[^.]+$/u, ''),
|
||||
sourceFormat: '.pdf',
|
||||
content,
|
||||
sections,
|
||||
pageCount,
|
||||
warnings
|
||||
}
|
||||
}
|
||||
|
||||
function nativePdfSections(pages: PdfTextPage[]): ParsedSection[] {
|
||||
return pages
|
||||
.filter((page) => page.content.length > 0)
|
||||
.map((page) => ({
|
||||
locator: `第 ${page.pageNumber} 页`,
|
||||
content: page.content,
|
||||
method: 'native' as const
|
||||
}))
|
||||
}
|
||||
|
||||
export class DocumentParsingService {
|
||||
constructor(
|
||||
private readonly settingsStore: DocumentParsingSettingsStore,
|
||||
private readonly modelManager: DocumentOcrModelManager,
|
||||
private readonly ocrBroker: DocumentOcrBroker
|
||||
) {}
|
||||
|
||||
async snapshot(): Promise<DocumentParsingSnapshot> {
|
||||
const settings = await this.settingsStore.get()
|
||||
const [localOcr, ocrModels] = await Promise.all([
|
||||
this.modelManager.getStatus(settings.localOcrModelId),
|
||||
this.modelManager.getSnapshot()
|
||||
])
|
||||
return documentParsingSnapshotSchema.parse({
|
||||
settings,
|
||||
status: {
|
||||
nativeParsingAvailable: true,
|
||||
conversionAvailable: false,
|
||||
localOcr
|
||||
},
|
||||
ocrModels
|
||||
})
|
||||
}
|
||||
|
||||
async update(input: unknown): Promise<DocumentParsingSnapshot> {
|
||||
await this.settingsStore.update(input)
|
||||
return this.snapshot()
|
||||
}
|
||||
|
||||
parse: ParseDocumentForPurpose = async (
|
||||
name,
|
||||
buffer,
|
||||
purpose,
|
||||
signal
|
||||
) => {
|
||||
ensureNotAborted(signal)
|
||||
if (extname(name).toLowerCase() !== '.pdf') {
|
||||
return parseDocument(name, buffer)
|
||||
}
|
||||
|
||||
const settings = await this.settingsStore.get()
|
||||
const pages = await extractPdfTextPages(buffer)
|
||||
ensureNotAborted(signal)
|
||||
const mode = effectiveOcrMode(settings, purpose)
|
||||
const pagesWithoutUsefulText = pages
|
||||
.filter((page) => !hasUsefulText(page.content))
|
||||
.map((page) => page.pageNumber)
|
||||
const ocrPageNumbers =
|
||||
mode === 'always'
|
||||
? pages.map((page) => page.pageNumber)
|
||||
: mode === 'auto'
|
||||
? pagesWithoutUsefulText
|
||||
: []
|
||||
|
||||
if (mode === 'disabled' || !settings.localOcrEnabled) {
|
||||
const native = nativePdfSections(pages)
|
||||
if (native.length > 0) {
|
||||
return buildPdfDocument(
|
||||
name,
|
||||
native,
|
||||
pages.length,
|
||||
pagesWithoutUsefulText.length > 0
|
||||
? ['部分页面没有有效文本，当前工作流未启用 OCR']
|
||||
: []
|
||||
)
|
||||
}
|
||||
throw new DocumentTextUnavailableError(
|
||||
'PDF 没有可用文本层，当前工作流未启用 OCR'
|
||||
)
|
||||
}
|
||||
if (ocrPageNumbers.length === 0) {
|
||||
return buildPdfDocument(
|
||||
name,
|
||||
nativePdfSections(pages),
|
||||
pages.length
|
||||
)
|
||||
}
|
||||
if (pages.length > settings.maximumPages) {
|
||||
throw new Error(
|
||||
`PDF 共 ${pages.length} 页，超过本地 OCR 的 ${settings.maximumPages} 页限制`
|
||||
)
|
||||
}
|
||||
const modelStatus = await this.modelManager.getStatus(
|
||||
settings.localOcrModelId
|
||||
)
|
||||
if (!modelStatus.available || !modelStatus.verified) {
|
||||
throw new Error(modelStatus.detail)
|
||||
}
|
||||
|
||||
const ocrRequest = {
|
||||
modelId: settings.localOcrModelId,
|
||||
fileName: name,
|
||||
mimeType: 'application/pdf' as const,
|
||||
data: Uint8Array.from(buffer).buffer,
|
||||
maximumPages: settings.maximumPages,
|
||||
pageNumbers: ocrPageNumbers,
|
||||
pageTimeoutSeconds: settings.pageTimeoutSeconds
|
||||
}
|
||||
const ocr = await (signal
|
||||
? this.ocrBroker.recognize(ocrRequest, signal)
|
||||
: this.ocrBroker.recognize(ocrRequest))
|
||||
ensureNotAborted(signal)
|
||||
const ocrByLocator = new Map(
|
||||
ocr.sections.map((section) => [section.locator, section])
|
||||
)
|
||||
const merged = pages.flatMap((page): ParsedSection[] => {
|
||||
const locator = `第 ${page.pageNumber} 页`
|
||||
const recognized = ocrByLocator.get(locator)
|
||||
if (
|
||||
recognized &&
|
||||
(mode === 'always' || !hasUsefulText(page.content))
|
||||
) {
|
||||
return [
|
||||
{
|
||||
locator,
|
||||
content: recognized.content,
|
||||
method: 'ocr',
|
||||
confidence: recognized.confidence
|
||||
}
|
||||
]
|
||||
}
|
||||
return page.content
|
||||
? [{ locator, content: page.content, method: 'native' }]
|
||||
: []
|
||||
})
|
||||
return buildPdfDocument(name, merged, pages.length, ocr.warnings)
|
||||
}
|
||||
|
||||
async diagnose(
|
||||
name: string,
|
||||
buffer: Buffer,
|
||||
purpose: DocumentParsingPurpose = 'diagnostic'
|
||||
): Promise<DocumentParsingDiagnostic> {
|
||||
const startedAt = Date.now()
|
||||
const parsed = await this.parse(name, buffer, purpose)
|
||||
const ocrPageCount = parsed.sections.filter(
|
||||
(section) => section.method === 'ocr'
|
||||
).length
|
||||
const nativePageCount = parsed.sections.filter(
|
||||
(section) => section.method !== 'ocr'
|
||||
).length
|
||||
return documentParsingDiagnosticSchema.parse({
|
||||
fileName: name,
|
||||
sourceFormat:
|
||||
parsed.sourceFormat.replace(/^\./u, '').toUpperCase() || 'UNKNOWN',
|
||||
pageCount:
|
||||
parsed.sourceFormat === '.pdf'
|
||||
? (parsed.pageCount ?? parsed.sections.length)
|
||||
: 0,
|
||||
ocrPageCount,
|
||||
characterCount: parsed.content.length,
|
||||
method:
|
||||
ocrPageCount > 0 && nativePageCount > 0
|
||||
? 'mixed'
|
||||
: ocrPageCount > 0
|
||||
? 'ocr'
|
||||
: 'native',
|
||||
durationMs: Date.now() - startedAt,
|
||||
preview: parsed.content.slice(0, 2_000),
|
||||
warnings: parsed.warnings
|
||||
})
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,120 @@
|
||||
import {
|
||||
mkdtemp,
|
||||
readFile,
|
||||
readdir,
|
||||
rm,
|
||||
writeFile
|
||||
} from 'node:fs/promises'
|
||||
import { tmpdir } from 'node:os'
|
||||
import { join } from 'node:path'
|
||||
import { afterEach, describe, expect, it } from 'vitest'
|
||||
import {
|
||||
defaultDocumentParsingSettings,
|
||||
DocumentParsingSettingsStore
|
||||
} from './document-parsing-settings-store'
|
||||
|
||||
const temporaryDirectories: string[] = []
|
||||
|
||||
async function createStore(): Promise<{
|
||||
directory: string
|
||||
filePath: string
|
||||
store: DocumentParsingSettingsStore
|
||||
}> {
|
||||
const directory = await mkdtemp(
|
||||
join(tmpdir(), 'goodbuddy-document-parsing-settings-')
|
||||
)
|
||||
temporaryDirectories.push(directory)
|
||||
const filePath = join(directory, 'document-parsing-settings.json')
|
||||
return {
|
||||
directory,
|
||||
filePath,
|
||||
store: new DocumentParsingSettingsStore(filePath)
|
||||
}
|
||||
}
|
||||
|
||||
afterEach(async () => {
|
||||
await Promise.all(
|
||||
temporaryDirectories.splice(0).map((directory) =>
|
||||
rm(directory, { recursive: true, force: true })
|
||||
)
|
||||
)
|
||||
})
|
||||
|
||||
describe('DocumentParsingSettingsStore', () => {
|
||||
it('returns local-first defaults without creating a file', async () => {
|
||||
const { directory, store } = await createStore()
|
||||
|
||||
await expect(store.get()).resolves.toEqual(
|
||||
defaultDocumentParsingSettings
|
||||
)
|
||||
await expect(readdir(directory)).resolves.toEqual([])
|
||||
})
|
||||
|
||||
it('persists a complete versioned settings document', async () => {
|
||||
const { filePath, store } = await createStore()
|
||||
const settings = {
|
||||
...defaultDocumentParsingSettings,
|
||||
chatWorkflow: 'fast-text' as const,
|
||||
maximumPages: 42
|
||||
}
|
||||
|
||||
await expect(store.update(settings)).resolves.toEqual(settings)
|
||||
expect(JSON.parse(await readFile(filePath, 'utf8'))).toEqual({
|
||||
version: 2,
|
||||
...settings
|
||||
})
|
||||
await expect(
|
||||
new DocumentParsingSettingsStore(filePath).get()
|
||||
).resolves.toEqual(settings)
|
||||
})
|
||||
|
||||
it('migrates legacy cloud permissions to the local OCR provider', async () => {
|
||||
const { filePath, store } = await createStore()
|
||||
const {
|
||||
ocrProvider: _ocrProvider,
|
||||
...legacySettings
|
||||
} = defaultDocumentParsingSettings
|
||||
void _ocrProvider
|
||||
await writeFile(
|
||||
filePath,
|
||||
JSON.stringify({
|
||||
version: 1,
|
||||
...legacySettings,
|
||||
chatCloudPermission: 'always',
|
||||
knowledgeCloudPermission: 'never'
|
||||
}),
|
||||
'utf8'
|
||||
)
|
||||
|
||||
await expect(store.get()).resolves.toEqual(
|
||||
defaultDocumentParsingSettings
|
||||
)
|
||||
})
|
||||
|
||||
it('rejects incomplete or out-of-range settings', async () => {
|
||||
const { directory, store } = await createStore()
|
||||
|
||||
await expect(store.update({})).rejects.toThrow()
|
||||
await expect(
|
||||
store.update({
|
||||
...defaultDocumentParsingSettings,
|
||||
maximumPages: 0
|
||||
})
|
||||
).rejects.toThrow()
|
||||
await expect(readdir(directory)).resolves.toEqual([])
|
||||
})
|
||||
|
||||
it('isolates corrupt settings and restores defaults', async () => {
|
||||
const { directory, filePath, store } = await createStore()
|
||||
await writeFile(filePath, '{not-json', 'utf8')
|
||||
|
||||
await expect(store.get()).resolves.toEqual(
|
||||
defaultDocumentParsingSettings
|
||||
)
|
||||
const entries = await readdir(directory)
|
||||
expect(entries).toHaveLength(1)
|
||||
expect(entries[0]).toMatch(
|
||||
/^document-parsing-settings\.json\.corrupt-\d+-[a-f0-9]{12}$/u
|
||||
)
|
||||
})
|
||||
})
|
||||
@@ -0,0 +1,181 @@
|
||||
import { randomBytes } from 'node:crypto'
|
||||
import {
|
||||
mkdir,
|
||||
readFile,
|
||||
rename,
|
||||
rm,
|
||||
writeFile
|
||||
} from 'node:fs/promises'
|
||||
import { dirname } from 'node:path'
|
||||
import { z } from 'zod'
|
||||
import {
|
||||
documentParsingSettingsSchema,
|
||||
documentParsingSettingsUpdateSchema,
|
||||
type DocumentParsingSettings
|
||||
} from '../shared/document-parsing-contracts'
|
||||
|
||||
const CURRENT_SETTINGS_VERSION = 2
|
||||
|
||||
const storedDocumentParsingSettingsSchema =
|
||||
documentParsingSettingsSchema
|
||||
.extend({
|
||||
version: z.literal(CURRENT_SETTINGS_VERSION)
|
||||
})
|
||||
.strict()
|
||||
|
||||
type StoredDocumentParsingSettings = z.infer<
|
||||
typeof storedDocumentParsingSettingsSchema
|
||||
>
|
||||
|
||||
const legacyDocumentParsingSettingsSchema =
|
||||
documentParsingSettingsSchema
|
||||
.omit({ ocrProvider: true })
|
||||
.extend({
|
||||
version: z.literal(1),
|
||||
chatCloudPermission: z.enum(['ask', 'always', 'never']),
|
||||
knowledgeCloudPermission: z.enum(['ask', 'always', 'never'])
|
||||
})
|
||||
.strict()
|
||||
|
||||
export const defaultDocumentParsingSettings: DocumentParsingSettings = {
|
||||
chatWorkflow: 'auto',
|
||||
knowledgeWorkflow: 'complete-index',
|
||||
pdfOcrMode: 'auto',
|
||||
ocrProvider: 'local',
|
||||
localOcrEnabled: true,
|
||||
localOcrModelId: 'pp-ocrv6-tiny',
|
||||
maximumPages: 100,
|
||||
ocrConcurrency: 1,
|
||||
pageTimeoutSeconds: 60
|
||||
}
|
||||
|
||||
function isMissingFile(error: unknown): boolean {
|
||||
return (
|
||||
error !== null &&
|
||||
typeof error === 'object' &&
|
||||
'code' in error &&
|
||||
error.code === 'ENOENT'
|
||||
)
|
||||
}
|
||||
|
||||
export class DocumentParsingSettingsStore {
|
||||
private settings?: StoredDocumentParsingSettings
|
||||
private updateQueue: Promise<void> = Promise.resolve()
|
||||
|
||||
constructor(private readonly filePath: string) {}
|
||||
|
||||
private async isolateCorruptFile(): Promise<void> {
|
||||
const isolatedPath =
|
||||
`${this.filePath}.corrupt-${Date.now()}-` +
|
||||
randomBytes(6).toString('hex')
|
||||
try {
|
||||
await rename(this.filePath, isolatedPath)
|
||||
} catch (error) {
|
||||
if (!isMissingFile(error)) {
|
||||
throw new Error('文档解析设置损坏且无法隔离', {
|
||||
cause: error
|
||||
})
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
private async loadStored(): Promise<StoredDocumentParsingSettings> {
|
||||
if (this.settings) {
|
||||
return this.settings
|
||||
}
|
||||
try {
|
||||
const contents = await readFile(this.filePath, 'utf8')
|
||||
let parsed: unknown
|
||||
try {
|
||||
parsed = JSON.parse(contents) as unknown
|
||||
} catch {
|
||||
await this.isolateCorruptFile()
|
||||
this.settings = {
|
||||
version: CURRENT_SETTINGS_VERSION,
|
||||
...defaultDocumentParsingSettings
|
||||
}
|
||||
return this.settings
|
||||
}
|
||||
const result =
|
||||
storedDocumentParsingSettingsSchema.safeParse(parsed)
|
||||
if (!result.success) {
|
||||
const legacy =
|
||||
legacyDocumentParsingSettingsSchema.safeParse(parsed)
|
||||
if (legacy.success) {
|
||||
const {
|
||||
version: _version,
|
||||
chatCloudPermission: _chatCloudPermission,
|
||||
knowledgeCloudPermission: _knowledgeCloudPermission,
|
||||
...settings
|
||||
} = legacy.data
|
||||
void _version
|
||||
void _chatCloudPermission
|
||||
void _knowledgeCloudPermission
|
||||
this.settings = {
|
||||
version: CURRENT_SETTINGS_VERSION,
|
||||
ocrProvider: 'local',
|
||||
...settings
|
||||
}
|
||||
return this.settings
|
||||
}
|
||||
await this.isolateCorruptFile()
|
||||
this.settings = {
|
||||
version: CURRENT_SETTINGS_VERSION,
|
||||
...defaultDocumentParsingSettings
|
||||
}
|
||||
return this.settings
|
||||
}
|
||||
this.settings = result.data
|
||||
} catch (error) {
|
||||
if (!isMissingFile(error)) {
|
||||
throw new Error('无法读取文档解析设置', { cause: error })
|
||||
}
|
||||
this.settings = {
|
||||
version: CURRENT_SETTINGS_VERSION,
|
||||
...defaultDocumentParsingSettings
|
||||
}
|
||||
}
|
||||
return this.settings
|
||||
}
|
||||
|
||||
async get(): Promise<DocumentParsingSettings> {
|
||||
const { version: _version, ...settings } = await this.loadStored()
|
||||
void _version
|
||||
return documentParsingSettingsSchema.parse(settings)
|
||||
}
|
||||
|
||||
update(input: unknown): Promise<DocumentParsingSettings> {
|
||||
const operation = this.updateQueue.then(async () => {
|
||||
const updates = documentParsingSettingsUpdateSchema.parse(input)
|
||||
const next: StoredDocumentParsingSettings = {
|
||||
version: CURRENT_SETTINGS_VERSION,
|
||||
...updates
|
||||
}
|
||||
await mkdir(dirname(this.filePath), { recursive: true })
|
||||
const temporaryPath =
|
||||
`${this.filePath}.${process.pid}.` +
|
||||
`${randomBytes(6).toString('hex')}.tmp`
|
||||
try {
|
||||
await writeFile(
|
||||
temporaryPath,
|
||||
`${JSON.stringify(next, null, 2)}\n`,
|
||||
{
|
||||
encoding: 'utf8',
|
||||
mode: 0o600,
|
||||
flag: 'wx'
|
||||
}
|
||||
)
|
||||
await rename(temporaryPath, this.filePath)
|
||||
} finally {
|
||||
await rm(temporaryPath, { force: true })
|
||||
}
|
||||
this.settings = next
|
||||
return this.get()
|
||||
})
|
||||
this.updateQueue = operation.then(
|
||||
() => undefined,
|
||||
() => undefined
|
||||
)
|
||||
return operation
|
||||
}
|
||||
}
|
||||
+32
-4
@@ -67,6 +67,10 @@ import { KnowledgeEmbeddingIndexRepository } from './knowledge/knowledge-embeddi
|
||||
import { GlobalTlsPolicy } from './global-tls-policy'
|
||||
import type { AgentRuntimeSelection } from '../shared/runtime-selection-contracts'
|
||||
import { waitForCleanup } from './shutdown'
|
||||
import { DocumentParsingSettingsStore } from './document-parsing-settings-store'
|
||||
import { DocumentOcrModelManager } from './document-ocr-model-manager'
|
||||
import { DocumentOcrBroker } from './document-ocr-broker'
|
||||
import { DocumentParsingService } from './document-parsing-service'
|
||||
|
||||
const shortcut = 'CommandOrControl+Shift+Space'
|
||||
const mainModuleDirectory = dirname(fileURLToPath(import.meta.url))
|
||||
@@ -98,6 +102,8 @@ let knowledgeGateway: KnowledgeMcpGateway | undefined
|
||||
let assistantDatabase: AssistantDatabase | undefined
|
||||
let browserService: BrowserService | undefined
|
||||
let globalTlsPolicy: GlobalTlsPolicy | undefined
|
||||
let documentOcrBroker: DocumentOcrBroker | undefined
|
||||
let documentOcrModelManager: DocumentOcrModelManager | undefined
|
||||
|
||||
function createEmbeddingProvider(
|
||||
settings: ResolvedRuntimeSettings
|
||||
@@ -340,6 +346,20 @@ if (hasSingleInstanceLock) {
|
||||
const applicationSettingsStore = new ApplicationSettingsStore(
|
||||
join(app.getPath('userData'), 'application-settings.json')
|
||||
)
|
||||
const documentParsingSettingsStore =
|
||||
new DocumentParsingSettingsStore(
|
||||
join(app.getPath('userData'), 'document-parsing-settings.json')
|
||||
)
|
||||
documentOcrModelManager = new DocumentOcrModelManager({
|
||||
userDataDirectory: app.getPath('userData'),
|
||||
fetch: globalThis.fetch
|
||||
})
|
||||
documentOcrBroker = new DocumentOcrBroker(mainWindow)
|
||||
const documentParsingService = new DocumentParsingService(
|
||||
documentParsingSettingsStore,
|
||||
documentOcrModelManager,
|
||||
documentOcrBroker
|
||||
)
|
||||
const versionChecker = new VersionChecker({
|
||||
fetch: globalThis.fetch,
|
||||
currentVersion: app.getVersion(),
|
||||
@@ -362,7 +382,8 @@ if (hasSingleInstanceLock) {
|
||||
knowledgeService = new KnowledgeService({
|
||||
databasePath: join(app.getPath('userData'), 'knowledge.sqlite'),
|
||||
managedRoot: join(app.getPath('userData'), 'knowledge'),
|
||||
extractStructured: createModelGraphExtractor(settingsStore)
|
||||
extractStructured: createModelGraphExtractor(settingsStore),
|
||||
parseDocument: documentParsingService.parse
|
||||
})
|
||||
await knowledgeService.initialize()
|
||||
const embeddingIndexCoordinator = new EmbeddingIndexCoordinator(
|
||||
@@ -459,7 +480,9 @@ if (hasSingleInstanceLock) {
|
||||
selectedRuntimeManager = new SelectedRuntimeManager(
|
||||
createSelectedRuntime
|
||||
)
|
||||
const contextManager = new ContextManager()
|
||||
const contextManager = new ContextManager({
|
||||
parseDocument: documentParsingService.parse
|
||||
})
|
||||
const approvalBroker = new ToolApprovalBroker()
|
||||
|
||||
const shortcutRegistered = globalShortcut.register(shortcut, () => {
|
||||
@@ -510,7 +533,10 @@ if (hasSingleInstanceLock) {
|
||||
selectedRuntimeManager,
|
||||
speechTranscriptionService,
|
||||
knowledgeGateway,
|
||||
launchWechatSidecar
|
||||
launchWechatSidecar,
|
||||
documentParsingService,
|
||||
documentOcrModelManager,
|
||||
documentOcrBroker
|
||||
)
|
||||
loadMainWindow(mainWindow)
|
||||
|
||||
@@ -550,7 +576,9 @@ app.on('before-quit', (event) => {
|
||||
Promise.resolve().then(() => knowledgeGateway?.dispose()),
|
||||
Promise.resolve().then(() => knowledgeService?.dispose()),
|
||||
Promise.resolve().then(() => browserService?.dispose()),
|
||||
Promise.resolve().then(() => globalTlsPolicy?.dispose())
|
||||
Promise.resolve().then(() => globalTlsPolicy?.dispose()),
|
||||
Promise.resolve().then(() => documentOcrModelManager?.dispose()),
|
||||
Promise.resolve().then(() => documentOcrBroker?.dispose())
|
||||
])
|
||||
globalShortcut.unregisterAll()
|
||||
tray?.destroy()
|
||||
|
||||
+180
-1
@@ -24,6 +24,10 @@ const electronMocks = vi.hoisted(() => {
|
||||
canceled: true,
|
||||
filePaths: [] as string[]
|
||||
})),
|
||||
showSaveDialog: vi.fn(async () => ({
|
||||
canceled: true,
|
||||
filePath: undefined as string | undefined
|
||||
})),
|
||||
openPath: vi.fn(async () => ''),
|
||||
showItemInFolder: vi.fn(),
|
||||
openExternal: vi.fn(async () => undefined)
|
||||
@@ -248,7 +252,8 @@ vi.mock('electron', () => ({
|
||||
},
|
||||
BrowserWindow: class {},
|
||||
dialog: {
|
||||
showOpenDialog: electronMocks.showOpenDialog
|
||||
showOpenDialog: electronMocks.showOpenDialog,
|
||||
showSaveDialog: electronMocks.showSaveDialog
|
||||
},
|
||||
ipcMain: {
|
||||
handle: electronMocks.handle,
|
||||
@@ -290,6 +295,180 @@ vi.mock('./channels/channel-env', () => ({
|
||||
)
|
||||
}))
|
||||
|
||||
describe('registerIpcHandlers model ZIP dialogs', () => {
|
||||
afterEach(() => {
|
||||
electronMocks.handlers.clear()
|
||||
vi.clearAllMocks()
|
||||
})
|
||||
|
||||
it('imports and exports speech and OCR ZIPs through trusted dialogs', async () => {
|
||||
const webContents = {
|
||||
mainFrame: { url: 'file:///goodbuddy/index.html' },
|
||||
getURL: vi.fn(() => 'file:///goodbuddy/index.html'),
|
||||
send: vi.fn()
|
||||
}
|
||||
const window = {
|
||||
webContents,
|
||||
isDestroyed: vi.fn(() => false),
|
||||
isMaximized: vi.fn(() => false),
|
||||
on: vi.fn(),
|
||||
removeListener: vi.fn()
|
||||
}
|
||||
const event = {
|
||||
sender: webContents,
|
||||
senderFrame: webContents.mainFrame
|
||||
}
|
||||
const speechSnapshot = {
|
||||
catalog: [],
|
||||
installed: [],
|
||||
operations: []
|
||||
}
|
||||
const speechModelManager = {
|
||||
rootDirectory: 'C:\\models\\speech',
|
||||
importArchive: vi.fn(async () => speechSnapshot),
|
||||
exportArchive: vi.fn(async () => undefined),
|
||||
getSnapshot: vi.fn(async () => speechSnapshot),
|
||||
cancel: vi.fn()
|
||||
}
|
||||
const ocrSnapshot = {
|
||||
settings: {},
|
||||
models: {
|
||||
catalog: [],
|
||||
installed: [],
|
||||
operations: []
|
||||
}
|
||||
}
|
||||
const documentParsingService = {
|
||||
snapshot: vi.fn(async () => ocrSnapshot)
|
||||
}
|
||||
const documentOcrModelManager = {
|
||||
importArchive: vi.fn(async () => undefined),
|
||||
exportArchive: vi.fn(async () => undefined)
|
||||
}
|
||||
const dispose = registerIpcHandlers(
|
||||
window as never,
|
||||
{ capability: 'text' } as never,
|
||||
'CommandOrControl+Shift+Space',
|
||||
{} as never,
|
||||
{} as never,
|
||||
{ clear: vi.fn() } as never,
|
||||
{} as never,
|
||||
{ claimDueSchedules: vi.fn(() => []) } as never,
|
||||
{ clear: vi.fn() } as never,
|
||||
{} as never,
|
||||
vi.fn(async () => undefined),
|
||||
undefined,
|
||||
undefined,
|
||||
undefined,
|
||||
undefined,
|
||||
undefined,
|
||||
undefined,
|
||||
speechModelManager as never,
|
||||
undefined,
|
||||
undefined,
|
||||
undefined,
|
||||
undefined,
|
||||
undefined,
|
||||
documentParsingService as never,
|
||||
documentOcrModelManager as never
|
||||
)
|
||||
|
||||
await expect(
|
||||
electronMocks.handlers.get(
|
||||
ipcChannels.speechModelsImportArchive
|
||||
)?.(event, { modelId: 'speech-model' })
|
||||
).resolves.toBeUndefined()
|
||||
expect(speechModelManager.importArchive).not.toHaveBeenCalled()
|
||||
|
||||
electronMocks.showOpenDialog.mockResolvedValueOnce({
|
||||
canceled: false,
|
||||
filePaths: ['C:\\transfer\\speech.zip']
|
||||
})
|
||||
await expect(
|
||||
electronMocks.handlers.get(
|
||||
ipcChannels.speechModelsImportArchive
|
||||
)?.(event, { modelId: 'speech-model' })
|
||||
).resolves.toBe(speechSnapshot)
|
||||
expect(electronMocks.showOpenDialog).toHaveBeenLastCalledWith(
|
||||
window,
|
||||
expect.objectContaining({
|
||||
properties: ['openFile'],
|
||||
filters: [
|
||||
{
|
||||
name: 'GoodBuddy 模型 ZIP',
|
||||
extensions: ['zip']
|
||||
}
|
||||
]
|
||||
})
|
||||
)
|
||||
expect(speechModelManager.importArchive).toHaveBeenCalledWith(
|
||||
'speech-model',
|
||||
'C:\\transfer\\speech.zip'
|
||||
)
|
||||
|
||||
electronMocks.showSaveDialog.mockResolvedValueOnce({
|
||||
canceled: false,
|
||||
filePath: 'C:\\transfer\\speech-model'
|
||||
})
|
||||
await expect(
|
||||
electronMocks.handlers.get(
|
||||
ipcChannels.speechModelsExportArchive
|
||||
)?.(event, { modelId: 'speech-model' })
|
||||
).resolves.toBe(speechSnapshot)
|
||||
expect(speechModelManager.exportArchive).toHaveBeenCalledWith(
|
||||
'speech-model',
|
||||
'C:\\transfer\\speech-model.zip'
|
||||
)
|
||||
|
||||
electronMocks.showOpenDialog.mockResolvedValueOnce({
|
||||
canceled: false,
|
||||
filePaths: ['C:\\transfer\\ocr.zip']
|
||||
})
|
||||
await expect(
|
||||
electronMocks.handlers.get(
|
||||
ipcChannels.documentOcrModelsImportArchive
|
||||
)?.(event, { modelId: 'ocr-model' })
|
||||
).resolves.toBe(ocrSnapshot)
|
||||
expect(documentOcrModelManager.importArchive).toHaveBeenCalledWith(
|
||||
'ocr-model',
|
||||
'C:\\transfer\\ocr.zip'
|
||||
)
|
||||
|
||||
electronMocks.showSaveDialog.mockResolvedValueOnce({
|
||||
canceled: false,
|
||||
filePath: 'C:\\transfer\\ocr-model.ZIP'
|
||||
})
|
||||
await expect(
|
||||
electronMocks.handlers.get(
|
||||
ipcChannels.documentOcrModelsExportArchive
|
||||
)?.(event, { modelId: 'ocr-model' })
|
||||
).resolves.toBe(ocrSnapshot)
|
||||
expect(documentOcrModelManager.exportArchive).toHaveBeenCalledWith(
|
||||
'ocr-model',
|
||||
'C:\\transfer\\ocr-model.ZIP'
|
||||
)
|
||||
|
||||
await expect(
|
||||
electronMocks.handlers.get(
|
||||
ipcChannels.speechModelsExportArchive
|
||||
)?.(
|
||||
{
|
||||
sender: {},
|
||||
senderFrame: webContents.mainFrame
|
||||
},
|
||||
{ modelId: 'speech-model' }
|
||||
)
|
||||
).rejects.toThrow('拒绝来自未知窗口的 IPC 请求')
|
||||
await expect(
|
||||
electronMocks.handlers.get(
|
||||
ipcChannels.documentOcrModelsImportArchive
|
||||
)?.(event, {})
|
||||
).rejects.toThrow()
|
||||
|
||||
await dispose()
|
||||
})
|
||||
})
|
||||
|
||||
describe('registerIpcHandlers connection tests', () => {
|
||||
afterEach(() => {
|
||||
electronMocks.handlers.clear()
|
||||
|
||||
+280
-6
@@ -73,6 +73,12 @@ import {
|
||||
embeddingIndexJobRequestSchema,
|
||||
embeddingSettingsSnapshotSchema
|
||||
} from '../shared/embedding-contracts'
|
||||
import {
|
||||
documentOcrModelActionInputSchema,
|
||||
documentOcrFailureSchema,
|
||||
documentOcrResultSchema,
|
||||
documentParsingSettingsUpdateSchema
|
||||
} from '../shared/document-parsing-contracts'
|
||||
import {
|
||||
agentRuntimeSelectionSchema,
|
||||
type AgentRuntimeSelection
|
||||
@@ -178,6 +184,9 @@ import type { VersionChecker } from './version-checker'
|
||||
import type { SpeechModelManager } from './speech/speech-model-manager'
|
||||
import type { SpeechTranscriptionService } from './speech/speech-transcription-service'
|
||||
import type { EmbeddingIndexCoordinator } from './knowledge/embedding-index-coordinator'
|
||||
import type { DocumentParsingService } from './document-parsing-service'
|
||||
import type { DocumentOcrModelManager } from './document-ocr-model-manager'
|
||||
import type { DocumentOcrBroker } from './document-ocr-broker'
|
||||
import { OpenAIEmbeddingClient } from './knowledge/openai-embedding-client'
|
||||
import {
|
||||
magicNotePlainText,
|
||||
@@ -321,6 +330,17 @@ const taskStatusRequestSchema = z
|
||||
status: z.enum(['completed', 'cancelled'])
|
||||
})
|
||||
.strict()
|
||||
|
||||
const modelArchiveDialogFilters = [
|
||||
{
|
||||
name: 'GoodBuddy 模型 ZIP',
|
||||
extensions: ['zip']
|
||||
}
|
||||
]
|
||||
|
||||
function ensureZipExtension(path: string): string {
|
||||
return extname(path).toLowerCase() === '.zip' ? path : `${path}.zip`
|
||||
}
|
||||
const expertUpdateRequestSchema = z
|
||||
.object({
|
||||
expertId: assistantIdSchema,
|
||||
@@ -575,7 +595,10 @@ export function registerIpcHandlers(
|
||||
selectedRuntimes?: SelectedRuntimeResolver,
|
||||
speechTranscriptionService?: SpeechTranscriptionService,
|
||||
knowledgeGateway?: KnowledgeMcpGateway,
|
||||
launchWechatSidecar?: WechatSidecarLauncher
|
||||
launchWechatSidecar?: WechatSidecarLauncher,
|
||||
documentParsingService?: DocumentParsingService,
|
||||
documentOcrModelManager?: DocumentOcrModelManager,
|
||||
documentOcrBroker?: DocumentOcrBroker
|
||||
): () => Promise<void> {
|
||||
const activeRequests = new Map<string, AbortController>()
|
||||
const pendingAgentQuestions = new Map<
|
||||
@@ -2463,6 +2486,233 @@ export function registerIpcHandlers(
|
||||
}
|
||||
)
|
||||
|
||||
ipcMain.handle(ipcChannels.documentParsingGet, (event) => {
|
||||
assertTrustedSender(event, window)
|
||||
if (!documentParsingService) {
|
||||
throw new Error('文档解析设置服务不可用')
|
||||
}
|
||||
return documentParsingService.snapshot()
|
||||
})
|
||||
|
||||
ipcMain.handle(
|
||||
ipcChannels.documentParsingUpdate,
|
||||
(event, input: unknown) => {
|
||||
assertTrustedSender(event, window)
|
||||
if (!documentParsingService) {
|
||||
throw new Error('文档解析设置服务不可用')
|
||||
}
|
||||
return documentParsingService.update(
|
||||
documentParsingSettingsUpdateSchema.parse(input)
|
||||
)
|
||||
}
|
||||
)
|
||||
|
||||
ipcMain.handle(
|
||||
ipcChannels.documentParsingTest,
|
||||
async (event) => {
|
||||
assertTrustedSender(event, window)
|
||||
if (!documentParsingService) {
|
||||
throw new Error('文档解析设置服务不可用')
|
||||
}
|
||||
const result = await dialog.showOpenDialog(window, {
|
||||
title: '选择测试文档',
|
||||
properties: ['openFile'],
|
||||
filters: [
|
||||
{
|
||||
name: '支持的文档',
|
||||
extensions: supportedDocumentExtensions.map((extension) =>
|
||||
extension.slice(1)
|
||||
)
|
||||
}
|
||||
]
|
||||
})
|
||||
const selectedPath = result.filePaths[0]
|
||||
if (result.canceled || !selectedPath) {
|
||||
return undefined
|
||||
}
|
||||
try {
|
||||
const canonicalPath = await realpath(selectedPath)
|
||||
const fileStat = await stat(canonicalPath)
|
||||
if (!fileStat.isFile() || fileStat.size > 20 * 1024 * 1024) {
|
||||
throw new Error('测试文档必须小于 20MB 且不能是目录')
|
||||
}
|
||||
return documentParsingService.diagnose(
|
||||
basename(canonicalPath),
|
||||
await readFile(canonicalPath)
|
||||
)
|
||||
} catch (error) {
|
||||
if (error instanceof Error && !('code' in error)) {
|
||||
throw error
|
||||
}
|
||||
throw new Error('无法读取测试文档，请检查文件权限和状态', {
|
||||
cause: error
|
||||
})
|
||||
}
|
||||
}
|
||||
)
|
||||
|
||||
ipcMain.handle(
|
||||
ipcChannels.documentOcrModelsInstall,
|
||||
(event, input: unknown) => {
|
||||
assertTrustedSender(event, window)
|
||||
if (!documentOcrModelManager || !documentParsingService) {
|
||||
throw new Error('本地 OCR 模型服务不可用')
|
||||
}
|
||||
const { modelId } =
|
||||
documentOcrModelActionInputSchema.parse(input)
|
||||
return trackExecution(
|
||||
documentOcrModelManager
|
||||
.install(modelId)
|
||||
.then(() => documentParsingService.snapshot())
|
||||
)
|
||||
}
|
||||
)
|
||||
|
||||
ipcMain.handle(
|
||||
ipcChannels.documentOcrModelsCancel,
|
||||
(event, input: unknown) => {
|
||||
assertTrustedSender(event, window)
|
||||
if (!documentOcrModelManager) {
|
||||
throw new Error('本地 OCR 模型服务不可用')
|
||||
}
|
||||
const { modelId } =
|
||||
documentOcrModelActionInputSchema.parse(input)
|
||||
return documentOcrModelManager.cancel(modelId)
|
||||
}
|
||||
)
|
||||
|
||||
ipcMain.handle(
|
||||
ipcChannels.documentOcrModelsRemove,
|
||||
async (event, input: unknown) => {
|
||||
assertTrustedSender(event, window)
|
||||
if (!documentOcrModelManager || !documentParsingService) {
|
||||
throw new Error('本地 OCR 模型服务不可用')
|
||||
}
|
||||
const { modelId } =
|
||||
documentOcrModelActionInputSchema.parse(input)
|
||||
await documentOcrModelManager.remove(modelId)
|
||||
return documentParsingService.snapshot()
|
||||
}
|
||||
)
|
||||
|
||||
ipcMain.handle(
|
||||
ipcChannels.documentOcrModelsImportArchive,
|
||||
async (event, input: unknown) => {
|
||||
assertTrustedSender(event, window)
|
||||
if (!documentOcrModelManager || !documentParsingService) {
|
||||
throw new Error('本地 OCR 模型服务不可用')
|
||||
}
|
||||
const { modelId } =
|
||||
documentOcrModelActionInputSchema.parse(input)
|
||||
const result = await dialog.showOpenDialog(window, {
|
||||
title: '导入 OCR 模型 ZIP',
|
||||
properties: ['openFile'],
|
||||
filters: modelArchiveDialogFilters
|
||||
})
|
||||
const archivePath = result.filePaths[0]
|
||||
if (result.canceled || !archivePath) {
|
||||
return undefined
|
||||
}
|
||||
return trackExecution(
|
||||
documentOcrModelManager
|
||||
.importArchive(modelId, archivePath)
|
||||
.then(() => documentParsingService.snapshot())
|
||||
)
|
||||
}
|
||||
)
|
||||
|
||||
ipcMain.handle(
|
||||
ipcChannels.documentOcrModelsExportArchive,
|
||||
async (event, input: unknown) => {
|
||||
assertTrustedSender(event, window)
|
||||
if (!documentOcrModelManager || !documentParsingService) {
|
||||
throw new Error('本地 OCR 模型服务不可用')
|
||||
}
|
||||
const { modelId } =
|
||||
documentOcrModelActionInputSchema.parse(input)
|
||||
const result = await dialog.showSaveDialog(window, {
|
||||
title: '导出 OCR 模型 ZIP',
|
||||
defaultPath: `${modelId}.zip`,
|
||||
filters: modelArchiveDialogFilters
|
||||
})
|
||||
if (result.canceled || !result.filePath) {
|
||||
return undefined
|
||||
}
|
||||
const destination = ensureZipExtension(result.filePath)
|
||||
await documentOcrModelManager.exportArchive(
|
||||
modelId,
|
||||
destination
|
||||
)
|
||||
return documentParsingService.snapshot()
|
||||
}
|
||||
)
|
||||
|
||||
ipcMain.handle(
|
||||
ipcChannels.documentOcrModelsOpenRepository,
|
||||
async (event, input: unknown) => {
|
||||
assertTrustedSender(event, window)
|
||||
if (!documentOcrModelManager) {
|
||||
throw new Error('本地 OCR 模型服务不可用')
|
||||
}
|
||||
const { modelId } =
|
||||
documentOcrModelActionInputSchema.parse(input)
|
||||
const snapshot = await documentOcrModelManager.getSnapshot()
|
||||
const entry = snapshot.catalog.find(
|
||||
(candidate) => candidate.id === modelId
|
||||
)
|
||||
if (!entry) {
|
||||
throw new Error('未知的 OCR 模型')
|
||||
}
|
||||
await shell.openExternal(entry.repositoryUrl)
|
||||
}
|
||||
)
|
||||
|
||||
ipcMain.handle(
|
||||
ipcChannels.documentOcrModelsOpenDirectory,
|
||||
async (event) => {
|
||||
assertTrustedSender(event, window)
|
||||
if (!documentOcrModelManager) {
|
||||
throw new Error('本地 OCR 模型服务不可用')
|
||||
}
|
||||
await documentOcrModelManager.getSnapshot()
|
||||
const error = await shell.openPath(
|
||||
documentOcrModelManager.rootDirectory
|
||||
)
|
||||
if (error) {
|
||||
throw new Error('无法打开 OCR 模型目录')
|
||||
}
|
||||
}
|
||||
)
|
||||
|
||||
ipcMain.handle(
|
||||
ipcChannels.documentParsingOcrAssets,
|
||||
(event, input: unknown) => {
|
||||
assertTrustedSender(event, window)
|
||||
if (!documentOcrModelManager) {
|
||||
throw new Error('本地 OCR 模型服务不可用')
|
||||
}
|
||||
const { modelId } =
|
||||
documentOcrModelActionInputSchema.parse(input)
|
||||
return documentOcrModelManager.getAssets(modelId)
|
||||
}
|
||||
)
|
||||
|
||||
ipcMain.handle(
|
||||
ipcChannels.documentParsingOcrRespond,
|
||||
(event, input: unknown) => {
|
||||
assertTrustedSender(event, window)
|
||||
if (!documentOcrBroker) {
|
||||
throw new Error('本地 OCR 任务服务不可用')
|
||||
}
|
||||
const result = documentOcrResultSchema.safeParse(input)
|
||||
documentOcrBroker.respond(
|
||||
result.success
|
||||
? result.data
|
||||
: documentOcrFailureSchema.parse(input)
|
||||
)
|
||||
}
|
||||
)
|
||||
|
||||
ipcMain.handle(ipcChannels.versionCheck, async (event) => {
|
||||
assertTrustedSender(event, window)
|
||||
if (!versionChecker) {
|
||||
@@ -2610,7 +2860,7 @@ export function registerIpcHandlers(
|
||||
)
|
||||
|
||||
ipcMain.handle(
|
||||
ipcChannels.speechModelsImportLocal,
|
||||
ipcChannels.speechModelsImportArchive,
|
||||
async (event, input: unknown) => {
|
||||
assertTrustedSender(event, window)
|
||||
if (!speechModelManager) {
|
||||
@@ -2618,20 +2868,44 @@ export function registerIpcHandlers(
|
||||
}
|
||||
const { modelId } = speechModelActionInputSchema.parse(input)
|
||||
const result = await dialog.showOpenDialog(window, {
|
||||
properties: ['openDirectory']
|
||||
title: '导入语音模型 ZIP',
|
||||
properties: ['openFile'],
|
||||
filters: modelArchiveDialogFilters
|
||||
})
|
||||
const directory = result.filePaths[0]
|
||||
if (result.canceled || !directory) {
|
||||
const archivePath = result.filePaths[0]
|
||||
if (result.canceled || !archivePath) {
|
||||
return undefined
|
||||
}
|
||||
return trackExecution(
|
||||
speechModelManager
|
||||
.registerLocalDirectory(modelId, directory)
|
||||
.importArchive(modelId, archivePath)
|
||||
.then(() => speechModelManager.getSnapshot())
|
||||
)
|
||||
}
|
||||
)
|
||||
|
||||
ipcMain.handle(
|
||||
ipcChannels.speechModelsExportArchive,
|
||||
async (event, input: unknown) => {
|
||||
assertTrustedSender(event, window)
|
||||
if (!speechModelManager) {
|
||||
throw new Error('语音模型服务不可用')
|
||||
}
|
||||
const { modelId } = speechModelActionInputSchema.parse(input)
|
||||
const result = await dialog.showSaveDialog(window, {
|
||||
title: '导出语音模型 ZIP',
|
||||
defaultPath: `${modelId}.zip`,
|
||||
filters: modelArchiveDialogFilters
|
||||
})
|
||||
if (result.canceled || !result.filePath) {
|
||||
return undefined
|
||||
}
|
||||
const destination = ensureZipExtension(result.filePath)
|
||||
await speechModelManager.exportArchive(modelId, destination)
|
||||
return speechModelManager.getSnapshot()
|
||||
}
|
||||
)
|
||||
|
||||
ipcMain.handle(
|
||||
ipcChannels.speechModelsOpenRepository,
|
||||
async (event, input: unknown) => {
|
||||
|
||||
@@ -0,0 +1,50 @@
|
||||
import { beforeEach, describe, expect, it, vi } from 'vitest'
|
||||
|
||||
const getDocument = vi.hoisted(() => vi.fn())
|
||||
|
||||
vi.mock('pdfjs-dist/legacy/build/pdf.mjs', () => ({
|
||||
getDocument
|
||||
}))
|
||||
|
||||
import { extractPdfTextPages } from './document-parser'
|
||||
|
||||
describe('PDF extraction in Electron main', () => {
|
||||
beforeEach(() => {
|
||||
getDocument.mockReset()
|
||||
})
|
||||
|
||||
it('disables PDF.js DOM factories for headless text extraction', async () => {
|
||||
const cleanup = vi.fn()
|
||||
const destroy = vi.fn(async () => undefined)
|
||||
getDocument.mockReturnValue({
|
||||
promise: Promise.resolve({
|
||||
numPages: 1,
|
||||
getPage: vi.fn(async () => ({
|
||||
getTextContent: vi.fn(async () => ({
|
||||
items: [{ str: 'PDF body text' }]
|
||||
})),
|
||||
cleanup
|
||||
}))
|
||||
}),
|
||||
destroy
|
||||
})
|
||||
|
||||
await expect(
|
||||
extractPdfTextPages(Buffer.from('synthetic PDF'))
|
||||
).resolves.toEqual([
|
||||
{
|
||||
pageNumber: 1,
|
||||
content: 'PDF body text'
|
||||
}
|
||||
])
|
||||
expect(getDocument).toHaveBeenCalledWith({
|
||||
data: expect.any(Uint8Array),
|
||||
disableFontFace: true,
|
||||
isOffscreenCanvasSupported: false,
|
||||
useSystemFonts: false,
|
||||
useWorkerFetch: false
|
||||
})
|
||||
expect(cleanup).toHaveBeenCalledOnce()
|
||||
expect(destroy).toHaveBeenCalledOnce()
|
||||
})
|
||||
})
|
||||
@@ -5,12 +5,17 @@ import { extname } from 'node:path'
|
||||
export type ParsedSection = {
|
||||
locator: string
|
||||
content: string
|
||||
method?: 'native' | 'ocr' | 'converted' | 'vision'
|
||||
confidence?: number
|
||||
}
|
||||
|
||||
export type ParsedDocument = {
|
||||
title: string
|
||||
sourceFormat: string
|
||||
content: string
|
||||
sections: ParsedSection[]
|
||||
warnings: string[]
|
||||
pageCount?: number
|
||||
}
|
||||
|
||||
export type DocumentChunk = {
|
||||
@@ -160,12 +165,43 @@ function parseOfficeArchive(
|
||||
}
|
||||
|
||||
async function parsePdf(buffer: Buffer): Promise<ParsedSection[]> {
|
||||
const pages = await extractPdfTextPages(buffer)
|
||||
return pages
|
||||
.filter((page) => page.content.length > 0)
|
||||
.map((page) => ({
|
||||
locator: `第 ${page.pageNumber} 页`,
|
||||
content: page.content
|
||||
}))
|
||||
}
|
||||
|
||||
export type PdfTextPage = {
|
||||
pageNumber: number
|
||||
content: string
|
||||
}
|
||||
|
||||
export class DocumentTextUnavailableError extends Error {
|
||||
constructor(message = '文档中没有可索引的文本内容') {
|
||||
super(message)
|
||||
this.name = 'DocumentTextUnavailableError'
|
||||
}
|
||||
}
|
||||
|
||||
export async function extractPdfTextPages(
|
||||
buffer: Buffer
|
||||
): Promise<PdfTextPage[]> {
|
||||
const pdfjs = await import('pdfjs-dist/legacy/build/pdf.mjs')
|
||||
const loadingTask = pdfjs.getDocument({
|
||||
data: new Uint8Array(buffer)
|
||||
data: new Uint8Array(buffer),
|
||||
// Electron's main process identifies itself as process.type ===
|
||||
// "browser", so PDF.js otherwise selects DOM font factories even
|
||||
// though no document exists there.
|
||||
disableFontFace: true,
|
||||
isOffscreenCanvasSupported: false,
|
||||
useSystemFonts: false,
|
||||
useWorkerFetch: false
|
||||
})
|
||||
const document = await loadingTask.promise
|
||||
const sections: ParsedSection[] = []
|
||||
const pages: PdfTextPage[] = []
|
||||
try {
|
||||
for (let pageNumber = 1; pageNumber <= document.numPages; pageNumber += 1) {
|
||||
const page = await document.getPage(pageNumber)
|
||||
@@ -175,18 +211,13 @@ async function parsePdf(buffer: Buffer): Promise<ParsedSection[]> {
|
||||
.join(' ')
|
||||
.replace(/\s+/g, ' ')
|
||||
.trim()
|
||||
if (content) {
|
||||
sections.push({
|
||||
locator: `第 ${pageNumber} 页`,
|
||||
content
|
||||
})
|
||||
}
|
||||
pages.push({ pageNumber, content })
|
||||
page.cleanup()
|
||||
}
|
||||
} finally {
|
||||
await loadingTask.destroy()
|
||||
}
|
||||
return sections
|
||||
return pages
|
||||
}
|
||||
|
||||
export async function parseDocument(
|
||||
@@ -227,12 +258,14 @@ export async function parseDocument(
|
||||
.join('\n\n')
|
||||
.slice(0, maximumExtractedCharacters)
|
||||
if (!content) {
|
||||
throw new Error('文档中没有可索引的文本内容')
|
||||
throw new DocumentTextUnavailableError()
|
||||
}
|
||||
return {
|
||||
title: name.replace(/\.[^.]+$/, ''),
|
||||
sourceFormat: extension || 'unknown',
|
||||
content,
|
||||
sections
|
||||
sections,
|
||||
warnings: []
|
||||
}
|
||||
}
|
||||
|
||||
|
||||
@@ -18,7 +18,12 @@ import {
|
||||
relative,
|
||||
resolve
|
||||
} from 'node:path'
|
||||
import { chunkDocument, parseDocument, supportedDocumentExtensions } from './document-parser'
|
||||
import {
|
||||
chunkDocument,
|
||||
parseDocument,
|
||||
supportedDocumentExtensions,
|
||||
type ParsedDocument
|
||||
} from './document-parser'
|
||||
import { classifyEmbeddingError } from './embedding-errors'
|
||||
import {
|
||||
extractKnowledgeGraph,
|
||||
@@ -99,6 +104,12 @@ export type KnowledgeServiceOptions = {
|
||||
urlImporter?: UrlImporter
|
||||
embeddingProvider?: EmbeddingProvider
|
||||
embeddingBatchSize?: number
|
||||
parseDocument?: (
|
||||
name: string,
|
||||
buffer: Buffer,
|
||||
purpose: 'knowledge-index',
|
||||
signal?: AbortSignal
|
||||
) => Promise<ParsedDocument>
|
||||
}
|
||||
|
||||
const supportedExtensions = new Set<string>(supportedDocumentExtensions)
|
||||
@@ -118,6 +129,9 @@ export class KnowledgeService {
|
||||
private readonly managedRoot: string
|
||||
private readonly extractStructured?: ExtractStructured
|
||||
private readonly urlImporter: UrlImporter
|
||||
private readonly documentParser: NonNullable<
|
||||
KnowledgeServiceOptions['parseDocument']
|
||||
>
|
||||
private embeddingProvider?: EmbeddingProvider
|
||||
private readonly embeddingBatchSize: number
|
||||
private readonly watchers = new Map<string, FSWatcher>()
|
||||
@@ -131,6 +145,9 @@ export class KnowledgeService {
|
||||
this.managedRoot = resolve(options.managedRoot)
|
||||
this.extractStructured = options.extractStructured
|
||||
this.urlImporter = options.urlImporter ?? new UrlImporter()
|
||||
this.documentParser =
|
||||
options.parseDocument ??
|
||||
((name, buffer) => parseDocument(name, buffer))
|
||||
this.embeddingProvider = options.embeddingProvider
|
||||
const embeddingBatchSize = options.embeddingBatchSize ?? 16
|
||||
if (
|
||||
@@ -839,9 +856,11 @@ export class KnowledgeService {
|
||||
})
|
||||
continue
|
||||
}
|
||||
const parsed = await parseDocument(
|
||||
const parsed = await this.documentParser(
|
||||
basename(file.absolutePath),
|
||||
buffer
|
||||
buffer,
|
||||
'knowledge-index',
|
||||
this.lifecycleController.signal
|
||||
)
|
||||
this.updateKnowledgeTask(parsingTask.id, {
|
||||
progress: 75,
|
||||
|
||||
@@ -0,0 +1,211 @@
|
||||
import { createHash } from 'node:crypto'
|
||||
import {
|
||||
mkdtemp,
|
||||
mkdir,
|
||||
readFile,
|
||||
rm,
|
||||
writeFile
|
||||
} from 'node:fs/promises'
|
||||
import { tmpdir } from 'node:os'
|
||||
import { join } from 'node:path'
|
||||
import { zipSync } from 'fflate'
|
||||
import { afterEach, describe, expect, it } from 'vitest'
|
||||
import {
|
||||
exportModelArchive,
|
||||
extractModelArchive
|
||||
} from './model-archive'
|
||||
|
||||
const temporaryDirectories: string[] = []
|
||||
|
||||
async function temporaryDirectory(): Promise<string> {
|
||||
const directory = await mkdtemp(
|
||||
join(tmpdir(), 'goodbuddy-model-archive-')
|
||||
)
|
||||
temporaryDirectories.push(directory)
|
||||
return directory
|
||||
}
|
||||
|
||||
function sha256(value: Uint8Array): string {
|
||||
return createHash('sha256').update(value).digest('hex')
|
||||
}
|
||||
|
||||
afterEach(async () => {
|
||||
await Promise.all(
|
||||
temporaryDirectories.splice(0).map((directory) =>
|
||||
rm(directory, { recursive: true, force: true })
|
||||
)
|
||||
)
|
||||
})
|
||||
|
||||
describe('model archive', () => {
|
||||
it('exports and extracts only declared verified model files', async () => {
|
||||
const directory = await temporaryDirectory()
|
||||
const source = join(directory, 'source')
|
||||
const extracted = join(directory, 'extracted')
|
||||
const archive = join(directory, 'model.zip')
|
||||
await Promise.all([mkdir(source), mkdir(extracted)])
|
||||
const model = Buffer.from('verified model bytes')
|
||||
const tokens = Buffer.from('verified tokens')
|
||||
await Promise.all([
|
||||
writeFile(join(source, 'model.onnx'), model),
|
||||
writeFile(join(source, 'tokens.txt'), tokens),
|
||||
writeFile(join(source, 'ignored.txt'), 'not exported'),
|
||||
writeFile(archive, 'archive selected for replacement')
|
||||
])
|
||||
|
||||
await exportModelArchive({
|
||||
destinationPath: archive,
|
||||
sourceDirectory: source,
|
||||
descriptor: {
|
||||
kind: 'speech',
|
||||
modelId: 'test-model',
|
||||
displayName: 'Test model',
|
||||
files: [
|
||||
{
|
||||
name: 'model.onnx',
|
||||
role: 'model',
|
||||
size: model.byteLength,
|
||||
sha256: sha256(model)
|
||||
},
|
||||
{
|
||||
name: 'tokens.txt',
|
||||
role: 'tokens',
|
||||
size: tokens.byteLength,
|
||||
sha256: sha256(tokens)
|
||||
}
|
||||
]
|
||||
}
|
||||
})
|
||||
|
||||
await expect(
|
||||
extractModelArchive({
|
||||
archivePath: archive,
|
||||
destinationDirectory: extracted,
|
||||
expectedKind: 'speech',
|
||||
expectedModelId: 'test-model',
|
||||
expectedFiles: [
|
||||
{ name: 'model.onnx', role: 'model' },
|
||||
{ name: 'tokens.txt', role: 'tokens' }
|
||||
],
|
||||
maximumArchiveBytes: 1024 * 1024,
|
||||
maximumFileBytes: 1024,
|
||||
maximumTotalBytes: 2048
|
||||
})
|
||||
).resolves.toMatchObject({
|
||||
kind: 'speech',
|
||||
modelId: 'test-model'
|
||||
})
|
||||
await expect(readFile(join(extracted, 'model.onnx'))).resolves.toEqual(
|
||||
model
|
||||
)
|
||||
await expect(readFile(join(extracted, 'tokens.txt'))).resolves.toEqual(
|
||||
tokens
|
||||
)
|
||||
})
|
||||
|
||||
it('preserves an existing archive when source verification fails', async () => {
|
||||
const directory = await temporaryDirectory()
|
||||
const source = join(directory, 'source')
|
||||
const archive = join(directory, 'model.zip')
|
||||
await mkdir(source)
|
||||
const model = Buffer.from('changed model')
|
||||
await Promise.all([
|
||||
writeFile(join(source, 'model.onnx'), model),
|
||||
writeFile(archive, 'existing archive')
|
||||
])
|
||||
|
||||
await expect(
|
||||
exportModelArchive({
|
||||
destinationPath: archive,
|
||||
sourceDirectory: source,
|
||||
descriptor: {
|
||||
kind: 'speech',
|
||||
modelId: 'test-model',
|
||||
displayName: 'Test model',
|
||||
files: [
|
||||
{
|
||||
name: 'model.onnx',
|
||||
role: 'model',
|
||||
size: model.byteLength,
|
||||
sha256: 'a'.repeat(64)
|
||||
}
|
||||
]
|
||||
}
|
||||
})
|
||||
).rejects.toThrow('模型文件校验失败')
|
||||
await expect(readFile(archive, 'utf8')).resolves.toBe(
|
||||
'existing archive'
|
||||
)
|
||||
})
|
||||
|
||||
it('rejects path traversal and undeclared archive entries', async () => {
|
||||
const directory = await temporaryDirectory()
|
||||
const archive = join(directory, 'unsafe.zip')
|
||||
const extracted = join(directory, 'extracted')
|
||||
await mkdir(extracted)
|
||||
await writeFile(
|
||||
archive,
|
||||
zipSync({
|
||||
'../model.onnx': Buffer.from('unsafe')
|
||||
})
|
||||
)
|
||||
|
||||
await expect(
|
||||
extractModelArchive({
|
||||
archivePath: archive,
|
||||
destinationDirectory: extracted,
|
||||
expectedKind: 'speech',
|
||||
expectedModelId: 'test-model',
|
||||
expectedFiles: [{ name: 'model.onnx', role: 'model' }],
|
||||
maximumArchiveBytes: 1024 * 1024,
|
||||
maximumFileBytes: 1024,
|
||||
maximumTotalBytes: 1024
|
||||
})
|
||||
).rejects.toThrow()
|
||||
})
|
||||
|
||||
it('rejects an archive whose manifest model ID does not match', async () => {
|
||||
const directory = await temporaryDirectory()
|
||||
const archive = join(directory, 'mismatch.zip')
|
||||
const extracted = join(directory, 'extracted')
|
||||
await mkdir(extracted)
|
||||
const model = Buffer.from('model')
|
||||
await writeFile(
|
||||
archive,
|
||||
zipSync({
|
||||
'goodbuddy-model.json': Buffer.from(
|
||||
JSON.stringify({
|
||||
format: 'goodbuddy-model-archive',
|
||||
version: 1,
|
||||
kind: 'speech',
|
||||
modelId: 'other-model',
|
||||
displayName: 'Other model',
|
||||
exportedAt: '2026-08-11T00:00:00.000Z',
|
||||
files: [
|
||||
{
|
||||
name: 'model.onnx',
|
||||
role: 'model',
|
||||
size: model.byteLength,
|
||||
sha256: sha256(model)
|
||||
}
|
||||
]
|
||||
})
|
||||
),
|
||||
'model.onnx': model
|
||||
})
|
||||
)
|
||||
|
||||
await expect(
|
||||
extractModelArchive({
|
||||
archivePath: archive,
|
||||
destinationDirectory: extracted,
|
||||
expectedKind: 'speech',
|
||||
expectedModelId: 'test-model',
|
||||
expectedFiles: [{ name: 'model.onnx', role: 'model' }],
|
||||
maximumArchiveBytes: 1024 * 1024,
|
||||
maximumFileBytes: 1024,
|
||||
maximumTotalBytes: 1024
|
||||
})
|
||||
).rejects.toThrow('模型 ID 不匹配')
|
||||
})
|
||||
})
|
||||
@@ -0,0 +1,617 @@
|
||||
import { createHash, randomUUID } from 'node:crypto'
|
||||
import {
|
||||
lstat,
|
||||
open,
|
||||
readFile,
|
||||
rename,
|
||||
rm,
|
||||
type FileHandle
|
||||
} from 'node:fs/promises'
|
||||
import { dirname, resolve } from 'node:path'
|
||||
import {
|
||||
Unzip,
|
||||
UnzipInflate,
|
||||
UnzipPassThrough,
|
||||
Zip,
|
||||
ZipPassThrough
|
||||
} from 'fflate'
|
||||
import { z } from 'zod'
|
||||
|
||||
const ARCHIVE_MANIFEST_NAME = 'goodbuddy-model.json'
|
||||
const ARCHIVE_FORMAT = 'goodbuddy-model-archive'
|
||||
const ARCHIVE_VERSION = 1
|
||||
const MAXIMUM_ARCHIVE_ENTRIES = 40
|
||||
const MAXIMUM_MANIFEST_BYTES = 256 * 1024
|
||||
|
||||
const archiveFileNameSchema = z
|
||||
.string()
|
||||
.min(1)
|
||||
.max(255)
|
||||
.regex(/^[^/\\:\0]+$/u)
|
||||
|
||||
const modelArchiveFileSchema = z
|
||||
.object({
|
||||
name: archiveFileNameSchema,
|
||||
role: z.string().trim().min(1).max(64),
|
||||
size: z.number().int().positive().safe(),
|
||||
sha256: z.string().regex(/^[a-f0-9]{64}$/u)
|
||||
})
|
||||
.strict()
|
||||
|
||||
const modelArchiveDescriptorSchema = z
|
||||
.object({
|
||||
kind: z.enum(['speech', 'document-ocr']),
|
||||
modelId: z
|
||||
.string()
|
||||
.min(1)
|
||||
.max(96)
|
||||
.regex(/^[a-z0-9]+(?:-[a-z0-9]+)*$/u),
|
||||
displayName: z.string().trim().min(1).max(120),
|
||||
files: z.array(modelArchiveFileSchema).min(1).max(32)
|
||||
})
|
||||
.strict()
|
||||
|
||||
const modelArchiveManifestSchema = modelArchiveDescriptorSchema
|
||||
.extend({
|
||||
format: z.literal(ARCHIVE_FORMAT),
|
||||
version: z.literal(ARCHIVE_VERSION),
|
||||
exportedAt: z.string().datetime()
|
||||
})
|
||||
.strict()
|
||||
.superRefine((manifest, context) => {
|
||||
if (
|
||||
new Set(manifest.files.map((file) => file.name.toLowerCase()))
|
||||
.size !== manifest.files.length
|
||||
) {
|
||||
context.addIssue({
|
||||
code: 'custom',
|
||||
path: ['files'],
|
||||
message: '模型 ZIP 清单包含重复文件'
|
||||
})
|
||||
}
|
||||
})
|
||||
|
||||
export type ModelArchiveKind = z.infer<
|
||||
typeof modelArchiveManifestSchema
|
||||
>['kind']
|
||||
|
||||
export type ModelArchiveFile = z.infer<typeof modelArchiveFileSchema>
|
||||
|
||||
export type ModelArchiveDescriptor = {
|
||||
kind: ModelArchiveKind
|
||||
modelId: string
|
||||
displayName: string
|
||||
files: ModelArchiveFile[]
|
||||
}
|
||||
|
||||
export type ModelArchiveExpectedFile = {
|
||||
name: string
|
||||
role: string
|
||||
}
|
||||
|
||||
type ExportModelArchiveOptions = {
|
||||
destinationPath: string
|
||||
sourceDirectory: string
|
||||
descriptor: ModelArchiveDescriptor
|
||||
}
|
||||
|
||||
type ExtractModelArchiveOptions = {
|
||||
archivePath: string
|
||||
destinationDirectory: string
|
||||
expectedKind: ModelArchiveKind
|
||||
expectedModelId: string
|
||||
expectedFiles: ModelArchiveExpectedFile[]
|
||||
maximumArchiveBytes: number
|
||||
maximumFileBytes: number
|
||||
maximumTotalBytes: number
|
||||
signal?: AbortSignal
|
||||
onProgress?: (completedBytes: number) => void
|
||||
}
|
||||
|
||||
function safeChild(parent: string, name: string): string {
|
||||
const child = resolve(parent, name)
|
||||
if (dirname(child) !== resolve(parent)) {
|
||||
throw new Error('模型 ZIP 路径超出临时目录')
|
||||
}
|
||||
return child
|
||||
}
|
||||
|
||||
function ensureArchiveName(name: string): string {
|
||||
return archiveFileNameSchema.parse(name)
|
||||
}
|
||||
|
||||
function ensureUniqueFiles(files: ModelArchiveExpectedFile[]): void {
|
||||
const names = files.map((file) => ensureArchiveName(file.name))
|
||||
if (new Set(names.map((name) => name.toLowerCase())).size !== names.length) {
|
||||
throw new Error('模型目录包含重复文件名')
|
||||
}
|
||||
}
|
||||
|
||||
async function hashFile(path: string): Promise<ModelArchiveFile['sha256']> {
|
||||
const handle = await open(path, 'r')
|
||||
const hash = createHash('sha256')
|
||||
const buffer = Buffer.allocUnsafe(64 * 1024)
|
||||
try {
|
||||
while (true) {
|
||||
const { bytesRead } = await handle.read(buffer, 0, buffer.length)
|
||||
if (bytesRead === 0) {
|
||||
break
|
||||
}
|
||||
hash.update(buffer.subarray(0, bytesRead))
|
||||
}
|
||||
} finally {
|
||||
await handle.close()
|
||||
}
|
||||
return hash.digest('hex')
|
||||
}
|
||||
|
||||
function checkedLimit(value: number, label: string): number {
|
||||
if (!Number.isSafeInteger(value) || value <= 0) {
|
||||
throw new RangeError(`${label}无效`)
|
||||
}
|
||||
return value
|
||||
}
|
||||
|
||||
function ensureNotAborted(signal?: AbortSignal): void {
|
||||
if (signal?.aborted) {
|
||||
throw signal.reason instanceof Error
|
||||
? signal.reason
|
||||
: new Error('模型 ZIP 导入已取消')
|
||||
}
|
||||
}
|
||||
|
||||
async function pushFileIntoArchive(
|
||||
archive: Zip,
|
||||
file: ModelArchiveFile,
|
||||
sourcePath: string,
|
||||
waitForOutput: () => Promise<void>
|
||||
): Promise<void> {
|
||||
const input = new ZipPassThrough(ensureArchiveName(file.name))
|
||||
archive.add(input)
|
||||
const sourceInfo = await lstat(sourcePath)
|
||||
if (!sourceInfo.isFile() || sourceInfo.isSymbolicLink()) {
|
||||
throw new Error(`模型文件不可导出：${file.name}`)
|
||||
}
|
||||
const handle = await open(sourcePath, 'r')
|
||||
const buffer = Buffer.allocUnsafe(64 * 1024)
|
||||
const hash = createHash('sha256')
|
||||
let size = 0
|
||||
try {
|
||||
const openedInfo = await handle.stat()
|
||||
if (
|
||||
!openedInfo.isFile() ||
|
||||
openedInfo.dev !== sourceInfo.dev ||
|
||||
openedInfo.ino !== sourceInfo.ino
|
||||
) {
|
||||
throw new Error(`模型文件在打开前已发生变化：${file.name}`)
|
||||
}
|
||||
while (true) {
|
||||
const { bytesRead } = await handle.read(buffer, 0, buffer.length)
|
||||
if (bytesRead === 0) {
|
||||
break
|
||||
}
|
||||
const chunk = buffer.subarray(0, bytesRead)
|
||||
hash.update(chunk)
|
||||
size += bytesRead
|
||||
input.push(Uint8Array.from(chunk))
|
||||
await waitForOutput()
|
||||
}
|
||||
if (size !== file.size || hash.digest('hex') !== file.sha256) {
|
||||
throw new Error(`模型文件校验失败：${file.name}`)
|
||||
}
|
||||
input.push(new Uint8Array(), true)
|
||||
await waitForOutput()
|
||||
} finally {
|
||||
await handle.close()
|
||||
}
|
||||
}
|
||||
|
||||
async function pushBytesIntoArchive(
|
||||
archive: Zip,
|
||||
name: string,
|
||||
value: Uint8Array,
|
||||
waitForOutput: () => Promise<void>
|
||||
): Promise<void> {
|
||||
const input = new ZipPassThrough(ensureArchiveName(name))
|
||||
archive.add(input)
|
||||
input.push(value, true)
|
||||
await waitForOutput()
|
||||
}
|
||||
|
||||
async function replaceArchiveFile(
|
||||
partialPath: string,
|
||||
destinationPath: string
|
||||
): Promise<void> {
|
||||
const backupPath = `${destinationPath}.${randomUUID()}.backup`
|
||||
let movedExistingFile = false
|
||||
try {
|
||||
try {
|
||||
await rename(destinationPath, backupPath)
|
||||
movedExistingFile = true
|
||||
const existingInfo = await lstat(backupPath)
|
||||
if (!existingInfo.isFile() || existingInfo.isSymbolicLink()) {
|
||||
throw new Error('模型 ZIP 导出目标必须是普通文件')
|
||||
}
|
||||
} catch (error) {
|
||||
if ((error as NodeJS.ErrnoException).code !== 'ENOENT') {
|
||||
throw error
|
||||
}
|
||||
}
|
||||
await rename(partialPath, destinationPath)
|
||||
if (movedExistingFile) {
|
||||
await rm(backupPath, { force: true }).catch(() => undefined)
|
||||
}
|
||||
} catch (error) {
|
||||
if (movedExistingFile) {
|
||||
await rm(destinationPath, { force: true }).catch(() => undefined)
|
||||
await rename(backupPath, destinationPath).catch(() => undefined)
|
||||
}
|
||||
throw error
|
||||
}
|
||||
}
|
||||
|
||||
export async function exportModelArchive(
|
||||
options: ExportModelArchiveOptions
|
||||
): Promise<void> {
|
||||
const descriptor = modelArchiveDescriptorSchema.parse(
|
||||
options.descriptor
|
||||
)
|
||||
ensureUniqueFiles(descriptor.files)
|
||||
const sourceDirectory = resolve(options.sourceDirectory)
|
||||
const destinationPath = resolve(options.destinationPath)
|
||||
const partialPath = `${destinationPath}.${randomUUID()}.partial`
|
||||
const output = await open(partialPath, 'wx')
|
||||
let writeChain = Promise.resolve()
|
||||
let archiveError: Error | undefined
|
||||
let resolveFinished: (() => void) | undefined
|
||||
let rejectFinished: ((error: Error) => void) | undefined
|
||||
const finished = new Promise<void>((resolvePromise, rejectPromise) => {
|
||||
resolveFinished = resolvePromise
|
||||
rejectFinished = rejectPromise
|
||||
})
|
||||
const archive = new Zip((error, data, final) => {
|
||||
if (error) {
|
||||
archiveError = error
|
||||
rejectFinished?.(error)
|
||||
return
|
||||
}
|
||||
writeChain = writeChain.then(async () => {
|
||||
if (data.byteLength > 0) {
|
||||
await output.write(data)
|
||||
}
|
||||
})
|
||||
if (final) {
|
||||
void writeChain.then(resolveFinished, rejectFinished)
|
||||
}
|
||||
})
|
||||
const waitForOutput = async (): Promise<void> => {
|
||||
await writeChain
|
||||
if (archiveError) {
|
||||
throw archiveError
|
||||
}
|
||||
}
|
||||
try {
|
||||
const manifest = modelArchiveManifestSchema.parse({
|
||||
format: ARCHIVE_FORMAT,
|
||||
version: ARCHIVE_VERSION,
|
||||
kind: descriptor.kind,
|
||||
modelId: descriptor.modelId,
|
||||
displayName: descriptor.displayName,
|
||||
exportedAt: new Date().toISOString(),
|
||||
files: descriptor.files
|
||||
})
|
||||
await pushBytesIntoArchive(
|
||||
archive,
|
||||
ARCHIVE_MANIFEST_NAME,
|
||||
Buffer.from(`${JSON.stringify(manifest, null, 2)}\n`, 'utf8'),
|
||||
waitForOutput
|
||||
)
|
||||
for (const file of descriptor.files) {
|
||||
await pushFileIntoArchive(
|
||||
archive,
|
||||
file,
|
||||
safeChild(sourceDirectory, file.name),
|
||||
waitForOutput
|
||||
)
|
||||
}
|
||||
archive.end()
|
||||
await finished
|
||||
await output.sync()
|
||||
await output.close()
|
||||
await replaceArchiveFile(partialPath, destinationPath)
|
||||
} catch (error) {
|
||||
archive.terminate()
|
||||
await output.close().catch(() => undefined)
|
||||
await rm(partialPath, { force: true })
|
||||
throw error
|
||||
}
|
||||
}
|
||||
|
||||
function closeHandle(handle: FileHandle): Promise<void> {
|
||||
return handle.close().catch(() => undefined)
|
||||
}
|
||||
|
||||
export async function extractModelArchive(
|
||||
options: ExtractModelArchiveOptions
|
||||
): Promise<ModelArchiveDescriptor> {
|
||||
ensureNotAborted(options.signal)
|
||||
const maximumArchiveBytes = checkedLimit(
|
||||
options.maximumArchiveBytes,
|
||||
'模型 ZIP 大小限制'
|
||||
)
|
||||
const maximumFileBytes = checkedLimit(
|
||||
options.maximumFileBytes,
|
||||
'模型文件大小限制'
|
||||
)
|
||||
const maximumTotalBytes = checkedLimit(
|
||||
options.maximumTotalBytes,
|
||||
'模型展开大小限制'
|
||||
)
|
||||
const expectedFiles = options.expectedFiles.map((file) => ({
|
||||
name: ensureArchiveName(file.name),
|
||||
role: file.role
|
||||
}))
|
||||
ensureUniqueFiles(expectedFiles)
|
||||
const allowedNames = new Set([
|
||||
ARCHIVE_MANIFEST_NAME,
|
||||
...expectedFiles.map((file) => file.name)
|
||||
])
|
||||
const source = resolve(options.archivePath)
|
||||
let sourceInfo
|
||||
try {
|
||||
sourceInfo = await lstat(source)
|
||||
} catch (error) {
|
||||
throw new Error('无法读取模型 ZIP', { cause: error })
|
||||
}
|
||||
if (
|
||||
!sourceInfo.isFile() ||
|
||||
sourceInfo.isSymbolicLink() ||
|
||||
sourceInfo.size <= 0 ||
|
||||
sourceInfo.size > maximumArchiveBytes
|
||||
) {
|
||||
throw new Error('模型 ZIP 必须是大小合规的普通文件')
|
||||
}
|
||||
|
||||
let input: FileHandle | undefined
|
||||
try {
|
||||
input = await open(source, 'r')
|
||||
const openedInfo = await input.stat()
|
||||
if (
|
||||
!openedInfo.isFile() ||
|
||||
openedInfo.size !== sourceInfo.size ||
|
||||
openedInfo.dev !== sourceInfo.dev ||
|
||||
openedInfo.ino !== sourceInfo.ino
|
||||
) {
|
||||
await input.close()
|
||||
throw new Error('模型 ZIP 在打开前已发生变化')
|
||||
}
|
||||
} catch (error) {
|
||||
await input?.close().catch(() => undefined)
|
||||
if (error instanceof Error && error.message.startsWith('模型 ZIP')) {
|
||||
throw error
|
||||
}
|
||||
throw new Error('无法读取模型 ZIP', { cause: error })
|
||||
}
|
||||
if (!input) {
|
||||
throw new Error('无法读取模型 ZIP')
|
||||
}
|
||||
|
||||
const destination = resolve(options.destinationDirectory)
|
||||
const seenNames = new Set<string>()
|
||||
const openHandles = new Set<FileHandle>()
|
||||
const completions: Promise<void>[] = []
|
||||
const pendingWrites = new Set<Promise<void>>()
|
||||
let entryCount = 0
|
||||
let totalBytes = 0
|
||||
let completedModelBytes = 0
|
||||
let fatalError: Error | undefined
|
||||
const fail = (error: unknown): Error => {
|
||||
const resolvedError =
|
||||
error instanceof Error ? error : new Error('模型 ZIP 已损坏')
|
||||
fatalError ??= resolvedError
|
||||
return resolvedError
|
||||
}
|
||||
const unzip = new Unzip((file) => {
|
||||
try {
|
||||
entryCount += 1
|
||||
if (
|
||||
entryCount > MAXIMUM_ARCHIVE_ENTRIES ||
|
||||
entryCount > allowedNames.size
|
||||
) {
|
||||
throw new Error('模型 ZIP 包含过多条目')
|
||||
}
|
||||
const name = ensureArchiveName(file.name)
|
||||
const key = name.toLowerCase()
|
||||
if (seenNames.has(key)) {
|
||||
throw new Error('模型 ZIP 包含重复条目')
|
||||
}
|
||||
seenNames.add(key)
|
||||
if (!allowedNames.has(name)) {
|
||||
throw new Error(`模型 ZIP 包含未声明文件：${name}`)
|
||||
}
|
||||
const entryMaximum =
|
||||
name === ARCHIVE_MANIFEST_NAME
|
||||
? MAXIMUM_MANIFEST_BYTES
|
||||
: maximumFileBytes
|
||||
if (
|
||||
file.originalSize !== undefined &&
|
||||
(file.originalSize <= 0 ||
|
||||
file.originalSize > entryMaximum ||
|
||||
totalBytes + file.originalSize > maximumTotalBytes)
|
||||
) {
|
||||
throw new Error(`模型 ZIP 条目大小超出限制：${name}`)
|
||||
}
|
||||
const handlePromise = open(
|
||||
safeChild(destination, name),
|
||||
'wx'
|
||||
).then((handle) => {
|
||||
openHandles.add(handle)
|
||||
return handle
|
||||
})
|
||||
let written = 0
|
||||
let writeChain = Promise.resolve()
|
||||
let resolveEntry: (() => void) | undefined
|
||||
let rejectEntry: ((error: Error) => void) | undefined
|
||||
const completion = new Promise<void>((resolveEntryPromise, rejectEntryPromise) => {
|
||||
resolveEntry = resolveEntryPromise
|
||||
rejectEntry = rejectEntryPromise
|
||||
})
|
||||
completions.push(completion)
|
||||
file.ondata = (error, data, final) => {
|
||||
if (error) {
|
||||
rejectEntry?.(fail(error))
|
||||
return
|
||||
}
|
||||
if (fatalError) {
|
||||
file.terminate()
|
||||
rejectEntry?.(fatalError)
|
||||
return
|
||||
}
|
||||
if (options.signal?.aborted) {
|
||||
file.terminate()
|
||||
rejectEntry?.(
|
||||
fail(
|
||||
options.signal.reason instanceof Error
|
||||
? options.signal.reason
|
||||
: new Error('模型 ZIP 导入已取消')
|
||||
)
|
||||
)
|
||||
return
|
||||
}
|
||||
written += data.byteLength
|
||||
totalBytes += data.byteLength
|
||||
if (name !== ARCHIVE_MANIFEST_NAME) {
|
||||
completedModelBytes += data.byteLength
|
||||
options.onProgress?.(completedModelBytes)
|
||||
}
|
||||
if (
|
||||
written > entryMaximum ||
|
||||
totalBytes > maximumTotalBytes
|
||||
) {
|
||||
file.terminate()
|
||||
rejectEntry?.(
|
||||
fail(new Error(`模型 ZIP 条目大小超出限制：${name}`))
|
||||
)
|
||||
return
|
||||
}
|
||||
writeChain = writeChain.then(async () => {
|
||||
const handle = await handlePromise
|
||||
if (data.byteLength > 0) {
|
||||
await handle.write(data)
|
||||
}
|
||||
})
|
||||
const pendingWrite = writeChain
|
||||
pendingWrites.add(pendingWrite)
|
||||
void pendingWrite.then(
|
||||
() => pendingWrites.delete(pendingWrite),
|
||||
() => pendingWrites.delete(pendingWrite)
|
||||
)
|
||||
if (final) {
|
||||
void writeChain.then(async () => {
|
||||
const handle = await handlePromise
|
||||
openHandles.delete(handle)
|
||||
await closeHandle(handle)
|
||||
resolveEntry?.()
|
||||
}, (writeError: unknown) => {
|
||||
rejectEntry?.(fail(writeError))
|
||||
})
|
||||
}
|
||||
}
|
||||
file.start()
|
||||
} catch (error) {
|
||||
file.terminate()
|
||||
fail(error)
|
||||
}
|
||||
})
|
||||
unzip.register(UnzipPassThrough)
|
||||
unzip.register(UnzipInflate)
|
||||
|
||||
const buffer = Buffer.allocUnsafe(16 * 1024)
|
||||
try {
|
||||
while (true) {
|
||||
ensureNotAborted(options.signal)
|
||||
if (fatalError) {
|
||||
throw fatalError
|
||||
}
|
||||
const { bytesRead } = await input.read(buffer, 0, buffer.length)
|
||||
if (bytesRead === 0) {
|
||||
unzip.push(new Uint8Array(), true)
|
||||
break
|
||||
}
|
||||
unzip.push(
|
||||
Uint8Array.from(buffer.subarray(0, bytesRead)),
|
||||
false
|
||||
)
|
||||
await Promise.all([...pendingWrites])
|
||||
}
|
||||
await Promise.all(completions)
|
||||
if (fatalError) {
|
||||
throw fatalError
|
||||
}
|
||||
} catch (error) {
|
||||
throw fail(error)
|
||||
} finally {
|
||||
await input.close()
|
||||
await Promise.all(
|
||||
[...openHandles].map((handle) => closeHandle(handle))
|
||||
)
|
||||
}
|
||||
|
||||
if (
|
||||
seenNames.size !== allowedNames.size ||
|
||||
[...allowedNames].some(
|
||||
(name) => !seenNames.has(name.toLowerCase())
|
||||
)
|
||||
) {
|
||||
throw new Error('模型 ZIP 缺少必需文件')
|
||||
}
|
||||
|
||||
let manifest
|
||||
try {
|
||||
manifest = modelArchiveManifestSchema.parse(
|
||||
JSON.parse(
|
||||
await readFile(
|
||||
safeChild(destination, ARCHIVE_MANIFEST_NAME),
|
||||
'utf8'
|
||||
)
|
||||
) as unknown
|
||||
)
|
||||
} catch {
|
||||
throw new Error('模型 ZIP 清单无效')
|
||||
}
|
||||
if (
|
||||
manifest.kind !== options.expectedKind ||
|
||||
manifest.modelId !== options.expectedModelId
|
||||
) {
|
||||
throw new Error('模型 ZIP 类型或模型 ID 不匹配')
|
||||
}
|
||||
if (
|
||||
manifest.files.length !== expectedFiles.length ||
|
||||
expectedFiles.some((expected) => {
|
||||
const archived = manifest.files.find(
|
||||
(file) => file.name === expected.name
|
||||
)
|
||||
return !archived || archived.role !== expected.role
|
||||
})
|
||||
) {
|
||||
throw new Error('模型 ZIP 清单与当前模型目录不匹配')
|
||||
}
|
||||
for (const archived of manifest.files) {
|
||||
const path = safeChild(destination, archived.name)
|
||||
const metadata = await lstat(path)
|
||||
if (
|
||||
!metadata.isFile() ||
|
||||
metadata.isSymbolicLink() ||
|
||||
metadata.size !== archived.size ||
|
||||
(await hashFile(path)) !== archived.sha256
|
||||
) {
|
||||
throw new Error(`模型 ZIP 文件校验失败：${archived.name}`)
|
||||
}
|
||||
}
|
||||
return {
|
||||
kind: manifest.kind,
|
||||
modelId: manifest.modelId,
|
||||
displayName: manifest.displayName,
|
||||
files: manifest.files
|
||||
}
|
||||
}
|
||||
@@ -373,6 +373,47 @@ describe('SpeechModelManager downloads', () => {
|
||||
operations: []
|
||||
})
|
||||
})
|
||||
|
||||
it('round-trips a verified model through an offline ZIP archive', async () => {
|
||||
const userData = await temporaryDirectory()
|
||||
const modelBytes = new TextEncoder().encode('verified model bytes')
|
||||
const tokenBytes = new TextEncoder().encode('verified tokens')
|
||||
const catalog = downloadableCatalog(modelBytes, tokenBytes)
|
||||
const manager = new SpeechModelManager({
|
||||
userDataDirectory: userData,
|
||||
catalog,
|
||||
fetch: vi.fn<typeof fetch>(async (input) => {
|
||||
const bytes = String(input).endsWith('model.onnx')
|
||||
? modelBytes
|
||||
: tokenBytes
|
||||
return new Response(bytes, {
|
||||
headers: { 'content-length': String(bytes.byteLength) }
|
||||
})
|
||||
})
|
||||
})
|
||||
const archive = join(userData, 'speech-model.zip')
|
||||
|
||||
await manager.install('download-test-model')
|
||||
await manager.exportArchive('download-test-model', archive)
|
||||
await manager.remove('download-test-model')
|
||||
|
||||
await expect(
|
||||
manager.importArchive('download-test-model', archive)
|
||||
).resolves.toMatchObject({
|
||||
id: 'download-test-model',
|
||||
source: 'local',
|
||||
files: [
|
||||
{
|
||||
name: 'model.onnx',
|
||||
sha256: sha256(modelBytes)
|
||||
},
|
||||
{
|
||||
name: 'tokens.txt',
|
||||
sha256: sha256(tokenBytes)
|
||||
}
|
||||
]
|
||||
})
|
||||
})
|
||||
})
|
||||
|
||||
describe('SpeechModelManager local import', () => {
|
||||
|
||||
@@ -25,12 +25,18 @@ import {
|
||||
type SpeechModelSnapshot
|
||||
} from '../../shared/speech-model-contracts'
|
||||
import { SPEECH_MODEL_CATALOG } from './speech-model-catalog'
|
||||
import {
|
||||
exportModelArchive,
|
||||
extractModelArchive
|
||||
} from '../model-archive'
|
||||
|
||||
const DEFAULT_MAX_FILE_BYTES = 2 * 1024 * 1024 * 1024
|
||||
const MAX_REDIRECTS = 3
|
||||
const MANIFEST_FILE_NAME = 'manifest.json'
|
||||
const SELECTION_FILE_NAME = '.selection.json'
|
||||
const PARTIAL_SUFFIX = '.partial'
|
||||
const MAXIMUM_ARCHIVE_BYTES = 4 * 1024 * 1024 * 1024 - 1
|
||||
const ARCHIVE_OVERHEAD_BYTES = 1024 * 1024
|
||||
|
||||
const selectionSchema = z
|
||||
.object({
|
||||
@@ -380,6 +386,143 @@ export class SpeechModelManager {
|
||||
}
|
||||
}
|
||||
|
||||
async exportArchive(
|
||||
modelId: string,
|
||||
destinationPath: string
|
||||
): Promise<void> {
|
||||
const entry = this.requireCatalogEntry(modelId)
|
||||
await this.ensureRoot()
|
||||
const installed = (await this.readInstalled()).find(
|
||||
(model) => model.id === entry.id
|
||||
)
|
||||
if (!installed) {
|
||||
throw new Error('只能导出已安装的语音模型')
|
||||
}
|
||||
const directory = this.modelDirectory(entry.id)
|
||||
const files = []
|
||||
for (const expected of entry.files) {
|
||||
const recorded = installed.files.find(
|
||||
(file) =>
|
||||
file.name === expected.name && file.role === expected.role
|
||||
)
|
||||
if (
|
||||
!recorded ||
|
||||
recorded.size <= 0 ||
|
||||
recorded.size > this.maxFileBytes ||
|
||||
(expected.download &&
|
||||
(recorded.size !== expected.download.size ||
|
||||
recorded.sha256 !== expected.download.sha256))
|
||||
) {
|
||||
throw new Error(`语音模型文件不可导出：${expected.name}`)
|
||||
}
|
||||
files.push({
|
||||
name: expected.name,
|
||||
role: expected.role,
|
||||
size: recorded.size,
|
||||
sha256: recorded.sha256
|
||||
})
|
||||
}
|
||||
await exportModelArchive({
|
||||
destinationPath,
|
||||
sourceDirectory: directory,
|
||||
descriptor: {
|
||||
kind: 'speech',
|
||||
modelId: entry.id,
|
||||
displayName: entry.displayName,
|
||||
files
|
||||
}
|
||||
})
|
||||
}
|
||||
|
||||
async importArchive(
|
||||
modelId: string,
|
||||
archivePath: string
|
||||
): Promise<InstalledSpeechModel> {
|
||||
const entry = this.requireCatalogEntry(modelId)
|
||||
const expectedTotal = entry.files.reduce(
|
||||
(total, file) =>
|
||||
total + (file.download?.size ?? this.maxFileBytes),
|
||||
0
|
||||
)
|
||||
const maximumTotalBytes = Math.min(
|
||||
MAXIMUM_ARCHIVE_BYTES,
|
||||
expectedTotal + ARCHIVE_OVERHEAD_BYTES
|
||||
)
|
||||
const operation = this.beginOperation(
|
||||
entry.id,
|
||||
'import',
|
||||
expectedTotal
|
||||
)
|
||||
let stagingDirectory: string | undefined
|
||||
try {
|
||||
await this.ensureRoot()
|
||||
await this.assertNotInstalled(entry.id)
|
||||
stagingDirectory = await this.createStagingDirectory(entry.id)
|
||||
operation.progress.phase = 'transferring'
|
||||
const descriptor = await extractModelArchive({
|
||||
archivePath,
|
||||
destinationDirectory: stagingDirectory,
|
||||
expectedKind: 'speech',
|
||||
expectedModelId: entry.id,
|
||||
expectedFiles: entry.files.map((file) => ({
|
||||
name: file.name,
|
||||
role: file.role
|
||||
})),
|
||||
maximumArchiveBytes: Math.min(
|
||||
MAXIMUM_ARCHIVE_BYTES,
|
||||
maximumTotalBytes + ARCHIVE_OVERHEAD_BYTES
|
||||
),
|
||||
maximumFileBytes: this.maxFileBytes,
|
||||
maximumTotalBytes,
|
||||
signal: operation.controller.signal,
|
||||
onProgress: (completedBytes) => {
|
||||
operation.progress.completedBytes = completedBytes
|
||||
}
|
||||
})
|
||||
for (const expected of entry.files) {
|
||||
const archived = descriptor.files.find(
|
||||
(file) =>
|
||||
file.name === expected.name &&
|
||||
file.role === expected.role
|
||||
)
|
||||
if (
|
||||
!archived ||
|
||||
archived.size > this.maxFileBytes ||
|
||||
(expected.download &&
|
||||
(archived.size !== expected.download.size ||
|
||||
archived.sha256 !== expected.download.sha256))
|
||||
) {
|
||||
throw new Error(
|
||||
`语音模型 ZIP 与当前模型目录不匹配：${expected.name}`
|
||||
)
|
||||
}
|
||||
}
|
||||
operation.progress.phase = 'installing'
|
||||
operation.progress.currentFile = null
|
||||
const installed = installedSpeechModelSchema.parse({
|
||||
id: entry.id,
|
||||
displayName: entry.displayName,
|
||||
source: 'local',
|
||||
installedAt: new Date().toISOString(),
|
||||
files: descriptor.files
|
||||
})
|
||||
await writeFile(
|
||||
safeChild(stagingDirectory, MANIFEST_FILE_NAME),
|
||||
`${JSON.stringify(installed, null, 2)}\n`,
|
||||
{ encoding: 'utf8', flag: 'wx' }
|
||||
)
|
||||
ensureNotAborted(operation.controller.signal)
|
||||
await rename(stagingDirectory, this.modelDirectory(entry.id))
|
||||
stagingDirectory = undefined
|
||||
return installed
|
||||
} finally {
|
||||
this.operations.delete(entry.id)
|
||||
if (stagingDirectory) {
|
||||
await rm(stagingDirectory, { recursive: true, force: true })
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
private async ensureRoot(): Promise<void> {
|
||||
await mkdir(this.rootDirectory, { recursive: true })
|
||||
}
|
||||
|
||||
Reference in New Issue
Block a user