[ContextAugmentedFunctionAgent] multi tool function call with retrieval

erniebot-agent/erniebot_agent/agents/__init__.py

@@ -15,6 +15,8 @@
 from .base import Agent
 from .functional_agent import FunctionalAgent
 from .functional_agent_with_retrieval import (
+    ContextAugmentedFunctionalAgent,
+    FunctionalAgentWithQueryPlanning,
     FunctionalAgentWithRetrieval,
     FunctionalAgentWithRetrievalScoreTool,
     FunctionalAgentWithRetrievalTool,

erniebot-agent/erniebot_agent/agents/functional_agent_with_retrieval.py

@@ -146,6 +146,20 @@ async def _async_run(self, prompt: str, files: Optional[List[File]] = None) -> A
 
             chat_history.append(HumanMessage(content=prompt))
 
+            step_input = HumanMessage(
+                content=self.rag_prompt.format(query=prompt, documents=results["documents"])
+            )
+            fake_chat_history: List[Message] = []
+            fake_chat_history.append(step_input)
+            llm_resp = await self._async_run_llm_without_hooks(
+                messages=fake_chat_history,
+                functions=None,
+                system=self.system_message.content if self.system_message is not None else None,
+            )
+
+            # Get RAG results
+            output_message = llm_resp.message
+
             outputs = []
             for item in results["documents"]:
                 outputs.append(
@@ -172,7 +186,8 @@ async def _async_run(self, prompt: str, files: Optional[List[File]] = None) -> A
             actions_taken.append(action)
             # return response
             tool_ret_json = json.dumps({"documents": outputs}, ensure_ascii=False)
-            next_step_input = FunctionMessage(name=action.tool_name, content=tool_ret_json)
+            # next_step_input = FunctionMessage(name=action.tool_name, content=tool_ret_json)
+            next_step_input = FunctionMessage(name=action.tool_name, content=output_message.content)
             tool_resp = ToolResponse(json=tool_ret_json, files=[])
             await self._callback_manager.on_tool_end(agent=self, tool=self.search_tool, response=tool_resp)
 
@@ -240,7 +255,7 @@ async def _async_run(self, prompt: str, files: Optional[List[File]] = None) -> A
             await self._callback_manager.on_tool_start(
                 agent=self, tool=self.search_tool, input_args=tool_args
             )
-            chat_history.append(HumanMessage(content=prompt))
+
             outputs = []
             for item in results["documents"]:
                 outputs.append(
@@ -251,25 +266,97 @@ async def _async_run(self, prompt: str, files: Optional[List[File]] = None) -> A
                     }
                 )
 
-            chat_history.append(
-                AIMessage(
-                    content="",
-                    function_call={
-                        "name": "KnowledgeBaseTool",
-                        "thoughts": "这是一个检索的需求，我需要在KnowledgeBaseTool知识库中检索出与输入的query相关的段落，并返回给用户。",
-                        "arguments": tool_args,
-                    },
+            # return response
+            tool_ret_json = json.dumps({"documents": outputs}, ensure_ascii=False)
+            # Direct Prompt
+            next_step_input = HumanMessage(content=f"问题：{prompt}，要求：请在第一步执行检索的操作,并且检索只允许调用一次")
+            tool_resp = ToolResponse(json=tool_ret_json, files=[])
+            await self._callback_manager.on_tool_end(agent=self, tool=self.search_tool, response=tool_resp)
+
+            num_steps_taken = 0
+            while num_steps_taken < self.max_steps:
+                curr_step_output = await self._async_step(
+                    next_step_input, chat_history, actions_taken, files_involved
                 )
+                if curr_step_output is None:
+                    response = self._create_finished_response(chat_history, actions_taken, files_involved)
+                    self.memory.add_message(chat_history[0])
+                    self.memory.add_message(chat_history[-1])
+                    return response
+                num_steps_taken += 1
+            response = self._create_stopped_response(chat_history, actions_taken, files_involved)
+            return response
+        else:
+            logger.info(
+                f"Irrelevant retrieval results. Fallbacking to FunctionalAgent for the query: {prompt}"
             )
+            return await super()._async_run(prompt)
+
+    async def _maybe_retrieval(
+        self,
+        step_input,
+    ):
+        documents = self.knowledge_base.search(step_input, top_k=self.top_k, filters=None)
+        documents = [item for item in documents if item["score"] > self.threshold]
+        results = {}
+        results["documents"] = documents
+        return results
+
+
+class ContextAugmentedFunctionalAgent(FunctionalAgent):
+    def __init__(self, knowledge_base: BaizhongSearch, top_k: int = 3, threshold: float = 0.1, **kwargs):
+        super().__init__(**kwargs)
+        self.knowledge_base = knowledge_base
+        self.top_k = top_k
+        self.threshold = threshold
+        self.rag_prompt = PromptTemplate(RAG_PROMPT, input_variables=["documents", "query"])
+        self.search_tool = KnowledgeBaseTool()
+
+    async def _async_run(self, prompt: str, files: Optional[List[File]] = None) -> AgentResponse:
+        results = await self._maybe_retrieval(prompt)
+        if len(results["documents"]) > 0:
+            # RAG
+            chat_history: List[Message] = []
+            actions_taken: List[AgentAction] = []
+            files_involved: List[AgentFile] = []
+
+            tool_args = json.dumps({"query": prompt}, ensure_ascii=False)
+            await self._callback_manager.on_tool_start(
+                agent=self, tool=self.search_tool, input_args=tool_args
+            )
+            step_input = HumanMessage(
+                content=self.rag_prompt.format(query=prompt, documents=results["documents"])
+            )
+            fake_chat_history: List[Message] = []
+            fake_chat_history.append(step_input)
+            llm_resp = await self._async_run_llm_without_hooks(
+                messages=fake_chat_history,
+                functions=None,
+                system=self.system_message.content if self.system_message is not None else None,
+            )
+
+            # Get RAG results
+            output_message = llm_resp.message
+
+            outputs = []
+            for item in results["documents"]:
+                outputs.append(
+                    {
+                        "id": item["id"],
+                        "title": item["title"],
+                        "document": item["content_se"],
+                    }
+                )
 
-            # Knowledge Retrieval Tool
-            action = AgentAction(tool_name="KnowledgeBaseTool", tool_args=tool_args)
-            actions_taken.append(action)
             # return response
             tool_ret_json = json.dumps({"documents": outputs}, ensure_ascii=False)
-            next_step_input = FunctionMessage(name=action.tool_name, content=tool_ret_json)
+            # 会有无限循环调用工具的问题
+            next_step_input = HumanMessage(
+                content=f"背景信息为：{output_message.content} \n 要求：选择相应的工具回答或者根据背景信息直接回答：{prompt}"
+            )
             tool_resp = ToolResponse(json=tool_ret_json, files=[])
             await self._callback_manager.on_tool_end(agent=self, tool=self.search_tool, response=tool_resp)
+
             num_steps_taken = 0
             while num_steps_taken < self.max_steps:
                 curr_step_output = await self._async_step(
@@ -298,3 +385,36 @@ async def _maybe_retrieval(
         results = {}
         results["documents"] = documents
         return results
+
+
+class FunctionalAgentWithQueryPlanning(FunctionalAgent):
+    def __init__(self, top_k: int = 3, threshold: float = 0.1, **kwargs):
+        super().__init__(**kwargs)
+        self.top_k = top_k
+        self.threshold = threshold
+        self.rag_prompt = PromptTemplate(RAG_PROMPT, input_variables=["documents", "query"])
+        # self.search_tool = KnowledgeBaseTool()
+
+    async def _async_run(self, prompt: str, files: Optional[List[File]] = None) -> AgentResponse:
+        # RAG
+        chat_history: List[Message] = []
+        actions_taken: List[AgentAction] = []
+        files_involved: List[AgentFile] = []
+        # 会有无限循环调用工具的问题
+        # next_step_input = HumanMessage(
+        #     content=f"请选择合适的工具来回答：{prompt}，如果需要的话，可以对把问题分解成子问题，然后每个子问题选择合适的工具回答。"
+        # )
+        next_step_input = HumanMessage(content=prompt)
+        num_steps_taken = 0
+        while num_steps_taken < self.max_steps:
+            curr_step_output = await self._async_step(
+                next_step_input, chat_history, actions_taken, files_involved
+            )
+            if curr_step_output is None:
+                response = self._create_finished_response(chat_history, actions_taken, files_involved)
+                self.memory.add_message(chat_history[0])
+                self.memory.add_message(chat_history[-1])
+                return response
+            num_steps_taken += 1
+        response = self._create_stopped_response(chat_history, actions_taken, files_involved)
+        return response

erniebot-agent/erniebot_agent/tools/baizhong_tool.py

@@ -1,6 +1,6 @@
 from __future__ import annotations
 
-from typing import Any, List, Optional, Type
+from typing import Any, Dict, List, Optional, Type
 
 from erniebot_agent.messages import AIMessage, HumanMessage
 from erniebot_agent.tools.schema import ToolParameterView
@@ -30,9 +30,10 @@ class BaizhongSearchTool(Tool):
     ouptut_type: Type[ToolParameterView] = BaizhongSearchToolOutputView
 
     def __init__(
-        self, description, db, threshold: float = 0.0, input_type=None, output_type=None, examples=None
+        self, name, description, db, threshold: float = 0.0, input_type=None, output_type=None, examples=None
     ) -> None:
         super().__init__()
+        self.name = name
         self.db = db
         self.description = description
         self.few_shot_examples = []
@@ -44,7 +45,7 @@ def __init__(
             self.few_shot_examples = examples
         self.threshold = threshold
 
-    async def __call__(self, query: str, top_k: int = 3, filters: Optional[dict[str, Any]] = None):
+    async def __call__(self, query: str, top_k: int = 3, filters: Optional[Dict[str, Any]] = None):
         documents = self.db.search(query, top_k, filters)
         documents = [item for item in documents if item["score"] > self.threshold]
         return {"documents": documents}

erniebot-agent/erniebot_agent/tools/openai_search_tool.py

@@ -0,0 +1,76 @@
+from __future__ import annotations
+
+from typing import Any, Dict, List, Optional, Type
+
+from erniebot_agent.messages import AIMessage, HumanMessage
+from erniebot_agent.tools.schema import ToolParameterView
+from pydantic import Field
+
+from .base import Tool
+
+
+class OpenAISearchToolInputView(ToolParameterView):
+    query: str = Field(description="查询语句")
+    top_k: int = Field(description="返回结果数量")
+
+
+class SearchResponseDocument(ToolParameterView):
+    title: str = Field(description="检索结果的标题")
+    document: str = Field(description="检索结果的内容")
+
+
+class OpenAISearchToolOutputView(ToolParameterView):
+    documents: List[SearchResponseDocument] = Field(description="检索结果，内容和用户输入query相关的段落")
+
+
+class OpenAISearchTool(Tool):
+    description: str = "在知识库中检索与用户输入query相关的段落"
+    input_type: Type[ToolParameterView] = OpenAISearchToolInputView
+    ouptut_type: Type[ToolParameterView] = OpenAISearchToolOutputView
+
+    def __init__(
+        self, name, description, db, threshold: float = 0.0, input_type=None, output_type=None, examples=None
+    ) -> None:
+        super().__init__()
+        self.name = name
+        self.db = db
+        self.description = description
+        self.few_shot_examples = []
+        if input_type is not None:
+            self.input_type = input_type
+        if output_type is not None:
+            self.ouptut_type = output_type
+        if examples is not None:
+            self.few_shot_examples = examples
+        self.threshold = threshold
+
+    async def __call__(self, query: str, top_k: int = 3, filters: Optional[Dict[str, Any]] = None):
+        documents = self.db.similarity_search_with_relevance_scores(query, top_k)
+        docs = []
+        for doc, score in documents:
+            if score > self.threshold:
+                docs.append(
+                    {"document": doc.page_content, "title": doc.metadata["source"], "meta": doc.metadata}
+                )
+
+        return {"documents": docs}
+
+    @property
+    def examples(
+        self,
+    ) -> List[Any]:
+        few_shot_objects: List[Any] = []
+        for item in self.few_shot_examples:
+            few_shot_objects.append(HumanMessage(item["user"]))
+            few_shot_objects.append(
+                AIMessage(
+                    "",
+                    function_call={
+                        "name": self.tool_name,
+                        "thoughts": item["thoughts"],
+                        "arguments": item["arguments"],
+                    },
+                )
+            )
+
+        return few_shot_objects