From 2016ba70629750ab47d7f7f663d27cb725f6def6 Mon Sep 17 00:00:00 2001 From: kangfenmao Date: Wed, 11 Sep 2024 11:54:37 +0800 Subject: [PATCH] feat: add attachment files --- package.json | 1 + .../migrations/001_create_files_table.sql | 1 + src/main/file.ts | 36 +++++++++---------- src/main/ipc.ts | 20 +++++++++-- src/main/utils/file.ts | 16 +++++++++ src/preload/index.d.ts | 17 +++++---- src/preload/index.ts | 17 +++++---- src/renderer/src/hooks/useAssistant.ts | 4 +-- src/renderer/src/pages/files/FilesPage.tsx | 2 +- .../pages/home/Inputbar/AttachmentButton.tsx | 27 +++++++------- .../src/pages/home/Inputbar/Inputbar.tsx | 8 ++--- .../home/Messages/MessageAttachments.tsx | 6 +++- .../src/pages/home/Messages/Messages.tsx | 8 ++++- src/renderer/src/providers/OpenAIProvider.ts | 7 ++-- src/renderer/src/services/messages.ts | 4 +++ src/renderer/src/services/storage.ts | 13 +++++++ src/renderer/src/types/index.ts | 11 +++++- src/renderer/src/utils/index.ts | 16 +-------- tsconfig.node.json | 1 + yarn.lock | 10 ++++++ 20 files changed, 149 insertions(+), 76 deletions(-) diff --git a/package.json b/package.json index c113832b..8f271888 100644 --- a/package.json +++ b/package.json @@ -70,6 +70,7 @@ "i18next": "^23.11.5", "localforage": "^1.10.0", "lodash": "^4.17.21", + "mime": "^4.0.4", "openai": "^4.52.1", "prettier": "^3.2.4", "react": "^18.2.0", diff --git a/resources/migrations/001_create_files_table.sql b/resources/migrations/001_create_files_table.sql index e5d08fd8..a744d77f 100644 --- a/resources/migrations/001_create_files_table.sql +++ b/resources/migrations/001_create_files_table.sql @@ -5,5 +5,6 @@ CREATE TABLE IF NOT EXISTS files ( path TEXT NOT NULL, size INTEGER NOT NULL, ext TEXT NOT NULL, + type TEXT NOT NULL, created_at TEXT NOT NULL ) diff --git a/src/main/file.ts b/src/main/file.ts index 556bdc6d..df3b0c67 100644 --- a/src/main/file.ts +++ b/src/main/file.ts @@ -5,15 +5,8 @@ import * as fs from 'fs' import * as path from 'path' import { v4 as uuidv4 } from 'uuid' -interface FileMetadata { - id: string - name: string - file_name: string - path: string - size: number - ext: string - created_at: Date -} +import { FileMetadata } from '../renderer/src/types' +import { getFileType } from './utils/file' export class File { private storageDir: string @@ -88,6 +81,7 @@ export class File { const fileMetadataPromises = result.filePaths.map(async (filePath) => { const stats = fs.statSync(filePath) const ext = path.extname(filePath) + const fileType = getFileType(ext) return { id: uuidv4(), @@ -96,27 +90,29 @@ export class File { path: filePath, created_at: stats.birthtime, size: stats.size, - ext: ext + ext: ext, + type: fileType } }) return Promise.all(fileMetadataPromises) } - async uploadFile(filePath: string): Promise { - const duplicateFile = await this.findDuplicateFile(filePath) + async uploadFile(file: FileMetadata): Promise { + const duplicateFile = await this.findDuplicateFile(file.path) if (duplicateFile) { return duplicateFile } const uuid = uuidv4() - const name = path.basename(filePath) + const name = path.basename(file.path) const ext = path.extname(name) const destPath = path.join(this.storageDir, uuid + ext) - await fs.promises.copyFile(filePath, destPath) + await fs.promises.copyFile(file.path, destPath) const stats = await fs.promises.stat(destPath) + const fileType = getFileType(ext) const fileMetadata: FileMetadata = { id: uuid, @@ -125,12 +121,13 @@ export class File { path: destPath, created_at: stats.birthtime, size: stats.size, - ext: ext + ext: ext, + type: fileType } const stmt = this.db.prepare(` - INSERT INTO files (id, name, file_name, path, size, ext, created_at) - VALUES (?, ?, ?, ?, ?, ?, ?) + INSERT INTO files (id, name, file_name, path, size, ext, type, created_at) + VALUES (?, ?, ?, ?, ?, ?, ?, ?) `) stmt.run( @@ -140,6 +137,7 @@ export class File { fileMetadata.path, fileMetadata.size, fileMetadata.ext, + fileMetadata.type, fileMetadata.created_at.toISOString() ) @@ -155,8 +153,8 @@ export class File { } } - async batchUploadFiles(filePaths: string[]): Promise { - const uploadPromises = filePaths.map((filePath) => this.uploadFile(filePath)) + async batchUploadFiles(files: FileMetadata[]): Promise { + const uploadPromises = files.map((file) => this.uploadFile(file)) return Promise.all(uploadPromises) } diff --git a/src/main/ipc.ts b/src/main/ipc.ts index 92d122a2..d6edc2e5 100644 --- a/src/main/ipc.ts +++ b/src/main/ipc.ts @@ -1,5 +1,9 @@ import { BrowserWindow, ipcMain, OpenDialogOptions, session, shell } from 'electron' +import Logger from 'electron-log' +import fs from 'fs' +import path from 'path' +import { FileMetadata } from '../renderer/src/types' import { appConfig, titleBarOverlayDark, titleBarOverlayLight } from './config' import { File } from './file' import AppUpdater from './updater' @@ -34,18 +38,28 @@ export function registerIpc(mainWindow: BrowserWindow, app: Electron.App) { ipcMain.handle('zip:compress', (_, text: string) => compress(text)) ipcMain.handle('zip:decompress', (_, text: Buffer) => decompress(text)) + ipcMain.handle('image:base64', async (_, filePath) => { + try { + const data = await fs.promises.readFile(filePath) + return `data:image/${path.extname(filePath).slice(1)};base64,${data.toString('base64')}` + } catch (error) { + Logger.error('Error reading file:', error) + return '' + } + }) + ipcMain.handle('file:select', async (_, options?: OpenDialogOptions) => await fileManager.selectFile(options)) - ipcMain.handle('file:upload', async (_, filePath: string) => await fileManager.uploadFile(filePath)) + ipcMain.handle('file:upload', async (_, file: FileMetadata) => await fileManager.uploadFile(file)) ipcMain.handle('file:delete', async (_, fileId: string) => { await fileManager.deleteFile(fileId) return { success: true } }) - ipcMain.handle('file:batchUpload', async (_, filePaths: string[]) => await fileManager.batchUploadFiles(filePaths)) + ipcMain.handle('file:batchUpload', async (_, files: FileMetadata[]) => await fileManager.batchUploadFiles(files)) ipcMain.handle('file:batchDelete', async (_, fileIds: string[]) => { await fileManager.batchDeleteFiles(fileIds) return { success: true } }) - ipcMain.handle('file:getAll', () => fileManager.getAllFiles()) + ipcMain.handle('file:all', () => fileManager.getAllFiles()) ipcMain.handle('minapp', (_, args) => { createMinappWindow({ diff --git a/src/main/utils/file.ts b/src/main/utils/file.ts index 31691ad9..580d138f 100644 --- a/src/main/utils/file.ts +++ b/src/main/utils/file.ts @@ -3,6 +3,8 @@ import logger from 'electron-log' import { writeFile } from 'fs' import { readFile } from 'fs/promises' +import { FileType } from '../../renderer/src/types' + export async function saveFile( _: Electron.IpcMainInvokeEvent, fileName: string, @@ -53,3 +55,17 @@ export async function openFile( return null } } + +export function getFileType(ext: string): FileType { + const imageExts = ['.jpg', '.jpeg', '.png', '.gif', '.bmp', '.webp'] + const videoExts = ['.mp4', '.avi', '.mov', '.wmv', '.flv', '.mkv'] + const audioExts = ['.mp3', '.wav', '.ogg', '.flac', '.aac'] + const documentExts = ['.pdf', '.doc', '.docx', '.xls', '.xlsx', '.ppt', '.pptx', '.txt'] + + ext = ext.toLowerCase() + if (imageExts.includes(ext)) return FileType.IMAGE + if (videoExts.includes(ext)) return FileType.VIDEO + if (audioExts.includes(ext)) return FileType.AUDIO + if (documentExts.includes(ext)) return FileType.DOCUMENT + return FileType.OTHER +} diff --git a/src/preload/index.d.ts b/src/preload/index.d.ts index 9d28f9b4..8b665dc7 100644 --- a/src/preload/index.d.ts +++ b/src/preload/index.d.ts @@ -22,12 +22,17 @@ declare global { reload: () => void compress: (text: string) => Promise decompress: (text: Buffer) => Promise - fileSelect: (options?: OpenDialogOptions) => Promise - fileUpload: (filePath: string) => Promise - fileDelete: (fileId: string) => Promise<{ success: boolean }> - fileBatchUpload: (filePaths: string[]) => Promise - fileBatchDelete: (fileIds: string[]) => Promise<{ success: boolean }> - fileGetAll: () => Promise + file: { + select: (options?: OpenDialogOptions) => Promise + upload: (file: FileMetadata) => Promise + delete: (fileId: string) => Promise<{ success: boolean }> + batchUpload: (files: FileMetadata[]) => Promise + batchDelete: (fileIds: string[]) => Promise<{ success: boolean }> + all: () => Promise + } + image: { + base64: (filePath: string) => Promise + } } } } diff --git a/src/preload/index.ts b/src/preload/index.ts index d0866f3e..0938959f 100644 --- a/src/preload/index.ts +++ b/src/preload/index.ts @@ -16,12 +16,17 @@ const api = { }, compress: (text: string) => ipcRenderer.invoke('zip:compress', text), decompress: (text: Buffer) => ipcRenderer.invoke('zip:decompress', text), - fileSelect: (options?: OpenDialogOptions) => ipcRenderer.invoke('file:select', options), - fileUpload: (filePath: string) => ipcRenderer.invoke('file:upload', filePath), - fileDelete: (fileId: string) => ipcRenderer.invoke('file:delete', fileId), - fileBatchUpload: (filePaths: string[]) => ipcRenderer.invoke('file:batchUpload', filePaths), - fileBatchDelete: (fileIds: string[]) => ipcRenderer.invoke('file:batchDelete', fileIds), - fileGetAll: () => ipcRenderer.invoke('file:getAll') + file: { + select: (options?: OpenDialogOptions) => ipcRenderer.invoke('file:select', options), + upload: (filePath: string) => ipcRenderer.invoke('file:upload', filePath), + delete: (fileId: string) => ipcRenderer.invoke('file:delete', fileId), + batchUpload: (filePaths: string[]) => ipcRenderer.invoke('file:batchUpload', filePaths), + batchDelete: (fileIds: string[]) => ipcRenderer.invoke('file:batchDelete', fileIds), + all: () => ipcRenderer.invoke('file:all') + }, + image: { + base64: (filePath: string) => ipcRenderer.invoke('image:base64', filePath) + } } // Use `contextBridge` APIs to expose Electron APIs to diff --git a/src/renderer/src/hooks/useAssistant.ts b/src/renderer/src/hooks/useAssistant.ts index 817a9835..93c27507 100644 --- a/src/renderer/src/hooks/useAssistant.ts +++ b/src/renderer/src/hooks/useAssistant.ts @@ -1,4 +1,5 @@ import { getDefaultTopic } from '@renderer/services/assistant' +import LocalStorage from '@renderer/services/storage' import { useAppDispatch, useAppSelector } from '@renderer/store' import { addAssistant, @@ -16,7 +17,6 @@ import { } from '@renderer/store/assistants' import { setDefaultModel, setTopicNamingModel, setTranslateModel } from '@renderer/store/llm' import { Assistant, AssistantSettings, Model, Topic } from '@renderer/types' -import localforage from 'localforage' export function useAssistants() { const { assistants } = useAppSelector((state) => state.assistants) @@ -30,7 +30,7 @@ export function useAssistants() { dispatch(removeAssistant({ id })) const assistant = assistants.find((a) => a.id === id) if (assistant) { - assistant.topics.forEach((id) => localforage.removeItem(`topic:${id}`)) + assistant.topics.forEach(({ id }) => LocalStorage.removeTopic(id)) } } } diff --git a/src/renderer/src/pages/files/FilesPage.tsx b/src/renderer/src/pages/files/FilesPage.tsx index 05e63899..f80f818f 100644 --- a/src/renderer/src/pages/files/FilesPage.tsx +++ b/src/renderer/src/pages/files/FilesPage.tsx @@ -11,7 +11,7 @@ const FilesPage: FC = () => { const [files, setFiles] = useState([]) useEffect(() => { - window.api.fileGetAll().then(setFiles) + window.api.file.all().then(setFiles) }, []) const dataSource = files.map((file) => ({ diff --git a/src/renderer/src/pages/home/Inputbar/AttachmentButton.tsx b/src/renderer/src/pages/home/Inputbar/AttachmentButton.tsx index 33c3f195..58c49aa0 100644 --- a/src/renderer/src/pages/home/Inputbar/AttachmentButton.tsx +++ b/src/renderer/src/pages/home/Inputbar/AttachmentButton.tsx @@ -1,29 +1,30 @@ import { PaperClipOutlined } from '@ant-design/icons' -import { Tooltip, Upload } from 'antd' +import { FileMetadata } from '@renderer/types' +import { Tooltip } from 'antd' import { FC } from 'react' import { useTranslation } from 'react-i18next' interface Props { - files: File[] - setFiles: (files: File[]) => void + files: FileMetadata[] + setFiles: (files: FileMetadata[]) => void ToolbarButton: any } const AttachmentButton: FC = ({ files, setFiles, ToolbarButton }) => { const { t } = useTranslation() + const onSelectFile = async () => { + const _files = await window.api.file.select({ + filters: [{ name: 'Files', extensions: ['jpg', 'png', 'jpeg'] }] + }) + _files && setFiles(_files) + } + return ( - {}} - accept="image/*" - itemRender={() => null} - maxCount={1} - onChange={async ({ file }) => file?.originFileObj && setFiles([file.originFileObj as File])}> - - - - + + + ) } diff --git a/src/renderer/src/pages/home/Inputbar/Inputbar.tsx b/src/renderer/src/pages/home/Inputbar/Inputbar.tsx index 49b37d3d..d8b8aab2 100644 --- a/src/renderer/src/pages/home/Inputbar/Inputbar.tsx +++ b/src/renderer/src/pages/home/Inputbar/Inputbar.tsx @@ -16,7 +16,7 @@ import { EVENT_NAMES, EventEmitter } from '@renderer/services/event' import { estimateInputTokenCount } from '@renderer/services/messages' import store, { useAppDispatch, useAppSelector } from '@renderer/store' import { setGenerating, setSearching } from '@renderer/store/runtime' -import { Assistant, Message, Topic } from '@renderer/types' +import { Assistant, FileMetadata, Message, Topic } from '@renderer/types' import { delay, uuid } from '@renderer/utils' import { Button, Popconfirm, Tooltip } from 'antd' import TextArea, { TextAreaRef } from 'antd/es/input/TextArea' @@ -47,7 +47,7 @@ const Inputbar: FC = ({ assistant, setActiveTopic }) => { const [contextCount, setContextCount] = useState(0) const generating = useAppSelector((state) => state.runtime.generating) const textareaRef = useRef(null) - const [files, setFiles] = useState([]) + const [files, setFiles] = useState([]) const { t } = useTranslation() const containerRef = useRef(null) const { showTopics, toggleShowTopics } = useShowTopics() @@ -56,7 +56,7 @@ const Inputbar: FC = ({ assistant, setActiveTopic }) => { _text = text - const sendMessage = useCallback(() => { + const sendMessage = useCallback(async () => { if (generating) { return } @@ -76,7 +76,7 @@ const Inputbar: FC = ({ assistant, setActiveTopic }) => { } if (files.length > 0) { - message.files = files + message.files = await window.api.file.batchUpload(files) } EventEmitter.emit(EVENT_NAMES.SEND_MESSAGE, message) diff --git a/src/renderer/src/pages/home/Messages/MessageAttachments.tsx b/src/renderer/src/pages/home/Messages/MessageAttachments.tsx index f05bd443..8c7044d9 100644 --- a/src/renderer/src/pages/home/Messages/MessageAttachments.tsx +++ b/src/renderer/src/pages/home/Messages/MessageAttachments.tsx @@ -8,7 +8,11 @@ interface Props { } const MessageAttachments: FC = ({ message }) => { - return {message.images?.map((image) => )} + return ( + + {message.files?.map((image) => )} + + ) } const Container = styled.div` diff --git a/src/renderer/src/pages/home/Messages/Messages.tsx b/src/renderer/src/pages/home/Messages/Messages.tsx index 3c77d36a..69c101ad 100644 --- a/src/renderer/src/pages/home/Messages/Messages.tsx +++ b/src/renderer/src/pages/home/Messages/Messages.tsx @@ -3,7 +3,12 @@ import { useProviderByAssistant } from '@renderer/hooks/useProvider' import { getTopic } from '@renderer/hooks/useTopic' import { fetchChatCompletion, fetchMessagesSummary } from '@renderer/services/api' import { EVENT_NAMES, EventEmitter } from '@renderer/services/event' -import { estimateHistoryTokenCount, filterMessages, getContextCount } from '@renderer/services/messages' +import { + deleteMessageFiles, + estimateHistoryTokenCount, + filterMessages, + getContextCount +} from '@renderer/services/messages' import LocalStorage from '@renderer/services/storage' import { Assistant, Message, Model, Topic } from '@renderer/types' import { getBriefInfo, runAsyncFunction, uuid } from '@renderer/utils' @@ -56,6 +61,7 @@ const Messages: FC = ({ assistant, topic, setActiveTopic }) => { const _messages = messages.filter((m) => m.id !== message.id) setMessages(_messages) localforage.setItem(`topic:${topic.id}`, { id: topic.id, messages: _messages }) + deleteMessageFiles(message) }, [messages, topic.id] ) diff --git a/src/renderer/src/providers/OpenAIProvider.ts b/src/renderer/src/providers/OpenAIProvider.ts index 755d56a1..bc570384 100644 --- a/src/renderer/src/providers/OpenAIProvider.ts +++ b/src/renderer/src/providers/OpenAIProvider.ts @@ -3,7 +3,7 @@ import { getAssistantSettings, getDefaultModel, getTopNamingModel } from '@rende import { EVENT_NAMES } from '@renderer/services/event' import { filterContextMessages, filterMessages } from '@renderer/services/messages' import { Assistant, Message, Provider, Suggestion } from '@renderer/types' -import { fileToBase64, removeQuotes } from '@renderer/utils' +import { removeQuotes } from '@renderer/utils' import { first, takeRight } from 'lodash' import OpenAI from 'openai' import { @@ -33,13 +33,13 @@ export default class OpenAIProvider extends BaseProvider { return message.content } - if (file.type.includes('image')) { + if (file.type === 'image') { return [ { type: 'text', text: message.content }, { type: 'image_url', image_url: { - url: await fileToBase64(file) + url: await window.api.image.base64(file.path) } } ] @@ -58,7 +58,6 @@ export default class OpenAIProvider extends BaseProvider { const { contextCount, maxTokens } = getAssistantSettings(assistant) const systemMessage = assistant.prompt ? { role: 'system', content: assistant.prompt } : undefined - const userMessages: ChatCompletionMessageParam[] = [] for (const message of filterMessages(filterContextMessages(takeRight(messages, contextCount + 1)))) { diff --git a/src/renderer/src/services/messages.ts b/src/renderer/src/services/messages.ts index 44f5dc73..88ecdd2e 100644 --- a/src/renderer/src/services/messages.ts +++ b/src/renderer/src/services/messages.ts @@ -59,3 +59,7 @@ export function estimateHistoryTokenCount(assistant: Assistant, msgs: Message[]) return all.usedTokens - 7 } + +export function deleteMessageFiles(message: Message) { + message.files && window.api.file.batchDelete(message.files.map((f) => f.id)) +} diff --git a/src/renderer/src/services/storage.ts b/src/renderer/src/services/storage.ts index fcb8e03f..54168533 100644 --- a/src/renderer/src/services/storage.ts +++ b/src/renderer/src/services/storage.ts @@ -2,6 +2,8 @@ import { Topic } from '@renderer/types' import { convertToBase64 } from '@renderer/utils' import localforage from 'localforage' +import { deleteMessageFiles } from './messages' + const IMAGE_PREFIX = 'image://' export default class LocalStorage { @@ -15,12 +17,23 @@ export default class LocalStorage { } static async removeTopic(id: string) { + const messages = await this.getTopicMessages(id) + + for (const message of messages) { + await deleteMessageFiles(message) + } + localforage.removeItem(`topic:${id}`) } static async clearTopicMessages(id: string) { const topic = await this.getTopic(id) + if (topic) { + for (const message of topic?.messages ?? []) { + await deleteMessageFiles(message) + } + topic.messages = [] await localforage.setItem(`topic:${id}`, topic) } diff --git a/src/renderer/src/types/index.ts b/src/renderer/src/types/index.ts index 715c1b61..4bc9a109 100644 --- a/src/renderer/src/types/index.ts +++ b/src/renderer/src/types/index.ts @@ -27,7 +27,7 @@ export type Message = { createdAt: string status: 'sending' | 'pending' | 'success' | 'paused' | 'error' modelId?: string - files?: File[] + files?: FileMetadata[] images?: string[] usage?: OpenAI.Completions.CompletionUsage type?: 'text' | '@' | 'clear' @@ -94,5 +94,14 @@ export interface FileMetadata { path: string size: number ext: string + type: FileType created_at: Date } + +export enum FileType { + IMAGE = 'image', + VIDEO = 'video', + AUDIO = 'audio', + DOCUMENT = 'document', + OTHER = 'other' +} diff --git a/src/renderer/src/utils/index.ts b/src/renderer/src/utils/index.ts index dbdbcc77..a8373097 100644 --- a/src/renderer/src/utils/index.ts +++ b/src/renderer/src/utils/index.ts @@ -1,5 +1,6 @@ import { Model } from '@renderer/types' import imageCompression from 'browser-image-compression' +// @ts-ignore next-line` import { v4 as uuidv4 } from 'uuid' export const runAsyncFunction = async (fn: () => void) => { @@ -223,18 +224,3 @@ export function getBriefInfo(text: string, maxLength: number = 50): string { // 截取前面的内容,并在末尾添加 "..." return truncatedText + '...' } - -export async function fileToBase64(file: File): Promise { - return new Promise((resolve, reject) => { - try { - const reader = new FileReader() - reader.onload = (e: ProgressEvent) => { - const result = e.target?.result - resolve(typeof result === 'string' ? result : '') - } - reader.readAsDataURL(file) - } catch (error: any) { - reject(error) - } - }) -} diff --git a/tsconfig.node.json b/tsconfig.node.json index dbea6885..69682c7f 100644 --- a/tsconfig.node.json +++ b/tsconfig.node.json @@ -5,6 +5,7 @@ "src/main/**/*", "src/preload/**/*", "src/main/env.d.ts", + "src/renderer/src/types/index.ts" ], "compilerOptions": { "composite": true, diff --git a/yarn.lock b/yarn.lock index 1439d7c4..a7c148ec 100644 --- a/yarn.lock +++ b/yarn.lock @@ -1944,6 +1944,7 @@ __metadata: i18next: "npm:^23.11.5" localforage: "npm:^1.10.0" lodash: "npm:^4.17.21" + mime: "npm:^4.0.4" openai: "npm:^4.52.1" prettier: "npm:^3.2.4" react: "npm:^18.2.0" @@ -6592,6 +6593,15 @@ __metadata: languageName: node linkType: hard +"mime@npm:^4.0.4": + version: 4.0.4 + resolution: "mime@npm:4.0.4" + bin: + mime: bin/cli.js + checksum: 10c0/3046e425ed616613af8c7f4b268ff33ab564baeb24a117ef00475cbb23fbae91369ff2d9918cc6408162b0016bde34ea8cc4041b830fc2c45a8ecaf5b7e3e26f + languageName: node + linkType: hard + "mimic-fn@npm:^2.1.0": version: 2.1.0 resolution: "mimic-fn@npm:2.1.0"