langgraph-store-postgres 0.1.0__py3-none-any.whl

langgraph_store_postgres/__init__.py

@@ -0,0 +1,164 @@
+"""PostgreSQL implementation of the LangGraph long-term-memory ``BaseStore``.
+
+``PostgresStore`` persists namespaced key/value memory in a single table and
+supports get / put / delete / search (namespace-prefix + filter) and
+``list_namespaces``. With an ``IndexConfig`` it also does **semantic search**
+natively via `pgvector <https://github.com/pgvector/pgvector>`_ (cosine
+distance, HNSW index). Built on ``langgraph-store-core``.
+"""
+
+from __future__ import annotations
+
+from typing import List, Optional, Tuple
+
+from sqlalchemy import (
+    Column,
+    Index,
+    MetaData,
+    String,
+    Table,
+    create_engine,
+    delete as sa_delete,
+    select,
+    text,
+)
+from sqlalchemy.dialects.postgresql import JSONB, insert as pg_insert
+from sqlalchemy.engine import Engine
+
+from langgraph_store_core import IndexConfig, KVStore
+
+__all__ = ["PostgresStore"]
+
+
+def _escape_like(value: str) -> str:
+    return value.replace("\\", "\\\\").replace("%", "\\%").replace("_", "\\_")
+
+
+class PostgresStore(KVStore):
+    """LangGraph ``BaseStore`` backed by PostgreSQL.
+
+    Example:
+        ```python
+        from langgraph_store_postgres import PostgresStore
+
+        store = PostgresStore("postgresql://user:pass@localhost:5432/db")
+        store.put(("users", "1"), "profile", {"name": "Kamal"})
+        item = store.get(("users", "1"), "profile")
+
+        # semantic search (needs the pgvector extension on the server)
+        store = PostgresStore(url, index={"dims": 1024, "embed": my_embedder, "fields": ["text"]})
+        hits = store.search(("users", "1"), query="what does the user like?")
+        ```
+
+    When ``index`` is given the table gains an ``embedding vector(dims)`` column
+    and an HNSW cosine index; ``CREATE EXTENSION IF NOT EXISTS vector`` is run
+    (requires the extension to be installed on the server).
+    """
+
+    supports_native_vector_search = True
+
+    def __init__(
+        self,
+        url: Optional[str] = None,
+        *,
+        table_name: str = "langgraph_store",
+        engine: Optional[Engine] = None,
+        index: Optional[IndexConfig] = None,
+    ) -> None:
+        super().__init__(index=index)
+        if engine is None and url is None:
+            raise ValueError("Provide either 'url' or 'engine'")
+        self._engine = engine or create_engine(url)  # type: ignore[arg-type]
+        self._metadata = MetaData()
+        columns = [
+            Column("prefix", String, primary_key=True),
+            Column("key", String, primary_key=True),
+            Column("value", JSONB, nullable=False),
+            Column("created_at", String, nullable=False),
+            Column("updated_at", String, nullable=False),
+        ]
+        if self.dims:
+            from pgvector.sqlalchemy import Vector  # optional dependency
+
+            columns.append(Column("embedding", Vector(self.dims), nullable=True))
+            with self._engine.begin() as conn:
+                conn.execute(text("CREATE EXTENSION IF NOT EXISTS vector"))
+        self._table = Table(table_name, self._metadata, *columns)
+        if self.dims:
+            Index(
+                f"{table_name}_embedding_hnsw",
+                self._table.c.embedding,
+                postgresql_using="hnsw",
+                postgresql_with={"m": 16, "ef_construction": 64},
+                postgresql_ops={"embedding": "vector_cosine_ops"},
+            )
+        self._metadata.create_all(self._engine)
+
+    def _row(self, r) -> dict:
+        emb = getattr(r, "embedding", None) if self.dims else None
+        return {
+            "prefix": r.prefix, "key": r.key, "value": r.value,
+            "created_at": r.created_at, "updated_at": r.updated_at,
+            "embedding": list(emb) if emb is not None else None,
+        }
+
+    def _read(self, prefix: str, key: str) -> Optional[dict]:
+        t = self._table
+        with self._engine.connect() as conn:
+            r = conn.execute(
+                select(t).where(t.c.prefix == prefix, t.c.key == key)
+            ).first()
+        return self._row(r) if r else None
+
+    def _write(self, row: dict) -> None:
+        t = self._table
+        values = {k: row[k] for k in ("prefix", "key", "value", "created_at", "updated_at")}
+        if self.dims:
+            values["embedding"] = row.get("embedding")
+        stmt = pg_insert(t).values(**values).on_conflict_do_update(
+            index_elements=[t.c.prefix, t.c.key],
+            set_={k: v for k, v in values.items() if k not in ("prefix", "key")},
+        )
+        with self._engine.begin() as conn:
+            conn.execute(stmt)
+
+    def _remove(self, prefix: str, key: str) -> None:
+        t = self._table
+        with self._engine.begin() as conn:
+            conn.execute(sa_delete(t).where(t.c.prefix == prefix, t.c.key == key))
+
+    def _prefix_clause(self, prefix: str):
+        t = self._table
+        if prefix == "":
+            return None
+        return t.c.prefix.like(f"{_escape_like(prefix)}%", escape="\\")
+
+    def _scan(self, prefix: str) -> List[dict]:
+        t = self._table
+        stmt = select(t)
+        clause = self._prefix_clause(prefix)
+        if clause is not None:
+            stmt = stmt.where(clause)
+        with self._engine.connect() as conn:
+            rows = conn.execute(stmt).all()
+        return [self._row(r) for r in rows]
+
+    def _vector_search(
+        self, prefix: str, vector: List[float], filter: Optional[dict], limit: int
+    ) -> List[Tuple[dict, float]]:
+        """Native pgvector ANN: cosine distance ordered, prefix + equality filters in SQL."""
+        t = self._table
+        distance = t.c.embedding.cosine_distance(vector).label("distance")
+        stmt = select(t, distance).where(t.c.embedding.isnot(None))
+        clause = self._prefix_clause(prefix)
+        if clause is not None:
+            stmt = stmt.where(clause)
+        # Push plain-equality filters into SQL as JSONB containment (type-safe);
+        # operator filters ($gt, $in, ...) are applied by core on the returned rows.
+        equality = {f: v for f, v in (filter or {}).items() if not isinstance(v, dict)}
+        if equality:
+            stmt = stmt.where(t.c.value.contains(equality))
+        stmt = stmt.order_by(distance).limit(limit)
+        with self._engine.connect() as conn:
+            rows = conn.execute(stmt).all()
+        return [(self._row(r), 1.0 - float(r.distance)) for r in rows]

langgraph_store_postgres-0.1.0.dist-info/METADATA

@@ -0,0 +1,68 @@
+Metadata-Version: 2.5
+Name: langgraph-store-postgres
+Version: 0.1.0
+Summary: PostgreSQL long-term-memory store (BaseStore) for LangGraph — namespaced key/value memory with filters, list_namespaces and native semantic (pgvector) search
+Project-URL: Homepage, https://github.com/skamalj/langgraph-store
+Project-URL: Repository, https://github.com/skamalj/langgraph-store.git
+Author-email: Kamal <skamalj@gmail.com>
+Keywords: agent-memory,basestore,langgraph,long-term-memory,postgres,postgresql,store
+Classifier: License :: OSI Approved :: MIT License
+Classifier: Operating System :: OS Independent
+Classifier: Programming Language :: Python :: 3
+Classifier: Topic :: Scientific/Engineering :: Artificial Intelligence
+Requires-Python: >=3.10
+Requires-Dist: langgraph-store-core>=0.1.0
+Requires-Dist: pgvector>=0.2
+Requires-Dist: psycopg2-binary
+Requires-Dist: sqlalchemy>=2.0
+Description-Content-Type: text/markdown
+
+# langgraph-store-postgres
+
+A **PostgreSQL** long-term-memory store (`BaseStore`) for [LangGraph](https://langchain-ai.github.io/langgraph/) — namespaced key/value memory with prefix search, filters, `list_namespaces`, and **native semantic search via pgvector**.
+
+```bash
+pip install langgraph-store-postgres
+```
+
+```python
+from langgraph_store_postgres import PostgresStore
+
+store = PostgresStore("postgresql://user:pass@localhost:5432/db")   # or engine=<Engine>
+
+store.put(("users", "1", "memories"), "food", {"text": "loves sushi", "kind": "pref"})
+item = store.get(("users", "1", "memories"), "food")
+hits = store.search(("users", "1"), filter={"kind": "pref"}, limit=10)
+spaces = store.list_namespaces(prefix=("users",))
+```
+
+Use it as a LangGraph store: `graph.compile(store=PostgresStore(...))`. Async methods (`aget`/`aput`/`asearch`/…) work too — the sync calls run in a thread.
+
+## Semantic search (pgvector)
+
+Pass a LangGraph `IndexConfig` and the store embeds the configured fields on `put` and ranks `search(query=...)` by cosine similarity using pgvector (`<=>`, HNSW index):
+
+```python
+from langgraph_store_core import bedrock_titan_embeddings
+from langgraph_store_postgres import PostgresStore
+
+store = PostgresStore(url, table_name="memory",
+                      index={"dims": 1024, "embed": bedrock_titan_embeddings(dimensions=1024), "fields": ["text"]})
+store.put(("memories", "kamal"), "k1", {"text": "the user loves sushi", "kind": "pref"})
+hits = store.search(("memories", "kamal"), query="what food does the user like?", filter={"kind": "pref"})
+print(hits[0].score, hits[0].value)
+```
+
+`embed` may be any LangChain `Embeddings`, a `list[str] -> list[list[float]]` callable, or a provider string. `fields` defaults to `["$"]` (whole value as JSON). `put(..., index=False)` skips embedding for one item.
+
+Requires the [pgvector](https://github.com/pgvector/pgvector) extension on the server (`apt install postgresql-16-pgvector`, or built in on RDS / Cloud SQL / Azure Database). The store runs `CREATE EXTENSION IF NOT EXISTS vector`, adds an `embedding vector(dims)` column and an HNSW cosine index. Namespace prefix and plain-equality filters run inside the SQL query (JSONB containment); operator filters (`$gt`, `$in`, …) are applied on the returned candidates.
+
+## Data model
+
+A single table `(prefix, key, value JSONB, created_at, updated_at[, embedding vector])` with primary key `(prefix, key)`, created automatically. Namespace tuples are joined with a unit separator into `prefix`; search is a prefix scan, filters and namespace matching are evaluated in the [core](https://pypi.org/project/langgraph-store-core/).
+
+Docs: <https://skamalj.github.io/agentstate-reducer/>
+
+## License
+
+MIT

langgraph_store_postgres-0.1.0.dist-info/RECORD

@@ -0,0 +1,4 @@
+langgraph_store_postgres/__init__.py,sha256=rVY1t69g2-4YzjlqSY6He0beKwx3T7PAjVV7faF3Hr4,6511
+langgraph_store_postgres-0.1.0.dist-info/METADATA,sha256=B0k8P8K32eTXOJu_H0pQlo-kcrlHJY24r3h7yDcKtP4,3649
+langgraph_store_postgres-0.1.0.dist-info/WHEEL,sha256=THafob7ofN-NsuMN7Mg4qZyHaQI7KkD-QlcQatYhXPo,87
+langgraph_store_postgres-0.1.0.dist-info/RECORD,,

langgraph_store_postgres-0.1.0.dist-info/WHEEL

@@ -0,0 +1,4 @@
+Wheel-Version: 1.0
+Generator: hatchling 1.32.3
+Root-Is-Purelib: true
+Tag: py3-none-any