feat(OCR): implement Mistral OCR provider and enhance file handling

- Introduced MistralOcrProvider for processing OCR tasks with Mistral API. - Updated BaseOcrProvider to return LocalFileSource in parseFile method. - Enhanced file upload and retrieval logic in MistralService. - Added support for local file source handling in various services. - Updated preload API and type definitions to accommodate new OCR provider functionalities.
2025-03-24 12:38:07 +08:00
parent 64838cb3fb
commit 5fe4bf3f19
19 changed files with 342 additions and 73 deletions
@@ -1,7 +1,7 @@
 import fs from 'node:fs'

 import { windowService } from '@main/services/WindowService'
-import { FileSource, OcrProvider } from '@types'
+import { FileSource, LocalFileSource, OcrProvider } from '@types'
 import Logger from 'electron-log'
 import pdfParse from 'pdf-parse'
 export default abstract class BaseOcrProvider {
@@ -12,7 +12,7 @@ export default abstract class BaseOcrProvider {
    }
    this.provider = provider
  }
-  abstract parseFile(sourceId: string, file: FileSource): Promise<{ processedFile: FileSource }>
+  abstract parseFile(sourceId: string, file: FileSource): Promise<{ processedFile: LocalFileSource }>
  /**
   * 辅助方法：延迟执行
   */
@@ -0,0 +1,203 @@
+import fs from 'node:fs'
+
+import { MistralService } from '@main/services/file/MistralService'
+import { MistralClientManager } from '@main/services/MistralClientManager'
+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 { FileSource, FileTypes, isLocalFile, LocalFileSource, OcrProvider } from '@types'
+import Logger from 'electron-log'
+import path from 'path'
+import { v4 as uuidv4 } from 'uuid'
+
+import BaseOcrProvider from './BaseOcrProvider'
+
+type PreuploadResponse = DocumentURLChunk | ImageURLChunk
+
+export default class MistralOcrProvider extends BaseOcrProvider {
+  private sdk: Mistral
+  private fileService: MistralService
+
+  constructor(provider: OcrProvider) {
+    super(provider)
+    const clientManager = MistralClientManager.getInstance()
+    clientManager.initializeClient(provider.apiKey!)
+    this.sdk = clientManager.getClient()
+    this.fileService = new MistralService(provider.apiKey!)
+  }
+
+  private async preupload(file: FileSource): Promise<PreuploadResponse> {
+    let document: PreuploadResponse
+    if (isLocalFile(file)) {
+      Logger.info(`OCR preupload started for local file: ${file.path}`)
+
+      const pdfInfo = await this.getPdfInfo(file.path)
+      if (pdfInfo.pageCount >= 1000) {
+        throw new Error(`PDF page count (${pdfInfo.pageCount}) exceeds the limit of 1000 pages`)
+      }
+      if (pdfInfo.fileSize >= 512 * 1024 * 1024) {
+        const fileSizeMB = Math.round(pdfInfo.fileSize / (1024 * 1024))
+        throw new Error(`PDF file size (${fileSizeMB}MB) exceeds the limit of 300MB`)
+      }
+
+      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)
+        }
+
+        const fileUrl = await this.sdk.files.getSignedUrl({
+          fileId: uploadResponse.fileId
+        })
+        Logger.info('Got signed URL:', fileUrl)
+
+        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}`
+        }
+      }
+    } else {
+      if (file.ext.toLowerCase() === '.pdf') {
+        document = {
+          type: 'document_url',
+          documentUrl: file.url
+        }
+      } else {
+        document = {
+          type: 'image_url',
+          imageUrl: file.url
+        }
+      }
+    }
+
+    if (!document) {
+      throw new Error('Unsupported file type')
+    }
+    return document
+  }
+
+  public async parseFile(sourceId: string, file: FileSource): Promise<{ processedFile: LocalFileSource }> {
+    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.sendOcrProgress(sourceId, 100)
+        const processedFile = this.convertFile(result, file)
+        return {
+          processedFile
+        }
+      } else {
+        throw new Error('OCR processing failed: OCR response is empty')
+      }
+    } catch (error) {
+      throw new Error('OCR processing failed: ' + error)
+    }
+  }
+
+  private convertFile(result: OCRResponse, file: FileSource): LocalFileSource {
+    // Create a unique directory for this conversion to store images
+    const conversionId = uuidv4()
+    let outputPath = ''
+    let outputFileName = ''
+    if (isLocalFile(file)) {
+      outputPath = path.join(path.dirname(file.path), conversionId)
+      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('![', imgStart)
+              const mdEnd = pageMarkdown.indexOf(')', imgStart)
+
+              if (mdStart >= 0 && mdEnd >= 0) {
+                // Replace just this specific image reference
+                pageMarkdown =
+                  pageMarkdown.substring(0, mdStart) +
+                  `![Image ${counter}](${relativeImagePath})` +
+                  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: mdFileName,
+      origin_name: mdFileName,
+      path: mdFilePath,
+      created_at: new Date().toISOString(),
+      type: FileTypes.DOCUMENT,
+      ext: '.md',
+      size: fs.statSync(mdFilePath).size,
+      count: result.pages.length,
+      source: 'local'
+    }
+  }
+}
@@ -1,12 +1,12 @@
-import { FileType, KnowledgeBaseParams } from '@types'
+import { FileType, OcrProvider as Provider } from '@types'

 import BaseOcrProvider from './BaseOcrProvider'
 import OcrProviderFactory from './OcrProviderFactory'

 export default class OcrProvider {
  private sdk: BaseOcrProvider
-  constructor(base: KnowledgeBaseParams) {
-    this.sdk = OcrProviderFactory.create(base)
+  constructor(provider: Provider) {
+    this.sdk = OcrProviderFactory.create(provider)
  }
  public async parseFile(sourceId: string, file: FileType): Promise<{ processedFile: FileType }> {
    return this.sdk.parseFile(sourceId, file)
@@ -73,7 +73,8 @@ class FileStorage {
            size: storedStats.size,
            ext,
            type: getFileType(ext),
-            count: 2
+            count: 2,
+            source: 'local'
          }
        }
      }
@@ -112,7 +113,8 @@ class FileStorage {
        size: stats.size,
        ext: ext,
        type: fileType,
-        count: 1
+        count: 1,
+        source: 'local' as const
      }
    })

@@ -177,7 +179,8 @@ class FileStorage {
      size: stats.size,
      ext: ext,
      type: fileType,
-      count: 1
+      count: 1,
+      source: 'local'
    }

    return fileMetadata
@@ -201,7 +204,8 @@ class FileStorage {
      size: stats.size,
      ext: ext,
      type: fileType,
-      count: 1
+      count: 1,
+      source: 'local'
    }

    return fileInfo
@@ -267,6 +271,16 @@ class FileStorage {
    }
  }

+  public base64File = async (
+    _: Electron.IpcMainInvokeEvent,
+    filePath: string
+  ): Promise<{ data: Buffer; mime: string }> => {
+    return {
+      data: await fs.promises.readFile(filePath),
+      mime: 'application/pdf'
+    }
+  }
+
  public binaryFile = async (_: Electron.IpcMainInvokeEvent, id: string): Promise<{ data: Buffer; mime: string }> => {
    const filePath = path.join(this.storageDir, id)
    const data = await fs.promises.readFile(filePath)
@@ -424,7 +438,8 @@ class FileStorage {
        size: stats.size,
        ext: ext,
        type: fileType,
-        count: 1
+        count: 1,
+        source: 'local'
      }

      return fileMetadata
@@ -177,9 +177,10 @@ class KnowledgeService {
          task: async () => {
            // 添加OCR预处理逻辑
            let fileToProcess: FileType = file
-            if (base.preprocessing && file.ext.toLowerCase() === '.pdf') {
+            if (base.preprocessing && base.ocrProvider && file.ext.toLowerCase() === '.pdf') {
              try {
-                const ocrProvider = new OcrProvider(base)
+                file.source = 'local'
+                const ocrProvider = new OcrProvider(base.ocrProvider)
                Logger.info(`Starting OCR processing for file: ${file.path}`)

                const { processedFile } = await ocrProvider.parseFile(item.id, file)
@@ -508,6 +509,10 @@ class KnowledgeService {
  ): Promise<ExtractChunkData[]> => {
    return await new Reranker(base).rerank(search, results)
  }
+
+  public getStorageDir = (): string => {
+    return this.storageDir
+  }
 }

 export default new KnowledgeService()
@@ -6,6 +6,7 @@ export class FileServiceManager {
  private static instance: FileServiceManager
  private services: Map<string, BaseFileService> = new Map()

+  // eslint-disable-next-line @typescript-eslint/no-empty-function
  private constructor() {}

  static getInstance(): FileServiceManager {
@@ -60,41 +60,26 @@ export class GeminiService extends BaseFileService {
      return cachedResponse
    }

-    const response = await this.fileManager.getFile(fileId)
+    const response = await this.fileManager.listFiles()

-    // 根据文件状态设置响应状态
-    let status: 'success' | 'processing' | 'failed' | 'unknown'
-    switch (response.state) {
-      case FileState.ACTIVE:
-        status = 'success'
-        break
-      case FileState.PROCESSING:
-        status = 'processing'
-        break
-      case FileState.FAILED:
-        status = 'failed'
-        break
-      default:
-        status = 'unknown'
+    if (response.files) {
+      const file = response.files.filter((file) => file.state === FileState.ACTIVE).find((file) => file.name === fileId)
+      if (file) {
+        return {
+          fileId: fileId,
+          displayName: file.displayName || '',
+          status: 'success',
+          originalFile: file
+        }
+      }
    }

-    const fileResponse: FileUploadResponse = {
-      fileId,
-      displayName: response.displayName || '',
-      status,
-      originalFile: response
+    return {
+      fileId: fileId,
+      displayName: '',
+      status: 'failed',
+      originalFile: undefined
    }
-
-    // 只缓存成功的文件
-    if (status === 'success') {
-      CacheService.set(
-        `${GeminiService.FILE_LIST_CACHE_KEY}_${fileId}`,
-        fileResponse,
-        GeminiService.FILE_CACHE_DURATION
-      )
-    }
-
-    return fileResponse
  }

  async listFiles(): Promise<FileListResponse> {
@@ -1,11 +1,14 @@
+import fs from 'node:fs/promises'
+
+import { Mistral } from '@mistralai/mistralai'
 import { FileListResponse, FileUploadResponse, LocalFileSource } from '@types'
-import { fileFrom } from 'fetch-blob/from.js'
+import Logger from 'electron-log'

 import { MistralClientManager } from '../MistralClientManager'
 import { BaseFileService } from './BaseFileService'

 export class MistralService extends BaseFileService {
-  private readonly client
+  private readonly client: Mistral

  constructor(apiKey: string) {
    super(apiKey)
@@ -16,9 +19,12 @@ export class MistralService extends BaseFileService {

  async uploadFile(file: LocalFileSource): Promise<FileUploadResponse> {
    try {
-      const blob = await fileFrom(file.path)
+      const fileBuffer = await fs.readFile(file.path)
      const response = await this.client.files.upload({
-        file: blob,
+        file: {
+          fileName: file.path,
+          content: new Uint8Array(fileBuffer)
+        },
        purpose: 'ocr'
      })

@@ -28,7 +34,7 @@ export class MistralService extends BaseFileService {
        status: 'success'
      }
    } catch (error) {
-      console.error('Error uploading file:', error)
+      Logger.error('Error uploading file:', error)
      return {
        fileId: '',
        displayName: file.origin_name,
@@ -78,7 +84,12 @@ export class MistralService extends BaseFileService {
      }
    } catch (error) {
      console.error('Error retrieving file:', error)
-      throw error
+      return {
+        fileId: fileId,
+        displayName: '',
+        status: 'failed',
+        originalFile: undefined
+      }
    }
  }
 }
@@ -57,7 +57,8 @@ export function getAllFiles(dirPath: string, arrayOfFiles: FileType[] = []): Fil
        count: 1,
        origin_name: name,
        type: fileType,
-        created_at: new Date().toISOString()
+        created_at: new Date().toISOString(),
+        source: 'local'
      }

      arrayOfFiles.push(fileItem)
@@ -1,8 +1,16 @@
 import { ElectronAPI } from '@electron-toolkit/preload'
-import type { FileMetadataResponse, ListFilesResponse, UploadFileResponse } from '@google/generative-ai/server'
 import { ExtractChunkData } from '@llm-tools/embedjs-interfaces'
 import type { MCPServer, MCPTool } from '@renderer/types'
-import { AppInfo, FileType, KnowledgeBaseParams, KnowledgeItem, LanguageVarious, WebDavConfig } from '@renderer/types'
+import {
+  AppInfo,
+  FileListResponse,
+  FileType,
+  FileUploadResponse,
+  KnowledgeBaseParams,
+  KnowledgeItem,
+  LanguageVarious,
+  WebDavConfig
+} from '@renderer/types'
 import type { LoaderReturn } from '@shared/config/types'
 import type { OpenDialogOptions } from 'electron'
 import type { UpdateInfo } from 'electron-updater'
@@ -66,6 +74,7 @@ declare global {
        ) => Promise<string | null>
        saveImage: (name: string, data: string) => void
        base64Image: (fileId: string) => Promise<{ mime: string; base64: string; data: string }>
+        base64File: (filePath: string) => Promise<{ mime: string; data: Buffer }>
        download: (url: string) => Promise<FileType | null>
        copy: (fileId: string, destPath: string) => Promise<void>
        binaryFile: (fileId: string) => Promise<{ data: Buffer; mime: string }>
@@ -118,9 +127,9 @@ declare global {
        resetMinimumSize: () => Promise<void>
      }
      fileService: {
-        upload: (type: string, apiKey: string, file: FileType) => Promise<UploadFileResponse>
-        retrieve: (type: string, apiKey: string, fileId: string) => Promise<FileMetadataResponse | undefined>
-        list: (type: string, apiKey: string) => Promise<ListFilesResponse>
+        upload: (type: string, apiKey: string, file: FileType) => Promise<FileUploadResponse>
+        retrieve: (type: string, apiKey: string, fileId: string) => Promise<FileUploadResponse>
+        list: (type: string, apiKey: string) => Promise<FileListResponse>
        delete: (type: string, apiKey: string, fileId: string) => Promise<void>
      }
      selectionMenu: {
@@ -53,6 +53,7 @@ const api = {
    selectFolder: () => ipcRenderer.invoke('file:selectFolder'),
    saveImage: (name: string, data: string) => ipcRenderer.invoke('file:saveImage', name, data),
    base64Image: (fileId: string) => ipcRenderer.invoke('file:base64Image', fileId),
+    base64File: (filePath: string) => ipcRenderer.invoke('file:base64File', filePath),
    download: (url: string) => ipcRenderer.invoke('file:download', url),
    copy: (fileId: string, destPath: string) => ipcRenderer.invoke('file:copy', fileId, destPath),
    binaryFile: (fileId: string) => ipcRenderer.invoke('file:binaryFile', fileId)
@@ -10,6 +10,7 @@ const OcrIcon: FC<React.SVGProps<SVGSVGElement>> = (props) => (
    width="16"
    height="16"
    strokeWidth="1"
+    fill="var(--color-text-2)"
    {...props}>
    <path
      d="M640.189 960.294l-0.131-64 254.651-0.522 1.285-255.968 64 0.322-1.605 319.515-318.2 0.653z m-256.239-0.075l-318.197-1.237 0.251-319.042 64 0.051-0.201 255.239 254.396 0.989-0.249 64zM66.004 383.837L65.53 64.399l318.728 1.829-0.367 63.999-254.265-1.459 0.378 254.975-64 0.094zM897.495 383l-0.661-252.989-254.683 0.217-0.055-64 318.569-0.271 0.829 316.876-63.999 0.167zM404.866 510.522c0 24.75-5.328 46.895-15.984 66.43-10.656 19.537-25.553 34.719-44.688 45.547-19.137 10.828-40.562 16.242-64.281 16.242-23.146 0-44.201-5.242-63.164-15.727-18.965-10.484-33.717-25.207-44.258-44.172-10.543-18.963-15.812-40.418-15.812-64.367 0-25.093 5.328-47.665 15.984-67.718 10.656-20.05 25.609-35.549 44.859-46.492 19.25-10.941 41.135-16.414 65.656-16.414 23.604 0 44.714 5.242 63.336 15.727 18.619 10.484 33 25.438 43.141 44.859s15.211 41.451 15.211 66.085z m-78.719 2.062c0-20.281-3.896-36.265-11.688-47.953-7.793-11.688-18.45-17.531-31.969-17.531-14.781 0-26.297 5.615-34.547 16.844-8.25 11.231-12.375 27.1-12.375 47.609 0 20.053 4.096 35.693 12.289 46.922 8.191 11.23 19.336 16.844 33.43 16.844 8.594 0 16.328-2.52 23.203-7.562 6.875-5.041 12.203-12.26 15.984-21.656 3.782-9.396 5.673-20.566 5.673-33.517zM623.147 627.741c-20.396 7.332-43.943 11-70.641 11-26.24 0-48.872-5.012-67.891-15.039-19.021-10.025-33.545-24.291-43.57-42.797-10.028-18.504-15.039-39.846-15.039-64.023 0-26.009 5.613-49.156 16.844-69.437 11.229-20.281 27.097-35.949 47.609-47.008 20.51-11.057 44.114-16.586 70.813-16.586 21.312 0 41.938 2.578 61.875 7.734v68.578c-6.875-4.125-15.068-7.332-24.578-9.625a122.892 122.892 0 0 0-28.875-3.438c-20.168 0-36.066 5.787-47.695 17.359-11.631 11.575-17.446 27.271-17.446 47.093 0 19.709 5.814 35.264 17.446 46.664 11.629 11.402 27.184 17.102 46.664 17.102 17.988 0 36.15-4.582 54.484-13.75v66.173zM789.351 634.444l-18.391-53.109c-3.553-10.426-8.164-18.619-13.836-24.578-5.672-5.957-11.832-8.938-18.477-8.938h-2.922v86.625h-74.25V387.975h98.656c34.488 0 59.955 5.645 76.398 16.93 16.441 11.287 24.664 28.217 24.664 50.789 0 16.959-4.785 31.168-14.352 42.625-9.568 11.458-23.805 19.765-42.711 24.921v0.688c10.426 3.209 19.105 8.422 26.039 15.641 6.932 7.219 13.148 17.934 18.648 32.141l24.234 62.734h-83.7z m-7.047-168.265c0-8.25-2.521-14.781-7.562-19.594-5.043-4.812-12.949-7.219-23.719-7.219h-15.297v56.375h13.406c9.969 0 17.988-2.807 24.062-8.422 6.073-5.613 9.11-12.659 9.11-21.14z"
@@ -8,6 +8,7 @@ const ToolIcon: FC<React.SVGProps<SVGSVGElement>> = (props) => (
    width="16"
    height="16"
    className="icon"
+    fill="var(--color-text-2)"
    {...props}>
    <path d="M732.16 450.048h-12.288V377.856c0-37.888-31.744-69.12-70.656-69.12h-73.216c1.024-7.168 1.024-14.848 1.024-22.016 0-75.776-62.976-137.728-140.288-137.728S296.448 210.944 296.448 286.72c0 7.168 0 14.848 1.024 22.016H224.256c-38.912 0-70.656 31.232-70.656 69.12v102.4c0 9.216 3.584 15.872 10.752 21.504 7.168 5.632 14.848 7.168 23.552 5.632 7.68-1.536 16.384-3.072 22.528-3.072 46.08 0 84.992 38.4 84.992 83.456 0 46.08-37.888 83.456-84.992 83.456-6.656 0-14.848-1.536-22.528-3.072-7.68-2.048-16.896 0.512-23.552 5.632-7.168 5.12-10.752 12.288-10.752 21.504v102.4c0 37.888 31.744 69.12 70.656 69.12H650.24c19.456 0 38.4-8.192 51.712-22.016 12.288-13.312 18.432-29.696 17.92-47.104v-72.192h12.288c77.312 0 140.288-61.952 140.288-137.728 0-75.776-62.976-137.728-140.288-137.728z m-66.048 239.616v105.984c0 8.192-7.168 15.36-15.872 15.36H224.768c-8.704 0-15.872-6.656-15.872-15.36v-70.656c0.512 0 1.536 0 2.048-0.512 77.312-0.512 139.776-61.952 139.776-137.728 0-37.376-14.848-72.704-41.984-98.816-27.136-26.112-62.976-39.936-100.352-38.912V378.368c0-8.192 7.168-15.36 15.872-15.36h110.08c8.704 0 17.92-4.608 23.552-12.288 4.608-6.656 5.632-15.36 2.048-25.088-5.12-11.264-8.192-24.064-8.192-36.352 0-46.08 37.888-83.456 84.992-83.456 46.08 0 84.992 38.4 84.992 83.456 0 10.752-1.024 23.04-7.68 35.84-4.608 7.68-3.584 18.432 1.536 25.6 5.12 6.144 13.824 12.288 23.552 12.288H650.24c8.704 0 15.872 6.656 15.872 15.36v107.52c0 9.216 3.584 15.872 10.752 21.504 7.68 5.632 16.896 6.656 25.6 3.072 8.192-4.096 17.408-6.144 28.672-6.144 46.08 0 84.992 38.4 84.992 83.456 0 46.08-37.888 83.456-84.992 83.456-8.192 0-17.408-3.072-26.112-5.632l-3.072-1.024c-10.24-2.56-18.432-1.536-25.088 3.584-6.656 3.072-10.752 11.264-10.752 21.504z" />
  </svg>
@@ -1,8 +1,12 @@
 import Doc2xLogo from '@renderer/assets/images/ocr/doc2x.svg'
+import MistralLogo from '@renderer/assets/images/providers/mistral.png'
+
 export function getOcrProviderLogo(providerId: string) {
  switch (providerId) {
    case 'doc2x':
      return Doc2xLogo
+    case 'mistral':
+      return MistralLogo
    default:
      return undefined
  }
@@ -14,5 +18,11 @@ export const OCR_PROVIDER_CONFIG = {
      official: 'https://doc2x.noedgeai.com',
      apiKey: 'https://open.noedgeai.com/apiKeys'
    }
+  },
+  mistral: {
+    websites: {
+      official: 'https://mistral.ai',
+      apiKey: 'https://mistral.ai/api-keys'
+    }
  }
 }
@@ -30,7 +30,17 @@ import {
  filterEmptyMessages,
  filterUserRoleStartMessages
 } from '@renderer/services/MessagesService'
-import { Assistant, FileType, FileTypes, MCPToolResponse, Message, Model, Provider, Suggestion } from '@renderer/types'
+import {
+  Assistant,
+  FileType,
+  FileTypes,
+  FileUploadResponse,
+  MCPToolResponse,
+  Message,
+  Model,
+  Provider,
+  Suggestion
+} from '@renderer/types'
 import { removeSpecialCharactersForTopicName } from '@renderer/utils'
 import { fileToBase64 } from '@renderer/utils/file'
 import {
@@ -84,25 +94,28 @@ export default class GeminiProvider extends BaseProvider {
      } as InlineDataPart
    }

-    // Retrieve file from Gemini uploaded files
-    const fileMetadata = await window.api.fileService.retrieve(this.provider.type, this.apiKey, file.id)
-
-    if (fileMetadata) {
+    // 尝试检索文件
+    const response: FileUploadResponse = await window.api.fileService.retrieve(this.provider.type, this.apiKey, file.id)
+    if (response && response.status === 'success') {
      return {
        fileData: {
-          fileUri: fileMetadata.uri,
-          mimeType: fileMetadata.mimeType
+          fileUri: response.originalFile.uri,
+          mimeType: response.originalFile.mimeType
        }
      } as FileDataPart
+    } else {
+      console.log('file not found', response)
    }
-
-    // If file is not found, upload it to Gemini
-    const uploadResult = await window.api.fileService.upload(this.provider.type, this.apiKey, file)
-
+    // 如果文件不存在，上传新文件
+    const uploadResponse: FileUploadResponse = await window.api.fileService.upload(
+      this.provider.type,
+      this.apiKey,
+      file
+    )
    return {
      fileData: {
-        fileUri: uploadResult.file.uri,
-        mimeType: uploadResult.file.mimeType
+        fileUri: uploadResponse.originalFile.file.uri,
+        mimeType: uploadResponse.originalFile.file.mimeType
      }
    } as FileDataPart
  }
@@ -804,6 +804,13 @@ const migrateConfig = {
          name: 'Doc2x',
          apiKey: '',
          apiHost: 'https://v2.doc2x.noedgeai.com'
+        },
+        {
+          id: 'mistral',
+          name: 'Mistral',
+          model: 'mistral-ocr-latest',
+          apiKey: '',
+          apiHost: 'https://api.mistral.ai'
        }
      ]
    }
@@ -13,6 +13,13 @@ const initialState: OcrState = {
      name: 'Doc2x',
      apiKey: '',
      apiHost: 'https://v2.doc2x.noedgeai.com'
+    },
+    {
+      id: 'mistral',
+      name: 'Mistral',
+      model: 'mistral-ocr-latest',
+      apiKey: '',
+      apiHost: 'https://api.mistral.ai'
    }
  ],
  defaultProvider: ''
@@ -340,6 +340,7 @@ export type KnowledgeBaseParams = {
  rerankModelProvider?: string
  topN?: number
  preprocessing?: boolean
+  ocrProvider?: OcrProvider
 }

 export interface OcrProvider {
@@ -1,9 +1,7 @@
-import fs from 'fs'
-
 export const fileToBase64 = async (filePath: string) => {
-  const buffer = await fs.promises.readFile(filePath)
+  const result = await window.api.file.base64Image(filePath)
  return {
-    data: buffer.toString('base64'),
-    mimeType: 'application/pdf'
+    data: result.base64,
+    mimeType: result.mime
  }
 }