wehub/langchain-ai--langgraph
1
0
Fork 0
Code Issues Pull Requests Actions Deployments Models Agent Notes Packages Projects Releases Wiki Activity
Files
main
langchain-ai--langgraph/libs/langgraph/tests/test_large_cases.py
T
wehub-resource-sync a7d6d88f6f
CI / changes (push) Has been cancelled
Details
CI / cd libs/checkpoint (push) Has been cancelled
Details
CI / cd libs/checkpoint-conformance (push) Has been cancelled
Details
CI / cd libs/checkpoint-postgres (push) Has been cancelled
Details
CI / cd libs/checkpoint-sqlite (push) Has been cancelled
Details
CI / cd libs/cli (push) Has been cancelled
Details
CI / cd libs/prebuilt (push) Has been cancelled
Details
CI / cd libs/sdk-py (push) Has been cancelled
Details
CI / cd libs/langgraph (push) Has been cancelled
Details
CI / Check SDK methods matching (push) Has been cancelled
Details
CI / Check CLI schema hasn't changed #3.13 (push) Has been cancelled
Details
CI / CLI integration test (push) Has been cancelled
Details
CI / sdk-py integration test (push) Has been cancelled
Details
CI / CI Success (push) Has been cancelled
Details
baseline / benchmark (push) Has been cancelled
Details
Deploy Redirects to GitHub Pages / deploy (push) Has been cancelled
Details
chore: import upstream snapshot with attribution
2026-07-13 12:37:18 +08:00

6987 lines
210 KiB
Python
Raw Permalink Blame History

import json
import operator
import re
import time
from dataclasses import replace
from typing import Annotated, Any, Literal, cast
import pytest
from langchain_core.messages import AIMessage, AnyMessage, ToolCall
from langchain_core.runnables import RunnableConfig, RunnableMap, RunnablePick
from langchain_core.tools import tool
from langchain_core.version import VERSION as LANGCHAIN_CORE_VERSION
from langgraph.checkpoint.base import BaseCheckpointSaver
from langgraph.checkpoint.memory import InMemorySaver
from langgraph.prebuilt.chat_agent_executor import create_react_agent
from langgraph.prebuilt.tool_node import ToolNode
from pytest_mock import MockerFixture
from syrupy import SnapshotAssertion
from typing_extensions import TypedDict
from langgraph._internal._constants import PULL, PUSH
from langgraph.channels.last_value import LastValue
from langgraph.channels.untracked_value import UntrackedValue
from langgraph.constants import END, START
from langgraph.graph import StateGraph
from langgraph.graph.message import MessagesState, add_messages
from langgraph.pregel import NodeBuilder, Pregel
from langgraph.types import (
Command,
Durability,
Interrupt,
PregelTask,
RetryPolicy,
Send,
StateSnapshot,
StreamWriter,
interrupt,
)
from tests.agents import AgentAction, AgentFinish
from tests.any_int import AnyInt
from tests.any_str import AnyDict, AnyStr, UnsortedSequence
from tests.fake_chat import FakeChatModel
from tests.fake_tracer import FakeTracer
from tests.messages import (
_AnyIdAIMessage,
_AnyIdAIMessageChunk,
_AnyIdHumanMessage,
_AnyIdToolMessage,
)
def test_invoke_two_processes_in_out_interrupt(
sync_checkpointer: BaseCheckpointSaver,
mocker: MockerFixture,
) -> None:
add_one = mocker.Mock(side_effect=lambda x: x + 1)
one = NodeBuilder().subscribe_only("input").do(add_one).write_to("inbox")
two = NodeBuilder().subscribe_only("inbox").do(add_one).write_to("output")
app = Pregel(
nodes={"one": one, "two": two},
channels={
"inbox": LastValue(int),
"output": LastValue(int),
"input": LastValue(int),
},
input_channels="input",
output_channels="output",
checkpointer=sync_checkpointer,
interrupt_after_nodes=["one"],
)
thread1 = {"configurable": {"thread_id": "1"}}
thread2 = {"configurable": {"thread_id": "2"}}
# start execution, stop at inbox
assert app.invoke(2, thread1, durability="async") is None
# inbox == 3
checkpoint = sync_checkpointer.get(thread1)
assert checkpoint is not None
assert checkpoint["channel_values"]["inbox"] == 3
# resume execution, finish
assert app.invoke(None, thread1, durability="async") == 4
# start execution again, stop at inbox
assert app.invoke(20, thread1, durability="async") is None
# inbox == 21
checkpoint = sync_checkpointer.get(thread1)
assert checkpoint is not None
assert checkpoint["channel_values"]["inbox"] == 21
# send a new value in, interrupting the previous execution
assert app.invoke(3, thread1, durability="async") is None
assert app.invoke(None, thread1, durability="async") == 5
# start execution again, stopping at inbox
assert app.invoke(20, thread2, durability="async") is None
# inbox == 21
snapshot = app.get_state(thread2)
assert snapshot.values["inbox"] == 21
assert snapshot.next == ("two",)
# update the state, resume
app.update_state(thread2, 25, as_node="one")
assert app.invoke(None, thread2) == 26
# no pending tasks
snapshot = app.get_state(thread2)
assert snapshot.next == ()
# list history
history = [c for c in app.get_state_history(thread1)]
assert len(history) == 8
assert history == [
StateSnapshot(
values={"inbox": 4, "output": 5, "input": 3},
tasks=(),
next=(),
config={
"configurable": {
"thread_id": "1",
"checkpoint_ns": "",
"checkpoint_id": AnyStr(),
}
},
metadata={
"parents": {},
"source": "loop",
"step": 6,
},
created_at=AnyStr(),
parent_config=history[1].config,
interrupts=(),
),
StateSnapshot(
values={"inbox": 4, "output": 4, "input": 3},
tasks=(PregelTask(AnyStr(), "two", (PULL, "two"), result={"output": 5}),),
next=("two",),
config={
"configurable": {
"thread_id": "1",
"checkpoint_ns": "",
"checkpoint_id": AnyStr(),
}
},
metadata={
"parents": {},
"source": "loop",
"step": 5,
},
created_at=AnyStr(),
parent_config=history[2].config,
interrupts=(),
),
StateSnapshot(
values={"inbox": 21, "output": 4, "input": 3},
tasks=(PregelTask(AnyStr(), "one", (PULL, "one"), result={"inbox": 4}),),
next=("one",),
config={
"configurable": {
"thread_id": "1",
"checkpoint_ns": "",
"checkpoint_id": AnyStr(),
}
},
metadata={
"parents": {},
"source": "input",
"step": 4,
},
created_at=AnyStr(),
parent_config=history[3].config,
interrupts=(),
),
StateSnapshot(
values={"inbox": 21, "output": 4, "input": 20},
tasks=(PregelTask(AnyStr(), "two", (PULL, "two")),),
next=("two",),
config={
"configurable": {
"thread_id": "1",
"checkpoint_ns": "",
"checkpoint_id": AnyStr(),
}
},
metadata={
"parents": {},
"source": "loop",
"step": 3,
},
created_at=AnyStr(),
parent_config=history[4].config,
interrupts=(),
),
StateSnapshot(
values={"inbox": 3, "output": 4, "input": 20},
tasks=(PregelTask(AnyStr(), "one", (PULL, "one"), result={"inbox": 21}),),
next=("one",),
config={
"configurable": {
"thread_id": "1",
"checkpoint_ns": "",
"checkpoint_id": AnyStr(),
}
},
metadata={
"parents": {},
"source": "input",
"step": 2,
},
created_at=AnyStr(),
parent_config=history[5].config,
interrupts=(),
),
StateSnapshot(
values={"inbox": 3, "output": 4, "input": 2},
tasks=(),
next=(),
config={
"configurable": {
"thread_id": "1",
"checkpoint_ns": "",
"checkpoint_id": AnyStr(),
}
},
metadata={
"parents": {},
"source": "loop",
"step": 1,
},
created_at=AnyStr(),
parent_config=history[6].config,
interrupts=(),
),
StateSnapshot(
values={"inbox": 3, "input": 2},
tasks=(PregelTask(AnyStr(), "two", (PULL, "two"), result={"output": 4}),),
next=("two",),
config={
"configurable": {
"thread_id": "1",
"checkpoint_ns": "",
"checkpoint_id": AnyStr(),
}
},
metadata={
"parents": {},
"source": "loop",
"step": 0,
},
created_at=AnyStr(),
parent_config=history[7].config,
interrupts=(),
),
StateSnapshot(
values={"input": 2},
tasks=(PregelTask(AnyStr(), "one", (PULL, "one"), result={"inbox": 3}),),
next=("one",),
config={
"configurable": {
"thread_id": "1",
"checkpoint_ns": "",
"checkpoint_id": AnyStr(),
}
},
metadata={
"parents": {},
"source": "input",
"step": -1,
},
created_at=AnyStr(),
parent_config=None,
interrupts=(),
),
]
# re-running from any previous checkpoint should re-run nodes
assert [c for c in app.stream(None, history[0].config, stream_mode="updates")] == []
assert [c for c in app.stream(None, history[1].config, stream_mode="updates")] == [
{"two": {"output": 5}},
]
assert [c for c in app.stream(None, history[2].config, stream_mode="updates")] == [
{"one": {"inbox": 4}},
{"__interrupt__": ()},
]
def test_fork_always_re_runs_nodes(
sync_checkpointer: BaseCheckpointSaver, mocker: MockerFixture
) -> None:
add_one = mocker.Mock(side_effect=lambda _: 1)
builder = StateGraph(Annotated[int, operator.add])
builder.add_node("add_one", add_one)
builder.add_edge(START, "add_one")
builder.add_conditional_edges("add_one", lambda cnt: "add_one" if cnt < 6 else END)
graph = builder.compile(checkpointer=sync_checkpointer)
thread1 = {"configurable": {"thread_id": "1"}}
# start execution, stop at inbox
assert [
*graph.stream(1, thread1, stream_mode=["values", "updates"], durability="async")
] == [
("values", 1),
("updates", {"add_one": 1}),
("values", 2),
("updates", {"add_one": 1}),
("values", 3),
("updates", {"add_one": 1}),
("values", 4),
("updates", {"add_one": 1}),
("values", 5),
("updates", {"add_one": 1}),
("values", 6),
]
# list history
history = [c for c in graph.get_state_history(thread1)]
assert history == [
StateSnapshot(
values=6,
tasks=(),
next=(),
config={
"configurable": {
"thread_id": "1",
"checkpoint_ns": "",
"checkpoint_id": AnyStr(),
}
},
metadata={
"parents": {},
"source": "loop",
"step": 5,
},
created_at=AnyStr(),
parent_config=history[1].config,
interrupts=(),
),
StateSnapshot(
values=5,
tasks=(PregelTask(AnyStr(), "add_one", (PULL, "add_one"), result=1),),
next=("add_one",),
config={
"configurable": {
"thread_id": "1",
"checkpoint_ns": "",
"checkpoint_id": AnyStr(),
}
},
metadata={
"parents": {},
"source": "loop",
"step": 4,
},
created_at=AnyStr(),
parent_config=history[2].config,
interrupts=(),
),
StateSnapshot(
values=4,
tasks=(PregelTask(AnyStr(), "add_one", (PULL, "add_one"), result=1),),
next=("add_one",),
config={
"configurable": {
"thread_id": "1",
"checkpoint_ns": "",
"checkpoint_id": AnyStr(),
}
},
metadata={
"parents": {},
"source": "loop",
"step": 3,
},
created_at=AnyStr(),
parent_config=history[3].config,
interrupts=(),
),
StateSnapshot(
values=3,
tasks=(PregelTask(AnyStr(), "add_one", (PULL, "add_one"), result=1),),
next=("add_one",),
config={
"configurable": {
"thread_id": "1",
"checkpoint_ns": "",
"checkpoint_id": AnyStr(),
}
},
metadata={
"parents": {},
"source": "loop",
"step": 2,
},
created_at=AnyStr(),
parent_config=history[4].config,
interrupts=(),
),
StateSnapshot(
values=2,
tasks=(PregelTask(AnyStr(), "add_one", (PULL, "add_one"), result=1),),
next=("add_one",),
config={
"configurable": {
"thread_id": "1",
"checkpoint_ns": "",
"checkpoint_id": AnyStr(),
}
},
metadata={
"parents": {},
"source": "loop",
"step": 1,
},
created_at=AnyStr(),
parent_config=history[5].config,
interrupts=(),
),
StateSnapshot(
values=1,
tasks=(PregelTask(AnyStr(), "add_one", (PULL, "add_one"), result=1),),
next=("add_one",),
config={
"configurable": {
"thread_id": "1",
"checkpoint_ns": "",
"checkpoint_id": AnyStr(),
}
},
metadata={
"parents": {},
"source": "loop",
"step": 0,
},
created_at=AnyStr(),
parent_config=history[6].config,
interrupts=(),
),
StateSnapshot(
values=0,
tasks=(PregelTask(AnyStr(), "__start__", (PULL, "__start__"), result=1),),
next=("__start__",),
config={
"configurable": {
"thread_id": "1",
"checkpoint_ns": "",
"checkpoint_id": AnyStr(),
}
},
metadata={
"parents": {},
"source": "input",
"step": -1,
},
created_at=AnyStr(),
parent_config=None,
interrupts=(),
),
]
# forking from any previous checkpoint should re-run nodes
assert [
c for c in graph.stream(None, history[0].config, stream_mode="updates")
] == []
assert [
c for c in graph.stream(None, history[1].config, stream_mode="updates")
] == [
{"add_one": 1},
]
assert [
c for c in graph.stream(None, history[2].config, stream_mode="updates")
] == [
{"add_one": 1},
{"add_one": 1},
]
def test_conditional_state_graph(
snapshot: SnapshotAssertion,
sync_checkpointer: BaseCheckpointSaver,
) -> None:
from langchain_core.language_models.fake import FakeStreamingListLLM
from langchain_core.prompts import PromptTemplate
from langchain_core.tools import tool
class AgentState(TypedDict, total=False):
input: Annotated[str, UntrackedValue]
agent_outcome: AgentAction | AgentFinish | None
intermediate_steps: Annotated[list[tuple[AgentAction, str]], operator.add]
class ToolState(TypedDict, total=False):
agent_outcome: AgentAction | AgentFinish
# Assemble the tools
@tool()
def search_api(query: str) -> str:
"""Searches the API for the query."""
return f"result for {query}"
tools = [search_api]
# Construct the agent
prompt = PromptTemplate.from_template("Hello!")
llm = FakeStreamingListLLM(
responses=[
"tool:search_api:query",
"tool:search_api:another",
"finish:answer",
]
)
def agent_parser(input: str) -> dict[str, AgentAction | AgentFinish]:
if input.startswith("finish"):
_, answer = input.split(":")
return {
"agent_outcome": AgentFinish(
return_values={"answer": answer}, log=input
)
}
else:
_, tool_name, tool_input = input.split(":")
return {
"agent_outcome": AgentAction(
tool=tool_name, tool_input=tool_input, log=input
)
}
agent = prompt | llm | agent_parser
# Define tool execution logic
def execute_tools(data: ToolState) -> dict:
# check session in data
assert "input" not in data
assert "intermediate_steps" not in data
# execute the tool
agent_action: AgentAction = data.pop("agent_outcome")
observation = {t.name: t for t in tools}[agent_action.tool].invoke(
agent_action.tool_input
)
return {"intermediate_steps": [[agent_action, observation]]}
# Define decision-making logic
def should_continue(data: AgentState) -> str:
# check session in data
# Logic to decide whether to continue in the loop or exit
if isinstance(data["agent_outcome"], AgentFinish):
return "exit"
else:
return "continue"
# Define a new graph
workflow = StateGraph(AgentState)
workflow.add_node("agent", agent)
workflow.add_node("tools", execute_tools, input_schema=ToolState)
workflow.set_entry_point("agent")
workflow.add_conditional_edges(
"agent", should_continue, {"continue": "tools", "exit": END}
)
workflow.add_edge("tools", "agent")
app = workflow.compile()
if isinstance(sync_checkpointer, InMemorySaver):
assert json.dumps(app.get_input_jsonschema()) == snapshot
assert json.dumps(app.get_output_jsonschema()) == 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({"input": "what is weather in sf"}) == {
"input": "what is weather in sf",
"intermediate_steps": [
[
AgentAction(
tool="search_api",
tool_input="query",
log="tool:search_api:query",
),
"result for query",
],
[
AgentAction(
tool="search_api",
tool_input="another",
log="tool:search_api:another",
),
"result for another",
],
],
"agent_outcome": AgentFinish(
return_values={"answer": "answer"}, log="finish:answer"
),
}
assert [*app.stream({"input": "what is weather in sf"})] == [
{
"agent": {
"agent_outcome": AgentAction(
tool="search_api",
tool_input="query",
log="tool:search_api:query",
),
}
},
{
"tools": {
"intermediate_steps": [
[
AgentAction(
tool="search_api",
tool_input="query",
log="tool:search_api:query",
),
"result for query",
]
],
}
},
{
"agent": {
"agent_outcome": AgentAction(
tool="search_api",
tool_input="another",
log="tool:search_api:another",
),
}
},
{
"tools": {
"intermediate_steps": [
[
AgentAction(
tool="search_api",
tool_input="another",
log="tool:search_api:another",
),
"result for another",
],
],
}
},
{
"agent": {
"agent_outcome": AgentFinish(
return_values={"answer": "answer"}, log="finish:answer"
),
}
},
]
# test state get/update methods with interrupt_after
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(
{"input": "what is weather in sf"}, config, durability="exit"
)
] == [
{
"agent": {
"agent_outcome": AgentAction(
tool="search_api",
tool_input="query",
log="tool:search_api:query",
),
}
},
{"__interrupt__": ()},
]
assert app_w_interrupt.get_state(config) == StateSnapshot(
values={
"agent_outcome": AgentAction(
tool="search_api", tool_input="query", log="tool:search_api:query"
),
"intermediate_steps": [],
},
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,
},
parent_config=None,
interrupts=(),
)
app_w_interrupt.update_state(
config,
{
"agent_outcome": AgentAction(
tool="search_api",
tool_input="query",
log="tool:search_api:a different query",
)
},
)
assert app_w_interrupt.get_state(config) == StateSnapshot(
values={
"agent_outcome": AgentAction(
tool="search_api",
tool_input="query",
log="tool:search_api:a different query",
),
"intermediate_steps": [],
},
tasks=(PregelTask(AnyStr(), "tools", (PULL, "tools")),),
next=("tools",),
config={
"configurable": {
"thread_id": "1",
"checkpoint_ns": "",
"checkpoint_id": AnyStr(),
}
},
created_at=AnyStr(),
metadata={
"parents": {},
"source": "update",
"step": 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": {
"intermediate_steps": [
[
AgentAction(
tool="search_api",
tool_input="query",
log="tool:search_api:a different query",
),
"result for query",
]
],
}
},
{
"agent": {
"agent_outcome": AgentAction(
tool="search_api",
tool_input="another",
log="tool:search_api:another",
),
}
},
{"__interrupt__": ()},
]
app_w_interrupt.update_state(
config,
{
"agent_outcome": AgentFinish(
return_values={"answer": "a really nice answer"},
log="finish:a really nice answer",
)
},
)
assert app_w_interrupt.get_state(config) == StateSnapshot(
values={
"agent_outcome": AgentFinish(
return_values={"answer": "a really nice answer"},
log="finish:a really nice answer",
),
"intermediate_steps": [
[
AgentAction(
tool="search_api",
tool_input="query",
log="tool:search_api:a different query",
),
"result for query",
]
],
},
tasks=(),
next=(),
config={
"configurable": {
"thread_id": "1",
"checkpoint_ns": "",
"checkpoint_id": AnyStr(),
}
},
created_at=AnyStr(),
metadata={
"parents": {},
"source": "update",
"step": 5,
},
parent_config=(
list(app_w_interrupt.checkpointer.list(config, limit=2))[-1].config
),
interrupts=(),
)
# test state get/update methods with interrupt_before
app_w_interrupt = workflow.compile(
checkpointer=sync_checkpointer,
interrupt_before=["tools"],
debug=True,
)
config = {"configurable": {"thread_id": "2"}}
llm.i = 0 # reset the llm
assert [
c
for c in app_w_interrupt.stream(
{"input": "what is weather in sf"}, config, durability="exit"
)
] == [
{
"agent": {
"agent_outcome": AgentAction(
tool="search_api", tool_input="query", log="tool:search_api:query"
),
}
},
{"__interrupt__": ()},
]
assert app_w_interrupt.get_state(config) == StateSnapshot(
values={
"agent_outcome": AgentAction(
tool="search_api", tool_input="query", log="tool:search_api:query"
),
"intermediate_steps": [],
},
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,
},
parent_config=None,
interrupts=(),
)
app_w_interrupt.update_state(
config,
{
"agent_outcome": AgentAction(
tool="search_api",
tool_input="query",
log="tool:search_api:a different query",
)
},
)
assert app_w_interrupt.get_state(config) == StateSnapshot(
values={
"agent_outcome": AgentAction(
tool="search_api",
tool_input="query",
log="tool:search_api:a different query",
),
"intermediate_steps": [],
},
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,
},
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": {
"intermediate_steps": [
[
AgentAction(
tool="search_api",
tool_input="query",
log="tool:search_api:a different query",
),
"result for query",
]
],
}
},
{
"agent": {
"agent_outcome": AgentAction(
tool="search_api",
tool_input="another",
log="tool:search_api:another",
),
}
},
{"__interrupt__": ()},
]
app_w_interrupt.update_state(
config,
{
"agent_outcome": AgentFinish(
return_values={"answer": "a really nice answer"},
log="finish:a really nice answer",
)
},
)
assert app_w_interrupt.get_state(config) == StateSnapshot(
values={
"agent_outcome": AgentFinish(
return_values={"answer": "a really nice answer"},
log="finish:a really nice answer",
),
"intermediate_steps": [
[
AgentAction(
tool="search_api",
tool_input="query",
log="tool:search_api:a different query",
),
"result for query",
]
],
},
tasks=(),
next=(),
config={
"configurable": {
"thread_id": "2",
"checkpoint_ns": "",
"checkpoint_id": AnyStr(),
}
},
created_at=AnyStr(),
metadata={
"parents": {},
"source": "update",
"step": 5,
},
parent_config=(
list(app_w_interrupt.checkpointer.list(config, limit=2))[-1].config
),
interrupts=(),
)
# test w interrupt before all
app_w_interrupt = workflow.compile(
checkpointer=sync_checkpointer,
interrupt_before="*",
debug=True,
)
config = {"configurable": {"thread_id": "3"}}
llm.i = 0 # reset the llm
assert [
c
for c in app_w_interrupt.stream(
{"input": "what is weather in sf"}, config, durability="exit"
)
] == [
{"__interrupt__": ()},
]
assert app_w_interrupt.get_state(config) == StateSnapshot(
values={
"intermediate_steps": [],
},
tasks=(PregelTask(AnyStr(), "agent", (PULL, "agent")),),
next=("agent",),
config={
"configurable": {
"thread_id": "3",
"checkpoint_ns": "",
"checkpoint_id": AnyStr(),
}
},
created_at=AnyStr(),
metadata={
"parents": {},
"source": "loop",
"step": 0,
},
parent_config=None,
interrupts=(),
)
assert [c for c in app_w_interrupt.stream(None, config)] == [
{
"agent": {
"agent_outcome": AgentAction(
tool="search_api", tool_input="query", log="tool:search_api:query"
),
}
},
{"__interrupt__": ()},
]
assert app_w_interrupt.get_state(config) == StateSnapshot(
values={
"agent_outcome": AgentAction(
tool="search_api", tool_input="query", log="tool:search_api:query"
),
"intermediate_steps": [],
},
tasks=(PregelTask(AnyStr(), "tools", (PULL, "tools")),),
next=("tools",),
config={
"configurable": {
"thread_id": "3",
"checkpoint_ns": "",
"checkpoint_id": AnyStr(),
}
},
created_at=AnyStr(),
metadata={
"parents": {},
"source": "loop",
"step": 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": {
"intermediate_steps": [
[
AgentAction(
tool="search_api",
tool_input="query",
log="tool:search_api:query",
),
"result for query",
]
],
}
},
{"__interrupt__": ()},
]
assert app_w_interrupt.get_state(config) == StateSnapshot(
values={
"agent_outcome": AgentAction(
tool="search_api", tool_input="query", log="tool:search_api:query"
),
"intermediate_steps": [
[
AgentAction(
tool="search_api",
tool_input="query",
log="tool:search_api:query",
),
"result for query",
]
],
},
tasks=(PregelTask(AnyStr(), "agent", (PULL, "agent")),),
next=("agent",),
config={
"configurable": {
"thread_id": "3",
"checkpoint_ns": "",
"checkpoint_id": AnyStr(),
}
},
created_at=AnyStr(),
metadata={
"parents": {},
"source": "loop",
"step": 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)] == [
{
"agent": {
"agent_outcome": AgentAction(
tool="search_api",
tool_input="another",
log="tool:search_api:another",
),
}
},
{"__interrupt__": ()},
]
# test w interrupt after all
app_w_interrupt = workflow.compile(
checkpointer=sync_checkpointer,
interrupt_after="*",
)
config = {"configurable": {"thread_id": "4"}}
llm.i = 0 # reset the llm
assert [
c
for c in app_w_interrupt.stream(
{"input": "what is weather in sf"}, config, durability="exit"
)
] == [
{
"agent": {
"agent_outcome": AgentAction(
tool="search_api", tool_input="query", log="tool:search_api:query"
),
}
},
{"__interrupt__": ()},
]
assert app_w_interrupt.get_state(config) == StateSnapshot(
values={
"agent_outcome": AgentAction(
tool="search_api", tool_input="query", log="tool:search_api:query"
),
"intermediate_steps": [],
},
tasks=(PregelTask(AnyStr(), "tools", (PULL, "tools")),),
next=("tools",),
config={
"configurable": {
"thread_id": "4",
"checkpoint_ns": "",
"checkpoint_id": AnyStr(),
}
},
created_at=AnyStr(),
metadata={
"parents": {},
"source": "loop",
"step": 1,
},
parent_config=None,
interrupts=(),
)
assert [c for c in app_w_interrupt.stream(None, config)] == [
{
"tools": {
"intermediate_steps": [
[
AgentAction(
tool="search_api",
tool_input="query",
log="tool:search_api:query",
),
"result for query",
]
],
}
},
{"__interrupt__": ()},
]
assert app_w_interrupt.get_state(config) == StateSnapshot(
values={
"agent_outcome": AgentAction(
tool="search_api", tool_input="query", log="tool:search_api:query"
),
"intermediate_steps": [
[
AgentAction(
tool="search_api",
tool_input="query",
log="tool:search_api:query",
),
"result for query",
]
],
},
tasks=(PregelTask(AnyStr(), "agent", (PULL, "agent")),),
next=("agent",),
config={
"configurable": {
"thread_id": "4",
"checkpoint_ns": "",
"checkpoint_id": AnyStr(),
}
},
created_at=AnyStr(),
metadata={
"parents": {},
"source": "loop",
"step": 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)] == [
{
"agent": {
"agent_outcome": AgentAction(
tool="search_api",
tool_input="another",
log="tool:search_api:another",
),
}
},
{"__interrupt__": ()},
]
def test_prebuilt_tool_chat(snapshot: SnapshotAssertion) -> None:
from langchain_core.messages import AIMessage, HumanMessage
from langchain_core.tools import tool
@tool()
def search_api(query: str) -> str:
"""Searches the API for the query."""
return f"result for {query}"
tools = [search_api]
model = FakeChatModel(
messages=[
AIMessage(
content="",
tool_calls=[
{
"id": "tool_call123",
"name": "search_api",
"args": {"query": "query"},
},
],
),
AIMessage(
content="",
tool_calls=[
{
"id": "tool_call234",
"name": "search_api",
"args": {"query": "another"},
},
{
"id": "tool_call567",
"name": "search_api",
"args": {"query": "a third one"},
},
],
),
AIMessage(content="answer"),
]
)
app = create_react_agent(model, tools)
assert json.dumps(app.get_input_jsonschema()) == snapshot
assert json.dumps(app.get_output_jsonschema()) == 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(
{"messages": [HumanMessage(content="what is weather in sf")]}
) == {
"messages": [
_AnyIdHumanMessage(content="what is weather in sf"),
_AnyIdAIMessage(
content="",
tool_calls=[
{
"id": "tool_call123",
"name": "search_api",
"args": {"query": "query"},
},
],
),
_AnyIdToolMessage(
content="result for query",
name="search_api",
tool_call_id="tool_call123",
),
_AnyIdAIMessage(
content="",
tool_calls=[
{
"id": "tool_call234",
"name": "search_api",
"args": {"query": "another"},
},
{
"id": "tool_call567",
"name": "search_api",
"args": {"query": "a third one"},
},
],
),
_AnyIdToolMessage(
content="result for another",
name="search_api",
tool_call_id="tool_call234",
),
_AnyIdToolMessage(
content="result for a third one",
name="search_api",
tool_call_id="tool_call567",
id=AnyStr(),
),
_AnyIdAIMessage(content="answer"),
]
}
events = [
c
for c in app.stream(
{"messages": [HumanMessage(content="what is weather in sf")]},
stream_mode="messages",
)
]
assert events[:3] == [
(
_AnyIdAIMessageChunk(
content="",
tool_calls=[
{
"name": "search_api",
"args": {"query": "query"},
"id": "tool_call123",
"type": "tool_call",
}
],
tool_call_chunks=[
{
"name": "search_api",
"args": '{"query": "query"}',
"id": "tool_call123",
"index": None,
"type": "tool_call_chunk",
}
],
chunk_position="last",
),
{
"langgraph_step": 1,
"langgraph_node": "agent",
"langgraph_triggers": ("branch:to:agent",),
"langgraph_path": (PULL, "agent"),
"langgraph_checkpoint_ns": AnyStr("agent:"),
"checkpoint_ns": AnyStr("agent:"),
"ls_provider": "fakechatmodel",
"ls_model_type": "chat",
"ls_integration": "langchain_chat_model",
"lc_versions": {"langchain-core": LANGCHAIN_CORE_VERSION},
},
),
(
_AnyIdToolMessage(
content="result for query",
name="search_api",
tool_call_id="tool_call123",
),
{
"ls_integration": "langgraph",
"langgraph_step": 2,
"langgraph_node": "tools",
"langgraph_triggers": (PUSH,),
"langgraph_path": (PUSH, AnyInt(), False),
"langgraph_checkpoint_ns": AnyStr("tools:"),
},
),
(
_AnyIdAIMessageChunk(
content="",
tool_calls=[
{
"name": "search_api",
"args": {"query": "another"},
"id": "tool_call234",
"type": "tool_call",
},
{
"name": "search_api",
"args": {"query": "a third one"},
"id": "tool_call567",
"type": "tool_call",
},
],
tool_call_chunks=[
{
"name": "search_api",
"args": '{"query": "another"}',
"id": "tool_call234",
"index": None,
"type": "tool_call_chunk",
},
{
"name": "search_api",
"args": '{"query": "a third one"}',
"id": "tool_call567",
"index": None,
"type": "tool_call_chunk",
},
],
chunk_position="last",
),
{
"langgraph_step": 3,
"langgraph_node": "agent",
"langgraph_triggers": ("branch:to:agent",),
"langgraph_path": (PULL, "agent"),
"langgraph_checkpoint_ns": AnyStr("agent:"),
"checkpoint_ns": AnyStr("agent:"),
"ls_provider": "fakechatmodel",
"ls_model_type": "chat",
"ls_integration": "langchain_chat_model",
"lc_versions": {"langchain-core": LANGCHAIN_CORE_VERSION},
},
),
]
assert events[3:5] == UnsortedSequence(
(
_AnyIdToolMessage(
content="result for another",
name="search_api",
tool_call_id="tool_call234",
),
{
"ls_integration": "langgraph",
"langgraph_step": 4,
"langgraph_node": "tools",
"langgraph_triggers": (PUSH,),
"langgraph_path": (PUSH, AnyInt(), False),
"langgraph_checkpoint_ns": AnyStr("tools:"),
},
),
(
_AnyIdToolMessage(
content="result for a third one",
name="search_api",
tool_call_id="tool_call567",
),
{
"ls_integration": "langgraph",
"langgraph_step": 4,
"langgraph_node": "tools",
"langgraph_triggers": (PUSH,),
"langgraph_path": (PUSH, AnyInt(), False),
"langgraph_checkpoint_ns": AnyStr("tools:"),
},
),
)
assert events[5:] == [
(
_AnyIdAIMessageChunk(
content="answer",
chunk_position="last",
),
{
"langgraph_step": 5,
"langgraph_node": "agent",
"langgraph_triggers": ("branch:to:agent",),
"langgraph_path": (PULL, "agent"),
"langgraph_checkpoint_ns": AnyStr("agent:"),
"checkpoint_ns": AnyStr("agent:"),
"ls_provider": "fakechatmodel",
"ls_model_type": "chat",
"ls_integration": "langchain_chat_model",
"lc_versions": {"langchain-core": LANGCHAIN_CORE_VERSION},
},
),
]
assert app.invoke(
{"messages": [HumanMessage(content="what is weather in sf")]},
{"recursion_limit": 2},
debug=True,
) == {
"messages": [
_AnyIdHumanMessage(content="what is weather in sf"),
_AnyIdAIMessage(content="Sorry, need more steps to process this request."),
]
}
model.i = 0 # reset the model
invoke_updates_events = app.invoke(
{"messages": [HumanMessage(content="what is weather in sf")]},
stream_mode="updates",
)
stream_updates_events = [
*app.stream({"messages": [HumanMessage(content="what is weather in sf")]})
]
for output in (invoke_updates_events, stream_updates_events):
assert output[:3] == [
{
"agent": {
"messages": [
_AnyIdAIMessage(
content="",
tool_calls=[
{
"id": "tool_call123",
"name": "search_api",
"args": {"query": "query"},
},
],
)
]
}
},
{
"tools": {
"messages": [
_AnyIdToolMessage(
content="result for query",
name="search_api",
tool_call_id="tool_call123",
)
]
}
},
{
"agent": {
"messages": [
_AnyIdAIMessage(
content="",
tool_calls=[
{
"id": "tool_call234",
"name": "search_api",
"args": {"query": "another"},
},
{
"id": "tool_call567",
"name": "search_api",
"args": {"query": "a third one"},
},
],
)
]
}
},
]
assert output[3:5] == UnsortedSequence(
{
"tools": {
"messages": [
_AnyIdToolMessage(
content="result for another",
name="search_api",
tool_call_id="tool_call234",
),
]
}
},
{
"tools": {
"messages": [
_AnyIdToolMessage(
content="result for a third one",
name="search_api",
tool_call_id="tool_call567",
),
]
}
},
)
assert output[5:] == [
{"agent": {"messages": [_AnyIdAIMessage(content="answer")]}}
]
def test_state_graph_packets(
sync_checkpointer: BaseCheckpointSaver, mocker: MockerFixture
) -> None:
from langchain_core.language_models.fake_chat_models import (
FakeMessagesListChatModel,
)
from langchain_core.messages import (
AIMessage,
BaseMessage,
HumanMessage,
ToolCall,
ToolMessage,
)
from langchain_core.tools import tool
class AgentState(TypedDict):
messages: Annotated[list[BaseMessage], add_messages]
@tool()
def search_api(query: str) -> str:
"""Searches the API for the query."""
time.sleep(0.1)
return f"result for {query}"
tools = [search_api]
tools_by_name = {t.name: t for t in tools}
model = FakeMessagesListChatModel(
responses=[
AIMessage(
id="ai1",
content="",
tool_calls=[
{
"id": "tool_call123",
"name": "search_api",
"args": {"query": "query"},
},
],
),
AIMessage(
id="ai2",
content="",
tool_calls=[
{
"id": "tool_call234",
"name": "search_api",
"args": {"query": "another", "idx": 0},
},
{
"id": "tool_call567",
"name": "search_api",
"args": {"query": "a third one", "idx": 1},
},
],
),
AIMessage(id="ai3", content="answer"),
]
)
def agent(data: AgentState) -> AgentState:
return {
"messages": model.invoke(data["messages"]),
"something_extra": "hi there",
}
# Define decision-making logic
def should_continue(data: dict) -> str:
assert data["something_extra"] == "hi there", (
"nodes can pass extra data to their cond edges, which isn't saved in state"
)
# Logic to decide whether to continue in the loop or exit
if tool_calls := data["messages"][-1].tool_calls:
return [Send("tools", tool_call) for tool_call in tool_calls]
else:
return END
def tools_node(input: ToolCall, config: RunnableConfig) -> AgentState:
time.sleep(input["args"].get("idx", 0) / 10)
output = tools_by_name[input["name"]].invoke(input["args"], config)
return {
"messages": ToolMessage(
content=output, name=input["name"], tool_call_id=input["id"]
)
}
# Define a new graph
workflow = StateGraph(AgentState)
# Define the two nodes we will cycle between
workflow.add_node("agent", agent)
workflow.add_node("tools", tools_node)
# 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("agent", should_continue)
# 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({"messages": HumanMessage(content="what is weather in sf")}) == {
"messages": [
_AnyIdHumanMessage(content="what is weather in sf"),
AIMessage(
id="ai1",
content="",
tool_calls=[
{
"id": "tool_call123",
"name": "search_api",
"args": {"query": "query"},
},
],
),
_AnyIdToolMessage(
content="result for query",
name="search_api",
tool_call_id="tool_call123",
),
AIMessage(
id="ai2",
content="",
tool_calls=[
{
"id": "tool_call234",
"name": "search_api",
"args": {"query": "another", "idx": 0},
},
{
"id": "tool_call567",
"name": "search_api",
"args": {"query": "a third one", "idx": 1},
},
],
),
_AnyIdToolMessage(
content="result for another",
name="search_api",
tool_call_id="tool_call234",
),
_AnyIdToolMessage(
content="result for a third one",
name="search_api",
tool_call_id="tool_call567",
),
AIMessage(content="answer", id="ai3"),
]
}
assert [
c
for c in app.stream(
{"messages": [HumanMessage(content="what is weather in sf")]}
)
] == [
{
"agent": {
"messages": AIMessage(
id="ai1",
content="",
tool_calls=[
{
"id": "tool_call123",
"name": "search_api",
"args": {"query": "query"},
},
],
)
},
},
{
"tools": {
"messages": _AnyIdToolMessage(
content="result for query",
name="search_api",
tool_call_id="tool_call123",
)
}
},
{
"agent": {
"messages": AIMessage(
id="ai2",
content="",
tool_calls=[
{
"id": "tool_call234",
"name": "search_api",
"args": {"query": "another", "idx": 0},
},
{
"id": "tool_call567",
"name": "search_api",
"args": {"query": "a third one", "idx": 1},
},
],
)
}
},
{
"tools": {
"messages": _AnyIdToolMessage(
content="result for another",
name="search_api",
tool_call_id="tool_call234",
)
},
},
{
"tools": {
"messages": _AnyIdToolMessage(
content="result for a third one",
name="search_api",
tool_call_id="tool_call567",
),
},
},
{"agent": {"messages": AIMessage(content="answer", id="ai3")}},
]
# interrupt after agent
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(
{"messages": HumanMessage(content="what is weather in sf")},
config,
durability="exit",
)
] == [
{
"agent": {
"messages": AIMessage(
id="ai1",
content="",
tool_calls=[
{
"id": "tool_call123",
"name": "search_api",
"args": {"query": "query"},
},
],
)
}
},
{"__interrupt__": ()},
]
assert app_w_interrupt.get_state(config) == StateSnapshot(
values={
"messages": [
_AnyIdHumanMessage(content="what is weather in sf"),
AIMessage(
id="ai1",
content="",
tool_calls=[
{
"id": "tool_call123",
"name": "search_api",
"args": {"query": "query"},
},
],
),
]
},
tasks=(PregelTask(AnyStr(), "tools", (PUSH, 0, False)),),
next=("tools",),
config={
"configurable": {
"thread_id": "1",
"checkpoint_ns": "",
"checkpoint_id": AnyStr(),
}
},
created_at=AnyStr(),
metadata={
"parents": {},
"source": "loop",
"step": 1,
},
parent_config=None,
interrupts=(),
)
# modify ai message
last_message = (app_w_interrupt.get_state(config)).values["messages"][-1]
last_message.tool_calls[0]["args"]["query"] = "a different query"
app_w_interrupt.update_state(
config, {"messages": last_message, "something_extra": "hi there"}
)
# message was replaced instead of appended
assert app_w_interrupt.get_state(config) == StateSnapshot(
values={
"messages": [
_AnyIdHumanMessage(content="what is weather in sf"),
AIMessage(
id="ai1",
content="",
tool_calls=[
{
"id": "tool_call123",
"name": "search_api",
"args": {"query": "a different query"},
},
],
),
]
},
tasks=(PregelTask(AnyStr(), "tools", (PUSH, 0, False)),),
next=("tools",),
config={
"configurable": {
"thread_id": "1",
"checkpoint_ns": "",
"checkpoint_id": AnyStr(),
}
},
created_at=AnyStr(),
metadata={
"parents": {},
"source": "update",
"step": 2,
},
parent_config=(
[*app_w_interrupt.checkpointer.list(config, limit=2)][-1].config
),
interrupts=(),
)
assert [c for c in app_w_interrupt.stream(None, config)] == [
{
"tools": {
"messages": _AnyIdToolMessage(
content="result for a different query",
name="search_api",
tool_call_id="tool_call123",
)
}
},
{
"agent": {
"messages": AIMessage(
id="ai2",
content="",
tool_calls=[
{
"id": "tool_call234",
"name": "search_api",
"args": {"query": "another", "idx": 0},
},
{
"id": "tool_call567",
"name": "search_api",
"args": {"query": "a third one", "idx": 1},
},
],
)
},
},
{"__interrupt__": ()},
]
assert app_w_interrupt.get_state(config) == StateSnapshot(
values={
"messages": [
_AnyIdHumanMessage(content="what is weather in sf"),
AIMessage(
id="ai1",
content="",
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(
id="ai2",
content="",
tool_calls=[
{
"id": "tool_call234",
"name": "search_api",
"args": {"query": "another", "idx": 0},
},
{
"id": "tool_call567",
"name": "search_api",
"args": {"query": "a third one", "idx": 1},
},
],
),
]
},
tasks=(
PregelTask(AnyStr(), "tools", (PUSH, 0, False)),
PregelTask(AnyStr(), "tools", (PUSH, 1, False)),
),
next=("tools", "tools"),
config={
"configurable": {
"thread_id": "1",
"checkpoint_ns": "",
"checkpoint_id": AnyStr(),
}
},
created_at=AnyStr(),
metadata={
"parents": {},
"source": "loop",
"step": 4,
},
parent_config=(
[*app_w_interrupt.checkpointer.list(config, limit=2)][-1].config
),
interrupts=(),
)
app_w_interrupt.update_state(
config,
{
"messages": AIMessage(content="answer", id="ai2"),
"something_extra": "hi there",
},
)
# replaces message even if object identity is different, as long as id is the same
assert app_w_interrupt.get_state(config) == StateSnapshot(
values={
"messages": [
_AnyIdHumanMessage(content="what is weather in sf"),
AIMessage(
id="ai1",
content="",
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,
},
parent_config=(
[*app_w_interrupt.checkpointer.list(config, limit=2)][-1].config
),
interrupts=(),
)
# interrupt before tools
app_w_interrupt = workflow.compile(
checkpointer=sync_checkpointer,
interrupt_before=["tools"],
)
config = {"configurable": {"thread_id": "2"}}
model.i = 0
assert [
c
for c in app_w_interrupt.stream(
{"messages": HumanMessage(content="what is weather in sf")},
config,
durability="exit",
)
] == [
{
"agent": {
"messages": AIMessage(
id="ai1",
content="",
tool_calls=[
{
"id": "tool_call123",
"name": "search_api",
"args": {"query": "query"},
},
],
)
}
},
{"__interrupt__": ()},
]
assert app_w_interrupt.get_state(config) == StateSnapshot(
values={
"messages": [
_AnyIdHumanMessage(content="what is weather in sf"),
AIMessage(
id="ai1",
content="",
tool_calls=[
{
"id": "tool_call123",
"name": "search_api",
"args": {"query": "query"},
},
],
),
]
},
tasks=(PregelTask(AnyStr(), "tools", (PUSH, 0, False)),),
next=("tools",),
config={
"configurable": {
"thread_id": "2",
"checkpoint_ns": "",
"checkpoint_id": AnyStr(),
}
},
created_at=AnyStr(),
metadata={
"parents": {},
"source": "loop",
"step": 1,
},
parent_config=None,
interrupts=(),
)
# modify ai message
last_message = (app_w_interrupt.get_state(config)).values["messages"][-1]
last_message.tool_calls[0]["args"]["query"] = "a different query"
app_w_interrupt.update_state(
config, {"messages": last_message, "something_extra": "hi there"}
)
# message was replaced instead of appended
assert app_w_interrupt.get_state(config) == StateSnapshot(
values={
"messages": [
_AnyIdHumanMessage(content="what is weather in sf"),
AIMessage(
id="ai1",
content="",
tool_calls=[
{
"id": "tool_call123",
"name": "search_api",
"args": {"query": "a different query"},
},
],
),
]
},
tasks=(PregelTask(AnyStr(), "tools", (PUSH, 0, False)),),
next=("tools",),
config=app_w_interrupt.checkpointer.get_tuple(config).config,
created_at=AnyStr(),
metadata={
"parents": {},
"source": "update",
"step": 2,
},
parent_config=(
[*app_w_interrupt.checkpointer.list(config, limit=2)][-1].config
),
interrupts=(),
)
assert [c for c in app_w_interrupt.stream(None, config)] == [
{
"tools": {
"messages": _AnyIdToolMessage(
content="result for a different query",
name="search_api",
tool_call_id="tool_call123",
)
}
},
{
"agent": {
"messages": AIMessage(
id="ai2",
content="",
tool_calls=[
{
"id": "tool_call234",
"name": "search_api",
"args": {"query": "another", "idx": 0},
},
{
"id": "tool_call567",
"name": "search_api",
"args": {"query": "a third one", "idx": 1},
},
],
)
},
},
{"__interrupt__": ()},
]
assert app_w_interrupt.get_state(config) == StateSnapshot(
values={
"messages": [
_AnyIdHumanMessage(content="what is weather in sf"),
AIMessage(
id="ai1",
content="",
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(
id="ai2",
content="",
tool_calls=[
{
"id": "tool_call234",
"name": "search_api",
"args": {"query": "another", "idx": 0},
},
{
"id": "tool_call567",
"name": "search_api",
"args": {"query": "a third one", "idx": 1},
},
],
),
]
},
tasks=(
PregelTask(AnyStr(), "tools", (PUSH, 0, False)),
PregelTask(AnyStr(), "tools", (PUSH, 1, False)),
),
next=("tools", "tools"),
config={
"configurable": {
"thread_id": "2",
"checkpoint_ns": "",
"checkpoint_id": AnyStr(),
}
},
created_at=AnyStr(),
metadata={
"parents": {},
"source": "loop",
"step": 4,
},
parent_config=(
[*app_w_interrupt.checkpointer.list(config, limit=2)][-1].config
),
interrupts=(),
)
app_w_interrupt.update_state(
config,
{
"messages": AIMessage(content="answer", id="ai2"),
"something_extra": "hi there",
},
)
# replaces message even if object identity is different, as long as id is the same
assert app_w_interrupt.get_state(config) == StateSnapshot(
values={
"messages": [
_AnyIdHumanMessage(content="what is weather in sf"),
AIMessage(
id="ai1",
content="",
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,
},
parent_config=(
[*app_w_interrupt.checkpointer.list(config, limit=2)][-1].config
),
interrupts=(),
)
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 FakeFunctionChatModel(FakeMessagesListChatModel):
def bind_functions(self, functions: list):
return self
def _generate(
self,
messages: list[BaseMessage],
stop: list[str] | None = None,
run_manager: CallbackManagerForLLMRun | None = 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 = FakeFunctionChatModel(
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 = StateGraph(state_schema=Annotated[list[AnyMessage], add_messages]) # type: ignore[arg-type]
# 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_jsonschema()) == snapshot
assert json.dumps(app.get_output_jsonschema()) == 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, durability="exit"
)
] == [
{
"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,
},
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,
},
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,
},
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,
},
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, durability="exit"
)
] == [
{
"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,
},
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,
},
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,
},
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,
},
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 BaseMessage 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,
},
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 FakeFunctionChatModel(FakeMessagesListChatModel):
def bind_functions(self, functions: list):
return self
def _generate(
self,
messages: list[BaseMessage],
stop: list[str] | None = None,
run_manager: CallbackManagerForLLMRun | None = 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 = FakeFunctionChatModel(
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-000000000004",
)
]
},
{
"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-000000000005",
)
]
},
{"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, durability="exit"
)
] == [
{
"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,
},
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,
},
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,
},
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,
},
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, durability="exit"
)
] == [
{
"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,
},
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,
},
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,
},
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,
},
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 BaseMessage 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,
},
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,
},
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-000000000008",
),
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-000000000010"
),
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))
class State(TypedDict, total=False):
query: str
answer: str
docs: Annotated[list[str], sorted_add]
def rewrite_query(data: State) -> State:
return {"query": f"query: {data['query']}"}
def retriever_one(data: State) -> State:
# timer ensures stream output order is stable
# also, it confirms that the update order is not dependent on finishing order
# instead being defined by the order of the nodes/edges in the graph definition
# ie. stable between invocations
time.sleep(0.1)
return {"docs": ["doc1", "doc2"]}
def retriever_two(data: State) -> State:
return {"docs": ["doc3", "doc4"]}
def qa(data: State) -> State:
return {"answer": ",".join(data["docs"])}
workflow = StateGraph(State)
workflow.add_node("rewrite_query", rewrite_query)
workflow.add_node("retriever_one", retriever_one)
workflow.add_node("retriever_two", retriever_two)
workflow.add_node("qa", qa)
workflow.set_entry_point("rewrite_query")
workflow.add_edge("rewrite_query", "retriever_one")
workflow.add_edge("rewrite_query", "retriever_two")
workflow.add_edge("retriever_one", "qa")
workflow.add_edge("retriever_two", "qa")
workflow.set_finish_point("qa")
app = workflow.compile()
assert app.invoke({"query": "what is weather in sf"}) == {
"query": "query: what is weather in sf",
"docs": ["doc1", "doc2", "doc3", "doc4"],
"answer": "doc1,doc2,doc3,doc4",
}
assert [*app.stream({"query": "what is weather in sf"})] == [
{"rewrite_query": {"query": "query: what is weather in sf"}},
{"retriever_two": {"docs": ["doc3", "doc4"]}},
{"retriever_one": {"docs": ["doc1", "doc2"]}},
{"qa": {"answer": "doc1,doc2,doc3,doc4"}},
]
assert [*app.stream({"query": "what is weather in sf"}, stream_mode="values")] == [
{"query": "what is weather in sf", "docs": []},
{"query": "query: what is weather in sf", "docs": []},
{
"query": "query: what is weather in sf",
"docs": ["doc1", "doc2", "doc3", "doc4"],
},
{
"query": "query: what is weather in sf",
"docs": ["doc1", "doc2", "doc3", "doc4"],
"answer": "doc1,doc2,doc3,doc4",
},
]
assert [
*app.stream(
{"query": "what is weather in sf"},
stream_mode=["values", "updates", "debug"],
)
] == [
("values", {"query": "what is weather in sf", "docs": []}),
(
"debug",
{
"type": "task",
"timestamp": AnyStr(),
"step": 1,
"payload": {
"id": AnyStr(),
"name": "rewrite_query",
"input": {"query": "what is weather in sf", "docs": []},
"triggers": ("branch:to:rewrite_query",),
},
},
),
("updates", {"rewrite_query": {"query": "query: what is weather in sf"}}),
(
"debug",
{
"type": "task_result",
"timestamp": AnyStr(),
"step": 1,
"payload": {
"id": AnyStr(),
"name": "rewrite_query",
"result": {
"query": "query: what is weather in sf",
},
"error": None,
"interrupts": [],
},
},
),
("values", {"query": "query: what is weather in sf", "docs": []}),
(
"debug",
{
"type": "task",
"timestamp": AnyStr(),
"step": 2,
"payload": {
"id": AnyStr(),
"name": "retriever_one",
"input": {"query": "query: what is weather in sf", "docs": []},
"triggers": ("branch:to:retriever_one",),
},
},
),
(
"debug",
{
"type": "task",
"timestamp": AnyStr(),
"step": 2,
"payload": {
"id": AnyStr(),
"name": "retriever_two",
"input": {"query": "query: what is weather in sf", "docs": []},
"triggers": ("branch:to:retriever_two",),
},
},
),
(
"updates",
{"retriever_two": {"docs": ["doc3", "doc4"]}},
),
(
"debug",
{
"type": "task_result",
"timestamp": AnyStr(),
"step": 2,
"payload": {
"id": AnyStr(),
"name": "retriever_two",
"result": {
"docs": ["doc3", "doc4"],
},
"error": None,
"interrupts": [],
},
},
),
(
"updates",
{"retriever_one": {"docs": ["doc1", "doc2"]}},
),
(
"debug",
{
"type": "task_result",
"timestamp": AnyStr(),
"step": 2,
"payload": {
"id": AnyStr(),
"name": "retriever_one",
"result": {
"docs": ["doc1", "doc2"],
},
"error": None,
"interrupts": [],
},
},
),
(
"values",
{
"query": "query: what is weather in sf",
"docs": ["doc1", "doc2", "doc3", "doc4"],
},
),
(
"debug",
{
"type": "task",
"timestamp": AnyStr(),
"step": 3,
"payload": {
"id": AnyStr(),
"name": "qa",
"input": {
"query": "query: what is weather in sf",
"docs": ["doc1", "doc2", "doc3", "doc4"],
},
"triggers": ("branch:to:qa",),
},
},
),
("updates", {"qa": {"answer": "doc1,doc2,doc3,doc4"}}),
(
"debug",
{
"type": "task_result",
"timestamp": AnyStr(),
"step": 3,
"payload": {
"id": AnyStr(),
"name": "qa",
"result": {
"answer": "doc1,doc2,doc3,doc4",
},
"error": None,
"interrupts": [],
},
},
),
(
"values",
{
"query": "query: what is weather in sf",
"answer": "doc1,doc2,doc3,doc4",
"docs": ["doc1", "doc2", "doc3", "doc4"],
},
),
]
def test_dynamic_interrupt(sync_checkpointer: BaseCheckpointSaver) -> None:
class State(TypedDict):
my_key: Annotated[str, operator.add]
market: str
tool_two_node_count = 0
def tool_two_node(s: State) -> State:
nonlocal tool_two_node_count
tool_two_node_count += 1
if s["market"] == "DE":
answer = interrupt("Just because...")
else:
answer = " all good"
return {"my_key": answer}
tool_two_graph = StateGraph(State)
tool_two_graph.add_node("tool_two", tool_two_node, retry_policy=RetryPolicy())
tool_two_graph.add_edge(START, "tool_two")
tool_two = tool_two_graph.compile()
tracer = FakeTracer()
assert tool_two.invoke(
{"my_key": "value", "market": "DE"}, {"callbacks": [tracer]}
) == {
"my_key": "value",
"market": "DE",
"__interrupt__": [Interrupt(value="Just because...", id=AnyStr())],
}
assert tool_two_node_count == 1, "interrupts aren't retried"
assert len(tracer.runs) == 1
run = tracer.runs[0]
assert run.end_time is not None
assert run.error is None
assert run.outputs == {"market": "DE", "my_key": "value"}
assert tool_two.invoke({"my_key": "value", "market": "US"}) == {
"my_key": "value all good",
"market": "US",
}
tool_two = tool_two_graph.compile(checkpointer=sync_checkpointer)
# missing thread_id
with pytest.raises(ValueError, match="thread_id"):
tool_two.invoke({"my_key": "value", "market": "DE"})
# flow: interrupt -> resume with answer
thread2 = {"configurable": {"thread_id": "2"}}
# stop when about to enter node
assert [
c for c in tool_two.stream({"my_key": "value ⛰️", "market": "DE"}, thread2)
] == [
{
"__interrupt__": (
Interrupt(
value="Just because...",
id=AnyStr(),
),
)
},
]
# resume with answer
assert [c for c in tool_two.stream(Command(resume=" my answer"), thread2)] == [
{"tool_two": {"my_key": " my answer"}},
]
# flow: interrupt -> clear tasks
thread1 = {"configurable": {"thread_id": "1"}}
# stop when about to enter node
assert tool_two.invoke(
{"my_key": "value ⛰️", "market": "DE"}, thread1, durability="exit"
) == {
"my_key": "value ⛰️",
"market": "DE",
"__interrupt__": [Interrupt(value="Just because...", id=AnyStr())],
}
assert [c.metadata for c in tool_two.checkpointer.list(thread1)] == [
{
"parents": {},
"source": "loop",
"step": 0,
},
]
assert tool_two.get_state(thread1) == StateSnapshot(
values={"my_key": "value ⛰️", "market": "DE"},
next=("tool_two",),
tasks=(
PregelTask(
AnyStr(),
"tool_two",
(PULL, "tool_two"),
interrupts=(
Interrupt(
value="Just because...",
id=AnyStr(),
),
),
),
),
config={
"configurable": {
"thread_id": "1",
"checkpoint_ns": "",
"checkpoint_id": AnyStr(),
}
},
created_at=AnyStr(),
metadata={
"parents": {},
"source": "loop",
"step": 0,
},
parent_config=None,
interrupts=(
Interrupt(
value="Just because...",
id=AnyStr(),
),
),
)
# clear the interrupt and next tasks
tool_two.update_state(thread1, None, as_node=END)
# interrupt and next tasks are cleared
assert tool_two.get_state(thread1) == StateSnapshot(
values={"my_key": "value ⛰️", "market": "DE"},
next=(),
tasks=(),
config={
"configurable": {
"thread_id": "1",
"checkpoint_ns": "",
"checkpoint_id": AnyStr(),
}
},
created_at=AnyStr(),
metadata={
"parents": {},
"source": "update",
"step": 1,
},
parent_config=(list(tool_two.checkpointer.list(thread1, limit=2))[-1].config),
interrupts=(),
)
def test_partial_pending_checkpoint(sync_checkpointer: BaseCheckpointSaver) -> None:
class State(TypedDict):
my_key: Annotated[str, operator.add]
market: str
def tool_one(s: State) -> State:
return {"my_key": " one"}
tool_two_node_count = 0
def tool_two_node(s: State) -> State:
nonlocal tool_two_node_count
tool_two_node_count += 1
if s["market"] == "DE":
time.sleep(0.1)
answer = interrupt("Just because...")
else:
answer = " all good"
return {"my_key": answer}
def start(state: State) -> list[Send | str]:
return ["tool_two", Send("tool_one", state)]
tool_two_graph = StateGraph(State)
tool_two_graph.add_node("tool_two", tool_two_node, retry_policy=RetryPolicy())
tool_two_graph.add_node("tool_one", tool_one)
tool_two_graph.set_conditional_entry_point(start)
tool_two = tool_two_graph.compile()
tracer = FakeTracer()
assert tool_two.invoke(
{"my_key": "value", "market": "DE"}, {"callbacks": [tracer]}
) == {
"my_key": "value one",
"market": "DE",
"__interrupt__": [Interrupt(value="Just because...", id=AnyStr())],
}
assert tool_two_node_count == 1, "interrupts aren't retried"
assert len(tracer.runs) == 1
run = tracer.runs[0]
assert run.end_time is not None
assert run.error is None
assert run.outputs == {"market": "DE", "my_key": "value one"}
assert tool_two.invoke({"my_key": "value", "market": "US"}) == {
"my_key": "value all good one",
"market": "US",
}
tool_two = tool_two_graph.compile(checkpointer=sync_checkpointer)
# missing thread_id
with pytest.raises(ValueError, match="thread_id"):
tool_two.invoke({"my_key": "value", "market": "DE"})
# flow: interrupt -> resume with answer
thread2 = {"configurable": {"thread_id": "2"}}
# stop when about to enter node
assert [
c for c in tool_two.stream({"my_key": "value ⛰️", "market": "DE"}, thread2)
] == [
{
"tool_one": {"my_key": " one"},
},
{
"__interrupt__": (
Interrupt(
value="Just because...",
id=AnyStr(),
),
)
},
]
# resume with answer
assert [c for c in tool_two.stream(Command(resume=" my answer"), thread2)] == [
{"tool_one": {"my_key": " one"}, "__metadata__": {"cached": True}},
{"tool_two": {"my_key": " my answer"}},
]
# flow: interrupt -> clear tasks
thread1 = {"configurable": {"thread_id": "1"}}
# stop when about to enter node
assert tool_two.invoke(
{"my_key": "value ⛰️", "market": "DE"}, thread1, durability="exit"
) == {
"my_key": "value ⛰️ one",
"market": "DE",
"__interrupt__": [Interrupt(value="Just because...", id=AnyStr())],
}
assert [c.metadata for c in tool_two.checkpointer.list(thread1)] == [
{
"parents": {},
"source": "loop",
"step": 0,
},
]
assert tool_two.get_state(thread1) == StateSnapshot(
values={"my_key": "value ⛰️ one", "market": "DE"},
next=("tool_two",),
tasks=(
PregelTask(
id=AnyStr(),
name="tool_one",
path=("__pregel_push", 0, False),
result={"my_key": " one"},
),
PregelTask(
AnyStr(),
"tool_two",
(PULL, "tool_two"),
interrupts=(
Interrupt(
value="Just because...",
id=AnyStr(),
),
),
),
),
config={
"configurable": {
"thread_id": "1",
"checkpoint_ns": "",
"checkpoint_id": AnyStr(),
}
},
created_at=AnyStr(),
metadata={
"parents": {},
"source": "loop",
"step": 0,
},
parent_config=None,
interrupts=(
Interrupt(
value="Just because...",
id=AnyStr(),
),
),
)
# clear the interrupt and next tasks
tool_two.update_state(thread1, None, as_node=END)
# interrupt and unresolved tasks are cleared, finished tasks are kept
assert tool_two.get_state(thread1) == StateSnapshot(
values={"my_key": "value ⛰️ one", "market": "DE"},
next=(),
tasks=(),
config={
"configurable": {
"thread_id": "1",
"checkpoint_ns": "",
"checkpoint_id": AnyStr(),
}
},
created_at=AnyStr(),
metadata={
"parents": {},
"source": "update",
"step": 1,
},
parent_config=([*tool_two.checkpointer.list(thread1, limit=2)][-1].config),
interrupts=(),
)
def test_dynamic_interrupt_subgraph(sync_checkpointer: BaseCheckpointSaver) -> None:
class SubgraphState(TypedDict):
my_key: str
market: str
tool_two_node_count = 0
def tool_two_node(s: SubgraphState) -> SubgraphState:
nonlocal tool_two_node_count
tool_two_node_count += 1
if s["market"] == "DE":
answer = interrupt("Just because...")
else:
answer = " all good"
return {"my_key": answer}
subgraph = StateGraph(SubgraphState)
subgraph.add_node("do", tool_two_node, retry_policy=RetryPolicy())
subgraph.add_edge(START, "do")
class State(TypedDict):
my_key: Annotated[str, operator.add]
market: str
tool_two_graph = StateGraph(State)
tool_two_graph.add_node("tool_two", subgraph.compile())
tool_two_graph.add_edge(START, "tool_two")
tool_two = tool_two_graph.compile()
tracer = FakeTracer()
assert tool_two.invoke(
{"my_key": "value", "market": "DE"}, {"callbacks": [tracer]}
) == {
"my_key": "value",
"market": "DE",
"__interrupt__": [
Interrupt(
value="Just because...",
id=AnyStr(),
)
],
}
assert tool_two_node_count == 1, "interrupts aren't retried"
assert len(tracer.runs) == 1
run = tracer.runs[0]
assert run.end_time is not None
assert run.error is None
assert run.outputs == {"market": "DE", "my_key": "value"}
assert tool_two.invoke({"my_key": "value", "market": "US"}) == {
"my_key": "value all good",
"market": "US",
}
tool_two = tool_two_graph.compile(checkpointer=sync_checkpointer)
# missing thread_id
with pytest.raises(ValueError, match="thread_id"):
tool_two.invoke({"my_key": "value", "market": "DE"})
# flow: interrupt -> resume with answer
thread2 = {"configurable": {"thread_id": "2"}}
# stop when about to enter node
assert [
c for c in tool_two.stream({"my_key": "value ⛰️", "market": "DE"}, thread2)
] == [
{
"__interrupt__": (
Interrupt(
value="Just because...",
id=AnyStr(),
),
)
},
]
# resume with answer
assert [c for c in tool_two.stream(Command(resume=" my answer"), thread2)] == [
{"tool_two": {"my_key": " my answer", "market": "DE"}},
]
# flow: interrupt -> clear tasks
thread1 = {"configurable": {"thread_id": "1"}}
# stop when about to enter node
assert tool_two.invoke(
{"my_key": "value ⛰️", "market": "DE"}, thread1, durability="exit"
) == {
"my_key": "value ⛰️",
"market": "DE",
"__interrupt__": [
Interrupt(
value="Just because...",
id=AnyStr(),
)
],
}
assert [
c.metadata
for c in tool_two.checkpointer.list(
{"configurable": {"thread_id": "1", "checkpoint_ns": ""}}
)
] == [
{
"parents": {},
"source": "loop",
"step": 0,
},
]
assert tool_two.get_state(thread1) == StateSnapshot(
values={"my_key": "value ⛰️", "market": "DE"},
next=("tool_two",),
tasks=(
PregelTask(
AnyStr(),
"tool_two",
(PULL, "tool_two"),
interrupts=(
Interrupt(
value="Just because...",
id=AnyStr(),
),
),
state={
"configurable": {
"thread_id": "1",
"checkpoint_ns": AnyStr("tool_two:"),
}
},
),
),
config={
"configurable": {
"thread_id": "1",
"checkpoint_ns": AnyStr(),
"checkpoint_id": AnyStr(),
}
},
created_at=AnyStr(),
metadata={
"parents": {},
"source": "loop",
"step": 0,
},
parent_config=None,
interrupts=(
Interrupt(
value="Just because...",
id=AnyStr(),
),
),
)
# clear the interrupt and next tasks
tool_two.update_state(thread1, None, as_node=END)
# interrupt and next tasks are cleared
assert tool_two.get_state(thread1) == StateSnapshot(
values={"my_key": "value ⛰️", "market": "DE"},
next=(),
tasks=(),
config={
"configurable": {
"thread_id": "1",
"checkpoint_ns": AnyStr(),
"checkpoint_id": AnyStr(),
}
},
created_at=AnyStr(),
metadata={
"parents": {},
"source": "update",
"step": 1,
},
parent_config=(
list(
tool_two.checkpointer.list(
{"configurable": {"thread_id": "1", "checkpoint_ns": ""}}, limit=2
)
)[-1].config
),
interrupts=(),
)
def test_send_dedupe_on_resume(
sync_checkpointer: BaseCheckpointSaver, durability: Durability
) -> None:
class InterruptOnce:
ticks: int = 0
def __call__(self, state):
self.ticks += 1
if self.ticks == 1:
interrupt("Bahh")
return ["|".join(("flaky", str(state)))]
class Node:
def __init__(self, name: str):
self.name = name
self.ticks = 0
setattr(self, "__name__", name)
def __call__(self, state):
time.sleep(0)
# sleep makes it more likely to trigger edge case where 1st task
# finishes before 2nd is registered in futures dict
self.ticks += 1
update = (
[self.name]
if isinstance(state, list)
else ["|".join((self.name, str(state)))]
)
if isinstance(state, Command):
return replace(state, update=update)
else:
return update
def send_for_fun(state):
return [
Send("2", Command(goto=Send("2", 3))),
Send("2", Command(goto=Send("flaky", 4))),
"3.1",
]
def route_to_three(state) -> Literal["3"]:
return "3"
builder = StateGraph(Annotated[list, operator.add])
builder.add_node(Node("1"))
builder.add_node(Node("2"))
builder.add_node(Node("3"))
builder.add_node(Node("3.1"))
builder.add_node("flaky", InterruptOnce())
builder.add_edge(START, "1")
builder.add_conditional_edges("1", send_for_fun)
builder.add_conditional_edges("2", route_to_three)
graph = builder.compile(checkpointer=sync_checkpointer)
thread1 = {"configurable": {"thread_id": "1"}}
assert graph.invoke(["0"], thread1, durability=durability) == {
"__interrupt__": [
Interrupt(
value="Bahh",
id=AnyStr(),
),
],
}
assert builder.nodes["2"].runnable.func.ticks == 3
assert builder.nodes["flaky"].runnable.func.ticks == 1
# check state
state = graph.get_state(thread1)
assert state.next == ("flaky",)
# check history
history = [c for c in graph.get_state_history(thread1)]
assert len(history) == (4 if durability != "exit" else 1)
# resume execution
assert graph.invoke(None, thread1, durability=durability) == [
"0",
"1",
"3.1",
"2|Command(goto=Send(node='2', arg=3))",
"2|Command(goto=Send(node='flaky', arg=4))",
"3",
"2|3",
"flaky|4",
"3",
]
# node "2" doesn't get called again, as we recover writes saved before
assert builder.nodes["2"].runnable.func.ticks == 3
# node "flaky" gets called again, as it was interrupted
assert builder.nodes["flaky"].runnable.func.ticks == 2
# check state
state = graph.get_state(thread1)
assert state.next == ()
# check history
history = [c for c in graph.get_state_history(thread1)]
assert len(history) == (6 if durability != "exit" else 2)
expected_history = [
StateSnapshot(
values=[
"0",
"1",
"3.1",
"2|Command(goto=Send(node='2', arg=3))",
"2|Command(goto=Send(node='flaky', arg=4))",
"3",
"2|3",
"flaky|4",
"3",
],
next=(),
config={
"configurable": {
"thread_id": "1",
"checkpoint_ns": "",
"checkpoint_id": AnyStr(),
}
},
metadata={
"source": "loop",
"step": 4,
"parents": {},
},
created_at=AnyStr(),
parent_config={
"configurable": {
"thread_id": "1",
"checkpoint_ns": "",
"checkpoint_id": AnyStr(),
}
},
tasks=(),
interrupts=(),
),
StateSnapshot(
values=[
"0",
"1",
"3.1",
"2|Command(goto=Send(node='2', arg=3))",
"2|Command(goto=Send(node='flaky', arg=4))",
"3",
"2|3",
"flaky|4",
],
next=("3",),
config={
"configurable": {
"thread_id": "1",
"checkpoint_ns": "",
"checkpoint_id": AnyStr(),
}
},
metadata={
"source": "loop",
"step": 3,
"parents": {},
},
created_at=AnyStr(),
parent_config={
"configurable": {
"thread_id": "1",
"checkpoint_ns": "",
"checkpoint_id": AnyStr(),
}
},
tasks=(
PregelTask(
id=AnyStr(),
name="3",
path=("__pregel_pull", "3"),
error=None,
interrupts=(),
state=None,
result=["3"],
),
),
interrupts=(),
),
StateSnapshot(
values=[
"0",
"1",
"3.1",
"2|Command(goto=Send(node='2', arg=3))",
"2|Command(goto=Send(node='flaky', arg=4))",
],
next=("2", "flaky", "3"),
config={
"configurable": {
"thread_id": "1",
"checkpoint_ns": "",
"checkpoint_id": AnyStr(),
}
},
metadata={
"source": "loop",
"step": 2,
"parents": {},
},
created_at=AnyStr(),
parent_config={
"configurable": {
"thread_id": "1",
"checkpoint_ns": "",
"checkpoint_id": AnyStr(),
}
},
tasks=(
PregelTask(
id=AnyStr(),
name="2",
path=("__pregel_push", 0, False),
error=None,
interrupts=(),
state=None,
result=["2|3"],
),
PregelTask(
id=AnyStr(),
name="flaky",
path=("__pregel_push", 1, False),
error=None,
interrupts=(Interrupt(value="Bahh", id=AnyStr()),),
state=None,
result=["flaky|4"] if durability != "exit" else None,
),
PregelTask(
id=AnyStr(),
name="3",
path=("__pregel_pull", "3"),
error=None,
interrupts=(),
state=None,
result=["3"],
),
),
interrupts=(Interrupt(value="Bahh", id=AnyStr()),),
),
StateSnapshot(
values=["0", "1"],
next=("2", "2", "3.1"),
config={
"configurable": {
"thread_id": "1",
"checkpoint_ns": "",
"checkpoint_id": AnyStr(),
}
},
metadata={
"source": "loop",
"step": 1,
"parents": {},
},
created_at=AnyStr(),
parent_config={
"configurable": {
"thread_id": "1",
"checkpoint_ns": "",
"checkpoint_id": AnyStr(),
}
},
tasks=(
PregelTask(
id=AnyStr(),
name="2",
path=("__pregel_push", 0, False),
error=None,
interrupts=(),
state=None,
result=["2|Command(goto=Send(node='2', arg=3))"],
),
PregelTask(
id=AnyStr(),
name="2",
path=("__pregel_push", 1, False),
error=None,
interrupts=(),
state=None,
result=["2|Command(goto=Send(node='flaky', arg=4))"],
),
PregelTask(
id=AnyStr(),
name="3.1",
path=("__pregel_pull", "3.1"),
error=None,
interrupts=(),
state=None,
result=["3.1"],
),
),
interrupts=(),
),
StateSnapshot(
values=["0"],
next=("1",),
config={
"configurable": {
"thread_id": "1",
"checkpoint_ns": "",
"checkpoint_id": AnyStr(),
}
},
metadata={
"source": "loop",
"step": 0,
"parents": {},
},
created_at=AnyStr(),
parent_config={
"configurable": {
"thread_id": "1",
"checkpoint_ns": "",
"checkpoint_id": AnyStr(),
}
},
tasks=(
PregelTask(
id=AnyStr(),
name="1",
path=("__pregel_pull", "1"),
error=None,
interrupts=(),
state=None,
result=["1"],
),
),
interrupts=(),
),
StateSnapshot(
values=[],
next=("__start__",),
config={
"configurable": {
"thread_id": "1",
"checkpoint_ns": "",
"checkpoint_id": AnyStr(),
}
},
metadata={
"source": "input",
"step": -1,
"parents": {},
},
created_at=AnyStr(),
parent_config=None,
interrupts=(),
tasks=(
PregelTask(
id=AnyStr(),
name="__start__",
path=("__pregel_pull", "__start__"),
error=None,
interrupts=(),
state=None,
result=["0"],
),
),
),
]
if durability != "exit":
assert history == expected_history
else:
assert history[0] == expected_history[0]._replace(
parent_config=history[1].config
)
assert history[1] == expected_history[2]._replace(parent_config=None)
def test_nested_graph_state(sync_checkpointer: BaseCheckpointSaver) -> None:
class InnerState(TypedDict):
my_key: str
my_other_key: str
def inner_1(state: InnerState):
return {
"my_key": state["my_key"] + " here",
"my_other_key": state["my_key"],
}
def inner_2(state: InnerState):
return {
"my_key": state["my_key"] + " and there",
"my_other_key": state["my_key"],
}
inner = StateGraph(InnerState)
inner.add_node("inner_1", inner_1)
inner.add_node("inner_2", inner_2)
inner.add_edge("inner_1", "inner_2")
inner.set_entry_point("inner_1")
inner.set_finish_point("inner_2")
class State(TypedDict):
my_key: str
other_parent_key: str
def outer_1(state: State):
return {"my_key": "hi " + state["my_key"]}
def outer_2(state: State):
return {"my_key": state["my_key"] + " and back again"}
graph = StateGraph(State)
graph.add_node("outer_1", outer_1)
graph.add_node(
"inner",
inner.compile(interrupt_before=["inner_2"]),
)
graph.add_node("outer_2", outer_2)
graph.set_entry_point("outer_1")
graph.add_edge("outer_1", "inner")
graph.add_edge("inner", "outer_2")
graph.set_finish_point("outer_2")
app = graph.compile(checkpointer=sync_checkpointer)
config = {"configurable": {"thread_id": "1"}}
app.invoke({"my_key": "my value"}, config, durability="exit")
# test state w/ nested subgraph state (right after interrupt)
# first get_state without subgraph state
expected = StateSnapshot(
values={"my_key": "hi my value"},
tasks=(
PregelTask(
AnyStr(),
"inner",
(PULL, "inner"),
state={"configurable": {"thread_id": "1", "checkpoint_ns": AnyStr()}},
),
),
next=("inner",),
config={
"configurable": {
"thread_id": "1",
"checkpoint_ns": "",
"checkpoint_id": AnyStr(),
}
},
metadata={
"parents": {},
"source": "loop",
"step": 1,
},
created_at=AnyStr(),
parent_config=None,
interrupts=(),
)
assert app.get_state(config) == expected
assert list(app.get_state_history(config)) == [expected]
# now, get_state with subgraphs state
assert app.get_state(config, subgraphs=True) == StateSnapshot(
values={"my_key": "hi my value"},
tasks=(
PregelTask(
AnyStr(),
"inner",
(PULL, "inner"),
state=StateSnapshot(
values={
"my_key": "hi my value here",
"my_other_key": "hi my value",
},
tasks=(
PregelTask(
AnyStr(),
"inner_2",
(PULL, "inner_2"),
),
),
next=("inner_2",),
config={
"configurable": {
"thread_id": "1",
"checkpoint_ns": AnyStr("inner:"),
"checkpoint_id": AnyStr(),
"checkpoint_map": AnyDict(
{"": AnyStr(), AnyStr("child:"): AnyStr()}
),
}
},
metadata={
"parents": {
"": AnyStr(),
},
"source": "loop",
"step": 1,
},
created_at=AnyStr(),
parent_config=None,
interrupts=(),
),
),
),
next=("inner",),
config={
"configurable": {
"thread_id": "1",
"checkpoint_ns": "",
"checkpoint_id": AnyStr(),
}
},
metadata={
"parents": {},
"source": "loop",
"step": 1,
},
created_at=AnyStr(),
parent_config=None,
interrupts=(),
)
# get_state_history for a subgraph returns its checkpoints
child_history = [*app.get_state_history(app.get_state(config).tasks[0].state)]
expected_child_history = [
StateSnapshot(
values={"my_key": "hi my value here", "my_other_key": "hi my value"},
next=("inner_2",),
config={
"configurable": {
"thread_id": "1",
"checkpoint_ns": AnyStr("inner:"),
"checkpoint_id": AnyStr(),
"checkpoint_map": AnyDict(
{"": AnyStr(), AnyStr("child:"): AnyStr()}
),
}
},
metadata={
"source": "loop",
"step": 1,
"parents": {"": AnyStr()},
},
created_at=AnyStr(),
parent_config=None,
interrupts=(),
tasks=(PregelTask(AnyStr(), "inner_2", (PULL, "inner_2")),),
),
]
assert child_history == expected_child_history
# resume
app.invoke(None, config, durability="exit")
# test state w/ nested subgraph state (after resuming from interrupt)
assert app.get_state(config) == StateSnapshot(
values={"my_key": "hi my value here and there and back again"},
tasks=(),
next=(),
config={
"configurable": {
"thread_id": "1",
"checkpoint_ns": "",
"checkpoint_id": AnyStr(),
}
},
metadata={
"parents": {},
"source": "loop",
"step": 3,
},
created_at=AnyStr(),
parent_config=(
{
"configurable": {
"thread_id": "1",
"checkpoint_ns": "",
"checkpoint_id": AnyStr(),
}
}
),
interrupts=(),
)
# test full history at the end
actual_history = list(app.get_state_history(config))
expected_history = [
StateSnapshot(
values={"my_key": "hi my value here and there and back again"},
tasks=(),
next=(),
config={
"configurable": {
"thread_id": "1",
"checkpoint_ns": "",
"checkpoint_id": AnyStr(),
}
},
metadata={
"parents": {},
"source": "loop",
"step": 3,
},
created_at=AnyStr(),
parent_config=(
{
"configurable": {
"thread_id": "1",
"checkpoint_ns": "",
"checkpoint_id": AnyStr(),
}
}
),
interrupts=(),
),
StateSnapshot(
values={"my_key": "hi my value"},
tasks=(
PregelTask(
AnyStr(),
"inner",
(PULL, "inner"),
state={
"configurable": {"thread_id": "1", "checkpoint_ns": AnyStr()}
},
result=None,
),
),
next=("inner",),
config={
"configurable": {
"thread_id": "1",
"checkpoint_ns": "",
"checkpoint_id": AnyStr(),
}
},
metadata={
"parents": {},
"source": "loop",
"step": 1,
},
created_at=AnyStr(),
parent_config=None,
interrupts=(),
),
]
assert actual_history == expected_history
# test looking up parent state by checkpoint ID
for actual_snapshot, expected_snapshot in zip(actual_history, expected_history):
assert app.get_state(actual_snapshot.config) == expected_snapshot
def test_doubly_nested_graph_state(
sync_checkpointer: BaseCheckpointSaver,
) -> None:
class State(TypedDict):
my_key: str
class ChildState(TypedDict):
my_key: str
class GrandChildState(TypedDict):
my_key: str
def grandchild_1(state: ChildState):
return {"my_key": state["my_key"] + " here"}
def grandchild_2(state: ChildState):
return {
"my_key": state["my_key"] + " and there",
}
grandchild = StateGraph(GrandChildState)
grandchild.add_node("grandchild_1", grandchild_1)
grandchild.add_node("grandchild_2", grandchild_2)
grandchild.add_edge("grandchild_1", "grandchild_2")
grandchild.set_entry_point("grandchild_1")
grandchild.set_finish_point("grandchild_2")
child = StateGraph(ChildState)
child.add_node(
"child_1",
grandchild.compile(interrupt_before=["grandchild_2"]),
)
child.set_entry_point("child_1")
child.set_finish_point("child_1")
def parent_1(state: State):
return {"my_key": "hi " + state["my_key"]}
def parent_2(state: State):
return {"my_key": state["my_key"] + " and back again"}
graph = StateGraph(State)
graph.add_node("parent_1", parent_1)
graph.add_node("child", child.compile())
graph.add_node("parent_2", parent_2)
graph.set_entry_point("parent_1")
graph.add_edge("parent_1", "child")
graph.add_edge("child", "parent_2")
graph.set_finish_point("parent_2")
app = graph.compile(checkpointer=sync_checkpointer)
# test invoke w/ nested interrupt
config = {"configurable": {"thread_id": "1"}}
assert [
c
for c in app.stream(
{"my_key": "my value"}, config, subgraphs=True, durability="exit"
)
] == [
((), {"parent_1": {"my_key": "hi my value"}}),
(
(AnyStr("child:"), AnyStr("child_1:")),
{"grandchild_1": {"my_key": "hi my value here"}},
),
((), {"__interrupt__": ()}),
]
# get state without subgraphs
outer_state = app.get_state(config)
assert outer_state == StateSnapshot(
values={"my_key": "hi my value"},
tasks=(
PregelTask(
AnyStr(),
"child",
(PULL, "child"),
state={
"configurable": {
"thread_id": "1",
"checkpoint_ns": AnyStr("child"),
}
},
),
),
next=("child",),
config={
"configurable": {
"thread_id": "1",
"checkpoint_ns": "",
"checkpoint_id": AnyStr(),
}
},
metadata={
"parents": {},
"source": "loop",
"step": 1,
},
created_at=AnyStr(),
parent_config=None,
interrupts=(),
)
child_state = app.get_state(outer_state.tasks[0].state)
assert child_state == StateSnapshot(
values={"my_key": "hi my value"},
tasks=(
PregelTask(
AnyStr(),
"child_1",
(PULL, "child_1"),
state={
"configurable": {
"thread_id": "1",
"checkpoint_ns": AnyStr(),
}
},
),
),
next=("child_1",),
config={
"configurable": {
"thread_id": "1",
"checkpoint_ns": AnyStr("child:"),
"checkpoint_id": AnyStr(),
"checkpoint_map": AnyDict(
{
"": AnyStr(),
AnyStr("child:"): AnyStr(),
}
),
}
},
metadata={
"parents": {"": AnyStr()},
"source": "loop",
"step": 0,
},
created_at=AnyStr(),
parent_config=None,
interrupts=(),
)
grandchild_state = app.get_state(child_state.tasks[0].state)
assert grandchild_state == StateSnapshot(
values={"my_key": "hi my value here"},
tasks=(
PregelTask(
AnyStr(),
"grandchild_2",
(PULL, "grandchild_2"),
),
),
next=("grandchild_2",),
config={
"configurable": {
"thread_id": "1",
"checkpoint_ns": AnyStr(),
"checkpoint_id": AnyStr(),
"checkpoint_map": AnyDict(
{
"": AnyStr(),
AnyStr("child:"): AnyStr(),
AnyStr(re.compile(r"child:.+|child1:")): AnyStr(),
}
),
}
},
metadata={
"parents": AnyDict(
{
"": AnyStr(),
AnyStr("child:"): AnyStr(),
}
),
"source": "loop",
"step": 1,
},
created_at=AnyStr(),
parent_config=None,
interrupts=(),
)
# get state with subgraphs
assert app.get_state(config, subgraphs=True) == StateSnapshot(
values={"my_key": "hi my value"},
tasks=(
PregelTask(
AnyStr(),
"child",
(PULL, "child"),
state=StateSnapshot(
values={"my_key": "hi my value"},
tasks=(
PregelTask(
AnyStr(),
"child_1",
(PULL, "child_1"),
state=StateSnapshot(
values={"my_key": "hi my value here"},
tasks=(
PregelTask(
AnyStr(),
"grandchild_2",
(PULL, "grandchild_2"),
),
),
next=("grandchild_2",),
config={
"configurable": {
"thread_id": "1",
"checkpoint_ns": AnyStr(),
"checkpoint_id": AnyStr(),
"checkpoint_map": AnyDict(
{
"": AnyStr(),
AnyStr("child:"): AnyStr(),
AnyStr(
re.compile(r"child:.+|child1:")
): AnyStr(),
}
),
}
},
metadata={
"parents": AnyDict(
{
"": AnyStr(),
AnyStr("child:"): AnyStr(),
}
),
"source": "loop",
"step": 1,
},
created_at=AnyStr(),
parent_config=None,
interrupts=(),
),
),
),
next=("child_1",),
config={
"configurable": {
"thread_id": "1",
"checkpoint_ns": AnyStr("child:"),
"checkpoint_id": AnyStr(),
"checkpoint_map": AnyDict(
{"": AnyStr(), AnyStr("child:"): AnyStr()}
),
}
},
metadata={
"parents": {"": AnyStr()},
"source": "loop",
"step": 0,
},
created_at=AnyStr(),
parent_config=None,
interrupts=(),
),
),
),
next=("child",),
config={
"configurable": {
"thread_id": "1",
"checkpoint_ns": "",
"checkpoint_id": AnyStr(),
}
},
metadata={
"parents": {},
"source": "loop",
"step": 1,
},
created_at=AnyStr(),
parent_config=None,
interrupts=(),
)
# # resume
assert [c for c in app.stream(None, config, subgraphs=True, durability="exit")] == [
(
(AnyStr("child:"), AnyStr("child_1:")),
{"grandchild_2": {"my_key": "hi my value here and there"}},
),
((AnyStr("child:"),), {"child_1": {"my_key": "hi my value here and there"}}),
((), {"child": {"my_key": "hi my value here and there"}}),
((), {"parent_2": {"my_key": "hi my value here and there and back again"}}),
]
# get state with and without subgraphs
assert (
app.get_state(config)
== app.get_state(config, subgraphs=True)
== StateSnapshot(
values={"my_key": "hi my value here and there and back again"},
tasks=(),
next=(),
config={
"configurable": {
"thread_id": "1",
"checkpoint_ns": "",
"checkpoint_id": AnyStr(),
}
},
metadata={
"parents": {},
"source": "loop",
"step": 3,
},
created_at=AnyStr(),
parent_config=(
{
"configurable": {
"thread_id": "1",
"checkpoint_ns": "",
"checkpoint_id": AnyStr(),
}
}
),
interrupts=(),
)
)
# get outer graph history
outer_history = list(app.get_state_history(config))
assert outer_history == [
StateSnapshot(
values={"my_key": "hi my value here and there and back again"},
tasks=(),
next=(),
config={
"configurable": {
"thread_id": "1",
"checkpoint_ns": "",
"checkpoint_id": AnyStr(),
}
},
metadata={
"parents": {},
"source": "loop",
"step": 3,
},
created_at=AnyStr(),
parent_config={
"configurable": {
"thread_id": "1",
"checkpoint_ns": "",
"checkpoint_id": AnyStr(),
}
},
interrupts=(),
),
StateSnapshot(
values={"my_key": "hi my value"},
tasks=(
PregelTask(
AnyStr(),
"child",
(PULL, "child"),
state={
"configurable": {
"thread_id": "1",
"checkpoint_ns": AnyStr("child"),
}
},
result=None,
),
),
next=("child",),
config={
"configurable": {
"thread_id": "1",
"checkpoint_ns": "",
"checkpoint_id": AnyStr(),
}
},
metadata={
"parents": {},
"source": "loop",
"step": 1,
},
created_at=AnyStr(),
parent_config=None,
interrupts=(),
),
]
# get child graph history
child_history = list(app.get_state_history(outer_history[1].tasks[0].state))
assert child_history == [
StateSnapshot(
values={"my_key": "hi my value"},
next=("child_1",),
config={
"configurable": {
"thread_id": "1",
"checkpoint_ns": AnyStr("child:"),
"checkpoint_id": AnyStr(),
"checkpoint_map": AnyDict(
{"": AnyStr(), AnyStr("child:"): AnyStr()}
),
}
},
metadata={
"source": "loop",
"step": 0,
"parents": {"": AnyStr()},
},
created_at=AnyStr(),
parent_config=None,
tasks=(
PregelTask(
id=AnyStr(),
name="child_1",
path=(PULL, "child_1"),
state={
"configurable": {
"thread_id": "1",
"checkpoint_ns": AnyStr("child:"),
}
},
result=None,
),
),
interrupts=(),
),
]
# get grandchild graph history
grandchild_history = list(app.get_state_history(child_history[0].tasks[0].state))
assert grandchild_history == [
StateSnapshot(
values={"my_key": "hi my value here"},
next=("grandchild_2",),
config={
"configurable": {
"thread_id": "1",
"checkpoint_ns": AnyStr(),
"checkpoint_id": AnyStr(),
"checkpoint_map": AnyDict(
{
"": AnyStr(),
AnyStr("child:"): AnyStr(),
AnyStr(re.compile(r"child:.+|child1:")): AnyStr(),
}
),
}
},
metadata={
"source": "loop",
"step": 1,
"parents": AnyDict(
{
"": AnyStr(),
AnyStr("child:"): AnyStr(),
}
),
},
created_at=AnyStr(),
parent_config=None,
tasks=(
PregelTask(
id=AnyStr(),
name="grandchild_2",
path=(PULL, "grandchild_2"),
result=None,
),
),
interrupts=(),
),
]
def test_send_to_nested_graphs(sync_checkpointer: BaseCheckpointSaver) -> None:
class OverallState(TypedDict):
subjects: list[str]
jokes: Annotated[list[str], operator.add]
def continue_to_jokes(state: OverallState):
return [Send("generate_joke", {"subject": s}) for s in state["subjects"]]
class JokeState(TypedDict):
subject: str
def edit(state: JokeState):
subject = state["subject"]
return {"subject": f"{subject} - hohoho"}
# subgraph
subgraph = StateGraph(JokeState, output_schema=OverallState)
subgraph.add_node("edit", edit)
subgraph.add_node(
"generate", lambda state: {"jokes": [f"Joke about {state['subject']}"]}
)
subgraph.set_entry_point("edit")
subgraph.add_edge("edit", "generate")
subgraph.set_finish_point("generate")
# parent graph
builder = StateGraph(OverallState)
builder.add_node(
"generate_joke",
subgraph.compile(interrupt_before=["generate"]),
)
builder.add_conditional_edges(START, continue_to_jokes)
builder.add_edge("generate_joke", END)
graph = builder.compile(checkpointer=sync_checkpointer)
config = {"configurable": {"thread_id": "1"}}
tracer = FakeTracer()
# invoke and pause at nested interrupt
assert graph.invoke(
{"subjects": ["cats", "dogs"]}, config={**config, "callbacks": [tracer]}
) == {
"subjects": ["cats", "dogs"],
"jokes": [],
}
assert len(tracer.runs) == 1, "Should produce exactly 1 root run"
# check state
outer_state = graph.get_state(config)
# update state of dogs joke graph
graph.update_state(outer_state.tasks[1].state, {"subject": "turtles - hohoho"})
# continue past interrupt
assert sorted(
graph.stream(None, config=config),
key=lambda d: d["generate_joke"]["jokes"][0],
) == [
{"generate_joke": {"jokes": ["Joke about cats - hohoho"]}},
{"generate_joke": {"jokes": ["Joke about turtles - hohoho"]}},
]
def test_send_react_interrupt(
sync_checkpointer: BaseCheckpointSaver,
) -> None:
from langchain_core.messages import AIMessage, HumanMessage, ToolCall, ToolMessage
ai_message = AIMessage(
"",
id="ai1",
tool_calls=[ToolCall(name="foo", args={"hi": [1, 2, 3]}, id=AnyStr())],
)
def agent(state):
return {"messages": ai_message}
def route(state):
if isinstance(state["messages"][-1], AIMessage):
return [
Send(call["name"], call) for call in state["messages"][-1].tool_calls
]
foo_called = 0
def foo(call: ToolCall):
nonlocal foo_called
foo_called += 1
return {"messages": ToolMessage(str(call["args"]), tool_call_id=call["id"])}
builder = StateGraph(MessagesState)
builder.add_node(agent)
builder.add_node(foo)
builder.add_edge(START, "agent")
builder.add_conditional_edges("agent", route)
graph = builder.compile()
assert graph.invoke({"messages": [HumanMessage("hello")]}) == {
"messages": [
_AnyIdHumanMessage(content="hello"),
_AnyIdAIMessage(
content="",
tool_calls=[
{
"name": "foo",
"args": {"hi": [1, 2, 3]},
"id": "",
"type": "tool_call",
}
],
),
_AnyIdToolMessage(
content="{'hi': [1, 2, 3]}",
tool_call_id=AnyStr(),
),
]
}
assert foo_called == 1
# simple interrupt-resume flow
foo_called = 0
graph = builder.compile(checkpointer=sync_checkpointer, interrupt_before=["foo"])
thread1 = {"configurable": {"thread_id": "1"}}
assert graph.invoke({"messages": [HumanMessage("hello")]}, thread1) == {
"messages": [
_AnyIdHumanMessage(content="hello"),
_AnyIdAIMessage(
content="",
tool_calls=[
{
"name": "foo",
"args": {"hi": [1, 2, 3]},
"id": "",
"type": "tool_call",
}
],
),
]
}
assert foo_called == 0
assert graph.invoke(None, thread1) == {
"messages": [
_AnyIdHumanMessage(content="hello"),
_AnyIdAIMessage(
content="",
tool_calls=[
{
"name": "foo",
"args": {"hi": [1, 2, 3]},
"id": "",
"type": "tool_call",
}
],
),
_AnyIdToolMessage(
content="{'hi': [1, 2, 3]}",
tool_call_id=AnyStr(),
),
]
}
assert foo_called == 1
# interrupt-update-resume flow
foo_called = 0
graph = builder.compile(checkpointer=sync_checkpointer, interrupt_before=["foo"])
thread1 = {"configurable": {"thread_id": "2"}}
assert graph.invoke(
{"messages": [HumanMessage("hello")]}, thread1, durability="exit"
) == {
"messages": [
_AnyIdHumanMessage(content="hello"),
_AnyIdAIMessage(
content="",
tool_calls=[
{
"name": "foo",
"args": {"hi": [1, 2, 3]},
"id": "",
"type": "tool_call",
}
],
),
]
}
assert foo_called == 0
# get state should show the pending task
state = graph.get_state(thread1)
assert state == StateSnapshot(
values={
"messages": [
_AnyIdHumanMessage(content="hello"),
_AnyIdAIMessage(
content="",
tool_calls=[
{
"name": "foo",
"args": {"hi": [1, 2, 3]},
"id": "",
"type": "tool_call",
}
],
),
]
},
next=("foo",),
config={
"configurable": {
"thread_id": "2",
"checkpoint_ns": "",
"checkpoint_id": AnyStr(),
}
},
metadata={
"step": 1,
"source": "loop",
"parents": {},
},
created_at=AnyStr(),
parent_config=None,
interrupts=(),
tasks=(
PregelTask(
id=AnyStr(),
name="foo",
path=("__pregel_push", 0, False),
error=None,
interrupts=(),
state=None,
result=None,
),
),
)
# remove the tool call, clearing the pending task
graph.update_state(
thread1, {"messages": AIMessage("Bye now", id=ai_message.id, tool_calls=[])}
)
# tool call no longer in pending tasks
assert graph.get_state(thread1) == StateSnapshot(
values={
"messages": [
_AnyIdHumanMessage(content="hello"),
_AnyIdAIMessage(
content="Bye now",
tool_calls=[],
),
]
},
next=(),
config={
"configurable": {
"thread_id": "2",
"checkpoint_ns": "",
"checkpoint_id": AnyStr(),
}
},
metadata={
"step": 2,
"source": "update",
"parents": {},
},
created_at=AnyStr(),
parent_config=(
{
"configurable": {
"thread_id": "2",
"checkpoint_ns": "",
"checkpoint_id": AnyStr(),
}
}
),
interrupts=(),
tasks=(),
)
# tool call not executed
assert graph.invoke(None, thread1) == {
"messages": [
_AnyIdHumanMessage(content="hello"),
_AnyIdAIMessage(content="Bye now"),
]
}
assert foo_called == 0
# interrupt-update-resume flow, creating new Send in update call
foo_called = 0
graph = builder.compile(checkpointer=sync_checkpointer, interrupt_before=["foo"])
thread1 = {"configurable": {"thread_id": "3"}}
assert graph.invoke(
{"messages": [HumanMessage("hello")]}, thread1, durability="exit"
) == {
"messages": [
_AnyIdHumanMessage(content="hello"),
_AnyIdAIMessage(
content="",
tool_calls=[
{
"name": "foo",
"args": {"hi": [1, 2, 3]},
"id": "",
"type": "tool_call",
}
],
),
]
}
assert foo_called == 0
# get state should show the pending task
state = graph.get_state(thread1)
assert state == StateSnapshot(
values={
"messages": [
_AnyIdHumanMessage(content="hello"),
_AnyIdAIMessage(
content="",
tool_calls=[
{
"name": "foo",
"args": {"hi": [1, 2, 3]},
"id": "",
"type": "tool_call",
}
],
),
]
},
next=("foo",),
config={
"configurable": {
"thread_id": "3",
"checkpoint_ns": "",
"checkpoint_id": AnyStr(),
}
},
metadata={
"step": 1,
"source": "loop",
"parents": {},
},
created_at=AnyStr(),
parent_config=None,
interrupts=(),
tasks=(
PregelTask(
id=AnyStr(),
name="foo",
path=("__pregel_push", 0, False),
error=None,
interrupts=(),
state=None,
result=None,
),
),
)
# replace the tool call, should clear previous send, create new one
graph.update_state(
thread1,
{
"messages": AIMessage(
"",
id=ai_message.id,
tool_calls=[
{
"name": "foo",
"args": {"hi": [4, 5, 6]},
"id": "tool1",
"type": "tool_call",
}
],
)
},
)
# prev tool call no longer in pending tasks, new tool call is
assert graph.get_state(thread1) == StateSnapshot(
values={
"messages": [
_AnyIdHumanMessage(content="hello"),
_AnyIdAIMessage(
content="",
tool_calls=[
{
"name": "foo",
"args": {"hi": [4, 5, 6]},
"id": "tool1",
"type": "tool_call",
}
],
),
]
},
next=("foo",),
config={
"configurable": {
"thread_id": "3",
"checkpoint_ns": "",
"checkpoint_id": AnyStr(),
}
},
metadata={
"step": 2,
"source": "update",
"parents": {},
},
created_at=AnyStr(),
parent_config=(
{
"configurable": {
"thread_id": "3",
"checkpoint_ns": "",
"checkpoint_id": AnyStr(),
}
}
),
interrupts=(),
tasks=(
PregelTask(
id=AnyStr(),
name="foo",
path=("__pregel_push", 0, False),
error=None,
interrupts=(),
state=None,
result=None,
),
),
)
# prev tool call not executed, new tool call is
assert graph.invoke(None, thread1) == {
"messages": [
_AnyIdHumanMessage(content="hello"),
AIMessage(
"",
id="ai1",
tool_calls=[
{
"name": "foo",
"args": {"hi": [4, 5, 6]},
"id": "tool1",
"type": "tool_call",
}
],
),
_AnyIdToolMessage(content="{'hi': [4, 5, 6]}", tool_call_id="tool1"),
]
}
assert foo_called == 1
def test_send_react_interrupt_control(
sync_checkpointer: BaseCheckpointSaver, snapshot: SnapshotAssertion
) -> None:
from langchain_core.messages import AIMessage, HumanMessage, ToolCall, ToolMessage
ai_message = AIMessage(
"",
id="ai1",
tool_calls=[ToolCall(name="foo", args={"hi": [1, 2, 3]}, id=AnyStr())],
)
def agent(state) -> Command[Literal["foo"]]:
return Command(
update={"messages": ai_message},
goto=[Send(call["name"], call) for call in ai_message.tool_calls],
)
foo_called = 0
def foo(call: ToolCall):
nonlocal foo_called
foo_called += 1
return {"messages": ToolMessage(str(call["args"]), tool_call_id=call["id"])}
builder = StateGraph(MessagesState)
builder.add_node(agent)
builder.add_node(foo)
builder.add_edge(START, "agent")
graph = builder.compile()
if isinstance(sync_checkpointer, InMemorySaver):
assert graph.get_graph().draw_mermaid() == snapshot
assert graph.invoke({"messages": [HumanMessage("hello")]}) == {
"messages": [
_AnyIdHumanMessage(content="hello"),
_AnyIdAIMessage(
content="",
tool_calls=[
{
"name": "foo",
"args": {"hi": [1, 2, 3]},
"id": "",
"type": "tool_call",
}
],
),
_AnyIdToolMessage(
content="{'hi': [1, 2, 3]}",
tool_call_id=AnyStr(),
),
]
}
assert foo_called == 1
# simple interrupt-resume flow
foo_called = 0
graph = builder.compile(checkpointer=sync_checkpointer, interrupt_before=["foo"])
thread1 = {"configurable": {"thread_id": "1"}}
assert graph.invoke({"messages": [HumanMessage("hello")]}, thread1) == {
"messages": [
_AnyIdHumanMessage(content="hello"),
_AnyIdAIMessage(
content="",
tool_calls=[
{
"name": "foo",
"args": {"hi": [1, 2, 3]},
"id": "",
"type": "tool_call",
}
],
),
]
}
assert foo_called == 0
assert graph.invoke(None, thread1) == {
"messages": [
_AnyIdHumanMessage(content="hello"),
_AnyIdAIMessage(
content="",
tool_calls=[
{
"name": "foo",
"args": {"hi": [1, 2, 3]},
"id": "",
"type": "tool_call",
}
],
),
_AnyIdToolMessage(
content="{'hi': [1, 2, 3]}",
tool_call_id=AnyStr(),
),
]
}
assert foo_called == 1
# interrupt-update-resume flow
foo_called = 0
graph = builder.compile(checkpointer=sync_checkpointer, interrupt_before=["foo"])
thread1 = {"configurable": {"thread_id": "2"}}
assert graph.invoke(
{"messages": [HumanMessage("hello")]}, thread1, durability="exit"
) == {
"messages": [
_AnyIdHumanMessage(content="hello"),
_AnyIdAIMessage(
content="",
tool_calls=[
{
"name": "foo",
"args": {"hi": [1, 2, 3]},
"id": "",
"type": "tool_call",
}
],
),
]
}
assert foo_called == 0
# get state should show the pending task
state = graph.get_state(thread1)
assert state == StateSnapshot(
values={
"messages": [
_AnyIdHumanMessage(content="hello"),
_AnyIdAIMessage(
content="",
tool_calls=[
{
"name": "foo",
"args": {"hi": [1, 2, 3]},
"id": "",
"type": "tool_call",
}
],
),
]
},
next=("foo",),
config={
"configurable": {
"thread_id": "2",
"checkpoint_ns": "",
"checkpoint_id": AnyStr(),
}
},
metadata={
"step": 1,
"source": "loop",
"parents": {},
},
created_at=AnyStr(),
parent_config=None,
tasks=(
PregelTask(
id=AnyStr(),
name="foo",
path=("__pregel_push", 0, False),
error=None,
interrupts=(),
state=None,
result=None,
),
),
interrupts=(),
)
# remove the tool call, clearing the pending task
graph.update_state(
thread1, {"messages": AIMessage("Bye now", id=ai_message.id, tool_calls=[])}
)
# tool call no longer in pending tasks
assert graph.get_state(thread1) == StateSnapshot(
values={
"messages": [
_AnyIdHumanMessage(content="hello"),
_AnyIdAIMessage(
content="Bye now",
tool_calls=[],
),
]
},
next=(),
config={
"configurable": {
"thread_id": "2",
"checkpoint_ns": "",
"checkpoint_id": AnyStr(),
}
},
metadata={
"step": 2,
"source": "update",
"parents": {},
},
created_at=AnyStr(),
parent_config=(
{
"configurable": {
"thread_id": "2",
"checkpoint_ns": "",
"checkpoint_id": AnyStr(),
}
}
),
interrupts=(),
tasks=(),
)
# tool call not executed
assert graph.invoke(None, thread1) == {
"messages": [
_AnyIdHumanMessage(content="hello"),
_AnyIdAIMessage(content="Bye now"),
]
}
assert foo_called == 0
# interrupt-update-resume flow, creating new Send in update call
# TODO add here test with invoke(Command())
def test_weather_subgraph(
sync_checkpointer: BaseCheckpointSaver, snapshot: SnapshotAssertion
) -> None:
from langchain_core.language_models.fake_chat_models import (
FakeMessagesListChatModel,
)
# setup subgraph
@tool
def get_weather(city: str):
"""Get the weather for a specific city"""
return f"I'ts sunny in {city}!"
weather_model = FakeMessagesListChatModel(
responses=[
AIMessage(
content="",
tool_calls=[
ToolCall(
id="tool_call123",
name="get_weather",
args={"city": "San Francisco"},
)
],
)
]
)
class SubGraphState(MessagesState):
city: str
def model_node(state: SubGraphState, writer: StreamWriter):
writer(" very")
result = weather_model.invoke(state["messages"])
return {"city": cast(AIMessage, result).tool_calls[0]["args"]["city"]}
def weather_node(state: SubGraphState, writer: StreamWriter):
writer(" good")
result = get_weather.invoke({"city": state["city"]})
return {"messages": [{"role": "assistant", "content": result}]}
subgraph = StateGraph(SubGraphState)
subgraph.add_node(model_node)
subgraph.add_node(weather_node)
subgraph.add_edge(START, "model_node")
subgraph.add_edge("model_node", "weather_node")
subgraph.add_edge("weather_node", END)
subgraph = subgraph.compile(interrupt_before=["weather_node"])
# setup main graph
class RouterState(MessagesState):
route: Literal["weather", "other"]
router_model = FakeMessagesListChatModel(
responses=[
AIMessage(
content="",
tool_calls=[
ToolCall(
id="tool_call123",
name="router",
args={"dest": "weather"},
)
],
)
]
)
def router_node(state: RouterState, writer: StreamWriter):
writer("I'm")
system_message = "Classify the incoming query as either about weather or not."
messages = [{"role": "system", "content": system_message}] + state["messages"]
route = router_model.invoke(messages)
return {"route": cast(AIMessage, route).tool_calls[0]["args"]["dest"]}
def normal_llm_node(state: RouterState):
return {"messages": [AIMessage("Hello!")]}
def route_after_prediction(state: RouterState):
if state["route"] == "weather":
return "weather_graph"
else:
return "normal_llm_node"
def weather_graph(state: RouterState):
return subgraph.invoke(state)
graph = StateGraph(RouterState)
graph.add_node(router_node)
graph.add_node(normal_llm_node)
graph.add_node("weather_graph", weather_graph)
graph.add_edge(START, "router_node")
graph.add_conditional_edges(
"router_node",
route_after_prediction,
path_map=["weather_graph", "normal_llm_node"],
)
graph.add_edge("normal_llm_node", END)
graph.add_edge("weather_graph", END)
graph = graph.compile(checkpointer=sync_checkpointer)
if isinstance(sync_checkpointer, InMemorySaver):
assert graph.get_graph(xray=1).draw_mermaid() == snapshot
config = {"configurable": {"thread_id": "1"}}
thread2 = {"configurable": {"thread_id": "2"}}
inputs = {"messages": [{"role": "user", "content": "what's the weather in sf"}]}
# run with custom output
assert [
c for c in graph.stream(inputs, thread2, stream_mode="custom", subgraphs=True)
] == [
((), "I'm"),
((AnyStr("weather_graph:"),), " very"),
]
assert [
c for c in graph.stream(None, thread2, stream_mode="custom", subgraphs=True)
] == [
((AnyStr("weather_graph:"),), " good"),
]
# run until interrupt
assert [
c
for c in graph.stream(
inputs,
config=config,
stream_mode="updates",
subgraphs=True,
durability="exit",
)
] == [
((), {"router_node": {"route": "weather"}}),
((AnyStr("weather_graph:"),), {"model_node": {"city": "San Francisco"}}),
((), {"__interrupt__": ()}),
]
# check current state
state = graph.get_state(config)
assert state == StateSnapshot(
values={
"messages": [_AnyIdHumanMessage(content="what's the weather in sf")],
"route": "weather",
},
next=("weather_graph",),
config={
"configurable": {
"thread_id": "1",
"checkpoint_ns": "",
"checkpoint_id": AnyStr(),
}
},
metadata={
"source": "loop",
"step": 1,
"parents": {},
},
created_at=AnyStr(),
parent_config=None,
tasks=(
PregelTask(
id=AnyStr(),
name="weather_graph",
path=(PULL, "weather_graph"),
state={
"configurable": {
"thread_id": "1",
"checkpoint_ns": AnyStr("weather_graph:"),
}
},
),
),
interrupts=(),
)
# update
graph.update_state(state.tasks[0].state, {"city": "la"})
# run after update
assert [
c
for c in graph.stream(
None, config=config, stream_mode="updates", subgraphs=True
)
] == [
(
(AnyStr("weather_graph:"),),
{
"weather_node": {
"messages": [{"role": "assistant", "content": "I'ts sunny in la!"}]
}
},
),
(
(),
{
"weather_graph": {
"messages": [
_AnyIdHumanMessage(content="what's the weather in sf"),
_AnyIdAIMessage(content="I'ts sunny in la!"),
]
}
},
),
]
# try updating acting as weather node
config = {"configurable": {"thread_id": "14"}}
inputs = {"messages": [{"role": "user", "content": "what's the weather in sf"}]}
assert [
c
for c in graph.stream(
inputs,
config=config,
stream_mode="updates",
subgraphs=True,
durability="exit",
)
] == [
((), {"router_node": {"route": "weather"}}),
((AnyStr("weather_graph:"),), {"model_node": {"city": "San Francisco"}}),
((), {"__interrupt__": ()}),
]
state = graph.get_state(config, subgraphs=True)
assert state == StateSnapshot(
values={
"messages": [_AnyIdHumanMessage(content="what's the weather in sf")],
"route": "weather",
},
next=("weather_graph",),
config={
"configurable": {
"thread_id": "14",
"checkpoint_ns": "",
"checkpoint_id": AnyStr(),
}
},
metadata={
"source": "loop",
"step": 1,
"parents": {},
},
created_at=AnyStr(),
parent_config=None,
tasks=(
PregelTask(
id=AnyStr(),
name="weather_graph",
path=(PULL, "weather_graph"),
state=StateSnapshot(
values={
"messages": [
_AnyIdHumanMessage(content="what's the weather in sf")
],
"city": "San Francisco",
},
next=("weather_node",),
config={
"configurable": {
"thread_id": "14",
"checkpoint_ns": AnyStr("weather_graph:"),
"checkpoint_id": AnyStr(),
"checkpoint_map": AnyDict(
{
"": AnyStr(),
AnyStr("weather_graph:"): AnyStr(),
}
),
}
},
metadata={
"source": "loop",
"step": 1,
"parents": {"": AnyStr()},
},
created_at=AnyStr(),
parent_config=None,
tasks=(
PregelTask(
id=AnyStr(),
name="weather_node",
path=(PULL, "weather_node"),
),
),
interrupts=(),
),
),
),
interrupts=(),
)
graph.update_state(
state.tasks[0].state.config,
{"messages": [{"role": "assistant", "content": "rainy"}]},
as_node="weather_node",
)
state = graph.get_state(config, subgraphs=True)
assert state == StateSnapshot(
values={
"messages": [_AnyIdHumanMessage(content="what's the weather in sf")],
"route": "weather",
},
next=("weather_graph",),
config={
"configurable": {
"thread_id": "14",
"checkpoint_ns": "",
"checkpoint_id": AnyStr(),
}
},
metadata={
"source": "loop",
"step": 1,
"parents": {},
},
created_at=AnyStr(),
parent_config=None,
tasks=(
PregelTask(
id=AnyStr(),
name="weather_graph",
path=(PULL, "weather_graph"),
state=StateSnapshot(
values={
"messages": [
_AnyIdHumanMessage(content="what's the weather in sf"),
_AnyIdAIMessage(content="rainy"),
],
"city": "San Francisco",
},
next=(),
config={
"configurable": {
"thread_id": "14",
"checkpoint_ns": AnyStr("weather_graph:"),
"checkpoint_id": AnyStr(),
"checkpoint_map": AnyDict(
{
"": AnyStr(),
AnyStr("weather_graph:"): AnyStr(),
}
),
}
},
metadata={
"step": 2,
"source": "update",
"parents": {"": AnyStr()},
},
created_at=AnyStr(),
parent_config=(
{
"configurable": {
"thread_id": "14",
"checkpoint_ns": AnyStr("weather_graph:"),
"checkpoint_id": AnyStr(),
"checkpoint_map": AnyDict(
{
"": AnyStr(),
AnyStr("weather_graph:"): AnyStr(),
}
),
}
}
),
interrupts=(),
tasks=(),
),
),
),
interrupts=(),
)
assert [
c
for c in graph.stream(
None, config=config, stream_mode="updates", subgraphs=True
)
] == [
(
(),
{
"weather_graph": {
"messages": [
_AnyIdHumanMessage(content="what's the weather in sf"),
_AnyIdAIMessage(content="rainy"),
]
}
},
),
]
# run with custom output, without subgraph streaming, should omit subgraph chunks
assert [
c
for c in graph.stream(
inputs, {"configurable": {"thread_id": "3"}}, stream_mode="custom"
)
] == [
"I'm",
]
# run with messages output, with subgraph streaming, should inc subgraph messages
assert [
c
for c in graph.stream(
inputs,
{"configurable": {"thread_id": "4"}},
stream_mode="messages",
subgraphs=True,
)
] == [
(
(),
(
_AnyIdAIMessage(
content="",
tool_calls=[
ToolCall(
id="tool_call123",
name="router",
args={"dest": "weather"},
)
],
),
{
"thread_id": "4",
"langgraph_step": 1,
"langgraph_node": "router_node",
"langgraph_triggers": ("branch:to:router_node",),
"langgraph_path": ("__pregel_pull", "router_node"),
"langgraph_checkpoint_ns": AnyStr("router_node:"),
"checkpoint_ns": AnyStr("router_node:"),
"ls_provider": "fakemessageslistchatmodel",
"ls_model_type": "chat",
"ls_integration": "langchain_chat_model",
"lc_versions": {"langchain-core": LANGCHAIN_CORE_VERSION},
},
),
),
(
(AnyStr("weather_graph:"),),
(
_AnyIdAIMessage(
content="",
tool_calls=[
ToolCall(
id="tool_call123",
name="get_weather",
args={"city": "San Francisco"},
)
],
),
{
"thread_id": "4",
"langgraph_step": 1,
"langgraph_node": "model_node",
"langgraph_triggers": ("branch:to:model_node",),
"langgraph_path": ("__pregel_pull", "model_node"),
"langgraph_checkpoint_ns": AnyStr("weather_graph:"),
"checkpoint_ns": AnyStr("weather_graph:"),
"ls_provider": "fakemessageslistchatmodel",
"ls_model_type": "chat",
"ls_integration": "langchain_chat_model",
"lc_versions": {"langchain-core": LANGCHAIN_CORE_VERSION},
},
),
),
]
# run with messages output, without subgraph streaming, should exc subgraph messages
assert [
c
for c in graph.stream(
inputs,
{"configurable": {"thread_id": "5"}},
stream_mode="messages",
)
] == [
(
_AnyIdAIMessage(
content="",
tool_calls=[
ToolCall(
id="tool_call123",
name="router",
args={"dest": "weather"},
)
],
),
{
"thread_id": "5",
"langgraph_step": 1,
"langgraph_node": "router_node",
"langgraph_triggers": ("branch:to:router_node",),
"langgraph_path": ("__pregel_pull", "router_node"),
"langgraph_checkpoint_ns": AnyStr("router_node:"),
"checkpoint_ns": AnyStr("router_node:"),
"ls_provider": "fakemessageslistchatmodel",
"ls_model_type": "chat",
"ls_integration": "langchain_chat_model",
"lc_versions": {"langchain-core": LANGCHAIN_CORE_VERSION},
},
),
]
def test_subgraph_to_end_does_not_warn() -> None:
"""Regression test for https://github.com/langchain-ai/langgraph/issues/5572."""
class State(TypedDict):
x: str
def update_x(state: State):
return Command(goto=END, update={"x": state["x"] + "!"})
# Subgraph
subgraph_builder = StateGraph(State)
subgraph_builder.add_node("update_x", update_x)
subgraph_builder.add_edge(START, "update_x")
subgraph_builder.add_edge("update_x", END)
subgraph = subgraph_builder.compile()
# Parent graph
builder = StateGraph(State)
builder.add_node("subgraph_node", subgraph)
builder.add_edge(START, "subgraph_node")
builder.add_edge("subgraph_node", END)
graph = builder.compile()
response = graph.invoke({"x": "hello"})
assert response == {"x": "hello!"}
Reference in New Issue View Git Blame Copy Permalink