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

Merge pull request #2105 from langchain-ai/nc/14oct/is-last-step-fix

Browse Source 
Fix IsLastStep counter for runs with checkpointers
This commit is contained in:
Nuno Campos
2024-10-14 17:21:31 -07:00
committed by GitHub
parent edec5c055e d48faecd42
commit e8b875906f
12 changed files with 211 additions and 111 deletions
Download Patch File Download Diff File Expand all files Collapse all files
libs/langgraph/langgraph/managed/__init__.py
+2 -2
View File
@@ -1,3 +1,3 @@
from langgraph.managed.is_last_step import IsLastStep
from langgraph.managed.is_last_step import IsLastStep, RemainingSteps
__all__ = ["IsLastStep"]
__all__ = ["IsLastStep", "RemainingSteps"]
libs/langgraph/langgraph/managed/base.py
+9 -8
View File
@@ -13,22 +13,23 @@ from typing import (
Union,
)
from langchain_core.runnables import RunnableConfig
from typing_extensions import Self, TypeGuard
from langgraph.types import LoopProtocol
V = TypeVar("V")
U = TypeVar("U")
class ManagedValue(ABC, Generic[V]):
def __init__(self, config: RunnableConfig) -> None:
self.config = config
def __init__(self, loop: LoopProtocol) -> None:
self.loop = loop
@classmethod
@contextmanager
def enter(cls, config: RunnableConfig, **kwargs: Any) -> Iterator[Self]:
def enter(cls, loop: LoopProtocol, **kwargs: Any) -> Iterator[Self]:
try:
value = cls(config, **kwargs)
value = cls(loop, **kwargs)
yield value
finally:
# because managed value and Pregel have reference to each other
@@ -40,9 +41,9 @@ class ManagedValue(ABC, Generic[V]):
@classmethod
@asynccontextmanager
async def aenter(cls, config: RunnableConfig, **kwargs: Any) -> AsyncIterator[Self]:
async def aenter(cls, loop: LoopProtocol, **kwargs: Any) -> AsyncIterator[Self]:
try:
value = cls(config, **kwargs)
value = cls(loop, **kwargs)
yield value
finally:
# because managed value and Pregel have reference to each other
@@ -53,7 +54,7 @@ class ManagedValue(ABC, Generic[V]):
pass
@abstractmethod
def __call__(self, step: int) -> V: ...
def __call__(self) -> V: ...
class WritableManagedValue(Generic[V, U], ManagedValue[V], ABC):
libs/langgraph/langgraph/managed/context.py
+10 -10
View File
@@ -13,10 +13,10 @@ from typing import (
Union,
)
from langchain_core.runnables import RunnableConfig
from typing_extensions import Self
from langgraph.managed.base import ConfiguredManagedValue, ManagedValue, V
from langgraph.types import LoopProtocol
class Context(ManagedValue[V], Generic[V]):
@@ -46,14 +46,14 @@ class Context(ManagedValue[V], Generic[V]):
@classmethod
@contextmanager
def enter(cls, config: RunnableConfig, **kwargs: Any) -> Iterator[Self]:
with super().enter(config, **kwargs) as self:
def enter(cls, loop: LoopProtocol, **kwargs: Any) -> Iterator[Self]:
with super().enter(loop, **kwargs) as self:
if self.ctx is None:
raise ValueError(
"Synchronous context manager not found. Please initialize Context value with a sync context manager, or invoke your graph asynchronously."
)
ctx = (
self.ctx(config) # type: ignore[call-arg]
self.ctx(loop.config) # type: ignore[call-arg]
if signature(self.ctx).parameters.get("config")
else self.ctx()
)
@@ -63,17 +63,17 @@ class Context(ManagedValue[V], Generic[V]):
@classmethod
@asynccontextmanager
async def aenter(cls, config: RunnableConfig, **kwargs: Any) -> AsyncIterator[Self]:
async with super().aenter(config, **kwargs) as self:
async def aenter(cls, loop: LoopProtocol, **kwargs: Any) -> AsyncIterator[Self]:
async with super().aenter(loop, **kwargs) as self:
if self.actx is not None:
ctx = (
self.actx(config) # type: ignore[call-arg]
self.actx(loop.config) # type: ignore[call-arg]
if signature(self.actx).parameters.get("config")
else self.actx()
)
elif self.ctx is not None:
ctx = (
self.ctx(config) # type: ignore
self.ctx(loop.config) # type: ignore
if signature(self.ctx).parameters.get("config")
else self.ctx()
)
@@ -96,7 +96,7 @@ class Context(ManagedValue[V], Generic[V]):
def __init__(
self,
config: RunnableConfig,
loop: LoopProtocol,
*,
ctx: Union[None, Type[ContextManager[V]], Type[AsyncContextManager[V]]] = None,
actx: Optional[Type[AsyncContextManager[V]]] = None,
@@ -104,5 +104,5 @@ class Context(ManagedValue[V], Generic[V]):
self.ctx = ctx
self.actx = actx
def __call__(self, step: int) -> V:
def __call__(self) -> V:
return self.value
libs/langgraph/langgraph/managed/is_last_step.py
+10 -2
View File
@@ -4,8 +4,16 @@ from langgraph.managed.base import ManagedValue
class IsLastStepManager(ManagedValue[bool]):
def __call__(self, step: int) -> bool:
return step == self.config.get("recursion_limit", 0) - 1
def __call__(self) -> bool:
return self.loop.step == self.loop.stop
IsLastStep = Annotated[bool, IsLastStepManager]
class RemainingStepsManager(ManagedValue[bool]):
def __call__(self) -> bool:
return self.loop.stop - self.loop.step
RemainingSteps = Annotated[bool, RemainingStepsManager]
libs/langgraph/langgraph/managed/shared_value.py
+21 -22
View File
@@ -7,13 +7,11 @@ from typing import (
Optional,
Sequence,
Type,
cast,
)
from langchain_core.runnables import RunnableConfig
from typing_extensions import NotRequired, Required, Self
from langgraph.constants import CONF, CONFIG_KEY_STORE
from langgraph.constants import CONF
from langgraph.errors import InvalidUpdateError
from langgraph.managed.base import (
ChannelKeyPlaceholder,
@@ -21,7 +19,8 @@ from langgraph.managed.base import (
ConfiguredManagedValue,
WritableManagedValue,
)
from langgraph.store.base import BaseStore, PutOp
from langgraph.store.base import PutOp
from langgraph.types import LoopProtocol
V = dict[str, Any]
@@ -55,25 +54,26 @@ class SharedValue(WritableManagedValue[Value, Update]):
@classmethod
@contextmanager
def enter(cls, config: RunnableConfig, **kwargs: Any) -> Iterator[Self]:
with super().enter(config, **kwargs) as value:
if value.store is not None:
saved = value.store.search(value.ns)
def enter(cls, loop: LoopProtocol, **kwargs: Any) -> Iterator[Self]:
with super().enter(loop, **kwargs) as value:
if loop.store is not None:
saved = loop.store.search(value.ns)
value.value = {it.key: it.value for it in saved}
yield value
@classmethod
@asynccontextmanager
async def aenter(cls, config: RunnableConfig, **kwargs: Any) -> AsyncIterator[Self]:
async with super().aenter(config, **kwargs) as value:
if value.store is not None:
saved = await value.store.asearch(value.ns)
async def aenter(cls, loop: LoopProtocol, **kwargs: Any) -> AsyncIterator[Self]:
async with super().aenter(loop, **kwargs) as value:
if loop.store is not None:
saved = await loop.store.asearch(value.ns)
value.value = {it.key: it.value for it in saved}
yield value
def __init__(
self, config: RunnableConfig, *, typ: Type[Any], scope: str, key: str
self, loop: LoopProtocol, *, typ: Type[Any], scope: str, key: str
) -> None:
super().__init__(loop)
if typ := _strip_extras(typ):
if typ not in (
dict,
@@ -83,18 +83,17 @@ class SharedValue(WritableManagedValue[Value, Update]):
raise ValueError("SharedValue must be a dict")
self.scope = scope
self.value: Value = {}
self.store = cast(BaseStore, config[CONF].get(CONFIG_KEY_STORE))
if self.store is None:
if self.loop.store is None:
pass
elif scope_value := config[CONF].get(self.scope):
elif scope_value := self.loop.config[CONF].get(self.scope):
self.ns = ("scoped", scope, key, scope_value)
else:
raise ValueError(
f"Scope {scope} for shared state key not in config.configurable"
)
def __call__(self, step: int) -> Value:
return self.value.copy()
def __call__(self) -> Value:
return self.value
def _process_update(self, values: Sequence[Update]) -> list[PutOp]:
writes: list[PutOp] = []
@@ -112,13 +111,13 @@ class SharedValue(WritableManagedValue[Value, Update]):
return writes
def update(self, values: Sequence[Update]) -> None:
if self.store is None:
if self.loop.store is None:
self._process_update(values)
else:
return self.store.batch(self._process_update(values))
return self.loop.store.batch(self._process_update(values))
async def aupdate(self, writes: Sequence[Update]) -> None:
if self.store is None:
if self.loop.store is None:
self._process_update(writes)
else:
return await self.store.abatch(self._process_update(writes))
return await self.loop.store.abatch(self._process_update(writes))
libs/langgraph/langgraph/prebuilt/chat_agent_executor.py
+45 -7
View File
@@ -14,7 +14,7 @@ from langgraph._api.deprecation import deprecated_parameter
from langgraph.graph import StateGraph
from langgraph.graph.graph import CompiledGraph
from langgraph.graph.message import add_messages
from langgraph.managed import IsLastStep
from langgraph.managed import IsLastStep, RemainingSteps
from langgraph.prebuilt.tool_executor import ToolExecutor
from langgraph.prebuilt.tool_node import ToolNode
from langgraph.store.base import BaseStore
@@ -33,6 +33,8 @@ class AgentState(TypedDict):
is_last_step: IsLastStep
remaining_steps: RemainingSteps
StateSchema = TypeVar("StateSchema", bound=AgentState)
StateSchemaType = Type[StateSchema]
@@ -529,10 +531,28 @@ def create_react_agent(
# Define the function that calls the model
def call_model(state: AgentState, config: RunnableConfig) -> AgentState:
response = model_runnable.invoke(state, config)
has_tool_calls = isinstance(response, AIMessage) and response.tool_calls
all_tools_return_direct = (
all(call["name"] in should_return_direct for call in response.tool_calls)
if isinstance(response, AIMessage)
else False
)
if (
state["is_last_step"]
and isinstance(response, AIMessage)
and response.tool_calls
(
"remaining_steps" not in state
and state["is_last_step"]
and has_tool_calls
)
or (
"remaining_steps" in state
and state["remaining_steps"] < 1
and all_tools_return_direct
)
or (
"remaining_steps" in state
and state["remaining_steps"] < 2
and has_tool_calls
)
):
return {
"messages": [
@@ -547,10 +567,28 @@ def create_react_agent(
async def acall_model(state: AgentState, config: RunnableConfig) -> AgentState:
response = await model_runnable.ainvoke(state, config)
has_tool_calls = isinstance(response, AIMessage) and response.tool_calls
all_tools_return_direct = (
all(call["name"] in should_return_direct for call in response.tool_calls)
if isinstance(response, AIMessage)
else False
)
if (
state["is_last_step"]
and isinstance(response, AIMessage)
and response.tool_calls
(
"remaining_steps" not in state
and state["is_last_step"]
and has_tool_calls
)
or (
"remaining_steps" in state
and state["remaining_steps"] < 1
and all_tools_return_direct
)
or (
"remaining_steps" in state
and state["remaining_steps"] < 2
and has_tool_calls
)
):
return {
"messages": [
libs/langgraph/langgraph/pregel/__init__.py
+27 -5
View File
@@ -88,7 +88,7 @@ from langgraph.pregel.utils import find_subgraph_pregel, get_new_channel_version
from langgraph.pregel.validate import validate_graph, validate_keys
from langgraph.pregel.write import ChannelWrite, ChannelWriteEntry
from langgraph.store.base import BaseStore
from langgraph.types import All, Checkpointer, StateSnapshot, StreamMode
from langgraph.types import All, Checkpointer, LoopProtocol, StateSnapshot, StreamMode
from langgraph.utils.config import (
ensure_config,
merge_configs,
@@ -433,7 +433,14 @@ class Pregel(Runnable[Union[dict[str, Any], Any], Union[dict[str, Any], Any]]):
)
with ChannelsManager(
self.channels, saved.checkpoint, saved.config, skip_context=True
self.channels,
saved.checkpoint,
LoopProtocol(
config=saved.config,
step=saved.metadata.get("step", -1) + 1,
stop=saved.metadata.get("step", -1) + 2,
),
skip_context=True,
) as (channels, managed):
# tasks for this checkpoint
next_tasks = prepare_next_tasks(
@@ -511,7 +518,14 @@ class Pregel(Runnable[Union[dict[str, Any], Any], Union[dict[str, Any], Any]]):
)
async with AsyncChannelsManager(
self.channels, saved.checkpoint, saved.config, skip_context=True
self.channels,
saved.checkpoint,
LoopProtocol(
config=saved.config,
step=saved.metadata.get("step", -1) + 1,
stop=saved.metadata.get("step", -1) + 2,
),
skip_context=True,
) as (
channels,
managed,
@@ -835,7 +849,11 @@ class Pregel(Runnable[Union[dict[str, Any], Any], Union[dict[str, Any], Any]]):
if as_node not in self.nodes:
raise InvalidUpdateError(f"Node {as_node} does not exist")
# update channels
with ChannelsManager(self.channels, checkpoint, config) as (
with ChannelsManager(
self.channels,
checkpoint,
LoopProtocol(config=config, step=step + 1, stop=step + 2),
) as (
channels,
managed,
):
@@ -981,7 +999,11 @@ class Pregel(Runnable[Union[dict[str, Any], Any], Union[dict[str, Any], Any]]):
if as_node not in self.nodes:
raise InvalidUpdateError(f"Node {as_node} does not exist")
# update channels, acting as the chosen node
async with AsyncChannelsManager(self.channels, checkpoint, config) as (
async with AsyncChannelsManager(
self.channels,
checkpoint,
LoopProtocol(config=config, step=step + 1, stop=step + 2),
) as (
channels,
managed,
):
libs/langgraph/langgraph/pregel/algo.py
+5 -8
View File
@@ -57,7 +57,7 @@ from langgraph.pregel.log import logger
from langgraph.pregel.manager import ChannelsManager
from langgraph.pregel.read import PregelNode
from langgraph.store.base import BaseStore
from langgraph.types import All, PregelExecutableTask, PregelTask
from langgraph.types import All, LoopProtocol, PregelExecutableTask, PregelTask
from langgraph.utils.config import merge_configs, patch_config
GetNextVersion = Callable[[Optional[V], BaseChannel], V]
@@ -148,7 +148,7 @@ def local_read(
with ChannelsManager(
{k: v for k, v in channels.items() if k in updated},
checkpoint,
config,
LoopProtocol(config=config, step=step, stop=step + 1),
skip_context=True,
) as (local_channels, _):
apply_writes(copy_checkpoint(checkpoint), local_channels, [task], None)
@@ -156,7 +156,7 @@ def local_read(
else:
values = read_channels(channels, select)
if managed_keys:
values.update({k: managed[k](step) for k in managed_keys})
values.update({k: managed[k]() for k in managed_keys})
return values
@@ -493,9 +493,7 @@ def prepare_single_task(
):
try:
val = next(
_proc_input(
step, proc, managed, channels, for_execution=for_execution
)
_proc_input(proc, managed, channels, for_execution=for_execution)
)
except StopIteration:
return
@@ -583,7 +581,6 @@ def prepare_single_task(
def _proc_input(
step: int,
proc: PregelNode,
managed: ManagedValueMapping,
channels: Mapping[str, BaseChannel],
@@ -605,7 +602,7 @@ def _proc_input(
except EmptyChannelError:
continue
else:
val[k] = managed[k](step)
val[k] = managed[k]()
except EmptyChannelError:
return
elif isinstance(proc.channels, list):
libs/langgraph/langgraph/pregel/loop.py
+11 -28
View File
@@ -100,7 +100,7 @@ from langgraph.pregel.manager import AsyncChannelsManager, ChannelsManager
from langgraph.pregel.read import PregelNode
from langgraph.pregel.utils import get_new_channel_versions
from langgraph.store.base import BaseStore
from langgraph.types import All, PregelExecutableTask, StreamMode
from langgraph.types import All, LoopProtocol, PregelExecutableTask, StreamProtocol
from langgraph.utils.config import patch_configurable
V = TypeVar("V")
@@ -112,22 +112,6 @@ INPUT_RESUMING = object()
SPECIAL_CHANNELS = (ERROR, INTERRUPT, SCHEDULED)
class StreamProtocol:
__slots__ = ("modes", "__call__")
modes: set[StreamMode]
__call__: Callable[[StreamChunk], None]
def __init__(
self,
__call__: Callable[[StreamChunk], None],
modes: set[StreamMode],
) -> None:
self.__call__ = __call__
self.modes = modes
def DuplexStream(*streams: StreamProtocol) -> StreamProtocol:
def __call__(value: StreamChunk) -> None:
for stream in streams:
@@ -137,16 +121,13 @@ def DuplexStream(*streams: StreamProtocol) -> StreamProtocol:
return StreamProtocol(__call__, {mode for s in streams for mode in s.modes})
class PregelLoop:
class PregelLoop(LoopProtocol):
input: Optional[Any]
config: RunnableConfig
store: Optional[BaseStore]
checkpointer: Optional[BaseCheckpointSaver]
nodes: Mapping[str, PregelNode]
specs: Mapping[str, Union[BaseChannel, ManagedValueSpec]]
output_keys: Union[str, Sequence[str]]
stream_keys: Union[str, Sequence[str]]
stream: Optional[StreamProtocol]
skip_done_tasks: bool
is_nested: bool
@@ -177,8 +158,6 @@ class PregelLoop:
checkpoint_previous_versions: dict[str, Union[str, float, int]]
prev_checkpoint_config: Optional[RunnableConfig]
step: int
stop: int
status: Literal[
"pending", "done", "interrupt_before", "interrupt_after", "out_of_steps"
]
@@ -202,10 +181,14 @@ class PregelLoop:
check_subgraphs: bool = True,
debug: bool = False,
) -> None:
self.stream = stream
super().__init__(
step=0,
stop=0,
config=config,
stream=stream,
store=store,
)
self.input = input
self.config = config
self.store = store
self.checkpointer = checkpointer
self.nodes = nodes
self.specs = specs
@@ -730,7 +713,7 @@ class SyncPregelLoop(PregelLoop, ContextManager):
self.submit = self.stack.enter_context(BackgroundExecutor(self.config))
self.channels, self.managed = self.stack.enter_context(
ChannelsManager(self.specs, self.checkpoint, self.config, self.store)
ChannelsManager(self.specs, self.checkpoint, self)
)
self.stack.push(self._suppress_interrupt)
self.status = "pending"
@@ -858,7 +841,7 @@ class AsyncPregelLoop(PregelLoop, AsyncContextManager):
self.submit = await self.stack.enter_async_context(AsyncBackgroundExecutor())
self.channels, self.managed = await self.stack.enter_async_context(
AsyncChannelsManager(self.specs, self.checkpoint, self.config, self.store)
AsyncChannelsManager(self.specs, self.checkpoint, self)
)
self.stack.push(self._suppress_interrupt)
self.status = "pending"
libs/langgraph/langgraph/pregel/manager.py
+8 -16
View File
@@ -1,33 +1,27 @@
import asyncio
from contextlib import AsyncExitStack, ExitStack, asynccontextmanager, contextmanager
from typing import AsyncIterator, Iterator, Mapping, Optional, Union
from langchain_core.runnables import RunnableConfig
from typing import AsyncIterator, Iterator, Mapping, Union
from langgraph.channels.base import BaseChannel
from langgraph.checkpoint.base import Checkpoint
from langgraph.constants import CONFIG_KEY_STORE
from langgraph.managed.base import (
ConfiguredManagedValue,
ManagedValueMapping,
ManagedValueSpec,
)
from langgraph.managed.context import Context
from langgraph.store.base import BaseStore
from langgraph.utils.config import patch_configurable
from langgraph.types import LoopProtocol
@contextmanager
def ChannelsManager(
specs: Mapping[str, Union[BaseChannel, ManagedValueSpec]],
checkpoint: Checkpoint,
config: RunnableConfig,
store: Optional[BaseStore] = None,
loop: LoopProtocol,
*,
skip_context: bool = False,
) -> Iterator[tuple[Mapping[str, BaseChannel], ManagedValueMapping]]:
"""Manage channels for the lifetime of a Pregel invocation (multiple steps)."""
config_for_managed = patch_configurable(config, {CONFIG_KEY_STORE: store})
channel_specs: dict[str, BaseChannel] = {}
managed_specs: dict[str, ManagedValueSpec] = {}
for k, v in specs.items():
@@ -48,9 +42,9 @@ def ChannelsManager(
ManagedValueMapping(
{
key: stack.enter_context(
value.cls.enter(config_for_managed, **value.kwargs)
value.cls.enter(loop, **value.kwargs)
if isinstance(value, ConfiguredManagedValue)
else value.enter(config_for_managed)
else value.enter(loop)
)
for key, value in managed_specs.items()
}
@@ -62,13 +56,11 @@ def ChannelsManager(
async def AsyncChannelsManager(
specs: Mapping[str, Union[BaseChannel, ManagedValueSpec]],
checkpoint: Checkpoint,
config: RunnableConfig,
store: Optional[BaseStore] = None,
loop: LoopProtocol,
*,
skip_context: bool = False,
) -> AsyncIterator[tuple[Mapping[str, BaseChannel], ManagedValueMapping]]:
"""Manage channels for the lifetime of a Pregel invocation (multiple steps)."""
config_for_managed = patch_configurable(config, {CONFIG_KEY_STORE: store})
channel_specs: dict[str, BaseChannel] = {}
managed_specs: dict[str, ManagedValueSpec] = {}
for k, v in specs.items():
@@ -85,9 +77,9 @@ async def AsyncChannelsManager(
if tasks := {
asyncio.create_task(
stack.enter_async_context(
value.cls.aenter(config_for_managed, **value.kwargs)
value.cls.aenter(loop, **value.kwargs)
if isinstance(value, ConfiguredManagedValue)
else value.aenter(config_for_managed)
else value.aenter(loop)
)
): key
for key, value in managed_specs.items()
libs/langgraph/langgraph/types.py
+46
View File
@@ -1,6 +1,7 @@
from collections import deque
from dataclasses import dataclass
from typing import (
TYPE_CHECKING,
Any,
Callable,
Literal,
@@ -15,6 +16,9 @@ from langchain_core.runnables import Runnable, RunnableConfig
from langgraph.checkpoint.base import BaseCheckpointSaver, CheckpointMetadata
if TYPE_CHECKING:
from langgraph.store.base import BaseStore
All = Literal["*"]
"""Special value to indicate that graph should interrupt on all nodes."""
@@ -213,3 +217,45 @@ class Send:
and self.node == value.node
and self.arg == value.arg
)
StreamChunk = tuple[tuple[str, ...], str, Any]
class StreamProtocol:
__slots__ = ("modes", "__call__")
modes: set[StreamMode]
__call__: Callable[[StreamChunk], None]
def __init__(
self,
__call__: Callable[[StreamChunk], None],
modes: set[StreamMode],
) -> None:
self.__call__ = __call__
self.modes = modes
class LoopProtocol:
config: RunnableConfig
store: Optional["BaseStore"]
stream: Optional[StreamProtocol]
step: int
stop: int
def __init__(
self,
*,
step: int,
stop: int,
config: RunnableConfig,
store: Optional["BaseStore"] = None,
stream: Optional[StreamProtocol] = None,
) -> None:
self.stream = stream
self.config = config
self.store = store
self.step = step
self.stop = stop
libs/scheduler-kafka/langgraph/scheduler/kafka/executor.py
+17 -3
View File
@@ -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,14 @@ 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,
),
) as (channels, managed), AsyncBackgroundExecutor() as submit:
if task := await asyncio.to_thread(
prepare_single_task,
@@ -379,7 +386,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"],