From 3fa3a586b5d4fd6b0507b0219b2023962e39aed1 Mon Sep 17 00:00:00 2001 From: Nuno Campos Date: Fri, 13 Jun 2025 14:21:05 -0700 Subject: [PATCH] Revert "Remove `MessageGraph` (#4875)" This reverts commit a5e6223569c6be5062fa3c22e63efa2468dc094f. --- docs/docs/cloud/deployment/graph_rebuild.md | 10 +- docs/docs/tutorials/extraction/retries.ipynb | 6 +- libs/langgraph/langgraph/graph/__init__.py | 3 +- libs/langgraph/langgraph/graph/message.py | 52 + libs/langgraph/tests/test_large_cases.py | 1570 ++++++++++++++++- .../langgraph/tests/test_large_cases_async.py | 413 ++++- libs/langgraph/tests/test_pregel.py | 20 +- libs/prebuilt/langgraph/prebuilt/tool_node.py | 2 +- .../langgraph/prebuilt/tool_validator.py | 4 +- 9 files changed, 2049 insertions(+), 31 deletions(-) diff --git a/docs/docs/cloud/deployment/graph_rebuild.md b/docs/docs/cloud/deployment/graph_rebuild.md index 75dc26377..eb5b8a322 100644 --- a/docs/docs/cloud/deployment/graph_rebuild.md +++ b/docs/docs/cloud/deployment/graph_rebuild.md @@ -20,7 +20,7 @@ my-app/ |-- openai_agent.py # code for your graph ``` -where the graph is defined in `openai_agent.py`. +where the graph is defined in `openai_agent.py`. ### No rebuild @@ -28,11 +28,11 @@ In the standard LangGraph API configuration, the server uses the compiled graph ```python from langchain_openai import ChatOpenAI -from langgraph.graph import END, START, StateGraph, MessagesState +from langgraph.graph import END, START, MessageGraph model = ChatOpenAI(temperature=0) -graph_workflow = StateGraph(MessagesState) +graph_workflow = MessageGraph() graph_workflow.add_node("agent", model) graph_workflow.add_edge("agent", END) @@ -61,7 +61,7 @@ To make your graph rebuild on each new run with custom configuration, you need t from typing import Annotated from typing_extensions import TypedDict from langchain_openai import ChatOpenAI -from langgraph.graph import END, START +from langgraph.graph import END, START, MessageGraph from langgraph.graph.state import StateGraph from langgraph.graph.message import add_messages from langgraph.prebuilt import ToolNode @@ -144,4 +144,4 @@ Finally, you need to specify the path to your graph-making function (`make_graph } ``` -See more info on LangGraph API configuration file [here](../reference/cli.md#configuration-file) +See more info on LangGraph API configuration file [here](../reference/cli.md#configuration-file) \ No newline at end of file diff --git a/docs/docs/tutorials/extraction/retries.ipynb b/docs/docs/tutorials/extraction/retries.ipynb index ccd552e3a..de081052c 100644 --- a/docs/docs/tutorials/extraction/retries.ipynb +++ b/docs/docs/tutorials/extraction/retries.ipynb @@ -89,7 +89,7 @@ }, { "cell_type": "code", - "execution_count": null, + "execution_count": 2, "id": "baf669a0-04ee-492d-80d8-8fcb658ed128", "metadata": {}, "outputs": [], @@ -313,8 +313,8 @@ "\n", " builder.add_edge(\"finalizer\", END)\n", "\n", - " # These functions let the step be used in a\n", - " # StateGraph with 'messages' as the key.\n", + " # These functions let the step be used in a MessageGraph\n", + " # or a StateGraph with 'messages' as the key.\n", " def encode(x: Union[Sequence[AnyMessage], PromptValue]) -> dict:\n", " \"\"\"Ensure the input is the correct format.\"\"\"\n", " if isinstance(x, PromptValue):\n", diff --git a/libs/langgraph/langgraph/graph/__init__.py b/libs/langgraph/langgraph/graph/__init__.py index 9b67c2a42..2581713c3 100644 --- a/libs/langgraph/langgraph/graph/__init__.py +++ b/libs/langgraph/langgraph/graph/__init__.py @@ -1,11 +1,12 @@ from langgraph.constants import END, START -from langgraph.graph.message import MessagesState, add_messages +from langgraph.graph.message import MessageGraph, MessagesState, add_messages from langgraph.graph.state import StateGraph __all__ = [ "END", "START", "StateGraph", + "MessageGraph", "add_messages", "MessagesState", ] diff --git a/libs/langgraph/langgraph/graph/message.py b/libs/langgraph/langgraph/graph/message.py index 059102027..e50bcfec2 100644 --- a/libs/langgraph/langgraph/graph/message.py +++ b/libs/langgraph/langgraph/graph/message.py @@ -25,6 +25,7 @@ from langchain_core.messages import ( from typing_extensions import TypedDict from langgraph.constants import CONF, CONFIG_KEY_SEND +from langgraph.graph.state import StateGraph Messages = Union[list[MessageLikeRepresentation], MessageLikeRepresentation] @@ -226,6 +227,57 @@ def add_messages( return merged +class MessageGraph(StateGraph): + """A StateGraph where every node receives a list of messages as input and returns one or more messages as output. + + MessageGraph is a subclass of StateGraph whose entire state is a single, append-only* list of messages. + Each node in a MessageGraph takes a list of messages as input and returns zero or more + messages as output. The `add_messages` function is used to merge the output messages from each node + into the existing list of messages in the graph's state. + + Examples: + ```pycon + >>> from langgraph.graph.message import MessageGraph + ... + >>> builder = MessageGraph() + >>> builder.add_node("chatbot", lambda state: [("assistant", "Hello!")]) + >>> builder.set_entry_point("chatbot") + >>> builder.set_finish_point("chatbot") + >>> builder.compile().invoke([("user", "Hi there.")]) + [HumanMessage(content="Hi there.", id='...'), AIMessage(content="Hello!", id='...')] + ``` + + ```pycon + >>> from langchain_core.messages import AIMessage, HumanMessage, ToolMessage + >>> from langgraph.graph.message import MessageGraph + ... + >>> builder = MessageGraph() + >>> builder.add_node( + ... "chatbot", + ... lambda state: [ + ... AIMessage( + ... content="Hello!", + ... tool_calls=[{"name": "search", "id": "123", "args": {"query": "X"}}], + ... ) + ... ], + ... ) + >>> builder.add_node( + ... "search", lambda state: [ToolMessage(content="Searching...", tool_call_id="123")] + ... ) + >>> builder.set_entry_point("chatbot") + >>> builder.add_edge("chatbot", "search") + >>> builder.set_finish_point("search") + >>> builder.compile().invoke([HumanMessage(content="Hi there. Can you search for X?")]) + {'messages': [HumanMessage(content="Hi there. Can you search for X?", id='b8b7d8f4-7f4d-4f4d-9c1d-f8b8d8f4d9c1'), + AIMessage(content="Hello!", id='f4d9c1d8-8d8f-4d9c-b8b7-d8f4f4d9c1d8'), + ToolMessage(content="Searching...", id='d8f4f4d9-c1d8-4f4d-b8b7-d8f4f4d9c1d8', tool_call_id="123")]} + ``` + """ + + def __init__(self) -> None: + super().__init__(Annotated[list[AnyMessage], add_messages]) # type: ignore[arg-type] + + class MessagesState(TypedDict): messages: Annotated[list[AnyMessage], add_messages] diff --git a/libs/langgraph/tests/test_large_cases.py b/libs/langgraph/tests/test_large_cases.py index 509059251..e706a6b02 100644 --- a/libs/langgraph/tests/test_large_cases.py +++ b/libs/langgraph/tests/test_large_cases.py @@ -3,10 +3,10 @@ import operator import re import time from dataclasses import replace -from typing import Annotated, Literal, Optional, Union, cast +from typing import Annotated, Any, Literal, Optional, Union, cast import pytest -from langchain_core.runnables import RunnableConfig +from langchain_core.runnables import RunnableConfig, RunnableMap, RunnablePick from pytest_mock import MockerFixture from syrupy import SnapshotAssertion from typing_extensions import TypedDict @@ -18,8 +18,9 @@ from langgraph.checkpoint.memory import InMemorySaver from langgraph.constants import END, PULL, PUSH, START from langgraph.errors import NodeInterrupt from langgraph.graph import StateGraph -from langgraph.graph.message import MessagesState, add_messages +from langgraph.graph.message import MessageGraph, MessagesState, add_messages from langgraph.prebuilt.chat_agent_executor import create_react_agent +from langgraph.prebuilt.tool_node import ToolNode from langgraph.pregel import NodeBuilder, Pregel from langgraph.types import ( Command, @@ -2379,6 +2380,1569 @@ def test_state_graph_packets( ) +def test_message_graph( + snapshot: SnapshotAssertion, + deterministic_uuids: MockerFixture, + sync_checkpointer: BaseCheckpointSaver, +) -> None: + from copy import deepcopy + + from langchain_core.callbacks import CallbackManagerForLLMRun + from langchain_core.language_models.fake_chat_models import ( + FakeMessagesListChatModel, + ) + from langchain_core.messages import AIMessage, BaseMessage, HumanMessage + from langchain_core.outputs import ChatGeneration, ChatResult + from langchain_core.tools import tool + + class FakeFuntionChatModel(FakeMessagesListChatModel): + def bind_functions(self, functions: list): + return self + + def _generate( + self, + messages: list[BaseMessage], + stop: Optional[list[str]] = None, + run_manager: Optional[CallbackManagerForLLMRun] = None, + **kwargs: Any, + ) -> ChatResult: + response = deepcopy(self.responses[self.i]) + if self.i < len(self.responses) - 1: + self.i += 1 + else: + self.i = 0 + generation = ChatGeneration(message=response) + return ChatResult(generations=[generation]) + + @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="", + tool_calls=[ + { + "id": "tool_call123", + "name": "search_api", + "args": {"query": "query"}, + } + ], + id="ai1", + ), + AIMessage( + content="", + tool_calls=[ + { + "id": "tool_call456", + "name": "search_api", + "args": {"query": "another"}, + } + ], + id="ai2", + ), + AIMessage(content="answer", id="ai3"), + ] + ) + + # 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 not last_message.tool_calls: + return "end" + # Otherwise if there is, we continue + else: + return "continue" + + # Define a new graph + workflow = MessageGraph() + + # Define the two nodes we will cycle between + workflow.add_node("agent", model) + workflow.add_node("tools", ToolNode(tools)) + + # 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": "tools", + # 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("tools", "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() + + if isinstance(sync_checkpointer, InMemorySaver): + assert json.dumps(app.get_input_schema().model_json_schema()) == snapshot + assert json.dumps(app.get_output_schema().model_json_schema()) == snapshot + assert json.dumps(app.get_graph().to_json(), indent=2) == snapshot + assert app.get_graph().draw_mermaid(with_styles=False) == snapshot + + assert app.invoke(HumanMessage(content="what is weather in sf")) == [ + _AnyIdHumanMessage( + content="what is weather in sf", + ), + AIMessage( + content="", + tool_calls=[ + { + "id": "tool_call123", + "name": "search_api", + "args": {"query": "query"}, + } + ], + id="ai1", # respects ids passed in + ), + _AnyIdToolMessage( + content="result for query", + name="search_api", + tool_call_id="tool_call123", + ), + AIMessage( + content="", + tool_calls=[ + { + "id": "tool_call456", + "name": "search_api", + "args": {"query": "another"}, + } + ], + id="ai2", + ), + _AnyIdToolMessage( + content="result for another", + name="search_api", + tool_call_id="tool_call456", + ), + AIMessage(content="answer", id="ai3"), + ] + + assert [*app.stream([HumanMessage(content="what is weather in sf")])] == [ + { + "agent": AIMessage( + content="", + tool_calls=[ + { + "id": "tool_call123", + "name": "search_api", + "args": {"query": "query"}, + } + ], + id="ai1", + ) + }, + { + "tools": [ + _AnyIdToolMessage( + content="result for query", + name="search_api", + tool_call_id="tool_call123", + ) + ] + }, + { + "agent": AIMessage( + content="", + tool_calls=[ + { + "id": "tool_call456", + "name": "search_api", + "args": {"query": "another"}, + } + ], + id="ai2", + ) + }, + { + "tools": [ + _AnyIdToolMessage( + content="result for another", + name="search_api", + tool_call_id="tool_call456", + ) + ] + }, + {"agent": AIMessage(content="answer", id="ai3")}, + ] + + app_w_interrupt = workflow.compile( + checkpointer=sync_checkpointer, + interrupt_after=["agent"], + ) + config = {"configurable": {"thread_id": "1"}} + + assert [ + c for c in app_w_interrupt.stream(("human", "what is weather in sf"), config) + ] == [ + { + "agent": AIMessage( + content="", + tool_calls=[ + { + "id": "tool_call123", + "name": "search_api", + "args": {"query": "query"}, + } + ], + id="ai1", + ) + }, + {"__interrupt__": ()}, + ] + + assert app_w_interrupt.get_state(config) == StateSnapshot( + values=[ + _AnyIdHumanMessage(content="what is weather in sf"), + AIMessage( + content="", + tool_calls=[ + { + "id": "tool_call123", + "name": "search_api", + "args": {"query": "query"}, + } + ], + id="ai1", + ), + ], + tasks=(PregelTask(AnyStr(), "tools", (PULL, "tools")),), + next=("tools",), + config={ + "configurable": { + "thread_id": "1", + "checkpoint_ns": "", + "checkpoint_id": AnyStr(), + } + }, + created_at=AnyStr(), + metadata={ + "parents": {}, + "source": "loop", + "step": 1, + "thread_id": "1", + }, + parent_config=None, + interrupts=(), + ) + + # modify ai message + last_message = app_w_interrupt.get_state(config).values[-1] + last_message.tool_calls[0]["args"] = {"query": "a different query"} + next_config = app_w_interrupt.update_state(config, last_message) + + # message was replaced instead of appended + assert app_w_interrupt.get_state(config) == StateSnapshot( + values=[ + _AnyIdHumanMessage(content="what is weather in sf"), + AIMessage( + content="", + id="ai1", + tool_calls=[ + { + "id": "tool_call123", + "name": "search_api", + "args": {"query": "a different query"}, + } + ], + ), + ], + tasks=(PregelTask(AnyStr(), "tools", (PULL, "tools")),), + next=("tools",), + config=next_config, + created_at=AnyStr(), + metadata={ + "parents": {}, + "source": "update", + "step": 2, + "thread_id": "1", + }, + parent_config=( + list(app_w_interrupt.checkpointer.list(config, limit=2))[-1].config + ), + interrupts=(), + ) + + assert [c for c in app_w_interrupt.stream(None, config)] == [ + { + "tools": [ + _AnyIdToolMessage( + content="result for a different query", + name="search_api", + tool_call_id="tool_call123", + ) + ] + }, + { + "agent": AIMessage( + content="", + tool_calls=[ + { + "id": "tool_call456", + "name": "search_api", + "args": {"query": "another"}, + } + ], + id="ai2", + ) + }, + {"__interrupt__": ()}, + ] + + assert app_w_interrupt.get_state(config) == StateSnapshot( + values=[ + _AnyIdHumanMessage(content="what is weather in sf"), + AIMessage( + content="", + id="ai1", + tool_calls=[ + { + "id": "tool_call123", + "name": "search_api", + "args": {"query": "a different query"}, + } + ], + ), + _AnyIdToolMessage( + content="result for a different query", + name="search_api", + tool_call_id="tool_call123", + ), + AIMessage( + content="", + tool_calls=[ + { + "id": "tool_call456", + "name": "search_api", + "args": {"query": "another"}, + } + ], + id="ai2", + ), + ], + tasks=(PregelTask(AnyStr(), "tools", (PULL, "tools")),), + next=("tools",), + config={ + "configurable": { + "thread_id": "1", + "checkpoint_ns": "", + "checkpoint_id": AnyStr(), + } + }, + created_at=AnyStr(), + metadata={ + "parents": {}, + "source": "loop", + "step": 4, + "thread_id": "1", + }, + parent_config=( + list(app_w_interrupt.checkpointer.list(config, limit=2))[-1].config + ), + interrupts=(), + ) + + app_w_interrupt.update_state( + config, + AIMessage(content="answer", id="ai2"), # replace existing message + ) + + # replaces message even if object identity is different, as long as id is the same + assert app_w_interrupt.get_state(config) == StateSnapshot( + values=[ + _AnyIdHumanMessage(content="what is weather in sf"), + AIMessage( + content="", + id="ai1", + tool_calls=[ + { + "id": "tool_call123", + "name": "search_api", + "args": {"query": "a different query"}, + } + ], + ), + _AnyIdToolMessage( + content="result for a different query", + name="search_api", + tool_call_id="tool_call123", + ), + AIMessage(content="answer", id="ai2"), + ], + tasks=(), + next=(), + config={ + "configurable": { + "thread_id": "1", + "checkpoint_ns": "", + "checkpoint_id": AnyStr(), + } + }, + created_at=AnyStr(), + metadata={ + "parents": {}, + "source": "update", + "step": 5, + "thread_id": "1", + }, + parent_config=( + list(app_w_interrupt.checkpointer.list(config, limit=2))[-1].config + ), + interrupts=(), + ) + + app_w_interrupt = workflow.compile( + checkpointer=sync_checkpointer, + interrupt_before=["tools"], + ) + config = {"configurable": {"thread_id": "2"}} + model.i = 0 # reset the llm + + assert [c for c in app_w_interrupt.stream("what is weather in sf", config)] == [ + { + "agent": AIMessage( + content="", + tool_calls=[ + { + "id": "tool_call123", + "name": "search_api", + "args": {"query": "query"}, + } + ], + id="ai1", + ) + }, + {"__interrupt__": ()}, + ] + + assert app_w_interrupt.get_state(config) == StateSnapshot( + values=[ + _AnyIdHumanMessage(content="what is weather in sf"), + AIMessage( + content="", + tool_calls=[ + { + "id": "tool_call123", + "name": "search_api", + "args": {"query": "query"}, + } + ], + id="ai1", + ), + ], + tasks=(PregelTask(AnyStr(), "tools", (PULL, "tools")),), + next=("tools",), + config={ + "configurable": { + "thread_id": "2", + "checkpoint_ns": "", + "checkpoint_id": AnyStr(), + } + }, + created_at=AnyStr(), + metadata={ + "parents": {}, + "source": "loop", + "step": 1, + "thread_id": "2", + }, + parent_config=None, + interrupts=(), + ) + + # modify ai message + last_message = app_w_interrupt.get_state(config).values[-1] + last_message.tool_calls[0]["args"] = {"query": "a different query"} + app_w_interrupt.update_state(config, last_message) + + # message was replaced instead of appended + assert app_w_interrupt.get_state(config) == StateSnapshot( + values=[ + _AnyIdHumanMessage(content="what is weather in sf"), + AIMessage( + content="", + id="ai1", + tool_calls=[ + { + "id": "tool_call123", + "name": "search_api", + "args": {"query": "a different query"}, + } + ], + ), + ], + tasks=(PregelTask(AnyStr(), "tools", (PULL, "tools")),), + next=("tools",), + config={ + "configurable": { + "thread_id": "2", + "checkpoint_ns": "", + "checkpoint_id": AnyStr(), + } + }, + created_at=AnyStr(), + metadata={ + "parents": {}, + "source": "update", + "step": 2, + "thread_id": "2", + }, + parent_config=( + list(app_w_interrupt.checkpointer.list(config, limit=2))[-1].config + ), + interrupts=(), + ) + + assert [c for c in app_w_interrupt.stream(None, config)] == [ + { + "tools": [ + _AnyIdToolMessage( + content="result for a different query", + name="search_api", + tool_call_id="tool_call123", + ) + ] + }, + { + "agent": AIMessage( + content="", + tool_calls=[ + { + "id": "tool_call456", + "name": "search_api", + "args": {"query": "another"}, + } + ], + id="ai2", + ) + }, + {"__interrupt__": ()}, + ] + + assert app_w_interrupt.get_state(config) == StateSnapshot( + values=[ + _AnyIdHumanMessage(content="what is weather in sf"), + AIMessage( + content="", + id="ai1", + tool_calls=[ + { + "id": "tool_call123", + "name": "search_api", + "args": {"query": "a different query"}, + } + ], + ), + _AnyIdToolMessage( + content="result for a different query", + name="search_api", + tool_call_id="tool_call123", + ), + AIMessage( + content="", + tool_calls=[ + { + "id": "tool_call456", + "name": "search_api", + "args": {"query": "another"}, + } + ], + id="ai2", + ), + ], + tasks=(PregelTask(AnyStr(), "tools", (PULL, "tools")),), + next=("tools",), + config={ + "configurable": { + "thread_id": "2", + "checkpoint_ns": "", + "checkpoint_id": AnyStr(), + } + }, + created_at=AnyStr(), + metadata={ + "parents": {}, + "source": "loop", + "step": 4, + "thread_id": "2", + }, + parent_config=( + list(app_w_interrupt.checkpointer.list(config, limit=2))[-1].config + ), + interrupts=(), + ) + + app_w_interrupt.update_state( + config, + AIMessage(content="answer", id="ai2"), + ) + + # replaces message even if object identity is different, as long as id is the same + assert app_w_interrupt.get_state(config) == StateSnapshot( + values=[ + _AnyIdHumanMessage(content="what is weather in sf"), + AIMessage( + content="", + id="ai1", + tool_calls=[ + { + "id": "tool_call123", + "name": "search_api", + "args": {"query": "a different query"}, + } + ], + ), + _AnyIdToolMessage( + content="result for a different query", + name="search_api", + tool_call_id="tool_call123", + id=AnyStr(), + ), + AIMessage(content="answer", id="ai2"), + ], + tasks=(), + next=(), + config={ + "configurable": { + "thread_id": "2", + "checkpoint_ns": "", + "checkpoint_id": AnyStr(), + } + }, + created_at=AnyStr(), + metadata={ + "parents": {}, + "source": "update", + "step": 5, + "thread_id": "2", + }, + parent_config=( + list(app_w_interrupt.checkpointer.list(config, limit=2))[-1].config + ), + interrupts=(), + ) + + # add an extra message as if it came from "tools" node + app_w_interrupt.update_state(config, ("ai", "an extra message"), as_node="tools") + + # extra message is coerced BaseMessge and appended + # now the next node is "agent" per the graph edges + assert app_w_interrupt.get_state(config) == StateSnapshot( + values=[ + _AnyIdHumanMessage(content="what is weather in sf"), + AIMessage( + content="", + id="ai1", + tool_calls=[ + { + "id": "tool_call123", + "name": "search_api", + "args": {"query": "a different query"}, + } + ], + ), + _AnyIdToolMessage( + content="result for a different query", + name="search_api", + tool_call_id="tool_call123", + id=AnyStr(), + ), + AIMessage(content="answer", id="ai2"), + _AnyIdAIMessage(content="an extra message"), + ], + tasks=(PregelTask(AnyStr(), "agent", (PULL, "agent")),), + next=("agent",), + config={ + "configurable": { + "thread_id": "2", + "checkpoint_ns": "", + "checkpoint_id": AnyStr(), + } + }, + created_at=AnyStr(), + metadata={ + "parents": {}, + "source": "update", + "step": 6, + "thread_id": "2", + }, + parent_config=( + list(app_w_interrupt.checkpointer.list(config, limit=2))[-1].config + ), + interrupts=(), + ) + + +def test_root_graph( + deterministic_uuids: MockerFixture, + sync_checkpointer: BaseCheckpointSaver, +) -> None: + from copy import deepcopy + + from langchain_core.callbacks import CallbackManagerForLLMRun + from langchain_core.language_models.fake_chat_models import ( + FakeMessagesListChatModel, + ) + from langchain_core.messages import ( + AIMessage, + BaseMessage, + HumanMessage, + ToolMessage, + ) + from langchain_core.outputs import ChatGeneration, ChatResult + from langchain_core.tools import tool + + class FakeFuntionChatModel(FakeMessagesListChatModel): + def bind_functions(self, functions: list): + return self + + def _generate( + self, + messages: list[BaseMessage], + stop: Optional[list[str]] = None, + run_manager: Optional[CallbackManagerForLLMRun] = None, + **kwargs: Any, + ) -> ChatResult: + response = deepcopy(self.responses[self.i]) + if self.i < len(self.responses) - 1: + self.i += 1 + else: + self.i = 0 + generation = ChatGeneration(message=response) + return ChatResult(generations=[generation]) + + @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="", + tool_calls=[ + { + "id": "tool_call123", + "name": "search_api", + "args": {"query": "query"}, + } + ], + id="ai1", + ), + AIMessage( + content="", + tool_calls=[ + { + "id": "tool_call456", + "name": "search_api", + "args": {"query": "another"}, + } + ], + id="ai2", + ), + AIMessage(content="answer", id="ai3"), + ] + ) + + # 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 not last_message.tool_calls: + return "end" + # Otherwise if there is, we continue + else: + return "continue" + + class State(TypedDict): + __root__: Annotated[list[BaseMessage], add_messages] + + # Define a new graph + workflow = StateGraph(State) + + # Define the two nodes we will cycle between + workflow.add_node("agent", model) + workflow.add_node("tools", ToolNode(tools)) + + # 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": "tools", + # 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("tools", "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")) == [ + _AnyIdHumanMessage( + content="what is weather in sf", + ), + AIMessage( + content="", + tool_calls=[ + { + "id": "tool_call123", + "name": "search_api", + "args": {"query": "query"}, + } + ], + id="ai1", # respects ids passed in + ), + _AnyIdToolMessage( + content="result for query", + name="search_api", + tool_call_id="tool_call123", + ), + AIMessage( + content="", + tool_calls=[ + { + "id": "tool_call456", + "name": "search_api", + "args": {"query": "another"}, + } + ], + id="ai2", + ), + _AnyIdToolMessage( + content="result for another", + name="search_api", + tool_call_id="tool_call456", + ), + AIMessage(content="answer", id="ai3"), + ] + + assert [*app.stream([HumanMessage(content="what is weather in sf")])] == [ + { + "agent": AIMessage( + content="", + tool_calls=[ + { + "id": "tool_call123", + "name": "search_api", + "args": {"query": "query"}, + } + ], + id="ai1", + ) + }, + { + "tools": [ + ToolMessage( + content="result for query", + name="search_api", + tool_call_id="tool_call123", + id="00000000-0000-4000-8000-000000000024", + ) + ] + }, + { + "agent": AIMessage( + content="", + tool_calls=[ + { + "id": "tool_call456", + "name": "search_api", + "args": {"query": "another"}, + } + ], + id="ai2", + ) + }, + { + "tools": [ + ToolMessage( + content="result for another", + name="search_api", + tool_call_id="tool_call456", + id="00000000-0000-4000-8000-000000000030", + ) + ] + }, + {"agent": AIMessage(content="answer", id="ai3")}, + ] + + app_w_interrupt = workflow.compile( + checkpointer=sync_checkpointer, + interrupt_after=["agent"], + ) + config = {"configurable": {"thread_id": "1"}} + + assert [ + c for c in app_w_interrupt.stream(("human", "what is weather in sf"), config) + ] == [ + { + "agent": AIMessage( + content="", + tool_calls=[ + { + "id": "tool_call123", + "name": "search_api", + "args": {"query": "query"}, + } + ], + id="ai1", + ) + }, + {"__interrupt__": ()}, + ] + + assert app_w_interrupt.get_state(config) == StateSnapshot( + values=[ + _AnyIdHumanMessage(content="what is weather in sf"), + AIMessage( + content="", + tool_calls=[ + { + "id": "tool_call123", + "name": "search_api", + "args": {"query": "query"}, + } + ], + id="ai1", + ), + ], + tasks=(PregelTask(AnyStr(), "tools", (PULL, "tools")),), + next=("tools",), + config={ + "configurable": { + "thread_id": "1", + "checkpoint_ns": "", + "checkpoint_id": AnyStr(), + } + }, + created_at=AnyStr(), + metadata={ + "parents": {}, + "source": "loop", + "step": 1, + "thread_id": "1", + }, + parent_config=None, + interrupts=(), + ) + + # modify ai message + last_message = app_w_interrupt.get_state(config).values[-1] + last_message.tool_calls[0]["args"] = {"query": "a different query"} + next_config = app_w_interrupt.update_state(config, last_message) + + # message was replaced instead of appended + assert app_w_interrupt.get_state(config) == StateSnapshot( + values=[ + _AnyIdHumanMessage(content="what is weather in sf"), + AIMessage( + content="", + id="ai1", + tool_calls=[ + { + "id": "tool_call123", + "name": "search_api", + "args": {"query": "a different query"}, + } + ], + ), + ], + tasks=(PregelTask(AnyStr(), "tools", (PULL, "tools")),), + next=("tools",), + config=next_config, + created_at=AnyStr(), + metadata={ + "parents": {}, + "source": "update", + "step": 2, + "thread_id": "1", + }, + parent_config=( + list(app_w_interrupt.checkpointer.list(config, limit=2))[-1].config + ), + interrupts=(), + ) + + assert [c for c in app_w_interrupt.stream(None, config)] == [ + { + "tools": [ + _AnyIdToolMessage( + content="result for a different query", + name="search_api", + tool_call_id="tool_call123", + ) + ] + }, + { + "agent": AIMessage( + content="", + tool_calls=[ + { + "id": "tool_call456", + "name": "search_api", + "args": {"query": "another"}, + } + ], + id="ai2", + ) + }, + {"__interrupt__": ()}, + ] + + assert app_w_interrupt.get_state(config) == StateSnapshot( + values=[ + _AnyIdHumanMessage(content="what is weather in sf"), + AIMessage( + content="", + id="ai1", + tool_calls=[ + { + "id": "tool_call123", + "name": "search_api", + "args": {"query": "a different query"}, + } + ], + ), + _AnyIdToolMessage( + content="result for a different query", + name="search_api", + tool_call_id="tool_call123", + id=AnyStr(), + ), + AIMessage( + content="", + tool_calls=[ + { + "id": "tool_call456", + "name": "search_api", + "args": {"query": "another"}, + } + ], + id="ai2", + ), + ], + tasks=(PregelTask(AnyStr(), "tools", (PULL, "tools")),), + next=("tools",), + config={ + "configurable": { + "thread_id": "1", + "checkpoint_ns": "", + "checkpoint_id": AnyStr(), + } + }, + created_at=AnyStr(), + metadata={ + "parents": {}, + "source": "loop", + "step": 4, + "thread_id": "1", + }, + parent_config=( + list(app_w_interrupt.checkpointer.list(config, limit=2))[-1].config + ), + interrupts=(), + ) + + app_w_interrupt.update_state( + config, + AIMessage(content="answer", id="ai2"), # replace existing message + ) + + # replaces message even if object identity is different, as long as id is the same + assert app_w_interrupt.get_state(config) == StateSnapshot( + values=[ + _AnyIdHumanMessage(content="what is weather in sf"), + AIMessage( + content="", + id="ai1", + tool_calls=[ + { + "id": "tool_call123", + "name": "search_api", + "args": {"query": "a different query"}, + } + ], + ), + _AnyIdToolMessage( + content="result for a different query", + name="search_api", + tool_call_id="tool_call123", + id=AnyStr(), + ), + AIMessage(content="answer", id="ai2"), + ], + tasks=(), + next=(), + config={ + "configurable": { + "thread_id": "1", + "checkpoint_ns": "", + "checkpoint_id": AnyStr(), + } + }, + created_at=AnyStr(), + metadata={ + "parents": {}, + "source": "update", + "step": 5, + "thread_id": "1", + }, + parent_config=( + list(app_w_interrupt.checkpointer.list(config, limit=2))[-1].config + ), + interrupts=(), + ) + + app_w_interrupt = workflow.compile( + checkpointer=sync_checkpointer, + interrupt_before=["tools"], + ) + config = {"configurable": {"thread_id": "2"}} + model.i = 0 # reset the llm + + assert [c for c in app_w_interrupt.stream("what is weather in sf", config)] == [ + { + "agent": AIMessage( + content="", + tool_calls=[ + { + "id": "tool_call123", + "name": "search_api", + "args": {"query": "query"}, + } + ], + id="ai1", + ) + }, + {"__interrupt__": ()}, + ] + + assert app_w_interrupt.get_state(config) == StateSnapshot( + values=[ + _AnyIdHumanMessage(content="what is weather in sf"), + AIMessage( + content="", + tool_calls=[ + { + "id": "tool_call123", + "name": "search_api", + "args": {"query": "query"}, + } + ], + id="ai1", + ), + ], + tasks=(PregelTask(AnyStr(), "tools", (PULL, "tools")),), + next=("tools",), + config={ + "configurable": { + "thread_id": "2", + "checkpoint_ns": "", + "checkpoint_id": AnyStr(), + } + }, + created_at=AnyStr(), + metadata={ + "parents": {}, + "source": "loop", + "step": 1, + "thread_id": "2", + }, + parent_config=None, + interrupts=(), + ) + + # modify ai message + last_message = app_w_interrupt.get_state(config).values[-1] + last_message.tool_calls[0]["args"] = {"query": "a different query"} + app_w_interrupt.update_state(config, last_message) + + # message was replaced instead of appended + assert app_w_interrupt.get_state(config) == StateSnapshot( + values=[ + _AnyIdHumanMessage(content="what is weather in sf"), + AIMessage( + content="", + id="ai1", + tool_calls=[ + { + "id": "tool_call123", + "name": "search_api", + "args": {"query": "a different query"}, + } + ], + ), + ], + tasks=(PregelTask(AnyStr(), "tools", (PULL, "tools")),), + next=("tools",), + config={ + "configurable": { + "thread_id": "2", + "checkpoint_ns": "", + "checkpoint_id": AnyStr(), + } + }, + created_at=AnyStr(), + metadata={ + "parents": {}, + "source": "update", + "step": 2, + "thread_id": "2", + }, + parent_config=( + list(app_w_interrupt.checkpointer.list(config, limit=2))[-1].config + ), + interrupts=(), + ) + + assert [c for c in app_w_interrupt.stream(None, config)] == [ + { + "tools": [ + _AnyIdToolMessage( + content="result for a different query", + name="search_api", + tool_call_id="tool_call123", + ) + ] + }, + { + "agent": AIMessage( + content="", + tool_calls=[ + { + "id": "tool_call456", + "name": "search_api", + "args": {"query": "another"}, + } + ], + id="ai2", + ) + }, + {"__interrupt__": ()}, + ] + + assert app_w_interrupt.get_state(config) == StateSnapshot( + values=[ + _AnyIdHumanMessage(content="what is weather in sf"), + AIMessage( + content="", + id="ai1", + tool_calls=[ + { + "id": "tool_call123", + "name": "search_api", + "args": {"query": "a different query"}, + } + ], + ), + _AnyIdToolMessage( + content="result for a different query", + name="search_api", + tool_call_id="tool_call123", + id=AnyStr(), + ), + AIMessage( + content="", + tool_calls=[ + { + "id": "tool_call456", + "name": "search_api", + "args": {"query": "another"}, + } + ], + id="ai2", + ), + ], + tasks=(PregelTask(AnyStr(), "tools", (PULL, "tools")),), + next=("tools",), + config={ + "configurable": { + "thread_id": "2", + "checkpoint_ns": "", + "checkpoint_id": AnyStr(), + } + }, + created_at=AnyStr(), + metadata={ + "parents": {}, + "source": "loop", + "step": 4, + "thread_id": "2", + }, + parent_config=( + list(app_w_interrupt.checkpointer.list(config, limit=2))[-1].config + ), + interrupts=(), + ) + + app_w_interrupt.update_state( + config, + AIMessage(content="answer", id="ai2"), + ) + + # replaces message even if object identity is different, as long as id is the same + assert app_w_interrupt.get_state(config) == StateSnapshot( + values=[ + _AnyIdHumanMessage(content="what is weather in sf"), + AIMessage( + content="", + id="ai1", + tool_calls=[ + { + "id": "tool_call123", + "name": "search_api", + "args": {"query": "a different query"}, + } + ], + ), + _AnyIdToolMessage( + content="result for a different query", + name="search_api", + tool_call_id="tool_call123", + ), + AIMessage(content="answer", id="ai2"), + ], + tasks=(), + next=(), + config={ + "configurable": { + "thread_id": "2", + "checkpoint_ns": "", + "checkpoint_id": AnyStr(), + } + }, + created_at=AnyStr(), + metadata={ + "parents": {}, + "source": "update", + "step": 5, + "thread_id": "2", + }, + parent_config=( + list(app_w_interrupt.checkpointer.list(config, limit=2))[-1].config + ), + interrupts=(), + ) + + # add an extra message as if it came from "tools" node + app_w_interrupt.update_state(config, ("ai", "an extra message"), as_node="tools") + + # extra message is coerced BaseMessge and appended + # now the next node is "agent" per the graph edges + assert app_w_interrupt.get_state(config) == StateSnapshot( + values=[ + _AnyIdHumanMessage(content="what is weather in sf"), + AIMessage( + content="", + id="ai1", + tool_calls=[ + { + "id": "tool_call123", + "name": "search_api", + "args": {"query": "a different query"}, + } + ], + ), + _AnyIdToolMessage( + content="result for a different query", + name="search_api", + tool_call_id="tool_call123", + id=AnyStr(), + ), + AIMessage(content="answer", id="ai2"), + _AnyIdAIMessage(content="an extra message"), + ], + tasks=(PregelTask(AnyStr(), "agent", (PULL, "agent")),), + next=("agent",), + config={ + "configurable": { + "thread_id": "2", + "checkpoint_ns": "", + "checkpoint_id": AnyStr(), + } + }, + created_at=AnyStr(), + metadata={ + "parents": {}, + "source": "update", + "step": 6, + "thread_id": "2", + }, + parent_config=( + list(app_w_interrupt.checkpointer.list(config, limit=2))[-1].config + ), + interrupts=(), + ) + + # create new graph with one more state key, reuse previous thread history + + def simple_add(left, right): + if not isinstance(right, list): + right = [right] + return left + right + + class MoreState(TypedDict): + __root__: Annotated[list[BaseMessage], simple_add] + something_else: str + + # Define a new graph + new_workflow = StateGraph(MoreState) + new_workflow.add_node( + "agent", RunnableMap(__root__=RunnablePick("__root__") | model) + ) + new_workflow.add_node( + "tools", RunnableMap(__root__=RunnablePick("__root__") | ToolNode(tools)) + ) + new_workflow.set_entry_point("agent") + new_workflow.add_conditional_edges( + "agent", + RunnablePick("__root__") | should_continue, + { + # If `tools`, then we call the tool node. + "continue": "tools", + # Otherwise we finish. + "end": END, + }, + ) + new_workflow.add_edge("tools", "agent") + new_app = new_workflow.compile(checkpointer=sync_checkpointer) + model.i = 0 # reset the llm + + # previous state is converted to new schema + assert new_app.get_state(config) == StateSnapshot( + values={ + "__root__": [ + _AnyIdHumanMessage(content="what is weather in sf"), + AIMessage( + content="", + id="ai1", + tool_calls=[ + { + "id": "tool_call123", + "name": "search_api", + "args": {"query": "a different query"}, + } + ], + ), + _AnyIdToolMessage( + content="result for a different query", + name="search_api", + tool_call_id="tool_call123", + ), + AIMessage(content="answer", id="ai2"), + _AnyIdAIMessage(content="an extra message"), + ] + }, + tasks=(PregelTask(AnyStr(), "agent", (PULL, "agent")),), + next=("agent",), + config={ + "configurable": { + "thread_id": "2", + "checkpoint_ns": "", + "checkpoint_id": AnyStr(), + } + }, + created_at=AnyStr(), + metadata={ + "parents": {}, + "source": "update", + "step": 6, + "thread_id": "2", + }, + parent_config=(list(new_app.checkpointer.list(config, limit=2))[-1].config), + interrupts=(), + ) + + # new input is merged to old state + assert new_app.invoke( + { + "__root__": [HumanMessage(content="what is weather in la")], + "something_else": "value", + }, + config, + interrupt_before=["agent"], + ) == { + "__root__": [ + HumanMessage( + content="what is weather in sf", + id="00000000-0000-4000-8000-000000000051", + ), + AIMessage( + content="", + id="ai1", + tool_calls=[ + { + "name": "search_api", + "args": {"query": "a different query"}, + "id": "tool_call123", + } + ], + ), + _AnyIdToolMessage( + content="result for a different query", + name="search_api", + tool_call_id="tool_call123", + ), + AIMessage(content="answer", id="ai2"), + AIMessage( + content="an extra message", id="00000000-0000-4000-8000-000000000066" + ), + HumanMessage(content="what is weather in la"), + ], + "something_else": "value", + } + + def test_in_one_fan_out_out_one_graph_state() -> None: def sorted_add(x: list[str], y: list[str]) -> list[str]: return sorted(operator.add(x, y)) diff --git a/libs/langgraph/tests/test_large_cases_async.py b/libs/langgraph/tests/test_large_cases_async.py index 90d1a0c33..6e16442db 100644 --- a/libs/langgraph/tests/test_large_cases_async.py +++ b/libs/langgraph/tests/test_large_cases_async.py @@ -20,9 +20,10 @@ from langgraph.channels.ephemeral_value import EphemeralValue from langgraph.channels.last_value import LastValue from langgraph.checkpoint.base import BaseCheckpointSaver from langgraph.constants import END, PULL, PUSH, START -from langgraph.graph.message import add_messages +from langgraph.graph.message import MessageGraph, add_messages from langgraph.graph.state import StateGraph from langgraph.prebuilt.chat_agent_executor import create_react_agent +from langgraph.prebuilt.tool_node import ToolNode from langgraph.pregel import NodeBuilder, Pregel from langgraph.types import PregelTask, Send, StateSnapshot, StreamWriter from tests.any_int import AnyInt @@ -2082,7 +2083,417 @@ async def test_state_graph_packets(async_checkpointer: BaseCheckpointSaver) -> N ) +async def test_message_graph(async_checkpointer: BaseCheckpointSaver) -> None: + from langchain_core.language_models.fake_chat_models import ( + FakeMessagesListChatModel, + ) + from langchain_core.messages import AIMessage, HumanMessage + from langchain_core.tools import tool + + 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="", + tool_calls=[ + { + "id": "tool_call123", + "name": "search_api", + "args": {"query": "query"}, + } + ], + id="ai1", + ), + AIMessage( + content="", + tool_calls=[ + { + "id": "tool_call456", + "name": "search_api", + "args": {"query": "another"}, + } + ], + id="ai2", + ), + AIMessage(content="answer", id="ai3"), + ] + ) + + # 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 not last_message.tool_calls: + return "end" + # Otherwise if there is, we continue + else: + return "continue" + + # Define a new graph + workflow = MessageGraph() + + # Define the two nodes we will cycle between + workflow.add_node("agent", model) + workflow.add_node("tools", ToolNode(tools)) + + # 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": "tools", + # 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("tools", "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")) == [ + _AnyIdHumanMessage( + content="what is weather in sf", + ), + AIMessage( + content="", + tool_calls=[ + { + "id": "tool_call123", + "name": "search_api", + "args": {"query": "query"}, + } + ], + id="ai1", # respects ids passed in + ), + _AnyIdToolMessage( + content="result for query", + name="search_api", + tool_call_id="tool_call123", + ), + AIMessage( + content="", + tool_calls=[ + { + "id": "tool_call456", + "name": "search_api", + "args": {"query": "another"}, + } + ], + id="ai2", + ), + _AnyIdToolMessage( + content="result for another", + name="search_api", + tool_call_id="tool_call456", + ), + AIMessage(content="answer", id="ai3"), + ] + + assert [ + c async for c in app.astream([HumanMessage(content="what is weather in sf")]) + ] == [ + { + "agent": AIMessage( + content="", + tool_calls=[ + { + "id": "tool_call123", + "name": "search_api", + "args": {"query": "query"}, + } + ], + id="ai1", + ) + }, + { + "tools": [ + _AnyIdToolMessage( + content="result for query", + name="search_api", + tool_call_id="tool_call123", + ) + ] + }, + { + "agent": AIMessage( + content="", + tool_calls=[ + { + "id": "tool_call456", + "name": "search_api", + "args": {"query": "another"}, + } + ], + id="ai2", + ) + }, + { + "tools": [ + _AnyIdToolMessage( + content="result for another", + name="search_api", + tool_call_id="tool_call456", + ) + ] + }, + {"agent": AIMessage(content="answer", id="ai3")}, + ] + + app_w_interrupt = workflow.compile( + checkpointer=async_checkpointer, + interrupt_after=["agent"], + ) + config = {"configurable": {"thread_id": "1"}} + + assert [ + c + async for c in app_w_interrupt.astream( + HumanMessage(content="what is weather in sf"), config + ) + ] == [ + { + "agent": AIMessage( + content="", + tool_calls=[ + { + "id": "tool_call123", + "name": "search_api", + "args": {"query": "query"}, + } + ], + id="ai1", + ) + }, + {"__interrupt__": ()}, + ] + + tup = await app_w_interrupt.checkpointer.aget_tuple(config) + assert await app_w_interrupt.aget_state(config) == StateSnapshot( + values=[ + _AnyIdHumanMessage(content="what is weather in sf"), + AIMessage( + content="", + tool_calls=[ + { + "id": "tool_call123", + "name": "search_api", + "args": {"query": "query"}, + } + ], + id="ai1", + ), + ], + tasks=(PregelTask(AnyStr(), "tools", (PULL, "tools")),), + next=("tools",), + config=tup.config, + created_at=tup.checkpoint["ts"], + metadata={ + "parents": {}, + "source": "loop", + "step": 1, + "thread_id": "1", + }, + parent_config=None, + interrupts=(), + ) + + # modify ai message + last_message = (await app_w_interrupt.aget_state(config)).values[-1] + last_message.tool_calls[0]["args"] = {"query": "a different query"} + await app_w_interrupt.aupdate_state(config, last_message) + + # message was replaced instead of appended + tup = await app_w_interrupt.checkpointer.aget_tuple(config) + assert await app_w_interrupt.aget_state(config) == StateSnapshot( + values=[ + _AnyIdHumanMessage(content="what is weather in sf"), + AIMessage( + content="", + id="ai1", + tool_calls=[ + { + "id": "tool_call123", + "name": "search_api", + "args": {"query": "a different query"}, + } + ], + ), + ], + tasks=(PregelTask(AnyStr(), "tools", (PULL, "tools")),), + next=("tools",), + config=tup.config, + created_at=tup.checkpoint["ts"], + metadata={ + "parents": {}, + "source": "update", + "step": 2, + "thread_id": "1", + }, + parent_config=( + [c async for c in app_w_interrupt.checkpointer.alist(config, limit=2)][ + -1 + ].config + ), + interrupts=(), + ) + + assert [c async for c in app_w_interrupt.astream(None, config)] == [ + { + "tools": [ + _AnyIdToolMessage( + content="result for a different query", + name="search_api", + tool_call_id="tool_call123", + ) + ] + }, + { + "agent": AIMessage( + content="", + tool_calls=[ + { + "id": "tool_call456", + "name": "search_api", + "args": {"query": "another"}, + } + ], + id="ai2", + ) + }, + {"__interrupt__": ()}, + ] + + tup = await app_w_interrupt.checkpointer.aget_tuple(config) + assert await app_w_interrupt.aget_state(config) == StateSnapshot( + values=[ + _AnyIdHumanMessage(content="what is weather in sf"), + AIMessage( + content="", + id="ai1", + tool_calls=[ + { + "id": "tool_call123", + "name": "search_api", + "args": {"query": "a different query"}, + } + ], + ), + _AnyIdToolMessage( + content="result for a different query", + name="search_api", + tool_call_id="tool_call123", + ), + AIMessage( + content="", + tool_calls=[ + { + "id": "tool_call456", + "name": "search_api", + "args": {"query": "another"}, + } + ], + id="ai2", + ), + ], + tasks=(PregelTask(AnyStr(), "tools", (PULL, "tools")),), + next=("tools",), + config=tup.config, + created_at=tup.checkpoint["ts"], + metadata={ + "parents": {}, + "source": "loop", + "step": 4, + "thread_id": "1", + }, + parent_config=( + [c async for c in app_w_interrupt.checkpointer.alist(config, limit=2)][ + -1 + ].config + ), + interrupts=(), + ) + + await app_w_interrupt.aupdate_state( + config, + AIMessage(content="answer", id="ai2"), + ) + + # replaces message even if object identity is different, as long as id is the same + tup = await app_w_interrupt.checkpointer.aget_tuple(config) + assert await app_w_interrupt.aget_state(config) == StateSnapshot( + values=[ + _AnyIdHumanMessage(content="what is weather in sf"), + AIMessage( + content="", + id="ai1", + tool_calls=[ + { + "id": "tool_call123", + "name": "search_api", + "args": {"query": "a different query"}, + } + ], + ), + _AnyIdToolMessage( + content="result for a different query", + name="search_api", + tool_call_id="tool_call123", + ), + AIMessage(content="answer", id="ai2"), + ], + tasks=(), + next=(), + config=tup.config, + created_at=tup.checkpoint["ts"], + metadata={ + "parents": {}, + "source": "update", + "step": 5, + "thread_id": "1", + }, + parent_config=( + [c async for c in app_w_interrupt.checkpointer.alist(config, limit=2)][ + -1 + ].config + ), + interrupts=(), + ) + + async def test_in_one_fan_out_out_one_graph_state() -> None: + def sorted_add(x: list[str], y: list[str]) -> list[str]: + return sorted(operator.add(x, y)) + class State(TypedDict, total=False): query: str answer: str diff --git a/libs/langgraph/tests/test_pregel.py b/libs/langgraph/tests/test_pregel.py index 2f173b105..05561bc78 100644 --- a/libs/langgraph/tests/test_pregel.py +++ b/libs/langgraph/tests/test_pregel.py @@ -46,7 +46,7 @@ from langgraph.constants import CONFIG_KEY_NODE_FINISHED, ERROR, PULL, START from langgraph.errors import InvalidUpdateError, ParentCommand from langgraph.func import entrypoint, task from langgraph.graph import END, StateGraph -from langgraph.graph.message import MessagesState, add_messages +from langgraph.graph.message import MessageGraph, MessagesState, add_messages from langgraph.prebuilt.tool_node import ToolNode from langgraph.pregel import ( GraphRecursionError, @@ -3994,14 +3994,9 @@ def test_checkpoint_metadata(sync_checkpointer: BaseCheckpointSaver) -> None: def test_remove_message_via_state_update( sync_checkpointer: BaseCheckpointSaver, ) -> None: - from langchain_core.messages import ( - AIMessage, - AnyMessage, - HumanMessage, - RemoveMessage, - ) + from langchain_core.messages import AIMessage, HumanMessage, RemoveMessage - workflow = StateGraph(Annotated[list[AnyMessage], add_messages]) + workflow = MessageGraph() workflow.add_node( "chatbot", lambda state: [ @@ -4032,14 +4027,9 @@ def test_remove_message_via_state_update( def test_remove_message_from_node(): - from langchain_core.messages import ( - AIMessage, - AnyMessage, - HumanMessage, - RemoveMessage, - ) + from langchain_core.messages import AIMessage, HumanMessage, RemoveMessage - workflow = StateGraph(Annotated[list[AnyMessage], add_messages]) + workflow = MessageGraph() workflow.add_node( "chatbot", lambda state: [ diff --git a/libs/prebuilt/langgraph/prebuilt/tool_node.py b/libs/prebuilt/langgraph/prebuilt/tool_node.py index 384e8acd0..c66835699 100644 --- a/libs/prebuilt/langgraph/prebuilt/tool_node.py +++ b/libs/prebuilt/langgraph/prebuilt/tool_node.py @@ -629,7 +629,7 @@ def tools_condition( Args: state: The state to check for - tool calls. Must have a list of messages or have the + tool calls. Must have a list of messages (MessageGraph) or have the "messages" key (StateGraph). Returns: diff --git a/libs/prebuilt/langgraph/prebuilt/tool_validator.py b/libs/prebuilt/langgraph/prebuilt/tool_validator.py index 7ef75a710..0e58c7d6c 100644 --- a/libs/prebuilt/langgraph/prebuilt/tool_validator.py +++ b/libs/prebuilt/langgraph/prebuilt/tool_validator.py @@ -2,7 +2,7 @@ in a langchain graph. It applies a pydantic schema to tool_calls in the models' outputs, and returns a ToolMessage with the validated content. If the schema is not valid, it returns a ToolMessage with the error message. The ValidationNode can be used in a -StateGraph with a "messages" key. If multiple tool calls are +StateGraph with a "messages" key or in a MessageGraph. If multiple tool calls are requested, they will be run in parallel. """ @@ -49,7 +49,7 @@ def _default_format_error( class ValidationNode(RunnableCallable): """A node that validates all tools requests from the last AIMessage. - It can be used in StateGraph with a "messages" key. + It can be used either in StateGraph with a "messages" key or in MessageGraph. !!! note