fix: dataset query when single dataset

api/core/agent/agent/multi_dataset_router_agent.py

-from typing import Tuple, List, Any, Union, Sequence, Optional
+from typing import Tuple, List, Any, Union, Sequence, Optional, cast
 
 from langchain.agents import OpenAIFunctionsAgent, BaseSingleActionAgent
 from langchain.callbacks.base import BaseCallbackManager
@@ -42,7 +42,9 @@ class MultiDatasetRouterAgent(OpenAIFunctionsAgent):
         if len(self.tools) == 0:
             return AgentFinish(return_values={"output": ''}, log='')
         elif len(self.tools) == 1:
-            rst = next(iter(self.tools)).run(kwargs['input'])
+            tool = next(iter(self.tools))
+            tool = cast(DatasetRetrieverTool, tool)
+            rst = tool.run(tool_input={'dataset_id': tool.dataset_id, 'query': kwargs['input']})
             return AgentFinish(return_values={"output": rst}, log=rst)
 
         if intermediate_steps:

api/core/completion.py

@@ -103,6 +103,7 @@ class Completion:
             #     # todo streaming flush the agent result to user, not call final llm
             #     pass
 
+            # todo or use fake llm
             cls.run_final_llm(
                 tenant_id=app.tenant_id,
                 mode=app.mode,

api/core/tool/dataset_retriever_tool.py

@@ -27,6 +27,7 @@ class DatasetRetrieverTool(BaseTool):
     description: str = "use this to retrieve a dataset. "
 
     tenant_id: str
+    dataset_id: str
     k: int = 3
 
     @classmethod
@@ -38,6 +39,7 @@ class DatasetRetrieverTool(BaseTool):
         description += '\nID of dataset MUST be ' + dataset.id
         return cls(
             tenant_id=dataset.tenant_id,
+            dataset_id=dataset.id,
             description=description,
             **kwargs
         )