WEBDIENSTE/langgraph
1
0
Fork 0
mirror of https://github.com/langchain-ai/langgraph.git synced 2026-08-20 06:35:46 +02:00
Code Issues Packages Projects Releases Wiki Activity
Files
a3b8d0694de458b370a0e44e743f658aaddb77cf
langgraph/tests/test_pregel_async.py
T

3319 lines
107 KiB
Python
Raw Blame History

import asyncio
import json
import operator
from contextlib import asynccontextmanager, contextmanager
from typing import (
Annotated,
Any,
AsyncGenerator,
AsyncIterator,
Generator,
Optional,
TypedDict,
Union,
)
from uuid import UUID
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.aiosqlite import AsyncSqliteSaver
from langgraph.checkpoint.base import CheckpointAt
from langgraph.graph import END, Graph, StateGraph
from langgraph.graph.message import MessageGraph
from langgraph.prebuilt.chat_agent_executor import (
create_function_calling_executor,
create_tool_calling_executor,
)
from langgraph.prebuilt.tool_executor import ToolExecutor
from langgraph.pregel import Channel, GraphRecursionError, Pregel, StateSnapshot
from tests.any_str import AnyStr
from tests.memory_assert import MemorySaverAssertImmutable
async 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"}
assert await app.ainvoke(2) == 3
assert await app.ainvoke(2, output_keys=["output"]) == {"output": 3}
assert await gapp.ainvoke(2) == 3
@pytest.mark.parametrize(
"falsy_value",
[None, False, 0, "", [], {}, set(), frozenset(), 0.0, 0j],
)
async 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 falsy_value == await gapp.ainvoke(1)
async def test_invoke_single_process_in_out_implicit_channels(
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})
assert app.input_schema.schema() == {"title": "LangGraphInput"}
assert app.output_schema.schema() == {"title": "LangGraphOutput"}
assert await app.ainvoke(2) == 3
async 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}, output_channels=["output", "fixed", "output_plus_one"]
)
assert app.input_schema.schema() == {"title": "LangGraphInput"}
assert app.output_schema.schema() == {
"title": "LangGraphOutput",
"type": "object",
"properties": {
"output": {"title": "Output"},
"fixed": {"title": "Fixed"},
"output_plus_one": {"title": "Output Plus One"},
},
}
assert await app.ainvoke(2) == {"output": 3, "fixed": 5, "output_plus_one": 4}
async 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},
output_channels=["output"],
)
assert app.input_schema.schema() == {"title": "LangGraphInput"}
assert app.output_schema.schema() == {
"title": "LangGraphOutput",
"type": "object",
"properties": {"output": {"title": "Output"}},
}
assert await app.ainvoke(2) == {"output": 3}
async 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,
},
input_channels=["input"],
output_channels=["output"],
)
assert app.input_schema.schema() == {
"title": "LangGraphInput",
"type": "object",
"properties": {"input": {"title": "Input"}},
}
assert app.output_schema.schema() == {
"title": "LangGraphOutput",
"type": "object",
"properties": {"output": {"title": "Output"}},
}
assert await app.ainvoke({"input": 2}) == {"output": 3}
async 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}, stream_channels=["inbox", "output"])
assert await app.ainvoke(2) == 4
assert await app.ainvoke(2, input_keys="inbox") == 3
with pytest.raises(GraphRecursionError):
await app.ainvoke(2, {"recursion_limit": 1})
step = 0
async for values in app.astream(2):
step += 1
if step == 1:
assert values == {
"inbox": 3,
}
elif step == 2:
assert values == {
"inbox": 3,
"output": 4,
}
assert step == 2
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 await gapp.ainvoke(2) == 4
step = 0
async for values in gapp.astream(2):
step += 1
if step == 1:
assert values == {
"add_one": 3,
}
elif step == 2:
assert values == {
"add_one_more": 4,
}
assert step == 2
@pytest.mark.parametrize(
"checkpoint_at", [CheckpointAt.END_OF_RUN, CheckpointAt.END_OF_STEP]
)
async def test_invoke_two_processes_in_out_interrupt(
mocker: MockerFixture, checkpoint_at: CheckpointAt
) -> 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(at=checkpoint_at)
app = Pregel(
nodes={"one": one, "two": two},
checkpointer=memory,
interrupt_after_nodes=["one"],
)
# start execution, stop at inbox
assert await app.ainvoke(2, {"configurable": {"thread_id": 1}}) is None
# inbox == 3
checkpoint = await memory.aget({"configurable": {"thread_id": 1}})
assert checkpoint is not None
assert checkpoint["channel_values"]["inbox"] == 3
# resume execution, finish
assert await app.ainvoke(None, {"configurable": {"thread_id": 1}}) == 4
# start execution again, stop at inbox
assert await app.ainvoke(20, {"configurable": {"thread_id": 1}}) is None
# inbox == 21
checkpoint = await memory.aget({"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 await app.ainvoke(3, {"configurable": {"thread_id": 1}}) is None
assert await app.ainvoke(None, {"configurable": {"thread_id": 1}}) == 5
# start execution again, stopping at inbox
assert await app.ainvoke(20, {"configurable": {"thread_id": 2}}) is None
# inbox == 21
snapshot = await app.aget_state({"configurable": {"thread_id": 2}})
assert snapshot.values["inbox"] == 21
assert snapshot.next == ("two",)
# update the state, resume
await app.aupdate_state({"configurable": {"thread_id": 2}}, 25, as_node="one")
assert await app.ainvoke(None, {"configurable": {"thread_id": 2}}) == 26
# no pending tasks
snapshot = await app.aget_state({"configurable": {"thread_id": 2}})
assert snapshot.next == ()
async 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).abatch
| Channel.write_to("output").abatch
)
app = Pregel(
nodes={"one": one, "two": two},
channels={"inbox": Topic(int)},
input_channels=["input", "inbox"],
stream_channels=["inbox", "output"],
)
# [12 + 1, 2 + 1 + 1]
assert [
c
async for c in app.astream(
{"input": 2, "inbox": 12}, output_keys="output", stream_mode="updates"
)
] == [
{"two": 13},
{"two": 4},
]
assert [
c async for c in app.astream({"input": 2, "inbox": 12}, output_keys="output")
] == [13, 4]
assert [
c async for c in app.astream({"input": 2, "inbox": 12}, stream_mode="updates")
] == [
{"one": {"inbox": 3}, "two": {"output": 13}},
{"two": {"output": 4}},
]
assert [c async for c in app.astream({"input": 2, "inbox": 12})] == [
{"inbox": [3], "output": 13},
{"inbox": [], "output": 4},
]
async def test_batch_two_processes_in_out() -> None:
async def add_one_with_delay(inp: int) -> int:
await asyncio.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)},
)
assert await app.abatch([3, 2, 1, 3, 5]) == [5, 4, 3, 5, 7]
assert await app.abatch([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 await gapp.abatch([3, 2, 1, 3, 5]) == [5, 4, 3, 5, 7]
async 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)
# No state is left over from previous invocations
for _ in range(10):
assert await app.ainvoke(2, {"recursion_limit": test_size}) == 2 + test_size
# Concurrent invocations do not interfere with each other
assert await asyncio.gather(
*(app.ainvoke(2, {"recursion_limit": test_size}) for _ in range(10))
) == [2 + test_size for _ in range(10)]
async 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)
# No state is left over from previous invocations
for _ in range(3):
# Then invoke pubsub
assert await app.abatch([2, 1, 3, 4, 5], {"recursion_limit": test_size}) == [
2 + test_size,
1 + test_size,
3 + test_size,
4 + test_size,
5 + test_size,
]
# Concurrent invocations do not interfere with each other
assert await asyncio.gather(
*(app.abatch([2, 1, 3, 4, 5], {"recursion_limit": test_size}) for _ in range(3))
) == [
[2 + test_size, 1 + test_size, 3 + test_size, 4 + test_size, 5 + test_size]
for _ in range(3)
]
async 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})
with pytest.raises(InvalidUpdateError):
# LastValue channels can only be updated once per iteration
await app.ainvoke(2)
async 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={"output": Topic(int)},
)
# An Topic channel accumulates updates into a sequence
assert await app.ainvoke(2) == [3, 3]
@pytest.mark.parametrize(
"checkpoint_at", [CheckpointAt.END_OF_RUN, CheckpointAt.END_OF_STEP]
)
async def test_invoke_checkpoint(
mocker: MockerFixture, checkpoint_at: CheckpointAt
) -> 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(at=checkpoint_at)
app = Pregel(
nodes={"one": one},
channels={"total": BinaryOperatorAggregate(int, operator.add)},
checkpointer=memory,
)
# total starts out as 0, so output is 0+2=2
assert await app.ainvoke(2, {"configurable": {"thread_id": "1"}}) == 2
checkpoint = await memory.aget({"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 await app.ainvoke(3, {"configurable": {"thread_id": "1"}}) == 5
checkpoint = await memory.aget({"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):
await app.ainvoke(4, {"configurable": {"thread_id": "1"}})
# checkpoint is not updated
checkpoint = await memory.aget({"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 await app.ainvoke(5, {"configurable": {"thread_id": "2"}}) == 5
checkpoint = await memory.aget({"configurable": {"thread_id": "1"}})
assert checkpoint is not None
assert checkpoint["channel_values"].get("total") == 7
checkpoint = await memory.aget({"configurable": {"thread_id": "2"}})
assert checkpoint is not None
assert checkpoint["channel_values"].get("total") == 5
@pytest.mark.parametrize(
"checkpoint_at", [CheckpointAt.END_OF_RUN, CheckpointAt.END_OF_STEP]
)
async def test_invoke_checkpoint_aiosqlite(
mocker: MockerFixture, checkpoint_at: CheckpointAt
) -> 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
)
async with AsyncSqliteSaver.from_conn_string(":memory:") as memory:
memory.at = checkpoint_at
app = Pregel(
nodes={"one": one},
channels={"total": BinaryOperatorAggregate(int, operator.add)},
checkpointer=memory,
debug=True,
)
thread_1 = {"configurable": {"thread_id": "1"}}
# total starts out as 0, so output is 0+2=2
assert await app.ainvoke(2, thread_1) == 2
state = await app.aget_state(thread_1)
assert state is not None
assert state.values.get("total") == 2
assert (
state.config["configurable"]["thread_ts"]
== (await memory.aget(thread_1))["ts"]
)
# total is now 2, so output is 2+3=5
assert await app.ainvoke(3, thread_1) == 5
state = await app.aget_state(thread_1)
assert state is not None
assert state.values.get("total") == 7
assert (
state.config["configurable"]["thread_ts"]
== (await memory.aget(thread_1))["ts"]
)
# total is now 2+5=7, so output would be 7+4=11, but raises ValueError
with pytest.raises(ValueError):
await app.ainvoke(4, thread_1)
# checkpoint is not updated
state = await app.aget_state(thread_1)
assert state is not None
assert state.values.get("total") == 7
thread_2 = {"configurable": {"thread_id": "2"}}
# on a new thread, total starts out as 0, so output is 0+5=5
assert await app.ainvoke(5, thread_2) == 5
state = await app.aget_state({"configurable": {"thread_id": "1"}})
assert state is not None
assert state.values.get("total") == 7
assert state.next == ()
state = await app.aget_state(thread_2)
assert state is not None
assert state.values.get("total") == 5
assert state.next == ()
# list all checkpoints for thread 1
thread_1_history = [c async for c in app.aget_state_history(thread_1)]
# there are 2: one for each successful ainvoke()
assert len(thread_1_history) == 2
# sorted descending
assert (
thread_1_history[0].config["configurable"]["thread_ts"]
> thread_1_history[1].config["configurable"]["thread_ts"]
)
# the second checkpoint
assert thread_1_history[0].values["total"] == 7
# the first checkpoint
assert thread_1_history[1].values["total"] == 2
# can get each checkpoint using aget with config
assert (await memory.aget(thread_1_history[0].config))[
"ts"
] == thread_1_history[0].config["configurable"]["thread_ts"]
assert (await memory.aget(thread_1_history[1].config))[
"ts"
] == thread_1_history[1].config["configurable"]["thread_ts"]
thread_1_next_config = await app.aupdate_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([h async for h in app.aget_state_history(thread_1)]) == 3
# the latest checkpoint is the updated one
assert await app.aget_state(thread_1) == await app.aget_state(
thread_1_next_config
)
async 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)},
)
# 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 await app.ainvoke(2) == [13, 13]
assert await asyncio.gather(*(app.ainvoke(2) for _ in range(100))) == [
[13, 13] for _ in range(100)
]
async def test_invoke_join_then_call_other_pubsub(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")
}
)
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),
},
)
# Then invoke pubsub
for _ in range(10):
assert await app.ainvoke([2, 3]) == 27
assert await asyncio.gather(*(app.ainvoke([2, 3]) for _ in range(10))) == [
27 for _ in range(10)
]
async 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}, stream_channels=["output", "between"])
# Then invoke pubsub
assert [c async for c in app.astream(2)] == [
{"between": 3, "output": 3},
{"between": 3, "output": 4},
]
async 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})
# It finishes executing (once no more messages being published)
# but returns nothing, as nothing was published to "output" topic
assert await app.ainvoke(2) is None
async def test_channel_enter_exit_timing(mocker: MockerFixture) -> None:
setup_sync = mocker.Mock()
cleanup_sync = mocker.Mock()
setup_async = mocker.Mock()
cleanup_async = mocker.Mock()
@contextmanager
def an_int() -> Generator[int, None, None]:
setup_sync()
try:
yield 5
finally:
cleanup_sync()
@asynccontextmanager
async def an_int_async() -> AsyncGenerator[int, None]:
setup_async()
try:
yield 5
finally:
cleanup_async()
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).abatch
| Channel.write_to("output").abatch
)
app = Pregel(
nodes={"one": one, "two": two},
channels={
"inbox": Topic(int),
"ctx": Context(an_int, an_int_async, typ=int),
},
output_channels=["inbox", "output"],
stream_channels=["inbox", "output"],
)
async def aenumerate(aiter: AsyncIterator[Any]) -> AsyncIterator[tuple[int, Any]]:
i = 0
async for chunk in aiter:
yield i, chunk
i += 1
assert setup_sync.call_count == 0
assert cleanup_sync.call_count == 0
assert setup_async.call_count == 0
assert cleanup_async.call_count == 0
async for i, chunk in aenumerate(app.astream(2)):
assert setup_sync.call_count == 0, "Sync context manager should not be used"
assert cleanup_sync.call_count == 0, "Sync context manager should not be used"
assert setup_async.call_count == 1, "Expected setup to be called once"
assert cleanup_async.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 setup_sync.call_count == 0
assert cleanup_sync.call_count == 0
assert setup_async.call_count == 1, "Expected setup to be called once"
assert cleanup_async.call_count == 1, "Expected cleanup to be called once"
@pytest.mark.parametrize(
"checkpoint_at", [CheckpointAt.END_OF_RUN, CheckpointAt.END_OF_STEP]
)
async def test_conditional_graph(checkpoint_at: CheckpointAt) -> 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",
]
)
async 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
async def execute_tools(data: dict) -> dict:
agent_action: AgentAction = data.pop("agent_outcome")
observation = await {t.name: t for t in tools}[agent_action.tool].ainvoke(
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
async 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 await app.ainvoke({"input": "what is weather in sf"}) == {
"input": "what is weather in sf",
"intermediate_steps": [
(
AgentAction(
tool="search_api",
tool_input="query",
log="tool:search_api:query",
),
"result for query",
),
(
AgentAction(
tool="search_api",
tool_input="another",
log="tool:search_api:another",
),
"result for another",
),
],
"agent_outcome": AgentFinish(
return_values={"answer": "answer"}, log="finish:answer"
),
}
# deepcopy because the nodes mutate the data
assert [
deepcopy(c) async for c in app.astream({"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"
),
}
},
]
patches = [c async for c in app.astream_log({"input": "what is weather in sf"})]
patch_paths = {op["path"] for log in patches for op in log.ops}
# Check that agent (one of the nodes) has its output streamed to the logs
assert "/logs/agent/streamed_output/-" in patch_paths
assert "/logs/agent:2/streamed_output/-" in patch_paths
assert "/logs/agent:3/streamed_output/-" in patch_paths
# Check that agent (one of the nodes) has its final output set in the logs
assert "/logs/agent/final_output" in patch_paths
assert "/logs/agent:2/final_output" in patch_paths
assert "/logs/agent:3/final_output" in patch_paths
assert [
p["value"]
for log in patches
for p in log.ops
if p["path"] == "/logs/agent/final_output"
or p["path"] == "/logs/agent:2/final_output"
or p["path"] == "/logs/agent:3/final_output"
] == [
{
"input": "what is weather in sf",
"agent_outcome": AgentAction(
tool="search_api", tool_input="query", log="tool:search_api:query"
),
},
{
"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",
),
},
{
"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(at=checkpoint_at),
interrupt_after=["agent"],
)
config = {"configurable": {"thread_id": "1"}}
assert [
c
async for c in app_w_interrupt.astream(
{"input": "what is weather in sf"}, config
)
] == [
{
"agent": {
"input": "what is weather in sf",
"agent_outcome": AgentAction(
tool="search_api", tool_input="query", log="tool:search_api:query"
),
}
}
]
assert await app_w_interrupt.aget_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=(await app_w_interrupt.checkpointer.aget_tuple(config)).config,
)
await app_w_interrupt.aupdate_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 await app_w_interrupt.aget_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=(await app_w_interrupt.checkpointer.aget_tuple(config)).config,
)
assert [c async for c in app_w_interrupt.astream(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",
),
}
},
]
await app_w_interrupt.aupdate_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 await app_w_interrupt.aget_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=(await app_w_interrupt.checkpointer.aget_tuple(config)).config,
)
# test state get/update methods with interrupt_before
app_w_interrupt = workflow.compile(
checkpointer=MemorySaverAssertImmutable(at=checkpoint_at),
interrupt_before=["tools"],
)
config = {"configurable": {"thread_id": "2"}}
llm.i = 0
assert [
c
async for c in app_w_interrupt.astream(
{"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 await app_w_interrupt.aget_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=(await app_w_interrupt.checkpointer.aget_tuple(config)).config,
)
await app_w_interrupt.aupdate_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 await app_w_interrupt.aget_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=(await app_w_interrupt.checkpointer.aget_tuple(config)).config,
)
assert [c async for c in app_w_interrupt.astream(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",
),
}
},
]
await app_w_interrupt.aupdate_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 await app_w_interrupt.aget_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=(await app_w_interrupt.checkpointer.aget_tuple(config)).config,
)
# test re-invoke to continue with interrupt_before
app_w_interrupt = workflow.compile(
checkpointer=MemorySaverAssertImmutable(at=checkpoint_at),
interrupt_before=["tools"],
)
config = {"configurable": {"thread_id": "2"}}
llm.i = 0 # reset the llm
assert [
c
async for c in app_w_interrupt.astream(
{"input": "what is weather in sf"}, config
)
] == [
{
"agent": {
"input": "what is weather in sf",
"agent_outcome": AgentAction(
tool="search_api", tool_input="query", log="tool:search_api:query"
),
}
}
]
assert await app_w_interrupt.aget_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=(await app_w_interrupt.checkpointer.aget_tuple(config)).config,
)
assert [c async for c in app_w_interrupt.astream(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 async for c in app_w_interrupt.astream(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"
),
}
},
]
@pytest.mark.parametrize(
"checkpoint_at", [CheckpointAt.END_OF_RUN, CheckpointAt.END_OF_STEP]
)
async def test_conditional_graph_state(checkpoint_at: CheckpointAt) -> 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):
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 await app.ainvoke({"input": "what is weather in sf"}) == {
"input": "what is weather in sf",
"intermediate_steps": [
(
AgentAction(
tool="search_api",
tool_input="query",
log="tool:search_api:query",
),
"result for query",
),
(
AgentAction(
tool="search_api",
tool_input="another",
log="tool:search_api:another",
),
"result for another",
),
],
"agent_outcome": AgentFinish(
return_values={"answer": "answer"}, log="finish:answer"
),
}
assert [c async for c in app.astream({"input": "what is weather in sf"})] == [
{
"agent": {
"agent_outcome": AgentAction(
tool="search_api", tool_input="query", log="tool:search_api:query"
),
}
},
{
"tools": {
"intermediate_steps": [
(
AgentAction(
tool="search_api",
tool_input="query",
log="tool:search_api:query",
),
"result for query",
)
],
}
},
{
"agent": {
"agent_outcome": AgentAction(
tool="search_api",
tool_input="another",
log="tool:search_api:another",
),
}
},
{
"tools": {
"intermediate_steps": [
(
AgentAction(
tool="search_api",
tool_input="another",
log="tool:search_api:another",
),
"result for another",
),
],
}
},
{
"agent": {
"agent_outcome": AgentFinish(
return_values={"answer": "answer"}, log="finish:answer"
),
}
},
]
patches = [c async for c in app.astream_log({"input": "what is weather in sf"})]
patch_paths = {op["path"] for log in patches for op in log.ops}
# Check that agent (one of the nodes) has its output streamed to the logs
assert "/logs/agent/streamed_output/-" in patch_paths
# Check that agent (one of the ndoes) has its final output set in the logs
assert "/logs/agent/final_output" in patch_paths
assert [
p["value"]
for log in patches
for p in log.ops
if p["path"] == "/logs/agent/final_output"
or p["path"] == "/logs/agent:2/final_output"
or p["path"] == "/logs/agent:3/final_output"
] == [
{
"agent_outcome": AgentAction(
tool="search_api", tool_input="query", log="tool:search_api:query"
)
},
{
"agent_outcome": AgentAction(
tool="search_api", tool_input="another", log="tool:search_api:another"
)
},
{
"agent_outcome": AgentFinish(
return_values={"answer": "answer"}, log="finish:answer"
),
},
]
# test state get/update methods with interrupt_after
app_w_interrupt = workflow.compile(
checkpointer=MemorySaverAssertImmutable(at=checkpoint_at),
interrupt_after=["agent"],
)
config = {"configurable": {"thread_id": "1"}}
assert [
c
async for c in app_w_interrupt.astream(
{"input": "what is weather in sf"}, config
)
] == [
{
"agent": {
"agent_outcome": AgentAction(
tool="search_api", tool_input="query", log="tool:search_api:query"
),
}
},
]
assert await app_w_interrupt.aget_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=(await app_w_interrupt.checkpointer.aget_tuple(config)).config,
)
await app_w_interrupt.aupdate_state(
config,
{
"agent_outcome": AgentAction(
tool="search_api",
tool_input="query",
log="tool:search_api:a different query",
)
},
)
assert await app_w_interrupt.aget_state(config) == StateSnapshot(
values={
"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=(await app_w_interrupt.checkpointer.aget_tuple(config)).config,
)
assert [c async for c in app_w_interrupt.astream(None, config)] == [
{
"tools": {
"intermediate_steps": [
(
AgentAction(
tool="search_api",
tool_input="query",
log="tool:search_api:a different query",
),
"result for query",
)
],
}
},
{
"agent": {
"agent_outcome": AgentAction(
tool="search_api",
tool_input="another",
log="tool:search_api:another",
),
}
},
]
await app_w_interrupt.aupdate_state(
config,
{
"agent_outcome": AgentFinish(
return_values={"answer": "a really nice answer"},
log="finish:a really nice answer",
)
},
)
assert await app_w_interrupt.aget_state(config) == StateSnapshot(
values={
"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=(await app_w_interrupt.checkpointer.aget_tuple(config)).config,
)
# test state get/update methods with interrupt_before
app_w_interrupt = workflow.compile(
checkpointer=MemorySaverAssertImmutable(at=checkpoint_at),
interrupt_before=["tools"],
)
config = {"configurable": {"thread_id": "2"}}
llm.i = 0 # reset the llm
assert [
c
async for c in app_w_interrupt.astream(
{"input": "what is weather in sf"}, config
)
] == [
{
"agent": {
"agent_outcome": AgentAction(
tool="search_api", tool_input="query", log="tool:search_api:query"
),
}
},
]
assert await app_w_interrupt.aget_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=(await app_w_interrupt.checkpointer.aget_tuple(config)).config,
)
await app_w_interrupt.aupdate_state(
config,
{
"agent_outcome": AgentAction(
tool="search_api",
tool_input="query",
log="tool:search_api:a different query",
)
},
)
assert await app_w_interrupt.aget_state(config) == StateSnapshot(
values={
"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=(await app_w_interrupt.checkpointer.aget_tuple(config)).config,
)
assert [c async for c in app_w_interrupt.astream(None, config)] == [
{
"tools": {
"intermediate_steps": [
(
AgentAction(
tool="search_api",
tool_input="query",
log="tool:search_api:a different query",
),
"result for query",
)
],
}
},
{
"agent": {
"agent_outcome": AgentAction(
tool="search_api",
tool_input="another",
log="tool:search_api:another",
),
}
},
]
await app_w_interrupt.aupdate_state(
config,
{
"agent_outcome": AgentFinish(
return_values={"answer": "a really nice answer"},
log="finish:a really nice answer",
)
},
)
assert await app_w_interrupt.aget_state(config) == StateSnapshot(
values={
"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=(await app_w_interrupt.checkpointer.aget_tuple(config)).config,
)
async def test_conditional_entrypoint_graph() -> None:
async def left(data: str) -> str:
return data + "->left"
async 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)
workflow.add_edge("right", END)
app = workflow.compile()
assert await app.ainvoke("what is weather in sf") == "what is weather in sf->right"
assert [c async for c in app.astream("what is weather in sf")] == [
{"right": "what is weather in sf->right"},
]
async def test_conditional_entrypoint_graph_state() -> None:
class AgentState(TypedDict, total=False):
input: str
output: str
steps: Annotated[list[str], operator.add]
async def left(data: AgentState) -> AgentState:
return {"output": data["input"] + "->left"}
async 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)
workflow.add_edge("right", END)
app = workflow.compile()
assert await app.ainvoke({"input": "what is weather in sf"}) == {
"input": "what is weather in sf",
"output": "what is weather in sf->right",
"steps": [],
}
assert [c async for c in app.astream({"input": "what is weather in sf"})] == [
{"right": {"output": "what is weather in sf->right"}},
]
async def test_prebuilt_tool_chat() -> 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 await app.ainvoke(
{"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 [
c
async for c in app.astream(
{"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",
tool_call_id="tool_call234",
name="search_api",
id=AnyStr(),
),
ToolMessage(
content="result for a third one",
tool_call_id="tool_call567",
name="search_api",
id=AnyStr(),
),
]
}
},
{"agent": {"messages": [AIMessage(content="answer", id=AnyStr())]}},
]
async def test_prebuilt_chat() -> 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 await app.ainvoke(
{"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 [
c
async for c in app.astream(
{"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())]}},
]
@pytest.mark.parametrize(
"checkpoint_at", [CheckpointAt.END_OF_RUN, CheckpointAt.END_OF_STEP]
)
async def test_message_graph(checkpoint_at: CheckpointAt) -> None:
from langchain.chat_models.fake import FakeMessagesListChatModel
from langchain_community.tools import tool
from langchain_core.agents import AgentAction
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]
model = FakeFuntionChatModel(
responses=[
AIMessage(
content="",
additional_kwargs={
"function_call": {
"name": "search_api",
"arguments": json.dumps("query"),
}
},
id="ai1",
),
AIMessage(
content="",
additional_kwargs={
"function_call": {
"name": "search_api",
"arguments": json.dumps("another"),
}
},
id="ai2",
),
AIMessage(content="answer", id="ai3"),
]
)
tool_executor = ToolExecutor(tools)
# 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 "function_call" not in last_message.additional_kwargs:
return "end"
# Otherwise if there is, we continue
else:
return "continue"
async def call_tool(messages):
# Based on the continue condition
# we know the last message involves a function call
last_message = messages[-1]
# We construct an AgentAction from the function_call
action = AgentAction(
tool=last_message.additional_kwargs["function_call"]["name"],
tool_input=json.loads(
last_message.additional_kwargs["function_call"]["arguments"]
),
log="",
)
# We call the tool_executor and get back a response
response = await tool_executor.ainvoke(action)
# We use the response to create a FunctionMessage
return FunctionMessage(content=str(response), name=action.tool)
# Define a new graph
workflow = MessageGraph()
# Define the two nodes we will cycle between
workflow.add_node("agent", model)
workflow.add_node("action", call_tool)
# 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 await app.ainvoke(HumanMessage(content="what is weather in sf")) == [
HumanMessage(content="what is weather in sf", id=AnyStr()),
AIMessage(
content="",
additional_kwargs={
"function_call": {"name": "search_api", "arguments": '"query"'}
},
id="ai1",
),
FunctionMessage(content="result for query", name="search_api", id=AnyStr()),
AIMessage(
content="",
additional_kwargs={
"function_call": {"name": "search_api", "arguments": '"another"'}
},
id="ai2",
),
FunctionMessage(content="result for another", name="search_api", id=AnyStr()),
AIMessage(content="answer", id="ai3"),
]
assert [
c async for c in app.astream([HumanMessage(content="what is weather in sf")])
] == [
{
"agent": AIMessage(
content="",
additional_kwargs={
"function_call": {"name": "search_api", "arguments": '"query"'}
},
id="ai1",
)
},
{
"action": FunctionMessage(
content="result for query", name="search_api", id=AnyStr()
)
},
{
"agent": AIMessage(
content="",
additional_kwargs={
"function_call": {"name": "search_api", "arguments": '"another"'}
},
id="ai2",
)
},
{
"action": FunctionMessage(
content="result for another", name="search_api", id=AnyStr()
)
},
{"agent": AIMessage(content="answer", id="ai3")},
]
app_w_interrupt = workflow.compile(
checkpointer=MemorySaverAssertImmutable(at=checkpoint_at),
interrupt_after=["agent"],
)
config = {"configurable": {"thread_id": "1"}}
assert [
c
async for c in app_w_interrupt.astream(
HumanMessage(content="what is weather in sf"), config
)
] == [
{
"agent": AIMessage(
content="",
additional_kwargs={
"function_call": {"name": "search_api", "arguments": '"query"'}
},
id="ai1",
)
},
]
assert await app_w_interrupt.aget_state(config) == StateSnapshot(
values=[
HumanMessage(
content="what is weather in sf",
id=AnyStr(),
),
AIMessage(
content="",
additional_kwargs={
"function_call": {"name": "search_api", "arguments": '"query"'}
},
id="ai1",
),
],
next=("action",),
config=(await app_w_interrupt.checkpointer.aget_tuple(config)).config,
)
# modify ai message
last_message = (await app_w_interrupt.aget_state(config)).values[-1]
last_message.additional_kwargs["function_call"]["arguments"] = '"a different query"'
await app_w_interrupt.aupdate_state(config, last_message)
# message was replaced instead of appended
assert await app_w_interrupt.aget_state(config) == StateSnapshot(
values=[
HumanMessage(
content="what is weather in sf",
id=AnyStr(),
),
AIMessage(
content="",
additional_kwargs={
"function_call": {
"name": "search_api",
"arguments": '"a different query"',
}
},
id="ai1",
),
],
next=("action",),
config=app_w_interrupt.checkpointer.get_tuple(config).config,
)
assert [c async for c in app_w_interrupt.astream(None, config)] == [
{
"action": FunctionMessage(
content="result for a different query",
name="search_api",
id=AnyStr(),
)
},
{
"agent": AIMessage(
content="",
additional_kwargs={
"function_call": {"name": "search_api", "arguments": '"another"'}
},
id="ai2",
)
},
]
assert await app_w_interrupt.aget_state(config) == StateSnapshot(
values=[
HumanMessage(
content="what is weather in sf",
id=AnyStr(),
),
AIMessage(
content="",
additional_kwargs={
"function_call": {
"name": "search_api",
"arguments": '"a different query"',
}
},
id="ai1",
),
FunctionMessage(
content="result for a different query",
name="search_api",
id=AnyStr(),
),
AIMessage(
content="",
additional_kwargs={
"function_call": {"name": "search_api", "arguments": '"another"'}
},
id="ai2",
),
],
next=("action",),
config=app_w_interrupt.checkpointer.get_tuple(config).config,
)
await app_w_interrupt.aupdate_state(
config,
AIMessage(content="answer", id="ai2"),
)
# replaces message even if object identity is different, as long as id is the same
assert await app_w_interrupt.aget_state(config) == StateSnapshot(
values=[
HumanMessage(
content="what is weather in sf",
id=AnyStr(),
),
AIMessage(
content="",
additional_kwargs={
"function_call": {
"name": "search_api",
"arguments": '"a different query"',
}
},
id="ai1",
),
FunctionMessage(
content="result for a different query",
name="search_api",
id=AnyStr(),
),
AIMessage(content="answer", id="ai2"),
],
next=(),
config=app_w_interrupt.checkpointer.get_tuple(config).config,
)
async def test_in_one_fan_out_out_one_graph_state() -> None:
def sorted_add(x: list[str], y: list[str]) -> list[str]:
return sorted(operator.add(x, y))
class State(TypedDict, total=False):
query: str
answer: str
docs: Annotated[list[str], sorted_add]
async def rewrite_query(data: State) -> State:
return {"query": f'query: {data["query"]}'}
async def retriever_one(data: State) -> State:
return {"docs": ["doc1", "doc2"]}
async def retriever_two(data: State) -> State:
return {"docs": ["doc3", "doc4"]}
async def qa(data: State) -> State:
return {"answer": ",".join(data["docs"])}
workflow = StateGraph(State)
workflow.add_node("rewrite_query", rewrite_query)
workflow.add_node("retriever_one", retriever_one)
workflow.add_node("retriever_two", retriever_two)
workflow.add_node("qa", qa)
workflow.set_entry_point("rewrite_query")
workflow.add_edge("rewrite_query", "retriever_one")
workflow.add_edge("rewrite_query", "retriever_two")
workflow.add_edge("retriever_one", "qa")
workflow.add_edge("retriever_two", "qa")
workflow.set_finish_point("qa")
app = workflow.compile()
assert await app.ainvoke({"query": "what is weather in sf"}) == {
"query": "query: what is weather in sf",
"docs": ["doc1", "doc2", "doc3", "doc4"],
"answer": "doc1,doc2,doc3,doc4",
}
assert [c async for c in app.astream({"query": "what is weather in sf"})] == [
{"rewrite_query": {"query": "query: what is weather in sf"}},
{
"retriever_two": {"docs": ["doc3", "doc4"]},
"retriever_one": {"docs": ["doc1", "doc2"]},
},
{"qa": {"answer": "doc1,doc2,doc3,doc4"}},
]
assert [
c
async for c in app.astream(
{"query": "what is weather in sf"}, stream_mode="values"
)
] == [
{"query": "what is weather in sf", "docs": []},
{"query": "query: what is weather in sf", "docs": []},
{
"query": "query: what is weather in sf",
"docs": ["doc1", "doc2", "doc3", "doc4"],
},
{
"query": "query: what is weather in sf",
"docs": ["doc1", "doc2", "doc3", "doc4"],
"answer": "doc1,doc2,doc3,doc4",
},
]
@pytest.mark.parametrize(
"checkpoint_at", [CheckpointAt.END_OF_RUN, CheckpointAt.END_OF_STEP]
)
async def test_in_one_fan_out_state_graph_waiting_edge(
checkpoint_at: CheckpointAt,
) -> 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]
async def rewrite_query(data: State) -> State:
return {"query": f'query: {data["query"]}'}
async def analyzer_one(data: State) -> State:
return {"query": f'analyzed: {data["query"]}'}
async def retriever_one(data: State) -> State:
return {"docs": ["doc1", "doc2"]}
async def retriever_two(data: State) -> State:
return {"docs": ["doc3", "doc4"]}
async def qa(data: State) -> State:
return {"answer": ",".join(data["docs"])}
workflow = StateGraph(State)
workflow.add_node("rewrite_query", rewrite_query)
workflow.add_node("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 await app.ainvoke({"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 [c async for c in app.astream({"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(at=checkpoint_at),
interrupt_after=["retriever_one"],
)
config = {"configurable": {"thread_id": "1"}}
assert [
c
async for c in app_w_interrupt.astream(
{"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 async for c in app_w_interrupt.astream(None, config)] == [
{"qa": {"answer": "doc1,doc2,doc3,doc4"}},
]
@pytest.mark.parametrize(
"checkpoint_at", [CheckpointAt.END_OF_RUN, CheckpointAt.END_OF_STEP]
)
async def test_in_one_fan_out_state_graph_waiting_edge_via_branch(
snapshot: SnapshotAssertion,
checkpoint_at: CheckpointAt,
) -> 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]
async def rewrite_query(data: State) -> State:
return {"query": f'query: {data["query"]}'}
async def analyzer_one(data: State) -> State:
return {"query": f'analyzed: {data["query"]}'}
async def retriever_one(data: State) -> State:
return {"docs": ["doc1", "doc2"]}
async def retriever_two(data: State) -> State:
return {"docs": ["doc3", "doc4"]}
async def qa(data: State) -> State:
return {"answer": ",".join(data["docs"])}
workflow = StateGraph(State)
workflow.add_node("rewrite_query", rewrite_query)
workflow.add_node("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", lambda _: "retriever_two", {"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_ascii() == snapshot
assert await app.ainvoke({"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 [c async for c in app.astream({"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(at=checkpoint_at),
interrupt_after=["retriever_one"],
)
config = {"configurable": {"thread_id": "1"}}
assert [
c
async for c in app_w_interrupt.astream(
{"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 async for c in app_w_interrupt.astream(None, config)] == [
{"qa": {"answer": "doc1,doc2,doc3,doc4"}},
]
@pytest.mark.parametrize(
"checkpoint_at", [CheckpointAt.END_OF_RUN, CheckpointAt.END_OF_STEP]
)
async def test_in_one_fan_out_state_graph_waiting_edge_custom_state_class(
snapshot: SnapshotAssertion,
checkpoint_at: CheckpointAt,
) -> 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]
async def rewrite_query(data: State) -> State:
return {"query": f"query: {data.query}"}
async def analyzer_one(data: State) -> State:
return {"query": f"analyzed: {data.query}"}
async def retriever_one(data: State) -> State:
return {"docs": ["doc1", "doc2"]}
async def retriever_two(data: State) -> State:
return {"docs": ["doc3", "doc4"]}
async def qa(data: State) -> State:
return {"answer": ",".join(data.docs)}
async 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_ascii() == snapshot
with pytest.raises(ValidationError):
await app.ainvoke({"query": {}})
assert await app.ainvoke({"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 [c async for c in app.astream({"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(at=checkpoint_at),
interrupt_after=["retriever_one"],
)
config = {"configurable": {"thread_id": "1"}}
assert [
c
async for c in app_w_interrupt.astream(
{"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 async for c in app_w_interrupt.astream(None, config)] == [
{"qa": {"answer": "doc1,doc2,doc3,doc4"}},
]
@pytest.mark.parametrize(
"checkpoint_at", [CheckpointAt.END_OF_RUN, CheckpointAt.END_OF_STEP]
)
async def test_in_one_fan_out_state_graph_waiting_edge_plus_regular(
checkpoint_at: CheckpointAt,
) -> 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]
async def rewrite_query(data: State) -> State:
return {"query": f'query: {data["query"]}'}
async def analyzer_one(data: State) -> State:
return {"query": f'analyzed: {data["query"]}'}
async def retriever_one(data: State) -> State:
return {"docs": ["doc1", "doc2"]}
async def retriever_two(data: State) -> State:
return {"docs": ["doc3", "doc4"]}
async def qa(data: State) -> State:
return {"answer": ",".join(data["docs"])}
workflow = StateGraph(State)
workflow.add_node("rewrite_query", rewrite_query)
workflow.add_node("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 await app.ainvoke({"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 [c async for c in app.astream({"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(at=checkpoint_at),
interrupt_after=["retriever_one"],
)
config = {"configurable": {"thread_id": "1"}}
assert [
c
async for c in app_w_interrupt.astream(
{"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 async for c in app_w_interrupt.astream(None, config)] == [
{"qa": {"answer": "doc1,doc2,doc3,doc4"}},
]
async 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]
async def rewrite_query(data: State) -> State:
return {"query": f'query: {data["query"]}'}
async def analyzer_one(data: State) -> State:
return {"query": f'analyzed: {data["query"]}'}
async def retriever_one(data: State) -> State:
return {"docs": ["doc1", "doc2"]}
async def retriever_two(data: State) -> State:
return {"docs": ["doc3", "doc4"]}
async def qa(data: State) -> State:
return {"answer": ",".join(data["docs"])}
async 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 await app.ainvoke({"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 [c async for c in app.astream({"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"}},
]
async 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]
async def rewrite_query(data: State) -> State:
return {"query": f'query: {data["query"]}'}
async def retriever_picker(data: State) -> list[str]:
return ["analyzer_one", "retriever_two"]
async def analyzer_one(data: State) -> State:
return {"query": f'analyzed: {data["query"]}'}
async def retriever_one(data: State) -> State:
return {"docs": ["doc1", "doc2"]}
async def retriever_two(data: State) -> State:
return {"docs": ["doc3", "doc4"]}
async def qa(data: State) -> State:
return {"answer": ",".join(data["docs"])}
async 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 await app.ainvoke({"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 [c async for c in app.astream({"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"}},
]
async def test_nested_graph(snapshot: SnapshotAssertion) -> None:
class State(TypedDict):
my_key: str
async def up(state: State):
return {"my_key": state["my_key"] + " there"}
inner = StateGraph(State)
inner.add_node("up", up)
inner.set_entry_point("up")
inner.set_finish_point("up")
async 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_ascii() == snapshot
assert await app.ainvoke({"my_key": "my value"}) == {
"my_key": "my value there and back again"
}
assert [chunk async for chunk in app.astream({"my_key": "my value"})] == [
{"inner": {"my_key": "my value there"}},
{"side": {"my_key": "my value there and back again"}},
]
assert [
chunk
async for chunk in app.astream({"my_key": "my value"}, stream_mode="values")
] == [
{"my_key": "my value"},
{"my_key": "my value there"},
{"my_key": "my value there and back again"},
]
times_called = 0
async for event in app.astream_events(
{"my_key": "my value"},
version="v1",
config={"run_id": UUID(int=0)},
stream_mode="values",
):
if event["event"] == "on_chain_end" and event["run_id"] == str(UUID(int=0)):
times_called += 1
assert event["data"] == {
"output": {"my_key": "my value there and back again"}
}
assert times_called == 1
times_called = 0
async for event in app.astream_events(
{"my_key": "my value"},
version="v1",
config={"run_id": UUID(int=0)},
):
if event["event"] == "on_chain_end" and event["run_id"] == str(UUID(int=0)):
times_called += 1
assert event["data"] == {
"output": [
{"inner": {"my_key": "my value there"}},
{"side": {"my_key": "my value there and back again"}},
]
}
assert times_called == 1
Reference in New Issue View Git Blame Copy Permalink