mirror of
https://github.com/langchain-ai/langgraph.git
synced 2026-09-07 18:27:52 +02:00
fix(langgraph): fix graph rendering for defer=True (#6130)
### Description Some graphs with `defer=True` nodes rendered incorrectly. E.g.: * edge C2 -> E1 is missing and edge C2 -> END should not appear in #5772 * edge E3 -> END is missing and edge E -> END should not appear in #5182 * extra edge #5369 Fix: * Record the destinations declared by get_static_writes for each node. Build step_sources as a union of the runtime writes and the static writes (instead of just runtime writes). * Label deferred nodes with 'deferred' ### https://github.com/langchain-ai/langgraph/issues/5772 'Before' is how they were rendered before this PR | No defer | Before (defer `E1`) | After (defer `E1`) | -------- | ------- | ------- | | <img height="400" alt="defer_after" src="https://github.com/user-attachments/assets/0a9fc992-1b6a-4c6d-8752-de54c703c329" /> | <img height="400" alt="defer_before" src="https://github.com/user-attachments/assets/825b09fc-3fb8-461a-9928-20c8d9cfc533" /> | <img height="400" alt="defer_after" src="https://github.com/user-attachments/assets/ce5334f7-b469-47b0-8f1e-35bda2544a4e" /> | Before: * For deferred joins (NamedBarrierValueAfterFinish), a writer from an upstream node may not produce a runtime task.writes entry until the barrier opens. draw_graph() builds edges from task.writes, so one side of the join (here C2) never gets recorded as a source, and C2 is seen as a sink, so there is an implicit edge: C2 -> END edge added. After: * C2's write to the join channel is recorded even if the barrier hasn’t opened. When E1 finally schedules, we correctly find both sources B2 and C2 for the same trigger and emit edges: B2 -> E1 and C2 -> E1. With C2 -> E1 present, C2 is no longer a terminus, so the unexpected edge: C2 -> END is not added. ### Other graphs Graphs for the most part remain unchanged. See: ### #5182 | No defer | Before (defer `d`) | After (defer `d`) | -------- | ------- | ------- | | <img height="400" alt="defer_after" src="https://github.com/user-attachments/assets/3509d25c-f3ad-473c-b877-c155b8008cd5" /> | <img height="400" alt="defer_before" src="https://github.com/user-attachments/assets/7af38e77-eb70-414d-b8fe-667da943f9e0" /> | <img height="400" alt="defer_after" src="https://github.com/user-attachments/assets/bc87a19f-b4fb-42d3-a6ee-5b0982d9af71" /> | ### https://github.com/langchain-ai/langgraph/issues/5369 | No defer | Before (defer `595577`, `52642`) | After (defer `595577`, `52642`) | -------- | ------- | ------- | | <img height="400" alt="defer_after" src="https://github.com/user-attachments/assets/7c0824ce-3921-4dce-bc16-278f64289d28" /> | <img height="400" alt="defer_before" src="https://github.com/user-attachments/assets/28661079-7502-4912-874b-c086c0204a87" /> | <img height="400" alt="defer_after" src="https://github.com/user-attachments/assets/2a04956a-ed79-40e5-98b8-f6ecb2597a2e" /> |
This commit is contained in:
@@ -2,7 +2,7 @@ from __future__ import annotations
|
||||
|
||||
from collections import defaultdict
|
||||
from collections.abc import Mapping, Sequence
|
||||
from typing import Any, cast
|
||||
from typing import Any, NamedTuple, cast
|
||||
|
||||
from langchain_core.runnables.config import RunnableConfig
|
||||
from langchain_core.runnables.graph import Graph, Node
|
||||
@@ -10,6 +10,7 @@ from langgraph.checkpoint.base import BaseCheckpointSaver
|
||||
|
||||
from langgraph._internal._constants import CONF, CONFIG_KEY_SEND, INPUT
|
||||
from langgraph.channels.base import BaseChannel
|
||||
from langgraph.channels.last_value import LastValueAfterFinish
|
||||
from langgraph.constants import END, START
|
||||
from langgraph.managed.base import ManagedValueSpec
|
||||
from langgraph.pregel._algo import (
|
||||
@@ -25,6 +26,19 @@ from langgraph.pregel._write import ChannelWrite
|
||||
from langgraph.types import All, Checkpointer
|
||||
|
||||
|
||||
class Edge(NamedTuple):
|
||||
source: str
|
||||
target: str
|
||||
conditional: bool
|
||||
data: str | None
|
||||
|
||||
|
||||
class TriggerEdge(NamedTuple):
|
||||
source: str
|
||||
conditional: bool
|
||||
data: str | None
|
||||
|
||||
|
||||
def draw_graph(
|
||||
config: RunnableConfig,
|
||||
*,
|
||||
@@ -49,7 +63,7 @@ def draw_graph(
|
||||
The graph for this Pregel instance.
|
||||
"""
|
||||
# (src, dest, is_conditional, label)
|
||||
edges: set[tuple[str, str, bool, str | None]] = set()
|
||||
edges: set[Edge] = set()
|
||||
|
||||
step = -1
|
||||
checkpoint = empty_checkpoint()
|
||||
@@ -63,8 +77,9 @@ def draw_graph(
|
||||
checkpoint,
|
||||
)
|
||||
static_seen: set[Any] = set()
|
||||
sources: dict[str, set[tuple[str, bool, str | None]]] = {}
|
||||
step_sources: dict[str, set[tuple[str, bool, str | None]]] = {}
|
||||
sources: dict[str, set[TriggerEdge]] = {}
|
||||
step_sources: dict[str, set[TriggerEdge]] = {}
|
||||
static_declared_writes: dict[str, set[TriggerEdge]] = defaultdict(set)
|
||||
# remove node mappers
|
||||
nodes = {
|
||||
k: v.copy(update={"mapper": None}) if v.mapper is not None else v
|
||||
@@ -123,32 +138,36 @@ def draw_graph(
|
||||
# END writes are not written, but become edges directly
|
||||
for t in writes:
|
||||
if t[0] == END:
|
||||
edges.add((task.name, t[0], True, t[2]))
|
||||
edges.add(Edge(task.name, t[0], True, t[2]))
|
||||
writes = [t for t in writes if t[0] != END]
|
||||
conditionals.update(
|
||||
{(task.name, t[0], t[1] or None): t[2] for t in writes}
|
||||
)
|
||||
# record static writes for edge creation
|
||||
for t in writes:
|
||||
static_declared_writes[task.name].add(
|
||||
TriggerEdge(t[0], True, t[2])
|
||||
)
|
||||
task.config[CONF][CONFIG_KEY_SEND]([t[:2] for t in writes])
|
||||
# collect sources
|
||||
step_sources = {
|
||||
task.name: {
|
||||
(
|
||||
step_sources = {}
|
||||
for task in tasks.values():
|
||||
task_edges = {
|
||||
TriggerEdge(
|
||||
w[0],
|
||||
(task.name, w[0], w[1] or None) in conditionals,
|
||||
conditionals.get((task.name, w[0], w[1] or None)),
|
||||
)
|
||||
for w in task.writes
|
||||
}
|
||||
for task in tasks.values()
|
||||
}
|
||||
task_edges |= static_declared_writes.get(task.name, set())
|
||||
step_sources[task.name] = task_edges
|
||||
sources.update(step_sources)
|
||||
# invert triggers
|
||||
trigger_to_sources: dict[str, set[tuple[str, bool, str | None]]] = defaultdict(
|
||||
set
|
||||
)
|
||||
trigger_to_sources: dict[str, set[TriggerEdge]] = defaultdict(set)
|
||||
for src, triggers in sources.items():
|
||||
for trigger, cond, label in triggers:
|
||||
trigger_to_sources[trigger].add((src, cond, label))
|
||||
trigger_to_sources[trigger].add(TriggerEdge(src, cond, label))
|
||||
# apply writes
|
||||
updated_channels = apply_writes(
|
||||
checkpoint, channels, tasks.values(), get_next_version, trigger_to_nodes
|
||||
@@ -170,26 +189,39 @@ def draw_graph(
|
||||
trigger_to_nodes=trigger_to_nodes,
|
||||
updated_channels=updated_channels,
|
||||
)
|
||||
# collect deferred nodes
|
||||
deferred_nodes: set[str] = set()
|
||||
edges_to_deferred_nodes: set[Edge] = set()
|
||||
for channel, item in channels.items():
|
||||
if isinstance(item, LastValueAfterFinish):
|
||||
deferred_node = channel.split(":", 2)[-1]
|
||||
deferred_nodes.add(deferred_node)
|
||||
# collect edges
|
||||
for task in tasks.values():
|
||||
added = False
|
||||
for trigger in task.triggers:
|
||||
for src, cond, label in sorted(trigger_to_sources[trigger]):
|
||||
edges.add((src, task.name, cond, label))
|
||||
# record edge to be reviewed later
|
||||
if task.name in deferred_nodes:
|
||||
edges_to_deferred_nodes.add(Edge(src, task.name, cond, label))
|
||||
edges.add(Edge(src, task.name, cond, label))
|
||||
# if the edge is from this step, skip adding the implicit edges
|
||||
if (trigger, cond, label) in step_sources.get(src, set()):
|
||||
added = True
|
||||
else:
|
||||
sources[src].discard((trigger, cond, label))
|
||||
sources[src].discard(TriggerEdge(trigger, cond, label))
|
||||
# if no edges from this step, add implicit edges from all previous tasks
|
||||
if not added:
|
||||
for src in step_sources:
|
||||
edges.add((src, task.name, True, None))
|
||||
edges.add(Edge(src, task.name, True, None))
|
||||
|
||||
# assemble the graph
|
||||
graph = Graph()
|
||||
# add nodes
|
||||
for name, node in nodes.items():
|
||||
metadata = dict(node.metadata or {})
|
||||
if name in deferred_nodes:
|
||||
metadata["defer"] = True
|
||||
if name in interrupt_before_nodes and name in interrupt_after_nodes:
|
||||
metadata["__interrupt"] = "before,after"
|
||||
elif name in interrupt_before_nodes:
|
||||
|
||||
@@ -307,6 +307,99 @@
|
||||
|
||||
'''
|
||||
# ---
|
||||
# name: test_get_graph_nonterminal_last_step_source
|
||||
'''
|
||||
{
|
||||
"edges": [
|
||||
{
|
||||
"source": "__start__",
|
||||
"target": "human"
|
||||
},
|
||||
{
|
||||
"conditional": true,
|
||||
"source": "chatbot",
|
||||
"target": "human"
|
||||
},
|
||||
{
|
||||
"conditional": true,
|
||||
"source": "chatbot",
|
||||
"target": "tools"
|
||||
},
|
||||
{
|
||||
"conditional": true,
|
||||
"source": "human",
|
||||
"target": "__end__"
|
||||
},
|
||||
{
|
||||
"conditional": true,
|
||||
"source": "human",
|
||||
"target": "chatbot"
|
||||
},
|
||||
{
|
||||
"source": "tools",
|
||||
"target": "chatbot"
|
||||
}
|
||||
],
|
||||
"nodes": [
|
||||
{
|
||||
"data": {
|
||||
"id": [
|
||||
"langgraph",
|
||||
"_internal",
|
||||
"_runnable",
|
||||
"RunnableCallable"
|
||||
],
|
||||
"name": "__start__"
|
||||
},
|
||||
"id": "__start__",
|
||||
"type": "runnable"
|
||||
},
|
||||
{
|
||||
"data": {
|
||||
"id": [
|
||||
"langgraph",
|
||||
"_internal",
|
||||
"_runnable",
|
||||
"RunnableCallable"
|
||||
],
|
||||
"name": "chatbot"
|
||||
},
|
||||
"id": "chatbot",
|
||||
"type": "runnable"
|
||||
},
|
||||
{
|
||||
"data": {
|
||||
"id": [
|
||||
"langgraph",
|
||||
"_internal",
|
||||
"_runnable",
|
||||
"RunnableCallable"
|
||||
],
|
||||
"name": "tools"
|
||||
},
|
||||
"id": "tools",
|
||||
"type": "runnable"
|
||||
},
|
||||
{
|
||||
"data": {
|
||||
"id": [
|
||||
"langgraph",
|
||||
"_internal",
|
||||
"_runnable",
|
||||
"RunnableCallable"
|
||||
],
|
||||
"name": "human"
|
||||
},
|
||||
"id": "human",
|
||||
"type": "runnable"
|
||||
},
|
||||
{
|
||||
"id": "__end__"
|
||||
}
|
||||
]
|
||||
}
|
||||
'''
|
||||
# ---
|
||||
# name: test_get_graph_root_channel
|
||||
'''
|
||||
{
|
||||
@@ -795,99 +888,6 @@
|
||||
|
||||
'''
|
||||
# ---
|
||||
# name: test_get_graph_nonterminal_last_step_source
|
||||
'''
|
||||
{
|
||||
"edges": [
|
||||
{
|
||||
"source": "__start__",
|
||||
"target": "human"
|
||||
},
|
||||
{
|
||||
"conditional": true,
|
||||
"source": "chatbot",
|
||||
"target": "human"
|
||||
},
|
||||
{
|
||||
"conditional": true,
|
||||
"source": "chatbot",
|
||||
"target": "tools"
|
||||
},
|
||||
{
|
||||
"conditional": true,
|
||||
"source": "human",
|
||||
"target": "__end__"
|
||||
},
|
||||
{
|
||||
"conditional": true,
|
||||
"source": "human",
|
||||
"target": "chatbot"
|
||||
},
|
||||
{
|
||||
"source": "tools",
|
||||
"target": "chatbot"
|
||||
}
|
||||
],
|
||||
"nodes": [
|
||||
{
|
||||
"data": {
|
||||
"id": [
|
||||
"langgraph",
|
||||
"_internal",
|
||||
"_runnable",
|
||||
"RunnableCallable"
|
||||
],
|
||||
"name": "__start__"
|
||||
},
|
||||
"id": "__start__",
|
||||
"type": "runnable"
|
||||
},
|
||||
{
|
||||
"data": {
|
||||
"id": [
|
||||
"langgraph",
|
||||
"_internal",
|
||||
"_runnable",
|
||||
"RunnableCallable"
|
||||
],
|
||||
"name": "chatbot"
|
||||
},
|
||||
"id": "chatbot",
|
||||
"type": "runnable"
|
||||
},
|
||||
{
|
||||
"data": {
|
||||
"id": [
|
||||
"langgraph",
|
||||
"_internal",
|
||||
"_runnable",
|
||||
"RunnableCallable"
|
||||
],
|
||||
"name": "tools"
|
||||
},
|
||||
"id": "tools",
|
||||
"type": "runnable"
|
||||
},
|
||||
{
|
||||
"data": {
|
||||
"id": [
|
||||
"langgraph",
|
||||
"_internal",
|
||||
"_runnable",
|
||||
"RunnableCallable"
|
||||
],
|
||||
"name": "human"
|
||||
},
|
||||
"id": "human",
|
||||
"type": "runnable"
|
||||
},
|
||||
{
|
||||
"id": "__end__"
|
||||
}
|
||||
]
|
||||
}
|
||||
'''
|
||||
# ---
|
||||
# name: test_repeat_condition
|
||||
'''
|
||||
graph TD;
|
||||
|
||||
Reference in New Issue
Block a user