mirror of
https://github.com/langchain-ai/langgraph.git
synced 2026-08-20 06:35:46 +02:00
Nuno Campos
f2803372ec
- Move all validation to compile() This allows adding edges before nodes - Detect more cases of missing edges with shorthand branches
3663 lines
116 KiB
Python
3663 lines
116 KiB
Python
import json
|
|
import operator
|
|
import time
|
|
import warnings
|
|
from concurrent.futures import ThreadPoolExecutor
|
|
from contextlib import contextmanager
|
|
from typing import Annotated, Any, Generator, 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.base import CheckpointAt
|
|
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_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 app.invoke(2) == 3
|
|
|
|
|
|
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 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,
|
|
},
|
|
output_channels=["output"],
|
|
)
|
|
|
|
assert app.input_schema.schema() == {"title": "LangGraphInput"}
|
|
assert app.output_schema.schema() == {
|
|
"title": "LangGraphOutput",
|
|
"type": "object",
|
|
"properties": {"output": {"title": "Output"}},
|
|
}
|
|
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,
|
|
},
|
|
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 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},
|
|
)
|
|
|
|
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
|
|
|
|
|
|
@pytest.mark.parametrize(
|
|
"checkpoint_at", [CheckpointAt.END_OF_RUN, CheckpointAt.END_OF_STEP]
|
|
)
|
|
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 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)},
|
|
input_channels=["input", "inbox"],
|
|
stream_channels=["output", "inbox"],
|
|
)
|
|
|
|
# [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},
|
|
]
|
|
|
|
|
|
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})
|
|
|
|
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)
|
|
|
|
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)
|
|
|
|
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})
|
|
|
|
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={"output": Topic(int)},
|
|
)
|
|
|
|
# An Inbox channel accumulates updates into a sequence
|
|
assert app.invoke(2) == [3, 3]
|
|
|
|
|
|
@pytest.mark.parametrize(
|
|
"checkpoint_at", [CheckpointAt.END_OF_RUN, CheckpointAt.END_OF_STEP]
|
|
)
|
|
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 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
|
|
|
|
|
|
@pytest.mark.parametrize(
|
|
"checkpoint_at", [CheckpointAt.END_OF_RUN, CheckpointAt.END_OF_STEP]
|
|
)
|
|
def test_invoke_checkpoint_sqlite(
|
|
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
|
|
)
|
|
|
|
with SqliteSaver.from_conn_string(":memory:") as memory:
|
|
memory.at = checkpoint_at
|
|
app = Pregel(
|
|
nodes={"one": one},
|
|
channels={"total": BinaryOperatorAggregate(int, operator.add)},
|
|
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) == 2
|
|
state = app.get_state(thread_1)
|
|
assert state is not None
|
|
assert state.values.get("total") == 2
|
|
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
|
|
|
|
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) == 5
|
|
state = app.get_state({"configurable": {"thread_id": "1"}})
|
|
assert state is not None
|
|
assert state.values.get("total") == 7
|
|
assert state.next == ()
|
|
state = app.get_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 for c in app.get_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 (
|
|
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))) == 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)},
|
|
)
|
|
|
|
# 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_app(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)},
|
|
)
|
|
|
|
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}, stream_channels=["output", "between"])
|
|
|
|
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})
|
|
|
|
# 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_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"
|
|
|
|
|
|
@pytest.mark.parametrize(
|
|
"checkpoint_at", [CheckpointAt.END_OF_RUN, CheckpointAt.END_OF_STEP]
|
|
)
|
|
def test_conditional_graph(
|
|
snapshot: SnapshotAssertion, 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",
|
|
]
|
|
)
|
|
|
|
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_ascii() == snapshot
|
|
assert (
|
|
app.get_graph(add_condition_nodes=False).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_ascii() == snapshot
|
|
assert (
|
|
app.get_graph(xray=True, add_condition_nodes=False).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(at=checkpoint_at),
|
|
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,
|
|
)
|
|
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,
|
|
)
|
|
|
|
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,
|
|
)
|
|
|
|
# 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 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,
|
|
)
|
|
|
|
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,
|
|
)
|
|
|
|
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,
|
|
)
|
|
|
|
# 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 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,
|
|
)
|
|
|
|
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_ascii() == 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"},
|
|
]
|
|
|
|
|
|
@pytest.mark.parametrize(
|
|
"checkpoint_at", [CheckpointAt.END_OF_RUN, CheckpointAt.END_OF_STEP]
|
|
)
|
|
def test_conditional_graph_state(
|
|
snapshot: SnapshotAssertion, 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, 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_ascii() == 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(at=checkpoint_at),
|
|
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,
|
|
)
|
|
|
|
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,
|
|
)
|
|
|
|
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,
|
|
)
|
|
|
|
# test state get/update methods with interrupt_before
|
|
|
|
app_w_interrupt = workflow.compile(
|
|
checkpointer=MemorySaverAssertImmutable(at=checkpoint_at),
|
|
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,
|
|
)
|
|
|
|
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,
|
|
)
|
|
|
|
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,
|
|
)
|
|
|
|
|
|
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_ascii() == 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_ascii() == 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_ascii() == 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())]}},
|
|
]
|
|
|
|
|
|
@pytest.mark.parametrize(
|
|
"checkpoint_at", [CheckpointAt.END_OF_RUN, CheckpointAt.END_OF_STEP]
|
|
)
|
|
def test_message_graph(
|
|
snapshot: SnapshotAssertion,
|
|
checkpoint_at: CheckpointAt,
|
|
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_ascii() == 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(at=checkpoint_at),
|
|
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,
|
|
)
|
|
|
|
# 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,
|
|
)
|
|
|
|
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,
|
|
)
|
|
|
|
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,
|
|
)
|
|
|
|
app_w_interrupt = workflow.compile(
|
|
checkpointer=MemorySaverAssertImmutable(at=checkpoint_at),
|
|
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,
|
|
)
|
|
|
|
# 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,
|
|
)
|
|
|
|
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,
|
|
)
|
|
|
|
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,
|
|
)
|
|
|
|
# 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,
|
|
)
|
|
|
|
|
|
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",
|
|
},
|
|
]
|
|
|
|
|
|
@pytest.mark.parametrize(
|
|
"checkpoint_at", [CheckpointAt.END_OF_RUN, CheckpointAt.END_OF_STEP]
|
|
)
|
|
def test_in_one_fan_out_state_graph_waiting_edge(
|
|
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]
|
|
|
|
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_ascii() == 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(at=checkpoint_at),
|
|
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"}},
|
|
]
|
|
|
|
|
|
@pytest.mark.parametrize(
|
|
"checkpoint_at", [CheckpointAt.END_OF_RUN, CheckpointAt.END_OF_STEP]
|
|
)
|
|
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]
|
|
|
|
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_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 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(at=checkpoint_at),
|
|
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"}},
|
|
]
|
|
|
|
|
|
@pytest.mark.parametrize(
|
|
"checkpoint_at", [CheckpointAt.END_OF_RUN, CheckpointAt.END_OF_STEP]
|
|
)
|
|
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]
|
|
|
|
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_ascii() == 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(at=checkpoint_at),
|
|
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"}},
|
|
]
|
|
|
|
|
|
@pytest.mark.parametrize(
|
|
"checkpoint_at", [CheckpointAt.END_OF_RUN, CheckpointAt.END_OF_STEP]
|
|
)
|
|
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]
|
|
|
|
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(at=checkpoint_at),
|
|
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_ascii() == snapshot
|
|
assert app.invoke({"my_key": "my_value"}) == {"my_key": "my_value"}
|
|
|
|
|
|
def test_nested_graph(snapshot: SnapshotAssertion) -> None:
|
|
class State(TypedDict):
|
|
my_key: str
|
|
|
|
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")
|
|
|
|
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 app.invoke({"my_key": "my value"}) == {
|
|
"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
|