Cosmos DB + AI

Modern AI applications demand more than a traditional database can deliver. When your RAG pipeline needs sub-10ms vector lookups across three continents, when your intelligent app must store documents alongside their embeddings without data synchronization headaches, and when your AI serving layer requires five-nines availability with guaranteed low latency — you need a database purpose-built for the AI era. Azure Cosmos DB has evolved into exactly that: a globally distributed, multi-model database with native vector search that eliminates the need for a separate vector store, unifying your operational and AI data in a single, planet-scale platform.

This guide dives deep into Cosmos DB’s vector search capabilities, walks through building production RAG applications, explores the change feed for real-time AI processing, and examines how global distribution transforms AI serving architectures. Whether you are migrating from a standalone vector database or designing a new AI-native application from scratch, this article covers the advanced patterns you need.

What Is Azure Cosmos DB?

Azure Cosmos DB is Microsoft’s fully managed, globally distributed, multi-model database service. Originally built to power services like Xbox, Skype, and Office 365 at planetary scale, it has evolved into the go-to database for AI-native applications that demand guaranteed low latency, elastic throughput, and turnkey global distribution.

Unlike traditional databases that require you to choose between relational, document, key-value, or graph paradigms, Cosmos DB supports multiple APIs — NoSQL (native), MongoDB, PostgreSQL, Apache Cassandra, Apache Gremlin, and Table — all on the same underlying globally distributed engine. For AI applications, this means you can store your operational data, user sessions, conversation history, and vector embeddings in a single service without stitching together multiple infrastructure components.

Why Cosmos DB is ideal for AI workloads

• Co-located vectors and data — store embeddings alongside the documents they represent, eliminating synchronization issues between a separate vector DB and your operational store.
• DiskANN-powered vector search — Microsoft Research’s DiskANN algorithm delivers high-recall approximate nearest neighbor (ANN) search with significantly lower memory overhead than HNSW-based alternatives.
• Sub-10ms latency at the 99th percentile — critical for real-time AI serving where user experience depends on response speed.
• Automatic and instant global distribution — replicate your vectors and documents to any Azure region with a single click, serving AI queries from the closest datacenter.
• Change feed for real-time AI — trigger embedding generation, re-indexing, or downstream AI processing the instant data changes.

Vector Search in Cosmos DB

Vector search in Cosmos DB is built on DiskANN, an approximate nearest neighbor (ANN) algorithm developed by Microsoft Research. Unlike HNSW indexes that require the entire graph to reside in memory, DiskANN stores the index on SSD with a compressed in-memory representation, enabling vector search over billions of vectors without proportional memory costs.

Configuring a vector embedding policy

Before storing vectors, you define a vector embedding policy on your container that tells Cosmos DB which properties contain vectors, their dimensions, the distance function to use, and the index type:

from azure.cosmos import CosmosClient, PartitionKey

# Initialize the client
client = CosmosClient("https://myaccount.documents.azure.com:443/", credential)
db = client.get_database_client("ai-app-db")

# Define the vector embedding policy
vector_embedding_policy = {
    "vectorEmbeddings": [
        {
            "path": "/embedding",
            "dataType": "float32",
            "distanceFunction": "cosine",
            "dimensions": 1536
        }
    ]
}

# Define the indexing policy with vector index
indexing_policy = {
    "includedPaths": [{"path": "/*"}],
    "excludedPaths": [{"path": "/embedding/*"}],
    "vectorIndexes": [
        {
            "path": "/embedding",
            "type": "diskANN"
        }
    ]
}

# Create the container with vector support
container = db.create_container(
    id="documents",
    partition_key=PartitionKey(path="/category"),
    vector_embedding_policy=vector_embedding_policy,
    indexing_policy=indexing_policy
)

Vector index types

Cosmos DB offers three vector index types, each optimized for different workload characteristics:

Performing vector search queries

Cosmos DB uses the VectorDistance built-in function in its NoSQL query language. You pass the query vector and retrieve documents ranked by similarity:

# Generate embedding for the user's query
query_embedding = openai_client.embeddings.create(
    input="How do I configure autoscaling in Kubernetes?",
    model="text-embedding-3-large"
).data[0].embedding

# Perform vector search in Cosmos DB
results = container.query_items(
    query="""
        SELECT TOP 5 c.title, c.content, c.category,
               VectorDistance(c.embedding, @queryVector) AS score
        FROM c
        ORDER BY VectorDistance(c.embedding, @queryVector)
    """,
    parameters=[
        {"name": "@queryVector", "value": query_embedding}
    ],
    enable_cross_partition_query=True
)

for doc in results:
    print(f"{doc['title']} (score: {doc['score']:.4f})")

Building a RAG Application with Cosmos DB

Retrieval-Augmented Generation (RAG) is the most common AI pattern for grounding LLM responses in your own data. Cosmos DB serves as both the document store and the vector store in this architecture, eliminating the need to synchronize data between separate systems. Here is a complete end-to-end implementation:

Step 1: Ingest documents with embeddings

import json
from openai import AzureOpenAI
from azure.cosmos import CosmosClient

openai_client = AzureOpenAI(
    azure_endpoint="https://my-openai.openai.azure.com/",
    api_version="2024-06-01",
    api_key=api_key
)

def ingest_document(container, doc_id, title, content, category):
    """Ingest a document with its embedding into Cosmos DB."""

    # Generate embedding from the content
    embedding = openai_client.embeddings.create(
        input=content,
        model="text-embedding-3-large"
    ).data[0].embedding

    # Store document + embedding together
    document = {
        "id": doc_id,
        "title": title,
        "content": content,
        "category": category,
        "embedding": embedding,  # 1536-dim float32 vector
        "metadata": {
            "source": "knowledge-base",
            "indexed_at": "2025-01-15T10:30:00Z",
            "chunk_index": 0
        }
    }

    container.upsert_item(document)
    return document["id"]

# Ingest a batch of documents
docs = [
    ("doc-001", "Kubernetes Autoscaling", "HPA scales pods based on...", "k8s"),
    ("doc-002", "Azure Functions Triggers", "Cosmos DB trigger enables...", "serverless"),
    ("doc-003", "Redis Cache Patterns", "Cache-aside pattern with...", "caching"),
]

for doc_id, title, content, cat in docs:
    ingest_document(container, doc_id, title, content, cat)

Step 2: Query by similarity and generate a grounded response

def rag_query(container, user_question, top_k=5):
    """Full RAG pipeline: embed query -> vector search -> LLM response."""

    # 1. Generate embedding for the user's question
    query_vector = openai_client.embeddings.create(
        input=user_question,
        model="text-embedding-3-large"
    ).data[0].embedding

    # 2. Vector search in Cosmos DB
    results = list(container.query_items(
        query=f"""
            SELECT TOP {top_k} c.title, c.content, c.category,
                   VectorDistance(c.embedding, @qv) AS similarity
            FROM c
            ORDER BY VectorDistance(c.embedding, @qv)
        """,
        parameters=[{"name": "@qv", "value": query_vector}],
        enable_cross_partition_query=True
    ))

    # 3. Build context from retrieved documents
    context = "\n\n".join([
        f"## {r['title']}\n{r['content']}"
        for r in results
    ])

    # 4. Generate grounded response with Azure OpenAI
    response = openai_client.chat.completions.create(
        model="gpt-4o",
        messages=[
            {"role": "system", "content": f"""Answer based only on
                the following context. Cite sources by title.
                \n\nContext:\n{context}"""},
            {"role": "user", "content": user_question}
        ],
        temperature=0.3
    )

    return {
        "answer": response.choices[0].message.content,
        "sources": [{"title": r["title"], "score": r["similarity"]} for r in results]
    }

# Usage
result = rag_query(container, "How do I set up autoscaling?")
print(result["answer"])

Integrated Vector Database: A New Paradigm

The AI industry has rapidly adopted standalone vector databases like Pinecone, Weaviate, and Milvus. However, Cosmos DB represents a fundamentally different approach: the integrated vector database, where vectors live alongside the data they represent in the same document, the same container, and the same globally distributed infrastructure.

Co-located vs. separated architectures

The integrated model particularly excels in scenarios where you need hybrid queries — combining vector similarity with metadata filters. For example, finding the most similar product descriptions but only within a specific category and price range, all in a single query with a single round-trip:

SELECT TOP 10 c.title, c.price, c.category,
       VectorDistance(c.embedding, @queryVector) AS score
FROM c
WHERE c.category = "electronics"
  AND c.price BETWEEN 100 AND 500
  AND c.inStock = true
ORDER BY VectorDistance(c.embedding, @queryVector)

Cosmos DB APIs for Vector Workloads

One of Cosmos DB’s distinctive strengths is its support for multiple API surfaces on the same underlying engine. Each API brings vector search to developers using the tools and query languages they already know.

NoSQL API (native)

The native API provides the richest vector search experience with the VectorDistance function, full support for hybrid queries, DiskANN indexing, and the tightest integration with the Cosmos DB SDK. This is the recommended API for new AI applications.

MongoDB API (vCore)

The MongoDB vCore offering supports the $search aggregation pipeline stage with vector search. If your team already uses MongoDB, you can run vector queries using familiar syntax:

# MongoDB vCore vector search via aggregation pipeline
pipeline = [
    {
        "$search": {
            "cosmosSearch": {
                "vector": query_embedding,
                "path": "embedding",
                "k": 5,
                "efSearch": 40
            },
            "returnStoredSource": True
        }
    },
    {
        "$project": {
            "title": 1,
            "content": 1,
            "score": {"$meta": "searchScore"}
        }
    }
]

results = collection.aggregate(pipeline)

PostgreSQL API (pgvector)

For teams with PostgreSQL expertise, Cosmos DB for PostgreSQL supports the pgvector extension. Vector operations use standard SQL with the <-> (L2 distance), <=> (cosine distance), and <#> (inner product) operators:

-- Enable pgvector extension
CREATE EXTENSION IF NOT EXISTS vector;

-- Create table with vector column
CREATE TABLE documents (
    id     SERIAL PRIMARY KEY,
    title  TEXT NOT NULL,
    content TEXT,
    embedding vector(1536)
);

-- Create an IVFFlat index for ANN search
CREATE INDEX ON documents
USING ivfflat (embedding vector_cosine_ops)
WITH (lists = 100);

-- Query nearest neighbors
SELECT title, content, embedding <=> '[0.12, -0.03, ...]' AS distance
FROM documents
ORDER BY embedding <=> '[0.12, -0.03, ...]'
LIMIT 5;

Change Feed + AI: Real-Time Intelligence

The Cosmos DB change feed is a persistent, ordered log of all inserts and updates to a container. For AI applications, it is a game-changer: instead of batch-processing data for embeddings or running periodic re-indexing jobs, you can react to data changes in real time.

Common AI patterns with change feed

• Automatic embedding generation — when a new document is inserted without an embedding, trigger an Azure Function that generates and writes back the embedding.
• Real-time content moderation — pipe every new user-generated content item through an AI content safety model the instant it is written.
• Incremental knowledge base updates — as knowledge articles change, automatically re-embed and update the vector index without full rebuilds.
• Event-driven AI pipelines — trigger multi-step AI workflows (summarize, classify, extract entities) on each new data item.

Azure Function with Cosmos DB trigger

import azure.functions as func
import json, os
from openai import AzureOpenAI

app = func.FunctionApp()
openai_client = AzureOpenAI(
    azure_endpoint=os.environ["OPENAI_ENDPOINT"],
    api_version="2024-06-01",
    api_key=os.environ["OPENAI_KEY"]
)

# Cosmos DB trigger — fires on every insert/update
@app.cosmos_db_trigger_v3(
    arg_name="documents",
    container_name="raw-content",
    database_name="ai-app-db",
    connection="CosmosDBConnection",
    lease_container_name="leases",
    create_lease_container_if_not_exists=True
)
# Output binding — writes enriched docs to another container
@app.cosmos_db_output_v3(
    arg_name="outputDoc",
    container_name="enriched-content",
    database_name="ai-app-db",
    connection="CosmosDBConnection"
)
def enrich_with_embedding(documents: func.DocumentList, outputDoc: func.Out[func.DocumentList]):
    enriched = []
    for doc in documents:
        doc_dict = doc.to_dict()

        # Skip if already has embedding
        if "embedding" in doc_dict:
            continue

        # Generate embedding from content
        content = doc_dict.get("content", "")
        embedding = openai_client.embeddings.create(
            input=content,
            model="text-embedding-3-large"
        ).data[0].embedding

        doc_dict["embedding"] = embedding
        doc_dict["enrichment_status"] = "completed"
        enriched.append(func.Document(json.dumps(doc_dict)))

    if enriched:
        outputDoc.set(enriched)

Global Distribution for AI Serving

Most vector databases are inherently single-region. When your users span the globe, every vector search query from a distant region incurs cross-continental latency. Cosmos DB’s global distribution changes this equation fundamentally.

Multi-region writes for AI

With multi-region writes enabled, your application can write new documents and their embeddings to the nearest Cosmos DB region, and those writes are automatically replicated to all other configured regions. This is transformative for AI applications:

• Edge AI serving — deploy your AI application in multiple regions and let each instance read vectors from its local Cosmos DB replica. A user in Singapore gets the same sub-10ms vector search latency as a user in Virginia.
• Distributed ingestion — if your data arrives from multiple geographies (e.g., IoT sensors, global user activity), each region can ingest and embed locally.
• Disaster recovery — automatic failover with zero data loss when using strong consistency, or near-zero data loss with session consistency.

Consistency levels and AI workloads

Cosmos DB offers five consistency levels. For AI applications, the right choice depends on your tolerance for stale vectors:

Key Capabilities at a Glance

Vector Database Comparison

Choosing the right vector store for your AI workload requires evaluating more than just search latency. Here is how Cosmos DB compares against popular alternatives across dimensions that matter for production AI systems:

Pricing Considerations

Cosmos DB pricing revolves around Request Units (RUs) — a normalized measure of compute that abstracts CPU, memory, and I/O. Understanding RU consumption for vector operations is essential to controlling costs.

Capacity models

Cost optimization strategies

• Exclude vector paths from regular indexing — this alone can reduce RU consumption for writes by 30-50%.
• Use quantizedFlat for dev/test — lower memory footprint means lower costs when you do not need production-grade throughput.
• Partition intelligently — co-locate frequently queried documents in the same partition to avoid cross-partition vector queries, which consume more RUs.
• Leverage hierarchical partitioning — for multi-tenant AI applications, use synthetic partition keys like /tenantId to isolate workloads and enable per-tenant cost tracking.
• Consider reserved capacity — 1-year and 3-year reservations offer up to 63% savings on provisioned throughput.

Architecture Best Practices

1. Co-locate embeddings with source documents. Store the vector in the same Cosmos DB document as the text it represents. This guarantees consistency and enables hybrid queries without joins or secondary lookups.
2. Use DiskANN indexing for production workloads. Reserve flat and quantizedFlat for development, testing, or datasets under 10,000 vectors. DiskANN scales to hundreds of millions of vectors with predictable latency.
3. Exclude vector paths from the standard index. Add /embedding/* to excludedPaths in your indexing policy. Vector properties have their own dedicated index and should not be double-indexed.
4. Choose the right consistency level. Default to session consistency for user-facing AI applications. Use bounded staleness when you need cross-session guarantees with minimal latency impact.
5. Implement chunking before embedding. Split large documents into 500-1000 token chunks. Store each chunk as a separate document with a parentDocId field for reassembly during retrieval.
6. Use change feed for embedding pipelines. Decouple document ingestion from embedding generation. Write raw documents first, then let an Azure Function triggered by the change feed handle embedding generation asynchronously.
7. Design partition keys around query patterns. If your AI queries always filter by tenant or category, use that as your partition key. Cross-partition vector queries work but cost significantly more RUs.
8. Enable multi-region writes for global AI apps. Serve vector queries from the nearest region to your users. The latency difference between a local read (2-5ms) and a cross-continent read (150-300ms) is make-or-break for real-time AI experiences.
9. Monitor RU consumption for vector operations. Use Azure Monitor and Cosmos DB’s built-in metrics to track RU usage per query. Set alerts when vector search RU consumption exceeds expected thresholds.
10. Plan for dimension flexibility. Newer embedding models like text-embedding-3-large support configurable dimensions. Start with the full 3072 dimensions for maximum quality, then experiment with reduced dimensions (1536, 256) to find the optimal quality/cost tradeoff for your use case.

Next Steps

Azure Cosmos DB’s evolution into a unified operational and vector database marks a significant shift in how AI applications manage data. Rather than assembling a patchwork of specialized services — one for documents, one for vectors, one for caching — you can build on a single, globally distributed platform that handles all three concerns with industry-leading SLAs.

To get started with Cosmos DB for your AI workloads:

• Explore the vector search quickstart — the official documentation includes interactive notebooks and sample datasets to get your first vector queries running in minutes.
• Try the AI sample gallery — Microsoft provides end-to-end sample applications including RAG chatbots, AI agents with memory, and recommendation engines built on Cosmos DB.
• Experiment with Semantic Kernel — if you are building AI agents, the Semantic Kernel Cosmos DB connector lets you use Cosmos DB as your agent’s persistent memory with minimal configuration.
• Evaluate the serverless tier — spin up a serverless Cosmos DB account at no minimum cost to prototype your vector search architecture before committing to provisioned throughput.
• Review the Azure AI reference architectures — the Azure Architecture Center features validated patterns for production RAG systems, multi-agent architectures, and global AI serving platforms built on Cosmos DB.