WEBDIENSTE/langgraph
1
0
Fork 0
mirror of https://github.com/langchain-ai/langgraph.git synced 2026-09-08 02:37:52 +02:00
Code Issues Packages Projects Releases Wiki Activity

Merge branch 'main' into vb/add-graph-command-docs

Browse Source
This commit is contained in:
vbarda
2024-12-04 18:37:20 -05:00
parent d52bb911a4 dcc2617396
commit 797b919cf9
18 changed files with 553 additions and 192 deletions
Download Patch File Download Diff File Expand all files Collapse all files
docs/docs/concepts/persistence.md
+4 -1
View File
@@ -276,9 +276,11 @@ The attributes it has are:
Beyond simple retrieval, the store also supports semantic search, allowing you to find memories based on meaning rather than exact matches. To enable this, configure the store with an embedding model:
```python
from langchain.embeddings import init_embeddings
store = InMemoryStore(
index={
"embed": "openai:text-embedding-3-small", # Embedding provider
"embed": init_embeddings("openai:text-embedding-3-small"), # Embedding provider
"dims": 1536, # Embedding dimensions
"fields": ["food_preference", "$"] # Fields to embed
}
@@ -289,6 +291,7 @@ Now when searching, you can use natural language queries to find relevant memori
```python
# Find memories about food preferences
# (This can be done after putting memories into the store)
memories = store.search(
namespace_for_memory,
query="What does the user like to eat?",
docs/docs/how-tos/memory/semantic-search.ipynb
-1
View File
@@ -297,7 +297,6 @@
}
],
"source": [
"embeddings = init_embeddings(\"openai:text-embedding-3-small\")\n",
"store = InMemoryStore(\n",
" index={\n",
" \"embed\": embeddings,\n",
libs/langgraph/langgraph/constants.py
+4 -2
View File
@@ -72,9 +72,11 @@ CONFIG_KEY_CHECKPOINT_ID = sys.intern("checkpoint_id")
CONFIG_KEY_CHECKPOINT_NS = sys.intern("checkpoint_ns")
# holds the current checkpoint_ns, "" for root graph
CONFIG_KEY_NODE_FINISHED = sys.intern("__pregel_node_finished")
# callback to be called when a node is finished
CONFIG_KEY_RESUME_VALUE = sys.intern("__pregel_resume_value")
# holds the value that "answers" an interrupt() call
CONFIG_KEY_WRITES = sys.intern("__pregel_writes")
# read-only list of existing task writes
CONFIG_KEY_SCRATCHPAD = sys.intern("__pregel_scratchpad")
# holds a mutable dict for temporary storage scoped to the current task
# --- Other constants ---
PUSH = sys.intern("__pregel_push")
libs/langgraph/langgraph/graph/__init__.py
+1 -2
View File
@@ -1,13 +1,12 @@
from langgraph.graph.graph import END, START, Graph
from langgraph.graph.message import MessageGraph, MessagesState, add_messages
from langgraph.graph.state import GraphCommand, StateGraph
from langgraph.graph.state import StateGraph
__all__ = [
"END",
"START",
"Graph",
"StateGraph",
"GraphCommand",
"MessageGraph",
"add_messages",
"MessagesState",
libs/langgraph/langgraph/graph/state.py
+12 -48
View File
@@ -1,4 +1,3 @@
import dataclasses
import inspect
import logging
import typing
@@ -9,7 +8,6 @@ from types import FunctionType
from typing import (
Any,
Callable,
Generic,
Literal,
NamedTuple,
Optional,
@@ -55,7 +53,7 @@ from langgraph.managed.base import (
from langgraph.pregel.read import ChannelRead, PregelNode
from langgraph.pregel.write import SKIP_WRITE, ChannelWrite, ChannelWriteEntry
from langgraph.store.base import BaseStore
from langgraph.types import _DC_KWARGS, All, Checkpointer, Command, N, RetryPolicy
from langgraph.types import All, Checkpointer, Command, RetryPolicy
from langgraph.utils.fields import get_field_default
from langgraph.utils.pydantic import create_model
from langgraph.utils.runnable import RunnableCallable, coerce_to_runnable
@@ -84,34 +82,6 @@ def _get_node_name(node: RunnableLike) -> str:
raise TypeError(f"Unsupported node type: {type(node)}")
@dataclasses.dataclass(**_DC_KWARGS)
class GraphCommand(Generic[N], Command[N]):
"""One or more commands to update a StateGraph's state and go to, or send messages to nodes.
Args:
graph: graph to send the command to. Supported values are:
- None: the current graph (default)
- GraphCommand.PARENT: closest parent graph
update: state update to apply to the graph's state at the current superstep.
send: list of `Send` objects to send to other nodes.
resume: value to resume execution with. Will be used when `interrupt()` is called.
goto: name of the node to navigate to next.
Can be any node that belongs to the specified `graph` (current or parent).
If `goto` not specified, the graph will halt after executing the current superstep.
"""
goto: Union[str, Sequence[str]] = ()
def __repr__(self) -> str:
# get all non-None values
contents = ", ".join(
f"{key}={value!r}"
for key, value in dataclasses.asdict(self).items()
if value
)
return f"Command({contents})"
class StateNodeSpec(NamedTuple):
runnable: Runnable
metadata: Optional[dict[str, Any]]
@@ -404,7 +374,7 @@ class StateGraph(Graph):
input = input_hint
if (
(rtn := hints.get("return"))
and get_origin(rtn) in (Command, GraphCommand)
and get_origin(rtn) is Command
and (rargs := get_args(rtn))
and get_origin(rargs[0]) is Literal
and (vals := get_args(rargs[0]))
@@ -846,15 +816,12 @@ def _control_branch(value: Any) -> Sequence[Union[str, Send]]:
if value.graph == Command.PARENT:
raise ParentCommand(value)
rtn: list[Union[str, Send]] = []
if isinstance(value, GraphCommand):
if isinstance(value.goto, str):
rtn.append(value.goto)
else:
rtn.extend(value.goto)
if isinstance(value.send, Send):
rtn.append(value.send)
if isinstance(value.goto, Send):
rtn.append(value.goto)
elif isinstance(value.goto, str):
rtn.append(value.goto)
else:
rtn.extend(value.send)
rtn.extend(value.goto)
return rtn
@@ -866,15 +833,12 @@ async def _acontrol_branch(value: Any) -> Sequence[Union[str, Send]]:
if value.graph == Command.PARENT:
raise ParentCommand(value)
rtn: list[Union[str, Send]] = []
if isinstance(value, GraphCommand):
if isinstance(value.goto, str):
rtn.append(value.goto)
else:
rtn.extend(value.goto)
if isinstance(value.send, Send):
rtn.append(value.send)
if isinstance(value.goto, Send):
rtn.append(value.goto)
elif isinstance(value.goto, str):
rtn.append(value.goto)
else:
rtn.extend(value.send)
rtn.extend(value.goto)
return rtn
libs/langgraph/langgraph/pregel/algo.py
+16 -19
View File
@@ -37,13 +37,13 @@ from langgraph.constants import (
CONFIG_KEY_CHECKPOINT_NS,
CONFIG_KEY_CHECKPOINTER,
CONFIG_KEY_READ,
CONFIG_KEY_RESUME_VALUE,
CONFIG_KEY_SCRATCHPAD,
CONFIG_KEY_SEND,
CONFIG_KEY_STORE,
CONFIG_KEY_TASK_ID,
CONFIG_KEY_WRITES,
EMPTY_SEQ,
INTERRUPT,
MISSING,
NO_WRITES,
NS_END,
NS_SEP,
@@ -589,14 +589,13 @@ def prepare_single_task(
},
CONFIG_KEY_CHECKPOINT_ID: None,
CONFIG_KEY_CHECKPOINT_NS: task_checkpoint_ns,
CONFIG_KEY_RESUME_VALUE: next(
(
v
for tid, c, v in pending_writes
if tid in (NULL_TASK_ID, task_id) and c == RESUME
),
configurable.get(CONFIG_KEY_RESUME_VALUE, MISSING),
),
CONFIG_KEY_WRITES: [
w
for w in pending_writes
+ configurable.get(CONFIG_KEY_WRITES, [])
if w[0] in (NULL_TASK_ID, task_id)
],
CONFIG_KEY_SCRATCHPAD: {},
},
),
triggers,
@@ -713,15 +712,13 @@ def prepare_single_task(
},
CONFIG_KEY_CHECKPOINT_ID: None,
CONFIG_KEY_CHECKPOINT_NS: task_checkpoint_ns,
CONFIG_KEY_RESUME_VALUE: next(
(
v
for tid, c, v in pending_writes
if tid in (NULL_TASK_ID, task_id)
and c == RESUME
),
configurable.get(CONFIG_KEY_RESUME_VALUE, MISSING),
),
CONFIG_KEY_WRITES: [
w
for w in pending_writes
+ configurable.get(CONFIG_KEY_WRITES, [])
if w[0] in (NULL_TASK_ID, task_id)
],
CONFIG_KEY_SCRATCHPAD: {},
},
),
triggers,
libs/langgraph/langgraph/pregel/io.py
+13 -7
View File
@@ -4,6 +4,7 @@ from uuid import UUID
from langchain_core.runnables.utils import AddableDict
from langgraph.channels.base import BaseChannel, EmptyChannelError
from langgraph.checkpoint.base import PendingWrite
from langgraph.constants import (
EMPTY_SEQ,
ERROR,
@@ -66,26 +67,31 @@ def read_channels(
def map_command(
cmd: Command,
cmd: Command, pending_writes: list[PendingWrite]
) -> Iterator[tuple[str, str, Any]]:
"""Map input chunk to a sequence of pending writes in the form (channel, value)."""
if cmd.graph == Command.PARENT:
raise InvalidUpdateError("There is not parent graph")
if cmd.send:
if isinstance(cmd.send, (tuple, list)):
sends = cmd.send
if cmd.goto:
if isinstance(cmd.goto, (tuple, list)):
sends = cmd.goto
else:
sends = [cmd.send]
sends = [cmd.goto]
for send in sends:
if not isinstance(send, Send):
raise TypeError(
f"In Command.send, expected Send, got {type(send).__name__}"
f"In Command.goto, expected Send, got {type(send).__name__}"
)
yield (NULL_TASK_ID, PUSH if FF_SEND_V2 else TASKS, send)
# TODO handle goto str for state graph
if cmd.resume:
if isinstance(cmd.resume, dict) and all(is_task_id(k) for k in cmd.resume):
for tid, resume in cmd.resume.items():
yield (tid, RESUME, resume)
existing: list[Any] = next(
(w[2] for w in pending_writes if w[0] == tid and w[1] == RESUME), []
)
existing.append(resume)
yield (tid, RESUME, existing)
else:
yield (NULL_TASK_ID, RESUME, cmd.resume)
if cmd.update:
libs/langgraph/langgraph/pregel/loop.py
+23 -2
View File
@@ -26,6 +26,7 @@ from typing_extensions import ParamSpec, Self
from langgraph.channels.base import BaseChannel
from langgraph.checkpoint.base import (
WRITES_IDX_MAP,
BaseCheckpointSaver,
ChannelVersions,
Checkpoint,
@@ -263,8 +264,28 @@ class PregelLoop(LoopProtocol):
"""Put writes for a task, to be read by the next tick."""
if not writes:
return
# deduplicate writes to special channels, last write wins
if all(w[0] in WRITES_IDX_MAP for w in writes):
writes = list({w[0]: w for w in writes}.values())
# save writes
self.checkpoint_pending_writes.extend((task_id, k, v) for k, v in writes)
for c, v in writes:
if (
c in WRITES_IDX_MAP
and (
idx := next(
(
i
for i, w in enumerate(self.checkpoint_pending_writes)
if w[0] == task_id and w[1] == c
),
None,
)
)
is not None
):
self.checkpoint_pending_writes[idx] = (task_id, c, v)
else:
self.checkpoint_pending_writes.append((task_id, c, v))
if self.checkpointer_put_writes is not None:
self.submit(
self.checkpointer_put_writes,
@@ -536,7 +557,7 @@ class PregelLoop(LoopProtocol):
elif isinstance(self.input, Command):
writes: defaultdict[str, list[tuple[str, Any]]] = defaultdict(list)
# group writes by task ID
for tid, c, v in map_command(self.input):
for tid, c, v in map_command(self.input, self.checkpoint_pending_writes):
writes[tid].append((c, v))
if not writes:
raise EmptyInputError("Received empty Command input")
libs/langgraph/langgraph/pregel/runner.py
+3
View File
@@ -21,6 +21,7 @@ from langgraph.constants import (
INTERRUPT,
NO_WRITES,
PUSH,
RESUME,
TAG_HIDDEN,
)
from langgraph.errors import GraphBubbleUp, GraphInterrupt
@@ -297,6 +298,8 @@ class PregelRunner:
if isinstance(exception, GraphInterrupt):
# save interrupt to checkpointer
if interrupts := [(INTERRUPT, i) for i in exception.args[0]]:
if resumes := [w for w in task.writes if w[0] == RESUME]:
interrupts.extend(resumes)
self.put_writes(task.id, interrupts)
elif isinstance(exception, GraphBubbleUp):
raise exception
libs/langgraph/langgraph/types.py
+54 -14
View File
@@ -13,6 +13,7 @@ from typing import (
Optional,
Sequence,
Type,
TypedDict,
TypeVar,
Union,
cast,
@@ -21,11 +22,16 @@ from typing import (
from langchain_core.runnables import Runnable, RunnableConfig
from typing_extensions import Self
from langgraph.checkpoint.base import BaseCheckpointSaver, CheckpointMetadata
from langgraph.checkpoint.base import (
BaseCheckpointSaver,
CheckpointMetadata,
PendingWrite,
)
if TYPE_CHECKING:
from langgraph.store.base import BaseStore
All = Literal["*"]
"""Special value to indicate that graph should interrupt on all nodes."""
@@ -252,8 +258,8 @@ class Command(Generic[N]):
graph: Optional[str] = None
update: Optional[dict[str, Any]] = None
send: Union[Send, Sequence[Send]] = ()
resume: Optional[Union[Any, dict[str, Any]]] = None
goto: Union[Send, Sequence[Union[Send, str]], str] = ()
def __repr__(self) -> str:
# get all non-None values
@@ -309,26 +315,60 @@ class LoopProtocol:
self.stop = stop
class PregelScratchpad(TypedDict, total=False):
interrupt_counter: int
used_null_resume: bool
resume: list[Any]
def interrupt(value: Any) -> Any:
from langgraph.constants import (
CONFIG_KEY_CHECKPOINT_NS,
CONFIG_KEY_RESUME_VALUE,
MISSING,
CONFIG_KEY_SCRATCHPAD,
CONFIG_KEY_SEND,
CONFIG_KEY_TASK_ID,
CONFIG_KEY_WRITES,
NS_SEP,
NULL_TASK_ID,
RESUME,
)
from langgraph.errors import GraphInterrupt
from langgraph.utils.config import get_configurable
conf = get_configurable()
if (resume := conf.get(CONFIG_KEY_RESUME_VALUE, MISSING)) and resume is not MISSING:
return resume
# track interrupt index
scratchpad: PregelScratchpad = conf[CONFIG_KEY_SCRATCHPAD]
if "interrupt_counter" not in scratchpad:
scratchpad["interrupt_counter"] = 0
else:
raise GraphInterrupt(
(
Interrupt(
value=value,
resumable=True,
ns=cast(str, conf[CONFIG_KEY_CHECKPOINT_NS]).split(NS_SEP),
),
)
scratchpad["interrupt_counter"] += 1
idx = scratchpad["interrupt_counter"]
# find previous resume values
task_id = conf[CONFIG_KEY_TASK_ID]
writes: list[PendingWrite] = conf[CONFIG_KEY_WRITES]
scratchpad.setdefault(
"resume", next((w[2] for w in writes if w[0] == task_id and w[1] == RESUME), [])
)
if scratchpad["resume"]:
if idx < len(scratchpad["resume"]):
return scratchpad["resume"][idx]
# find current resume value
if not scratchpad.get("used_null_resume"):
scratchpad["used_null_resume"] = True
for tid, c, v in sorted(writes, key=lambda x: x[0], reverse=True):
if tid == NULL_TASK_ID and c == RESUME:
assert len(scratchpad["resume"]) == idx, (scratchpad["resume"], idx)
scratchpad["resume"].append(v)
print("saving:", scratchpad["resume"])
conf[CONFIG_KEY_SEND]([(RESUME, scratchpad["resume"])])
return v
# no resume value found
raise GraphInterrupt(
(
Interrupt(
value=value,
resumable=True,
ns=cast(str, conf[CONFIG_KEY_CHECKPOINT_NS]).split(NS_SEP),
),
)
)
libs/langgraph/tests/test_pregel.py
+169 -29
View File
@@ -65,7 +65,7 @@ from langgraph.constants import (
START,
)
from langgraph.errors import InvalidUpdateError, MultipleSubgraphsError, NodeInterrupt
from langgraph.graph import END, Graph, GraphCommand, StateGraph
from langgraph.graph import END, Graph, StateGraph
from langgraph.graph.message import MessageGraph, MessagesState, add_messages
from langgraph.managed.shared_value import SharedValue
from langgraph.prebuilt.chat_agent_executor import create_tool_calling_executor
@@ -270,10 +270,10 @@ def test_graph_validation_with_command() -> None:
bar: str
def node_a(state: State):
return GraphCommand(goto="b", update={"foo": "bar"})
return Command(goto="b", update={"foo": "bar"})
def node_b(state: State):
return GraphCommand(goto=END, update={"bar": "baz"})
return Command(goto=END, update={"bar": "baz"})
builder = StateGraph(State)
builder.add_node("a", node_a)
@@ -1925,8 +1925,8 @@ def test_send_sequences() -> None:
def send_for_fun(state):
return [
Send("2", Command(send=Send("2", 3))),
Send("2", GraphCommand(send=Send("2", 4))),
Send("2", Command(goto=Send("2", 3))),
Send("2", Command(goto=Send("2", 4))),
"3.1",
]
@@ -1947,8 +1947,8 @@ def test_send_sequences() -> None:
== [
"0",
"1",
"2|Command(send=Send(node='2', arg=3))",
"2|Command(send=Send(node='2', arg=4))",
"2|Command(goto=Send(node='2', arg=3))",
"2|Command(goto=Send(node='2', arg=4))",
"2|3",
"2|4",
"3",
@@ -1959,8 +1959,8 @@ def test_send_sequences() -> None:
"0",
"1",
"3.1",
"2|Command(send=Send(node='2', arg=3))",
"2|Command(send=Send(node='2', arg=4))",
"2|Command(goto=Send(node='2', arg=3))",
"2|Command(goto=Send(node='2', arg=4))",
"3",
"2|3",
"2|4",
@@ -2000,15 +2000,15 @@ def test_send_dedupe_on_resume(
if isinstance(state, list)
else ["|".join((self.name, str(state)))]
)
if isinstance(state, GraphCommand):
if isinstance(state, Command):
return replace(state, update=update)
else:
return update
def send_for_fun(state):
return [
Send("2", GraphCommand(send=Send("2", 3))),
Send("2", GraphCommand(send=Send("flaky", 4))),
Send("2", Command(goto=Send("2", 3))),
Send("2", Command(goto=Send("flaky", 4))),
"3.1",
]
@@ -2030,8 +2030,8 @@ def test_send_dedupe_on_resume(
assert graph.invoke(["0"], thread1, debug=1) == [
"0",
"1",
"2|Command(send=Send(node='2', arg=3))",
"2|Command(send=Send(node='flaky', arg=4))",
"2|Command(goto=Send(node='2', arg=3))",
"2|Command(goto=Send(node='flaky', arg=4))",
"2|3",
]
assert builder.nodes["2"].runnable.func.ticks == 3
@@ -2046,8 +2046,8 @@ def test_send_dedupe_on_resume(
assert graph.invoke(None, thread1, debug=1) == [
"0",
"1",
"2|Command(send=Send(node='2', arg=3))",
"2|Command(send=Send(node='flaky', arg=4))",
"2|Command(goto=Send(node='2', arg=3))",
"2|Command(goto=Send(node='flaky', arg=4))",
"2|3",
"flaky|4",
"3",
@@ -2069,8 +2069,8 @@ def test_send_dedupe_on_resume(
values=[
"0",
"1",
"2|Command(send=Send(node='2', arg=3))",
"2|Command(send=Send(node='flaky', arg=4))",
"2|Command(goto=Send(node='2', arg=3))",
"2|Command(goto=Send(node='flaky', arg=4))",
"2|3",
"flaky|4",
"3",
@@ -2105,8 +2105,8 @@ def test_send_dedupe_on_resume(
values=[
"0",
"1",
"2|Command(send=Send(node='2', arg=3))",
"2|Command(send=Send(node='flaky', arg=4))",
"2|Command(goto=Send(node='2', arg=3))",
"2|Command(goto=Send(node='flaky', arg=4))",
"2|3",
"flaky|4",
],
@@ -2123,8 +2123,8 @@ def test_send_dedupe_on_resume(
"writes": {
"1": ["1"],
"2": [
["2|Command(send=Send(node='2', arg=3))"],
["2|Command(send=Send(node='flaky', arg=4))"],
["2|Command(goto=Send(node='2', arg=3))"],
["2|Command(goto=Send(node='flaky', arg=4))"],
["2|3"],
],
"flaky": ["flaky|4"],
@@ -2209,7 +2209,7 @@ def test_send_dedupe_on_resume(
error=None,
interrupts=(),
state=None,
result=["2|Command(send=Send(node='2', arg=3))"],
result=["2|Command(goto=Send(node='2', arg=3))"],
),
PregelTask(
id=AnyStr(),
@@ -2223,7 +2223,7 @@ def test_send_dedupe_on_resume(
error=None,
interrupts=(),
state=None,
result=["2|Command(send=Send(node='flaky', arg=4))"],
result=["2|Command(goto=Send(node='flaky', arg=4))"],
),
PregelTask(
id=AnyStr(),
@@ -2786,10 +2786,10 @@ def test_send_react_interrupt_control(
tool_calls=[ToolCall(name="foo", args={"hi": [1, 2, 3]}, id=AnyStr())],
)
def agent(state) -> GraphCommand[Literal["foo"]]:
return GraphCommand(
def agent(state) -> Command[Literal["foo"]]:
return Command(
update={"messages": ai_message},
send=[Send(call["name"], call) for call in ai_message.tool_calls],
goto=[Send(call["name"], call) for call in ai_message.tool_calls],
)
foo_called = 0
@@ -14580,9 +14580,9 @@ def test_parent_command(request: pytest.FixtureRequest, checkpointer_name: str)
from langchain_core.tools import tool
@tool(return_direct=True)
def get_user_name() -> GraphCommand:
def get_user_name() -> Command:
"""Retrieve user name"""
return GraphCommand(update={"user_name": "Meow"}, graph=GraphCommand.PARENT)
return Command(update={"user_name": "Meow"}, graph=Command.PARENT)
subgraph_builder = StateGraph(MessagesState)
subgraph_builder.add_node("tool", get_user_name)
@@ -14680,3 +14680,143 @@ def test_interrupt_subgraph(request: pytest.FixtureRequest, checkpointer_name: s
assert graph.invoke({"baz": ""}, thread1)
# Resume with answer
assert graph.invoke(Command(resume="bar"), thread1)
@pytest.mark.parametrize("checkpointer_name", ALL_CHECKPOINTERS_SYNC)
def test_interrupt_multiple(request: pytest.FixtureRequest, checkpointer_name: str):
checkpointer = request.getfixturevalue(f"checkpointer_{checkpointer_name}")
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=checkpointer)
thread1 = {"configurable": {"thread_id": "1"}}
assert [e for e in graph.stream({"my_key": "DE", "market": "DE"}, thread1)] == [
{
"__interrupt__": (
Interrupt(
value={"value": 1},
resumable=True,
ns=[AnyStr("node:")],
when="during",
),
)
}
]
assert [
event
for event in graph.stream(
Command(resume="answer 1", update={"my_key": "foofoo"}), thread1
)
] == [
{
"__interrupt__": (
Interrupt(
value={"value": 2},
resumable=True,
ns=[AnyStr("node:")],
when="during",
),
)
}
]
assert [event for event in graph.stream(Command(resume="answer 2"), thread1)] == [
{"node": {"my_key": "answer 1 answer 2"}},
]
@pytest.mark.parametrize("checkpointer_name", ALL_CHECKPOINTERS_SYNC)
def test_interrupt_loop(request: pytest.FixtureRequest, checkpointer_name: str):
checkpointer = request.getfixturevalue(f"checkpointer_{checkpointer_name}")
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=checkpointer)
thread1 = {"configurable": {"thread_id": "1"}}
assert [e for e in graph.stream({"other": ""}, thread1)] == [
{
"__interrupt__": (
Interrupt(
value="How old are you?",
resumable=True,
ns=[AnyStr("node:")],
when="during",
),
)
}
]
assert [
event
for event in graph.stream(
Command(resume="13"),
thread1,
)
] == [
{
"__interrupt__": (
Interrupt(
value="invalid response",
resumable=True,
ns=[AnyStr("node:")],
when="during",
),
)
}
]
assert [
event
for event in graph.stream(
Command(resume="15"),
thread1,
)
] == [
{
"__interrupt__": (
Interrupt(
value="invalid response",
resumable=True,
ns=[AnyStr("node:")],
when="during",
),
)
}
]
assert [event for event in graph.stream(Command(resume="19"), thread1)] == [
{"node": {"age": 19}},
]
libs/langgraph/tests/test_pregel_async.py
+191 -34
View File
@@ -62,7 +62,7 @@ from langgraph.constants import (
START,
)
from langgraph.errors import InvalidUpdateError, MultipleSubgraphsError, NodeInterrupt
from langgraph.graph import END, Graph, GraphCommand, StateGraph
from langgraph.graph import END, Graph, StateGraph
from langgraph.graph.message import MessageGraph, MessagesState, add_messages
from langgraph.managed.shared_value import SharedValue
from langgraph.prebuilt.chat_agent_executor import create_tool_calling_executor
@@ -2580,8 +2580,8 @@ async def test_send_sequences(checkpointer_name: str) -> None:
async def send_for_fun(state):
return [
Send("2", Command(send=Send("2", 3))),
Send("2", GraphCommand(send=Send("2", 4))),
Send("2", Command(goto=Send("2", 3))),
Send("2", Command(goto=Send("2", 4))),
"3.1",
]
@@ -2602,8 +2602,8 @@ async def test_send_sequences(checkpointer_name: str) -> None:
== [
"0",
"1",
"2|Command(send=Send(node='2', arg=3))",
"2|Command(send=Send(node='2', arg=4))",
"2|Command(goto=Send(node='2', arg=3))",
"2|Command(goto=Send(node='2', arg=4))",
"2|3",
"2|4",
"3",
@@ -2614,8 +2614,8 @@ async def test_send_sequences(checkpointer_name: str) -> None:
"0",
"1",
"3.1",
"2|Command(send=Send(node='2', arg=3))",
"2|Command(send=Send(node='2', arg=4))",
"2|Command(goto=Send(node='2', arg=3))",
"2|Command(goto=Send(node='2', arg=4))",
"3",
"2|3",
"2|4",
@@ -2632,16 +2632,16 @@ async def test_send_sequences(checkpointer_name: str) -> None:
assert await graph.ainvoke(["0"], thread1) == [
"0",
"1",
"2|Command(send=Send(node='2', arg=3))",
"2|Command(send=Send(node='2', arg=4))",
"2|Command(goto=Send(node='2', arg=3))",
"2|Command(goto=Send(node='2', arg=4))",
"2|3",
"2|4",
]
assert await graph.ainvoke(None, thread1) == [
"0",
"1",
"2|Command(send=Send(node='2', arg=3))",
"2|Command(send=Send(node='2', arg=4))",
"2|Command(goto=Send(node='2', arg=3))",
"2|Command(goto=Send(node='2', arg=4))",
"2|3",
"2|4",
"3",
@@ -2677,15 +2677,15 @@ async def test_send_dedupe_on_resume(checkpointer_name: str) -> None:
if isinstance(state, list)
else ["|".join((self.name, str(state)))]
)
if isinstance(state, GraphCommand):
if isinstance(state, Command):
return replace(state, update=update)
else:
return update
def send_for_fun(state):
return [
Send("2", GraphCommand(send=Send("2", 3))),
Send("2", GraphCommand(send=Send("flaky", 4))),
Send("2", Command(goto=Send("2", 3))),
Send("2", Command(goto=Send("flaky", 4))),
"3.1",
]
@@ -2708,8 +2708,8 @@ async def test_send_dedupe_on_resume(checkpointer_name: str) -> None:
assert await graph.ainvoke(["0"], thread1, debug=1) == [
"0",
"1",
"2|Command(send=Send(node='2', arg=3))",
"2|Command(send=Send(node='flaky', arg=4))",
"2|Command(goto=Send(node='2', arg=3))",
"2|Command(goto=Send(node='flaky', arg=4))",
"2|3",
]
assert builder.nodes["2"].runnable.func.ticks == 3
@@ -2718,8 +2718,8 @@ async def test_send_dedupe_on_resume(checkpointer_name: str) -> None:
assert await graph.ainvoke(None, thread1, debug=1) == [
"0",
"1",
"2|Command(send=Send(node='2', arg=3))",
"2|Command(send=Send(node='flaky', arg=4))",
"2|Command(goto=Send(node='2', arg=3))",
"2|Command(goto=Send(node='flaky', arg=4))",
"2|3",
"flaky|4",
"3",
@@ -2736,8 +2736,8 @@ async def test_send_dedupe_on_resume(checkpointer_name: str) -> None:
values=[
"0",
"1",
"2|Command(send=Send(node='2', arg=3))",
"2|Command(send=Send(node='flaky', arg=4))",
"2|Command(goto=Send(node='2', arg=3))",
"2|Command(goto=Send(node='flaky', arg=4))",
"2|3",
"flaky|4",
"3",
@@ -2772,8 +2772,8 @@ async def test_send_dedupe_on_resume(checkpointer_name: str) -> None:
values=[
"0",
"1",
"2|Command(send=Send(node='2', arg=3))",
"2|Command(send=Send(node='flaky', arg=4))",
"2|Command(goto=Send(node='2', arg=3))",
"2|Command(goto=Send(node='flaky', arg=4))",
"2|3",
"flaky|4",
],
@@ -2790,8 +2790,8 @@ async def test_send_dedupe_on_resume(checkpointer_name: str) -> None:
"writes": {
"1": ["1"],
"2": [
["2|Command(send=Send(node='2', arg=3))"],
["2|Command(send=Send(node='flaky', arg=4))"],
["2|Command(goto=Send(node='2', arg=3))"],
["2|Command(goto=Send(node='flaky', arg=4))"],
["2|3"],
],
"flaky": ["flaky|4"],
@@ -2876,7 +2876,7 @@ async def test_send_dedupe_on_resume(checkpointer_name: str) -> None:
error=None,
interrupts=(),
state=None,
result=["2|Command(send=Send(node='2', arg=3))"],
result=["2|Command(goto=Send(node='2', arg=3))"],
),
PregelTask(
id=AnyStr(),
@@ -2890,7 +2890,7 @@ async def test_send_dedupe_on_resume(checkpointer_name: str) -> None:
error=None,
interrupts=(),
state=None,
result=["2|Command(send=Send(node='flaky', arg=4))"],
result=["2|Command(goto=Send(node='flaky', arg=4))"],
),
PregelTask(
id=AnyStr(),
@@ -3448,9 +3448,9 @@ async def test_send_react_interrupt_control(
)
async def agent(state) -> Command[Literal["foo"]]:
return GraphCommand(
return Command(
update={"messages": ai_message},
send=[Send(call["name"], call) for call in ai_message.tool_calls],
goto=[Send(call["name"], call) for call in ai_message.tool_calls],
)
foo_called = 0
@@ -3761,13 +3761,13 @@ async def test_max_concurrency(checkpointer_name: str) -> None:
@pytest.mark.parametrize("checkpointer_name", ALL_CHECKPOINTERS_ASYNC)
async def test_max_concurrency_control(checkpointer_name: str) -> None:
async def node1(state) -> GraphCommand[Literal["2"]]:
return GraphCommand(update=["1"], send=[Send("2", idx) for idx in range(100)])
async def node1(state) -> Command[Literal["2"]]:
return Command(update=["1"], goto=[Send("2", idx) for idx in range(100)])
node2_currently = 0
node2_max_currently = 0
async def node2(state) -> GraphCommand[Literal["3"]]:
async def node2(state) -> Command[Literal["3"]]:
nonlocal node2_currently, node2_max_currently
node2_currently += 1
if node2_currently > node2_max_currently:
@@ -3775,7 +3775,7 @@ async def test_max_concurrency_control(checkpointer_name: str) -> None:
await asyncio.sleep(0.1)
node2_currently -= 1
return GraphCommand(update=[state], goto="3")
return Command(update=[state], goto="3")
async def node3(state) -> Literal["3"]:
return ["3"]
@@ -12788,9 +12788,9 @@ async def test_parent_command(checkpointer_name: str) -> None:
from langchain_core.tools import tool
@tool(return_direct=True)
def get_user_name() -> GraphCommand:
def get_user_name() -> Command:
"""Retrieve user name"""
return GraphCommand(update={"user_name": "Meow"}, graph=GraphCommand.PARENT)
return Command(update={"user_name": "Meow"}, graph=Command.PARENT)
subgraph_builder = StateGraph(MessagesState)
subgraph_builder.add_node("tool", get_user_name)
@@ -12896,3 +12896,160 @@ async def test_interrupt_subgraph(checkpointer_name: str):
assert await graph.ainvoke({"baz": ""}, thread1)
# Resume with answer
assert await graph.ainvoke(Command(resume="bar"), thread1)
@pytest.mark.skipif(
sys.version_info < (3, 11),
reason="Python 3.11+ is required for async contextvars support",
)
@pytest.mark.parametrize("checkpointer_name", ALL_CHECKPOINTERS_ASYNC)
async def test_interrupt_multiple(checkpointer_name: str):
class State(TypedDict):
my_key: Annotated[str, operator.add]
async 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")
async with awith_checkpointer(checkpointer_name) as checkpointer:
graph = builder.compile(checkpointer=checkpointer)
thread1 = {"configurable": {"thread_id": "1"}}
assert [
e async for e in graph.astream({"my_key": "DE", "market": "DE"}, thread1)
] == [
{
"__interrupt__": (
Interrupt(
value={"value": 1},
resumable=True,
ns=[AnyStr("node:")],
when="during",
),
)
}
]
assert [
event
async for event in graph.astream(
Command(resume="answer 1", update={"my_key": "foofoo"}),
thread1,
stream_mode="updates",
)
] == [
{
"__interrupt__": (
Interrupt(
value={"value": 2},
resumable=True,
ns=[AnyStr("node:")],
when="during",
),
)
}
]
assert [
event
async for event in graph.astream(
Command(resume="answer 2"), thread1, stream_mode="updates"
)
] == [
{"node": {"my_key": "answer 1 answer 2"}},
]
@pytest.mark.skipif(
sys.version_info < (3, 11),
reason="Python 3.11+ is required for async contextvars support",
)
@pytest.mark.parametrize("checkpointer_name", ALL_CHECKPOINTERS_ASYNC)
async def test_interrupt_loop(checkpointer_name: str):
class State(TypedDict):
age: int
other: str
async 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")
async with awith_checkpointer(checkpointer_name) as checkpointer:
graph = builder.compile(checkpointer=checkpointer)
thread1 = {"configurable": {"thread_id": "1"}}
assert [e async for e in graph.astream({"other": ""}, thread1)] == [
{
"__interrupt__": (
Interrupt(
value="How old are you?",
resumable=True,
ns=[AnyStr("node:")],
when="during",
),
)
}
]
assert [
event
async for event in graph.astream(
Command(resume="13"),
thread1,
)
] == [
{
"__interrupt__": (
Interrupt(
value="invalid response",
resumable=True,
ns=[AnyStr("node:")],
when="during",
),
)
}
]
assert [
event
async for event in graph.astream(
Command(resume="15"),
thread1,
)
] == [
{
"__interrupt__": (
Interrupt(
value="invalid response",
resumable=True,
ns=[AnyStr("node:")],
when="during",
),
)
}
]
assert [
event async for event in graph.astream(Command(resume="19"), thread1)
] == [
{"node": {"age": 19}},
]
libs/scheduler-kafka/tests/any.py
+18
View File
@@ -35,3 +35,21 @@ class AnyDict(dict):
return False
else:
return True
class AnyList(list):
def __init__(self, *args, **kwargs) -> None:
super().__init__(*args, **kwargs)
def __eq__(self, other: object) -> bool:
if not self and isinstance(other, list):
return True
if not isinstance(other, list) or len(self) != len(other):
return False
for i, v in enumerate(self):
if v == other[i]:
continue
else:
return False
else:
return True
libs/scheduler-kafka/tests/test_push.py
+9 -9
View File
@@ -11,10 +11,10 @@ from aiokafka import AIOKafkaProducer
from langgraph.checkpoint.base import BaseCheckpointSaver
from langgraph.constants import FF_SEND_V2, START
from langgraph.errors import NodeInterrupt
from langgraph.graph.state import CompiledStateGraph, GraphCommand, StateGraph
from langgraph.graph.state import CompiledStateGraph, StateGraph
from langgraph.scheduler.kafka import serde
from langgraph.scheduler.kafka.types import MessageToOrchestrator, Topics
from langgraph.types import Send
from langgraph.types import Command, Send
from tests.any import AnyDict
from tests.drain import drain_topics_async
@@ -48,15 +48,15 @@ def mk_push_graph(
if isinstance(state, list)
else ["|".join((self.name, str(state)))]
)
if isinstance(state, GraphCommand):
if isinstance(state, Command):
return state.copy(update=update)
else:
return update
def send_for_fun(state):
return [
Send("2", GraphCommand(send=Send("2", 3))),
Send("2", GraphCommand(send=Send("flaky", 4))),
Send("2", Command(goto=Send("2", 3))),
Send("2", Command(goto=Send("flaky", 4))),
"3.1",
]
@@ -105,8 +105,8 @@ async def test_push_graph(topics: Topics, acheckpointer: BaseCheckpointSaver) ->
== [
"0",
"1",
"2|Control(send=Send(node='2', arg=3))",
"2|Control(send=Send(node='flaky', arg=4))",
"2|Control(goto=Send(node='2', arg=3))",
"2|Control(goto=Send(node='flaky', arg=4))",
"2|3",
]
)
@@ -182,8 +182,8 @@ async def test_push_graph(topics: Topics, acheckpointer: BaseCheckpointSaver) ->
== [
"0",
"1",
"2|Control(send=Send(node='2', arg=3))",
"2|Control(send=Send(node='flaky', arg=4))",
"2|Control(goto=Send(node='2', arg=3))",
"2|Control(goto=Send(node='flaky', arg=4))",
"2|3",
"flaky|4",
"3",
libs/scheduler-kafka/tests/test_push_sync.py
+9 -9
View File
@@ -10,11 +10,11 @@ import pytest
from langgraph.checkpoint.base import BaseCheckpointSaver
from langgraph.constants import FF_SEND_V2, START
from langgraph.errors import NodeInterrupt
from langgraph.graph.state import CompiledStateGraph, GraphCommand, StateGraph
from langgraph.graph.state import CompiledStateGraph, StateGraph
from langgraph.scheduler.kafka import serde
from langgraph.scheduler.kafka.default_sync import DefaultProducer
from langgraph.scheduler.kafka.types import MessageToOrchestrator, Topics
from langgraph.types import Send
from langgraph.types import Command, Send
from tests.any import AnyDict
from tests.drain import drain_topics
@@ -48,15 +48,15 @@ def mk_push_graph(
if isinstance(state, list)
else ["|".join((self.name, str(state)))]
)
if isinstance(state, GraphCommand):
if isinstance(state, Command):
return state.copy(update=update)
else:
return update
def send_for_fun(state):
return [
Send("2", GraphCommand(send=Send("2", 3))),
Send("2", GraphCommand(send=Send("flaky", 4))),
Send("2", Command(goto=Send("2", 3))),
Send("2", Command(goto=Send("flaky", 4))),
"3.1",
]
@@ -106,8 +106,8 @@ def test_push_graph(topics: Topics, acheckpointer: BaseCheckpointSaver) -> None:
== [
"0",
"1",
"2|Control(send=Send(node='2', arg=3))",
"2|Control(send=Send(node='flaky', arg=4))",
"2|Control(goto=Send(node='2', arg=3))",
"2|Control(goto=Send(node='flaky', arg=4))",
"2|3",
]
)
@@ -184,8 +184,8 @@ def test_push_graph(topics: Topics, acheckpointer: BaseCheckpointSaver) -> None:
== [
"0",
"1",
"2|Control(send=Send(node='2', arg=3))",
"2|Control(send=Send(node='flaky', arg=4))",
"2|Control(goto=Send(node='2', arg=3))",
"2|Control(goto=Send(node='flaky', arg=4))",
"2|3",
"flaky|4",
"3",
libs/scheduler-kafka/tests/test_subgraph.py
+13 -7
View File
@@ -15,7 +15,7 @@ from langgraph.graph.state import StateGraph
from langgraph.pregel import Pregel
from langgraph.scheduler.kafka import serde
from langgraph.scheduler.kafka.types import MessageToOrchestrator, Topics
from tests.any import AnyDict
from tests.any import AnyDict, AnyList
from tests.drain import drain_topics_async
from tests.messages import _AnyIdAIMessage, _AnyIdHumanMessage
@@ -196,7 +196,8 @@ async def test_subgraph_w_interrupt(
"__pregel_resuming": False,
"__pregel_store": None,
"__pregel_task_id": history[0].tasks[0].id,
"__pregel_resume_value": None,
"__pregel_scratchpad": {},
"__pregel_writes": AnyList(),
"checkpoint_id": None,
"checkpoint_map": {
"": history[0].config["configurable"]["checkpoint_id"]
@@ -261,7 +262,8 @@ async def test_subgraph_w_interrupt(
"__pregel_resuming": False,
"__pregel_store": None,
"__pregel_task_id": history[0].tasks[0].id,
"__pregel_resume_value": None,
"__pregel_scratchpad": {},
"__pregel_writes": AnyList(),
"checkpoint_id": c.config["configurable"]["checkpoint_id"],
"checkpoint_map": {
"": history[0].config["configurable"]["checkpoint_id"]
@@ -356,7 +358,8 @@ async def test_subgraph_w_interrupt(
"__pregel_resuming": False,
"__pregel_store": None,
"__pregel_task_id": history[0].tasks[0].id,
"__pregel_resume_value": None,
"__pregel_scratchpad": {},
"__pregel_writes": AnyList(),
"checkpoint_id": c.config["configurable"]["checkpoint_id"],
"checkpoint_map": {
"": history[0].config["configurable"]["checkpoint_id"]
@@ -461,7 +464,8 @@ async def test_subgraph_w_interrupt(
"__pregel_resuming": True,
"__pregel_store": None,
"__pregel_task_id": history[1].tasks[0].id,
"__pregel_resume_value": None,
"__pregel_scratchpad": {},
"__pregel_writes": AnyList(),
"checkpoint_id": None,
"checkpoint_map": {
"": history[1].config["configurable"]["checkpoint_id"]
@@ -521,7 +525,8 @@ async def test_subgraph_w_interrupt(
"__pregel_resuming": True,
"__pregel_store": None,
"__pregel_task_id": history[1].tasks[0].id,
"__pregel_resume_value": None,
"__pregel_scratchpad": {},
"__pregel_writes": AnyList(),
"checkpoint_id": c.config["configurable"]["checkpoint_id"],
"checkpoint_map": {
"": history[1].config["configurable"]["checkpoint_id"]
@@ -637,7 +642,8 @@ async def test_subgraph_w_interrupt(
"__pregel_resuming": True,
"__pregel_store": None,
"__pregel_task_id": history[1].tasks[0].id,
"__pregel_resume_value": None,
"__pregel_scratchpad": {},
"__pregel_writes": AnyList(),
"checkpoint_id": c.config["configurable"]["checkpoint_id"],
"checkpoint_map": {
"": history[1].config["configurable"]["checkpoint_id"]
libs/scheduler-kafka/tests/test_subgraph_sync.py
+13 -7
View File
@@ -15,7 +15,7 @@ from langgraph.pregel import Pregel
from langgraph.scheduler.kafka import serde
from langgraph.scheduler.kafka.default_sync import DefaultProducer
from langgraph.scheduler.kafka.types import MessageToOrchestrator, Topics
from tests.any import AnyDict
from tests.any import AnyDict, AnyList
from tests.drain import drain_topics
from tests.messages import _AnyIdAIMessage, _AnyIdHumanMessage
@@ -195,7 +195,8 @@ def test_subgraph_w_interrupt(
"__pregel_resuming": False,
"__pregel_store": None,
"__pregel_task_id": history[0].tasks[0].id,
"__pregel_resume_value": None,
"__pregel_scratchpad": {},
"__pregel_writes": AnyList(),
"checkpoint_id": None,
"checkpoint_map": {
"": history[0].config["configurable"]["checkpoint_id"]
@@ -260,7 +261,8 @@ def test_subgraph_w_interrupt(
"__pregel_dedupe_tasks": True,
"__pregel_resuming": False,
"__pregel_task_id": history[0].tasks[0].id,
"__pregel_resume_value": None,
"__pregel_scratchpad": {},
"__pregel_writes": AnyList(),
"checkpoint_id": c.config["configurable"]["checkpoint_id"],
"checkpoint_map": {
"": history[0].config["configurable"]["checkpoint_id"]
@@ -355,7 +357,8 @@ def test_subgraph_w_interrupt(
"__pregel_store": None,
"__pregel_resuming": False,
"__pregel_task_id": history[0].tasks[0].id,
"__pregel_resume_value": None,
"__pregel_scratchpad": {},
"__pregel_writes": AnyList(),
"checkpoint_id": c.config["configurable"]["checkpoint_id"],
"checkpoint_map": {
"": history[0].config["configurable"]["checkpoint_id"]
@@ -459,7 +462,8 @@ def test_subgraph_w_interrupt(
"__pregel_store": None,
"__pregel_resuming": True,
"__pregel_task_id": history[1].tasks[0].id,
"__pregel_resume_value": None,
"__pregel_scratchpad": {},
"__pregel_writes": AnyList(),
"checkpoint_id": None,
"checkpoint_map": {
"": history[1].config["configurable"]["checkpoint_id"]
@@ -519,7 +523,8 @@ def test_subgraph_w_interrupt(
"__pregel_store": None,
"__pregel_resuming": True,
"__pregel_task_id": history[1].tasks[0].id,
"__pregel_resume_value": None,
"__pregel_scratchpad": {},
"__pregel_writes": AnyList(),
"checkpoint_id": c.config["configurable"]["checkpoint_id"],
"checkpoint_map": {
"": history[1].config["configurable"]["checkpoint_id"]
@@ -635,7 +640,8 @@ def test_subgraph_w_interrupt(
"__pregel_resuming": True,
"__pregel_store": None,
"__pregel_task_id": history[1].tasks[0].id,
"__pregel_resume_value": None,
"__pregel_scratchpad": {},
"__pregel_writes": AnyList(),
"checkpoint_id": c.config["configurable"]["checkpoint_id"],
"checkpoint_map": {
"": history[1].config["configurable"]["checkpoint_id"]
libs/sdk-py/langgraph_sdk/schema.py
+1 -1
View File
@@ -373,6 +373,6 @@ class Send(TypedDict):
class Command(TypedDict, total=False):
send: Union[Send, Sequence[Send]]
goto: Union[Send, str, Sequence[Union[Send, str]]]
update: dict[str, Any]
resume: Any