diff --git a/libs/langgraph/langgraph/pregel/__init__.py b/libs/langgraph/langgraph/pregel/__init__.py index 3d7414494..98d82712b 100644 --- a/libs/langgraph/langgraph/pregel/__init__.py +++ b/libs/langgraph/langgraph/pregel/__init__.py @@ -26,7 +26,6 @@ from uuid import UUID, uuid5 from langchain_core.globals import get_debug from langchain_core.runnables import ( Runnable, - RunnableLambda, RunnableSequence, ) from langchain_core.runnables.base import Input, Output @@ -37,7 +36,6 @@ from langchain_core.runnables.config import ( ) from langchain_core.runnables.utils import ( ConfigurableFieldSpec, - get_function_nonlocals, get_unique_config_specs, ) from langchain_core.tracers._streaming import _StreamingCallbackHandler @@ -86,7 +84,7 @@ from langgraph.pregel.messages import StreamMessagesHandler from langgraph.pregel.read import PregelNode from langgraph.pregel.retry import RetryPolicy from langgraph.pregel.runner import PregelRunner -from langgraph.pregel.utils import get_new_channel_versions +from langgraph.pregel.utils import find_subgraph_pregel, get_new_channel_versions from langgraph.pregel.validate import validate_graph, validate_keys from langgraph.pregel.write import ChannelWrite, ChannelWriteEntry from langgraph.store.base import BaseStore @@ -100,7 +98,6 @@ from langgraph.utils.config import ( ) from langgraph.utils.pydantic import create_model from langgraph.utils.queue import AsyncQueue, SyncQueue # type: ignore[attr-defined] -from langgraph.utils.runnable import RunnableCallable WriteValue = Union[Callable[[Input], Output], Any] @@ -391,32 +388,10 @@ class Pregel(Runnable[Union[dict[str, Any], Any], Union[dict[str, Any], Any]]): if namespace is not None: if not namespace.startswith(name): continue + # find the subgraph, if any - graph: Optional[Pregel] = None - candidates = [node.bound] - for candidate in candidates: - if ( - isinstance(candidate, Pregel) - # subgraphs that disabled checkpointing are not considered - and candidate.checkpointer is not False - ): - graph = candidate - break - elif isinstance(candidate, RunnableSequence): - candidates.extend(candidate.steps) - elif isinstance(candidate, RunnableLambda): - candidates.extend(candidate.deps) - elif isinstance(candidate, RunnableCallable): - if candidate.func is not None: - candidates.extend( - nl.__self__ if hasattr(nl, "__self__") else nl - for nl in get_function_nonlocals(candidate.func) - ) - if candidate.afunc is not None: - candidates.extend( - nl.__self__ if hasattr(nl, "__self__") else nl - for nl in get_function_nonlocals(candidate.afunc) - ) + graph = cast(Optional[Pregel], find_subgraph_pregel(node.bound)) + # if found, yield recursively if graph: if name == namespace: diff --git a/libs/langgraph/langgraph/pregel/debug.py b/libs/langgraph/langgraph/pregel/debug.py index 7e8f9546f..53d49f7e1 100644 --- a/libs/langgraph/langgraph/pregel/debug.py +++ b/libs/langgraph/langgraph/pregel/debug.py @@ -20,8 +20,17 @@ from langchain_core.utils.input import get_bolded_text, get_colored_text from langgraph.channels.base import BaseChannel from langgraph.checkpoint.base import Checkpoint, CheckpointMetadata, PendingWrite -from langgraph.constants import ERROR, INTERRUPT, TAG_HIDDEN +from langgraph.constants import ( + CONF, + CONFIG_KEY_CHECKPOINT_NS, + ERROR, + INTERRUPT, + NS_END, + NS_SEP, + TAG_HIDDEN, +) from langgraph.pregel.io import read_channels +from langgraph.pregel.utils import find_subgraph_pregel from langgraph.types import PregelExecutableTask, PregelTask, StateSnapshot @@ -45,6 +54,7 @@ class CheckpointTask(TypedDict): name: str error: Optional[str] interrupts: list[dict] + state: Optional[RunnableConfig] class CheckpointPayload(TypedDict): @@ -140,6 +150,27 @@ def map_debug_checkpoint( parent_config: Optional[RunnableConfig], ) -> Iterator[DebugOutputCheckpoint]: """Produce "checkpoint" events for stream_mode=debug.""" + + parent_ns = config[CONF].get(CONFIG_KEY_CHECKPOINT_NS, "") + task_states: dict[str, Union[RunnableConfig, StateSnapshot]] = {} + + for task in tasks: + if not find_subgraph_pregel(task.proc): + continue + + # assemble checkpoint_ns for this task + task_ns = f"{task.name}{NS_END}{task.id}" + if parent_ns: + task_ns = f"{parent_ns}{NS_SEP}{task_ns}" + + # set config as signal that subgraph checkpoints exist + task_states[task.id] = { + CONF: { + "thread_id": config[CONF]["thread_id"], + CONFIG_KEY_CHECKPOINT_NS: task_ns, + } + } + yield { "type": "checkpoint", "timestamp": checkpoint["ts"], @@ -155,14 +186,16 @@ def map_debug_checkpoint( "id": t.id, "name": t.name, "error": t.error, + "state": t.state, } if t.error else { "id": t.id, "name": t.name, "interrupts": tuple(asdict(i) for i in t.interrupts), + "state": t.state, } - for t in tasks_w_writes(tasks, pending_writes, None) + for t in tasks_w_writes(tasks, pending_writes, task_states) ], }, } diff --git a/libs/langgraph/langgraph/pregel/utils.py b/libs/langgraph/langgraph/pregel/utils.py index 3a29e5ed1..2b09f8f75 100644 --- a/libs/langgraph/langgraph/pregel/utils.py +++ b/libs/langgraph/langgraph/pregel/utils.py @@ -1,4 +1,10 @@ +from typing import Optional + +from langchain_core.runnables import RunnableLambda, RunnableSequence +from langchain_core.runnables.utils import get_function_nonlocals + from langgraph.checkpoint.base import ChannelVersions +from langgraph.utils.runnable import Runnable, RunnableCallable, RunnableSeq def get_new_channel_versions( @@ -17,3 +23,34 @@ def get_new_channel_versions( new_versions = current_versions return new_versions + + +def find_subgraph_pregel(candidate: Runnable) -> Optional[Runnable]: + from langgraph.pregel import Pregel + + candidates: list[Runnable] = [candidate] + + for c in candidates: + if ( + isinstance(c, Pregel) + # subgraphs that disabled checkpointing are not considered + and c.checkpointer is not False + ): + return c + elif isinstance(c, RunnableSequence) or isinstance(c, RunnableSeq): + candidates.extend(c.steps) + elif isinstance(c, RunnableLambda): + candidates.extend(c.deps) + elif isinstance(c, RunnableCallable): + if c.func is not None: + candidates.extend( + nl.__self__ if hasattr(nl, "__self__") else nl + for nl in get_function_nonlocals(c.func) + ) + if c.afunc is not None: + candidates.extend( + nl.__self__ if hasattr(nl, "__self__") else nl + for nl in get_function_nonlocals(c.afunc) + ) + + return None diff --git a/libs/langgraph/pyproject.toml b/libs/langgraph/pyproject.toml index 4a90b63d3..8c5efd8d9 100644 --- a/libs/langgraph/pyproject.toml +++ b/libs/langgraph/pyproject.toml @@ -1,6 +1,6 @@ [tool.poetry] name = "langgraph" -version = "0.2.34" +version = "0.2.35" description = "Building stateful, multi-actor applications with LLMs" authors = [] license = "MIT" diff --git a/libs/langgraph/tests/test_pregel.py b/libs/langgraph/tests/test_pregel.py index 196fe4223..e75df2245 100644 --- a/libs/langgraph/tests/test_pregel.py +++ b/libs/langgraph/tests/test_pregel.py @@ -6996,11 +6996,14 @@ def test_branch_then( # test stream_mode=debug tool_two = tool_two_graph.compile(checkpointer=checkpointer) thread10 = {"configurable": {"thread_id": "10"}} - assert [ + + res = [ *tool_two.stream( {"my_key": "value", "market": "DE"}, thread10, stream_mode="debug" ) - ] == [ + ] + + assert res == [ { "type": "checkpoint", "timestamp": AnyStr(), @@ -7026,7 +7029,14 @@ def test_branch_then( }, "parent_config": None, "next": ["__start__"], - "tasks": [{"id": AnyStr(), "name": "__start__", "interrupts": ()}], + "tasks": [ + { + "id": AnyStr(), + "name": "__start__", + "interrupts": (), + "state": None, + } + ], }, }, { @@ -7067,7 +7077,9 @@ def test_branch_then( }, }, "next": ["prepare"], - "tasks": [{"id": AnyStr(), "name": "prepare", "interrupts": ()}], + "tasks": [ + {"id": AnyStr(), "name": "prepare", "interrupts": (), "state": None} + ], }, }, { @@ -7131,7 +7143,14 @@ def test_branch_then( }, }, "next": ["tool_two_slow"], - "tasks": [{"id": AnyStr(), "name": "tool_two_slow", "interrupts": ()}], + "tasks": [ + { + "id": AnyStr(), + "name": "tool_two_slow", + "interrupts": (), + "state": None, + } + ], }, }, { @@ -7195,7 +7214,9 @@ def test_branch_then( }, }, "next": ["finish"], - "tasks": [{"id": AnyStr(), "name": "finish", "interrupts": ()}], + "tasks": [ + {"id": AnyStr(), "name": "finish", "interrupts": (), "state": None} + ], }, }, { @@ -11561,3 +11582,167 @@ def test_enum_node_names(): graph = graph.compile() assert graph.invoke({"foo": "hello"}) == {"foo": "hello", "bar": "hello!"} + + +def test_debug_subgraphs(): + class State(TypedDict): + messages: Annotated[list[str], operator.add] + + def node(name): + def _node(state: State): + return {"messages": [f"entered {name} node"]} + + return _node + + parent = StateGraph(State) + child = StateGraph(State) + + child.add_node("c_one", node("c_one")) + child.add_node("c_two", node("c_two")) + child.add_edge(START, "c_one") + child.add_edge("c_one", "c_two") + child.add_edge("c_two", END) + + parent.add_node("p_one", node("p_one")) + parent.add_node("p_two", child.compile()) + parent.add_edge(START, "p_one") + parent.add_edge("p_one", "p_two") + parent.add_edge("p_two", END) + + graph = parent.compile(checkpointer=MemorySaver()) + + config = {"configurable": {"thread_id": "1"}} + events = [ + *graph.stream( + {"messages": []}, + config=config, + stream_mode="debug", + ) + ] + + checkpoint_events = list( + reversed([e["payload"] for e in events if e["type"] == "checkpoint"]) + ) + checkpoint_history = list(graph.get_state_history(config)) + + assert len(checkpoint_events) == len(checkpoint_history) + + def normalize_config(config: Optional[dict]) -> Optional[dict]: + if config is None: + return None + return config["configurable"] + + for stream, history in zip(checkpoint_events, checkpoint_history): + assert stream["values"] == history.values + assert stream["next"] == list(history.next) + assert normalize_config(stream["config"]) == normalize_config(history.config) + assert normalize_config(stream["parent_config"]) == normalize_config( + history.parent_config + ) + + assert len(stream["tasks"]) == len(history.tasks) + for stream_task, history_task in zip(stream["tasks"], history.tasks): + assert stream_task["id"] == history_task.id + assert stream_task["name"] == history_task.name + assert stream_task["interrupts"] == history_task.interrupts + assert stream_task.get("error") == history_task.error + assert stream_task.get("state") == history_task.state + + +def test_debug_nested_subgraphs(): + from collections import defaultdict + + class State(TypedDict): + messages: Annotated[list[str], operator.add] + + def node(name): + def _node(state: State): + return {"messages": [f"entered {name} node"]} + + return _node + + grand_parent = StateGraph(State) + parent = StateGraph(State) + child = StateGraph(State) + + child.add_node("c_one", node("c_one")) + child.add_node("c_two", node("c_two")) + child.add_edge(START, "c_one") + child.add_edge("c_one", "c_two") + child.add_edge("c_two", END) + + parent.add_node("p_one", node("p_one")) + parent.add_node("p_two", child.compile()) + parent.add_edge(START, "p_one") + parent.add_edge("p_one", "p_two") + parent.add_edge("p_two", END) + + grand_parent.add_node("gp_one", node("gp_one")) + grand_parent.add_node("gp_two", parent.compile()) + grand_parent.add_edge(START, "gp_one") + grand_parent.add_edge("gp_one", "gp_two") + grand_parent.add_edge("gp_two", END) + + graph = grand_parent.compile(checkpointer=MemorySaver()) + + config = {"configurable": {"thread_id": "1"}} + events = [ + *graph.stream( + {"messages": []}, + config=config, + stream_mode="debug", + subgraphs=True, + ) + ] + + stream_ns: dict[tuple, dict] = defaultdict(list) + for ns, e in events: + if e["type"] == "checkpoint": + stream_ns[ns].append(e["payload"]) + + assert list(stream_ns.keys()) == [ + (), + (AnyStr("gp_two:"),), + (AnyStr("gp_two:"), AnyStr("p_two:")), + ] + + history_ns = { + ns: list( + graph.get_state_history( + {"configurable": {"thread_id": "1", "checkpoint_ns": "|".join(ns)}} + ) + )[::-1] + for ns in stream_ns.keys() + } + + def normalize_config(config: Optional[dict]) -> Optional[dict]: + if config is None: + return None + + clean_config = {} + clean_config["thread_id"] = config["configurable"]["thread_id"] + clean_config["checkpoint_id"] = config["configurable"]["checkpoint_id"] + clean_config["checkpoint_ns"] = config["configurable"]["checkpoint_ns"] + + return clean_config + + for checkpoint_events, checkpoint_history in zip( + stream_ns.values(), history_ns.values() + ): + for stream, history in zip(checkpoint_events, checkpoint_history): + assert stream["values"] == history.values + assert stream["next"] == list(history.next) + assert normalize_config(stream["config"]) == normalize_config( + history.config + ) + assert normalize_config(stream["parent_config"]) == normalize_config( + history.parent_config + ) + + assert len(stream["tasks"]) == len(history.tasks) + for stream_task, history_task in zip(stream["tasks"], history.tasks): + assert stream_task["id"] == history_task.id + assert stream_task["name"] == history_task.name + assert stream_task["interrupts"] == history_task.interrupts + assert stream_task.get("error") == history_task.error + assert stream_task.get("state") == history_task.state diff --git a/libs/langgraph/tests/test_pregel_async.py b/libs/langgraph/tests/test_pregel_async.py index 0290bf069..3be52643b 100644 --- a/libs/langgraph/tests/test_pregel_async.py +++ b/libs/langgraph/tests/test_pregel_async.py @@ -5655,7 +5655,14 @@ async def test_branch_then(checkpointer_name: str) -> None: }, "parent_config": None, "next": ["__start__"], - "tasks": [{"id": AnyStr(), "name": "__start__", "interrupts": ()}], + "tasks": [ + { + "id": AnyStr(), + "name": "__start__", + "interrupts": (), + "state": None, + } + ], }, }, { @@ -5696,7 +5703,14 @@ async def test_branch_then(checkpointer_name: str) -> None: }, }, "next": ["prepare"], - "tasks": [{"id": AnyStr(), "name": "prepare", "interrupts": ()}], + "tasks": [ + { + "id": AnyStr(), + "name": "prepare", + "interrupts": (), + "state": None, + } + ], }, }, { @@ -5761,7 +5775,12 @@ async def test_branch_then(checkpointer_name: str) -> None: }, "next": ["tool_two_slow"], "tasks": [ - {"id": AnyStr(), "name": "tool_two_slow", "interrupts": ()} + { + "id": AnyStr(), + "name": "tool_two_slow", + "interrupts": (), + "state": None, + } ], }, }, @@ -5826,7 +5845,14 @@ async def test_branch_then(checkpointer_name: str) -> None: }, }, "next": ["finish"], - "tasks": [{"id": AnyStr(), "name": "finish", "interrupts": ()}], + "tasks": [ + { + "id": AnyStr(), + "name": "finish", + "interrupts": (), + "state": None, + } + ], }, }, { @@ -5937,7 +5963,14 @@ async def test_branch_then(checkpointer_name: str) -> None: }, "parent_config": None, "next": ["__start__"], - "tasks": [{"id": AnyStr(), "name": "__start__", "interrupts": ()}], + "tasks": [ + { + "id": AnyStr(), + "name": "__start__", + "interrupts": (), + "state": None, + } + ], }, }, { @@ -5978,7 +6011,14 @@ async def test_branch_then(checkpointer_name: str) -> None: }, }, "next": ["prepare"], - "tasks": [{"id": AnyStr(), "name": "prepare", "interrupts": ()}], + "tasks": [ + { + "id": AnyStr(), + "name": "prepare", + "interrupts": (), + "state": None, + } + ], }, }, { @@ -6043,7 +6083,12 @@ async def test_branch_then(checkpointer_name: str) -> None: }, "next": ["tool_two_slow"], "tasks": [ - {"id": AnyStr(), "name": "tool_two_slow", "interrupts": ()} + { + "id": AnyStr(), + "name": "tool_two_slow", + "interrupts": (), + "state": None, + } ], }, }, @@ -9773,3 +9818,174 @@ async def test_store_injected_async(checkpointer_name: str, store_name: str) -> assert ( len((await the_store.asearch(("foo", "bar")))) == 1 ) # still overwriting the same one + + +async def test_debug_subgraphs(): + class State(TypedDict): + messages: Annotated[list[str], operator.add] + + def node(name): + async def _node(state: State): + return {"messages": [f"entered {name} node"]} + + return _node + + parent = StateGraph(State) + child = StateGraph(State) + + child.add_node("c_one", node("c_one")) + child.add_node("c_two", node("c_two")) + child.add_edge(START, "c_one") + child.add_edge("c_one", "c_two") + child.add_edge("c_two", END) + + parent.add_node("p_one", node("p_one")) + parent.add_node("p_two", child.compile()) + parent.add_edge(START, "p_one") + parent.add_edge("p_one", "p_two") + parent.add_edge("p_two", END) + + graph = parent.compile(checkpointer=MemorySaver()) + + config = {"configurable": {"thread_id": "1"}} + events = [ + c + async for c in graph.astream( + {"messages": []}, + config=config, + stream_mode="debug", + ) + ] + + checkpoint_events = list( + reversed([e["payload"] for e in events if e["type"] == "checkpoint"]) + ) + checkpoint_history = [c async for c in graph.aget_state_history(config)] + + assert len(checkpoint_events) == len(checkpoint_history) + + def normalize_config(config: Optional[dict]) -> Optional[dict]: + if config is None: + return None + return config["configurable"] + + for stream, history in zip(checkpoint_events, checkpoint_history): + assert stream["values"] == history.values + assert stream["next"] == list(history.next) + assert normalize_config(stream["config"]) == normalize_config(history.config) + assert normalize_config(stream["parent_config"]) == normalize_config( + history.parent_config + ) + + assert len(stream["tasks"]) == len(history.tasks) + for stream_task, history_task in zip(stream["tasks"], history.tasks): + assert stream_task["id"] == history_task.id + assert stream_task["name"] == history_task.name + assert stream_task["interrupts"] == history_task.interrupts + assert stream_task.get("error") == history_task.error + assert stream_task.get("state") == history_task.state + + +async def test_debug_nested_subgraphs(): + from collections import defaultdict + + class State(TypedDict): + messages: Annotated[list[str], operator.add] + + def node(name): + async def _node(state: State): + return {"messages": [f"entered {name} node"]} + + return _node + + grand_parent = StateGraph(State) + parent = StateGraph(State) + child = StateGraph(State) + + child.add_node("c_one", node("c_one")) + child.add_node("c_two", node("c_two")) + child.add_edge(START, "c_one") + child.add_edge("c_one", "c_two") + child.add_edge("c_two", END) + + parent.add_node("p_one", node("p_one")) + parent.add_node("p_two", child.compile()) + parent.add_edge(START, "p_one") + parent.add_edge("p_one", "p_two") + parent.add_edge("p_two", END) + + grand_parent.add_node("gp_one", node("gp_one")) + grand_parent.add_node("gp_two", parent.compile()) + grand_parent.add_edge(START, "gp_one") + grand_parent.add_edge("gp_one", "gp_two") + grand_parent.add_edge("gp_two", END) + + graph = grand_parent.compile(checkpointer=MemorySaver()) + + config = {"configurable": {"thread_id": "1"}} + events = [ + c + async for c in graph.astream( + {"messages": []}, + config=config, + stream_mode="debug", + subgraphs=True, + ) + ] + + stream_ns: dict[tuple, dict] = defaultdict(list) + for ns, e in events: + if e["type"] == "checkpoint": + stream_ns[ns].append(e["payload"]) + + assert list(stream_ns.keys()) == [ + (), + (AnyStr("gp_two:"),), + (AnyStr("gp_two:"), AnyStr("p_two:")), + ] + + history_ns = {} + for ns in stream_ns.keys(): + + async def get_history(): + history = [ + c + async for c in graph.aget_state_history( + {"configurable": {"thread_id": "1", "checkpoint_ns": "|".join(ns)}} + ) + ] + return history[::-1] + + history_ns[ns] = await get_history() + + def normalize_config(config: Optional[dict]) -> Optional[dict]: + if config is None: + return None + + clean_config = {} + clean_config["thread_id"] = config["configurable"]["thread_id"] + clean_config["checkpoint_id"] = config["configurable"]["checkpoint_id"] + clean_config["checkpoint_ns"] = config["configurable"]["checkpoint_ns"] + + return clean_config + + for checkpoint_events, checkpoint_history in zip( + stream_ns.values(), history_ns.values() + ): + for stream, history in zip(checkpoint_events, checkpoint_history): + assert stream["values"] == history.values + assert stream["next"] == list(history.next) + assert normalize_config(stream["config"]) == normalize_config( + history.config + ) + assert normalize_config(stream["parent_config"]) == normalize_config( + history.parent_config + ) + + assert len(stream["tasks"]) == len(history.tasks) + for stream_task, history_task in zip(stream["tasks"], history.tasks): + assert stream_task["id"] == history_task.id + assert stream_task["name"] == history_task.name + assert stream_task["interrupts"] == history_task.interrupts + assert stream_task.get("error") == history_task.error + assert stream_task.get("state") == history_task.state