diff --git a/libs/checkpoint-postgres/langgraph/store/postgres/aio.py b/libs/checkpoint-postgres/langgraph/store/postgres/aio.py index 981cc91b4..71d497012 100644 --- a/libs/checkpoint-postgres/langgraph/store/postgres/aio.py +++ b/libs/checkpoint-postgres/langgraph/store/postgres/aio.py @@ -18,6 +18,7 @@ from psycopg.rows import dict_row from langgraph.store.base import GetOp, ListNamespacesOp, Op, PutOp, Result, SearchOp +from langgraph.store.base.batch import AsyncBatchedBaseStore from langgraph.store.postgres.base import ( BasePostgresStore, Row, @@ -29,7 +30,9 @@ logger = logging.getLogger(__name__) -class AsyncPostgresStore(BasePostgresStore[AsyncConnection]): +class AsyncPostgresStore(AsyncBatchedBaseStore, BasePostgresStore[AsyncConnection]): + __slots__ = ("_deserializer",) + def __init__( self, conn: AsyncConnection[Any], @@ -38,7 +41,8 @@ def __init__( Callable[[Union[bytes, orjson.Fragment]], dict[str, Any]] ] = None, ) -> None: - super().__init__(deserializer=deserializer) + super().__init__() + self._deserializer = deserializer self.conn = conn self.conn = conn self.loop = asyncio.get_running_loop() diff --git a/libs/checkpoint-postgres/langgraph/store/postgres/base.py b/libs/checkpoint-postgres/langgraph/store/postgres/base.py index d979c5fbc..8bd2b8279 100644 --- a/libs/checkpoint-postgres/langgraph/store/postgres/base.py +++ b/libs/checkpoint-postgres/langgraph/store/postgres/base.py @@ -59,20 +59,10 @@ C = TypeVar("C", bound=BaseConnection) -class BasePostgresStore(BaseStore, Generic[C]): +class BasePostgresStore(Generic[C]): MIGRATIONS = MIGRATIONS conn: C - __slots__ = ("_deserializer",) - - def __init__( - self, - *, - deserializer: Optional[ - Callable[[Union[bytes, orjson.Fragment]], dict[str, Any]] - ] = None, - ) -> None: - super().__init__() - self._deserializer = deserializer + _deserializer: Optional[Callable[[Union[bytes, orjson.Fragment]], dict[str, Any]]] def _get_batch_GET_ops_queries( self, @@ -166,7 +156,9 @@ def _get_batch_search_queries( params.extend([key, json.dumps(value)]) query += " AND " + " AND ".join(filter_conditions) - query += " LIMIT %s OFFSET %s" + # Note: we will need to not do this if sim/keyword search + # is used + query += " ORDER BY updated_at DESC LIMIT %s OFFSET %s" params.extend([op.limit, op.offset]) queries.append((query, params)) @@ -227,7 +219,9 @@ def _get_batch_list_namespaces_queries( return queries -class PostgresStore(BasePostgresStore[Connection]): +class PostgresStore(BaseStore, BasePostgresStore[Connection]): + __slots__ = ("_deserializer",) + def __init__( self, conn: Connection[Any], @@ -236,7 +230,8 @@ def __init__( Callable[[Union[bytes, orjson.Fragment]], dict[str, Any]] ] = None, ) -> None: - super().__init__(deserializer=deserializer) + super().__init__() + self._deserializer = deserializer self.conn = conn def batch(self, ops: Iterable[Op]) -> list[Result]: diff --git a/libs/checkpoint-postgres/pyproject.toml b/libs/checkpoint-postgres/pyproject.toml index d98c154a7..d687b9a42 100644 --- a/libs/checkpoint-postgres/pyproject.toml +++ b/libs/checkpoint-postgres/pyproject.toml @@ -1,6 +1,6 @@ [tool.poetry] name = "langgraph-checkpoint-postgres" -version = "2.0.0" +version = "2.0.1" description = "Library with a Postgres implementation of LangGraph checkpoint saver." authors = [] license = "MIT" diff --git a/libs/checkpoint/langgraph/store/base/__init__.py b/libs/checkpoint/langgraph/store/base/__init__.py index 8d11a3015..098462339 100644 --- a/libs/checkpoint/langgraph/store/base/__init__.py +++ b/libs/checkpoint/langgraph/store/base/__init__.py @@ -182,7 +182,11 @@ def _validate_namespace(namespace: tuple[str, ...]) -> None: class BaseStore(ABC): - """Abstract base class for key-value stores.""" + """Abstract base class for persistent key-value stores. + + Stores enable persistence and memory that can be shared across threads, + scoped to user IDs, assistant IDs, or other arbitrary namespaces. + """ __slots__ = ("__weakref__",)