feat: add gemini files support

package.json

@@ -50,6 +50,7 @@
     "@electron-toolkit/preload": "^3.0.0",
     "@electron-toolkit/utils": "^3.0.0",
     "@electron/notarize": "^2.5.0",
+    "@google/generative-ai": "^0.21.0",
     "@llm-tools/embedjs": "patch:@llm-tools/embedjs@npm%3A0.1.25#~/.yarn/patches/@llm-tools-embedjs-npm-0.1.25-ec5645cf36.patch",
     "@llm-tools/embedjs-libsql": "patch:@llm-tools/embedjs-libsql@npm%3A0.1.25#~/.yarn/patches/@llm-tools-embedjs-libsql-npm-0.1.25-fad000d74c.patch",
     "@llm-tools/embedjs-loader-csv": "^0.1.25",
@@ -81,7 +82,6 @@
     "@electron-toolkit/eslint-config-prettier": "^2.0.0",
     "@electron-toolkit/eslint-config-ts": "^1.0.1",
     "@electron-toolkit/tsconfig": "^1.0.1",
-    "@google/generative-ai": "^0.21.0",
     "@hello-pangea/dnd": "^16.6.0",
     "@kangfenmao/keyv-storage": "^0.1.0",
     "@reduxjs/toolkit": "^2.2.5",

src/main/ipc.ts

@@ -11,6 +11,7 @@ import BackupManager from './services/BackupManager'
 import { configManager } from './services/ConfigManager'
 import { ExportService } from './services/ExportService'
 import FileStorage from './services/FileStorage'
+import { GeminiService } from './services/GeminiService'
 import KnowledgeService from './services/KnowledgeService'
 import { registerShortcuts, unregisterAllShortcuts } from './services/ShortcutService'
 import { windowService } from './services/WindowService'
@@ -167,4 +168,9 @@ export function registerIpc(mainWindow: BrowserWindow, app: Electron.App) {
       mainWindow?.setSize(1080, height)
     }
   })
+
+  // gemini
+  ipcMain.handle('gemini:upload-file', GeminiService.uploadFile)
+  ipcMain.handle('gemini:base64-file', GeminiService.base64File)
+  ipcMain.handle('gemini:retrieve-file', GeminiService.retrieveFile)
 }

src/main/services/CacheService.ts

@@ -0,0 +1,74 @@
+interface CacheItem<T> {
+  data: T
+  timestamp: number
+  duration: number
+}
+
+export class CacheService {
+  private static cache: Map<string, CacheItem<any>> = new Map()
+
+  /**
+   * Set cache
+   * @param key Cache key
+   * @param data Cache data
+   * @param duration Cache duration (in milliseconds)
+   */
+  static set<T>(key: string, data: T, duration: number): void {
+    this.cache.set(key, {
+      data,
+      timestamp: Date.now(),
+      duration
+    })
+  }
+
+  /**
+   * Get cache
+   * @param key Cache key
+   * @returns Returns data if cache exists and not expired, otherwise returns null
+   */
+  static get<T>(key: string): T | null {
+    const item = this.cache.get(key)
+    if (!item) return null
+
+    const now = Date.now()
+    if (now - item.timestamp > item.duration) {
+      this.remove(key)
+      return null
+    }
+
+    return item.data
+  }
+
+  /**
+   * Remove specific cache
+   * @param key Cache key
+   */
+  static remove(key: string): void {
+    this.cache.delete(key)
+  }
+
+  /**
+   * Clear all cache
+   */
+  static clear(): void {
+    this.cache.clear()
+  }
+
+  /**
+   * Check if cache exists and is valid
+   * @param key Cache key
+   * @returns boolean
+   */
+  static has(key: string): boolean {
+    const item = this.cache.get(key)
+    if (!item) return false
+
+    const now = Date.now()
+    if (now - item.timestamp > item.duration) {
+      this.remove(key)
+      return false
+    }
+
+    return true
+  }
+}

src/main/services/GeminiService.ts

@@ -0,0 +1,53 @@
+import { FileMetadataResponse, FileState, GoogleAIFileManager } from '@google/generative-ai/server'
+import { FileType } from '@types'
+import fs from 'fs'
+
+import { CacheService } from './CacheService'
+
+export class GeminiService {
+  private static readonly FILE_LIST_CACHE_KEY = 'gemini_file_list'
+  private static readonly CACHE_DURATION = 3000
+
+  static async uploadFile(_: Electron.IpcMainInvokeEvent, file: FileType, apiKey: string) {
+    const fileManager = new GoogleAIFileManager(apiKey)
+    const uploadResult = await fileManager.uploadFile(file.path, {
+      mimeType: 'application/pdf',
+      displayName: file.origin_name
+    })
+    return uploadResult
+  }
+
+  static async base64File(_: Electron.IpcMainInvokeEvent, file: FileType) {
+    return {
+      data: Buffer.from(fs.readFileSync(file.path)).toString('base64'),
+      mimeType: 'application/pdf'
+    }
+  }
+
+  static async retrieveFile(
+    _: Electron.IpcMainInvokeEvent,
+    file: FileType,
+    apiKey: string
+  ): Promise<FileMetadataResponse | undefined> {
+    const fileManager = new GoogleAIFileManager(apiKey)
+
+    const cachedResponse = CacheService.get<any>(GeminiService.FILE_LIST_CACHE_KEY)
+    if (cachedResponse) {
+      return GeminiService.processResponse(cachedResponse, file)
+    }
+
+    const response = await fileManager.listFiles()
+    CacheService.set(GeminiService.FILE_LIST_CACHE_KEY, response, GeminiService.CACHE_DURATION)
+
+    return GeminiService.processResponse(response, file)
+  }
+
+  private static processResponse(response: any, file: FileType) {
+    if (response.files) {
+      return response.files
+        .filter((file) => file.state === FileState.ACTIVE)
+        .find((i) => i.displayName === file.origin_name && Number(i.sizeBytes) === file.size)
+    }
+    return undefined
+  }
+}

src/preload/index.d.ts

@@ -1,4 +1,6 @@
 import { ElectronAPI } from '@electron-toolkit/preload'
+import type { FileMetadataResponse } from '@google/generative-ai/server'
+import { UploadFileResponse } from '@google/generative-ai/server'
 import { AddLoaderReturn, ExtractChunkData } from '@llm-tools/embedjs-interfaces'
 import { FileType } from '@renderer/types'
 import { WebDavConfig } from '@renderer/types'
@@ -80,6 +82,11 @@ declare global {
         setMinimumSize: (width: number, height: number) => Promise<void>
         resetMinimumSize: () => Promise<void>
       }
+      gemini: {
+        uploadFile: (file: FileType, apiKey: string) => Promise<UploadFileResponse>
+        retrieveFile: (file: FileType, apiKey: string) => Promise<FileMetadataResponse | undefined>
+        base64File: (file: FileType) => Promise<{ data: string; mimeType: string }>
+      }
     }
   }
 }

src/preload/index.ts

@@ -1,5 +1,5 @@
 import { electronAPI } from '@electron-toolkit/preload'
-import { KnowledgeBaseParams, KnowledgeItem, Shortcut, WebDavConfig } from '@types'
+import { FileType, KnowledgeBaseParams, KnowledgeItem, Shortcut, WebDavConfig } from '@types'
 import { contextBridge, ipcRenderer, OpenDialogOptions } from 'electron'
 
 // Custom APIs for renderer
@@ -74,6 +74,11 @@ const api = {
   window: {
     setMinimumSize: (width: number, height: number) => ipcRenderer.invoke('window:set-minimum-size', width, height),
     resetMinimumSize: () => ipcRenderer.invoke('window:reset-minimum-size')
+  },
+  gemini: {
+    uploadFile: (file: FileType, apiKey: string) => ipcRenderer.invoke('gemini:upload-file', file, apiKey),
+    base64File: (file: FileType) => ipcRenderer.invoke('gemini:base64-file', file),
+    retrieveFile: (file: FileType, apiKey: string) => ipcRenderer.invoke('gemini:retrieve-file', file, apiKey)
   }
 }
 

src/renderer/src/components/TopView/index.tsx

@@ -28,10 +28,7 @@ const TopViewContainer: React.FC<Props> = ({ children }) => {
   const elementsRef = useRef<ElementItem[]>([])
   elementsRef.current = elements
 
-  // 消息提示默认为 1s 后关闭，使用方法 window.message 代替 antd message
-  const [messageApi, messageContextHolder] = message.useMessage({
-    duration: 1
-  })
+  const [messageApi, messageContextHolder] = message.useMessage()
   const [modal, modalContextHolder] = Modal.useModal()
 
   useAppInit()

src/renderer/src/i18n/locales/en-us.json

@@ -183,6 +183,7 @@
       "name": "Name",
       "open": "Open",
       "size": "Size",
+      "type": "Type",
       "text": "Text",
       "title": "Files",
       "edit": "Edit",

src/renderer/src/i18n/locales/ja-jp.json

@@ -183,6 +183,7 @@
       "name": "名前",
       "open": "開く",
       "size": "サイズ",
+      "type": "タイプ",
       "text": "テキスト",
       "title": "ファイル",
       "edit": "編集",

src/renderer/src/i18n/locales/ru-ru.json

@@ -183,6 +183,7 @@
       "name": "Имя",
       "open": "Открыть",
       "size": "Размер",
+      "type": "Тип",
       "text": "Текст",
       "title": "Файлы",
       "edit": "Редактировать",

src/renderer/src/i18n/locales/zh-cn.json

@@ -184,6 +184,7 @@
       "name": "文件名",
       "open": "打开",
       "size": "大小",
+      "type": "类型",
       "text": "文本",
       "title": "文件",
       "edit": "编辑",

src/renderer/src/i18n/locales/zh-tw.json

@@ -183,6 +183,7 @@
       "name": "名稱",
       "open": "打開",
       "size": "大小",
+      "type": "類型",
       "text": "文本",
       "title": "檔案",
       "edit": "編輯",

src/renderer/src/pages/files/ContentView.tsx

@@ -0,0 +1,129 @@
+import FileManager from '@renderer/services/FileManager'
+import { FileType, FileTypes } from '@renderer/types'
+import { formatFileSize } from '@renderer/utils'
+import { Col, Image, Row, Spin, Table } from 'antd'
+import React, { memo } from 'react'
+import styled from 'styled-components'
+
+import GeminiFiles from './GeminiFiles'
+
+interface ContentViewProps {
+  id: FileTypes | 'all' | string
+  files?: FileType[]
+  dataSource?: any[]
+  columns: any[]
+}
+
+const ContentView: React.FC<ContentViewProps> = ({ id, files, dataSource, columns }) => {
+  if (id === FileTypes.IMAGE && files?.length && files?.length > 0) {
+    return (
+      <Image.PreviewGroup>
+        <Row gutter={[16, 16]}>
+          {files?.map((file) => (
+            <Col key={file.id} xs={24} sm={12} md={8} lg={4} xl={3}>
+              <ImageWrapper>
+                <LoadingWrapper>
+                  <Spin />
+                </LoadingWrapper>
+                <Image
+                  src={FileManager.getFileUrl(file)}
+                  style={{ height: '100%', objectFit: 'cover', cursor: 'pointer' }}
+                  preview={{ mask: false }}
+                  onLoad={(e) => {
+                    const img = e.target as HTMLImageElement
+                    img.parentElement?.classList.add('loaded')
+                  }}
+                />
+                <ImageInfo>
+                  <div>{formatFileSize(file)}</div>
+                </ImageInfo>
+              </ImageWrapper>
+            </Col>
+          ))}
+        </Row>
+      </Image.PreviewGroup>
+    )
+  }
+
+  if (id.startsWith('gemini_')) {
+    return <GeminiFiles id={id.replace('gemini_', '') as string} />
+  }
+
+  return (
+    <Table
+      dataSource={dataSource}
+      columns={columns}
+      style={{ width: '100%' }}
+      size="small"
+      pagination={{ pageSize: 100 }}
+    />
+  )
+}
+
+const ImageWrapper = styled.div`
+  position: relative;
+  aspect-ratio: 1;
+  overflow: hidden;
+  border-radius: 8px;
+  background-color: var(--color-background-soft);
+  display: flex;
+  align-items: center;
+  justify-content: center;
+  border: 0.5px solid var(--color-border);
+
+  .ant-image {
+    height: 100%;
+    width: 100%;
+    opacity: 0;
+    transition:
+      opacity 0.3s ease,
+      transform 0.3s ease;
+
+    &.loaded {
+      opacity: 1;
+    }
+  }
+
+  &:hover {
+    .ant-image.loaded {
+      transform: scale(1.05);
+    }
+
+    div:last-child {
+      opacity: 1;
+    }
+  }
+`
+
+const LoadingWrapper = styled.div`
+  position: absolute;
+  top: 0;
+  left: 0;
+  right: 0;
+  bottom: 0;
+  display: flex;
+  align-items: center;
+  justify-content: center;
+  background-color: var(--color-background-soft);
+`
+
+const ImageInfo = styled.div`
+  position: absolute;
+  bottom: 0;
+  left: 0;
+  right: 0;
+  background: rgba(0, 0, 0, 0.6);
+  color: white;
+  padding: 5px 8px;
+  opacity: 0;
+  transition: opacity 0.3s ease;
+  font-size: 12px;
+
+  > div:first-child {
+    white-space: nowrap;
+    overflow: hidden;
+    text-overflow: ellipsis;
+  }
+`
+
+export default memo(ContentView)

src/renderer/src/pages/files/FilesPage.tsx

@@ -10,21 +10,27 @@ import { Navbar, NavbarCenter } from '@renderer/components/app/Navbar'
 import TextEditPopup from '@renderer/components/Popups/TextEditPopup'
 import Scrollbar from '@renderer/components/Scrollbar'
 import db from '@renderer/databases'
+import { useProviders } from '@renderer/hooks/useProvider'
 import FileManager from '@renderer/services/FileManager'
 import store from '@renderer/store'
 import { FileType, FileTypes } from '@renderer/types'
 import { formatFileSize } from '@renderer/utils'
 import type { MenuProps } from 'antd'
-import { Button, Col, Dropdown, Image, Menu, Row, Spin, Table } from 'antd'
+import { Button, Dropdown, Menu } from 'antd'
 import dayjs from 'dayjs'
 import { useLiveQuery } from 'dexie-react-hooks'
-import { FC, useState } from 'react'
+import { FC, useMemo, useState } from 'react'
 import { useTranslation } from 'react-i18next'
 import styled from 'styled-components'
 
+import ContentView from './ContentView'
+
 const FilesPage: FC = () => {
   const { t } = useTranslation()
-  const [fileType, setFileType] = useState<FileTypes | 'all'>('all')
+  const [fileType, setFileType] = useState<FileTypes | 'all' | 'gemini'>('all')
+  const { providers } = useProviders()
+
+  const geminiProviders = providers.filter((provider) => provider.type === 'gemini')
 
   const files = useLiveQuery<FileType[]>(() => {
     if (fileType === 'all') {
@@ -111,58 +117,68 @@ const FilesPage: FC = () => {
       created_at: dayjs(file.created_at).format('MM-DD HH:mm'),
       created_at_unix: dayjs(file.created_at).unix(),
       actions: (
-        <Dropdown menu={{ items: getActionMenu(file.id) }} trigger={['click']}>
+        <Dropdown menu={{ items: getActionMenu(file.id) }} trigger={['click']} placement="bottom" arrow>
           <Button type="text" size="small" icon={<EllipsisOutlined />} />
         </Dropdown>
       )
     }
   })
 
-  const columns = [
-    {
-      title: t('files.name'),
-      dataIndex: 'file',
-      key: 'file',
-      width: '300px'
-    },
-    {
-      title: t('files.size'),
-      dataIndex: 'size',
-      key: 'size',
-      width: '80px',
-      sorter: (a: { size_bytes: number }, b: { size_bytes: number }) => b.size_bytes - a.size_bytes,
-      align: 'center'
-    },
-    {
-      title: t('files.count'),
-      dataIndex: 'count',
-      key: 'count',
-      width: '60px',
-      sorter: (a: { count: number }, b: { count: number }) => b.count - a.count,
-      align: 'center'
-    },
-    {
-      title: t('files.created_at'),
-      dataIndex: 'created_at',
-      key: 'created_at',
-      width: '120px',
-      align: 'center',
-      sorter: (a: { created_at_unix: number }, b: { created_at_unix: number }) => b.created_at_unix - a.created_at_unix
-    },
-    {
-      title: t('files.actions'),
-      dataIndex: 'actions',
-      key: 'actions',
-      width: '50px'
-    }
-  ]
+  const columns = useMemo(
+    () => [
+      {
+        title: t('files.name'),
+        dataIndex: 'file',
+        key: 'file',
+        width: '300px'
+      },
+      {
+        title: t('files.size'),
+        dataIndex: 'size',
+        key: 'size',
+        width: '80px',
+        sorter: (a: { size_bytes: number }, b: { size_bytes: number }) => b.size_bytes - a.size_bytes,
+        align: 'center'
+      },
+      {
+        title: t('files.count'),
+        dataIndex: 'count',
+        key: 'count',
+        width: '60px',
+        sorter: (a: { count: number }, b: { count: number }) => b.count - a.count,
+        align: 'center'
+      },
+      {
+        title: t('files.created_at'),
+        dataIndex: 'created_at',
+        key: 'created_at',
+        width: '120px',
+        align: 'center',
+        sorter: (a: { created_at_unix: number }, b: { created_at_unix: number }) =>
+          b.created_at_unix - a.created_at_unix
+      },
+      {
+        title: t('files.actions'),
+        dataIndex: 'actions',
+        key: 'actions',
+        width: '80px',
+        align: 'center'
+      }
+    ],
+    [t]
+  )
 
   const menuItems = [
     { key: 'all', label: t('files.all'), icon: <FileTextOutlined /> },
     { key: FileTypes.IMAGE, label: t('files.image'), icon: <FileImageOutlined /> },
     { key: FileTypes.TEXT, label: t('files.text'), icon: <FileTextOutlined /> },
-    { key: FileTypes.DOCUMENT, label: t('files.document'), icon: <FilePdfOutlined /> }
-  ]
+    { key: FileTypes.DOCUMENT, label: t('files.document'), icon: <FilePdfOutlined /> },
+    ...geminiProviders.map((provider) => ({
+      key: 'gemini_' + provider.id,
+      label: provider.name,
+      icon: <FilePdfOutlined />
+    }))
+  ].filter(Boolean) as MenuProps['items']
 
   return (
     <Container>
@@ -174,41 +190,7 @@ const FilesPage: FC = () => {
           <Menu selectedKeys={[fileType]} items={menuItems} onSelect={({ key }) => setFileType(key as FileTypes)} />
         </SideNav>
         <TableContainer right>
-          {fileType === FileTypes.IMAGE && files?.length && files?.length > 0 ? (
-            <Image.PreviewGroup>
-              <Row gutter={[16, 16]}>
-                {files?.map((file) => (
-                  <Col key={file.id} xs={24} sm={12} md={8} lg={4} xl={3}>
-                    <ImageWrapper>
-                      <LoadingWrapper>
-                        <Spin />
-                      </LoadingWrapper>
-                      <Image
-                        src={FileManager.getFileUrl(file)}
-                        style={{ height: '100%', objectFit: 'cover', cursor: 'pointer' }}
-                        preview={{ mask: false }}
-                        onLoad={(e) => {
-                          const img = e.target as HTMLImageElement
-                          img.parentElement?.classList.add('loaded')
-                        }}
-                      />
-                      <ImageInfo>
-                        <div>{formatFileSize(file)}</div>
-                      </ImageInfo>
-                    </ImageWrapper>
-                  </Col>
-                ))}
-              </Row>
-            </Image.PreviewGroup>
-          ) : (
-            <Table
-              dataSource={dataSource}
-              columns={columns as any}
-              style={{ width: '100%' }}
-              size="small"
-              pagination={{ pageSize: 100 }}
-            />
-          )}
+          <ContentView id={fileType} files={files} dataSource={dataSource} columns={columns} />
         </TableContainer>
       </ContentContainer>
     </Container>
@@ -242,72 +224,6 @@ const FileNameText = styled.div`
   cursor: pointer;
 `
 
-const ImageWrapper = styled.div`
-  position: relative;
-  aspect-ratio: 1;
-  overflow: hidden;
-  border-radius: 8px;
-  background-color: var(--color-background-soft);
-  display: flex;
-  align-items: center;
-  justify-content: center;
-  border: 0.5px solid var(--color-border);
-
-  .ant-image {
-    height: 100%;
-    width: 100%;
-    opacity: 0;
-    transition:
-      opacity 0.3s ease,
-      transform 0.3s ease;
-
-    &.loaded {
-      opacity: 1;
-    }
-  }
-
-  &:hover {
-    .ant-image.loaded {
-      transform: scale(1.05);
-    }
-
-    div:last-child {
-      opacity: 1;
-    }
-  }
-`
-
-const LoadingWrapper = styled.div`
-  position: absolute;
-  top: 0;
-  left: 0;
-  right: 0;
-  bottom: 0;
-  display: flex;
-  align-items: center;
-  justify-content: center;
-  background-color: var(--color-background-soft);
-`
-
-const ImageInfo = styled.div`
-  position: absolute;
-  bottom: 0;
-  left: 0;
-  right: 0;
-  background: rgba(0, 0, 0, 0.6);
-  color: white;
-  padding: 5px 8px;
-  opacity: 0;
-  transition: opacity 0.3s ease;
-  font-size: 12px;
-
-  > div:first-child {
-    white-space: nowrap;
-    overflow: hidden;
-    text-overflow: ellipsis;
-  }
-`
-
 const SideNav = styled.div`
   width: var(--assistants-width);
   border-right: 0.5px solid var(--color-border);

src/renderer/src/pages/files/GeminiFiles.tsx

@@ -0,0 +1,101 @@
+import { DeleteOutlined } from '@ant-design/icons'
+import { FileMetadataResponse, FileState } from '@google/generative-ai/server'
+import { useProvider } from '@renderer/hooks/useProvider'
+import GeminiProvider from '@renderer/providers/GeminiProvider'
+import { runAsyncFunction } from '@renderer/utils'
+import { Table } from 'antd'
+import type { ColumnsType } from 'antd/es/table'
+import { FC, useCallback, useEffect, useState } from 'react'
+import { useTranslation } from 'react-i18next'
+import styled from 'styled-components'
+
+interface GeminiFilesProps {
+  id: string
+}
+
+const GeminiFiles: FC<GeminiFilesProps> = ({ id }) => {
+  const { provider } = useProvider(id)
+  const [files, setFiles] = useState<FileMetadataResponse[]>([])
+  const { t } = useTranslation()
+  const [loading, setLoading] = useState(false)
+
+  const fetchFiles = useCallback(async () => {
+    const geminiProvider = new GeminiProvider(provider)
+    const { files } = await geminiProvider.listFiles()
+    files && setFiles(files.filter((file) => file.state === FileState.ACTIVE))
+  }, [provider])
+
+  const columns: ColumnsType<FileMetadataResponse> = [
+    {
+      title: t('files.name'),
+      dataIndex: 'displayName',
+      key: 'displayName'
+    },
+    {
+      title: t('files.type'),
+      dataIndex: 'mimeType',
+      key: 'mimeType'
+    },
+    {
+      title: t('files.size'),
+      dataIndex: 'sizeBytes',
+      key: 'sizeBytes',
+      render: (size: string) => `${(parseInt(size) / 1024 / 1024).toFixed(2)} MB`
+    },
+    {
+      title: t('files.created_at'),
+      dataIndex: 'createTime',
+      key: 'createTime',
+      render: (time: string) => new Date(time).toLocaleString()
+    },
+    {
+      title: t('files.actions'),
+      dataIndex: 'actions',
+      key: 'actions',
+      align: 'center',
+      render: (_, record) => {
+        const geminiProvider = new GeminiProvider(provider)
+        return (
+          <DeleteOutlined
+            style={{ cursor: 'pointer', color: 'var(--color-error)' }}
+            onClick={() => {
+              setFiles(files.filter((file) => file.name !== record.name))
+              geminiProvider.deleteFile(record.name).catch((error) => {
+                console.error('Failed to delete file:', error)
+                setFiles((prev) => [...prev, record])
+              })
+            }}
+          />
+        )
+      }
+    }
+  ]
+
+  useEffect(() => {
+    runAsyncFunction(async () => {
+      try {
+        setLoading(true)
+        await fetchFiles()
+        setLoading(false)
+      } catch (error: any) {
+        console.error('Failed to fetch files:', error)
+        window.message.error(error.message)
+        setLoading(false)
+      }
+    })
+  }, [fetchFiles])
+
+  useEffect(() => {
+    setFiles([])
+  }, [id])
+
+  return (
+    <Container>
+      <Table columns={columns} dataSource={files} rowKey="name" loading={loading} />
+    </Container>
+  )
+}
+
+const Container = styled.div``
+
+export default GeminiFiles

src/renderer/src/providers/GeminiProvider.ts

@@ -1,5 +1,6 @@
 import {
   Content,
+  FileDataPart,
   GoogleGenerativeAI,
   HarmBlockThreshold,
   HarmCategory,
@@ -8,13 +9,14 @@ import {
   RequestOptions,
   TextPart
 } from '@google/generative-ai'
+import { GoogleAIFileManager, ListFilesResponse } from '@google/generative-ai/server'
 import { isEmbeddingModel, isWebSearchModel } from '@renderer/config/models'
 import { getStoreSetting } from '@renderer/hooks/useSettings'
 import i18n from '@renderer/i18n'
 import { getAssistantSettings, getDefaultModel, getTopNamingModel } from '@renderer/services/AssistantService'
 import { EVENT_NAMES } from '@renderer/services/EventService'
 import { filterContextMessages } from '@renderer/services/MessagesService'
-import { Assistant, FileTypes, Message, Model, Provider, Suggestion } from '@renderer/types'
+import { Assistant, FileType, FileTypes, Message, Model, Provider, Suggestion } from '@renderer/types'
 import { removeSpecialCharacters } from '@renderer/utils'
 import axios from 'axios'
 import { first, isEmpty, last, takeRight } from 'lodash'
@@ -39,6 +41,43 @@ export default class GeminiProvider extends BaseProvider {
     return this.provider.apiHost
   }
 
+  private async handlePdfFile(file: FileType): Promise<Part> {
+    const smallFileSize = 20 * 1024 * 1024
+    const isSmallFile = file.size < smallFileSize
+
+    if (isSmallFile) {
+      const { data, mimeType } = await window.api.gemini.base64File(file)
+      return {
+        inlineData: {
+          data,
+          mimeType
+        }
+      } as InlineDataPart
+    }
+
+    // Retrieve file from Gemini uploaded files
+    const fileMetadata = await window.api.gemini.retrieveFile(file, this.apiKey)
+
+    if (fileMetadata) {
+      return {
+        fileData: {
+          fileUri: fileMetadata.uri,
+          mimeType: fileMetadata.mimeType
+        }
+      } as FileDataPart
+    }
+
+    // If file is not found, upload it to Gemini
+    const uploadResult = await window.api.gemini.uploadFile(file, this.apiKey)
+
+    return {
+      fileData: {
+        fileUri: uploadResult.file.uri,
+        mimeType: uploadResult.file.mimeType
+      }
+    } as FileDataPart
+  }
+
   private async getMessageContents(message: Message): Promise<Content> {
     const role = message.role === 'user' ? 'user' : 'model'
 
@@ -54,6 +93,12 @@ export default class GeminiProvider extends BaseProvider {
           }
         } as InlineDataPart)
       }
+
+      if (file.ext === '.pdf') {
+        parts.push(await this.handlePdfFile(file))
+        continue
+      }
+
       if ([FileTypes.TEXT, FileTypes.DOCUMENT].includes(file.type)) {
         const fileContent = await (await window.api.file.read(file.id + file.ext)).trim()
         parts.push({
@@ -93,7 +138,7 @@ export default class GeminiProvider extends BaseProvider {
         model: model.id,
         systemInstruction: assistant.prompt,
         // @ts-ignore googleSearch is not a valid tool for Gemini
-        tools: assistant.enableWebSearch && isWebSearchModel(model) ? [{ googleSearch: {} }] : [],
+        tools: assistant.enableWebSearch && isWebSearchModel(model) ? [{ googleSearch: {} }] : undefined,
         generationConfig: {
           maxOutputTokens: maxTokens,
           temperature: assistant?.settings?.temperature,
@@ -300,4 +345,14 @@ export default class GeminiProvider extends BaseProvider {
     const data = await this.sdk.getGenerativeModel({ model: model.id }, this.requestOptions).embedContent('hi')
     return data.embedding.values.length
   }
+
+  public async listFiles(): Promise<ListFilesResponse> {
+    const fileManager = new GoogleAIFileManager(this.apiKey)
+    return await fileManager.listFiles()
+  }
+
+  public async deleteFile(fileId: string): Promise<void> {
+    const fileManager = new GoogleAIFileManager(this.apiKey)
+    await fileManager.deleteFile(fileId)
+  }
 }

src/renderer/src/services/ApiService.ts

@@ -230,16 +230,8 @@ export async function fetchModels(provider: Provider) {
 
 function formatErrorMessage(error: any): string {
   try {
-    return (
-      '```json\n' +
-      JSON.stringify(
-        error?.error?.message || error?.response?.data || error?.response || error?.request || error,
-        null,
-        2
-      ) +
-      '\n```'
-    )
+    return '```json\n' + JSON.stringify(error, null, 2) + '\n```'
   } catch (e) {
-    return 'Error: ' + error.message
+    return 'Error: ' + error?.message
   }
 }