langchain 0.3.26__py3-none-any.whl → 0.4.0.dev0__py3-none-any.whl

langchain/agents/agent.py

@@ -4,6 +4,7 @@ from __future__ import annotations
 
 import asyncio
 import builtins
+import contextlib
 import json
 import logging
 import time
@@ -41,7 +42,7 @@ from langchain_core.runnables.utils import AddableDict
 from langchain_core.tools import BaseTool
 from langchain_core.utils.input import get_color_mapping
 from pydantic import BaseModel, ConfigDict, model_validator
-from typing_extensions import Self
+from typing_extensions import Self, override
 
 from langchain._api.deprecation import AGENT_DEPRECATION_WARNING
 from langchain.agents.agent_iterator import AgentExecutorIterator
@@ -63,6 +64,7 @@ class BaseSingleActionAgent(BaseModel):
         return ["output"]
 
     def get_allowed_tools(self) -> Optional[list[str]]:
+        """Get allowed tools."""
         return None
 
     @abstractmethod
@@ -114,8 +116,8 @@ class BaseSingleActionAgent(BaseModel):
     def return_stopped_response(
         self,
         early_stopping_method: str,
-        intermediate_steps: list[tuple[AgentAction, str]],
-        **kwargs: Any,
+        intermediate_steps: list[tuple[AgentAction, str]],  # noqa: ARG002
+        **_: Any,
     ) -> AgentFinish:
         """Return response when agent has been stopped due to max iterations.
 
@@ -123,7 +125,6 @@ class BaseSingleActionAgent(BaseModel):
             early_stopping_method: Method to use for early stopping.
             intermediate_steps: Steps the LLM has taken to date,
                 along with observations.
-            **kwargs: User inputs.
 
         Returns:
             AgentFinish: Agent finish object.
@@ -134,12 +135,11 @@ class BaseSingleActionAgent(BaseModel):
         if early_stopping_method == "force":
             # `force` just returns a constant string
             return AgentFinish(
-                {"output": "Agent stopped due to iteration limit or time limit."}, ""
-            )
-        else:
-            raise ValueError(
-                f"Got unsupported early_stopping_method `{early_stopping_method}`"
+                {"output": "Agent stopped due to iteration limit or time limit."},
+                "",
             )
+        msg = f"Got unsupported early_stopping_method `{early_stopping_method}`"
+        raise ValueError(msg)
 
     @classmethod
     def from_llm_and_tools(
@@ -167,6 +167,7 @@ class BaseSingleActionAgent(BaseModel):
         """Return Identifier of an agent type."""
         raise NotImplementedError
 
+    @override
     def dict(self, **kwargs: Any) -> builtins.dict:
         """Return dictionary representation of agent.
 
@@ -195,12 +196,10 @@ class BaseSingleActionAgent(BaseModel):
 
             # If working with agent executor
             agent.agent.save(file_path="path/agent.yaml")
+
         """
         # Convert file to Path object.
-        if isinstance(file_path, str):
-            save_path = Path(file_path)
-        else:
-            save_path = file_path
+        save_path = Path(file_path) if isinstance(file_path, str) else file_path
 
         directory_path = save_path.parent
         directory_path.mkdir(parents=True, exist_ok=True)
@@ -208,16 +207,18 @@ class BaseSingleActionAgent(BaseModel):
         # Fetch dictionary to save
         agent_dict = self.dict()
         if "_type" not in agent_dict:
-            raise NotImplementedError(f"Agent {self} does not support saving")
+            msg = f"Agent {self} does not support saving"
+            raise NotImplementedError(msg)
 
         if save_path.suffix == ".json":
-            with open(file_path, "w") as f:
+            with save_path.open("w") as f:
                 json.dump(agent_dict, f, indent=4)
         elif save_path.suffix.endswith((".yaml", ".yml")):
-            with open(file_path, "w") as f:
+            with save_path.open("w") as f:
                 yaml.dump(agent_dict, f, default_flow_style=False)
         else:
-            raise ValueError(f"{save_path} must be json or yaml")
+            msg = f"{save_path} must be json or yaml"
+            raise ValueError(msg)
 
     def tool_run_logging_kwargs(self) -> builtins.dict:
         """Return logging kwargs for tool run."""
@@ -289,8 +290,8 @@ class BaseMultiActionAgent(BaseModel):
     def return_stopped_response(
         self,
         early_stopping_method: str,
-        intermediate_steps: list[tuple[AgentAction, str]],
-        **kwargs: Any,
+        intermediate_steps: list[tuple[AgentAction, str]],  # noqa: ARG002
+        **_: Any,
     ) -> AgentFinish:
         """Return response when agent has been stopped due to max iterations.
 
@@ -298,7 +299,6 @@ class BaseMultiActionAgent(BaseModel):
             early_stopping_method: Method to use for early stopping.
             intermediate_steps: Steps the LLM has taken to date,
                 along with observations.
-            **kwargs: User inputs.
 
         Returns:
             AgentFinish: Agent finish object.
@@ -309,23 +309,20 @@ class BaseMultiActionAgent(BaseModel):
         if early_stopping_method == "force":
             # `force` just returns a constant string
             return AgentFinish({"output": "Agent stopped due to max iterations."}, "")
-        else:
-            raise ValueError(
-                f"Got unsupported early_stopping_method `{early_stopping_method}`"
-            )
+        msg = f"Got unsupported early_stopping_method `{early_stopping_method}`"
+        raise ValueError(msg)
 
     @property
     def _agent_type(self) -> str:
         """Return Identifier of an agent type."""
         raise NotImplementedError
 
+    @override
     def dict(self, **kwargs: Any) -> builtins.dict:
         """Return dictionary representation of agent."""
         _dict = super().model_dump()
-        try:
+        with contextlib.suppress(NotImplementedError):
             _dict["_type"] = str(self._agent_type)
-        except NotImplementedError:
-            pass
         return _dict
 
     def save(self, file_path: Union[Path, str]) -> None:
@@ -343,29 +340,29 @@ class BaseMultiActionAgent(BaseModel):
 
             # If working with agent executor
             agent.agent.save(file_path="path/agent.yaml")
+
         """
         # Convert file to Path object.
-        if isinstance(file_path, str):
-            save_path = Path(file_path)
-        else:
-            save_path = file_path
+        save_path = Path(file_path) if isinstance(file_path, str) else file_path
 
         # Fetch dictionary to save
         agent_dict = self.dict()
         if "_type" not in agent_dict:
-            raise NotImplementedError(f"Agent {self} does not support saving.")
+            msg = f"Agent {self} does not support saving."
+            raise NotImplementedError(msg)
 
         directory_path = save_path.parent
         directory_path.mkdir(parents=True, exist_ok=True)
 
         if save_path.suffix == ".json":
-            with open(file_path, "w") as f:
+            with save_path.open("w") as f:
                 json.dump(agent_dict, f, indent=4)
         elif save_path.suffix.endswith((".yaml", ".yml")):
-            with open(file_path, "w") as f:
+            with save_path.open("w") as f:
                 yaml.dump(agent_dict, f, default_flow_style=False)
         else:
-            raise ValueError(f"{save_path} must be json or yaml")
+            msg = f"{save_path} must be json or yaml"
+            raise ValueError(msg)
 
     def tool_run_logging_kwargs(self) -> builtins.dict:
         """Return logging kwargs for tool run."""
@@ -382,7 +379,7 @@ class AgentOutputParser(BaseOutputParser[Union[AgentAction, AgentFinish]]):
 
 
 class MultiActionAgentOutputParser(
-    BaseOutputParser[Union[list[AgentAction], AgentFinish]]
+    BaseOutputParser[Union[list[AgentAction], AgentFinish]],
 ):
     """Base class for parsing agent output into agent actions/finish.
 
@@ -449,7 +446,7 @@ class RunnableAgent(BaseSingleActionAgent):
         Returns:
             Action specifying what tool to use.
         """
-        inputs = {**kwargs, **{"intermediate_steps": intermediate_steps}}
+        inputs = {**kwargs, "intermediate_steps": intermediate_steps}
         final_output: Any = None
         if self.stream_runnable:
             # Use streaming to make sure that the underlying LLM is invoked in a
@@ -488,7 +485,7 @@ class RunnableAgent(BaseSingleActionAgent):
         Returns:
             Action specifying what tool to use.
         """
-        inputs = {**kwargs, **{"intermediate_steps": intermediate_steps}}
+        inputs = {**kwargs, "intermediate_steps": intermediate_steps}
         final_output: Any = None
         if self.stream_runnable:
             # Use streaming to make sure that the underlying LLM is invoked in a
@@ -498,7 +495,8 @@ class RunnableAgent(BaseSingleActionAgent):
             # Because the response from the plan is not a generator, we need to
             # accumulate the output into final output and return that.
             async for chunk in self.runnable.astream(
-                inputs, config={"callbacks": callbacks}
+                inputs,
+                config={"callbacks": callbacks},
             ):
                 if final_output is None:
                     final_output = chunk
@@ -506,7 +504,8 @@ class RunnableAgent(BaseSingleActionAgent):
                     final_output += chunk
         else:
             final_output = await self.runnable.ainvoke(
-                inputs, config={"callbacks": callbacks}
+                inputs,
+                config={"callbacks": callbacks},
             )
         return final_output
 
@@ -565,7 +564,7 @@ class RunnableMultiActionAgent(BaseMultiActionAgent):
         Returns:
             Action specifying what tool to use.
         """
-        inputs = {**kwargs, **{"intermediate_steps": intermediate_steps}}
+        inputs = {**kwargs, "intermediate_steps": intermediate_steps}
         final_output: Any = None
         if self.stream_runnable:
             # Use streaming to make sure that the underlying LLM is invoked in a
@@ -604,7 +603,7 @@ class RunnableMultiActionAgent(BaseMultiActionAgent):
         Returns:
             Action specifying what tool to use.
         """
-        inputs = {**kwargs, **{"intermediate_steps": intermediate_steps}}
+        inputs = {**kwargs, "intermediate_steps": intermediate_steps}
         final_output: Any = None
         if self.stream_runnable:
             # Use streaming to make sure that the underlying LLM is invoked in a
@@ -614,7 +613,8 @@ class RunnableMultiActionAgent(BaseMultiActionAgent):
             # Because the response from the plan is not a generator, we need to
             # accumulate the output into final output and return that.
             async for chunk in self.runnable.astream(
-                inputs, config={"callbacks": callbacks}
+                inputs,
+                config={"callbacks": callbacks},
             ):
                 if final_output is None:
                     final_output = chunk
@@ -622,7 +622,8 @@ class RunnableMultiActionAgent(BaseMultiActionAgent):
                     final_output += chunk
         else:
             final_output = await self.runnable.ainvoke(
-                inputs, config={"callbacks": callbacks}
+                inputs,
+                config={"callbacks": callbacks},
             )
 
         return final_output
@@ -652,6 +653,7 @@ class LLMSingleActionAgent(BaseSingleActionAgent):
         """
         return list(set(self.llm_chain.input_keys) - {"intermediate_steps"})
 
+    @override
     def dict(self, **kwargs: Any) -> builtins.dict:
         """Return dictionary representation of agent."""
         _dict = super().dict()
@@ -736,6 +738,7 @@ class Agent(BaseSingleActionAgent):
     allowed_tools: Optional[list[str]] = None
     """Allowed tools for the agent. If None, all tools are allowed."""
 
+    @override
     def dict(self, **kwargs: Any) -> builtins.dict:
         """Return dictionary representation of agent."""
         _dict = super().dict()
@@ -751,17 +754,6 @@ class Agent(BaseSingleActionAgent):
         """Return values of the agent."""
         return ["output"]
 
-    def _fix_text(self, text: str) -> str:
-        """Fix the text.
-
-        Args:
-            text: Text to fix.
-
-        Returns:
-            str: Fixed text.
-        """
-        raise ValueError("fix_text not implemented for this agent.")
-
     @property
     def _stop(self) -> list[str]:
         return [
@@ -770,7 +762,8 @@ class Agent(BaseSingleActionAgent):
         ]
 
     def _construct_scratchpad(
-        self, intermediate_steps: list[tuple[AgentAction, str]]
+        self,
+        intermediate_steps: list[tuple[AgentAction, str]],
     ) -> Union[str, list[BaseMessage]]:
         """Construct the scratchpad that lets the agent continue its thought process."""
         thoughts = ""
@@ -819,11 +812,12 @@ class Agent(BaseSingleActionAgent):
         """
         full_inputs = self.get_full_inputs(intermediate_steps, **kwargs)
         full_output = await self.llm_chain.apredict(callbacks=callbacks, **full_inputs)
-        agent_output = await self.output_parser.aparse(full_output)
-        return agent_output
+        return await self.output_parser.aparse(full_output)
 
     def get_full_inputs(
-        self, intermediate_steps: list[tuple[AgentAction, str]], **kwargs: Any
+        self,
+        intermediate_steps: list[tuple[AgentAction, str]],
+        **kwargs: Any,
     ) -> builtins.dict[str, Any]:
         """Create the full inputs for the LLMChain from intermediate steps.
 
@@ -837,8 +831,7 @@ class Agent(BaseSingleActionAgent):
         """
         thoughts = self._construct_scratchpad(intermediate_steps)
         new_inputs = {"agent_scratchpad": thoughts, "stop": self._stop}
-        full_inputs = {**kwargs, **new_inputs}
-        return full_inputs
+        return {**kwargs, **new_inputs}
 
     @property
     def input_keys(self) -> list[str]:
@@ -866,7 +859,7 @@ class Agent(BaseSingleActionAgent):
         if "agent_scratchpad" not in prompt.input_variables:
             logger.warning(
                 "`agent_scratchpad` should be a variable in prompt.input_variables."
-                " Did not find it, so adding it at the end."
+                " Did not find it, so adding it at the end.",
             )
             prompt.input_variables.append("agent_scratchpad")
             if isinstance(prompt, PromptTemplate):
@@ -874,7 +867,8 @@ class Agent(BaseSingleActionAgent):
             elif isinstance(prompt, FewShotPromptTemplate):
                 prompt.suffix += "\n{agent_scratchpad}"
             else:
-                raise ValueError(f"Got unexpected prompt type {type(prompt)}")
+                msg = f"Got unexpected prompt type {type(prompt)}"
+                raise ValueError(msg)
         return self
 
     @property
@@ -907,8 +901,6 @@ class Agent(BaseSingleActionAgent):
             tools: Tools to use.
         """
 
-        pass
-
     @classmethod
     @abstractmethod
     def _get_default_output_parser(cls, **kwargs: Any) -> AgentOutputParser:
@@ -973,9 +965,10 @@ class Agent(BaseSingleActionAgent):
         if early_stopping_method == "force":
             # `force` just returns a constant string
             return AgentFinish(
-                {"output": "Agent stopped due to iteration limit or time limit."}, ""
+                {"output": "Agent stopped due to iteration limit or time limit."},
+                "",
             )
-        elif early_stopping_method == "generate":
+        if early_stopping_method == "generate":
             # Generate does one final forward pass
             thoughts = ""
             for action, observation in intermediate_steps:
@@ -995,15 +988,14 @@ class Agent(BaseSingleActionAgent):
             if isinstance(parsed_output, AgentFinish):
                 # If we can extract, we send the correct stuff
                 return parsed_output
-            else:
-                # If we can extract, but the tool is not the final tool,
-                # we just return the full output
-                return AgentFinish({"output": full_output}, full_output)
-        else:
-            raise ValueError(
-                "early_stopping_method should be one of `force` or `generate`, "
-                f"got {early_stopping_method}"
-            )
+            # If we can extract, but the tool is not the final tool,
+            # we just return the full output
+            return AgentFinish({"output": full_output}, full_output)
+        msg = (
+            "early_stopping_method should be one of `force` or `generate`, "
+            f"got {early_stopping_method}"
+        )
+        raise ValueError(msg)
 
     def tool_run_logging_kwargs(self) -> builtins.dict:
         """Return logging kwargs for tool run."""
@@ -1021,6 +1013,7 @@ class ExceptionTool(BaseTool):
     description: str = "Exception tool"
     """Description of the tool."""
 
+    @override
     def _run(
         self,
         query: str,
@@ -1028,6 +1021,7 @@ class ExceptionTool(BaseTool):
     ) -> str:
         return query
 
+    @override
     async def _arun(
         self,
         query: str,
@@ -1082,7 +1076,8 @@ class AgentExecutor(Chain):
       as an observation.
     """
     trim_intermediate_steps: Union[
-        int, Callable[[list[tuple[AgentAction, str]]], list[tuple[AgentAction, str]]]
+        int,
+        Callable[[list[tuple[AgentAction, str]]], list[tuple[AgentAction, str]]],
     ] = -1
     """How to trim the intermediate steps before returning them.
     Defaults to -1, which means no trimming.
@@ -1130,12 +1125,14 @@ class AgentExecutor(Chain):
         agent = self.agent
         tools = self.tools
         allowed_tools = agent.get_allowed_tools()  # type: ignore[union-attr]
-        if allowed_tools is not None:
-            if set(allowed_tools) != set([tool.name for tool in tools]):
-                raise ValueError(
-                    f"Allowed tools ({allowed_tools}) different than "
-                    f"provided tools ({[tool.name for tool in tools]})"
-                )
+        if allowed_tools is not None and set(allowed_tools) != {
+            tool.name for tool in tools
+        }:
+            msg = (
+                f"Allowed tools ({allowed_tools}) different than "
+                f"provided tools ({[tool.name for tool in tools]})"
+            )
+            raise ValueError(msg)
         return self
 
     @model_validator(mode="before")
@@ -1161,11 +1158,13 @@ class AgentExecutor(Chain):
             stream_runnable = values.pop("stream_runnable", True)
             if multi_action:
                 values["agent"] = RunnableMultiActionAgent(
-                    runnable=agent, stream_runnable=stream_runnable
+                    runnable=agent,
+                    stream_runnable=stream_runnable,
                 )
             else:
                 values["agent"] = RunnableAgent(
-                    runnable=agent, stream_runnable=stream_runnable
+                    runnable=agent,
+                    stream_runnable=stream_runnable,
                 )
         return values
 
@@ -1180,10 +1179,10 @@ class AgentExecutor(Chain):
         to reflect the changes made in the root_validator.
         """
         if isinstance(self.agent, Runnable):
-            return cast(RunnableAgentType, self.agent)
-        else:
-            return self.agent
+            return cast("RunnableAgentType", self.agent)
+        return self.agent
 
+    @override
     def save(self, file_path: Union[Path, str]) -> None:
         """Raise error - saving not supported for Agent Executors.
 
@@ -1193,11 +1192,12 @@ class AgentExecutor(Chain):
         Raises:
             ValueError: Saving not supported for agent executors.
         """
-        raise ValueError(
+        msg = (
             "Saving not supported for agent executors. "
             "If you are trying to save the agent, please use the "
             "`.save_agent(...)`"
         )
+        raise ValueError(msg)
 
     def save_agent(self, file_path: Union[Path, str]) -> None:
         """Save the underlying agent.
@@ -1213,7 +1213,7 @@ class AgentExecutor(Chain):
         callbacks: Callbacks = None,
         *,
         include_run_info: bool = False,
-        async_: bool = False,  # arg kept for backwards compat, but ignored
+        async_: bool = False,  # noqa: ARG002 arg kept for backwards compat, but ignored
     ) -> AgentExecutorIterator:
         """Enables iteration over steps taken to reach final output.
 
@@ -1249,9 +1249,8 @@ class AgentExecutor(Chain):
         :meta private:
         """
         if self.return_intermediate_steps:
-            return self._action_agent.return_values + ["intermediate_steps"]
-        else:
-            return self._action_agent.return_values
+            return [*self._action_agent.return_values, "intermediate_steps"]
+        return self._action_agent.return_values
 
     def lookup_tool(self, name: str) -> BaseTool:
         """Lookup tool by name.
@@ -1267,13 +1266,7 @@ class AgentExecutor(Chain):
     def _should_continue(self, iterations: int, time_elapsed: float) -> bool:
         if self.max_iterations is not None and iterations >= self.max_iterations:
             return False
-        if (
-            self.max_execution_time is not None
-            and time_elapsed >= self.max_execution_time
-        ):
-            return False
-
-        return True
+        return self.max_execution_time is None or time_elapsed < self.max_execution_time
 
     def _return(
         self,
@@ -1296,7 +1289,9 @@ class AgentExecutor(Chain):
     ) -> dict[str, Any]:
         if run_manager:
             await run_manager.on_agent_finish(
-                output, color="green", verbose=self.verbose
+                output,
+                color="green",
+                verbose=self.verbose,
             )
         final_output = output.return_values
         if self.return_intermediate_steps:
@@ -1304,15 +1299,15 @@ class AgentExecutor(Chain):
         return final_output
 
     def _consume_next_step(
-        self, values: NextStepOutput
+        self,
+        values: NextStepOutput,
     ) -> Union[AgentFinish, list[tuple[AgentAction, str]]]:
         if isinstance(values[-1], AgentFinish):
-            assert len(values) == 1
+            if len(values) != 1:
+                msg = "Expected a single AgentFinish output, but got multiple values."
+                raise ValueError(msg)
             return values[-1]
-        else:
-            return [
-                (a.action, a.observation) for a in values if isinstance(a, AgentStep)
-            ]
+        return [(a.action, a.observation) for a in values if isinstance(a, AgentStep)]
 
     def _take_next_step(
         self,
@@ -1323,16 +1318,15 @@ class AgentExecutor(Chain):
         run_manager: Optional[CallbackManagerForChainRun] = None,
     ) -> Union[AgentFinish, list[tuple[AgentAction, str]]]:
         return self._consume_next_step(
-            [
-                a
-                for a in self._iter_next_step(
+            list(
+                self._iter_next_step(
                     name_to_tool_map,
                     color_mapping,
                     inputs,
                     intermediate_steps,
                     run_manager,
-                )
-            ]
+                ),
+            ),
         )
 
     def _iter_next_step(
@@ -1362,12 +1356,13 @@ class AgentExecutor(Chain):
             else:
                 raise_error = False
             if raise_error:
-                raise ValueError(
+                msg = (
                     "An output parsing error occurred. "
                     "In order to pass this error back to the agent and have it try "
                     "again, pass `handle_parsing_errors=True` to the AgentExecutor. "
-                    f"This is the error: {str(e)}"
+                    f"This is the error: {e!s}"
                 )
+                raise ValueError(msg) from e
             text = str(e)
             if isinstance(self.handle_parsing_errors, bool):
                 if e.send_to_llm:
@@ -1380,7 +1375,8 @@ class AgentExecutor(Chain):
             elif callable(self.handle_parsing_errors):
                 observation = self.handle_parsing_errors(e)
             else:
-                raise ValueError("Got unexpected type of `handle_parsing_errors`")
+                msg = "Got unexpected type of `handle_parsing_errors`"
+                raise ValueError(msg) from e  # noqa: TRY004
             output = AgentAction("_Exception", observation, text)
             if run_manager:
                 run_manager.on_agent_action(output, color="green")
@@ -1401,15 +1397,15 @@ class AgentExecutor(Chain):
             return
 
         actions: list[AgentAction]
-        if isinstance(output, AgentAction):
-            actions = [output]
-        else:
-            actions = output
+        actions = [output] if isinstance(output, AgentAction) else output
         for agent_action in actions:
             yield agent_action
         for agent_action in actions:
             yield self._perform_agent_action(
-                name_to_tool_map, color_mapping, agent_action, run_manager
+                name_to_tool_map,
+                color_mapping,
+                agent_action,
+                run_manager,
             )
 
     def _perform_agent_action(
@@ -1469,7 +1465,7 @@ class AgentExecutor(Chain):
                     intermediate_steps,
                     run_manager,
                 )
-            ]
+            ],
         )
 
     async def _aiter_next_step(
@@ -1499,12 +1495,13 @@ class AgentExecutor(Chain):
             else:
                 raise_error = False
             if raise_error:
-                raise ValueError(
+                msg = (
                     "An output parsing error occurred. "
                     "In order to pass this error back to the agent and have it try "
                     "again, pass `handle_parsing_errors=True` to the AgentExecutor. "
-                    f"This is the error: {str(e)}"
+                    f"This is the error: {e!s}"
                 )
+                raise ValueError(msg) from e
             text = str(e)
             if isinstance(self.handle_parsing_errors, bool):
                 if e.send_to_llm:
@@ -1517,7 +1514,8 @@ class AgentExecutor(Chain):
             elif callable(self.handle_parsing_errors):
                 observation = self.handle_parsing_errors(e)
             else:
-                raise ValueError("Got unexpected type of `handle_parsing_errors`")
+                msg = "Got unexpected type of `handle_parsing_errors`"
+                raise ValueError(msg) from e  # noqa: TRY004
             output = AgentAction("_Exception", observation, text)
             tool_run_kwargs = self._action_agent.tool_run_logging_kwargs()
             observation = await ExceptionTool().arun(
@@ -1536,10 +1534,7 @@ class AgentExecutor(Chain):
             return
 
         actions: list[AgentAction]
-        if isinstance(output, AgentAction):
-            actions = [output]
-        else:
-            actions = output
+        actions = [output] if isinstance(output, AgentAction) else output
         for agent_action in actions:
             yield agent_action
 
@@ -1547,7 +1542,10 @@ class AgentExecutor(Chain):
         result = await asyncio.gather(
             *[
                 self._aperform_agent_action(
-                    name_to_tool_map, color_mapping, agent_action, run_manager
+                    name_to_tool_map,
+                    color_mapping,
+                    agent_action,
+                    run_manager,
                 )
                 for agent_action in actions
             ],
@@ -1566,7 +1564,9 @@ class AgentExecutor(Chain):
     ) -> AgentStep:
         if run_manager:
             await run_manager.on_agent_action(
-                agent_action, verbose=self.verbose, color="green"
+                agent_action,
+                verbose=self.verbose,
+                color="green",
             )
         # Otherwise we lookup the tool
         if agent_action.tool in name_to_tool_map:
@@ -1608,7 +1608,8 @@ class AgentExecutor(Chain):
         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", "red"]
+            [tool.name for tool in self.tools],
+            excluded_colors=["green", "red"],
         )
         intermediate_steps: list[tuple[AgentAction, str]] = []
         # Let's start tracking the number of iterations and time elapsed
@@ -1626,7 +1627,9 @@ class AgentExecutor(Chain):
             )
             if isinstance(next_step_output, AgentFinish):
                 return self._return(
-                    next_step_output, intermediate_steps, run_manager=run_manager
+                    next_step_output,
+                    intermediate_steps,
+                    run_manager=run_manager,
                 )
 
             intermediate_steps.extend(next_step_output)
@@ -1636,12 +1639,16 @@ class AgentExecutor(Chain):
                 tool_return = self._get_tool_return(next_step_action)
                 if tool_return is not None:
                     return self._return(
-                        tool_return, intermediate_steps, run_manager=run_manager
+                        tool_return,
+                        intermediate_steps,
+                        run_manager=run_manager,
                     )
             iterations += 1
             time_elapsed = time.time() - start_time
         output = self._action_agent.return_stopped_response(
-            self.early_stopping_method, intermediate_steps, **inputs
+            self.early_stopping_method,
+            intermediate_steps,
+            **inputs,
         )
         return self._return(output, intermediate_steps, run_manager=run_manager)
 
@@ -1655,7 +1662,8 @@ class AgentExecutor(Chain):
         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"]
+            [tool.name for tool in self.tools],
+            excluded_colors=["green"],
         )
         intermediate_steps: list[tuple[AgentAction, str]] = []
         # Let's start tracking the number of iterations and time elapsed
@@ -1687,28 +1695,39 @@ class AgentExecutor(Chain):
                         tool_return = self._get_tool_return(next_step_action)
                         if tool_return is not None:
                             return await self._areturn(
-                                tool_return, intermediate_steps, run_manager=run_manager
+                                tool_return,
+                                intermediate_steps,
+                                run_manager=run_manager,
                             )
 
                     iterations += 1
                     time_elapsed = time.time() - start_time
                 output = self._action_agent.return_stopped_response(
-                    self.early_stopping_method, intermediate_steps, **inputs
+                    self.early_stopping_method,
+                    intermediate_steps,
+                    **inputs,
                 )
                 return await self._areturn(
-                    output, intermediate_steps, run_manager=run_manager
+                    output,
+                    intermediate_steps,
+                    run_manager=run_manager,
                 )
         except (TimeoutError, asyncio.TimeoutError):
             # stop early when interrupted by the async timeout
             output = self._action_agent.return_stopped_response(
-                self.early_stopping_method, intermediate_steps, **inputs
+                self.early_stopping_method,
+                intermediate_steps,
+                **inputs,
             )
             return await self._areturn(
-                output, intermediate_steps, run_manager=run_manager
+                output,
+                intermediate_steps,
+                run_manager=run_manager,
             )
 
     def _get_tool_return(
-        self, next_step_output: tuple[AgentAction, str]
+        self,
+        next_step_output: tuple[AgentAction, str],
     ) -> Optional[AgentFinish]:
         """Check if the tool is a returning tool."""
         agent_action, observation = next_step_output
@@ -1717,27 +1736,30 @@ class AgentExecutor(Chain):
         if len(self._action_agent.return_values) > 0:
             return_value_key = self._action_agent.return_values[0]
         # Invalid tools won't be in the map, so we return False.
-        if agent_action.tool in name_to_tool_map:
-            if name_to_tool_map[agent_action.tool].return_direct:
-                return AgentFinish(
-                    {return_value_key: observation},
-                    "",
-                )
+        if (
+            agent_action.tool in name_to_tool_map
+            and name_to_tool_map[agent_action.tool].return_direct
+        ):
+            return AgentFinish(
+                {return_value_key: observation},
+                "",
+            )
         return None
 
     def _prepare_intermediate_steps(
-        self, intermediate_steps: list[tuple[AgentAction, str]]
+        self,
+        intermediate_steps: list[tuple[AgentAction, str]],
     ) -> list[tuple[AgentAction, str]]:
         if (
             isinstance(self.trim_intermediate_steps, int)
             and self.trim_intermediate_steps > 0
         ):
             return intermediate_steps[-self.trim_intermediate_steps :]
-        elif callable(self.trim_intermediate_steps):
+        if callable(self.trim_intermediate_steps):
             return self.trim_intermediate_steps(intermediate_steps)
-        else:
-            return intermediate_steps
+        return intermediate_steps
 
+    @override
     def stream(
         self,
         input: Union[dict[str, Any], Any],
@@ -1768,6 +1790,7 @@ class AgentExecutor(Chain):
         )
         yield from iterator
 
+    @override
     async def astream(
         self,
         input: Union[dict[str, Any], Any],