langchain[patch]: Extract _aperform_agent_action from _aiter_next_step from AgentExecutor (#15707)

- **Description:** extreact the _aperform_agent_action in the
AgentExecutor class to allow for easier overriding. Extracted logic from
_iter_next_step into a new method _perform_agent_action for consistency
and easier overriding.
- **Issue:** #15706

Closes #15706
pull/16487/head
Gianfranco Demarco 5 months ago committed by GitHub
parent 95ee69a301
commit c69f599594
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194

@ -1179,6 +1179,17 @@ class AgentExecutor(Chain):
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
)
def _perform_agent_action(
self,
name_to_tool_map: Dict[str, BaseTool],
color_mapping: Dict[str, str],
agent_action: AgentAction,
run_manager: Optional[CallbackManagerForChainRun] = None,
) -> AgentStep:
if run_manager:
run_manager.on_agent_action(agent_action, color="green")
# Otherwise we lookup the tool
@ -1209,7 +1220,7 @@ class AgentExecutor(Chain):
callbacks=run_manager.get_child() if run_manager else None,
**tool_run_kwargs,
)
yield AgentStep(action=agent_action, observation=observation)
return AgentStep(action=agent_action, observation=observation)
async def _atake_next_step(
self,
@ -1303,8 +1314,26 @@ class AgentExecutor(Chain):
for agent_action in actions:
yield agent_action
# Use asyncio.gather to run multiple tool.arun() calls concurrently
result = await asyncio.gather(
*[
self._aperform_agent_action(
name_to_tool_map, color_mapping, agent_action, run_manager
)
for agent_action in actions
],
)
# TODO This could yield each result as it becomes available
for chunk in result:
yield chunk
async def _aperform_agent_action(
self,
name_to_tool_map: Dict[str, BaseTool],
color_mapping: Dict[str, str],
agent_action: AgentAction,
run_manager: Optional[AsyncCallbackManagerForChainRun] = None,
) -> AgentStep:
if run_manager:
await run_manager.on_agent_action(
@ -1340,15 +1369,6 @@ class AgentExecutor(Chain):
)
return AgentStep(action=agent_action, observation=observation)
# Use asyncio.gather to run multiple tool.arun() calls concurrently
result = await asyncio.gather(
*[_aperform_agent_action(agent_action) for agent_action in actions]
)
# TODO This could yield each result as it becomes available
for chunk in result:
yield chunk
def _call(
self,
inputs: Dict[str, str],

Loading…
Cancel
Save