Add async for chat agent exec

langgraph/prebuilt/chat_agent_executor.py

@@ -5,6 +5,7 @@ from typing import Annotated, Sequence, TypedDict
 from langchain.tools.render import format_tool_to_openai_function
 from langchain_core.agents import AgentAction
 from langchain_core.messages import BaseMessage, FunctionMessage
+from langchain_core.runnables import RunnableLambda
 
 from langgraph.graph import END, StateGraph
 from langgraph.prebuilt.tool_executor import ToolExecutor
@@ -17,11 +18,13 @@ def create_function_calling_executor(model, tools):
     else:
         tool_executor = ToolExecutor(tools)
         tool_classes = tools
-    model = model.bind_functions([format_tool_to_openai_function(t) for t in tool_classes])
+    model = model.bind_functions(
+        [format_tool_to_openai_function(t) for t in tool_classes]
+    )
 
     # Define the function that determines whether to continue or not
     def should_continue(state):
-        messages = state['messages']
+        messages = state["messages"]
         last_message = messages[-1]
         # If there is no function call, then we finish
         if "function_call" not in last_message.additional_kwargs:
@@ -32,23 +35,34 @@ def create_function_calling_executor(model, tools):
 
     # Define the function that calls the model
     def call_model(state):
-        messages = state['messages']
+        messages = state["messages"]
         response = model.invoke(messages)
         # We return a list, because this will get added to the existing list
         return {"messages": [response]}
 
+    async def acall_model(state):
+        messages = state["messages"]
+        response = await model.ainvoke(messages)
+        # We return a list, because this will get added to the existing list
+        return {"messages": [response]}
+
     # Define the function to execute tools
-    def call_tool(state):
-        messages = state['messages']
+    def _get_action(state):
+        messages = state["messages"]
         # Based on the continue condition
         # we know the last message involves a function call
         last_message = messages[-1]
         # We construct an AgentAction from the function_call
-        action = AgentAction(
+        return AgentAction(
             tool=last_message.additional_kwargs["function_call"]["name"],
-            tool_input=json.loads(last_message.additional_kwargs["function_call"]["arguments"]),
+            tool_input=json.loads(
+                last_message.additional_kwargs["function_call"]["arguments"]
+            ),
             log="",
         )
+
+    def call_tool(state):
+        action = _get_action(state)
         # We call the tool_executor and get back a response
         response = tool_executor.invoke(action)
         # We use the response to create a FunctionMessage
@@ -56,6 +70,15 @@ def create_function_calling_executor(model, tools):
         # We return a list, because this will get added to the existing list
         return {"messages": [function_message]}
 
+    async def acall_tool(state):
+        action = _get_action(state)
+        # We call the tool_executor and get back a response
+        response = await tool_executor.ainvoke(action)
+        # We use the response to create a FunctionMessage
+        function_message = FunctionMessage(content=str(response), name=action.tool)
+        # We return a list, because this will get added to the existing list
+        return {"messages": [function_message]}
+
     # We create the AgentState that we will pass around
     # This simply involves a list of messages
     # We want steps to return messages to append to the list
@@ -67,8 +90,8 @@ def create_function_calling_executor(model, tools):
     workflow = StateGraph(AgentState)
 
     # Define the two nodes we will cycle between
-    workflow.add_node("agent", call_model)
-    workflow.add_node("action", call_tool)
+    workflow.add_node("agent", RunnableLambda(call_model, acall_model))
+    workflow.add_node("action", RunnableLambda(call_tool, acall_tool))
 
     # Set the entrypoint as `agent`
     # This means that this node is the first one called
@@ -91,13 +114,13 @@ def create_function_calling_executor(model, tools):
             # If `tools`, then we call the tool node.
             "continue": "action",
             # Otherwise we finish.
-            "end": END
-        }
+            "end": END,
+        },
     )
 
     # We now add a normal edge from `tools` to `agent`.
     # This means that after `tools` is called, `agent` node is called next.
-    workflow.add_edge('action', 'agent')
+    workflow.add_edge("action", "agent")
 
     # Finally, we compile it!
     # This compiles it into a LangChain Runnable,

langgraph/prebuilt/tool_executor.py

@@ -1,9 +1,8 @@
-from typing import Any, Sequence
+from typing import Any, Sequence, Union
 
-from typing import Union
+from langchain_core.load.serializable import Serializable
 from langchain_core.runnables import RunnableBinding, RunnableLambda
 from langchain_core.tools import BaseTool
-from langchain_core.load.serializable import Serializable
 
 INVALID_TOOL_MSG_TEMPLATE = (
     "{requested_tool_name} is not a valid tool, "
@@ -13,6 +12,7 @@ INVALID_TOOL_MSG_TEMPLATE = (
 
 class ToolInvocationInterface:
     """Interface for invoking a tool"""
+
     tool: str
     tool_input: Union[str, dict]
 
@@ -25,21 +25,32 @@ class ToolInvocation(Serializable):
     tool_input: Union[str, dict]
     """The input to pass in to the Tool."""
 
-class ToolExecutor(RunnableBinding):
 
+class ToolExecutor(RunnableBinding):
     tools: Sequence[BaseTool]
     tool_map: dict
     invalid_tool_msg_template: str
-    def __init__(self, tools: Sequence[BaseTool], invalid_tool_msg_template: str = INVALID_TOOL_MSG_TEMPLATE, **kwargs: Any) -> None:
 
+    def __init__(
+        self,
+        tools: Sequence[BaseTool],
+        invalid_tool_msg_template: str = INVALID_TOOL_MSG_TEMPLATE,
+        **kwargs: Any,
+    ) -> None:
         bound = RunnableLambda(self._execute, afunc=self._aexecute)
-        super().__init__(bound=bound, tools=tools, tool_map ={t.name: t for t in tools}, invalid_tool_msg_template=invalid_tool_msg_template, **kwargs)
+        super().__init__(
+            bound=bound,
+            tools=tools,
+            tool_map={t.name: t for t in tools},
+            invalid_tool_msg_template=invalid_tool_msg_template,
+            **kwargs,
+        )
 
     def _execute(self, tool_invocation: ToolInvocationInterface) -> Any:
         if tool_invocation.tool not in self.tool_map:
             return self.invalid_tool_msg_template.format(
                 requested_tool_name=tool_invocation.tool,
-                available_tool_names_str=", ".join([t.name for t in self.tools])
+                available_tool_names_str=", ".join([t.name for t in self.tools]),
             )
         else:
             tool = self.tool_map[tool_invocation.tool]
@@ -50,9 +61,9 @@ class ToolExecutor(RunnableBinding):
         if tool_invocation.tool not in self.tool_map:
             return self.invalid_tool_msg_template.format(
                 requested_tool_name=tool_invocation.tool,
-                available_tool_names_str=", ".join([t.name for t in self.tools])
+                available_tool_names_str=", ".join([t.name for t in self.tools]),
             )
         else:
             tool = self.tool_map[tool_invocation.tool]
             output = await tool.ainvoke(tool_invocation.tool_input)
-            return output
+            return output