diff --git a/package.json b/package.json index 067c2ed79..e76601eb2 100644 --- a/package.json +++ b/package.json @@ -70,6 +70,7 @@ "@modelcontextprotocol/sdk": "patch:@modelcontextprotocol/sdk@npm%3A1.6.1#~/.yarn/patches/@modelcontextprotocol-sdk-npm-1.6.1-b46313efe7.patch", "@notionhq/client": "^2.2.15", "@tryfabric/martian": "^1.2.4", + "@types/pdf-parse": "^1.1.4", "@types/react-infinite-scroll-component": "^5.0.0", "adm-zip": "^0.5.16", "apache-arrow": "^18.1.0", @@ -85,6 +86,7 @@ "npx-scope-finder": "^1.2.0", "officeparser": "^4.1.1", "p-queue": "^8.1.0", + "pdf-parse": "^1.1.1", "proxy-agent": "^6.5.0", "tar": "^7.4.3", "tokenx": "^0.4.1", diff --git a/src/main/ipc.ts b/src/main/ipc.ts index 10255e563..91842c56a 100644 --- a/src/main/ipc.ts +++ b/src/main/ipc.ts @@ -154,6 +154,7 @@ export function registerIpc(mainWindow: BrowserWindow, app: Electron.App) { ipcMain.handle('file:clear', fileManager.clear) ipcMain.handle('file:read', fileManager.readFile) ipcMain.handle('file:delete', fileManager.deleteFile) + ipcMain.handle('file:deleteDir', fileManager.deleteDir) ipcMain.handle('file:get', fileManager.getFile) ipcMain.handle('file:selectFolder', fileManager.selectFolder) ipcMain.handle('file:create', fileManager.createTempFile) diff --git a/src/main/ocr/BaseOcrProvider.ts b/src/main/ocr/BaseOcrProvider.ts index 8fa4c3276..7c5837e23 100644 --- a/src/main/ocr/BaseOcrProvider.ts +++ b/src/main/ocr/BaseOcrProvider.ts @@ -1,5 +1,9 @@ -import { KnowledgeBaseParams } from '@types' +import fs from 'node:fs' +import { windowService } from '@main/services/WindowService' +import { FileType, KnowledgeBaseParams } from '@types' +import Logger from 'electron-log' +import pdfParse from 'pdf-parse' export default abstract class BaseOcrProvider { protected base: KnowledgeBaseParams constructor(base: KnowledgeBaseParams) { @@ -11,12 +15,48 @@ export default abstract class BaseOcrProvider { } this.base = base } - abstract parseFile(filePath: string): Promise<{ uid: string }> - abstract exportFile(filePath: string, uid: string): Promise + abstract parseFile(sourceId: string, file: FileType): Promise<{ processedFile: FileType }> /** * 辅助方法：延迟执行 */ public delay = (ms: number): Promise => { return new Promise((resolve) => setTimeout(resolve, ms)) } + + /** + * 获取PDF文件的页数和文件大小 + * @param filePath PDF文件路径 + * @returns 包含页数和文件大小(字节)的对象 + */ + public async getPdfInfo(filePath: string): Promise<{ pageCount: number; fileSize: number }> { + try { + Logger.info(`Getting PDF info for: ${filePath}`) + + if (!filePath.toLowerCase().endsWith('.pdf')) { + throw new Error('File is not a PDF') + } + + const stats = fs.statSync(filePath) + const fileSize = stats.size + const dataBuffer = fs.readFileSync(filePath) + const pdfData = await pdfParse(dataBuffer) + + Logger.info(`File ${filePath} has ${pdfData.numpages} pages and size: ${fileSize} bytes`) + return { + pageCount: pdfData.numpages, + fileSize: fileSize + } + } catch (error) { + Logger.error(`Failed to get PDF info for ${filePath}: ${error instanceof Error ? error.message : String(error)}`) + throw new Error('Failed to get PDF information') + } + } + + public async sendOcrProgress(sourceId: string, progress: number): Promise { + const mainWindow = windowService.getMainWindow() + mainWindow?.webContents.send('file-ocr-progress', { + itemId: sourceId, + progress: progress + }) + } } diff --git a/src/main/ocr/DefaultOcrProvider.ts b/src/main/ocr/DefaultOcrProvider.ts index 58259cc7b..9f5e892cb 100644 --- a/src/main/ocr/DefaultOcrProvider.ts +++ b/src/main/ocr/DefaultOcrProvider.ts @@ -1,4 +1,4 @@ -import { KnowledgeBaseParams } from '@types' +import { FileType, KnowledgeBaseParams } from '@types' import BaseOcrProvider from './BaseOcrProvider' @@ -6,10 +6,7 @@ export default class DefaultOcrProvider extends BaseOcrProvider { constructor(base: KnowledgeBaseParams) { super(base) } - parseFile(): Promise<{ uid: string }> { - throw new Error('Method not implemented.') - } - exportFile(): Promise { + public parseFile(): Promise<{ processedFile: FileType }> { throw new Error('Method not implemented.') } } diff --git a/src/main/ocr/Doc2xOcrProvider.ts b/src/main/ocr/Doc2xOcrProvider.ts index fae4eabda..90f820aa9 100644 --- a/src/main/ocr/Doc2xOcrProvider.ts +++ b/src/main/ocr/Doc2xOcrProvider.ts @@ -1,125 +1,266 @@ import fs from 'node:fs' +import path from 'node:path' -import { KnowledgeBaseParams } from '@types' +import { FileType, KnowledgeBaseParams } from '@types' +import AdmZip from 'adm-zip' import axios, { AxiosRequestConfig } from 'axios' +import Logger from 'electron-log' import BaseOcrProvider from './BaseOcrProvider' +type ApiResponse = { + code: string + data: T + message?: string +} + +type PreuploadResponse = { + uid: string + url: string +} + +type StatusResponse = { + status: string + progress: number +} + +type ParsedFileResponse = { + status: string + url: string +} + export default class Doc2xOcrProvider extends BaseOcrProvider { constructor(base: KnowledgeBaseParams) { super(base) } - public async parseFile(filePath: string): Promise<{ uid: string }> { - console.log('Parsing file:', filePath) - const { uid, url } = await this.preupload() - console.log('uid:', uid, 'url:', url) - await this.putFile(filePath, url) - while (true) { - // 可以修改为根据progress延迟 - await this.delay(1000) - const { status, progress } = await this.getStatus(uid) - console.log('status:', status, 'progress', progress) - if (status === 'success') { - break - } else if (status === 'failed') { - throw new Error('parse failed') + + public async parseFile(sourceId: string, file: FileType): Promise<{ processedFile: FileType }> { + try { + Logger.info(`OCR processing started: ${file.path}`) + + const pdfInfo = await this.getPdfInfo(file.path) + + // 文件页数小于1000页 + if (pdfInfo.pageCount >= 1000) { + throw new Error(`PDF page count (${pdfInfo.pageCount}) exceeds the limit of 1000 pages`) } - } - console.log('Parsing file success') - return { - uid: uid + // 文件大小小于300MB + if (pdfInfo.fileSize >= 300 * 1024 * 1024) { + const fileSizeMB = Math.round(pdfInfo.fileSize / (1024 * 1024)) + throw new Error(`PDF file size (${fileSizeMB}MB) exceeds the limit of 300MB`) + } + + // 步骤1: 准备上传 + const { uid, url } = await this.preupload() + Logger.info(`OCR preupload completed: uid=${uid}`) + + // 步骤2: 上传文件 + await this.putFile(file.path, url) + + // 步骤3: 等待处理完成 + await this.waitForProcessing(sourceId, uid) + Logger.info(`OCR parsing completed successfully for: ${file.path}`) + + // 步骤4: 导出文件 + const { path: outputPath } = await this.exportFile(file.path, uid) + + // 步骤5: 创建处理后的文件信息 + return { + processedFile: this.createProcessedFileInfo(file, outputPath) + } + } catch (error) { + Logger.error(`OCR processing failed for ${file.path}: ${error instanceof Error ? error.message : String(error)}`) + throw error } } - public async exportFile(filePath: string, uid: string): Promise { - console.log('Exporting file:', filePath) + + private createProcessedFileInfo(file: FileType, outputPath: string): FileType { + const outputFilePath = `${outputPath}/${file.id}.md` + return { + ...file, + name: file.name.replace('.pdf', '.md'), + path: outputFilePath, + ext: '.md', + size: fs.statSync(outputFilePath).size + } + } + + public async exportFile(filePath: string, uid: string): Promise<{ path: string }> { + Logger.info(`Exporting file: ${filePath}`) + + // 步骤1: 转换文件 await this.convertFile(uid, filePath) - console.log('Converted file success') - let exportUrl = '' + Logger.info(`File conversion completed for: ${filePath}`) + + // 步骤2: 等待导出并获取URL + const exportUrl = await this.waitForExport(uid) + + // 步骤3: 下载并解压文件 + return this.downloadFile(exportUrl, filePath) + } + + private async waitForProcessing(sourceId: string, uid: string): Promise { + while (true) { + await this.delay(1000) + const { status, progress } = await this.getStatus(uid) + await this.sendOcrProgress(sourceId, progress) + Logger.info(`OCR processing status: ${status}, progress: ${progress}%`) + + if (status === 'success') { + return + } else if (status === 'failed') { + throw new Error('OCR processing failed') + } + } + } + + private async waitForExport(uid: string): Promise { while (true) { await this.delay(1000) const { status, url } = await this.getParsedFile(uid) - console.log('status:', status, 'url:', url) - if (status === 'success') { - exportUrl = url - break + Logger.info(`Export status: ${status}`) + + if (status === 'success' && url) { + return url } else if (status === 'failed') { - throw new Error('export failed') + throw new Error('Export failed') } } - await this.downloadFile(exportUrl, filePath) - console.log('Exported file success') } - private async preupload(): Promise<{ uid: string; url: string }> { - const config: AxiosRequestConfig = { headers: { Authorization: `Bearer ${this.base.ocrProvider?.apiKey}` } } - const { data } = await axios.post(`${this.base.ocrProvider?.apiHost}/api/v2/parse/preupload`, null, config) - if (data.code === 'success') { - return { - uid: data.data.uid, - url: data.data.url + private async preupload(): Promise { + const config = this.createAuthConfig() + const endpoint = `${this.base.ocrProvider?.apiHost}/api/v2/parse/preupload` + + try { + const { data } = await axios.post>(endpoint, null, config) + + if (data.code === 'success' && data.data) { + return data.data + } else { + throw new Error(`API returned error: ${data.message || JSON.stringify(data)}`) } - } else { - throw new Error(`get preupload url failed: ${data}`) + } catch (error) { + Logger.error(`Failed to get preupload URL: ${error instanceof Error ? error.message : String(error)}`) + throw new Error('Failed to get preupload URL') } } private async putFile(filePath: string, url: string): Promise { - const fileContent = fs.readFileSync(filePath) - const response = await axios.put(url, fileContent) - if (response.status !== 200) { - throw new Error(`put file failed: ${response.statusText}`) + try { + const fileContent = fs.readFileSync(filePath) + const response = await axios.put(url, fileContent) + + if (response.status !== 200) { + throw new Error(`HTTP status ${response.status}: ${response.statusText}`) + } + } catch (error) { + Logger.error(`Failed to upload file ${filePath}: ${error instanceof Error ? error.message : String(error)}`) + throw new Error('Failed to upload file') } } - private async getStatus(uid: string): Promise<{ status: string; progress: number }> { - const config: AxiosRequestConfig = { headers: { Authorization: `Bearer ${this.base.ocrProvider?.apiKey}` } } - const response = await axios.get(`${this.base.ocrProvider?.apiHost}/api/v2/parse/status?uid=${uid}`, config) - if (response.status !== 200) { - throw new Error(`get status failed: ${response.statusText}`) - } - if (response.data.code !== 'success') { - throw new Error(`get status failed: ${response.data}`) - } - return { - status: response.data.data.status, - progress: response.data.data.progress + private async getStatus(uid: string): Promise { + const config = this.createAuthConfig() + const endpoint = `${this.base.ocrProvider?.apiHost}/api/v2/parse/status?uid=${uid}` + + try { + const response = await axios.get>(endpoint, config) + + if (response.data.code === 'success' && response.data.data) { + return response.data.data + } else { + throw new Error(`API returned error: ${response.data.message || JSON.stringify(response.data)}`) + } + } catch (error) { + Logger.error(`Failed to get status for uid ${uid}: ${error instanceof Error ? error.message : String(error)}`) + throw new Error('Failed to get processing status') } } private async convertFile(uid: string, filePath: string): Promise { - const fileName = filePath.split('/').pop()?.split('.')[0] - const config: AxiosRequestConfig = { - headers: { Authorization: `Bearer ${this.base.ocrProvider?.apiKey}`, 'Content-Type': 'application/json' } + const fileName = path.basename(filePath).split('.')[0] + const config = { + ...this.createAuthConfig(), + headers: { + ...this.createAuthConfig().headers, + 'Content-Type': 'application/json' + } } - const data = { + + const payload = { uid, to: 'md', formula_mode: 'normal', filename: fileName } - const response = await axios.post(`${this.base.ocrProvider?.apiHost}/api/v2/convert/parse`, data, config) - if (response.data.code !== 'success') { - throw new Error(`convert file failed: ${response.statusText}`) + + const endpoint = `${this.base.ocrProvider?.apiHost}/api/v2/convert/parse` + + try { + const response = await axios.post>(endpoint, payload, config) + + if (response.data.code !== 'success') { + throw new Error(`API returned error: ${response.data.message || JSON.stringify(response.data)}`) + } + } catch (error) { + Logger.error(`Failed to convert file ${filePath}: ${error instanceof Error ? error.message : String(error)}`) + throw new Error('Failed to convert file') } } - private async getParsedFile(uid: string): Promise<{ status: string; url: string }> { - const config: AxiosRequestConfig = { headers: { Authorization: `Bearer ${this.base.ocrProvider?.apiKey}` } } - const response = await axios.get(`${this.base.ocrProvider?.apiHost}/api/v2/convert/parse/result?uid=${uid}`, config) - if (response.status !== 200) { - throw new Error(`get parsed file failed: ${response.statusText}`) + private async getParsedFile(uid: string): Promise { + const config = this.createAuthConfig() + const endpoint = `${this.base.ocrProvider?.apiHost}/api/v2/convert/parse/result?uid=${uid}` + + try { + const response = await axios.get>(endpoint, config) + + if (response.status === 200 && response.data.data) { + return response.data.data + } else { + throw new Error(`HTTP status ${response.status}: ${response.statusText}`) + } + } catch (error) { + Logger.error( + `Failed to get parsed file for uid ${uid}: ${error instanceof Error ? error.message : String(error)}` + ) + throw new Error('Failed to get parsed file information') } + } + + private async downloadFile(url: string, filePath: string): Promise<{ path: string }> { + const dirPath = path.dirname(filePath) + const baseName = path.basename(filePath, path.extname(filePath)) + const zipPath = path.join(dirPath, `${baseName}.zip`) + const extractPath = path.join(dirPath, baseName) + + Logger.info(`Downloading to export path: ${zipPath}`) + + try { + // 下载文件 + const response = await axios.get(url, { responseType: 'arraybuffer' }) + fs.writeFileSync(zipPath, response.data) + + // 解压文件 + const zip = new AdmZip(zipPath) + zip.extractAllTo(extractPath, true) + + // 删除临时ZIP文件 + fs.unlinkSync(zipPath) + + return { path: extractPath } + } catch (error) { + Logger.error(`Failed to download and extract file: ${error instanceof Error ? error.message : String(error)}`) + throw new Error('Failed to download and extract file') + } + } + + private createAuthConfig(): AxiosRequestConfig { return { - status: response.data.data.status, - url: response.data.data.url + headers: { + Authorization: `Bearer ${this.base.ocrProvider?.apiKey}` + } } } - - private async downloadFile(url: string, filePath: string): Promise { - const fileName = filePath.split('/').pop()?.split('.')[0] + '.zip' - const exportPath = `${filePath.split('/').slice(0, -1).join('/')}/${fileName}` - console.log('exportPath:', exportPath) - const response = await axios.get(url, { responseType: 'arraybuffer' }) - fs.writeFileSync(exportPath, response.data) - } } diff --git a/src/main/ocr/OcrProvider.ts b/src/main/ocr/OcrProvider.ts index 5e28a898f..032c3062c 100644 --- a/src/main/ocr/OcrProvider.ts +++ b/src/main/ocr/OcrProvider.ts @@ -1,4 +1,4 @@ -import { KnowledgeBaseParams } from '@types' +import { FileType, KnowledgeBaseParams } from '@types' import BaseOcrProvider from './BaseOcrProvider' import OcrProviderFactory from './OcrProviderFactory' @@ -8,10 +8,7 @@ export default class OcrProvider { constructor(base: KnowledgeBaseParams) { this.sdk = OcrProviderFactory.create(base) } - public async parseFile(filePath: string): Promise<{ uid: string }> { - return this.sdk.parseFile(filePath) - } - public async exportFile(filePath: string, uid: string): Promise { - return this.sdk.exportFile(filePath, uid) + public async parseFile(sourceId: string, file: FileType): Promise<{ processedFile: FileType }> { + return this.sdk.parseFile(sourceId, file) } } diff --git a/src/main/services/FileStorage.ts b/src/main/services/FileStorage.ts index 9557e17c9..252435740 100644 --- a/src/main/services/FileStorage.ts +++ b/src/main/services/FileStorage.ts @@ -211,6 +211,10 @@ class FileStorage { await fs.promises.unlink(path.join(this.storageDir, id)) } + public deleteDir = async (_: Electron.IpcMainInvokeEvent, id: string): Promise => { + await fs.promises.rm(path.join(this.storageDir, id), { recursive: true }) + } + public readFile = async (_: Electron.IpcMainInvokeEvent, id: string): Promise => { const filePath = path.join(this.storageDir, id) diff --git a/src/main/services/KnowledgeService.ts b/src/main/services/KnowledgeService.ts index eb4c11717..d0bd5e6d8 100644 --- a/src/main/services/KnowledgeService.ts +++ b/src/main/services/KnowledgeService.ts @@ -132,7 +132,8 @@ class KnowledgeService { model, apiKey, dimensions, - batchSize + batchSize, + configuration: { baseURL } }) ) .setVectorDatabase(new LibSqlDb({ path: path.join(this.storageDir, id) })) @@ -149,6 +150,7 @@ class KnowledgeService { } public delete = async (_: Electron.IpcMainInvokeEvent, id: string): Promise => { + console.log('id', id) const dbPath = path.join(this.storageDir, id) if (fs.existsSync(dbPath)) { fs.rmSync(dbPath, { recursive: true }) @@ -161,7 +163,6 @@ class KnowledgeService { this.workload >= KnowledgeService.MAXIMUM_WORKLOAD ) } - private fileTask( ragApplication: RAGApplication, options: KnowledgeBaseAddItemOptionsNonNullableAttribute @@ -173,8 +174,27 @@ class KnowledgeService { loaderTasks: [ { state: LoaderTaskItemState.PENDING, - task: () => - addFileLoader(ragApplication, file, base, forceReload) + task: async () => { + // 添加OCR预处理逻辑 + let fileToProcess: FileType = file + if (base.preprocessing && file.ext.toLowerCase() === '.pdf') { + try { + const ocrProvider = new OcrProvider(base) + Logger.info(`Starting OCR processing for file: ${file.path}`) + + const { processedFile } = await ocrProvider.parseFile(item.id, file) + + fileToProcess = processedFile + Logger.info(`OCR processing completed: ${fileToProcess.path}`) + } catch (err) { + Logger.error(`OCR processing failed: ${err}`) + // 如果OCR失败，使用原始文件 + fileToProcess = file + } + } + + // 使用处理后的文件进行加载 + return addFileLoader(ragApplication, fileToProcess, base, forceReload) .then((result) => { loaderTask.loaderDoneReturn = result return result @@ -182,7 +202,8 @@ class KnowledgeService { .catch((err) => { Logger.error(err) return KnowledgeService.ERROR_LOADER_RETURN - }), + }) + }, evaluateTaskWorkload: { workload: file.size } } ], @@ -424,12 +445,6 @@ class KnowledgeService { } public add = async (_: Electron.IpcMainInvokeEvent, options: KnowledgeBaseAddItemOptions): Promise => { - const { base, item } = options - if (base.preprocessing) { - const ocrProvider = new OcrProvider(base) - const { uid } = await ocrProvider.parseFile((item.content as FileType).path) - await ocrProvider.exportFile((item.content as FileType).path, uid) - } return new Promise((resolve) => { const { base, item, forceReload = false } = options const optionsNonNullableAttribute = { base, item, forceReload } diff --git a/src/preload/index.d.ts b/src/preload/index.d.ts index 4d0b88048..ba14b8210 100644 --- a/src/preload/index.d.ts +++ b/src/preload/index.d.ts @@ -47,6 +47,7 @@ declare global { select: (options?: OpenDialogOptions) => Promise upload: (file: FileType) => Promise delete: (fileId: string) => Promise + deleteDir: (dirPath: string) => Promise read: (fileId: string) => Promise clear: () => Promise get: (filePath: string) => Promise diff --git a/src/preload/index.ts b/src/preload/index.ts index 831b59a7f..5072f4a45 100644 --- a/src/preload/index.ts +++ b/src/preload/index.ts @@ -37,6 +37,7 @@ const api = { select: (options?: OpenDialogOptions) => ipcRenderer.invoke('file:select', options), upload: (filePath: string) => ipcRenderer.invoke('file:upload', filePath), delete: (fileId: string) => ipcRenderer.invoke('file:delete', fileId), + deleteDir: (dirPath: string) => ipcRenderer.invoke('file:deleteDir', dirPath), read: (fileId: string) => ipcRenderer.invoke('file:read', fileId), clear: () => ipcRenderer.invoke('file:clear'), get: (filePath: string) => ipcRenderer.invoke('file:get', filePath), diff --git a/src/renderer/src/hooks/useKnowledge.ts b/src/renderer/src/hooks/useKnowledge.ts index 01a396772..51f145b91 100644 --- a/src/renderer/src/hooks/useKnowledge.ts +++ b/src/renderer/src/hooks/useKnowledge.ts @@ -149,6 +149,7 @@ export const useKnowledge = (baseId: string) => { } if (item.type === 'file' && typeof item.content === 'object') { await FileManager.deleteFile(item.content.id) + await window.api.file.deleteDir(item.content.id) } } // 刷新项目 @@ -225,6 +226,27 @@ export const useKnowledge = (baseId: string) => { return percent } + // 获取文件ocr处理进度 + const getFileOcrProgress = (itemId: string) => { + const [progress, setProgress] = useState(0) + useEffect(() => { + const cleanup = window.electron.ipcRenderer.on( + 'file-ocr-progress', + (_, { itemId: id, progress }: { itemId: string; progress: number }) => { + if (itemId === id) { + setProgress(progress) + } + } + ) + + return () => { + cleanup() + } + }, [itemId]) + + return progress + } + // 清除已完成的项目 const clearCompleted = () => { dispatch(clearCompletedProcessing({ baseId })) @@ -308,6 +330,7 @@ export const useKnowledge = (baseId: string) => { getProcessingStatus, getProcessingItemsByType, getDirectoryProcessingPercent, + getFileOcrProgress, clearCompleted, clearAll, removeItem, diff --git a/src/renderer/src/pages/knowledge/KnowledgeContent.tsx b/src/renderer/src/pages/knowledge/KnowledgeContent.tsx index 67b963edd..2dab9434e 100644 --- a/src/renderer/src/pages/knowledge/KnowledgeContent.tsx +++ b/src/renderer/src/pages/knowledge/KnowledgeContent.tsx @@ -55,6 +55,7 @@ const KnowledgeContent: FC = ({ selectedBase }) => { removeItem, getProcessingStatus, getDirectoryProcessingPercent, + getFileOcrProgress, addNote, addDirectory, updateItem @@ -69,6 +70,11 @@ const KnowledgeContent: FC = ({ selectedBase }) => { const getProgressingPercentForItem = (itemId: string) => getDirectoryProcessingPercent(itemId) + const getFileOcrProgressForItem = (itemId: string) => { + console.log('[KnowledgeContent] getFileOcrProgressForItem:', itemId) + return getFileOcrProgress(itemId) + } + const handleAddFile = () => { if (disabled) { return @@ -256,6 +262,7 @@ const KnowledgeContent: FC = ({ selectedBase }) => { {fileItems.reverse().map((item) => { const file = item.content as FileType + console.log('item', item) return ( @@ -270,7 +277,13 @@ const KnowledgeContent: FC = ({ selectedBase }) => { {item.uniqueId &&