langgraph-store-cosmosdb 0.1.0__tar.gz

langgraph_store_cosmosdb-0.1.0/.gitignore

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

langgraph_store_cosmosdb-0.1.0/PKG-INFO

@@ -0,0 +1,67 @@
+Metadata-Version: 2.5
+Name: langgraph-store-cosmosdb
+Version: 0.1.0
+Summary: Azure CosmosDB long-term-memory store (BaseStore) for LangGraph — namespaced key/value memory with filters, list_namespaces and native semantic (Cosmos DB vector) 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,azure,basestore,cosmosdb,langgraph,long-term-memory,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: azure-cosmos
+Requires-Dist: langgraph-store-core>=0.1.0
+Description-Content-Type: text/markdown
+
+# langgraph-store-cosmosdb
+
+An **Azure Cosmos DB (NoSQL)** 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 Cosmos DB vector search**.
+
+```bash
+pip install langgraph-store-cosmosdb
+```
+
+```python
+from langgraph_store_cosmosdb import CosmosDBStore
+
+store = CosmosDBStore(
+    endpoint="https://<acct>.documents.azure.com:443/", key="<key>",
+    database_name="langgraph", container_name="store",
+)
+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"})
+```
+
+Use it as a LangGraph store: `graph.compile(store=CosmosDBStore(...))`. Async methods work too (sync calls run in a thread).
+
+## Semantic search (Cosmos DB vector search)
+
+Pass a LangGraph `IndexConfig` and the store embeds the configured fields on `put` and ranks `search(query=...)` by cosine similarity using `VectorDistance` in the container's SQL:
+
+```python
+from langgraph_store_core import bedrock_titan_embeddings
+from langgraph_store_cosmosdb import CosmosDBStore
+
+store = CosmosDBStore(endpoint=..., key=..., database_name="langgraph", container_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.
+
+With `index` the container is created with a vector embedding policy on `/embedding` (cosine, `dims`) and a `diskANN` vector index (`vector_index_type="quantizedFlat"` / `"flat"` to change). The policy must be set at container creation, so use a **new container name** when enabling semantic search on an existing store. The account needs the **Vector Search for NoSQL** capability enabled (`az cosmosdb update ... --capabilities EnableNoSQLVectorSearch`, keeping any existing capabilities). Namespace prefix and plain-equality filters run inside the query; operator filters (`$gt`, `$in`, …) are applied on the returned candidates.
+
+## Data model
+
+A container partitioned by `/prefix` (the namespace), each item `{id, prefix, key, value, created_at, updated_at[, embedding]}` (`id` is the URL-encoded key). Database and container are auto-created under key-based auth. Search is a `STARTSWITH` prefix query; 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_cosmosdb-0.1.0/README.md

@@ -0,0 +1,50 @@
+# langgraph-store-cosmosdb
+
+An **Azure Cosmos DB (NoSQL)** 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 Cosmos DB vector search**.
+
+```bash
+pip install langgraph-store-cosmosdb
+```
+
+```python
+from langgraph_store_cosmosdb import CosmosDBStore
+
+store = CosmosDBStore(
+    endpoint="https://<acct>.documents.azure.com:443/", key="<key>",
+    database_name="langgraph", container_name="store",
+)
+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"})
+```
+
+Use it as a LangGraph store: `graph.compile(store=CosmosDBStore(...))`. Async methods work too (sync calls run in a thread).
+
+## Semantic search (Cosmos DB vector search)
+
+Pass a LangGraph `IndexConfig` and the store embeds the configured fields on `put` and ranks `search(query=...)` by cosine similarity using `VectorDistance` in the container's SQL:
+
+```python
+from langgraph_store_core import bedrock_titan_embeddings
+from langgraph_store_cosmosdb import CosmosDBStore
+
+store = CosmosDBStore(endpoint=..., key=..., database_name="langgraph", container_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.
+
+With `index` the container is created with a vector embedding policy on `/embedding` (cosine, `dims`) and a `diskANN` vector index (`vector_index_type="quantizedFlat"` / `"flat"` to change). The policy must be set at container creation, so use a **new container name** when enabling semantic search on an existing store. The account needs the **Vector Search for NoSQL** capability enabled (`az cosmosdb update ... --capabilities EnableNoSQLVectorSearch`, keeping any existing capabilities). Namespace prefix and plain-equality filters run inside the query; operator filters (`$gt`, `$in`, …) are applied on the returned candidates.
+
+## Data model
+
+A container partitioned by `/prefix` (the namespace), each item `{id, prefix, key, value, created_at, updated_at[, embedding]}` (`id` is the URL-encoded key). Database and container are auto-created under key-based auth. Search is a `STARTSWITH` prefix query; 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_cosmosdb-0.1.0/pyproject.toml

@@ -0,0 +1,35 @@
+[build-system]
+requires = ["hatchling"]
+build-backend = "hatchling.build"
+
+[project]
+name = "langgraph-store-cosmosdb"
+version = "0.1.0"
+description = "Azure CosmosDB long-term-memory store (BaseStore) for LangGraph — namespaced key/value memory with filters, list_namespaces and native semantic (Cosmos DB vector) search"
+authors = [{name = "Kamal", email = "skamalj@gmail.com"}]
+readme = "README.md"
+requires-python = ">=3.10"
+dependencies = [
+    "langgraph-store-core>=0.1.0",
+    "azure-cosmos",
+]
+keywords = ["langgraph", "store", "basestore", "cosmosdb", "azure", "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_cosmosdb"]

langgraph_store_cosmosdb-0.1.0/src/langgraph_store_cosmosdb/__init__.py

@@ -0,0 +1,161 @@
+"""Azure CosmosDB implementation of the LangGraph long-term-memory ``BaseStore``.
+
+``CosmosDBStore`` persists namespaced key/value memory in a container partitioned
+by namespace, and supports get / put / delete / search (namespace-prefix + filter)
+and ``list_namespaces``. With an ``IndexConfig`` it also does **semantic search**
+natively via Cosmos DB NoSQL vector search (``VectorDistance``, cosine).
+Built on ``langgraph-store-core``.
+"""
+
+from __future__ import annotations
+
+from typing import List, Optional, Tuple
+from urllib.parse import quote
+
+from azure.cosmos import CosmosClient, PartitionKey
+from azure.cosmos.exceptions import CosmosResourceNotFoundError
+
+from langgraph_store_core import IndexConfig, KVStore
+
+__all__ = ["CosmosDBStore"]
+
+_FIELDS = ("prefix", "key", "value", "created_at", "updated_at")
+_PROJECTION = ", ".join(f'c["{f}"]' for f in _FIELDS)  # value/key are reserved words
+
+
+class CosmosDBStore(KVStore):
+    """LangGraph ``BaseStore`` backed by Azure CosmosDB.
+
+    Example:
+        ```python
+        from langgraph_store_cosmosdb import CosmosDBStore
+
+        store = CosmosDBStore(
+            endpoint="https://<acct>.documents.azure.com:443/", key="<key>",
+            database_name="langgraph", container_name="store",
+        )
+        store.put(("users", "1"), "profile", {"name": "Kamal"})
+
+        # semantic search
+        store = CosmosDBStore(..., index={"dims": 1024, "embed": my_embedder, "fields": ["text"]})
+        hits = store.search(("users", "1"), query="what does the user like?")
+        ```
+
+    The container is partitioned by ``/prefix`` (the namespace) and auto-created
+    under key-based auth. With ``index`` the container is created with a vector
+    embedding policy on ``/embedding`` (cosine, ``dims``) and a ``diskANN``
+    vector index; the policy must be set at container creation, so use a new
+    container name when enabling semantic search on an existing store.
+    """
+
+    supports_native_vector_search = True
+
+    def __init__(
+        self,
+        *,
+        endpoint: str,
+        key: str,
+        database_name: str,
+        container_name: str,
+        index: Optional[IndexConfig] = None,
+        vector_index_type: str = "diskANN",
+    ) -> None:
+        super().__init__(index=index)
+        client = CosmosClient(endpoint, credential=key)
+        db = client.create_database_if_not_exists(database_name)
+        kwargs = {}
+        if self.dims:
+            kwargs["vector_embedding_policy"] = {
+                "vectorEmbeddings": [
+                    {
+                        "path": "/embedding",
+                        "dataType": "float32",
+                        "distanceFunction": "cosine",
+                        "dimensions": self.dims,
+                    }
+                ]
+            }
+            kwargs["indexing_policy"] = {
+                "automatic": True,
+                "indexingMode": "consistent",
+                "includedPaths": [{"path": "/*"}],
+                # vector paths must be excluded from the range index
+                "excludedPaths": [{"path": "/embedding/*"}, {"path": '/"_etag"/?'}],
+                "vectorIndexes": [{"path": "/embedding", "type": vector_index_type}],
+            }
+        self._container = db.create_container_if_not_exists(
+            id=container_name, partition_key=PartitionKey(path="/prefix"), **kwargs
+        )
+
+    @staticmethod
+    def _id(key: str) -> str:
+        return quote(key, safe="")  # Cosmos ids may not contain # / \ ?
+
+    def _row(self, item: dict) -> dict:
+        row = {f: item[f] for f in _FIELDS}
+        row["embedding"] = item.get("embedding")
+        return row
+
+    def _read(self, prefix: str, key: str) -> Optional[dict]:
+        try:
+            item = self._container.read_item(item=self._id(key), partition_key=prefix)
+        except CosmosResourceNotFoundError:
+            return None
+        return self._row(item)
+
+    def _write(self, row: dict) -> None:
+        doc = {"id": self._id(row["key"]), **{f: row[f] for f in _FIELDS}}
+        if self.dims and row.get("embedding") is not None:
+            doc["embedding"] = row["embedding"]
+        self._container.upsert_item(doc)
+
+    def _remove(self, prefix: str, key: str) -> None:
+        try:
+            self._container.delete_item(item=self._id(key), partition_key=prefix)
+        except CosmosResourceNotFoundError:
+            pass
+
+    def _scan(self, prefix: str) -> List[dict]:
+        proj = _PROJECTION + (', c["embedding"]' if self.dims else "")
+        if prefix == "":
+            query = f"SELECT {proj} FROM c"
+            params: list = []
+        else:
+            query = f'SELECT {proj} FROM c WHERE STARTSWITH(c["prefix"], @p)'
+            params = [{"name": "@p", "value": prefix}]
+        items = self._container.query_items(
+            query=query, parameters=params, enable_cross_partition_query=True
+        )
+        return [self._row(i) for i in items]
+
+    def _vector_search(
+        self, prefix: str, vector: List[float], filter: Optional[dict], limit: int
+    ) -> List[Tuple[dict, float]]:
+        """Native Cosmos vector search: ORDER BY VectorDistance, prefix + equality filters in SQL."""
+        where = ['IS_DEFINED(c["embedding"])']
+        params: list = [{"name": "@q", "value": vector}, {"name": "@k", "value": int(limit)}]
+        if prefix != "":
+            where.append('STARTSWITH(c["prefix"], @p)')
+            params.append({"name": "@p", "value": prefix})
+        # Push plain-equality filters into SQL; operator filters are applied by core.
+        n = 0
+        for field, expected in (filter or {}).items():
+            if not isinstance(expected, dict):
+                n += 1
+                where.append(f'c["value"][@f{n}] = @v{n}')
+                params += [{"name": f"@f{n}", "value": field}, {"name": f"@v{n}", "value": expected}]
+        query = (
+            f"SELECT TOP @k {_PROJECTION}, "
+            f'VectorDistance(c["embedding"], @q) AS score '
+            f"FROM c WHERE {' AND '.join(where)} "
+            f'ORDER BY VectorDistance(c["embedding"], @q)'
+        )
+        items = self._container.query_items(
+            query=query, parameters=params, enable_cross_partition_query=True
+        )
+        out: List[Tuple[dict, float]] = []
+        for i in items:
+            row = {f: i[f] for f in _FIELDS}
+            row["embedding"] = None  # not projected; not needed after ranking
+            out.append((row, float(i["score"])))  # cosine: similarity, higher is better
+        return out

langgraph_store_cosmosdb-0.1.0/tests/test_cosmosdb_semantic.py

@@ -0,0 +1,107 @@
+"""Semantic-search tests for CosmosDBStore (native VectorDistance path) against a real Cosmos DB account."""
+import os
+import time
+import uuid
+
+import pytest
+
+from langgraph_store_core.testing import FakeEmbeddings
+from langgraph_store_cosmosdb import CosmosDBStore
+
+EP = os.environ.get("COSMOS_ENDPOINT")
+KEY = os.environ.get("COSMOS_KEY")
+DB = os.environ.get("LG_COSMOS_DB", "langgraph_store_test")
+DIMS = 64
+
+pytestmark = pytest.mark.skipif(not (EP and KEY), reason="COSMOS_ENDPOINT/COSMOS_KEY not set")
+
+
+@pytest.fixture(scope="module")
+def store():
+    cname = f"vstore_{uuid.uuid4().hex[:8]}"
+    s = CosmosDBStore(endpoint=EP, key=KEY, database_name=DB, container_name=cname,
+                      index={"dims": DIMS, "embed": FakeEmbeddings(DIMS), "fields": ["text"]})
+    yield s
+    s._container.database_link  # noqa: B018 — keep proxy alive
+    from azure.cosmos import CosmosClient
+    CosmosClient(EP, credential=KEY).get_database_client(DB).delete_container(cname)
+
+
+@pytest.fixture(autouse=True)
+def _clean(store):
+    yield
+    for row in store._scan(""):
+        store._remove(row["prefix"], row["key"])
+
+
+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 _search(s, *a, **k):
+    # vector index population is asynchronous; retry briefly until results appear
+    for _ in range(12):
+        hits = s.search(*a, **k)
+        if hits:
+            return hits
+        time.sleep(2)
+    return s.search(*a, **k)
+
+
+def test_container_has_vector_policy(store):
+    props = store._container.read()
+    assert props.get("vectorEmbeddingPolicy", {}).get("vectorEmbeddings", [{}])[0].get("path") == "/embedding"
+    assert any(v.get("path") == "/embedding" for v in props["indexingPolicy"].get("vectorIndexes", []))
+
+
+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 = _search(store, ("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 _search(store, ("memories",), query="sushi")}
+    assert (("memories", "kamal"), "food") in keys and (("memories", "priya"), "food") in keys
+    facts = _search(store, ("memories", "kamal"), query="user", filter={"kind": "fact"})
+    assert facts and all(h.value["kind"] == "fact" for h in facts)
+
+
+def test_query_limit_offset(store):
+    _seed(store)
+    first = _search(store, ("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._container.query_items
+
+    def spy(query, **kw):
+        called.append(query)
+        return orig(query=query, **kw)
+
+    monkeypatch.setattr(store._container, "query_items", spy)
+    _search(store, ("memories",), query="sushi")
+    assert any("VectorDistance" in q for q in called)

langgraph_store_cosmosdb-0.1.0/tests/test_cosmosdb_store.py

@@ -0,0 +1,83 @@
+"""Tests for CosmosDBStore (LangGraph BaseStore) against real CosmosDB.
+
+Requires COSMOS_ENDPOINT + COSMOS_KEY. A unique container per test gives clean
+isolation (deleted on teardown).
+"""
+import os
+import uuid
+
+import pytest
+
+from langgraph_store_cosmosdb import CosmosDBStore
+
+EP = os.environ.get("COSMOS_ENDPOINT")
+KEY = os.environ.get("COSMOS_KEY")
+DB = os.environ.get("LG_COSMOS_DB", "langgraph_store_test")
+
+pytestmark = pytest.mark.skipif(not (EP and KEY), reason="COSMOS_ENDPOINT/COSMOS_KEY not set")
+
+
+@pytest.fixture()
+def store():
+    cname = f"test_{uuid.uuid4().hex[:8]}"
+    s = CosmosDBStore(endpoint=EP, key=KEY, database_name=DB, container_name=cname)
+    yield s
+    from azure.cosmos import CosmosClient
+    CosmosClient(EP, credential=KEY).get_database_client(DB).delete_container(cname)
+
+
+def test_put_get(store):
+    store.put(("users", "1"), "food", {"text": "sushi"})
+    it = store.get(("users", "1"), "food")
+    assert it is not None and it.value == {"text": "sushi"}
+    assert it.namespace == ("users", "1") and 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
+    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)
+    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"]
+
+
+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",)))
+    depth1 = store.list_namespaces(max_depth=1)
+    assert ("users",) in depth1 and ("orgs",) in depth1
+
+
+async def test_async_surface(store):
+    await store.aput(("u",), "k", {"v": 1})
+    it = await store.aget(("u",), "k")
+    assert it.value == {"v": 1}
+    await store.adelete(("u",), "k")
+    assert await store.aget(("u",), "k") is None