from __future__ import annotations import asyncio import logging import time from typing import ( TYPE_CHECKING, Any, AsyncIterator, Dict, Iterator, List, Optional, Tuple, Union, ) from uuid import UUID from langchain_core.agents import ( AgentAction, AgentFinish, AgentStep, ) from langchain_core.callbacks import ( AsyncCallbackManager, AsyncCallbackManagerForChainRun, CallbackManager, CallbackManagerForChainRun, Callbacks, ) from langchain_core.load.dump import dumpd from langchain_core.outputs import RunInfo from langchain_core.runnables.utils import AddableDict from langchain_core.tools import BaseTool from langchain_core.utils.input import get_color_mapping from langchain.schema import RUN_KEY from langchain.utilities.asyncio import asyncio_timeout if TYPE_CHECKING: from langchain.agents.agent import AgentExecutor, NextStepOutput logger = logging.getLogger(__name__) class AgentExecutorIterator: """Iterator for AgentExecutor.""" def __init__( self, agent_executor: AgentExecutor, inputs: Any, callbacks: Callbacks = None, *, tags: Optional[list[str]] = None, metadata: Optional[Dict[str, Any]] = None, run_name: Optional[str] = None, run_id: Optional[UUID] = None, include_run_info: bool = False, yield_actions: bool = False, ): """ Initialize the AgentExecutorIterator with the given AgentExecutor, inputs, and optional callbacks. Args: agent_executor (AgentExecutor): The AgentExecutor to iterate over. inputs (Any): The inputs to the AgentExecutor. callbacks (Callbacks, optional): The callbacks to use during iteration. Defaults to None. tags (Optional[list[str]], optional): The tags to use during iteration. Defaults to None. metadata (Optional[Dict[str, Any]], optional): The metadata to use during iteration. Defaults to None. run_name (Optional[str], optional): The name of the run. Defaults to None. run_id (Optional[UUID], optional): The ID of the run. Defaults to None. include_run_info (bool, optional): Whether to include run info in the output. Defaults to False. yield_actions (bool, optional): Whether to yield actions as they are generated. Defaults to False. """ self._agent_executor = agent_executor self.inputs = inputs self.callbacks = callbacks self.tags = tags self.metadata = metadata self.run_name = run_name self.run_id = run_id self.include_run_info = include_run_info self.yield_actions = yield_actions self.reset() _inputs: Dict[str, str] callbacks: Callbacks tags: Optional[list[str]] metadata: Optional[Dict[str, Any]] run_name: Optional[str] run_id: Optional[UUID] include_run_info: bool yield_actions: bool @property def inputs(self) -> Dict[str, str]: """The inputs to the AgentExecutor.""" return self._inputs @inputs.setter def inputs(self, inputs: Any) -> None: self._inputs = self.agent_executor.prep_inputs(inputs) @property def agent_executor(self) -> AgentExecutor: """The AgentExecutor to iterate over.""" return self._agent_executor @agent_executor.setter def agent_executor(self, agent_executor: AgentExecutor) -> None: self._agent_executor = agent_executor # force re-prep inputs in case agent_executor's prep_inputs fn changed self.inputs = self.inputs @property def name_to_tool_map(self) -> Dict[str, BaseTool]: """A mapping of tool names to tools.""" return {tool.name: tool for tool in self.agent_executor.tools} @property def color_mapping(self) -> Dict[str, str]: """A mapping of tool names to colors.""" return get_color_mapping( [tool.name for tool in self.agent_executor.tools], excluded_colors=["green", "red"], ) def reset(self) -> None: """ Reset the iterator to its initial state, clearing intermediate steps, iterations, and time elapsed. """ logger.debug("(Re)setting AgentExecutorIterator to fresh state") self.intermediate_steps: list[tuple[AgentAction, str]] = [] self.iterations = 0 # maybe better to start these on the first __anext__ call? self.time_elapsed = 0.0 self.start_time = time.time() def update_iterations(self) -> None: """ Increment the number of iterations and update the time elapsed. """ self.iterations += 1 self.time_elapsed = time.time() - self.start_time logger.debug( f"Agent Iterations: {self.iterations} ({self.time_elapsed:.2f}s elapsed)" ) def make_final_outputs( self, outputs: Dict[str, Any], run_manager: Union[CallbackManagerForChainRun, AsyncCallbackManagerForChainRun], ) -> AddableDict: # have access to intermediate steps by design in iterator, # so return only outputs may as well always be true. prepared_outputs = AddableDict( self.agent_executor.prep_outputs( self.inputs, outputs, return_only_outputs=True ) ) if self.include_run_info: prepared_outputs[RUN_KEY] = RunInfo(run_id=run_manager.run_id) return prepared_outputs def __iter__(self: "AgentExecutorIterator") -> Iterator[AddableDict]: logger.debug("Initialising AgentExecutorIterator") self.reset() callback_manager = CallbackManager.configure( self.callbacks, self.agent_executor.callbacks, self.agent_executor.verbose, self.tags, self.agent_executor.tags, self.metadata, self.agent_executor.metadata, ) run_manager = callback_manager.on_chain_start( dumpd(self.agent_executor), self.inputs, self.run_id, name=self.run_name, ) try: while self.agent_executor._should_continue( self.iterations, self.time_elapsed ): # take the next step: this plans next action, executes it, # yielding action and observation as they are generated next_step_seq: NextStepOutput = [] for chunk in self.agent_executor._iter_next_step( self.name_to_tool_map, self.color_mapping, self.inputs, self.intermediate_steps, run_manager, ): next_step_seq.append(chunk) # if we're yielding actions, yield them as they come # do not yield AgentFinish, which will be handled below if self.yield_actions: if isinstance(chunk, AgentAction): yield AddableDict(actions=[chunk], messages=chunk.messages) elif isinstance(chunk, AgentStep): yield AddableDict(steps=[chunk], messages=chunk.messages) # convert iterator output to format handled by _process_next_step_output next_step = self.agent_executor._consume_next_step(next_step_seq) # update iterations and time elapsed self.update_iterations() # decide if this is the final output output = self._process_next_step_output(next_step, run_manager) is_final = "intermediate_step" not in output # yield the final output always # for backwards compat, yield int. output if not yielding actions if not self.yield_actions or is_final: yield output # if final output reached, stop iteration if is_final: return except BaseException as e: run_manager.on_chain_error(e) raise # if we got here means we exhausted iterations or time yield self._stop(run_manager) async def __aiter__(self) -> AsyncIterator[AddableDict]: """ N.B. __aiter__ must be a normal method, so need to initialize async run manager on first __anext__ call where we can await it """ logger.debug("Initialising AgentExecutorIterator (async)") self.reset() callback_manager = AsyncCallbackManager.configure( self.callbacks, self.agent_executor.callbacks, self.agent_executor.verbose, self.tags, self.agent_executor.tags, self.metadata, self.agent_executor.metadata, ) run_manager = await callback_manager.on_chain_start( dumpd(self.agent_executor), self.inputs, self.run_id, name=self.run_name, ) try: async with asyncio_timeout(self.agent_executor.max_execution_time): while self.agent_executor._should_continue( self.iterations, self.time_elapsed ): # take the next step: this plans next action, executes it, # yielding action and observation as they are generated next_step_seq: NextStepOutput = [] async for chunk in self.agent_executor._aiter_next_step( self.name_to_tool_map, self.color_mapping, self.inputs, self.intermediate_steps, run_manager, ): next_step_seq.append(chunk) # if we're yielding actions, yield them as they come # do not yield AgentFinish, which will be handled below if self.yield_actions: if isinstance(chunk, AgentAction): yield AddableDict( actions=[chunk], messages=chunk.messages ) elif isinstance(chunk, AgentStep): yield AddableDict( steps=[chunk], messages=chunk.messages ) # convert iterator output to format handled by _process_next_step next_step = self.agent_executor._consume_next_step(next_step_seq) # update iterations and time elapsed self.update_iterations() # decide if this is the final output output = await self._aprocess_next_step_output( next_step, run_manager ) is_final = "intermediate_step" not in output # yield the final output always # for backwards compat, yield int. output if not yielding actions if not self.yield_actions or is_final: yield output # if final output reached, stop iteration if is_final: return except (TimeoutError, asyncio.TimeoutError): yield await self._astop(run_manager) return except BaseException as e: await run_manager.on_chain_error(e) raise # if we got here means we exhausted iterations or time yield await self._astop(run_manager) def _process_next_step_output( self, next_step_output: Union[AgentFinish, List[Tuple[AgentAction, str]]], run_manager: CallbackManagerForChainRun, ) -> AddableDict: """ Process the output of the next step, handling AgentFinish and tool return cases. """ logger.debug("Processing output of Agent loop step") if isinstance(next_step_output, AgentFinish): logger.debug( "Hit AgentFinish: _return -> on_chain_end -> run final output logic" ) return self._return(next_step_output, run_manager=run_manager) self.intermediate_steps.extend(next_step_output) logger.debug("Updated intermediate_steps with step output") # Check for tool return if len(next_step_output) == 1: next_step_action = next_step_output[0] tool_return = self.agent_executor._get_tool_return(next_step_action) if tool_return is not None: return self._return(tool_return, run_manager=run_manager) return AddableDict(intermediate_step=next_step_output) async def _aprocess_next_step_output( self, next_step_output: Union[AgentFinish, List[Tuple[AgentAction, str]]], run_manager: AsyncCallbackManagerForChainRun, ) -> AddableDict: """ Process the output of the next async step, handling AgentFinish and tool return cases. """ logger.debug("Processing output of async Agent loop step") if isinstance(next_step_output, AgentFinish): logger.debug( "Hit AgentFinish: _areturn -> on_chain_end -> run final output logic" ) return await self._areturn(next_step_output, run_manager=run_manager) self.intermediate_steps.extend(next_step_output) logger.debug("Updated intermediate_steps with step output") # Check for tool return if len(next_step_output) == 1: next_step_action = next_step_output[0] tool_return = self.agent_executor._get_tool_return(next_step_action) if tool_return is not None: return await self._areturn(tool_return, run_manager=run_manager) return AddableDict(intermediate_step=next_step_output) def _stop(self, run_manager: CallbackManagerForChainRun) -> AddableDict: """ Stop the iterator and raise a StopIteration exception with the stopped response. """ logger.warning("Stopping agent prematurely due to triggering stop condition") # this manually constructs agent finish with output key output = self.agent_executor._action_agent.return_stopped_response( self.agent_executor.early_stopping_method, self.intermediate_steps, **self.inputs, ) return self._return(output, run_manager=run_manager) async def _astop(self, run_manager: AsyncCallbackManagerForChainRun) -> AddableDict: """ Stop the async iterator and raise a StopAsyncIteration exception with the stopped response. """ logger.warning("Stopping agent prematurely due to triggering stop condition") output = self.agent_executor._action_agent.return_stopped_response( self.agent_executor.early_stopping_method, self.intermediate_steps, **self.inputs, ) return await self._areturn(output, run_manager=run_manager) def _return( self, output: AgentFinish, run_manager: CallbackManagerForChainRun ) -> AddableDict: """ Return the final output of the iterator. """ returned_output = self.agent_executor._return( output, self.intermediate_steps, run_manager=run_manager ) returned_output["messages"] = output.messages run_manager.on_chain_end(returned_output) return self.make_final_outputs(returned_output, run_manager) async def _areturn( self, output: AgentFinish, run_manager: AsyncCallbackManagerForChainRun ) -> AddableDict: """ Return the final output of the async iterator. """ returned_output = await self.agent_executor._areturn( output, self.intermediate_steps, run_manager=run_manager ) returned_output["messages"] = output.messages await run_manager.on_chain_end(returned_output) return self.make_final_outputs(returned_output, run_manager)