from __future__ import annotations import asyncio import logging from collections.abc import AsyncIterator, Callable, Iterable, Sequence from contextlib import asynccontextmanager from types import TracebackType from typing import Any, cast import orjson from langgraph.store.base import ( GetOp, ListNamespacesOp, Op, PutOp, Result, SearchOp, ) from langgraph.store.base.batch import AsyncBatchedBaseStore from psycopg import AsyncConnection, AsyncCursor, AsyncPipeline, Capabilities from psycopg.rows import DictRow, dict_row from psycopg_pool import AsyncConnectionPool from langgraph.checkpoint.postgres import _ainternal from langgraph.store.postgres.base import ( PLACEHOLDER, BasePostgresStore, PoolConfig, PostgresIndexConfig, Row, TTLConfig, _decode_ns_bytes, _ensure_index_config, _group_ops, _row_to_item, _row_to_search_item, ) logger = logging.getLogger(__name__) class AsyncPostgresStore(AsyncBatchedBaseStore, BasePostgresStore[_ainternal.Conn]): """Asynchronous Postgres-backed store with optional vector search using pgvector. !!! example "Examples" Basic setup and usage: ```python from langgraph.store.postgres import AsyncPostgresStore conn_string = "postgresql://user:pass@localhost:5432/dbname" async with AsyncPostgresStore.from_conn_string(conn_string) as store: await store.setup() # Run migrations. Done once # Store and retrieve data await store.aput(("users", "123"), "prefs", {"theme": "dark"}) item = await store.aget(("users", "123"), "prefs") ``` Vector search using LangChain embeddings: ```python from langchain.embeddings import init_embeddings from langgraph.store.postgres import AsyncPostgresStore conn_string = "postgresql://user:pass@localhost:5432/dbname" async with AsyncPostgresStore.from_conn_string( conn_string, index={ "dims": 1536, "embed": init_embeddings("openai:text-embedding-3-small"), "fields": ["text"] # specify which fields to embed. Default is the whole serialized value } ) as store: await store.setup() # Run migrations. Done once # Store documents await store.aput(("docs",), "doc1", {"text": "Python tutorial"}) await store.aput(("docs",), "doc2", {"text": "TypeScript guide"}) await store.aput(("docs",), "doc3", {"text": "Other guide"}, index=False) # don't index # Search by similarity results = await store.asearch(("docs",), query="programming guides", limit=2) ``` Using connection pooling for better performance: ```python from langgraph.store.postgres import AsyncPostgresStore, PoolConfig conn_string = "postgresql://user:pass@localhost:5432/dbname" async with AsyncPostgresStore.from_conn_string( conn_string, pool_config=PoolConfig( min_size=5, max_size=20 ) ) as store: await store.setup() # Run migrations. Done once # Use store with connection pooling... ``` Warning: Make sure to: 1. Call `setup()` before first use to create necessary tables and indexes 2. Have the pgvector extension available to use vector search 3. Use Python 3.10+ for async functionality Note: Semantic search is disabled by default. You can enable it by providing an `index` configuration when creating the store. Without this configuration, all `index` arguments passed to `put` or `aput` will have no effect. Note: If you provide a TTL configuration, you must explicitly call `start_ttl_sweeper()` to begin the background task that removes expired items. Call `stop_ttl_sweeper()` to properly clean up resources when you're done with the store. """ __slots__ = ( "_deserializer", "pipe", "lock", "supports_pipeline", "index_config", "embeddings", "ttl_config", "_ttl_sweeper_task", "_ttl_stop_event", ) supports_ttl: bool = True def __init__( self, conn: _ainternal.Conn, *, pipe: AsyncPipeline | None = None, deserializer: Callable[[bytes | orjson.Fragment], dict[str, Any]] | None = None, index: PostgresIndexConfig | None = None, ttl: TTLConfig | None = None, ) -> None: if isinstance(conn, AsyncConnectionPool) and pipe is not None: raise ValueError( "Pipeline should be used only with a single AsyncConnection, not AsyncConnectionPool." ) super().__init__() self._deserializer = deserializer self.conn = conn self.pipe = pipe self.lock = asyncio.Lock() self.loop = asyncio.get_running_loop() self.supports_pipeline = Capabilities().has_pipeline() self.index_config = index if self.index_config: self.embeddings, self.index_config = _ensure_index_config(self.index_config) else: self.embeddings = None self.ttl_config = ttl self._ttl_sweeper_task: asyncio.Task[None] | None = None self._ttl_stop_event = asyncio.Event() async def abatch(self, ops: Iterable[Op]) -> list[Result]: grouped_ops, num_ops = _group_ops(ops) results: list[Result] = [None] * num_ops if self.pipe: async with self.pipe: await self._execute_batch(grouped_ops, results) else: await self._execute_batch(grouped_ops, results) return results @classmethod @asynccontextmanager async def from_conn_string( cls, conn_string: str, *, pipeline: bool = False, pool_config: PoolConfig | None = None, index: PostgresIndexConfig | None = None, ttl: TTLConfig | None = None, ) -> AsyncIterator[AsyncPostgresStore]: """Create a new AsyncPostgresStore instance from a connection string. Args: conn_string: The Postgres connection info string. pipeline: Whether to use AsyncPipeline (only for single connections) pool_config: Configuration for the connection pool. If provided, will create a connection pool and use it instead of a single connection. This overrides the `pipeline` argument. index: The embedding config. Returns: AsyncPostgresStore: A new AsyncPostgresStore instance. """ if pool_config is not None: pc = pool_config.copy() async with cast( AsyncConnectionPool[AsyncConnection[DictRow]], AsyncConnectionPool( conn_string, min_size=pc.pop("min_size", 1), max_size=pc.pop("max_size", None), kwargs={ "autocommit": True, "prepare_threshold": 0, "row_factory": dict_row, **(pc.pop("kwargs", None) or {}), }, **cast(dict, pc), ), ) as pool: yield cls(conn=pool, index=index, ttl=ttl) else: async with await AsyncConnection.connect( conn_string, autocommit=True, prepare_threshold=0, row_factory=dict_row ) as conn: if pipeline: async with conn.pipeline() as pipe: yield cls(conn=conn, pipe=pipe, index=index, ttl=ttl) else: yield cls(conn=conn, index=index, ttl=ttl) async def setup(self) -> None: """Set up the store database asynchronously. This method creates the necessary tables in the Postgres database if they don't already exist and runs database migrations. It MUST be called directly by the user the first time the store is used. """ async def _get_version(cur: AsyncCursor[DictRow], table: str) -> int: await cur.execute( f""" CREATE TABLE IF NOT EXISTS {table} ( v INTEGER PRIMARY KEY ) """ ) await cur.execute(f"SELECT v FROM {table} ORDER BY v DESC LIMIT 1") row = cast(dict, await cur.fetchone()) if row is None: version = -1 else: version = row["v"] return version async with self._cursor() as cur: version = await _get_version(cur, table="store_migrations") for v, sql in enumerate(self.MIGRATIONS[version + 1 :], start=version + 1): await cur.execute(sql) await cur.execute("INSERT INTO store_migrations (v) VALUES (%s)", (v,)) if self.index_config: version = await _get_version(cur, table="vector_migrations") for v, migration in enumerate( self.VECTOR_MIGRATIONS[version + 1 :], start=version + 1 ): sql = migration.sql if migration.params: params = { k: v(self) if v is not None and callable(v) else v for k, v in migration.params.items() } if "dims" in params: try: params["dims"] = int(params["dims"]) except Exception as e: raise ValueError( f"Invalid dims for vector index: {params['dims']}" ) from e if "vector_type" in params: vt = str(params["vector_type"]) if vt not in ("vector", "halfvec"): raise ValueError( f"Invalid vector_type for pgvector: {vt}" ) params["vector_type"] = vt if "index_type" in params: it = str(params["index_type"]) if it not in ("hnsw", "ivfflat"): raise ValueError( f"Invalid index_type for pgvector: {it}" ) params["index_type"] = it sql = sql % params await cur.execute(sql) await cur.execute( "INSERT INTO vector_migrations (v) VALUES (%s)", (v,) ) async def sweep_ttl(self) -> int: """Delete expired store items based on TTL. Returns: int: The number of deleted items. """ async with self._cursor() as cur: await cur.execute( """ DELETE FROM store WHERE expires_at IS NOT NULL AND expires_at < NOW() """ ) deleted_count = cur.rowcount return deleted_count async def start_ttl_sweeper( self, sweep_interval_minutes: int | None = None ) -> asyncio.Task[None]: """Periodically delete expired store items based on TTL. Returns: Task that can be awaited or cancelled. """ if not self.ttl_config: return asyncio.create_task(asyncio.sleep(0)) if self._ttl_sweeper_task is not None and not self._ttl_sweeper_task.done(): return self._ttl_sweeper_task self._ttl_stop_event.clear() interval = float( sweep_interval_minutes or self.ttl_config.get("sweep_interval_minutes") or 5 ) logger.info(f"Starting store TTL sweeper with interval {interval} minutes") async def _sweep_loop() -> None: while not self._ttl_stop_event.is_set(): try: try: await asyncio.wait_for( self._ttl_stop_event.wait(), timeout=interval * 60, ) break except asyncio.TimeoutError: pass expired_items = await self.sweep_ttl() if expired_items > 0: logger.info(f"Store swept {expired_items} expired items") except asyncio.CancelledError: break except Exception as exc: logger.exception("Store TTL sweep iteration failed", exc_info=exc) task = asyncio.create_task(_sweep_loop()) task.set_name("ttl_sweeper") self._ttl_sweeper_task = task return task async def stop_ttl_sweeper(self, timeout: float | None = None) -> bool: """Stop the TTL sweeper task if it's running. Args: timeout: Maximum time to wait for the task to stop, in seconds. If `None`, wait indefinitely. Returns: bool: True if the task was successfully stopped or wasn't running, False if the timeout was reached before the task stopped. """ if self._ttl_sweeper_task is None or self._ttl_sweeper_task.done(): return True logger.info("Stopping TTL sweeper task") self._ttl_stop_event.set() if timeout is not None: try: await asyncio.wait_for(self._ttl_sweeper_task, timeout=timeout) success = True except asyncio.TimeoutError: success = False else: await self._ttl_sweeper_task success = True if success: self._ttl_sweeper_task = None logger.info("TTL sweeper task stopped") else: logger.warning("Timed out waiting for TTL sweeper task to stop") return success async def __aenter__(self) -> AsyncPostgresStore: return self async def __aexit__( self, exc_type: type[BaseException] | None, exc_val: BaseException | None, exc_tb: TracebackType | None, ) -> None: # Ensure the TTL sweeper task is stopped when exiting the context if hasattr(self, "_ttl_sweeper_task") and self._ttl_sweeper_task is not None: # Set the event to signal the task to stop self._ttl_stop_event.set() # We don't wait for the task to complete here to avoid blocking # The task will clean up itself gracefully async def _execute_batch( self, grouped_ops: dict, results: list[Result], conn: AsyncConnection[DictRow] | None = None, ) -> None: # Keep `conn` for compatibility with subclasses overriding this private hook. # All database I/O goes through `_cursor()`, which owns connection acquisition. async with self._cursor(pipeline=True) as cur: if GetOp in grouped_ops: await self._batch_get_ops( cast(Sequence[tuple[int, GetOp]], grouped_ops[GetOp]), results, cur, ) if SearchOp in grouped_ops: await self._batch_search_ops( cast(Sequence[tuple[int, SearchOp]], grouped_ops[SearchOp]), results, cur, ) if ListNamespacesOp in grouped_ops: await self._batch_list_namespaces_ops( cast( Sequence[tuple[int, ListNamespacesOp]], grouped_ops[ListNamespacesOp], ), results, cur, ) if PutOp in grouped_ops: await self._batch_put_ops( cast(Sequence[tuple[int, PutOp]], grouped_ops[PutOp]), cur, ) async def _batch_get_ops( self, get_ops: Sequence[tuple[int, GetOp]], results: list[Result], cur: AsyncCursor[DictRow], ) -> None: for query, params, namespace, items in self._get_batch_GET_ops_queries(get_ops): await cur.execute(query, params) rows = cast(list[Row], await cur.fetchall()) key_to_row = {row["key"]: row for row in rows} for idx, key in items: row = key_to_row.get(key) if row: results[idx] = _row_to_item( namespace, row, loader=self._deserializer ) else: results[idx] = None async def _batch_put_ops( self, put_ops: Sequence[tuple[int, PutOp]], cur: AsyncCursor[DictRow], ) -> None: queries, embedding_request = self._prepare_batch_PUT_queries(put_ops) if embedding_request: if self.embeddings is None: # Should not get here since the embedding config is required # to return an embedding_request above raise ValueError( "Embedding configuration is required for vector operations " f"(for semantic search). " f"Please provide an EmbeddingConfig when initializing the {self.__class__.__name__}." ) query, txt_params = embedding_request vectors = await self.embeddings.aembed_documents( [param[-1] for param in txt_params] ) queries.append( ( query, [ p for (ns, k, pathname, _), vector in zip( txt_params, vectors, strict=False ) for p in (ns, k, pathname, vector) ], ) ) for query, params in queries: await cur.execute(query, params) async def _batch_search_ops( self, search_ops: Sequence[tuple[int, SearchOp]], results: list[Result], cur: AsyncCursor[DictRow], ) -> None: queries, embedding_requests = self._prepare_batch_search_queries(search_ops) if embedding_requests and self.embeddings: vectors = await self.embeddings.aembed_documents( [query for _, query in embedding_requests] ) for (idx, _), vector in zip(embedding_requests, vectors, strict=False): _paramslist = queries[idx][1] for i in range(len(_paramslist)): if _paramslist[i] is PLACEHOLDER: _paramslist[i] = vector for (idx, _), (query, params) in zip(search_ops, queries, strict=False): await cur.execute(query, params) rows = cast(list[Row], await cur.fetchall()) items = [ _row_to_search_item( _decode_ns_bytes(row["prefix"]), row, loader=self._deserializer ) for row in rows ] results[idx] = items async def _batch_list_namespaces_ops( self, list_ops: Sequence[tuple[int, ListNamespacesOp]], results: list[Result], cur: AsyncCursor[DictRow], ) -> None: queries = self._get_batch_list_namespaces_queries(list_ops) for (query, params), (idx, _) in zip(queries, list_ops, strict=False): await cur.execute(query, params) rows = cast(list[dict], await cur.fetchall()) namespaces = [_decode_ns_bytes(row["truncated_prefix"]) for row in rows] results[idx] = namespaces @asynccontextmanager async def _cursor( self, *, pipeline: bool = False ) -> AsyncIterator[AsyncCursor[DictRow]]: """Create a database cursor as a context manager. Args: pipeline: whether to use pipeline for the DB operations inside the context manager. Will be applied regardless of whether the PostgresStore instance was initialized with a pipeline. If pipeline mode is not supported, will fall back to using transaction context manager. """ is_pooled_conn = isinstance(self.conn, AsyncConnectionPool) # With AsyncConnectionPool, each _cursor() call checks out its own connection. # The pool does not hand out the same connection concurrently, so a shared lock # across calls is unnecessary here. lock = asyncio.Lock() if is_pooled_conn else self.lock async with _ainternal.get_connection(self.conn) as conn: if self.pipe: # a connection in pipeline mode can be used concurrently # in multiple threads/coroutines, but only one cursor can be # used at a time try: async with conn.cursor(binary=True, row_factory=dict_row) as cur: yield cur finally: if pipeline: await self.pipe.sync() elif pipeline: # a connection not in pipeline mode can only be used by one # thread/coroutine at a time, so we acquire a lock if self.supports_pipeline: async with ( lock, conn.pipeline(), conn.cursor(binary=True, row_factory=dict_row) as cur, ): yield cur else: async with ( lock, conn.transaction(), conn.cursor(binary=True, row_factory=dict_row) as cur, ): yield cur else: async with ( lock, conn.cursor(binary=True, row_factory=dict_row) as cur, ): yield cur