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_async.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

4057 lines
125 KiB
Python
Raw Permalink Blame History

import asyncio
import operator
import re
import sys
from typing import (
Annotated,
Literal,
cast,
)
import pytest
from langchain_core.messages import AnyMessage, ToolCall
from langchain_core.runnables import RunnableConfig, RunnablePick
from langchain_core.version import VERSION as LANGCHAIN_CORE_VERSION
from langgraph.checkpoint.base import BaseCheckpointSaver
from langgraph.prebuilt.chat_agent_executor import create_react_agent
from langgraph.prebuilt.tool_node import ToolNode
from pytest_mock import MockerFixture
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.message import add_messages
from langgraph.graph.state import StateGraph
from langgraph.pregel import NodeBuilder, Pregel
from langgraph.types import PregelTask, Send, StateSnapshot, StreamWriter
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,
)
pytestmark = pytest.mark.anyio
async def test_invoke_two_processes_in_out_interrupt(
async_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=async_checkpointer,
interrupt_after_nodes=["one"],
)
thread1 = {"configurable": {"thread_id": "1"}}
thread2 = {"configurable": {"thread_id": "2"}}
# start execution, stop at inbox
assert await app.ainvoke(2, thread1, durability="async") is None
# inbox == 3
checkpoint = await async_checkpointer.aget(thread1)
assert checkpoint is not None
assert checkpoint["channel_values"]["inbox"] == 3
# resume execution, finish
assert await app.ainvoke(None, thread1, durability="async") == 4
# start execution again, stop at inbox
assert await app.ainvoke(20, thread1, durability="async") is None
# inbox == 21
checkpoint = await async_checkpointer.aget(thread1)
assert checkpoint is not None
assert checkpoint["channel_values"]["inbox"] == 21
# send a new value in, interrupting the previous execution
assert await app.ainvoke(3, thread1, durability="async") is None
assert await app.ainvoke(None, thread1, durability="async") == 5
# start execution again, stopping at inbox
assert await app.ainvoke(20, thread2, durability="async") is None
# inbox == 21
snapshot = await app.aget_state(thread2)
assert snapshot.values["inbox"] == 21
assert snapshot.next == ("two",)
# update the state, resume
await app.aupdate_state(thread2, 25, as_node="one")
assert await app.ainvoke(None, thread2) == 26
# no pending tasks
snapshot = await app.aget_state(thread2)
assert snapshot.next == ()
# list history
history = [c async for c in app.aget_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=(),
),
]
# forking from any previous checkpoint should re-run nodes
assert [
c async for c in app.astream(None, history[0].config, stream_mode="updates")
] == []
assert [
c async for c in app.astream(None, history[1].config, stream_mode="updates")
] == [
{"two": {"output": 5}},
]
assert [
c async for c in app.astream(None, history[2].config, stream_mode="updates")
] == [
{"one": {"inbox": 4}},
{"__interrupt__": ()},
]
async def test_fork_always_re_runs_nodes(
async_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=async_checkpointer)
thread1 = {"configurable": {"thread_id": "1"}}
# start execution, stop at inbox
assert [
c
async for c in graph.astream(
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 async for c in graph.aget_state_history(thread1)]
assert history == [
StateSnapshot(
values=6,
next=(),
tasks=(),
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 async for c in graph.astream(None, history[0].config, stream_mode="updates")
] == []
assert [
c async for c in graph.astream(None, history[1].config, stream_mode="updates")
] == [
{"add_one": 1},
]
assert [
c async for c in graph.astream(None, history[2].config, stream_mode="updates")
] == [
{"add_one": 1},
{"add_one": 1},
]
async def test_conditional_graph_state(async_checkpointer: BaseCheckpointSaver) -> None:
from langchain_core.agents import AgentAction, AgentFinish
from langchain_core.language_models.fake import FakeStreamingListLLM
from langchain_core.prompts import PromptTemplate
from langchain_core.tools import tool
class AgentState(TypedDict):
input: Annotated[str, UntrackedValue]
agent_outcome: AgentAction | AgentFinish | None
intermediate_steps: Annotated[list[tuple[AgentAction, str]], operator.add]
# 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: AgentState) -> dict:
# 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:
# 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)
workflow.set_entry_point("agent")
workflow.add_conditional_edges(
"agent", should_continue, {"continue": "tools", "exit": END}
)
workflow.add_edge("tools", "agent")
app = workflow.compile()
assert await app.ainvoke({"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 [c async for c in app.astream({"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"
),
}
},
]
patches = [c async for c in app.astream_log({"input": "what is weather in sf"})]
patch_paths = {op["path"] for log in patches for op in log.ops}
# Check that agent (one of the nodes) has its output streamed to the logs
assert "/logs/agent/streamed_output/-" in patch_paths
# Check that agent (one of the nodes) has its final output set in the logs
assert "/logs/agent/final_output" in patch_paths
assert [
p["value"]
for log in patches
for p in log.ops
if p["path"] == "/logs/agent/final_output"
or p["path"] == "/logs/agent:2/final_output"
or p["path"] == "/logs/agent:3/final_output"
] == [
{
"agent_outcome": AgentAction(
tool="search_api", tool_input="query", log="tool:search_api:query"
)
},
{
"agent_outcome": AgentAction(
tool="search_api", tool_input="another", log="tool:search_api:another"
)
},
{
"agent_outcome": AgentFinish(
return_values={"answer": "answer"}, log="finish:answer"
),
},
]
# test state get/update methods with interrupt_after
app_w_interrupt = workflow.compile(
checkpointer=async_checkpointer,
interrupt_after=["agent"],
)
config = {"configurable": {"thread_id": "1"}}
assert [
c
async for c in app_w_interrupt.astream(
{"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 await app_w_interrupt.aget_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=(),
)
await app_w_interrupt.aupdate_state(
config,
{
"agent_outcome": AgentAction(
tool="search_api",
tool_input="query",
log="tool:search_api:a different query",
)
},
)
assert await app_w_interrupt.aget_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=(
[c async for c in app_w_interrupt.checkpointer.alist(config, limit=2)][
-1
].config
),
interrupts=(),
)
assert [c async for c in app_w_interrupt.astream(None, config)] == [
{
"tools": {
"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__": ()},
]
await app_w_interrupt.aupdate_state(
config,
{
"agent_outcome": AgentFinish(
return_values={"answer": "a really nice answer"},
log="finish:a really nice answer",
)
},
)
assert await app_w_interrupt.aget_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=(
[c async for c in app_w_interrupt.checkpointer.alist(config, limit=2)][
-1
].config
),
interrupts=(),
)
# test state get/update methods with interrupt_before
app_w_interrupt = workflow.compile(
checkpointer=async_checkpointer,
interrupt_before=["tools"],
)
config = {"configurable": {"thread_id": "2"}}
llm.i = 0 # reset the llm
assert [
c
async for c in app_w_interrupt.astream(
{"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 await app_w_interrupt.aget_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=(),
)
await app_w_interrupt.aupdate_state(
config,
{
"agent_outcome": AgentAction(
tool="search_api",
tool_input="query",
log="tool:search_api:a different query",
)
},
)
assert await app_w_interrupt.aget_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=[
c async for c in app_w_interrupt.checkpointer.alist(config, limit=2)
][-1].config,
interrupts=(),
)
assert [c async for c in app_w_interrupt.astream(None, config)] == [
{
"tools": {
"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__": ()},
]
await app_w_interrupt.aupdate_state(
config,
{
"agent_outcome": AgentFinish(
return_values={"answer": "a really nice answer"},
log="finish:a really nice answer",
)
},
)
assert await app_w_interrupt.aget_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=[
c async for c in app_w_interrupt.checkpointer.alist(config, limit=2)
][-1].config,
interrupts=(),
)
async def test_prebuilt_tool_chat() -> None:
from langchain_core.messages import AIMessage, HumanMessage
from langchain_core.tools import tool
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"),
]
)
@tool()
def search_api(query: str) -> str:
"""Searches the API for the query."""
return f"result for {query}"
tools = [search_api]
app = create_react_agent(model, tools)
assert await app.ainvoke(
{"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
async for c in app.astream(
{"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},
},
),
]
stream_updates_events = [
c
async for c in app.astream(
{"messages": [HumanMessage(content="what is weather in sf")]}
)
]
assert stream_updates_events[: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 stream_updates_events[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 stream_updates_events[5:] == [
{"agent": {"messages": [_AnyIdAIMessage(content="answer")]}}
]
async def test_state_graph_packets(async_checkpointer: BaseCheckpointSaver) -> None:
from langchain_core.language_models.fake_chat_models import (
FakeMessagesListChatModel,
)
from langchain_core.messages import (
AIMessage,
BaseMessage,
HumanMessage,
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."""
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"),
]
)
# Define decision-making logic
def should_continue(data: AgentState) -> str:
# 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
async def tools_node(input: ToolCall, config: RunnableConfig) -> AgentState:
await asyncio.sleep(input["args"].get("idx", 0) / 10)
output = await tools_by_name[input["name"]].ainvoke(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", {"messages": RunnablePick("messages") | model})
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 await app.ainvoke(
{"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
async for c in app.astream(
{"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=async_checkpointer,
interrupt_after=["agent"],
)
config = {"configurable": {"thread_id": "1"}}
assert [
c
async for c in app_w_interrupt.astream(
{"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 await app_w_interrupt.aget_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 = (await app_w_interrupt.aget_state(config)).values["messages"][-1]
last_message.tool_calls[0]["args"]["query"] = "a different query"
await app_w_interrupt.aupdate_state(config, {"messages": last_message})
# message was replaced instead of appended
tup = await app_w_interrupt.checkpointer.aget_tuple(config)
assert await app_w_interrupt.aget_state(config) == StateSnapshot(
values={
"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=tup.config,
created_at=tup.checkpoint["ts"],
metadata={
"parents": {},
"source": "update",
"step": 2,
},
parent_config=(
[c async for c in app_w_interrupt.checkpointer.alist(config, limit=2)][
-1
].config
),
interrupts=(),
)
assert [c async for c in app_w_interrupt.astream(None, config)] == [
{
"tools": {
"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__": ()},
]
tup = await app_w_interrupt.checkpointer.aget_tuple(config)
assert await app_w_interrupt.aget_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=tup.config,
created_at=tup.checkpoint["ts"],
metadata={
"parents": {},
"source": "loop",
"step": 4,
},
parent_config=(
[c async for c in app_w_interrupt.checkpointer.alist(config, limit=2)][
-1
].config
),
interrupts=(),
)
await app_w_interrupt.aupdate_state(
config,
{"messages": AIMessage(content="answer", id="ai2")},
)
# replaces message even if object identity is different, as long as id is the same
tup = await app_w_interrupt.checkpointer.aget_tuple(config)
assert await app_w_interrupt.aget_state(config) == StateSnapshot(
values={
"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=tup.config,
created_at=tup.checkpoint["ts"],
metadata={
"parents": {},
"source": "update",
"step": 5,
},
parent_config=(
[c async for c in app_w_interrupt.checkpointer.alist(config, limit=2)][
-1
].config
),
interrupts=(),
)
# interrupt before tools
app_w_interrupt = workflow.compile(
checkpointer=async_checkpointer,
interrupt_before=["tools"],
)
config = {"configurable": {"thread_id": "2"}}
model.i = 0
assert [
c
async for c in app_w_interrupt.astream(
{"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__": ()},
]
tup = await app_w_interrupt.checkpointer.aget_tuple(config)
assert await app_w_interrupt.aget_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=tup.config,
created_at=tup.checkpoint["ts"],
metadata={
"parents": {},
"source": "loop",
"step": 1,
},
parent_config=None,
interrupts=(),
)
# modify ai message
last_message = (await app_w_interrupt.aget_state(config)).values["messages"][-1]
last_message.tool_calls[0]["args"]["query"] = "a different query"
await app_w_interrupt.aupdate_state(config, {"messages": last_message})
# message was replaced instead of appended
tup = await app_w_interrupt.checkpointer.aget_tuple(config)
assert await app_w_interrupt.aget_state(config) == StateSnapshot(
values={
"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=tup.config,
created_at=tup.checkpoint["ts"],
metadata={
"parents": {},
"source": "update",
"step": 2,
},
parent_config=(
[c async for c in app_w_interrupt.checkpointer.alist(config, limit=2)][
-1
].config
),
interrupts=(),
)
assert [c async for c in app_w_interrupt.astream(None, config)] == [
{
"tools": {
"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__": ()},
]
tup = await app_w_interrupt.checkpointer.aget_tuple(config)
assert await app_w_interrupt.aget_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=tup.config,
created_at=tup.checkpoint["ts"],
metadata={
"parents": {},
"source": "loop",
"step": 4,
},
parent_config=(
[c async for c in app_w_interrupt.checkpointer.alist(config, limit=2)][
-1
].config
),
interrupts=(),
)
await app_w_interrupt.aupdate_state(
config,
{"messages": AIMessage(content="answer", id="ai2")},
)
# replaces message even if object identity is different, as long as id is the same
tup = await app_w_interrupt.checkpointer.aget_tuple(config)
assert await app_w_interrupt.aget_state(config) == StateSnapshot(
values={
"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=tup.config,
created_at=tup.checkpoint["ts"],
metadata={
"parents": {},
"source": "update",
"step": 5,
},
parent_config=(
[c async for c in app_w_interrupt.checkpointer.alist(config, limit=2)][
-1
].config
),
interrupts=(),
)
async def test_message_graph(async_checkpointer: BaseCheckpointSaver) -> None:
from langchain_core.language_models.fake_chat_models import (
FakeMessagesListChatModel,
)
from langchain_core.messages import AIMessage, HumanMessage
from langchain_core.tools import tool
class FakeFunctionChatModel(FakeMessagesListChatModel):
def bind_functions(self, functions: list):
return self
@tool()
def search_api(query: str) -> str:
"""Searches the API for the query."""
return f"result for {query}"
tools = [search_api]
model = 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()
assert await app.ainvoke([HumanMessage(content="what is weather in sf")]) == [
_AnyIdHumanMessage(
content="what is weather in sf",
),
AIMessage(
content="",
tool_calls=[
{
"id": "tool_call123",
"name": "search_api",
"args": {"query": "query"},
}
],
id="ai1", # respects ids passed in
),
_AnyIdToolMessage(
content="result for query",
name="search_api",
tool_call_id="tool_call123",
),
AIMessage(
content="",
tool_calls=[
{
"id": "tool_call456",
"name": "search_api",
"args": {"query": "another"},
}
],
id="ai2",
),
_AnyIdToolMessage(
content="result for another",
name="search_api",
tool_call_id="tool_call456",
),
AIMessage(content="answer", id="ai3"),
]
assert [
c async for c in app.astream([HumanMessage(content="what is weather in sf")])
] == [
{
"agent": AIMessage(
content="",
tool_calls=[
{
"id": "tool_call123",
"name": "search_api",
"args": {"query": "query"},
}
],
id="ai1",
)
},
{
"tools": [
_AnyIdToolMessage(
content="result for query",
name="search_api",
tool_call_id="tool_call123",
)
]
},
{
"agent": AIMessage(
content="",
tool_calls=[
{
"id": "tool_call456",
"name": "search_api",
"args": {"query": "another"},
}
],
id="ai2",
)
},
{
"tools": [
_AnyIdToolMessage(
content="result for another",
name="search_api",
tool_call_id="tool_call456",
)
]
},
{"agent": AIMessage(content="answer", id="ai3")},
]
app_w_interrupt = workflow.compile(
checkpointer=async_checkpointer,
interrupt_after=["agent"],
)
config = {"configurable": {"thread_id": "1"}}
assert [
c
async for c in app_w_interrupt.astream(
HumanMessage(content="what is weather in sf"),
config,
durability="exit",
)
] == [
{
"agent": AIMessage(
content="",
tool_calls=[
{
"id": "tool_call123",
"name": "search_api",
"args": {"query": "query"},
}
],
id="ai1",
)
},
{"__interrupt__": ()},
]
tup = await app_w_interrupt.checkpointer.aget_tuple(config)
assert await app_w_interrupt.aget_state(config) == StateSnapshot(
values=[
_AnyIdHumanMessage(content="what is weather in sf"),
AIMessage(
content="",
tool_calls=[
{
"id": "tool_call123",
"name": "search_api",
"args": {"query": "query"},
}
],
id="ai1",
),
],
tasks=(PregelTask(AnyStr(), "tools", (PULL, "tools")),),
next=("tools",),
config=tup.config,
created_at=tup.checkpoint["ts"],
metadata={
"parents": {},
"source": "loop",
"step": 1,
},
parent_config=None,
interrupts=(),
)
# modify ai message
last_message = (await app_w_interrupt.aget_state(config)).values[-1]
last_message.tool_calls[0]["args"] = {"query": "a different query"}
await app_w_interrupt.aupdate_state(config, last_message)
# message was replaced instead of appended
tup = await app_w_interrupt.checkpointer.aget_tuple(config)
assert await app_w_interrupt.aget_state(config) == StateSnapshot(
values=[
_AnyIdHumanMessage(content="what is weather in sf"),
AIMessage(
content="",
id="ai1",
tool_calls=[
{
"id": "tool_call123",
"name": "search_api",
"args": {"query": "a different query"},
}
],
),
],
tasks=(PregelTask(AnyStr(), "tools", (PULL, "tools")),),
next=("tools",),
config=tup.config,
created_at=tup.checkpoint["ts"],
metadata={
"parents": {},
"source": "update",
"step": 2,
},
parent_config=(
[c async for c in app_w_interrupt.checkpointer.alist(config, limit=2)][
-1
].config
),
interrupts=(),
)
assert [c async for c in app_w_interrupt.astream(None, config)] == [
{
"tools": [
_AnyIdToolMessage(
content="result for a different query",
name="search_api",
tool_call_id="tool_call123",
)
]
},
{
"agent": AIMessage(
content="",
tool_calls=[
{
"id": "tool_call456",
"name": "search_api",
"args": {"query": "another"},
}
],
id="ai2",
)
},
{"__interrupt__": ()},
]
tup = await app_w_interrupt.checkpointer.aget_tuple(config)
assert await app_w_interrupt.aget_state(config) == StateSnapshot(
values=[
_AnyIdHumanMessage(content="what is weather in sf"),
AIMessage(
content="",
id="ai1",
tool_calls=[
{
"id": "tool_call123",
"name": "search_api",
"args": {"query": "a different query"},
}
],
),
_AnyIdToolMessage(
content="result for a different query",
name="search_api",
tool_call_id="tool_call123",
),
AIMessage(
content="",
tool_calls=[
{
"id": "tool_call456",
"name": "search_api",
"args": {"query": "another"},
}
],
id="ai2",
),
],
tasks=(PregelTask(AnyStr(), "tools", (PULL, "tools")),),
next=("tools",),
config=tup.config,
created_at=tup.checkpoint["ts"],
metadata={
"parents": {},
"source": "loop",
"step": 4,
},
parent_config=(
[c async for c in app_w_interrupt.checkpointer.alist(config, limit=2)][
-1
].config
),
interrupts=(),
)
await app_w_interrupt.aupdate_state(
config,
AIMessage(content="answer", id="ai2"),
)
# replaces message even if object identity is different, as long as id is the same
tup = await app_w_interrupt.checkpointer.aget_tuple(config)
assert await app_w_interrupt.aget_state(config) == StateSnapshot(
values=[
_AnyIdHumanMessage(content="what is weather in sf"),
AIMessage(
content="",
id="ai1",
tool_calls=[
{
"id": "tool_call123",
"name": "search_api",
"args": {"query": "a different query"},
}
],
),
_AnyIdToolMessage(
content="result for a different query",
name="search_api",
tool_call_id="tool_call123",
),
AIMessage(content="answer", id="ai2"),
],
tasks=(),
next=(),
config=tup.config,
created_at=tup.checkpoint["ts"],
metadata={
"parents": {},
"source": "update",
"step": 5,
},
parent_config=(
[c async for c in app_w_interrupt.checkpointer.alist(config, limit=2)][
-1
].config
),
interrupts=(),
)
async def test_in_one_fan_out_out_one_graph_state() -> None:
def sorted_add(x: list[str], y: list[str]) -> list[str]:
return sorted(operator.add(x, y))
class State(TypedDict, total=False):
query: str
answer: str
docs: Annotated[list[str], operator.add]
async def rewrite_query(data: State) -> State:
return {"query": f"query: {data['query']}"}
async def retriever_one(data: State) -> State:
await asyncio.sleep(0.1)
return {"docs": ["doc1", "doc2"]}
async def retriever_two(data: State) -> State:
return {"docs": ["doc3", "doc4"]}
async 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 await app.ainvoke({"query": "what is weather in sf"}) == {
"query": "query: what is weather in sf",
"docs": ["doc1", "doc2", "doc3", "doc4"],
"answer": "doc1,doc2,doc3,doc4",
}
assert [c async for c in app.astream({"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 [
c
async for c in app.astream(
{"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 [
c
async for c in app.astream(
{"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"],
},
),
]
async def test_nested_graph_state(async_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=async_checkpointer)
config = {"configurable": {"thread_id": "1"}}
await app.ainvoke({"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 await app.aget_state(config) == expected
# now, get_state with subgraphs state
assert await app.aget_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 returns outer graph checkpoints
assert [c async for c in app.aget_state_history(config)] == [expected]
# get_state_history for a subgraph returns its checkpoints
child_history = [
c
async for c in app.aget_state_history(
(await app.aget_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,
tasks=(PregelTask(AnyStr(), "inner_2", (PULL, "inner_2")),),
interrupts=(),
),
]
assert child_history == expected_child_history
# resume
await app.ainvoke(None, config, durability="exit")
# test state w/ nested subgraph state (after resuming from interrupt)
assert await app.aget_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 = [c async for c in app.aget_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 await app.aget_state(actual_snapshot.config) == expected_snapshot
async def test_doubly_nested_graph_state(
async_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=async_checkpointer)
# test invoke w/ nested interrupt
config = {"configurable": {"thread_id": "1"}}
assert [
c
async for c in app.astream(
{"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 = await app.aget_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 = await app.aget_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 = await app.aget_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 await app.aget_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 async for c in app.astream(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 (
await app.aget_state(config)
== await app.aget_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 = [c async for c in app.aget_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"),
}
},
),
),
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 = [
c async for c in app.aget_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 = [
c async for c in app.aget_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=(),
),
]
async def test_send_to_nested_graphs(async_checkpointer: BaseCheckpointSaver) -> None:
class OverallState(TypedDict):
subjects: list[str]
jokes: Annotated[list[str], operator.add]
async def continue_to_jokes(state: OverallState):
return [Send("generate_joke", {"subject": s}) for s in state["subjects"]]
class JokeState(TypedDict):
subject: str
async 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=async_checkpointer)
config = {"configurable": {"thread_id": "1"}}
tracer = FakeTracer()
# invoke and pause at nested interrupt
assert await graph.ainvoke(
{"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 = await graph.aget_state(config)
# update state of dogs joke graph
await graph.aupdate_state(
outer_state.tasks[1].state, {"subject": "turtles - hohoho"}
)
# continue past interrupt
assert await graph.ainvoke(None, config=config) == {
"subjects": ["cats", "dogs"],
"jokes": ["Joke about cats - hohoho", "Joke about turtles - hohoho"],
}
@pytest.mark.skipif(
sys.version_info < (3, 11),
reason="Python 3.11+ is required for async contextvars support",
)
async def test_weather_subgraph(
async_checkpointer: BaseCheckpointSaver,
) -> None:
from langchain_core.language_models.fake_chat_models import (
FakeMessagesListChatModel,
)
from langchain_core.messages import AIMessage, ToolCall
from langchain_core.tools import tool
from langgraph.graph import MessagesState
# 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):
# this tests that all async checkpointers tested also implement sync methods
# as the subgraph called with sync invoke will use sync checkpointer methods
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)
def get_first_in_list():
return [*graph.get_state_history(config, limit=1)][0]
graph = graph.compile(checkpointer=async_checkpointer)
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
async for c in graph.astream(
inputs, thread2, stream_mode="custom", subgraphs=True
)
] == [
((), "I'm"),
((AnyStr("weather_graph:"),), " very"),
]
assert [
c
async for c in graph.astream(
None, thread2, stream_mode="custom", subgraphs=True
)
] == [
((AnyStr("weather_graph:"),), " good"),
]
# run until interrupt
assert [
c
async for c in graph.astream(
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 = await graph.aget_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=(),
)
# confirm that list() delegates to alist() correctly
assert await asyncio.to_thread(get_first_in_list) == state
# update
await graph.aupdate_state(state.tasks[0].state, {"city": "la"})
# run after update
assert [
c
async for c in graph.astream(
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
async for c in graph.astream(
inputs,
config=config,
stream_mode="updates",
subgraphs=True,
durability="exit",
)
] == [
((), {"router_node": {"route": "weather"}}),
((AnyStr("weather_graph:"),), {"model_node": {"city": "San Francisco"}}),
((), {"__interrupt__": ()}),
]
state = await graph.aget_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,
interrupts=(),
tasks=(
PregelTask(
id=AnyStr(),
name="weather_node",
path=(PULL, "weather_node"),
),
),
),
),
),
interrupts=(),
)
await graph.aupdate_state(
state.tasks[0].state.config,
{"messages": [{"role": "assistant", "content": "rainy"}]},
as_node="weather_node",
)
state = await graph.aget_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,
interrupts=(),
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(),
}
),
}
}
),
tasks=(),
interrupts=(),
),
),
),
)
assert [
c
async for c in graph.astream(
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
async for c in graph.astream(
inputs, {"configurable": {"thread_id": "3"}}, stream_mode="custom"
)
] == [
"I'm",
]
# run with messages output, with subgraph streaming, should inc subgraph messages
assert [
c
async for c in graph.astream(
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
async for c in graph.astream(
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},
},
),
]
Reference in New Issue View Git Blame Copy Permalink