diff --git a/libs/checkpoint/langgraph/checkpoint/memory/__init__.py b/libs/checkpoint/langgraph/checkpoint/memory/__init__.py index e30c082c7..b11f6e21c 100644 --- a/libs/checkpoint/langgraph/checkpoint/memory/__init__.py +++ b/libs/checkpoint/langgraph/checkpoint/memory/__init__.py @@ -1,4 +1,3 @@ -import asyncio import logging import os import pickle @@ -6,7 +5,6 @@ import random import shutil from collections import defaultdict from contextlib import AbstractAsyncContextManager, AbstractContextManager, ExitStack -from functools import partial from types import TracebackType from typing import Any, AsyncIterator, Dict, Iterator, Optional, Sequence, Tuple, Type @@ -395,9 +393,7 @@ class MemorySaver( Returns: Optional[CheckpointTuple]: The retrieved checkpoint tuple, or None if no matching checkpoint was found. """ - return await asyncio.get_running_loop().run_in_executor( - None, self.get_tuple, config - ) + return self.get_tuple(config) async def alist( self, @@ -418,24 +414,8 @@ class MemorySaver( Yields: AsyncIterator[CheckpointTuple]: An asynchronous iterator of checkpoint tuples. """ - loop = asyncio.get_running_loop() - iter = await loop.run_in_executor( - None, - partial( - self.list, - before=before, - limit=limit, - filter=filter, - ), - config, - ) - while True: - # handling StopIteration exception inside coroutine won't work - # as expected, so using next() with default value to break the loop - if item := await loop.run_in_executor(None, next, iter, None): - yield item - else: - break + for item in self.list(config, filter=filter, before=before, limit=limit): + yield item async def aput( self, @@ -455,9 +435,7 @@ class MemorySaver( Returns: RunnableConfig: The updated config containing the saved checkpoint's timestamp. """ - return await asyncio.get_running_loop().run_in_executor( - None, self.put, config, checkpoint, metadata, new_versions - ) + return self.put(config, checkpoint, metadata, new_versions) async def aput_writes( self, @@ -474,10 +452,8 @@ class MemorySaver( config (RunnableConfig): The config to associate with the writes. writes (List[Tuple[str, Any]]): The writes to save, each as a (channel, value) pair. task_id (str): Identifier for the task creating the writes. + return self.put_writes(config, writes, task_id) """ - return await asyncio.get_running_loop().run_in_executor( - None, self.put_writes, config, writes, task_id - ) def get_next_version(self, current: Optional[str], channel: ChannelProtocol) -> str: if current is None: diff --git a/libs/langgraph/langgraph/pregel/loop.py b/libs/langgraph/langgraph/pregel/loop.py index a8e945edd..cf0716e7c 100644 --- a/libs/langgraph/langgraph/pregel/loop.py +++ b/libs/langgraph/langgraph/pregel/loop.py @@ -311,7 +311,9 @@ class PregelLoop(LoopProtocol): ) -> Optional[PregelExecutableTask]: """Accept a PUSH from a task, potentially returning a new task to start.""" # don't start if we should interrupt *after* the original task - if should_interrupt(self.checkpoint, self.interrupt_after, [task]): + if self.interrupt_after and should_interrupt( + self.checkpoint, self.interrupt_after, [task] + ): self.to_interrupt.append(task) return if pushed := cast( @@ -333,7 +335,9 @@ class PregelLoop(LoopProtocol): ), ): # don't start if we should interrupt *before* the new task - if should_interrupt(self.checkpoint, self.interrupt_before, [pushed]): + if self.interrupt_before and should_interrupt( + self.checkpoint, self.interrupt_before, [pushed] + ): self.to_interrupt.append(pushed) return # produce debug output @@ -409,7 +413,7 @@ class PregelLoop(LoopProtocol): } ) # after execution, check if we should interrupt - if should_interrupt( + if self.interrupt_after and should_interrupt( self.checkpoint, self.interrupt_after, self.tasks.values() ): self.status = "interrupt_after" @@ -481,7 +485,7 @@ class PregelLoop(LoopProtocol): return self.tick(input_keys=input_keys) # before execution, check if we should interrupt - if should_interrupt( + if self.interrupt_before and should_interrupt( self.checkpoint, self.interrupt_before, self.tasks.values() ): self.status = "interrupt_before" diff --git a/libs/langgraph/langgraph/utils/runnable.py b/libs/langgraph/langgraph/utils/runnable.py index ccebba862..7cd6a85b9 100644 --- a/libs/langgraph/langgraph/utils/runnable.py +++ b/libs/langgraph/langgraph/utils/runnable.py @@ -404,12 +404,10 @@ class RunnableSeq(Runnable): config = patch_config( config, callbacks=run_manager.get_child(f"seq:step:{i+1}") ) - context = copy_context() - context.run(_set_config_context, config) if i == 0: - input = context.run(step.invoke, input, config, **kwargs) + input = step.invoke(input, config, **kwargs) else: - input = context.run(step.invoke, input, config) + input = step.invoke(input, config) # finish the root run except BaseException as e: run_manager.on_chain_error(e) @@ -443,16 +441,10 @@ class RunnableSeq(Runnable): config = patch_config( config, callbacks=run_manager.get_child(f"seq:step:{i+1}") ) - context = copy_context() - context.run(_set_config_context, config) if i == 0: - coro = step.ainvoke(input, config, **kwargs) + input = await step.ainvoke(input, config, **kwargs) else: - coro = step.ainvoke(input, config) - if ASYNCIO_ACCEPTS_CONTEXT: - input = await asyncio.create_task(coro, context=context) - else: - input = await asyncio.create_task(coro) + input = await step.ainvoke(input, config) # finish the root run except BaseException as e: await run_manager.on_chain_error(e)