langgraph-checkpoint-postgres 2.0.6__tar.gz → 2.0.8__tar.gz

{langgraph_checkpoint_postgres-2.0.6 → langgraph_checkpoint_postgres-2.0.8}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: langgraph-checkpoint-postgres
-Version: 2.0.6
+Version: 2.0.8
 Summary: Library with a Postgres implementation of LangGraph checkpoint saver.
 Home-page: https://www.github.com/langchain-ai/langgraph
 License: MIT

{langgraph_checkpoint_postgres-2.0.6 → langgraph_checkpoint_postgres-2.0.8}/langgraph/checkpoint/postgres/__init__.py

@@ -5,7 +5,6 @@ from typing import Any, Optional
 
 from langchain_core.runnables import RunnableConfig
 from psycopg import Capabilities, Connection, Cursor, Pipeline
-from psycopg.errors import UndefinedTable
 from psycopg.rows import DictRow, dict_row
 from psycopg.types.json import Jsonb
 from psycopg_pool import ConnectionPool
@@ -76,16 +75,15 @@ class PostgresSaver(BasePostgresSaver):
         the first time checkpointer is used.
         """
         with self._cursor() as cur:
-            try:
-                row = cur.execute(
-                    "SELECT v FROM checkpoint_migrations ORDER BY v DESC LIMIT 1"
-                ).fetchone()
-                if row is None:
-                    version = -1
-                else:
-                    version = row["v"]
-            except UndefinedTable:
+            cur.execute(self.MIGRATIONS[0])
+            results = cur.execute(
+                "SELECT v FROM checkpoint_migrations ORDER BY v DESC LIMIT 1"
+            )
+            row = results.fetchone()
+            if row is None:
                 version = -1
+            else:
+                version = row["v"]
             for v, migration in zip(
                 range(version + 1, len(self.MIGRATIONS)),
                 self.MIGRATIONS[version + 1 :],

{langgraph_checkpoint_postgres-2.0.6 → langgraph_checkpoint_postgres-2.0.8}/langgraph/checkpoint/postgres/aio.py

@@ -5,7 +5,6 @@ from typing import Any, Optional
 
 from langchain_core.runnables import RunnableConfig
 from psycopg import AsyncConnection, AsyncCursor, AsyncPipeline, Capabilities
-from psycopg.errors import UndefinedTable
 from psycopg.rows import DictRow, dict_row
 from psycopg.types.json import Jsonb
 from psycopg_pool import AsyncConnectionPool
@@ -55,7 +54,7 @@ class AsyncPostgresSaver(BasePostgresSaver):
         pipeline: bool = False,
         serde: Optional[SerializerProtocol] = None,
     ) -> AsyncIterator["AsyncPostgresSaver"]:
-        """Create a new PostgresSaver instance from a connection string.
+        """Create a new AsyncPostgresSaver instance from a connection string.
 
         Args:
             conn_string (str): The Postgres connection info string.
@@ -81,17 +80,15 @@ class AsyncPostgresSaver(BasePostgresSaver):
         the first time checkpointer is used.
         """
         async with self._cursor() as cur:
-            try:
-                results = await cur.execute(
-                    "SELECT v FROM checkpoint_migrations ORDER BY v DESC LIMIT 1"
-                )
-                row = await results.fetchone()
-                if row is None:
-                    version = -1
-                else:
-                    version = row["v"]
-            except UndefinedTable:
+            await cur.execute(self.MIGRATIONS[0])
+            results = await cur.execute(
+                "SELECT v FROM checkpoint_migrations ORDER BY v DESC LIMIT 1"
+            )
+            row = await results.fetchone()
+            if row is None:
                 version = -1
+            else:
+                version = row["v"]
             for v, migration in zip(
                 range(version + 1, len(self.MIGRATIONS)),
                 self.MIGRATIONS[version + 1 :],

{langgraph_checkpoint_postgres-2.0.6 → langgraph_checkpoint_postgres-2.0.8}/langgraph/checkpoint/postgres/base.py

@@ -57,6 +57,17 @@ MIGRATIONS = [
     PRIMARY KEY (thread_id, checkpoint_ns, checkpoint_id, task_id, idx)
 );""",
     "ALTER TABLE checkpoint_blobs ALTER COLUMN blob DROP not null;",
+    """
+    """,
+    """
+    CREATE INDEX CONCURRENTLY IF NOT EXISTS checkpoints_thread_id_idx ON checkpoints(thread_id);
+    """,
+    """
+    CREATE INDEX CONCURRENTLY IF NOT EXISTS checkpoint_blobs_thread_id_idx ON checkpoint_blobs(thread_id);
+    """,
+    """
+    CREATE INDEX CONCURRENTLY IF NOT EXISTS checkpoint_writes_thread_id_idx ON checkpoint_writes(thread_id);
+    """,
 ]
 
 SELECT_SQL = f"""

{langgraph_checkpoint_postgres-2.0.6 → langgraph_checkpoint_postgres-2.0.8}/langgraph/store/postgres/aio.py

@@ -6,7 +6,6 @@ from typing import Any, Callable, Optional, Union, cast
 
 import orjson
 from psycopg import AsyncConnection, AsyncCursor, AsyncPipeline, Capabilities
-from psycopg.errors import UndefinedTable
 from psycopg.rows import DictRow, dict_row
 from psycopg_pool import AsyncConnectionPool
 
@@ -37,6 +36,75 @@ 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 key-value storage:
+        ```python
+        from langgraph.store.postgres import AsyncPostgresStore
+
+        async with AsyncPostgresStore.from_conn_string(
+            "postgresql://user:pass@localhost:5432/dbname"
+        ) as store:
+            await store.setup()
+
+            # 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
+
+        async with AsyncPostgresStore.from_conn_string(
+            "postgresql://user:pass@localhost:5432/dbname",
+            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()  # Do this once to run migrations
+
+            # Store documents
+            await store.aput(("docs",), "doc1", {"text": "Python tutorial"})
+            await store.aput(("docs",), "doc2", {"text": "TypeScript guide"})
+            # Don't index the following
+            await store.aput(("docs",), "doc3", {"text": "Other guide"}, index=False)
+
+            # Search by similarity
+            results = await store.asearch(("docs",), query="python programming")
+        ```
+
+        Using connection pooling for better performance:
+        ```python
+        from langgraph.store.postgres import AsyncPostgresStore, PoolConfig
+
+        async with AsyncPostgresStore.from_conn_string(
+            "postgresql://user:pass@localhost:5432/dbname",
+            pool_config=PoolConfig(
+                min_size=5,
+                max_size=20
+            )
+        ) as store:
+            await store.setup()
+            # 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.
+    """
+
     __slots__ = (
         "_deserializer",
         "pipe",
@@ -150,22 +218,19 @@ class AsyncPostgresStore(AsyncBatchedBaseStore, BasePostgresStore[_ainternal.Con
         """
 
         async def _get_version(cur: AsyncCursor[DictRow], table: str) -> int:
-            try:
-                await cur.execute(f"SELECT v FROM {table} ORDER BY v DESC LIMIT 1")
-                row = await cur.fetchone()
-                if row is None:
-                    version = -1
-                else:
-                    version = row["v"]
-            except UndefinedTable:
-                version = -1
-                await cur.execute(
-                    f"""
-                    CREATE TABLE IF NOT EXISTS {table} (
-                        v INTEGER PRIMARY KEY
-                    )
-                """
+            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:

{langgraph_checkpoint_postgres-2.0.6 → langgraph_checkpoint_postgres-2.0.8}/langgraph/store/postgres/base.py

@@ -21,7 +21,6 @@ from typing import (
 
 import orjson
 from psycopg import Capabilities, Connection, Cursor, Pipeline
-from psycopg.errors import UndefinedTable
 from psycopg.rows import DictRow, dict_row
 from psycopg.types.json import Jsonb
 from psycopg_pool import ConnectionPool
@@ -56,6 +55,7 @@ class Migration(NamedTuple):
 
     sql: str
     params: Optional[dict[str, Any]] = None
+    condition: Optional[Callable[["BasePostgresStore"], bool]] = None
 
 
 MIGRATIONS: Sequence[str] = [
@@ -72,7 +72,7 @@ CREATE TABLE IF NOT EXISTS store (
 """,
     """
 -- For faster lookups by prefix
-CREATE INDEX IF NOT EXISTS store_prefix_idx ON store USING btree (prefix text_pattern_ops);
+CREATE INDEX CONCURRENTLY IF NOT EXISTS store_prefix_idx ON store USING btree (prefix text_pattern_ops);
 """,
 ]
 
@@ -104,11 +104,29 @@ CREATE TABLE IF NOT EXISTS store_vectors (
             ),
         },
     ),
-    # TODO: Add an HNSW or IVFFlat index depending on config
-    # First must improve the search query when filtering by
-    # namespace
+    Migration(
+        """
+CREATE INDEX CONCURRENTLY IF NOT EXISTS store_vectors_embedding_idx ON store_vectors 
+    USING %(index_type)s (embedding %(ops)s)%(index_params)s;
+""",
+        condition=lambda store: bool(
+            store.index_config and _get_index_params(store)[0] != "flat"
+        ),
+        params={
+            "index_type": lambda store: _get_index_params(store)[0],
+            "ops": lambda store: _get_vector_type_ops(store),
+            "index_params": lambda store: (
+                " WITH ("
+                + ", ".join(f"{k}={v}" for k, v in _get_index_params(store)[1].items())
+                + ")"
+                if _get_index_params(store)[1]
+                else ""
+            ),
+        },
+    ),
 ]
 
+
 C = TypeVar("C", bound=Union[_pg_internal.Conn, _ainternal.Conn])
 
 
@@ -140,6 +158,8 @@ class PoolConfig(TypedDict, total=False):
 class ANNIndexConfig(TypedDict, total=False):
     """Configuration for vector index in PostgreSQL store."""
 
+    kind: Literal["hnsw", "ivfflat", "flat"]
+    """Type of index to use: 'hnsw' for Hierarchical Navigable Small World, or 'ivfflat' for Inverted File Flat."""
     vector_type: Literal["vector", "halfvec"]
     """Type of vector storage to use.
     Options:
@@ -148,6 +168,35 @@ class ANNIndexConfig(TypedDict, total=False):
     """
 
 
+class HNSWConfig(ANNIndexConfig, total=False):
+    """Configuration for HNSW (Hierarchical Navigable Small World) index."""
+
+    kind: Literal["hnsw"]  # type: ignore[misc]
+    m: int
+    """Maximum number of connections per layer. Default is 16."""
+    ef_construction: int
+    """Size of dynamic candidate list for index construction. Default is 64."""
+
+
+class IVFFlatConfig(ANNIndexConfig, total=False):
+    """IVFFlat index divides vectors into lists, and then searches a subset of those lists that are closest to the query vector. It has faster build times and uses less memory than HNSW, but has lower query performance (in terms of speed-recall tradeoff).
+
+    Three keys to achieving good recall are:
+    1. Create the index after the table has some data
+    2. Choose an appropriate number of lists - a good place to start is rows / 1000 for up to 1M rows and sqrt(rows) for over 1M rows
+    3. When querying, specify an appropriate number of probes (higher is better for recall, lower is better for speed) - a good place to start is sqrt(lists)
+    """
+
+    kind: Literal["ivfflat"]  # type: ignore[misc]
+    nlist: int
+    """Number of inverted lists (clusters) for IVF index.
+    
+    Determines the number of clusters used in the index structure.
+    Higher values can improve search speed but increase index size and build time.
+    Typically set to the square root of the number of vectors in the index.
+    """
+
+
 class PostgresIndexConfig(IndexConfig, total=False):
     """Configuration for vector embeddings in PostgreSQL store with pgvector-specific options.
 
@@ -321,7 +370,7 @@ class BasePostgresStore(Generic[C]):
             if op.query and self.index_config:
                 embedding_requests.append((idx, op.query))
 
-                score_operator = _get_distance_operator(self)
+                score_operator, post_operator = _get_distance_operator(self)
                 vector_type = (
                     cast(PostgresIndexConfig, self.index_config)
                     .get("ann_index_config", {})
@@ -351,18 +400,28 @@ class BasePostgresStore(Generic[C]):
                     if not filter_conditions
                     else " AND " + " AND ".join(filter_conditions)
                 )
+                if op.namespace_prefix:
+                    prefix_filter_str = f"WHERE s.prefix LIKE %s {filter_str} "
+                    ns_args: Sequence = (f"{_namespace_to_text(op.namespace_prefix)}%",)
+                else:
+                    ns_args = ()
+                    if filter_str:
+                        prefix_filter_str = f"WHERE {filter_str} "
+                    else:
+                        prefix_filter_str = ""
+
                 base_query = f"""
                     WITH scored AS (
-                        SELECT s.prefix, s.key, s.value, s.created_at, s.updated_at, {score_operator} AS score
+                        SELECT s.prefix, s.key, s.value, s.created_at, s.updated_at, {score_operator} AS neg_score
                         FROM store s
                         JOIN store_vectors sv ON s.prefix = sv.prefix AND s.key = sv.key
-                        WHERE s.prefix LIKE %s {filter_str}
-                        ORDER BY {score_operator} DESC 
+                        {prefix_filter_str}
+                        ORDER BY {score_operator} ASC 
                         LIMIT %s
                     )
                     SELECT * FROM (
                         SELECT DISTINCT ON (prefix, key) 
-                            prefix, key, value, created_at, updated_at, score 
+                            prefix, key, value, created_at, updated_at, {post_operator} as score 
                         FROM scored 
                         ORDER BY prefix, key, score DESC
                     ) AS unique_docs
@@ -372,7 +431,7 @@ class BasePostgresStore(Generic[C]):
                 """
                 params = [
                     _PLACEHOLDER,  # Vector placeholder
-                    f"{_namespace_to_text(op.namespace_prefix)}%",
+                    *ns_args,
                     *filter_params,
                     _PLACEHOLDER,
                     expanded_limit,
@@ -474,6 +533,58 @@ class BasePostgresStore(Generic[C]):
 
 
 class PostgresStore(BaseStore, BasePostgresStore[_pg_internal.Conn]):
+    """Postgres-backed store with optional vector search using pgvector.
+
+    !!! example "Examples"
+        Basic setup and key-value storage:
+        ```python
+        from langgraph.store.postgres import PostgresStore
+
+        store = PostgresStore(
+            connection_string="postgresql://user:pass@localhost:5432/dbname"
+        )
+        store.setup()
+
+        # Store and retrieve data
+        store.put(("users", "123"), "prefs", {"theme": "dark"})
+        item = store.get(("users", "123"), "prefs")
+        ```
+
+        Vector search using LangChain embeddings:
+        ```python
+        from langchain.embeddings import init_embeddings
+        from langgraph.store.postgres import PostgresStore
+
+        store = PostgresStore(
+            connection_string="postgresql://user:pass@localhost:5432/dbname",
+            index={
+                "dims": 1536,
+                "embed": init_embeddings("openai:text-embedding-3-small"),
+                "fields": ["text"]  # specify which fields to embed. Default is the whole serialized value
+            }
+        )
+        store.setup() # Do this once to run migrations
+
+        # Store documents
+        store.put(("docs",), "doc1", {"text": "Python tutorial"})
+        store.put(("docs",), "doc2", {"text": "TypeScript guide"})
+        store.put(("docs",), "doc2", {"text": "Other guide"}, index=False) # don't index
+
+        # Search by similarity
+        results = store.search(("docs",), query="python programming")
+        ```
+
+    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.
+
+    Warning:
+        Make sure to call `setup()` before first use to create necessary tables and indexes.
+        The pgvector extension must be available to use vector search.
+
+    """
+
     __slots__ = (
         "_deserializer",
         "pipe",
@@ -702,7 +813,6 @@ class PostgresStore(BaseStore, BasePostgresStore[_pg_internal.Conn]):
                         _paramslist[i] = embedding
 
         for (idx, _), (query, params) in zip(search_ops, queries):
-            # Execute the actual query
             cur.execute(query, params)
             rows = cast(list[Row], cur.fetchall())
             results[idx] = [
@@ -736,22 +846,19 @@ class PostgresStore(BaseStore, BasePostgresStore[_pg_internal.Conn]):
         """
 
         def _get_version(cur: Cursor[dict[str, Any]], table: str) -> int:
-            try:
-                cur.execute(f"SELECT v FROM {table} ORDER BY v DESC LIMIT 1")
-                row = cast(dict, cur.fetchone())
-                if row is None:
-                    version = -1
-                else:
-                    version = row["v"]
-            except UndefinedTable:
-                version = -1
-                cur.execute(
-                    f"""
-                    CREATE TABLE IF NOT EXISTS {table} (
-                        v INTEGER PRIMARY KEY
-                    )
-                """
+            cur.execute(
+                f"""
+                CREATE TABLE IF NOT EXISTS {table} (
+                    v INTEGER PRIMARY KEY
                 )
+            """
+            )
+            cur.execute(f"SELECT v FROM {table} ORDER BY v DESC LIMIT 1")
+            row = cast(dict, cur.fetchone())
+            if row is None:
+                version = -1
+            else:
+                version = row["v"]
             return version
 
         with self._cursor() as cur:
@@ -765,6 +872,8 @@ class PostgresStore(BaseStore, BasePostgresStore[_pg_internal.Conn]):
                 for v, migration in enumerate(
                     self.VECTOR_MIGRATIONS[version + 1 :], start=version + 1
                 ):
+                    if migration.condition and not migration.condition(self):
+                        continue
                     sql = migration.sql
                     if migration.params:
                         params = {
@@ -823,6 +932,18 @@ def _get_vector_type_ops(store: BasePostgresStore) -> str:
     return f"{type_prefix}_{distance_suffix}"
 
 
+def _get_index_params(store: Any) -> tuple[str, dict[str, Any]]:
+    """Get the index type and configuration based on config."""
+    if not store.index_config:
+        return "hnsw", {}
+
+    config = cast(PostgresIndexConfig, store.index_config)
+    index_config = config.get("ann_index_config", _DEFAULT_ANN_CONFIG).copy()
+    kind = index_config.pop("kind", "hnsw")
+    index_config.pop("vector_type", None)
+    return kind, index_config
+
+
 def _namespace_to_text(
     namespace: tuple[str, ...], handle_wildcards: bool = False
 ) -> str:
@@ -915,7 +1036,7 @@ def _decode_ns_bytes(namespace: Union[str, bytes, list]) -> tuple[str, ...]:
     return tuple(namespace.split("."))
 
 
-def _get_distance_operator(store: Any) -> str:
+def _get_distance_operator(store: Any) -> tuple[str, str]:
     """Get the distance operator and score expression based on config."""
     # Note: Today, we are not using ANN indices due to restrictions
     # on PGVector's support for mixing vector and non-vector filters
@@ -936,12 +1057,22 @@ def _get_distance_operator(store: Any) -> str:
     config = cast(PostgresIndexConfig, store.index_config)
     distance_type = config.get("distance_type", "cosine")
 
+    # Return the operator and the score expression
+    # The operator is used in the CTE and will be compatible with an ASCENDING ORDER
+    # sort clause.
+    # The score expression is used in the final query and will be compatible with
+    # a DESCENDING ORDER sort clause and the user's expectations of what the similarity score
+    # should be.
     if distance_type == "l2":
-        return "1 - (sv.embedding <-> %s::%s)"
+        # Final: "-(sv.embedding <-> %s::%s)"
+        # We return the "l2 similarity" so that the sorting order is the same
+        return "sv.embedding <-> %s::%s", "-scored.neg_score"
     elif distance_type == "inner_product":
-        return "-(sv.embedding <#> %s::%s)"
-    else:  # cosine
-        return "1 - (sv.embedding <=> %s::%s)"
+        # Final: "-(sv.embedding <#> %s::%s)"
+        return "sv.embedding <#> %s::%s", "-(scored.neg_score)"
+    else:  # cosine similarity
+        # Final:  "1 - (sv.embedding <=> %s::%s)"
+        return "sv.embedding <=> %s::%s", "1 - scored.neg_score"
 
 
 def _ensure_index_config(

{langgraph_checkpoint_postgres-2.0.6 → langgraph_checkpoint_postgres-2.0.8}/pyproject.toml

@@ -1,6 +1,6 @@
 [tool.poetry]
 name = "langgraph-checkpoint-postgres"
-version = "2.0.6"
+version = "2.0.8"
 description = "Library with a Postgres implementation of LangGraph checkpoint saver."
 authors = []
 license = "MIT"