diff --git a/pkg/persistence/migrations/dbm010_pipeline_multi_knowledge_base.py b/pkg/persistence/migrations/dbm010_pipeline_multi_knowledge_base.py new file mode 100644 index 00000000..c28b64ed --- /dev/null +++ b/pkg/persistence/migrations/dbm010_pipeline_multi_knowledge_base.py @@ -0,0 +1,88 @@ +from .. import migration + +import sqlalchemy + +from ...entity.persistence import pipeline as persistence_pipeline + + +@migration.migration_class(10) +class DBMigratePipelineMultiKnowledgeBase(migration.DBMigration): + """Pipeline support multiple knowledge base binding""" + + async def upgrade(self): + """Upgrade""" + # read all pipelines + pipelines = await self.ap.persistence_mgr.execute_async(sqlalchemy.select(persistence_pipeline.LegacyPipeline)) + + for pipeline in pipelines: + serialized_pipeline = self.ap.persistence_mgr.serialize_model(persistence_pipeline.LegacyPipeline, pipeline) + + config = serialized_pipeline['config'] + + # Convert knowledge-base from string to array + if 'local-agent' in config['ai']: + current_kb = config['ai']['local-agent'].get('knowledge-base', '') + + # If it's already a list, skip + if isinstance(current_kb, list): + continue + + # Convert string to list + if current_kb and current_kb != '__none__': + config['ai']['local-agent']['knowledge-bases'] = [current_kb] + else: + config['ai']['local-agent']['knowledge-bases'] = [] + + # Remove old field + if 'knowledge-base' in config['ai']['local-agent']: + del config['ai']['local-agent']['knowledge-base'] + + await self.ap.persistence_mgr.execute_async( + sqlalchemy.update(persistence_pipeline.LegacyPipeline) + .where(persistence_pipeline.LegacyPipeline.uuid == serialized_pipeline['uuid']) + .values( + { + 'config': config, + 'for_version': self.ap.ver_mgr.get_current_version(), + } + ) + ) + + async def downgrade(self): + """Downgrade""" + # read all pipelines + pipelines = await self.ap.persistence_mgr.execute_async(sqlalchemy.select(persistence_pipeline.LegacyPipeline)) + + for pipeline in pipelines: + serialized_pipeline = self.ap.persistence_mgr.serialize_model(persistence_pipeline.LegacyPipeline, pipeline) + + config = serialized_pipeline['config'] + + # Convert knowledge-bases from array back to string + if 'local-agent' in config['ai']: + current_kbs = config['ai']['local-agent'].get('knowledge-bases', []) + + # If it's already a string, skip + if isinstance(current_kbs, str): + continue + + # Convert list to string (take first one or empty) + if current_kbs and len(current_kbs) > 0: + config['ai']['local-agent']['knowledge-base'] = current_kbs[0] + else: + config['ai']['local-agent']['knowledge-base'] = '' + + # Remove new field + if 'knowledge-bases' in config['ai']['local-agent']: + del config['ai']['local-agent']['knowledge-bases'] + + await self.ap.persistence_mgr.execute_async( + sqlalchemy.update(persistence_pipeline.LegacyPipeline) + .where(persistence_pipeline.LegacyPipeline.uuid == serialized_pipeline['uuid']) + .values( + { + 'config': config, + 'for_version': self.ap.ver_mgr.get_current_version(), + } + ) + ) diff --git a/pkg/provider/runners/localagent.py b/pkg/provider/runners/localagent.py index 7ab1e739..f89a4e1c 100644 --- a/pkg/provider/runners/localagent.py +++ b/pkg/provider/runners/localagent.py @@ -40,10 +40,14 @@ class LocalAgentRunner(runner.RequestRunner): """运行请求""" pending_tool_calls = [] - kb_uuid = query.pipeline_config['ai']['local-agent']['knowledge-base'] - - if kb_uuid == '__none__': - kb_uuid = None + # Get knowledge bases list (new field) + kb_uuids = query.pipeline_config['ai']['local-agent'].get('knowledge-bases', []) + + # Fallback to old field for backward compatibility + if not kb_uuids: + old_kb_uuid = query.pipeline_config['ai']['local-agent'].get('knowledge-base', '') + if old_kb_uuid and old_kb_uuid != '__none__': + kb_uuids = [old_kb_uuid] user_message = copy.deepcopy(query.user_message) @@ -57,21 +61,28 @@ class LocalAgentRunner(runner.RequestRunner): user_message_text += ce.text break - if kb_uuid and user_message_text: + if kb_uuids and user_message_text: # only support text for now - kb = await self.ap.rag_mgr.get_knowledge_base_by_uuid(kb_uuid) + all_results = [] + + # Retrieve from each knowledge base + for kb_uuid in kb_uuids: + kb = await self.ap.rag_mgr.get_knowledge_base_by_uuid(kb_uuid) - if not kb: - self.ap.logger.warning(f'Knowledge base {kb_uuid} not found') - raise ValueError(f'Knowledge base {kb_uuid} not found') + if not kb: + self.ap.logger.warning(f'Knowledge base {kb_uuid} not found, skipping') + continue - result = await kb.retrieve(user_message_text, kb.knowledge_base_entity.top_k) + result = await kb.retrieve(user_message_text, kb.knowledge_base_entity.top_k) + + if result: + all_results.extend(result) final_user_message_text = '' - if result: + if all_results: rag_context = '\n\n'.join( - f'[{i + 1}] {entry.metadata.get("text", "")}' for i, entry in enumerate(result) + f'[{i + 1}] {entry.metadata.get("text", "")}' for i, entry in enumerate(all_results) ) final_user_message_text = rag_combined_prompt_template.format( rag_context=rag_context, user_message=user_message_text diff --git a/pkg/utils/constants.py b/pkg/utils/constants.py index ee0fe9c9..5a7ab4fa 100644 --- a/pkg/utils/constants.py +++ b/pkg/utils/constants.py @@ -1,6 +1,6 @@ semantic_version = 'v4.4.1' -required_database_version = 9 +required_database_version = 10 """Tag the version of the database schema, used to check if the database needs to be migrated""" debug_mode = False diff --git a/templates/default-pipeline-config.json b/templates/default-pipeline-config.json index c5398e76..efbb9c3f 100644 --- a/templates/default-pipeline-config.json +++ b/templates/default-pipeline-config.json @@ -45,7 +45,7 @@ "content": "You are a helpful assistant." } ], - "knowledge-base": "" + "knowledge-bases": [] }, "dify-service-api": { "base-url": "https://api.dify.ai/v1", diff --git a/templates/metadata/pipeline/ai.yaml b/templates/metadata/pipeline/ai.yaml index e4d16a95..f6d54ee6 100644 --- a/templates/metadata/pipeline/ai.yaml +++ b/templates/metadata/pipeline/ai.yaml @@ -80,16 +80,16 @@ stages: zh_Hans: 除非您了解消息结构，否则请只使用 system 单提示词 type: prompt-editor required: true - - name: knowledge-base + - name: knowledge-bases label: - en_US: Knowledge Base + en_US: Knowledge Bases zh_Hans: 知识库 description: - en_US: Configure the knowledge base to use for the agent, if not selected, the agent will directly use the LLM to reply + en_US: Configure the knowledge bases to use for the agent, if not selected, the agent will directly use the LLM to reply zh_Hans: 配置用于提升回复质量的知识库，若不选择，则直接使用大模型回复 - type: knowledge-base-selector + type: knowledge-base-multi-selector required: false - default: '' + default: [] - name: tbox-app-api label: en_US: Tbox App API diff --git a/web/src/app/home/components/dynamic-form/DynamicFormComponent.tsx b/web/src/app/home/components/dynamic-form/DynamicFormComponent.tsx index 5cdd2ff7..dd2178f2 100644 --- a/web/src/app/home/components/dynamic-form/DynamicFormComponent.tsx +++ b/web/src/app/home/components/dynamic-form/DynamicFormComponent.tsx @@ -58,6 +58,9 @@ export default function DynamicFormComponent({ case 'knowledge-base-selector': fieldSchema = z.string(); break; + case 'knowledge-base-multi-selector': + fieldSchema = z.array(z.string()); + break; case 'bot-selector': fieldSchema = z.string(); break; diff --git a/web/src/app/home/components/dynamic-form/DynamicFormItemComponent.tsx b/web/src/app/home/components/dynamic-form/DynamicFormItemComponent.tsx index 03603c26..e078adb2 100644 --- a/web/src/app/home/components/dynamic-form/DynamicFormItemComponent.tsx +++ b/web/src/app/home/components/dynamic-form/DynamicFormItemComponent.tsx @@ -29,6 +29,15 @@ import { useTranslation } from 'react-i18next'; import { extractI18nObject } from '@/i18n/I18nProvider'; import { Textarea } from '@/components/ui/textarea'; import { Card, CardContent } from '@/components/ui/card'; +import { + Dialog, + DialogContent, + DialogHeader, + DialogTitle, + DialogFooter, +} from '@/components/ui/dialog'; +import { Checkbox } from '@/components/ui/checkbox'; +import { Plus, X } from 'lucide-react'; export default function DynamicFormItemComponent({ config, @@ -44,6 +53,8 @@ export default function DynamicFormItemComponent({ const [knowledgeBases, setKnowledgeBases] = useState([]); const [bots, setBots] = useState([]); const [uploading, setUploading] = useState(false); + const [kbDialogOpen, setKbDialogOpen] = useState(false); + const [tempSelectedKBIds, setTempSelectedKBIds] = useState([]); const { t } = useTranslation(); const handleFileUpload = async (file: File): Promise => { @@ -90,7 +101,10 @@ export default function DynamicFormItemComponent({ }, [config.type]); useEffect(() => { - if (config.type === DynamicFormItemType.KNOWLEDGE_BASE_SELECTOR) { + if ( + config.type === DynamicFormItemType.KNOWLEDGE_BASE_SELECTOR || + config.type === DynamicFormItemType.KNOWLEDGE_BASE_MULTI_SELECTOR + ) { httpClient .getKnowledgeBases() .then((resp) => { @@ -336,6 +350,128 @@ export default function DynamicFormItemComponent({ ); + case DynamicFormItemType.KNOWLEDGE_BASE_MULTI_SELECTOR: + return ( + <> +
+ {field.value && field.value.length > 0 ? ( +
+ {field.value.map((kbId: string) => { + const kb = knowledgeBases.find((base) => base.uuid === kbId); + if (!kb) return null; + return ( +
+
+
{kb.name}
+ {kb.description && ( +
+ {kb.description} +
+ )} +
+ +
+ ); + })} +
+ ) : ( +
+

+ {t('knowledge.noKnowledgeBaseSelected')} +

+
+ )} +
+ + + + {/* Knowledge Base Selection Dialog */} + + + + {t('knowledge.selectKnowledgeBases')} + +
+ {knowledgeBases.map((base) => { + const isSelected = tempSelectedKBIds.includes( + base.uuid ?? '', + ); + return ( +
{ + const kbId = base.uuid ?? ''; + setTempSelectedKBIds((prev) => + prev.includes(kbId) + ? prev.filter((id) => id !== kbId) + : [...prev, kbId], + ); + }} + > + +
+
{base.name}
+ {base.description && ( +
+ {base.description} +
+ )} +
+
+ ); + })} +
+ + + + + + + + ); + case DynamicFormItemType.BOT_SELECTOR: return (