From 0337c6649b063cb1ace6cdc962cf82a4a98e7956 Mon Sep 17 00:00:00 2001 From: kangfenmao Date: Wed, 11 Sep 2024 15:20:12 +0800 Subject: [PATCH] feat: Added tracking column to files table and updated FileMetadata interface. - Added a "count" column with default value 1 to the "files" table for tracking purposes. - Improved file duplication and deletion handling. - Updated regular expression for vision models to include additional providers. - Improved removal of topics for assistants from local storage. - Added support for human-readable date formats in file metadata. - Improved handling of messages with image attachments to include base64 encoded images in the response. - Added new 'count' property to the FileMetadata interface. --- .../migrations/001_create_files_table.sql | 3 +- src/main/file.ts | 46 +++++++++++-------- src/renderer/src/config/models.ts | 2 +- src/renderer/src/hooks/useAssistant.ts | 10 ++-- src/renderer/src/pages/files/FilesPage.tsx | 3 +- .../src/providers/AnthropicProvider.ts | 37 +++++++++++++-- src/renderer/src/types/index.ts | 1 + 7 files changed, 70 insertions(+), 32 deletions(-) diff --git a/resources/migrations/001_create_files_table.sql b/resources/migrations/001_create_files_table.sql index a744d77f..837de796 100644 --- a/resources/migrations/001_create_files_table.sql +++ b/resources/migrations/001_create_files_table.sql @@ -6,5 +6,6 @@ CREATE TABLE IF NOT EXISTS files ( size INTEGER NOT NULL, ext TEXT NOT NULL, type TEXT NOT NULL, - created_at TEXT NOT NULL + created_at TEXT NOT NULL, + count INTEGER DEFAULT 1 ) diff --git a/src/main/file.ts b/src/main/file.ts index df3b0c67..4eac4348 100644 --- a/src/main/file.ts +++ b/src/main/file.ts @@ -91,7 +91,8 @@ export class File { created_at: stats.birthtime, size: stats.size, ext: ext, - type: fileType + type: fileType, + count: 1 } }) @@ -102,7 +103,12 @@ export class File { const duplicateFile = await this.findDuplicateFile(file.path) if (duplicateFile) { - return duplicateFile + // Increment the count for the duplicate file + const updateStmt = this.db.prepare('UPDATE files SET count = count + 1 WHERE id = ?') + updateStmt.run(duplicateFile.id) + + // Fetch the updated file metadata + return this.getFile(duplicateFile.id)! } const uuid = uuidv4() @@ -122,12 +128,13 @@ export class File { created_at: stats.birthtime, size: stats.size, ext: ext, - type: fileType + type: fileType, + count: 1 } const stmt = this.db.prepare(` - INSERT INTO files (id, name, file_name, path, size, ext, type, created_at) - VALUES (?, ?, ?, ?, ?, ?, ?, ?) + INSERT INTO files (id, name, file_name, path, size, ext, type, created_at, count) + VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?) `) stmt.run( @@ -138,7 +145,8 @@ export class File { fileMetadata.size, fileMetadata.ext, fileMetadata.type, - fileMetadata.created_at.toISOString() + fileMetadata.created_at.toISOString(), + fileMetadata.count ) return fileMetadata @@ -147,9 +155,16 @@ export class File { async deleteFile(fileId: string): Promise { const fileMetadata = this.getFile(fileId) if (fileMetadata) { - await fs.promises.unlink(fileMetadata.path) - const stmt = this.db.prepare('DELETE FROM files WHERE id = ?') - stmt.run(fileId) + if (fileMetadata.count > 1) { + // Decrement the count if there are multiple references + const updateStmt = this.db.prepare('UPDATE files SET count = count - 1 WHERE id = ?') + updateStmt.run(fileId) + } else { + // Delete the file and database entry if this is the last reference + await fs.promises.unlink(fileMetadata.path) + const deleteStmt = this.db.prepare('DELETE FROM files WHERE id = ?') + deleteStmt.run(fileId) + } } } @@ -166,21 +181,12 @@ export class File { getFile(id: string): FileMetadata | null { const stmt = this.db.prepare('SELECT * FROM files WHERE id = ?') const row = stmt.get(id) as any - if (row) { - return { - ...row, - created_at: new Date(row.created_at) - } - } - return null + return row ? { ...row, created_at: new Date(row.created_at) } : null } getAllFiles(): FileMetadata[] { const stmt = this.db.prepare('SELECT * FROM files') const rows = stmt.all() as any[] - return rows.map((row) => ({ - ...row, - created_at: new Date(row.created_at) - })) + return rows.map((row) => ({ ...row, created_at: new Date(row.created_at) })) } } diff --git a/src/renderer/src/config/models.ts b/src/renderer/src/config/models.ts index d89b8a45..3dbaac26 100644 --- a/src/renderer/src/config/models.ts +++ b/src/renderer/src/config/models.ts @@ -1,7 +1,7 @@ import { Model } from '@renderer/types' const TEXT_TO_IMAGE_REGEX = /flux|diffusion|stabilityai|sd-turbo|dall|cogview/i -const VISION_REGEX = /llava|moondream|minicpm|gemini/i +const VISION_REGEX = /llava|moondream|minicpm|gemini|claude|vision/i const EMBEDDING_REGEX = /embedding/i export const SYSTEM_MODELS: Record = { diff --git a/src/renderer/src/hooks/useAssistant.ts b/src/renderer/src/hooks/useAssistant.ts index 93c27507..e807cf04 100644 --- a/src/renderer/src/hooks/useAssistant.ts +++ b/src/renderer/src/hooks/useAssistant.ts @@ -29,9 +29,8 @@ export function useAssistants() { removeAssistant: (id: string) => { dispatch(removeAssistant({ id })) const assistant = assistants.find((a) => a.id === id) - if (assistant) { - assistant.topics.forEach(({ id }) => LocalStorage.removeTopic(id)) - } + const topics = assistant?.topics || [] + topics.forEach(({ id }) => LocalStorage.removeTopic(id)) } } } @@ -45,7 +44,10 @@ export function useAssistant(id: string) { assistant, model: assistant?.model ?? defaultModel, addTopic: (topic: Topic) => dispatch(addTopic({ assistantId: assistant.id, topic })), - removeTopic: (topic: Topic) => dispatch(removeTopic({ assistantId: assistant.id, topic })), + removeTopic: (topic: Topic) => { + LocalStorage.removeTopic(topic.id) + dispatch(removeTopic({ assistantId: assistant.id, topic })) + }, updateTopic: (topic: Topic) => dispatch(updateTopic({ assistantId: assistant.id, topic })), updateTopics: (topics: Topic[]) => dispatch(updateTopics({ assistantId: assistant.id, topics })), removeAllTopics: () => dispatch(removeAllTopics({ assistantId: assistant.id })), diff --git a/src/renderer/src/pages/files/FilesPage.tsx b/src/renderer/src/pages/files/FilesPage.tsx index 4d363152..467f9ec3 100644 --- a/src/renderer/src/pages/files/FilesPage.tsx +++ b/src/renderer/src/pages/files/FilesPage.tsx @@ -2,6 +2,7 @@ import { Navbar, NavbarCenter } from '@renderer/components/app/Navbar' import { VStack } from '@renderer/components/Layout' import { FileMetadata } from '@renderer/types' import { Image, Table } from 'antd' +import dayjs from 'dayjs' import { FC, useEffect, useState } from 'react' import { useTranslation } from 'react-i18next' import styled from 'styled-components' @@ -18,7 +19,7 @@ const FilesPage: FC = () => { file: , name: {file.name}, size: `${(file.size / 1024 / 1024).toFixed(2)} MB`, - created_at: file.created_at.toISOString().split('T')[0] + created_at: dayjs(file.created_at).format('MM-DD HH:mm') })) const columns = [ diff --git a/src/renderer/src/providers/AnthropicProvider.ts b/src/renderer/src/providers/AnthropicProvider.ts index c8f3bb05..abd4b5be 100644 --- a/src/renderer/src/providers/AnthropicProvider.ts +++ b/src/renderer/src/providers/AnthropicProvider.ts @@ -18,6 +18,31 @@ export default class AnthropicProvider extends BaseProvider { this.sdk = new Anthropic({ apiKey: provider.apiKey, baseURL: this.getBaseURL() }) } + private async getMessageContent(message: Message): Promise { + const file = first(message.files) + + if (!file) { + return message.content + } + + if (file.type === 'image') { + const base64Data = await window.api.image.base64(file.path) + return [ + { type: 'text', text: message.content }, + { + type: 'image', + source: { + data: base64Data.base64, + media_type: base64Data.mime.replace('jpg', 'jpeg') as any, + type: 'base64' + } + } + ] + } + + return message.content + } + public async completions( messages: Message[], assistant: Assistant, @@ -27,12 +52,14 @@ export default class AnthropicProvider extends BaseProvider { const model = assistant.model || defaultModel const { contextCount, maxTokens } = getAssistantSettings(assistant) - const userMessages = filterMessages(filterContextMessages(takeRight(messages, contextCount + 2))).map((message) => { - return { + const userMessages: MessageParam[] = [] + + for (const message of filterMessages(filterContextMessages(takeRight(messages, contextCount + 2)))) { + userMessages.push({ role: message.role, - content: message.content - } - }) + content: await this.getMessageContent(message) + }) + } if (first(userMessages)?.role === 'assistant') { userMessages.shift() diff --git a/src/renderer/src/types/index.ts b/src/renderer/src/types/index.ts index 4bc9a109..5f46c8d8 100644 --- a/src/renderer/src/types/index.ts +++ b/src/renderer/src/types/index.ts @@ -96,6 +96,7 @@ export interface FileMetadata { ext: string type: FileType created_at: Date + count: number } export enum FileType {