From 38ef3d5218bd175c967760227006306ea0c6145d Mon Sep 17 00:00:00 2001 From: midas8181919 Date: Sat, 10 Feb 2024 23:10:32 +0000 Subject: [PATCH] replace function_call with tool_call --- langgraph/prebuilt/chat_agent_executor.py | 116 ++++++- tests/test_pregel.py | 361 +++++++++++++++++++++- 2 files changed, 475 insertions(+), 2 deletions(-) diff --git a/langgraph/prebuilt/chat_agent_executor.py b/langgraph/prebuilt/chat_agent_executor.py index bc7208107..83e48a392 100644 --- a/langgraph/prebuilt/chat_agent_executor.py +++ b/langgraph/prebuilt/chat_agent_executor.py @@ -5,7 +5,7 @@ from typing import Annotated, Sequence, TypedDict from langchain_core.agents import AgentAction from langchain_core.messages import BaseMessage, FunctionMessage from langchain_core.runnables import RunnableLambda -from langchain_core.utils.function_calling import convert_to_openai_function +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 @@ -124,3 +124,117 @@ def create_function_calling_executor(model, tools): # This compiles it into a LangChain Runnable, # meaning you can use it as you would any other runnable return workflow.compile() + +def create_tool_calling_executor(model, tools): + 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]) + + # Define the function that determines whether to continue or not + def should_continue(state): + messages = state["messages"] + last_message = messages[-1] + # If there is no function call, then we finish + if "tool_call" not in last_message.additional_kwargs: + return "end" + # Otherwise if there is, we continue + else: + return "continue" + + # Define the function that calls the model + def call_model(state): + 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 _get_action(state): + 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="", + ) + + 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 + 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]} + + 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 + # 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) + + # Define the two nodes we will cycle between + 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 + 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 + return workflow.compile() diff --git a/tests/test_pregel.py b/tests/test_pregel.py index f9fef277b..696555f61 100644 --- a/tests/test_pregel.py +++ b/tests/test_pregel.py @@ -10,6 +10,7 @@ import pytest from langchain_core.runnables import RunnablePassthrough from pytest_mock import MockerFixture +from langchain_core._api import deprecated from langgraph.channels.base import InvalidUpdateError from langgraph.channels.binop import BinaryOperatorAggregate from langgraph.channels.context import Context @@ -20,7 +21,7 @@ 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 +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 @@ -1061,7 +1062,365 @@ 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 + from langchain_core.messages import AIMessage, FunctionMessage, 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] + + app = create_tool_calling_executor( + 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"), + ] + ), + tools, + ) + + assert app.invoke( + {"messages": [HumanMessage(content="what is weather in sf")]} + ) == { + "messages": [ + HumanMessage(content="what is weather in sf"), + AIMessage( + content="", + additional_kwargs={ + "tool_calls": [{ + "id": "tool_call123", + "type": "function", + "function":{ + "name": "search_api", + "arguments": "query", + } + }] + }, + ), + FunctionMessage(content="result for query", name="search_api"), + AIMessage( + content="", + additional_kwargs={ + "tool_calls": [{ + "id": "tool_call234", + "type": "function", + "function":{ + "name": "search_api", + "arguments": "another", + } + }] + }, + ), + FunctionMessage(content="result for another", name="search_api"), + AIMessage(content="answer"), + ] + } + + assert [ + *app.stream({"messages": [HumanMessage(content="what is weather in sf")]}) + ] == [ + { + "agent": { + "messages": [ + AIMessage( + content="", + additional_kwargs={ + "tool_calls": [{ + "id": "tool_call123", + "type": "function", + "function":{ + "name": "search_api", + "arguments": "query", + } + }] + }, + ) + ] + } + }, + { + "action": { + "messages": [ + FunctionMessage(content="result for query", name="search_api") + ] + } + }, + { + "agent": { + "messages": [ + AIMessage( + content="", + additional_kwargs={ + "tool_calls": [{ + "id": "tool_call234", + "type": "function", + "function":{ + "name": "search_api", + "arguments": "another", + } + }] + }, + ) + ] + } + }, + { + "action": { + "messages": [ + FunctionMessage(content="result for another", name="search_api") + ] + } + }, + {"agent": {"messages": [AIMessage(content="answer")]}}, + { + "__end__": { + "messages": [ + HumanMessage(content="what is weather in sf"), + AIMessage( + content="", + additional_kwargs={ + "tool_calls": [{ + "id": "tool_call123", + "type": "function", + "function":{ + "name": "search_api", + "arguments": "query", + } + }] + }, + ), + FunctionMessage(content="result for query", name="search_api"), + AIMessage( + content="", + additional_kwargs={ + "tool_calls": [{ + "id": "tool_call234", + "type": "function", + "function":{ + "name": "search_api", + "arguments": "another", + } + }] + }, + ), + FunctionMessage(content="result for another", name="search_api"), + AIMessage(content="answer"), + ] + } + }, + ] + + +def test_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, FunctionMessage, 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={ + "function_call": { + "name": "search_api", + "arguments": json.dumps("query"), + } + }, + ), + AIMessage( + content="", + additional_kwargs={ + "function_call": { + "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 "function_call" 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["function_call"]["name"], + tool_input=json.loads( + last_message.additional_kwargs["function_call"]["arguments"] + ), + log="", + ) + # We call the tool_executor and get back a response + response = tool_executor.invoke(action) + # We use the response to create a FunctionMessage + return FunctionMessage(content=str(response), name=action.tool) + + # 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={ + "function_call": {"name": "search_api", "arguments": '"query"'} + }, + ), + FunctionMessage(content="result for query", name="search_api"), + AIMessage( + content="", + additional_kwargs={ + "function_call": {"name": "search_api", "arguments": '"another"'} + }, + ), + FunctionMessage(content="result for another", name="search_api"), + AIMessage(content="answer"), + ] + + assert [*app.stream([HumanMessage(content="what is weather in sf")])] == [ + { + "agent": AIMessage( + content="", + additional_kwargs={ + "function_call": {"name": "search_api", "arguments": '"query"'} + }, + ) + }, + {"action": FunctionMessage(content="result for query", name="search_api")}, + { + "agent": AIMessage( + content="", + additional_kwargs={ + "function_call": {"name": "search_api", "arguments": '"another"'} + }, + ) + }, + {"action": FunctionMessage(content="result for another", name="search_api")}, + {"agent": AIMessage(content="answer")}, + { + "__end__": [ + HumanMessage(content="what is weather in sf"), + AIMessage( + content="", + additional_kwargs={ + "function_call": {"name": "search_api", "arguments": '"query"'} + }, + ), + FunctionMessage(content="result for query", name="search_api"), + AIMessage( + content="", + additional_kwargs={ + "function_call": { + "name": "search_api", + "arguments": '"another"', + } + }, + ), + FunctionMessage(content="result for another", name="search_api"), + AIMessage(content="answer"), + ] + }, + ] + +@deprecated("*") def test_prebuilt_chat() -> None: from langchain.chat_models.fake import FakeMessagesListChatModel from langchain_community.tools import tool