mirror of
https://github.com/infiniflow/ragflow.git
synced 2025-12-08 12:32:30 +08:00
Feature (agent): Allow Retrieval kb_ids param use kb_id,and allow list kb_name or kb_id (#9531)
### 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)
This commit is contained in:
@ -86,10 +86,16 @@ class Retrieval(ToolBase, ABC):
|
|||||||
kb_ids.append(id)
|
kb_ids.append(id)
|
||||||
continue
|
continue
|
||||||
kb_nm = self._canvas.get_variable_value(id)
|
kb_nm = self._canvas.get_variable_value(id)
|
||||||
e, kb = KnowledgebaseService.get_by_name(kb_nm, self._canvas._tenant_id)
|
# if kb_nm is a list
|
||||||
if not e:
|
kb_nm_list = kb_nm if isinstance(kb_nm, list) else [kb_nm]
|
||||||
raise Exception(f"Dataset({kb_nm}) does not exist.")
|
for nm_or_id in kb_nm_list:
|
||||||
kb_ids.append(kb.id)
|
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]))
|
filtered_kb_ids: list[str] = list(set([kb_id for kb_id in kb_ids if kb_id]))
|
||||||
|
|
||||||
|
|||||||
Reference in New Issue
Block a user