fix: rename knowledage -> knowledge
This commit is contained in:
kangfenmao
@@ -0,0 +1,126 @@
|
||||
import fs from 'node:fs'
|
||||
import path from 'node:path'
|
||||
|
||||
import { windowService } from '@main/services/WindowService'
|
||||
import { getFileExt } from '@main/utils/file'
|
||||
import { FileMetadata, PreprocessProvider } from '@types'
|
||||
import { app } from 'electron'
|
||||
import pdfjs from 'pdfjs-dist'
|
||||
import { TypedArray } from 'pdfjs-dist/types/src/display/api'
|
||||
|
||||
export default abstract class BasePreprocessProvider {
|
||||
protected provider: PreprocessProvider
|
||||
protected userId?: string
|
||||
public storageDir = path.join(app.getPath('userData'), 'Data', 'Files')
|
||||
|
||||
constructor(provider: PreprocessProvider, userId?: string) {
|
||||
if (!provider) {
|
||||
throw new Error('Preprocess provider is not set')
|
||||
}
|
||||
this.provider = provider
|
||||
this.userId = userId
|
||||
}
|
||||
abstract parseFile(sourceId: string, file: FileMetadata): Promise<{ processedFile: FileMetadata; quota?: number }>
|
||||
|
||||
abstract checkQuota(): Promise<number>
|
||||
|
||||
/**
|
||||
* 检查文件是否已经被预处理过
|
||||
* 统一检测方法：如果 Data/Files/{file.id} 是目录，说明已被预处理
|
||||
* @param file 文件信息
|
||||
* @returns 如果已处理返回处理后的文件信息，否则返回null
|
||||
*/
|
||||
public async checkIfAlreadyProcessed(file: FileMetadata): Promise<FileMetadata | null> {
|
||||
try {
|
||||
// 检查 Data/Files/{file.id} 是否是目录
|
||||
const preprocessDirPath = path.join(this.storageDir, file.id)
|
||||
|
||||
if (fs.existsSync(preprocessDirPath)) {
|
||||
const stats = await fs.promises.stat(preprocessDirPath)
|
||||
|
||||
// 如果是目录，说明已经被预处理过
|
||||
if (stats.isDirectory()) {
|
||||
// 查找目录中的处理结果文件
|
||||
const files = await fs.promises.readdir(preprocessDirPath)
|
||||
|
||||
// 查找主要的处理结果文件（.md 或 .txt）
|
||||
const processedFile = files.find((fileName) => fileName.endsWith('.md') || fileName.endsWith('.txt'))
|
||||
|
||||
if (processedFile) {
|
||||
const processedFilePath = path.join(preprocessDirPath, processedFile)
|
||||
const processedStats = await fs.promises.stat(processedFilePath)
|
||||
const ext = getFileExt(processedFile)
|
||||
|
||||
return {
|
||||
...file,
|
||||
name: file.name.replace(file.ext, ext),
|
||||
path: processedFilePath,
|
||||
ext: ext,
|
||||
size: processedStats.size,
|
||||
created_at: processedStats.birthtime.toISOString()
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
return null
|
||||
} catch (error) {
|
||||
// 如果检查过程中出现错误，返回null表示未处理
|
||||
return null
|
||||
}
|
||||
}
|
||||
|
||||
/**
|
||||
* 辅助方法：延迟执行
|
||||
*/
|
||||
public delay = (ms: number): Promise<void> => {
|
||||
return new Promise((resolve) => setTimeout(resolve, ms))
|
||||
}
|
||||
|
||||
public async readPdf(
|
||||
source: string | URL | TypedArray,
|
||||
passwordCallback?: (fn: (password: string) => void, reason: string) => string
|
||||
) {
|
||||
const documentLoadingTask = pdfjs.getDocument(source)
|
||||
if (passwordCallback) {
|
||||
documentLoadingTask.onPassword = passwordCallback
|
||||
}
|
||||
|
||||
const document = await documentLoadingTask.promise
|
||||
return document
|
||||
}
|
||||
|
||||
public async sendPreprocessProgress(sourceId: string, progress: number): Promise<void> {
|
||||
const mainWindow = windowService.getMainWindow()
|
||||
mainWindow?.webContents.send('file-preprocess-progress', {
|
||||
itemId: sourceId,
|
||||
progress: progress
|
||||
})
|
||||
}
|
||||
|
||||
/**
|
||||
* 将文件移动到附件目录
|
||||
* @param fileId 文件id
|
||||
* @param filePaths 需要移动的文件路径数组
|
||||
* @returns 移动后的文件路径数组
|
||||
*/
|
||||
public moveToAttachmentsDir(fileId: string, filePaths: string[]): string[] {
|
||||
const attachmentsPath = path.join(this.storageDir, fileId)
|
||||
if (!fs.existsSync(attachmentsPath)) {
|
||||
fs.mkdirSync(attachmentsPath, { recursive: true })
|
||||
}
|
||||
|
||||
const movedPaths: string[] = []
|
||||
|
||||
for (const filePath of filePaths) {
|
||||
if (fs.existsSync(filePath)) {
|
||||
const fileName = path.basename(filePath)
|
||||
const destPath = path.join(attachmentsPath, fileName)
|
||||
fs.copyFileSync(filePath, destPath)
|
||||
fs.unlinkSync(filePath) // 删除原文件，实现"移动"
|
||||
movedPaths.push(destPath)
|
||||
}
|
||||
}
|
||||
return movedPaths
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,16 @@
|
||||
import { FileMetadata, PreprocessProvider } from '@types'
|
||||
|
||||
import BasePreprocessProvider from './BasePreprocessProvider'
|
||||
|
||||
export default class DefaultPreprocessProvider extends BasePreprocessProvider {
|
||||
constructor(provider: PreprocessProvider) {
|
||||
super(provider)
|
||||
}
|
||||
public parseFile(): Promise<{ processedFile: FileMetadata }> {
|
||||
throw new Error('Method not implemented.')
|
||||
}
|
||||
|
||||
public checkQuota(): Promise<number> {
|
||||
throw new Error('Method not implemented.')
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,331 @@
|
||||
import fs from 'node:fs'
|
||||
import path from 'node:path'
|
||||
|
||||
import { loggerService } from '@logger'
|
||||
import { FileMetadata, PreprocessProvider } from '@types'
|
||||
import AdmZip from 'adm-zip'
|
||||
import axios, { AxiosRequestConfig } from 'axios'
|
||||
|
||||
import BasePreprocessProvider from './BasePreprocessProvider'
|
||||
|
||||
const logger = loggerService.withContext('Doc2xPreprocessProvider')
|
||||
|
||||
type ApiResponse<T> = {
|
||||
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 Doc2xPreprocessProvider extends BasePreprocessProvider {
|
||||
constructor(provider: PreprocessProvider) {
|
||||
super(provider)
|
||||
}
|
||||
|
||||
private async validateFile(filePath: string): Promise<void> {
|
||||
const pdfBuffer = await fs.promises.readFile(filePath)
|
||||
|
||||
const doc = await this.readPdf(new Uint8Array(pdfBuffer))
|
||||
|
||||
// 文件页数小于1000页
|
||||
if (doc.numPages >= 1000) {
|
||||
throw new Error(`PDF page count (${doc.numPages}) exceeds the limit of 1000 pages`)
|
||||
}
|
||||
// 文件大小小于300MB
|
||||
if (pdfBuffer.length >= 300 * 1024 * 1024) {
|
||||
const fileSizeMB = Math.round(pdfBuffer.length / (1024 * 1024))
|
||||
throw new Error(`PDF file size (${fileSizeMB}MB) exceeds the limit of 300MB`)
|
||||
}
|
||||
}
|
||||
|
||||
public async parseFile(sourceId: string, file: FileMetadata): Promise<{ processedFile: FileMetadata }> {
|
||||
try {
|
||||
logger.info(`Preprocess processing started: ${file.path}`)
|
||||
|
||||
// 步骤1: 准备上传
|
||||
const { uid, url } = await this.preupload()
|
||||
logger.info(`Preprocess preupload completed: uid=${uid}`)
|
||||
|
||||
await this.validateFile(file.path)
|
||||
|
||||
// 步骤2: 上传文件
|
||||
await this.putFile(file.path, url)
|
||||
|
||||
// 步骤3: 等待处理完成
|
||||
await this.waitForProcessing(sourceId, uid)
|
||||
logger.info(`Preprocess parsing completed successfully for: ${file.path}`)
|
||||
|
||||
// 步骤4: 导出文件
|
||||
const { path: outputPath } = await this.exportFile(file, uid)
|
||||
|
||||
// 步骤5: 创建处理后的文件信息
|
||||
return {
|
||||
processedFile: this.createProcessedFileInfo(file, outputPath)
|
||||
}
|
||||
} catch (error) {
|
||||
logger.error(
|
||||
`Preprocess processing failed for ${file.path}: ${error instanceof Error ? error.message : String(error)}`
|
||||
)
|
||||
throw error
|
||||
}
|
||||
}
|
||||
|
||||
private createProcessedFileInfo(file: FileMetadata, outputPath: string): FileMetadata {
|
||||
const outputFilePath = `${outputPath}/${file.name.split('.').slice(0, -1).join('.')}.md`
|
||||
return {
|
||||
...file,
|
||||
name: file.name.replace('.pdf', '.md'),
|
||||
path: outputFilePath,
|
||||
ext: '.md',
|
||||
size: fs.statSync(outputFilePath).size
|
||||
}
|
||||
}
|
||||
|
||||
/**
|
||||
* 导出文件
|
||||
* @param file 文件信息
|
||||
* @param uid 预上传响应的uid
|
||||
* @returns 导出文件的路径
|
||||
*/
|
||||
public async exportFile(file: FileMetadata, uid: string): Promise<{ path: string }> {
|
||||
logger.info(`Exporting file: ${file.path}`)
|
||||
|
||||
// 步骤1: 转换文件
|
||||
await this.convertFile(uid, file.path)
|
||||
logger.info(`File conversion completed for: ${file.path}`)
|
||||
|
||||
// 步骤2: 等待导出并获取URL
|
||||
const exportUrl = await this.waitForExport(uid)
|
||||
|
||||
// 步骤3: 下载并解压文件
|
||||
return this.downloadFile(exportUrl, file)
|
||||
}
|
||||
|
||||
/**
|
||||
* 等待处理完成
|
||||
* @param sourceId 源文件ID
|
||||
* @param uid 预上传响应的uid
|
||||
*/
|
||||
private async waitForProcessing(sourceId: string, uid: string): Promise<void> {
|
||||
while (true) {
|
||||
await this.delay(1000)
|
||||
const { status, progress } = await this.getStatus(uid)
|
||||
await this.sendPreprocessProgress(sourceId, progress)
|
||||
logger.info(`Preprocess processing status: ${status}, progress: ${progress}%`)
|
||||
|
||||
if (status === 'success') {
|
||||
return
|
||||
} else if (status === 'failed') {
|
||||
throw new Error('Preprocess processing failed')
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/**
|
||||
* 等待导出完成
|
||||
* @param uid 预上传响应的uid
|
||||
* @returns 导出文件的url
|
||||
*/
|
||||
private async waitForExport(uid: string): Promise<string> {
|
||||
while (true) {
|
||||
await this.delay(1000)
|
||||
const { status, url } = await this.getParsedFile(uid)
|
||||
logger.info(`Export status: ${status}`)
|
||||
|
||||
if (status === 'success' && url) {
|
||||
return url
|
||||
} else if (status === 'failed') {
|
||||
throw new Error('Export failed')
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/**
|
||||
* 预上传文件
|
||||
* @returns 预上传响应的url和uid
|
||||
*/
|
||||
private async preupload(): Promise<PreuploadResponse> {
|
||||
const config = this.createAuthConfig()
|
||||
const endpoint = `${this.provider.apiHost}/api/v2/parse/preupload`
|
||||
|
||||
try {
|
||||
const { data } = await axios.post<ApiResponse<PreuploadResponse>>(endpoint, null, config)
|
||||
|
||||
if (data.code === 'success' && data.data) {
|
||||
return data.data
|
||||
} else {
|
||||
throw new Error(`API returned error: ${data.message || JSON.stringify(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')
|
||||
}
|
||||
}
|
||||
|
||||
/**
|
||||
* 上传文件
|
||||
* @param filePath 文件路径
|
||||
* @param url 预上传响应的url
|
||||
*/
|
||||
private async putFile(filePath: string, url: string): Promise<void> {
|
||||
try {
|
||||
const fileStream = fs.createReadStream(filePath)
|
||||
const response = await axios.put(url, fileStream)
|
||||
|
||||
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<StatusResponse> {
|
||||
const config = this.createAuthConfig()
|
||||
const endpoint = `${this.provider.apiHost}/api/v2/parse/status?uid=${uid}`
|
||||
|
||||
try {
|
||||
const response = await axios.get<ApiResponse<StatusResponse>>(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')
|
||||
}
|
||||
}
|
||||
|
||||
/**
|
||||
* Preprocess文件
|
||||
* @param uid 预上传响应的uid
|
||||
* @param filePath 文件路径
|
||||
*/
|
||||
private async convertFile(uid: string, filePath: string): Promise<void> {
|
||||
const fileName = path.parse(filePath).name
|
||||
const config = {
|
||||
...this.createAuthConfig(),
|
||||
headers: {
|
||||
...this.createAuthConfig().headers,
|
||||
'Content-Type': 'application/json'
|
||||
}
|
||||
}
|
||||
|
||||
const payload = {
|
||||
uid,
|
||||
to: 'md',
|
||||
formula_mode: 'normal',
|
||||
filename: fileName
|
||||
}
|
||||
|
||||
const endpoint = `${this.provider.apiHost}/api/v2/convert/parse`
|
||||
|
||||
try {
|
||||
const response = await axios.post<ApiResponse<any>>(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')
|
||||
}
|
||||
}
|
||||
|
||||
/**
|
||||
* 获取解析后的文件信息
|
||||
* @param uid 预上传响应的uid
|
||||
* @returns 解析后的文件信息
|
||||
*/
|
||||
private async getParsedFile(uid: string): Promise<ParsedFileResponse> {
|
||||
const config = this.createAuthConfig()
|
||||
const endpoint = `${this.provider.apiHost}/api/v2/convert/parse/result?uid=${uid}`
|
||||
|
||||
try {
|
||||
const response = await axios.get<ApiResponse<ParsedFileResponse>>(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')
|
||||
}
|
||||
}
|
||||
|
||||
/**
|
||||
* 下载文件
|
||||
* @param url 导出文件的url
|
||||
* @param file 文件信息
|
||||
* @returns 下载文件的路径
|
||||
*/
|
||||
private async downloadFile(url: string, file: FileMetadata): Promise<{ path: string }> {
|
||||
const dirPath = this.storageDir
|
||||
// 使用统一的存储路径：Data/Files/{file.id}/
|
||||
const extractPath = path.join(dirPath, file.id)
|
||||
const zipPath = path.join(dirPath, `${file.id}.zip`)
|
||||
|
||||
// 确保目录存在
|
||||
fs.mkdirSync(dirPath, { recursive: true })
|
||||
fs.mkdirSync(extractPath, { recursive: true })
|
||||
|
||||
logger.info(`Downloading to export path: ${zipPath}`)
|
||||
|
||||
try {
|
||||
// 下载文件
|
||||
const response = await axios.get(url, { responseType: 'arraybuffer' })
|
||||
fs.writeFileSync(zipPath, response.data)
|
||||
|
||||
// 确保提取目录存在
|
||||
if (!fs.existsSync(extractPath)) {
|
||||
fs.mkdirSync(extractPath, { recursive: true })
|
||||
}
|
||||
|
||||
// 解压文件
|
||||
const zip = new AdmZip(zipPath)
|
||||
zip.extractAllTo(extractPath, true)
|
||||
logger.info(`Extracted files to: ${extractPath}`)
|
||||
|
||||
// 删除临时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 {
|
||||
headers: {
|
||||
Authorization: `Bearer ${this.provider.apiKey}`
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
public checkQuota(): Promise<number> {
|
||||
throw new Error('Method not implemented.')
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,396 @@
|
||||
import fs from 'node:fs'
|
||||
import path from 'node:path'
|
||||
|
||||
import { loggerService } from '@logger'
|
||||
import { FileMetadata, PreprocessProvider } from '@types'
|
||||
import AdmZip from 'adm-zip'
|
||||
import axios from 'axios'
|
||||
|
||||
import BasePreprocessProvider from './BasePreprocessProvider'
|
||||
|
||||
const logger = loggerService.withContext('MineruPreprocessProvider')
|
||||
|
||||
type ApiResponse<T> = {
|
||||
code: number
|
||||
data: T
|
||||
msg?: string
|
||||
trace_id?: string
|
||||
}
|
||||
|
||||
type BatchUploadResponse = {
|
||||
batch_id: string
|
||||
file_urls: string[]
|
||||
}
|
||||
|
||||
type ExtractProgress = {
|
||||
extracted_pages: number
|
||||
total_pages: number
|
||||
start_time: string
|
||||
}
|
||||
|
||||
type ExtractFileResult = {
|
||||
file_name: string
|
||||
state: 'done' | 'waiting-file' | 'pending' | 'running' | 'converting' | 'failed'
|
||||
err_msg: string
|
||||
full_zip_url?: string
|
||||
extract_progress?: ExtractProgress
|
||||
}
|
||||
|
||||
type ExtractResultResponse = {
|
||||
batch_id: string
|
||||
extract_result: ExtractFileResult[]
|
||||
}
|
||||
|
||||
type QuotaResponse = {
|
||||
code: number
|
||||
data: {
|
||||
user_left_quota: number
|
||||
total_left_quota: number
|
||||
}
|
||||
msg?: string
|
||||
trace_id?: string
|
||||
}
|
||||
|
||||
export default class MineruPreprocessProvider extends BasePreprocessProvider {
|
||||
constructor(provider: PreprocessProvider, userId?: string) {
|
||||
super(provider, userId)
|
||||
// todo：免费期结束后删除
|
||||
this.provider.apiKey = this.provider.apiKey || import.meta.env.MAIN_VITE_MINERU_API_KEY
|
||||
}
|
||||
|
||||
public async parseFile(
|
||||
sourceId: string,
|
||||
file: FileMetadata
|
||||
): Promise<{ processedFile: FileMetadata; quota: number }> {
|
||||
try {
|
||||
logger.info(`MinerU preprocess processing started: ${file.path}`)
|
||||
await this.validateFile(file.path)
|
||||
|
||||
// 1. 获取上传URL并上传文件
|
||||
const batchId = await this.uploadFile(file)
|
||||
logger.info(`MinerU file upload completed: batch_id=${batchId}`)
|
||||
|
||||
// 2. 等待处理完成并获取结果
|
||||
const extractResult = await this.waitForCompletion(sourceId, batchId, file.origin_name)
|
||||
logger.info(`MinerU processing completed for batch: ${batchId}`)
|
||||
|
||||
// 3. 下载并解压文件
|
||||
const { path: outputPath } = await this.downloadAndExtractFile(extractResult.full_zip_url!, file)
|
||||
|
||||
// 4. check quota
|
||||
const quota = await this.checkQuota()
|
||||
|
||||
// 5. 创建处理后的文件信息
|
||||
return {
|
||||
processedFile: this.createProcessedFileInfo(file, outputPath),
|
||||
quota
|
||||
}
|
||||
} catch (error: any) {
|
||||
logger.error(`MinerU preprocess processing failed for ${file.path}: ${error.message}`)
|
||||
throw new Error(error.message)
|
||||
}
|
||||
}
|
||||
|
||||
public async checkQuota() {
|
||||
try {
|
||||
const quota = await fetch(`${this.provider.apiHost}/api/v4/quota`, {
|
||||
method: 'GET',
|
||||
headers: {
|
||||
'Content-Type': 'application/json',
|
||||
Authorization: `Bearer ${this.provider.apiKey}`,
|
||||
token: this.userId ?? ''
|
||||
}
|
||||
})
|
||||
if (!quota.ok) {
|
||||
throw new Error(`HTTP ${quota.status}: ${quota.statusText}`)
|
||||
}
|
||||
const response: QuotaResponse = await quota.json()
|
||||
return response.data.user_left_quota
|
||||
} catch (error) {
|
||||
logger.error('Error checking quota:', error)
|
||||
throw error
|
||||
}
|
||||
}
|
||||
|
||||
private async validateFile(filePath: string): Promise<void> {
|
||||
const pdfBuffer = await fs.promises.readFile(filePath)
|
||||
|
||||
const doc = await this.readPdf(new Uint8Array(pdfBuffer))
|
||||
|
||||
// 文件页数小于600页
|
||||
if (doc.numPages >= 600) {
|
||||
throw new Error(`PDF page count (${doc.numPages}) exceeds the limit of 600 pages`)
|
||||
}
|
||||
// 文件大小小于200MB
|
||||
if (pdfBuffer.length >= 200 * 1024 * 1024) {
|
||||
const fileSizeMB = Math.round(pdfBuffer.length / (1024 * 1024))
|
||||
throw new Error(`PDF file size (${fileSizeMB}MB) exceeds the limit of 200MB`)
|
||||
}
|
||||
}
|
||||
|
||||
private createProcessedFileInfo(file: FileMetadata, outputPath: string): FileMetadata {
|
||||
// 查找解压后的主要文件
|
||||
let finalPath = ''
|
||||
let finalName = file.origin_name.replace('.pdf', '.md')
|
||||
|
||||
try {
|
||||
const files = fs.readdirSync(outputPath)
|
||||
|
||||
const mdFile = files.find((f) => f.endsWith('.md'))
|
||||
if (mdFile) {
|
||||
const originalMdPath = path.join(outputPath, mdFile)
|
||||
const newMdPath = path.join(outputPath, finalName)
|
||||
|
||||
// 重命名文件为原始文件名
|
||||
try {
|
||||
fs.renameSync(originalMdPath, newMdPath)
|
||||
finalPath = newMdPath
|
||||
logger.info(`Renamed markdown file from ${mdFile} to ${finalName}`)
|
||||
} catch (renameError) {
|
||||
logger.warn(`Failed to rename file ${mdFile} to ${finalName}: ${renameError}`)
|
||||
// 如果重命名失败，使用原文件
|
||||
finalPath = originalMdPath
|
||||
finalName = mdFile
|
||||
}
|
||||
}
|
||||
} catch (error) {
|
||||
logger.warn(`Failed to read output directory ${outputPath}: ${error}`)
|
||||
finalPath = path.join(outputPath, `${file.id}.md`)
|
||||
}
|
||||
|
||||
return {
|
||||
...file,
|
||||
name: finalName,
|
||||
path: finalPath,
|
||||
ext: '.md',
|
||||
size: fs.existsSync(finalPath) ? fs.statSync(finalPath).size : 0
|
||||
}
|
||||
}
|
||||
|
||||
private async downloadAndExtractFile(zipUrl: string, file: FileMetadata): Promise<{ path: string }> {
|
||||
const dirPath = this.storageDir
|
||||
|
||||
const zipPath = path.join(dirPath, `${file.id}.zip`)
|
||||
const extractPath = path.join(dirPath, `${file.id}`)
|
||||
|
||||
logger.info(`Downloading MinerU result to: ${zipPath}`)
|
||||
|
||||
try {
|
||||
// 下载ZIP文件
|
||||
const response = await axios.get(zipUrl, { responseType: 'arraybuffer' })
|
||||
fs.writeFileSync(zipPath, response.data)
|
||||
logger.info(`Downloaded ZIP file: ${zipPath}`)
|
||||
|
||||
// 确保提取目录存在
|
||||
if (!fs.existsSync(extractPath)) {
|
||||
fs.mkdirSync(extractPath, { recursive: true })
|
||||
}
|
||||
|
||||
// 解压文件
|
||||
const zip = new AdmZip(zipPath)
|
||||
zip.extractAllTo(extractPath, true)
|
||||
logger.info(`Extracted files to: ${extractPath}`)
|
||||
|
||||
// 删除临时ZIP文件
|
||||
fs.unlinkSync(zipPath)
|
||||
|
||||
return { path: extractPath }
|
||||
} catch (error: any) {
|
||||
logger.error(`Failed to download and extract file: ${error.message}`)
|
||||
throw new Error(error.message)
|
||||
}
|
||||
}
|
||||
|
||||
private async uploadFile(file: FileMetadata): Promise<string> {
|
||||
try {
|
||||
// 步骤1: 获取上传URL
|
||||
const { batchId, fileUrls } = await this.getBatchUploadUrls(file)
|
||||
logger.debug(`Got upload URLs for batch: ${batchId}`)
|
||||
|
||||
logger.debug(`batchId: ${batchId}, fileurls: ${fileUrls}`)
|
||||
// 步骤2: 上传文件到获取的URL
|
||||
await this.putFileToUrl(file.path, fileUrls[0])
|
||||
logger.info(`File uploaded successfully: ${file.path}`)
|
||||
|
||||
return batchId
|
||||
} catch (error: any) {
|
||||
logger.error(`Failed to upload file ${file.path}: ${error.message}`)
|
||||
throw new Error(error.message)
|
||||
}
|
||||
}
|
||||
|
||||
private async getBatchUploadUrls(file: FileMetadata): Promise<{ batchId: string; fileUrls: string[] }> {
|
||||
const endpoint = `${this.provider.apiHost}/api/v4/file-urls/batch`
|
||||
|
||||
const payload = {
|
||||
language: 'auto',
|
||||
enable_formula: true,
|
||||
enable_table: true,
|
||||
files: [
|
||||
{
|
||||
name: file.origin_name,
|
||||
is_ocr: true,
|
||||
data_id: file.id
|
||||
}
|
||||
]
|
||||
}
|
||||
|
||||
try {
|
||||
const response = await fetch(endpoint, {
|
||||
method: 'POST',
|
||||
headers: {
|
||||
'Content-Type': 'application/json',
|
||||
Authorization: `Bearer ${this.provider.apiKey}`,
|
||||
token: this.userId ?? '',
|
||||
Accept: '*/*'
|
||||
},
|
||||
body: JSON.stringify(payload)
|
||||
})
|
||||
|
||||
if (response.ok) {
|
||||
const data: ApiResponse<BatchUploadResponse> = await response.json()
|
||||
if (data.code === 0 && data.data) {
|
||||
const { batch_id, file_urls } = data.data
|
||||
return {
|
||||
batchId: batch_id,
|
||||
fileUrls: file_urls
|
||||
}
|
||||
} else {
|
||||
throw new Error(`API returned error: ${data.msg || JSON.stringify(data)}`)
|
||||
}
|
||||
} else {
|
||||
throw new Error(`HTTP ${response.status}: ${response.statusText}`)
|
||||
}
|
||||
} catch (error: any) {
|
||||
logger.error(`Failed to get batch upload URLs: ${error.message}`)
|
||||
throw new Error(error.message)
|
||||
}
|
||||
}
|
||||
|
||||
private async putFileToUrl(filePath: string, uploadUrl: string): Promise<void> {
|
||||
try {
|
||||
const fileBuffer = await fs.promises.readFile(filePath)
|
||||
|
||||
const response = await fetch(uploadUrl, {
|
||||
method: 'PUT',
|
||||
body: fileBuffer,
|
||||
headers: {
|
||||
'Content-Type': 'application/pdf'
|
||||
}
|
||||
// headers: {
|
||||
// 'Content-Length': fileBuffer.length.toString()
|
||||
// }
|
||||
})
|
||||
|
||||
if (!response.ok) {
|
||||
// 克隆 response 以避免消费 body stream
|
||||
const responseClone = response.clone()
|
||||
|
||||
try {
|
||||
const responseBody = await responseClone.text()
|
||||
const errorInfo = {
|
||||
status: response.status,
|
||||
statusText: response.statusText,
|
||||
url: response.url,
|
||||
type: response.type,
|
||||
redirected: response.redirected,
|
||||
headers: Object.fromEntries(response.headers.entries()),
|
||||
body: responseBody
|
||||
}
|
||||
|
||||
logger.error('Response details:', errorInfo)
|
||||
throw new Error(`Upload failed with status ${response.status}: ${responseBody}`)
|
||||
} catch (parseError) {
|
||||
throw new Error(`Upload failed with status ${response.status}. Could not parse response body.`)
|
||||
}
|
||||
}
|
||||
|
||||
logger.info(`File uploaded successfully to: ${uploadUrl}`)
|
||||
} catch (error: any) {
|
||||
logger.error(`Failed to upload file to URL ${uploadUrl}: ${error}`)
|
||||
throw new Error(error.message)
|
||||
}
|
||||
}
|
||||
|
||||
private async getExtractResults(batchId: string): Promise<ExtractResultResponse> {
|
||||
const endpoint = `${this.provider.apiHost}/api/v4/extract-results/batch/${batchId}`
|
||||
|
||||
try {
|
||||
const response = await fetch(endpoint, {
|
||||
method: 'GET',
|
||||
headers: {
|
||||
'Content-Type': 'application/json',
|
||||
Authorization: `Bearer ${this.provider.apiKey}`,
|
||||
token: this.userId ?? ''
|
||||
}
|
||||
})
|
||||
|
||||
if (response.ok) {
|
||||
const data: ApiResponse<ExtractResultResponse> = await response.json()
|
||||
if (data.code === 0 && data.data) {
|
||||
return data.data
|
||||
} else {
|
||||
throw new Error(`API returned error: ${data.msg || JSON.stringify(data)}`)
|
||||
}
|
||||
} else {
|
||||
throw new Error(`HTTP ${response.status}: ${response.statusText}`)
|
||||
}
|
||||
} catch (error: any) {
|
||||
logger.error(`Failed to get extract results for batch ${batchId}: ${error.message}`)
|
||||
throw new Error(error.message)
|
||||
}
|
||||
}
|
||||
|
||||
private async waitForCompletion(
|
||||
sourceId: string,
|
||||
batchId: string,
|
||||
fileName: string,
|
||||
maxRetries: number = 60,
|
||||
intervalMs: number = 5000
|
||||
): Promise<ExtractFileResult> {
|
||||
let retries = 0
|
||||
|
||||
while (retries < maxRetries) {
|
||||
try {
|
||||
const result = await this.getExtractResults(batchId)
|
||||
|
||||
// 查找对应文件的处理结果
|
||||
const fileResult = result.extract_result.find((item) => item.file_name === fileName)
|
||||
if (!fileResult) {
|
||||
throw new Error(`File ${fileName} not found in batch results`)
|
||||
}
|
||||
|
||||
// 检查处理状态
|
||||
if (fileResult.state === 'done' && fileResult.full_zip_url) {
|
||||
logger.info(`Processing completed for file: ${fileName}`)
|
||||
return fileResult
|
||||
} else if (fileResult.state === 'failed') {
|
||||
throw new Error(`Processing failed for file: ${fileName}, error: ${fileResult.err_msg}`)
|
||||
} else if (fileResult.state === 'running') {
|
||||
// 发送进度更新
|
||||
if (fileResult.extract_progress) {
|
||||
const progress = Math.round(
|
||||
(fileResult.extract_progress.extracted_pages / fileResult.extract_progress.total_pages) * 100
|
||||
)
|
||||
await this.sendPreprocessProgress(sourceId, progress)
|
||||
logger.info(`File ${fileName} processing progress: ${progress}%`)
|
||||
} else {
|
||||
// 如果没有具体进度信息，发送一个通用进度
|
||||
await this.sendPreprocessProgress(sourceId, 50)
|
||||
logger.info(`File ${fileName} is still processing...`)
|
||||
}
|
||||
}
|
||||
} catch (error) {
|
||||
logger.warn(`Failed to check status for batch ${batchId}, retry ${retries + 1}/${maxRetries}`)
|
||||
if (retries === maxRetries - 1) {
|
||||
throw error
|
||||
}
|
||||
}
|
||||
|
||||
retries++
|
||||
await new Promise((resolve) => setTimeout(resolve, intervalMs))
|
||||
}
|
||||
|
||||
throw new Error(`Processing timeout for batch: ${batchId}`)
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,189 @@
|
||||
import fs from 'node:fs'
|
||||
|
||||
import { loggerService } from '@logger'
|
||||
import { MistralClientManager } from '@main/services/MistralClientManager'
|
||||
import { MistralService } from '@main/services/remotefile/MistralService'
|
||||
import { Mistral } from '@mistralai/mistralai'
|
||||
import { DocumentURLChunk } from '@mistralai/mistralai/models/components/documenturlchunk'
|
||||
import { ImageURLChunk } from '@mistralai/mistralai/models/components/imageurlchunk'
|
||||
import { OCRResponse } from '@mistralai/mistralai/models/components/ocrresponse'
|
||||
import { FileMetadata, FileTypes, PreprocessProvider, Provider } from '@types'
|
||||
import path from 'path'
|
||||
|
||||
import BasePreprocessProvider from './BasePreprocessProvider'
|
||||
|
||||
type PreuploadResponse = DocumentURLChunk | ImageURLChunk
|
||||
|
||||
const logger = loggerService.withContext('MistralPreprocessProvider')
|
||||
|
||||
export default class MistralPreprocessProvider extends BasePreprocessProvider {
|
||||
private sdk: Mistral
|
||||
private fileService: MistralService
|
||||
|
||||
constructor(provider: PreprocessProvider) {
|
||||
super(provider)
|
||||
const clientManager = MistralClientManager.getInstance()
|
||||
const aiProvider: Provider = {
|
||||
id: provider.id,
|
||||
type: 'mistral',
|
||||
name: provider.name,
|
||||
apiKey: provider.apiKey!,
|
||||
apiHost: provider.apiHost!,
|
||||
models: []
|
||||
}
|
||||
clientManager.initializeClient(aiProvider)
|
||||
this.sdk = clientManager.getClient()
|
||||
this.fileService = new MistralService(aiProvider)
|
||||
}
|
||||
|
||||
private async preupload(file: FileMetadata): Promise<PreuploadResponse> {
|
||||
let document: PreuploadResponse
|
||||
logger.info(`preprocess preupload started for local file: ${file.path}`)
|
||||
|
||||
if (file.ext.toLowerCase() === '.pdf') {
|
||||
const uploadResponse = await this.fileService.uploadFile(file)
|
||||
|
||||
if (uploadResponse.status === 'failed') {
|
||||
logger.error('File upload failed:', uploadResponse)
|
||||
throw new Error('Failed to upload file: ' + uploadResponse.displayName)
|
||||
}
|
||||
await this.sendPreprocessProgress(file.id, 15)
|
||||
const fileUrl = await this.sdk.files.getSignedUrl({
|
||||
fileId: uploadResponse.fileId
|
||||
})
|
||||
logger.info('Got signed URL:', fileUrl)
|
||||
await this.sendPreprocessProgress(file.id, 20)
|
||||
document = {
|
||||
type: 'document_url',
|
||||
documentUrl: fileUrl.url
|
||||
}
|
||||
} else {
|
||||
const base64Image = Buffer.from(fs.readFileSync(file.path)).toString('base64')
|
||||
document = {
|
||||
type: 'image_url',
|
||||
imageUrl: `data:image/png;base64,${base64Image}`
|
||||
}
|
||||
}
|
||||
|
||||
if (!document) {
|
||||
throw new Error('Unsupported file type')
|
||||
}
|
||||
return document
|
||||
}
|
||||
|
||||
public async parseFile(sourceId: string, file: FileMetadata): Promise<{ processedFile: FileMetadata }> {
|
||||
try {
|
||||
const document = await this.preupload(file)
|
||||
const result = await this.sdk.ocr.process({
|
||||
model: this.provider.model!,
|
||||
document: document,
|
||||
includeImageBase64: true
|
||||
})
|
||||
if (result) {
|
||||
await this.sendPreprocessProgress(sourceId, 100)
|
||||
const processedFile = this.convertFile(result, file)
|
||||
return {
|
||||
processedFile
|
||||
}
|
||||
} else {
|
||||
throw new Error('preprocess processing failed: OCR response is empty')
|
||||
}
|
||||
} catch (error) {
|
||||
throw new Error('preprocess processing failed: ' + error)
|
||||
}
|
||||
}
|
||||
|
||||
private convertFile(result: OCRResponse, file: FileMetadata): FileMetadata {
|
||||
// 使用统一的存储路径：Data/Files/{file.id}/
|
||||
const conversionId = file.id
|
||||
const outputPath = path.join(this.storageDir, file.id)
|
||||
// const outputPath = this.storageDir
|
||||
const outputFileName = path.basename(file.path, path.extname(file.path))
|
||||
fs.mkdirSync(outputPath, { recursive: true })
|
||||
|
||||
const markdownParts: string[] = []
|
||||
let counter = 0
|
||||
|
||||
// Process each page
|
||||
result.pages.forEach((page) => {
|
||||
let pageMarkdown = page.markdown
|
||||
|
||||
// Process images from this page
|
||||
page.images.forEach((image) => {
|
||||
if (image.imageBase64) {
|
||||
let imageFormat = 'jpeg' // default format
|
||||
let imageBase64Data = image.imageBase64
|
||||
|
||||
// Check for data URL prefix more efficiently
|
||||
const prefixEnd = image.imageBase64.indexOf(';base64,')
|
||||
if (prefixEnd > 0) {
|
||||
const prefix = image.imageBase64.substring(0, prefixEnd)
|
||||
const formatIndex = prefix.indexOf('image/')
|
||||
if (formatIndex >= 0) {
|
||||
imageFormat = prefix.substring(formatIndex + 6)
|
||||
}
|
||||
imageBase64Data = image.imageBase64.substring(prefixEnd + 8)
|
||||
}
|
||||
|
||||
const imageFileName = `img-${counter}.${imageFormat}`
|
||||
const imagePath = path.join(outputPath, imageFileName)
|
||||
|
||||
// Save image file
|
||||
try {
|
||||
fs.writeFileSync(imagePath, Buffer.from(imageBase64Data, 'base64'))
|
||||
|
||||
// Update image reference in markdown
|
||||
// Use relative path for better portability
|
||||
const relativeImagePath = `./${imageFileName}`
|
||||
|
||||
// Find the start and end of the image markdown
|
||||
const imgStart = pageMarkdown.indexOf(image.imageBase64)
|
||||
if (imgStart >= 0) {
|
||||
// Find the markdown image syntax around this base64
|
||||
const mdStart = pageMarkdown.lastIndexOf('` +
|
||||
pageMarkdown.substring(mdEnd + 1)
|
||||
}
|
||||
}
|
||||
|
||||
counter++
|
||||
} catch (error) {
|
||||
logger.error(`Failed to save image ${imageFileName}:`, error)
|
||||
}
|
||||
}
|
||||
})
|
||||
|
||||
markdownParts.push(pageMarkdown)
|
||||
})
|
||||
|
||||
// Combine all markdown content with double newlines for readability
|
||||
const combinedMarkdown = markdownParts.join('\n\n')
|
||||
|
||||
// Write the markdown content to a file
|
||||
const mdFileName = `${outputFileName}.md`
|
||||
const mdFilePath = path.join(outputPath, mdFileName)
|
||||
fs.writeFileSync(mdFilePath, combinedMarkdown)
|
||||
|
||||
return {
|
||||
id: conversionId,
|
||||
name: file.name.replace(/\.[^/.]+$/, '.md'),
|
||||
origin_name: file.origin_name,
|
||||
path: mdFilePath,
|
||||
created_at: new Date().toISOString(),
|
||||
type: FileTypes.DOCUMENT,
|
||||
ext: '.md',
|
||||
size: fs.statSync(mdFilePath).size,
|
||||
count: 1
|
||||
} as FileMetadata
|
||||
}
|
||||
|
||||
public checkQuota(): Promise<number> {
|
||||
throw new Error('Method not implemented.')
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,30 @@
|
||||
import { FileMetadata, PreprocessProvider as Provider } from '@types'
|
||||
|
||||
import BasePreprocessProvider from './BasePreprocessProvider'
|
||||
import PreprocessProviderFactory from './PreprocessProviderFactory'
|
||||
|
||||
export default class PreprocessProvider {
|
||||
private sdk: BasePreprocessProvider
|
||||
constructor(provider: Provider, userId?: string) {
|
||||
this.sdk = PreprocessProviderFactory.create(provider, userId)
|
||||
}
|
||||
public async parseFile(
|
||||
sourceId: string,
|
||||
file: FileMetadata
|
||||
): Promise<{ processedFile: FileMetadata; quota?: number }> {
|
||||
return this.sdk.parseFile(sourceId, file)
|
||||
}
|
||||
|
||||
public async checkQuota(): Promise<number> {
|
||||
return this.sdk.checkQuota()
|
||||
}
|
||||
|
||||
/**
|
||||
* 检查文件是否已经被预处理过
|
||||
* @param file 文件信息
|
||||
* @returns 如果已处理返回处理后的文件信息，否则返回null
|
||||
*/
|
||||
public async checkIfAlreadyProcessed(file: FileMetadata): Promise<FileMetadata | null> {
|
||||
return this.sdk.checkIfAlreadyProcessed(file)
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,21 @@
|
||||
import { PreprocessProvider } from '@types'
|
||||
|
||||
import BasePreprocessProvider from './BasePreprocessProvider'
|
||||
import DefaultPreprocessProvider from './DefaultPreprocessProvider'
|
||||
import Doc2xPreprocessProvider from './Doc2xPreprocessProvider'
|
||||
import MineruPreprocessProvider from './MineruPreprocessProvider'
|
||||
import MistralPreprocessProvider from './MistralPreprocessProvider'
|
||||
export default class PreprocessProviderFactory {
|
||||
static create(provider: PreprocessProvider, userId?: string): BasePreprocessProvider {
|
||||
switch (provider.id) {
|
||||
case 'doc2x':
|
||||
return new Doc2xPreprocessProvider(provider)
|
||||
case 'mistral':
|
||||
return new MistralPreprocessProvider(provider)
|
||||
case 'mineru':
|
||||
return new MineruPreprocessProvider(provider, userId)
|
||||
default:
|
||||
return new DefaultPreprocessProvider(provider)
|
||||
}
|
||||
}
|
||||
}
|
||||
Reference in New Issue
Block a user