WEBDIENSTE/langgraph
1
0
Fork 0
mirror of https://github.com/langchain-ai/langgraph.git synced 2026-09-13 13:17:52 +02:00
Code Issues Packages Projects Releases Wiki Activity
Files
48df2d10847bd636d2f7ad6657cb3b406d523fb9
langgraph/tests/test_pregel.py
T
Nuno Campos 611ecdb1cd Checkpoint inputs before starting the first step for easier error recovery 
- this enables easier retrying, for any error just do .invoke(None, config) no matter which step the error happened on
2024-05-06 11:53:20 -07:00

4547 lines
147 KiB
Python
Raw Blame History

import json
import operator
import time
import warnings
from collections import Counter
from concurrent.futures import ThreadPoolExecutor
from contextlib import contextmanager
from typing import Annotated, Any, Generator, Literal, Optional, TypedDict, Union
import pytest
from langchain_core.runnables import RunnableLambda, RunnablePassthrough
from pytest_mock import MockerFixture
from syrupy import SnapshotAssertion
from langgraph.channels.base import InvalidUpdateError
from langgraph.channels.binop import BinaryOperatorAggregate
from langgraph.channels.context import Context
from langgraph.channels.last_value import LastValue
from langgraph.channels.topic import Topic
from langgraph.checkpoint.sqlite import SqliteSaver
from langgraph.graph import END, Graph
from langgraph.graph.message import MessageGraph
from langgraph.graph.state import StateGraph
from langgraph.prebuilt.chat_agent_executor import (
create_function_calling_executor,
create_tool_calling_executor,
)
from langgraph.prebuilt.tool_node import ToolNode
from langgraph.pregel import Channel, GraphRecursionError, Pregel, StateSnapshot
from tests.any_str import AnyStr
from tests.memory_assert import MemorySaverAssertImmutable
def test_graph_validation() -> None:
def logic(inp: str) -> str:
return ""
workflow = Graph()
workflow.add_node("agent", logic)
workflow.set_entry_point("agent")
workflow.set_finish_point("agent")
assert workflow.compile(), "valid graph"
workflow = Graph()
workflow.add_node("agent", logic)
workflow.set_entry_point("agent")
with pytest.raises(ValueError, match="dead-end"):
workflow.compile()
workflow = Graph()
workflow.add_node("agent", logic)
workflow.set_finish_point("agent")
with pytest.raises(ValueError, match="not reachable"):
workflow.compile()
workflow = Graph()
workflow.add_node("agent", logic)
workflow.add_node("tools", logic)
workflow.set_entry_point("agent")
workflow.add_conditional_edges("agent", logic, {"continue": "tools", "exit": END})
workflow.add_edge("tools", "agent")
assert workflow.compile(), "valid graph"
workflow = Graph()
workflow.add_node("agent", logic)
workflow.add_node("tools", logic)
workflow.set_entry_point("tools")
workflow.add_conditional_edges("agent", logic, {"continue": "tools", "exit": END})
workflow.add_edge("tools", "agent")
assert workflow.compile(), "valid graph"
workflow = Graph()
workflow.set_entry_point("tools")
workflow.add_conditional_edges("agent", logic, {"continue": "tools", "exit": END})
workflow.add_edge("tools", "agent")
workflow.add_node("agent", logic)
workflow.add_node("tools", logic)
assert workflow.compile(), "valid graph"
workflow = Graph()
workflow.set_entry_point("tools")
workflow.add_conditional_edges(
"agent", logic, {"continue": "tools", "exit": END, "hmm": "extra"}
)
workflow.add_edge("tools", "agent")
workflow.add_node("agent", logic)
workflow.add_node("tools", logic)
with pytest.raises(ValueError, match="unknown"): # extra is not defined
workflow.compile()
workflow = Graph()
workflow.set_entry_point("agent")
workflow.add_conditional_edges("agent", logic, {"continue": "tools", "exit": END})
workflow.add_edge("tools", "extra")
workflow.add_node("agent", logic)
workflow.add_node("tools", logic)
with pytest.raises(ValueError, match="unknown"): # extra is not defined
workflow.compile()
workflow = Graph()
workflow.add_node("agent", logic)
workflow.add_node("tools", logic)
workflow.add_node("extra", logic)
workflow.set_entry_point("agent")
workflow.add_conditional_edges("agent", logic, {"continue": "tools", "exit": END})
workflow.add_edge("tools", "agent")
with pytest.raises(ValueError): # extra is dead-end / not reachable
workflow.compile()
workflow = Graph()
workflow.add_node("agent", logic)
workflow.add_node("tools", logic)
workflow.add_node("extra", logic)
workflow.set_entry_point("agent")
workflow.add_conditional_edges("agent", logic)
workflow.add_edge("tools", "agent")
with pytest.raises(ValueError): # extra is dead-end
workflow.compile()
def test_invoke_single_process_in_out(mocker: MockerFixture) -> None:
add_one = mocker.Mock(side_effect=lambda x: x + 1)
chain = Channel.subscribe_to("input") | add_one | Channel.write_to("output")
app = Pregel(
nodes={
"one": chain,
},
channels={
"input": LastValue(int),
"output": LastValue(int),
},
input_channels="input",
output_channels="output",
)
graph = Graph()
graph.add_node("add_one", add_one)
graph.set_entry_point("add_one")
graph.set_finish_point("add_one")
gapp = graph.compile()
assert app.input_schema.schema() == {"title": "LangGraphInput", "type": "integer"}
assert app.output_schema.schema() == {"title": "LangGraphOutput", "type": "integer"}
with warnings.catch_warnings():
warnings.simplefilter("error") # raise warnings as errors
assert app.config_schema().schema() == {
"properties": {},
"title": "LangGraphConfig",
"type": "object",
}
assert app.invoke(2) == 3
assert app.invoke(2, output_keys=["output"]) == {"output": 3}
assert repr(app), "does not raise recursion error"
assert gapp.invoke(2, debug=True) == 3
@pytest.mark.parametrize(
"falsy_value",
[None, False, 0, "", [], {}, set(), frozenset(), 0.0, 0j],
)
def test_invoke_single_process_in_out_falsy_values(falsy_value: Any) -> None:
graph = Graph()
graph.add_node("return_falsy_const", lambda *args, **kwargs: falsy_value)
graph.set_entry_point("return_falsy_const")
graph.set_finish_point("return_falsy_const")
gapp = graph.compile()
assert gapp.invoke(1) == falsy_value
def test_invoke_single_process_in_write_kwargs(mocker: MockerFixture) -> None:
add_one = mocker.Mock(side_effect=lambda x: x + 1)
chain = (
Channel.subscribe_to("input")
| add_one
| Channel.write_to("output", fixed=5, output_plus_one=lambda x: x + 1)
)
app = Pregel(
nodes={"one": chain},
channels={
"input": LastValue(int),
"output": LastValue(int),
"fixed": LastValue(int),
"output_plus_one": LastValue(int),
},
output_channels=["output", "fixed", "output_plus_one"],
input_channels="input",
)
assert app.input_schema.schema() == {"title": "LangGraphInput", "type": "integer"}
assert app.output_schema.schema() == {
"title": "LangGraphOutput",
"type": "object",
"properties": {
"output": {"title": "Output", "type": "integer"},
"fixed": {"title": "Fixed", "type": "integer"},
"output_plus_one": {"title": "Output Plus One", "type": "integer"},
},
}
assert app.invoke(2) == {"output": 3, "fixed": 5, "output_plus_one": 4}
def test_invoke_single_process_in_out_dict(mocker: MockerFixture) -> None:
add_one = mocker.Mock(side_effect=lambda x: x + 1)
chain = Channel.subscribe_to("input") | add_one | Channel.write_to("output")
app = Pregel(
nodes={"one": chain},
channels={"input": LastValue(int), "output": LastValue(int)},
input_channels="input",
output_channels=["output"],
)
assert app.input_schema.schema() == {"title": "LangGraphInput", "type": "integer"}
assert app.output_schema.schema() == {
"title": "LangGraphOutput",
"type": "object",
"properties": {"output": {"title": "Output", "type": "integer"}},
}
assert app.invoke(2) == {"output": 3}
def test_invoke_single_process_in_dict_out_dict(mocker: MockerFixture) -> None:
add_one = mocker.Mock(side_effect=lambda x: x + 1)
chain = Channel.subscribe_to("input") | add_one | Channel.write_to("output")
app = Pregel(
nodes={"one": chain},
channels={"input": LastValue(int), "output": LastValue(int)},
input_channels=["input"],
output_channels=["output"],
)
assert app.input_schema.schema() == {
"title": "LangGraphInput",
"type": "object",
"properties": {"input": {"title": "Input", "type": "integer"}},
}
assert app.output_schema.schema() == {
"title": "LangGraphOutput",
"type": "object",
"properties": {"output": {"title": "Output", "type": "integer"}},
}
assert app.invoke({"input": 2}) == {"output": 3}
def test_invoke_two_processes_in_out(mocker: MockerFixture) -> None:
add_one = mocker.Mock(side_effect=lambda x: x + 1)
one = Channel.subscribe_to("input") | add_one | Channel.write_to("inbox")
two = Channel.subscribe_to("inbox") | add_one | Channel.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",
)
assert app.invoke(2) == 4
assert app.invoke(2, input_keys="inbox") == 3
with pytest.raises(GraphRecursionError):
app.invoke(2, {"recursion_limit": 1})
graph = Graph()
graph.add_node("add_one", add_one)
graph.add_node("add_one_more", add_one)
graph.set_entry_point("add_one")
graph.set_finish_point("add_one_more")
graph.add_edge("add_one", "add_one_more")
gapp = graph.compile()
assert gapp.invoke(2) == 4
for step, values in enumerate(gapp.stream(2), start=1):
if step == 1:
assert values == {
"add_one": 3,
}
elif step == 2:
assert values == {
"add_one_more": 4,
}
else:
assert 0, f"{step}:{values}"
assert step == 2
def test_invoke_two_processes_in_out_interrupt(mocker: MockerFixture) -> None:
add_one = mocker.Mock(side_effect=lambda x: x + 1)
one = Channel.subscribe_to("input") | add_one | Channel.write_to("inbox")
two = Channel.subscribe_to("inbox") | add_one | Channel.write_to("output")
memory = MemorySaverAssertImmutable()
app = Pregel(
nodes={"one": one, "two": two},
channels={
"inbox": LastValue(int),
"output": LastValue(int),
"input": LastValue(int),
},
input_channels="input",
output_channels="output",
checkpointer=memory,
interrupt_after_nodes=["one"],
)
# start execution, stop at inbox
assert app.invoke(2, {"configurable": {"thread_id": 1}}) is None
# inbox == 3
checkpoint = memory.get({"configurable": {"thread_id": 1}})
assert checkpoint is not None
assert checkpoint["channel_values"]["inbox"] == 3
# resume execution, finish
assert app.invoke(None, {"configurable": {"thread_id": 1}}) == 4
# start execution again, stop at inbox
assert app.invoke(20, {"configurable": {"thread_id": 1}}) is None
# inbox == 21
checkpoint = memory.get({"configurable": {"thread_id": 1}})
assert checkpoint is not None
assert checkpoint["channel_values"]["inbox"] == 21
# send a new value in, interrupting the previous execution
assert app.invoke(3, {"configurable": {"thread_id": 1}}) is None
assert app.invoke(None, {"configurable": {"thread_id": 1}}) == 5
# start execution again, stopping at inbox
assert app.invoke(20, {"configurable": {"thread_id": 2}}) is None
# inbox == 21
snapshot = app.get_state({"configurable": {"thread_id": 2}})
assert snapshot.values["inbox"] == 21
assert snapshot.next == ("two",)
# update the state, resume
app.update_state({"configurable": {"thread_id": 2}}, 25, as_node="one")
assert app.invoke(None, {"configurable": {"thread_id": 2}}) == 26
# no pending tasks
snapshot = app.get_state({"configurable": {"thread_id": 2}})
assert snapshot.next == ()
def test_invoke_two_processes_in_dict_out(mocker: MockerFixture) -> None:
add_one = mocker.Mock(side_effect=lambda x: x + 1)
one = Channel.subscribe_to("input") | add_one | Channel.write_to("inbox")
two = (
Channel.subscribe_to("inbox")
| RunnableLambda(add_one).batch
| Channel.write_to("output").batch
)
app = Pregel(
nodes={"one": one, "two": two},
channels={
"inbox": Topic(int),
"output": LastValue(int),
"input": LastValue(int),
},
input_channels=["input", "inbox"],
stream_channels=["output", "inbox"],
output_channels=["output"],
)
# [12 + 1, 2 + 1 + 1]
assert [
*app.stream(
{"input": 2, "inbox": 12}, output_keys="output", stream_mode="updates"
)
] == [
{"two": 13},
{"two": 4},
]
assert [*app.stream({"input": 2, "inbox": 12}, output_keys="output")] == [
13,
4,
]
assert [*app.stream({"input": 2, "inbox": 12}, stream_mode="updates")] == [
{"one": {"inbox": 3}, "two": {"output": 13}},
{"two": {"output": 4}},
]
assert [*app.stream({"input": 2, "inbox": 12})] == [
{"inbox": [3], "output": 13},
{"inbox": [], "output": 4},
]
assert [*app.stream({"input": 2, "inbox": 12}, stream_mode="debug")] == [
{
"type": "task",
"timestamp": AnyStr(),
"step": 0,
"payload": {
"id": "7a3cc398-2e02-5023-ad7b-e4848d3b67fa",
"name": "one",
"input": 2,
"triggers": ["input"],
},
},
{
"type": "task",
"timestamp": AnyStr(),
"step": 0,
"payload": {
"id": "34e90af0-f97e-54e0-a159-691da37f175f",
"name": "two",
"input": [12],
"triggers": ["inbox"],
},
},
{
"type": "task_result",
"timestamp": AnyStr(),
"step": 0,
"payload": {
"id": "7a3cc398-2e02-5023-ad7b-e4848d3b67fa",
"name": "one",
"result": [("inbox", 3)],
},
},
{
"type": "task_result",
"timestamp": AnyStr(),
"step": 0,
"payload": {
"id": "34e90af0-f97e-54e0-a159-691da37f175f",
"name": "two",
"result": [("output", 13)],
},
},
{
"type": "checkpoint",
"timestamp": AnyStr(),
"step": 0,
"payload": {"config": None, "values": {"output": 13, "inbox": [3]}},
},
{
"type": "task",
"timestamp": AnyStr(),
"step": 1,
"payload": {
"id": "cf7cf374-2a2a-556f-8561-91737af89d2f",
"name": "two",
"input": [3],
"triggers": ["inbox"],
},
},
{
"type": "task_result",
"timestamp": AnyStr(),
"step": 1,
"payload": {
"id": "cf7cf374-2a2a-556f-8561-91737af89d2f",
"name": "two",
"result": [("output", 4)],
},
},
{
"type": "checkpoint",
"timestamp": AnyStr(),
"step": 1,
"payload": {"config": None, "values": {"output": 4, "inbox": []}},
},
]
def test_batch_two_processes_in_out() -> None:
def add_one_with_delay(inp: int) -> int:
time.sleep(inp / 10)
return inp + 1
one = Channel.subscribe_to("input") | add_one_with_delay | Channel.write_to("one")
two = Channel.subscribe_to("one") | add_one_with_delay | Channel.write_to("output")
app = Pregel(
nodes={"one": one, "two": two},
channels={
"one": LastValue(int),
"output": LastValue(int),
"input": LastValue(int),
},
input_channels="input",
output_channels="output",
)
assert app.batch([3, 2, 1, 3, 5]) == [5, 4, 3, 5, 7]
assert app.batch([3, 2, 1, 3, 5], output_keys=["output"]) == [
{"output": 5},
{"output": 4},
{"output": 3},
{"output": 5},
{"output": 7},
]
graph = Graph()
graph.add_node("add_one", add_one_with_delay)
graph.add_node("add_one_more", add_one_with_delay)
graph.set_entry_point("add_one")
graph.set_finish_point("add_one_more")
graph.add_edge("add_one", "add_one_more")
gapp = graph.compile()
assert gapp.batch([3, 2, 1, 3, 5]) == [5, 4, 3, 5, 7]
def test_invoke_many_processes_in_out(mocker: MockerFixture) -> None:
test_size = 100
add_one = mocker.Mock(side_effect=lambda x: x + 1)
nodes = {"-1": Channel.subscribe_to("input") | add_one | Channel.write_to("-1")}
for i in range(test_size - 2):
nodes[str(i)] = (
Channel.subscribe_to(str(i - 1)) | add_one | Channel.write_to(str(i))
)
nodes["last"] = Channel.subscribe_to(str(i)) | add_one | Channel.write_to("output")
app = Pregel(
nodes=nodes,
channels={str(i): LastValue(int) for i in range(-1, test_size - 2)}
| {"input": LastValue(int), "output": LastValue(int)},
input_channels="input",
output_channels="output",
)
for _ in range(10):
assert app.invoke(2, {"recursion_limit": test_size}) == 2 + test_size
with ThreadPoolExecutor() as executor:
assert [
*executor.map(app.invoke, [2] * 10, [{"recursion_limit": test_size}] * 10)
] == [2 + test_size] * 10
def test_batch_many_processes_in_out(mocker: MockerFixture) -> None:
test_size = 100
add_one = mocker.Mock(side_effect=lambda x: x + 1)
nodes = {"-1": Channel.subscribe_to("input") | add_one | Channel.write_to("-1")}
for i in range(test_size - 2):
nodes[str(i)] = (
Channel.subscribe_to(str(i - 1)) | add_one | Channel.write_to(str(i))
)
nodes["last"] = Channel.subscribe_to(str(i)) | add_one | Channel.write_to("output")
app = Pregel(
nodes=nodes,
channels={str(i): LastValue(int) for i in range(-1, test_size - 2)}
| {"input": LastValue(int), "output": LastValue(int)},
input_channels="input",
output_channels="output",
)
for _ in range(3):
assert app.batch([2, 1, 3, 4, 5], {"recursion_limit": test_size}) == [
2 + test_size,
1 + test_size,
3 + test_size,
4 + test_size,
5 + test_size,
]
with ThreadPoolExecutor() as executor:
assert [
*executor.map(
app.batch, [[2, 1, 3, 4, 5]] * 3, [{"recursion_limit": test_size}] * 3
)
] == [
[2 + test_size, 1 + test_size, 3 + test_size, 4 + test_size, 5 + test_size]
] * 3
def test_invoke_two_processes_two_in_two_out_invalid(mocker: MockerFixture) -> None:
add_one = mocker.Mock(side_effect=lambda x: x + 1)
one = Channel.subscribe_to("input") | add_one | Channel.write_to("output")
two = Channel.subscribe_to("input") | add_one | Channel.write_to("output")
app = Pregel(
nodes={"one": one, "two": two},
channels={"output": LastValue(int), "input": LastValue(int)},
input_channels="input",
output_channels="output",
)
with pytest.raises(InvalidUpdateError):
# LastValue channels can only be updated once per iteration
app.invoke(2)
def test_invoke_two_processes_two_in_two_out_valid(mocker: MockerFixture) -> None:
add_one = mocker.Mock(side_effect=lambda x: x + 1)
one = Channel.subscribe_to("input") | add_one | Channel.write_to("output")
two = Channel.subscribe_to("input") | add_one | Channel.write_to("output")
app = Pregel(
nodes={"one": one, "two": two},
channels={
"input": LastValue(int),
"output": Topic(int),
},
input_channels="input",
output_channels="output",
)
# An Inbox channel accumulates updates into a sequence
assert app.invoke(2) == [3, 3]
def test_invoke_checkpoint(mocker: MockerFixture) -> None:
add_one = mocker.Mock(side_effect=lambda x: x["total"] + x["input"])
def raise_if_above_10(input: int) -> int:
if input > 10:
raise ValueError("Input is too large")
return input
one = (
Channel.subscribe_to(["input"]).join(["total"])
| add_one
| Channel.write_to("output", "total")
| raise_if_above_10
)
memory = MemorySaverAssertImmutable()
app = Pregel(
nodes={"one": one},
channels={
"total": BinaryOperatorAggregate(int, operator.add),
"input": LastValue(int),
"output": LastValue(int),
},
input_channels="input",
output_channels="output",
checkpointer=memory,
)
# total starts out as 0, so output is 0+2=2
assert app.invoke(2, {"configurable": {"thread_id": "1"}}) == 2
checkpoint = memory.get({"configurable": {"thread_id": "1"}})
assert checkpoint is not None
assert checkpoint["channel_values"].get("total") == 2
# total is now 2, so output is 2+3=5
assert app.invoke(3, {"configurable": {"thread_id": "1"}}) == 5
checkpoint = memory.get({"configurable": {"thread_id": "1"}})
assert checkpoint is not None
assert checkpoint["channel_values"].get("total") == 7
# total is now 2+5=7, so output would be 7+4=11, but raises ValueError
with pytest.raises(ValueError):
app.invoke(4, {"configurable": {"thread_id": "1"}})
# checkpoint is not updated
checkpoint = memory.get({"configurable": {"thread_id": "1"}})
assert checkpoint is not None
assert checkpoint["channel_values"].get("total") == 7
# on a new thread, total starts out as 0, so output is 0+5=5
assert app.invoke(5, {"configurable": {"thread_id": "2"}}) == 5
checkpoint = memory.get({"configurable": {"thread_id": "1"}})
assert checkpoint is not None
assert checkpoint["channel_values"].get("total") == 7
checkpoint = memory.get({"configurable": {"thread_id": "2"}})
assert checkpoint is not None
assert checkpoint["channel_values"].get("total") == 5
def test_invoke_checkpoint_sqlite(mocker: MockerFixture) -> None:
adder = mocker.Mock(side_effect=lambda x: x["total"] + x["input"])
def raise_if_above_10(input: int) -> int:
if input > 10:
raise ValueError("Input is too large")
return input
one = (
Channel.subscribe_to(["input"]).join(["total"])
| adder
| Channel.write_to("output", "total")
| raise_if_above_10
)
with SqliteSaver.from_conn_string(":memory:") as memory:
app = Pregel(
nodes={"one": one},
channels={
"total": BinaryOperatorAggregate(int, operator.add),
"input": LastValue(int),
"output": LastValue(int),
},
input_channels="input",
output_channels="output",
checkpointer=memory,
)
thread_1 = {"configurable": {"thread_id": "1"}}
# total starts out as 0, so output is 0+2=2
assert app.invoke(2, thread_1, debug=1) == 2
state = app.get_state(thread_1)
assert state is not None
assert state.values.get("total") == 2
assert state.next == ()
assert state.config["configurable"]["thread_ts"] == memory.get(thread_1)["ts"]
# total is now 2, so output is 2+3=5
assert app.invoke(3, thread_1) == 5
state = app.get_state(thread_1)
assert state is not None
assert state.values.get("total") == 7
assert state.config["configurable"]["thread_ts"] == memory.get(thread_1)["ts"]
# total is now 2+5=7, so output would be 7+4=11, but raises ValueError
with pytest.raises(ValueError):
app.invoke(4, thread_1)
# checkpoint is not updated
state = app.get_state(thread_1)
assert state is not None
assert state.values.get("total") == 7
assert state.next == ("one",)
"""we checkpoint inputs and it failed on "one", so the next node is one"""
# we can recover from error by sending new inputs
assert app.invoke(2, thread_1) == 9
state = app.get_state(thread_1)
assert state is not None
assert state.values.get("total") == 16, "total is now 7+9=16"
assert state.next == ()
thread_2 = {"configurable": {"thread_id": "2"}}
# on a new thread, total starts out as 0, so output is 0+5=5
assert app.invoke(5, thread_2, debug=True) == 5
state = app.get_state({"configurable": {"thread_id": "1"}})
assert state is not None
assert state.values.get("total") == 16
assert state.next == (), "checkpoint of other thread not touched"
state = app.get_state(thread_2)
assert state is not None
assert state.values.get("total") == 5
assert state.next == ()
assert len(list(app.get_state_history(thread_1, limit=1))) == 1
# list all checkpoints for thread 1
thread_1_history = [c for c in app.get_state_history(thread_1)]
# there are 7 checkpoints
assert len(thread_1_history) == 7
assert Counter(c.metadata["source"] for c in thread_1_history) == {
"input": 4,
"loop": 3,
}
# sorted descending
assert (
thread_1_history[0].config["configurable"]["thread_ts"]
> thread_1_history[1].config["configurable"]["thread_ts"]
)
# cursor pagination
cursored = list(
app.get_state_history(thread_1, limit=1, before=thread_1_history[0].config)
)
assert len(cursored) == 1
assert cursored[0].config == thread_1_history[1].config
# the last checkpoint
assert thread_1_history[0].values["total"] == 16
# the first "loop" checkpoint
assert thread_1_history[-2].values["total"] == 2
# can get each checkpoint using aget with config
assert (
memory.get(thread_1_history[0].config)["ts"]
== thread_1_history[0].config["configurable"]["thread_ts"]
)
assert (
memory.get(thread_1_history[1].config)["ts"]
== thread_1_history[1].config["configurable"]["thread_ts"]
)
thread_1_next_config = app.update_state(thread_1_history[1].config, 10)
# update creates a new checkpoint
assert (
thread_1_next_config["configurable"]["thread_ts"]
> thread_1_history[0].config["configurable"]["thread_ts"]
)
# 1 more checkpoint in history
assert len(list(app.get_state_history(thread_1))) == 8
assert Counter(
c.metadata["source"] for c in app.get_state_history(thread_1)
) == {
"update": 1,
"input": 4,
"loop": 3,
}
# the latest checkpoint is the updated one
assert app.get_state(thread_1) == app.get_state(thread_1_next_config)
def test_invoke_two_processes_two_in_join_two_out(mocker: MockerFixture) -> None:
add_one = mocker.Mock(side_effect=lambda x: x + 1)
add_10_each = mocker.Mock(side_effect=lambda x: sorted(y + 10 for y in x))
one = Channel.subscribe_to("input") | add_one | Channel.write_to("inbox")
chain_three = Channel.subscribe_to("input") | add_one | Channel.write_to("inbox")
chain_four = (
Channel.subscribe_to("inbox") | add_10_each | Channel.write_to("output")
)
app = Pregel(
nodes={
"one": one,
"chain_three": chain_three,
"chain_four": chain_four,
},
channels={
"inbox": Topic(int),
"output": LastValue(int),
"input": LastValue(int),
},
input_channels="input",
output_channels="output",
)
# Then invoke app
# We get a single array result as chain_four waits for all publishers to finish
# before operating on all elements published to topic_two as an array
for _ in range(100):
assert app.invoke(2) == [13, 13]
with ThreadPoolExecutor() as executor:
assert [*executor.map(app.invoke, [2] * 100)] == [[13, 13]] * 100
def test_invoke_join_then_call_other_pregel(mocker: MockerFixture) -> None:
add_one = mocker.Mock(side_effect=lambda x: x + 1)
add_10_each = mocker.Mock(side_effect=lambda x: [y + 10 for y in x])
inner_app = Pregel(
nodes={
"one": Channel.subscribe_to("input") | add_one | Channel.write_to("output")
},
channels={
"output": LastValue(int),
"input": LastValue(int),
},
input_channels="input",
output_channels="output",
)
one = (
Channel.subscribe_to("input")
| add_10_each
| Channel.write_to("inbox_one").map()
)
two = (
Channel.subscribe_to("inbox_one")
| inner_app.map()
| sorted
| Channel.write_to("outbox_one")
)
chain_three = Channel.subscribe_to("outbox_one") | sum | Channel.write_to("output")
app = Pregel(
nodes={
"one": one,
"two": two,
"chain_three": chain_three,
},
channels={
"inbox_one": Topic(int),
"outbox_one": LastValue(int),
"output": LastValue(int),
"input": LastValue(int),
},
input_channels="input",
output_channels="output",
)
for _ in range(10):
assert app.invoke([2, 3]) == 27
with ThreadPoolExecutor() as executor:
assert [*executor.map(app.invoke, [[2, 3]] * 10)] == [27] * 10
def test_invoke_two_processes_one_in_two_out(mocker: MockerFixture) -> None:
add_one = mocker.Mock(side_effect=lambda x: x + 1)
one = (
Channel.subscribe_to("input")
| add_one
| Channel.write_to(output=RunnablePassthrough(), between=RunnablePassthrough())
)
two = Channel.subscribe_to("between") | add_one | Channel.write_to("output")
app = Pregel(
nodes={"one": one, "two": two},
channels={
"input": LastValue(int),
"between": LastValue(int),
"output": LastValue(int),
},
stream_channels=["output", "between"],
input_channels="input",
output_channels="output",
)
assert [c for c in app.stream(2, stream_mode="updates")] == [
{"one": {"between": 3, "output": 3}},
{"two": {"output": 4}},
]
assert [c for c in app.stream(2)] == [
{"between": 3, "output": 3},
{"between": 3, "output": 4},
]
def test_invoke_two_processes_no_out(mocker: MockerFixture) -> None:
add_one = mocker.Mock(side_effect=lambda x: x + 1)
one = Channel.subscribe_to("input") | add_one | Channel.write_to("between")
two = Channel.subscribe_to("between") | add_one
app = Pregel(
nodes={"one": one, "two": two},
channels={
"input": LastValue(int),
"between": LastValue(int),
"output": LastValue(int),
},
input_channels="input",
output_channels="output",
)
# It finishes executing (once no more messages being published)
# but returns nothing, as nothing was published to OUT topic
assert app.invoke(2) is None
def test_invoke_two_processes_no_in(mocker: MockerFixture) -> None:
add_one = mocker.Mock(side_effect=lambda x: x + 1)
one = Channel.subscribe_to("between") | add_one | Channel.write_to("output")
two = Channel.subscribe_to("between") | add_one
with pytest.raises(ValueError):
Pregel(nodes={"one": one, "two": two})
def test_channel_enter_exit_timing(mocker: MockerFixture) -> None:
setup = mocker.Mock()
cleanup = mocker.Mock()
@contextmanager
def an_int() -> Generator[int, None, None]:
setup()
try:
yield 5
finally:
cleanup()
add_one = mocker.Mock(side_effect=lambda x: x + 1)
one = Channel.subscribe_to("input") | add_one | Channel.write_to("inbox")
two = (
Channel.subscribe_to("inbox")
| RunnableLambda(add_one).batch
| Channel.write_to("output").batch
)
app = Pregel(
nodes={"one": one, "two": two},
channels={
"inbox": Topic(int),
"ctx": Context(an_int, typ=int),
"output": LastValue(int),
"input": LastValue(int),
},
input_channels="input",
output_channels=["inbox", "output"],
stream_channels=["inbox", "output"],
)
assert setup.call_count == 0
assert cleanup.call_count == 0
for i, chunk in enumerate(app.stream(2)):
assert setup.call_count == 1, "Expected setup to be called once"
assert cleanup.call_count == 0, "Expected cleanup to not be called yet"
if i == 0:
assert chunk == {"inbox": [3]}
elif i == 1:
assert chunk == {"inbox": [], "output": 4}
else:
assert False, "Expected only two chunks"
assert cleanup.call_count == 1, "Expected cleanup to be called once"
def test_conditional_graph(snapshot: SnapshotAssertion) -> None:
from copy import deepcopy
from langchain.llms.fake import FakeStreamingListLLM
from langchain_community.tools import tool
from langchain_core.agents import AgentAction, AgentFinish
from langchain_core.prompts import PromptTemplate
from langchain_core.runnables import RunnablePassthrough
# 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) -> Union[AgentAction, AgentFinish]:
if input.startswith("finish"):
_, answer = input.split(":")
return AgentFinish(return_values={"answer": answer}, log=input)
else:
_, tool_name, tool_input = input.split(":")
return AgentAction(tool=tool_name, tool_input=tool_input, log=input)
agent = RunnablePassthrough.assign(agent_outcome=prompt | llm | agent_parser)
# Define tool execution logic
def execute_tools(data: dict) -> dict:
agent_action: AgentAction = data.pop("agent_outcome")
observation = {t.name: t for t in tools}[agent_action.tool].invoke(
agent_action.tool_input
)
if data.get("intermediate_steps") is None:
data["intermediate_steps"] = []
data["intermediate_steps"].append((agent_action, observation))
return data
# Define decision-making logic
def should_continue(data: dict) -> 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 = Graph()
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 json.dumps(app.get_graph().to_json(), indent=2) == snapshot
assert app.get_graph().draw_mermaid(with_styles=False) == snapshot
assert json.dumps(app.get_graph(xray=True).to_json(), indent=2) == snapshot
assert app.get_graph(xray=True).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"
),
}
# deepcopy because the nodes mutate the data
assert [deepcopy(c) for c in app.stream({"input": "what is weather in sf"})] == [
{
"agent": {
"input": "what is weather in sf",
"agent_outcome": AgentAction(
tool="search_api", tool_input="query", log="tool:search_api:query"
),
}
},
{
"tools": {
"input": "what is weather in sf",
"intermediate_steps": [
(
AgentAction(
tool="search_api",
tool_input="query",
log="tool:search_api:query",
),
"result for query",
)
],
}
},
{
"agent": {
"input": "what is weather in sf",
"intermediate_steps": [
(
AgentAction(
tool="search_api",
tool_input="query",
log="tool:search_api:query",
),
"result for query",
)
],
"agent_outcome": AgentAction(
tool="search_api",
tool_input="another",
log="tool:search_api:another",
),
}
},
{
"tools": {
"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": {
"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"
),
}
},
]
# test state get/update methods with interrupt_after
app_w_interrupt = workflow.compile(
checkpointer=MemorySaverAssertImmutable(),
interrupt_after=["agent"],
)
config = {"configurable": {"thread_id": "1"}}
assert [
c for c in app_w_interrupt.stream({"input": "what is weather in sf"}, config)
] == [
{
"agent": {
"input": "what is weather in sf",
"agent_outcome": AgentAction(
tool="search_api", tool_input="query", log="tool:search_api:query"
),
}
}
]
assert app_w_interrupt.get_state(config) == StateSnapshot(
values={
"agent": {
"input": "what is weather in sf",
"agent_outcome": AgentAction(
tool="search_api", tool_input="query", log="tool:search_api:query"
),
},
},
next=("tools",),
config=app_w_interrupt.checkpointer.get_tuple(config).config,
metadata={"source": "loop", "step": 0},
)
assert (
app_w_interrupt.checkpointer.get_tuple(config).config["configurable"][
"thread_ts"
]
is not None
)
app_w_interrupt.update_state(
config,
{
"agent_outcome": AgentAction(
tool="search_api",
tool_input="query",
log="tool:search_api:a different query",
),
"input": "what is weather in sf",
},
)
assert app_w_interrupt.get_state(config) == StateSnapshot(
values={
"agent": {
"agent_outcome": AgentAction(
tool="search_api",
tool_input="query",
log="tool:search_api:a different query",
),
"input": "what is weather in sf",
},
},
next=("tools",),
config=app_w_interrupt.checkpointer.get_tuple(config).config,
metadata={"source": "update", "step": 1},
)
assert [c for c in app_w_interrupt.stream(None, config)] == [
{
"tools": {
"input": "what is weather in sf",
"intermediate_steps": [
(
AgentAction(
tool="search_api",
tool_input="query",
log="tool:search_api:a different query",
),
"result for query",
)
],
}
},
{
"agent": {
"input": "what is weather in sf",
"intermediate_steps": [
(
AgentAction(
tool="search_api",
tool_input="query",
log="tool:search_api:a different query",
),
"result for query",
)
],
"agent_outcome": AgentAction(
tool="search_api",
tool_input="another",
log="tool:search_api:another",
),
}
},
]
app_w_interrupt.update_state(
config,
{
"input": "what is weather in sf",
"intermediate_steps": [
(
AgentAction(
tool="search_api",
tool_input="query",
log="tool:search_api:a different query",
),
"result for query",
)
],
"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": {
"input": "what is weather in sf",
"intermediate_steps": [
(
AgentAction(
tool="search_api",
tool_input="query",
log="tool:search_api:a different query",
),
"result for query",
)
],
"agent_outcome": AgentFinish(
return_values={"answer": "a really nice answer"},
log="finish:a really nice answer",
),
},
},
next=(),
config=app_w_interrupt.checkpointer.get_tuple(config).config,
metadata={"source": "update", "step": 4},
)
# test state get/update methods with interrupt_before
app_w_interrupt = workflow.compile(
checkpointer=MemorySaverAssertImmutable(),
interrupt_before=["tools"],
)
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)
] == [
{
"agent": {
"input": "what is weather in sf",
"agent_outcome": AgentAction(
tool="search_api", tool_input="query", log="tool:search_api:query"
),
}
}
]
assert app_w_interrupt.get_state(config) == StateSnapshot(
values={
"agent": {
"input": "what is weather in sf",
"agent_outcome": AgentAction(
tool="search_api", tool_input="query", log="tool:search_api:query"
),
},
},
next=("tools",),
config=app_w_interrupt.checkpointer.get_tuple(config).config,
metadata={"source": "loop", "step": 0},
)
app_w_interrupt.update_state(
config,
{
"agent_outcome": AgentAction(
tool="search_api",
tool_input="query",
log="tool:search_api:a different query",
),
"input": "what is weather in sf",
},
)
assert app_w_interrupt.get_state(config) == StateSnapshot(
values={
"agent": {
"agent_outcome": AgentAction(
tool="search_api",
tool_input="query",
log="tool:search_api:a different query",
),
"input": "what is weather in sf",
},
},
next=("tools",),
config=app_w_interrupt.checkpointer.get_tuple(config).config,
metadata={"source": "update", "step": 1},
)
assert [c for c in app_w_interrupt.stream(None, config)] == [
{
"tools": {
"input": "what is weather in sf",
"intermediate_steps": [
(
AgentAction(
tool="search_api",
tool_input="query",
log="tool:search_api:a different query",
),
"result for query",
)
],
}
},
{
"agent": {
"input": "what is weather in sf",
"intermediate_steps": [
(
AgentAction(
tool="search_api",
tool_input="query",
log="tool:search_api:a different query",
),
"result for query",
)
],
"agent_outcome": AgentAction(
tool="search_api",
tool_input="another",
log="tool:search_api:another",
),
}
},
]
app_w_interrupt.update_state(
config,
{
"input": "what is weather in sf",
"intermediate_steps": [
(
AgentAction(
tool="search_api",
tool_input="query",
log="tool:search_api:a different query",
),
"result for query",
)
],
"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": {
"input": "what is weather in sf",
"intermediate_steps": [
(
AgentAction(
tool="search_api",
tool_input="query",
log="tool:search_api:a different query",
),
"result for query",
)
],
"agent_outcome": AgentFinish(
return_values={"answer": "a really nice answer"},
log="finish:a really nice answer",
),
},
},
next=(),
config=app_w_interrupt.checkpointer.get_tuple(config).config,
metadata={"source": "update", "step": 4},
)
# test re-invoke to continue with interrupt_before
app_w_interrupt = workflow.compile(
checkpointer=MemorySaverAssertImmutable(),
interrupt_before=["tools"],
)
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)
] == [
{
"agent": {
"input": "what is weather in sf",
"agent_outcome": AgentAction(
tool="search_api", tool_input="query", log="tool:search_api:query"
),
}
}
]
assert app_w_interrupt.get_state(config) == StateSnapshot(
values={
"agent": {
"input": "what is weather in sf",
"agent_outcome": AgentAction(
tool="search_api", tool_input="query", log="tool:search_api:query"
),
},
},
next=("tools",),
config=app_w_interrupt.checkpointer.get_tuple(config).config,
metadata={"source": "loop", "step": 0},
)
assert [c for c in app_w_interrupt.stream(None, config)] == [
{
"tools": {
"input": "what is weather in sf",
"intermediate_steps": [
(
AgentAction(
tool="search_api",
tool_input="query",
log="tool:search_api:query",
),
"result for query",
)
],
}
},
{
"agent": {
"input": "what is weather in sf",
"intermediate_steps": [
(
AgentAction(
tool="search_api",
tool_input="query",
log="tool:search_api:query",
),
"result for query",
)
],
"agent_outcome": AgentAction(
tool="search_api",
tool_input="another",
log="tool:search_api:another",
),
}
},
]
assert [c for c in app_w_interrupt.stream(None, config)] == [
{
"tools": {
"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": {
"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"
),
}
},
]
def test_conditional_entrypoint_graph(snapshot: SnapshotAssertion) -> None:
def left(data: str) -> str:
return data + "->left"
def right(data: str) -> str:
return data + "->right"
def should_start(data: str) -> str:
# Logic to decide where to start
if len(data) > 10:
return "go-right"
else:
return "go-left"
# Define a new graph
workflow = Graph()
workflow.add_node("left", left)
workflow.add_node("right", right)
workflow.set_conditional_entry_point(
should_start, {"go-left": "left", "go-right": "right"}
)
workflow.add_conditional_edges("left", lambda data: END, {END: END})
workflow.add_edge("right", END)
app = workflow.compile()
assert app.get_input_schema().schema_json() == snapshot
assert app.get_output_schema().schema_json() == 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("what is weather in sf", debug=True)
== "what is weather in sf->right"
)
assert [*app.stream("what is weather in sf")] == [
{"right": "what is weather in sf->right"},
]
def test_conditional_state_graph(snapshot: SnapshotAssertion) -> None:
from langchain.llms.fake import FakeStreamingListLLM
from langchain_community.tools import tool
from langchain_core.agents import AgentAction, AgentFinish
from langchain_core.prompts import PromptTemplate
class AgentState(TypedDict, total=False):
input: str
agent_outcome: Optional[Union[AgentAction, AgentFinish]]
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, Union[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:
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 app.get_input_schema().schema_json() == snapshot
assert app.get_output_schema().schema_json() == 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=MemorySaverAssertImmutable(),
interrupt_after=["agent"],
)
config = {"configurable": {"thread_id": "1"}}
assert [
c for c in app_w_interrupt.stream({"input": "what is weather in sf"}, config)
] == [
{
"agent": {
"agent_outcome": AgentAction(
tool="search_api", tool_input="query", log="tool:search_api:query"
),
}
},
]
assert app_w_interrupt.get_state(config) == StateSnapshot(
values={
"input": "what is weather in sf",
"agent_outcome": AgentAction(
tool="search_api", tool_input="query", log="tool:search_api:query"
),
"intermediate_steps": [],
},
next=("tools",),
config=app_w_interrupt.checkpointer.get_tuple(config).config,
metadata={"source": "loop", "step": 1},
)
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={
"input": "what is weather in sf",
"agent_outcome": AgentAction(
tool="search_api",
tool_input="query",
log="tool:search_api:a different query",
),
"intermediate_steps": [],
},
next=("tools",),
config=app_w_interrupt.checkpointer.get_tuple(config).config,
metadata={"source": "update", "step": 2},
)
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",
),
}
},
]
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={
"input": "what is weather in sf",
"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",
)
],
},
next=(),
config=app_w_interrupt.checkpointer.get_tuple(config).config,
metadata={"source": "update", "step": 5},
)
# test state get/update methods with interrupt_before
app_w_interrupt = workflow.compile(
checkpointer=MemorySaverAssertImmutable(),
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)
] == [
{
"agent": {
"agent_outcome": AgentAction(
tool="search_api", tool_input="query", log="tool:search_api:query"
),
}
},
]
assert app_w_interrupt.get_state(config) == StateSnapshot(
values={
"input": "what is weather in sf",
"agent_outcome": AgentAction(
tool="search_api", tool_input="query", log="tool:search_api:query"
),
"intermediate_steps": [],
},
next=("tools",),
config=app_w_interrupt.checkpointer.get_tuple(config).config,
metadata={"source": "loop", "step": 1},
)
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={
"input": "what is weather in sf",
"agent_outcome": AgentAction(
tool="search_api",
tool_input="query",
log="tool:search_api:a different query",
),
"intermediate_steps": [],
},
next=("tools",),
config=app_w_interrupt.checkpointer.get_tuple(config).config,
metadata={"source": "update", "step": 2},
)
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",
),
}
},
]
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={
"input": "what is weather in sf",
"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",
)
],
},
next=(),
config=app_w_interrupt.checkpointer.get_tuple(config).config,
metadata={"source": "update", "step": 5},
)
# test w interrupt before all
app_w_interrupt = workflow.compile(
checkpointer=MemorySaverAssertImmutable(),
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)
] == []
assert app_w_interrupt.get_state(config) == StateSnapshot(
values={
"input": "what is weather in sf",
"intermediate_steps": [],
},
next=("agent",),
config=app_w_interrupt.checkpointer.get_tuple(config).config,
metadata={"source": "loop", "step": 0},
)
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"
),
}
},
]
assert app_w_interrupt.get_state(config) == StateSnapshot(
values={
"input": "what is weather in sf",
"agent_outcome": AgentAction(
tool="search_api", tool_input="query", log="tool:search_api:query"
),
"intermediate_steps": [],
},
next=("tools",),
config=app_w_interrupt.checkpointer.get_tuple(config).config,
metadata={"source": "loop", "step": 1},
)
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",
)
],
}
},
]
assert app_w_interrupt.get_state(config) == StateSnapshot(
values={
"input": "what is weather in sf",
"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",
)
],
},
next=("agent",),
config=app_w_interrupt.checkpointer.get_tuple(config).config,
metadata={"source": "loop", "step": 2},
)
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",
),
}
},
]
# test w interrupt after all
app_w_interrupt = workflow.compile(
checkpointer=MemorySaverAssertImmutable(),
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)
] == [
{
"agent": {
"agent_outcome": AgentAction(
tool="search_api", tool_input="query", log="tool:search_api:query"
),
}
},
]
assert app_w_interrupt.get_state(config) == StateSnapshot(
values={
"input": "what is weather in sf",
"agent_outcome": AgentAction(
tool="search_api", tool_input="query", log="tool:search_api:query"
),
"intermediate_steps": [],
},
next=("tools",),
config=app_w_interrupt.checkpointer.get_tuple(config).config,
metadata={"source": "loop", "step": 1},
)
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",
)
],
}
},
]
assert app_w_interrupt.get_state(config) == StateSnapshot(
values={
"input": "what is weather in sf",
"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",
)
],
},
next=("agent",),
config=app_w_interrupt.checkpointer.get_tuple(config).config,
metadata={"source": "loop", "step": 2},
)
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",
),
}
},
]
def test_state_graph_w_config(snapshot: SnapshotAssertion) -> None:
from langchain.llms.fake import FakeStreamingListLLM
from langchain_community.tools import tool
from langchain_core.agents import AgentAction, AgentFinish
from langchain_core.prompts import PromptTemplate
class AgentState(TypedDict, total=False):
input: str
agent_outcome: Optional[Union[AgentAction, AgentFinish]]
intermediate_steps: Annotated[list[tuple[AgentAction, str]], operator.add]
class Config(TypedDict, total=False):
tools: list[str]
# 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, Union[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:
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, Config)
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 app.config_schema().schema_json() == snapshot
def test_conditional_entrypoint_graph_state(snapshot: SnapshotAssertion) -> None:
class AgentState(TypedDict, total=False):
input: str
output: str
steps: Annotated[list[str], operator.add]
def left(data: AgentState) -> AgentState:
return {"output": data["input"] + "->left"}
def right(data: AgentState) -> AgentState:
return {"output": data["input"] + "->right"}
def should_start(data: AgentState) -> str:
assert data["steps"] == [], "Expected input to be read from the state"
# Logic to decide where to start
if len(data["input"]) > 10:
return "go-right"
else:
return "go-left"
# Define a new graph
workflow = StateGraph(AgentState)
workflow.add_node("left", left)
workflow.add_node("right", right)
workflow.set_conditional_entry_point(
should_start, {"go-left": "left", "go-right": "right"}
)
workflow.add_conditional_edges("left", lambda data: END, {END: END})
workflow.add_edge("right", END)
app = workflow.compile()
assert app.get_input_schema().schema_json() == snapshot
assert app.get_output_schema().schema_json() == 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",
"output": "what is weather in sf->right",
"steps": [],
}
assert [*app.stream({"input": "what is weather in sf"})] == [
{"right": {"output": "what is weather in sf->right"}},
]
def test_prebuilt_tool_chat(snapshot: SnapshotAssertion) -> None:
from langchain.chat_models.fake import FakeMessagesListChatModel
from langchain_community.tools import tool
from langchain_core.messages import AIMessage, HumanMessage, ToolMessage
class FakeFuntionChatModel(FakeMessagesListChatModel):
def bind_tools(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]
app = create_tool_calling_executor(
FakeFuntionChatModel(
responses=[
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"),
]
),
tools,
)
assert app.get_input_schema().schema_json() == snapshot
assert app.get_output_schema().schema_json() == 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": [
HumanMessage(content="what is weather in sf", id=AnyStr()),
AIMessage(
id=AnyStr(),
content="",
tool_calls=[
{
"id": "tool_call123",
"name": "search_api",
"args": {"query": "query"},
},
],
),
ToolMessage(
content="result for query",
name="search_api",
tool_call_id="tool_call123",
id=AnyStr(),
),
AIMessage(
id=AnyStr(),
content="",
tool_calls=[
{
"id": "tool_call234",
"name": "search_api",
"args": {"query": "another"},
},
{
"id": "tool_call567",
"name": "search_api",
"args": {"query": "a third one"},
},
],
),
ToolMessage(
content="result for another",
name="search_api",
tool_call_id="tool_call234",
id=AnyStr(),
),
ToolMessage(
content="result for a third one",
name="search_api",
tool_call_id="tool_call567",
id=AnyStr(),
),
AIMessage(content="answer", id=AnyStr()),
]
}
assert app.invoke(
{"messages": [HumanMessage(content="what is weather in sf")]},
stream_mode="updates",
) == [
{
"agent": {
"messages": [
AIMessage(
content="",
tool_calls=[
{
"id": "tool_call123",
"name": "search_api",
"args": {"query": "query"},
},
],
id=AnyStr(),
)
]
}
},
{
"action": {
"messages": [
ToolMessage(
content="result for query",
name="search_api",
tool_call_id="tool_call123",
id=AnyStr(),
)
]
}
},
{
"agent": {
"messages": [
AIMessage(
content="",
tool_calls=[
{
"id": "tool_call234",
"name": "search_api",
"args": {"query": "another"},
},
{
"id": "tool_call567",
"name": "search_api",
"args": {"query": "a third one"},
},
],
id=AnyStr(),
)
]
}
},
{
"action": {
"messages": [
ToolMessage(
content="result for another",
name="search_api",
tool_call_id="tool_call234",
id=AnyStr(),
),
ToolMessage(
content="result for a third one",
name="search_api",
tool_call_id="tool_call567",
id=AnyStr(),
),
]
}
},
{
"agent": {
"messages": [
AIMessage(
content="answer",
id=AnyStr(),
)
]
}
},
]
assert [
*app.stream({"messages": [HumanMessage(content="what is weather in sf")]})
] == [
{
"agent": {
"messages": [
AIMessage(
id=AnyStr(),
content="",
tool_calls=[
{
"id": "tool_call123",
"name": "search_api",
"args": {"query": "query"},
},
],
)
]
}
},
{
"action": {
"messages": [
ToolMessage(
content="result for query",
name="search_api",
tool_call_id="tool_call123",
id=AnyStr(),
)
]
}
},
{
"agent": {
"messages": [
AIMessage(
id=AnyStr(),
content="",
tool_calls=[
{
"id": "tool_call234",
"name": "search_api",
"args": {"query": "another"},
},
{
"id": "tool_call567",
"name": "search_api",
"args": {"query": "a third one"},
},
],
)
]
}
},
{
"action": {
"messages": [
ToolMessage(
content="result for another",
name="search_api",
tool_call_id="tool_call234",
id=AnyStr(),
),
ToolMessage(
content="result for a third one",
name="search_api",
tool_call_id="tool_call567",
id=AnyStr(),
),
]
}
},
{"agent": {"messages": [AIMessage(content="answer", id=AnyStr())]}},
]
def test_prebuilt_chat(snapshot: SnapshotAssertion) -> None:
from langchain.chat_models.fake import FakeMessagesListChatModel
from langchain_community.tools import tool
from langchain_core.messages import AIMessage, FunctionMessage, HumanMessage
class FakeFuntionChatModel(FakeMessagesListChatModel):
def bind_functions(self, functions: list):
return self
@tool()
def search_api(query: str) -> str:
"""Searches the API for the query."""
return f"result for {query}"
tools = [search_api]
app = create_function_calling_executor(
FakeFuntionChatModel(
responses=[
AIMessage(
content="",
additional_kwargs={
"function_call": {
"name": "search_api",
"arguments": json.dumps("query"),
}
},
),
AIMessage(
content="",
additional_kwargs={
"function_call": {
"name": "search_api",
"arguments": json.dumps("another"),
}
},
),
AIMessage(content="answer"),
]
),
tools,
)
assert app.get_input_schema().schema_json() == snapshot
assert app.get_output_schema().schema_json() == 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": [
HumanMessage(content="what is weather in sf", id=AnyStr()),
AIMessage(
id=AnyStr(),
content="",
additional_kwargs={
"function_call": {"name": "search_api", "arguments": '"query"'}
},
),
FunctionMessage(content="result for query", name="search_api", id=AnyStr()),
AIMessage(
id=AnyStr(),
content="",
additional_kwargs={
"function_call": {"name": "search_api", "arguments": '"another"'}
},
),
FunctionMessage(
content="result for another", name="search_api", id=AnyStr()
),
AIMessage(content="answer", id=AnyStr()),
]
}
assert [
*app.stream({"messages": [HumanMessage(content="what is weather in sf")]})
] == [
{
"agent": {
"messages": [
AIMessage(
id=AnyStr(),
content="",
additional_kwargs={
"function_call": {
"name": "search_api",
"arguments": '"query"',
}
},
)
]
}
},
{
"action": {
"messages": [
FunctionMessage(
content="result for query", name="search_api", id=AnyStr()
)
]
}
},
{
"agent": {
"messages": [
AIMessage(
id=AnyStr(),
content="",
additional_kwargs={
"function_call": {
"name": "search_api",
"arguments": '"another"',
}
},
)
]
}
},
{
"action": {
"messages": [
FunctionMessage(
content="result for another", name="search_api", id=AnyStr()
)
]
}
},
{"agent": {"messages": [AIMessage(content="answer", id=AnyStr())]}},
]
def test_message_graph(
snapshot: SnapshotAssertion,
deterministic_uuids: MockerFixture,
) -> None:
from copy import deepcopy
from langchain.chat_models.fake import FakeMessagesListChatModel
from langchain_community.tools import tool
from langchain_core.callbacks import CallbackManagerForLLMRun
from langchain_core.messages import (
AIMessage,
BaseMessage,
HumanMessage,
ToolMessage,
)
from langchain_core.outputs import ChatGeneration, ChatResult
class FakeFuntionChatModel(FakeMessagesListChatModel):
def bind_functions(self, functions: list):
return self
def _generate(
self,
messages: list[BaseMessage],
stop: Optional[list[str]] = None,
run_manager: Optional[CallbackManagerForLLMRun] = None,
**kwargs: Any,
) -> ChatResult:
response = deepcopy(self.responses[self.i])
if self.i < len(self.responses) - 1:
self.i += 1
else:
self.i = 0
generation = ChatGeneration(message=response)
return ChatResult(generations=[generation])
@tool()
def search_api(query: str) -> str:
"""Searches the API for the query."""
return f"result for {query}"
tools = [search_api]
model = FakeFuntionChatModel(
responses=[
AIMessage(
content="",
tool_calls=[
{
"id": "tool_call123",
"name": "search_api",
"args": {"query": "query"},
}
],
id="ai1",
),
AIMessage(
content="",
tool_calls=[
{
"id": "tool_call456",
"name": "search_api",
"args": {"query": "another"},
}
],
id="ai2",
),
AIMessage(content="answer", id="ai3"),
]
)
# Define the function that determines whether to continue or not
def should_continue(messages):
last_message = messages[-1]
# If there is no function call, then we finish
if not last_message.tool_calls:
return "end"
# Otherwise if there is, we continue
else:
return "continue"
# Define a new graph
workflow = MessageGraph()
# Define the two nodes we will cycle between
workflow.add_node("agent", model)
workflow.add_node("action", 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": "action",
# 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("action", "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.get_input_schema().schema_json() == snapshot
assert app.get_output_schema().schema_json() == 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")) == [
HumanMessage(
content="what is weather in sf",
id="00000000-0000-4000-8000-000000000002", # adds missing ids
),
AIMessage(
content="",
tool_calls=[
{
"id": "tool_call123",
"name": "search_api",
"args": {"query": "query"},
}
],
id="ai1", # respects ids passed in
),
ToolMessage(
content="result for query",
name="search_api",
tool_call_id="tool_call123",
id="00000000-0000-4000-8000-000000000011",
),
AIMessage(
content="",
tool_calls=[
{
"id": "tool_call456",
"name": "search_api",
"args": {"query": "another"},
}
],
id="ai2",
),
ToolMessage(
content="result for another",
name="search_api",
tool_call_id="tool_call456",
id="00000000-0000-4000-8000-000000000020",
),
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",
)
},
{
"action": [
ToolMessage(
content="result for query",
name="search_api",
tool_call_id="tool_call123",
id="00000000-0000-4000-8000-000000000036",
)
]
},
{
"agent": AIMessage(
content="",
tool_calls=[
{
"id": "tool_call456",
"name": "search_api",
"args": {"query": "another"},
}
],
id="ai2",
)
},
{
"action": [
ToolMessage(
content="result for another",
name="search_api",
tool_call_id="tool_call456",
id="00000000-0000-4000-8000-000000000045",
)
]
},
{"agent": AIMessage(content="answer", id="ai3")},
]
app_w_interrupt = workflow.compile(
checkpointer=MemorySaverAssertImmutable(),
interrupt_after=["agent"],
)
config = {"configurable": {"thread_id": "1"}}
assert [
c for c in app_w_interrupt.stream(("human", "what is weather in sf"), config)
] == [
{
"agent": AIMessage(
content="",
tool_calls=[
{
"id": "tool_call123",
"name": "search_api",
"args": {"query": "query"},
}
],
id="ai1",
)
},
]
assert app_w_interrupt.get_state(config) == StateSnapshot(
values=[
HumanMessage(content="what is weather in sf", id=AnyStr()),
AIMessage(
content="",
tool_calls=[
{
"id": "tool_call123",
"name": "search_api",
"args": {"query": "query"},
}
],
id="ai1",
),
],
next=("action",),
config=app_w_interrupt.checkpointer.get_tuple(config).config,
metadata={"source": "loop", "step": 1},
)
# 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=[
HumanMessage(content="what is weather in sf", id=AnyStr()),
AIMessage(
content="",
id="ai1",
tool_calls=[
{
"id": "tool_call123",
"name": "search_api",
"args": {"query": "a different query"},
}
],
),
],
next=("action",),
config=next_config,
metadata={"source": "update", "step": 2},
)
assert [c for c in app_w_interrupt.stream(None, config)] == [
{
"action": [
ToolMessage(
content="result for a different query",
name="search_api",
tool_call_id="tool_call123",
id=AnyStr(),
)
]
},
{
"agent": AIMessage(
content="",
tool_calls=[
{
"id": "tool_call456",
"name": "search_api",
"args": {"query": "another"},
}
],
id="ai2",
)
},
]
assert app_w_interrupt.get_state(config) == StateSnapshot(
values=[
HumanMessage(
content="what is weather in sf",
id=AnyStr(),
),
AIMessage(
content="",
id="ai1",
tool_calls=[
{
"id": "tool_call123",
"name": "search_api",
"args": {"query": "a different query"},
}
],
),
ToolMessage(
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",
),
],
next=("action",),
config=app_w_interrupt.checkpointer.get_tuple(config).config,
metadata={"source": "loop", "step": 4},
)
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=[
HumanMessage(
content="what is weather in sf",
id=AnyStr(),
),
AIMessage(
content="",
id="ai1",
tool_calls=[
{
"id": "tool_call123",
"name": "search_api",
"args": {"query": "a different query"},
}
],
),
ToolMessage(
content="result for a different query",
name="search_api",
tool_call_id="tool_call123",
id=AnyStr(),
),
AIMessage(content="answer", id="ai2"),
],
next=(),
config=app_w_interrupt.checkpointer.get_tuple(config).config,
metadata={"source": "update", "step": 5},
)
app_w_interrupt = workflow.compile(
checkpointer=MemorySaverAssertImmutable(),
interrupt_before=["action"],
)
config = {"configurable": {"thread_id": "2"}}
model.i = 0 # reset the llm
assert [c for c in app_w_interrupt.stream("what is weather in sf", config)] == [
{
"agent": AIMessage(
content="",
tool_calls=[
{
"id": "tool_call123",
"name": "search_api",
"args": {"query": "query"},
}
],
id="ai1",
)
},
]
assert app_w_interrupt.get_state(config) == StateSnapshot(
values=[
HumanMessage(
content="what is weather in sf",
id=AnyStr(),
),
AIMessage(
content="",
tool_calls=[
{
"id": "tool_call123",
"name": "search_api",
"args": {"query": "query"},
}
],
id="ai1",
),
],
next=("action",),
config=app_w_interrupt.checkpointer.get_tuple(config).config,
metadata={"source": "loop", "step": 1},
)
# 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=[
HumanMessage(
content="what is weather in sf",
id=AnyStr(),
),
AIMessage(
content="",
id="ai1",
tool_calls=[
{
"id": "tool_call123",
"name": "search_api",
"args": {"query": "a different query"},
}
],
),
],
next=("action",),
config=app_w_interrupt.checkpointer.get_tuple(config).config,
metadata={"source": "update", "step": 2},
)
assert [c for c in app_w_interrupt.stream(None, config)] == [
{
"action": [
ToolMessage(
content="result for a different query",
name="search_api",
tool_call_id="tool_call123",
id=AnyStr(),
)
]
},
{
"agent": AIMessage(
content="",
tool_calls=[
{
"id": "tool_call456",
"name": "search_api",
"args": {"query": "another"},
}
],
id="ai2",
)
},
]
assert app_w_interrupt.get_state(config) == StateSnapshot(
values=[
HumanMessage(
content="what is weather in sf",
id=AnyStr(),
),
AIMessage(
content="",
id="ai1",
tool_calls=[
{
"id": "tool_call123",
"name": "search_api",
"args": {"query": "a different query"},
}
],
),
ToolMessage(
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",
),
],
next=("action",),
config=app_w_interrupt.checkpointer.get_tuple(config).config,
metadata={"source": "loop", "step": 4},
)
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=[
HumanMessage(
content="what is weather in sf",
id=AnyStr(),
),
AIMessage(
content="",
id="ai1",
tool_calls=[
{
"id": "tool_call123",
"name": "search_api",
"args": {"query": "a different query"},
}
],
),
ToolMessage(
content="result for a different query",
name="search_api",
tool_call_id="tool_call123",
id=AnyStr(),
),
AIMessage(content="answer", id="ai2"),
],
next=(),
config=app_w_interrupt.checkpointer.get_tuple(config).config,
metadata={"source": "update", "step": 5},
)
# add an extra message as if it came from "action" node
app_w_interrupt.update_state(config, ("ai", "an extra message"), as_node="action")
# extra message is coerced BaseMessge and appended
# now the next node is "agent" per the graph edges
assert app_w_interrupt.get_state(config) == StateSnapshot(
values=[
HumanMessage(
content="what is weather in sf",
id=AnyStr(),
),
AIMessage(
content="",
id="ai1",
tool_calls=[
{
"id": "tool_call123",
"name": "search_api",
"args": {"query": "a different query"},
}
],
),
ToolMessage(
content="result for a different query",
name="search_api",
tool_call_id="tool_call123",
id=AnyStr(),
),
AIMessage(content="answer", id="ai2"),
AIMessage(content="an extra message", id=AnyStr()),
],
next=("agent",),
config=app_w_interrupt.checkpointer.get_tuple(config).config,
metadata={"source": "update", "step": 6},
)
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:
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",
},
]
def test_start_branch_then(snapshot: SnapshotAssertion) -> None:
class State(TypedDict):
my_key: Annotated[str, operator.add]
market: str
# this graph is invalid because there is no path to END
invalid_graph = StateGraph(State)
invalid_graph.add_node("tool_two_slow", lambda s: {"my_key": "slow"})
invalid_graph.add_node("tool_two_fast", lambda s: {"my_key": "fast"})
invalid_graph.set_conditional_entry_point(
lambda s: "tool_two_slow" if s["market"] == "DE" else "tool_two_fast"
)
with pytest.raises(ValueError):
invalid_graph.compile()
tool_two_graph = StateGraph(State)
tool_two_graph.add_node("tool_two_slow", lambda s: {"my_key": " slow"})
tool_two_graph.add_node("tool_two_fast", lambda s: {"my_key": " fast"})
tool_two_graph.set_conditional_entry_point(
lambda s: "tool_two_slow" if s["market"] == "DE" else "tool_two_fast", then=END
)
tool_two = tool_two_graph.compile()
assert tool_two.get_graph().draw_mermaid() == snapshot
assert tool_two.invoke({"my_key": "value", "market": "DE"}) == {
"my_key": "value slow",
"market": "DE",
}
assert tool_two.invoke({"my_key": "value", "market": "US"}) == {
"my_key": "value fast",
"market": "US",
}
with SqliteSaver.from_conn_string(":memory:") as saver:
tool_two = tool_two_graph.compile(
checkpointer=saver, interrupt_before=["tool_two_fast", "tool_two_slow"]
)
# missing thread_id
with pytest.raises(ValueError, match="thread_id"):
tool_two.invoke({"my_key": "value", "market": "DE"})
thread1 = {"configurable": {"thread_id": "1"}}
# stop when about to enter node
assert tool_two.invoke({"my_key": "value", "market": "DE"}, thread1) == {
"my_key": "value",
"market": "DE",
}
assert tool_two.get_state(thread1) == StateSnapshot(
values={"my_key": "value", "market": "DE"},
next=("tool_two_slow",),
config=tool_two.checkpointer.get_tuple(thread1).config,
metadata={"source": "loop", "step": 0},
)
# resume, for same result as above
assert tool_two.invoke(None, thread1, debug=1) == {
"my_key": "value slow",
"market": "DE",
}
assert tool_two.get_state(thread1) == StateSnapshot(
values={"my_key": "value slow", "market": "DE"},
next=(),
config=tool_two.checkpointer.get_tuple(thread1).config,
metadata={"source": "loop", "step": 1},
)
thread2 = {"configurable": {"thread_id": "2"}}
# stop when about to enter node
assert tool_two.invoke({"my_key": "value", "market": "US"}, thread2) == {
"my_key": "value",
"market": "US",
}
assert tool_two.get_state(thread2) == StateSnapshot(
values={"my_key": "value", "market": "US"},
next=("tool_two_fast",),
config=tool_two.checkpointer.get_tuple(thread2).config,
metadata={"source": "loop", "step": 0},
)
# resume, for same result as above
assert tool_two.invoke(None, thread2, debug=1) == {
"my_key": "value fast",
"market": "US",
}
assert tool_two.get_state(thread2) == StateSnapshot(
values={"my_key": "value fast", "market": "US"},
next=(),
config=tool_two.checkpointer.get_tuple(thread2).config,
metadata={"source": "loop", "step": 1},
)
def test_branch_then(snapshot: SnapshotAssertion) -> None:
class State(TypedDict):
my_key: Annotated[str, operator.add]
market: str
# this graph is invalid because there is no path to "finish"
invalid_graph = StateGraph(State)
invalid_graph.set_entry_point("prepare")
invalid_graph.set_finish_point("finish")
invalid_graph.add_conditional_edges(
source="prepare",
path=lambda s: "tool_two_slow" if s["market"] == "DE" else "tool_two_fast",
path_map=["tool_two_slow", "tool_two_fast"],
)
invalid_graph.add_node("prepare", lambda s: {"my_key": " prepared"})
invalid_graph.add_node("tool_two_slow", lambda s: {"my_key": " slow"})
invalid_graph.add_node("tool_two_fast", lambda s: {"my_key": " fast"})
invalid_graph.add_node("finish", lambda s: {"my_key": " finished"})
with pytest.raises(ValueError):
invalid_graph.compile()
tool_two_graph = StateGraph(State)
tool_two_graph.set_entry_point("prepare")
tool_two_graph.set_finish_point("finish")
tool_two_graph.add_conditional_edges(
source="prepare",
path=lambda s: "tool_two_slow" if s["market"] == "DE" else "tool_two_fast",
then="finish",
)
tool_two_graph.add_node("prepare", lambda s: {"my_key": " prepared"})
tool_two_graph.add_node("tool_two_slow", lambda s: {"my_key": " slow"})
tool_two_graph.add_node("tool_two_fast", lambda s: {"my_key": " fast"})
tool_two_graph.add_node("finish", lambda s: {"my_key": " finished"})
tool_two = tool_two_graph.compile()
assert tool_two.get_graph().draw_mermaid(with_styles=False) == snapshot
assert tool_two.get_graph().draw_mermaid() == snapshot
assert tool_two.invoke({"my_key": "value", "market": "DE"}, debug=1) == {
"my_key": "value prepared slow finished",
"market": "DE",
}
assert tool_two.invoke({"my_key": "value", "market": "US"}) == {
"my_key": "value prepared fast finished",
"market": "US",
}
with SqliteSaver.from_conn_string(":memory:") as saver:
# test stream_mode=debug
tool_two = tool_two_graph.compile(checkpointer=saver)
thread10 = {"configurable": {"thread_id": "10"}}
assert [
*tool_two.stream(
{"my_key": "value", "market": "DE"}, thread10, stream_mode="debug"
)
] == [
{
"type": "checkpoint",
"timestamp": AnyStr(),
"step": 0,
"payload": {
"config": {
"configurable": {"thread_id": "10", "thread_ts": AnyStr()}
},
"values": {"my_key": "value", "market": "DE"},
},
},
{
"type": "task",
"timestamp": AnyStr(),
"step": 1,
"payload": {
"id": "e7879e70-6335-5867-9ec6-957fbb3da6fa",
"name": "prepare",
"input": {"my_key": "value", "market": "DE"},
"triggers": ["start:prepare"],
},
},
{
"type": "task_result",
"timestamp": AnyStr(),
"step": 1,
"payload": {
"id": "e7879e70-6335-5867-9ec6-957fbb3da6fa",
"name": "prepare",
"result": [("my_key", " prepared")],
},
},
{
"type": "checkpoint",
"timestamp": AnyStr(),
"step": 1,
"payload": {
"config": {
"configurable": {"thread_id": "10", "thread_ts": AnyStr()}
},
"values": {"my_key": "value prepared", "market": "DE"},
},
},
{
"type": "task",
"timestamp": AnyStr(),
"step": 2,
"payload": {
"id": "122f31bd-0e14-5b8f-91e7-4f241047a3fd",
"name": "tool_two_slow",
"input": {"my_key": "value prepared", "market": "DE"},
"triggers": ["branch:prepare:condition:tool_two_slow"],
},
},
{
"type": "task_result",
"timestamp": AnyStr(),
"step": 2,
"payload": {
"id": "122f31bd-0e14-5b8f-91e7-4f241047a3fd",
"name": "tool_two_slow",
"result": [("my_key", " slow")],
},
},
{
"type": "checkpoint",
"timestamp": AnyStr(),
"step": 2,
"payload": {
"config": {
"configurable": {"thread_id": "10", "thread_ts": AnyStr()}
},
"values": {"my_key": "value prepared slow", "market": "DE"},
},
},
{
"type": "task",
"timestamp": AnyStr(),
"step": 3,
"payload": {
"id": "48a16051-2c14-5ff5-9cfe-e8c7c32d5c83",
"name": "finish",
"input": {"my_key": "value prepared slow", "market": "DE"},
"triggers": ["branch:prepare:condition:then"],
},
},
{
"type": "task_result",
"timestamp": AnyStr(),
"step": 3,
"payload": {
"id": "48a16051-2c14-5ff5-9cfe-e8c7c32d5c83",
"name": "finish",
"result": [("my_key", " finished")],
},
},
{
"type": "checkpoint",
"timestamp": AnyStr(),
"step": 3,
"payload": {
"config": {
"configurable": {"thread_id": "10", "thread_ts": AnyStr()}
},
"values": {
"my_key": "value prepared slow finished",
"market": "DE",
},
},
},
]
tool_two = tool_two_graph.compile(
checkpointer=saver, interrupt_before=["tool_two_fast", "tool_two_slow"]
)
# missing thread_id
with pytest.raises(ValueError, match="thread_id"):
tool_two.invoke({"my_key": "value", "market": "DE"})
thread1 = {"configurable": {"thread_id": "1"}}
# stop when about to enter node
assert tool_two.invoke({"my_key": "value", "market": "DE"}, thread1) == {
"my_key": "value prepared",
"market": "DE",
}
assert tool_two.get_state(thread1) == StateSnapshot(
values={"my_key": "value prepared", "market": "DE"},
next=("tool_two_slow",),
config=tool_two.checkpointer.get_tuple(thread1).config,
metadata={"source": "loop", "step": 1},
)
# resume, for same result as above
assert tool_two.invoke(None, thread1, debug=1) == {
"my_key": "value prepared slow finished",
"market": "DE",
}
assert tool_two.get_state(thread1) == StateSnapshot(
values={"my_key": "value prepared slow finished", "market": "DE"},
next=(),
config=tool_two.checkpointer.get_tuple(thread1).config,
metadata={"source": "loop", "step": 3},
)
thread2 = {"configurable": {"thread_id": "2"}}
# stop when about to enter node
assert tool_two.invoke({"my_key": "value", "market": "US"}, thread2) == {
"my_key": "value prepared",
"market": "US",
}
assert tool_two.get_state(thread2) == StateSnapshot(
values={"my_key": "value prepared", "market": "US"},
next=("tool_two_fast",),
config=tool_two.checkpointer.get_tuple(thread2).config,
metadata={"source": "loop", "step": 1},
)
# resume, for same result as above
assert tool_two.invoke(None, thread2, debug=1) == {
"my_key": "value prepared fast finished",
"market": "US",
}
assert tool_two.get_state(thread2) == StateSnapshot(
values={"my_key": "value prepared fast finished", "market": "US"},
next=(),
config=tool_two.checkpointer.get_tuple(thread2).config,
metadata={"source": "loop", "step": 3},
)
with SqliteSaver.from_conn_string(":memory:") as saver:
tool_two = tool_two_graph.compile(
checkpointer=saver, interrupt_after=["prepare"]
)
# missing thread_id
with pytest.raises(ValueError, match="thread_id"):
tool_two.invoke({"my_key": "value", "market": "DE"})
thread1 = {"configurable": {"thread_id": "1"}}
# stop when about to enter node
assert tool_two.invoke({"my_key": "value", "market": "DE"}, thread1) == {
"my_key": "value prepared",
"market": "DE",
}
assert tool_two.get_state(thread1) == StateSnapshot(
values={"my_key": "value prepared", "market": "DE"},
next=("tool_two_slow",),
config=tool_two.checkpointer.get_tuple(thread1).config,
metadata={"source": "loop", "step": 1},
)
# resume, for same result as above
assert tool_two.invoke(None, thread1, debug=1) == {
"my_key": "value prepared slow finished",
"market": "DE",
}
assert tool_two.get_state(thread1) == StateSnapshot(
values={"my_key": "value prepared slow finished", "market": "DE"},
next=(),
config=tool_two.checkpointer.get_tuple(thread1).config,
metadata={"source": "loop", "step": 3},
)
thread2 = {"configurable": {"thread_id": "2"}}
# stop when about to enter node
assert tool_two.invoke({"my_key": "value", "market": "US"}, thread2) == {
"my_key": "value prepared",
"market": "US",
}
assert tool_two.get_state(thread2) == StateSnapshot(
values={"my_key": "value prepared", "market": "US"},
next=("tool_two_fast",),
config=tool_two.checkpointer.get_tuple(thread2).config,
metadata={"source": "loop", "step": 1},
)
# resume, for same result as above
assert tool_two.invoke(None, thread2, debug=1) == {
"my_key": "value prepared fast finished",
"market": "US",
}
assert tool_two.get_state(thread2) == StateSnapshot(
values={"my_key": "value prepared fast finished", "market": "US"},
next=(),
config=tool_two.checkpointer.get_tuple(thread2).config,
metadata={"source": "loop", "step": 3},
)
def test_in_one_fan_out_state_graph_waiting_edge(snapshot: SnapshotAssertion) -> None:
def sorted_add(
x: list[str], y: Union[list[str], list[tuple[str, str]]]
) -> list[str]:
if isinstance(y[0], tuple):
for rem, _ in y:
x.remove(rem)
y = [t[1] for t in y]
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 analyzer_one(data: State) -> State:
return {"query": f'analyzed: {data["query"]}'}
def retriever_one(data: State) -> State:
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("analyzer_one", analyzer_one)
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", "analyzer_one")
workflow.add_edge("analyzer_one", "retriever_one")
workflow.add_edge("rewrite_query", "retriever_two")
workflow.add_edge(["retriever_one", "retriever_two"], "qa")
workflow.set_finish_point("qa")
app = workflow.compile()
assert app.get_graph().draw_mermaid(with_styles=False) == snapshot
assert app.invoke({"query": "what is weather in sf"}) == {
"query": "analyzed: 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"}},
{
"analyzer_one": {"query": "analyzed: query: what is weather in sf"},
"retriever_two": {"docs": ["doc3", "doc4"]},
},
{"retriever_one": {"docs": ["doc1", "doc2"]}},
{"qa": {"answer": "doc1,doc2,doc3,doc4"}},
]
app_w_interrupt = workflow.compile(
checkpointer=MemorySaverAssertImmutable(),
interrupt_after=["retriever_one"],
)
config = {"configurable": {"thread_id": "1"}}
assert [
c for c in app_w_interrupt.stream({"query": "what is weather in sf"}, config)
] == [
{"rewrite_query": {"query": "query: what is weather in sf"}},
{
"analyzer_one": {"query": "analyzed: query: what is weather in sf"},
"retriever_two": {"docs": ["doc3", "doc4"]},
},
{"retriever_one": {"docs": ["doc1", "doc2"]}},
]
assert [c for c in app_w_interrupt.stream(None, config)] == [
{"qa": {"answer": "doc1,doc2,doc3,doc4"}},
]
def test_in_one_fan_out_state_graph_waiting_edge_via_branch(
snapshot: SnapshotAssertion,
) -> None:
def sorted_add(
x: list[str], y: Union[list[str], list[tuple[str, str]]]
) -> list[str]:
if isinstance(y[0], tuple):
for rem, _ in y:
x.remove(rem)
y = [t[1] for t in y]
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 analyzer_one(data: State) -> State:
return {"query": f'analyzed: {data["query"]}'}
def retriever_one(data: State) -> State:
return {"docs": ["doc1", "doc2"]}
def retriever_two(data: State) -> State:
return {"docs": ["doc3", "doc4"]}
def qa(data: State) -> State:
return {"answer": ",".join(data["docs"])}
def rewrite_query_then(data: State) -> Literal["retriever_two"]:
return "retriever_two"
workflow = StateGraph(State)
workflow.add_node("rewrite_query", rewrite_query)
workflow.add_node("analyzer_one", analyzer_one)
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", "analyzer_one")
workflow.add_edge("analyzer_one", "retriever_one")
workflow.add_conditional_edges("rewrite_query", rewrite_query_then)
workflow.add_edge(["retriever_one", "retriever_two"], "qa")
workflow.set_finish_point("qa")
app = workflow.compile()
assert app.get_graph().draw_mermaid(with_styles=False) == snapshot
assert app.invoke({"query": "what is weather in sf"}, debug=True) == {
"query": "analyzed: 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"}},
{
"analyzer_one": {"query": "analyzed: query: what is weather in sf"},
"retriever_two": {"docs": ["doc3", "doc4"]},
},
{"retriever_one": {"docs": ["doc1", "doc2"]}},
{"qa": {"answer": "doc1,doc2,doc3,doc4"}},
]
app_w_interrupt = workflow.compile(
checkpointer=MemorySaverAssertImmutable(),
interrupt_after=["retriever_one"],
)
config = {"configurable": {"thread_id": "1"}}
assert [
c for c in app_w_interrupt.stream({"query": "what is weather in sf"}, config)
] == [
{"rewrite_query": {"query": "query: what is weather in sf"}},
{
"analyzer_one": {"query": "analyzed: query: what is weather in sf"},
"retriever_two": {"docs": ["doc3", "doc4"]},
},
{"retriever_one": {"docs": ["doc1", "doc2"]}},
]
assert [c for c in app_w_interrupt.stream(None, config)] == [
{"qa": {"answer": "doc1,doc2,doc3,doc4"}},
]
def test_in_one_fan_out_state_graph_waiting_edge_custom_state_class(
snapshot: SnapshotAssertion,
) -> None:
from langchain_core.pydantic_v1 import BaseModel, ValidationError
def sorted_add(
x: list[str], y: Union[list[str], list[tuple[str, str]]]
) -> list[str]:
if isinstance(y[0], tuple):
for rem, _ in y:
x.remove(rem)
y = [t[1] for t in y]
return sorted(operator.add(x, y))
class State(BaseModel):
query: str
answer: Optional[str] = None
docs: Annotated[list[str], sorted_add]
def rewrite_query(data: State) -> State:
return {"query": f"query: {data.query}"}
def analyzer_one(data: State) -> State:
return {"query": f"analyzed: {data.query}"}
def retriever_one(data: State) -> State:
return {"docs": ["doc1", "doc2"]}
def retriever_two(data: State) -> State:
return {"docs": ["doc3", "doc4"]}
def qa(data: State) -> State:
return {"answer": ",".join(data.docs)}
def decider(data: State) -> str:
assert isinstance(data, State)
return "retriever_two"
workflow = StateGraph(State)
workflow.add_node("rewrite_query", rewrite_query)
workflow.add_node("analyzer_one", analyzer_one)
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", "analyzer_one")
workflow.add_edge("analyzer_one", "retriever_one")
workflow.add_conditional_edges(
"rewrite_query", decider, {"retriever_two": "retriever_two"}
)
workflow.add_edge(["retriever_one", "retriever_two"], "qa")
workflow.set_finish_point("qa")
app = workflow.compile()
assert app.get_graph().draw_mermaid(with_styles=False) == snapshot
with pytest.raises(ValidationError):
app.invoke({"query": {}})
assert app.invoke({"query": "what is weather in sf"}) == {
"query": "analyzed: 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"}},
{
"analyzer_one": {"query": "analyzed: query: what is weather in sf"},
"retriever_two": {"docs": ["doc3", "doc4"]},
},
{"retriever_one": {"docs": ["doc1", "doc2"]}},
{"qa": {"answer": "doc1,doc2,doc3,doc4"}},
]
app_w_interrupt = workflow.compile(
checkpointer=MemorySaverAssertImmutable(),
interrupt_after=["retriever_one"],
)
config = {"configurable": {"thread_id": "1"}}
assert [
c for c in app_w_interrupt.stream({"query": "what is weather in sf"}, config)
] == [
{"rewrite_query": {"query": "query: what is weather in sf"}},
{
"analyzer_one": {"query": "analyzed: query: what is weather in sf"},
"retriever_two": {"docs": ["doc3", "doc4"]},
},
{"retriever_one": {"docs": ["doc1", "doc2"]}},
]
assert [c for c in app_w_interrupt.stream(None, config)] == [
{"qa": {"answer": "doc1,doc2,doc3,doc4"}},
]
def test_in_one_fan_out_state_graph_waiting_edge_plus_regular() -> None:
def sorted_add(
x: list[str], y: Union[list[str], list[tuple[str, str]]]
) -> list[str]:
if isinstance(y[0], tuple):
for rem, _ in y:
x.remove(rem)
y = [t[1] for t in y]
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 analyzer_one(data: State) -> State:
return {"query": f'analyzed: {data["query"]}'}
def retriever_one(data: State) -> State:
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("analyzer_one", analyzer_one)
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", "analyzer_one")
workflow.add_edge("analyzer_one", "retriever_one")
workflow.add_edge("rewrite_query", "retriever_two")
workflow.add_edge(["retriever_one", "retriever_two"], "qa")
workflow.set_finish_point("qa")
# silly edge, to make sure having been triggered before doesn't break
# semantics of named barrier (== waiting edges)
workflow.add_edge("rewrite_query", "qa")
app = workflow.compile()
assert app.invoke({"query": "what is weather in sf"}) == {
"query": "analyzed: 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"}},
{
"analyzer_one": {"query": "analyzed: query: what is weather in sf"},
"retriever_two": {"docs": ["doc3", "doc4"]},
"qa": {"answer": ""},
},
{"retriever_one": {"docs": ["doc1", "doc2"]}},
{"qa": {"answer": "doc1,doc2,doc3,doc4"}},
]
app_w_interrupt = workflow.compile(
checkpointer=MemorySaverAssertImmutable(),
interrupt_after=["retriever_one"],
)
config = {"configurable": {"thread_id": "1"}}
assert [
c for c in app_w_interrupt.stream({"query": "what is weather in sf"}, config)
] == [
{"rewrite_query": {"query": "query: what is weather in sf"}},
{
"analyzer_one": {"query": "analyzed: query: what is weather in sf"},
"retriever_two": {"docs": ["doc3", "doc4"]},
"qa": {"answer": ""},
},
{"retriever_one": {"docs": ["doc1", "doc2"]}},
]
assert [c for c in app_w_interrupt.stream(None, config)] == [
{"qa": {"answer": "doc1,doc2,doc3,doc4"}},
]
def test_in_one_fan_out_state_graph_waiting_edge_multiple() -> None:
def sorted_add(
x: list[str], y: Union[list[str], list[tuple[str, str]]]
) -> list[str]:
if isinstance(y[0], tuple):
for rem, _ in y:
x.remove(rem)
y = [t[1] for t in y]
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 analyzer_one(data: State) -> State:
return {"query": f'analyzed: {data["query"]}'}
def retriever_one(data: State) -> State:
return {"docs": ["doc1", "doc2"]}
def retriever_two(data: State) -> State:
return {"docs": ["doc3", "doc4"]}
def qa(data: State) -> State:
return {"answer": ",".join(data["docs"])}
def decider(data: State) -> None:
return None
def decider_cond(data: State) -> str:
if data["query"].count("analyzed") > 1:
return "qa"
else:
return "rewrite_query"
workflow = StateGraph(State)
workflow.add_node("rewrite_query", rewrite_query)
workflow.add_node("analyzer_one", analyzer_one)
workflow.add_node("retriever_one", retriever_one)
workflow.add_node("retriever_two", retriever_two)
workflow.add_node("decider", decider)
workflow.add_node("qa", qa)
workflow.set_entry_point("rewrite_query")
workflow.add_edge("rewrite_query", "analyzer_one")
workflow.add_edge("analyzer_one", "retriever_one")
workflow.add_edge("rewrite_query", "retriever_two")
workflow.add_edge(["retriever_one", "retriever_two"], "decider")
workflow.add_conditional_edges("decider", decider_cond)
workflow.set_finish_point("qa")
app = workflow.compile()
assert app.invoke({"query": "what is weather in sf"}) == {
"query": "analyzed: query: analyzed: query: what is weather in sf",
"answer": "doc1,doc1,doc2,doc2,doc3,doc3,doc4,doc4",
"docs": ["doc1", "doc1", "doc2", "doc2", "doc3", "doc3", "doc4", "doc4"],
}
assert [*app.stream({"query": "what is weather in sf"})] == [
{"rewrite_query": {"query": "query: what is weather in sf"}},
{
"analyzer_one": {"query": "analyzed: query: what is weather in sf"},
"retriever_two": {"docs": ["doc3", "doc4"]},
},
{"retriever_one": {"docs": ["doc1", "doc2"]}},
{"rewrite_query": {"query": "query: analyzed: query: what is weather in sf"}},
{
"analyzer_one": {
"query": "analyzed: query: analyzed: query: what is weather in sf"
},
"retriever_two": {"docs": ["doc3", "doc4"]},
},
{
"retriever_one": {"docs": ["doc1", "doc2"]},
},
{"qa": {"answer": "doc1,doc1,doc2,doc2,doc3,doc3,doc4,doc4"}},
]
def test_in_one_fan_out_state_graph_waiting_edge_multiple_cond_edge() -> None:
def sorted_add(
x: list[str], y: Union[list[str], list[tuple[str, str]]]
) -> list[str]:
if isinstance(y[0], tuple):
for rem, _ in y:
x.remove(rem)
y = [t[1] for t in y]
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_picker(data: State) -> list[str]:
return ["analyzer_one", "retriever_two"]
def analyzer_one(data: State) -> State:
return {"query": f'analyzed: {data["query"]}'}
def retriever_one(data: State) -> State:
return {"docs": ["doc1", "doc2"]}
def retriever_two(data: State) -> State:
return {"docs": ["doc3", "doc4"]}
def qa(data: State) -> State:
return {"answer": ",".join(data["docs"])}
def decider(data: State) -> None:
return None
def decider_cond(data: State) -> str:
if data["query"].count("analyzed") > 1:
return "qa"
else:
return "rewrite_query"
workflow = StateGraph(State)
workflow.add_node("rewrite_query", rewrite_query)
workflow.add_node("analyzer_one", analyzer_one)
workflow.add_node("retriever_one", retriever_one)
workflow.add_node("retriever_two", retriever_two)
workflow.add_node("decider", decider)
workflow.add_node("qa", qa)
workflow.set_entry_point("rewrite_query")
workflow.add_conditional_edges("rewrite_query", retriever_picker)
workflow.add_edge("analyzer_one", "retriever_one")
workflow.add_edge(["retriever_one", "retriever_two"], "decider")
workflow.add_conditional_edges("decider", decider_cond)
workflow.set_finish_point("qa")
app = workflow.compile()
assert app.invoke({"query": "what is weather in sf"}) == {
"query": "analyzed: query: analyzed: query: what is weather in sf",
"answer": "doc1,doc1,doc2,doc2,doc3,doc3,doc4,doc4",
"docs": ["doc1", "doc1", "doc2", "doc2", "doc3", "doc3", "doc4", "doc4"],
}
assert [*app.stream({"query": "what is weather in sf"})] == [
{"rewrite_query": {"query": "query: what is weather in sf"}},
{
"analyzer_one": {"query": "analyzed: query: what is weather in sf"},
"retriever_two": {"docs": ["doc3", "doc4"]},
},
{"retriever_one": {"docs": ["doc1", "doc2"]}},
{"rewrite_query": {"query": "query: analyzed: query: what is weather in sf"}},
{
"analyzer_one": {
"query": "analyzed: query: analyzed: query: what is weather in sf"
},
"retriever_two": {"docs": ["doc3", "doc4"]},
},
{
"retriever_one": {"docs": ["doc1", "doc2"]},
},
{"qa": {"answer": "doc1,doc1,doc2,doc2,doc3,doc3,doc4,doc4"}},
]
def test_simple_multi_edge(snapshot: SnapshotAssertion) -> None:
class State(TypedDict):
my_key: Annotated[str, operator.add]
def up(state: State):
pass
def side(state: State):
pass
def down(state: State):
pass
graph = StateGraph(State)
graph.add_node("up", up)
graph.add_node("side", side)
graph.add_node("down", down)
graph.set_entry_point("up")
graph.add_edge("up", "side")
graph.add_edge(["up", "side"], "down")
graph.set_finish_point("down")
app = graph.compile()
assert app.get_graph().draw_mermaid(with_styles=False) == snapshot
assert app.invoke({"my_key": "my_value"}) == {"my_key": "my_value"}
def test_nested_graph_xray(snapshot: SnapshotAssertion) -> None:
class State(TypedDict):
my_key: Annotated[str, operator.add]
market: str
def logic(state: State):
pass
tool_two_graph = StateGraph(State)
tool_two_graph.add_node("tool_two_slow", logic)
tool_two_graph.add_node("tool_two_fast", logic)
tool_two_graph.set_conditional_entry_point(
lambda s: "tool_two_slow" if s["market"] == "DE" else "tool_two_fast",
then=END,
)
tool_two = tool_two_graph.compile()
graph = StateGraph(State)
graph.add_node("tool_one", logic)
graph.add_node("tool_two", tool_two)
graph.add_node("tool_three", logic)
graph.set_conditional_entry_point(lambda s: "tool_one", then=END)
app = graph.compile()
assert app.get_graph(xray=True).to_json() == snapshot
assert app.get_graph(xray=True).draw_mermaid() == snapshot
def test_nested_graph(snapshot: SnapshotAssertion) -> None:
def never_called_fn(state: Any):
assert 0, "This function should never be called"
never_called = RunnableLambda(never_called_fn)
class InnerState(TypedDict):
my_key: str
my_other_key: str
def up(state: InnerState):
return {"my_key": state["my_key"] + " there", "my_other_key": state["my_key"]}
inner = StateGraph(InnerState)
inner.add_node("up", up)
inner.set_entry_point("up")
inner.set_finish_point("up")
class State(TypedDict):
my_key: str
never_called: Any
def side(state: State):
return {"my_key": state["my_key"] + " and back again"}
graph = StateGraph(State)
graph.add_node("inner", inner.compile())
graph.add_node("side", side)
graph.set_entry_point("inner")
graph.add_edge("inner", "side")
graph.set_finish_point("side")
app = graph.compile()
assert app.get_graph().draw_mermaid(with_styles=False) == snapshot
assert app.get_graph(xray=True).draw_mermaid() == snapshot
assert app.invoke(
{"my_key": "my value", "never_called": never_called}, debug=True
) == {
"my_key": "my value there and back again",
"never_called": never_called,
}
assert [*app.stream({"my_key": "my value", "never_called": never_called})] == [
{"inner": {"my_key": "my value there"}},
{"side": {"my_key": "my value there and back again"}},
]
assert [
*app.stream(
{"my_key": "my value", "never_called": never_called}, stream_mode="values"
)
] == [
{
"my_key": "my value",
"never_called": never_called,
},
{
"my_key": "my value there",
"never_called": never_called,
},
{
"my_key": "my value there and back again",
"never_called": never_called,
},
]
chain = app | RunnablePassthrough()
assert chain.invoke({"my_key": "my value", "never_called": never_called}) == {
"my_key": "my value there and back again",
"never_called": never_called,
}
assert [*chain.stream({"my_key": "my value", "never_called": never_called})] == [
{"inner": {"my_key": "my value there"}},
{"side": {"my_key": "my value there and back again"}},
]
def test_repeat_condition(snapshot: SnapshotAssertion) -> None:
class AgentState(TypedDict):
hello: str
def router(state: AgentState) -> str:
return "hmm"
workflow = StateGraph(AgentState)
workflow.add_node("Researcher", lambda x: x)
workflow.add_node("Chart Generator", lambda x: x)
workflow.add_node("Call Tool", lambda x: x)
workflow.add_conditional_edges(
"Researcher",
router,
{"continue": "Chart Generator", "call_tool": "Call Tool", "end": END},
)
workflow.add_conditional_edges(
"Chart Generator",
router,
{"continue": "Researcher", "call_tool": "Call Tool", "end": END},
)
workflow.add_conditional_edges(
"Call Tool",
# Each agent node updates the 'sender' field
# the tool calling node does not, meaning
# this edge will route back to the original agent
# who invoked the tool
lambda x: x["sender"],
{
"Researcher": "Researcher",
"Chart Generator": "Chart Generator",
},
)
workflow.set_entry_point("Researcher")
app = workflow.compile()
assert app.get_graph().draw_mermaid(with_styles=False) == snapshot
Reference in New Issue View Git Blame Copy Permalink