langgraph-store-cosmosdb

An Azure Cosmos DB (NoSQL) long-term-memory store (BaseStore) for LangGraph — namespaced key/value memory with prefix search, filters, list_namespaces, and native semantic search via Cosmos DB vector search.

pip install langgraph-store-cosmosdb

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:

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.

Docs: https://skamalj.github.io/agentstate-reducer/

License

MIT