From a5e6223569c6be5062fa3c22e63efa2468dc094f Mon Sep 17 00:00:00 2001 From: Sydney Runkle <54324534+sydney-runkle@users.noreply.github.com> Date: Thu, 29 May 2025 21:50:23 -0400 Subject: [PATCH] Remove `MessageGraph` (#4875) Remove MessageGraph Co-authored-by: Nuno Campos --- 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, 31 insertions(+), 2049 deletions(-) diff --git a/docs/docs/cloud/deployment/graph_rebuild.md b/docs/docs/cloud/deployment/graph_rebuild.md index eb5b8a322..75dc26377 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, 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 @@ -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) \ No newline at end of file +See more info on LangGraph API configuration file [here](../reference/cli.md#configuration-file) diff --git a/docs/docs/tutorials/extraction/retries.ipynb b/docs/docs/tutorials/extraction/retries.ipynb index de081052c..ccd552e3a 100644 --- a/docs/docs/tutorials/extraction/retries.ipynb +++ b/docs/docs/tutorials/extraction/retries.ipynb @@ -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", - " # or a StateGraph with 'messages' as the key.\n", + " # These functions let the step be used in a\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", diff --git a/libs/langgraph/langgraph/graph/__init__.py b/libs/langgraph/langgraph/graph/__init__.py index 2581713c3..9b67c2a42 100644 --- a/libs/langgraph/langgraph/graph/__init__.py +++ b/libs/langgraph/langgraph/graph/__init__.py @@ -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", ] diff --git a/libs/langgraph/langgraph/graph/message.py b/libs/langgraph/langgraph/graph/message.py index 2e2682fe8..2a752e0c5 100644 --- a/libs/langgraph/langgraph/graph/message.py +++ b/libs/langgraph/langgraph/graph/message.py @@ -24,7 +24,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] @@ -226,57 +225,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] diff --git a/libs/langgraph/tests/test_large_cases.py b/libs/langgraph/tests/test_large_cases.py index 1993d154b..d0a6e657c 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, Any, Literal, Optional, Union, cast +from typing import Annotated, Literal, Optional, Union, cast import pytest -from langchain_core.runnables import RunnableConfig, RunnableMap, RunnablePick +from langchain_core.runnables import RunnableConfig from pytest_mock import MockerFixture from syrupy import SnapshotAssertion from typing_extensions import TypedDict @@ -18,9 +18,8 @@ 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 MessageGraph, MessagesState, add_messages +from langgraph.graph.message import 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.store.base import BaseStore from langgraph.store.memory import InMemoryStore @@ -2382,1569 +2381,6 @@ 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 9c2d7935c..0c2c1ff5a 100644 --- a/libs/langgraph/tests/test_large_cases_async.py +++ b/libs/langgraph/tests/test_large_cases_async.py @@ -20,10 +20,9 @@ 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 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.store.base import BaseStore from langgraph.types import PregelTask, Send, StateSnapshot, StreamWriter @@ -2084,417 +2083,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 - ) - ] == [ - { - "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 d41999d1c..0f0cce810 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 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, @@ -4211,9 +4211,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: [ @@ -4244,9 +4249,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: [ diff --git a/libs/prebuilt/langgraph/prebuilt/tool_node.py b/libs/prebuilt/langgraph/prebuilt/tool_node.py index 03bf1c3e9..b2c7cf8d0 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 (MessageGraph) or have the + tool calls. Must have a list of messages 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 0e58c7d6c..7ef75a710 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 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