diff --git a/libs/checkpoint/langgraph/store/base/batch.py b/libs/checkpoint/langgraph/store/base/batch.py index 6cfc11419..94189ad72 100644 --- a/libs/checkpoint/langgraph/store/base/batch.py +++ b/libs/checkpoint/langgraph/store/base/batch.py @@ -1,7 +1,18 @@ import asyncio import functools import weakref -from typing import Any, Callable, Iterable, Literal, Optional, TypeVar, Union +import sys +from typing import ( + Any, + Callable, + Iterable, + Literal, + Optional, + TypeVar, + Union, + AsyncIterator, +) +from contextlib import asynccontextmanager from langgraph.store.base import ( BaseStore, @@ -17,9 +28,15 @@ from langgraph.store.base import ( SearchOp, _validate_namespace, ) +import asyncio +import logging + +logger = logging.getLogger(__name__) F = TypeVar("F", bound=Callable) +SUPPORTS_EXC_NOTES = sys.version_info >= (3, 11) + def _check_loop(func: F) -> F: @functools.wraps(func) @@ -67,7 +84,8 @@ class AsyncBatchedBaseStore(BaseStore): ) -> Optional[Item]: fut = self._loop.create_future() self._aqueue[fut] = GetOp(namespace, key) - return await fut + async with _raise_with_queue_stats(self._aqueue): + return await fut async def asearch( self, @@ -81,7 +99,8 @@ class AsyncBatchedBaseStore(BaseStore): ) -> list[SearchItem]: fut = self._loop.create_future() self._aqueue[fut] = SearchOp(namespace_prefix, filter, limit, offset, query) - return await fut + async with _raise_with_queue_stats(self._aqueue): + return await fut async def aput( self, @@ -93,7 +112,8 @@ class AsyncBatchedBaseStore(BaseStore): _validate_namespace(namespace) fut = self._loop.create_future() self._aqueue[fut] = PutOp(namespace, key, value, index) - return await fut + async with _raise_with_queue_stats(self._aqueue): + return await fut async def adelete( self, @@ -102,7 +122,8 @@ class AsyncBatchedBaseStore(BaseStore): ) -> None: fut = self._loop.create_future() self._aqueue[fut] = PutOp(namespace, key, None) - return await fut + async with _raise_with_queue_stats(self._aqueue): + return await fut async def alist_namespaces( self, @@ -127,7 +148,8 @@ class AsyncBatchedBaseStore(BaseStore): offset=offset, ) self._aqueue[fut] = op - return await fut + async with _raise_with_queue_stats(self._aqueue): + return await fut @_check_loop def batch(self, ops: Iterable[Op]) -> list[Result]: @@ -281,3 +303,17 @@ async def _run( break # remove strong ref to store del s + + +@asynccontextmanager +async def _raise_with_queue_stats( + queue: dict[asyncio.Future, Op] +) -> AsyncIterator[None]: + try: + yield + except asyncio.CancelledError as e: + if SUPPORTS_EXC_NOTES: + e.add_note(f"Queue size: {len(queue)}") + else: + logger.warning(f"Queue size: {len(queue)}") + raise