diff --git a/libs/checkpoint/langgraph/store/base/__init__.py b/libs/checkpoint/langgraph/store/base/__init__.py index b784d3aa5..cf1f75a41 100644 --- a/libs/checkpoint/langgraph/store/base/__init__.py +++ b/libs/checkpoint/langgraph/store/base/__init__.py @@ -6,11 +6,15 @@ scoped to user IDs, assistant IDs, or other arbitrary namespaces. from abc import ABC, abstractmethod from datetime import datetime -from typing import ( Any, Iterable, Literal, NamedTuple, - Optional, TypedDict, Union, cast) +from typing import Any, Iterable, Literal, NamedTuple, Optional, TypedDict, Union, cast from langchain_core.embeddings import Embeddings -from langgraph.store.base._embed import AEmbeddingsFunc, EmbeddingsFunc, ensure_embeddings + +from langgraph.store.base._embed import ( + AEmbeddingsFunc, + EmbeddingsFunc, + ensure_embeddings, +) class Item: @@ -528,6 +532,7 @@ class BaseStore(ABC): ) return (await self.abatch([op]))[0] + __all__ = [ "BaseStore", "Item", @@ -541,4 +546,4 @@ __all__ = [ "NamespaceMatchType", "Embeddings", "ensure_embeddings", -] \ No newline at end of file +] diff --git a/libs/checkpoint/langgraph/store/base/_embed.py b/libs/checkpoint/langgraph/store/base/_embed.py index a2b4b4ae4..1c82c86e6 100644 --- a/libs/checkpoint/langgraph/store/base/_embed.py +++ b/libs/checkpoint/langgraph/store/base/_embed.py @@ -7,8 +7,7 @@ asynchronous operations. """ import asyncio -from typing import (Any, Awaitable, Callable, List, Optional, Sequence, - TypeGuard, Union) +from typing import Any, Awaitable, Callable, List, Optional, Sequence, Union from langchain_core.embeddings import Embeddings @@ -30,7 +29,7 @@ Similar to EmbeddingsFunc, but returns an awaitable that resolves to the embeddi def ensure_embeddings( embed: Union[Embeddings, EmbeddingsFunc, AEmbeddingsFunc, None], *, - aembed: Optional[AEmbeddingsFunc] = None + aembed: Optional[AEmbeddingsFunc] = None, ) -> Embeddings: """Ensure that an embedding function conforms to LangChain's Embeddings interface. @@ -180,14 +179,14 @@ class EmbeddingsLambda(Embeddings): def _is_async_callable( func: Any, -) -> TypeGuard[Callable[..., Awaitable]]: +) -> bool: """Check if a function is async. - + This includes both async def functions and classes with async __call__ methods. - + Args: func: Function or callable object to check. - + Returns: True if the function is async, False otherwise. """ diff --git a/libs/checkpoint/langgraph/store/base/batch.py b/libs/checkpoint/langgraph/store/base/batch.py index 545ab1d1c..90f837622 100644 --- a/libs/checkpoint/langgraph/store/base/batch.py +++ b/libs/checkpoint/langgraph/store/base/batch.py @@ -6,6 +6,9 @@ from langgraph.store.base import ( BaseStore, GetOp, Item, + ListNamespacesOp, + MatchCondition, + NameSpacePath, Op, PutOp, SearchOp, @@ -69,6 +72,74 @@ class AsyncBatchedBaseStore(BaseStore): self._aqueue[fut] = PutOp(namespace, key, None) return await fut + async def alist_namespaces( + self, + *, + prefix: Optional[NameSpacePath] = None, + suffix: Optional[NameSpacePath] = None, + max_depth: Optional[int] = None, + limit: int = 100, + offset: int = 0, + ) -> list[tuple[str, ...]]: + fut = self._loop.create_future() + match_conditions = [] + if prefix: + match_conditions.append(MatchCondition(match_type="prefix", path=prefix)) + if suffix: + match_conditions.append(MatchCondition(match_type="suffix", path=suffix)) + + op = ListNamespacesOp( + match_conditions=tuple(match_conditions), + max_depth=max_depth, + limit=limit, + offset=offset, + ) + self._aqueue[fut] = op + return await fut + + +def _dedupe_ops(values: list[Op]) -> tuple[Optional[list[int]], list[Op]]: + """Dedupe operations while preserving order for results. + + Args: + values: List of operations to dedupe + + Returns: + Tuple of (listen indices, deduped operations) + where listen indices map deduped operation results back to original positions + """ + if len(values) <= 1: + return None, list(values) + + dedupped: list[Op] = [] + listen: list[int] = [] + puts: dict[tuple[tuple[str, ...], str], int] = {} + + for op in values: + if isinstance(op, (GetOp, SearchOp, ListNamespacesOp)): + try: + listen.append(dedupped.index(op)) + except ValueError: + listen.append(len(dedupped)) + dedupped.append(op) + elif isinstance(op, PutOp): + putkey = (op.namespace, op.key) + if putkey in puts: + # Overwrite previous put + ix = puts[putkey] + dedupped[ix] = op + listen.append(ix) + else: + puts[putkey] = len(dedupped) + listen.append(len(dedupped)) + dedupped.append(op) + + else: # Any new ops will be treated regularly + listen.append(len(dedupped)) + dedupped.append(op) + + return listen, dedupped + async def _run( aqueue: dict[asyncio.Future, Op], store: weakref.ReferenceType[BaseStore] @@ -82,7 +153,12 @@ async def _run( taken = aqueue.copy() # action each operation try: - results = await s.abatch(taken.values()) + values = list(taken.values()) + listen, dedupped = _dedupe_ops(values) + results = await s.abatch(dedupped) + if listen is not None: + results = [results[ix] for ix in listen] + # set the results of each operation for fut, result in zip(taken, results): fut.set_result(result) diff --git a/libs/checkpoint/tests/test_store.py b/libs/checkpoint/tests/test_store.py index 9d06281d0..0ecd4bd84 100644 --- a/libs/checkpoint/tests/test_store.py +++ b/libs/checkpoint/tests/test_store.py @@ -10,6 +10,18 @@ from langgraph.store.base.batch import AsyncBatchedBaseStore from langgraph.store.memory import InMemoryStore +class MockAsyncBatchedStore(AsyncBatchedBaseStore): + def __init__(self) -> None: + super().__init__() + self._store = InMemoryStore() + + def batch(self, ops: Iterable[Op]) -> list[Result]: + return self._store.batch(ops) + + async def abatch(self, ops: Iterable[Op]) -> list[Result]: + return self._store.batch(ops) + + async def test_async_batch_store(mocker: MockerFixture) -> None: abatch = mocker.stub() @@ -313,17 +325,6 @@ async def test_cannot_put_empty_namespace() -> None: store.delete(("langgraph", "foo"), "bar") assert store.get(("langgraph", "foo"), "bar") is None - class MockAsyncBatchedStore(AsyncBatchedBaseStore): - def __init__(self) -> None: - super().__init__() - self._store = InMemoryStore() - - def batch(self, ops: Iterable[Op]) -> list[Result]: - return self._store.batch(ops) - - async def abatch(self, ops: Iterable[Op]) -> list[Result]: - return self._store.batch(ops) - async_store = MockAsyncBatchedStore() doc = {"foo": "bar"} @@ -354,3 +355,68 @@ async def test_cannot_put_empty_namespace() -> None: assert (await async_store.asearch(("valid", "namespace")))[0].value == doc await async_store.adelete(("valid", "namespace"), "key") assert (await async_store.aget(("valid", "namespace"), "key")) is None + + +async def test_async_batch_store_deduplication(mocker: MockerFixture) -> None: + abatch = mocker.spy(InMemoryStore, "batch") + store = MockAsyncBatchedStore() + + same_doc = {"value": "same"} + diff_doc = {"value": "different"} + await asyncio.gather( + store.aput(namespace=("test",), key="same", value=same_doc), + store.aput(namespace=("test",), key="different", value=diff_doc), + ) + abatch.reset_mock() + + results = await asyncio.gather( + store.aget(namespace=("test",), key="same"), + store.aget(namespace=("test",), key="same"), + store.aget(namespace=("test",), key="different"), + ) + + assert len(results) == 3 + assert results[0] == results[1] + assert results[0] != results[2] + assert results[0].value == same_doc # type: ignore + assert results[2].value == diff_doc # type: ignore + assert len(abatch.call_args_list) == 1 + ops = list(abatch.call_args_list[0].args[1]) + assert len(ops) == 2 + assert GetOp(("test",), "same") in ops + assert GetOp(("test",), "different") in ops + + abatch.reset_mock() + + doc1 = {"value": 1} + doc2 = {"value": 2} + results = await asyncio.gather( + store.aput(namespace=("test",), key="key", value=doc1), + store.aput(namespace=("test",), key="key", value=doc2), + ) + assert len(abatch.call_args_list) == 1 + ops = list(abatch.call_args_list[0].args[1]) + assert len(ops) == 1 + assert ops[0] == PutOp(("test",), "key", doc2) + assert len(results) == 2 + assert all(result is None for result in results) + + result = await store.aget(namespace=("test",), key="key") + assert result is not None + assert result.value == doc2 + + abatch.reset_mock() + + results = await asyncio.gather( + store.asearch(("test",), filter={"value": 2}), + store.asearch(("test",), filter={"value": 2}), + ) + assert len(abatch.call_args_list) == 1 + ops = list(abatch.call_args_list[0].args[1]) + assert len(ops) == 1 + assert len(results) == 2 + assert results[0] == results[1] + assert len(results[0]) == 1 + assert results[0][0].value == doc2 + + abatch.reset_mock() diff --git a/libs/langgraph/langgraph/errors.py b/libs/langgraph/langgraph/errors.py index 2450b42b1..0737a31d0 100644 --- a/libs/langgraph/langgraph/errors.py +++ b/libs/langgraph/langgraph/errors.py @@ -2,7 +2,7 @@ from enum import Enum from typing import Any, Sequence from langgraph.checkpoint.base import EmptyChannelError # noqa: F401 -from langgraph.types import Interrupt +from langgraph.types import Command, Interrupt # EmptyChannelError re-exported for backwards compatibility @@ -58,7 +58,11 @@ class InvalidUpdateError(Exception): pass -class GraphInterrupt(Exception): +class GraphBubbleUp(Exception): + pass + + +class GraphInterrupt(GraphBubbleUp): """Raised when a subgraph is interrupted, suppressed by the root graph. Never raised directly, or surfaced to the user.""" @@ -73,13 +77,20 @@ class NodeInterrupt(GraphInterrupt): super().__init__([Interrupt(value=value)]) -class GraphDelegate(Exception): +class GraphDelegate(GraphBubbleUp): """Raised when a graph is delegated (for distributed mode).""" def __init__(self, *args: dict[str, Any]) -> None: super().__init__(*args) +class ParentCommand(GraphBubbleUp): + args: tuple[Command] + + def __init__(self, command: Command) -> None: + super().__init__(command) + + class EmptyInputError(Exception): """Raised when graph receives an empty input.""" diff --git a/libs/langgraph/langgraph/graph/state.py b/libs/langgraph/langgraph/graph/state.py index cce742911..0684fb29c 100644 --- a/libs/langgraph/langgraph/graph/state.py +++ b/libs/langgraph/langgraph/graph/state.py @@ -37,7 +37,12 @@ from langgraph.channels.ephemeral_value import EphemeralValue from langgraph.channels.last_value import LastValue from langgraph.channels.named_barrier_value import NamedBarrierValue from langgraph.constants import EMPTY_SEQ, NS_END, NS_SEP, SELF, TAG_HIDDEN -from langgraph.errors import ErrorCode, InvalidUpdateError, create_error_message +from langgraph.errors import ( + ErrorCode, + InvalidUpdateError, + ParentCommand, + create_error_message, +) from langgraph.graph.graph import END, START, Branch, CompiledGraph, Graph, Send from langgraph.managed.base import ( ChannelKeyPlaceholder, @@ -623,6 +628,8 @@ class CompiledStateGraph(CompiledGraph): def _get_root(input: Any) -> Any: if isinstance(input, Command): + if input.graph == Command.PARENT: + return SKIP_WRITE return input.update else: return input @@ -640,6 +647,8 @@ class CompiledStateGraph(CompiledGraph): ) return input.get(key, SKIP_WRITE) elif isinstance(input, Command): + if input.graph == Command.PARENT: + return SKIP_WRITE return _get_state_key(input.update, key=key) elif get_type_hints(type(input)): value = getattr(input, key, SKIP_WRITE) @@ -822,6 +831,8 @@ def _control_branch(value: Any) -> Sequence[Union[str, Send]]: return [value] if not isinstance(value, GraphCommand): return EMPTY_SEQ + if value.graph == Command.PARENT: + raise ParentCommand(value) rtn: list[Union[str, Send]] = [] if isinstance(value.goto, str): rtn.append(value.goto) @@ -839,6 +850,8 @@ async def _acontrol_branch(value: Any) -> Sequence[Union[str, Send]]: return [value] if not isinstance(value, GraphCommand): return EMPTY_SEQ + if value.graph == Command.PARENT: + raise ParentCommand(value) rtn: list[Union[str, Send]] = [] if isinstance(value.goto, str): rtn.append(value.goto) diff --git a/libs/langgraph/langgraph/prebuilt/tool_node.py b/libs/langgraph/langgraph/prebuilt/tool_node.py index cdcbdd819..1ea0dd56c 100644 --- a/libs/langgraph/langgraph/prebuilt/tool_node.py +++ b/libs/langgraph/langgraph/prebuilt/tool_node.py @@ -37,7 +37,7 @@ from langchain_core.tools import tool as create_tool from langchain_core.tools.base import get_all_basemodel_annotations from typing_extensions import Annotated, get_args, get_origin -from langgraph.errors import GraphInterrupt +from langgraph.errors import GraphBubbleUp from langgraph.store.base import BaseStore from langgraph.utils.runnable import RunnableCallable @@ -275,7 +275,7 @@ class ToolNode(RunnableCallable): # (2) a NodeInterrupt is raised inside a graph node for a graph called as a tool # (3) a GraphInterrupt is raised when a subgraph is interrupted inside a graph called as a tool # (2 and 3 can happen in a "supervisor w/ tools" multi-agent architecture) - except GraphInterrupt as e: + except GraphBubbleUp as e: raise e except Exception as e: if isinstance(self.handle_tool_errors, tuple): @@ -316,7 +316,7 @@ class ToolNode(RunnableCallable): # (2) a NodeInterrupt is raised inside a graph node for a graph called as a tool # (3) a GraphInterrupt is raised when a subgraph is interrupted inside a graph called as a tool # (2 and 3 can happen in a "supervisor w/ tools" multi-agent architecture) - except GraphInterrupt as e: + except GraphBubbleUp as e: raise e except Exception as e: if isinstance(self.handle_tool_errors, tuple): diff --git a/libs/langgraph/langgraph/pregel/algo.py b/libs/langgraph/langgraph/pregel/algo.py index 564c53022..1410e432f 100644 --- a/libs/langgraph/langgraph/pregel/algo.py +++ b/libs/langgraph/langgraph/pregel/algo.py @@ -602,6 +602,7 @@ def prepare_single_task( None, task_id, task_path, + writers=proc.flat_writers, ) else: @@ -720,6 +721,7 @@ def prepare_single_task( None, task_id, task_path, + writers=proc.flat_writers, ) else: return PregelTask(task_id, name, task_path) diff --git a/libs/langgraph/langgraph/pregel/executor.py b/libs/langgraph/langgraph/pregel/executor.py index 246510fb4..70aea29e3 100644 --- a/libs/langgraph/langgraph/pregel/executor.py +++ b/libs/langgraph/langgraph/pregel/executor.py @@ -20,7 +20,7 @@ from langchain_core.runnables import RunnableConfig from langchain_core.runnables.config import get_executor_for_config from typing_extensions import ParamSpec -from langgraph.errors import GraphInterrupt +from langgraph.errors import GraphBubbleUp P = ParamSpec("P") T = TypeVar("T") @@ -68,7 +68,7 @@ class BackgroundExecutor(ContextManager): def done(self, task: concurrent.futures.Future) -> None: try: task.result() - except GraphInterrupt: + except GraphBubbleUp: # This exception is an interruption signal, not an error # so we don't want to re-raise it on exit self.tasks.pop(task) @@ -155,7 +155,7 @@ class AsyncBackgroundExecutor(AsyncContextManager): if exc := task.exception(): # This exception is an interruption signal, not an error # so we don't want to re-raise it on exit - if isinstance(exc, GraphInterrupt): + if isinstance(exc, GraphBubbleUp): self.tasks.pop(task) else: self.tasks.pop(task) diff --git a/libs/langgraph/langgraph/pregel/io.py b/libs/langgraph/langgraph/pregel/io.py index 6695e1ce0..693dffce2 100644 --- a/libs/langgraph/langgraph/pregel/io.py +++ b/libs/langgraph/langgraph/pregel/io.py @@ -15,6 +15,7 @@ from langgraph.constants import ( TAG_HIDDEN, TASKS, ) +from langgraph.errors import InvalidUpdateError from langgraph.pregel.log import logger from langgraph.types import Command, PregelExecutableTask, Send @@ -68,6 +69,8 @@ def map_command( cmd: Command, ) -> 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 diff --git a/libs/langgraph/langgraph/pregel/retry.py b/libs/langgraph/langgraph/pregel/retry.py index ea9162dc2..6e52a7c41 100644 --- a/libs/langgraph/langgraph/pregel/retry.py +++ b/libs/langgraph/langgraph/pregel/retry.py @@ -2,6 +2,7 @@ import asyncio import logging import random import time +from dataclasses import replace from functools import partial from typing import Any, Callable, Optional, Sequence @@ -10,9 +11,10 @@ from langgraph.constants import ( CONFIG_KEY_CHECKPOINT_NS, CONFIG_KEY_RESUMING, CONFIG_KEY_SEND, + NS_SEP, ) -from langgraph.errors import _SEEN_CHECKPOINT_NS, GraphInterrupt -from langgraph.types import PregelExecutableTask, RetryPolicy +from langgraph.errors import _SEEN_CHECKPOINT_NS, GraphBubbleUp, ParentCommand +from langgraph.types import Command, PregelExecutableTask, RetryPolicy from langgraph.utils.config import patch_configurable logger = logging.getLogger(__name__) @@ -40,7 +42,21 @@ def run_with_retry( task.proc.invoke(task.input, config) # if successful, end break - except GraphInterrupt: + except ParentCommand as exc: + ns: str = config[CONF][CONFIG_KEY_CHECKPOINT_NS] + cmd = exc.args[0] + if cmd.graph == ns: + # this command is for the current graph, handle it + for w in task.writers: + w.invoke(cmd, config) + break + elif cmd.graph == Command.PARENT: + # this command is for the parent graph, assign it to the parent + parent_ns = NS_SEP.join(ns.split(NS_SEP)[:-1]) + exc.args = (replace(cmd, graph=parent_ns),) + # bubble up + raise + except GraphBubbleUp: # if interrupted, end raise except Exception as exc: @@ -118,7 +134,21 @@ async def arun_with_retry( await task.proc.ainvoke(task.input, config) # if successful, end break - except GraphInterrupt: + except ParentCommand as exc: + ns: str = config[CONF][CONFIG_KEY_CHECKPOINT_NS] + cmd = exc.args[0] + if cmd.graph == ns: + # this command is for the current graph, handle it + for w in task.writers: + w.invoke(cmd, config) + break + elif cmd.graph == Command.PARENT: + # this command is for the parent graph, assign it to the parent + parent_ns = NS_SEP.join(ns.split(NS_SEP)[:-1]) + exc.args = (replace(cmd, graph=parent_ns),) + # bubble up + raise + except GraphBubbleUp: # if interrupted, end raise except Exception as exc: diff --git a/libs/langgraph/langgraph/pregel/runner.py b/libs/langgraph/langgraph/pregel/runner.py index 64e5c8d3c..9e3879b0f 100644 --- a/libs/langgraph/langgraph/pregel/runner.py +++ b/libs/langgraph/langgraph/pregel/runner.py @@ -23,7 +23,7 @@ from langgraph.constants import ( PUSH, TAG_HIDDEN, ) -from langgraph.errors import GraphDelegate, GraphInterrupt +from langgraph.errors import GraphBubbleUp, GraphInterrupt from langgraph.pregel.executor import Submit from langgraph.pregel.retry import arun_with_retry, run_with_retry from langgraph.types import PregelExecutableTask, RetryPolicy @@ -298,7 +298,7 @@ class PregelRunner: # save interrupt to checkpointer if interrupts := [(INTERRUPT, i) for i in exception.args[0]]: self.put_writes(task.id, interrupts) - elif isinstance(exception, GraphDelegate): + elif isinstance(exception, GraphBubbleUp): raise exception else: # save error to checkpointer @@ -324,7 +324,7 @@ def _should_stop_others( if fut.cancelled(): return True if exc := fut.exception(): - return not isinstance(exc, GraphInterrupt) + return not isinstance(exc, GraphBubbleUp) else: return False diff --git a/libs/langgraph/langgraph/types.py b/libs/langgraph/langgraph/types.py index 104412d8e..7bf9148c5 100644 --- a/libs/langgraph/langgraph/types.py +++ b/libs/langgraph/langgraph/types.py @@ -5,6 +5,7 @@ from typing import ( TYPE_CHECKING, Any, Callable, + ClassVar, Generic, Hashable, Literal, @@ -140,6 +141,7 @@ class PregelExecutableTask(NamedTuple): id: str path: tuple[Union[str, int, tuple], ...] scheduled: bool = False + writers: Sequence[Runnable] = () class StateSnapshot(NamedTuple): @@ -239,6 +241,7 @@ N = TypeVar("N", bound=Hashable) class Command(Generic[N]): """One or more commands to update the graph's state and send messages to nodes.""" + graph: Optional[str] = None update: Optional[dict[str, Any]] = None send: Union[Send, Sequence[Send]] = () resume: Optional[Union[Any, dict[str, Any]]] = None @@ -252,6 +255,8 @@ class Command(Generic[N]): ) return f"Command({contents})" + PARENT: ClassVar[Literal["__parent__"]] = "__parent__" + StreamChunk = tuple[tuple[str, ...], str, Any] diff --git a/libs/langgraph/tests/test_pregel.py b/libs/langgraph/tests/test_pregel.py index fcee1fa20..c2ed63d28 100644 --- a/libs/langgraph/tests/test_pregel.py +++ b/libs/langgraph/tests/test_pregel.py @@ -14395,3 +14395,79 @@ def test_runnable_passthrough_node_graph() -> None: graph = graph_builder.compile() assert graph.get_graph(xray=True).to_json() == graph.get_graph(xray=False).to_json() + + +@pytest.mark.parametrize("checkpointer_name", ALL_CHECKPOINTERS_SYNC) +def test_parent_command(request: pytest.FixtureRequest, checkpointer_name: str) -> None: + from langchain_core.messages import BaseMessage + from langchain_core.tools import tool + + @tool(return_direct=True) + def get_user_name() -> GraphCommand: + """Retrieve user name""" + return GraphCommand(update={"user_name": "Meow"}, graph=GraphCommand.PARENT) + + subgraph_builder = StateGraph(MessagesState) + subgraph_builder.add_node("tool", get_user_name) + subgraph_builder.add_edge(START, "tool") + subgraph = subgraph_builder.compile() + + 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") + checkpointer = request.getfixturevalue(f"checkpointer_{checkpointer_name}") + graph = builder.compile(checkpointer=checkpointer) + + config = {"configurable": {"thread_id": "1"}} + + assert graph.invoke({"messages": [("user", "get user name")]}, config) == { + "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", + "writes": { + "alice": { + "user_name": "Meow", + } + }, + "thread_id": "1", + "step": 1, + "parents": {}, + }, + created_at=AnyStr(), + parent_config={ + "configurable": { + "thread_id": "1", + "checkpoint_ns": "", + "checkpoint_id": AnyStr(), + } + }, + tasks=(), + ) diff --git a/libs/langgraph/tests/test_pregel_async.py b/libs/langgraph/tests/test_pregel_async.py index c5812e896..813fa8240 100644 --- a/libs/langgraph/tests/test_pregel_async.py +++ b/libs/langgraph/tests/test_pregel_async.py @@ -12597,3 +12597,83 @@ async def test_debug_nested_subgraphs(): assert stream_task["interrupts"] == history_task.interrupts assert stream_task.get("error") == history_task.error assert stream_task.get("state") == history_task.state + + +@pytest.mark.parametrize("checkpointer_name", ALL_CHECKPOINTERS_ASYNC) +async def test_parent_command(checkpointer_name: str) -> None: + from langchain_core.messages import BaseMessage + from langchain_core.tools import tool + + @tool(return_direct=True) + def get_user_name() -> GraphCommand: + """Retrieve user name""" + return GraphCommand(update={"user_name": "Meow"}, graph=GraphCommand.PARENT) + + subgraph_builder = StateGraph(MessagesState) + subgraph_builder.add_node("tool", get_user_name) + subgraph_builder.add_edge(START, "tool") + subgraph = subgraph_builder.compile() + + 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") + async with awith_checkpointer(checkpointer_name) as checkpointer: + graph = builder.compile(checkpointer=checkpointer) + + config = {"configurable": {"thread_id": "1"}} + + assert await graph.ainvoke( + {"messages": [("user", "get user name")]}, config + ) == { + "messages": [ + _AnyIdHumanMessage( + content="get user name", additional_kwargs={}, response_metadata={} + ), + ], + "user_name": "Meow", + } + assert await graph.aget_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", + "writes": { + "alice": { + "user_name": "Meow", + } + }, + "thread_id": "1", + "step": 1, + "parents": {}, + }, + created_at=AnyStr(), + parent_config={ + "configurable": { + "thread_id": "1", + "checkpoint_ns": "", + "checkpoint_id": AnyStr(), + } + }, + tasks=(), + )