From ccb9f0b0d71f5360cedc28d6c5d2af6e2d9d0f1c Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?=E5=A4=A9=E6=B5=B7=E8=92=BC=E7=81=86?= Date: Tue, 19 Aug 2025 09:42:39 +0800 Subject: [PATCH] Feature (agent): Allow Retrieval kb_ids param use kb_id,and allow list kb_name or kb_id (#9531) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit ### What problem does this PR solve? Allow Retrieval kb_ids param use kb_id,and allow list kb_name or kb_id。 - Add judgment on whether the knowledge base name is a list and support batch queries -When the knowledge base name does not exist, try using the ID for querying -If both query methods fail, throw an exception ### Type of change - [x] New Feature (non-breaking change which adds functionality) --- agent/tools/retrieval.py | 14 ++++++++++---- 1 file changed, 10 insertions(+), 4 deletions(-) diff --git a/agent/tools/retrieval.py b/agent/tools/retrieval.py index 655f1e264..27ddfe4b2 100644 --- a/agent/tools/retrieval.py +++ b/agent/tools/retrieval.py @@ -86,10 +86,16 @@ class Retrieval(ToolBase, ABC): kb_ids.append(id) continue kb_nm = self._canvas.get_variable_value(id) - e, kb = KnowledgebaseService.get_by_name(kb_nm, self._canvas._tenant_id) - if not e: - raise Exception(f"Dataset({kb_nm}) does not exist.") - kb_ids.append(kb.id) + # if kb_nm is a list + kb_nm_list = kb_nm if isinstance(kb_nm, list) else [kb_nm] + for nm_or_id in kb_nm_list: + e, kb = KnowledgebaseService.get_by_name(nm_or_id, + self._canvas._tenant_id) + if not e: + e, kb = KnowledgebaseService.get_by_id(nm_or_id) + if not e: + raise Exception(f"Dataset({nm_or_id}) does not exist.") + kb_ids.append(kb.id) filtered_kb_ids: list[str] = list(set([kb_id for kb_id in kb_ids if kb_id]))