From 5e175e098b01a72caaab26918d51534fee5f560f Mon Sep 17 00:00:00 2001 From: Nuno Campos Date: Mon, 14 Oct 2024 17:10:10 -0700 Subject: [PATCH] Update kafka --- .../langgraph/scheduler/kafka/executor.py | 21 ++++++++++++++++--- 1 file changed, 18 insertions(+), 3 deletions(-) diff --git a/libs/scheduler-kafka/langgraph/scheduler/kafka/executor.py b/libs/scheduler-kafka/langgraph/scheduler/kafka/executor.py index c2ad752e4..da2c2bafb 100644 --- a/libs/scheduler-kafka/langgraph/scheduler/kafka/executor.py +++ b/libs/scheduler-kafka/langgraph/scheduler/kafka/executor.py @@ -37,7 +37,7 @@ from langgraph.scheduler.kafka.types import ( Sendable, Topics, ) -from langgraph.types import RetryPolicy +from langgraph.types import LoopProtocol, RetryPolicy from langgraph.utils.config import patch_configurable @@ -183,7 +183,15 @@ class AsyncKafkaExecutor(AbstractAsyncContextManager): if saved.checkpoint["id"] != msg["config"]["configurable"]["checkpoint_id"]: raise CheckpointNotLatest() async with AsyncChannelsManager( - graph.channels, saved.checkpoint, msg["config"], self.graph.store + graph.channels, + saved.checkpoint, + LoopProtocol( + config=msg["config"], + store=self.graph.store, + step=saved.metadata["step"] + 1, + stop=saved.metadata["step"] + 2, + ), + self.graph.store, ) as (channels, managed), AsyncBackgroundExecutor() as submit: if task := await asyncio.to_thread( prepare_single_task, @@ -379,7 +387,14 @@ class KafkaExecutor(AbstractContextManager): if saved.checkpoint["id"] != msg["config"]["configurable"]["checkpoint_id"]: raise CheckpointNotLatest() with ChannelsManager( - graph.channels, saved.checkpoint, msg["config"], self.graph.store + 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"],