Reapply "Remove MessageGraph (#4875 )"

This reverts commit 3fa3a586b5.
2026-09-08 18:57:52 +02:00 · 2025-06-29 11:20:08 -04:00
9 changed files with 31 additions and 2044 deletions
@@ -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, MessageGraph
+from langgraph.graph import END, START, StateGraph, MessagesState
 model = ChatOpenAI(temperature=0)
-graph_workflow = MessageGraph()
+graph_workflow = StateGraph(MessagesState)
 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, MessageGraph
+from langgraph.graph import END, START
 from langgraph.graph.state import StateGraph
 from langgraph.graph.message import add_messages
 from langgraph.prebuilt import ToolNode
@@ -89,7 +89,7 @@
  },
  {
   "cell_type": "code",
-   "execution_count": 2,
+   "execution_count": null,
   "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 MessageGraph\n",
+    "    # These functions let the step be used in a\n",
-    "    # or a StateGraph with 'messages' as the key.\n",
+    "    # 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",
@@ -1,12 +1,11 @@
 from langgraph.constants import END, START
-from langgraph.graph.message import MessageGraph, MessagesState, add_messages
+from langgraph.graph.message import MessagesState, add_messages
 from langgraph.graph.state import StateGraph
 __all__ = [
    "END",
    "START",
    "StateGraph",
    "MessageGraph",
    "add_messages",
    "MessagesState",
 ]
@@ -25,7 +25,6 @@ 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]
@@ -227,57 +226,6 @@ 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]
@@ -20,10 +20,9 @@ from langgraph.channels.last_value import LastValue
 from langgraph.channels.untracked_value import UntrackedValue
 from langgraph.checkpoint.base import BaseCheckpointSaver
 from langgraph.constants import END, PULL, PUSH, START
-from langgraph.graph.message import MessageGraph, add_messages
+from langgraph.graph.message import 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
@@ -2059,415 +2058,7 @@ 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,
            checkpoint_during=False,
        )
    ] == [
        {
            "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,
        },
        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,
        },
        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,
        },
        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,
        },
        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
@@ -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 MessageGraph, MessagesState, add_messages
+from langgraph.graph.message import MessagesState, add_messages
 from langgraph.prebuilt.tool_node import ToolNode
 from langgraph.pregel import (
    GraphRecursionError,
@@ -3984,9 +3984,14 @@ 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, HumanMessage, RemoveMessage
+    from langchain_core.messages import (
        AIMessage,
        AnyMessage,
        HumanMessage,
        RemoveMessage,
    )
-    workflow = MessageGraph()
+    workflow = StateGraph(Annotated[list[AnyMessage], add_messages])
    workflow.add_node(
        "chatbot",
        lambda state: [
@@ -4017,9 +4022,14 @@ def test_remove_message_via_state_update(
 def test_remove_message_from_node():
-    from langchain_core.messages import AIMessage, HumanMessage, RemoveMessage
+    from langchain_core.messages import (
        AIMessage,
        AnyMessage,
        HumanMessage,
        RemoveMessage,
    )
-    workflow = MessageGraph()
+    workflow = StateGraph(Annotated[list[AnyMessage], add_messages])
    workflow.add_node(
        "chatbot",
        lambda state: [
@@ -629,7 +629,7 @@ def tools_condition(
    Args:
        state: The state to check for
-            tool calls. Must have a list of messages (MessageGraph) or have the
+            tool calls. Must have a list of messages or have the
            "messages" key (StateGraph).
    Returns:
@@ -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 or in a MessageGraph. If multiple tool calls are
+StateGraph with a "messages" key. 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 either in StateGraph with a "messages" key or in MessageGraph.
+    It can be used in StateGraph with a "messages" key.
    !!! note