Revert "Remove MessageGraph (#4875)"

This reverts commit a5e6223569.
This commit is contained in:
Nuno Campos
2025-06-13 14:21:05 -07:00
parent 053b606b46
commit 3fa3a586b5
9 changed files with 2049 additions and 31 deletions
+5 -5
View File
@@ -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)
+3 -3
View File
@@ -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",
+2 -1
View File
@@ -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",
]
+52
View File
@@ -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]
File diff suppressed because it is too large Load Diff
+412 -1
View File
@@ -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
+5 -15
View File
@@ -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: [
@@ -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:
@@ -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