agent refactors (#997)

langchain/agents/agent.py

@@ -375,35 +375,22 @@ class AgentExecutor(Chain, BaseModel):
             final_output["intermediate_steps"] = intermediate_steps
         return final_output
 
-    def _call(self, inputs: Dict[str, str]) -> Dict[str, Any]:
-        """Run text through and get agent response."""
-        # Make sure that every tool is synchronous (not a coroutine)
-        for tool in self.tools:
-            if asyncio.iscoroutinefunction(tool.func):
-                raise ValueError(
-                    "Tools cannot be asynchronous for `run` method. "
-                    "Please use `arun` instead."
-                )
+    def _take_next_step(
+        self,
+        name_to_tool_map: Dict[str, Tool],
+        color_mapping: Dict[str, str],
+        inputs: Dict[str, str],
+        intermediate_steps: List[Tuple[AgentAction, str]],
+    ) -> Union[AgentFinish, Tuple[AgentAction, str]]:
+        """Take a single step in the thought-action-observation loop.
 
-        # Do any preparation necessary when receiving a new input.
-        self.agent.prepare_for_new_call()
-        # Construct a mapping of tool name to tool for easy lookup
-        name_to_tool_map = {tool.name: tool for tool in self.tools}
-        # We construct a mapping from each tool to a color, used for logging.
-        color_mapping = get_color_mapping(
-            [tool.name for tool in self.tools], excluded_colors=["green"]
-        )
-        intermediate_steps: List[Tuple[AgentAction, str]] = []
-        # Let's start tracking the iterations the agent has gone through
-        iterations = 0
-        # We now enter the agent loop (until it returns something).
-        while self._should_continue(iterations):
+        Override this to take control of how the agent makes and acts on choices.
+        """
         # Call the LLM to see what to do.
         output = self.agent.plan(intermediate_steps, **inputs)
         # If the tool chosen is the finishing tool, then we end and return.
         if isinstance(output, AgentFinish):
-                return self._return(output, intermediate_steps)
-
+            return output
         # Otherwise we lookup the tool
         if output.tool in name_to_tool_map:
             tool = name_to_tool_map[output.tool]
@@ -436,11 +423,41 @@ class AgentExecutor(Chain, BaseModel):
             llm_prefix=llm_prefix,
             verbose=self.verbose,
         )
-            intermediate_steps.append((output, observation))
         if return_direct:
             # Set the log to "" because we do not want to log it.
-                output = AgentFinish({self.agent.return_values[0]: observation}, "")
-                return self._return(output, intermediate_steps)
+            return AgentFinish({self.agent.return_values[0]: observation}, "")
+        return output, observation
+
+    def _call(self, inputs: Dict[str, str]) -> Dict[str, Any]:
+        """Run text through and get agent response."""
+        # Make sure that every tool is synchronous (not a coroutine)
+        for tool in self.tools:
+            if asyncio.iscoroutinefunction(tool.func):
+                raise ValueError(
+                    "Tools cannot be asynchronous for `run` method. "
+                    "Please use `arun` instead."
+                )
+
+        # Do any preparation necessary when receiving a new input.
+        self.agent.prepare_for_new_call()
+        # Construct a mapping of tool name to tool for easy lookup
+        name_to_tool_map = {tool.name: tool for tool in self.tools}
+        # We construct a mapping from each tool to a color, used for logging.
+        color_mapping = get_color_mapping(
+            [tool.name for tool in self.tools], excluded_colors=["green"]
+        )
+        intermediate_steps: List[Tuple[AgentAction, str]] = []
+        # Let's start tracking the iterations the agent has gone through
+        iterations = 0
+        # We now enter the agent loop (until it returns something).
+        while self._should_continue(iterations):
+            next_step_output = self._take_next_step(
+                name_to_tool_map, color_mapping, inputs, intermediate_steps
+            )
+            if isinstance(next_step_output, AgentFinish):
+                return self._return(next_step_output, intermediate_steps)
+
+            intermediate_steps.append(next_step_output)
             iterations += 1
         output = self.agent.return_stopped_response(
             self.early_stopping_method, intermediate_steps, **inputs