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."""