WEBDIENSTE/langgraph
1
0
Fork 0
mirror of https://github.com/langchain-ai/langgraph.git synced 2026-09-03 08:18:47 +02:00
Code Issues Packages Projects Releases Wiki Activity

Add support for multiple subgraphs called in a single node (#3056)

Browse Source
This commit is contained in:
Nuno Campos
2025-01-15 18:01:31 -08:00
committed by GitHub
parent 1763dd69a7 6b9369876f
commit 6bfc6be307
14 changed files with 164 additions and 173 deletions
Download Patch File Download Diff File Expand all files Collapse all files
libs/langgraph/langgraph/errors.py
-15
View File
@@ -107,18 +107,3 @@ class CheckpointNotLatest(Exception):
"""Raised when the checkpoint is not the latest version (for distributed mode)."""
pass
class MultipleSubgraphsError(Exception):
"""Raised when multiple subgraphs are called inside the same node.
Troubleshooting guides:
- [MULTIPLE_SUBGRAPHS](https://python.langchain.com/docs/troubleshooting/errors/MULTIPLE_SUBGRAPHS)
"""
pass
_SEEN_CHECKPOINT_NS: set[str] = set()
"""Used for subgraph detection."""
libs/langgraph/langgraph/pregel/__init__.py
+19 -44
View File
@@ -115,6 +115,7 @@ from langgraph.utils.config import (
patch_checkpoint_map,
patch_config,
patch_configurable,
recast_checkpoint_ns,
)
from langgraph.utils.fields import get_enhanced_type_hints
from langgraph.utils.pydantic import create_model
@@ -694,19 +695,15 @@ class Pregel(PregelProtocol):
checkpoint_ns := config[CONF].get(CONFIG_KEY_CHECKPOINT_NS, "")
) and CONFIG_KEY_CHECKPOINTER not in config[CONF]:
# remove task_ids from checkpoint_ns
recast_checkpoint_ns = NS_SEP.join(
part.split(NS_END)[0] for part in checkpoint_ns.split(NS_SEP)
)
recast = recast_checkpoint_ns(checkpoint_ns)
# find the subgraph with the matching name
for _, pregel in self.get_subgraphs(
namespace=recast_checkpoint_ns, recurse=True
):
for _, pregel in self.get_subgraphs(namespace=recast, recurse=True):
return pregel.get_state(
patch_configurable(config, {CONFIG_KEY_CHECKPOINTER: checkpointer}),
subgraphs=subgraphs,
)
else:
raise ValueError(f"Subgraph {recast_checkpoint_ns} not found")
raise ValueError(f"Subgraph {recast} not found")
config = merge_configs(self.config, config) if self.config else config
saved = checkpointer.get_tuple(config)
@@ -731,19 +728,15 @@ class Pregel(PregelProtocol):
checkpoint_ns := config[CONF].get(CONFIG_KEY_CHECKPOINT_NS, "")
) and CONFIG_KEY_CHECKPOINTER not in config[CONF]:
# remove task_ids from checkpoint_ns
recast_checkpoint_ns = NS_SEP.join(
part.split(NS_END)[0] for part in checkpoint_ns.split(NS_SEP)
)
recast = recast_checkpoint_ns(checkpoint_ns)
# find the subgraph with the matching name
async for _, pregel in self.aget_subgraphs(
namespace=recast_checkpoint_ns, recurse=True
):
async for _, pregel in self.aget_subgraphs(namespace=recast, recurse=True):
return await pregel.aget_state(
patch_configurable(config, {CONFIG_KEY_CHECKPOINTER: checkpointer}),
subgraphs=subgraphs,
)
else:
raise ValueError(f"Subgraph {recast_checkpoint_ns} not found")
raise ValueError(f"Subgraph {recast} not found")
config = merge_configs(self.config, config) if self.config else config
saved = await checkpointer.aget_tuple(config)
@@ -774,13 +767,9 @@ class Pregel(PregelProtocol):
checkpoint_ns := config[CONF].get(CONFIG_KEY_CHECKPOINT_NS, "")
) and CONFIG_KEY_CHECKPOINTER not in config[CONF]:
# remove task_ids from checkpoint_ns
recast_checkpoint_ns = NS_SEP.join(
part.split(NS_END)[0] for part in checkpoint_ns.split(NS_SEP)
)
recast = recast_checkpoint_ns(checkpoint_ns)
# find the subgraph with the matching name
for _, pregel in self.get_subgraphs(
namespace=recast_checkpoint_ns, recurse=True
):
for _, pregel in self.get_subgraphs(namespace=recast, recurse=True):
yield from pregel.get_state_history(
patch_configurable(config, {CONFIG_KEY_CHECKPOINTER: checkpointer}),
filter=filter,
@@ -789,7 +778,7 @@ class Pregel(PregelProtocol):
)
return
else:
raise ValueError(f"Subgraph {recast_checkpoint_ns} not found")
raise ValueError(f"Subgraph {recast} not found")
config = merge_configs(
self.config,
@@ -824,13 +813,9 @@ class Pregel(PregelProtocol):
checkpoint_ns := config[CONF].get(CONFIG_KEY_CHECKPOINT_NS, "")
) and CONFIG_KEY_CHECKPOINTER not in config[CONF]:
# remove task_ids from checkpoint_ns
recast_checkpoint_ns = NS_SEP.join(
part.split(NS_END)[0] for part in checkpoint_ns.split(NS_SEP)
)
recast = recast_checkpoint_ns(checkpoint_ns)
# find the subgraph with the matching name
async for _, pregel in self.aget_subgraphs(
namespace=recast_checkpoint_ns, recurse=True
):
async for _, pregel in self.aget_subgraphs(namespace=recast, recurse=True):
async for state in pregel.aget_state_history(
patch_configurable(config, {CONFIG_KEY_CHECKPOINTER: checkpointer}),
filter=filter,
@@ -840,7 +825,7 @@ class Pregel(PregelProtocol):
yield state
return
else:
raise ValueError(f"Subgraph {recast_checkpoint_ns} not found")
raise ValueError(f"Subgraph {recast} not found")
config = merge_configs(
self.config,
@@ -879,20 +864,16 @@ class Pregel(PregelProtocol):
checkpoint_ns := config[CONF].get(CONFIG_KEY_CHECKPOINT_NS, "")
) and CONFIG_KEY_CHECKPOINTER not in config[CONF]:
# remove task_ids from checkpoint_ns
recast_checkpoint_ns = NS_SEP.join(
part.split(NS_END)[0] for part in checkpoint_ns.split(NS_SEP)
)
recast = recast_checkpoint_ns(checkpoint_ns)
# find the subgraph with the matching name
for _, pregel in self.get_subgraphs(
namespace=recast_checkpoint_ns, recurse=True
):
for _, pregel in self.get_subgraphs(namespace=recast, recurse=True):
return pregel.update_state(
patch_configurable(config, {CONFIG_KEY_CHECKPOINTER: checkpointer}),
values,
as_node,
)
else:
raise ValueError(f"Subgraph {recast_checkpoint_ns} not found")
raise ValueError(f"Subgraph {recast} not found")
# get last checkpoint
config = ensure_config(self.config, config)
@@ -1163,20 +1144,16 @@ class Pregel(PregelProtocol):
checkpoint_ns := config[CONF].get(CONFIG_KEY_CHECKPOINT_NS, "")
) and CONFIG_KEY_CHECKPOINTER not in config[CONF]:
# remove task_ids from checkpoint_ns
recast_checkpoint_ns = NS_SEP.join(
part.split(NS_END)[0] for part in checkpoint_ns.split(NS_SEP)
)
recast = recast_checkpoint_ns(checkpoint_ns)
# find the subgraph with the matching name
async for _, pregel in self.aget_subgraphs(
namespace=recast_checkpoint_ns, recurse=True
):
async for _, pregel in self.aget_subgraphs(namespace=recast, recurse=True):
return await pregel.aupdate_state(
patch_configurable(config, {CONFIG_KEY_CHECKPOINTER: checkpointer}),
values,
as_node,
)
else:
raise ValueError(f"Subgraph {recast_checkpoint_ns} not found")
raise ValueError(f"Subgraph {recast} not found")
# get last checkpoint
config = ensure_config(self.config, config)
@@ -1642,7 +1619,6 @@ class Pregel(PregelProtocol):
interrupt_after=interrupt_after_,
manager=run_manager,
debug=debug,
check_subgraphs=self.checkpointer is not True,
) as loop:
# create runner
runner = PregelRunner(
@@ -1870,7 +1846,6 @@ class Pregel(PregelProtocol):
interrupt_after=interrupt_after_,
manager=run_manager,
debug=debug,
check_subgraphs=self.checkpointer is not True,
) as loop:
# create runner
runner = PregelRunner(
libs/langgraph/langgraph/pregel/algo.py
+3 -1
View File
@@ -680,7 +680,7 @@ def prepare_single_task(
"langgraph_checkpoint_ns": task_checkpoint_ns,
}
if task_id_checksum is not None:
assert task_id == task_id_checksum
assert task_id == task_id_checksum, f"{task_id} != {task_id_checksum}"
if for_execution:
if node := proc.node:
if proc.metadata:
@@ -766,6 +766,8 @@ def _scratchpad(
(w[2] for w in pending_writes if w[0] == NULL_TASK_ID and w[1] == RESUME),
MISSING,
),
# subgraph
subgraph_counter=0,
)
libs/langgraph/langgraph/pregel/loop.py
+16 -16
View File
@@ -63,12 +63,10 @@ from langgraph.constants import (
TAG_HIDDEN,
)
from langgraph.errors import (
_SEEN_CHECKPOINT_NS,
CheckpointNotLatest,
EmptyInputError,
GraphDelegate,
GraphInterrupt,
MultipleSubgraphsError,
)
from langgraph.managed.base import (
ManagedValueMapping,
@@ -116,6 +114,7 @@ from langgraph.types import (
Command,
LoopProtocol,
PregelExecutableTask,
PregelScratchpad,
StreamChunk,
StreamProtocol,
)
@@ -203,7 +202,6 @@ class PregelLoop(LoopProtocol):
interrupt_after: Union[All, Sequence[str]] = EMPTY_SEQ,
interrupt_before: Union[All, Sequence[str]] = EMPTY_SEQ,
manager: Union[None, AsyncParentRunManager, ParentRunManager] = None,
check_subgraphs: bool = True,
debug: bool = False,
) -> None:
super().__init__(
@@ -230,20 +228,26 @@ class PregelLoop(LoopProtocol):
self.debug = debug
if self.stream is not None and CONFIG_KEY_STREAM in config[CONF]:
self.stream = DuplexStream(self.stream, config[CONF][CONFIG_KEY_STREAM])
scratchpad: Optional[PregelScratchpad] = config[CONF].get(CONFIG_KEY_SCRATCHPAD)
if not self.config[CONF].get(CONFIG_KEY_DELEGATE) and scratchpad is not None:
if scratchpad["subgraph_counter"]:
self.config = patch_configurable(
self.config,
{
CONFIG_KEY_CHECKPOINT_NS: NS_SEP.join(
(
config[CONF][CONFIG_KEY_CHECKPOINT_NS],
str(scratchpad["subgraph_counter"]),
)
)
},
)
scratchpad["subgraph_counter"] += 1
if not self.is_nested and config[CONF].get(CONFIG_KEY_CHECKPOINT_NS):
self.config = patch_configurable(
self.config,
{CONFIG_KEY_CHECKPOINT_NS: "", CONFIG_KEY_CHECKPOINT_ID: None},
)
if check_subgraphs and self.is_nested and self.checkpointer is not None:
if self.config[CONF][CONFIG_KEY_CHECKPOINT_NS] in _SEEN_CHECKPOINT_NS:
raise MultipleSubgraphsError(
"Multiple subgraphs called inside the same node\n\n"
"Troubleshooting URL: https://python.langchain.com/docs"
"/troubleshooting/errors/MULTIPLE_SUBGRAPHS/"
)
else:
_SEEN_CHECKPOINT_NS.add(self.config[CONF][CONFIG_KEY_CHECKPOINT_NS])
if (
CONFIG_KEY_CHECKPOINT_MAP in self.config[CONF]
and self.config[CONF].get(CONFIG_KEY_CHECKPOINT_NS)
@@ -817,7 +821,6 @@ class SyncPregelLoop(PregelLoop, ContextManager):
interrupt_before: Union[All, Sequence[str]] = EMPTY_SEQ,
output_keys: Union[str, Sequence[str]] = EMPTY_SEQ,
stream_keys: Union[str, Sequence[str]] = EMPTY_SEQ,
check_subgraphs: bool = True,
debug: bool = False,
) -> None:
super().__init__(
@@ -832,7 +835,6 @@ class SyncPregelLoop(PregelLoop, ContextManager):
stream_keys=stream_keys,
interrupt_after=interrupt_after,
interrupt_before=interrupt_before,
check_subgraphs=check_subgraphs,
manager=manager,
debug=debug,
)
@@ -954,7 +956,6 @@ class AsyncPregelLoop(PregelLoop, AsyncContextManager):
manager: Union[None, AsyncParentRunManager, ParentRunManager] = None,
output_keys: Union[str, Sequence[str]] = EMPTY_SEQ,
stream_keys: Union[str, Sequence[str]] = EMPTY_SEQ,
check_subgraphs: bool = True,
debug: bool = False,
) -> None:
super().__init__(
@@ -969,7 +970,6 @@ class AsyncPregelLoop(PregelLoop, AsyncContextManager):
stream_keys=stream_keys,
interrupt_after=interrupt_after,
interrupt_before=interrupt_before,
check_subgraphs=check_subgraphs,
manager=manager,
debug=debug,
)
libs/langgraph/langgraph/pregel/retry.py
+9 -17
View File
@@ -12,7 +12,7 @@ from langgraph.constants import (
CONFIG_KEY_RESUMING,
NS_SEP,
)
from langgraph.errors import _SEEN_CHECKPOINT_NS, GraphBubbleUp, ParentCommand
from langgraph.errors import GraphBubbleUp, ParentCommand
from langgraph.types import Command, PregelExecutableTask, RetryPolicy
from langgraph.utils.config import patch_configurable
@@ -48,7 +48,10 @@ def run_with_retry(
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])
parts = ns.split(NS_SEP)
if parts[-1].isdigit():
parts.pop()
parent_ns = NS_SEP.join(parts[:-1])
exc.args = (replace(cmd, graph=parent_ns),)
# bubble up
raise
@@ -96,13 +99,6 @@ def run_with_retry(
)
# signal subgraphs to resume (if available)
config = patch_configurable(config, {CONFIG_KEY_RESUMING: True})
# clear checkpoint_ns seen (for subgraph detection)
if checkpoint_ns := config[CONF].get(CONFIG_KEY_CHECKPOINT_NS):
_SEEN_CHECKPOINT_NS.discard(checkpoint_ns)
finally:
# clear checkpoint_ns seen (for subgraph detection)
if checkpoint_ns := config[CONF].get(CONFIG_KEY_CHECKPOINT_NS):
_SEEN_CHECKPOINT_NS.discard(checkpoint_ns)
async def arun_with_retry(
@@ -140,7 +136,10 @@ async def arun_with_retry(
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])
parts = ns.split(NS_SEP)
if parts[-1].isdigit():
parts.pop()
parent_ns = NS_SEP.join(parts[:-1])
exc.args = (replace(cmd, graph=parent_ns),)
# bubble up
raise
@@ -188,10 +187,3 @@ async def arun_with_retry(
)
# signal subgraphs to resume (if available)
config = patch_configurable(config, {CONFIG_KEY_RESUMING: True})
# clear checkpoint_ns seen (for subgraph detection)
if checkpoint_ns := config[CONF].get(CONFIG_KEY_CHECKPOINT_NS):
_SEEN_CHECKPOINT_NS.discard(checkpoint_ns)
finally:
# clear checkpoint_ns seen (for subgraph detection)
if checkpoint_ns := config[CONF].get(CONFIG_KEY_CHECKPOINT_NS):
_SEEN_CHECKPOINT_NS.discard(checkpoint_ns)
libs/langgraph/langgraph/types.py
+2
View File
@@ -346,6 +346,8 @@ class PregelScratchpad(TypedDict):
interrupt_counter: int
resume: list[Any]
null_resume: Any
# subgraph
subgraph_counter: int
def interrupt(value: Any) -> Any:
libs/langgraph/langgraph/utils/config.py
+16
View File
@@ -23,9 +23,25 @@ from langgraph.constants import (
CONFIG_KEY_CHECKPOINT_ID,
CONFIG_KEY_CHECKPOINT_MAP,
CONFIG_KEY_CHECKPOINT_NS,
NS_END,
NS_SEP,
)
def recast_checkpoint_ns(ns: str) -> str:
"""Remove task IDs from checkpoint namespace.
Args:
ns (str): The checkpoint namespace with task IDs.
Returns:
str: The checkpoint namespace without task IDs.
"""
return NS_SEP.join(
part.split(NS_END)[0] for part in ns.split(NS_SEP) if not part.isdigit()
)
def patch_configurable(
config: Optional[RunnableConfig], patch: dict[str, Any]
) -> RunnableConfig:
libs/langgraph/tests/test_pregel.py
+3 -4
View File
@@ -51,7 +51,7 @@ from langgraph.checkpoint.base import (
)
from langgraph.checkpoint.memory import MemorySaver
from langgraph.constants import CONFIG_KEY_NODE_FINISHED, ERROR, PULL, START
from langgraph.errors import InvalidUpdateError, MultipleSubgraphsError
from langgraph.errors import InvalidUpdateError
from langgraph.func import entrypoint, task
from langgraph.graph import END, Graph, StateGraph
from langgraph.graph.message import MessageGraph, MessagesState, add_messages
@@ -1745,9 +1745,8 @@ def test_invoke_join_then_call_other_pregel(
# add checkpointer
app.checkpointer = checkpointer
# subgraph is called twice in the same node, through .map(), so raises
with pytest.raises(MultipleSubgraphsError):
app.invoke([2, 3], {"configurable": {"thread_id": "1"}})
# subgraph is called twice in the same node, but that works
assert app.invoke([2, 3], {"configurable": {"thread_id": "1"}}) == 27
# set inner graph checkpointer NeverCheckpoint
inner_app.checkpointer = False
libs/langgraph/tests/test_pregel_async.py
+3 -4
View File
@@ -48,7 +48,7 @@ from langgraph.checkpoint.base import (
)
from langgraph.checkpoint.memory import MemorySaver
from langgraph.constants import CONFIG_KEY_NODE_FINISHED, ERROR, PULL, PUSH, START
from langgraph.errors import InvalidUpdateError, MultipleSubgraphsError, NodeInterrupt
from langgraph.errors import InvalidUpdateError, NodeInterrupt
from langgraph.func import entrypoint, task
from langgraph.graph import END, Graph, StateGraph
from langgraph.graph.message import MessagesState, add_messages
@@ -4068,9 +4068,8 @@ async def test_invoke_join_then_call_other_pregel(
async with awith_checkpointer(checkpointer_name) as checkpointer:
# add checkpointer
app.checkpointer = checkpointer
# subgraph is called twice in the same node, through .map(), so raises
with pytest.raises(MultipleSubgraphsError):
await app.ainvoke([2, 3], {"configurable": {"thread_id": "1"}})
# subgraph is called twice, and that works
assert await app.ainvoke([2, 3], {"configurable": {"thread_id": "1"}}) == 27
# set inner graph checkpointer NeverCheckpoint
inner_app.checkpointer = False
libs/scheduler-kafka/langgraph/scheduler/kafka/executor.py
+38 -35
View File
@@ -1,5 +1,6 @@
import asyncio
import concurrent.futures
from collections.abc import Sequence
from contextlib import (
AbstractAsyncContextManager,
AbstractContextManager,
@@ -7,7 +8,7 @@ from contextlib import (
ExitStack,
)
from functools import partial
from typing import Any, Optional, Sequence
from typing import Any, Optional
from uuid import UUID
import orjson
@@ -15,7 +16,7 @@ from langchain_core.runnables import RunnableConfig
from typing_extensions import Self
import langgraph.scheduler.kafka.serde as serde
from langgraph.constants import CONFIG_KEY_DELEGATE, ERROR, NS_END, NS_SEP
from langgraph.constants import CONFIG_KEY_DELEGATE, ERROR
from langgraph.errors import CheckpointNotLatest, GraphDelegate, TaskNotFound
from langgraph.pregel import Pregel
from langgraph.pregel.algo import prepare_single_task
@@ -39,7 +40,7 @@ from langgraph.scheduler.kafka.types import (
Topics,
)
from langgraph.types import LoopProtocol, PregelExecutableTask, RetryPolicy
from langgraph.utils.config import patch_configurable
from langgraph.utils.config import patch_configurable, recast_checkpoint_ns
class AsyncKafkaExecutor(AbstractAsyncContextManager):
@@ -165,14 +166,12 @@ class AsyncKafkaExecutor(AbstractAsyncContextManager):
# find graph
if checkpoint_ns := msg["config"]["configurable"].get("checkpoint_ns"):
# remove task_ids from checkpoint_ns
recast_checkpoint_ns = NS_SEP.join(
part.split(NS_END)[0] for part in checkpoint_ns.split(NS_SEP)
)
recast = recast_checkpoint_ns(checkpoint_ns)
# find the subgraph with the matching name
if recast_checkpoint_ns in self.subgraphs:
graph = self.subgraphs[recast_checkpoint_ns]
if recast in self.subgraphs:
graph = self.subgraphs[recast]
else:
raise ValueError(f"Subgraph {recast_checkpoint_ns} not found")
raise ValueError(f"Subgraph {recast} not found")
else:
graph = self.graph
# process message
@@ -183,16 +182,19 @@ class AsyncKafkaExecutor(AbstractAsyncContextManager):
raise RuntimeError("Checkpoint not found")
if saved.checkpoint["id"] != msg["config"]["configurable"]["checkpoint_id"]:
raise CheckpointNotLatest()
async with AsyncChannelsManager(
graph.channels,
saved.checkpoint,
LoopProtocol(
config=msg["config"],
store=self.graph.store,
step=saved.metadata["step"] + 1,
stop=saved.metadata["step"] + 2,
),
) as (channels, managed), AsyncBackgroundExecutor(msg["config"]) as submit:
async with (
AsyncChannelsManager(
graph.channels,
saved.checkpoint,
LoopProtocol(
config=msg["config"],
store=self.graph.store,
step=saved.metadata["step"] + 1,
stop=saved.metadata["step"] + 2,
),
) as (channels, managed),
AsyncBackgroundExecutor(msg["config"]) as submit,
):
if task := await asyncio.to_thread(
prepare_single_task,
msg["task"]["path"],
@@ -378,14 +380,12 @@ class KafkaExecutor(AbstractContextManager):
# find graph
if checkpoint_ns := msg["config"]["configurable"].get("checkpoint_ns"):
# remove task_ids from checkpoint_ns
recast_checkpoint_ns = NS_SEP.join(
part.split(NS_END)[0] for part in checkpoint_ns.split(NS_SEP)
)
recast = recast_checkpoint_ns(checkpoint_ns)
# find the subgraph with the matching name
if recast_checkpoint_ns in self.subgraphs:
graph = self.subgraphs[recast_checkpoint_ns]
if recast in self.subgraphs:
graph = self.subgraphs[recast]
else:
raise ValueError(f"Subgraph {recast_checkpoint_ns} not found")
raise ValueError(f"Subgraph {recast} not found")
else:
graph = self.graph
# process message
@@ -396,16 +396,19 @@ class KafkaExecutor(AbstractContextManager):
raise RuntimeError("Checkpoint not found")
if saved.checkpoint["id"] != msg["config"]["configurable"]["checkpoint_id"]:
raise CheckpointNotLatest()
with ChannelsManager(
graph.channels,
saved.checkpoint,
LoopProtocol(
config=msg["config"],
store=self.graph.store,
step=saved.metadata["step"] + 1,
stop=saved.metadata["step"] + 2,
),
) as (channels, managed), BackgroundExecutor({}) as submit:
with (
ChannelsManager(
graph.channels,
saved.checkpoint,
LoopProtocol(
config=msg["config"],
store=self.graph.store,
step=saved.metadata["step"] + 1,
stop=saved.metadata["step"] + 2,
),
) as (channels, managed),
BackgroundExecutor({}) as submit,
):
if task := prepare_single_task(
msg["task"]["path"],
msg["task"]["id"],
libs/scheduler-kafka/langgraph/scheduler/kafka/orchestrator.py
+33 -35
View File
@@ -13,11 +13,11 @@ from typing_extensions import Self
import langgraph.scheduler.kafka.serde as serde
from langgraph.constants import (
CONF,
CONFIG_KEY_DEDUPE_TASKS,
CONFIG_KEY_ENSURE_LATEST,
CONFIG_KEY_SCRATCHPAD,
INTERRUPT,
NS_END,
NS_SEP,
SCHEDULED,
)
from langgraph.errors import CheckpointNotLatest, GraphInterrupt
@@ -37,7 +37,7 @@ from langgraph.scheduler.kafka.types import (
Topics,
)
from langgraph.types import RetryPolicy
from langgraph.utils.config import patch_configurable
from langgraph.utils.config import patch_configurable, recast_checkpoint_ns
class AsyncKafkaOrchestrator(AbstractAsyncContextManager):
@@ -140,14 +140,12 @@ class AsyncKafkaOrchestrator(AbstractAsyncContextManager):
# find graph
if checkpoint_ns := msg["config"]["configurable"].get("checkpoint_ns"):
# remove task_ids from checkpoint_ns
recast_checkpoint_ns = NS_SEP.join(
part.split(NS_END)[0] for part in checkpoint_ns.split(NS_SEP)
)
recast = recast_checkpoint_ns(checkpoint_ns)
# find the subgraph with the matching name
if recast_checkpoint_ns in self.subgraphs:
graph = self.subgraphs[recast_checkpoint_ns]
if recast in self.subgraphs:
graph = self.subgraphs[recast]
else:
raise ValueError(f"Subgraph {recast_checkpoint_ns} not found")
raise ValueError(f"Subgraph {recast} not found")
else:
graph = self.graph
# process message
@@ -163,7 +161,6 @@ class AsyncKafkaOrchestrator(AbstractAsyncContextManager):
stream_keys=graph.stream_channels,
interrupt_after=graph.interrupt_after_nodes,
interrupt_before=graph.interrupt_before_nodes,
check_subgraphs=False,
) as loop:
if loop.tick(input_keys=graph.input_channels):
# wait for checkpoint to be saved
@@ -173,6 +170,16 @@ class AsyncKafkaOrchestrator(AbstractAsyncContextManager):
if new_tasks := [
t for t in loop.tasks.values() if not t.scheduled and not t.writes
]:
config = patch_configurable(
loop.config,
{
**loop.checkpoint_config["configurable"],
CONFIG_KEY_DEDUPE_TASKS: True,
CONFIG_KEY_ENSURE_LATEST: True,
},
)
if CONFIG_KEY_SCRATCHPAD in config[CONF]:
config[CONF][CONFIG_KEY_SCRATCHPAD]["subgraph_counter"] = 0
# send messages to executor
futures = await asyncio.gather(
*(
@@ -180,16 +187,7 @@ class AsyncKafkaOrchestrator(AbstractAsyncContextManager):
self.topics.executor,
value=serde.dumps(
MessageToExecutor(
config=patch_configurable(
loop.config,
{
**loop.checkpoint_config[
"configurable"
],
CONFIG_KEY_DEDUPE_TASKS: True,
CONFIG_KEY_ENSURE_LATEST: True,
},
),
config=config,
task=ExecutorTask(id=task.id, path=task.path),
finally_send=msg.get("finally_send"),
)
@@ -330,14 +328,12 @@ class KafkaOrchestrator(AbstractContextManager):
# find graph
if checkpoint_ns := msg["config"]["configurable"].get("checkpoint_ns"):
# remove task_ids from checkpoint_ns
recast_checkpoint_ns = NS_SEP.join(
part.split(NS_END)[0] for part in checkpoint_ns.split(NS_SEP)
)
recast = recast_checkpoint_ns(checkpoint_ns)
# find the subgraph with the matching name
if recast_checkpoint_ns in self.subgraphs:
graph = self.subgraphs[recast_checkpoint_ns]
if recast in self.subgraphs:
graph = self.subgraphs[recast]
else:
raise ValueError(f"Subgraph {recast_checkpoint_ns} not found")
raise ValueError(f"Subgraph {recast} not found")
else:
graph = self.graph
# process message
@@ -353,7 +349,6 @@ class KafkaOrchestrator(AbstractContextManager):
stream_keys=graph.stream_channels,
interrupt_after=graph.interrupt_after_nodes,
interrupt_before=graph.interrupt_before_nodes,
check_subgraphs=False,
) as loop:
if loop.tick(input_keys=graph.input_channels):
# wait for checkpoint to be saved
@@ -363,20 +358,23 @@ class KafkaOrchestrator(AbstractContextManager):
if new_tasks := [
t for t in loop.tasks.values() if not t.scheduled and not t.writes
]:
config = patch_configurable(
loop.config,
{
**loop.checkpoint_config["configurable"],
CONFIG_KEY_DEDUPE_TASKS: True,
CONFIG_KEY_ENSURE_LATEST: True,
},
)
if CONFIG_KEY_SCRATCHPAD in config[CONF]:
config[CONF][CONFIG_KEY_SCRATCHPAD]["subgraph_counter"] = 0
# send messages to executor
futures = [
self.producer.send(
self.topics.executor,
value=serde.dumps(
MessageToExecutor(
config=patch_configurable(
loop.config,
{
**loop.checkpoint_config["configurable"],
CONFIG_KEY_DEDUPE_TASKS: True,
CONFIG_KEY_ENSURE_LATEST: True,
},
),
config=config,
task=ExecutorTask(id=task.id, path=task.path),
finally_send=msg.get("finally_send"),
)
libs/scheduler-kafka/tests/any.py
+8
View File
@@ -53,3 +53,11 @@ class AnyList(list):
return False
else:
return True
class AnyInt(int):
def __init__(self) -> None:
super().__init__()
def __eq__(self, other: object) -> bool:
return isinstance(other, int)
libs/scheduler-kafka/tests/test_subgraph.py
+7 -1
View File
@@ -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, AnyInt
from tests.drain import drain_topics_async
from tests.messages import _AnyIdAIMessage, _AnyIdHumanMessage
@@ -198,6 +198,7 @@ async def test_subgraph_w_interrupt(
"__pregel_store": None,
"__pregel_task_id": history[0].tasks[0].id,
"__pregel_scratchpad": {
"subgraph_counter": AnyInt(),
"call_counter": 0,
"interrupt_counter": -1,
"null_resume": None,
@@ -269,6 +270,7 @@ async def test_subgraph_w_interrupt(
"__pregel_store": None,
"__pregel_task_id": history[0].tasks[0].id,
"__pregel_scratchpad": {
"subgraph_counter": AnyInt(),
"call_counter": 0,
"interrupt_counter": -1,
"null_resume": None,
@@ -370,6 +372,7 @@ async def test_subgraph_w_interrupt(
"__pregel_store": None,
"__pregel_task_id": history[0].tasks[0].id,
"__pregel_scratchpad": {
"subgraph_counter": AnyInt(),
"call_counter": 0,
"interrupt_counter": -1,
"null_resume": None,
@@ -481,6 +484,7 @@ async def test_subgraph_w_interrupt(
"__pregel_store": None,
"__pregel_task_id": history[1].tasks[0].id,
"__pregel_scratchpad": {
"subgraph_counter": AnyInt(),
"call_counter": 0,
"interrupt_counter": -1,
"null_resume": None,
@@ -547,6 +551,7 @@ async def test_subgraph_w_interrupt(
"__pregel_store": None,
"__pregel_task_id": history[1].tasks[0].id,
"__pregel_scratchpad": {
"subgraph_counter": AnyInt(),
"call_counter": 0,
"interrupt_counter": -1,
"null_resume": None,
@@ -669,6 +674,7 @@ async def test_subgraph_w_interrupt(
"__pregel_store": None,
"__pregel_task_id": history[1].tasks[0].id,
"__pregel_scratchpad": {
"subgraph_counter": AnyInt(),
"call_counter": 0,
"interrupt_counter": -1,
"null_resume": None,
libs/scheduler-kafka/tests/test_subgraph_sync.py
+7 -1
View File
@@ -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, AnyInt
from tests.drain import drain_topics
from tests.messages import _AnyIdAIMessage, _AnyIdHumanMessage
@@ -197,6 +197,7 @@ def test_subgraph_w_interrupt(
"__pregel_store": None,
"__pregel_task_id": history[0].tasks[0].id,
"__pregel_scratchpad": {
"subgraph_counter": AnyInt(),
"call_counter": 0,
"interrupt_counter": -1,
"null_resume": None,
@@ -268,6 +269,7 @@ def test_subgraph_w_interrupt(
"__pregel_resuming": False,
"__pregel_task_id": history[0].tasks[0].id,
"__pregel_scratchpad": {
"subgraph_counter": AnyInt(),
"call_counter": 0,
"interrupt_counter": -1,
"null_resume": None,
@@ -369,6 +371,7 @@ def test_subgraph_w_interrupt(
"__pregel_resuming": False,
"__pregel_task_id": history[0].tasks[0].id,
"__pregel_scratchpad": {
"subgraph_counter": AnyInt(),
"call_counter": 0,
"interrupt_counter": -1,
"null_resume": None,
@@ -479,6 +482,7 @@ def test_subgraph_w_interrupt(
"__pregel_resuming": True,
"__pregel_task_id": history[1].tasks[0].id,
"__pregel_scratchpad": {
"subgraph_counter": AnyInt(),
"call_counter": 0,
"interrupt_counter": -1,
"null_resume": None,
@@ -545,6 +549,7 @@ def test_subgraph_w_interrupt(
"__pregel_resuming": True,
"__pregel_task_id": history[1].tasks[0].id,
"__pregel_scratchpad": {
"subgraph_counter": AnyInt(),
"call_counter": 0,
"interrupt_counter": -1,
"null_resume": None,
@@ -667,6 +672,7 @@ def test_subgraph_w_interrupt(
"__pregel_store": None,
"__pregel_task_id": history[1].tasks[0].id,
"__pregel_scratchpad": {
"subgraph_counter": AnyInt(),
"call_counter": 0,
"interrupt_counter": -1,
"null_resume": None,