langraph-cosmos 0.1.1

Library with a Cosmos DB implementation of LangGraph checkpoint saver.

LangGraph Cosmos DB

Azure Cosmos DB implementation of LangGraph checkpoint saver and store for building stateful, multi-actor applications with LLMs.

Overview

This library provides a production-ready integration between LangGraph and Azure Cosmos DB, enabling:

• Durable checkpoint storage for LangGraph agents (coming soon)
• Persistent memory store with support for both sync and async operations
• Vector search capabilities using Cosmos DB's native vector indexing
• Automatic TTL (Time-To-Live) for managing data lifecycle
• Global distribution with Azure Cosmos DB's multi-region replication

Why Cosmos DB for LangGraph?

Azure Cosmos DB is an ideal backend for LangGraph applications:

• 🌍 Global Distribution: Deploy your AI agents globally with multi-region writes and reads
• ⚡ Low Latency: Single-digit millisecond response times for state access
• 🔍 Vector Search: Native support for vector embeddings with DiskANN, quantized flat, and flat indexes
• 📈 Elastic Scale: Automatic scaling from zero to planet-scale throughput
• 💾 Flexible Schema: Store complex agent state with JSON documents
• 🔒 Enterprise-Ready: Built-in security, compliance, and 99.999% SLA

Installation

pip install langraph-cosmos

Optional Dependencies

For vector search functionality:

pip install langchain-openai  # or your preferred embeddings provider

Quick Start

Synchronous Store

from azure.cosmos import CosmosClient
from langraph_cosmos.store import CosmosStore

# Create Cosmos DB client
client = CosmosClient(
    url="https://your-account.documents.azure.com:443/",
    credential="your-primary-key"
)

# Initialize store
store = CosmosStore(
    client=client,
    database_name="langgraph_db",
    container_name="memory_store"
)

# Setup container (creates if doesn't exist)
store.setup()

# Store data
store.put(
    namespace=("users", "user123"),
    key="preferences",
    value={"theme": "dark", "language": "en"}
)

# Retrieve data
item = store.get(
    namespace=("users", "user123"),
    key="preferences"
)
print(item.value)  # {'theme': 'dark', 'language': 'en'}

# List namespaces
namespaces = list(store.list_namespaces(prefix=("users",)))
print(namespaces)  # [('users', 'user123')]

# Search within namespace
results = list(store.search(
    namespace_prefix=("users",),
    limit=10
))

Asynchronous Store

import asyncio
from azure.cosmos.aio import CosmosClient
from langraph_cosmos.store import aio

async def main():
    # Create async client
    client = CosmosClient(
        url="https://your-account.documents.azure.com:443/",
        credential="your-primary-key"
    )

    # Initialize async store
    store = aio.CosmosStore(
        client=client,
        database_name="langgraph_db",
        container_name="memory_store"
    )

    await store.setup()

    # Store data asynchronously
    await store.aput(
        namespace=("agents", "agent_1"),
        key="state",
        value={"step": 5, "context": "processing"}
    )

    # Retrieve data
    item = await store.aget(
        namespace=("agents", "agent_1"),
        key="state"
    )
    print(item.value)

    # Batch operations
    items = [
        ("agents", "agent_1", "config", {"temperature": 0.7}),
        ("agents", "agent_2", "config", {"temperature": 0.9}),
    ]

    await store.abatch([
        {"namespace": ns, "key": k, "value": v}
        for ns, _, k, v in [(tuple(i[0].split(",")), i[1], i[2], i[3]) for i in items]
    ])

    await client.close()

asyncio.run(main())

Connection String

Both sync and async stores support connection strings:

from langraph_cosmos.store import CosmosStore

with CosmosStore.from_connection_string(
    connection_string="AccountEndpoint=https://...;AccountKey=...;",
    database_name="langgraph_db",
    container_name="memory_store"
) as store:
    store.setup()
    store.put(("test",), "key1", {"data": "value"})

Advanced Features

Vector Search

Enable semantic search capabilities with vector embeddings:

from azure.cosmos import CosmosClient
from langchain_openai import OpenAIEmbeddings
from langraph_cosmos.store import CosmosStore, CosmosIndexConfig

# Configure vector indexing
index_config = CosmosIndexConfig(
    embed=OpenAIEmbeddings(model="text-embedding-3-small"),
    fields=["text", "content"],  # Fields to embed
    dims=1536,  # Embedding dimensions
    vector_index_type="diskANN",  # Options: flat, diskANN, quantizedFlat
    distance_type="cosine"  # Options: cosine, euclidean, dotproduct
)

store = CosmosStore(
    client=client,
    database_name="langgraph_db",
    container_name="vector_store",
    index=index_config
)

store.setup()

# Store documents with automatic embedding
store.put(
    namespace=("documents",),
    key="doc1",
    value={
        "text": "LangGraph is a framework for building stateful agents",
        "metadata": {"source": "docs"}
    }
)

store.put(
    namespace=("documents",),
    key="doc2",
    value={
        "text": "Azure Cosmos DB provides global distribution and low latency",
        "metadata": {"source": "azure"}
    }
)

# Semantic search
results = list(store.search(
    namespace_prefix=("documents",),
    query="agent frameworks",
    limit=5
))

for item in results:
    print(f"Score: {item.score}, Text: {item.value['text']}")

Quantized Vector Search

For reduced memory footprint:

index_config = CosmosIndexConfig(
    embed=OpenAIEmbeddings(),
    fields=["text"],
    dims=1536,
    vector_index_type="quantizedFlat",
    distance_type="cosine",
    quantization_byte_size=2  # 1, 2, or 4 bytes
)

TTL (Time-To-Live)

Automatically expire old data:

from langraph_cosmos.store import CosmosStore

store = CosmosStore(
    client=client,
    database_name="langgraph_db",
    container_name="ephemeral_store",
    ttl={"default_ttl": 60}  # 60 minutes
)

store.setup()

# Start TTL sweeper thread (required for TTL to work)
store.start_ttl_sweeper()

# Store with custom TTL
store.put(
    namespace=("sessions",),
    key="session_123",
    value={"user": "alice", "active": True}
    # Will expire after 60 minutes
)

# Later, stop the sweeper
store.stop_ttl_sweeper()

For async:

# Start TTL sweeper task
await store.start_ttl_sweeper()

# Later, stop it
await store.stop_ttl_sweeper()

Filtering

Search with metadata filters:

# Store items with filterable metadata
store.put(
    namespace=("products",),
    key="prod1",
    value={"name": "Widget", "price": 19.99, "category": "tools"}
)

store.put(
    namespace=("products",),
    key="prod2",
    value={"name": "Gadget", "price": 29.99, "category": "electronics"}
)

# Search with filters
results = list(store.search(
    namespace_prefix=("products",),
    filter={"category": {"$eq": "tools"}}
))

Supported filter operators:

• $eq: Equal to
• $ne: Not equal to
• $gt: Greater than
• $gte: Greater than or equal
• $lt: Less than
• $lte: Less than or equal

Resources

• LangGraph: Documentation | GitHub
• Azure Cosmos DB: Documentation | Python SDK
• Vector Search: Cosmos DB Vector Search Guide

Contributing

Contributions are welcome! Please feel free to submit a Pull Request.

License

This project is licensed under the MIT License - see the LICENSE file for details.

Acknowledgments

Built on top of:

• LangGraph - The LangChain orchestration framework
• Azure Cosmos DB - Microsoft's globally distributed database
• LangChain - For embeddings integration

Support

For issues and questions:

• GitHub Issues: Report bugs or request features
• LangChain Forum: Community discussions
• Azure Support: Azure Cosmos DB support

---

Built with ❤️ for the LangGraph community