feat: allow knowledge base multiple search #1346 (#1773)

* feat: agent can select multiple knowledge bases * feat: basic search multiple knowledge base * fix bug: knowledge base is delete, assistants and agents sync delete * fix bug: assistant and knowledge base button sync * feat: allow to search multiple knowledge base * chore: finish rebase to upstream/main
2025-02-17 16:36:25 +08:00 · 2025-02-17 16:36:25 +08:00 · 266f909045
commit 266f909045
parent bad2f15c1f
8 changed files with 83 additions and 78 deletions
--- a/src/renderer/src/hooks/useKnowledge.ts
+++ b/src/renderer/src/hooks/useKnowledge.ts
@ -307,16 +307,22 @@ export const useKnowledgeBases = () => {

    // remove assistant knowledge_base
    const _assistants = assistants.map((assistant) => {
-      if (assistant.knowledge_base?.id === baseId) {
-        return { ...assistant, knowledge_base: undefined }
+      if (assistant.knowledge_bases?.find((kb) => kb.id === baseId)) {
+        return {
+          ...assistant,
+          knowledge_bases: assistant.knowledge_bases.filter((kb) => kb.id !== baseId)
+        }
      }
      return assistant
    })

    // remove agent knowledge_base
    const _agents = agents.map((agent) => {
-      if (agent.knowledge_base?.id === baseId) {
-        return { ...agent, knowledge_base: undefined }
+      if (agent.knowledge_bases?.find((kb) => kb.id === baseId)) {
+        return {
+          ...agent,
+          knowledge_bases: agent.knowledge_bases.filter((kb) => kb.id !== baseId)
+        }
      }
      return agent
    })
--- a/src/renderer/src/pages/agents/components/AddAgentPopup.tsx
+++ b/src/renderer/src/pages/agents/components/AddAgentPopup.tsx
@ -9,7 +9,7 @@ import { useSidebarIconShow } from '@renderer/hooks/useSidebarIcon'
 import { fetchGenerate } from '@renderer/services/ApiService'
 import { getDefaultModel } from '@renderer/services/AssistantService'
 import { useAppSelector } from '@renderer/store'
-import { Agent } from '@renderer/types'
+import { Agent, KnowledgeBase } from '@renderer/types'
 import { getLeadingEmoji, uuid } from '@renderer/utils'
 import { Button, Form, FormInstance, Input, Modal, Popover, Select, SelectProps } from 'antd'
 import TextArea from 'antd/es/input/TextArea'
@ -25,7 +25,7 @@ type FieldType = {
  id: string
  name: string
  prompt: string
-  knowledge_base_id: string
+  knowledge_base_id: string[]
 }

 const PopupContainer: React.FC<Props> = ({ resolve }) => {
@ -57,7 +57,9 @@ const PopupContainer: React.FC<Props> = ({ resolve }) => {
    const _agent: Agent = {
      id: uuid(),
      name: values.name,
-      knowledge_base: knowledgeState.bases.find((t) => t.id === values.knowledge_base_id),
+      knowledge_bases: values.knowledge_base_id
+        .map((id) => knowledgeState.bases.find((t) => t.id === id))
+        .filter((base): base is KnowledgeBase => base !== undefined),
      emoji: _emoji,
      prompt: values.prompt,
      defaultModel: getDefaultModel(),
@ -156,6 +158,7 @@ const PopupContainer: React.FC<Props> = ({ resolve }) => {
        {showKnowledgeIcon && (
          <Form.Item name="knowledge_base_id" label={t('agents.add.knowledge_base')} rules={[{ required: false }]}>
            <Select
+              mode="multiple"
              allowClear
              placeholder={t('agents.add.knowledge_base.placeholder')}
              menuItemSelectedIcon={<CheckOutlined />}
--- a/src/renderer/src/pages/home/Inputbar/Inputbar.tsx
+++ b/src/renderer/src/pages/home/Inputbar/Inputbar.tsx
@ -52,7 +52,6 @@ interface Props {

 let _text = ''
 let _files: FileType[] = []
-let _base: KnowledgeBase | undefined

 const Inputbar: FC<Props> = ({ assistant: _assistant, setActiveTopic }) => {
  const [text, setText] = useState(_text)
@ -83,7 +82,7 @@ const Inputbar: FC<Props> = ({ assistant: _assistant, setActiveTopic }) => {
  const [spaceClickCount, setSpaceClickCount] = useState(0)
  const spaceClickTimer = useRef<NodeJS.Timeout>()
  const [isTranslating, setIsTranslating] = useState(false)
-  const [selectedKnowledgeBase, setSelectedKnowledgeBase] = useState<KnowledgeBase | undefined>(_base)
+  const [selectedKnowledgeBases, setSelectedKnowledgeBases] = useState<KnowledgeBase[]>([])
  const [mentionModels, setMentionModels] = useState<Model[]>([])
  const [isMentionPopupOpen, setIsMentionPopupOpen] = useState(false)

@ -104,7 +103,6 @@ const Inputbar: FC<Props> = ({ assistant: _assistant, setActiveTopic }) => {

  _text = text
  _files = files
-  _base = selectedKnowledgeBase

  const sendMessage = useCallback(async () => {
    await modelGenerating()
@ -124,8 +122,8 @@ const Inputbar: FC<Props> = ({ assistant: _assistant, setActiveTopic }) => {
      status: 'success'
    }

-    if (selectedKnowledgeBase) {
-      message.knowledgeBaseIds = [selectedKnowledgeBase.id]
+    if (selectedKnowledgeBases) {
+      message.knowledgeBaseIds = selectedKnowledgeBases.map((base) => base.id)
    }

    if (files.length > 0) {
@ -144,7 +142,7 @@ const Inputbar: FC<Props> = ({ assistant: _assistant, setActiveTopic }) => {
    setTimeout(() => resizeTextArea(), 0)

    setExpend(false)
-  }, [inputEmpty, text, assistant.id, assistant.topics, selectedKnowledgeBase, files, mentionModels])
+  }, [inputEmpty, text, assistant.id, assistant.topics, selectedKnowledgeBases, files, mentionModels])

  const translate = async () => {
    if (isTranslating) {
@ -458,14 +456,15 @@ const Inputbar: FC<Props> = ({ assistant: _assistant, setActiveTopic }) => {
  }, [])

  useEffect(() => {
-    setSelectedKnowledgeBase(showKnowledgeIcon ? assistant.knowledge_base : undefined)
-  }, [assistant.id, assistant.knowledge_base, showKnowledgeIcon])
+    // if assistant knowledge bases are undefined return []
+    setSelectedKnowledgeBases(showKnowledgeIcon ? (assistant.knowledge_bases ?? []) : [])
+  }, [assistant.id, assistant.knowledge_bases, showKnowledgeIcon])

  const textareaRows = window.innerHeight >= 1000 || isBubbleStyle ? 2 : 1

-  const handleKnowledgeBaseSelect = (base?: KnowledgeBase) => {
-    updateAssistant({ ...assistant, knowledge_base: base })
-    setSelectedKnowledgeBase(base)
+  const handleKnowledgeBaseSelect = (bases?: KnowledgeBase[]) => {
+    updateAssistant({ ...assistant, knowledge_bases: bases })
+    setSelectedKnowledgeBases(bases ?? [])
  }

  const onMentionModel = (model: Model) => {
@ -573,7 +572,7 @@ const Inputbar: FC<Props> = ({ assistant: _assistant, setActiveTopic }) => {
              </Tooltip>
              {showKnowledgeIcon && (
                <KnowledgeBaseButton
-                  selectedBase={selectedKnowledgeBase}
+                  selectedBases={selectedKnowledgeBases}
                  onSelect={handleKnowledgeBaseSelect}
                  ToolbarButton={ToolbarButton}
                  disabled={files.length > 0}
--- a/src/renderer/src/pages/home/Inputbar/KnowledgeBaseButton.tsx
+++ b/src/renderer/src/pages/home/Inputbar/KnowledgeBaseButton.tsx
@ -1,71 +1,63 @@
-import { FileSearchOutlined } from '@ant-design/icons'
+import { CheckOutlined, FileSearchOutlined } from '@ant-design/icons'
 import { useAppSelector } from '@renderer/store'
 import { KnowledgeBase } from '@renderer/types'
-import { Button, Popover, Tooltip } from 'antd'
+import { Popover, Select, SelectProps, Tooltip } from 'antd'
 import { FC } from 'react'
 import { useTranslation } from 'react-i18next'
 import styled from 'styled-components'

 interface Props {
-  selectedBase?: KnowledgeBase
-  onSelect: (base?: KnowledgeBase) => void
+  selectedBases?: KnowledgeBase[]
+  onSelect: (bases: KnowledgeBase[]) => void
  disabled?: boolean
  ToolbarButton?: any
 }

-const KnowledgeBaseSelector: FC<Props> = ({ selectedBase, onSelect }) => {
+const KnowledgeBaseSelector: FC<Props> = ({ selectedBases, onSelect }) => {
  const { t } = useTranslation()
  const knowledgeState = useAppSelector((state) => state.knowledge)
+  const knowledgeOptions: SelectProps['options'] = knowledgeState.bases.map((base) => ({
+    label: base.name,
+    value: base.id
+  }))

  return (
    <SelectorContainer>
      {knowledgeState.bases.length === 0 ? (
        <EmptyMessage>{t('knowledge.no_bases')}</EmptyMessage>
      ) : (
-        <>
-          {selectedBase && (
-            <Button type="link" block onClick={() => onSelect(undefined)} style={{ textAlign: 'left' }}>
-              {t('knowledge.clear_selection')}
-            </Button>
-          )}
-          {knowledgeState.bases.map((base) => (
-            <Button
-              key={base.id}
-              type={selectedBase?.id === base.id ? 'primary' : 'text'}
-              block
-              onClick={() => onSelect(base)}
-              style={{ textAlign: 'left' }}>
-              {base.name}
-            </Button>
-          ))}
-        </>
+        <Select
+          mode="multiple"
+          value={selectedBases?.map((base) => base.id)}
+          allowClear
+          placeholder={t('agents.add.knowledge_base.placeholder')}
+          menuItemSelectedIcon={<CheckOutlined />}
+          options={knowledgeOptions}
+          onChange={(ids) => {
+            const newSelected = knowledgeState.bases.filter((base) => ids.includes(base.id))
+            onSelect(newSelected)
+          }}
+          style={{ width: '200px' }}
+        />
      )}
    </SelectorContainer>
  )
 }

-const KnowledgeBaseButton: FC<Props> = ({ selectedBase, onSelect, disabled, ToolbarButton }) => {
+const KnowledgeBaseButton: FC<Props> = ({ selectedBases, onSelect, disabled, ToolbarButton }) => {
  const { t } = useTranslation()

-  if (selectedBase) {
-    return (
-      <Tooltip placement="top" title={selectedBase.name} arrow>
-        <ToolbarButton type="text" onClick={() => onSelect(undefined)}>
-          <FileSearchOutlined style={{ color: selectedBase ? 'var(--color-link)' : 'var(--color-icon)' }} />
-        </ToolbarButton>
-      </Tooltip>
-    )
-  }
-
  return (
    <Tooltip placement="top" title={t('chat.input.knowledge_base')} arrow>
      <Popover
        placement="top"
-        content={<KnowledgeBaseSelector selectedBase={selectedBase} onSelect={onSelect} />}
+        content={<KnowledgeBaseSelector selectedBases={selectedBases} onSelect={onSelect} />}
        overlayStyle={{ maxWidth: 400 }}
        trigger="click">
-        <ToolbarButton type="text" onClick={() => selectedBase && onSelect(undefined)} disabled={disabled}>
-          <FileSearchOutlined style={{ color: selectedBase ? 'var(--color-link)' : 'var(--color-icon)' }} />
+        <ToolbarButton type="text" disabled={disabled}>
+          <FileSearchOutlined
+            style={{ color: selectedBases && selectedBases?.length > 0 ? 'var(--color-link)' : 'var(--color-icon)' }}
+          />
        </ToolbarButton>
      </Popover>
    </Tooltip>
--- a/src/renderer/src/pages/settings/AssistantSettings/AssistantKnowledgeBaseSettings.tsx
+++ b/src/renderer/src/pages/settings/AssistantSettings/AssistantKnowledgeBaseSettings.tsx
@ -26,8 +26,8 @@ const AssistantKnowledgeBaseSettings: React.FC<Props> = ({ assistant, updateAssi
  })

  const onUpdate = (value) => {
-    const knowledge_base = knowledgeState.bases.find((t) => t.id === value)
-    const _assistant = { ...assistant, knowledge_base }
+    const knowledge_bases = value.map((id) => knowledgeState.bases.find((b) => b.id === id))
+    const _assistant = { ...assistant, knowledge_bases }
    updateAssistant(_assistant)
  }

@ -37,8 +37,9 @@ const AssistantKnowledgeBaseSettings: React.FC<Props> = ({ assistant, updateAssi
        {t('common.knowledge_base')}
      </Box>
      <Select
+        mode="multiple"
        allowClear
-        defaultValue={assistant.knowledge_base?.id}
+        value={assistant.knowledge_bases?.map((b) => b.id)}
        placeholder={t('agents.add.knowledge_base.placeholder')}
        menuItemSelectedIcon={<CheckOutlined />}
        options={knowledgeOptions}
--- a/src/renderer/src/providers/BaseProvider.ts
+++ b/src/renderer/src/providers/BaseProvider.ts
@ -5,6 +5,7 @@ import { getKnowledgeReferences } from '@renderer/services/KnowledgeService'
 import store from '@renderer/store'
 import { Assistant, GenerateImageParams, Message, Model, Provider, Suggestion } from '@renderer/types'
 import { delay, isJSON, parseJSON } from '@renderer/utils'
+import { t } from 'i18next'
 import OpenAI from 'openai'

 import { CompletionsParams } from '.'
@ -83,21 +84,35 @@ export default abstract class BaseProvider {
      return message.content
    }

-    const knowledgeId = message.knowledgeBaseIds[0]
-    const base = store.getState().knowledge.bases.find((kb) => kb.id === knowledgeId)
+    const bases = store.getState().knowledge.bases.filter((kb) => message.knowledgeBaseIds?.includes(kb.id))

-    if (!base) {
+    if (!bases || bases.length === 0) {
      return message.content
    }

-    const { referencesContent, referencesCount } = await getKnowledgeReferences(base, message)
+    const allReferencesPromises = bases.map(async (base) => {
+      const references = await getKnowledgeReferences(base, message)

-    // 如果知识库中未检索到内容则使用通用逻辑
-    if (referencesCount === 0) {
+      return {
+        knowledgeBaseId: base.id,
+        references
+      }
+    })
+    const allReferences = (await Promise.all(allReferencesPromises))
+      .filter((result) => result.references && result.references.length > 0)
+      .flat()
+
+    if (allReferences.length === 0) {
+      window.message.info({
+        content: t('knowledge.no_match'),
+        duration: 4,
+        key: 'knowledge-base-no-match-info'
+      })
      return message.content
    }
+    const allReferencesContent = `\`\`\`json\n${JSON.stringify(allReferences, null, 2)}\n\`\`\``

-    return REFERENCE_PROMPT.replace('{question}', message.content).replace('{references}', referencesContent)
+    return REFERENCE_PROMPT.replace('{question}', message.content).replace('{references}', allReferencesContent)
  }

  protected getCustomParameters(assistant: Assistant) {
--- a/src/renderer/src/services/KnowledgeService.ts
+++ b/src/renderer/src/services/KnowledgeService.ts
@ -3,7 +3,6 @@ import { DEFAULT_KNOWLEDGE_DOCUMENT_COUNT, DEFAULT_KNOWLEDGE_THRESHOLD } from '@
 import { getEmbeddingMaxContext } from '@renderer/config/embedings'
 import AiProvider from '@renderer/providers/AiProvider'
 import { FileType, KnowledgeBase, KnowledgeBaseParams, Message } from '@renderer/types'
-import { t } from 'i18next'
 import { take } from 'lodash'

 import { getProviderByModel } from './AssistantService'
@ -91,14 +90,6 @@ export const getKnowledgeReferences = async (base: KnowledgeBase, message: Messa
        return item.score >= threshold
      })
    )
-  if (searchResults.length === 0) {
-    window.message.info({
-      content: t('knowledge.no_match'),
-      duration: 4,
-      key: 'knowledge-base-no-match-info'
-    })
-    return { referencesContent: '', referencesCount: 0 }
-  }

  const _searchResults = await Promise.all(
    searchResults.map(async (item) => {
@ -121,7 +112,5 @@ export const getKnowledgeReferences = async (base: KnowledgeBase, message: Messa
    })
  )

-  const referencesContent = `\`\`\`json\n${JSON.stringify(references, null, 2)}\n\`\`\``
-
-  return { referencesContent, referencesCount: references.length }
+  return references
 }
--- a/src/renderer/src/types/index.ts
+++ b/src/renderer/src/types/index.ts
@ -5,7 +5,7 @@ export type Assistant = {
  id: string
  name: string
  prompt: string
-  knowledge_base?: KnowledgeBase
+  knowledge_bases?: KnowledgeBase[]
  topics: Topic[]
  type: string
  emoji?: string