From 3c58d96db548e8efee89df6f1595eecff12a4c18 Mon Sep 17 00:00:00 2001
From: Soulter <905617992@qq.com>
Date: Fri, 24 Oct 2025 14:45:07 +0800
Subject: [PATCH] feat: add configuration for final knowledge base retrieval
 count and update related components

---
 astrbot/core/config/default.py                |  16 +-
 astrbot/core/pipeline/process_stage/utils.py  |   3 +-
 .../shared/KnowledgeBaseSelector.vue          | 178 +++++++++---------
 3 files changed, 103 insertions(+), 94 deletions(-)

diff --git a/astrbot/core/config/default.py b/astrbot/core/config/default.py
index d23364b19..36ac616ae 100644
--- a/astrbot/core/config/default.py
+++ b/astrbot/core/config/default.py
@@ -137,6 +137,7 @@ DEFAULT_CONFIG = {
     "default_kb_collection": "",  # 默认知识库名称, 已经过时
     "plugin_set": ["*"],  # "*" 表示使用所有可用的插件, 空列表表示不使用任何插件
     "kb_names": [],  # 默认知识库名称列表
+    "kb_final_top_k": 5,  # 知识库检索的最终返回结果数量
 }
 
 
@@ -2001,6 +2002,8 @@ CONFIG_METADATA_2 = {
             "default_kb_collection": {
                 "type": "string",
             },
+            "kb_names": {"type": "list", "items": {"type": "string"}},
+            "kb_final_top_k": {"type": "int", "default": 5},
         },
     },
 }
@@ -2079,10 +2082,17 @@ CONFIG_METADATA_3 = {
                 "description": "知识库",
                 "type": "object",
                 "items": {
-                    "default_kb_collection": {
-                        "description": "默认使用的知识库",
-                        "type": "string",
+                    "kb_names": {
+                        "description": "知识库列表",
+                        "type": "list",
+                        "items": {"type": "string"},
                         "_special": "select_knowledgebase",
+                        "hint": "支持多选",
+                    },
+                    "kb_final_top_k": {
+                        "description": "最终返回结果数",
+                        "type": "int",
+                        "hint": "从知识库中检索到的结果数量，越大可能获得越多相关信息，但也可能引入噪音。建议根据实际需求调整",
                     },
                 },
             },
diff --git a/astrbot/core/pipeline/process_stage/utils.py b/astrbot/core/pipeline/process_stage/utils.py
index bf40b3cae..8fe34e3ee 100644
--- a/astrbot/core/pipeline/process_stage/utils.py
+++ b/astrbot/core/pipeline/process_stage/utils.py
@@ -7,7 +7,6 @@ async def inject_kb_context(
     umo: str,
     p_ctx: PipelineContext,
     req: ProviderRequest,
-    top_k: int = 5,
 ) -> None:
     """inject knowledge base context into the provider request
 
@@ -17,6 +16,7 @@ async def inject_kb_context(
     """
     kb_mgr = p_ctx.plugin_manager.context.kb_manager
     kb_names = p_ctx.astrbot_config.get("kb_names", [])
+    top_k = p_ctx.astrbot_config.get("kb_final_top_k", 5)
 
     if not kb_names:
         return
@@ -24,6 +24,7 @@ async def inject_kb_context(
     kb_context = await kb_mgr.retrieve(
         query=req.prompt,
         kb_names=kb_names,
+        top_m_final=top_k,
     )
     if not kb_context:
         return
diff --git a/dashboard/src/components/shared/KnowledgeBaseSelector.vue b/dashboard/src/components/shared/KnowledgeBaseSelector.vue
index 3e0a8be15..a90e6f5e2 100644
--- a/dashboard/src/components/shared/KnowledgeBaseSelector.vue
+++ b/dashboard/src/components/shared/KnowledgeBaseSelector.vue
@@ -1,11 +1,21 @@
 <template>
   <div class="d-flex align-center justify-space-between">
-    <span v-if="!modelValue" style="color: rgb(var(--v-theme-primaryText));">
+    <span v-if="!modelValue || (Array.isArray(modelValue) && modelValue.length === 0)" 
+          style="color: rgb(var(--v-theme-primaryText));">
       未选择
     </span>
-    <span v-else>
-      {{ modelValue }}
-    </span>
+    <div v-else class="d-flex flex-wrap gap-1">
+      <v-chip 
+        v-for="name in modelValue" 
+        :key="name" 
+        size="small" 
+        color="primary" 
+        variant="tonal"
+        closable
+        @click:close="removeKnowledgeBase(name)">
+        {{ name }}
+      </v-chip>
+    </div>
     <v-btn size="small" color="primary" variant="tonal" @click="openDialog">
       {{ buttonText }}
     </v-btn>
@@ -21,86 +31,57 @@
       <v-card-text class="pa-0" style="max-height: 400px; overflow-y: auto;">
         <v-progress-linear v-if="loading" indeterminate color="primary"></v-progress-linear>
         
-        <!-- 插件未安装提示 -->
-        <div v-if="!loading && !pluginInstalled" class="text-center py-8">
-          <v-icon size="64" color="grey-lighten-1">mdi-puzzle-outline</v-icon>
-          <p class="text-grey mt-4 mb-4">知识库插件未安装</p>
-          <v-btn color="primary" variant="tonal" @click="goToKnowledgeBasePage">
-            前往知识库页面
-          </v-btn>
-        </div>
-        
         <!-- 知识库列表 -->
-        <v-list v-else-if="!loading && pluginInstalled" density="compact">
-          <!-- 不使用选项 -->
-          <v-list-item
-            :value="''"
-            @click="selectKnowledgeBase({ collection_name: '' })"
-            :active="selectedKnowledgeBase === ''"
-            rounded="md"
-            class="ma-1">
-            <template v-slot:prepend>
-              <v-icon color="grey-lighten-1">mdi-close-circle-outline</v-icon>
-            </template>
-            <v-list-item-title>不使用</v-list-item-title>
-            <v-list-item-subtitle>不使用任何知识库</v-list-item-subtitle>
-            
-            <template v-slot:append>
-              <v-icon v-if="selectedKnowledgeBase === ''" color="primary">mdi-check-circle</v-icon>
-            </template>
-          </v-list-item>
-          
-          <v-divider v-if="knowledgeBaseList.length > 0" class="my-2"></v-divider>
-          
+        <v-list v-if="!loading" density="compact">
           <!-- 知识库选项 -->
           <v-list-item
             v-for="kb in knowledgeBaseList"
-            :key="kb.collection_name"
-            :value="kb.collection_name"
-            @click="selectKnowledgeBase(kb)"
-            :active="selectedKnowledgeBase === kb.collection_name"
+            :key="kb.kb_id"
+            :value="kb.kb_name"
+            @click="selectKnowledgeBase(kb.kb_name)"
+            :active="isSelected(kb.kb_name)"
             rounded="md"
             class="ma-1">
             <template v-slot:prepend>
-              <span class="emoji-icon">{{ kb.emoji || '🙂' }}</span>
+              <span class="emoji-icon">{{ kb.emoji || '📚' }}</span>
             </template>
-            <v-list-item-title>{{ kb.collection_name }}</v-list-item-title>
+            <v-list-item-title>{{ kb.kb_name }}</v-list-item-title>
             <v-list-item-subtitle>
               {{ kb.description || '无描述' }}
-              <span v-if="kb.count !== undefined"> - {{ kb.count }} 项知识</span>
+              <span v-if="kb.doc_count !== undefined"> - {{ kb.doc_count }} 个文档</span>
+              <span v-if="kb.chunk_count !== undefined"> - {{ kb.chunk_count }} 个块</span>
             </v-list-item-subtitle>
             
             <template v-slot:append>
-              <v-icon v-if="selectedKnowledgeBase === kb.collection_name" color="primary">mdi-check-circle</v-icon>
+              <v-icon v-if="isSelected(kb.kb_name)" color="primary">
+                mdi-checkbox-marked
+              </v-icon>
+              <v-icon v-else color="grey-lighten-1">
+                mdi-checkbox-blank-outline
+              </v-icon>
             </template>
           </v-list-item>
           
           <!-- 当没有知识库时显示创建提示 -->
-          <div v-if="knowledgeBaseList.length === 0" class="text-center py-4">
-            <p class="text-grey mb-4">暂无知识库</p>
-            <v-btn color="primary" variant="tonal" size="small" @click="goToKnowledgeBasePage">
+          <div v-if="knowledgeBaseList.length === 0" class="text-center py-8">
+            <v-icon size="64" color="grey-lighten-1">mdi-database-off</v-icon>
+            <p class="text-grey mt-4 mb-4">暂无知识库</p>
+            <v-btn color="primary" variant="tonal" @click="goToKnowledgeBasePage">
               创建知识库
             </v-btn>
           </div>
         </v-list>
-        
-        <!-- 空状态（插件未安装时保留原有逻辑） -->
-        <div v-else-if="!loading && !pluginInstalled && knowledgeBaseList.length === 0" class="text-center py-8">
-          <v-icon size="64" color="grey-lighten-1">mdi-database-off</v-icon>
-          <p class="text-grey mt-4 mb-4">暂无知识库</p>
-          <v-btn color="primary" variant="tonal" @click="goToKnowledgeBasePage">
-            创建知识库
-          </v-btn>
-        </div>
       </v-card-text>
       
       <v-card-actions class="pa-4">
+        <div v-if="selectedKnowledgeBases.length > 0" class="text-caption text-grey">
+          已选择 {{ selectedKnowledgeBases.length }} 个知识库
+        </div>
         <v-spacer></v-spacer>
         <v-btn variant="text" @click="cancelSelection">取消</v-btn>
         <v-btn 
           color="primary" 
-          @click="confirmSelection"
-          :disabled="selectedKnowledgeBase === null || selectedKnowledgeBase === undefined">
+          @click="confirmSelection">
           确认选择
         </v-btn>
       </v-card-actions>
@@ -115,8 +96,8 @@ import { useRouter } from 'vue-router'
 
 const props = defineProps({
   modelValue: {
-    type: String,
-    default: ''
+    type: Array,
+    default: () => []
   },
   buttonText: {
     type: String,
@@ -130,74 +111,87 @@ const router = useRouter()
 const dialog = ref(false)
 const knowledgeBaseList = ref([])
 const loading = ref(false)
-const selectedKnowledgeBase = ref('')
-const pluginInstalled = ref(false)
+const selectedKnowledgeBases = ref([])
 
-// 监听 modelValue 变化，同步到 selectedKnowledgeBase
+// 监听 modelValue 变化，同步到 selectedKnowledgeBases
 watch(() => props.modelValue, (newValue) => {
-  selectedKnowledgeBase.value = newValue || ''
+  selectedKnowledgeBases.value = Array.isArray(newValue) ? [...newValue] : []
 }, { immediate: true })
 
 async function openDialog() {
-  selectedKnowledgeBase.value = props.modelValue || ''
+  // 初始化选中状态
+  selectedKnowledgeBases.value = Array.isArray(props.modelValue) 
+    ? [...props.modelValue] 
+    : []
+  
   dialog.value = true
-  await checkPluginAndLoadKnowledgeBases()
+  await loadKnowledgeBases()
 }
 
-async function checkPluginAndLoadKnowledgeBases() {
+async function loadKnowledgeBases() {
   loading.value = true
   try {
-    // 首先检查插件是否安装
-    const pluginResponse = await axios.get('/api/plugin/get?name=astrbot_plugin_knowledge_base')
+    const response = await axios.get('/api/kb/list', {
+      params: {
+        page: 1,
+        page_size: 100
+      }
+    })
     
-    if (pluginResponse.data.status === 'ok' && pluginResponse.data.data.length > 0) {
-      pluginInstalled.value = true
-      // 插件已安装，获取知识库列表
-      await loadKnowledgeBases()
+    if (response.data.status === 'ok') {
+      knowledgeBaseList.value = response.data.data.items || []
     } else {
-      pluginInstalled.value = false
+      console.error('加载知识库列表失败:', response.data.message)
       knowledgeBaseList.value = []
     }
   } catch (error) {
-    console.error('检查知识库插件失败:', error)
-    pluginInstalled.value = false
+    console.error('加载知识库列表失败:', error)
     knowledgeBaseList.value = []
   } finally {
     loading.value = false
   }
 }
 
-async function loadKnowledgeBases() {
-  try {
-    const response = await axios.get('/api/plug/alkaid/kb/collections')
-    if (response.data.status === 'ok') {
-      knowledgeBaseList.value = response.data.data || []
-    } else {
-      knowledgeBaseList.value = []
-    }
-  } catch (error) {
-    console.error('加载知识库列表失败:', error)
-    knowledgeBaseList.value = []
+function isSelected(kbName) {
+  return selectedKnowledgeBases.value.includes(kbName)
+}
+
+function selectKnowledgeBase(kbName) {
+  // 多选模式：切换选中状态
+  const index = selectedKnowledgeBases.value.indexOf(kbName)
+  if (index > -1) {
+    selectedKnowledgeBases.value.splice(index, 1)
+  } else {
+    selectedKnowledgeBases.value.push(kbName)
   }
 }
 
-function selectKnowledgeBase(kb) {
-  selectedKnowledgeBase.value = kb.collection_name
+function removeKnowledgeBase(kbName) {
+  const index = selectedKnowledgeBases.value.indexOf(kbName)
+  if (index > -1) {
+    selectedKnowledgeBases.value.splice(index, 1)
+  }
+  
+  // 立即更新父组件
+  emit('update:modelValue', [...selectedKnowledgeBases.value])
 }
 
 function confirmSelection() {
-  emit('update:modelValue', selectedKnowledgeBase.value)
+  emit('update:modelValue', [...selectedKnowledgeBases.value])
   dialog.value = false
 }
 
 function cancelSelection() {
-  selectedKnowledgeBase.value = props.modelValue || ''
+  // 恢复到原始值
+  selectedKnowledgeBases.value = Array.isArray(props.modelValue) 
+    ? [...props.modelValue] 
+    : []
   dialog.value = false
 }
 
 function goToKnowledgeBasePage() {
   dialog.value = false
-  router.push('/alkaid/knowledge-base')
+  router.push('/knowledge-base')
 }
 </script>
 
@@ -222,4 +216,8 @@ function goToKnowledgeBasePage() {
   align-items: center;
   justify-content: center;
 }
+
+.gap-1 {
+  gap: 4px;
+}
 </style>