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

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) {
+      if (assistant.knowledge_bases?.find((kb) => kb.id === baseId)) {
-        return { ...assistant, knowledge_base: undefined }
+        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) {
+      if (agent.knowledge_bases?.find((kb) => kb.id === baseId)) {
-        return { ...agent, knowledge_base: undefined }
+        return {
+          ...agent,
+          knowledge_bases: agent.knowledge_bases.filter((kb) => kb.id !== baseId)
+        }
       }
       return agent
     })

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 />}

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) {
+    if (selectedKnowledgeBases) {
-      message.knowledgeBaseIds = [selectedKnowledgeBase.id]
+      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)
+    // if assistant knowledge bases are undefined return []
-  }, [assistant.id, assistant.knowledge_base, showKnowledgeIcon])
+    setSelectedKnowledgeBases(showKnowledgeIcon ? (assistant.knowledge_bases ?? []) : [])
+  }, [assistant.id, assistant.knowledge_bases, showKnowledgeIcon])
 
   const textareaRows = window.innerHeight >= 1000 || isBubbleStyle ? 2 : 1
 
-  const handleKnowledgeBaseSelect = (base?: KnowledgeBase) => {
+  const handleKnowledgeBaseSelect = (bases?: KnowledgeBase[]) => {
-    updateAssistant({ ...assistant, knowledge_base: base })
+    updateAssistant({ ...assistant, knowledge_bases: bases })
-    setSelectedKnowledgeBase(base)
+    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}

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
+  selectedBases?: KnowledgeBase[]
-  onSelect: (base?: KnowledgeBase) => void
+  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>
       ) : (
-        <>
+        <Select
-          {selectedBase && (
+          mode="multiple"
-            <Button type="link" block onClick={() => onSelect(undefined)} style={{ textAlign: 'left' }}>
+          value={selectedBases?.map((base) => base.id)}
-              {t('knowledge.clear_selection')}
+          allowClear
-            </Button>
+          placeholder={t('agents.add.knowledge_base.placeholder')}
-          )}
+          menuItemSelectedIcon={<CheckOutlined />}
-          {knowledgeState.bases.map((base) => (
+          options={knowledgeOptions}
-            <Button
+          onChange={(ids) => {
-              key={base.id}
+            const newSelected = knowledgeState.bases.filter((base) => ids.includes(base.id))
-              type={selectedBase?.id === base.id ? 'primary' : 'text'}
+            onSelect(newSelected)
-              block
+          }}
-              onClick={() => onSelect(base)}
+          style={{ width: '200px' }}
-              style={{ textAlign: 'left' }}>
+        />
-              {base.name}
-            </Button>
-          ))}
-        </>
       )}
     </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}>
+        <ToolbarButton type="text" disabled={disabled}>
-          <FileSearchOutlined style={{ color: selectedBase ? 'var(--color-link)' : 'var(--color-icon)' }} />
+          <FileSearchOutlined
+            style={{ color: selectedBases && selectedBases?.length > 0 ? 'var(--color-link)' : 'var(--color-icon)' }}
+          />
         </ToolbarButton>
       </Popover>
     </Tooltip>

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 knowledge_bases = value.map((id) => knowledgeState.bases.find((b) => b.id === id))
-    const _assistant = { ...assistant, knowledge_base }
+    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}

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 bases = store.getState().knowledge.bases.filter((kb) => message.knowledgeBaseIds?.includes(kb.id))
-    const base = store.getState().knowledge.bases.find((kb) => kb.id === knowledgeId)
 
-    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)
 
-    // 如果知识库中未检索到内容则使用通用逻辑
+      return {
-    if (referencesCount === 0) {
+        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) {

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 references
-
-  return { referencesContent, referencesCount: references.length }
 }

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