wehub-resource-sync
a7d6d88f6f
CI / changes (push) Has been cancelled
CI / cd libs/checkpoint (push) Has been cancelled
CI / cd libs/checkpoint-conformance (push) Has been cancelled
CI / cd libs/checkpoint-postgres (push) Has been cancelled
CI / cd libs/checkpoint-sqlite (push) Has been cancelled
CI / cd libs/cli (push) Has been cancelled
CI / cd libs/prebuilt (push) Has been cancelled
CI / cd libs/sdk-py (push) Has been cancelled
CI / cd libs/langgraph (push) Has been cancelled
CI / Check SDK methods matching (push) Has been cancelled
CI / Check CLI schema hasn't changed #3.13 (push) Has been cancelled
CI / CLI integration test (push) Has been cancelled
CI / sdk-py integration test (push) Has been cancelled
CI / CI Success (push) Has been cancelled
baseline / benchmark (push) Has been cancelled
Deploy Redirects to GitHub Pages / deploy (push) Has been cancelled
9678 lines
303 KiB
Python
9678 lines
303 KiB
Python
import enum
|
|
import functools
|
|
import gc
|
|
import json
|
|
import logging
|
|
import operator
|
|
import threading
|
|
import time
|
|
import uuid
|
|
from collections import Counter, deque
|
|
from collections.abc import Sequence
|
|
from concurrent.futures import ThreadPoolExecutor
|
|
from dataclasses import dataclass, field
|
|
from random import randrange
|
|
from typing import Annotated, Any, Literal, get_type_hints
|
|
|
|
import pytest
|
|
from langchain_core.language_models import GenericFakeChatModel
|
|
from langchain_core.messages import AIMessage, AnyMessage, HumanMessage, RemoveMessage
|
|
from langchain_core.runnables import (
|
|
RunnableConfig,
|
|
RunnableLambda,
|
|
RunnablePassthrough,
|
|
)
|
|
from langchain_core.runnables.graph import Edge
|
|
from langchain_core.version import VERSION as LANGCHAIN_CORE_VERSION
|
|
from langgraph.cache.base import BaseCache
|
|
from langgraph.checkpoint.base import (
|
|
BaseCheckpointSaver,
|
|
Checkpoint,
|
|
CheckpointMetadata,
|
|
CheckpointTuple,
|
|
)
|
|
from langgraph.checkpoint.memory import InMemorySaver
|
|
from langgraph.prebuilt.tool_node import ToolNode
|
|
from langgraph.store.base import BaseStore
|
|
from langsmith import traceable
|
|
from pydantic import BaseModel, ConfigDict, Field, ValidationError
|
|
from pytest_mock import MockerFixture
|
|
from syrupy import SnapshotAssertion
|
|
from typing_extensions import NotRequired, TypedDict
|
|
|
|
from langgraph._internal._constants import CONFIG_KEY_NODE_FINISHED, ERROR, PULL
|
|
from langgraph.channels.binop import BinaryOperatorAggregate
|
|
from langgraph.channels.delta import DeltaChannel
|
|
from langgraph.channels.ephemeral_value import EphemeralValue
|
|
from langgraph.channels.last_value import LastValue
|
|
from langgraph.channels.topic import Topic
|
|
from langgraph.channels.untracked_value import UntrackedValue
|
|
from langgraph.config import get_stream_writer
|
|
from langgraph.errors import GraphRecursionError, InvalidUpdateError, ParentCommand
|
|
from langgraph.func import entrypoint, task
|
|
from langgraph.graph import END, START, StateGraph
|
|
from langgraph.graph.message import MessagesState, _messages_delta_reducer, add_messages
|
|
from langgraph.pregel import (
|
|
NodeBuilder,
|
|
Pregel,
|
|
)
|
|
from langgraph.pregel._loop import SyncPregelLoop
|
|
from langgraph.pregel._runner import PregelRunner
|
|
from langgraph.types import (
|
|
CachePolicy,
|
|
Command,
|
|
Durability,
|
|
Interrupt,
|
|
Overwrite,
|
|
PregelTask,
|
|
RetryPolicy,
|
|
Send,
|
|
StateSnapshot,
|
|
StateUpdate,
|
|
StreamWriter,
|
|
interrupt,
|
|
)
|
|
from tests.agents import AgentAction, AgentFinish
|
|
from tests.any_str import AnyStr, AnyVersion, FloatBetween, UnsortedSequence
|
|
from tests.messages import (
|
|
_AnyIdAIMessage,
|
|
_AnyIdAIMessageChunk,
|
|
_AnyIdHumanMessage,
|
|
_AnyIdToolMessage,
|
|
)
|
|
|
|
pytestmark = pytest.mark.anyio
|
|
|
|
logger = logging.getLogger(__name__)
|
|
|
|
|
|
def test_graph_validation() -> None:
|
|
class State(TypedDict):
|
|
hello: str
|
|
|
|
graph = StateGraph(State)
|
|
graph.add_node("start", lambda x: x)
|
|
graph.add_edge("__start__", "start")
|
|
graph.add_edge("unknown", "start")
|
|
graph.add_edge("start", "__end__")
|
|
with pytest.raises(ValueError, match="Found edge starting at unknown node "):
|
|
graph.compile()
|
|
|
|
def bad_reducer(a): ...
|
|
|
|
class BadReducerState(TypedDict):
|
|
hello: Annotated[str, bad_reducer]
|
|
|
|
with pytest.raises(ValueError, match="Invalid reducer"):
|
|
StateGraph(BadReducerState)
|
|
|
|
def node_b(state: State) -> State:
|
|
return {"hello": "world"}
|
|
|
|
builder = StateGraph(State)
|
|
builder.add_node("a", node_b)
|
|
builder.add_node("b", node_b)
|
|
builder.add_node("c", node_b)
|
|
builder.set_entry_point("a")
|
|
builder.add_edge("a", "b")
|
|
builder.add_edge("a", "c")
|
|
graph = builder.compile()
|
|
|
|
with pytest.raises(InvalidUpdateError, match="At key 'hello'"):
|
|
graph.invoke({"hello": "there"})
|
|
|
|
|
|
def test_request_drain_allows_inflight_call_scheduling(
|
|
sync_checkpointer: BaseCheckpointSaver,
|
|
) -> None:
|
|
from langgraph.runtime import RunControl
|
|
|
|
@task
|
|
def child(x: int) -> int:
|
|
return x + 1
|
|
|
|
control = RunControl()
|
|
|
|
@entrypoint(checkpointer=sync_checkpointer)
|
|
def graph(x: int) -> int:
|
|
control.request_drain()
|
|
fut = child(x)
|
|
return fut.result()
|
|
|
|
config = {"configurable": {"thread_id": "drain-call-sync"}}
|
|
|
|
assert graph.invoke(1, config=config, control=control) == 2
|
|
assert control.drain_requested
|
|
|
|
|
|
def test_invalid_checkpointer_type() -> None:
|
|
class State(TypedDict):
|
|
foo: str
|
|
|
|
builder = StateGraph(State)
|
|
builder.add_node("start", lambda state: state)
|
|
builder.set_entry_point("start")
|
|
builder.set_finish_point("start")
|
|
|
|
class NotACheckpointer:
|
|
pass
|
|
|
|
with pytest.raises(TypeError, match="Invalid checkpointer provided"):
|
|
builder.compile(checkpointer=NotACheckpointer())
|
|
|
|
|
|
def test_graph_validation_with_command() -> None:
|
|
class State(TypedDict):
|
|
foo: str
|
|
bar: str
|
|
|
|
def node_a(state: State):
|
|
return Command(goto="b", update={"foo": "bar"})
|
|
|
|
def node_b(state: State):
|
|
return Command(goto=END, update={"bar": "baz"})
|
|
|
|
builder = StateGraph(State)
|
|
builder.add_node("a", node_a)
|
|
builder.add_node("b", node_b)
|
|
builder.add_edge(START, "a")
|
|
graph = builder.compile()
|
|
assert graph.invoke({"foo": ""}) == {"foo": "bar", "bar": "baz"}
|
|
|
|
|
|
def test_checkpoint_errors() -> None:
|
|
class FaultyGetCheckpointer(InMemorySaver):
|
|
def get_tuple(self, config: RunnableConfig) -> CheckpointTuple | None:
|
|
raise ValueError("Faulty get_tuple")
|
|
|
|
class FaultyPutCheckpointer(InMemorySaver):
|
|
def put(
|
|
self,
|
|
config: RunnableConfig,
|
|
checkpoint: Checkpoint,
|
|
metadata: CheckpointMetadata,
|
|
new_versions: dict[str, str | int | float] | None = None,
|
|
) -> RunnableConfig:
|
|
raise ValueError("Faulty put")
|
|
|
|
class FaultyPutWritesCheckpointer(InMemorySaver):
|
|
def put_writes(
|
|
self, config: RunnableConfig, writes: list[tuple[str, Any]], task_id: str
|
|
) -> RunnableConfig:
|
|
raise ValueError("Faulty put_writes")
|
|
|
|
class FaultyVersionCheckpointer(InMemorySaver):
|
|
def get_next_version(self, current: int | None, channel: None) -> int:
|
|
raise ValueError("Faulty get_next_version")
|
|
|
|
def logic(inp: str) -> str:
|
|
return ""
|
|
|
|
builder = StateGraph(Annotated[str, operator.add])
|
|
builder.add_node("agent", logic)
|
|
builder.add_edge(START, "agent")
|
|
|
|
graph = builder.compile(checkpointer=FaultyGetCheckpointer())
|
|
with pytest.raises(ValueError, match="Faulty get_tuple"):
|
|
graph.invoke("", {"configurable": {"thread_id": "thread-1"}})
|
|
|
|
graph = builder.compile(checkpointer=FaultyPutCheckpointer())
|
|
with pytest.raises(ValueError, match="Faulty put"):
|
|
graph.invoke("", {"configurable": {"thread_id": "thread-1"}})
|
|
|
|
graph = builder.compile(checkpointer=FaultyVersionCheckpointer())
|
|
with pytest.raises(ValueError, match="Faulty get_next_version"):
|
|
graph.invoke("", {"configurable": {"thread_id": "thread-1"}})
|
|
|
|
# add parallel node
|
|
builder.add_node("parallel", logic)
|
|
builder.add_edge(START, "parallel")
|
|
graph = builder.compile(checkpointer=FaultyPutWritesCheckpointer())
|
|
with pytest.raises(ValueError, match="Faulty put_writes"):
|
|
graph.invoke(
|
|
"", {"configurable": {"thread_id": "thread-1"}}, durability="async"
|
|
)
|
|
|
|
|
|
def test_context_json_schema() -> None:
|
|
"""Test that config json schema is generated properly."""
|
|
chain = NodeBuilder().subscribe_only("input").write_to("output")
|
|
|
|
@dataclass
|
|
class Foo:
|
|
x: int
|
|
y: str = field(default="foo")
|
|
|
|
app = Pregel(
|
|
nodes={
|
|
"one": chain,
|
|
},
|
|
channels={
|
|
"ephemeral": EphemeralValue(Any),
|
|
"input": LastValue(int),
|
|
"output": LastValue(int),
|
|
},
|
|
input_channels=["input", "ephemeral"],
|
|
output_channels="output",
|
|
context_schema=Foo,
|
|
)
|
|
|
|
assert app.get_context_jsonschema() == {
|
|
"properties": {
|
|
"x": {
|
|
"title": "X",
|
|
"type": "integer",
|
|
},
|
|
"y": {
|
|
"default": "foo",
|
|
"title": "Y",
|
|
"type": "string",
|
|
},
|
|
},
|
|
"required": [
|
|
"x",
|
|
],
|
|
"title": "Foo",
|
|
"type": "object",
|
|
}
|
|
|
|
|
|
def test_node_schemas_custom_output() -> None:
|
|
class State(TypedDict):
|
|
hello: str
|
|
bye: str
|
|
messages: Annotated[list[str], add_messages]
|
|
|
|
class Output(TypedDict):
|
|
messages: list[str]
|
|
|
|
class StateForA(TypedDict):
|
|
hello: str
|
|
messages: Annotated[list[str], add_messages]
|
|
|
|
def node_a(state: StateForA) -> State:
|
|
assert state == {
|
|
"hello": "there",
|
|
"messages": [_AnyIdHumanMessage(content="hello")],
|
|
}
|
|
|
|
class StateForB(TypedDict):
|
|
bye: str
|
|
now: int
|
|
|
|
def node_b(state: StateForB):
|
|
assert state == {
|
|
"bye": "world",
|
|
}
|
|
return {
|
|
"now": 123,
|
|
"hello": "again",
|
|
}
|
|
|
|
class StateForC(TypedDict):
|
|
hello: str
|
|
now: int
|
|
|
|
def node_c(state: StateForC) -> StateForC:
|
|
assert state == {
|
|
"hello": "again",
|
|
"now": 123,
|
|
}
|
|
|
|
builder = StateGraph(State, output_schema=Output)
|
|
builder.add_node("a", node_a)
|
|
builder.add_node("b", node_b)
|
|
builder.add_node("c", node_c)
|
|
builder.add_edge(START, "a")
|
|
builder.add_edge("a", "b")
|
|
builder.add_edge("b", "c")
|
|
graph = builder.compile()
|
|
|
|
assert graph.invoke({"hello": "there", "bye": "world", "messages": "hello"}) == {
|
|
"messages": [_AnyIdHumanMessage(content="hello")],
|
|
}
|
|
|
|
builder = StateGraph(State, output_schema=Output)
|
|
builder.add_node("a", node_a)
|
|
builder.add_node("b", node_b)
|
|
builder.add_node("c", node_c)
|
|
builder.add_edge(START, "a")
|
|
builder.add_edge("a", "b")
|
|
builder.add_edge("b", "c")
|
|
graph = builder.compile()
|
|
|
|
assert graph.invoke(
|
|
{
|
|
"hello": "there",
|
|
"bye": "world",
|
|
"messages": "hello",
|
|
"now": 345, # ignored because not in input schema
|
|
}
|
|
) == {
|
|
"messages": [_AnyIdHumanMessage(content="hello")],
|
|
}
|
|
|
|
assert [
|
|
c
|
|
for c in graph.stream(
|
|
{
|
|
"hello": "there",
|
|
"bye": "world",
|
|
"messages": "hello",
|
|
"now": 345, # ignored because not in input schema
|
|
}
|
|
)
|
|
] == [
|
|
{"a": None},
|
|
{"b": {"hello": "again", "now": 123}},
|
|
{"c": None},
|
|
]
|
|
|
|
|
|
def test_reducer_before_first_node() -> None:
|
|
class State(TypedDict):
|
|
hello: str
|
|
messages: Annotated[list[str], add_messages]
|
|
|
|
def node_a(state: State) -> State:
|
|
assert state == {
|
|
"hello": "there",
|
|
"messages": [_AnyIdHumanMessage(content="hello")],
|
|
}
|
|
|
|
builder = StateGraph(State)
|
|
builder.add_node("a", node_a)
|
|
builder.set_entry_point("a")
|
|
builder.set_finish_point("a")
|
|
graph = builder.compile()
|
|
assert graph.invoke({"hello": "there", "messages": "hello"}) == {
|
|
"hello": "there",
|
|
"messages": [_AnyIdHumanMessage(content="hello")],
|
|
}
|
|
|
|
class State(TypedDict):
|
|
hello: str
|
|
messages: Annotated[list[str], add_messages]
|
|
|
|
def node_a(state: State) -> State:
|
|
assert state == {
|
|
"hello": "there",
|
|
"messages": [_AnyIdHumanMessage(content="hello")],
|
|
}
|
|
|
|
builder = StateGraph(State)
|
|
builder.add_node("a", node_a)
|
|
builder.set_entry_point("a")
|
|
builder.set_finish_point("a")
|
|
graph = builder.compile()
|
|
assert graph.invoke({"hello": "there", "messages": "hello"}) == {
|
|
"hello": "there",
|
|
"messages": [_AnyIdHumanMessage(content="hello")],
|
|
}
|
|
|
|
class State(TypedDict):
|
|
hello: str
|
|
messages: Annotated[Sequence[str], add_messages]
|
|
|
|
def node_a(state: State) -> State:
|
|
assert state == {
|
|
"hello": "there",
|
|
"messages": [_AnyIdHumanMessage(content="hello")],
|
|
}
|
|
|
|
builder = StateGraph(State)
|
|
builder.add_node("a", node_a)
|
|
builder.set_entry_point("a")
|
|
builder.set_finish_point("a")
|
|
graph = builder.compile()
|
|
assert graph.invoke({"hello": "there", "messages": "hello"}) == {
|
|
"hello": "there",
|
|
"messages": [_AnyIdHumanMessage(content="hello")],
|
|
}
|
|
|
|
|
|
def test_invoke_single_process_in_out(mocker: MockerFixture) -> None:
|
|
add_one = mocker.Mock(side_effect=lambda x: x + 1)
|
|
chain = NodeBuilder().subscribe_only("input").do(add_one).write_to("output")
|
|
|
|
app = Pregel(
|
|
nodes={
|
|
"one": chain,
|
|
},
|
|
channels={
|
|
"input": LastValue(int),
|
|
"output": LastValue(int),
|
|
},
|
|
input_channels="input",
|
|
output_channels="output",
|
|
)
|
|
|
|
assert app.input_schema.model_json_schema() == {
|
|
"title": "LangGraphInput",
|
|
"type": "integer",
|
|
}
|
|
assert app.output_schema.model_json_schema() == {
|
|
"title": "LangGraphOutput",
|
|
"type": "integer",
|
|
}
|
|
assert app.get_context_jsonschema() is None
|
|
|
|
assert app.invoke(2) == 3
|
|
assert app.invoke(2, output_keys=["output"]) == {"output": 3}
|
|
assert repr(app), "does not raise recursion error"
|
|
|
|
|
|
def test_invoke_single_process_in_write_kwargs(mocker: MockerFixture) -> None:
|
|
add_one = mocker.Mock(side_effect=lambda x: x + 1)
|
|
chain = (
|
|
NodeBuilder()
|
|
.subscribe_only("input")
|
|
.do(add_one)
|
|
.write_to("output", fixed=5, output_plus_one=lambda x: x + 1)
|
|
)
|
|
|
|
app = Pregel(
|
|
nodes={"one": chain},
|
|
channels={
|
|
"input": LastValue(int),
|
|
"output": LastValue(int),
|
|
"fixed": LastValue(int),
|
|
"output_plus_one": LastValue(int),
|
|
},
|
|
output_channels=["output", "fixed", "output_plus_one"],
|
|
input_channels="input",
|
|
)
|
|
|
|
assert app.input_schema.model_json_schema() == {
|
|
"title": "LangGraphInput",
|
|
"type": "integer",
|
|
}
|
|
assert app.output_schema.model_json_schema() == {
|
|
"title": "LangGraphOutput",
|
|
"type": "object",
|
|
"properties": {
|
|
"output": {"title": "Output", "type": "integer", "default": None},
|
|
"fixed": {"title": "Fixed", "type": "integer", "default": None},
|
|
"output_plus_one": {
|
|
"title": "Output Plus One",
|
|
"type": "integer",
|
|
"default": None,
|
|
},
|
|
},
|
|
}
|
|
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 = NodeBuilder().subscribe_only("input").do(add_one).write_to("output")
|
|
|
|
app = Pregel(
|
|
nodes={"one": chain},
|
|
channels={"input": LastValue(int), "output": LastValue(int)},
|
|
input_channels="input",
|
|
output_channels=["output"],
|
|
)
|
|
|
|
assert app.input_schema.model_json_schema() == {
|
|
"title": "LangGraphInput",
|
|
"type": "integer",
|
|
}
|
|
assert app.output_schema.model_json_schema() == {
|
|
"title": "LangGraphOutput",
|
|
"type": "object",
|
|
"properties": {
|
|
"output": {"title": "Output", "type": "integer", "default": None}
|
|
},
|
|
}
|
|
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 = NodeBuilder().subscribe_only("input").do(add_one).write_to("output")
|
|
|
|
app = Pregel(
|
|
nodes={"one": chain},
|
|
channels={"input": LastValue(int), "output": LastValue(int)},
|
|
input_channels=["input"],
|
|
output_channels=["output"],
|
|
)
|
|
assert app.input_schema.model_json_schema() == {
|
|
"title": "LangGraphInput",
|
|
"type": "object",
|
|
"properties": {"input": {"title": "Input", "type": "integer", "default": None}},
|
|
}
|
|
assert app.output_schema.model_json_schema() == {
|
|
"title": "LangGraphOutput",
|
|
"type": "object",
|
|
"properties": {
|
|
"output": {"title": "Output", "type": "integer", "default": None}
|
|
},
|
|
}
|
|
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 = NodeBuilder().subscribe_only("input").do(add_one).write_to("inbox")
|
|
two = NodeBuilder().subscribe_only("inbox").do(add_one).write_to("output")
|
|
|
|
app = Pregel(
|
|
nodes={"one": one, "two": two},
|
|
channels={
|
|
"inbox": LastValue(int),
|
|
"output": LastValue(int),
|
|
"input": LastValue(int),
|
|
},
|
|
input_channels="input",
|
|
output_channels="output",
|
|
)
|
|
|
|
assert app.invoke(2) == 4
|
|
|
|
with pytest.raises(GraphRecursionError):
|
|
app.invoke(2, {"recursion_limit": 1}, debug=1)
|
|
|
|
|
|
def test_run_from_checkpoint_id_retains_previous_writes(
|
|
sync_checkpointer: BaseCheckpointSaver,
|
|
) -> None:
|
|
class MyState(TypedDict):
|
|
myval: Annotated[int, operator.add]
|
|
otherval: bool
|
|
|
|
class Anode:
|
|
def __init__(self):
|
|
self.switch = False
|
|
|
|
def __call__(self, state: MyState):
|
|
self.switch = not self.switch
|
|
return {"myval": 2 if self.switch else 1, "otherval": self.switch}
|
|
|
|
builder = StateGraph(MyState)
|
|
thenode = Anode() # Fun.
|
|
builder.add_node("node_one", thenode)
|
|
builder.add_node("node_two", thenode)
|
|
builder.add_edge(START, "node_one")
|
|
|
|
def _getedge(src: str):
|
|
swap = "node_one" if src == "node_two" else "node_two"
|
|
|
|
def _edge(st: MyState) -> Literal["__end__", "node_one", "node_two"]:
|
|
if st["myval"] > 3:
|
|
return END
|
|
if st["otherval"]:
|
|
return swap
|
|
return src
|
|
|
|
return _edge
|
|
|
|
builder.add_conditional_edges("node_one", _getedge("node_one"))
|
|
builder.add_conditional_edges("node_two", _getedge("node_two"))
|
|
graph = builder.compile(checkpointer=sync_checkpointer)
|
|
|
|
thread_id = uuid.uuid4()
|
|
thread1 = {"configurable": {"thread_id": str(thread_id)}}
|
|
|
|
result = graph.invoke({"myval": 1}, thread1, durability="async")
|
|
assert result["myval"] == 4
|
|
history = [c for c in graph.get_state_history(thread1)]
|
|
|
|
assert len(history) == 4
|
|
assert history[-1].values == {"myval": 0}
|
|
assert history[0].values == {"myval": 4, "otherval": False}
|
|
|
|
second_run_config = {
|
|
**thread1,
|
|
"configurable": {
|
|
**thread1["configurable"],
|
|
"checkpoint_id": history[1].config["configurable"]["checkpoint_id"],
|
|
},
|
|
}
|
|
second_result = graph.invoke(None, second_run_config)
|
|
assert second_result == {"myval": 5, "otherval": True}
|
|
|
|
new_history = [
|
|
c
|
|
for c in graph.get_state_history(
|
|
{"configurable": {"thread_id": str(thread_id), "checkpoint_ns": ""}}
|
|
)
|
|
]
|
|
|
|
# +2: one fork checkpoint from time travel, one from the new execution
|
|
assert len(new_history) == len(history) + 2
|
|
# new_history[0] is the new execution result, new_history[1] is the fork
|
|
assert new_history[1].metadata["source"] == "fork"
|
|
for original, new in zip(history, new_history[2:]):
|
|
assert original.values == new.values
|
|
assert original.next == new.next
|
|
assert original.metadata["step"] == new.metadata["step"]
|
|
|
|
def _get_tasks(hist: list, start: int):
|
|
return [h.tasks for h in hist[start:]]
|
|
|
|
assert _get_tasks(new_history, 2) == _get_tasks(history, 0)
|
|
|
|
|
|
def test_batch_two_processes_in_out() -> None:
|
|
def add_one_with_delay(inp: int) -> int:
|
|
time.sleep(inp / 10)
|
|
return inp + 1
|
|
|
|
one = NodeBuilder().subscribe_only("input").do(add_one_with_delay).write_to("one")
|
|
two = NodeBuilder().subscribe_only("one").do(add_one_with_delay).write_to("output")
|
|
|
|
app = Pregel(
|
|
nodes={"one": one, "two": two},
|
|
channels={
|
|
"one": LastValue(int),
|
|
"output": LastValue(int),
|
|
"input": LastValue(int),
|
|
},
|
|
input_channels="input",
|
|
output_channels="output",
|
|
)
|
|
|
|
assert app.batch([3, 2, 1, 3, 5]) == [5, 4, 3, 5, 7]
|
|
assert app.batch([3, 2, 1, 3, 5], output_keys=["output"]) == [
|
|
{"output": 5},
|
|
{"output": 4},
|
|
{"output": 3},
|
|
{"output": 5},
|
|
{"output": 7},
|
|
]
|
|
|
|
|
|
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": NodeBuilder().subscribe_only("input").do(add_one).write_to("-1")}
|
|
for i in range(test_size - 2):
|
|
nodes[str(i)] = (
|
|
NodeBuilder().subscribe_only(str(i - 1)).do(add_one).write_to(str(i))
|
|
)
|
|
nodes["last"] = NodeBuilder().subscribe_only(str(i)).do(add_one).write_to("output")
|
|
|
|
app = Pregel(
|
|
nodes=nodes,
|
|
channels={str(i): LastValue(int) for i in range(-1, test_size - 2)}
|
|
| {"input": LastValue(int), "output": LastValue(int)},
|
|
input_channels="input",
|
|
output_channels="output",
|
|
)
|
|
|
|
for _ in range(10):
|
|
assert app.invoke(2, {"recursion_limit": test_size}) == 2 + test_size
|
|
|
|
with ThreadPoolExecutor() as executor:
|
|
assert [
|
|
*executor.map(app.invoke, [2] * 10, [{"recursion_limit": test_size}] * 10)
|
|
] == [2 + test_size] * 10
|
|
|
|
|
|
def test_batch_many_processes_in_out(mocker: MockerFixture) -> None:
|
|
test_size = 100
|
|
add_one = mocker.Mock(side_effect=lambda x: x + 1)
|
|
|
|
nodes = {"-1": NodeBuilder().subscribe_only("input").do(add_one).write_to("-1")}
|
|
for i in range(test_size - 2):
|
|
nodes[str(i)] = (
|
|
NodeBuilder().subscribe_only(str(i - 1)).do(add_one).write_to(str(i))
|
|
)
|
|
nodes["last"] = NodeBuilder().subscribe_only(str(i)).do(add_one).write_to("output")
|
|
|
|
app = Pregel(
|
|
nodes=nodes,
|
|
channels={str(i): LastValue(int) for i in range(-1, test_size - 2)}
|
|
| {"input": LastValue(int), "output": LastValue(int)},
|
|
input_channels="input",
|
|
output_channels="output",
|
|
)
|
|
|
|
for _ in range(3):
|
|
assert app.batch([2, 1, 3, 4, 5], {"recursion_limit": test_size}) == [
|
|
2 + test_size,
|
|
1 + test_size,
|
|
3 + test_size,
|
|
4 + test_size,
|
|
5 + test_size,
|
|
]
|
|
|
|
with ThreadPoolExecutor() as executor:
|
|
assert [
|
|
*executor.map(
|
|
app.batch, [[2, 1, 3, 4, 5]] * 3, [{"recursion_limit": test_size}] * 3
|
|
)
|
|
] == [
|
|
[2 + test_size, 1 + test_size, 3 + test_size, 4 + test_size, 5 + test_size]
|
|
] * 3
|
|
|
|
|
|
def test_invoke_two_processes_two_in_two_out_invalid(mocker: MockerFixture) -> None:
|
|
add_one = mocker.Mock(side_effect=lambda x: x + 1)
|
|
|
|
one = NodeBuilder().subscribe_only("input").do(add_one).write_to("output")
|
|
two = NodeBuilder().subscribe_only("input").do(add_one).write_to("output")
|
|
|
|
app = Pregel(
|
|
nodes={"one": one, "two": two},
|
|
channels={"output": LastValue(int), "input": LastValue(int)},
|
|
input_channels="input",
|
|
output_channels="output",
|
|
)
|
|
|
|
with pytest.raises(InvalidUpdateError):
|
|
# LastValue channels can only be updated once per iteration
|
|
app.invoke(2)
|
|
|
|
class State(TypedDict):
|
|
hello: str
|
|
|
|
def my_node(input: State) -> State:
|
|
return {"hello": "world"}
|
|
|
|
builder = StateGraph(State)
|
|
builder.add_node("one", my_node)
|
|
builder.add_node("two", my_node)
|
|
builder.set_conditional_entry_point(lambda _: ["one", "two"])
|
|
|
|
graph = builder.compile()
|
|
with pytest.raises(InvalidUpdateError, match="At key 'hello'"):
|
|
graph.invoke({"hello": "there"}, debug=True)
|
|
|
|
|
|
def test_invoke_two_processes_two_in_two_out_valid(mocker: MockerFixture) -> None:
|
|
add_one = mocker.Mock(side_effect=lambda x: x + 1)
|
|
|
|
one = NodeBuilder().subscribe_only("input").do(add_one).write_to("output")
|
|
two = NodeBuilder().subscribe_only("input").do(add_one).write_to("output")
|
|
|
|
app = Pregel(
|
|
nodes={"one": one, "two": two},
|
|
channels={
|
|
"input": LastValue(int),
|
|
"output": Topic(int),
|
|
},
|
|
input_channels="input",
|
|
output_channels="output",
|
|
)
|
|
|
|
# An Inbox channel accumulates updates into a sequence
|
|
assert app.invoke(2) == [3, 3]
|
|
|
|
|
|
def test_invoke_checkpoint_two(
|
|
mocker: MockerFixture, sync_checkpointer: BaseCheckpointSaver
|
|
) -> None:
|
|
add_one = mocker.Mock(side_effect=lambda x: x["total"] + x["input"])
|
|
errored_once = False
|
|
|
|
def raise_if_above_10(input: int) -> int:
|
|
nonlocal errored_once
|
|
if input > 4:
|
|
if errored_once:
|
|
pass
|
|
else:
|
|
errored_once = True
|
|
raise ConnectionError("I will be retried")
|
|
if input > 10:
|
|
raise ValueError("Input is too large")
|
|
return input
|
|
|
|
one = (
|
|
NodeBuilder()
|
|
.subscribe_to("input")
|
|
.read_from("total")
|
|
.do(add_one)
|
|
.write_to("output", "total")
|
|
.do(raise_if_above_10)
|
|
)
|
|
|
|
app = Pregel(
|
|
nodes={"one": one},
|
|
channels={
|
|
"total": BinaryOperatorAggregate(int, operator.add),
|
|
"input": LastValue(int),
|
|
"output": LastValue(int),
|
|
},
|
|
input_channels="input",
|
|
output_channels="output",
|
|
checkpointer=sync_checkpointer,
|
|
retry_policy=RetryPolicy(),
|
|
)
|
|
|
|
# total starts out as 0, so output is 0+2=2
|
|
assert app.invoke(2, {"configurable": {"thread_id": "1"}}) == 2
|
|
checkpoint = sync_checkpointer.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
|
|
assert errored_once, "errored and retried"
|
|
checkpoint_tup = sync_checkpointer.get_tuple({"configurable": {"thread_id": "1"}})
|
|
assert checkpoint_tup is not None
|
|
assert checkpoint_tup.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, error is recorded
|
|
checkpoint_tup = sync_checkpointer.get_tuple({"configurable": {"thread_id": "1"}})
|
|
assert checkpoint_tup is not None
|
|
assert checkpoint_tup.checkpoint["channel_values"].get("total") == 7
|
|
assert checkpoint_tup.pending_writes == [
|
|
(AnyStr(), ERROR, "ValueError('Input is too large')")
|
|
]
|
|
# 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 = sync_checkpointer.get({"configurable": {"thread_id": "1"}})
|
|
assert checkpoint is not None
|
|
assert checkpoint["channel_values"].get("total") == 7
|
|
checkpoint = sync_checkpointer.get({"configurable": {"thread_id": "2"}})
|
|
assert checkpoint is not None
|
|
assert checkpoint["channel_values"].get("total") == 5
|
|
|
|
|
|
def test_pending_writes_resume(
|
|
sync_checkpointer: BaseCheckpointSaver, durability: Durability
|
|
) -> None:
|
|
class State(TypedDict):
|
|
value: Annotated[int, operator.add]
|
|
|
|
class AwhileMaker:
|
|
def __init__(self, sleep: float, rtn: dict | Exception) -> None:
|
|
self.sleep = sleep
|
|
self.rtn = rtn
|
|
self.reset()
|
|
|
|
def __call__(self, input: State) -> Any:
|
|
self.calls += 1
|
|
time.sleep(self.sleep)
|
|
if isinstance(self.rtn, Exception):
|
|
raise self.rtn
|
|
else:
|
|
return self.rtn
|
|
|
|
def reset(self):
|
|
self.calls = 0
|
|
|
|
one = AwhileMaker(0.1, {"value": 2})
|
|
two = AwhileMaker(0.2, ConnectionError("I'm not good"))
|
|
builder = StateGraph(State)
|
|
builder.add_node("one", one)
|
|
builder.add_node(
|
|
"two",
|
|
two,
|
|
retry_policy=RetryPolicy(max_attempts=2, initial_interval=0, jitter=False),
|
|
)
|
|
builder.add_edge(START, "one")
|
|
builder.add_edge(START, "two")
|
|
graph = builder.compile(checkpointer=sync_checkpointer)
|
|
|
|
thread1: RunnableConfig = {"configurable": {"thread_id": "1"}}
|
|
with pytest.raises(ConnectionError, match="I'm not good"):
|
|
graph.invoke({"value": 1}, thread1, durability=durability)
|
|
|
|
# both nodes should have been called once
|
|
assert one.calls == 1
|
|
assert two.calls == 2 # two attempts
|
|
|
|
# latest checkpoint should be before nodes "one", "two"
|
|
# but we should have applied the write from "one"
|
|
state = graph.get_state(thread1)
|
|
assert state is not None
|
|
assert state.values == {"value": 3}
|
|
assert state.next == ("two",)
|
|
assert state.tasks == (
|
|
PregelTask(AnyStr(), "one", (PULL, "one"), result={"value": 2}),
|
|
PregelTask(AnyStr(), "two", (PULL, "two"), 'ConnectionError("I\'m not good")'),
|
|
)
|
|
assert state.metadata == {
|
|
"parents": {},
|
|
"source": "loop",
|
|
"step": 0,
|
|
}
|
|
# get_state with checkpoint_id should not apply any pending writes
|
|
state = graph.get_state(state.config)
|
|
assert state is not None
|
|
assert state.values == {"value": 1}
|
|
assert state.next == ("one", "two")
|
|
# should contain pending write of "one"
|
|
checkpoint = sync_checkpointer.get_tuple(thread1)
|
|
assert checkpoint is not None
|
|
# should contain error from "two"
|
|
expected_writes = [
|
|
(AnyStr(), "value", 2),
|
|
(AnyStr(), ERROR, 'ConnectionError("I\'m not good")'),
|
|
]
|
|
assert len(checkpoint.pending_writes) == 2
|
|
assert all(w in expected_writes for w in checkpoint.pending_writes)
|
|
# both non-error pending writes come from same task
|
|
non_error_writes = [w for w in checkpoint.pending_writes if w[1] != ERROR]
|
|
# error write is from the other task
|
|
error_write = next(w for w in checkpoint.pending_writes if w[1] == ERROR)
|
|
assert error_write[0] != non_error_writes[0][0]
|
|
|
|
# resume execution
|
|
with pytest.raises(ConnectionError, match="I'm not good"):
|
|
graph.invoke(None, thread1, durability=durability)
|
|
|
|
# node "one" succeeded previously, so shouldn't be called again
|
|
assert one.calls == 1
|
|
# node "two" should have been called once again
|
|
assert two.calls == 4 # two attempts before + two attempts now
|
|
|
|
# confirm no new checkpoints saved
|
|
state_two = graph.get_state(thread1)
|
|
assert state_two.metadata == state.metadata
|
|
|
|
# resume execution, without exception
|
|
two.rtn = {"value": 3}
|
|
# both the pending write and the new write were applied, 1 + 2 + 3 = 6
|
|
assert graph.invoke(None, thread1, durability=durability) == {"value": 6}
|
|
|
|
# check all final checkpoints
|
|
checkpoints = [c for c in sync_checkpointer.list(thread1)]
|
|
# we should have 3
|
|
assert len(checkpoints) == (3 if durability != "exit" else 2)
|
|
# the last one not too interesting for this test
|
|
assert checkpoints[0] == CheckpointTuple(
|
|
config={
|
|
"configurable": {
|
|
"thread_id": "1",
|
|
"checkpoint_ns": "",
|
|
"checkpoint_id": AnyStr(),
|
|
}
|
|
},
|
|
checkpoint={
|
|
"v": 4,
|
|
"id": AnyStr(),
|
|
"ts": AnyStr(),
|
|
"versions_seen": {
|
|
"one": {
|
|
"branch:to:one": AnyVersion(),
|
|
},
|
|
"two": {
|
|
"branch:to:two": AnyVersion(),
|
|
},
|
|
"__input__": {},
|
|
"__start__": {
|
|
"__start__": AnyVersion(),
|
|
},
|
|
"__interrupt__": {
|
|
"value": AnyVersion(),
|
|
"__start__": AnyVersion(),
|
|
"branch:to:one": AnyVersion(),
|
|
"branch:to:two": AnyVersion(),
|
|
},
|
|
},
|
|
"channel_versions": {
|
|
"value": AnyVersion(),
|
|
"__start__": AnyVersion(),
|
|
"branch:to:one": AnyVersion(),
|
|
"branch:to:two": AnyVersion(),
|
|
},
|
|
"channel_values": {"value": 6},
|
|
"updated_channels": ["value"],
|
|
},
|
|
metadata={
|
|
"parents": {},
|
|
"step": 1,
|
|
"source": "loop",
|
|
},
|
|
parent_config={
|
|
"configurable": {
|
|
"thread_id": "1",
|
|
"checkpoint_ns": "",
|
|
"checkpoint_id": checkpoints[1].config["configurable"]["checkpoint_id"],
|
|
}
|
|
},
|
|
pending_writes=[],
|
|
)
|
|
# the previous one we assert that pending writes contains both
|
|
# - original error
|
|
# - successful writes from resuming after preventing error
|
|
assert checkpoints[1] == CheckpointTuple(
|
|
config={
|
|
"configurable": {
|
|
"thread_id": "1",
|
|
"checkpoint_ns": "",
|
|
"checkpoint_id": AnyStr(),
|
|
}
|
|
},
|
|
checkpoint={
|
|
"v": 4,
|
|
"id": AnyStr(),
|
|
"ts": AnyStr(),
|
|
"versions_seen": {
|
|
"__input__": {},
|
|
"__start__": {
|
|
"__start__": AnyVersion(),
|
|
},
|
|
},
|
|
"channel_versions": {
|
|
"value": AnyVersion(),
|
|
"__start__": AnyVersion(),
|
|
"branch:to:one": AnyVersion(),
|
|
"branch:to:two": AnyVersion(),
|
|
},
|
|
"channel_values": {
|
|
"value": 1,
|
|
"branch:to:one": None,
|
|
"branch:to:two": None,
|
|
},
|
|
"updated_channels": ["branch:to:one", "branch:to:two", "value"],
|
|
},
|
|
metadata={
|
|
"parents": {},
|
|
"step": 0,
|
|
"source": "loop",
|
|
},
|
|
parent_config={
|
|
"configurable": {
|
|
"thread_id": "1",
|
|
"checkpoint_ns": "",
|
|
"checkpoint_id": (
|
|
checkpoints[2].config["configurable"]["checkpoint_id"]
|
|
),
|
|
}
|
|
}
|
|
if durability != "exit"
|
|
else None,
|
|
pending_writes=(
|
|
UnsortedSequence(
|
|
(AnyStr(), "value", 2),
|
|
(AnyStr(), "__error__", 'ConnectionError("I\'m not good")'),
|
|
(AnyStr(), "value", 3),
|
|
)
|
|
if durability != "exit"
|
|
else UnsortedSequence(
|
|
(AnyStr(), "value", 2),
|
|
(AnyStr(), "__error__", 'ConnectionError("I\'m not good")'),
|
|
# the write against the previous checkpoint is not saved, as it is
|
|
# produced in a run where only the next checkpoint (the last) is saved
|
|
)
|
|
),
|
|
)
|
|
if durability == "exit":
|
|
return
|
|
assert checkpoints[2] == CheckpointTuple(
|
|
config={
|
|
"configurable": {
|
|
"thread_id": "1",
|
|
"checkpoint_ns": "",
|
|
"checkpoint_id": AnyStr(),
|
|
}
|
|
},
|
|
checkpoint={
|
|
"v": 4,
|
|
"id": AnyStr(),
|
|
"ts": AnyStr(),
|
|
"versions_seen": {"__input__": {}},
|
|
"channel_versions": {
|
|
"__start__": AnyVersion(),
|
|
},
|
|
"channel_values": {"__start__": {"value": 1}},
|
|
"updated_channels": ["__start__"],
|
|
},
|
|
metadata={
|
|
"parents": {},
|
|
"step": -1,
|
|
"source": "input",
|
|
},
|
|
parent_config=None,
|
|
pending_writes=UnsortedSequence(
|
|
(AnyStr(), "value", 1),
|
|
(AnyStr(), "branch:to:one", None),
|
|
(AnyStr(), "branch:to:two", None),
|
|
),
|
|
)
|
|
|
|
|
|
def test_cond_edge_after_send() -> None:
|
|
class Node:
|
|
def __init__(self, name: str):
|
|
self.name = name
|
|
setattr(self, "__name__", name)
|
|
|
|
def __call__(self, state):
|
|
return [self.name]
|
|
|
|
def send_for_fun(state):
|
|
return [Send("2", state), Send("2", state)]
|
|
|
|
def route_to_three(state) -> Literal["3"]:
|
|
return "3"
|
|
|
|
builder = StateGraph(Annotated[list, operator.add])
|
|
builder.add_node(Node("1"))
|
|
builder.add_node(Node("2"))
|
|
builder.add_node(Node("3"))
|
|
builder.add_edge(START, "1")
|
|
builder.add_conditional_edges("1", send_for_fun)
|
|
builder.add_conditional_edges("2", route_to_three)
|
|
graph = builder.compile()
|
|
assert graph.invoke(["0"]) == ["0", "1", "2", "2", "3"]
|
|
|
|
|
|
def test_concurrent_emit_sends() -> None:
|
|
class Node:
|
|
def __init__(self, name: str):
|
|
self.name = name
|
|
setattr(self, "__name__", name)
|
|
|
|
def __call__(self, state):
|
|
return (
|
|
[self.name]
|
|
if isinstance(state, list)
|
|
else ["|".join((self.name, str(state)))]
|
|
)
|
|
|
|
def send_for_fun(state):
|
|
return [Send("2", 1), Send("2", 2), "3.1"]
|
|
|
|
def send_for_profit(state):
|
|
return [Send("2", 3), Send("2", 4)]
|
|
|
|
def route_to_three(state) -> Literal["3"]:
|
|
return "3"
|
|
|
|
builder = StateGraph(Annotated[list, operator.add])
|
|
builder.add_node(Node("1"))
|
|
builder.add_node(Node("1.1"))
|
|
builder.add_node(Node("2"))
|
|
builder.add_node(Node("3"))
|
|
builder.add_node(Node("3.1"))
|
|
builder.add_edge(START, "1")
|
|
builder.add_edge(START, "1.1")
|
|
builder.add_conditional_edges("1", send_for_fun)
|
|
builder.add_conditional_edges("1.1", send_for_profit)
|
|
builder.add_conditional_edges("2", route_to_three)
|
|
graph = builder.compile()
|
|
assert graph.invoke(["0"]) == [
|
|
"0",
|
|
"1",
|
|
"1.1",
|
|
"3.1",
|
|
"2|1",
|
|
"2|2",
|
|
"2|3",
|
|
"2|4",
|
|
"3",
|
|
]
|
|
|
|
|
|
def test_send_sequences() -> None:
|
|
class Node:
|
|
def __init__(self, name: str):
|
|
self.name = name
|
|
setattr(self, "__name__", name)
|
|
|
|
def __call__(self, state):
|
|
update = (
|
|
[self.name]
|
|
if isinstance(state, list)
|
|
else ["|".join((self.name, str(state)))]
|
|
)
|
|
if isinstance(state, Command):
|
|
return [state, Command(update=update)]
|
|
else:
|
|
return update
|
|
|
|
def send_for_fun(state):
|
|
return [
|
|
Send("2", Command(goto=Send("2", 3))),
|
|
Send("2", Command(goto=Send("2", 4))),
|
|
"3.1",
|
|
]
|
|
|
|
def route_to_three(state) -> Literal["3"]:
|
|
return "3"
|
|
|
|
builder = StateGraph(Annotated[list, operator.add])
|
|
builder.add_node(Node("1"))
|
|
builder.add_node(Node("2"))
|
|
builder.add_node(Node("3"))
|
|
builder.add_node(Node("3.1"))
|
|
builder.add_edge(START, "1")
|
|
builder.add_conditional_edges("1", send_for_fun)
|
|
builder.add_conditional_edges("2", route_to_three)
|
|
graph = builder.compile()
|
|
assert graph.invoke(["0"]) == [
|
|
"0",
|
|
"1",
|
|
"3.1",
|
|
"2|Command(goto=Send(node='2', arg=3))",
|
|
"2|Command(goto=Send(node='2', arg=4))",
|
|
"3",
|
|
"2|3",
|
|
"2|4",
|
|
"3",
|
|
]
|
|
|
|
|
|
def test_imp_task(
|
|
sync_checkpointer: BaseCheckpointSaver, durability: Durability
|
|
) -> None:
|
|
mapper_calls = 0
|
|
|
|
class Context(TypedDict):
|
|
model: str
|
|
|
|
@task()
|
|
def mapper(input: int) -> str:
|
|
nonlocal mapper_calls
|
|
mapper_calls += 1
|
|
time.sleep(input / 100)
|
|
return str(input) * 2
|
|
|
|
@entrypoint(checkpointer=sync_checkpointer, context_schema=Context)
|
|
def graph(input: list[int]) -> list[str]:
|
|
futures = [mapper(i) for i in input]
|
|
mapped = [f.result() for f in futures]
|
|
answer = interrupt("question")
|
|
return [m + answer for m in mapped]
|
|
|
|
assert graph.get_input_jsonschema() == {
|
|
"type": "array",
|
|
"items": {"type": "integer"},
|
|
"title": "LangGraphInput",
|
|
}
|
|
assert graph.get_output_jsonschema() == {
|
|
"type": "array",
|
|
"items": {"type": "string"},
|
|
"title": "LangGraphOutput",
|
|
}
|
|
assert graph.get_context_jsonschema() == {
|
|
"properties": {"model": {"title": "Model", "type": "string"}},
|
|
"required": ["model"],
|
|
"title": "Context",
|
|
"type": "object",
|
|
}
|
|
|
|
thread1 = {"configurable": {"thread_id": "1"}}
|
|
result = [*graph.stream([0, 1], thread1, durability=durability)]
|
|
# mapper tasks run concurrently so output order is non-deterministic
|
|
assert sorted(result[:-1], key=lambda d: str(d)) == [
|
|
{"mapper": "00"},
|
|
{"mapper": "11"},
|
|
]
|
|
assert result[-1] == {
|
|
"__interrupt__": (
|
|
Interrupt(
|
|
value="question",
|
|
id=AnyStr(),
|
|
),
|
|
)
|
|
}
|
|
assert mapper_calls == 2
|
|
|
|
assert graph.invoke(Command(resume="answer"), thread1, durability=durability) == [
|
|
"00answer",
|
|
"11answer",
|
|
]
|
|
assert mapper_calls == 2
|
|
|
|
|
|
def test_imp_nested(
|
|
sync_checkpointer: BaseCheckpointSaver, durability: Durability
|
|
) -> None:
|
|
def mynode(input: list[str]) -> list[str]:
|
|
return [it + "a" for it in input]
|
|
|
|
builder = StateGraph(list[str])
|
|
builder.add_node(mynode)
|
|
builder.add_edge(START, "mynode")
|
|
add_a = builder.compile()
|
|
|
|
@task
|
|
def submapper(input: int) -> str:
|
|
time.sleep(input / 100)
|
|
return str(input)
|
|
|
|
@task()
|
|
def mapper(input: int) -> str:
|
|
sub = submapper(input)
|
|
time.sleep(input / 100)
|
|
return sub.result() * 2
|
|
|
|
@entrypoint(checkpointer=sync_checkpointer)
|
|
def graph(input: list[int]) -> list[str]:
|
|
futures = [mapper(i) for i in input]
|
|
mapped = [f.result() for f in futures]
|
|
answer = interrupt("question")
|
|
final = [m + answer for m in mapped]
|
|
return add_a.invoke(final)
|
|
|
|
assert graph.get_input_jsonschema() == {
|
|
"type": "array",
|
|
"items": {"type": "integer"},
|
|
"title": "LangGraphInput",
|
|
}
|
|
assert graph.get_output_jsonschema() == {
|
|
"type": "array",
|
|
"items": {"type": "string"},
|
|
"title": "LangGraphOutput",
|
|
}
|
|
|
|
thread1 = {"configurable": {"thread_id": "1"}}
|
|
result = [*graph.stream([0, 1], thread1, durability=durability)]
|
|
# nested tasks run concurrently so output order is non-deterministic
|
|
assert sorted(result[:-1], key=lambda d: str(d)) == [
|
|
{"mapper": "00"},
|
|
{"mapper": "11"},
|
|
{"submapper": "0"},
|
|
{"submapper": "1"},
|
|
]
|
|
assert result[-1] == {
|
|
"__interrupt__": (
|
|
Interrupt(
|
|
value="question",
|
|
id=AnyStr(),
|
|
),
|
|
)
|
|
}
|
|
|
|
assert graph.invoke(Command(resume="answer"), thread1, durability=durability) == [
|
|
"00answera",
|
|
"11answera",
|
|
]
|
|
|
|
|
|
def test_imp_stream_order(
|
|
sync_checkpointer: BaseCheckpointSaver, durability: Durability
|
|
) -> None:
|
|
@task()
|
|
def foo(state: dict) -> tuple:
|
|
return state["a"] + "foo", "bar"
|
|
|
|
@task
|
|
def bar(a: str, b: str, c: str | None = None) -> dict:
|
|
return {"a": a + b, "c": (c or "") + "bark"}
|
|
|
|
@task
|
|
def baz(state: dict) -> dict:
|
|
return {"a": state["a"] + "baz", "c": "something else"}
|
|
|
|
@entrypoint(checkpointer=sync_checkpointer)
|
|
def graph(state: dict) -> dict:
|
|
fut_foo = foo(state)
|
|
fut_bar = bar(*fut_foo.result())
|
|
fut_baz = baz(fut_bar.result())
|
|
return fut_baz.result()
|
|
|
|
thread1 = {"configurable": {"thread_id": "1"}}
|
|
assert [c for c in graph.stream({"a": "0"}, thread1, durability=durability)] == [
|
|
{
|
|
"foo": (
|
|
"0foo",
|
|
"bar",
|
|
)
|
|
},
|
|
{"bar": {"a": "0foobar", "c": "bark"}},
|
|
{"baz": {"a": "0foobarbaz", "c": "something else"}},
|
|
{"graph": {"a": "0foobarbaz", "c": "something else"}},
|
|
]
|
|
|
|
assert graph.get_state(thread1).values == {"a": "0foobarbaz", "c": "something else"}
|
|
|
|
|
|
def test_invoke_checkpoint_three(
|
|
mocker: MockerFixture, sync_checkpointer: BaseCheckpointSaver
|
|
) -> None:
|
|
adder = mocker.Mock(side_effect=lambda x: x["total"] + x["input"])
|
|
|
|
def raise_if_above_10(input: int) -> int:
|
|
if input > 10:
|
|
raise ValueError("Input is too large")
|
|
return input
|
|
|
|
one = (
|
|
NodeBuilder()
|
|
.subscribe_to("input")
|
|
.read_from("total")
|
|
.do(adder)
|
|
.write_to("output", "total")
|
|
.do(raise_if_above_10)
|
|
)
|
|
|
|
app = Pregel(
|
|
nodes={"one": one},
|
|
channels={
|
|
"total": BinaryOperatorAggregate(int, operator.add),
|
|
"input": LastValue(int),
|
|
"output": LastValue(int),
|
|
},
|
|
input_channels="input",
|
|
output_channels="output",
|
|
checkpointer=sync_checkpointer,
|
|
)
|
|
|
|
thread_1 = {"configurable": {"thread_id": "1"}}
|
|
# total starts out as 0, so output is 0+2=2
|
|
assert app.invoke(2, thread_1, durability="async") == 2
|
|
state = app.get_state(thread_1)
|
|
assert state is not None
|
|
assert state.values.get("total") == 2
|
|
assert state.next == ()
|
|
assert (
|
|
state.config["configurable"]["checkpoint_id"]
|
|
== sync_checkpointer.get(thread_1)["id"]
|
|
)
|
|
# total is now 2, so output is 2+3=5
|
|
assert app.invoke(3, thread_1, durability="async") == 5
|
|
state = app.get_state(thread_1)
|
|
assert state is not None
|
|
assert state.values.get("total") == 7
|
|
assert (
|
|
state.config["configurable"]["checkpoint_id"]
|
|
== sync_checkpointer.get(thread_1)["id"]
|
|
)
|
|
# 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, durability="async")
|
|
# checkpoint is updated with new input
|
|
state = app.get_state(thread_1)
|
|
assert state is not None
|
|
assert state.values.get("total") == 7
|
|
assert state.next == ("one",)
|
|
"""we checkpoint inputs and it failed on "one", so the next node is one"""
|
|
# we can recover from error by sending new inputs
|
|
assert app.invoke(2, thread_1, durability="async") == 9
|
|
state = app.get_state(thread_1)
|
|
assert state is not None
|
|
assert state.values.get("total") == 16, "total is now 7+9=16"
|
|
assert state.next == ()
|
|
|
|
thread_2 = {"configurable": {"thread_id": "2"}}
|
|
# on a new thread, total starts out as 0, so output is 0+5=5
|
|
assert app.invoke(5, thread_2) == 5
|
|
state = app.get_state(thread_1)
|
|
assert state is not None
|
|
assert state.values.get("total") == 16
|
|
assert state.next == (), "checkpoint of other thread not touched"
|
|
state = app.get_state(thread_2)
|
|
assert state is not None
|
|
assert state.values.get("total") == 5
|
|
assert state.next == ()
|
|
|
|
assert len(list(app.get_state_history(thread_1, limit=1))) == 1
|
|
# list all checkpoints for thread 1
|
|
thread_1_history = [c for c in app.get_state_history(thread_1)]
|
|
# there are 7 checkpoints
|
|
assert len(thread_1_history) == 7
|
|
assert Counter(c.metadata["source"] for c in thread_1_history) == {
|
|
"input": 4,
|
|
"loop": 3,
|
|
}
|
|
# sorted descending
|
|
assert (
|
|
thread_1_history[0].config["configurable"]["checkpoint_id"]
|
|
> thread_1_history[1].config["configurable"]["checkpoint_id"]
|
|
)
|
|
# cursor pagination
|
|
cursored = list(
|
|
app.get_state_history(thread_1, limit=1, before=thread_1_history[0].config)
|
|
)
|
|
assert len(cursored) == 1
|
|
assert cursored[0].config == thread_1_history[1].config
|
|
# the last checkpoint
|
|
assert thread_1_history[0].values["total"] == 16
|
|
# the first "loop" checkpoint
|
|
assert thread_1_history[-2].values["total"] == 2
|
|
# can get each checkpoint using aget with config
|
|
assert (
|
|
sync_checkpointer.get(thread_1_history[0].config)["id"]
|
|
== thread_1_history[0].config["configurable"]["checkpoint_id"]
|
|
)
|
|
assert (
|
|
sync_checkpointer.get(thread_1_history[1].config)["id"]
|
|
== thread_1_history[1].config["configurable"]["checkpoint_id"]
|
|
)
|
|
|
|
thread_1_next_config = app.update_state(thread_1_history[1].config, 10)
|
|
# update creates a new checkpoint
|
|
assert (
|
|
thread_1_next_config["configurable"]["checkpoint_id"]
|
|
> thread_1_history[0].config["configurable"]["checkpoint_id"]
|
|
)
|
|
# update makes new checkpoint child of the previous one
|
|
assert (
|
|
app.get_state(thread_1_next_config).parent_config == thread_1_history[1].config
|
|
)
|
|
# 1 more checkpoint in history
|
|
assert len(list(app.get_state_history(thread_1))) == 8
|
|
assert Counter(c.metadata["source"] for c in app.get_state_history(thread_1)) == {
|
|
"update": 1,
|
|
"input": 4,
|
|
"loop": 3,
|
|
}
|
|
# the latest checkpoint is the updated one
|
|
assert app.get_state(thread_1) == app.get_state(thread_1_next_config)
|
|
|
|
|
|
def test_invoke_two_processes_two_in_join_two_out(mocker: MockerFixture) -> None:
|
|
add_one = mocker.Mock(side_effect=lambda x: x + 1)
|
|
add_10_each = mocker.Mock(side_effect=lambda x: sorted(y + 10 for y in x))
|
|
|
|
one = NodeBuilder().subscribe_only("input").do(add_one).write_to("inbox")
|
|
chain_three = NodeBuilder().subscribe_only("input").do(add_one).write_to("inbox")
|
|
chain_four = (
|
|
NodeBuilder().subscribe_only("inbox").do(add_10_each).write_to("output")
|
|
)
|
|
|
|
app = Pregel(
|
|
nodes={
|
|
"one": one,
|
|
"chain_three": chain_three,
|
|
"chain_four": chain_four,
|
|
},
|
|
channels={
|
|
"inbox": Topic(int),
|
|
"output": LastValue(int),
|
|
"input": LastValue(int),
|
|
},
|
|
input_channels="input",
|
|
output_channels="output",
|
|
)
|
|
|
|
# Then invoke app
|
|
# We get a single array result as chain_four waits for all publishers to finish
|
|
# before operating on all elements published to topic_two as an array
|
|
for _ in range(100):
|
|
assert app.invoke(2) == [13, 13]
|
|
|
|
with ThreadPoolExecutor() as executor:
|
|
assert [*executor.map(app.invoke, [2] * 100)] == [[13, 13]] * 100
|
|
|
|
|
|
def test_invoke_join_then_call_other_pregel(
|
|
mocker: MockerFixture, sync_checkpointer: BaseCheckpointSaver
|
|
) -> 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": NodeBuilder().subscribe_only("input").do(add_one).write_to("output")
|
|
},
|
|
channels={
|
|
"output": LastValue(int),
|
|
"input": LastValue(int),
|
|
},
|
|
input_channels="input",
|
|
output_channels="output",
|
|
)
|
|
|
|
one = NodeBuilder().subscribe_only("input").do(add_10_each).write_to("inbox_one")
|
|
two = (
|
|
NodeBuilder()
|
|
.subscribe_only("inbox_one")
|
|
.do(inner_app.map())
|
|
.write_to("outbox_one")
|
|
)
|
|
chain_three = NodeBuilder().subscribe_only("outbox_one").do(sum).write_to("output")
|
|
|
|
app = Pregel(
|
|
nodes={
|
|
"one": one,
|
|
"two": two,
|
|
"chain_three": chain_three,
|
|
},
|
|
channels={
|
|
"inbox_one": Topic(int),
|
|
"outbox_one": LastValue(int),
|
|
"output": LastValue(int),
|
|
"input": LastValue(int),
|
|
},
|
|
input_channels="input",
|
|
output_channels="output",
|
|
)
|
|
|
|
for _ in range(10):
|
|
assert app.invoke([2, 3]) == 27
|
|
|
|
with ThreadPoolExecutor() as executor:
|
|
assert [*executor.map(app.invoke, [[2, 3]] * 10)] == [27] * 10
|
|
|
|
# add checkpointer
|
|
app.checkpointer = sync_checkpointer
|
|
# subgraph is called twice in the same node, but that works
|
|
assert app.invoke([2, 3], {"configurable": {"thread_id": "1"}}) == 27
|
|
|
|
# set inner graph checkpointer NeverCheckpoint
|
|
inner_app.checkpointer = False
|
|
# subgraph still called twice, but checkpointing for inner graph is disabled
|
|
assert app.invoke([2, 3], {"configurable": {"thread_id": "1"}}) == 27
|
|
|
|
|
|
def test_invoke_two_processes_one_in_two_out(mocker: MockerFixture) -> None:
|
|
add_one = mocker.Mock(side_effect=lambda x: x + 1)
|
|
|
|
one = (
|
|
NodeBuilder().subscribe_only("input").do(add_one).write_to("output", "between")
|
|
)
|
|
two = NodeBuilder().subscribe_only("between").do(add_one).write_to("output")
|
|
|
|
app = Pregel(
|
|
nodes={"one": one, "two": two},
|
|
channels={
|
|
"input": LastValue(int),
|
|
"between": LastValue(int),
|
|
"output": LastValue(int),
|
|
},
|
|
stream_channels=["output", "between"],
|
|
input_channels="input",
|
|
output_channels="output",
|
|
)
|
|
|
|
assert [c for c in app.stream(2, stream_mode="updates")] == [
|
|
{"one": {"between": 3, "output": 3}},
|
|
{"two": {"output": 4}},
|
|
]
|
|
assert [c for c in app.stream(2)] == [
|
|
{"between": 3, "output": 3},
|
|
{"between": 3, "output": 4},
|
|
]
|
|
|
|
|
|
def test_invoke_two_processes_no_out(mocker: MockerFixture) -> None:
|
|
add_one = mocker.Mock(side_effect=lambda x: x + 1)
|
|
one = NodeBuilder().subscribe_only("input").do(add_one).write_to("between")
|
|
two = NodeBuilder().subscribe_only("between").do(add_one)
|
|
|
|
app = Pregel(
|
|
nodes={"one": one, "two": two},
|
|
channels={
|
|
"input": LastValue(int),
|
|
"between": LastValue(int),
|
|
"output": LastValue(int),
|
|
},
|
|
input_channels="input",
|
|
output_channels="output",
|
|
)
|
|
|
|
# It finishes executing (once no more messages being published)
|
|
# but returns nothing, as nothing was published to OUT topic
|
|
assert app.invoke(2) is None
|
|
|
|
|
|
def test_invoke_two_processes_no_in(mocker: MockerFixture) -> None:
|
|
add_one = mocker.Mock(side_effect=lambda x: x + 1)
|
|
|
|
one = NodeBuilder().subscribe_only("between").do(add_one).write_to("output")
|
|
two = NodeBuilder().subscribe_only("between").do(add_one)
|
|
|
|
with pytest.raises(TypeError):
|
|
Pregel(nodes={"one": one, "two": two})
|
|
|
|
|
|
def test_conditional_entrypoint_to_multiple_state_graph(
|
|
snapshot: SnapshotAssertion,
|
|
) -> None:
|
|
class OverallState(TypedDict):
|
|
locations: list[str]
|
|
results: Annotated[list[str], operator.add]
|
|
|
|
def get_weather(state: OverallState) -> OverallState:
|
|
location = state["location"]
|
|
weather = "sunny" if len(location) > 2 else "cloudy"
|
|
return {"results": [f"It's {weather} in {location}"]}
|
|
|
|
def continue_to_weather(state: OverallState) -> list[Send]:
|
|
return [
|
|
Send("get_weather", {"location": location})
|
|
for location in state["locations"]
|
|
]
|
|
|
|
workflow = StateGraph(OverallState)
|
|
|
|
workflow.add_node("get_weather", get_weather)
|
|
workflow.add_edge("get_weather", END)
|
|
workflow.set_conditional_entry_point(continue_to_weather, path_map=["get_weather"])
|
|
|
|
app = workflow.compile()
|
|
|
|
assert json.dumps(app.get_input_jsonschema()) == snapshot
|
|
assert json.dumps(app.get_output_jsonschema()) == snapshot
|
|
assert json.dumps(app.get_graph().to_json(), indent=2) == snapshot
|
|
assert app.get_graph().draw_mermaid(with_styles=False) == snapshot
|
|
|
|
assert app.invoke({"locations": ["sf", "nyc"]}, debug=True) == {
|
|
"locations": ["sf", "nyc"],
|
|
"results": ["It's cloudy in sf", "It's sunny in nyc"],
|
|
}
|
|
|
|
assert [*app.stream({"locations": ["sf", "nyc"]}, stream_mode="values")][-1] == {
|
|
"locations": ["sf", "nyc"],
|
|
"results": ["It's cloudy in sf", "It's sunny in nyc"],
|
|
}
|
|
|
|
|
|
def test_conditional_state_graph_with_list_edge_inputs(snapshot: SnapshotAssertion):
|
|
class State(TypedDict):
|
|
foo: Annotated[list[str], operator.add]
|
|
|
|
graph_builder = StateGraph(State)
|
|
graph_builder.add_node("A", lambda x: {"foo": ["A"]})
|
|
graph_builder.add_node("B", lambda x: {"foo": ["B"]})
|
|
graph_builder.add_edge(START, "A")
|
|
graph_builder.add_edge(START, "B")
|
|
graph_builder.add_edge(["A", "B"], END)
|
|
|
|
app = graph_builder.compile()
|
|
assert app.invoke({"foo": []}) == {"foo": ["A", "B"]}
|
|
|
|
assert json.dumps(app.get_graph().to_json(), indent=2) == snapshot
|
|
assert app.get_graph().draw_mermaid(with_styles=False) == snapshot
|
|
|
|
|
|
def test_state_graph_w_config_inherited_state_keys(snapshot: SnapshotAssertion) -> None:
|
|
from langchain_core.language_models.fake import FakeStreamingListLLM
|
|
from langchain_core.prompts import PromptTemplate
|
|
from langchain_core.tools import tool
|
|
|
|
class BaseState(TypedDict):
|
|
input: str
|
|
agent_outcome: AgentAction | AgentFinish | None
|
|
|
|
class AgentState(BaseState, total=False):
|
|
intermediate_steps: Annotated[list[tuple[AgentAction, str]], operator.add]
|
|
|
|
assert get_type_hints(AgentState).keys() == {
|
|
"input",
|
|
"agent_outcome",
|
|
"intermediate_steps",
|
|
}
|
|
|
|
class Context(TypedDict, total=False):
|
|
tools: list[str]
|
|
|
|
# Assemble the tools
|
|
@tool()
|
|
def search_api(query: str) -> str:
|
|
"""Searches the API for the query."""
|
|
return f"result for {query}"
|
|
|
|
tools = [search_api]
|
|
|
|
# Construct the agent
|
|
prompt = PromptTemplate.from_template("Hello!")
|
|
|
|
llm = FakeStreamingListLLM(
|
|
responses=[
|
|
"tool:search_api:query",
|
|
"tool:search_api:another",
|
|
"finish:answer",
|
|
]
|
|
)
|
|
|
|
def agent_parser(input: str) -> dict[str, 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
|
|
builder = StateGraph(AgentState, Context)
|
|
|
|
builder.add_node("agent", agent)
|
|
builder.add_node("tools", execute_tools)
|
|
|
|
builder.set_entry_point("agent")
|
|
|
|
builder.add_conditional_edges(
|
|
"agent", should_continue, {"continue": "tools", "exit": END}
|
|
)
|
|
|
|
builder.add_edge("tools", "agent")
|
|
|
|
app = builder.compile()
|
|
|
|
assert json.dumps(app.get_context_jsonschema()) == snapshot
|
|
assert json.dumps(app.get_input_jsonschema()) == snapshot
|
|
assert json.dumps(app.get_output_jsonschema()) == snapshot
|
|
|
|
assert builder.channels.keys() == {"input", "agent_outcome", "intermediate_steps"}
|
|
|
|
assert app.invoke({"input": "what is weather in sf"}) == {
|
|
"agent_outcome": AgentFinish(
|
|
return_values={"answer": "answer"}, log="finish:answer"
|
|
),
|
|
"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",
|
|
),
|
|
],
|
|
}
|
|
|
|
|
|
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 json.dumps(app.get_input_jsonschema()) == snapshot
|
|
assert json.dumps(app.get_output_jsonschema()) == snapshot
|
|
assert json.dumps(app.get_graph().to_json(), indent=2) == snapshot
|
|
assert app.get_graph().draw_mermaid(with_styles=False) == snapshot
|
|
|
|
assert app.invoke({"input": "what is weather in sf"}) == {
|
|
"input": "what is weather in sf",
|
|
"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_in_one_fan_out_state_graph_waiting_edge(
|
|
snapshot: SnapshotAssertion, sync_checkpointer: BaseCheckpointSaver
|
|
) -> None:
|
|
def sorted_add(x: list[str], y: 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]
|
|
|
|
workflow = StateGraph(State)
|
|
|
|
@workflow.add_node
|
|
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:
|
|
time.sleep(0.1) # to ensure stream order
|
|
return {"docs": ["doc3", "doc4"]}
|
|
|
|
def qa(data: State) -> State:
|
|
return {"answer": ",".join(data["docs"])}
|
|
|
|
workflow.add_node(analyzer_one)
|
|
workflow.add_node(retriever_one)
|
|
workflow.add_node(retriever_two)
|
|
workflow.add_node(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()
|
|
|
|
if isinstance(sync_checkpointer, InMemorySaver):
|
|
assert app.get_graph().draw_mermaid(with_styles=False) == snapshot
|
|
|
|
assert app.invoke({"query": "what is weather in sf"}) == {
|
|
"query": "analyzed: query: what is weather in sf",
|
|
"docs": ["doc1", "doc2", "doc3", "doc4"],
|
|
"answer": "doc1,doc2,doc3,doc4",
|
|
}
|
|
|
|
assert [*app.stream({"query": "what is weather in sf"})] == [
|
|
{"rewrite_query": {"query": "query: what is weather in sf"}},
|
|
{"analyzer_one": {"query": "analyzed: query: what is weather in sf"}},
|
|
{"retriever_two": {"docs": ["doc3", "doc4"]}},
|
|
{"retriever_one": {"docs": ["doc1", "doc2"]}},
|
|
{"qa": {"answer": "doc1,doc2,doc3,doc4"}},
|
|
]
|
|
|
|
app_w_interrupt = workflow.compile(
|
|
checkpointer=sync_checkpointer,
|
|
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"]}},
|
|
{"__interrupt__": ()},
|
|
]
|
|
|
|
assert [c for c in app_w_interrupt.stream(None, config)] == [
|
|
{"qa": {"answer": "doc1,doc2,doc3,doc4"}},
|
|
]
|
|
|
|
app_w_interrupt = workflow.compile(
|
|
checkpointer=sync_checkpointer,
|
|
interrupt_before=["qa"],
|
|
)
|
|
config = {"configurable": {"thread_id": "2"}}
|
|
|
|
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"]}},
|
|
{"__interrupt__": ()},
|
|
]
|
|
|
|
app_w_interrupt.update_state(config, {"docs": ["doc5"]})
|
|
expected_parent_config = list(app_w_interrupt.checkpointer.list(config, limit=2))[
|
|
-1
|
|
].config
|
|
assert app_w_interrupt.get_state(config) == StateSnapshot(
|
|
values={
|
|
"query": "analyzed: query: what is weather in sf",
|
|
"docs": ["doc1", "doc2", "doc3", "doc4", "doc5"],
|
|
},
|
|
tasks=(PregelTask(AnyStr(), "qa", (PULL, "qa")),),
|
|
next=("qa",),
|
|
config={
|
|
"configurable": {
|
|
"thread_id": "2",
|
|
"checkpoint_ns": "",
|
|
"checkpoint_id": AnyStr(),
|
|
}
|
|
},
|
|
created_at=AnyStr(),
|
|
metadata={
|
|
"parents": {},
|
|
"source": "update",
|
|
"step": 4,
|
|
},
|
|
parent_config=expected_parent_config,
|
|
interrupts=(),
|
|
)
|
|
|
|
assert [c for c in app_w_interrupt.stream(None, config, debug=1)] == [
|
|
{"qa": {"answer": "doc1,doc2,doc3,doc4,doc5"}},
|
|
]
|
|
|
|
|
|
@pytest.mark.parametrize("use_waiting_edge", (True, False))
|
|
def test_in_one_fan_out_state_graph_defer_node(
|
|
snapshot: SnapshotAssertion,
|
|
sync_checkpointer: BaseCheckpointSaver,
|
|
use_waiting_edge: bool,
|
|
) -> None:
|
|
def sorted_add(x: list[str], y: 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]
|
|
|
|
workflow = StateGraph(State)
|
|
|
|
@workflow.add_node
|
|
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:
|
|
time.sleep(0.1) # to ensure stream order
|
|
return {"docs": ["doc3", "doc4"]}
|
|
|
|
def qa(data: State) -> State:
|
|
return {"answer": ",".join(data["docs"])}
|
|
|
|
workflow.add_node(analyzer_one)
|
|
workflow.add_node(retriever_one)
|
|
workflow.add_node(retriever_two)
|
|
workflow.add_node(qa, defer=True)
|
|
|
|
workflow.set_entry_point("rewrite_query")
|
|
workflow.add_edge("rewrite_query", "retriever_one")
|
|
workflow.add_edge("retriever_one", "analyzer_one")
|
|
workflow.add_edge("rewrite_query", "retriever_two")
|
|
if use_waiting_edge:
|
|
workflow.add_edge(["retriever_one", "retriever_two"], "qa")
|
|
else:
|
|
workflow.add_edge("retriever_one", "qa")
|
|
workflow.add_edge("retriever_two", "qa")
|
|
workflow.set_finish_point("qa")
|
|
|
|
app = workflow.compile()
|
|
|
|
if isinstance(sync_checkpointer, InMemorySaver):
|
|
assert app.get_graph().draw_mermaid(with_styles=False) == snapshot
|
|
|
|
assert app.invoke({"query": "what is weather in sf"}) == {
|
|
"query": "analyzed: query: what is weather in sf",
|
|
"docs": ["doc1", "doc2", "doc3", "doc4"],
|
|
"answer": "doc1,doc2,doc3,doc4",
|
|
}
|
|
|
|
assert [*app.stream({"query": "what is weather in sf"})] == [
|
|
{"rewrite_query": {"query": "query: what is weather in sf"}},
|
|
{"retriever_one": {"docs": ["doc1", "doc2"]}},
|
|
{"retriever_two": {"docs": ["doc3", "doc4"]}},
|
|
{"analyzer_one": {"query": "analyzed: query: what is weather in sf"}},
|
|
{"qa": {"answer": "doc1,doc2,doc3,doc4"}},
|
|
]
|
|
|
|
assert [*app.stream({"query": "what is weather in sf"}, stream_mode="debug")] == [
|
|
{
|
|
"type": "task",
|
|
"timestamp": AnyStr(),
|
|
"step": 1,
|
|
"payload": {
|
|
"id": AnyStr(),
|
|
"name": "rewrite_query",
|
|
"input": {"query": "what is weather in sf", "docs": []},
|
|
"triggers": ("branch:to:rewrite_query",),
|
|
},
|
|
},
|
|
{
|
|
"type": "task_result",
|
|
"timestamp": AnyStr(),
|
|
"step": 1,
|
|
"payload": {
|
|
"id": AnyStr(),
|
|
"name": "rewrite_query",
|
|
"error": None,
|
|
"result": {
|
|
"query": "query: what is weather in sf",
|
|
},
|
|
"interrupts": [],
|
|
},
|
|
},
|
|
{
|
|
"type": "task",
|
|
"timestamp": AnyStr(),
|
|
"step": 2,
|
|
"payload": {
|
|
"id": AnyStr(),
|
|
"name": "retriever_one",
|
|
"input": {"query": "query: what is weather in sf", "docs": []},
|
|
"triggers": ("branch:to:retriever_one",),
|
|
},
|
|
},
|
|
{
|
|
"type": "task",
|
|
"timestamp": AnyStr(),
|
|
"step": 2,
|
|
"payload": {
|
|
"id": AnyStr(),
|
|
"name": "retriever_two",
|
|
"input": {"query": "query: what is weather in sf", "docs": []},
|
|
"triggers": ("branch:to:retriever_two",),
|
|
},
|
|
},
|
|
{
|
|
"type": "task_result",
|
|
"timestamp": AnyStr(),
|
|
"step": 2,
|
|
"payload": {
|
|
"id": AnyStr(),
|
|
"name": "retriever_one",
|
|
"error": None,
|
|
"result": {
|
|
"docs": ["doc1", "doc2"],
|
|
},
|
|
"interrupts": [],
|
|
},
|
|
},
|
|
{
|
|
"type": "task_result",
|
|
"timestamp": AnyStr(),
|
|
"step": 2,
|
|
"payload": {
|
|
"id": AnyStr(),
|
|
"name": "retriever_two",
|
|
"error": None,
|
|
"result": {
|
|
"docs": ["doc3", "doc4"],
|
|
},
|
|
"interrupts": [],
|
|
},
|
|
},
|
|
{
|
|
"type": "task",
|
|
"timestamp": AnyStr(),
|
|
"step": 3,
|
|
"payload": {
|
|
"id": AnyStr(),
|
|
"name": "analyzer_one",
|
|
"input": {
|
|
"query": "query: what is weather in sf",
|
|
"docs": ["doc1", "doc2", "doc3", "doc4"],
|
|
},
|
|
"triggers": ("branch:to:analyzer_one",),
|
|
},
|
|
},
|
|
{
|
|
"type": "task_result",
|
|
"timestamp": AnyStr(),
|
|
"step": 3,
|
|
"payload": {
|
|
"id": AnyStr(),
|
|
"name": "analyzer_one",
|
|
"error": None,
|
|
"result": {
|
|
"query": "analyzed: query: what is weather in sf",
|
|
},
|
|
"interrupts": [],
|
|
},
|
|
},
|
|
{
|
|
"type": "task",
|
|
"timestamp": AnyStr(),
|
|
"step": 4,
|
|
"payload": {
|
|
"id": AnyStr(),
|
|
"name": "qa",
|
|
"input": {
|
|
"query": "analyzed: query: what is weather in sf",
|
|
"docs": ["doc1", "doc2", "doc3", "doc4"],
|
|
},
|
|
"triggers": ("branch:to:qa", "join:retriever_one+retriever_two:qa")
|
|
if use_waiting_edge
|
|
else ("branch:to:qa",),
|
|
},
|
|
},
|
|
{
|
|
"type": "task_result",
|
|
"timestamp": AnyStr(),
|
|
"step": 4,
|
|
"payload": {
|
|
"id": AnyStr(),
|
|
"name": "qa",
|
|
"error": None,
|
|
"result": {
|
|
"answer": "doc1,doc2,doc3,doc4",
|
|
},
|
|
"interrupts": [],
|
|
},
|
|
},
|
|
]
|
|
|
|
app_w_interrupt = workflow.compile(
|
|
checkpointer=sync_checkpointer,
|
|
interrupt_after=["analyzer_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"}},
|
|
{"retriever_one": {"docs": ["doc1", "doc2"]}},
|
|
{"retriever_two": {"docs": ["doc3", "doc4"]}},
|
|
{"analyzer_one": {"query": "analyzed: query: what is weather in sf"}},
|
|
{"__interrupt__": ()},
|
|
]
|
|
|
|
assert [c for c in app_w_interrupt.stream(None, config)] == [
|
|
{"qa": {"answer": "doc1,doc2,doc3,doc4"}},
|
|
]
|
|
|
|
app_w_interrupt = workflow.compile(
|
|
checkpointer=sync_checkpointer,
|
|
interrupt_before=["qa"],
|
|
)
|
|
config = {"configurable": {"thread_id": "2"}}
|
|
|
|
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"}},
|
|
{"retriever_one": {"docs": ["doc1", "doc2"]}},
|
|
{"retriever_two": {"docs": ["doc3", "doc4"]}},
|
|
{"analyzer_one": {"query": "analyzed: query: what is weather in sf"}},
|
|
{"__interrupt__": ()},
|
|
]
|
|
|
|
app_w_interrupt.update_state(config, {"docs": ["doc5"]})
|
|
expected_parent_config = list(app_w_interrupt.checkpointer.list(config, limit=2))[
|
|
-1
|
|
].config
|
|
assert app_w_interrupt.get_state(config) == StateSnapshot(
|
|
values={
|
|
"query": "analyzed: query: what is weather in sf",
|
|
"docs": ["doc1", "doc2", "doc3", "doc4", "doc5"],
|
|
},
|
|
tasks=(PregelTask(AnyStr(), "qa", (PULL, "qa")),),
|
|
next=("qa",),
|
|
config={
|
|
"configurable": {
|
|
"thread_id": "2",
|
|
"checkpoint_ns": "",
|
|
"checkpoint_id": AnyStr(),
|
|
}
|
|
},
|
|
created_at=AnyStr(),
|
|
metadata={
|
|
"parents": {},
|
|
"source": "update",
|
|
"step": 4,
|
|
},
|
|
parent_config=expected_parent_config,
|
|
interrupts=(),
|
|
)
|
|
|
|
assert [c for c in app_w_interrupt.stream(None, config, debug=1)] == [
|
|
{"qa": {"answer": "doc1,doc2,doc3,doc4,doc5"}},
|
|
]
|
|
|
|
|
|
def test_in_one_fan_out_state_graph_waiting_edge_via_branch(
|
|
snapshot: SnapshotAssertion, sync_checkpointer: BaseCheckpointSaver
|
|
) -> None:
|
|
def sorted_add(x: list[str], y: 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:
|
|
time.sleep(0.1)
|
|
return {"docs": ["doc3", "doc4"]}
|
|
|
|
def qa(data: State) -> State:
|
|
return {"answer": ",".join(data["docs"])}
|
|
|
|
def rewrite_query_then(data: State) -> Literal["retriever_two"]:
|
|
return "retriever_two"
|
|
|
|
workflow = StateGraph(State)
|
|
|
|
workflow.add_node("rewrite_query", rewrite_query)
|
|
workflow.add_node("analyzer_one", analyzer_one)
|
|
workflow.add_node("retriever_one", retriever_one)
|
|
workflow.add_node("retriever_two", retriever_two)
|
|
workflow.add_node("qa", qa)
|
|
|
|
workflow.set_entry_point("rewrite_query")
|
|
workflow.add_edge("rewrite_query", "analyzer_one")
|
|
workflow.add_edge("analyzer_one", "retriever_one")
|
|
workflow.add_conditional_edges("rewrite_query", rewrite_query_then)
|
|
workflow.add_edge(["retriever_one", "retriever_two"], "qa")
|
|
workflow.set_finish_point("qa")
|
|
|
|
app = workflow.compile()
|
|
|
|
if isinstance(sync_checkpointer, InMemorySaver):
|
|
assert app.get_graph().draw_mermaid(with_styles=False) == snapshot
|
|
|
|
assert app.invoke({"query": "what is weather in sf"}, debug=True) == {
|
|
"query": "analyzed: query: what is weather in sf",
|
|
"docs": ["doc1", "doc2", "doc3", "doc4"],
|
|
"answer": "doc1,doc2,doc3,doc4",
|
|
}
|
|
|
|
assert [*app.stream({"query": "what is weather in sf"})] == [
|
|
{"rewrite_query": {"query": "query: what is weather in sf"}},
|
|
{"analyzer_one": {"query": "analyzed: query: what is weather in sf"}},
|
|
{"retriever_two": {"docs": ["doc3", "doc4"]}},
|
|
{"retriever_one": {"docs": ["doc1", "doc2"]}},
|
|
{"qa": {"answer": "doc1,doc2,doc3,doc4"}},
|
|
]
|
|
|
|
app_w_interrupt = workflow.compile(
|
|
checkpointer=sync_checkpointer,
|
|
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"]}},
|
|
{"__interrupt__": ()},
|
|
]
|
|
|
|
assert [c for c in app_w_interrupt.stream(None, config)] == [
|
|
{"qa": {"answer": "doc1,doc2,doc3,doc4"}},
|
|
]
|
|
|
|
|
|
def test_in_one_fan_out_state_graph_waiting_edge_custom_state_class_pydantic2(
|
|
snapshot: SnapshotAssertion,
|
|
sync_checkpointer: BaseCheckpointSaver,
|
|
) -> None:
|
|
def sorted_add(x: list[str], y: 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 InnerObject(BaseModel):
|
|
yo: int
|
|
|
|
class State(BaseModel):
|
|
model_config = ConfigDict(arbitrary_types_allowed=True)
|
|
|
|
query: str
|
|
inner: Annotated[InnerObject, lambda x, y: y]
|
|
answer: str | None = None
|
|
docs: Annotated[list[str], sorted_add]
|
|
|
|
class StateUpdate(BaseModel):
|
|
query: str | None = None
|
|
answer: str | None = None
|
|
docs: list[str] | None = None
|
|
|
|
class UpdateDocs34(BaseModel):
|
|
docs: list[str] = Field(default_factory=lambda: ["doc3", "doc4"])
|
|
|
|
class Input(BaseModel):
|
|
query: str
|
|
inner: InnerObject
|
|
|
|
class Output(BaseModel):
|
|
answer: str
|
|
docs: list[str]
|
|
|
|
def rewrite_query(data: State) -> State:
|
|
assert isinstance(data.inner, InnerObject)
|
|
return {"query": f"query: {data.query}"}
|
|
|
|
def analyzer_one(data: State) -> State:
|
|
assert isinstance(data.inner, InnerObject)
|
|
return StateUpdate(query=f"analyzed: {data.query}")
|
|
|
|
def retriever_one(data: State) -> State:
|
|
return {"docs": ["doc1", "doc2"]}
|
|
|
|
def retriever_two(data: State) -> State:
|
|
time.sleep(0.1)
|
|
return UpdateDocs34()
|
|
|
|
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, input_schema=Input, output_schema=Output)
|
|
|
|
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()
|
|
|
|
if isinstance(sync_checkpointer, InMemorySaver):
|
|
assert app.get_graph().draw_mermaid(with_styles=False) == snapshot
|
|
assert app.get_input_jsonschema() == snapshot
|
|
assert app.get_output_jsonschema() == snapshot
|
|
|
|
with pytest.raises(ValidationError):
|
|
app.invoke({"query": {}})
|
|
|
|
assert app.invoke({"query": "what is weather in sf", "inner": {"yo": 1}}) == {
|
|
"docs": ["doc1", "doc2", "doc3", "doc4"],
|
|
"answer": "doc1,doc2,doc3,doc4",
|
|
}
|
|
|
|
assert [*app.stream({"query": "what is weather in sf", "inner": {"yo": 1}})] == [
|
|
{"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=sync_checkpointer,
|
|
interrupt_after=["retriever_one"],
|
|
)
|
|
config = {"configurable": {"thread_id": "1"}}
|
|
|
|
assert [
|
|
c
|
|
for c in app_w_interrupt.stream(
|
|
{"query": "what is weather in sf", "inner": {"yo": 1}}, 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"]}},
|
|
{"__interrupt__": ()},
|
|
]
|
|
|
|
assert [c for c in app_w_interrupt.stream(None, config)] == [
|
|
{"qa": {"answer": "doc1,doc2,doc3,doc4"}},
|
|
]
|
|
|
|
assert app_w_interrupt.update_state(
|
|
config, {"docs": ["doc5"]}, as_node="rewrite_query"
|
|
) == {
|
|
"configurable": {
|
|
"thread_id": "1",
|
|
"checkpoint_id": AnyStr(),
|
|
"checkpoint_ns": "",
|
|
}
|
|
}
|
|
|
|
|
|
def test_in_one_fan_out_state_graph_waiting_edge_custom_state_class_pydantic_input(
|
|
sync_checkpointer: BaseCheckpointSaver,
|
|
) -> None:
|
|
def sorted_add(x: list[str], y: 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 InnerObject(BaseModel):
|
|
yo: int
|
|
|
|
class QueryModel(BaseModel):
|
|
query: str
|
|
|
|
class State(QueryModel):
|
|
inner: InnerObject
|
|
answer: str | None = None
|
|
docs: Annotated[list[str], sorted_add]
|
|
|
|
class StateUpdate(BaseModel):
|
|
query: str | None = None
|
|
answer: str | None = None
|
|
docs: list[str] | None = None
|
|
|
|
class Input(QueryModel):
|
|
inner: InnerObject
|
|
|
|
class Output(BaseModel):
|
|
answer: str
|
|
docs: list[str]
|
|
|
|
def rewrite_query(data: State) -> State:
|
|
return {"query": f"query: {data.query}"}
|
|
|
|
def analyzer_one(data: State) -> State:
|
|
return StateUpdate(query=f"analyzed: {data.query}")
|
|
|
|
def retriever_one(data: State) -> State:
|
|
return {"docs": ["doc1", "doc2"]}
|
|
|
|
def retriever_two(data: State) -> State:
|
|
time.sleep(0.1)
|
|
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, input_schema=Input, output_schema=Output)
|
|
|
|
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.invoke(
|
|
Input(query="what is weather in sf", inner=InnerObject(yo=1))
|
|
) == {
|
|
"docs": ["doc1", "doc2", "doc3", "doc4"],
|
|
"answer": "doc1,doc2,doc3,doc4",
|
|
}
|
|
|
|
assert [
|
|
*app.stream(Input(query="what is weather in sf", inner=InnerObject(yo=1)))
|
|
] == [
|
|
{"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=sync_checkpointer,
|
|
interrupt_after=["retriever_one"],
|
|
)
|
|
config = {"configurable": {"thread_id": "1"}}
|
|
|
|
assert [
|
|
c
|
|
for c in app_w_interrupt.stream(
|
|
Input(query="what is weather in sf", inner=InnerObject(yo=1)), 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"]}},
|
|
{"__interrupt__": ()},
|
|
]
|
|
|
|
assert [c for c in app_w_interrupt.stream(None, config)] == [
|
|
{"qa": {"answer": "doc1,doc2,doc3,doc4"}},
|
|
]
|
|
|
|
assert app_w_interrupt.update_state(
|
|
config, {"docs": ["doc5"]}, as_node="rewrite_query"
|
|
) == {
|
|
"configurable": {
|
|
"thread_id": "1",
|
|
"checkpoint_id": AnyStr(),
|
|
"checkpoint_ns": "",
|
|
}
|
|
}
|
|
|
|
|
|
def test_in_one_fan_out_state_graph_waiting_edge_plus_regular(
|
|
sync_checkpointer: BaseCheckpointSaver,
|
|
) -> None:
|
|
def sorted_add(x: list[str], y: 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:
|
|
time.sleep(0.1)
|
|
return {"query": f"analyzed: {data['query']}"}
|
|
|
|
def retriever_one(data: State) -> State:
|
|
return {"docs": ["doc1", "doc2"]}
|
|
|
|
def retriever_two(data: State) -> State:
|
|
time.sleep(0.2)
|
|
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",
|
|
}
|
|
|
|
rewrite = {"rewrite_query": {"query": "query: what is weather in sf"}}
|
|
analyzer = {"analyzer_one": {"query": "analyzed: query: what is weather in sf"}}
|
|
empty_qa = {"qa": {"answer": ""}}
|
|
retriever_one_result = {"retriever_one": {"docs": ["doc1", "doc2"]}}
|
|
retriever_two_result = {"retriever_two": {"docs": ["doc3", "doc4"]}}
|
|
|
|
def assert_valid_stream_order(
|
|
chunks: list[dict[str, Any]], terminal: dict[str, Any]
|
|
) -> None:
|
|
assert chunks[0] == rewrite
|
|
assert chunks[-1] == terminal
|
|
middle = chunks[1:-1]
|
|
assert sorted(middle, key=repr) == sorted(
|
|
[empty_qa, analyzer, retriever_one_result, retriever_two_result], key=repr
|
|
)
|
|
assert middle.index(analyzer) < middle.index(retriever_one_result)
|
|
|
|
assert_valid_stream_order(
|
|
[*app.stream({"query": "what is weather in sf"})],
|
|
{"qa": {"answer": "doc1,doc2,doc3,doc4"}},
|
|
)
|
|
|
|
app_w_interrupt = workflow.compile(
|
|
checkpointer=sync_checkpointer,
|
|
interrupt_after=["retriever_one"],
|
|
)
|
|
config = {"configurable": {"thread_id": "1"}}
|
|
|
|
assert_valid_stream_order(
|
|
[c for c in app_w_interrupt.stream({"query": "what is weather in sf"}, config)],
|
|
{"__interrupt__": ()},
|
|
)
|
|
|
|
assert [c for c in app_w_interrupt.stream(None, config)] == [
|
|
{"qa": {"answer": "doc1,doc2,doc3,doc4"}},
|
|
]
|
|
|
|
|
|
@pytest.mark.parametrize("with_cache", [True, False])
|
|
def test_in_one_fan_out_state_graph_waiting_edge_multiple(
|
|
with_cache: bool, cache: BaseCache
|
|
) -> None:
|
|
def sorted_add(x: list[str], y: 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]
|
|
|
|
rewrite_query_count = 0
|
|
|
|
def rewrite_query(data: State) -> State:
|
|
nonlocal rewrite_query_count
|
|
rewrite_query_count += 1
|
|
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:
|
|
time.sleep(0.1)
|
|
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,
|
|
cache_policy=CachePolicy() if with_cache else None,
|
|
)
|
|
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(cache=cache)
|
|
|
|
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"},
|
|
"__metadata__": {"cached": True},
|
|
}
|
|
if with_cache
|
|
else {"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"]}},
|
|
{"decider": None},
|
|
{
|
|
"rewrite_query": {"query": "query: analyzed: query: what is weather in sf"},
|
|
"__metadata__": {"cached": True},
|
|
}
|
|
if with_cache
|
|
else {
|
|
"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"]}},
|
|
{"decider": None},
|
|
{"qa": {"answer": "doc1,doc1,doc2,doc2,doc3,doc3,doc4,doc4"}},
|
|
]
|
|
assert rewrite_query_count == 2 if with_cache else 4
|
|
|
|
# clear the cache
|
|
if with_cache:
|
|
app.clear_cache()
|
|
|
|
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 rewrite_query_count == 4
|
|
|
|
|
|
def test_callable_in_conditional_edges_with_no_path_map() -> None:
|
|
class State(TypedDict, total=False):
|
|
query: str
|
|
|
|
def rewrite(data: State) -> State:
|
|
return {"query": f"query: {data['query']}"}
|
|
|
|
def analyze(data: State) -> State:
|
|
return {"query": f"analyzed: {data['query']}"}
|
|
|
|
class ChooseAnalyzer:
|
|
def __call__(self, data: State) -> str:
|
|
return "analyzer"
|
|
|
|
workflow = StateGraph(State)
|
|
workflow.add_node("rewriter", rewrite)
|
|
workflow.add_node("analyzer", analyze)
|
|
workflow.add_conditional_edges("rewriter", ChooseAnalyzer())
|
|
workflow.set_entry_point("rewriter")
|
|
app = workflow.compile()
|
|
|
|
assert app.invoke({"query": "what is weather in sf"}) == {
|
|
"query": "analyzed: query: what is weather in sf",
|
|
}
|
|
|
|
|
|
def test_function_in_conditional_edges_with_no_path_map() -> None:
|
|
class State(TypedDict, total=False):
|
|
query: str
|
|
|
|
def rewrite(data: State) -> State:
|
|
return {"query": f"query: {data['query']}"}
|
|
|
|
def analyze(data: State) -> State:
|
|
return {"query": f"analyzed: {data['query']}"}
|
|
|
|
def choose_analyzer(data: State) -> str:
|
|
return "analyzer"
|
|
|
|
workflow = StateGraph(State)
|
|
workflow.add_node("rewriter", rewrite)
|
|
workflow.add_node("analyzer", analyze)
|
|
workflow.add_conditional_edges("rewriter", choose_analyzer)
|
|
workflow.set_entry_point("rewriter")
|
|
app = workflow.compile()
|
|
|
|
assert app.invoke({"query": "what is weather in sf"}) == {
|
|
"query": "analyzed: query: what is weather in sf",
|
|
}
|
|
|
|
|
|
def test_in_one_fan_out_state_graph_waiting_edge_multiple_cond_edge() -> None:
|
|
def sorted_add(x: list[str], y: 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:
|
|
time.sleep(0.1)
|
|
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"]}},
|
|
{"decider": None},
|
|
{"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"]}},
|
|
{"decider": None},
|
|
{"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 other(state: State):
|
|
return {"my_key": "_more"}
|
|
|
|
def down(state: State):
|
|
pass
|
|
|
|
graph = StateGraph(State)
|
|
|
|
graph.add_node("up", up)
|
|
graph.add_node("side", side)
|
|
graph.add_node("other", other)
|
|
graph.add_node("down", down)
|
|
|
|
graph.set_entry_point("up")
|
|
graph.add_edge("up", "side")
|
|
graph.add_edge("up", "other")
|
|
graph.add_edge(["up", "side"], "down")
|
|
graph.set_finish_point("down")
|
|
|
|
app = graph.compile()
|
|
|
|
assert app.get_graph().draw_mermaid(with_styles=False) == snapshot
|
|
assert app.invoke({"my_key": "my_value"}) == {"my_key": "my_value_more"}
|
|
assert [*app.stream({"my_key": "my_value"})] in (
|
|
[
|
|
{"up": None},
|
|
{"side": None},
|
|
{"other": {"my_key": "_more"}},
|
|
{"down": None},
|
|
],
|
|
[
|
|
{"up": None},
|
|
{"other": {"my_key": "_more"}},
|
|
{"side": None},
|
|
{"down": None},
|
|
],
|
|
)
|
|
|
|
|
|
def test_nested_graph_xray(snapshot: SnapshotAssertion) -> None:
|
|
class State(TypedDict):
|
|
my_key: Annotated[str, operator.add]
|
|
market: str
|
|
|
|
def logic(state: State):
|
|
pass
|
|
|
|
tool_two_graph = StateGraph(State)
|
|
tool_two_graph.add_node("tool_two_slow", logic)
|
|
tool_two_graph.add_node("tool_two_fast", logic)
|
|
tool_two_graph.set_conditional_entry_point(
|
|
lambda s: "tool_two_slow" if s["market"] == "DE" else "tool_two_fast",
|
|
["tool_two_slow", "tool_two_fast"],
|
|
)
|
|
tool_two = tool_two_graph.compile()
|
|
|
|
graph = StateGraph(State)
|
|
graph.add_node("tool_one", logic)
|
|
graph.add_node("tool_two", tool_two)
|
|
graph.add_node("tool_three", logic)
|
|
graph.set_conditional_entry_point(
|
|
lambda s: "tool_one", ["tool_one", "tool_two", "tool_three"]
|
|
)
|
|
app = graph.compile()
|
|
|
|
assert app.get_graph(xray=True).to_json() == snapshot
|
|
assert app.get_graph(xray=True).draw_mermaid() == snapshot
|
|
|
|
|
|
def test_nested_graph(snapshot: SnapshotAssertion) -> None:
|
|
def never_called_fn(state: Any):
|
|
assert 0, "This function should never be called"
|
|
|
|
never_called = RunnableLambda(never_called_fn)
|
|
|
|
class InnerState(TypedDict):
|
|
my_key: str
|
|
my_other_key: str
|
|
|
|
def up(state: InnerState):
|
|
return {"my_key": state["my_key"] + " there", "my_other_key": state["my_key"]}
|
|
|
|
inner = StateGraph(InnerState)
|
|
inner.add_node("up", up)
|
|
inner.set_entry_point("up")
|
|
inner.set_finish_point("up")
|
|
|
|
class State(TypedDict):
|
|
my_key: str
|
|
never_called: Any
|
|
|
|
def side(state: State):
|
|
return {"my_key": state["my_key"] + " and back again"}
|
|
|
|
graph = StateGraph(State)
|
|
graph.add_node("inner", inner.compile())
|
|
graph.add_node("side", side)
|
|
graph.set_entry_point("inner")
|
|
graph.add_edge("inner", "side")
|
|
graph.set_finish_point("side")
|
|
|
|
app = graph.compile()
|
|
|
|
assert app.get_graph().draw_mermaid(with_styles=False) == snapshot
|
|
assert app.get_graph(xray=True).draw_mermaid() == snapshot
|
|
assert app.invoke(
|
|
{"my_key": "my value", "never_called": never_called},
|
|
print_mode=["values", "updates"],
|
|
) == {
|
|
"my_key": "my value there and back again",
|
|
"never_called": never_called,
|
|
}
|
|
assert [*app.stream({"my_key": "my value", "never_called": never_called})] == [
|
|
{"inner": {"my_key": "my value there"}},
|
|
{"side": {"my_key": "my value there and back again"}},
|
|
]
|
|
assert [
|
|
*app.stream(
|
|
{"my_key": "my value", "never_called": never_called}, stream_mode="values"
|
|
)
|
|
] == [
|
|
{
|
|
"my_key": "my value",
|
|
"never_called": never_called,
|
|
},
|
|
{
|
|
"my_key": "my value there",
|
|
"never_called": never_called,
|
|
},
|
|
{
|
|
"my_key": "my value there and back again",
|
|
"never_called": never_called,
|
|
},
|
|
]
|
|
|
|
chain = app | RunnablePassthrough()
|
|
|
|
assert chain.invoke({"my_key": "my value", "never_called": never_called}) == {
|
|
"my_key": "my value there and back again",
|
|
"never_called": never_called,
|
|
}
|
|
assert [*chain.stream({"my_key": "my value", "never_called": never_called})] == [
|
|
{"inner": {"my_key": "my value there"}},
|
|
{"side": {"my_key": "my value there and back again"}},
|
|
]
|
|
|
|
|
|
def test_subgraph_checkpoint_true(
|
|
sync_checkpointer: BaseCheckpointSaver, durability: Durability
|
|
) -> None:
|
|
class InnerState(TypedDict):
|
|
my_key: Annotated[str, operator.add]
|
|
my_other_key: str
|
|
|
|
def inner_1(state: InnerState):
|
|
return {"my_key": " got here", "my_other_key": state["my_key"]}
|
|
|
|
def inner_2(state: InnerState):
|
|
return {"my_key": " and there"}
|
|
|
|
inner = StateGraph(InnerState)
|
|
inner.add_node("inner_1", inner_1)
|
|
inner.add_node("inner_2", inner_2)
|
|
inner.add_edge("inner_1", "inner_2")
|
|
inner.set_entry_point("inner_1")
|
|
inner.set_finish_point("inner_2")
|
|
|
|
class State(TypedDict):
|
|
my_key: str
|
|
|
|
graph = StateGraph(State)
|
|
graph.add_node("inner", inner.compile(checkpointer=True))
|
|
graph.add_edge(START, "inner")
|
|
graph.add_conditional_edges(
|
|
"inner", lambda s: "inner" if s["my_key"].count("there") < 2 else END
|
|
)
|
|
app = graph.compile(checkpointer=sync_checkpointer)
|
|
|
|
config = {"configurable": {"thread_id": "2"}}
|
|
assert [
|
|
c
|
|
for c in app.stream(
|
|
{"my_key": ""}, config, subgraphs=True, durability=durability
|
|
)
|
|
] == [
|
|
(("inner",), {"inner_1": {"my_key": " got here", "my_other_key": ""}}),
|
|
(("inner",), {"inner_2": {"my_key": " and there"}}),
|
|
((), {"inner": {"my_key": " got here and there"}}),
|
|
(
|
|
("inner",),
|
|
{
|
|
"inner_1": {
|
|
"my_key": " got here",
|
|
"my_other_key": " got here and there got here and there",
|
|
}
|
|
},
|
|
),
|
|
(("inner",), {"inner_2": {"my_key": " and there"}}),
|
|
(
|
|
(),
|
|
{
|
|
"inner": {
|
|
"my_key": " got here and there got here and there got here and there"
|
|
}
|
|
},
|
|
),
|
|
]
|
|
|
|
checkpoints = list(app.get_state_history(config))
|
|
if durability != "exit":
|
|
assert len(checkpoints) == 4
|
|
else:
|
|
assert len(checkpoints) == 1
|
|
|
|
|
|
def test_subgraph_durability_inherited(durability: Durability) -> None:
|
|
sync_checkpointer = InMemorySaver()
|
|
|
|
class InnerState(TypedDict):
|
|
my_key: Annotated[str, operator.add]
|
|
my_other_key: str
|
|
|
|
def inner_1(state: InnerState):
|
|
return {"my_key": " got here", "my_other_key": state["my_key"]}
|
|
|
|
def inner_2(state: InnerState):
|
|
return {"my_key": " and there"}
|
|
|
|
inner = StateGraph(InnerState)
|
|
inner.add_node("inner_1", inner_1)
|
|
inner.add_node("inner_2", inner_2)
|
|
inner.add_edge("inner_1", "inner_2")
|
|
inner.set_entry_point("inner_1")
|
|
inner.set_finish_point("inner_2")
|
|
|
|
class State(TypedDict):
|
|
my_key: str
|
|
|
|
inner_app = inner.compile(checkpointer=sync_checkpointer)
|
|
graph = StateGraph(State)
|
|
graph.add_node("inner", inner_app)
|
|
graph.add_edge(START, "inner")
|
|
graph.add_conditional_edges(
|
|
"inner", lambda s: "inner" if s["my_key"].count("there") < 2 else END
|
|
)
|
|
app = graph.compile(checkpointer=sync_checkpointer)
|
|
thread_id = str(uuid.uuid4())
|
|
config = {"configurable": {"thread_id": thread_id}}
|
|
app.invoke({"my_key": ""}, config, subgraphs=True, durability=durability)
|
|
if durability != "exit":
|
|
checkpoints = list(sync_checkpointer.list(config))
|
|
assert len(checkpoints) == 12
|
|
else:
|
|
checkpoints = list(sync_checkpointer.list(config))
|
|
assert len(checkpoints) == 1
|
|
|
|
|
|
def test_subgraph_checkpoint_true_interrupt(
|
|
sync_checkpointer: BaseCheckpointSaver, durability: Durability
|
|
) -> None:
|
|
# Define subgraph
|
|
class SubgraphState(TypedDict):
|
|
# note that none of these keys are shared with the parent graph state
|
|
bar: str
|
|
baz: str
|
|
|
|
def subgraph_node_1(state: SubgraphState):
|
|
baz_value = interrupt("Provide baz value")
|
|
return {"baz": baz_value}
|
|
|
|
def subgraph_node_2(state: SubgraphState):
|
|
return {"bar": state["bar"] + state["baz"]}
|
|
|
|
subgraph_builder = StateGraph(SubgraphState)
|
|
subgraph_builder.add_node(subgraph_node_1)
|
|
subgraph_builder.add_node(subgraph_node_2)
|
|
subgraph_builder.add_edge(START, "subgraph_node_1")
|
|
subgraph_builder.add_edge("subgraph_node_1", "subgraph_node_2")
|
|
subgraph = subgraph_builder.compile(checkpointer=True)
|
|
|
|
class ParentState(TypedDict):
|
|
foo: str
|
|
|
|
def node_1(state: ParentState):
|
|
return {"foo": "hi! " + state["foo"]}
|
|
|
|
def node_2(state: ParentState):
|
|
response = subgraph.invoke({"bar": state["foo"]})
|
|
return {"foo": response["bar"]}
|
|
|
|
builder = StateGraph(ParentState)
|
|
builder.add_node("node_1", node_1)
|
|
builder.add_node("node_2", node_2)
|
|
builder.add_edge(START, "node_1")
|
|
builder.add_edge("node_1", "node_2")
|
|
|
|
graph = builder.compile(checkpointer=sync_checkpointer)
|
|
config = {"configurable": {"thread_id": "1"}}
|
|
|
|
assert graph.invoke({"foo": "foo"}, config, durability=durability) == {
|
|
"foo": "hi! foo",
|
|
"__interrupt__": [
|
|
Interrupt(
|
|
value="Provide baz value",
|
|
id=AnyStr(),
|
|
)
|
|
],
|
|
}
|
|
assert graph.get_state(config, subgraphs=True).tasks[0].state.values == {
|
|
"bar": "hi! foo"
|
|
}
|
|
assert graph.invoke(Command(resume="baz"), config, durability=durability) == {
|
|
"foo": "hi! foobaz"
|
|
}
|
|
|
|
|
|
def test_stream_subgraphs_during_execution(
|
|
sync_checkpointer: BaseCheckpointSaver,
|
|
) -> None:
|
|
class InnerState(TypedDict):
|
|
my_key: Annotated[str, operator.add]
|
|
my_other_key: str
|
|
|
|
def inner_1(state: InnerState):
|
|
return {"my_key": "got here", "my_other_key": state["my_key"]}
|
|
|
|
def inner_2(state: InnerState):
|
|
time.sleep(0.5)
|
|
return {
|
|
"my_key": " and there",
|
|
"my_other_key": state["my_key"],
|
|
}
|
|
|
|
inner = StateGraph(InnerState)
|
|
inner.add_node("inner_1", inner_1)
|
|
inner.add_node("inner_2", inner_2)
|
|
inner.add_edge("inner_1", "inner_2")
|
|
inner.set_entry_point("inner_1")
|
|
inner.set_finish_point("inner_2")
|
|
|
|
class State(TypedDict):
|
|
my_key: Annotated[str, operator.add]
|
|
|
|
def outer_1(state: State):
|
|
time.sleep(0.2)
|
|
return {"my_key": " and parallel"}
|
|
|
|
def outer_2(state: State):
|
|
return {"my_key": " and back again"}
|
|
|
|
graph = StateGraph(State)
|
|
graph.add_node("inner", inner.compile())
|
|
graph.add_node("outer_1", outer_1)
|
|
graph.add_node("outer_2", outer_2)
|
|
|
|
graph.add_edge(START, "inner")
|
|
graph.add_edge(START, "outer_1")
|
|
graph.add_edge(["inner", "outer_1"], "outer_2")
|
|
graph.add_edge("outer_2", END)
|
|
|
|
app = graph.compile(checkpointer=sync_checkpointer)
|
|
|
|
start = time.perf_counter()
|
|
chunks: list[tuple[float, Any]] = []
|
|
config = {"configurable": {"thread_id": "2"}}
|
|
for c in app.stream({"my_key": ""}, config, subgraphs=True):
|
|
chunks.append((round(time.perf_counter() - start, 1), c))
|
|
for idx in range(len(chunks)):
|
|
elapsed, c = chunks[idx]
|
|
chunks[idx] = (round(elapsed - chunks[0][0], 1), c)
|
|
|
|
assert chunks == [
|
|
# arrives before "inner" finishes
|
|
(
|
|
FloatBetween(0.0, 0.1),
|
|
(
|
|
(AnyStr("inner:"),),
|
|
{"inner_1": {"my_key": "got here", "my_other_key": ""}},
|
|
),
|
|
),
|
|
(FloatBetween(0.2, 0.3), ((), {"outer_1": {"my_key": " and parallel"}})),
|
|
(
|
|
FloatBetween(0.5, 0.8),
|
|
(
|
|
(AnyStr("inner:"),),
|
|
{"inner_2": {"my_key": " and there", "my_other_key": "got here"}},
|
|
),
|
|
),
|
|
(FloatBetween(0.5, 0.8), ((), {"inner": {"my_key": "got here and there"}})),
|
|
(FloatBetween(0.5, 0.8), ((), {"outer_2": {"my_key": " and back again"}})),
|
|
]
|
|
|
|
|
|
def test_stream_buffering_single_node(sync_checkpointer: BaseCheckpointSaver) -> None:
|
|
class State(TypedDict):
|
|
my_key: Annotated[str, operator.add]
|
|
|
|
def node(state: State, writer: StreamWriter):
|
|
writer("Before sleep")
|
|
time.sleep(0.2)
|
|
writer("After sleep")
|
|
return {"my_key": "got here"}
|
|
|
|
builder = StateGraph(State)
|
|
builder.add_node("node", node)
|
|
builder.add_edge(START, "node")
|
|
builder.add_edge("node", END)
|
|
graph = builder.compile(checkpointer=sync_checkpointer)
|
|
|
|
start = time.perf_counter()
|
|
chunks: list[tuple[float, Any]] = []
|
|
config = {"configurable": {"thread_id": "2"}}
|
|
for c in graph.stream({"my_key": ""}, config, stream_mode="custom"):
|
|
chunks.append((round(time.perf_counter() - start, 1), c))
|
|
|
|
assert chunks == [
|
|
(FloatBetween(0.0, 0.1), "Before sleep"),
|
|
(FloatBetween(0.2, 0.3), "After sleep"),
|
|
]
|
|
|
|
|
|
def test_nested_graph_interrupts_parallel(
|
|
sync_checkpointer: BaseCheckpointSaver, durability: Durability
|
|
) -> None:
|
|
class InnerState(TypedDict):
|
|
my_key: Annotated[str, operator.add]
|
|
my_other_key: str
|
|
|
|
def inner_1(state: InnerState):
|
|
time.sleep(0.1)
|
|
return {"my_key": "got here", "my_other_key": state["my_key"]}
|
|
|
|
def inner_2(state: InnerState):
|
|
return {
|
|
"my_key": " and there",
|
|
"my_other_key": state["my_key"],
|
|
}
|
|
|
|
inner = StateGraph(InnerState)
|
|
inner.add_node("inner_1", inner_1)
|
|
inner.add_node("inner_2", inner_2)
|
|
inner.add_edge("inner_1", "inner_2")
|
|
inner.set_entry_point("inner_1")
|
|
inner.set_finish_point("inner_2")
|
|
|
|
class State(TypedDict):
|
|
my_key: Annotated[str, operator.add]
|
|
|
|
def outer_1(state: State):
|
|
return {"my_key": " and parallel"}
|
|
|
|
def outer_2(state: State):
|
|
return {"my_key": " and back again"}
|
|
|
|
graph = StateGraph(State)
|
|
graph.add_node("inner", inner.compile(interrupt_before=["inner_2"]))
|
|
graph.add_node("outer_1", outer_1)
|
|
graph.add_node("outer_2", outer_2)
|
|
|
|
graph.add_edge(START, "inner")
|
|
graph.add_edge(START, "outer_1")
|
|
graph.add_edge(["inner", "outer_1"], "outer_2")
|
|
graph.set_finish_point("outer_2")
|
|
|
|
app = graph.compile(checkpointer=sync_checkpointer)
|
|
|
|
# test invoke w/ nested interrupt
|
|
config = {"configurable": {"thread_id": "1"}}
|
|
assert app.invoke({"my_key": ""}, config, durability=durability) == {
|
|
"my_key": " and parallel",
|
|
}
|
|
|
|
assert app.invoke(None, config, durability=durability) == {
|
|
"my_key": "got here and there and parallel and back again",
|
|
}
|
|
|
|
# below combo of assertions is asserting two things
|
|
# - outer_1 finishes before inner interrupts (because we see its output in stream, which only happens after node finishes)
|
|
# - the writes of outer are persisted in 1st call and used in 2nd call, ie outer isn't called again (because we dont see outer_1 output again in 2nd stream)
|
|
# test stream updates w/ nested interrupt
|
|
config = {"configurable": {"thread_id": "2"}}
|
|
assert [
|
|
*app.stream({"my_key": ""}, config, subgraphs=True, durability=durability)
|
|
] == [
|
|
# we got to parallel node first
|
|
((), {"outer_1": {"my_key": " and parallel"}}),
|
|
((AnyStr("inner:"),), {"inner_1": {"my_key": "got here", "my_other_key": ""}}),
|
|
((), {"__interrupt__": ()}),
|
|
]
|
|
assert [*app.stream(None, config, durability=durability)] == [
|
|
{"outer_1": {"my_key": " and parallel"}, "__metadata__": {"cached": True}},
|
|
{"inner": {"my_key": "got here and there"}},
|
|
{"outer_2": {"my_key": " and back again"}},
|
|
]
|
|
|
|
# test stream values w/ nested interrupt
|
|
config = {"configurable": {"thread_id": "3"}}
|
|
assert [
|
|
*app.stream(
|
|
{"my_key": ""},
|
|
config,
|
|
stream_mode="values",
|
|
durability=durability,
|
|
)
|
|
] == [
|
|
{"my_key": ""},
|
|
{"my_key": " and parallel"},
|
|
]
|
|
assert [*app.stream(None, config, stream_mode="values", durability=durability)] == [
|
|
{"my_key": ""},
|
|
{"my_key": "got here and there and parallel"},
|
|
{"my_key": "got here and there and parallel and back again"},
|
|
]
|
|
|
|
# test interrupts BEFORE the parallel node
|
|
app = graph.compile(checkpointer=sync_checkpointer, interrupt_before=["outer_1"])
|
|
config = {"configurable": {"thread_id": "4"}}
|
|
assert [
|
|
*app.stream(
|
|
{"my_key": ""},
|
|
config,
|
|
stream_mode="values",
|
|
durability=durability,
|
|
)
|
|
] == [{"my_key": ""}]
|
|
# while we're waiting for the node w/ interrupt inside to finish
|
|
assert [*app.stream(None, config, stream_mode="values", durability=durability)] == [
|
|
{"my_key": ""},
|
|
{"my_key": " and parallel"},
|
|
]
|
|
assert [*app.stream(None, config, stream_mode="values", durability=durability)] == [
|
|
{"my_key": ""},
|
|
{"my_key": "got here and there and parallel"},
|
|
{"my_key": "got here and there and parallel and back again"},
|
|
]
|
|
|
|
# test interrupts AFTER the parallel node
|
|
app = graph.compile(checkpointer=sync_checkpointer, interrupt_after=["outer_1"])
|
|
config = {"configurable": {"thread_id": "5"}}
|
|
assert [
|
|
*app.stream(
|
|
{"my_key": ""},
|
|
config,
|
|
stream_mode="values",
|
|
durability=durability,
|
|
)
|
|
] == [
|
|
{"my_key": ""},
|
|
{"my_key": " and parallel"},
|
|
]
|
|
assert [*app.stream(None, config, stream_mode="values", durability=durability)] == [
|
|
{"my_key": ""},
|
|
{"my_key": "got here and there and parallel"},
|
|
]
|
|
assert [*app.stream(None, config, stream_mode="values", durability=durability)] == [
|
|
{"my_key": "got here and there and parallel"},
|
|
{"my_key": "got here and there and parallel and back again"},
|
|
]
|
|
|
|
|
|
def test_doubly_nested_graph_interrupts(
|
|
sync_checkpointer: BaseCheckpointSaver, durability: Durability
|
|
) -> None:
|
|
class State(TypedDict):
|
|
my_key: str
|
|
|
|
class ChildState(TypedDict):
|
|
my_key: str
|
|
|
|
class GrandChildState(TypedDict):
|
|
my_key: str
|
|
|
|
def grandchild_1(state: ChildState):
|
|
return {"my_key": state["my_key"] + " here"}
|
|
|
|
def grandchild_2(state: ChildState):
|
|
return {
|
|
"my_key": state["my_key"] + " and there",
|
|
}
|
|
|
|
grandchild = StateGraph(GrandChildState)
|
|
grandchild.add_node("grandchild_1", grandchild_1)
|
|
grandchild.add_node("grandchild_2", grandchild_2)
|
|
grandchild.add_edge("grandchild_1", "grandchild_2")
|
|
grandchild.set_entry_point("grandchild_1")
|
|
grandchild.set_finish_point("grandchild_2")
|
|
|
|
child = StateGraph(ChildState)
|
|
child.add_node(
|
|
"child_1",
|
|
grandchild.compile(interrupt_before=["grandchild_2"]),
|
|
)
|
|
child.set_entry_point("child_1")
|
|
child.set_finish_point("child_1")
|
|
|
|
def parent_1(state: State):
|
|
return {"my_key": "hi " + state["my_key"]}
|
|
|
|
def parent_2(state: State):
|
|
return {"my_key": state["my_key"] + " and back again"}
|
|
|
|
graph = StateGraph(State)
|
|
graph.add_node("parent_1", parent_1)
|
|
graph.add_node("child", child.compile())
|
|
graph.add_node("parent_2", parent_2)
|
|
graph.set_entry_point("parent_1")
|
|
graph.add_edge("parent_1", "child")
|
|
graph.add_edge("child", "parent_2")
|
|
graph.set_finish_point("parent_2")
|
|
|
|
app = graph.compile(checkpointer=sync_checkpointer)
|
|
|
|
# test invoke w/ nested interrupt
|
|
config = {"configurable": {"thread_id": "1"}}
|
|
assert app.invoke({"my_key": "my value"}, config, durability=durability) == {
|
|
"my_key": "hi my value",
|
|
}
|
|
|
|
assert app.invoke(None, config, durability=durability) == {
|
|
"my_key": "hi my value here and there and back again",
|
|
}
|
|
|
|
# test stream updates w/ nested interrupt
|
|
nodes: list[str] = []
|
|
config = {
|
|
"configurable": {"thread_id": "2", CONFIG_KEY_NODE_FINISHED: nodes.append}
|
|
}
|
|
assert [*app.stream({"my_key": "my value"}, config, durability=durability)] == [
|
|
{"parent_1": {"my_key": "hi my value"}},
|
|
{"__interrupt__": ()},
|
|
]
|
|
assert nodes == ["parent_1", "grandchild_1"]
|
|
assert [*app.stream(None, config, durability=durability)] == [
|
|
{"child": {"my_key": "hi my value here and there"}},
|
|
{"parent_2": {"my_key": "hi my value here and there and back again"}},
|
|
]
|
|
assert nodes == [
|
|
"parent_1",
|
|
"grandchild_1",
|
|
"grandchild_2",
|
|
"child_1",
|
|
"child",
|
|
"parent_2",
|
|
]
|
|
|
|
# test stream values w/ nested interrupt
|
|
config = {"configurable": {"thread_id": "3"}}
|
|
assert [
|
|
*app.stream(
|
|
{"my_key": "my value"},
|
|
config,
|
|
stream_mode="values",
|
|
durability=durability,
|
|
)
|
|
] == [
|
|
{"my_key": "my value"},
|
|
{"my_key": "hi my value"},
|
|
]
|
|
assert [*app.stream(None, config, stream_mode="values", durability=durability)] == [
|
|
{"my_key": "hi my value"},
|
|
{"my_key": "hi my value here and there"},
|
|
{"my_key": "hi my value here and 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,
|
|
{
|
|
"redo": "Researcher",
|
|
"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
|
|
|
|
|
|
def test_checkpoint_metadata(sync_checkpointer: BaseCheckpointSaver) -> None:
|
|
"""This test verifies that a run's configurable fields are merged with the
|
|
previous checkpoint config for each step in the run.
|
|
"""
|
|
# set up test
|
|
from langchain_core.language_models.fake_chat_models import (
|
|
FakeMessagesListChatModel,
|
|
)
|
|
from langchain_core.messages import AIMessage, AnyMessage
|
|
from langchain_core.prompts import ChatPromptTemplate
|
|
from langchain_core.tools import tool
|
|
|
|
# graph state
|
|
class BaseState(TypedDict):
|
|
messages: Annotated[list[AnyMessage], add_messages]
|
|
|
|
# initialize graph nodes
|
|
@tool()
|
|
def search_api(query: str) -> str:
|
|
"""Searches the API for the query."""
|
|
return f"result for {query}"
|
|
|
|
tools = [search_api]
|
|
|
|
prompt = ChatPromptTemplate.from_messages(
|
|
[
|
|
("system", "You are a nice assistant."),
|
|
("placeholder", "{messages}"),
|
|
]
|
|
)
|
|
|
|
model = FakeMessagesListChatModel(
|
|
responses=[
|
|
AIMessage(
|
|
content="",
|
|
tool_calls=[
|
|
{
|
|
"id": "tool_call123",
|
|
"name": "search_api",
|
|
"args": {"query": "query"},
|
|
},
|
|
],
|
|
),
|
|
AIMessage(content="answer"),
|
|
]
|
|
)
|
|
|
|
@traceable(run_type="llm")
|
|
def agent(state: BaseState) -> BaseState:
|
|
formatted = prompt.invoke(state)
|
|
response = model.invoke(formatted)
|
|
return {"messages": response, "usage_metadata": {"total_tokens": 123}}
|
|
|
|
def should_continue(data: BaseState) -> str:
|
|
# Logic to decide whether to continue in the loop or exit
|
|
if not data["messages"][-1].tool_calls:
|
|
return "exit"
|
|
else:
|
|
return "continue"
|
|
|
|
# define graphs w/ and w/o interrupt
|
|
workflow = StateGraph(BaseState)
|
|
workflow.add_node("agent", agent)
|
|
workflow.add_node("tools", ToolNode(tools))
|
|
workflow.set_entry_point("agent")
|
|
workflow.add_conditional_edges(
|
|
"agent", should_continue, {"continue": "tools", "exit": END}
|
|
)
|
|
workflow.add_edge("tools", "agent")
|
|
|
|
# graph w/o interrupt
|
|
app = workflow.compile(checkpointer=sync_checkpointer)
|
|
|
|
# graph w/ interrupt
|
|
app_w_interrupt = workflow.compile(
|
|
checkpointer=sync_checkpointer, interrupt_before=["tools"]
|
|
)
|
|
|
|
# assertions
|
|
|
|
# invoke graph w/o interrupt
|
|
assert app.invoke(
|
|
{"messages": ["what is weather in sf"]},
|
|
{
|
|
"configurable": {
|
|
"thread_id": "1",
|
|
"test_config_1": "foo",
|
|
"test_config_2": "bar",
|
|
},
|
|
},
|
|
) == {
|
|
"messages": [
|
|
_AnyIdHumanMessage(content="what is weather in sf"),
|
|
_AnyIdAIMessage(
|
|
content="",
|
|
tool_calls=[
|
|
{
|
|
"name": "search_api",
|
|
"args": {"query": "query"},
|
|
"id": "tool_call123",
|
|
"type": "tool_call",
|
|
}
|
|
],
|
|
),
|
|
_AnyIdToolMessage(
|
|
content="result for query",
|
|
name="search_api",
|
|
tool_call_id="tool_call123",
|
|
),
|
|
_AnyIdAIMessage(content="answer"),
|
|
]
|
|
}
|
|
|
|
config = {"configurable": {"thread_id": "1"}}
|
|
|
|
# assert that checkpoint metadata contains the run's configurable fields
|
|
chkpnt_metadata_1 = sync_checkpointer.get_tuple(config).metadata
|
|
assert chkpnt_metadata_1["test_config_1"] == "foo"
|
|
assert chkpnt_metadata_1["test_config_2"] == "bar"
|
|
|
|
# Verify that all checkpoint metadata have the expected keys. This check
|
|
# is needed because a run may have an arbitrary number of steps depending
|
|
# on how the graph is constructed.
|
|
chkpnt_tuples_1 = sync_checkpointer.list(config)
|
|
for chkpnt_tuple in chkpnt_tuples_1:
|
|
assert chkpnt_tuple.metadata["test_config_1"] == "foo"
|
|
assert chkpnt_tuple.metadata["test_config_2"] == "bar"
|
|
|
|
# invoke graph, but interrupt before tool call
|
|
app_w_interrupt.invoke(
|
|
{"messages": ["what is weather in sf"]},
|
|
{
|
|
"configurable": {
|
|
"thread_id": "2",
|
|
"test_config_3": "foo",
|
|
"test_config_4": "bar",
|
|
},
|
|
},
|
|
)
|
|
|
|
config = {"configurable": {"thread_id": "2"}}
|
|
|
|
# assert that checkpoint metadata contains the run's configurable fields
|
|
chkpnt_metadata_2 = sync_checkpointer.get_tuple(config).metadata
|
|
assert chkpnt_metadata_2["test_config_3"] == "foo"
|
|
assert chkpnt_metadata_2["test_config_4"] == "bar"
|
|
|
|
# resume graph execution
|
|
app_w_interrupt.invoke(
|
|
input=None,
|
|
config={
|
|
"configurable": {
|
|
"thread_id": "2",
|
|
"test_config_3": "foo",
|
|
"test_config_4": "bar",
|
|
}
|
|
},
|
|
)
|
|
|
|
# assert that checkpoint metadata contains the run's configurable fields
|
|
chkpnt_metadata_3 = sync_checkpointer.get_tuple(config).metadata
|
|
assert chkpnt_metadata_3["test_config_3"] == "foo"
|
|
assert chkpnt_metadata_3["test_config_4"] == "bar"
|
|
|
|
# Verify that all checkpoint metadata have the expected keys. This check
|
|
# is needed because a run may have an arbitrary number of steps depending
|
|
# on how the graph is constructed.
|
|
chkpnt_tuples_2 = sync_checkpointer.list(config)
|
|
for chkpnt_tuple in chkpnt_tuples_2:
|
|
assert chkpnt_tuple.metadata["test_config_3"] == "foo"
|
|
assert chkpnt_tuple.metadata["test_config_4"] == "bar"
|
|
|
|
|
|
def test_remove_message_via_state_update(
|
|
sync_checkpointer: BaseCheckpointSaver,
|
|
) -> None:
|
|
from langchain_core.messages import AIMessage, HumanMessage, RemoveMessage
|
|
|
|
workflow = StateGraph(state_schema=Annotated[list[AnyMessage], add_messages]) # type: ignore[arg-type]
|
|
workflow.add_node(
|
|
"chatbot",
|
|
lambda state: [
|
|
AIMessage(
|
|
content="Hello! How can I help you",
|
|
)
|
|
],
|
|
)
|
|
|
|
workflow.set_entry_point("chatbot")
|
|
workflow.add_edge("chatbot", END)
|
|
|
|
app = workflow.compile(checkpointer=sync_checkpointer)
|
|
config = {"configurable": {"thread_id": "1"}}
|
|
output = app.invoke([HumanMessage(content="Hi")], config=config)
|
|
app.update_state(config, values=[RemoveMessage(id=output[-1].id)])
|
|
|
|
updated_state = app.get_state(config)
|
|
|
|
assert len(updated_state.values) == 1
|
|
assert updated_state.values[-1].content == "Hi"
|
|
|
|
app.checkpointer.delete_thread(config["configurable"]["thread_id"])
|
|
|
|
# Verify that the message was removed from the checkpointer
|
|
assert app.checkpointer.get_tuple(config) is None
|
|
assert [*app.get_state_history(config)] == []
|
|
|
|
|
|
def test_remove_message_from_node():
|
|
from langchain_core.messages import AIMessage, HumanMessage, RemoveMessage
|
|
|
|
workflow = StateGraph(state_schema=Annotated[list[AnyMessage], add_messages]) # type: ignore[arg-type]
|
|
workflow.add_node(
|
|
"chatbot",
|
|
lambda state: [
|
|
AIMessage(
|
|
content="Hello!",
|
|
),
|
|
AIMessage(
|
|
content="How can I help you?",
|
|
),
|
|
],
|
|
)
|
|
workflow.add_node("delete_messages", lambda state: [RemoveMessage(id=state[-2].id)])
|
|
workflow.set_entry_point("chatbot")
|
|
workflow.add_edge("chatbot", "delete_messages")
|
|
workflow.add_edge("delete_messages", END)
|
|
|
|
app = workflow.compile()
|
|
output = app.invoke([HumanMessage(content="Hi")])
|
|
assert len(output) == 2
|
|
assert output[-1].content == "How can I help you?"
|
|
|
|
|
|
def test_xray_lance(snapshot: SnapshotAssertion):
|
|
from langchain_core.messages import AnyMessage, HumanMessage
|
|
|
|
class Analyst(BaseModel):
|
|
affiliation: str = Field(
|
|
description="Primary affiliation of the investment analyst.",
|
|
)
|
|
name: str = Field(
|
|
description="Name of the investment analyst.",
|
|
pattern=r"^[a-zA-Z0-9_-]{1,64}$",
|
|
)
|
|
role: str = Field(
|
|
description="Role of the investment analyst in the context of the topic.",
|
|
)
|
|
description: str = Field(
|
|
description="Description of the investment analyst focus, concerns, and motives.",
|
|
)
|
|
|
|
@property
|
|
def persona(self) -> str:
|
|
return f"Name: {self.name}\nRole: {self.role}\nAffiliation: {self.affiliation}\nDescription: {self.description}\n"
|
|
|
|
class Perspectives(BaseModel):
|
|
analysts: list[Analyst] = Field(
|
|
description="Comprehensive list of investment analysts with their roles and affiliations.",
|
|
)
|
|
|
|
class Section(BaseModel):
|
|
section_title: str = Field(..., title="Title of the section")
|
|
context: str = Field(
|
|
..., title="Provide a clear summary of the focus area that you researched."
|
|
)
|
|
findings: str = Field(
|
|
...,
|
|
title="Give a clear and detailed overview of your findings based upon the expert interview.",
|
|
)
|
|
thesis: str = Field(
|
|
...,
|
|
title="Give a clear and specific investment thesis based upon these findings.",
|
|
)
|
|
|
|
class InterviewState(TypedDict):
|
|
messages: Annotated[list[AnyMessage], add_messages]
|
|
analyst: Analyst
|
|
section: Section
|
|
|
|
class ResearchGraphState(TypedDict):
|
|
analysts: list[Analyst]
|
|
topic: str
|
|
max_analysts: int
|
|
sections: list[Section]
|
|
interviews: Annotated[list, operator.add]
|
|
|
|
# Conditional edge
|
|
def route_messages(state):
|
|
return "ask_question"
|
|
|
|
def generate_question(state):
|
|
return ...
|
|
|
|
def generate_answer(state):
|
|
return ...
|
|
|
|
# Add nodes and edges
|
|
interview_builder = StateGraph(InterviewState)
|
|
interview_builder.add_node("ask_question", generate_question)
|
|
interview_builder.add_node("answer_question", generate_answer)
|
|
|
|
# Flow
|
|
interview_builder.add_edge(START, "ask_question")
|
|
interview_builder.add_edge("ask_question", "answer_question")
|
|
interview_builder.add_conditional_edges(
|
|
"answer_question", route_messages, ["ask_question", END]
|
|
)
|
|
|
|
# Interview
|
|
interview_graph = interview_builder.compile().with_config(
|
|
run_name="Conduct Interviews"
|
|
)
|
|
|
|
# View
|
|
assert interview_graph.get_graph().to_json() == snapshot
|
|
|
|
def run_all_interviews(state: ResearchGraphState):
|
|
"""Edge to run the interview sub-graph using Send"""
|
|
return [
|
|
Send(
|
|
"conduct_interview",
|
|
{
|
|
"analyst": Analyst(),
|
|
"messages": [
|
|
HumanMessage(
|
|
content="So you said you were writing an article on ...?"
|
|
)
|
|
],
|
|
},
|
|
)
|
|
for s in state["analysts"]
|
|
]
|
|
|
|
def generate_sections(state: ResearchGraphState):
|
|
return ...
|
|
|
|
def generate_analysts(state: ResearchGraphState):
|
|
return ...
|
|
|
|
builder = StateGraph(ResearchGraphState)
|
|
builder.add_node("generate_analysts", generate_analysts)
|
|
builder.add_node("conduct_interview", interview_builder.compile())
|
|
builder.add_node("generate_sections", generate_sections)
|
|
|
|
builder.add_edge(START, "generate_analysts")
|
|
builder.add_conditional_edges(
|
|
"generate_analysts", run_all_interviews, ["conduct_interview"]
|
|
)
|
|
builder.add_edge("conduct_interview", "generate_sections")
|
|
builder.add_edge("generate_sections", END)
|
|
|
|
graph = builder.compile()
|
|
|
|
# View
|
|
assert graph.get_graph().to_json() == snapshot
|
|
assert graph.get_graph(xray=1).to_json() == snapshot
|
|
|
|
|
|
def test_channel_values(sync_checkpointer: BaseCheckpointSaver) -> None:
|
|
config = {"configurable": {"thread_id": "1"}}
|
|
chain = NodeBuilder().subscribe_only("input").write_to("output")
|
|
app = Pregel(
|
|
nodes={
|
|
"one": chain,
|
|
},
|
|
channels={
|
|
"ephemeral": EphemeralValue(Any),
|
|
"input": LastValue(int),
|
|
"output": LastValue(int),
|
|
},
|
|
input_channels=["input", "ephemeral"],
|
|
output_channels="output",
|
|
checkpointer=sync_checkpointer,
|
|
)
|
|
app.invoke({"input": 1, "ephemeral": "meow"}, config)
|
|
assert sync_checkpointer.get(config)["channel_values"] == {"input": 1, "output": 1}
|
|
|
|
|
|
def test_xray_issue(snapshot: SnapshotAssertion) -> None:
|
|
class State(TypedDict):
|
|
messages: Annotated[list, add_messages]
|
|
|
|
def node(name):
|
|
def _node(state: State):
|
|
return {"messages": [("human", f"entered {name} node")]}
|
|
|
|
return _node
|
|
|
|
parent = StateGraph(State)
|
|
child = StateGraph(State)
|
|
|
|
child.add_node("c_one", node("c_one"))
|
|
child.add_node("c_two", node("c_two"))
|
|
|
|
child.add_edge("__start__", "c_one")
|
|
child.add_edge("c_two", "c_one")
|
|
|
|
child.add_conditional_edges(
|
|
"c_one", lambda x: str(randrange(0, 2)), {"0": "c_two", "1": "__end__"}
|
|
)
|
|
|
|
parent.add_node("p_one", node("p_one"))
|
|
parent.add_node("p_two", child.compile())
|
|
|
|
parent.add_edge("__start__", "p_one")
|
|
parent.add_edge("p_two", "p_one")
|
|
|
|
parent.add_conditional_edges(
|
|
"p_one", lambda x: str(randrange(0, 2)), {"0": "p_two", "1": "__end__"}
|
|
)
|
|
|
|
app = parent.compile()
|
|
|
|
assert app.get_graph(xray=True).draw_mermaid() == snapshot
|
|
|
|
|
|
def test_xray_bool(snapshot: SnapshotAssertion) -> None:
|
|
class State(TypedDict):
|
|
messages: Annotated[list, add_messages]
|
|
|
|
def node(name):
|
|
def _node(state: State):
|
|
return {"messages": [("human", f"entered {name} node")]}
|
|
|
|
return _node
|
|
|
|
grand_parent = StateGraph(State)
|
|
|
|
child = StateGraph(State)
|
|
|
|
child.add_node("c_one", node("c_one"))
|
|
child.add_node("c_two", node("c_two"))
|
|
|
|
child.add_edge("__start__", "c_one")
|
|
child.add_edge("c_two", "c_one")
|
|
|
|
child.add_conditional_edges(
|
|
"c_one", lambda x: str(randrange(0, 2)), {"0": "c_two", "1": "__end__"}
|
|
)
|
|
|
|
parent = StateGraph(State)
|
|
parent.add_node("p_one", node("p_one"))
|
|
parent.add_node("p_two", child.compile())
|
|
parent.add_edge("__start__", "p_one")
|
|
parent.add_edge("p_two", "p_one")
|
|
parent.add_conditional_edges(
|
|
"p_one", lambda x: str(randrange(0, 2)), {"0": "p_two", "1": "__end__"}
|
|
)
|
|
|
|
grand_parent.add_node("gp_one", node("gp_one"))
|
|
grand_parent.add_node("gp_two", parent.compile())
|
|
grand_parent.add_edge("__start__", "gp_one")
|
|
grand_parent.add_edge("gp_two", "gp_one")
|
|
grand_parent.add_conditional_edges(
|
|
"gp_one", lambda x: str(randrange(0, 2)), {"0": "gp_two", "1": "__end__"}
|
|
)
|
|
|
|
app = grand_parent.compile()
|
|
assert app.get_graph(xray=True).draw_mermaid() == snapshot
|
|
|
|
|
|
def test_multiple_sinks_subgraphs(snapshot: SnapshotAssertion) -> None:
|
|
class State(TypedDict):
|
|
messages: Annotated[list, add_messages]
|
|
|
|
subgraph_builder = StateGraph(State)
|
|
subgraph_builder.add_node("one", lambda x: x)
|
|
subgraph_builder.add_node("two", lambda x: x)
|
|
subgraph_builder.add_node("three", lambda x: x)
|
|
subgraph_builder.add_edge("__start__", "one")
|
|
subgraph_builder.add_conditional_edges("one", lambda x: "two", ["two", "three"])
|
|
subgraph = subgraph_builder.compile()
|
|
|
|
builder = StateGraph(State)
|
|
builder.add_node("uno", lambda x: x)
|
|
builder.add_node("dos", lambda x: x)
|
|
builder.add_node("subgraph", subgraph)
|
|
builder.add_edge("__start__", "uno")
|
|
builder.add_conditional_edges("uno", lambda x: "dos", ["dos", "subgraph"])
|
|
|
|
app = builder.compile()
|
|
assert app.get_graph(xray=True).draw_mermaid() == snapshot
|
|
|
|
|
|
def test_store_injected(
|
|
sync_checkpointer: BaseCheckpointSaver, sync_store: BaseStore
|
|
) -> None:
|
|
class State(TypedDict):
|
|
count: Annotated[int, operator.add]
|
|
|
|
doc_id = str(uuid.uuid4())
|
|
doc = {"some-key": "this-is-a-val"}
|
|
uid = uuid.uuid4().hex
|
|
namespace = (f"foo-{uid}", "bar")
|
|
thread_1 = str(uuid.uuid4())
|
|
thread_2 = str(uuid.uuid4())
|
|
|
|
class Node:
|
|
def __init__(self, i: int | None = None):
|
|
self.i = i
|
|
|
|
def __call__(self, inputs: State, config: RunnableConfig, store: BaseStore):
|
|
assert isinstance(store, BaseStore)
|
|
store.put(
|
|
(
|
|
namespace
|
|
if self.i is not None
|
|
and config["configurable"]["thread_id"] in (thread_1, thread_2)
|
|
else (f"foo_{self.i}", "bar")
|
|
),
|
|
doc_id,
|
|
{
|
|
**doc,
|
|
"from_thread": config["configurable"]["thread_id"],
|
|
"some_val": inputs["count"],
|
|
},
|
|
)
|
|
return {"count": 1}
|
|
|
|
builder = StateGraph(State)
|
|
builder.add_node("node", Node())
|
|
builder.add_edge("__start__", "node")
|
|
N = 50
|
|
M = 1
|
|
|
|
for i in range(N):
|
|
builder.add_node(f"node_{i}", Node(i))
|
|
builder.add_edge("__start__", f"node_{i}")
|
|
|
|
graph = builder.compile(store=sync_store, checkpointer=sync_checkpointer)
|
|
|
|
results = graph.batch(
|
|
[{"count": 0}] * M,
|
|
([{"configurable": {"thread_id": str(uuid.uuid4())}}] * (M - 1))
|
|
+ [{"configurable": {"thread_id": thread_1}}],
|
|
)
|
|
result = results[-1]
|
|
assert result == {"count": N + 1}
|
|
returned_doc = sync_store.get(namespace, doc_id).value
|
|
assert returned_doc == {**doc, "from_thread": thread_1, "some_val": 0}
|
|
assert len(sync_store.search(namespace)) == 1
|
|
# Check results after another turn of the same thread
|
|
result = graph.invoke({"count": 0}, {"configurable": {"thread_id": thread_1}})
|
|
assert result == {"count": (N + 1) * 2}
|
|
returned_doc = sync_store.get(namespace, doc_id).value
|
|
assert returned_doc == {**doc, "from_thread": thread_1, "some_val": N + 1}
|
|
assert len(sync_store.search(namespace)) == 1
|
|
|
|
result = graph.invoke({"count": 0}, {"configurable": {"thread_id": thread_2}})
|
|
assert result == {"count": N + 1}
|
|
returned_doc = sync_store.get(namespace, doc_id).value
|
|
assert returned_doc == {
|
|
**doc,
|
|
"from_thread": thread_2,
|
|
"some_val": 0,
|
|
} # Overwrites the whole doc
|
|
assert len(sync_store.search(namespace)) == 1 # still overwriting the same one
|
|
|
|
|
|
def test_enum_node_names():
|
|
class NodeName(str, enum.Enum):
|
|
BAZ = "baz"
|
|
|
|
class State(TypedDict):
|
|
foo: str
|
|
bar: str
|
|
|
|
def baz(state: State):
|
|
return {"bar": state["foo"] + "!"}
|
|
|
|
graph = StateGraph(State)
|
|
graph.add_node(NodeName.BAZ, baz)
|
|
graph.add_edge(START, NodeName.BAZ)
|
|
graph.add_edge(NodeName.BAZ, END)
|
|
graph = graph.compile()
|
|
|
|
assert graph.invoke({"foo": "hello"}) == {"foo": "hello", "bar": "hello!"}
|
|
|
|
|
|
def test_debug_retry(sync_checkpointer: BaseCheckpointSaver):
|
|
class State(TypedDict):
|
|
messages: Annotated[list[str], operator.add]
|
|
|
|
def node(name):
|
|
def _node(state: State):
|
|
return {"messages": [f"entered {name} node"]}
|
|
|
|
return _node
|
|
|
|
builder = StateGraph(State)
|
|
builder.add_node("one", node("one"))
|
|
builder.add_node("two", node("two"))
|
|
builder.add_edge(START, "one")
|
|
builder.add_edge("one", "two")
|
|
builder.add_edge("two", END)
|
|
|
|
graph = builder.compile(checkpointer=sync_checkpointer)
|
|
|
|
config = {"configurable": {"thread_id": "1"}}
|
|
graph.invoke({"messages": []}, config=config, durability="async")
|
|
|
|
# re-run step: 1
|
|
target_config = next(
|
|
c.parent_config
|
|
for c in sync_checkpointer.list(config)
|
|
if c.metadata["step"] == 1
|
|
)
|
|
update_config = graph.update_state(target_config, values=None)
|
|
|
|
events = [
|
|
*graph.stream(
|
|
None, config=update_config, stream_mode="debug", durability="async"
|
|
)
|
|
]
|
|
|
|
checkpoint_events = list(
|
|
reversed([e["payload"] for e in events if e["type"] == "checkpoint"])
|
|
)
|
|
|
|
checkpoint_history = {
|
|
c.config["configurable"]["checkpoint_id"]: c
|
|
for c in graph.get_state_history(config)
|
|
}
|
|
|
|
def lax_normalize_config(config: dict | None) -> dict | None:
|
|
if config is None:
|
|
return None
|
|
return config["configurable"]
|
|
|
|
for stream in checkpoint_events:
|
|
stream_conf = lax_normalize_config(stream["config"])
|
|
stream_parent_conf = lax_normalize_config(stream["parent_config"])
|
|
assert stream_conf != stream_parent_conf
|
|
|
|
# ensure the streamed checkpoint == checkpoint from checkpointer.list()
|
|
history = checkpoint_history[stream["config"]["configurable"]["checkpoint_id"]]
|
|
history_conf = lax_normalize_config(history.config)
|
|
assert stream_conf == history_conf
|
|
|
|
history_parent_conf = lax_normalize_config(history.parent_config)
|
|
assert stream_parent_conf == history_parent_conf
|
|
|
|
|
|
def test_debug_subgraphs(
|
|
sync_checkpointer: BaseCheckpointSaver, durability: Durability
|
|
):
|
|
class State(TypedDict):
|
|
messages: Annotated[list[str], operator.add]
|
|
|
|
def node(name):
|
|
def _node(state: State):
|
|
return {"messages": [f"entered {name} node"]}
|
|
|
|
return _node
|
|
|
|
parent = StateGraph(State)
|
|
child = StateGraph(State)
|
|
|
|
child.add_node("c_one", node("c_one"))
|
|
child.add_node("c_two", node("c_two"))
|
|
child.add_edge(START, "c_one")
|
|
child.add_edge("c_one", "c_two")
|
|
child.add_edge("c_two", END)
|
|
|
|
parent.add_node("p_one", node("p_one"))
|
|
parent.add_node("p_two", child.compile())
|
|
parent.add_edge(START, "p_one")
|
|
parent.add_edge("p_one", "p_two")
|
|
parent.add_edge("p_two", END)
|
|
|
|
graph = parent.compile(checkpointer=sync_checkpointer)
|
|
|
|
config = {"configurable": {"thread_id": "1"}}
|
|
events = [
|
|
*graph.stream(
|
|
{"messages": []},
|
|
config=config,
|
|
stream_mode="debug",
|
|
durability=durability,
|
|
)
|
|
]
|
|
|
|
checkpoint_events = list(
|
|
reversed([e["payload"] for e in events if e["type"] == "checkpoint"])
|
|
)
|
|
if durability == "exit":
|
|
checkpoint_events = checkpoint_events[:1]
|
|
checkpoint_history = list(graph.get_state_history(config))
|
|
|
|
assert len(checkpoint_events) == len(checkpoint_history)
|
|
|
|
def lax_normalize_config(config: dict | None) -> dict | None:
|
|
if config is None:
|
|
return None
|
|
return config["configurable"]
|
|
|
|
for stream, history in zip(checkpoint_events, checkpoint_history):
|
|
assert stream["values"] == history.values
|
|
assert stream["next"] == list(history.next)
|
|
assert lax_normalize_config(stream["config"]) == lax_normalize_config(
|
|
history.config
|
|
)
|
|
assert lax_normalize_config(stream["parent_config"]) == lax_normalize_config(
|
|
history.parent_config
|
|
)
|
|
|
|
assert len(stream["tasks"]) == len(history.tasks)
|
|
for stream_task, history_task in zip(stream["tasks"], history.tasks):
|
|
assert stream_task["id"] == history_task.id
|
|
assert stream_task["name"] == history_task.name
|
|
assert stream_task["interrupts"] == history_task.interrupts
|
|
assert stream_task.get("error") == history_task.error
|
|
assert stream_task.get("state") == history_task.state
|
|
|
|
|
|
def test_debug_nested_subgraphs(
|
|
sync_checkpointer: BaseCheckpointSaver, durability: Durability
|
|
):
|
|
from collections import defaultdict
|
|
|
|
class State(TypedDict):
|
|
messages: Annotated[list[str], operator.add]
|
|
|
|
def node(name):
|
|
def _node(state: State):
|
|
return {"messages": [f"entered {name} node"]}
|
|
|
|
return _node
|
|
|
|
grand_parent = StateGraph(State)
|
|
parent = StateGraph(State)
|
|
child = StateGraph(State)
|
|
|
|
child.add_node("c_one", node("c_one"))
|
|
child.add_node("c_two", node("c_two"))
|
|
child.add_edge(START, "c_one")
|
|
child.add_edge("c_one", "c_two")
|
|
child.add_edge("c_two", END)
|
|
|
|
parent.add_node("p_one", node("p_one"))
|
|
parent.add_node("p_two", child.compile())
|
|
parent.add_edge(START, "p_one")
|
|
parent.add_edge("p_one", "p_two")
|
|
parent.add_edge("p_two", END)
|
|
|
|
grand_parent.add_node("gp_one", node("gp_one"))
|
|
grand_parent.add_node("gp_two", parent.compile())
|
|
grand_parent.add_edge(START, "gp_one")
|
|
grand_parent.add_edge("gp_one", "gp_two")
|
|
grand_parent.add_edge("gp_two", END)
|
|
|
|
graph = grand_parent.compile(checkpointer=sync_checkpointer)
|
|
|
|
config = {"configurable": {"thread_id": "1"}}
|
|
events = [
|
|
*graph.stream(
|
|
{"messages": []},
|
|
config=config,
|
|
stream_mode="debug",
|
|
subgraphs=True,
|
|
durability=durability,
|
|
)
|
|
]
|
|
|
|
stream_ns: dict[tuple, dict] = defaultdict(list)
|
|
for ns, e in events:
|
|
if e["type"] == "checkpoint":
|
|
stream_ns[ns].append(e["payload"])
|
|
|
|
assert list(stream_ns.keys()) == [
|
|
(),
|
|
(AnyStr("gp_two:"),),
|
|
(AnyStr("gp_two:"), AnyStr("p_two:")),
|
|
]
|
|
|
|
history_ns = {
|
|
ns: list(
|
|
graph.get_state_history(
|
|
{"configurable": {"thread_id": "1", "checkpoint_ns": "|".join(ns)}}
|
|
)
|
|
)[::-1]
|
|
for ns in stream_ns.keys()
|
|
}
|
|
|
|
def normalize_config(config: dict | None) -> dict | None:
|
|
if config is None:
|
|
return None
|
|
|
|
clean_config = {}
|
|
clean_config["thread_id"] = config["configurable"]["thread_id"]
|
|
clean_config["checkpoint_id"] = config["configurable"]["checkpoint_id"]
|
|
clean_config["checkpoint_ns"] = config["configurable"]["checkpoint_ns"]
|
|
if "checkpoint_map" in config["configurable"]:
|
|
clean_config["checkpoint_map"] = config["configurable"]["checkpoint_map"]
|
|
|
|
return clean_config
|
|
|
|
for checkpoint_events, checkpoint_history, ns in zip(
|
|
stream_ns.values(), history_ns.values(), stream_ns.keys()
|
|
):
|
|
if durability == "exit":
|
|
checkpoint_events = checkpoint_events[-1:]
|
|
if ns: # Save no checkpoints for subgraphs when durability="exit"
|
|
assert not checkpoint_history
|
|
continue
|
|
assert len(checkpoint_events) == len(checkpoint_history)
|
|
for stream, history in zip(checkpoint_events, checkpoint_history):
|
|
assert stream["values"] == history.values
|
|
assert stream["next"] == list(history.next)
|
|
assert normalize_config(stream["config"]) == normalize_config(
|
|
history.config
|
|
)
|
|
assert normalize_config(stream["parent_config"]) == normalize_config(
|
|
history.parent_config
|
|
)
|
|
|
|
assert len(stream["tasks"]) == len(history.tasks)
|
|
for stream_task, history_task in zip(stream["tasks"], history.tasks):
|
|
assert stream_task["id"] == history_task.id
|
|
assert stream_task["name"] == history_task.name
|
|
assert stream_task["interrupts"] == history_task.interrupts
|
|
assert stream_task.get("error") == history_task.error
|
|
assert stream_task.get("state") == history_task.state
|
|
|
|
|
|
def test_add_sequence():
|
|
class State(TypedDict):
|
|
foo: Annotated[list[str], operator.add]
|
|
bar: str
|
|
|
|
def step1(state: State):
|
|
return {"foo": ["step1"], "bar": "baz"}
|
|
|
|
def step2(state: State):
|
|
return {"foo": ["step2"]}
|
|
|
|
# test raising if less than 1 steps
|
|
with pytest.raises(ValueError):
|
|
StateGraph(State).add_sequence([])
|
|
|
|
# test raising if duplicate step names
|
|
with pytest.raises(ValueError):
|
|
StateGraph(State).add_sequence([step1, step1])
|
|
|
|
with pytest.raises(ValueError):
|
|
StateGraph(State).add_sequence([("foo", step1), ("foo", step1)])
|
|
|
|
# test unnamed steps
|
|
builder = StateGraph(State)
|
|
builder.add_sequence([step1, step2])
|
|
builder.add_edge(START, "step1")
|
|
graph = builder.compile()
|
|
result = graph.invoke({"foo": []})
|
|
assert result == {"foo": ["step1", "step2"], "bar": "baz"}
|
|
stream_chunks = list(graph.stream({"foo": []}))
|
|
assert stream_chunks == [
|
|
{"step1": {"foo": ["step1"], "bar": "baz"}},
|
|
{"step2": {"foo": ["step2"]}},
|
|
]
|
|
|
|
# test named steps
|
|
builder_named_steps = StateGraph(State)
|
|
builder_named_steps.add_sequence([("meow1", step1), ("meow2", step2)])
|
|
builder_named_steps.add_edge(START, "meow1")
|
|
graph_named_steps = builder_named_steps.compile()
|
|
result = graph_named_steps.invoke({"foo": []})
|
|
stream_chunks = list(graph_named_steps.stream({"foo": []}))
|
|
assert result == {"foo": ["step1", "step2"], "bar": "baz"}
|
|
assert stream_chunks == [
|
|
{"meow1": {"foo": ["step1"], "bar": "baz"}},
|
|
{"meow2": {"foo": ["step2"]}},
|
|
]
|
|
|
|
builder_named_steps = StateGraph(State)
|
|
builder_named_steps.add_sequence(
|
|
[
|
|
("meow1", lambda state: {"foo": ["foo"]}),
|
|
("meow2", lambda state: {"bar": state["foo"][0] + "bar"}),
|
|
],
|
|
)
|
|
builder_named_steps.add_edge(START, "meow1")
|
|
graph_named_steps = builder_named_steps.compile()
|
|
result = graph_named_steps.invoke({"foo": []})
|
|
stream_chunks = list(graph_named_steps.stream({"foo": []}))
|
|
# filtered by output schema
|
|
assert result == {"bar": "foobar", "foo": ["foo"]}
|
|
assert stream_chunks == [
|
|
{"meow1": {"foo": ["foo"]}},
|
|
{"meow2": {"bar": "foobar"}},
|
|
]
|
|
|
|
# test two sequences
|
|
|
|
def a(state: State):
|
|
return {"foo": ["a"]}
|
|
|
|
def b(state: State):
|
|
return {"foo": ["b"]}
|
|
|
|
builder_two_sequences = StateGraph(State)
|
|
builder_two_sequences.add_sequence([a])
|
|
builder_two_sequences.add_sequence([b])
|
|
builder_two_sequences.add_edge(START, "a")
|
|
builder_two_sequences.add_edge("a", "b")
|
|
graph_two_sequences = builder_two_sequences.compile()
|
|
|
|
result = graph_two_sequences.invoke({"foo": []})
|
|
assert result == {"foo": ["a", "b"]}
|
|
|
|
stream_chunks = list(graph_two_sequences.stream({"foo": []}))
|
|
assert stream_chunks == [
|
|
{"a": {"foo": ["a"]}},
|
|
{"b": {"foo": ["b"]}},
|
|
]
|
|
|
|
# test mixed nodes and sequences
|
|
|
|
def c(state: State):
|
|
return {"foo": ["c"]}
|
|
|
|
def d(state: State):
|
|
return {"foo": ["d"]}
|
|
|
|
def e(state: State):
|
|
return {"foo": ["e"]}
|
|
|
|
def foo(state: State):
|
|
if state["foo"][0] == "a":
|
|
return "d"
|
|
else:
|
|
return "c"
|
|
|
|
builder_complex = StateGraph(State)
|
|
builder_complex.add_sequence([a, b])
|
|
builder_complex.add_conditional_edges("b", foo)
|
|
builder_complex.add_node(c)
|
|
builder_complex.add_sequence([d, e])
|
|
builder_complex.add_edge(START, "a")
|
|
graph_complex = builder_complex.compile()
|
|
|
|
result = graph_complex.invoke({"foo": []})
|
|
assert result == {"foo": ["a", "b", "d", "e"]}
|
|
|
|
result = graph_complex.invoke({"foo": ["start"]})
|
|
assert result == {"foo": ["start", "a", "b", "c"]}
|
|
|
|
stream_chunks = list(graph_complex.stream({"foo": []}))
|
|
assert stream_chunks == [
|
|
{"a": {"foo": ["a"]}},
|
|
{"b": {"foo": ["b"]}},
|
|
{"d": {"foo": ["d"]}},
|
|
{"e": {"foo": ["e"]}},
|
|
]
|
|
|
|
|
|
def test_runnable_passthrough_node_graph() -> None:
|
|
class State(TypedDict):
|
|
changeme: str
|
|
|
|
async def dummy(state):
|
|
return state
|
|
|
|
agent = dummy | RunnablePassthrough.assign(prediction=RunnableLambda(lambda x: x))
|
|
|
|
graph_builder = StateGraph(State)
|
|
|
|
graph_builder.add_node("agent", agent)
|
|
graph_builder.add_edge(START, "agent")
|
|
|
|
graph = graph_builder.compile()
|
|
|
|
assert graph.get_graph(xray=True).to_json() == graph.get_graph(xray=False).to_json()
|
|
|
|
|
|
@pytest.mark.parametrize("subgraph_persist", [True, False])
|
|
def test_parent_command(
|
|
sync_checkpointer: BaseCheckpointSaver, subgraph_persist: bool
|
|
) -> None:
|
|
from langchain_core.messages import BaseMessage
|
|
from langchain_core.tools import tool
|
|
|
|
@tool(return_direct=True)
|
|
def get_user_name() -> Command:
|
|
"""Retrieve user name"""
|
|
return Command(update={"user_name": "Meow"}, graph=Command.PARENT)
|
|
|
|
subgraph_builder = StateGraph(MessagesState)
|
|
subgraph_builder.add_node("tool", get_user_name)
|
|
subgraph_builder.add_edge(START, "tool")
|
|
subgraph = subgraph_builder.compile(checkpointer=subgraph_persist)
|
|
|
|
class CustomParentState(TypedDict):
|
|
messages: Annotated[list[BaseMessage], add_messages]
|
|
# this key is not available to the child graph
|
|
user_name: str
|
|
|
|
builder = StateGraph(CustomParentState)
|
|
builder.add_node("alice", subgraph)
|
|
builder.add_edge(START, "alice")
|
|
graph = builder.compile(checkpointer=sync_checkpointer)
|
|
|
|
config = {"configurable": {"thread_id": "1"}}
|
|
|
|
assert graph.invoke(
|
|
{"messages": [("user", "get user name")]}, config, durability="exit"
|
|
) == {
|
|
"messages": [
|
|
_AnyIdHumanMessage(
|
|
content="get user name", additional_kwargs={}, response_metadata={}
|
|
),
|
|
],
|
|
"user_name": "Meow",
|
|
}
|
|
assert graph.get_state(config) == StateSnapshot(
|
|
values={
|
|
"messages": [
|
|
_AnyIdHumanMessage(
|
|
content="get user name", additional_kwargs={}, response_metadata={}
|
|
),
|
|
],
|
|
"user_name": "Meow",
|
|
},
|
|
next=(),
|
|
config={
|
|
"configurable": {
|
|
"thread_id": "1",
|
|
"checkpoint_ns": "",
|
|
"checkpoint_id": AnyStr(),
|
|
}
|
|
},
|
|
metadata={
|
|
"source": "loop",
|
|
"step": 1,
|
|
"parents": {},
|
|
},
|
|
created_at=AnyStr(),
|
|
parent_config=None,
|
|
tasks=(),
|
|
interrupts=(),
|
|
)
|
|
|
|
|
|
def test_interrupt_subgraph(sync_checkpointer: BaseCheckpointSaver):
|
|
class State(TypedDict):
|
|
baz: str
|
|
|
|
def foo(state):
|
|
return {"baz": "foo"}
|
|
|
|
def bar(state):
|
|
value = interrupt("Please provide baz value:")
|
|
return {"baz": value}
|
|
|
|
child_builder = StateGraph(State)
|
|
child_builder.add_node(bar)
|
|
child_builder.add_edge(START, "bar")
|
|
|
|
builder = StateGraph(State)
|
|
builder.add_node(foo)
|
|
builder.add_node("bar", child_builder.compile())
|
|
builder.add_edge(START, "foo")
|
|
builder.add_edge("foo", "bar")
|
|
graph = builder.compile(checkpointer=sync_checkpointer)
|
|
|
|
thread1 = {"configurable": {"thread_id": "1"}}
|
|
# First run, interrupted at bar
|
|
assert graph.invoke({"baz": ""}, thread1)
|
|
# Resume with answer
|
|
assert graph.invoke(Command(resume="bar"), thread1)
|
|
|
|
|
|
@pytest.mark.parametrize("resume_style", ["null", "map"])
|
|
def test_interrupt_multiple(
|
|
sync_checkpointer: BaseCheckpointSaver, resume_style: Literal["null", "map"]
|
|
):
|
|
class State(TypedDict):
|
|
my_key: Annotated[str, operator.add]
|
|
|
|
def node(s: State) -> State:
|
|
answer = interrupt({"value": 1})
|
|
answer2 = interrupt({"value": 2})
|
|
return {"my_key": answer + " " + answer2}
|
|
|
|
builder = StateGraph(State)
|
|
builder.add_node("node", node)
|
|
builder.add_edge(START, "node")
|
|
|
|
graph = builder.compile(checkpointer=sync_checkpointer)
|
|
thread1 = {"configurable": {"thread_id": "1"}}
|
|
|
|
result = [e for e in graph.stream({"my_key": "DE", "market": "DE"}, thread1)]
|
|
assert result == [
|
|
{
|
|
"__interrupt__": (
|
|
Interrupt(
|
|
value={"value": 1},
|
|
id=AnyStr(),
|
|
),
|
|
)
|
|
}
|
|
]
|
|
|
|
result = [
|
|
event
|
|
for event in graph.stream(
|
|
Command(
|
|
resume="answer 1"
|
|
if resume_style == "null"
|
|
else {result[0]["__interrupt__"][0].id: "answer 1"},
|
|
update={"my_key": " foofoo "},
|
|
),
|
|
thread1,
|
|
)
|
|
]
|
|
assert result == [
|
|
{
|
|
"__interrupt__": (
|
|
Interrupt(
|
|
value={"value": 2},
|
|
id=AnyStr(),
|
|
),
|
|
)
|
|
}
|
|
]
|
|
|
|
assert [
|
|
event
|
|
for event in graph.stream(
|
|
Command(
|
|
resume="answer 2"
|
|
if resume_style == "null"
|
|
else {result[0]["__interrupt__"][0].id: "answer 2"}
|
|
),
|
|
thread1,
|
|
stream_mode="values",
|
|
)
|
|
] == [
|
|
{"my_key": "DE foofoo "},
|
|
{"my_key": "DE foofoo answer 1 answer 2"},
|
|
]
|
|
|
|
|
|
def test_interrupt_loop(sync_checkpointer: BaseCheckpointSaver):
|
|
class State(TypedDict):
|
|
age: int
|
|
other: str
|
|
|
|
def ask_age(s: State):
|
|
"""Ask an expert for help."""
|
|
question = "How old are you?"
|
|
value = None
|
|
for _ in range(10):
|
|
value: str = interrupt(question)
|
|
if not value.isdigit() or int(value) < 18:
|
|
question = "invalid response"
|
|
value = None
|
|
else:
|
|
break
|
|
|
|
return {"age": int(value)}
|
|
|
|
builder = StateGraph(State)
|
|
builder.add_node("node", ask_age)
|
|
builder.add_edge(START, "node")
|
|
|
|
graph = builder.compile(checkpointer=sync_checkpointer)
|
|
thread1 = {"configurable": {"thread_id": "1"}}
|
|
|
|
assert [e for e in graph.stream({"other": ""}, thread1)] == [
|
|
{
|
|
"__interrupt__": (
|
|
Interrupt(
|
|
value="How old are you?",
|
|
id=AnyStr(),
|
|
),
|
|
)
|
|
}
|
|
]
|
|
|
|
assert [
|
|
event
|
|
for event in graph.stream(
|
|
Command(resume="13"),
|
|
thread1,
|
|
)
|
|
] == [
|
|
{
|
|
"__interrupt__": (
|
|
Interrupt(
|
|
value="invalid response",
|
|
id=AnyStr(),
|
|
),
|
|
)
|
|
}
|
|
]
|
|
|
|
assert [
|
|
event
|
|
for event in graph.stream(
|
|
Command(resume="15"),
|
|
thread1,
|
|
)
|
|
] == [
|
|
{
|
|
"__interrupt__": (
|
|
Interrupt(
|
|
value="invalid response",
|
|
id=AnyStr(),
|
|
),
|
|
)
|
|
}
|
|
]
|
|
|
|
assert [event for event in graph.stream(Command(resume="19"), thread1)] == [
|
|
{"node": {"age": 19}},
|
|
]
|
|
|
|
|
|
def test_interrupt_functional(
|
|
sync_checkpointer: BaseCheckpointSaver, snapshot: SnapshotAssertion
|
|
) -> None:
|
|
@task
|
|
def foo(state: dict) -> dict:
|
|
return {"a": state["a"] + "foo"}
|
|
|
|
@task
|
|
def bar(state: dict) -> dict:
|
|
return {"a": state["a"] + "bar", "b": state["b"]}
|
|
|
|
@entrypoint(checkpointer=sync_checkpointer)
|
|
def graph(inputs: dict) -> dict:
|
|
fut_foo = foo(inputs)
|
|
value = interrupt("Provide value for bar:")
|
|
bar_input = {**fut_foo.result(), "b": value}
|
|
fut_bar = bar(bar_input)
|
|
return fut_bar.result()
|
|
|
|
config = {"configurable": {"thread_id": "1"}}
|
|
# First run, interrupted at bar
|
|
assert graph.invoke({"a": ""}, config) == {
|
|
"__interrupt__": [
|
|
Interrupt(
|
|
value="Provide value for bar:",
|
|
id=AnyStr(),
|
|
)
|
|
]
|
|
}
|
|
# Resume with an answer
|
|
res = graph.invoke(Command(resume="bar"), config)
|
|
assert res == {"a": "foobar", "b": "bar"}
|
|
|
|
|
|
def test_interrupt_task_functional(
|
|
sync_checkpointer: BaseCheckpointSaver, snapshot: SnapshotAssertion
|
|
) -> None:
|
|
@task
|
|
def foo(state: dict) -> dict:
|
|
return {"a": state["a"] + "foo"}
|
|
|
|
@task
|
|
def bar(state: dict) -> dict:
|
|
value = interrupt("Provide value for bar:")
|
|
return {"a": state["a"] + value}
|
|
|
|
@entrypoint(checkpointer=sync_checkpointer)
|
|
def graph(inputs: dict) -> dict:
|
|
fut_foo = foo(inputs)
|
|
fut_bar = bar(fut_foo.result())
|
|
return fut_bar.result()
|
|
|
|
config = {"configurable": {"thread_id": "1"}}
|
|
# First run, interrupted at bar
|
|
assert graph.invoke({"a": ""}, config) == {
|
|
"__interrupt__": [
|
|
Interrupt(
|
|
value="Provide value for bar:",
|
|
id=AnyStr(),
|
|
),
|
|
]
|
|
}
|
|
# Resume with an answer
|
|
res = graph.invoke(Command(resume="bar"), config)
|
|
assert res == {"a": "foobar"}
|
|
|
|
# Test that we can interrupt the same task multiple times
|
|
config = {"configurable": {"thread_id": "2"}}
|
|
|
|
@entrypoint(checkpointer=sync_checkpointer)
|
|
def graph(inputs: dict) -> dict:
|
|
foo_result = foo(inputs).result()
|
|
bar_result = bar(foo_result).result()
|
|
baz_result = bar(bar_result).result()
|
|
return baz_result
|
|
|
|
# First run, interrupted at bar
|
|
assert graph.invoke({"a": ""}, config) == {
|
|
"__interrupt__": [
|
|
Interrupt(
|
|
value="Provide value for bar:",
|
|
id=AnyStr(),
|
|
),
|
|
]
|
|
}
|
|
# Provide resumes
|
|
graph.invoke(Command(resume="bar"), config)
|
|
assert graph.invoke(Command(resume="baz"), config) == {"a": "foobarbaz"}
|
|
|
|
|
|
def test_root_mixed_return() -> None:
|
|
def my_node(state: list[str]):
|
|
return [Command(update=["a"]), ["b"]]
|
|
|
|
graph = StateGraph(Annotated[list[str], operator.add])
|
|
|
|
graph.add_node(my_node)
|
|
graph.add_edge(START, "my_node")
|
|
graph = graph.compile()
|
|
|
|
assert graph.invoke([]) == ["a", "b"]
|
|
|
|
|
|
def test_dict_mixed_return() -> None:
|
|
class State(TypedDict):
|
|
foo: Annotated[str, operator.add]
|
|
|
|
def my_node(state: State):
|
|
return [Command(update={"foo": "a"}), {"foo": "b"}]
|
|
|
|
graph = StateGraph(State)
|
|
graph.add_node(my_node)
|
|
graph.add_edge(START, "my_node")
|
|
graph = graph.compile()
|
|
|
|
assert graph.invoke({"foo": ""}) == {"foo": "ab"}
|
|
|
|
|
|
def test_command_pydantic_dataclass() -> None:
|
|
class PydanticState(BaseModel):
|
|
foo: str
|
|
|
|
@dataclass
|
|
class DataclassState:
|
|
foo: str
|
|
|
|
for State in (PydanticState, DataclassState):
|
|
|
|
def node_a(state) -> Command[Literal["node_b"]]:
|
|
return Command(
|
|
update=State(foo="foo"),
|
|
goto="node_b",
|
|
)
|
|
|
|
def node_b(state):
|
|
return {"foo": state.foo + "bar"}
|
|
|
|
builder = StateGraph(State)
|
|
builder.add_edge(START, "node_a")
|
|
builder.add_node(node_a)
|
|
builder.add_node(node_b)
|
|
graph = builder.compile()
|
|
assert graph.invoke(State(foo="")) == {"foo": "foobar"}
|
|
|
|
|
|
def test_command_with_static_breakpoints(
|
|
sync_checkpointer: BaseCheckpointSaver,
|
|
) -> None:
|
|
"""Test that we can use Command to resume and update with static breakpoints."""
|
|
|
|
class State(TypedDict):
|
|
"""The graph state."""
|
|
|
|
foo: str
|
|
|
|
def node1(state: State):
|
|
return {
|
|
"foo": state["foo"] + "|node-1",
|
|
}
|
|
|
|
def node2(state: State):
|
|
return {
|
|
"foo": state["foo"] + "|node-2",
|
|
}
|
|
|
|
builder = StateGraph(State)
|
|
builder.add_node("node1", node1)
|
|
builder.add_node("node2", node2)
|
|
builder.add_edge(START, "node1")
|
|
builder.add_edge("node1", "node2")
|
|
|
|
graph = builder.compile(checkpointer=sync_checkpointer, interrupt_before=["node1"])
|
|
config = {"configurable": {"thread_id": str(uuid.uuid4())}}
|
|
|
|
# Start the graph and interrupt at the first node
|
|
graph.invoke({"foo": "abc"}, config)
|
|
result = graph.invoke(Command(resume="node1"), config)
|
|
assert result == {"foo": "abc|node-1|node-2"}
|
|
|
|
|
|
def test_multistep_plan(sync_checkpointer: BaseCheckpointSaver):
|
|
from langchain_core.messages import AnyMessage
|
|
|
|
class State(TypedDict, total=False):
|
|
plan: list[str | list[str]]
|
|
messages: Annotated[list[AnyMessage], add_messages]
|
|
|
|
def planner(state: State):
|
|
if state.get("plan") is None:
|
|
# create plan somehow
|
|
plan = ["step1", ["step2", "step3"], "step4"]
|
|
# pick the first step to execute next
|
|
first_step, *plan = plan
|
|
# put the rest of plan in state
|
|
return Command(goto=first_step, update={"plan": plan})
|
|
elif state["plan"]:
|
|
# go to the next step of the plan
|
|
next_step, *next_plan = state["plan"]
|
|
return Command(goto=next_step, update={"plan": next_plan})
|
|
else:
|
|
# the end of the plan
|
|
pass
|
|
|
|
def step1(state: State):
|
|
return Command(goto="planner", update={"messages": [("human", "step1")]})
|
|
|
|
def step2(state: State):
|
|
return Command(goto="planner", update={"messages": [("human", "step2")]})
|
|
|
|
def step3(state: State):
|
|
return Command(goto="planner", update={"messages": [("human", "step3")]})
|
|
|
|
def step4(state: State):
|
|
return Command(goto="planner", update={"messages": [("human", "step4")]})
|
|
|
|
builder = StateGraph(State)
|
|
builder.add_node(planner)
|
|
builder.add_node(step1)
|
|
builder.add_node(step2)
|
|
builder.add_node(step3)
|
|
builder.add_node(step4)
|
|
builder.add_edge(START, "planner")
|
|
graph = builder.compile(checkpointer=sync_checkpointer)
|
|
|
|
config = {"configurable": {"thread_id": "1"}}
|
|
|
|
assert graph.invoke({"messages": [("human", "start")]}, config) == {
|
|
"messages": [
|
|
_AnyIdHumanMessage(content="start"),
|
|
_AnyIdHumanMessage(content="step1"),
|
|
_AnyIdHumanMessage(content="step2"),
|
|
_AnyIdHumanMessage(content="step3"),
|
|
_AnyIdHumanMessage(content="step4"),
|
|
],
|
|
"plan": [],
|
|
}
|
|
|
|
|
|
def test_command_goto_with_static_breakpoints(
|
|
sync_checkpointer: BaseCheckpointSaver,
|
|
) -> None:
|
|
"""Use Command goto with static breakpoints."""
|
|
|
|
class State(TypedDict):
|
|
"""The graph state."""
|
|
|
|
foo: Annotated[str, operator.add]
|
|
|
|
def node1(state: State):
|
|
return {
|
|
"foo": "|node-1",
|
|
}
|
|
|
|
def node2(state: State):
|
|
return {
|
|
"foo": "|node-2",
|
|
}
|
|
|
|
builder = StateGraph(State)
|
|
builder.add_node("node1", node1)
|
|
builder.add_node("node2", node2)
|
|
builder.add_edge(START, "node1")
|
|
builder.add_edge("node1", "node2")
|
|
|
|
graph = builder.compile(checkpointer=sync_checkpointer, interrupt_before=["node1"])
|
|
|
|
config = {"configurable": {"thread_id": str(uuid.uuid4())}}
|
|
|
|
# Start the graph and interrupt at the first node
|
|
graph.invoke({"foo": "abc"}, config)
|
|
result = graph.invoke(Command(goto=["node2"]), config)
|
|
assert result == {"foo": "abc|node-1|node-2|node-2"}
|
|
|
|
|
|
def test_parallel_node_execution():
|
|
"""Test that parallel nodes execute concurrently."""
|
|
|
|
class State(TypedDict):
|
|
results: Annotated[list[str], operator.add]
|
|
|
|
def slow_node(state: State):
|
|
time.sleep(1)
|
|
return {"results": ["slow"]}
|
|
|
|
def fast_node(state: State):
|
|
time.sleep(2)
|
|
return {"results": ["fast"]}
|
|
|
|
builder = StateGraph(State)
|
|
builder.add_node("slow", slow_node)
|
|
builder.add_node("fast", fast_node)
|
|
builder.add_edge(START, "slow")
|
|
builder.add_edge(START, "fast")
|
|
|
|
graph = builder.compile()
|
|
|
|
start = time.perf_counter()
|
|
result = graph.invoke({"results": []})
|
|
duration = time.perf_counter() - start
|
|
|
|
# Fast node result should be available first
|
|
assert "fast" in result["results"][0]
|
|
|
|
# Total duration should be less than sum of both nodes
|
|
assert duration < 3.0
|
|
|
|
|
|
def test_multiple_interrupt_state_persistence(
|
|
sync_checkpointer: BaseCheckpointSaver,
|
|
) -> None:
|
|
"""Test that state is preserved correctly across multiple interrupts."""
|
|
|
|
class State(TypedDict):
|
|
steps: Annotated[list[str], operator.add]
|
|
|
|
def interruptible_node(state: State):
|
|
first = interrupt("First interrupt")
|
|
second = interrupt("Second interrupt")
|
|
return {"steps": [first, second]}
|
|
|
|
builder = StateGraph(State)
|
|
builder.add_node("node", interruptible_node)
|
|
builder.add_edge(START, "node")
|
|
|
|
app = builder.compile(checkpointer=sync_checkpointer)
|
|
config = {"configurable": {"thread_id": "1"}}
|
|
|
|
# First execution - should hit first interrupt
|
|
app.invoke({"steps": []}, config)
|
|
|
|
# State should still be empty since node hasn't returned
|
|
state = app.get_state(config)
|
|
assert state.values == {"steps": []}
|
|
|
|
# Resume after first interrupt - should hit second interrupt
|
|
app.invoke(Command(resume="step1"), config)
|
|
|
|
# State should still be empty since node hasn't returned
|
|
state = app.get_state(config)
|
|
assert state.values == {"steps": []}
|
|
|
|
# Resume after second interrupt - node should complete
|
|
result = app.invoke(Command(resume="step2"), config)
|
|
|
|
# Now state should contain both steps since node returned
|
|
assert result["steps"] == ["step1", "step2"]
|
|
state = app.get_state(config)
|
|
assert state.values["steps"] == ["step1", "step2"]
|
|
|
|
|
|
def test_concurrent_execution_thread_safety():
|
|
"""Test thread safety during concurrent execution."""
|
|
|
|
class State(TypedDict):
|
|
counter: Annotated[int, operator.add]
|
|
|
|
results = deque() # thread-safe queue
|
|
threads: list[threading.Thread] = []
|
|
|
|
def slow_node(state: State):
|
|
time.sleep(0.1)
|
|
return {"counter": 1}
|
|
|
|
builder = StateGraph(State)
|
|
builder.add_node("node", slow_node)
|
|
builder.add_edge(START, "node")
|
|
graph = builder.compile()
|
|
|
|
def run_graph():
|
|
result = graph.invoke({"counter": 0})
|
|
results.append(result)
|
|
|
|
# Start multiple threads
|
|
for _ in range(10):
|
|
thread = threading.Thread(target=run_graph)
|
|
thread.start()
|
|
threads.append(thread)
|
|
|
|
# Wait for all threads
|
|
for thread in threads:
|
|
thread.join()
|
|
|
|
# Verify results are independent
|
|
assert len(results) == 10
|
|
for result in results:
|
|
assert result["counter"] == 1
|
|
|
|
|
|
def test_checkpoint_recovery(
|
|
sync_checkpointer: BaseCheckpointSaver, durability: Durability
|
|
):
|
|
"""Test recovery from checkpoints after failures."""
|
|
|
|
class State(TypedDict):
|
|
steps: Annotated[list[str], operator.add]
|
|
attempt: int # Track number of attempts
|
|
|
|
def failing_node(state: State):
|
|
# Fail on first attempt, succeed on retry
|
|
if state["attempt"] == 1:
|
|
raise RuntimeError("Simulated failure")
|
|
return {"steps": ["node1"]}
|
|
|
|
def second_node(state: State):
|
|
return {"steps": ["node2"]}
|
|
|
|
builder = StateGraph(State)
|
|
builder.add_node("node1", failing_node)
|
|
builder.add_node("node2", second_node)
|
|
builder.add_edge(START, "node1")
|
|
builder.add_edge("node1", "node2")
|
|
|
|
graph = builder.compile(checkpointer=sync_checkpointer)
|
|
config = {"configurable": {"thread_id": "1"}}
|
|
|
|
# First attempt should fail
|
|
with pytest.raises(RuntimeError):
|
|
graph.invoke(
|
|
{"steps": ["start"], "attempt": 1},
|
|
config,
|
|
durability=durability,
|
|
)
|
|
|
|
# Verify checkpoint state
|
|
state = graph.get_state(config)
|
|
assert state is not None
|
|
assert state.values == {"steps": ["start"], "attempt": 1} # input state saved
|
|
assert state.next == ("node1",) # Should retry failed node
|
|
assert "RuntimeError('Simulated failure')" in state.tasks[0].error
|
|
|
|
# Retry with updated attempt count
|
|
result = graph.invoke({"steps": [], "attempt": 2}, config, durability=durability)
|
|
assert result == {"steps": ["start", "node1", "node2"], "attempt": 2}
|
|
|
|
# Verify checkpoint history shows both attempts
|
|
history = list(graph.get_state_history(config))
|
|
if durability != "exit":
|
|
assert len(history) == 6 # Initial + failed attempt + successful attempt
|
|
else:
|
|
assert len(history) == 2 # error + success
|
|
|
|
# Verify the error was recorded in checkpoint
|
|
failed_checkpoint = next(c for c in history if c.tasks and c.tasks[0].error)
|
|
assert "RuntimeError('Simulated failure')" in failed_checkpoint.tasks[0].error
|
|
|
|
# Verify delete leaves it empty
|
|
graph.checkpointer.delete_thread(config["configurable"]["thread_id"])
|
|
assert graph.checkpointer.get_tuple(config) is None
|
|
assert [*graph.get_state_history(config)] == []
|
|
|
|
|
|
def test_multiple_updates_root() -> None:
|
|
def node_a(state):
|
|
return [Command(update="a1"), Command(update="a2")]
|
|
|
|
def node_b(state):
|
|
return "b"
|
|
|
|
graph = (
|
|
StateGraph(Annotated[str, operator.add])
|
|
.add_sequence([node_a, node_b])
|
|
.add_edge(START, "node_a")
|
|
.compile()
|
|
)
|
|
|
|
assert graph.invoke("") == "a1a2b"
|
|
|
|
# only streams the last update from node_a
|
|
assert [c for c in graph.stream("", stream_mode="updates")] == [
|
|
{"node_a": ["a1", "a2"]},
|
|
{"node_b": "b"},
|
|
]
|
|
|
|
|
|
def test_multiple_updates() -> None:
|
|
class State(TypedDict):
|
|
foo: Annotated[str, operator.add]
|
|
|
|
def node_a(state):
|
|
return [Command(update={"foo": "a1"}), Command(update={"foo": "a2"})]
|
|
|
|
def node_b(state):
|
|
return {"foo": "b"}
|
|
|
|
graph = (
|
|
StateGraph(State)
|
|
.add_sequence([node_a, node_b])
|
|
.add_edge(START, "node_a")
|
|
.compile()
|
|
)
|
|
|
|
assert graph.invoke({"foo": ""}) == {
|
|
"foo": "a1a2b",
|
|
}
|
|
|
|
# only streams the last update from node_a
|
|
assert [c for c in graph.stream({"foo": ""}, stream_mode="updates")] == [
|
|
{"node_a": [{"foo": "a1"}, {"foo": "a2"}]},
|
|
{"node_b": {"foo": "b"}},
|
|
]
|
|
|
|
|
|
def test_falsy_return_from_task(sync_checkpointer: BaseCheckpointSaver):
|
|
"""Test with a falsy return from a task."""
|
|
|
|
@task
|
|
def falsy_task() -> bool:
|
|
return False
|
|
|
|
@entrypoint(checkpointer=sync_checkpointer)
|
|
def graph(state: dict) -> dict:
|
|
"""React tool."""
|
|
falsy_task().result()
|
|
interrupt("test")
|
|
|
|
configurable = {"configurable": {"thread_id": uuid.uuid4()}}
|
|
assert [
|
|
chunk
|
|
for chunk in graph.stream(
|
|
{"a": 5}, configurable, stream_mode="debug", durability="exit"
|
|
)
|
|
] == [
|
|
{
|
|
"payload": {
|
|
"config": {
|
|
"configurable": {
|
|
"checkpoint_id": AnyStr(),
|
|
"checkpoint_ns": "",
|
|
"thread_id": AnyStr(),
|
|
},
|
|
},
|
|
"metadata": {
|
|
"parents": {},
|
|
"source": "input",
|
|
"step": -1,
|
|
},
|
|
"next": [
|
|
"graph",
|
|
],
|
|
"parent_config": None,
|
|
"tasks": [
|
|
{
|
|
"id": AnyStr(),
|
|
"interrupts": (),
|
|
"name": "graph",
|
|
"state": None,
|
|
},
|
|
],
|
|
"values": None,
|
|
},
|
|
"step": -1,
|
|
"timestamp": AnyStr(),
|
|
"type": "checkpoint",
|
|
},
|
|
{
|
|
"payload": {
|
|
"id": AnyStr(),
|
|
"input": {
|
|
"a": 5,
|
|
},
|
|
"name": "graph",
|
|
"triggers": ("__start__",),
|
|
},
|
|
"step": 0,
|
|
"timestamp": AnyStr(),
|
|
"type": "task",
|
|
},
|
|
{
|
|
"payload": {
|
|
"id": AnyStr(),
|
|
"input": (
|
|
(),
|
|
{},
|
|
),
|
|
"name": "falsy_task",
|
|
"triggers": ("__pregel_push",),
|
|
},
|
|
"step": 0,
|
|
"timestamp": AnyStr(),
|
|
"type": "task",
|
|
},
|
|
{
|
|
"payload": {
|
|
"error": None,
|
|
"id": AnyStr(),
|
|
"interrupts": [],
|
|
"name": "falsy_task",
|
|
"result": {
|
|
"__return__": False,
|
|
},
|
|
},
|
|
"step": 0,
|
|
"timestamp": AnyStr(),
|
|
"type": "task_result",
|
|
},
|
|
{
|
|
"payload": {
|
|
"error": None,
|
|
"id": AnyStr(),
|
|
"interrupts": [
|
|
{
|
|
"id": AnyStr(),
|
|
"value": "test",
|
|
},
|
|
],
|
|
"name": "graph",
|
|
"result": {},
|
|
},
|
|
"step": 0,
|
|
"timestamp": AnyStr(),
|
|
"type": "task_result",
|
|
},
|
|
]
|
|
assert [
|
|
c
|
|
for c in graph.stream(
|
|
Command(resume="123"),
|
|
configurable,
|
|
stream_mode="debug",
|
|
durability="exit",
|
|
)
|
|
] == [
|
|
{
|
|
"payload": {
|
|
"config": {
|
|
"configurable": {
|
|
"checkpoint_id": AnyStr(),
|
|
"checkpoint_ns": "",
|
|
"thread_id": AnyStr(),
|
|
},
|
|
},
|
|
"metadata": {
|
|
"parents": {},
|
|
"source": "input",
|
|
"step": -1,
|
|
},
|
|
"next": [
|
|
"graph",
|
|
],
|
|
"parent_config": None,
|
|
"tasks": [
|
|
{
|
|
"id": AnyStr(),
|
|
"interrupts": (
|
|
{
|
|
"id": AnyStr(),
|
|
"value": "test",
|
|
},
|
|
),
|
|
"name": "graph",
|
|
"state": None,
|
|
},
|
|
],
|
|
"values": None,
|
|
},
|
|
"step": -1,
|
|
"timestamp": AnyStr(),
|
|
"type": "checkpoint",
|
|
},
|
|
{
|
|
"payload": {
|
|
"id": AnyStr(),
|
|
"input": {
|
|
"a": 5,
|
|
},
|
|
"name": "graph",
|
|
"triggers": ("__start__",),
|
|
},
|
|
"step": 0,
|
|
"timestamp": AnyStr(),
|
|
"type": "task",
|
|
},
|
|
{
|
|
"payload": {
|
|
"id": AnyStr(),
|
|
"input": (
|
|
(),
|
|
{},
|
|
),
|
|
"name": "falsy_task",
|
|
"triggers": ("__pregel_push",),
|
|
},
|
|
"step": 0,
|
|
"timestamp": AnyStr(),
|
|
"type": "task",
|
|
},
|
|
{
|
|
"payload": {
|
|
"error": None,
|
|
"id": AnyStr(),
|
|
"interrupts": [],
|
|
"name": "graph",
|
|
"result": {
|
|
"__end__": None,
|
|
},
|
|
},
|
|
"step": 0,
|
|
"timestamp": AnyStr(),
|
|
"type": "task_result",
|
|
},
|
|
{
|
|
"payload": {
|
|
"config": {
|
|
"configurable": {
|
|
"checkpoint_id": AnyStr(),
|
|
"checkpoint_ns": "",
|
|
"thread_id": AnyStr(),
|
|
},
|
|
},
|
|
"metadata": {
|
|
"parents": {},
|
|
"source": "loop",
|
|
"step": 0,
|
|
},
|
|
"next": [],
|
|
"parent_config": None,
|
|
"tasks": [],
|
|
"values": None,
|
|
},
|
|
"step": 0,
|
|
"timestamp": AnyStr(),
|
|
"type": "checkpoint",
|
|
},
|
|
]
|
|
|
|
|
|
def test_multiple_interrupts_functional(sync_checkpointer: BaseCheckpointSaver):
|
|
"""Test multiple interrupts with functional API."""
|
|
|
|
counter = 0
|
|
|
|
@task
|
|
def double(x: int) -> int:
|
|
"""Increment the counter."""
|
|
nonlocal counter
|
|
counter += 1
|
|
return 2 * x
|
|
|
|
@entrypoint(checkpointer=sync_checkpointer)
|
|
def graph(state: dict) -> dict:
|
|
"""React tool."""
|
|
|
|
values = []
|
|
|
|
for idx in [1, 2, 3]:
|
|
values.extend([double(idx).result(), interrupt({"a": "boo"})])
|
|
|
|
return {"values": values}
|
|
|
|
configurable = {"configurable": {"thread_id": str(uuid.uuid4())}}
|
|
graph.invoke({}, configurable)
|
|
graph.invoke(Command(resume="a"), configurable)
|
|
graph.invoke(Command(resume="b"), configurable)
|
|
result = graph.invoke(Command(resume="c"), configurable)
|
|
# `double` value should be cached appropriately when used w/ `interrupt`
|
|
assert result == {
|
|
"values": [2, "a", 4, "b", 6, "c"],
|
|
}
|
|
assert counter == 3
|
|
|
|
|
|
def test_multiple_interrupts_functional_cache(
|
|
sync_checkpointer: BaseCheckpointSaver, cache: BaseCache
|
|
):
|
|
"""Test multiple interrupts with functional API."""
|
|
|
|
counter = 0
|
|
|
|
@task(cache_policy=CachePolicy())
|
|
def double(x: int) -> int:
|
|
"""Increment the counter."""
|
|
nonlocal counter
|
|
counter += 1
|
|
return 2 * x
|
|
|
|
@entrypoint(checkpointer=sync_checkpointer, cache=cache)
|
|
def graph(state: dict) -> dict:
|
|
"""React tool."""
|
|
|
|
values = []
|
|
|
|
for idx in [1, 1, 2, 2, 3, 3]:
|
|
values.extend([double(idx).result(), interrupt({"a": "boo"})])
|
|
|
|
return {"values": values}
|
|
|
|
configurable = {"configurable": {"thread_id": str(uuid.uuid4())}}
|
|
graph.invoke({}, configurable)
|
|
graph.invoke(Command(resume="a"), configurable)
|
|
graph.invoke(Command(resume="b"), configurable)
|
|
graph.invoke(Command(resume="c"), configurable)
|
|
graph.invoke(Command(resume="d"), configurable)
|
|
graph.invoke(Command(resume="e"), configurable)
|
|
result = graph.invoke(Command(resume="f"), configurable)
|
|
# `double` value should be cached appropriately when used w/ `interrupt`
|
|
assert result == {
|
|
"values": [2, "a", 2, "b", 4, "c", 4, "d", 6, "e", 6, "f"],
|
|
}
|
|
assert counter == 3
|
|
|
|
# should all be cached now
|
|
configurable = {"configurable": {"thread_id": str(uuid.uuid4())}}
|
|
graph.invoke({}, configurable)
|
|
graph.invoke(Command(resume="a"), configurable)
|
|
graph.invoke(Command(resume="b"), configurable)
|
|
graph.invoke(Command(resume="c"), configurable)
|
|
graph.invoke(Command(resume="d"), configurable)
|
|
graph.invoke(Command(resume="e"), configurable)
|
|
result = graph.invoke(Command(resume="f"), configurable)
|
|
# `double` value should be cached appropriately when used w/ `interrupt`
|
|
assert result == {
|
|
"values": [2, "a", 2, "b", 4, "c", 4, "d", 6, "e", 6, "f"],
|
|
}
|
|
assert counter == 3
|
|
|
|
# clear cache
|
|
double.clear_cache(cache)
|
|
|
|
# should recompute now
|
|
configurable = {"configurable": {"thread_id": str(uuid.uuid4())}}
|
|
graph.invoke({}, configurable)
|
|
graph.invoke(Command(resume="a"), configurable)
|
|
graph.invoke(Command(resume="b"), configurable)
|
|
graph.invoke(Command(resume="c"), configurable)
|
|
graph.invoke(Command(resume="d"), configurable)
|
|
graph.invoke(Command(resume="e"), configurable)
|
|
result = graph.invoke(Command(resume="f"), configurable)
|
|
# `double` value should be cached appropriately when used w/ `interrupt`
|
|
assert result == {
|
|
"values": [2, "a", 2, "b", 4, "c", 4, "d", 6, "e", 6, "f"],
|
|
}
|
|
assert counter == 6
|
|
|
|
|
|
def test_task_before_interrupt_resume(
|
|
sync_checkpointer: BaseCheckpointSaver,
|
|
) -> None:
|
|
"""Test that Command(resume=value) works correctly when a @task runs
|
|
before interrupt-producing tasks in an @entrypoint.
|
|
|
|
The @task wrapper on both setup and ask is essential to reproduce the bug:
|
|
- @task on setup triggers a mid-step put_writes (creating a new pending_writes list)
|
|
- @task on ask means interrupt() runs in a child scratchpad that must
|
|
delegate to the parent for null resume consumption tracking
|
|
"""
|
|
|
|
@entrypoint(checkpointer=sync_checkpointer)
|
|
def workflow(number_of_topics: int) -> dict:
|
|
@task
|
|
def setup() -> int:
|
|
return number_of_topics
|
|
|
|
@task
|
|
def ask(question: str) -> str:
|
|
return interrupt(question)
|
|
|
|
n = setup().result()
|
|
|
|
answers = []
|
|
for i in range(n):
|
|
q = f"Whats the answer for topic {i + 1}?"
|
|
answers.append(ask(q).result())
|
|
|
|
return {"answers": answers}
|
|
|
|
config = {"configurable": {"thread_id": "1"}}
|
|
|
|
# First invocation - should get first interrupt
|
|
result = workflow.invoke(2, config=config)
|
|
assert "__interrupt__" in result
|
|
assert len(result["__interrupt__"]) == 1
|
|
assert result["__interrupt__"][0].value == "Whats the answer for topic 1?"
|
|
|
|
# Resume with answer for topic 1 - should get second interrupt
|
|
result = workflow.invoke(Command(resume="answer1"), config=config)
|
|
assert "__interrupt__" in result, f"Expected interrupt for topic 2, got: {result}"
|
|
assert len(result["__interrupt__"]) == 1
|
|
assert result["__interrupt__"][0].value == "Whats the answer for topic 2?"
|
|
|
|
# Resume with answer for topic 2 - should get final result
|
|
result = workflow.invoke(Command(resume="answer2"), config=config)
|
|
assert result == {"answers": ["answer1", "answer2"]}
|
|
|
|
|
|
def test_multiple_tasks_before_interrupt_resume(
|
|
sync_checkpointer: BaseCheckpointSaver,
|
|
) -> None:
|
|
"""Test that Command(resume=value) works correctly when multiple @tasks
|
|
run before an interrupt-producing task in an @entrypoint."""
|
|
|
|
@entrypoint(checkpointer=sync_checkpointer)
|
|
def workflow(inputs: dict) -> dict:
|
|
@task
|
|
def step_a(x: int) -> int:
|
|
return x + 1
|
|
|
|
@task
|
|
def step_b(x: int) -> int:
|
|
return x * 2
|
|
|
|
@task
|
|
def ask(question: str) -> str:
|
|
return interrupt(question)
|
|
|
|
a = step_a(inputs["x"]).result()
|
|
b = step_b(a).result()
|
|
|
|
answer = ask(f"Result so far is {b}. What next?").result()
|
|
|
|
return {"computed": b, "answer": answer}
|
|
|
|
config = {"configurable": {"thread_id": "1"}}
|
|
|
|
# First invocation - should get interrupt
|
|
result = workflow.invoke({"x": 5}, config=config)
|
|
assert "__interrupt__" in result
|
|
assert result["__interrupt__"][0].value == "Result so far is 12. What next?"
|
|
|
|
# Resume
|
|
result = workflow.invoke(Command(resume="continue"), config=config)
|
|
assert result == {"computed": 12, "answer": "continue"}
|
|
|
|
|
|
def test_no_redundant_put_writes_for_cached_task(
|
|
sync_checkpointer: BaseCheckpointSaver,
|
|
) -> None:
|
|
"""Cached @tasks on resume must not trigger redundant put_writes."""
|
|
from unittest.mock import patch
|
|
|
|
from langgraph.pregel._loop import PregelLoop
|
|
|
|
@task
|
|
def setup(x: int) -> int:
|
|
return x
|
|
|
|
@task
|
|
def ask(question: str) -> str:
|
|
return interrupt(question)
|
|
|
|
@entrypoint(checkpointer=sync_checkpointer)
|
|
def workflow(x: int) -> dict:
|
|
n = setup(x).result()
|
|
answer = ask(f"q{n}").result()
|
|
return {"answer": answer}
|
|
|
|
config = {"configurable": {"thread_id": "1"}}
|
|
result = workflow.invoke(1, config=config)
|
|
assert "__interrupt__" in result
|
|
|
|
put_writes_task_ids: list[str] = []
|
|
orig = PregelLoop.put_writes
|
|
|
|
def spy(self, task_id, writes):
|
|
put_writes_task_ids.append(task_id)
|
|
return orig(self, task_id, writes)
|
|
|
|
with patch.object(PregelLoop, "put_writes", spy):
|
|
result = workflow.invoke(Command(resume="ans"), config=config)
|
|
|
|
assert result == {"answer": "ans"}
|
|
# Count unique non-null task IDs that got put_writes.
|
|
# Should be exactly 2: the ask task and the entrypoint task.
|
|
# If 3, the cached setup task is being redundantly re-committed.
|
|
non_null = set(tid for tid in put_writes_task_ids if not tid.startswith("00000000"))
|
|
assert len(non_null) == 2, (
|
|
f"Expected 2 task IDs in put_writes (ask + entrypoint), got {len(non_null)}"
|
|
)
|
|
|
|
|
|
def test_node_before_interrupt_resume_graph_api(
|
|
sync_checkpointer: BaseCheckpointSaver,
|
|
) -> None:
|
|
"""Test that Command(resume=value) works correctly in a StateGraph when a
|
|
node runs before a node that calls interrupt(). This is the graph-API
|
|
analog of test_task_before_interrupt_resume (entrypoint API)."""
|
|
|
|
class State(TypedDict):
|
|
topics: list[str]
|
|
answers: Annotated[list[str], operator.add]
|
|
|
|
def setup(state: State) -> dict:
|
|
return {"topics": [f"topic {i + 1}" for i in range(len(state["topics"]))]}
|
|
|
|
def ask(state: State) -> dict:
|
|
answers = []
|
|
for topic in state["topics"]:
|
|
answer = interrupt(f"Whats the answer for {topic}?")
|
|
answers.append(answer)
|
|
return {"answers": answers}
|
|
|
|
graph = (
|
|
StateGraph(State)
|
|
.add_node("setup", setup)
|
|
.add_node("ask", ask)
|
|
.add_edge(START, "setup")
|
|
.add_edge("setup", "ask")
|
|
.add_edge("ask", END)
|
|
.compile(checkpointer=sync_checkpointer)
|
|
)
|
|
|
|
config = {"configurable": {"thread_id": "1"}}
|
|
|
|
# First invocation - setup runs, then ask interrupts on the first topic
|
|
result = graph.invoke({"topics": ["a", "b"], "answers": []}, config=config)
|
|
assert "__interrupt__" in result
|
|
assert len(result["__interrupt__"]) == 1
|
|
assert result["__interrupt__"][0].value == "Whats the answer for topic 1?"
|
|
|
|
# Resume with answer for topic 1 - should get second interrupt
|
|
result = graph.invoke(Command(resume="answer1"), config=config)
|
|
assert "__interrupt__" in result, f"Expected interrupt for topic 2, got: {result}"
|
|
assert len(result["__interrupt__"]) == 1
|
|
assert result["__interrupt__"][0].value == "Whats the answer for topic 2?"
|
|
|
|
# Resume with answer for topic 2 - should complete
|
|
result = graph.invoke(Command(resume="answer2"), config=config)
|
|
assert result == {
|
|
"topics": ["topic 1", "topic 2"],
|
|
"answers": ["answer1", "answer2"],
|
|
}
|
|
|
|
|
|
def test_multiple_nodes_before_interrupt_resume_graph_api(
|
|
sync_checkpointer: BaseCheckpointSaver,
|
|
) -> None:
|
|
"""Test that Command(resume=value) works correctly in a StateGraph when
|
|
multiple nodes run before a node that calls interrupt(). This is the
|
|
graph-API analog of test_multiple_tasks_before_interrupt_resume."""
|
|
|
|
class State(TypedDict):
|
|
value: int
|
|
answer: str
|
|
|
|
def step_a(state: State) -> dict:
|
|
return {"value": state["value"] + 1}
|
|
|
|
def step_b(state: State) -> dict:
|
|
return {"value": state["value"] * 2}
|
|
|
|
def ask(state: State) -> dict:
|
|
answer = interrupt(f"Result so far is {state['value']}. What next?")
|
|
return {"answer": answer}
|
|
|
|
graph = (
|
|
StateGraph(State)
|
|
.add_node("step_a", step_a)
|
|
.add_node("step_b", step_b)
|
|
.add_node("ask", ask)
|
|
.add_edge(START, "step_a")
|
|
.add_edge("step_a", "step_b")
|
|
.add_edge("step_b", "ask")
|
|
.add_edge("ask", END)
|
|
.compile(checkpointer=sync_checkpointer)
|
|
)
|
|
|
|
config = {"configurable": {"thread_id": "1"}}
|
|
|
|
# First invocation - step_a and step_b run, then ask interrupts
|
|
result = graph.invoke({"value": 5, "answer": ""}, config=config)
|
|
assert "__interrupt__" in result
|
|
assert result["__interrupt__"][0].value == "Result so far is 12. What next?"
|
|
|
|
# Resume - should complete
|
|
result = graph.invoke(Command(resume="continue"), config=config)
|
|
assert result == {"value": 12, "answer": "continue"}
|
|
|
|
|
|
def test_node_before_multiple_interrupt_cycles_graph_api(
|
|
sync_checkpointer: BaseCheckpointSaver,
|
|
) -> None:
|
|
"""Test that a node running before an interrupt node does not interfere
|
|
with multiple interrupt/resume cycles in a StateGraph."""
|
|
|
|
class State(TypedDict):
|
|
count: int
|
|
data: str
|
|
|
|
def prepare(state: State) -> dict:
|
|
return {"count": state["count"] + 10}
|
|
|
|
def multi_interrupt(state: State) -> dict:
|
|
first = interrupt("First question?")
|
|
second = interrupt("Second question?")
|
|
return {"data": f"{first},{second}"}
|
|
|
|
graph = (
|
|
StateGraph(State)
|
|
.add_node("prepare", prepare)
|
|
.add_node("multi_interrupt", multi_interrupt)
|
|
.add_edge(START, "prepare")
|
|
.add_edge("prepare", "multi_interrupt")
|
|
.add_edge("multi_interrupt", END)
|
|
.compile(checkpointer=sync_checkpointer)
|
|
)
|
|
|
|
config = {"configurable": {"thread_id": "1"}}
|
|
|
|
# First invocation - prepare runs, multi_interrupt hits first interrupt
|
|
result = graph.invoke({"count": 0, "data": ""}, config=config)
|
|
assert "__interrupt__" in result
|
|
assert result["__interrupt__"][0].value == "First question?"
|
|
|
|
# Resume first interrupt - hits second interrupt
|
|
result = graph.invoke(Command(resume="first_answer"), config=config)
|
|
assert "__interrupt__" in result
|
|
assert result["__interrupt__"][0].value == "Second question?"
|
|
|
|
# Resume second interrupt - completes
|
|
result = graph.invoke(Command(resume="second_answer"), config=config)
|
|
assert result == {"count": 10, "data": "first_answer,second_answer"}
|
|
|
|
|
|
def test_double_interrupt_subgraph(sync_checkpointer: BaseCheckpointSaver) -> None:
|
|
class AgentState(TypedDict):
|
|
input: str
|
|
|
|
def node_1(state: AgentState):
|
|
result = interrupt("interrupt node 1")
|
|
return {"input": result}
|
|
|
|
def node_2(state: AgentState):
|
|
result = interrupt("interrupt node 2")
|
|
return {"input": result}
|
|
|
|
subgraph_builder = (
|
|
StateGraph(AgentState)
|
|
.add_node("node_1", node_1)
|
|
.add_node("node_2", node_2)
|
|
.add_edge(START, "node_1")
|
|
.add_edge("node_1", "node_2")
|
|
.add_edge("node_2", END)
|
|
)
|
|
|
|
# invoke the sub graph
|
|
subgraph = subgraph_builder.compile(checkpointer=sync_checkpointer)
|
|
thread = {"configurable": {"thread_id": str(uuid.uuid4())}}
|
|
assert [c for c in subgraph.stream({"input": "test"}, thread)] == [
|
|
{
|
|
"__interrupt__": (
|
|
Interrupt(
|
|
value="interrupt node 1",
|
|
id=AnyStr(),
|
|
),
|
|
)
|
|
},
|
|
]
|
|
# resume from the first interrupt
|
|
assert [c for c in subgraph.stream(Command(resume="123"), thread)] == [
|
|
{
|
|
"node_1": {"input": "123"},
|
|
},
|
|
{
|
|
"__interrupt__": (
|
|
Interrupt(
|
|
value="interrupt node 2",
|
|
id=AnyStr(),
|
|
),
|
|
)
|
|
},
|
|
]
|
|
# resume from the second interrupt
|
|
assert [c for c in subgraph.stream(Command(resume="123"), thread)] == [
|
|
{
|
|
"node_2": {"input": "123"},
|
|
},
|
|
]
|
|
|
|
subgraph = subgraph_builder.compile()
|
|
|
|
def invoke_sub_agent(state: AgentState):
|
|
return subgraph.invoke(state)
|
|
|
|
thread = {"configurable": {"thread_id": str(uuid.uuid4())}}
|
|
parent_agent = (
|
|
StateGraph(AgentState)
|
|
.add_node("invoke_sub_agent", invoke_sub_agent)
|
|
.add_edge(START, "invoke_sub_agent")
|
|
.add_edge("invoke_sub_agent", END)
|
|
.compile(checkpointer=sync_checkpointer)
|
|
)
|
|
|
|
assert [c for c in parent_agent.stream({"input": "test"}, thread)] == [
|
|
{
|
|
"__interrupt__": (
|
|
Interrupt(
|
|
value="interrupt node 1",
|
|
id=AnyStr(),
|
|
),
|
|
)
|
|
},
|
|
]
|
|
|
|
# resume from the first interrupt
|
|
assert [c for c in parent_agent.stream(Command(resume=True), thread)] == [
|
|
{
|
|
"__interrupt__": (
|
|
Interrupt(
|
|
value="interrupt node 2",
|
|
id=AnyStr(),
|
|
),
|
|
)
|
|
}
|
|
]
|
|
|
|
# resume from 2nd interrupt
|
|
assert [c for c in parent_agent.stream(Command(resume=True), thread)] == [
|
|
{
|
|
"invoke_sub_agent": {"input": True},
|
|
},
|
|
]
|
|
|
|
|
|
def test_multi_resume(sync_checkpointer: BaseCheckpointSaver) -> None:
|
|
class ChildState(TypedDict):
|
|
prompt: str
|
|
human_input: str
|
|
human_inputs: list[str]
|
|
|
|
def get_human_input(state: ChildState):
|
|
human_input = interrupt(state["prompt"])
|
|
|
|
return {
|
|
"human_input": human_input,
|
|
"human_inputs": [human_input],
|
|
}
|
|
|
|
child_graph = (
|
|
StateGraph(ChildState)
|
|
.add_node("get_human_input", get_human_input)
|
|
.add_edge(START, "get_human_input")
|
|
.add_edge("get_human_input", END)
|
|
.compile(checkpointer=sync_checkpointer)
|
|
)
|
|
|
|
class ParentState(TypedDict):
|
|
prompts: list[str]
|
|
human_inputs: Annotated[list[str], operator.add]
|
|
|
|
def assign_workers(state: ParentState) -> list[Send]:
|
|
return [
|
|
Send(
|
|
"child_graph",
|
|
{"prompt": prompt},
|
|
)
|
|
for prompt in state["prompts"]
|
|
]
|
|
|
|
def cleanup(state: ParentState):
|
|
assert len(state["human_inputs"]) == len(state["prompts"])
|
|
|
|
parent_graph = (
|
|
StateGraph(ParentState)
|
|
.add_node("child_graph", child_graph)
|
|
.add_node("cleanup", cleanup)
|
|
.add_conditional_edges(START, assign_workers, ["child_graph"])
|
|
.add_edge("child_graph", "cleanup")
|
|
.add_edge("cleanup", END)
|
|
.compile(checkpointer=sync_checkpointer)
|
|
)
|
|
|
|
thread_config: RunnableConfig = {
|
|
"configurable": {
|
|
"thread_id": uuid.uuid4(),
|
|
},
|
|
}
|
|
|
|
prompts = ["a", "b", "c", "d", "e"]
|
|
|
|
events = parent_graph.invoke(
|
|
{"prompts": prompts}, thread_config, stream_mode="values"
|
|
)
|
|
|
|
assert len(events["__interrupt__"]) == len(prompts)
|
|
interrupt_values = {i.value for i in events["__interrupt__"]}
|
|
assert interrupt_values == set(prompts)
|
|
|
|
resume_map: dict[str, str] = {
|
|
i.id: f"human input for prompt {i.value}"
|
|
for i in parent_graph.get_state(thread_config).interrupts
|
|
}
|
|
|
|
result = parent_graph.invoke(Command(resume=resume_map), thread_config)
|
|
assert result == {
|
|
"prompts": prompts,
|
|
"human_inputs": [f"human input for prompt {prompt}" for prompt in prompts],
|
|
}
|
|
|
|
|
|
def test_sync_streaming_with_functional_api() -> None:
|
|
"""Test streaming with functional API.
|
|
|
|
This test verifies that we're able to stream results as they're being generated
|
|
rather than have all the results arrive at once after the graph has completed.
|
|
|
|
The time of arrival between the two updates corresponding to the two `slow` tasks
|
|
should be roughly the task delay. If results are buffered until graph completion,
|
|
the two updates arrive back-to-back instead.
|
|
"""
|
|
|
|
time_delay = 0.05
|
|
|
|
@task()
|
|
def slow() -> dict:
|
|
time.sleep(time_delay)
|
|
return {"tic": time.monotonic()}
|
|
|
|
@entrypoint()
|
|
def graph(inputs: dict) -> list:
|
|
first = slow().result()
|
|
second = slow().result()
|
|
return [first, second]
|
|
|
|
arrival_times = []
|
|
|
|
for chunk in graph.stream({}):
|
|
if "slow" not in chunk: # We'll just look at the updates from `slow`
|
|
continue
|
|
arrival_times.append(time.monotonic())
|
|
|
|
assert len(arrival_times) == 2
|
|
delta = arrival_times[1] - arrival_times[0]
|
|
# Allow a small amount of scheduler jitter while still verifying the chunks
|
|
# arrived separately rather than back-to-back after graph completion.
|
|
assert delta > time_delay * 0.8
|
|
|
|
|
|
def test_entrypoint_without_checkpointer() -> None:
|
|
"""Test no checkpointer."""
|
|
states = []
|
|
config = {"configurable": {"thread_id": "1"}}
|
|
|
|
# Test without previous
|
|
@entrypoint()
|
|
def foo(inputs: Any) -> Any:
|
|
states.append(inputs)
|
|
return inputs
|
|
|
|
assert foo.invoke({"a": "1"}, config) == {"a": "1"}
|
|
|
|
@entrypoint()
|
|
def foo(inputs: Any, *, previous: Any) -> Any:
|
|
states.append(previous)
|
|
return {"previous": previous, "current": inputs}
|
|
|
|
assert foo.invoke({"a": "1"}, config) == {"current": {"a": "1"}, "previous": None}
|
|
assert foo.invoke({"a": "1"}, config) == {"current": {"a": "1"}, "previous": None}
|
|
|
|
|
|
def test_entrypoint_stateful(sync_checkpointer: BaseCheckpointSaver) -> None:
|
|
"""Test stateful entrypoint invoke."""
|
|
|
|
# Test invoke
|
|
states = []
|
|
|
|
@entrypoint(checkpointer=sync_checkpointer)
|
|
def foo(inputs, *, previous: Any) -> Any:
|
|
states.append(previous)
|
|
return {"previous": previous, "current": inputs}
|
|
|
|
config = {"configurable": {"thread_id": "1"}}
|
|
|
|
assert foo.invoke({"a": "1"}, config) == {"current": {"a": "1"}, "previous": None}
|
|
assert foo.invoke({"a": "2"}, config) == {
|
|
"current": {"a": "2"},
|
|
"previous": {"current": {"a": "1"}, "previous": None},
|
|
}
|
|
assert foo.invoke({"a": "3"}, config) == {
|
|
"current": {"a": "3"},
|
|
"previous": {
|
|
"current": {"a": "2"},
|
|
"previous": {"current": {"a": "1"}, "previous": None},
|
|
},
|
|
}
|
|
assert states == [
|
|
None,
|
|
{"current": {"a": "1"}, "previous": None},
|
|
{"current": {"a": "2"}, "previous": {"current": {"a": "1"}, "previous": None}},
|
|
]
|
|
|
|
# Test stream
|
|
@entrypoint(checkpointer=sync_checkpointer)
|
|
def foo(inputs, *, previous: Any) -> Any:
|
|
return {"previous": previous, "current": inputs}
|
|
|
|
config = {"configurable": {"thread_id": "2"}}
|
|
items = [item for item in foo.stream({"a": "1"}, config)]
|
|
assert items == [{"foo": {"current": {"a": "1"}, "previous": None}}]
|
|
|
|
|
|
def test_entrypoint_stateful_update_state(
|
|
sync_checkpointer: BaseCheckpointSaver,
|
|
) -> None:
|
|
"""Test stateful entrypoint invoke."""
|
|
|
|
# Test invoke
|
|
states = []
|
|
|
|
@entrypoint(checkpointer=sync_checkpointer)
|
|
def foo(inputs, *, previous: Any) -> Any:
|
|
states.append(previous)
|
|
return {"previous": previous, "current": inputs}
|
|
|
|
config = {"configurable": {"thread_id": "1"}}
|
|
|
|
# assert print(foo.input_channels)
|
|
foo.update_state(config, {"a": "-1"})
|
|
assert foo.invoke({"a": "1"}, config) == {
|
|
"current": {"a": "1"},
|
|
"previous": {"a": "-1"},
|
|
}
|
|
assert foo.invoke({"a": "2"}, config) == {
|
|
"current": {"a": "2"},
|
|
"previous": {"current": {"a": "1"}, "previous": {"a": "-1"}},
|
|
}
|
|
assert foo.invoke({"a": "3"}, config) == {
|
|
"current": {"a": "3"},
|
|
"previous": {
|
|
"current": {"a": "2"},
|
|
"previous": {"current": {"a": "1"}, "previous": {"a": "-1"}},
|
|
},
|
|
}
|
|
|
|
# update state
|
|
foo.update_state(config, {"a": "3"})
|
|
|
|
# Test stream
|
|
assert [item for item in foo.stream({"a": "1"}, config)] == [
|
|
{"foo": {"current": {"a": "1"}, "previous": {"a": "3"}}}
|
|
]
|
|
assert states == [
|
|
{"a": "-1"},
|
|
{"current": {"a": "1"}, "previous": {"a": "-1"}},
|
|
{
|
|
"current": {"a": "2"},
|
|
"previous": {"current": {"a": "1"}, "previous": {"a": "-1"}},
|
|
},
|
|
{"a": "3"},
|
|
]
|
|
|
|
|
|
def test_entrypoint_from_sync_generator() -> None:
|
|
"""@entrypoint does not support sync generators."""
|
|
previous_return_values = []
|
|
|
|
with pytest.raises(NotImplementedError):
|
|
|
|
@entrypoint()
|
|
def foo(inputs, previous=None) -> Any:
|
|
previous_return_values.append(previous)
|
|
yield "a"
|
|
yield "b"
|
|
|
|
|
|
def test_multiple_subgraphs(sync_checkpointer: BaseCheckpointSaver) -> None:
|
|
class State(TypedDict):
|
|
a: int
|
|
b: int
|
|
|
|
class Output(TypedDict):
|
|
result: int
|
|
|
|
# Define the subgraphs
|
|
def add(state):
|
|
return {"result": state["a"] + state["b"]}
|
|
|
|
add_subgraph = (
|
|
StateGraph(State, output_schema=Output)
|
|
.add_node(add)
|
|
.add_edge(START, "add")
|
|
.compile()
|
|
)
|
|
|
|
def multiply(state):
|
|
return {"result": state["a"] * state["b"]}
|
|
|
|
multiply_subgraph = (
|
|
StateGraph(State, output_schema=Output)
|
|
.add_node(multiply)
|
|
.add_edge(START, "multiply")
|
|
.compile()
|
|
)
|
|
|
|
# Test calling the same subgraph multiple times
|
|
def call_same_subgraph(state):
|
|
result = add_subgraph.invoke(state)
|
|
another_result = add_subgraph.invoke({"a": result["result"], "b": 10})
|
|
return another_result
|
|
|
|
parent_call_same_subgraph = (
|
|
StateGraph(State, output_schema=Output)
|
|
.add_node(call_same_subgraph)
|
|
.add_edge(START, "call_same_subgraph")
|
|
.compile(checkpointer=sync_checkpointer)
|
|
)
|
|
config = {"configurable": {"thread_id": "1"}}
|
|
assert parent_call_same_subgraph.invoke({"a": 2, "b": 3}, config) == {"result": 15}
|
|
|
|
# Test calling multiple subgraphs
|
|
class Output(TypedDict):
|
|
add_result: int
|
|
multiply_result: int
|
|
|
|
def call_multiple_subgraphs(state):
|
|
add_result = add_subgraph.invoke(state)
|
|
multiply_result = multiply_subgraph.invoke(state)
|
|
return {
|
|
"add_result": add_result["result"],
|
|
"multiply_result": multiply_result["result"],
|
|
}
|
|
|
|
parent_call_multiple_subgraphs = (
|
|
StateGraph(State, output_schema=Output)
|
|
.add_node(call_multiple_subgraphs)
|
|
.add_edge(START, "call_multiple_subgraphs")
|
|
.compile(checkpointer=sync_checkpointer)
|
|
)
|
|
config = {"configurable": {"thread_id": "2"}}
|
|
assert parent_call_multiple_subgraphs.invoke({"a": 2, "b": 3}, config) == {
|
|
"add_result": 5,
|
|
"multiply_result": 6,
|
|
}
|
|
|
|
|
|
def test_multiple_subgraphs_functional(sync_checkpointer: BaseCheckpointSaver) -> None:
|
|
# Define addition subgraph
|
|
@entrypoint()
|
|
def add(inputs: tuple[int, int]):
|
|
a, b = inputs
|
|
return a + b
|
|
|
|
# Define multiplication subgraph using tasks
|
|
@task
|
|
def multiply_task(a, b):
|
|
return a * b
|
|
|
|
@entrypoint()
|
|
def multiply(inputs: tuple[int, int]):
|
|
return multiply_task(*inputs).result()
|
|
|
|
# Test calling the same subgraph multiple times
|
|
@task
|
|
def call_same_subgraph(a, b):
|
|
result = add.invoke([a, b])
|
|
another_result = add.invoke([result, 10])
|
|
return another_result
|
|
|
|
@entrypoint(checkpointer=sync_checkpointer)
|
|
def parent_call_same_subgraph(inputs):
|
|
return call_same_subgraph(*inputs).result()
|
|
|
|
config = {"configurable": {"thread_id": "1"}}
|
|
assert parent_call_same_subgraph.invoke([2, 3], config) == 15
|
|
|
|
# Test calling multiple subgraphs
|
|
@task
|
|
def call_multiple_subgraphs(a, b):
|
|
add_result = add.invoke([a, b])
|
|
multiply_result = multiply.invoke([a, b])
|
|
return [add_result, multiply_result]
|
|
|
|
@entrypoint(checkpointer=sync_checkpointer)
|
|
def parent_call_multiple_subgraphs(inputs):
|
|
return call_multiple_subgraphs(*inputs).result()
|
|
|
|
config = {"configurable": {"thread_id": "2"}}
|
|
assert parent_call_multiple_subgraphs.invoke([2, 3], config) == [5, 6]
|
|
|
|
|
|
def test_multiple_subgraphs_mixed_entrypoint(
|
|
sync_checkpointer: BaseCheckpointSaver,
|
|
) -> None:
|
|
"""Test calling multiple StateGraph subgraphs from an entrypoint."""
|
|
|
|
class State(TypedDict):
|
|
a: int
|
|
b: int
|
|
|
|
class Output(TypedDict):
|
|
result: int
|
|
|
|
# Define the subgraphs
|
|
def add(state):
|
|
return {"result": state["a"] + state["b"]}
|
|
|
|
add_subgraph = (
|
|
StateGraph(State, output_schema=Output)
|
|
.add_node(add)
|
|
.add_edge(START, "add")
|
|
.compile()
|
|
)
|
|
|
|
def multiply(state):
|
|
return {"result": state["a"] * state["b"]}
|
|
|
|
multiply_subgraph = (
|
|
StateGraph(State, output_schema=Output)
|
|
.add_node(multiply)
|
|
.add_edge(START, "multiply")
|
|
.compile()
|
|
)
|
|
|
|
# Test calling the same subgraph multiple times
|
|
@task
|
|
def call_same_subgraph(a, b):
|
|
result = add_subgraph.invoke({"a": a, "b": b})["result"]
|
|
another_result = add_subgraph.invoke({"a": result, "b": 10})["result"]
|
|
return another_result
|
|
|
|
@entrypoint(checkpointer=sync_checkpointer)
|
|
def parent_call_same_subgraph(inputs):
|
|
return call_same_subgraph(*inputs).result()
|
|
|
|
config = {"configurable": {"thread_id": "1"}}
|
|
assert parent_call_same_subgraph.invoke([2, 3], config) == 15
|
|
|
|
# Test calling multiple subgraphs
|
|
@task
|
|
def call_multiple_subgraphs(a, b):
|
|
add_result = add_subgraph.invoke({"a": a, "b": b})["result"]
|
|
multiply_result = multiply_subgraph.invoke({"a": a, "b": b})["result"]
|
|
return [add_result, multiply_result]
|
|
|
|
@entrypoint(checkpointer=sync_checkpointer)
|
|
def parent_call_multiple_subgraphs(inputs):
|
|
return call_multiple_subgraphs(*inputs).result()
|
|
|
|
config = {"configurable": {"thread_id": "2"}}
|
|
assert parent_call_multiple_subgraphs.invoke([2, 3], config) == [5, 6]
|
|
|
|
|
|
def test_multiple_subgraphs_mixed_state_graph(
|
|
sync_checkpointer: BaseCheckpointSaver,
|
|
) -> None:
|
|
"""Test calling multiple entrypoint "subgraphs" from a StateGraph."""
|
|
|
|
class State(TypedDict):
|
|
a: int
|
|
b: int
|
|
|
|
class Output(TypedDict):
|
|
result: int
|
|
|
|
# Define addition subgraph
|
|
@entrypoint()
|
|
def add(inputs: tuple[int, int]):
|
|
a, b = inputs
|
|
return a + b
|
|
|
|
# Define multiplication subgraph using tasks
|
|
@task
|
|
def multiply_task(a, b):
|
|
return a * b
|
|
|
|
@entrypoint()
|
|
def multiply(inputs: tuple[int, int]):
|
|
return multiply_task(*inputs).result()
|
|
|
|
# Test calling the same subgraph multiple times
|
|
def call_same_subgraph(state):
|
|
result = add.invoke([state["a"], state["b"]])
|
|
another_result = add.invoke([result, 10])
|
|
return {"result": another_result}
|
|
|
|
parent_call_same_subgraph = (
|
|
StateGraph(State, output_schema=Output)
|
|
.add_node(call_same_subgraph)
|
|
.add_edge(START, "call_same_subgraph")
|
|
.compile(checkpointer=sync_checkpointer)
|
|
)
|
|
config = {"configurable": {"thread_id": "1"}}
|
|
assert parent_call_same_subgraph.invoke({"a": 2, "b": 3}, config) == {"result": 15}
|
|
|
|
# Test calling multiple subgraphs
|
|
class Output(TypedDict):
|
|
add_result: int
|
|
multiply_result: int
|
|
|
|
def call_multiple_subgraphs(state):
|
|
add_result = add.invoke([state["a"], state["b"]])
|
|
multiply_result = multiply.invoke([state["a"], state["b"]])
|
|
return {
|
|
"add_result": add_result,
|
|
"multiply_result": multiply_result,
|
|
}
|
|
|
|
parent_call_multiple_subgraphs = (
|
|
StateGraph(State, output_schema=Output)
|
|
.add_node(call_multiple_subgraphs)
|
|
.add_edge(START, "call_multiple_subgraphs")
|
|
.compile(checkpointer=sync_checkpointer)
|
|
)
|
|
config = {"configurable": {"thread_id": "2"}}
|
|
assert parent_call_multiple_subgraphs.invoke({"a": 2, "b": 3}, config) == {
|
|
"add_result": 5,
|
|
"multiply_result": 6,
|
|
}
|
|
|
|
|
|
def test_multiple_subgraphs_checkpointer(
|
|
sync_checkpointer: BaseCheckpointSaver,
|
|
) -> None:
|
|
class SubgraphState(TypedDict):
|
|
sub_counter: Annotated[int, operator.add]
|
|
|
|
def subgraph_node(state):
|
|
return {"sub_counter": 2}
|
|
|
|
sub_graph_1 = (
|
|
StateGraph(SubgraphState)
|
|
.add_node(subgraph_node)
|
|
.add_edge(START, "subgraph_node")
|
|
.compile(checkpointer=True)
|
|
)
|
|
|
|
class OtherSubgraphState(TypedDict):
|
|
other_sub_counter: Annotated[int, operator.add]
|
|
|
|
def other_subgraph_node(state):
|
|
return {"other_sub_counter": 3}
|
|
|
|
sub_graph_2 = (
|
|
StateGraph(OtherSubgraphState)
|
|
.add_node(other_subgraph_node)
|
|
.add_edge(START, "other_subgraph_node")
|
|
.compile()
|
|
)
|
|
|
|
class ParentState(TypedDict):
|
|
parent_counter: int
|
|
|
|
def parent_node(state):
|
|
result = sub_graph_1.invoke({"sub_counter": state["parent_counter"]})
|
|
other_result = sub_graph_2.invoke({"other_sub_counter": result["sub_counter"]})
|
|
return {"parent_counter": other_result["other_sub_counter"]}
|
|
|
|
parent_graph = (
|
|
StateGraph(ParentState)
|
|
.add_node(parent_node)
|
|
.add_edge(START, "parent_node")
|
|
.compile(checkpointer=sync_checkpointer)
|
|
)
|
|
|
|
config = {"configurable": {"thread_id": "1"}}
|
|
assert parent_graph.invoke({"parent_counter": 0}, config) == {"parent_counter": 5}
|
|
assert parent_graph.invoke({"parent_counter": 0}, config) == {"parent_counter": 7}
|
|
config = {"configurable": {"thread_id": "2"}}
|
|
assert [
|
|
c
|
|
for c in parent_graph.stream(
|
|
{"parent_counter": 0}, config, subgraphs=True, stream_mode="updates"
|
|
)
|
|
] == [
|
|
(("parent_node",), {"subgraph_node": {"sub_counter": 2}}),
|
|
(
|
|
(AnyStr("parent_node:"), "1"),
|
|
{"other_subgraph_node": {"other_sub_counter": 3}},
|
|
),
|
|
((), {"parent_node": {"parent_counter": 5}}),
|
|
]
|
|
assert [
|
|
c
|
|
for c in parent_graph.stream(
|
|
{"parent_counter": 0}, config, subgraphs=True, stream_mode="updates"
|
|
)
|
|
] == [
|
|
(("parent_node",), {"subgraph_node": {"sub_counter": 2}}),
|
|
(
|
|
(AnyStr("parent_node:"), "1"),
|
|
{"other_subgraph_node": {"other_sub_counter": 3}},
|
|
),
|
|
((), {"parent_node": {"parent_counter": 7}}),
|
|
]
|
|
|
|
|
|
def test_entrypoint_output_schema_with_return_and_save() -> None:
|
|
"""Test output schema inference with entrypoint.final."""
|
|
|
|
# Un-parameterized entrypoint.final is interpreted as entrypoint.final[Any, Any]
|
|
@entrypoint()
|
|
def foo2(inputs, *, previous: Any) -> entrypoint.final:
|
|
return entrypoint.final(value="foo", save=1)
|
|
|
|
assert foo2.get_output_jsonschema() == {
|
|
"title": "LangGraphOutput",
|
|
}
|
|
|
|
@entrypoint()
|
|
def foo(inputs, *, previous: Any) -> entrypoint.final[str, int]:
|
|
return entrypoint.final(value="foo", save=1)
|
|
|
|
assert foo.get_output_jsonschema() == {
|
|
"title": "LangGraphOutput",
|
|
"type": "string",
|
|
}
|
|
|
|
with pytest.raises(TypeError):
|
|
# Raise an exception on an improperly parameterized entrypoint.final
|
|
# User is attempting to parameterize in this case, so we'll offer
|
|
# a bit of help if it's not done correctly.
|
|
@entrypoint()
|
|
def foo(inputs, *, previous: Any) -> entrypoint.final[int]:
|
|
return entrypoint.final(value=1, save=1) # type: ignore
|
|
|
|
|
|
def test_entrypoint_with_return_and_save(
|
|
sync_checkpointer: BaseCheckpointSaver,
|
|
) -> None:
|
|
"""Test entrypoint with return and save."""
|
|
previous_ = None
|
|
|
|
@entrypoint(checkpointer=sync_checkpointer)
|
|
def foo(msg: str, *, previous: Any) -> entrypoint.final[int, list[str]]:
|
|
nonlocal previous_
|
|
previous_ = previous
|
|
previous = previous or []
|
|
return entrypoint.final(value=len(previous), save=previous + [msg])
|
|
|
|
assert foo.get_output_jsonschema() == {
|
|
"title": "LangGraphOutput",
|
|
"type": "integer",
|
|
}
|
|
|
|
config = {"configurable": {"thread_id": "1"}}
|
|
assert foo.invoke("hello", config) == 0
|
|
assert previous_ is None
|
|
assert foo.invoke("goodbye", config) == 1
|
|
assert previous_ == ["hello"]
|
|
assert foo.invoke("definitely", config) == 2
|
|
assert previous_ == ["hello", "goodbye"]
|
|
|
|
|
|
def test_overriding_injectable_args_with_tasks(sync_store: BaseStore) -> None:
|
|
"""Test overriding injectable args in tasks."""
|
|
|
|
@task
|
|
def foo(store: BaseStore, writer: StreamWriter, value: Any) -> None:
|
|
assert store is value
|
|
assert writer is value
|
|
|
|
@entrypoint(store=sync_store)
|
|
def main(inputs, store: BaseStore) -> str:
|
|
assert store is not None
|
|
foo(store=None, writer=None, value=None).result()
|
|
foo(store="hello", writer="hello", value="hello").result()
|
|
return "OK"
|
|
|
|
assert main.invoke({}) == "OK"
|
|
|
|
|
|
def test_named_tasks_functional() -> None:
|
|
class Foo:
|
|
def foo(self, value: str) -> dict:
|
|
return value + "foo"
|
|
|
|
f = Foo()
|
|
|
|
# class method task
|
|
foo = task(f.foo, name="custom_foo")
|
|
other_foo = task(f.foo, name="other_foo")
|
|
|
|
# regular function task
|
|
@task(name="custom_bar")
|
|
def bar(value: str) -> dict:
|
|
return value + "|bar"
|
|
|
|
def baz(update: str, value: str) -> dict:
|
|
return value + f"|{update}"
|
|
|
|
# partial function task (unnamed)
|
|
baz_task = task(functools.partial(baz, "baz"))
|
|
# partial function task (named_)
|
|
custom_baz_task = task(functools.partial(baz, "custom_baz"), name="custom_baz")
|
|
|
|
class Qux:
|
|
def __call__(self, value: str) -> dict:
|
|
return value + "|qux"
|
|
|
|
qux_task = task(Qux(), name="qux")
|
|
|
|
@entrypoint()
|
|
def workflow(inputs: dict) -> dict:
|
|
foo_result = foo(inputs).result()
|
|
other_foo(inputs).result()
|
|
fut_bar = bar(foo_result)
|
|
fut_baz = baz_task(fut_bar.result())
|
|
fut_custom_baz = custom_baz_task(fut_baz.result())
|
|
fut_qux = qux_task(fut_custom_baz.result())
|
|
return fut_qux.result()
|
|
|
|
assert list(workflow.stream("", stream_mode="updates")) == [
|
|
{"custom_foo": "foo"},
|
|
{"other_foo": "foo"},
|
|
{"custom_bar": "foo|bar"},
|
|
{"baz": "foo|bar|baz"},
|
|
{"custom_baz": "foo|bar|baz|custom_baz"},
|
|
{"qux": "foo|bar|baz|custom_baz|qux"},
|
|
{"workflow": "foo|bar|baz|custom_baz|qux"},
|
|
]
|
|
|
|
|
|
def test_tags_stream_mode_messages() -> None:
|
|
model = GenericFakeChatModel(messages=iter(["foo"]), tags=["meow"])
|
|
graph = (
|
|
StateGraph(MessagesState)
|
|
.add_node(
|
|
"call_model", lambda state: {"messages": model.invoke(state["messages"])}
|
|
)
|
|
.add_edge(START, "call_model")
|
|
.compile()
|
|
)
|
|
assert list(
|
|
graph.stream(
|
|
{
|
|
"messages": "hi",
|
|
},
|
|
stream_mode="messages",
|
|
)
|
|
) == [
|
|
(
|
|
_AnyIdAIMessageChunk(content="foo", chunk_position="last"),
|
|
{
|
|
"langgraph_step": 1,
|
|
"langgraph_node": "call_model",
|
|
"langgraph_triggers": ("branch:to:call_model",),
|
|
"langgraph_path": ("__pregel_pull", "call_model"),
|
|
"langgraph_checkpoint_ns": AnyStr("call_model:"),
|
|
"checkpoint_ns": AnyStr("call_model:"),
|
|
"ls_provider": "genericfakechatmodel",
|
|
"ls_model_type": "chat",
|
|
"ls_integration": "langchain_chat_model",
|
|
"lc_versions": {"langchain-core": LANGCHAIN_CORE_VERSION},
|
|
"tags": ["meow"],
|
|
},
|
|
)
|
|
]
|
|
|
|
|
|
def test_configurable_propagates_to_stream_metadata() -> None:
|
|
"""Regression: thread_id, run_id, assistant_id, graph_id,
|
|
and langgraph_auth_user_id from configurable must appear
|
|
in stream_mode='messages' metadata."""
|
|
|
|
def my_node(state):
|
|
return {"messages": HumanMessage(content="hello")}
|
|
|
|
graph = (
|
|
StateGraph(MessagesState)
|
|
.add_node("my_node", my_node)
|
|
.add_edge(START, "my_node")
|
|
.compile()
|
|
)
|
|
|
|
config = {
|
|
"configurable": {
|
|
"thread_id": "th-123",
|
|
"checkpoint_id": "ckpt-1",
|
|
"checkpoint_ns": "ns-1",
|
|
"task_id": "task-1",
|
|
"run_id": "run-456",
|
|
"assistant_id": "asst-789",
|
|
"graph_id": "graph-0",
|
|
"model": "gpt-4o",
|
|
"user_id": "uid-1",
|
|
"cron_id": "cron-1",
|
|
"langgraph_auth_user_id": "user-1",
|
|
# these should NOT be propagated into metadata
|
|
"some_api_key": "secret",
|
|
"custom_setting": {"nested": True},
|
|
},
|
|
}
|
|
results = list(graph.stream({"messages": []}, config, stream_mode="messages"))
|
|
assert len(results) == 1
|
|
_, metadata = results[0]
|
|
# propagated keys
|
|
assert metadata["thread_id"] == "th-123"
|
|
assert metadata["checkpoint_id"] == "ckpt-1"
|
|
assert metadata["checkpoint_ns"] == "ns-1"
|
|
assert metadata["task_id"] == "task-1"
|
|
assert metadata["run_id"] == "run-456"
|
|
assert metadata["assistant_id"] == "asst-789"
|
|
assert metadata["graph_id"] == "graph-0"
|
|
# These are only present in trace metadata by default as of langgraph 1.2
|
|
# assert metadata["model"] == "gpt-4o"
|
|
# assert metadata["user_id"] == "uid-1"
|
|
# assert metadata["cron_id"] == "cron-1"
|
|
# assert metadata["langgraph_auth_user_id"] == "user-1"
|
|
# non-allowlisted keys must not appear
|
|
assert "some_api_key" not in metadata
|
|
assert "custom_setting" not in metadata
|
|
|
|
|
|
def test_stream_mode_messages_command() -> None:
|
|
from langchain_core.messages import HumanMessage
|
|
|
|
def my_node(state):
|
|
return {"messages": HumanMessage(content="foo")}
|
|
|
|
def my_other_node(state):
|
|
return Command(update={"messages": HumanMessage(content="bar")})
|
|
|
|
def my_last_node(state):
|
|
return [Command(update={"messages": HumanMessage(content="baz")})]
|
|
|
|
graph = (
|
|
StateGraph(MessagesState)
|
|
.add_sequence([my_node, my_other_node, my_last_node])
|
|
.add_edge(START, "my_node")
|
|
.compile()
|
|
)
|
|
assert list(
|
|
graph.stream(
|
|
{
|
|
"messages": [],
|
|
},
|
|
stream_mode="messages",
|
|
)
|
|
) == [
|
|
(
|
|
_AnyIdHumanMessage(content="foo"),
|
|
{
|
|
"ls_integration": "langgraph",
|
|
"langgraph_step": 1,
|
|
"langgraph_node": "my_node",
|
|
"langgraph_triggers": ("branch:to:my_node",),
|
|
"langgraph_path": ("__pregel_pull", "my_node"),
|
|
"langgraph_checkpoint_ns": AnyStr("my_node:"),
|
|
},
|
|
),
|
|
(
|
|
_AnyIdHumanMessage(content="bar"),
|
|
{
|
|
"ls_integration": "langgraph",
|
|
"langgraph_step": 2,
|
|
"langgraph_node": "my_other_node",
|
|
"langgraph_triggers": ("branch:to:my_other_node",),
|
|
"langgraph_path": ("__pregel_pull", "my_other_node"),
|
|
"langgraph_checkpoint_ns": AnyStr("my_other_node:"),
|
|
},
|
|
),
|
|
(
|
|
_AnyIdHumanMessage(content="baz"),
|
|
{
|
|
"ls_integration": "langgraph",
|
|
"langgraph_step": 3,
|
|
"langgraph_node": "my_last_node",
|
|
"langgraph_triggers": ("branch:to:my_last_node",),
|
|
"langgraph_path": ("__pregel_pull", "my_last_node"),
|
|
"langgraph_checkpoint_ns": AnyStr("my_last_node:"),
|
|
},
|
|
),
|
|
]
|
|
|
|
|
|
def test_node_destinations() -> None:
|
|
class State(TypedDict):
|
|
foo: Annotated[str, operator.add]
|
|
|
|
def node_a(state: State):
|
|
value = state["foo"]
|
|
if value == "a":
|
|
goto = "node_b"
|
|
else:
|
|
goto = "node_c"
|
|
|
|
return Command(
|
|
update={"foo": value},
|
|
goto=goto,
|
|
graph=Command.PARENT,
|
|
)
|
|
|
|
subgraph = StateGraph(State).add_node(node_a).add_edge(START, "node_a").compile()
|
|
|
|
# test calling subgraph inside a node function
|
|
def call_subgraph(state: State):
|
|
return subgraph.invoke(state)
|
|
|
|
def node_b(state: State):
|
|
return {"foo": "b"}
|
|
|
|
def node_c(state: State):
|
|
return {"foo": "c"}
|
|
|
|
for subgraph_node in (subgraph, call_subgraph):
|
|
# destinations w/ tuples
|
|
builder = StateGraph(State)
|
|
builder.add_edge(START, "child")
|
|
builder.add_node("child", subgraph_node, destinations=("node_b", "node_c"))
|
|
builder.add_node(node_b)
|
|
builder.add_node(node_c)
|
|
compiled_graph = builder.compile()
|
|
assert compiled_graph.invoke({"foo": ""}) == {"foo": "c"}
|
|
|
|
graph = compiled_graph.get_graph()
|
|
assert [
|
|
Edge(source="__start__", target="child", data=None, conditional=False),
|
|
Edge(source="child", target="node_b", data=None, conditional=True),
|
|
Edge(source="child", target="node_c", data=None, conditional=True),
|
|
Edge(source="node_b", target="__end__", data=None, conditional=False),
|
|
Edge(source="node_c", target="__end__", data=None, conditional=False),
|
|
] == graph.edges
|
|
|
|
# destinations w/ dicts
|
|
builder = StateGraph(State)
|
|
builder.add_edge(START, "child")
|
|
builder.add_node(
|
|
"child", subgraph_node, destinations={"node_b": "foo", "node_c": "bar"}
|
|
)
|
|
builder.add_node(node_b)
|
|
builder.add_node(node_c)
|
|
compiled_graph = builder.compile()
|
|
assert compiled_graph.invoke({"foo": ""}) == {"foo": "c"}
|
|
|
|
graph = compiled_graph.get_graph()
|
|
assert [
|
|
Edge(source="__start__", target="child", data=None, conditional=False),
|
|
Edge(source="child", target="node_b", data="foo", conditional=True),
|
|
Edge(source="child", target="node_c", data="bar", conditional=True),
|
|
Edge(source="node_b", target="__end__", data=None, conditional=False),
|
|
Edge(source="node_c", target="__end__", data=None, conditional=False),
|
|
] == graph.edges
|
|
|
|
|
|
def test_pydantic_none_state_update() -> None:
|
|
class State(BaseModel):
|
|
foo: str | None
|
|
|
|
def node_a(state: State) -> State:
|
|
return State(foo=None)
|
|
|
|
graph = StateGraph(State).add_node(node_a).add_edge(START, "node_a").compile()
|
|
assert graph.invoke({"foo": ""}) == {"foo": None}
|
|
|
|
|
|
def test_pydantic_state_update_command() -> None:
|
|
class State(BaseModel):
|
|
foo: str | None
|
|
|
|
def node_a(state: State) -> State:
|
|
return Command(update=State(foo=None))
|
|
|
|
graph = StateGraph(State).add_node(node_a).add_edge(START, "node_a").compile()
|
|
assert graph.invoke({"foo": ""}) == {"foo": None}
|
|
|
|
class State(BaseModel):
|
|
foo: str | None = None
|
|
bar: str | None = None
|
|
|
|
def node_a(state: State):
|
|
return State(foo="foo")
|
|
|
|
def node_b(state: State):
|
|
return Command(update=State(bar="bar"))
|
|
|
|
builder = StateGraph(State)
|
|
builder.add_node(node_a)
|
|
builder.add_node(node_b)
|
|
builder.add_edge(START, "node_a")
|
|
builder.add_edge("node_a", "node_b")
|
|
builder.add_edge("node_b", END)
|
|
graph = builder.compile()
|
|
|
|
assert graph.invoke(State()) == {"foo": "foo", "bar": "bar"}
|
|
|
|
|
|
def test_pydantic_state_mutation() -> None:
|
|
class Inner(BaseModel):
|
|
a: int = 0
|
|
|
|
class State(BaseModel):
|
|
inner: Inner = Inner()
|
|
outer: int = 0
|
|
|
|
def my_node(state: State) -> State:
|
|
state.inner.a = 5
|
|
state.outer = 10
|
|
return state
|
|
|
|
graph = StateGraph(State).add_node(my_node).add_edge(START, "my_node").compile()
|
|
|
|
assert graph.invoke({"outer": 1}) == {"outer": 10, "inner": Inner(a=5)}
|
|
|
|
# test w/ default_factory
|
|
class State(BaseModel):
|
|
inner: Inner = Field(default_factory=Inner)
|
|
outer: int = 0
|
|
|
|
def my_node(state: State) -> State:
|
|
state.inner.a = 5
|
|
state.outer = 10
|
|
return state
|
|
|
|
graph = StateGraph(State).add_node(my_node).add_edge(START, "my_node").compile()
|
|
|
|
assert graph.invoke({"outer": 1}) == {"outer": 10, "inner": Inner(a=5)}
|
|
|
|
|
|
def test_pydantic_state_mutation_command() -> None:
|
|
class Inner(BaseModel):
|
|
a: int = 0
|
|
|
|
class State(BaseModel):
|
|
inner: Inner = Inner()
|
|
outer: int = 0
|
|
|
|
def my_node(state: State) -> State:
|
|
state.inner.a = 5
|
|
state.outer = 10
|
|
return Command(update=state)
|
|
|
|
graph = StateGraph(State).add_node(my_node).add_edge(START, "my_node").compile()
|
|
|
|
assert graph.invoke({"outer": 1}) == {"outer": 10, "inner": Inner(a=5)}
|
|
|
|
# test w/ default_factory
|
|
class State(BaseModel):
|
|
inner: Inner = Field(default_factory=Inner)
|
|
outer: int = 0
|
|
|
|
def my_node(state: State) -> State:
|
|
state.inner.a = 5
|
|
state.outer = 10
|
|
return Command(update=state)
|
|
|
|
graph = StateGraph(State).add_node(my_node).add_edge(START, "my_node").compile()
|
|
|
|
assert graph.invoke({"outer": 1}) == {"outer": 10, "inner": Inner(a=5)}
|
|
|
|
|
|
def test_get_stream_writer() -> None:
|
|
class State(TypedDict):
|
|
foo: str
|
|
|
|
def my_node(state):
|
|
writer = get_stream_writer()
|
|
writer("custom!")
|
|
return state
|
|
|
|
graph = StateGraph(State).add_node(my_node).add_edge(START, "my_node").compile()
|
|
assert list(graph.stream({"foo": "bar"}, stream_mode="custom")) == ["custom!"]
|
|
assert list(graph.stream({"foo": "bar"}, stream_mode="values")) == [
|
|
{"foo": "bar"},
|
|
{"foo": "bar"},
|
|
]
|
|
assert list(graph.stream({"foo": "bar"}, stream_mode=["custom", "updates"])) == [
|
|
(
|
|
"custom",
|
|
"custom!",
|
|
),
|
|
(
|
|
"updates",
|
|
{
|
|
"my_node": {
|
|
"foo": "bar",
|
|
},
|
|
},
|
|
),
|
|
]
|
|
|
|
|
|
def test_stream_messages_dedupe_inputs() -> None:
|
|
from langchain_core.messages import AIMessage
|
|
|
|
def call_model(state):
|
|
return {"messages": AIMessage("hi", id="1")}
|
|
|
|
def route(state):
|
|
return Command(goto="node_2", graph=Command.PARENT)
|
|
|
|
subgraph = (
|
|
StateGraph(MessagesState)
|
|
.add_node(call_model)
|
|
.add_node(route)
|
|
.add_edge(START, "call_model")
|
|
.add_edge("call_model", "route")
|
|
.compile()
|
|
)
|
|
|
|
graph = (
|
|
StateGraph(MessagesState)
|
|
.add_node("node_1", subgraph)
|
|
.add_node("node_2", lambda state: state)
|
|
.add_edge(START, "node_1")
|
|
.compile()
|
|
)
|
|
|
|
chunks = [
|
|
chunk
|
|
for ns, chunk in graph.stream(
|
|
{"messages": "hi"}, stream_mode="messages", subgraphs=True
|
|
)
|
|
]
|
|
|
|
assert len(chunks) == 1
|
|
assert chunks[0][0] == AIMessage("hi", id="1")
|
|
assert chunks[0][1]["langgraph_node"] == "call_model"
|
|
|
|
|
|
def test_stream_messages_dedupe_state(sync_checkpointer: BaseCheckpointSaver) -> None:
|
|
from langchain_core.messages import AIMessage
|
|
|
|
to_emit = [AIMessage("bye", id="1"), AIMessage("bye again", id="2")]
|
|
|
|
def call_model(state):
|
|
return {"messages": to_emit.pop(0)}
|
|
|
|
def route(state):
|
|
return Command(goto="node_2", graph=Command.PARENT)
|
|
|
|
subgraph = (
|
|
StateGraph(MessagesState)
|
|
.add_node(call_model)
|
|
.add_node(route)
|
|
.add_edge(START, "call_model")
|
|
.add_edge("call_model", "route")
|
|
.compile()
|
|
)
|
|
|
|
graph = (
|
|
StateGraph(MessagesState)
|
|
.add_node("node_1", subgraph)
|
|
.add_node("node_2", lambda state: state)
|
|
.add_edge(START, "node_1")
|
|
.compile(checkpointer=sync_checkpointer)
|
|
)
|
|
|
|
thread1 = {"configurable": {"thread_id": "1"}}
|
|
|
|
chunks = [
|
|
chunk
|
|
for ns, chunk in graph.stream(
|
|
{"messages": "hi"}, thread1, stream_mode="messages", subgraphs=True
|
|
)
|
|
]
|
|
|
|
assert len(chunks) == 1
|
|
assert chunks[0][0] == AIMessage("bye", id="1")
|
|
assert chunks[0][1]["langgraph_node"] == "call_model"
|
|
|
|
chunks = [
|
|
chunk
|
|
for ns, chunk in graph.stream(
|
|
{"messages": "hi again"},
|
|
thread1,
|
|
stream_mode="messages",
|
|
subgraphs=True,
|
|
)
|
|
]
|
|
|
|
assert len(chunks) == 1
|
|
assert chunks[0][0] == AIMessage("bye again", id="2")
|
|
assert chunks[0][1]["langgraph_node"] == "call_model"
|
|
|
|
|
|
def test_stream_messages_dedupe_pydantic_subgraph_interrupt(
|
|
sync_checkpointer: BaseCheckpointSaver,
|
|
) -> None:
|
|
"""Pydantic BaseModel state should not cause duplicate messages when
|
|
streaming from subgraphs that use interrupts. Regression test for a bug
|
|
where ``on_chain_start`` only populated the ``seen`` set for dict inputs,
|
|
skipping Pydantic model inputs entirely."""
|
|
|
|
class PydanticState(BaseModel):
|
|
messages: Annotated[list[AnyMessage], add_messages] = Field(
|
|
default_factory=list
|
|
)
|
|
|
|
def subgraph_proposal(state) -> Command[Literal["subgraph_approval"]]:
|
|
return Command(
|
|
goto="subgraph_approval",
|
|
update={"messages": [AIMessage(content="Proposal", id="proposal_msg")]},
|
|
)
|
|
|
|
def subgraph_approval(state) -> Command[Literal["__end__"]]:
|
|
resume_value = interrupt({"message": "Waiting for approval"})
|
|
user_msg = resume_value.get("user_message", "")
|
|
msgs = [HumanMessage(content=user_msg)] if user_msg else []
|
|
return Command(goto="__end__", update={"messages": msgs})
|
|
|
|
subgraph = (
|
|
StateGraph(PydanticState)
|
|
.add_node("proposal", subgraph_proposal)
|
|
.add_node("subgraph_approval", subgraph_approval)
|
|
.add_edge(START, "proposal")
|
|
.compile(checkpointer=sync_checkpointer)
|
|
)
|
|
|
|
def finalize(state) -> Command[Literal["__end__"]]:
|
|
return Command(
|
|
goto="__end__",
|
|
update={"messages": [AIMessage(content="Finalized", id="finalize_msg")]},
|
|
)
|
|
|
|
graph = (
|
|
StateGraph(PydanticState)
|
|
.add_node("subgraph", subgraph)
|
|
.add_node("finalize", finalize)
|
|
.add_edge(START, "subgraph")
|
|
.add_edge("subgraph", "finalize")
|
|
.compile(checkpointer=sync_checkpointer)
|
|
)
|
|
|
|
thread1 = {"configurable": {"thread_id": "1"}}
|
|
|
|
# First stream: should hit interrupt after proposal
|
|
chunks_req0 = [
|
|
(ns, chunk)
|
|
for ns, chunk in graph.stream(
|
|
{"messages": [HumanMessage(content="Create a proposal")]},
|
|
thread1,
|
|
stream_mode="messages",
|
|
subgraphs=True,
|
|
)
|
|
]
|
|
|
|
msg_ids_req0 = {chunk[0].id for _, chunk in chunks_req0}
|
|
assert "proposal_msg" in msg_ids_req0
|
|
|
|
# Verify interrupted
|
|
state = graph.get_state(thread1)
|
|
assert state.next
|
|
|
|
# Second stream: resume — should NOT duplicate messages from first stream
|
|
chunks_req1 = [
|
|
(ns, chunk)
|
|
for ns, chunk in graph.stream(
|
|
Command(resume={"user_message": "Yes"}),
|
|
thread1,
|
|
stream_mode="messages",
|
|
subgraphs=True,
|
|
)
|
|
]
|
|
|
|
msg_ids_req1 = {chunk[0].id for _, chunk in chunks_req1}
|
|
assert "finalize_msg" in msg_ids_req1
|
|
|
|
# The key assertion: no message IDs from request 0 should appear in request 1
|
|
duplicates = msg_ids_req0 & msg_ids_req1
|
|
assert not duplicates, f"Duplicate message IDs across requests: {duplicates}"
|
|
|
|
|
|
def test_interrupt_subgraph_reenter_checkpointer_true(
|
|
sync_checkpointer: BaseCheckpointSaver,
|
|
) -> None:
|
|
class SubgraphState(TypedDict):
|
|
foo: str
|
|
bar: str
|
|
|
|
class ParentState(TypedDict):
|
|
foo: str
|
|
counter: int
|
|
|
|
called = []
|
|
bar_values = []
|
|
|
|
def subnode_1(state: SubgraphState):
|
|
called.append("subnode_1")
|
|
bar_values.append(state.get("bar"))
|
|
return {"foo": "subgraph_1"}
|
|
|
|
def subnode_2(state: SubgraphState):
|
|
called.append("subnode_2")
|
|
value = interrupt("Provide value")
|
|
value += "baz"
|
|
return {"foo": "subgraph_2", "bar": value}
|
|
|
|
subgraph = (
|
|
StateGraph(SubgraphState)
|
|
.add_node(subnode_1)
|
|
.add_node(subnode_2)
|
|
.add_edge(START, "subnode_1")
|
|
.add_edge("subnode_1", "subnode_2")
|
|
.compile(checkpointer=True)
|
|
)
|
|
|
|
def call_subgraph(state: ParentState):
|
|
called.append("call_subgraph")
|
|
return subgraph.invoke(state)
|
|
|
|
def node(state: ParentState):
|
|
called.append("parent")
|
|
if state["counter"] < 1:
|
|
return Command(
|
|
goto="call_subgraph", update={"counter": state["counter"] + 1}
|
|
)
|
|
|
|
return {"foo": state["foo"] + "|" + "parent"}
|
|
|
|
parent = (
|
|
StateGraph(ParentState)
|
|
.add_node(call_subgraph)
|
|
.add_node(node)
|
|
.add_edge(START, "call_subgraph")
|
|
.add_edge("call_subgraph", "node")
|
|
.compile(checkpointer=sync_checkpointer)
|
|
)
|
|
|
|
config = {"configurable": {"thread_id": "1"}}
|
|
assert parent.invoke({"foo": "", "counter": 0}, config) == {
|
|
"foo": "",
|
|
"counter": 0,
|
|
"__interrupt__": [
|
|
Interrupt(
|
|
value="Provide value",
|
|
id=AnyStr(),
|
|
)
|
|
],
|
|
}
|
|
assert parent.invoke(Command(resume="bar"), config) == {
|
|
"foo": "subgraph_2",
|
|
"counter": 1,
|
|
"__interrupt__": [
|
|
Interrupt(
|
|
value="Provide value",
|
|
id=AnyStr(),
|
|
)
|
|
],
|
|
}
|
|
assert parent.invoke(Command(resume="qux"), config) == {
|
|
"foo": "subgraph_2|parent",
|
|
"counter": 1,
|
|
}
|
|
assert called == [
|
|
"call_subgraph",
|
|
"subnode_1",
|
|
"subnode_2",
|
|
"call_subgraph",
|
|
"subnode_2",
|
|
"parent",
|
|
"call_subgraph",
|
|
"subnode_1",
|
|
"subnode_2",
|
|
"call_subgraph",
|
|
"subnode_2",
|
|
"parent",
|
|
]
|
|
|
|
# invoke parent again (new turn)
|
|
assert parent.invoke({"foo": "meow", "counter": 0}, config) == {
|
|
"foo": "meow",
|
|
"counter": 0,
|
|
"__interrupt__": [
|
|
Interrupt(
|
|
value="Provide value",
|
|
id=AnyStr(),
|
|
)
|
|
],
|
|
}
|
|
# confirm that we preserve the state values from the previous invocation
|
|
assert bar_values == [None, "barbaz", "quxbaz"]
|
|
|
|
|
|
def test_empty_invoke() -> None:
|
|
def reducer_merge_dicts(
|
|
dict1: dict[Any, Any], dict2: dict[Any, Any]
|
|
) -> dict[Any, Any]:
|
|
merged = {**dict1, **dict2}
|
|
return merged
|
|
|
|
class SimpleGraphState(BaseModel):
|
|
x1: Annotated[list[str], operator.add] = []
|
|
x2: Annotated[dict[str, Any], reducer_merge_dicts] = {}
|
|
|
|
def update_x1_1(state: SimpleGraphState):
|
|
print(state)
|
|
return {"x1": ["111"]}
|
|
|
|
def update_x1_2(state: SimpleGraphState):
|
|
print(state)
|
|
state.x1.append("222")
|
|
return {"x1": ["222"]}
|
|
|
|
def update_x2_1(state: SimpleGraphState):
|
|
print(state)
|
|
return {"x2": {"111": 111}}
|
|
|
|
def update_x2_2(state: SimpleGraphState):
|
|
print(state)
|
|
return {"x2": {"222": 222}}
|
|
|
|
graph = StateGraph(SimpleGraphState)
|
|
graph.add_node("x1_1_node", update_x1_1)
|
|
graph.add_node("x1_2_node", update_x1_2)
|
|
graph.add_node("x2_1_node", update_x2_1)
|
|
graph.add_node("x2_2_node", update_x2_2)
|
|
graph.add_edge("x1_1_node", "x1_2_node")
|
|
graph.add_edge("x1_2_node", "x2_1_node")
|
|
graph.add_edge("x2_1_node", "x2_2_node")
|
|
|
|
graph.add_edge(START, "x1_1_node")
|
|
graph.add_edge("x2_2_node", END)
|
|
|
|
compiled = graph.compile()
|
|
|
|
assert compiled.invoke(SimpleGraphState()).get("x2") == {
|
|
"111": 111,
|
|
"222": 222,
|
|
}
|
|
|
|
|
|
def test_parallel_interrupts(sync_checkpointer: BaseCheckpointSaver) -> None:
|
|
# --- CHILD GRAPH ---
|
|
|
|
class ChildState(BaseModel):
|
|
prompt: str = Field(..., description="What is going to be asked to the user?")
|
|
human_input: str | None = Field(None, description="What the human said")
|
|
human_inputs: Annotated[list[str], operator.add] = Field(
|
|
default_factory=list, description="All of my messages"
|
|
)
|
|
|
|
def get_human_input(state: ChildState):
|
|
human_input = interrupt(state.prompt)
|
|
|
|
return dict(
|
|
human_input=human_input, # update child state
|
|
human_inputs=[human_input], # update parent state
|
|
)
|
|
|
|
child_graph_builder = StateGraph(ChildState)
|
|
child_graph_builder.add_node("get_human_input", get_human_input)
|
|
child_graph_builder.add_edge(START, "get_human_input")
|
|
child_graph_builder.add_edge("get_human_input", END)
|
|
child_graph = child_graph_builder.compile()
|
|
|
|
# --- PARENT GRAPH ---
|
|
|
|
class ParentState(BaseModel):
|
|
prompts: list[str] = Field(
|
|
..., description="What is going to be asked to the user?"
|
|
)
|
|
human_inputs: Annotated[list[str], operator.add] = Field(
|
|
default_factory=list, description="All of my messages"
|
|
)
|
|
|
|
def assign_workers(state: ParentState):
|
|
return [
|
|
Send(
|
|
"child_graph",
|
|
dict(
|
|
prompt=prompt,
|
|
),
|
|
)
|
|
for prompt in state.prompts
|
|
]
|
|
|
|
def cleanup(state: ParentState):
|
|
assert len(state.human_inputs) == len(state.prompts)
|
|
|
|
parent_graph_builder = StateGraph(ParentState)
|
|
parent_graph_builder.add_node("child_graph", child_graph)
|
|
parent_graph_builder.add_node("cleanup", cleanup)
|
|
|
|
parent_graph_builder.add_conditional_edges(START, assign_workers, ["child_graph"])
|
|
parent_graph_builder.add_edge("child_graph", "cleanup")
|
|
parent_graph_builder.add_edge("cleanup", END)
|
|
|
|
parent_graph = parent_graph_builder.compile(checkpointer=sync_checkpointer)
|
|
|
|
# --- CLIENT INVOCATION ---
|
|
|
|
thread_config = dict(
|
|
configurable=dict(
|
|
thread_id=str(uuid.uuid4()),
|
|
)
|
|
)
|
|
current_input = dict(
|
|
prompts=["a", "b"],
|
|
)
|
|
|
|
invokes = 0
|
|
events: dict[int, list[dict]] = {}
|
|
while invokes < 10:
|
|
# reset interrupt
|
|
invokes += 1
|
|
events[invokes] = []
|
|
current_interrupts: list[Interrupt] = []
|
|
|
|
# start / resume the graph
|
|
for event in parent_graph.stream(
|
|
input=current_input,
|
|
config=thread_config,
|
|
stream_mode="updates",
|
|
):
|
|
events[invokes].append(event)
|
|
# handle the interrupt
|
|
if "__interrupt__" in event:
|
|
current_interrupts.extend(event["__interrupt__"])
|
|
# assume that it breaks here, because it is an interrupt
|
|
|
|
# get human input and resume
|
|
if len(current_interrupts) > 0:
|
|
# we resume one at a time to preserve original test behavior,
|
|
# but we could also resume all at once if we wanted
|
|
# with a single dict mapping of interrupt ids to resume values
|
|
resume = {current_interrupts[0].id: f"Resume #{invokes}"}
|
|
current_input = Command(resume=resume)
|
|
|
|
# not more human input required, must be completed
|
|
else:
|
|
break
|
|
else:
|
|
assert False, "Detected infinite loop"
|
|
|
|
assert invokes == 3
|
|
assert len(events) == 3
|
|
|
|
assert events[1] == UnsortedSequence(
|
|
{
|
|
"__interrupt__": (
|
|
Interrupt(
|
|
value="a",
|
|
id=AnyStr(),
|
|
),
|
|
)
|
|
},
|
|
{
|
|
"__interrupt__": (
|
|
Interrupt(
|
|
value="b",
|
|
id=AnyStr(),
|
|
),
|
|
)
|
|
},
|
|
)
|
|
assert events[2] in (
|
|
UnsortedSequence(
|
|
{
|
|
"__interrupt__": (
|
|
Interrupt(
|
|
value="a",
|
|
id=AnyStr(),
|
|
),
|
|
)
|
|
},
|
|
{"child_graph": {"human_inputs": ["Resume #1"]}},
|
|
),
|
|
UnsortedSequence(
|
|
{
|
|
"__interrupt__": (
|
|
Interrupt(
|
|
value="b",
|
|
id=AnyStr(),
|
|
),
|
|
)
|
|
},
|
|
{"child_graph": {"human_inputs": ["Resume #1"]}},
|
|
),
|
|
)
|
|
assert events[3] == UnsortedSequence(
|
|
{
|
|
"child_graph": {"human_inputs": ["Resume #1"]},
|
|
"__metadata__": {"cached": True},
|
|
},
|
|
{"child_graph": {"human_inputs": ["Resume #2"]}},
|
|
{"cleanup": None},
|
|
)
|
|
|
|
|
|
def test_parallel_interrupts_double(sync_checkpointer: BaseCheckpointSaver) -> None:
|
|
# --- CHILD GRAPH ---
|
|
|
|
class ChildState(BaseModel):
|
|
prompt: str = Field(..., description="What is going to be asked to the user?")
|
|
human_input: str | None = Field(None, description="What the human said")
|
|
human_inputs: Annotated[list[str], operator.add] = Field(
|
|
default_factory=list, description="All of my messages"
|
|
)
|
|
|
|
def get_human_input(state: ChildState):
|
|
human_input = interrupt(state.prompt)
|
|
|
|
return dict(
|
|
human_inputs=[human_input], # update parent state
|
|
)
|
|
|
|
def get_dolphin_input(state: ChildState):
|
|
human_input = interrupt(state.prompt)
|
|
|
|
return dict(
|
|
human_inputs=[human_input], # update parent state
|
|
)
|
|
|
|
child_graph_builder = StateGraph(ChildState)
|
|
child_graph_builder.add_node("get_human_input", get_human_input)
|
|
child_graph_builder.add_node("get_dolphin_input", get_dolphin_input)
|
|
child_graph_builder.add_edge(START, "get_human_input")
|
|
child_graph_builder.add_edge(START, "get_dolphin_input")
|
|
child_graph = child_graph_builder.compile()
|
|
|
|
# --- PARENT GRAPH ---
|
|
|
|
class ParentState(BaseModel):
|
|
prompts: list[str] = Field(
|
|
..., description="What is going to be asked to the user?"
|
|
)
|
|
human_inputs: Annotated[list[str], operator.add] = Field(
|
|
default_factory=list, description="All of my messages"
|
|
)
|
|
|
|
def assign_workers(state: ParentState):
|
|
return [
|
|
Send(
|
|
"child_graph",
|
|
dict(
|
|
prompt=prompt,
|
|
),
|
|
)
|
|
for prompt in state.prompts
|
|
]
|
|
|
|
def cleanup(state: ParentState):
|
|
assert len(state.human_inputs) == len(state.prompts) * 2
|
|
|
|
parent_graph_builder = StateGraph(ParentState)
|
|
parent_graph_builder.add_node("child_graph", child_graph)
|
|
parent_graph_builder.add_node("cleanup", cleanup)
|
|
|
|
parent_graph_builder.add_conditional_edges(START, assign_workers, ["child_graph"])
|
|
parent_graph_builder.add_edge("child_graph", "cleanup")
|
|
parent_graph_builder.add_edge("cleanup", END)
|
|
|
|
parent_graph = parent_graph_builder.compile(checkpointer=sync_checkpointer)
|
|
|
|
# --- CLIENT INVOCATION ---
|
|
|
|
thread_config = dict(
|
|
configurable=dict(
|
|
thread_id=str(uuid.uuid4()),
|
|
)
|
|
)
|
|
current_input = dict(
|
|
prompts=["a", "b"],
|
|
)
|
|
|
|
invokes = 0
|
|
events: dict[int, list[dict]] = {}
|
|
while invokes < 10:
|
|
# reset interrupt
|
|
invokes += 1
|
|
events[invokes] = []
|
|
current_interrupts: list[Interrupt] = []
|
|
|
|
# start / resume the graph
|
|
for event in parent_graph.stream(
|
|
input=current_input,
|
|
config=thread_config,
|
|
stream_mode="updates",
|
|
):
|
|
events[invokes].append(event)
|
|
# handle the interrupt
|
|
if "__interrupt__" in event:
|
|
current_interrupts.extend(event["__interrupt__"])
|
|
# assume that it breaks here, because it is an interrupt
|
|
|
|
# get human input and resume
|
|
if len(current_interrupts) > 0:
|
|
# we resume one at a time to preserve original test behavior,
|
|
# but we could also resume all at once if we wanted
|
|
# with a single dict mapping of interrupt ids to resume values
|
|
resume = {current_interrupts[0].id: f"Resume #{invokes}"}
|
|
current_input = Command(resume=resume)
|
|
|
|
# not more human input required, must be completed
|
|
else:
|
|
break
|
|
else:
|
|
assert False, "Detected infinite loop"
|
|
|
|
assert invokes == 5
|
|
assert len(events) == 5
|
|
|
|
|
|
def test_pregel_loop_refcount():
|
|
gc.collect()
|
|
try:
|
|
gc.disable()
|
|
|
|
class State(TypedDict):
|
|
messages: Annotated[list, add_messages]
|
|
|
|
graph_builder = StateGraph(State)
|
|
|
|
def chatbot(state: State):
|
|
return {"messages": [("ai", "HIYA")]}
|
|
|
|
graph_builder.add_node("chatbot", chatbot)
|
|
graph_builder.set_entry_point("chatbot")
|
|
graph_builder.set_finish_point("chatbot")
|
|
graph = graph_builder.compile()
|
|
|
|
for _ in range(5):
|
|
graph.invoke({"messages": [{"role": "user", "content": "hi"}]})
|
|
assert (
|
|
len(
|
|
[obj for obj in gc.get_objects() if isinstance(obj, SyncPregelLoop)]
|
|
)
|
|
== 0
|
|
)
|
|
assert (
|
|
len([obj for obj in gc.get_objects() if isinstance(obj, PregelRunner)])
|
|
== 0
|
|
)
|
|
finally:
|
|
gc.enable()
|
|
|
|
|
|
def test_bulk_state_updates(
|
|
sync_checkpointer: BaseCheckpointSaver,
|
|
) -> None:
|
|
class State(TypedDict):
|
|
foo: str
|
|
baz: str
|
|
|
|
def node_a(state: State) -> State:
|
|
return {"foo": "bar"}
|
|
|
|
def node_b(state: State) -> State:
|
|
return {"baz": "qux"}
|
|
|
|
graph = (
|
|
StateGraph(State)
|
|
.add_node("node_a", node_a)
|
|
.add_node("node_b", node_b)
|
|
.add_edge(START, "node_a")
|
|
.add_edge("node_a", "node_b")
|
|
.compile(checkpointer=sync_checkpointer)
|
|
)
|
|
|
|
config = {"configurable": {"thread_id": "1"}}
|
|
|
|
# First update with node_a
|
|
graph.bulk_update_state(
|
|
config,
|
|
[
|
|
[
|
|
StateUpdate(values={"foo": "bar"}, as_node="node_a"),
|
|
]
|
|
],
|
|
)
|
|
|
|
# Then bulk update with both nodes
|
|
graph.bulk_update_state(
|
|
config,
|
|
[
|
|
[
|
|
StateUpdate(values={"foo": "updated"}, as_node="node_a"),
|
|
StateUpdate(values={"baz": "new"}, as_node="node_b"),
|
|
]
|
|
],
|
|
)
|
|
|
|
state = graph.get_state(config)
|
|
assert state.values == {"foo": "updated", "baz": "new"}
|
|
|
|
# Check if there are only two checkpoints
|
|
checkpoints = list(sync_checkpointer.list(config))
|
|
assert len(checkpoints) == 2
|
|
|
|
# perform multiple steps at the same time
|
|
config = {"configurable": {"thread_id": "2"}}
|
|
|
|
graph.bulk_update_state(
|
|
config,
|
|
[
|
|
[
|
|
StateUpdate(values={"foo": "bar"}, as_node="node_a"),
|
|
],
|
|
[
|
|
StateUpdate(values={"foo": "updated"}, as_node="node_a"),
|
|
StateUpdate(values={"baz": "new"}, as_node="node_b"),
|
|
],
|
|
],
|
|
)
|
|
|
|
state = graph.get_state(config)
|
|
assert state.values == {"foo": "updated", "baz": "new"}
|
|
|
|
checkpoints = list(sync_checkpointer.list(config))
|
|
assert len(checkpoints) == 2
|
|
|
|
# Should raise error if updating without as_node
|
|
with pytest.raises(InvalidUpdateError):
|
|
graph.bulk_update_state(
|
|
config,
|
|
[
|
|
[
|
|
StateUpdate(values={"foo": "error"}, as_node=None),
|
|
StateUpdate(values={"bar": "error"}, as_node=None),
|
|
]
|
|
],
|
|
)
|
|
|
|
# Should raise if no updates are provided
|
|
with pytest.raises(ValueError, match="No supersteps provided"):
|
|
graph.bulk_update_state(config, [])
|
|
|
|
# Should raise if no updates are provided
|
|
with pytest.raises(ValueError, match="No updates provided"):
|
|
graph.bulk_update_state(config, [[], []])
|
|
|
|
# Should raise if __end__ or __copy__ update is applied in bulk
|
|
with pytest.raises(InvalidUpdateError):
|
|
graph.bulk_update_state(
|
|
config,
|
|
[
|
|
[
|
|
StateUpdate(values=None, as_node="__end__"),
|
|
StateUpdate(values=None, as_node="__copy__"),
|
|
],
|
|
],
|
|
)
|
|
|
|
|
|
def test_pregel_node_copy() -> None:
|
|
class State(TypedDict):
|
|
foo: str
|
|
|
|
def agent(state: State) -> State:
|
|
return {"foo": "agent"}
|
|
|
|
def tool(state: State) -> State:
|
|
return {"foo": "tool"}
|
|
|
|
graph = (
|
|
StateGraph(State)
|
|
.add_node("agent", agent)
|
|
.add_node("tool", tool)
|
|
.add_edge(START, "agent")
|
|
.add_edge("agent", "tool")
|
|
.compile()
|
|
)
|
|
|
|
graph.invoke({"foo": "input"}, {"configurable": {"thread_id": "1"}})
|
|
graph.copy()
|
|
graph.nodes["agent"].copy({})
|
|
|
|
|
|
def test_update_as_input(
|
|
sync_checkpointer: BaseCheckpointSaver, durability: Durability
|
|
) -> None:
|
|
class State(TypedDict):
|
|
foo: str
|
|
|
|
def agent(state: State) -> State:
|
|
return {"foo": "agent"}
|
|
|
|
def tool(state: State) -> State:
|
|
return {"foo": "tool"}
|
|
|
|
graph = (
|
|
StateGraph(State)
|
|
.add_node("agent", agent)
|
|
.add_node("tool", tool)
|
|
.add_edge(START, "agent")
|
|
.add_edge("agent", "tool")
|
|
.compile(checkpointer=sync_checkpointer)
|
|
)
|
|
|
|
assert graph.invoke(
|
|
{"foo": "input"},
|
|
{"configurable": {"thread_id": "1"}},
|
|
durability=durability,
|
|
) == {"foo": "tool"}
|
|
|
|
assert graph.invoke(
|
|
{"foo": "input"},
|
|
{"configurable": {"thread_id": "1"}},
|
|
durability=durability,
|
|
) == {"foo": "tool"}
|
|
|
|
def map_snapshot(i: StateSnapshot) -> dict:
|
|
return {
|
|
"values": i.values,
|
|
"next": i.next,
|
|
"step": i.metadata.get("step"),
|
|
}
|
|
|
|
history = [
|
|
map_snapshot(s)
|
|
for s in graph.get_state_history({"configurable": {"thread_id": "1"}})
|
|
]
|
|
|
|
graph.bulk_update_state(
|
|
{"configurable": {"thread_id": "2"}},
|
|
[
|
|
# First turn
|
|
[StateUpdate({"foo": "input"}, "__input__")],
|
|
[StateUpdate({"foo": "input"}, "__start__")],
|
|
[StateUpdate({"foo": "agent"}, "agent")],
|
|
[StateUpdate({"foo": "tool"}, "tool")],
|
|
# Second turn
|
|
[StateUpdate({"foo": "input"}, "__input__")],
|
|
[StateUpdate({"foo": "input"}, "__start__")],
|
|
[StateUpdate({"foo": "agent"}, "agent")],
|
|
[StateUpdate({"foo": "tool"}, "tool")],
|
|
],
|
|
)
|
|
|
|
state = graph.get_state({"configurable": {"thread_id": "2"}})
|
|
assert state.values == {"foo": "tool"}
|
|
|
|
new_history = [
|
|
map_snapshot(s)
|
|
for s in graph.get_state_history({"configurable": {"thread_id": "2"}})
|
|
]
|
|
|
|
if durability != "exit":
|
|
assert new_history == history
|
|
else:
|
|
assert [new_history[0], new_history[4]] == history
|
|
|
|
|
|
def test_batch_update_as_input(
|
|
sync_checkpointer: BaseCheckpointSaver, durability: Durability
|
|
) -> None:
|
|
class State(TypedDict):
|
|
foo: str
|
|
tasks: Annotated[list[int], operator.add]
|
|
|
|
def agent(state: State) -> State:
|
|
return {"foo": "agent"}
|
|
|
|
def map(state: State) -> Command["task"]:
|
|
return Command(
|
|
goto=[
|
|
Send("task", {"index": 0}),
|
|
Send("task", {"index": 1}),
|
|
Send("task", {"index": 2}),
|
|
],
|
|
update={"foo": "map"},
|
|
)
|
|
|
|
def task(state: dict) -> State:
|
|
return {"tasks": [state["index"]]}
|
|
|
|
graph = (
|
|
StateGraph(State)
|
|
.add_node("agent", agent)
|
|
.add_node("map", map)
|
|
.add_node("task", task)
|
|
.add_edge(START, "agent")
|
|
.add_edge("agent", "map")
|
|
.compile(checkpointer=sync_checkpointer)
|
|
)
|
|
|
|
assert graph.invoke(
|
|
{"foo": "input"},
|
|
{"configurable": {"thread_id": "1"}},
|
|
durability=durability,
|
|
) == {
|
|
"foo": "map",
|
|
"tasks": [0, 1, 2],
|
|
}
|
|
|
|
def map_snapshot(i: StateSnapshot) -> dict:
|
|
return {
|
|
"values": i.values,
|
|
"next": i.next,
|
|
"step": i.metadata.get("step"),
|
|
"tasks": [t.name for t in i.tasks],
|
|
}
|
|
|
|
history = [
|
|
map_snapshot(s)
|
|
for s in graph.get_state_history({"configurable": {"thread_id": "1"}})
|
|
]
|
|
|
|
graph.bulk_update_state(
|
|
{"configurable": {"thread_id": "2"}},
|
|
[
|
|
[StateUpdate({"foo": "input"}, "__input__")],
|
|
[StateUpdate({"foo": "input"}, "__start__")],
|
|
[StateUpdate({"foo": "agent", "tasks": []}, "agent")],
|
|
[
|
|
StateUpdate(
|
|
Command(
|
|
goto=[
|
|
Send("task", {"index": 0}),
|
|
Send("task", {"index": 1}),
|
|
Send("task", {"index": 2}),
|
|
],
|
|
update={"foo": "map"},
|
|
),
|
|
"map",
|
|
)
|
|
],
|
|
[
|
|
StateUpdate({"tasks": [0]}, "task"),
|
|
StateUpdate({"tasks": [1]}, "task"),
|
|
StateUpdate({"tasks": [2]}, "task"),
|
|
],
|
|
],
|
|
)
|
|
|
|
state = graph.get_state({"configurable": {"thread_id": "2"}})
|
|
assert state.values == {"foo": "map", "tasks": [0, 1, 2]}
|
|
|
|
new_history = [
|
|
map_snapshot(s)
|
|
for s in graph.get_state_history({"configurable": {"thread_id": "2"}})
|
|
]
|
|
|
|
if durability != "exit":
|
|
assert new_history == history
|
|
else:
|
|
assert new_history[:1] == history
|
|
|
|
|
|
def test_migration_graph(snapshot: SnapshotAssertion) -> None:
|
|
class DummyState(BaseModel):
|
|
pass_count: int = 0
|
|
|
|
def increment_pass_count(state: DummyState):
|
|
state.pass_count += 1
|
|
return state
|
|
|
|
def route_b(state: DummyState):
|
|
if state.pass_count == 0:
|
|
return "X"
|
|
else:
|
|
return "Y"
|
|
|
|
migration_graph = StateGraph(DummyState)
|
|
|
|
migration_graph.add_node("B", increment_pass_count)
|
|
migration_graph.add_node("C", increment_pass_count)
|
|
migration_graph.add_node("D", increment_pass_count)
|
|
|
|
migration_graph.add_edge(START, "B")
|
|
|
|
migration_graph.add_conditional_edges(
|
|
"B",
|
|
route_b,
|
|
{
|
|
"X": "C",
|
|
"Y": "D",
|
|
},
|
|
)
|
|
|
|
migration_graph.add_edge("D", "B")
|
|
migration_graph.add_edge("C", END)
|
|
|
|
app = migration_graph.compile()
|
|
|
|
assert app.get_graph().draw_mermaid(with_styles=False) == snapshot
|
|
|
|
|
|
def test_get_graph_loop(snapshot: SnapshotAssertion) -> None:
|
|
class State(TypedDict):
|
|
foo: str
|
|
|
|
def human_node(state: State) -> State:
|
|
value = interrupt()
|
|
return {"foo": value}
|
|
|
|
def agent_node(state: State) -> State:
|
|
return {"foo": "Hi " + state["foo"]}
|
|
|
|
workflow = StateGraph(State)
|
|
workflow.add_node("human", human_node)
|
|
workflow.add_node("agent", agent_node)
|
|
workflow.add_edge(START, "human")
|
|
workflow.add_edge("human", "agent")
|
|
workflow.add_edge("agent", "human")
|
|
|
|
app = workflow.compile()
|
|
assert json.dumps(app.get_graph().to_json(), indent=2) == snapshot
|
|
assert app.get_graph().draw_mermaid(with_styles=False) == snapshot
|
|
|
|
|
|
def test_get_graph_self_loop(snapshot: SnapshotAssertion) -> None:
|
|
import random
|
|
|
|
subgraph_builder = StateGraph(MessagesState)
|
|
subgraph_builder.add_node("agent", lambda x: x)
|
|
subgraph_builder.add_edge(START, "agent")
|
|
subgraph = subgraph_builder.compile()
|
|
|
|
def worker_node(state: MessagesState) -> Command[Literal["worker_node", "__end__"]]:
|
|
subgraph_result = subgraph.invoke(state)
|
|
|
|
if random.choice([True, False]):
|
|
next_node_name = "worker_node"
|
|
else:
|
|
next_node_name = END
|
|
|
|
return Command(update=subgraph_result, goto=next_node_name)
|
|
|
|
self_loop_builder = StateGraph(MessagesState)
|
|
self_loop_builder.add_node("worker_node", worker_node)
|
|
self_loop_builder.add_edge(START, "worker_node")
|
|
self_loop_graph = self_loop_builder.compile()
|
|
|
|
assert json.dumps(self_loop_graph.get_graph().to_json(), indent=2) == snapshot
|
|
assert self_loop_graph.get_graph().draw_mermaid(with_styles=False) == snapshot
|
|
|
|
|
|
def test_get_graph_root_channel(snapshot: SnapshotAssertion) -> None:
|
|
child_builder = StateGraph(list)
|
|
child_builder.add_node("child_node", lambda x: x)
|
|
child_builder.add_edge(START, "child_node")
|
|
child_graph = child_builder.compile()
|
|
|
|
graph_builder = StateGraph(list)
|
|
graph_builder.add_node("child", child_graph)
|
|
graph_builder.add_edge(START, "child")
|
|
graph = graph_builder.compile()
|
|
|
|
assert json.dumps(graph.get_graph().to_json(), indent=2) == snapshot
|
|
assert graph.get_graph().draw_mermaid(with_styles=False) == snapshot
|
|
|
|
|
|
def test_imp_exception(
|
|
sync_checkpointer: BaseCheckpointSaver,
|
|
) -> None:
|
|
@task()
|
|
def my_task(number: int):
|
|
time.sleep(0.1)
|
|
return number * 2
|
|
|
|
@task()
|
|
def task_with_exception(number: int):
|
|
time.sleep(0.1)
|
|
raise Exception("This is a test exception")
|
|
|
|
@entrypoint(checkpointer=sync_checkpointer)
|
|
def my_workflow(number: int):
|
|
my_task(number).result()
|
|
try:
|
|
task_with_exception(number).result()
|
|
except Exception as e:
|
|
print(f"Exception caught: {e}")
|
|
my_task(number).result()
|
|
return "done"
|
|
|
|
thread1 = {"configurable": {"thread_id": "1"}}
|
|
assert my_workflow.invoke(1, thread1) == "done"
|
|
|
|
assert [c for c in my_workflow.stream(1, thread1)] == [
|
|
{"my_task": 2},
|
|
{"my_task": 2},
|
|
{"my_workflow": "done"},
|
|
]
|
|
|
|
|
|
@pytest.mark.parametrize("with_timeout", [False, "inner", "outer", "both"])
|
|
@pytest.mark.parametrize("subgraph_persist", [True, False])
|
|
def test_parent_command_goto(
|
|
sync_checkpointer: BaseCheckpointSaver, subgraph_persist: bool, with_timeout: bool
|
|
) -> None:
|
|
class State(TypedDict):
|
|
dialog_state: Annotated[list[str], operator.add]
|
|
|
|
def node_a_child(state):
|
|
return {"dialog_state": ["a_child_state"]}
|
|
|
|
def node_b_child(state):
|
|
return Command(
|
|
graph=Command.PARENT,
|
|
goto="node_b_parent",
|
|
update={"dialog_state": ["b_child_state"]},
|
|
)
|
|
|
|
sub_builder = StateGraph(State)
|
|
sub_builder.add_node(node_a_child)
|
|
sub_builder.add_node(node_b_child)
|
|
sub_builder.add_edge(START, "node_a_child")
|
|
sub_builder.add_edge("node_a_child", "node_b_child")
|
|
sub_graph = sub_builder.compile(checkpointer=subgraph_persist)
|
|
if with_timeout in ("inner", "both"):
|
|
sub_graph.step_timeout = 1
|
|
|
|
def node_b_parent(state):
|
|
return {"dialog_state": ["node_b_parent"]}
|
|
|
|
main_builder = StateGraph(State)
|
|
main_builder.add_node(node_b_parent)
|
|
main_builder.add_edge(START, "subgraph_node")
|
|
main_builder.add_node("subgraph_node", sub_graph, destinations=("node_b_parent",))
|
|
main_graph = main_builder.compile(sync_checkpointer, name="parent")
|
|
if with_timeout in ("outer", "both"):
|
|
main_graph.step_timeout = 1
|
|
|
|
config = {"configurable": {"thread_id": 1}}
|
|
|
|
assert main_graph.invoke(input={"dialog_state": ["init_state"]}, config=config) == {
|
|
"dialog_state": ["init_state", "b_child_state", "node_b_parent"]
|
|
}
|
|
|
|
|
|
@pytest.mark.parametrize("subgraph_persist", [True, False])
|
|
def test_parent_command_goto_deeply_nested(
|
|
sync_checkpointer: BaseCheckpointSaver,
|
|
subgraph_persist: bool,
|
|
) -> None:
|
|
"""Test Command.PARENT in a 3-level nested subgraph.
|
|
|
|
Command.PARENT should jump to sub_child_3 in the immediate parent (sub_graph).
|
|
|
|
Note: With operator.add, subgraph state (including its input) is merged with
|
|
parent state, causing the input to appear multiple times. This is expected.
|
|
"""
|
|
|
|
class State(TypedDict):
|
|
dialog_state: Annotated[list[str], operator.add]
|
|
|
|
# Level 3: Deepest subgraph that issues Command.PARENT
|
|
def sub_sub_child_node(state):
|
|
# Jump to immediate parent (sub_graph)
|
|
return Command(
|
|
graph=Command.PARENT,
|
|
goto="sub_child_3",
|
|
update={"dialog_state": ["sub_sub_child"]},
|
|
)
|
|
|
|
sub_sub_builder = StateGraph(State)
|
|
sub_sub_builder.add_node("sub_sub_child", sub_sub_child_node)
|
|
sub_sub_builder.add_edge(START, "sub_sub_child")
|
|
sub_sub_graph = sub_sub_builder.compile(
|
|
name="sub_sub_graph", checkpointer=subgraph_persist
|
|
)
|
|
|
|
# Level 2: Middle subgraph containing Level 3
|
|
def sub_child_1(state):
|
|
return {"dialog_state": ["sub_child_1"]}
|
|
|
|
def sub_child_3(state):
|
|
return {"dialog_state": ["sub_child_3"]}
|
|
|
|
sub_builder = StateGraph(State)
|
|
sub_builder.add_node("sub_child_1", sub_child_1)
|
|
sub_builder.add_node("sub_child_2", sub_sub_graph, destinations=("sub_child_3",))
|
|
sub_builder.add_node("sub_child_3", sub_child_3)
|
|
sub_builder.add_edge(START, "sub_child_1")
|
|
sub_builder.add_edge("sub_child_1", "sub_child_2")
|
|
sub_graph = sub_builder.compile(name="sub_graph", checkpointer=subgraph_persist)
|
|
|
|
# Level 1: Main graph containing Level 2
|
|
def child_1(state):
|
|
return {"dialog_state": ["child_1"]}
|
|
|
|
builder = StateGraph(State)
|
|
builder.add_node("child_1", child_1)
|
|
builder.add_node("child_2", sub_graph)
|
|
builder.add_edge(START, "child_1")
|
|
builder.add_edge("child_1", "child_2")
|
|
graph = builder.compile(name="main_graph", checkpointer=sync_checkpointer)
|
|
|
|
config = {"configurable": {"thread_id": 1}}
|
|
|
|
result = graph.invoke(input={"dialog_state": ["init"]}, config=config)
|
|
|
|
# Command.PARENT from sub_sub_child jumps to sub_child_3 in immediate parent
|
|
# State duplication occurs due to operator.add merging behavior
|
|
assert result == {
|
|
"dialog_state": [
|
|
"init",
|
|
"child_1",
|
|
"init",
|
|
"child_1",
|
|
"sub_child_1",
|
|
"sub_sub_child",
|
|
"sub_child_3",
|
|
]
|
|
}
|
|
|
|
|
|
@pytest.mark.parametrize("with_timeout", [True, False])
|
|
def test_timeout_with_parent_command(
|
|
sync_checkpointer: BaseCheckpointSaver, with_timeout: bool
|
|
) -> None:
|
|
"""Test that parent commands are properly propagated during timeouts."""
|
|
|
|
class State(TypedDict):
|
|
value: str
|
|
|
|
def parent_command_node(state: State) -> State:
|
|
time.sleep(0.1) # Add some delay before raising
|
|
return Command(graph=Command.PARENT, goto="test_cmd", update={"key": "value"})
|
|
|
|
builder = StateGraph(State)
|
|
builder.add_node("parent_cmd", parent_command_node)
|
|
builder.set_entry_point("parent_cmd")
|
|
graph = builder.compile(checkpointer=sync_checkpointer)
|
|
if with_timeout:
|
|
graph.step_timeout = 1
|
|
|
|
# Should propagate parent command, not timeout
|
|
thread1 = {"configurable": {"thread_id": "1"}}
|
|
with pytest.raises(ParentCommand) as exc_info:
|
|
graph.invoke({"value": "start"}, thread1)
|
|
assert exc_info.value.args[0].goto == "test_cmd"
|
|
assert exc_info.value.args[0].update == {"key": "value"}
|
|
|
|
|
|
def test_fork_and_update_task_results(sync_checkpointer: BaseCheckpointSaver) -> None:
|
|
"""Test forking and updating task results with state history."""
|
|
|
|
def checkpoint(values: dict[str, Any]):
|
|
return ("checkpoint", {"values": values})
|
|
|
|
def task(name: str, result: Any):
|
|
return ("task", {"name": name, "result": result})
|
|
|
|
def get_tree(history: list[StateSnapshot]) -> list:
|
|
"""Build a tree structure from state history for comparison."""
|
|
if not history:
|
|
return []
|
|
|
|
# Build a tree structure similar to renderForks
|
|
node_map: dict[str, dict] = {}
|
|
root_nodes: list[dict] = []
|
|
|
|
# Second pass: establish parent-child relationships
|
|
for item in reversed(history):
|
|
checkpoint_id = item.config["configurable"]["checkpoint_id"]
|
|
parent_checkpoint_id = (
|
|
item.parent_config["configurable"]["checkpoint_id"]
|
|
if item.parent_config
|
|
else None
|
|
)
|
|
node_map[checkpoint_id] = {"item": item, "children": []}
|
|
|
|
parent = node_map.get(parent_checkpoint_id)
|
|
(parent["children"] if parent else root_nodes).append(
|
|
node_map[checkpoint_id]
|
|
)
|
|
|
|
def node_to_tree(node: dict) -> list:
|
|
"""Convert a node to tree structure."""
|
|
result = [
|
|
checkpoint(node["item"].values),
|
|
] + [
|
|
task(task_info.name, task_info.result)
|
|
for task_info in node["item"].tasks
|
|
]
|
|
|
|
if len(node["children"]) > 1:
|
|
branches = [node_to_tree(child) for child in node["children"]]
|
|
return result + [branches]
|
|
elif len(node["children"]) == 1:
|
|
return result + node_to_tree(node["children"][0])
|
|
else:
|
|
return result
|
|
|
|
if len(root_nodes) == 1:
|
|
# Process all root nodes
|
|
return node_to_tree(root_nodes[0])
|
|
|
|
elif len(root_nodes) > 1:
|
|
# Multiple root nodes - treat as branches
|
|
branches = [node_to_tree(node) for node in root_nodes]
|
|
return branches
|
|
else:
|
|
return []
|
|
|
|
class State(TypedDict):
|
|
name: Annotated[str, lambda a, b: " > ".join([a, b]) if a else b]
|
|
|
|
# Define the graph with a sequence of nodes
|
|
def one(state: State) -> Command:
|
|
return Command(goto=[Send("two", {})], update={"name": "one"})
|
|
|
|
def two(state: State) -> State:
|
|
return {"name": "two"}
|
|
|
|
def three(state: State) -> State:
|
|
return {"name": "three"}
|
|
|
|
graph = (
|
|
StateGraph(State)
|
|
.add_node("one", one)
|
|
.add_node("two", two)
|
|
.add_node("three", three)
|
|
.add_edge(START, "one")
|
|
.add_edge("one", "two")
|
|
.add_edge("two", "three")
|
|
.compile(checkpointer=sync_checkpointer)
|
|
)
|
|
|
|
config = {"configurable": {"thread_id": "1"}}
|
|
history: list[StateSnapshot] = []
|
|
|
|
# Initial run
|
|
graph.invoke({"name": "start"}, config)
|
|
history = list(graph.get_state_history(config))
|
|
|
|
assert get_tree(history) == [
|
|
checkpoint({"name": ""}),
|
|
task("__start__", {"name": "start"}),
|
|
checkpoint({"name": "start"}),
|
|
task("one", {"name": "one"}),
|
|
checkpoint({"name": "start > one"}),
|
|
task("two", {"name": "two"}),
|
|
task("two", {"name": "two"}),
|
|
checkpoint({"name": "start > one > two > two"}),
|
|
task("three", {"name": "three"}),
|
|
checkpoint({"name": "start > one > two > two > three"}),
|
|
]
|
|
|
|
# Update the start state
|
|
graph.invoke(
|
|
None,
|
|
graph.update_state(
|
|
history[4].config,
|
|
values=[StateUpdate(values={"name": "start*"}, as_node="__start__")],
|
|
as_node="__copy__",
|
|
),
|
|
)
|
|
|
|
history = list(graph.get_state_history(config))
|
|
assert get_tree(history) == [
|
|
[
|
|
checkpoint({"name": ""}),
|
|
task("__start__", {"name": "start"}),
|
|
checkpoint({"name": "start"}),
|
|
task("one", {"name": "one"}),
|
|
checkpoint({"name": "start > one"}),
|
|
task("two", {"name": "two"}),
|
|
task("two", {"name": "two"}),
|
|
checkpoint({"name": "start > one > two > two"}),
|
|
task("three", {"name": "three"}),
|
|
checkpoint({"name": "start > one > two > two > three"}),
|
|
],
|
|
[
|
|
checkpoint({"name": ""}),
|
|
task("__start__", {"name": "start*"}),
|
|
checkpoint({"name": "start*"}),
|
|
task("one", {"name": "one"}),
|
|
checkpoint({"name": "start* > one"}),
|
|
task("two", {"name": "two"}),
|
|
task("two", {"name": "two"}),
|
|
checkpoint({"name": "start* > one > two > two"}),
|
|
task("three", {"name": "three"}),
|
|
checkpoint({"name": "start* > one > two > two > three"}),
|
|
],
|
|
]
|
|
|
|
# Fork from task "one"
|
|
# Start from the checkpoint that has the task "one"
|
|
assert history[3].values == {"name": "start*"}
|
|
assert len(history[3].tasks) == 1
|
|
assert history[3].tasks[0].name == "one"
|
|
|
|
graph.invoke(
|
|
None,
|
|
graph.update_state(
|
|
history[3].config,
|
|
[StateUpdate(values={"name": "one*"}, as_node="one")],
|
|
"__copy__",
|
|
),
|
|
)
|
|
|
|
history = list(graph.get_state_history(config))
|
|
assert get_tree(history) == [
|
|
[
|
|
checkpoint({"name": ""}),
|
|
task("__start__", {"name": "start"}),
|
|
checkpoint({"name": "start"}),
|
|
task("one", {"name": "one"}),
|
|
checkpoint({"name": "start > one"}),
|
|
task("two", {"name": "two"}),
|
|
task("two", {"name": "two"}),
|
|
checkpoint({"name": "start > one > two > two"}),
|
|
task("three", {"name": "three"}),
|
|
checkpoint({"name": "start > one > two > two > three"}),
|
|
],
|
|
[
|
|
checkpoint({"name": ""}),
|
|
task("__start__", {"name": "start*"}),
|
|
[
|
|
[
|
|
checkpoint({"name": "start*"}),
|
|
task("one", {"name": "one"}),
|
|
checkpoint({"name": "start* > one"}),
|
|
task("two", {"name": "two"}),
|
|
task("two", {"name": "two"}),
|
|
checkpoint({"name": "start* > one > two > two"}),
|
|
task("three", {"name": "three"}),
|
|
checkpoint({"name": "start* > one > two > two > three"}),
|
|
],
|
|
[
|
|
checkpoint({"name": "start*"}),
|
|
task("one", {"name": "one*"}),
|
|
checkpoint({"name": "start* > one*"}),
|
|
task("two", {"name": "two"}),
|
|
checkpoint({"name": "start* > one* > two"}),
|
|
task("three", {"name": "three"}),
|
|
checkpoint({"name": "start* > one* > two > three"}),
|
|
],
|
|
],
|
|
],
|
|
]
|
|
|
|
config = {"configurable": {"thread_id": "2"}}
|
|
|
|
# Initialize the thread once again
|
|
graph.invoke({"name": "start"}, config)
|
|
history = list(graph.get_state_history(config))
|
|
|
|
# Fork from task "two"
|
|
# Start from the checkpoint that has the task "two"
|
|
assert history[2].values == {"name": "start > one"}
|
|
|
|
graph.invoke(
|
|
None,
|
|
graph.update_state(
|
|
history[2].config,
|
|
[
|
|
StateUpdate(values={"name": "two"}, as_node="two"),
|
|
StateUpdate(values={"name": "two"}, as_node="two"),
|
|
],
|
|
"__copy__",
|
|
),
|
|
)
|
|
|
|
history = list(graph.get_state_history(config))
|
|
assert get_tree(history) == [
|
|
checkpoint({"name": ""}),
|
|
task("__start__", {"name": "start"}),
|
|
checkpoint({"name": "start"}),
|
|
task("one", {"name": "one"}),
|
|
[
|
|
[
|
|
checkpoint({"name": "start > one"}),
|
|
task("two", {"name": "two"}),
|
|
task("two", {"name": "two"}),
|
|
checkpoint({"name": "start > one > two > two"}),
|
|
task("three", {"name": "three"}),
|
|
checkpoint({"name": "start > one > two > two > three"}),
|
|
],
|
|
[
|
|
checkpoint({"name": "start > one"}),
|
|
task("two", {"name": "two"}),
|
|
task("two", {"name": "two"}),
|
|
checkpoint({"name": "start > one > two > two"}),
|
|
task("three", {"name": "three"}),
|
|
checkpoint({"name": "start > one > two > two > three"}),
|
|
],
|
|
],
|
|
]
|
|
|
|
# Fork task three
|
|
assert history[1].values == {"name": "start > one > two > two"}
|
|
assert len(history[1].tasks) == 1
|
|
assert history[1].tasks[0].name == "three"
|
|
|
|
graph.invoke(
|
|
None,
|
|
graph.update_state(
|
|
history[1].config,
|
|
[StateUpdate(values={"name": "three*"}, as_node="three")],
|
|
"__copy__",
|
|
),
|
|
)
|
|
|
|
history = list(graph.get_state_history(config))
|
|
assert get_tree(history) == [
|
|
checkpoint({"name": ""}),
|
|
task("__start__", {"name": "start"}),
|
|
checkpoint({"name": "start"}),
|
|
task("one", {"name": "one"}),
|
|
[
|
|
[
|
|
checkpoint({"name": "start > one"}),
|
|
task("two", {"name": "two"}),
|
|
task("two", {"name": "two"}),
|
|
checkpoint({"name": "start > one > two > two"}),
|
|
task("three", {"name": "three"}),
|
|
checkpoint({"name": "start > one > two > two > three"}),
|
|
],
|
|
[
|
|
checkpoint({"name": "start > one"}),
|
|
task("two", {"name": "two"}),
|
|
task("two", {"name": "two"}),
|
|
[
|
|
[
|
|
checkpoint({"name": "start > one > two > two"}),
|
|
task("three", {"name": "three"}),
|
|
checkpoint({"name": "start > one > two > two > three"}),
|
|
],
|
|
[
|
|
checkpoint({"name": "start > one > two > two"}),
|
|
task("three", {"name": "three*"}),
|
|
checkpoint({"name": "start > one > two > two > three*"}),
|
|
],
|
|
],
|
|
],
|
|
],
|
|
]
|
|
|
|
# Regenerate task three
|
|
assert history[3].values == {"name": "start > one > two > two"}
|
|
assert len(history[3].tasks) == 1
|
|
assert history[3].tasks[0].name == "three"
|
|
|
|
graph.invoke(None, graph.update_state(history[3].config, None, "__copy__"))
|
|
|
|
history = list(graph.get_state_history(config))
|
|
assert get_tree(history) == [
|
|
checkpoint({"name": ""}),
|
|
task("__start__", {"name": "start"}),
|
|
checkpoint({"name": "start"}),
|
|
task("one", {"name": "one"}),
|
|
[
|
|
[
|
|
checkpoint({"name": "start > one"}),
|
|
task("two", {"name": "two"}),
|
|
task("two", {"name": "two"}),
|
|
checkpoint({"name": "start > one > two > two"}),
|
|
task("three", {"name": "three"}),
|
|
checkpoint({"name": "start > one > two > two > three"}),
|
|
],
|
|
[
|
|
checkpoint({"name": "start > one"}),
|
|
task("two", {"name": "two"}),
|
|
task("two", {"name": "two"}),
|
|
[
|
|
[
|
|
checkpoint({"name": "start > one > two > two"}),
|
|
task("three", {"name": "three"}),
|
|
checkpoint({"name": "start > one > two > two > three"}),
|
|
],
|
|
[
|
|
checkpoint({"name": "start > one > two > two"}),
|
|
task("three", {"name": "three*"}),
|
|
checkpoint({"name": "start > one > two > two > three*"}),
|
|
],
|
|
[
|
|
checkpoint({"name": "start > one > two > two"}),
|
|
task("three", {"name": "three"}),
|
|
checkpoint({"name": "start > one > two > two > three"}),
|
|
],
|
|
],
|
|
],
|
|
],
|
|
]
|
|
|
|
|
|
def test_subgraph_streaming_sync() -> None:
|
|
"""Test subgraph streaming when used as a node in sync version"""
|
|
|
|
# Create a fake chat model that returns a simple response
|
|
model = GenericFakeChatModel(messages=iter(["The weather is sunny today."]))
|
|
|
|
# Create a subgraph that uses the fake chat model
|
|
def call_model_node(state: MessagesState, config: RunnableConfig) -> MessagesState:
|
|
"""Node that calls the model with the last message."""
|
|
messages = state["messages"]
|
|
last_message = messages[-1].content if messages else ""
|
|
response = model.invoke([("user", last_message)], config)
|
|
return {"messages": [response]}
|
|
|
|
# Build the subgraph
|
|
subgraph = StateGraph(MessagesState)
|
|
subgraph.add_node("call_model", call_model_node)
|
|
subgraph.add_edge(START, "call_model")
|
|
compiled_subgraph = subgraph.compile()
|
|
|
|
class SomeCustomState(TypedDict):
|
|
last_chunk: NotRequired[str]
|
|
num_chunks: NotRequired[int]
|
|
|
|
# Will invoke a subgraph as a function
|
|
def parent_node(state: SomeCustomState, config: RunnableConfig) -> dict:
|
|
"""Node that runs the subgraph."""
|
|
msgs = {"messages": [("user", "What is the weather in Tokyo?")]}
|
|
events = []
|
|
for event in compiled_subgraph.stream(msgs, config, stream_mode="messages"):
|
|
events.append(event)
|
|
ai_msg_chunks = [ai_msg_chunk for ai_msg_chunk, _ in events]
|
|
return {
|
|
"last_chunk": ai_msg_chunks[-1],
|
|
"num_chunks": len(ai_msg_chunks),
|
|
}
|
|
|
|
# Build the main workflow
|
|
workflow = StateGraph(SomeCustomState)
|
|
workflow.add_node("subgraph", parent_node)
|
|
workflow.add_edge(START, "subgraph")
|
|
compiled_workflow = workflow.compile()
|
|
|
|
# Test the basic functionality
|
|
result = compiled_workflow.invoke({})
|
|
|
|
assert result["last_chunk"].content == "today."
|
|
assert result["num_chunks"] == 9
|
|
|
|
|
|
def test_get_graph_nonterminal_last_step_source(snapshot: SnapshotAssertion) -> None:
|
|
class State(TypedDict):
|
|
messages: list[str]
|
|
|
|
def chatbot_node(state: State) -> State:
|
|
return {"messages": state["messages"] + ["chatbot"]}
|
|
|
|
def tools_node(state: State) -> State:
|
|
return {"messages": state["messages"] + ["tools"]}
|
|
|
|
def human_node(state: State) -> State:
|
|
return {"messages": state["messages"] + ["human"]}
|
|
|
|
def tools_condition(_: State) -> str:
|
|
return "tools"
|
|
|
|
def end_condition(_: State) -> str:
|
|
return "chatbot"
|
|
|
|
workflow = StateGraph(State)
|
|
workflow.add_node("chatbot", chatbot_node)
|
|
workflow.add_node("tools", tools_node)
|
|
workflow.add_node("human", human_node)
|
|
|
|
workflow.add_edge(START, "human")
|
|
workflow.add_edge("tools", "chatbot")
|
|
|
|
workflow.add_conditional_edges(
|
|
"chatbot", tools_condition, {"tools": "tools", "human": "human"}
|
|
)
|
|
workflow.add_conditional_edges(
|
|
"human", end_condition, {"chatbot": "chatbot", END: END}
|
|
)
|
|
|
|
app = workflow.compile()
|
|
graph = app.get_graph()
|
|
graph_json = graph.to_json()
|
|
|
|
assert json.dumps(graph_json, indent=2, sort_keys=True) == snapshot
|
|
|
|
|
|
def test_null_resume_disallowed_with_multiple_interrupts(
|
|
sync_checkpointer: BaseCheckpointSaver,
|
|
) -> None:
|
|
class State(TypedDict):
|
|
text_1: str
|
|
text_2: str
|
|
|
|
def human_node_1(state: State):
|
|
value = interrupt({"text_to_revise": state["text_1"]})
|
|
return {"text_1": value}
|
|
|
|
def human_node_2(state: State):
|
|
value = interrupt({"text_to_revise": state["text_2"]})
|
|
return {"text_2": value}
|
|
|
|
graph_builder = StateGraph(State)
|
|
graph_builder.add_node("human_node_1", human_node_1)
|
|
graph_builder.add_node("human_node_2", human_node_2)
|
|
|
|
# Add both nodes in parallel from START
|
|
graph_builder.add_edge(START, "human_node_1")
|
|
graph_builder.add_edge(START, "human_node_2")
|
|
|
|
checkpointer = InMemorySaver()
|
|
graph = graph_builder.compile(checkpointer=checkpointer)
|
|
|
|
thread_id = str(uuid.uuid4())
|
|
config: RunnableConfig = {"configurable": {"thread_id": thread_id}}
|
|
graph.invoke(
|
|
{"text_1": "original text 1", "text_2": "original text 2"}, config=config
|
|
)
|
|
|
|
resume_map = {
|
|
i.id: f"resume for prompt: {i.value['text_to_revise']}"
|
|
for i in graph.get_state(config).interrupts
|
|
}
|
|
with pytest.raises(
|
|
RuntimeError,
|
|
match="When there are multiple pending interrupts, you must specify the interrupt id when resuming.",
|
|
):
|
|
graph.invoke(Command(resume="singular resume"), config=config)
|
|
|
|
assert graph.invoke(Command(resume=resume_map), config=config) == {
|
|
"text_1": "resume for prompt: original text 1",
|
|
"text_2": "resume for prompt: original text 2",
|
|
}
|
|
|
|
|
|
def test_interrupt_stream_mode_values(sync_checkpointer: BaseCheckpointSaver):
|
|
"""Test that interrupts are surfaced on 'values' stream mode"""
|
|
|
|
class State(TypedDict):
|
|
robot_input: str
|
|
human_input: str
|
|
|
|
def robot_input_node(state: State) -> State:
|
|
return {"robot_input": "beep boop i am a robot"}
|
|
|
|
def human_input_node(state: State) -> Command:
|
|
human_input = interrupt("interrupt")
|
|
return Command(update={"human_input": human_input})
|
|
|
|
builder = StateGraph(State)
|
|
builder.add_node(robot_input_node)
|
|
builder.add_node(human_input_node)
|
|
builder.add_edge(START, "robot_input_node")
|
|
builder.add_edge("robot_input_node", "human_input_node")
|
|
app = builder.compile(checkpointer=sync_checkpointer)
|
|
config = {"configurable": {"thread_id": str(uuid.uuid4())}}
|
|
|
|
result = [*app.stream(State(), config, stream_mode=["updates", "values"])]
|
|
assert len(result) == 4
|
|
assert result == [
|
|
("updates", {"robot_input_node": {"robot_input": "beep boop i am a robot"}}),
|
|
("values", {"robot_input": "beep boop i am a robot"}),
|
|
("updates", {"__interrupt__": (Interrupt(value="interrupt", id=AnyStr()),)}),
|
|
(
|
|
"values",
|
|
{
|
|
"robot_input": "beep boop i am a robot",
|
|
"__interrupt__": (Interrupt(value="interrupt", id=AnyStr()),),
|
|
},
|
|
),
|
|
]
|
|
resume_result = [
|
|
*app.stream(
|
|
Command(resume="i am a human"), config, stream_mode=["updates", "values"]
|
|
)
|
|
]
|
|
assert resume_result == [
|
|
("values", {"robot_input": "beep boop i am a robot"}),
|
|
("updates", {"human_input_node": {"human_input": "i am a human"}}),
|
|
(
|
|
"values",
|
|
{"robot_input": "beep boop i am a robot", "human_input": "i am a human"},
|
|
),
|
|
]
|
|
|
|
|
|
def test_supersteps_populate_task_results(
|
|
sync_checkpointer: BaseCheckpointSaver,
|
|
) -> None:
|
|
class State(TypedDict):
|
|
num: int
|
|
text: str
|
|
|
|
def double(state: State) -> State:
|
|
return {"num": state["num"] * 2, "text": state["text"] * 2}
|
|
|
|
graph = (
|
|
StateGraph(State)
|
|
.add_node("double", double)
|
|
.add_edge(START, "double")
|
|
.add_edge("double", END)
|
|
.compile(checkpointer=sync_checkpointer)
|
|
)
|
|
|
|
def first_task_result(history: list[StateSnapshot], node: str) -> Any:
|
|
for s in history:
|
|
for t in s.tasks:
|
|
if t.name == node:
|
|
return t.result
|
|
return None
|
|
|
|
# reference run with invoke
|
|
ref_cfg = {"configurable": {"thread_id": "ref"}}
|
|
graph.invoke({"num": 1, "text": "one"}, ref_cfg)
|
|
ref_history = list(graph.get_state_history(ref_cfg))
|
|
|
|
ref_start_result = first_task_result(ref_history, "__start__")
|
|
ref_double_result = first_task_result(ref_history, "double")
|
|
assert ref_start_result == {"num": 1, "text": "one"}
|
|
assert ref_double_result == {"num": 2, "text": "oneone"}
|
|
|
|
# using supersteps
|
|
bulk_cfg = {"configurable": {"thread_id": "bulk"}}
|
|
graph.bulk_update_state(
|
|
bulk_cfg,
|
|
[
|
|
[StateUpdate(values={}, as_node="__input__")],
|
|
[StateUpdate(values={"num": 1, "text": "one"}, as_node="__start__")],
|
|
[StateUpdate(values={"num": 2, "text": "oneone"}, as_node="double")],
|
|
],
|
|
)
|
|
bulk_history = list(graph.get_state_history(bulk_cfg))
|
|
|
|
bulk_start_result = first_task_result(bulk_history, "__start__")
|
|
bulk_double_result = first_task_result(bulk_history, "double")
|
|
|
|
assert bulk_start_result == ref_start_result == {"num": 1, "text": "one"}
|
|
assert bulk_double_result == ref_double_result == {"num": 2, "text": "oneone"}
|
|
|
|
|
|
def test_multiple_writes_same_channel_from_same_node(
|
|
sync_checkpointer: BaseCheckpointSaver,
|
|
) -> None:
|
|
"""Test that a node can write multiple times to the same channel and that writes are ordered, reduced, and reflected in streamed events and state history."""
|
|
|
|
class State(TypedDict):
|
|
foo: Annotated[str, lambda a, b: ", ".join([x for x in [a, b] if x])]
|
|
|
|
def one(_: State) -> Command:
|
|
return Command(update=[("foo", "one.0"), ("foo", "one.1")])
|
|
|
|
def two(_: State) -> State:
|
|
return {"foo": "two"}
|
|
|
|
graph = (
|
|
StateGraph(State)
|
|
.add_node("one", one)
|
|
.add_node("two", two)
|
|
.add_edge(START, "one")
|
|
.add_edge("one", "two")
|
|
.add_edge("two", END)
|
|
.compile(checkpointer=sync_checkpointer)
|
|
)
|
|
|
|
config = {"configurable": {"thread_id": "1"}}
|
|
|
|
events = [
|
|
(ns, ev)
|
|
for ns, ev in graph.stream(
|
|
{"foo": "input"}, config, stream_mode=["updates", "tasks"]
|
|
)
|
|
]
|
|
|
|
assert events == [
|
|
(
|
|
"tasks",
|
|
{
|
|
"id": AnyStr(),
|
|
"name": "one",
|
|
"input": {"foo": "input"},
|
|
"triggers": ("branch:to:one",),
|
|
},
|
|
),
|
|
("updates", {"one": [{"foo": "one.0"}, {"foo": "one.1"}]}),
|
|
(
|
|
"tasks",
|
|
{
|
|
"id": AnyStr(),
|
|
"name": "one",
|
|
"error": None,
|
|
"result": {"foo": {"$writes": ["one.0", "one.1"]}},
|
|
"interrupts": [],
|
|
},
|
|
),
|
|
(
|
|
"tasks",
|
|
{
|
|
"id": AnyStr(),
|
|
"name": "two",
|
|
"input": {"foo": "input, one.0, one.1"},
|
|
"triggers": ("branch:to:two",),
|
|
},
|
|
),
|
|
("updates", {"two": {"foo": "two"}}),
|
|
(
|
|
"tasks",
|
|
{
|
|
"id": AnyStr(),
|
|
"name": "two",
|
|
"error": None,
|
|
"result": {"foo": "two"},
|
|
"interrupts": [],
|
|
},
|
|
),
|
|
]
|
|
|
|
def map_snapshot(s: StateSnapshot) -> dict:
|
|
return {
|
|
"tasks": [{"name": t.name, "result": t.result} for t in s.tasks],
|
|
"values": s.values,
|
|
}
|
|
|
|
history = [map_snapshot(s) for s in graph.get_state_history(config)]
|
|
|
|
assert history == [
|
|
{
|
|
"tasks": [],
|
|
"values": {"foo": "input, one.0, one.1, two"},
|
|
},
|
|
{
|
|
"tasks": [{"name": "two", "result": {"foo": "two"}}],
|
|
"values": {"foo": "input, one.0, one.1"},
|
|
},
|
|
{
|
|
"tasks": [
|
|
{"name": "one", "result": {"foo": {"$writes": ["one.0", "one.1"]}}}
|
|
],
|
|
"values": {"foo": "input"},
|
|
},
|
|
{
|
|
"tasks": [{"name": "__start__", "result": {"foo": "input"}}],
|
|
"values": {"foo": ""},
|
|
},
|
|
]
|
|
|
|
|
|
def test_send_with_untracked_value(sync_checkpointer: BaseCheckpointSaver):
|
|
"""Test that Send objects work correctly with untracked values in state."""
|
|
|
|
class UnserializableResource:
|
|
def __init__(self, name: str):
|
|
self.name = name
|
|
self.lock = threading.Lock()
|
|
|
|
class State(TypedDict):
|
|
messages: Annotated[list[str], operator.add]
|
|
session_resource: Annotated[UnserializableResource, UntrackedValue]
|
|
|
|
def setup_node(state: State) -> State:
|
|
resource = UnserializableResource("test_session")
|
|
return {"messages": ["setup complete"], "session_resource": resource}
|
|
|
|
def send_to_tool(state: State):
|
|
return [Send("tool_node", state)]
|
|
|
|
def tool_node(state: State) -> State:
|
|
resource = state["session_resource"]
|
|
assert isinstance(resource, UnserializableResource)
|
|
assert resource.name == "test_session"
|
|
|
|
new_resource = UnserializableResource("new_session")
|
|
|
|
return {
|
|
"messages": [f"tool used resource: {resource.name}"],
|
|
"session_resource": new_resource,
|
|
}
|
|
|
|
graph = StateGraph(State)
|
|
graph.add_node("setup", setup_node)
|
|
graph.add_node("tool_node", tool_node)
|
|
graph.add_edge(START, "setup")
|
|
graph.add_conditional_edges("setup", send_to_tool)
|
|
|
|
app = graph.compile(checkpointer=sync_checkpointer)
|
|
config = {"configurable": {"thread_id": "1"}}
|
|
result = app.invoke({}, config)
|
|
|
|
assert len(result["messages"]) == 2
|
|
assert result["messages"][0] == "setup complete"
|
|
assert result["messages"][1] == "tool used resource: test_session"
|
|
assert result["session_resource"].name == "new_session"
|
|
|
|
state = app.get_state(config)
|
|
assert "session_resource" not in state.values
|
|
|
|
|
|
def test_send_with_untracked_value_overlapping_keys(
|
|
sync_checkpointer: BaseCheckpointSaver,
|
|
):
|
|
"""Test that Send objects work correctly with untracked values in state."""
|
|
|
|
class State(TypedDict):
|
|
dictionary: dict
|
|
session_resource: Annotated[str, UntrackedValue]
|
|
|
|
def setup_node(state: State) -> State:
|
|
return {}
|
|
|
|
def send_to_tool(state: State):
|
|
return [
|
|
Send(
|
|
"tool_node",
|
|
{
|
|
"dictionary": {"session_resource": "legal_value"},
|
|
"session_resource": "illegal_value",
|
|
},
|
|
)
|
|
]
|
|
|
|
def tool_node(state: State) -> State:
|
|
print(f"STATE: {state}")
|
|
assert state["dictionary"] == {"session_resource": "legal_value"}
|
|
assert state["session_resource"] == "illegal_value"
|
|
|
|
return {
|
|
"dictionary": state["dictionary"],
|
|
"session_resource": "new_illegal_value",
|
|
}
|
|
|
|
graph = StateGraph(State)
|
|
graph.add_node("setup", setup_node)
|
|
graph.add_node("tool_node", tool_node)
|
|
graph.add_edge(START, "setup")
|
|
graph.add_conditional_edges("setup", send_to_tool)
|
|
|
|
app = graph.compile(checkpointer=sync_checkpointer)
|
|
config = {"configurable": {"thread_id": "1"}}
|
|
result = app.invoke({}, config)
|
|
|
|
assert result["session_resource"] == "new_illegal_value"
|
|
state = app.get_state(config)
|
|
assert "session_resource" not in state.values
|
|
assert state.values.get("dictionary") == {"session_resource": "legal_value"}
|
|
|
|
|
|
@pytest.mark.parametrize("as_json", [False, True])
|
|
def test_overwrite_sequential(
|
|
sync_checkpointer: BaseCheckpointSaver, as_json: bool
|
|
) -> None:
|
|
"""Test a sequential chain of nodes where the last node uses Overwrite to bypass a reducer and write a value directly to the channel."""
|
|
|
|
class State(TypedDict):
|
|
messages: Annotated[list, operator.add]
|
|
|
|
def node_a(state: State):
|
|
return {"messages": ["a"]}
|
|
|
|
def node_b(state: State):
|
|
overwrite = {"__overwrite__": ["b"]} if as_json else Overwrite(["b"])
|
|
return {"messages": overwrite}
|
|
|
|
builder = StateGraph(State)
|
|
builder.add_node("node_a", node_a)
|
|
builder.add_node("node_b", node_b)
|
|
builder.add_edge(START, "node_a")
|
|
builder.add_edge("node_a", "node_b")
|
|
|
|
graph = builder.compile(checkpointer=sync_checkpointer)
|
|
config = {"configurable": {"thread_id": "1"}}
|
|
result = graph.invoke({"messages": ["START"]}, config)
|
|
# a is overwritten by b
|
|
assert result == {"messages": ["b"]}
|
|
|
|
|
|
@pytest.mark.parametrize("as_json", [False, True])
|
|
def test_overwrite_parallel(
|
|
sync_checkpointer: BaseCheckpointSaver, as_json: bool
|
|
) -> None:
|
|
"""Test parallel nodes where max one node uses Overwrite to bypass a reducer and write a value directly to the channel."""
|
|
|
|
class State(TypedDict):
|
|
messages: Annotated[list, operator.add]
|
|
|
|
def node_a(state: State):
|
|
return {"messages": ["a"]}
|
|
|
|
def node_b(state: State):
|
|
overwrite = {"__overwrite__": ["b"]} if as_json else Overwrite(["b"])
|
|
return {"messages": overwrite}
|
|
|
|
def node_c(state: State):
|
|
return {"messages": ["c"]}
|
|
|
|
def node_d(state: State):
|
|
return {"messages": ["d"]}
|
|
|
|
builder = StateGraph(State)
|
|
builder.add_node("node_a", node_a)
|
|
builder.add_node("node_b", node_b)
|
|
builder.add_node("node_c", node_c)
|
|
builder.add_node("node_d", node_d)
|
|
builder.add_edge(START, "node_a")
|
|
builder.add_edge("node_a", "node_b")
|
|
builder.add_edge("node_a", "node_c")
|
|
builder.add_edge("node_b", "node_d")
|
|
builder.add_edge("node_c", "node_d")
|
|
|
|
graph = builder.compile(checkpointer=sync_checkpointer)
|
|
config = {"configurable": {"thread_id": "1"}}
|
|
result = graph.invoke({"messages": ["START"]}, config)
|
|
# a, c are overwritten by b, then d is written
|
|
assert result == {"messages": ["b", "d"]}
|
|
|
|
|
|
@pytest.mark.parametrize("as_json", [False, True])
|
|
def test_overwrite_parallel_error(
|
|
sync_checkpointer: BaseCheckpointSaver, as_json: bool
|
|
) -> None:
|
|
"""Test parallel nodes where more than one node uses Overwrite to bypass a reducer and write a value directly to the channel. In this case, InvalidUpdateError should be raised."""
|
|
|
|
class State(TypedDict):
|
|
messages: Annotated[list, operator.add]
|
|
|
|
def node_a(state: State):
|
|
return {"messages": ["a"]}
|
|
|
|
def node_b(state: State):
|
|
overwrite = {"__overwrite__": ["b"]} if as_json else Overwrite(["b"])
|
|
return {"messages": overwrite}
|
|
|
|
def node_c(state: State):
|
|
overwrite = {"__overwrite__": ["c"]} if as_json else Overwrite(["c"])
|
|
return {"messages": overwrite}
|
|
|
|
builder = StateGraph(State)
|
|
builder.add_node("node_a", node_a)
|
|
builder.add_node("node_b", node_b)
|
|
builder.add_node("node_c", node_c)
|
|
builder.add_edge(START, "node_a")
|
|
builder.add_edge("node_a", "node_b")
|
|
builder.add_edge("node_a", "node_c")
|
|
builder.add_edge("node_b", END)
|
|
builder.add_edge("node_c", END)
|
|
|
|
graph = builder.compile(checkpointer=sync_checkpointer)
|
|
config = {"configurable": {"thread_id": "1"}}
|
|
with pytest.raises(
|
|
InvalidUpdateError, match="Can receive only one Overwrite value per super-step."
|
|
):
|
|
graph.invoke({"messages": ["START"]}, config)
|
|
|
|
|
|
def test_fork_does_not_apply_pending_writes(
|
|
sync_checkpointer: BaseCheckpointSaver,
|
|
) -> None:
|
|
"""Test that forking with update_state does not apply pending writes from original execution."""
|
|
|
|
class State(TypedDict):
|
|
value: Annotated[int, operator.add]
|
|
|
|
def node_a(state: State) -> State:
|
|
return {"value": 10}
|
|
|
|
def node_b(state: State) -> State:
|
|
return {"value": 100}
|
|
|
|
graph = (
|
|
StateGraph(State)
|
|
.add_node("node_a", node_a)
|
|
.add_node("node_b", node_b)
|
|
.add_edge(START, "node_a")
|
|
.add_edge("node_a", "node_b")
|
|
.compile(checkpointer=sync_checkpointer)
|
|
)
|
|
|
|
thread1 = {"configurable": {"thread_id": "1"}}
|
|
graph.invoke({"value": 1}, thread1)
|
|
|
|
history = list(graph.get_state_history(thread1))
|
|
checkpoint_before_a = next(s for s in history if s.next == ("node_a",))
|
|
|
|
fork_config = graph.update_state(
|
|
checkpoint_before_a.config, {"value": 20}, as_node="node_a"
|
|
)
|
|
|
|
# Continue from fork (should run node_b)
|
|
result = graph.invoke(None, fork_config)
|
|
|
|
# Should be: 1 (input) + 20 (forked node_a) + 100 (node_b) = 121
|
|
assert result == {"value": 121}
|
|
|
|
|
|
async def test_delta_channel_end_to_end_inmemory() -> None:
|
|
"""Full graph run: DeltaChannel accumulates correctly across multiple turns."""
|
|
|
|
class State(TypedDict):
|
|
messages: Annotated[list, DeltaChannel(_messages_delta_reducer)]
|
|
|
|
def respond(state: State) -> dict:
|
|
n = len(state["messages"])
|
|
return {"messages": [AIMessage(content=f"reply-{n}", id=f"ai-{n}")]}
|
|
|
|
builder = StateGraph(State)
|
|
builder.add_node("respond", respond)
|
|
builder.add_edge(START, "respond")
|
|
graph = builder.compile(checkpointer=InMemorySaver())
|
|
|
|
config = {"configurable": {"thread_id": "diff-test-1"}}
|
|
|
|
# Turn 1
|
|
graph.invoke({"messages": [HumanMessage(content="hello", id="h1")]}, config)
|
|
# Turn 2
|
|
graph.invoke({"messages": [HumanMessage(content="world", id="h2")]}, config)
|
|
# Turn 3
|
|
graph.invoke({"messages": [HumanMessage(content="bye", id="h3")]}, config)
|
|
|
|
state = graph.get_state(config)
|
|
msgs = state.values["messages"]
|
|
# 3 human + 3 AI = 6 total
|
|
assert len(msgs) == 6, f"expected 6 messages, got {len(msgs)}: {msgs}"
|
|
assert msgs[0].content == "hello"
|
|
assert msgs[2].content == "world"
|
|
assert msgs[4].content == "bye"
|
|
assert msgs[1].content == "reply-1"
|
|
assert msgs[3].content == "reply-3"
|
|
assert msgs[5].content == "reply-5"
|
|
|
|
|
|
async def test_delta_channel_time_travel() -> None:
|
|
"""Time-travel back to turn-1 checkpoint and resume; continuation must not include turn-2 deltas."""
|
|
|
|
class State(TypedDict):
|
|
messages: Annotated[list, DeltaChannel(_messages_delta_reducer)]
|
|
|
|
counter = {"n": 0}
|
|
|
|
def respond(state: State) -> dict:
|
|
counter["n"] += 1
|
|
return {
|
|
"messages": [
|
|
AIMessage(content=f"ai-{counter['n']}", id=f"ai-{counter['n']}")
|
|
]
|
|
}
|
|
|
|
builder = StateGraph(State)
|
|
builder.add_node("respond", respond)
|
|
builder.add_edge(START, "respond")
|
|
saver = InMemorySaver()
|
|
graph = builder.compile(checkpointer=saver)
|
|
|
|
config = {"configurable": {"thread_id": "diff-time-travel"}}
|
|
|
|
# Run 2 turns: h1→ai-1, h2→ai-2
|
|
graph.invoke({"messages": [HumanMessage(content="h1", id="h1")]}, config)
|
|
graph.invoke({"messages": [HumanMessage(content="h2", id="h2")]}, config)
|
|
|
|
# Find the checkpoint after turn 1 (2 messages: h1 + ai-1)
|
|
history = list(graph.get_state_history(config))
|
|
after_turn1 = next(h for h in history if len(h.values.get("messages", [])) == 2)
|
|
|
|
assert len(after_turn1.values["messages"]) == 2
|
|
assert after_turn1.values["messages"][0].content == "h1"
|
|
assert after_turn1.values["messages"][1].content == "ai-1"
|
|
|
|
# Resume from turn-1 checkpoint: inject h3, expect 3 messages total (h1, ai-1, ai-N)
|
|
# NOT 5 messages (turn-2 deltas must not bleed into the resumed run)
|
|
result = graph.invoke(
|
|
{"messages": [HumanMessage(content="h3", id="h3")]},
|
|
after_turn1.config,
|
|
)
|
|
msgs = result["messages"]
|
|
# Should be: h1, ai-1, h3, ai-N — 4 messages total
|
|
assert len(msgs) == 4, (
|
|
f"expected 4 messages after time-travel resume, got {len(msgs)}: {msgs}"
|
|
)
|
|
assert msgs[0].content == "h1"
|
|
assert msgs[1].content == "ai-1"
|
|
assert msgs[2].content == "h3"
|
|
|
|
|
|
async def test_delta_channel_remove_message_end_to_end() -> None:
|
|
"""RemoveMessage inside a DeltaChannel graph must persist and reload correctly."""
|
|
|
|
class State(TypedDict):
|
|
messages: Annotated[list, DeltaChannel(_messages_delta_reducer)]
|
|
|
|
def respond(state: State) -> dict:
|
|
return {"messages": [AIMessage(content="reply", id="ai-1")]}
|
|
|
|
def delete_first(state: State) -> dict:
|
|
# removes the first message
|
|
return {"messages": [RemoveMessage(id=state["messages"][0].id)]}
|
|
|
|
builder = StateGraph(State)
|
|
builder.add_node("respond", respond)
|
|
builder.add_node("delete_first", delete_first)
|
|
builder.add_edge(START, "respond")
|
|
builder.add_edge("respond", "delete_first")
|
|
graph = builder.compile(checkpointer=InMemorySaver())
|
|
|
|
config = {"configurable": {"thread_id": "diff-remove-test"}}
|
|
graph.invoke({"messages": [HumanMessage(content="hello", id="h1")]}, config)
|
|
|
|
state = graph.get_state(config)
|
|
msgs = state.values["messages"]
|
|
# h1 was removed, only ai-1 should remain
|
|
assert len(msgs) == 1, f"expected 1 message, got {len(msgs)}: {msgs}"
|
|
assert msgs[0].id == "ai-1"
|
|
|
|
# A subsequent turn must reconstruct from the checkpoint correctly
|
|
graph.invoke({"messages": [HumanMessage(content="again", id="h2")]}, config)
|
|
state = graph.get_state(config)
|
|
msgs = state.values["messages"]
|
|
# ai-1 + h2 + ai-1(second reply, same id overwrites) + h2 removed
|
|
# more simply: after second run we expect ai-1 updated + h2 remaining minus deleted h2
|
|
# just assert h1 is still gone
|
|
assert all(m.id != "h1" for m in msgs), (
|
|
"h1 should still be absent after second turn"
|
|
)
|
|
|
|
|
|
async def test_delta_channel_update_by_id_end_to_end() -> None:
|
|
"""Updating a message by ID via DeltaChannel must persist and reload correctly."""
|
|
|
|
class State(TypedDict):
|
|
messages: Annotated[list, DeltaChannel(_messages_delta_reducer)]
|
|
|
|
def update_msg(state: State) -> dict:
|
|
# re-send h1 with updated content
|
|
return {"messages": [HumanMessage(content="updated", id="h1")]}
|
|
|
|
builder = StateGraph(State)
|
|
builder.add_node("update_msg", update_msg)
|
|
builder.add_edge(START, "update_msg")
|
|
graph = builder.compile(checkpointer=InMemorySaver())
|
|
|
|
config = {"configurable": {"thread_id": "diff-update-id-test"}}
|
|
graph.invoke({"messages": [HumanMessage(content="original", id="h1")]}, config)
|
|
|
|
state = graph.get_state(config)
|
|
msgs = state.values["messages"]
|
|
assert len(msgs) == 1, f"expected 1 message, got {len(msgs)}: {msgs}"
|
|
assert msgs[0].content == "updated"
|
|
assert msgs[0].id == "h1"
|
|
|
|
# Second turn: verify the updated state is the base for further accumulation
|
|
graph.invoke({"messages": [HumanMessage(content="new", id="h2")]}, config)
|
|
state = graph.get_state(config)
|
|
msgs = state.values["messages"]
|
|
ids = [m.id for m in msgs]
|
|
assert "h1" in ids # h1 persists (updated, not duplicated)
|
|
assert "h2" in ids
|
|
assert ids.count("h1") == 1, "h1 must not be duplicated"
|
|
|
|
|
|
async def test_delta_channel_durability_exit_stores_snapshot() -> None:
|
|
"""DeltaChannel must reload from a durability='exit' checkpoint."""
|
|
|
|
class State(TypedDict):
|
|
messages: Annotated[list, DeltaChannel(_messages_delta_reducer)]
|
|
|
|
def respond(state: State) -> dict:
|
|
return {"messages": [AIMessage(content="reply", id="ai1")]}
|
|
|
|
builder = StateGraph(State)
|
|
builder.add_node("respond", respond)
|
|
builder.add_edge(START, "respond")
|
|
graph = builder.compile(checkpointer=InMemorySaver())
|
|
config = {"configurable": {"thread_id": "delta-exit-test"}}
|
|
|
|
result = graph.invoke(
|
|
{"messages": [HumanMessage(content="hello", id="h1")]},
|
|
config,
|
|
durability="exit",
|
|
)
|
|
assert [m.content for m in result["messages"]] == ["hello", "reply"]
|
|
|
|
state = graph.get_state(config)
|
|
assert [m.content for m in state.values["messages"]] == ["hello", "reply"]
|
|
|
|
|
|
async def test_delta_channel_async_write_ordering() -> None:
|
|
"""In async mode, DeltaChannel write futures are awaited before the checkpoint
|
|
is committed, so aput_writes always precedes aput for delta-channel
|
|
checkpoints (those where the delta channel had a versioned write but
|
|
is absent from `channel_values`, i.e. no snapshot fired this step)."""
|
|
|
|
class State(TypedDict):
|
|
messages: Annotated[list, DeltaChannel(_messages_delta_reducer)]
|
|
|
|
def respond(state: State) -> dict:
|
|
i = len(state["messages"])
|
|
return {"messages": [AIMessage(content=f"r{i}", id=f"ai{i}")]}
|
|
|
|
order: list[str] = []
|
|
original_aput_writes = InMemorySaver.aput_writes
|
|
original_aput = InMemorySaver.aput
|
|
|
|
async def tracked_aput_writes(self, config, writes, task_id, task_path=""):
|
|
result = await original_aput_writes(self, config, writes, task_id, task_path)
|
|
order.append("aput_writes")
|
|
return result
|
|
|
|
async def tracked_aput(self, config, checkpoint, metadata, new_versions):
|
|
# A "delta" checkpoint here = `messages` versioned but absent from
|
|
# `channel_values` (no snapshot fired). When a snapshot does fire,
|
|
# `channel_values["messages"]` is a `_DeltaSnapshot` — also a delta
|
|
# checkpoint shape, since the writes still have to be persisted
|
|
# before the parent checkpoint commits.
|
|
channel_values = checkpoint.get("channel_values", {})
|
|
is_delta_step = (
|
|
"messages" in checkpoint.get("channel_versions", {})
|
|
and "messages" not in channel_values
|
|
)
|
|
order.append("aput_delta" if is_delta_step else "aput_other")
|
|
return await original_aput(self, config, checkpoint, metadata, new_versions)
|
|
|
|
InMemorySaver.aput_writes = tracked_aput_writes
|
|
InMemorySaver.aput = tracked_aput
|
|
try:
|
|
builder = StateGraph(State)
|
|
builder.add_node("respond", respond)
|
|
builder.add_edge(START, "respond")
|
|
graph = builder.compile(checkpointer=InMemorySaver())
|
|
config = {"configurable": {"thread_id": "async-ordering-test"}}
|
|
|
|
for i in range(3):
|
|
await graph.ainvoke(
|
|
{"messages": [HumanMessage(content=f"h{i}", id=f"h{i}")]}, config
|
|
)
|
|
|
|
# Every aput_delta must be preceded by at least one aput_writes
|
|
for i, event in enumerate(order):
|
|
if event == "aput_delta":
|
|
preceding = order[:i]
|
|
assert "aput_writes" in preceding, (
|
|
f"aput_delta at {i} had no preceding aput_writes: {order}"
|
|
)
|
|
last_write_idx = max(
|
|
j for j, e in enumerate(order[:i]) if e == "aput_writes"
|
|
)
|
|
assert last_write_idx < i, (
|
|
f"aput_writes at {last_write_idx} should precede aput_delta at {i}: {order}"
|
|
)
|
|
finally:
|
|
InMemorySaver.aput_writes = original_aput_writes
|
|
InMemorySaver.aput = original_aput
|
|
|
|
state = await graph.aget_state(config)
|
|
assert len(state.values["messages"]) == 6 # 3 human + 3 AI
|