langgraph-store-postgres 0.1.0__tar.gz

langgraph_store_postgres-0.1.0/.gitignore

@@ -0,0 +1,9 @@
+__pycache__/
+*.py[cod]
+*.egg-info/
+.pytest_cache/
+.venv/
+dist/
+build/
+uv.lock
+.env

langgraph_store_postgres-0.1.0/PKG-INFO

@@ -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/README.md

@@ -0,0 +1,49 @@
+# 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/pyproject.toml

@@ -0,0 +1,37 @@
+[build-system]
+requires = ["hatchling"]
+build-backend = "hatchling.build"
+
+[project]
+name = "langgraph-store-postgres"
+version = "0.1.0"
+description = "PostgreSQL long-term-memory store (BaseStore) for LangGraph — namespaced key/value memory with filters, list_namespaces and native semantic (pgvector) search"
+authors = [{name = "Kamal", email = "skamalj@gmail.com"}]
+readme = "README.md"
+requires-python = ">=3.10"
+dependencies = [
+    "langgraph-store-core>=0.1.0",
+    "sqlalchemy>=2.0",
+    "psycopg2-binary",
+    "pgvector>=0.2",
+]
+keywords = ["langgraph", "store", "basestore", "postgres", "postgresql", "long-term-memory", "agent-memory"]
+classifiers = [
+    "Programming Language :: Python :: 3",
+    "License :: OSI Approved :: MIT License",
+    "Operating System :: OS Independent",
+    "Topic :: Scientific/Engineering :: Artificial Intelligence",
+]
+
+[project.urls]
+Homepage = "https://github.com/skamalj/langgraph-store"
+Repository = "https://github.com/skamalj/langgraph-store.git"
+
+[dependency-groups]
+dev = ["pytest>=7.0", "pytest-asyncio"]
+
+[tool.pytest.ini_options]
+asyncio_mode = "auto"
+
+[tool.hatch.build.targets.wheel]
+packages = ["src/langgraph_store_postgres"]

langgraph_store_postgres-0.1.0/src/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/tests/test_postgres_semantic.py

@@ -0,0 +1,103 @@
+"""Semantic-search tests for PostgresStore (native pgvector path) against a real PostgreSQL with pgvector."""
+import os
+import uuid
+
+import pytest
+from sqlalchemy import create_engine, text
+
+from langgraph_store_core.testing import FakeEmbeddings
+from langgraph_store_postgres import PostgresStore
+
+URL = os.environ.get("SQL_TEST_URL", "postgresql+psycopg2://postgres:postgres@localhost:5433/postgres")
+DIMS = 64
+
+
+def _pgvector_available() -> bool:
+    try:
+        with create_engine(URL).begin() as c:
+            c.execute(text("CREATE EXTENSION IF NOT EXISTS vector"))
+        return True
+    except Exception:
+        return False
+
+
+pytestmark = pytest.mark.skipif(not _pgvector_available(), reason=f"PostgreSQL+pgvector not reachable: {URL}")
+
+
+@pytest.fixture()
+def store():
+    tname = f"lg_vstore_test_{uuid.uuid4().hex[:8]}"
+    s = PostgresStore(url=URL, table_name=tname, index={"dims": DIMS, "embed": FakeEmbeddings(DIMS), "fields": ["text"]})
+    yield s
+    with s._engine.begin() as c:
+        c.execute(text(f'DROP TABLE IF EXISTS "{tname}"'))
+    s._engine.dispose()
+
+
+def _seed(s):
+    s.put(("memories", "kamal"), "food", {"text": "the user loves sushi and japanese food", "kind": "pref"})
+    s.put(("memories", "kamal"), "lang", {"text": "the user writes python every day", "kind": "fact"})
+    s.put(("memories", "kamal"), "city", {"text": "the user lives by the sea in a big city", "kind": "fact"})
+    s.put(("memories", "priya"), "food", {"text": "loves sushi", "kind": "pref"})
+
+
+def test_table_has_vector_column_and_index(store):
+    with store._engine.connect() as c:
+        cols = {r[0]: r[1] for r in c.execute(text(
+            f"select column_name, udt_name from information_schema.columns where table_name='{store._table.name}'"))}
+        assert cols.get("embedding") == "vector"
+        idx = c.execute(text(f"select indexdef from pg_indexes where tablename='{store._table.name}'")).all()
+        assert any("hnsw" in r[0] for r in idx)
+
+
+def test_put_embeds_and_index_false_skips(store):
+    store.put(("n",), "k", {"text": "hello world"})
+    assert len(store._read("n", "k")["embedding"]) == DIMS
+    store.put(("n",), "k2", {"text": "hello"}, index=False)
+    assert store._read("n", "k2")["embedding"] is None
+
+
+def test_query_ranks_and_scores(store):
+    _seed(store)
+    hits = store.search(("memories", "kamal"), query="sushi and japanese food")
+    assert hits[0].key == "food" and hits[0].score is not None
+    assert hits[0].score >= hits[-1].score
+    assert all(h.namespace == ("memories", "kamal") for h in hits)
+
+
+def test_query_prefix_children_and_filter(store):
+    _seed(store)
+    keys = {(h.namespace, h.key) for h in store.search(("memories",), query="sushi")}
+    assert (("memories", "kamal"), "food") in keys and (("memories", "priya"), "food") in keys
+    facts = store.search(("memories", "kamal"), query="user", filter={"kind": "fact"})
+    assert facts and all(h.value["kind"] == "fact" for h in facts)
+    gt = store.search(("memories", "kamal"), query="user", filter={"kind": {"$in": ["pref"]}})
+    assert [h.key for h in gt] == ["food"]
+
+
+def test_query_limit_offset(store):
+    _seed(store)
+    first = store.search(("memories", "kamal"), query="user", limit=2)
+    rest = store.search(("memories", "kamal"), query="user", limit=2, offset=2)
+    assert len(first) == 2 and len(rest) == 1 and not {h.key for h in first} & {h.key for h in rest}
+
+
+def test_no_query_is_filter_only(store):
+    _seed(store)
+    hits = store.search(("memories", "kamal"))
+    assert [h.key for h in hits] == ["city", "food", "lang"] and all(h.score is None for h in hits)
+
+
+def test_native_path_used(store, monkeypatch):
+    _seed(store)
+    called = []
+    orig = store._vector_search
+    monkeypatch.setattr(store, "_vector_search", lambda *a, **k: (called.append(1), orig(*a, **k))[1])
+    store.search(("memories",), query="sushi")
+    assert called
+
+
+async def test_async_query(store):
+    _seed(store)
+    hits = await store.asearch(("memories", "kamal"), query="sushi")
+    assert hits[0].key == "food"

langgraph_store_postgres-0.1.0/tests/test_postgres_store.py

@@ -0,0 +1,119 @@
+"""Tests for PostgresStore (LangGraph BaseStore) against real local PostgreSQL."""
+import os
+import uuid
+
+import pytest
+from sqlalchemy import create_engine, text
+
+from langgraph_store_postgres import PostgresStore
+
+URL = os.environ.get(
+    "SQL_TEST_URL", "postgresql+psycopg2://postgres:postgres@localhost:5432/postgres"
+)
+
+
+def _db_available() -> bool:
+    try:
+        with create_engine(URL).connect() as c:
+            c.execute(text("select 1"))
+        return True
+    except Exception:
+        return False
+
+
+pytestmark = pytest.mark.skipif(not _db_available(), reason=f"PostgreSQL not reachable: {URL}")
+
+
+@pytest.fixture()
+def store():
+    tname = f"lg_store_test_{uuid.uuid4().hex[:8]}"
+    s = PostgresStore(url=URL, table_name=tname)
+    yield s
+    with s._engine.begin() as c:
+        c.execute(text(f'DROP TABLE IF EXISTS "{tname}"'))
+    s._engine.dispose()
+
+
+def test_put_get(store):
+    store.put(("users", "1"), "food", {"text": "sushi"})
+    it = store.get(("users", "1"), "food")
+    assert it is not None
+    assert it.value == {"text": "sushi"}
+    assert it.namespace == ("users", "1")
+    assert it.key == "food"
+
+
+def test_get_missing_is_none(store):
+    assert store.get(("nope",), "x") is None
+
+
+def test_update_preserves_created_at(store):
+    store.put(("u",), "k", {"v": 1})
+    a = store.get(("u",), "k")
+    store.put(("u",), "k", {"v": 2})
+    b = store.get(("u",), "k")
+    assert b.value == {"v": 2}
+    assert b.created_at == a.created_at        # created_at preserved across updates
+    assert b.updated_at >= a.updated_at
+
+
+def test_delete_both_ways(store):
+    store.put(("u",), "k", {"v": 1})
+    store.delete(("u",), "k")
+    assert store.get(("u",), "k") is None
+    store.put(("u",), "k2", {"v": 1})
+    store.put(("u",), "k2", None)              # put None == delete
+    assert store.get(("u",), "k2") is None
+
+
+def test_search_prefix_and_filter(store):
+    store.put(("users", "1", "mem"), "a", {"kind": "pref"})
+    store.put(("users", "1", "mem"), "b", {"kind": "fact"})
+    store.put(("users", "2", "mem"), "c", {"kind": "pref"})
+    assert sorted(i.key for i in store.search(("users", "1"))) == ["a", "b"]
+    assert [i.key for i in store.search(("users", "1"), filter={"kind": "pref"})] == ["a"]
+    assert store.search(("users", "1"))[0].namespace == ("users", "1", "mem")
+
+
+def test_search_operator_filter(store):
+    store.put(("n",), "a", {"score": 5})
+    store.put(("n",), "b", {"score": 15})
+    hits = store.search(("n",), filter={"score": {"$gt": 10}})
+    assert [i.key for i in hits] == ["b"]
+
+
+def test_search_limit_offset(store):
+    for i in range(5):
+        store.put(("n",), f"k{i}", {"i": i})
+    assert [i.key for i in store.search(("n",), limit=2)] == ["k0", "k1"]
+    assert [i.key for i in store.search(("n",), limit=2, offset=2)] == ["k2", "k3"]
+
+
+def test_list_namespaces(store):
+    store.put(("users", "1", "mem"), "a", {})
+    store.put(("users", "2", "mem"), "b", {})
+    store.put(("orgs", "x"), "c", {})
+    allns = store.list_namespaces()
+    assert ("users", "1", "mem") in allns and ("orgs", "x") in allns
+    assert all(n[0] == "users" for n in store.list_namespaces(prefix=("users",)))
+    assert all(n[-1] == "mem" for n in store.list_namespaces(suffix=("mem",)))
+    depth1 = store.list_namespaces(max_depth=1)
+    assert ("users",) in depth1 and ("orgs",) in depth1
+
+
+def test_list_namespaces_wildcard(store):
+    store.put(("a", "1", "x"), "k", {})
+    store.put(("a", "2", "x"), "k", {})
+    store.put(("a", "1", "y"), "k", {})
+    hits = store.list_namespaces(prefix=("a", "*", "x"))
+    assert set(hits) == {("a", "1", "x"), ("a", "2", "x")}
+
+
+async def test_async_surface(store):
+    await store.aput(("u",), "k", {"v": 1})
+    it = await store.aget(("u",), "k")
+    assert it.value == {"v": 1}
+    res = await store.asearch(("u",))
+    assert [i.key for i in res] == ["k"]
+    await store.adelete(("u",), "k")
+    assert await store.aget(("u",), "k") is None