diff --git a/package.json b/package.json index 545fd965..d348d017 100644 --- a/package.json +++ b/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", diff --git a/src/main/ipc.ts b/src/main/ipc.ts index 2c65fb42..1d5600f3 100644 --- a/src/main/ipc.ts +++ b/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) } diff --git a/src/main/services/CacheService.ts b/src/main/services/CacheService.ts new file mode 100644 index 00000000..d2984a99 --- /dev/null +++ b/src/main/services/CacheService.ts @@ -0,0 +1,74 @@ +interface CacheItem { + data: T + timestamp: number + duration: number +} + +export class CacheService { + private static cache: Map> = new Map() + + /** + * Set cache + * @param key Cache key + * @param data Cache data + * @param duration Cache duration (in milliseconds) + */ + static set(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(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 + } +} diff --git a/src/main/services/GeminiService.ts b/src/main/services/GeminiService.ts new file mode 100644 index 00000000..698d6b5b --- /dev/null +++ b/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 { + const fileManager = new GoogleAIFileManager(apiKey) + + const cachedResponse = CacheService.get(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 + } +} diff --git a/src/preload/index.d.ts b/src/preload/index.d.ts index 0333f1b0..ce635fee 100644 --- a/src/preload/index.d.ts +++ b/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 resetMinimumSize: () => Promise } + gemini: { + uploadFile: (file: FileType, apiKey: string) => Promise + retrieveFile: (file: FileType, apiKey: string) => Promise + base64File: (file: FileType) => Promise<{ data: string; mimeType: string }> + } } } } diff --git a/src/preload/index.ts b/src/preload/index.ts index 696d066e..8b6017be 100644 --- a/src/preload/index.ts +++ b/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) } } diff --git a/src/renderer/src/components/TopView/index.tsx b/src/renderer/src/components/TopView/index.tsx index 7fd59043..868a1c39 100644 --- a/src/renderer/src/components/TopView/index.tsx +++ b/src/renderer/src/components/TopView/index.tsx @@ -28,10 +28,7 @@ const TopViewContainer: React.FC = ({ children }) => { const elementsRef = useRef([]) 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() diff --git a/src/renderer/src/i18n/locales/en-us.json b/src/renderer/src/i18n/locales/en-us.json index 6fd31194..85ee3582 100644 --- a/src/renderer/src/i18n/locales/en-us.json +++ b/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", diff --git a/src/renderer/src/i18n/locales/ja-jp.json b/src/renderer/src/i18n/locales/ja-jp.json index 90965062..ca2ec747 100644 --- a/src/renderer/src/i18n/locales/ja-jp.json +++ b/src/renderer/src/i18n/locales/ja-jp.json @@ -183,6 +183,7 @@ "name": "名前", "open": "開く", "size": "サイズ", + "type": "タイプ", "text": "テキスト", "title": "ファイル", "edit": "編集", diff --git a/src/renderer/src/i18n/locales/ru-ru.json b/src/renderer/src/i18n/locales/ru-ru.json index 312d63a6..f40f3715 100644 --- a/src/renderer/src/i18n/locales/ru-ru.json +++ b/src/renderer/src/i18n/locales/ru-ru.json @@ -183,6 +183,7 @@ "name": "Имя", "open": "Открыть", "size": "Размер", + "type": "Тип", "text": "Текст", "title": "Файлы", "edit": "Редактировать", diff --git a/src/renderer/src/i18n/locales/zh-cn.json b/src/renderer/src/i18n/locales/zh-cn.json index 2a09d522..91a09dde 100644 --- a/src/renderer/src/i18n/locales/zh-cn.json +++ b/src/renderer/src/i18n/locales/zh-cn.json @@ -184,6 +184,7 @@ "name": "文件名", "open": "打开", "size": "大小", + "type": "类型", "text": "文本", "title": "文件", "edit": "编辑", diff --git a/src/renderer/src/i18n/locales/zh-tw.json b/src/renderer/src/i18n/locales/zh-tw.json index 3af6ba51..98601a11 100644 --- a/src/renderer/src/i18n/locales/zh-tw.json +++ b/src/renderer/src/i18n/locales/zh-tw.json @@ -183,6 +183,7 @@ "name": "名稱", "open": "打開", "size": "大小", + "type": "類型", "text": "文本", "title": "檔案", "edit": "編輯", diff --git a/src/renderer/src/pages/files/ContentView.tsx b/src/renderer/src/pages/files/ContentView.tsx new file mode 100644 index 00000000..1e75990c --- /dev/null +++ b/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 = ({ id, files, dataSource, columns }) => { + if (id === FileTypes.IMAGE && files?.length && files?.length > 0) { + return ( + + + {files?.map((file) => ( + + + + + + { + const img = e.target as HTMLImageElement + img.parentElement?.classList.add('loaded') + }} + /> + +
{formatFileSize(file)}
+
+
+ + ))} +
+
+ ) + } + + if (id.startsWith('gemini_')) { + return + } + + return ( + + ) +} + +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) diff --git a/src/renderer/src/pages/files/FilesPage.tsx b/src/renderer/src/pages/files/FilesPage.tsx index bfc066e7..a248ef64 100644 --- a/src/renderer/src/pages/files/FilesPage.tsx +++ b/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('all') + const [fileType, setFileType] = useState('all') + const { providers } = useProviders() + + const geminiProviders = providers.filter((provider) => provider.type === 'gemini') const files = useLiveQuery(() => { 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: ( - + - - - - - { - const img = e.target as HTMLImageElement - img.parentElement?.classList.add('loaded') - }} - /> - -
{formatFileSize(file)}
-
-
- - ))} - - - ) : ( -
- )} + @@ -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); diff --git a/src/renderer/src/pages/files/GeminiFiles.tsx b/src/renderer/src/pages/files/GeminiFiles.tsx new file mode 100644 index 00000000..1856d66c --- /dev/null +++ b/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 = ({ id }) => { + const { provider } = useProvider(id) + const [files, setFiles] = useState([]) + 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 = [ + { + 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 ( + { + 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 ( + +
+ + ) +} + +const Container = styled.div`` + +export default GeminiFiles diff --git a/src/renderer/src/providers/GeminiProvider.ts b/src/renderer/src/providers/GeminiProvider.ts index cc220e8e..e0baaa1b 100644 --- a/src/renderer/src/providers/GeminiProvider.ts +++ b/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 { + 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 { 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 { + const fileManager = new GoogleAIFileManager(this.apiKey) + return await fileManager.listFiles() + } + + public async deleteFile(fileId: string): Promise { + const fileManager = new GoogleAIFileManager(this.apiKey) + await fileManager.deleteFile(fileId) + } } diff --git a/src/renderer/src/services/ApiService.ts b/src/renderer/src/services/ApiService.ts index 046a5127..1e242441 100644 --- a/src/renderer/src/services/ApiService.ts +++ b/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 } }