From dca21078ad6fa6bbaf528580ef507c5b853f440f Mon Sep 17 00:00:00 2001 From: Ankush Gola <9536492+agola11@users.noreply.github.com> Date: Fri, 7 Apr 2023 16:23:03 +0200 Subject: [PATCH] Run tools concurrently in `_atake_next_step` (#2537) small refactor to allow this --- langchain/agents/agent.py | 16 ++++++++++++---- 1 file changed, 12 insertions(+), 4 deletions(-) diff --git a/langchain/agents/agent.py b/langchain/agents/agent.py index 447f6a10..f33a065f 100644 --- a/langchain/agents/agent.py +++ b/langchain/agents/agent.py @@ -1,6 +1,7 @@ """Chain that takes in an input and produces an action and action input.""" from __future__ import annotations +import asyncio import json import logging import time @@ -749,8 +750,10 @@ class AgentExecutor(Chain): actions = [output] else: actions = output - result = [] - for agent_action in actions: + + async def _aperform_agent_action( + agent_action: AgentAction, + ) -> Tuple[AgentAction, str]: if self.callback_manager.is_async: await self.callback_manager.on_agent_action( agent_action, verbose=self.verbose, color="green" @@ -782,9 +785,14 @@ class AgentExecutor(Chain): color=None, **tool_run_kwargs, ) - result.append((agent_action, observation)) + return agent_action, observation - return result + # Use asyncio.gather to run multiple tool.arun() calls concurrently + result = await asyncio.gather( + *[_aperform_agent_action(agent_action) for agent_action in actions] + ) + + return list(result) def _call(self, inputs: Dict[str, str]) -> Dict[str, Any]: """Run text through and get agent response."""