diff --git a/langgraph/prebuilt/chat_agent_executor.py b/langgraph/prebuilt/chat_agent_executor.py index eebb6998c..253473416 100644 --- a/langgraph/prebuilt/chat_agent_executor.py +++ b/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, diff --git a/langgraph/prebuilt/tool_executor.py b/langgraph/prebuilt/tool_executor.py index 82fe5e88a..b82e15653 100644 --- a/langgraph/prebuilt/tool_executor.py +++ b/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 \ No newline at end of file + return output