From ea5ccd7a80e11e53002486859ca6e0004d289518 Mon Sep 17 00:00:00 2001 From: Nuno Campos Date: Wed, 4 Dec 2024 14:15:30 -0800 Subject: [PATCH 1/6] lib: Add support for multiple interrupts per node - Includes support for interrupt loops --- libs/langgraph/langgraph/constants.py | 4 + libs/langgraph/langgraph/pregel/algo.py | 35 +++-- libs/langgraph/langgraph/pregel/io.py | 9 +- libs/langgraph/langgraph/pregel/loop.py | 25 +++- libs/langgraph/langgraph/pregel/runner.py | 3 + libs/langgraph/langgraph/types.py | 66 +++++++-- libs/langgraph/tests/test_pregel.py | 140 +++++++++++++++++++ libs/langgraph/tests/test_pregel_async.py | 157 ++++++++++++++++++++++ 8 files changed, 403 insertions(+), 36 deletions(-) diff --git a/libs/langgraph/langgraph/constants.py b/libs/langgraph/langgraph/constants.py index 478f23d68..efdf65143 100644 --- a/libs/langgraph/langgraph/constants.py +++ b/libs/langgraph/langgraph/constants.py @@ -75,6 +75,10 @@ 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") diff --git a/libs/langgraph/langgraph/pregel/algo.py b/libs/langgraph/langgraph/pregel/algo.py index 5d104a85f..0885f12aa 100644 --- a/libs/langgraph/langgraph/pregel/algo.py +++ b/libs/langgraph/langgraph/pregel/algo.py @@ -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, diff --git a/libs/langgraph/langgraph/pregel/io.py b/libs/langgraph/langgraph/pregel/io.py index 693dffce2..918f3d899 100644 --- a/libs/langgraph/langgraph/pregel/io.py +++ b/libs/langgraph/langgraph/pregel/io.py @@ -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,7 +67,7 @@ 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: @@ -85,7 +86,11 @@ def map_command( 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 = next( + (w 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: diff --git a/libs/langgraph/langgraph/pregel/loop.py b/libs/langgraph/langgraph/pregel/loop.py index 2a68b00f2..d9af9279e 100644 --- a/libs/langgraph/langgraph/pregel/loop.py +++ b/libs/langgraph/langgraph/pregel/loop.py @@ -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") diff --git a/libs/langgraph/langgraph/pregel/runner.py b/libs/langgraph/langgraph/pregel/runner.py index 9e3879b0f..f46210459 100644 --- a/libs/langgraph/langgraph/pregel/runner.py +++ b/libs/langgraph/langgraph/pregel/runner.py @@ -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 diff --git a/libs/langgraph/langgraph/types.py b/libs/langgraph/langgraph/types.py index 7bf9148c5..4047a28f7 100644 --- a/libs/langgraph/langgraph/types.py +++ b/libs/langgraph/langgraph/types.py @@ -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.""" @@ -300,26 +306,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), + ), ) + ) diff --git a/libs/langgraph/tests/test_pregel.py b/libs/langgraph/tests/test_pregel.py index 4ff56ca3e..4f768badd 100644 --- a/libs/langgraph/tests/test_pregel.py +++ b/libs/langgraph/tests/test_pregel.py @@ -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}}, + ] diff --git a/libs/langgraph/tests/test_pregel_async.py b/libs/langgraph/tests/test_pregel_async.py index cc244c363..addb18b2b 100644 --- a/libs/langgraph/tests/test_pregel_async.py +++ b/libs/langgraph/tests/test_pregel_async.py @@ -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}}, + ] From fb01d65dc07af9a46ff05a68dabb5ec57f1fb9d6 Mon Sep 17 00:00:00 2001 From: Nuno Campos Date: Wed, 4 Dec 2024 14:31:22 -0800 Subject: [PATCH 2/6] Lint --- libs/langgraph/langgraph/pregel/io.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/libs/langgraph/langgraph/pregel/io.py b/libs/langgraph/langgraph/pregel/io.py index 918f3d899..ed2c28938 100644 --- a/libs/langgraph/langgraph/pregel/io.py +++ b/libs/langgraph/langgraph/pregel/io.py @@ -86,8 +86,8 @@ def map_command( 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(): - existing = next( - (w for w in pending_writes if w[0] == tid and w[1] == 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) From 5c7a6689af406faeaabf251a2f6009256eb9521a Mon Sep 17 00:00:00 2001 From: Nuno Campos Date: Wed, 4 Dec 2024 14:41:30 -0800 Subject: [PATCH 3/6] Update tests --- libs/langgraph/langgraph/constants.py | 2 -- libs/scheduler-kafka/tests/any.py | 18 +++++++++++++++++ libs/scheduler-kafka/tests/test_subgraph.py | 20 ++++++++++++------- .../tests/test_subgraph_sync.py | 20 ++++++++++++------- 4 files changed, 44 insertions(+), 16 deletions(-) diff --git a/libs/langgraph/langgraph/constants.py b/libs/langgraph/langgraph/constants.py index efdf65143..e2d9f069a 100644 --- a/libs/langgraph/langgraph/constants.py +++ b/libs/langgraph/langgraph/constants.py @@ -72,8 +72,6 @@ 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 diff --git a/libs/scheduler-kafka/tests/any.py b/libs/scheduler-kafka/tests/any.py index 73744a1e8..3ea224173 100644 --- a/libs/scheduler-kafka/tests/any.py +++ b/libs/scheduler-kafka/tests/any.py @@ -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 diff --git a/libs/scheduler-kafka/tests/test_subgraph.py b/libs/scheduler-kafka/tests/test_subgraph.py index ebaaea580..4ab92676c 100644 --- a/libs/scheduler-kafka/tests/test_subgraph.py +++ b/libs/scheduler-kafka/tests/test_subgraph.py @@ -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"] diff --git a/libs/scheduler-kafka/tests/test_subgraph_sync.py b/libs/scheduler-kafka/tests/test_subgraph_sync.py index 75b9d6e73..5fa43998a 100644 --- a/libs/scheduler-kafka/tests/test_subgraph_sync.py +++ b/libs/scheduler-kafka/tests/test_subgraph_sync.py @@ -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"] From d457ad3cc272b8884a3a9e6cd829323748cc8ae5 Mon Sep 17 00:00:00 2001 From: William FH <13333726+hinthornw@users.noreply.github.com> Date: Wed, 4 Dec 2024 14:50:23 -0800 Subject: [PATCH 4/6] Clean up code snippet (#2637) --- docs/docs/concepts/persistence.md | 5 ++++- docs/docs/how-tos/memory/semantic-search.ipynb | 1 - 2 files changed, 4 insertions(+), 2 deletions(-) diff --git a/docs/docs/concepts/persistence.md b/docs/docs/concepts/persistence.md index 8f4aa993a..0ec126316 100644 --- a/docs/docs/concepts/persistence.md +++ b/docs/docs/concepts/persistence.md @@ -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?", diff --git a/docs/docs/how-tos/memory/semantic-search.ipynb b/docs/docs/how-tos/memory/semantic-search.ipynb index 24905e625..658e4bb29 100644 --- a/docs/docs/how-tos/memory/semantic-search.ipynb +++ b/docs/docs/how-tos/memory/semantic-search.ipynb @@ -297,7 +297,6 @@ } ], "source": [ - "embeddings = init_embeddings(\"openai:text-embedding-3-small\")\n", "store = InMemoryStore(\n", " index={\n", " \"embed\": embeddings,\n", From b4b3ac6f57adad2b0458026622c1e0ceb07c6c9f Mon Sep 17 00:00:00 2001 From: Nuno Campos Date: Wed, 4 Dec 2024 15:12:03 -0800 Subject: [PATCH 5/6] lib: Merge GraphCommand and Command - Now we have only Command - Command(goto=) combines the previous functionality of Command(send=) and Command(goto=) --- libs/langgraph/langgraph/graph/__init__.py | 3 +- libs/langgraph/langgraph/graph/state.py | 48 ++++---------- libs/langgraph/langgraph/pregel/io.py | 9 +-- libs/langgraph/langgraph/types.py | 2 +- libs/langgraph/tests/test_pregel.py | 58 ++++++++--------- libs/langgraph/tests/test_pregel_async.py | 68 ++++++++++---------- libs/scheduler-kafka/tests/test_push.py | 18 +++--- libs/scheduler-kafka/tests/test_push_sync.py | 18 +++--- libs/sdk-py/langgraph_sdk/schema.py | 2 +- 9 files changed, 101 insertions(+), 125 deletions(-) diff --git a/libs/langgraph/langgraph/graph/__init__.py b/libs/langgraph/langgraph/graph/__init__.py index 241106a3a..c81ad9903 100644 --- a/libs/langgraph/langgraph/graph/__init__.py +++ b/libs/langgraph/langgraph/graph/__init__.py @@ -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", diff --git a/libs/langgraph/langgraph/graph/state.py b/libs/langgraph/langgraph/graph/state.py index 1a7208a2a..e63f25111 100644 --- a/libs/langgraph/langgraph/graph/state.py +++ b/libs/langgraph/langgraph/graph/state.py @@ -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,22 +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.""" - - 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]] @@ -392,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])) @@ -834,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 @@ -854,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 diff --git a/libs/langgraph/langgraph/pregel/io.py b/libs/langgraph/langgraph/pregel/io.py index ed2c28938..c1fed349a 100644 --- a/libs/langgraph/langgraph/pregel/io.py +++ b/libs/langgraph/langgraph/pregel/io.py @@ -72,17 +72,18 @@ def map_command( """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 cmd.goto: if isinstance(cmd.send, (tuple, list)): - sends = cmd.send + 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(): diff --git a/libs/langgraph/langgraph/types.py b/libs/langgraph/langgraph/types.py index 4047a28f7..67c7e53f8 100644 --- a/libs/langgraph/langgraph/types.py +++ b/libs/langgraph/langgraph/types.py @@ -249,8 +249,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 diff --git a/libs/langgraph/tests/test_pregel.py b/libs/langgraph/tests/test_pregel.py index 4f768badd..68071594b 100644 --- a/libs/langgraph/tests/test_pregel.py +++ b/libs/langgraph/tests/test_pregel.py @@ -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) diff --git a/libs/langgraph/tests/test_pregel_async.py b/libs/langgraph/tests/test_pregel_async.py index addb18b2b..538730c78 100644 --- a/libs/langgraph/tests/test_pregel_async.py +++ b/libs/langgraph/tests/test_pregel_async.py @@ -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) diff --git a/libs/scheduler-kafka/tests/test_push.py b/libs/scheduler-kafka/tests/test_push.py index 15e9211a2..3d2e4d43d 100644 --- a/libs/scheduler-kafka/tests/test_push.py +++ b/libs/scheduler-kafka/tests/test_push.py @@ -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", diff --git a/libs/scheduler-kafka/tests/test_push_sync.py b/libs/scheduler-kafka/tests/test_push_sync.py index 27cd96cb7..ee33d613e 100644 --- a/libs/scheduler-kafka/tests/test_push_sync.py +++ b/libs/scheduler-kafka/tests/test_push_sync.py @@ -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", diff --git a/libs/sdk-py/langgraph_sdk/schema.py b/libs/sdk-py/langgraph_sdk/schema.py index 1ccae3e89..6237ea5bd 100644 --- a/libs/sdk-py/langgraph_sdk/schema.py +++ b/libs/sdk-py/langgraph_sdk/schema.py @@ -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 From df70e91daecac6b7d2b187b7c67a5c725e7dbe11 Mon Sep 17 00:00:00 2001 From: Nuno Campos Date: Wed, 4 Dec 2024 15:13:55 -0800 Subject: [PATCH 6/6] Lint --- libs/langgraph/langgraph/pregel/io.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/libs/langgraph/langgraph/pregel/io.py b/libs/langgraph/langgraph/pregel/io.py index c1fed349a..b2596d3ad 100644 --- a/libs/langgraph/langgraph/pregel/io.py +++ b/libs/langgraph/langgraph/pregel/io.py @@ -73,7 +73,7 @@ def map_command( if cmd.graph == Command.PARENT: raise InvalidUpdateError("There is not parent graph") if cmd.goto: - if isinstance(cmd.send, (tuple, list)): + if isinstance(cmd.goto, (tuple, list)): sends = cmd.goto else: sends = [cmd.goto]