Store and retrieve embeddings in Azure Cosmos DB

Vector embeddings are the foundation of semantic search. By storing them as document properties in Cosmos DB, you unify your architecture (no separate vector DB) and can seamlessly filter metadata alongside vector similarity.

1. What are Vector Embeddings?

Embeddings are arrays of floating-point numbers generated by machine learning models (like Azure OpenAI, Hugging Face). They capture the semantic meaning of text, images, or audio. Similar concepts result in vectors that are mathematically close to each other.

• Dimensions: The model dictates the array length. For example, Azure OpenAI's text-embedding-ada-002 produces 1,536 dimensions. text-embedding-3-large produces 3,072. Larger dimensions capture more nuance but consume more storage.

2. Designing the Document Structure

Embeddings are simply stored as an array property ([]) inside your JSON document, right next to your standard business data.

{
  "id": "doc-12345",
  "title": "Troubleshooting wireless network connections",
  "content": "This guide covers common WiFi connectivity issues...",
  "category": "networking",
  "productId": "router-x100",
  "createdDate": "2024-06-15T10:30:00Z",
  "embedding": [0.023, -0.041, 0.067, ...] // 1,536 floats
}

3. Pre-requisite: Enable Vector Search

Before you can use this feature, it must be enabled on your Azure Cosmos DB account.

• Portal: Settings > Features > Vector Search for NoSQL API
• CLI: az cosmosdb update --capabilities EnableNoSQLVectorSearch (Note: It can take up to 15 minutes to take effect.)

Under the Hood: The Raw Data Problem

Let's strip away all the database jargon and look at this as a pure technical data-routing problem.

When you use an AI model, you feed it text, and it returns a massive array of numbers. You save that into your database looking exactly like this:

{
  "category": "Finance",
  "embedding": [0.123, -0.045, 0.982, ...] // 1,536 total numbers
}

If you don't configure the database first, it sees that embedding array as just a massive, useless block of text. It has no idea those numbers represent AI meaning.

Here is how the configuration policies solve this hardware-level data problem.

1. Vector Policy = "What type of data is this?"

This rule tells the database's memory manager how to allocate space for that specific JSON key.

• path: /embedding: It tells the engine, "Look inside the JSON for a key named 'embedding'."
• dimensions: 1536: It allocates a strict memory slot for exactly 1,536 numbers. If your AI array has 1,535 or 1,537 numbers, the database will throw a data validation error and reject the write.
• distanceFunction: cosine: It loads a specific mathematical algorithm into the CPU. When you search, the CPU will run a math formula on these 1,536 numbers to see how closely their geometric angles align.

2. Indexing Policy = "How should we store it on the hard drive?"

This tells the database indexer how to organize the data on the physical solid-state drive (SSD) for fast retrieval.

• includedPaths: /*: It builds a standard inverted index (B-Tree) for normal text and numbers (like "category": "Finance"). This allows the database to find matching text rows instantly using simple binary search.
• excludedPaths: /embedding/*: This is the critical step. You are telling the standard text indexer, "Do NOT touch the array of numbers." If a standard database engine tries to index 1,536 individual numbers per row, the indexing engine will run out of RAM and freeze.
• type: diskANN: Because you excluded it from the normal index, you tell the database to route the /embedding path to a specialized engine called DiskANN. This engine builds a physical spatial graph index on the SSD. It links documents together with pointers based on how mathematically close their 1,536 numbers are, allowing lightning-fast search speeds instead of scanning every single row sequentially.

3. Container Creation = "Where does it physically live?"

When you create the container, you physically provision the hardware table in the cloud.

• partition_key: /category: This tells the database how to shard (physically split) your data across different server blades. If a document has "category": "Finance", it goes to Server Rack A. If it has "category": "HR", it goes to Server Rack B.

4. Container Vector Policies

CRITICAL RULE: Vector policies and indexing policies for vectors must be defined at container creation time. They cannot be added or modified later.

A vector policy defines how the database engine handles your embeddings:

• path: JSON path to the property (e.g., /embedding)
• dataType: float32 (default), float16, int8, or uint8
• dimensions: Number of elements (default: 1536)
• distanceFunction: Metric for calculating similarity (cosine, dotproduct, euclidean)

Understanding Distance Functions

1. Cosine: Measures the angle between vectors. Ideal for text embeddings (like Azure OpenAI) where magnitude shouldn't matter. Scores range from -1 (least similar) to +1 (most similar).
2. Dot Product: Measures angle and magnitude. If vectors are normalized (like OpenAI's), this yields the exact same result as Cosine but computes slightly faster.
3. Euclidean: Measures straight-line distance. Lower score = more similar. Rarely used for text.

Understanding Data Types

• float32: Full precision, best accuracy, uses the most storage. Recommended starting point.
• float16: 50% storage reduction with barely noticeable impact on search quality.
• int8 / uint8: Heavily compressed (quantized). Used for massive scale, but requires careful tuning.

Multiple Vector Types

You can store multiple embeddings in one document (e.g., one for the title, one for the body). Just add them to the policy:

vector_embedding_policy = {
    "vectorEmbeddings": [
        {"path": "/titleEmbedding", "dataType": "float32", "distanceFunction": "cosine", "dimensions": 1536},
        {"path": "/contentEmbedding", "dataType": "float32", "distanceFunction": "cosine", "dimensions": 1536}
    ]
}

5. Vector Indexing Policies

Without an index, the database performs a "full scan" (comparing the query to every single document), which is slow and expensive (high RUs).

Cosmos DB offers 3 vector index types:

1. flat: Exact, brute-force search. 100% recall. Limit: 505 dimensions. Best for small datasets.
2. quantizedFlat: Compresses vectors before indexing. Faster brute-force search. Up to 4,096 dimensions. Best for datasets up to ~50k vectors per partition.
3. diskANN: Microsoft Research's algorithm for fast, approximate search (ANN). Up to 4,096 dimensions. Best for large datasets (millions of vectors).

Note: quantizedFlat and diskANN require at least 1,000 vectors to become active. Below that, they silently fall back to full scans.

Crucial Indexing Rule: You must exclude the vector path from standard range indexes (includedPaths), otherwise Cosmos DB wastes storage indexing an array of 1,500 floats that you will never do standard > or < queries on.

6. Implementation Code: Creating a Container

from azure.cosmos import CosmosClient, PartitionKey

# 1. Define Vector Policy
vector_embedding_policy = {
    "vectorEmbeddings": [
        {
            "path": "/embedding",
            "dataType": "float32",
            "distanceFunction": "cosine",
            "dimensions": 1536
        }
    ]
}

# 2. Define Indexing Policy
indexing_policy = {
    "indexingMode": "consistent",
    "automatic": True,
    "includedPaths": [
        {"path": "/*"}
    ],
    "excludedPaths": [
        {"path": "/\"_etag\"/?"},
        {"path": "/embedding/*"} # EXCLUDE standard indexing for the vector array!
    ],
    "vectorIndexes": [
        {"path": "/embedding", "type": "diskANN"} # DEFINE the vector index
    ]
}

# 3. Create Container
container = database.create_container(
    id="knowledge-base",
    partition_key=PartitionKey(path="/category"),
    indexing_policy=indexing_policy,
    vector_embedding_policy=vector_embedding_policy
)

7. Implementation Code: Inserting Documents

To insert a document, you first call the Azure OpenAI API to generate the vector, then you upsert it into Cosmos DB.

from openai import AzureOpenAI
from azure.cosmos import CosmosClient

# Initialize clients
openai_client = AzureOpenAI(
    api_key=api_key,
    api_version="2024-02-01",
    azure_endpoint=endpoint
)
cosmos_client = CosmosClient(cosmos_endpoint, credential=cosmos_key)
container = cosmos_client.get_database_client("support-db").get_container_client("knowledge-base")

# 1. Generate embedding from document content
document_text = f"{title} {content}"
response = openai_client.embeddings.create(
    input=document_text,
    model="text-embedding-ada-002"
)
embedding = response.data[0].embedding # The 1,536 float array

# 2. Create the document object
document = {
    "id": document_id,
    "title": title,
    "content": content,
    "category": category,
    "productId": product_id,
    "createdDate": created_date,
    "embedding": embedding
}

# 3. Insert or update document
# We use upsert so that if the text changes later, we can just overwrite the whole doc
container.upsert_item(document)