Merge pull request #2048 from langchain-ai/dqbd/debug-tasks-state

feat(debug): send tasks info

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:

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)
             ],
         },
     }

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

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"

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

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