fix: knowledge bugs

.github/workflows/release.yml

@@ -41,6 +41,20 @@ jobs:
           node-version: 20
           arch: ${{ matrix.arch }}
 
+      - name: Get yarn cache directory path
+        id: yarn-cache-dir-path
+        run: echo "dir=$(yarn config get cacheFolder)" >> $GITHUB_OUTPUT
+
+      - name: Cache yarn dependencies
+        uses: actions/cache@v3
+        with:
+          path: |
+            ${{ steps.yarn-cache-dir-path.outputs.dir }}
+            node_modules
+          key: ${{ runner.os }}-yarn-${{ hashFiles('**/yarn.lock') }}
+          restore-keys: |
+            ${{ runner.os }}-yarn-
+
       - name: Install corepack
         run: corepack enable && corepack prepare yarn@4.3.1 --activate
 

.gitignore

@@ -36,6 +36,7 @@ node_modules
 dist
 out
 build/icons
+stats.html
 
 # ENV
 .env

electron-builder.yml

@@ -11,6 +11,16 @@ files:
   - '!src'
   - '!scripts'
   - '!local'
+  - '!docs'
+  - '!packages'
+  - '!stats.html'
+  - '!*.md'
+  - '!node_modules/rollup-plugin-visualizer'
+  - '!node_modules/js-tiktoken'
+  - '!node_modules/node_modules/pdf-parse/lib/pdf.js/v1.9.426'
+  - '!node_modules/node_modules/pdf-parse/lib/pdf.js/v1.10.88'
+  - '!node_modules/node_modules/pdf-parse/lib/pdf.js/v2.0.550'
+
 asarUnpack:
   - resources/**
   - '**/*.{node,dll,metal,exp,lib}'

electron.vite.config.ts

@@ -1,6 +1,11 @@
 import react from '@vitejs/plugin-react'
 import { defineConfig, externalizeDepsPlugin } from 'electron-vite'
 import { resolve } from 'path'
+import { visualizer } from 'rollup-plugin-visualizer'
+
+const visualizerPlugin = (type: 'renderer' | 'main') => {
+  return process.env[`VISUALIZER_${type.toUpperCase()}`] ? [visualizer({ open: true })] : []
+}
 
 export default defineConfig({
   main: {
@@ -8,8 +13,6 @@ export default defineConfig({
       externalizeDepsPlugin({
         exclude: [
           '@llm-tools/embedjs',
-          '@llm-tools/embedjs-lancedb',
-          '@llm-tools/embedjs-ollama',
           '@llm-tools/embedjs-openai',
           '@llm-tools/embedjs-loader-web',
           '@llm-tools/embedjs-loader-markdown',
@@ -17,9 +20,10 @@ export default defineConfig({
           '@llm-tools/embedjs-loader-xml',
           '@llm-tools/embedjs-loader-pdf',
           '@llm-tools/embedjs-loader-sitemap',
-          '@lancedb/lancedb'
+          '@llm-tools/embedjs-libsql'
         ]
-      })
+      }),
+      ...visualizerPlugin('main')
     ],
     resolve: {
       alias: {
@@ -30,7 +34,7 @@ export default defineConfig({
     },
     build: {
       rollupOptions: {
-        external: ['@lancedb/lancedb']
+        external: ['@libsql/client']
       }
     }
   },
@@ -38,7 +42,7 @@ export default defineConfig({
     plugins: [externalizeDepsPlugin()]
   },
   renderer: {
-    plugins: [react()],
+    plugins: [react(), ...visualizerPlugin('renderer')],
     resolve: {
       alias: {
         '@renderer': resolve('src/renderer/src'),

package.json

@@ -25,6 +25,8 @@
     "typecheck": "npm run typecheck:node && npm run typecheck:web",
     "start": "electron-vite preview",
     "dev": "electron-vite dev",
+    "analyze:renderer": "VISUALIZER_RENDERER=true yarn build",
+    "analyze:main": "VISUALIZER_MAIN=true yarn build",
     "build": "npm run typecheck && electron-vite build",
     "postinstall": "electron-builder install-app-deps",
     "build:unpack": "dotenv npm run build && electron-builder --dir",
@@ -46,17 +48,16 @@
     "@electron-toolkit/preload": "^3.0.0",
     "@electron-toolkit/utils": "^3.0.0",
     "@electron/notarize": "^2.5.0",
-    "@llm-tools/embedjs": "^0.1.24",
+    "@llm-tools/embedjs": "^0.1.25",
-    "@llm-tools/embedjs-lancedb": "^0.1.24",
+    "@llm-tools/embedjs-libsql": "^0.1.25",
-    "@llm-tools/embedjs-loader-csv": "^0.1.24",
+    "@llm-tools/embedjs-loader-csv": "^0.1.25",
-    "@llm-tools/embedjs-loader-markdown": "^0.1.24",
+    "@llm-tools/embedjs-loader-markdown": "^0.1.25",
-    "@llm-tools/embedjs-loader-msoffice": "^0.1.24",
+    "@llm-tools/embedjs-loader-msoffice": "^0.1.25",
-    "@llm-tools/embedjs-loader-pdf": "^0.1.24",
+    "@llm-tools/embedjs-loader-pdf": "^0.1.25",
-    "@llm-tools/embedjs-loader-sitemap": "^0.1.24",
+    "@llm-tools/embedjs-loader-sitemap": "^0.1.25",
-    "@llm-tools/embedjs-loader-web": "^0.1.24",
+    "@llm-tools/embedjs-loader-web": "^0.1.25",
-    "@llm-tools/embedjs-loader-xml": "^0.1.24",
+    "@llm-tools/embedjs-loader-xml": "^0.1.25",
-    "@llm-tools/embedjs-ollama": "^0.1.24",
+    "@llm-tools/embedjs-openai": "^0.1.25",
-    "@llm-tools/embedjs-openai": "^0.1.24",
     "@types/react-infinite-scroll-component": "^5.0.0",
     "adm-zip": "^0.5.16",
     "apache-arrow": "^18.0.0",
@@ -69,6 +70,7 @@
     "html2canvas": "^1.4.1",
     "markdown-it": "^14.1.0",
     "officeparser": "^4.1.1",
+    "tokenx": "^0.4.1",
     "webdav": "4.11.4"
   },
   "devDependencies": {
@@ -109,7 +111,6 @@
     "eslint-plugin-react-hooks": "^4.6.2",
     "eslint-plugin-simple-import-sort": "^12.1.1",
     "eslint-plugin-unused-imports": "^4.0.0",
-    "gpt-tokens": "^1.3.10",
     "i18next": "^23.11.5",
     "lodash": "^4.17.21",
     "mime": "^4.0.4",
@@ -132,6 +133,7 @@
     "rehype-raw": "^7.0.0",
     "remark-gfm": "^4.0.0",
     "remark-math": "^6.0.0",
+    "rollup-plugin-visualizer": "^5.12.0",
     "sass": "^1.77.2",
     "shiki": "^1.22.2",
     "styled-components": "^6.1.11",

src/main/ipc.ts

@@ -101,6 +101,7 @@ export function registerIpc(mainWindow: BrowserWindow, app: Electron.App) {
 
   // file
   ipcMain.handle('file:open', fileManager.open)
+  ipcMain.handle('file:openPath', fileManager.openPath)
   ipcMain.handle('file:save', fileManager.save)
   ipcMain.handle('file:select', fileManager.selectFile)
   ipcMain.handle('file:upload', fileManager.uploadFile)

src/main/services/FileStorage.ts

@@ -8,7 +8,8 @@ import {
   OpenDialogOptions,
   OpenDialogReturnValue,
   SaveDialogOptions,
-  SaveDialogReturnValue
+  SaveDialogReturnValue,
+  shell
 } from 'electron'
 import logger from 'electron-log'
 import * as fs from 'fs'
@@ -300,6 +301,10 @@ class FileStorage {
     }
   }
 
+  public openPath = async (_: Electron.IpcMainInvokeEvent, path: string): Promise<void> => {
+    shell.openPath(path).catch((err) => logger.error('[IPC - Error] Failed to open file:', err))
+  }
+
   public save = async (
     _: Electron.IpcMainInvokeEvent,
     fileName: string,

src/main/services/KnowledgeService.ts

@@ -1,11 +1,11 @@
 import * as fs from 'node:fs'
 import path from 'node:path'
 
-import { RAGApplication, RAGApplicationBuilder, TextLoader } from '@llm-tools/embedjs'
+import { LocalPathLoader, RAGApplication, RAGApplicationBuilder, TextLoader } from '@llm-tools/embedjs'
 import { AddLoaderReturn, ExtractChunkData } from '@llm-tools/embedjs-interfaces'
-import { LanceDb } from '@llm-tools/embedjs-lancedb'
+import { LibSqlDb } from '@llm-tools/embedjs-libsql'
 import { MarkdownLoader } from '@llm-tools/embedjs-loader-markdown'
-import { DocxLoader } from '@llm-tools/embedjs-loader-msoffice'
+import { DocxLoader, ExcelLoader, PptLoader } from '@llm-tools/embedjs-loader-msoffice'
 import { PdfLoader } from '@llm-tools/embedjs-loader-pdf'
 import { SitemapLoader } from '@llm-tools/embedjs-loader-sitemap'
 import { WebLoader } from '@llm-tools/embedjs-loader-web'
@@ -34,10 +34,11 @@ class KnowledgeService {
           model,
           apiKey,
           configuration: { baseURL },
-          dimensions: 1024
+          dimensions: 1024,
+          batchSize: 10
         })
       )
-      .setVectorDatabase(new LanceDb({ path: path.join(this.storageDir, id) }))
+      .setVectorDatabase(new LibSqlDb({ path: path.join(this.storageDir, id) }))
       .build()
   }
 
@@ -62,41 +63,58 @@ class KnowledgeService {
 
   public add = async (
     _: Electron.IpcMainInvokeEvent,
-    { base, item }: { base: KnowledgeBaseParams; item: KnowledgeItem }
+    { base, item, forceReload = false }: { base: KnowledgeBaseParams; item: KnowledgeItem; forceReload: boolean }
   ): Promise<AddLoaderReturn> => {
     const ragApplication = await this.getRagApplication(base)
 
+    if (item.type === 'directory') {
+      const directory = item.content as string
+      return await ragApplication.addLoader(new LocalPathLoader({ path: directory }), forceReload)
+    }
+
     if (item.type === 'url') {
       const content = item.content as string
       if (content.startsWith('http')) {
-        return await ragApplication.addLoader(new WebLoader({ urlOrContent: content }))
+        return await ragApplication.addLoader(new WebLoader({ urlOrContent: content }), forceReload)
       }
     }
 
     if (item.type === 'sitemap') {
       const content = item.content as string
-      return await ragApplication.addLoader(new SitemapLoader({ url: content }))
+      return await ragApplication.addLoader(new SitemapLoader({ url: content }), forceReload)
     }
 
     if (item.type === 'note') {
       const content = item.content as string
-      return await ragApplication.addLoader(new TextLoader({ text: content }))
+      return await ragApplication.addLoader(new TextLoader({ text: content }), forceReload)
     }
 
     if (item.type === 'file') {
       const file = item.content as FileType
 
       if (file.ext === '.pdf') {
-        return await ragApplication.addLoader(new PdfLoader({ filePathOrUrl: file.path }) as any)
+        return await ragApplication.addLoader(new PdfLoader({ filePathOrUrl: file.path }) as any, forceReload)
       }
 
       if (file.ext === '.docx') {
-        return await ragApplication.addLoader(new DocxLoader({ filePathOrUrl: file.path }) as any)
+        return await ragApplication.addLoader(new DocxLoader({ filePathOrUrl: file.path }) as any, forceReload)
       }
 
-      if (file.ext.startsWith('.md')) {
+      if (file.ext === '.pptx') {
-        return await ragApplication.addLoader(new MarkdownLoader({ filePathOrUrl: file.path }) as any)
+        return await ragApplication.addLoader(new PptLoader({ filePathOrUrl: file.path }) as any, forceReload)
       }
+
+      if (file.ext === '.xlsx') {
+        return await ragApplication.addLoader(new ExcelLoader({ filePathOrUrl: file.path }) as any, forceReload)
+      }
+
+      if (['.md', '.mdx'].includes(file.ext)) {
+        return await ragApplication.addLoader(new MarkdownLoader({ filePathOrUrl: file.path }) as any, forceReload)
+      }
+
+      const fileContent = fs.readFileSync(file.path, 'utf-8')
+
+      return await ragApplication.addLoader(new TextLoader({ text: fileContent }), forceReload)
     }
 
     return { entriesAdded: 0, uniqueId: '', loaderType: '' }

src/main/services/WindowService.ts

@@ -1,6 +1,7 @@
 import { is } from '@electron-toolkit/utils'
 import { isLinux, isWin } from '@main/constant'
 import { app, BrowserWindow, Menu, MenuItem, shell } from 'electron'
+import Logger from 'electron-log'
 import windowStateKeeper from 'electron-window-state'
 import { join } from 'path'
 
@@ -132,7 +133,16 @@ export class WindowService {
     })
 
     mainWindow.webContents.setWindowOpenHandler((details) => {
-      shell.openExternal(details.url)
+      const { url } = details
+
+      if (url.includes('http://file/')) {
+        const fileUrl = url.replace('http://file/', '')
+        const filePath = decodeURIComponent(fileUrl)
+        shell.openPath(filePath).catch((err) => Logger.error('Failed to open file:', err))
+      } else {
+        shell.openExternal(details.url)
+      }
+
       return { action: 'deny' }
     })
 

src/preload/index.d.ts

@@ -42,6 +42,7 @@ declare global {
         create: (fileName: string) => Promise<string>
         write: (filePath: string, data: Uint8Array | string) => Promise<void>
         open: (options?: OpenDialogOptions) => Promise<{ fileName: string; filePath: string; content: Buffer } | null>
+        openPath: (path: string) => Promise<void>
         save: (
           path: string,
           content: string | NodeJS.ArrayBufferView,
@@ -63,7 +64,15 @@ declare global {
         create: ({ id, model, apiKey, baseURL }: KnowledgeBaseParams) => Promise<void>
         reset: ({ base }: { base: KnowledgeBaseParams }) => Promise<void>
         delete: (id: string) => Promise<void>
-        add: ({ base, item }: { base: KnowledgeBaseParams; item: KnowledgeItem }) => Promise<AddLoaderReturn>
+        add: ({
+          base,
+          item,
+          forceReload = false
+        }: {
+          base: KnowledgeBaseParams
+          item: KnowledgeItem
+          forceReload?: boolean
+        }) => Promise<AddLoaderReturn>
         remove: ({ uniqueId, base }: { uniqueId: string; base: KnowledgeBaseParams }) => Promise<void>
         search: ({ search, base }: { search: string; base: KnowledgeBaseParams }) => Promise<ExtractChunkData[]>
       }

src/preload/index.ts

@@ -36,6 +36,7 @@ const api = {
     create: (fileName: string) => ipcRenderer.invoke('file:create', fileName),
     write: (filePath: string, data: Uint8Array | string) => ipcRenderer.invoke('file:write', filePath, data),
     open: (options?: { decompress: boolean }) => ipcRenderer.invoke('file:open', options),
+    openPath: (path: string) => ipcRenderer.invoke('file:openPath', path),
     save: (path: string, content: string, options?: { compress: boolean }) =>
       ipcRenderer.invoke('file:save', path, content, options),
     selectFolder: () => ipcRenderer.invoke('file:selectFolder'),
@@ -56,8 +57,15 @@ const api = {
       ipcRenderer.invoke('knowledge-base:create', { id, model, apiKey, baseURL }),
     reset: ({ base }: { base: KnowledgeBaseParams }) => ipcRenderer.invoke('knowledge-base:reset', { base }),
     delete: (id: string) => ipcRenderer.invoke('knowledge-base:delete', id),
-    add: ({ base, item }: { base: KnowledgeBaseParams; item: KnowledgeItem }) =>
+    add: ({
-      ipcRenderer.invoke('knowledge-base:add', { base, item }),
+      base,
+      item,
+      forceReload = false
+    }: {
+      base: KnowledgeBaseParams
+      item: KnowledgeItem
+      forceReload?: boolean
+    }) => ipcRenderer.invoke('knowledge-base:add', { base, item, forceReload }),
     remove: ({ uniqueId, base }: { uniqueId: string; base: KnowledgeBaseParams }) =>
       ipcRenderer.invoke('knowledge-base:remove', { uniqueId, base }),
     search: ({ search, base }: { search: string; base: KnowledgeBaseParams }) =>

src/renderer/src/assets/styles/index.scss

@@ -42,6 +42,7 @@
   --color-active: rgba(55, 55, 55, 1);
   --color-frame-border: #333;
   --color-group-background: var(--color-background-soft);
+  --color-reference-background: #0b0e12;
 
   --navbar-background-mac: rgba(30, 30, 30, 0.6);
   --navbar-background: rgba(30, 30, 30);
@@ -99,6 +100,7 @@ body[theme-mode='light'] {
   --color-active: var(--color-white-soft);
   --color-frame-border: #ddd;
   --color-group-background: var(--color-white);
+  --color-reference-background: #f1f7ff;
 
   --navbar-background-mac: rgba(255, 255, 255, 0.6);
   --navbar-background: rgba(255, 255, 255);

src/renderer/src/assets/styles/markdown.scss

@@ -229,11 +229,24 @@
 
   .footnotes {
     margin-top: 1em;
+    margin-bottom: 1em;
     padding-top: 1em;
-    border-top: 1px solid var(--color-border);
+
+    background-color: var(--color-reference-background);
+    border-radius: 8px;
+    padding: 8px 12px;
+
+    h4 {
+      margin-bottom: 5px;
+      font-size: 12px;
+    }
 
     ol {
       padding-left: 1em;
+      margin: 0;
+      li:last-child {
+        margin-bottom: 0;
+      }
     }
 
     li {

src/renderer/src/components/AssistantSettings/AssistantModelSettings.tsx

@@ -241,7 +241,7 @@ const AssistantModelSettings: FC<Props> = ({ assistant, updateAssistant, updateA
       )}
       <Divider style={{ margin: '10px 0' }} />
       <SettingRow style={{ minHeight: 30 }}>
-        <Label>{t('model.stream_output')}</Label>
+        <Label>{t('models.stream_output')}</Label>
         <Switch
           checked={streamOutput}
           onChange={(checked) => {

src/renderer/src/components/Popups/SelectModelPopup.tsx

@@ -1,7 +1,7 @@
 import { PushpinOutlined, SearchOutlined } from '@ant-design/icons'
 import VisionIcon from '@renderer/components/Icons/VisionIcon'
 import { TopView } from '@renderer/components/TopView'
-import { getModelLogo, isVisionModel, isWebSearchModel } from '@renderer/config/models'
+import { getModelLogo, isEmbeddingModel, isVisionModel, isWebSearchModel } from '@renderer/config/models'
 import db from '@renderer/databases'
 import { useProviders } from '@renderer/hooks/useProvider'
 import { getModelUniqId } from '@renderer/services/ModelService'
@@ -66,6 +66,7 @@ const PopupContainer: React.FC<PopupContainerProps> = ({ model, resolve }) => {
     .filter((p) => p.models && p.models.length > 0)
     .map((p) => {
       const filteredModels = sortBy(p.models, ['group', 'name'])
+        .filter((m) => !isEmbeddingModel(m))
         .filter((m) =>
           [m.name + m.provider + t('provider.' + p.id)].join('').toLowerCase().includes(searchText.toLowerCase())
         )
@@ -142,7 +143,7 @@ const PopupContainer: React.FC<PopupContainerProps> = ({ model, resolve }) => {
     if (pinnedItems.length > 0) {
       filteredItems.unshift({
         key: 'pinned',
-        label: t('model.pinned'),
+        label: t('models.pinned'),
         type: 'group',
         children: pinnedItems
       } as MenuItem)
@@ -188,7 +189,7 @@ const PopupContainer: React.FC<PopupContainerProps> = ({ model, resolve }) => {
             </SearchIcon>
           }
           ref={inputRef}
-          placeholder={t('model.search')}
+          placeholder={t('models.search')}
           value={searchText}
           onChange={(e) => setSearchText(e.target.value)}
           allowClear

src/renderer/src/config/minapps.ts

@@ -221,7 +221,8 @@ const _apps: MinAppType[] = [
     id: 'thinkany',
     name: 'ThinkAny',
     logo: ThinkAnyLogo,
-    url: 'https://thinkany.ai/'
+    url: 'https://thinkany.ai/',
+    bodered: true
   }
 ]
 

src/renderer/src/config/models.ts

@@ -151,9 +151,9 @@ export const VISION_REGEX = new RegExp(
   'i'
 )
 
-const TEXT_TO_IMAGE_REGEX = /flux|diffusion|stabilityai|sd-|dall|cogview/i
+export const TEXT_TO_IMAGE_REGEX = /flux|diffusion|stabilityai|sd-|dall|cogview/i
-const EMBEDDING_REGEX = /(?:^text-|embed|rerank|davinci|babbage|bge-|base|retrieval|uae-)/i
+export const EMBEDDING_REGEX = /(?:^text-|embed|rerank|davinci|babbage|bge-|base|retrieval|uae-|gte-)/i
-const NOT_SUPPORTED_REGEX = /(?:^tts|rerank|whisper|speech)/i
+export const NOT_SUPPORTED_REGEX = /(?:^tts|rerank|whisper|speech)/i
 
 export function getModelLogo(modelId: string) {
   const isLight = true
@@ -1047,18 +1047,38 @@ export function isTextToImageModel(model: Model): boolean {
 }
 
 export function isEmbeddingModel(model: Model): boolean {
-  return EMBEDDING_REGEX.test(model.id)
+  if (!model) {
+    return false
+  }
+
+  if (['anthropic'].includes(model?.provider)) {
+    return false
+  }
+
+  return EMBEDDING_REGEX.test(model.id) || model.type?.includes('embedding') || false
 }
 
 export function isVisionModel(model: Model): boolean {
+  if (!model) {
+    return false
+  }
+
   return VISION_REGEX.test(model.id) || model.type?.includes('vision') || false
 }
 
 export function isSupportedModel(model: OpenAI.Models.Model): boolean {
+  if (!model) {
+    return false
+  }
+
   return !NOT_SUPPORTED_REGEX.test(model.id)
 }
 
 export function isWebSearchModel(model: Model): boolean {
+  if (!model) {
+    return false
+  }
+
   const provider = getProviderByModel(model)
 
   if (!provider) {

src/renderer/src/config/prompts.ts

@@ -49,3 +49,30 @@ export const SUMMARIZE_PROMPT =
 
 export const TRANSLATE_PROMPT =
   'You are a translation expert. Translate from input language to {{target_language}}, provide the translation result directly without any explanation and keep original format. Do not translate if the target language is the same as the source language.'
+
+export const REFERENCE_PROMPT = `请根据参考资料回答问题，并使用脚注格式引用数据来源。参考资料可能和问题无关，请忽略无关的参考资料。
+
+## 脚注格式：
+
+1. **脚注标记**：在正文中使用 [^数字] 的形式标记脚注，例如 [^1]。
+2. **脚注内容**：在文档末尾使用 [^数字]: 脚注内容 的形式定义脚注的具体内容。
+
+## 脚注示例和要求：
+
+1. type 为 file 时：[^1]: [__name__](http://file/__url__)
+2. type 为 directory 时：[^1]: [__name__](http://file/__url__)
+3. type 为 url,sitemap 时：[^1]: [__name__](__url__)
+4. type 为 note 时：[^1]: __note__
+
+__url__ 替换成参考资料的 url
+__name__ 请根据参考资料的 url 进行解析和替换
+__note__ 请根据参考资料的 content 进行总结和替换
+
+## 我的问题是：
+
+{question}
+
+## 参考资料：
+
+{references}
+`

src/renderer/src/config/providers.ts

@@ -355,11 +355,11 @@ export const PROVIDER_CONFIG = {
   },
   aihubmix: {
     api: {
-      url: 'https://aihubmix.com'
+      url: 'https://aihubmix.com?aff=SJyh'
     },
     websites: {
       official: 'https://aihubmix.com/',
-      apiKey: 'https://aihubmix.com/token',
+      apiKey: 'https://aihubmix.com?aff=SJyh',
       docs: 'https://doc.aihubmix.com/',
       models: 'https://aihubmix.com/models'
     }

src/renderer/src/hooks/useKnowledge.ts

@@ -0,0 +1,279 @@
+/* eslint-disable react-hooks/rules-of-hooks */
+import { db } from '@renderer/databases/index'
+import KnowledgeQueue from '@renderer/queue/KnowledgeQueue'
+import FileManager from '@renderer/services/FileManager'
+import { getKnowledgeBaseParams } from '@renderer/services/KnowledgeService'
+import { RootState } from '@renderer/store'
+import {
+  addBase,
+  addFiles as addFilesAction,
+  addItem,
+  clearAllProcessing,
+  clearCompletedProcessing,
+  deleteBase,
+  removeItem as removeItemAction,
+  renameBase,
+  updateBase,
+  updateBases,
+  updateItemProcessingStatus,
+  updateNotes
+} from '@renderer/store/knowledge'
+import { FileType, KnowledgeBase, ProcessingStatus } from '@renderer/types'
+import { KnowledgeItem } from '@renderer/types'
+import { runAsyncFunction } from '@renderer/utils'
+import { useEffect, useState } from 'react'
+import { useDispatch, useSelector } from 'react-redux'
+import { v4 as uuidv4 } from 'uuid'
+
+export const useKnowledge = (baseId: string) => {
+  const dispatch = useDispatch()
+  const base = useSelector((state: RootState) => state.knowledge.bases.find((b) => b.id === baseId))
+
+  // 重命名知识库
+  const renameKnowledgeBase = (name: string) => {
+    dispatch(renameBase({ baseId, name }))
+  }
+
+  // 更新知识库
+  const updateKnowledgeBase = (base: KnowledgeBase) => {
+    dispatch(updateBase(base))
+  }
+
+  // 批量添加文件
+  const addFiles = (files: FileType[]) => {
+    const filesItems: KnowledgeItem[] = files.map((file) => ({
+      id: uuidv4(),
+      type: 'file' as const,
+      content: file,
+      created_at: Date.now(),
+      updated_at: Date.now(),
+      processingStatus: 'pending',
+      processingProgress: 0,
+      processingError: '',
+      retryCount: 0
+    }))
+    dispatch(addFilesAction({ baseId, items: filesItems }))
+    setTimeout(() => KnowledgeQueue.checkAllBases(), 0)
+  }
+
+  // 添加URL
+  const addUrl = (url: string) => {
+    const newUrlItem: KnowledgeItem = {
+      id: uuidv4(),
+      type: 'url' as const,
+      content: url,
+      created_at: Date.now(),
+      updated_at: Date.now(),
+      processingStatus: 'pending',
+      processingProgress: 0,
+      processingError: '',
+      retryCount: 0
+    }
+    dispatch(addItem({ baseId, item: newUrlItem }))
+    setTimeout(() => KnowledgeQueue.checkAllBases(), 0)
+  }
+
+  // 添加笔记
+  const addNote = async (content: string) => {
+    const noteId = uuidv4()
+    const note: KnowledgeItem = {
+      id: noteId,
+      type: 'note',
+      content,
+      created_at: Date.now(),
+      updated_at: Date.now()
+    }
+
+    // 存储完整笔记到数据库
+    await db.knowledge_notes.add(note)
+
+    // 在 store 中只存储引用
+    const noteRef: KnowledgeItem = {
+      id: noteId,
+      baseId,
+      type: 'note',
+      content: '', // store中不需要存储实际内容
+      created_at: Date.now(),
+      updated_at: Date.now(),
+      processingStatus: 'pending',
+      processingProgress: 0,
+      processingError: '',
+      retryCount: 0
+    }
+
+    dispatch(updateNotes({ baseId, item: noteRef }))
+    setTimeout(() => KnowledgeQueue.checkAllBases(), 0)
+  }
+
+  // 更新笔记内容
+  const updateNoteContent = async (noteId: string, content: string) => {
+    const note = await db.knowledge_notes.get(noteId)
+    if (note) {
+      const updatedNote = {
+        ...note,
+        content,
+        updated_at: Date.now()
+      }
+      await db.knowledge_notes.put(updatedNote)
+      dispatch(updateNotes({ baseId, item: updatedNote }))
+    }
+    setTimeout(() => KnowledgeQueue.checkAllBases(), 0)
+  }
+
+  // 获取笔记内容
+  const getNoteContent = async (noteId: string) => {
+    return await db.knowledge_notes.get(noteId)
+  }
+
+  // 移除项目
+  const removeItem = async (item: KnowledgeItem) => {
+    dispatch(removeItemAction({ baseId, item }))
+    if (base) {
+      if (item?.uniqueId) {
+        await window.api.knowledgeBase.remove({ uniqueId: item.uniqueId, base: getKnowledgeBaseParams(base) })
+      }
+      if (item.type === 'file' && typeof item.content === 'object') {
+        await FileManager.deleteFile(item.content.id)
+      }
+    }
+  }
+
+  // 更新处理状态
+  const updateItemStatus = (itemId: string, status: ProcessingStatus, progress?: number, error?: string) => {
+    dispatch(
+      updateItemProcessingStatus({
+        baseId,
+        itemId,
+        status,
+        progress,
+        error
+      })
+    )
+  }
+
+  // 获取特定项目的处理状态
+  const getProcessingStatus = (itemId: string) => {
+    return base?.items.find((item) => item.id === itemId)?.processingStatus
+  }
+
+  // 获取特定类型的所有处理项
+  const getProcessingItemsByType = (type: 'file' | 'url' | 'note') => {
+    return base?.items.filter((item) => item.type === type && item.processingStatus !== undefined) || []
+  }
+
+  // 清除已完成的项目
+  const clearCompleted = () => {
+    dispatch(clearCompletedProcessing({ baseId }))
+  }
+
+  // 清除所有处理状态
+  const clearAll = () => {
+    dispatch(clearAllProcessing({ baseId }))
+  }
+
+  // 添加 Sitemap
+  const addSitemap = (url: string) => {
+    const newSitemapItem: KnowledgeItem = {
+      id: uuidv4(),
+      type: 'sitemap' as const,
+      content: url,
+      created_at: Date.now(),
+      updated_at: Date.now(),
+      processingStatus: 'pending',
+      processingProgress: 0,
+      processingError: '',
+      retryCount: 0
+    }
+    dispatch(addItem({ baseId, item: newSitemapItem }))
+    setTimeout(() => KnowledgeQueue.checkAllBases(), 0)
+  }
+
+  // Add directory support
+  const addDirectory = (path: string) => {
+    const newDirectoryItem: KnowledgeItem = {
+      id: uuidv4(),
+      type: 'directory',
+      content: path,
+      created_at: Date.now(),
+      updated_at: Date.now(),
+      processingStatus: 'pending',
+      processingProgress: 0,
+      processingError: '',
+      retryCount: 0
+    }
+    dispatch(addItem({ baseId, item: newDirectoryItem }))
+    setTimeout(() => KnowledgeQueue.checkAllBases(), 0)
+  }
+
+  const fileItems = base?.items.filter((item) => item.type === 'file') || []
+  const directoryItems = base?.items.filter((item) => item.type === 'directory') || []
+  const urlItems = base?.items.filter((item) => item.type === 'url') || []
+  const sitemapItems = base?.items.filter((item) => item.type === 'sitemap') || []
+  const [noteItems, setNoteItems] = useState<KnowledgeItem[]>([])
+
+  useEffect(() => {
+    const notes = base?.items.filter((item) => item.type === 'note') || []
+    runAsyncFunction(async () => {
+      const newNoteItems = await Promise.all(
+        notes.map(async (item) => {
+          const note = await db.knowledge_notes.get(item.id)
+          return { ...item, content: note?.content || '' }
+        })
+      )
+      setNoteItems(newNoteItems.filter((note) => note !== undefined) as KnowledgeItem[])
+    })
+  }, [base?.items])
+
+  return {
+    base,
+    fileItems,
+    urlItems,
+    sitemapItems,
+    noteItems,
+    renameKnowledgeBase,
+    updateKnowledgeBase,
+    addFiles,
+    addUrl,
+    addSitemap,
+    addNote,
+    updateNoteContent,
+    getNoteContent,
+    updateItemStatus,
+    getProcessingStatus,
+    getProcessingItemsByType,
+    clearCompleted,
+    clearAll,
+    removeItem,
+    directoryItems,
+    addDirectory
+  }
+}
+
+export const useKnowledgeBases = () => {
+  const dispatch = useDispatch()
+  const bases = useSelector((state: RootState) => state.knowledge.bases)
+
+  const addKnowledgeBase = (base: KnowledgeBase) => {
+    dispatch(addBase(base))
+  }
+
+  const renameKnowledgeBase = (baseId: string, name: string) => {
+    dispatch(renameBase({ baseId, name }))
+  }
+
+  const deleteKnowledgeBase = (baseId: string) => {
+    dispatch(deleteBase({ baseId }))
+  }
+
+  const updateKnowledgeBases = (bases: KnowledgeBase[]) => {
+    dispatch(updateBases(bases))
+  }
+
+  return {
+    bases,
+    addKnowledgeBase,
+    renameKnowledgeBase,
+    deleteKnowledgeBase,
+    updateKnowledgeBases
+  }
+}

src/renderer/src/hooks/useShortcuts.ts

@@ -37,8 +37,10 @@ export const useShortcut = (
 
   const shortcutConfig = shortcuts.find((s) => s.key === shortcutKey)
 
+  console.log(shortcutConfig)
+
   useHotkeys(
-    shortcutConfig?.enabled ? formatShortcut(shortcutConfig.shortcut) : '',
+    shortcutConfig?.enabled ? formatShortcut(shortcutConfig.shortcut) : 'none',
     (e) => {
       if (options.preventDefault) {
         e.preventDefault()
@@ -49,7 +51,8 @@ export const useShortcut = (
     },
     {
       enableOnFormTags: options.enableOnFormTags,
-      description: options.description || shortcutConfig?.key
+      description: options.description || shortcutConfig?.key,
+      enabled: !!shortcutConfig?.enabled
     }
   )
 }

src/renderer/src/hooks/useknowledge.ts

@@ -6,6 +6,7 @@ import { getKnowledgeBaseParams } from '@renderer/services/KnowledgeService'
 import { RootState } from '@renderer/store'
 import {
   addBase,
+  addFiles as addFilesAction,
   addItem,
   clearAllProcessing,
   clearCompletedProcessing,
@@ -13,7 +14,7 @@ import {
   removeItem as removeItemAction,
   renameBase,
   updateBase,
-  updateFiles as updateFilesAction,
+  updateBases,
   updateItemProcessingStatus,
   updateNotes
 } from '@renderer/store/knowledge'
@@ -38,22 +39,20 @@ export const useKnowledge = (baseId: string) => {
     dispatch(updateBase(base))
   }
 
-  // 添加文件列表
+  // 批量添加文件
   const addFiles = (files: FileType[]) => {
-    for (const file of files) {
+    const filesItems: KnowledgeItem[] = files.map((file) => ({
-      const newItem: KnowledgeItem = {
+      id: uuidv4(),
-        id: uuidv4(),
+      type: 'file' as const,
-        type: 'file' as const,
+      content: file,
-        content: file,
+      created_at: Date.now(),
-        created_at: Date.now(),
+      updated_at: Date.now(),
-        updated_at: Date.now(),
+      processingStatus: 'pending',
-        processingStatus: 'pending',
+      processingProgress: 0,
-        processingProgress: 0,
+      processingError: '',
-        processingError: '',
+      retryCount: 0
-        retryCount: 0
+    }))
-      }
+    dispatch(addFilesAction({ baseId, items: filesItems }))
-      dispatch(addItem({ baseId, item: newItem }))
-    }
     setTimeout(() => KnowledgeQueue.checkAllBases(), 0)
   }
 
@@ -106,19 +105,6 @@ export const useKnowledge = (baseId: string) => {
     setTimeout(() => KnowledgeQueue.checkAllBases(), 0)
   }
 
-  // 更新文件列表
-  const updateFiles = (files: FileType[]) => {
-    const newItems = files.map((file) => ({
-      id: uuidv4(),
-      type: 'file' as const,
-      content: file,
-      created_at: Date.now(),
-      updated_at: Date.now()
-    }))
-    dispatch(updateFilesAction({ baseId, items: newItems }))
-    setTimeout(() => KnowledgeQueue.checkAllBases(), 0)
-  }
-
   // 更新笔记内容
   const updateNoteContent = async (noteId: string, content: string) => {
     const note = await db.knowledge_notes.get(noteId)
@@ -202,7 +188,25 @@ export const useKnowledge = (baseId: string) => {
     setTimeout(() => KnowledgeQueue.checkAllBases(), 0)
   }
 
+  // Add directory support
+  const addDirectory = (path: string) => {
+    const newDirectoryItem: KnowledgeItem = {
+      id: uuidv4(),
+      type: 'directory',
+      content: path,
+      created_at: Date.now(),
+      updated_at: Date.now(),
+      processingStatus: 'pending',
+      processingProgress: 0,
+      processingError: '',
+      retryCount: 0
+    }
+    dispatch(addItem({ baseId, item: newDirectoryItem }))
+    setTimeout(() => KnowledgeQueue.checkAllBases(), 0)
+  }
+
   const fileItems = base?.items.filter((item) => item.type === 'file') || []
+  const directoryItems = base?.items.filter((item) => item.type === 'directory') || []
   const urlItems = base?.items.filter((item) => item.type === 'url') || []
   const sitemapItems = base?.items.filter((item) => item.type === 'sitemap') || []
   const [noteItems, setNoteItems] = useState<KnowledgeItem[]>([])
@@ -232,7 +236,6 @@ export const useKnowledge = (baseId: string) => {
     addUrl,
     addSitemap,
     addNote,
-    updateFiles,
     updateNoteContent,
     getNoteContent,
     updateItemStatus,
@@ -240,7 +243,9 @@ export const useKnowledge = (baseId: string) => {
     getProcessingItemsByType,
     clearCompleted,
     clearAll,
-    removeItem
+    removeItem,
+    directoryItems,
+    addDirectory
   }
 }
 
@@ -260,10 +265,15 @@ export const useKnowledgeBases = () => {
     dispatch(deleteBase({ baseId }))
   }
 
+  const updateKnowledgeBases = (bases: KnowledgeBase[]) => {
+    dispatch(updateBases(bases))
+  }
+
   return {
     bases,
     addKnowledgeBase,
     renameKnowledgeBase,
-    deleteKnowledgeBase
+    deleteKnowledgeBase,
+    updateKnowledgeBases
   }
 }

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

@@ -252,16 +252,6 @@
     "minapp": {
       "title": "MinApp"
     },
-    "model": {
-      "pinned": "Pinned",
-      "search": "Search models...",
-      "stream_output": "Stream output",
-      "type": {
-        "select": "Select Model Types",
-        "text": "Text",
-        "vision": "Vision"
-      }
-    },
     "ollama": {
       "keep_alive_time.description": "The time in minutes to keep the connection alive, default is 5 minutes.",
       "keep_alive_time.placeholder": "Minutes",
@@ -532,7 +522,7 @@
       "search": "Search knowledge base",
       "empty": "No knowledge base found",
       "drag_file": "Drag file here",
-      "file_hint": "Support pdf, docx, txt and md",
+      "file_hint": "Support {{file_types}}",
       "add": {
         "title": "Add Knowledge Base"
       },
@@ -562,7 +552,26 @@
       "delete_confirm": "Are you sure you want to delete this knowledge base?",
       "sitemaps": "Websites",
       "add_sitemap": "Website Map",
-      "sitemap_placeholder": "Enter Website Map URL"
+      "sitemap_placeholder": "Enter Website Map URL",
+      "directories": "Directories",
+      "add_directory": "Add Directory",
+      "directory_placeholder": "Enter Directory Path"
+    },
+    "models": {
+      "pinned": "Pinned",
+      "search": "Search models...",
+      "stream_output": "Stream output",
+      "type": {
+        "select": "Select Model Types",
+        "text": "Text",
+        "vision": "Vision",
+        "embedding": "Embedding"
+      },
+      "all": "All",
+      "vision": "Vision",
+      "websearch": "WebSearch",
+      "free": "Free",
+      "embedding": "Embedding"
     }
   }
 }

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

@@ -252,16 +252,6 @@
     "minapp": {
       "title": "Встроенные приложения"
     },
-    "model": {
-      "pinned": "Закреплено",
-      "search": "Поиск моделей...",
-      "stream_output": "Потоковый вывод",
-      "type": {
-        "select": "Выберите тип модели",
-        "text": "Текст",
-        "vision": "Изображение"
-      }
-    },
     "ollama": {
       "keep_alive_time.description": "Время в минутах, в течение которого модель остается активной, по умолчанию 5 минут.",
       "keep_alive_time.placeholder": "Минуты",
@@ -532,7 +522,7 @@
       "search": "Поиск в базе знаний",
       "empty": "База знаний не найдена",
       "drag_file": "Перетащите файл сюда",
-      "file_hint": "Поддерживаются pdf, docx, txt и md",
+      "file_hint": "Поддерживаются {{file_types}}",
       "add": {
         "title": "Добавить базу знаний"
       },
@@ -562,7 +552,26 @@
       "delete_confirm": "Вы уверены, что хотите удалить эту базу знаний?",
       "sitemaps": "Сайты",
       "add_sitemap": "Карта сайта",
-      "sitemap_placeholder": "Введите URL карты сайта"
+      "sitemap_placeholder": "Введите URL карты сайта",
+      "directories": "Директории",
+      "add_directory": "Добавить директорию",
+      "directory_placeholder": "Введите путь к директории"
+    },
+    "models": {
+      "pinned": "Закреплено",
+      "search": "Поиск моделей...",
+      "stream_output": "Потоковый вывод",
+      "type": {
+        "select": "Выберите тип модели",
+        "text": "Текст",
+        "vision": "Изображение",
+        "embedding": "Встраиваемые"
+      },
+      "all": "Все",
+      "vision": "Визуальные модели",
+      "websearch": "Веб-поисковые модели",
+      "free": "Бесплатные модели",
+      "embedding": "Встраиваемые модели"
     }
   }
 }

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

@@ -253,16 +253,6 @@
     "minapp": {
       "title": "小程序"
     },
-    "model": {
-      "pinned": "已固定",
-      "search": "搜索模型...",
-      "stream_output": "流式输出",
-      "type": {
-        "select": "选择模型类型",
-        "text": "文本",
-        "vision": "图像"
-      }
-    },
     "ollama": {
       "keep_alive_time.description": "对话后模型在内存中保持的时间（默认：5分钟）",
       "keep_alive_time.placeholder": "分钟",
@@ -521,7 +511,7 @@
       "search": "搜索知识库",
       "empty": "暂无知识库",
       "drag_file": "拖拽文件到这里",
-      "file_hint": "支持 pdf, docx, txt, md 格式",
+      "file_hint": "支持 {{file_types}} 格式",
       "add": {
         "title": "添加知识库"
       },
@@ -551,7 +541,26 @@
       "delete_confirm": "确定要删除此知识库吗?",
       "sitemaps": "网站",
       "add_sitemap": "站点地图",
-      "sitemap_placeholder": "请输入站点地图 URL"
+      "sitemap_placeholder": "请输入站点地图 URL",
+      "directories": "目录",
+      "add_directory": "添加目录",
+      "directory_placeholder": "请输入目录路径"
+    },
+    "models": {
+      "pinned": "已固定",
+      "search": "搜索模型...",
+      "stream_output": "流式输出",
+      "type": {
+        "select": "选择模型类型",
+        "text": "文本",
+        "vision": "图像",
+        "embedding": "嵌入"
+      },
+      "all": "全部",
+      "vision": "视觉模型",
+      "websearch": "网络搜索模型",
+      "free": "免费模型",
+      "embedding": "嵌入模型"
     }
   }
 }

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

@@ -252,16 +252,6 @@
     "minapp": {
       "title": "小程序"
     },
-    "model": {
-      "pinned": "已固定",
-      "search": "搜尋模型...",
-      "stream_output": "串流輸出",
-      "type": {
-        "select": "選擇模型類型",
-        "text": "文字",
-        "vision": "圖像"
-      }
-    },
     "ollama": {
       "keep_alive_time.description": "對話後模型在記憶體中保持的時間（預設為 5 分鐘）。",
       "keep_alive_time.placeholder": "分鐘",
@@ -520,7 +510,7 @@
       "search": "搜尋知識庫",
       "empty": "暫無知識庫",
       "drag_file": "拖拽文件到這裡",
-      "file_hint": "支持 pdf, docx, txt, md 格式",
+      "file_hint": "支持 {{file_types}} 格式",
       "add": {
         "title": "添加知識庫"
       },
@@ -550,7 +540,26 @@
       "delete_confirm": "確定要刪除此知識庫嗎?",
       "sitemaps": "網站",
       "add_sitemap": "網站地圖",
-      "sitemap_placeholder": "請輸入網站地圖 URL"
+      "sitemap_placeholder": "請輸入網站地圖 URL",
+      "directories": "目錄",
+      "add_directory": "添加目錄",
+      "directory_placeholder": "請輸入目錄路徑"
+    },
+    "models": {
+      "pinned": "已固定",
+      "search": "搜尋模型...",
+      "stream_output": "串流輸出",
+      "type": {
+        "select": "選擇模型類型",
+        "text": "文字",
+        "vision": "圖像",
+        "embedding": "嵌入"
+      },
+      "all": "全部",
+      "vision": "視覺模型",
+      "websearch": "網路搜索模型",
+      "free": "免費模型",
+      "embedding": "嵌入模型"
     }
   }
 }

src/renderer/src/pages/home/Navbar.tsx

@@ -1,4 +1,4 @@
-import { SearchOutlined } from '@ant-design/icons'
+import { FormOutlined, SearchOutlined } from '@ant-design/icons'
 import { Navbar, NavbarLeft, NavbarRight } from '@renderer/components/app/Navbar'
 import AssistantSettingsPopup from '@renderer/components/AssistantSettings'
 import { HStack } from '@renderer/components/Layout'
@@ -47,8 +47,8 @@ const HeaderNavbar: FC<Props> = ({ activeAssistant }) => {
           <NavbarIcon onClick={toggleShowAssistants} style={{ marginLeft: isMac ? 8 : 0 }}>
             <i className="iconfont icon-hide-sidebar" />
           </NavbarIcon>
-          <NavbarIcon onClick={() => SearchPopup.show()}>
+          <NavbarIcon onClick={() => EventEmitter.emit(EVENT_NAMES.ADD_NEW_TOPIC)}>
-            <SearchOutlined />
+            <FormOutlined />
           </NavbarIcon>
         </NavbarLeft>
       )}
@@ -70,6 +70,9 @@ const HeaderNavbar: FC<Props> = ({ activeAssistant }) => {
           <SelectModelButton assistant={assistant} />
         </HStack>
         <HStack alignItems="center">
+          <NavbarIcon onClick={() => SearchPopup.show()}>
+            <SearchOutlined />
+          </NavbarIcon>
           <AppStorePopover>
             <NavbarIcon>
               <i className="iconfont icon-appstore" />

src/renderer/src/pages/home/Tabs/SettingsTab.tsx

@@ -162,7 +162,7 @@ const SettingsTab: FC<Props> = (props) => {
           </Col>
         </Row>
         <SettingRow>
-          <SettingRowTitleSmall>{t('model.stream_output')}</SettingRowTitleSmall>
+          <SettingRowTitleSmall>{t('models.stream_output')}</SettingRowTitleSmall>
           <Switch
             size="small"
             checked={streamOutput}

src/renderer/src/pages/knowledge/KnowledgeContent.tsx

@@ -2,6 +2,7 @@ import {
   DeleteOutlined,
   EditOutlined,
   FileTextOutlined,
+  FolderOutlined,
   GlobalOutlined,
   LinkOutlined,
   PlusOutlined,
@@ -10,7 +11,7 @@ import {
 import PromptPopup from '@renderer/components/Popups/PromptPopup'
 import TextEditPopup from '@renderer/components/Popups/TextEditPopup'
 import Scrollbar from '@renderer/components/Scrollbar'
-import { useKnowledge } from '@renderer/hooks/useknowledge'
+import { useKnowledge } from '@renderer/hooks/useKnowledge'
 import FileManager from '@renderer/services/FileManager'
 import { FileType, FileTypes, KnowledgeBase } from '@renderer/types'
 import { Button, Card, message, Typography, Upload } from 'antd'
@@ -28,6 +29,32 @@ interface KnowledgeContentProps {
   selectedBase: KnowledgeBase
 }
 
+const fileTypes = ['.pdf', '.docx', '.pptx', '.xlsx', '.txt', '.md', '.mdx']
+
+const FlexColumn = styled.div`
+  display: flex;
+  flex-direction: column;
+  gap: 8px;
+`
+
+const FlexAlignCenter = styled.div`
+  display: flex;
+  align-items: center;
+  gap: 16px;
+`
+
+const ClickableSpan = styled.span`
+  cursor: pointer;
+`
+
+const FileIcon = styled(FileTextOutlined)`
+  font-size: 16px;
+`
+
+const BottomSpacer = styled.div`
+  min-height: 20px;
+`
+
 const KnowledgeContent: FC<KnowledgeContentProps> = ({ selectedBase }) => {
   const { t } = useTranslation()
   const {
@@ -36,13 +63,15 @@ const KnowledgeContent: FC<KnowledgeContentProps> = ({ selectedBase }) => {
     fileItems,
     urlItems,
     sitemapItems,
+    directoryItems,
     addFiles,
     updateNoteContent,
     addUrl,
     addSitemap,
     removeItem,
     getProcessingStatus,
-    addNote
+    addNote,
+    addDirectory
   } = useKnowledge(selectedBase.id || '')
 
   if (!base) {
@@ -53,12 +82,10 @@ const KnowledgeContent: FC<KnowledgeContentProps> = ({ selectedBase }) => {
     const input = document.createElement('input')
     input.type = 'file'
     input.multiple = true
-    input.accept = '.pdf,.txt,.docx,.md'
+    input.accept = fileTypes.join(',')
     input.onchange = (e) => {
       const files = (e.target as HTMLInputElement).files
-      if (files) {
+      files && handleDrop(Array.from(files))
-        handleDrop(Array.from(files))
-      }
     }
     input.click()
   }
@@ -142,6 +169,12 @@ const KnowledgeContent: FC<KnowledgeContentProps> = ({ selectedBase }) => {
     editedText && updateNoteContent(note.id, editedText)
   }
 
+  const handleAddDirectory = async () => {
+    const path = await window.api.file.selectFolder()
+    console.log('[KnowledgeContent] Selected directory:', path)
+    path && addDirectory(path)
+  }
+
   return (
     <MainContent>
       <FileSection>
@@ -155,10 +188,12 @@ const KnowledgeContent: FC<KnowledgeContentProps> = ({ selectedBase }) => {
           showUploadList={false}
           customRequest={({ file }) => handleDrop([file as File])}
           multiple={true}
-          accept=".pdf,.txt,.docx,.md"
+          accept={fileTypes.join(',')}
           style={{ marginTop: 10, background: 'transparent' }}>
           <p className="ant-upload-text">{t('knowledge_base.drag_file')}</p>
-          <p className="ant-upload-hint">{t('knowledge_base.file_hint')}</p>
+          <p className="ant-upload-hint">
+            {t('knowledge_base.file_hint', { file_types: fileTypes.join(', ').replaceAll('.', '') })}
+          </p>
         </Dragger>
       </FileSection>
 
@@ -169,19 +204,46 @@ const KnowledgeContent: FC<KnowledgeContentProps> = ({ selectedBase }) => {
             <ItemCard key={item.id}>
               <ItemContent>
                 <ItemInfo>
-                  <FileTextOutlined style={{ fontSize: '16px' }} />
+                  <FileIcon />
-                  <span>{file.origin_name}</span>
+                  <ClickableSpan onClick={() => window.api.file.openPath(file.path)}>{file.origin_name}</ClickableSpan>
                 </ItemInfo>
-                <div style={{ display: 'flex', alignItems: 'center', gap: '16px' }}>
+                <FlexAlignCenter>
                   <StatusIcon sourceId={item.id} base={base} getProcessingStatus={getProcessingStatus} />
                   <Button type="text" danger onClick={() => removeItem(item)} icon={<DeleteOutlined />} />
-                </div>
+                </FlexAlignCenter>
               </ItemContent>
             </ItemCard>
           )
         })}
       </FileListSection>
 
+      <ContentSection>
+        <TitleWrapper>
+          <Title level={5}>{t('knowledge_base.directories')}</Title>
+          <Button icon={<PlusOutlined />} onClick={handleAddDirectory}>
+            {t('knowledge_base.add_directory')}
+          </Button>
+        </TitleWrapper>
+        <FlexColumn>
+          {directoryItems.map((item) => (
+            <ItemCard key={item.id}>
+              <ItemContent>
+                <ItemInfo>
+                  <FolderOutlined />
+                  <ClickableSpan onClick={() => window.api.file.openPath(item.content as string)}>
+                    {item.content as string}
+                  </ClickableSpan>
+                </ItemInfo>
+                <FlexAlignCenter>
+                  <StatusIcon sourceId={item.id} base={base} getProcessingStatus={getProcessingStatus} />
+                  <Button type="text" danger onClick={() => removeItem(item)} icon={<DeleteOutlined />} />
+                </FlexAlignCenter>
+              </ItemContent>
+            </ItemCard>
+          ))}
+        </FlexColumn>
+      </ContentSection>
+
       <ContentSection>
         <TitleWrapper>
           <Title level={5}>{t('knowledge_base.urls')}</Title>
@@ -189,24 +251,24 @@ const KnowledgeContent: FC<KnowledgeContentProps> = ({ selectedBase }) => {
             {t('knowledge_base.add_url')}
           </Button>
         </TitleWrapper>
-        <div style={{ display: 'flex', flexDirection: 'column', gap: '8px' }}>
+        <FlexColumn>
           {urlItems.map((item) => (
             <ItemCard key={item.id}>
               <ItemContent>
                 <ItemInfo>
-                  <LinkOutlined style={{ fontSize: '16px' }} />
+                  <LinkOutlined />
                   <a href={item.content as string} target="_blank" rel="noopener noreferrer">
                     {item.content as string}
                   </a>
                 </ItemInfo>
-                <div style={{ display: 'flex', alignItems: 'center', gap: '16px' }}>
+                <FlexAlignCenter>
                   <StatusIcon sourceId={item.id} base={base} getProcessingStatus={getProcessingStatus} />
                   <Button type="text" danger onClick={() => removeItem(item)} icon={<DeleteOutlined />} />
-                </div>
+                </FlexAlignCenter>
               </ItemContent>
             </ItemCard>
           ))}
-        </div>
+        </FlexColumn>
       </ContentSection>
 
       <ContentSection>
@@ -216,24 +278,24 @@ const KnowledgeContent: FC<KnowledgeContentProps> = ({ selectedBase }) => {
             {t('knowledge_base.add_sitemap')}
           </Button>
         </TitleWrapper>
-        <div style={{ display: 'flex', flexDirection: 'column', gap: '8px' }}>
+        <FlexColumn>
           {sitemapItems.map((item) => (
             <ItemCard key={item.id}>
               <ItemContent>
                 <ItemInfo>
-                  <GlobalOutlined style={{ fontSize: '16px' }} />
+                  <GlobalOutlined />
                   <a href={item.content as string} target="_blank" rel="noopener noreferrer">
                     {item.content as string}
                   </a>
                 </ItemInfo>
-                <div style={{ display: 'flex', alignItems: 'center', gap: '16px' }}>
+                <FlexAlignCenter>
                   <StatusIcon sourceId={item.id} base={base} getProcessingStatus={getProcessingStatus} />
                   <Button type="text" danger onClick={() => removeItem(item)} icon={<DeleteOutlined />} />
-                </div>
+                </FlexAlignCenter>
               </ItemContent>
             </ItemCard>
           ))}
-        </div>
+        </FlexColumn>
       </ContentSection>
 
       <ContentSection>
@@ -243,29 +305,31 @@ const KnowledgeContent: FC<KnowledgeContentProps> = ({ selectedBase }) => {
             {t('knowledge_base.add_note')}
           </Button>
         </TitleWrapper>
-        <div style={{ display: 'flex', flexDirection: 'column', gap: '12px' }}>
+        <FlexColumn>
           {noteItems.map((note) => (
             <ItemCard key={note.id}>
               <ItemContent>
                 <ItemInfo onClick={() => handleEditNote(note)} style={{ cursor: 'pointer' }}>
                   <span>{(note.content as string).slice(0, 50)}...</span>
                 </ItemInfo>
-                <div style={{ display: 'flex', alignItems: 'center', gap: '16px' }}>
+                <FlexAlignCenter>
                   <Button type="text" onClick={() => handleEditNote(note)} icon={<EditOutlined />} />
                   <StatusIcon sourceId={note.id} base={base} getProcessingStatus={getProcessingStatus} />
                   <Button type="text" danger onClick={() => removeItem(note)} icon={<DeleteOutlined />} />
-                </div>
+                </FlexAlignCenter>
               </ItemContent>
             </ItemCard>
           ))}
-        </div>
+        </FlexColumn>
       </ContentSection>
+
       <IndexSection>
         <Button type="primary" onClick={() => KnowledgeSearchPopup.show({ base })} icon={<SearchOutlined />}>
           {t('knowledge_base.search')}
         </Button>
       </IndexSection>
-      <div style={{ minHeight: '20px' }} />
+
+      <BottomSpacer />
     </MainContent>
   )
 }

src/renderer/src/pages/knowledge/KnowledgePage.tsx

@@ -1,9 +1,10 @@
 import { DeleteOutlined, EditOutlined, FileTextOutlined, PlusOutlined } from '@ant-design/icons'
 import { Navbar, NavbarCenter } from '@renderer/components/app/Navbar'
+import DragableList from '@renderer/components/DragableList'
 import ListItem from '@renderer/components/ListItem'
 import PromptPopup from '@renderer/components/Popups/PromptPopup'
 import Scrollbar from '@renderer/components/Scrollbar'
-import { useKnowledgeBases } from '@renderer/hooks/useknowledge'
+import { useKnowledgeBases } from '@renderer/hooks/useKnowledge'
 import { KnowledgeBase } from '@renderer/types'
 import { Dropdown, Empty, MenuProps } from 'antd'
 import { FC, useCallback, useEffect, useState } from 'react'
@@ -15,8 +16,9 @@ import KnowledgeContent from './KnowledgeContent'
 
 const KnowledgePage: FC = () => {
   const { t } = useTranslation()
-  const { bases, renameKnowledgeBase, deleteKnowledgeBase } = useKnowledgeBases()
+  const { bases, renameKnowledgeBase, deleteKnowledgeBase, updateKnowledgeBases } = useKnowledgeBases()
   const [selectedBase, setSelectedBase] = useState<KnowledgeBase>()
+  const [isDragging, setIsDragging] = useState(false)
 
   const handleAddKnowledge = async () => {
     await AddKnowledgePopup.show({ title: t('knowledge_base.add.title') })
@@ -82,24 +84,33 @@ const KnowledgePage: FC = () => {
       <ContentContainer id="content-container">
         <SideNav>
           <ScrollContainer>
-            {bases.map((base) => (
+            <DragableList
-              <Dropdown menu={{ items: getMenuItems(base) }} trigger={['contextMenu']} key={base.id}>
+              list={bases}
-                <div>
+              onUpdate={updateKnowledgeBases}
-                  <ListItem
+              style={{ marginBottom: 0, paddingBottom: isDragging ? 50 : 0 }}
-                    active={selectedBase?.id === base.id}
+              onDragStart={() => setIsDragging(true)}
-                    icon={<FileTextOutlined />}
+              onDragEnd={() => setIsDragging(false)}>
-                    title={base.name}
+              {(base) => (
-                    onClick={() => setSelectedBase(base)}
+                <Dropdown menu={{ items: getMenuItems(base) }} trigger={['contextMenu']} key={base.id}>
-                  />
+                  <div>
-                </div>
+                    <ListItem
-              </Dropdown>
+                      active={selectedBase?.id === base.id}
-            ))}
+                      icon={<FileTextOutlined />}
-            <AddKnowledgeItem onClick={handleAddKnowledge}>
+                      title={base.name}
-              <AddKnowledgeName>
+                      onClick={() => setSelectedBase(base)}
-                <PlusOutlined style={{ color: 'var(--color-text-2)', marginRight: 4 }} />
+                    />
-                {t('button.add')}
+                  </div>
-              </AddKnowledgeName>
+                </Dropdown>
-            </AddKnowledgeItem>
+              )}
+            </DragableList>
+            {!isDragging && (
+              <AddKnowledgeItem onClick={handleAddKnowledge}>
+                <AddKnowledgeName>
+                  <PlusOutlined style={{ color: 'var(--color-text-2)', marginRight: 4 }} />
+                  {t('button.add')}
+                </AddKnowledgeName>
+              </AddKnowledgeItem>
+            )}
             <div style={{ minHeight: '10px' }}></div>
           </ScrollContainer>
         </SideNav>

src/renderer/src/pages/knowledge/components/AddKnowledgePopup.tsx

@@ -1,6 +1,6 @@
 import { TopView } from '@renderer/components/TopView'
 import { isEmbeddingModel } from '@renderer/config/models'
-import { useKnowledgeBases } from '@renderer/hooks/useknowledge'
+import { useKnowledgeBases } from '@renderer/hooks/useKnowledge'
 import { useProviders } from '@renderer/hooks/useProvider'
 import { getKnowledgeBaseParams } from '@renderer/services/KnowledgeService'
 import { getModelUniqId } from '@renderer/services/ModelService'

src/renderer/src/pages/knowledge/components/KnowledgeSearchPopup.tsx

@@ -1,4 +1,4 @@
-import { ExtractChunkData } from '@llm-tools/embedjs-interfaces'
+import type { ExtractChunkData } from '@llm-tools/embedjs-interfaces'
 import { TopView } from '@renderer/components/TopView'
 import { getKnowledgeBaseParams } from '@renderer/services/KnowledgeService'
 import { KnowledgeBase } from '@renderer/types'

src/renderer/src/pages/settings/ProviderSettings/EditModelsPopup.tsx

@@ -1,12 +1,13 @@
 import { LoadingOutlined, MinusOutlined, PlusOutlined, QuestionCircleOutlined } from '@ant-design/icons'
 import VisionIcon from '@renderer/components/Icons/VisionIcon'
 import WebSearchIcon from '@renderer/components/Icons/WebSearchIcon'
+import { Center } from '@renderer/components/Layout'
 import { getModelLogo, isEmbeddingModel, isVisionModel, isWebSearchModel, SYSTEM_MODELS } from '@renderer/config/models'
 import { useProvider } from '@renderer/hooks/useProvider'
 import { fetchModels } from '@renderer/services/ApiService'
 import { Model, Provider } from '@renderer/types'
 import { getDefaultGroupName, isFreeModel, runAsyncFunction } from '@renderer/utils'
-import { Avatar, Button, Empty, Flex, Modal, Popover, Tag } from 'antd'
+import { Avatar, Button, Empty, Flex, Modal, Popover, Radio, Tag } from 'antd'
 import Search from 'antd/es/input/Search'
 import { groupBy, isEmpty, uniqBy } from 'lodash'
 import { useEffect, useState } from 'react'
@@ -29,14 +30,29 @@ const PopupContainer: React.FC<Props> = ({ provider: _provider, resolve }) => {
   const [listModels, setListModels] = useState<Model[]>([])
   const [loading, setLoading] = useState(false)
   const [searchText, setSearchText] = useState('')
-  const { t } = useTranslation()
+  const [filterType, setFilterType] = useState<string>('all')
+  const { t, i18n } = useTranslation()
 
   const systemModels = SYSTEM_MODELS[_provider.id] || []
   const allModels = uniqBy([...systemModels, ...listModels, ...models], 'id')
 
-  const list = searchText
+  const list = allModels.filter((model) => {
-    ? allModels.filter((model) => model.id.toLocaleLowerCase().includes(searchText.toLocaleLowerCase()))
+    if (searchText && !model.id.toLocaleLowerCase().includes(searchText.toLocaleLowerCase())) {
-    : allModels
+      return false
+    }
+    switch (filterType) {
+      case 'vision':
+        return isVisionModel(model)
+      case 'websearch':
+        return isWebSearchModel(model)
+      case 'free':
+        return isFreeModel(model)
+      case 'embedding':
+        return isEmbeddingModel(model)
+      default:
+        return true
+    }
+  })
 
   const modelGroups = groupBy(list, 'group')
 
@@ -89,7 +105,9 @@ const PopupContainer: React.FC<Props> = ({ provider: _provider, resolve }) => {
     return (
       <Flex>
         <ModelHeaderTitle>
-          {provider.isSystem ? t(`provider.${provider.id}`) : provider.name} {t('common.models')}
+          {provider.isSystem ? t(`provider.${provider.id}`) : provider.name}
+          {i18n.language.startsWith('zh') ? '' : ' '}
+          {t('common.models')}
         </ModelHeaderTitle>
         {loading && <LoadingOutlined size={20} />}
       </Flex>
@@ -111,6 +129,15 @@ const PopupContainer: React.FC<Props> = ({ provider: _provider, resolve }) => {
       }}
       centered>
       <SearchContainer>
+        <Center>
+          <Radio.Group value={filterType} onChange={(e) => setFilterType(e.target.value)} buttonStyle="solid">
+            <Radio.Button value="all">{t('models.all')}</Radio.Button>
+            <Radio.Button value="vision">{t('models.vision')}</Radio.Button>
+            <Radio.Button value="websearch">{t('models.websearch')}</Radio.Button>
+            <Radio.Button value="free">{t('models.free')}</Radio.Button>
+            <Radio.Button value="embedding">{t('models.embedding')}</Radio.Button>
+          </Radio.Group>
+        </Center>
         <Search placeholder={t('settings.provider.search_placeholder')} allowClear onSearch={setSearchText} />
       </SearchContainer>
       <ListContainer>
@@ -131,12 +158,12 @@ const PopupContainer: React.FC<Props> = ({ provider: _provider, resolve }) => {
                       {isWebSearchModel(model) && <WebSearchIcon />}
                       {isFreeModel(model) && (
                         <Tag style={{ marginLeft: 10 }} color="green">
-                          Free
+                          {t('models.free')}
                         </Tag>
                       )}
                       {isEmbeddingModel(model) && (
                         <Tag style={{ marginLeft: 10 }} color="orange">
-                          Embed
+                          {t('models.embedding')}
                         </Tag>
                       )}
                       {!isEmpty(model.description) && (
@@ -168,11 +195,16 @@ const PopupContainer: React.FC<Props> = ({ provider: _provider, resolve }) => {
 
 const SearchContainer = styled.div`
   display: flex;
-  flex-direction: row;
+  flex-direction: column;
-  align-items: center;
+  gap: 12px;
-  justify-content: space-between;
   padding: 0 22px;
   padding-bottom: 20px;
+  margin-top: -10px;
+
+  .ant-radio-group {
+    display: flex;
+    flex-wrap: wrap;
+  }
 `
 
 const ListContainer = styled.div`

src/renderer/src/pages/settings/ProviderSettings/ProviderSetting.tsx

@@ -9,7 +9,7 @@ import {
 } from '@ant-design/icons'
 import VisionIcon from '@renderer/components/Icons/VisionIcon'
 import WebSearchIcon from '@renderer/components/Icons/WebSearchIcon'
-import { getModelLogo, isVisionModel, isWebSearchModel, VISION_REGEX } from '@renderer/config/models'
+import { EMBEDDING_REGEX, getModelLogo, isVisionModel, isWebSearchModel, VISION_REGEX } from '@renderer/config/models'
 import { PROVIDER_CONFIG } from '@renderer/config/providers'
 import { useTheme } from '@renderer/context/ThemeProvider'
 import { useAssistants, useDefaultModel } from '@renderer/hooks/useAssistant'
@@ -165,7 +165,10 @@ const ProviderSetting: FC<Props> = ({ provider: _provider }) => {
       <Checkbox.Group
         value={model.type}
         onChange={(types) => onUpdateModelTypes(model, types as ModelType[])}
-        options={[{ label: t('model.type.vision'), value: 'vision', disabled: VISION_REGEX.test(model.id) }]}
+        options={[
+          { label: t('models.type.vision'), value: 'vision', disabled: VISION_REGEX.test(model.id) },
+          { label: t('models.type.embedding'), value: 'embedding', disabled: EMBEDDING_REGEX.test(model.id) }
+        ]}
       />
     </div>
   )
@@ -270,7 +273,7 @@ const ProviderSetting: FC<Props> = ({ provider: _provider }) => {
                 </Avatar>
                 {model.name} {isVisionModel(model) && <VisionIcon />}
                 {isWebSearchModel(model) && <WebSearchIcon />}
-                <Popover content={modelTypeContent(model)} title={t('model.type.select')} trigger="click">
+                <Popover content={modelTypeContent(model)} title={t('models.type.select')} trigger="click">
                   <SettingIcon />
                 </Popover>
               </ModelListHeader>

src/renderer/src/providers/BaseProvider.ts

@@ -1,7 +1,8 @@
+import { REFERENCE_PROMPT } from '@renderer/config/prompts'
 import { getOllamaKeepAliveTime } from '@renderer/hooks/useOllama'
 import { getKnowledgeBaseParams } from '@renderer/services/KnowledgeService'
 import store from '@renderer/store'
-import { Assistant, FileType, Message, Provider, Suggestion } from '@renderer/types'
+import { Assistant, Message, Provider, Suggestion } from '@renderer/types'
 import { delay } from '@renderer/utils'
 import { take } from 'lodash'
 import OpenAI from 'openai'
@@ -88,7 +89,6 @@ export default abstract class BaseProvider {
 
     const knowledgeId = message.knowledgeBaseIds[0]
     const base = store.getState().knowledge.bases.find((kb) => kb.id === knowledgeId)
-    console.debug('knowledge', base)
 
     if (!base) {
       return message.content
@@ -99,43 +99,20 @@ export default abstract class BaseProvider {
       base: getKnowledgeBaseParams(base)
     })
 
-    const references = take(searchResults, 5)
+    const references = take(searchResults, 6).map((item, index) => {
-      .map((item, index) => {
+      const sourceUrl = item.metadata.source
-        let sourceUrl = ''
+      const baseItem = base.items.find((i) => i.uniqueId === item.metadata.uniqueLoaderId)
-        let sourceName = ''
 
-        const baseItem = base.items.find((i) => i.uniqueId === item.metadata.uniqueLoaderId)
+      return {
+        id: index,
+        content: item.pageContent,
+        url: sourceUrl,
+        type: baseItem?.type
+      }
+    })
 
-        if (baseItem) {
+    const referencesContent = JSON.stringify(references, null, 2)
-          switch (baseItem.type) {
-            case 'file':
-              // sourceUrl = `file://${encodeURIComponent((baseItem?.content as FileType).path)}`
-              sourceName = (baseItem?.content as FileType).origin_name
-              break
-            case 'url':
-              sourceUrl = baseItem.content as string
-              sourceName = baseItem.content as string
-              break
-            case 'note':
-              sourceName = baseItem.content as string
-              break
-          }
-        }
 
-        return `
+    return REFERENCE_PROMPT.replace('{question}', message.content).replace('{references}', referencesContent)
----
-id: ${index}
-content: ${item.pageContent}
-source_type: ${baseItem?.type}
-source_name: ${sourceName}
-source_url: ${sourceUrl}
-`
-      })
-      .join('\n\n')
-
-    const prompt =
-      '回答问题请参考以下内容，并使用类似 [^1]: source 的脚注格式引用数据来源, source 根据 source_type 决定'
-
-    return [message.content, prompt, references].join('\n\n')
   }
 }

src/renderer/src/queue/KnowledgeQueue.ts

@@ -1,4 +1,4 @@
-import { AddLoaderReturn } from '@llm-tools/embedjs-interfaces'
+import type { AddLoaderReturn } from '@llm-tools/embedjs-interfaces'
 import db from '@renderer/databases'
 import { getKnowledgeBaseParams } from '@renderer/services/KnowledgeService'
 import store from '@renderer/store'
@@ -98,8 +98,7 @@ class KnowledgeQueue {
           break
         }
 
-        console.log(`[KnowledgeQueue] Processing item ${item.id} (${item.type}) in base ${baseId}`)
+        this.processItem(baseId, item)
-        await this.processItem(baseId, item)
       }
     } finally {
       console.log(`[KnowledgeQueue] Finished processing queue for base ${baseId}`)
@@ -152,25 +151,18 @@ class KnowledgeQueue {
       let result: AddLoaderReturn | null = null
       let note, content
 
+      console.log(`[KnowledgeQueue] Processing item: ${sourceItem.content}`)
+
       switch (item.type) {
-        case 'file':
-          console.log(`[KnowledgeQueue] Processing file: ${sourceItem.content}`)
-          result = await window.api.knowledgeBase.add({ base: baseParams, item: sourceItem })
-          break
-        case 'url':
-          console.log(`[KnowledgeQueue] Processing URL: ${sourceItem.content}`)
-          result = await window.api.knowledgeBase.add({ base: baseParams, item: sourceItem })
-          break
-        case 'sitemap':
-          console.log(`[KnowledgeQueue] Processing Sitemap: ${sourceItem.content}`)
-          result = await window.api.knowledgeBase.add({ base: baseParams, item: sourceItem })
-          break
         case 'note':
-          console.log(`[KnowledgeQueue] Processing note: ${sourceItem.content}`)
           note = await db.knowledge_notes.get(item.id)
-          if (!note) throw new Error(`Source note ${item.id} not found`)
+          if (note) {
-          content = note.content as string
+            content = note.content as string
-          result = await window.api.knowledgeBase.add({ base: baseParams, item: { ...sourceItem, content } })
+            result = await window.api.knowledgeBase.add({ base: baseParams, item: { ...sourceItem, content } })
+          }
+          break
+        default:
+          result = await window.api.knowledgeBase.add({ base: baseParams, item: sourceItem })
           break
       }
 

src/renderer/src/services/TokenService.ts

@@ -1,7 +1,7 @@
 import { Assistant, FileType, FileTypes, Message } from '@renderer/types'
-import { GPTTokens } from 'gpt-tokens'
 import { flatten, takeRight } from 'lodash'
 import { CompletionUsage } from 'openai/resources'
+import { approximateTokenSize } from 'tokenx'
 
 import { getAssistantSettings } from './AssistantService'
 import { filterContextMessages, filterMessages } from './MessagesService'
@@ -45,12 +45,7 @@ async function getMessageParam(message: Message): Promise<MessageItem[]> {
 }
 
 export function estimateTextTokens(text: string) {
-  const { usedTokens } = new GPTTokens({
+  return approximateTokenSize(text)
-    model: 'gpt-4o',
-    messages: [{ role: 'user', content: text }]
-  })
-
-  return usedTokens - 7
 }
 
 export function estimateImageTokens(file: FileType) {
@@ -58,11 +53,6 @@ export function estimateImageTokens(file: FileType) {
 }
 
 export async function estimateMessageUsage(message: Message): Promise<CompletionUsage> {
-  const { usedTokens, promptUsedTokens, completionUsedTokens } = new GPTTokens({
-    model: 'gpt-4o',
-    messages: await getMessageParam(message)
-  })
-
   let imageTokens = 0
 
   if (message.files) {
@@ -74,10 +64,12 @@ export async function estimateMessageUsage(message: Message): Promise<Completion
     }
   }
 
+  const tokens = estimateTextTokens(message.content)
+
   return {
-    prompt_tokens: promptUsedTokens,
+    prompt_tokens: tokens,
-    completion_tokens: completionUsedTokens,
+    completion_tokens: tokens,
-    total_tokens: usedTokens + (imageTokens ? imageTokens - 7 : 0)
+    total_tokens: tokens + (imageTokens ? imageTokens - 7 : 0)
   }
 }
 
@@ -121,16 +113,10 @@ export async function estimateHistoryTokens(assistant: Assistant, msgs: Message[
     allMessages = allMessages.concat(items)
   }
 
-  const { usedTokens } = new GPTTokens({
+  const prompt = assistant.prompt
-    model: 'gpt-4o',
+  const input = flatten(allMessages)
-    messages: [
+    .map((m) => m.content)
-      {
+    .join('\n')
-        role: 'system',
-        content: assistant.prompt
-      },
-      ...flatten(allMessages)
-    ]
-  })
 
-  return usedTokens - 7 + uasageTokens
+  return estimateTextTokens(prompt + input) + uasageTokens
 }

src/renderer/src/store/knowledge.ts

@@ -43,12 +43,24 @@ const knowledgeSlice = createSlice({
       }
     },
 
+    updateBases(state, action: PayloadAction<KnowledgeBase[]>) {
+      state.bases = action.payload
+    },
+
     addItem(state, action: PayloadAction<{ baseId: string; item: KnowledgeItem }>) {
       const base = state.bases.find((b) => b.id === action.payload.baseId)
       if (base) {
-        if (action.payload.item.type === 'note') {
+        if (action.payload.item.type === 'file') {
+          action.payload.item.created_at = new Date(action.payload.item.created_at).getTime()
+          action.payload.item.updated_at = new Date(action.payload.item.updated_at).getTime()
           base.items.push(action.payload.item)
         }
+        if (action.payload.item.type === 'directory') {
+          const directoryExists = base.items.some((item) => item.content === action.payload.item.content)
+          if (!directoryExists) {
+            base.items.push(action.payload.item)
+          }
+        }
         if (action.payload.item.type === 'url') {
           const urlExists = base.items.some((item) => item.content === action.payload.item.content)
           if (!urlExists) {
@@ -61,9 +73,7 @@ const knowledgeSlice = createSlice({
             base.items.push(action.payload.item)
           }
         }
-        if (action.payload.item.type === 'file') {
+        if (action.payload.item.type === 'note') {
-          action.payload.item.created_at = new Date(action.payload.item.created_at).getTime()
-          action.payload.item.updated_at = new Date(action.payload.item.updated_at).getTime()
           base.items.push(action.payload.item)
         }
         base.updated_at = Date.now()
@@ -79,12 +89,10 @@ const knowledgeSlice = createSlice({
       }
     },
 
-    updateFiles(state, action: PayloadAction<{ baseId: string; items: KnowledgeItem[] }>) {
+    addFiles(state, action: PayloadAction<{ baseId: string; items: KnowledgeItem[] }>) {
       const base = state.bases.find((b) => b.id === action.payload.baseId)
       if (base) {
-        // 保留非文件类型的项目
+        base.items = [...base.items, ...action.payload.items]
-        const nonFileItems = base.items.filter((item) => item.type !== 'file')
-        base.items = [...nonFileItems, ...action.payload.items]
         base.updated_at = Date.now()
       }
     },
@@ -170,8 +178,9 @@ export const {
   deleteBase,
   renameBase,
   updateBase,
+  updateBases,
   addItem,
-  updateFiles,
+  addFiles,
   updateNotes,
   removeItem,
   updateItemProcessingStatus,

src/renderer/src/types/index.ts

@@ -89,7 +89,7 @@ export type Provider = {
 
 export type ProviderType = 'openai' | 'anthropic' | 'gemini'
 
-export type ModelType = 'text' | 'vision'
+export type ModelType = 'text' | 'vision' | 'embedding'
 
 export type Model = {
   id: string
@@ -183,11 +183,13 @@ export interface Shortcut {
 
 export type ProcessingStatus = 'pending' | 'processing' | 'completed' | 'failed'
 
+export type KnowledgeItemType = 'file' | 'url' | 'note' | 'sitemap' | 'directory'
+
 export type KnowledgeItem = {
   id: string
   baseId?: string
   uniqueId?: string
-  type: 'file' | 'url' | 'note' | 'sitemap'
+  type: KnowledgeItemType
   content: string | FileType
   created_at: number
   updated_at: number
@@ -197,8 +199,6 @@ export type KnowledgeItem = {
   retryCount?: number
 }
 
-export type KnowledgeItemType = 'file' | 'url' | 'note' | 'sitemap'
-
 export interface KnowledgeBase {
   id: string
   name: string