From e4bfa8603d602a184cb0a00135ff8ff6ca5ae11d Mon Sep 17 00:00:00 2001 From: Nuno Campos Date: Mon, 19 Feb 2024 09:12:08 -0800 Subject: [PATCH] Support multiple tool calls, Lint --- langgraph/prebuilt/chat_agent_executor.py | 120 +++--- tests/test_pregel.py | 408 +++++++-------------- tests/test_pregel_async.py | 423 +++++++--------------- 3 files changed, 328 insertions(+), 623 deletions(-) diff --git a/langgraph/prebuilt/chat_agent_executor.py b/langgraph/prebuilt/chat_agent_executor.py index c99d56e65..fcd4d6806 100644 --- a/langgraph/prebuilt/chat_agent_executor.py +++ b/langgraph/prebuilt/chat_agent_executor.py @@ -1,17 +1,23 @@ import json import operator -from typing import Annotated, Sequence, TypedDict +from typing import Annotated, Sequence, TypedDict, Union -from langchain_core.agents import AgentAction +from langchain_core.language_models import LanguageModelLike from langchain_core.messages import BaseMessage, FunctionMessage, ToolMessage from langchain_core.runnables import RunnableLambda -from langchain_core.utils.function_calling import convert_to_openai_function, convert_to_openai_tool +from langchain_core.tools import BaseTool +from langchain_core.utils.function_calling import ( + convert_to_openai_function, + convert_to_openai_tool, +) from langgraph.graph import END, StateGraph -from langgraph.prebuilt.tool_executor import ToolExecutor +from langgraph.prebuilt.tool_executor import ToolExecutor, ToolInvocation -def create_function_calling_executor(model, tools): +def create_function_calling_executor( + model: LanguageModelLike, tools: Union[ToolExecutor, Sequence[BaseTool]] +): if isinstance(tools, ToolExecutor): tool_executor = tools tool_classes = tools.tools @@ -20,8 +26,15 @@ def create_function_calling_executor(model, tools): tool_classes = tools model = model.bind(functions=[convert_to_openai_function(t) for t in tool_classes]) + # 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 + # So we annotate the messages attribute with operator.add + class AgentState(TypedDict): + messages: Annotated[Sequence[BaseMessage], operator.add] + # Define the function that determines whether to continue or not - def should_continue(state): + def should_continue(state: AgentState): messages = state["messages"] last_message = messages[-1] # If there is no function call, then we finish @@ -32,34 +45,33 @@ def create_function_calling_executor(model, tools): return "continue" # Define the function that calls the model - def call_model(state): + def call_model(state: AgentState): 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): + async def acall_model(state: AgentState): 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 _get_action(state): + def _get_action(state: AgentState): 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 - return AgentAction( + return ToolInvocation( tool=last_message.additional_kwargs["function_call"]["name"], tool_input=json.loads( last_message.additional_kwargs["function_call"]["arguments"] ), - log="", ) - def call_tool(state): + def call_tool(state: AgentState): action = _get_action(state) # We call the tool_executor and get back a response response = tool_executor.invoke(action) @@ -68,7 +80,7 @@ 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): + async def acall_tool(state: AgentState): action = _get_action(state) # We call the tool_executor and get back a response response = await tool_executor.ainvoke(action) @@ -77,13 +89,6 @@ def create_function_calling_executor(model, tools): # 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 - # So we annotate the messages attribute with operator.add - class AgentState(TypedDict): - messages: Annotated[Sequence[BaseMessage], operator.add] - # Define a new graph workflow = StateGraph(AgentState) @@ -125,17 +130,27 @@ def create_function_calling_executor(model, tools): # meaning you can use it as you would any other runnable return workflow.compile() -def create_tool_calling_executor(model, tools): + +def create_tool_calling_executor( + model: LanguageModelLike, tools: Union[ToolExecutor, Sequence[BaseTool]] +): if isinstance(tools, ToolExecutor): tool_executor = tools tool_classes = tools.tools else: tool_executor = ToolExecutor(tools) tool_classes = tools - model = model.bind(functions=[convert_to_openai_tool(t) for t in tool_classes]) + model = model.bind(tools=[convert_to_openai_tool(t) for t in tool_classes]) + + # 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 + # So we annotate the messages attribute with operator.add + class AgentState(TypedDict): + messages: Annotated[Sequence[BaseMessage], operator.add] # Define the function that determines whether to continue or not - def should_continue(state): + def should_continue(state: AgentState): messages = state["messages"] last_message = messages[-1] # If there is no function call, then we finish @@ -146,57 +161,62 @@ def create_tool_calling_executor(model, tools): return "continue" # Define the function that calls the model - def call_model(state): + def call_model(state: AgentState): 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): + async def acall_model(state: AgentState): 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 _get_action(state): + def _get_actions(state: AgentState): messages = state["messages"] # Based on the continue condition # we know the last message involves a tool call last_message = messages[-1] - # We construct an AgentAction from the tool_calls - return AgentAction( - tool=last_message.additional_kwargs["tool_calls"][0]["function"]["name"], - tool_input=json.loads( - last_message.additional_kwargs["tool_calls"][0]["function"]["arguments"] - ), - log=last_message.additional_kwargs["tool_calls"][0]["id"], + # We construct an AgentAction from each of the tool_calls + return ( + [ + ToolInvocation( + tool=tool_call["function"]["name"], + tool_input=json.loads(tool_call["function"]["arguments"]), + ) + for tool_call in last_message.additional_kwargs["tool_calls"] + ], + [ + tool_call["id"] + for tool_call in last_message.additional_kwargs["tool_calls"] + ], ) - def call_tool(state): - action = _get_action(state) + def call_tool(state: AgentState): + actions, ids = _get_actions(state) # We call the tool_executor and get back a response - response = tool_executor.invoke(action) + responses = tool_executor.batch(actions) # We use the response to create a FunctionMessage - tool_message = ToolMessage(content=str(response), tool_call_id=action.log) + tool_messages = [ + ToolMessage(content=str(response), tool_call_id=id) + for response, id in zip(responses, ids) + ] # We return a list, because this will get added to the existing list - return {"messages": [tool_message]} + return {"messages": tool_messages} - async def acall_tool(state): - action = _get_action(state) + async def acall_tool(state: AgentState): + actions, ids = _get_actions(state) # We call the tool_executor and get back a response - response = await tool_executor.ainvoke(action) + responses = await tool_executor.abatch(actions) # We use the response to create a FunctionMessage - tool_message = ToolMessage(content=str(response), tool_call_id=action.log) + tool_messages = [ + ToolMessage(content=str(response), tool_call_id=id) + for response, id in zip(responses, ids) + ] # We return a list, because this will get added to the existing list - return {"messages": [tool_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 - # So we annotate the messages attribute with operator.add - class AgentState(TypedDict): - messages: Annotated[Sequence[BaseMessage], operator.add] + return {"messages": tool_messages} # Define a new graph workflow = StateGraph(AgentState) diff --git a/tests/test_pregel.py b/tests/test_pregel.py index 4b1a78659..13427b9c6 100644 --- a/tests/test_pregel.py +++ b/tests/test_pregel.py @@ -20,11 +20,15 @@ from langgraph.checkpoint.sqlite import SqliteSaver from langgraph.graph import END, Graph from langgraph.graph.message import MessageGraph from langgraph.graph.state import StateGraph -from langgraph.prebuilt.chat_agent_executor import create_function_calling_executor, create_tool_calling_executor +from langgraph.prebuilt.chat_agent_executor import ( + create_function_calling_executor, + create_tool_calling_executor, +) from langgraph.prebuilt.tool_executor import ToolExecutor from langgraph.pregel import Channel, GraphRecursionError, Pregel from langgraph.pregel.reserved import ReservedChannels + def test_invoke_single_process_in_out(mocker: MockerFixture) -> None: add_one = mocker.Mock(side_effect=lambda x: x + 1) chain = Channel.subscribe_to("input") | add_one | Channel.write_to("output") @@ -1060,6 +1064,7 @@ def test_conditional_graph_state() -> None: }, ] + def test_prebuilt_tool_chat() -> None: from langchain.chat_models.fake import FakeMessagesListChatModel from langchain_community.tools import tool @@ -1082,27 +1087,39 @@ def test_prebuilt_tool_chat() -> None: AIMessage( content="", additional_kwargs={ - "tool_calls": [{ - "id": "tool_call123", - "type": "function", - "function":{ - "name": "search_api", - "arguments": json.dumps("query"), + "tool_calls": [ + { + "id": "tool_call123", + "type": "function", + "function": { + "name": "search_api", + "arguments": json.dumps("query"), + }, } - }] + ] }, ), AIMessage( content="", additional_kwargs={ - "tool_calls": [{ - "id": "tool_call234", - "type": "function", - "function":{ - "name": "search_api", - "arguments": json.dumps("another"), - } - }] + "tool_calls": [ + { + "id": "tool_call234", + "type": "function", + "function": { + "name": "search_api", + "arguments": json.dumps("another"), + }, + }, + { + "id": "tool_call567", + "type": "function", + "function": { + "name": "search_api", + "arguments": '"a third one"', + }, + }, + ] }, ), AIMessage(content="answer"), @@ -1110,7 +1127,7 @@ def test_prebuilt_tool_chat() -> None: ), tools, ) - + assert app.invoke( {"messages": [HumanMessage(content="what is weather in sf")]} ) == { @@ -1119,31 +1136,44 @@ def test_prebuilt_tool_chat() -> None: AIMessage( content="", additional_kwargs={ - "tool_calls": [{ - "id": "tool_call123", - "type": "function", - "function":{ - "name": "search_api", - "arguments": "\"query\"", + "tool_calls": [ + { + "id": "tool_call123", + "type": "function", + "function": { + "name": "search_api", + "arguments": '"query"', + }, } - }] + ] }, ), ToolMessage(content="result for query", tool_call_id="tool_call123"), AIMessage( content="", additional_kwargs={ - "tool_calls": [{ - "id": "tool_call234", - "type": "function", - "function":{ - "name": "search_api", - "arguments": "\"another\"", - } - }] - }, + "tool_calls": [ + { + "id": "tool_call234", + "type": "function", + "function": { + "name": "search_api", + "arguments": '"another"', + }, + }, + { + "id": "tool_call567", + "type": "function", + "function": { + "name": "search_api", + "arguments": '"a third one"', + }, + }, + ] + }, ), ToolMessage(content="result for another", tool_call_id="tool_call234"), + ToolMessage(content="result for a third one", tool_call_id="tool_call567"), AIMessage(content="answer"), ] } @@ -1157,14 +1187,16 @@ def test_prebuilt_tool_chat() -> None: AIMessage( content="", additional_kwargs={ - "tool_calls": [{ + "tool_calls": [ + { "id": "tool_call123", "type": "function", - "function":{ + "function": { "name": "search_api", - "arguments": "\"query\"", - } - }] + "arguments": '"query"', + }, + } + ] }, ) ] @@ -1182,16 +1214,26 @@ def test_prebuilt_tool_chat() -> None: "messages": [ AIMessage( content="", - additional_kwargs={ - "tool_calls": [{ + additional_kwargs={ + "tool_calls": [ + { "id": "tool_call234", "type": "function", - "function":{ + "function": { "name": "search_api", - "arguments": "\"another\"", - } - }] - }, + "arguments": '"another"', + }, + }, + { + "id": "tool_call567", + "type": "function", + "function": { + "name": "search_api", + "arguments": '"a third one"', + }, + }, + ] + }, ) ] } @@ -1199,7 +1241,12 @@ def test_prebuilt_tool_chat() -> None: { "action": { "messages": [ - ToolMessage(content="result for another", tool_call_id="tool_call234") + ToolMessage( + content="result for another", tool_call_id="tool_call234" + ), + ToolMessage( + content="result for a third one", tool_call_id="tool_call567" + ), ] } }, @@ -1211,258 +1258,56 @@ def test_prebuilt_tool_chat() -> None: AIMessage( content="", additional_kwargs={ - "tool_calls": [{ + "tool_calls": [ + { "id": "tool_call123", "type": "function", - "function":{ + "function": { "name": "search_api", - "arguments": "\"query\"", - } - }] - }, + "arguments": '"query"', + }, + } + ] + }, + ), + ToolMessage( + content="result for query", tool_call_id="tool_call123" ), - ToolMessage(content="result for query", tool_call_id="tool_call123"), AIMessage( content="", additional_kwargs={ - "tool_calls": [{ + "tool_calls": [ + { "id": "tool_call234", "type": "function", - "function":{ + "function": { "name": "search_api", - "arguments": "\"another\"", - } - }] + "arguments": '"another"', + }, + }, + { + "id": "tool_call567", + "type": "function", + "function": { + "name": "search_api", + "arguments": '"a third one"', + }, + }, + ] }, ), - ToolMessage(content="result for another", tool_call_id="tool_call234"), + ToolMessage( + content="result for another", tool_call_id="tool_call234" + ), + ToolMessage( + content="result for a third one", tool_call_id="tool_call567" + ), AIMessage(content="answer"), ] } }, ] -def test_tool_message_graph() -> None: - from langchain.chat_models.fake import FakeMessagesListChatModel - from langchain_community.tools import tool - from langchain_core.agents import AgentAction - from langchain_core.messages import AIMessage, ToolMessage, HumanMessage - - class FakeFuntionChatModel(FakeMessagesListChatModel): - def bind_functions(self, functions: list): - return self - - @tool() - def search_api(query: str) -> str: - """Searches the API for the query.""" - return f"result for {query}" - - tools = [search_api] - - model = FakeFuntionChatModel( - responses=[ - AIMessage( - content="", - additional_kwargs={ - "tool_calls": [{ - "id": "tool_call123", - "type": "function", - "function":{ - "name": "search_api", - "arguments": json.dumps("query"), - } - }] - }, - ), - AIMessage( - content="", - additional_kwargs={ - "tool_calls": [{ - "id": "tool_call234", - "type": "function", - "function":{ - "name": "search_api", - "arguments": json.dumps("another"), - } - }] - }, - ), - AIMessage(content="answer"), - ] - ) - - tool_executor = ToolExecutor(tools) - - # Define the function that determines whether to continue or not - def should_continue(messages): - last_message = messages[-1] - # If there is no function call, then we finish - if "tool_calls" not in last_message.additional_kwargs: - return "end" - # Otherwise if there is, we continue - else: - return "continue" - - def call_tool(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( - tool=last_message.additional_kwargs["tool_calls"][0]["function"]["name"], - tool_input=json.loads( - last_message.additional_kwargs["tool_calls"][0]["function"]["arguments"] - ), - log=last_message.additional_kwargs["tool_calls"][0]["id"], - ) - # We call the tool_executor and get back a response - response = tool_executor.invoke(action) - # We use the response to create a ToolMessage - return ToolMessage(content=str(response), tool_call_id=action.log) - - # Define a new graph - workflow = MessageGraph() - - # Define the two nodes we will cycle between - workflow.add_node("agent", model) - workflow.add_node("action", call_tool) - - # Set the entrypoint as `agent` - # This means that this node is the first one called - workflow.set_entry_point("agent") - - # We now add a conditional edge - workflow.add_conditional_edges( - # First, we define the start node. We use `agent`. - # This means these are the edges taken after the `agent` node is called. - "agent", - # Next, we pass in the function that will determine which node is called next. - should_continue, - # Finally we pass in a mapping. - # The keys are strings, and the values are other nodes. - # END is a special node marking that the graph should finish. - # What will happen is we will call `should_continue`, and then the output of that - # will be matched against the keys in this mapping. - # Based on which one it matches, that node will then be called. - { - # If `tools`, then we call the tool node. - "continue": "action", - # Otherwise we finish. - "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") - - # Finally, we compile it! - # This compiles it into a LangChain Runnable, - # meaning you can use it as you would any other runnable - app = workflow.compile() - - assert app.invoke(HumanMessage(content="what is weather in sf")) == [ - HumanMessage(content="what is weather in sf"), - AIMessage( - content="", - additional_kwargs={ - "tool_calls": [{ - "id": "tool_call123", - "type": "function", - "function":{ - "name": "search_api", - "arguments": "\"query\"", - } - }] - }, - ), - ToolMessage(content="result for query", tool_call_id="tool_call123"), - AIMessage( - content="", - additional_kwargs={ - "tool_calls": [{ - "id": "tool_call234", - "type": "function", - "function":{ - "name": "search_api", - "arguments": "\"another\"", - } - }] - }, - ), - ToolMessage(content="result for another", tool_call_id="tool_call234"), - AIMessage(content="answer"), - ] - - assert [*app.stream([HumanMessage(content="what is weather in sf")])] == [ - { - "agent": AIMessage( - content="", - additional_kwargs={ - "tool_calls": [{ - "id": "tool_call123", - "type": "function", - "function":{ - "name": "search_api", - "arguments": "\"query\"", - } - }] - }, - ) - }, - {"action": ToolMessage(content="result for query", tool_call_id="tool_call123")}, - { - "agent": AIMessage( - content="", - additional_kwargs={ - "tool_calls": [{ - "id": "tool_call234", - "type": "function", - "function":{ - "name": "search_api", - "arguments": "\"another\"", - } - }] - }, - ) - }, - {"action": ToolMessage(content="result for another", tool_call_id="tool_call234")}, - {"agent": AIMessage(content="answer")}, - { - "__end__": [ - HumanMessage(content="what is weather in sf"), - AIMessage( - content="", - additional_kwargs={ - "tool_calls": [{ - "id": "tool_call123", - "type": "function", - "function":{ - "name": "search_api", - "arguments": "\"query\"", - } - }] - }, - ), - ToolMessage(content="result for query", tool_call_id="tool_call123"), - AIMessage( - content="", - additional_kwargs={ - "tool_calls": [{ - "id": "tool_call234", - "type": "function", - "function":{ - "name": "search_api", - "arguments": "\"another\"", - } - }] - }, - ), - ToolMessage(content="result for another", tool_call_id="tool_call234"), - AIMessage(content="answer"), - ] - }, - ] def test_prebuilt_chat() -> None: from langchain.chat_models.fake import FakeMessagesListChatModel @@ -1608,6 +1453,7 @@ def test_prebuilt_chat() -> None: }, ] + def test_message_graph() -> None: from langchain.chat_models.fake import FakeMessagesListChatModel from langchain_community.tools import tool @@ -1781,4 +1627,4 @@ def test_message_graph() -> None: AIMessage(content="answer"), ] }, - ] \ No newline at end of file + ] diff --git a/tests/test_pregel_async.py b/tests/test_pregel_async.py index 4f09e372a..30a344a95 100644 --- a/tests/test_pregel_async.py +++ b/tests/test_pregel_async.py @@ -26,7 +26,10 @@ from langgraph.checkpoint.aiosqlite import AsyncSqliteSaver from langgraph.checkpoint.memory import MemorySaver from langgraph.graph import END, Graph, StateGraph from langgraph.graph.message import MessageGraph -from langgraph.prebuilt.chat_agent_executor import create_function_calling_executor, create_tool_calling_executor +from langgraph.prebuilt.chat_agent_executor import ( + create_function_calling_executor, + create_tool_calling_executor, +) from langgraph.prebuilt.tool_executor import ToolExecutor from langgraph.pregel import Channel, GraphRecursionError, Pregel from langgraph.pregel.reserved import ReservedChannels @@ -1114,7 +1117,7 @@ async def test_conditional_graph_state() -> None: async def test_prebuilt_tool_chat() -> None: from langchain.chat_models.fake import FakeMessagesListChatModel from langchain_community.tools import tool - from langchain_core.messages import AIMessage, ToolMessage, HumanMessage + from langchain_core.messages import AIMessage, HumanMessage, ToolMessage class FakeFuntionChatModel(FakeMessagesListChatModel): def bind_functions(self, functions: list): @@ -1133,27 +1136,39 @@ async def test_prebuilt_tool_chat() -> None: AIMessage( content="", additional_kwargs={ - "tool_calls": [{ - "id": "tool_call123", - "type": "function", - "function":{ - "name": "search_api", - "arguments": json.dumps("query"), + "tool_calls": [ + { + "id": "tool_call123", + "type": "function", + "function": { + "name": "search_api", + "arguments": json.dumps("query"), + }, } - }] + ] }, ), AIMessage( content="", additional_kwargs={ - "tool_calls": [{ - "id": "tool_call234", - "type": "function", - "function":{ - "name": "search_api", - "arguments": json.dumps("another"), - } - }] + "tool_calls": [ + { + "id": "tool_call234", + "type": "function", + "function": { + "name": "search_api", + "arguments": json.dumps("another"), + }, + }, + { + "id": "tool_call567", + "type": "function", + "function": { + "name": "search_api", + "arguments": '"a third one"', + }, + }, + ] }, ), AIMessage(content="answer"), @@ -1170,31 +1185,44 @@ async def test_prebuilt_tool_chat() -> None: AIMessage( content="", additional_kwargs={ - "tool_calls": [{ - "id": "tool_call123", - "type": "function", - "function":{ - "name": "search_api", - "arguments": "\"query\"", + "tool_calls": [ + { + "id": "tool_call123", + "type": "function", + "function": { + "name": "search_api", + "arguments": '"query"', + }, } - }] - }, + ] + }, ), ToolMessage(content="result for query", tool_call_id="tool_call123"), AIMessage( content="", additional_kwargs={ - "tool_calls": [{ - "id": "tool_call234", - "type": "function", - "function":{ - "name": "search_api", - "arguments": "\"another\"", - } - }] + "tool_calls": [ + { + "id": "tool_call234", + "type": "function", + "function": { + "name": "search_api", + "arguments": '"another"', + }, + }, + { + "id": "tool_call567", + "type": "function", + "function": { + "name": "search_api", + "arguments": '"a third one"', + }, + }, + ] }, ), ToolMessage(content="result for another", tool_call_id="tool_call234"), + ToolMessage(content="result for a third one", tool_call_id="tool_call567"), AIMessage(content="answer"), ] } @@ -1211,15 +1239,17 @@ async def test_prebuilt_tool_chat() -> None: AIMessage( content="", additional_kwargs={ - "tool_calls": [{ - "id": "tool_call123", - "type": "function", - "function":{ - "name": "search_api", - "arguments": "\"query\"", + "tool_calls": [ + { + "id": "tool_call123", + "type": "function", + "function": { + "name": "search_api", + "arguments": '"query"', + }, } - }] - }, + ] + }, ) ] } @@ -1237,14 +1267,24 @@ async def test_prebuilt_tool_chat() -> None: AIMessage( content="", additional_kwargs={ - "tool_calls": [{ - "id": "tool_call234", - "type": "function", - "function":{ - "name": "search_api", - "arguments": "\"another\"", - } - }] + "tool_calls": [ + { + "id": "tool_call234", + "type": "function", + "function": { + "name": "search_api", + "arguments": '"another"', + }, + }, + { + "id": "tool_call567", + "type": "function", + "function": { + "name": "search_api", + "arguments": '"a third one"', + }, + }, + ] }, ) ] @@ -1253,7 +1293,12 @@ async def test_prebuilt_tool_chat() -> None: { "action": { "messages": [ - ToolMessage(content="result for another", tool_call_id="tool_call234") + ToolMessage( + content="result for another", tool_call_id="tool_call234" + ), + ToolMessage( + content="result for a third one", tool_call_id="tool_call567" + ), ] } }, @@ -1265,31 +1310,50 @@ async def test_prebuilt_tool_chat() -> None: AIMessage( content="", additional_kwargs={ - "tool_calls": [{ - "id": "tool_call123", - "type": "function", - "function":{ - "name": "search_api", - "arguments": "\"query\"", + "tool_calls": [ + { + "id": "tool_call123", + "type": "function", + "function": { + "name": "search_api", + "arguments": '"query"', + }, } - }] + ] }, ), - ToolMessage(content="result for query", tool_call_id="tool_call123"), + ToolMessage( + content="result for query", tool_call_id="tool_call123" + ), AIMessage( content="", additional_kwargs={ - "tool_calls": [{ - "id": "tool_call234", - "type": "function", - "function":{ - "name": "search_api", - "arguments": "\"another\"", - } - }] + "tool_calls": [ + { + "id": "tool_call234", + "type": "function", + "function": { + "name": "search_api", + "arguments": '"another"', + }, + }, + { + "id": "tool_call567", + "type": "function", + "function": { + "name": "search_api", + "arguments": '"a third one"', + }, + }, + ] }, ), - ToolMessage(content="result for another", tool_call_id="tool_call234"), + ToolMessage( + content="result for another", tool_call_id="tool_call234" + ), + ToolMessage( + content="result for a third one", tool_call_id="tool_call567" + ), AIMessage(content="answer"), ] } @@ -1297,231 +1361,6 @@ async def test_prebuilt_tool_chat() -> None: ] -async def test_message_tool_graph() -> None: - from langchain.chat_models.fake import FakeMessagesListChatModel - from langchain_community.tools import tool - from langchain_core.agents import AgentAction - from langchain_core.messages import AIMessage, ToolMessage, HumanMessage - - class FakeFuntionChatModel(FakeMessagesListChatModel): - def bind_functions(self, functions: list): - return self - - @tool() - def search_api(query: str) -> str: - """Searches the API for the query.""" - return f"result for {query}" - - tools = [search_api] - - model = FakeFuntionChatModel( - responses=[ - AIMessage( - content="", - additional_kwargs={ - "tool_calls": [{ - "id": "tool_call123", - "type": "function", - "function":{ - "name": "search_api", - "arguments": json.dumps("query"), - } - }] - }, - ), - AIMessage( - content="", - additional_kwargs={ - "tool_calls": [{ - "id": "tool_call234", - "type": "function", - "function":{ - "name": "search_api", - "arguments": json.dumps("another"), - } - }] - }, - ), - AIMessage(content="answer"), - ] - ) - - tool_executor = ToolExecutor(tools) - - # Define the function that determines whether to continue or not - def should_continue(messages): - last_message = messages[-1] - # If there is no function call, then we finish - if "tool_calls" not in last_message.additional_kwargs: - return "end" - # Otherwise if there is, we continue - else: - return "continue" - - async def call_tool(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( - tool=last_message.additional_kwargs["tool_calls"][0]["function"]["name"], - tool_input=json.loads( - last_message.additional_kwargs["tool_calls"][0]["function"]["arguments"] - ), - log=last_message.additional_kwargs["tool_calls"][0]["id"], - ) - # We call the tool_executor and get back a response - response = await tool_executor.ainvoke(action) - # We use the response to create a FunctionMessage - return ToolMessage(content=str(response), tool_call_id=action.log) - - # Define a new graph - workflow = MessageGraph() - - # Define the two nodes we will cycle between - workflow.add_node("agent", model) - workflow.add_node("action", call_tool) - - # Set the entrypoint as `agent` - # This means that this node is the first one called - workflow.set_entry_point("agent") - - # We now add a conditional edge - workflow.add_conditional_edges( - # First, we define the start node. We use `agent`. - # This means these are the edges taken after the `agent` node is called. - "agent", - # Next, we pass in the function that will determine which node is called next. - should_continue, - # Finally we pass in a mapping. - # The keys are strings, and the values are other nodes. - # END is a special node marking that the graph should finish. - # What will happen is we will call `should_continue`, and then the output of that - # will be matched against the keys in this mapping. - # Based on which one it matches, that node will then be called. - { - # If `tools`, then we call the tool node. - "continue": "action", - # Otherwise we finish. - "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") - - # Finally, we compile it! - # This compiles it into a LangChain Runnable, - # meaning you can use it as you would any other runnable - app = workflow.compile() - - assert await app.ainvoke(HumanMessage(content="what is weather in sf")) == [ - HumanMessage(content="what is weather in sf"), - AIMessage( - content="", - additional_kwargs={ - "tool_calls": [{ - "id": "tool_call123", - "type": "function", - "function":{ - "name": "search_api", - "arguments": "\"query\"", - } - }] - }, - ), - ToolMessage(content="result for query", tool_call_id="tool_call123"), - AIMessage( - content="", - additional_kwargs={ - "tool_calls": [{ - "id": "tool_call234", - "type": "function", - "function":{ - "name": "search_api", - "arguments": "\"another\"", - } - }] - }, - ), - ToolMessage(content="result for another", tool_call_id="tool_call234"), - AIMessage(content="answer"), - ] - - assert [ - c async for c in app.astream([HumanMessage(content="what is weather in sf")]) - ] == [ - { - "agent": AIMessage( - content="", - additional_kwargs={ - "tool_calls": [{ - "id": "tool_call123", - "type": "function", - "function":{ - "name": "search_api", - "arguments": "\"query\"", - } - }] - }, - ) - }, - {"action": ToolMessage(content="result for query", tool_call_id="tool_call123")}, - { - "agent": AIMessage( - content="", - additional_kwargs={ - "tool_calls": [{ - "id": "tool_call234", - "type": "function", - "function":{ - "name": "search_api", - "arguments": "\"another\"", - } - }] - }, - ) - }, - {"action": ToolMessage(content="result for another", tool_call_id="tool_call234")}, - {"agent": AIMessage(content="answer")}, - { - "__end__": [ - HumanMessage(content="what is weather in sf"), - AIMessage( - content="", - additional_kwargs={ - "tool_calls": [{ - "id": "tool_call123", - "type": "function", - "function":{ - "name": "search_api", - "arguments": "\"query\"", - } - }] - }, - ), - ToolMessage(content="result for query", tool_call_id="tool_call123"), - AIMessage( - content="", - additional_kwargs={ - "tool_calls": [{ - "id": "tool_call234", - "type": "function", - "function":{ - "name": "search_api", - "arguments": "\"another\"", - } - }] - }, - ), - ToolMessage(content="result for another", tool_call_id="tool_call234"), - AIMessage(content="answer"), - ] - }, - ] - - async def test_prebuilt_chat() -> None: from langchain.chat_models.fake import FakeMessagesListChatModel from langchain_community.tools import tool