diff --git a/langgraph/graph/message.py b/langgraph/graph/message.py new file mode 100644 index 000000000..d41a2e747 --- /dev/null +++ b/langgraph/graph/message.py @@ -0,0 +1,24 @@ +from typing import Annotated, Union + +from langchain_core.messages import AnyMessage + +from langgraph.graph.state import StateGraph + +Messages = Union[list[AnyMessage], AnyMessage] + + +def add_messages(left: Messages, right: Messages) -> Messages: + if not isinstance(left, list): + left = [left] + if not isinstance(right, list): + right = [right] + return left + right + + +class MessageGraph(StateGraph): + """A StateGraph where every node + - receives a list of messages as input + - returns one or more messages as output.""" + + def __init__(self) -> None: + super().__init__(Annotated[list[AnyMessage], add_messages]) diff --git a/tests/test_pregel.py b/tests/test_pregel.py index 3557a526d..ac198490b 100644 --- a/tests/test_pregel.py +++ b/tests/test_pregel.py @@ -17,8 +17,10 @@ from langgraph.channels.last_value import LastValue from langgraph.channels.topic import Topic from langgraph.checkpoint.memory import MemorySaver 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.tool_executor import ToolExecutor from langgraph.pregel import Channel, GraphRecursionError, Pregel from langgraph.pregel.reserved import ReservedChannels @@ -1118,3 +1120,179 @@ def test_prebuilt_chat() -> None: } }, ] + + +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) -> Self: + 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"), + ] + }, + ] diff --git a/tests/test_pregel_async.py b/tests/test_pregel_async.py index 3adc9e5a0..ce50a2271 100644 --- a/tests/test_pregel_async.py +++ b/tests/test_pregel_async.py @@ -25,7 +25,9 @@ from langgraph.channels.last_value import LastValue from langgraph.channels.topic import Topic 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 +from langgraph.prebuilt.tool_executor import ToolExecutor from langgraph.pregel import Channel, GraphRecursionError, Pregel from langgraph.pregel.reserved import ReservedChannels @@ -1168,3 +1170,181 @@ async def test_prebuilt_chat() -> None: } }, ] + + +async 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) -> Self: + 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" + + 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["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 = await tool_executor.ainvoke(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 await app.ainvoke(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 [ + c async for c in app.astream([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"), + ] + }, + ]