ai-assisted

Azure Cosmos DB Agent Kit Just Became Platform Material

How to Build an AI-Ready Data Platform with Azure Cosmos DB Agent Kit

Frank Garofalo

03 Jun 2026 — 14 min read

Azure Cosmos DB Agent Kit is increasingly a platform-level choice for Azure-native AI systems, not just an SDK experiment.

If you are still treating it like a clever demo, you may be optimizing for the wrong outcome.

Microsoft’s Azure guidance now clearly supports AI application patterns that combine Azure Cosmos DB, Azure OpenAI, governed access paths, and operational architecture patterns. That is not a universal mandate for every workload, but it is a strong signal that these designs are supported and increasingly common for Azure-native systems (Azure Cosmos DB docs: https://learn.microsoft.com/en-us/azure/cosmos-db/, Azure Architecture Center: https://learn.microsoft.com/en-us/azure/architecture/).

The more useful question is no longer “Can we make an agent work?”

It is:

“Can we run an AI application with predictable latency, governed access, repeatable delivery, and clear operational ownership?”

That is the right framing for Azure Cosmos DB Agent Kit.

Last quarter, a 14-person platform team I worked with hit a wall after their third internal agent because every workflow stored memory differently, one hot partition started returning 429s, and nobody could explain which prompts had touched customer escalation data.

This tutorial shows how to avoid that outcome with one reference scenario:

an internal support operations assistant for a multi-tenant SaaS platform.

By the end, you will have built a minimal Azure-native slice that:

stores conversation state in Azure Cosmos DB
stores tenant-scoped grounding metadata in a separate Cosmos DB container
uses managed identity for Cosmos DB and Azure OpenAI
retrieves context for one tenant
sends a grounded request to Azure OpenAI
persists the assistant response
emits basic telemetry you can expand later

Prerequisites

Before you start, have these ready:

an Azure subscription
permission to create Azure Cosmos DB, Azure OpenAI, Azure Functions or another app host, and Application Insights
permission to assign managed identity roles
an Azure OpenAI deployment already created
Python packages for azure-identity, azure-cosmos, and openai
Azure CLI or PowerShell if you want to provision from scripts

Expected Azure access:

rights to create resource groups and resources
rights to assign Cosmos DB Built-in Data Contributor
rights to assign Cognitive Services OpenAI User

The build order in this tutorial is deliberate:

define the reference architecture
model the data
provision Cosmos DB
wire identity
validate state and context reads
add Azure OpenAI
run one end-to-end request

Step 1: Start from a production-shaped reference architecture

For this tutorial, keep the architecture intentionally small:

client application
Python orchestration layer using Cosmos DB Agent Kit-style state patterns
Azure Cosmos DB for low-latency operational state and simple grounding metadata
Azure OpenAI for inference
Data API builder where you want governed REST or GraphQL access paths
Azure monitoring stack for logs, traces, and metrics

This is not a vector-search tutorial. Here, Cosmos DB is used for operational state plus tenant-scoped grounding metadata. If you need vector indexing and similarity search, that is a valid adjacent pattern, but it should be designed explicitly rather than implied.

The diagram below shows the hot path: user request, app orchestration, Cosmos DB state/context, Azure OpenAI, and telemetry.

What to observe: the critical path is intentionally small. Managed identity sits on the path to both Cosmos DB and Azure OpenAI, which reduces secret sprawl and simplifies hardening.

Limitation: this is still a minimal slice. It does not yet show private networking, multi-region failover, or advanced retrieval.

Validation checkpoint: confirm your team can name each component’s responsibility in one sentence. If not, the design is still too blurry.

Step 2: Model Cosmos DB data for the reference scenario

For the multi-tenant support assistant, use two containers:

conversation-state for session-scoped turns and workflow state
knowledge for tenant-scoped support metadata such as runbook summaries, incident notes, and policy snippets

That split matters because the access patterns differ.

A useful baseline:

conversation-state partition key: /sessionId
knowledge partition key: /tenantId

Why this matters operationally: partition keys are not a storage detail; they determine scale behavior, query efficiency, and where hot spots appear.

Limitation: no single partition key is universally correct. /sessionId is efficient for ordered history reads, while /tenantId is better for tenant-scoped grounding and governance.

Typical trade-offs:

/sessionId: efficient for one active conversation, but chatty sessions can skew load
/tenantId: efficient for tenant-scoped retrieval, but very large tenants can become hot
/workflowId: useful for long-running task state
/userId: useful when user-centric history dominates

Also decide lifecycle rules early:

ephemeral memory: short TTL
operator diagnostics: medium retention
audit records: explicit governed retention
long-term context: summarized, not raw transcript forever

Validation checkpoint: write down your primary read path for each container before you create it. If you cannot describe the dominant query shape, you are not ready to choose a partition key.

Step 3: Define the minimal app configuration

The first code sample sets the shape of the app configuration: Cosmos endpoint, database and containers, Azure OpenAI endpoint and model, and Application Insights connection string.

# Minimal config for an AI-ready Cosmos DB Agent Kit app
from dataclasses import dataclass
import os

@dataclass(frozen=True)
class Settings:
    cosmos_endpoint: str = os.environ["COSMOS_ENDPOINT"]
    cosmos_database: str = os.environ.get("COSMOS_DATABASE", "ai-platform")
    state_container: str = os.environ.get("STATE_CONTAINER", "conversation-state")
    context_container: str = os.environ.get("CONTEXT_CONTAINER", "knowledge")
    openai_endpoint: str = os.environ["AZURE_OPENAI_ENDPOINT"]
    openai_model: str = os.environ.get("AZURE_OPENAI_MODEL", "gpt-4o-mini")
    app_insights_connection_string: str = os.environ.get("APPLICATIONINSIGHTS_CONNECTION_STRING", "")

settings = Settings()
print({"database": settings.cosmos_database, "model": settings.openai_model})

Why this matters operationally: explicit configuration boundaries make environments reproducible and easier to audit.

Limitation: this is only configuration shape, not secret rotation, policy enforcement, or deployment-time validation.

Validation checkpoint: run it and confirm the expected database and model names print correctly for your environment.

Step 4: Provision Cosmos DB with a clear capacity choice

For this tutorial, we are using Azure Cosmos DB for NoSQL because it fits JSON operational state well.

Capacity planning depends on workload shape:

serverless can be attractive for low-volume or bursty early workloads
provisioned throughput or autoscale is usually the better fit once request patterns become more predictable

The Bicep below is intentionally more production-shaped than a bare minimum example. It uses autoscale throughput for the database and standard containers. If you choose serverless instead, treat that as an explicit trade-off for early-stage or low-volume workloads, not as a default production recommendation.

// Cosmos DB account and SQL database for agent state and context
param location string = resourceGroup().location
param accountName string
param databaseName string = 'ai-platform'
param maxThroughput int = 4000

resource cosmos 'Microsoft.DocumentDB/databaseAccounts@2024-05-15' = {
  name: accountName
  location: location
  kind: 'GlobalDocumentDB'
  properties: {
    databaseAccountOfferType: 'Standard'
    locations: [{ locationName: location, failoverPriority: 0 }]
    consistencyPolicy: { defaultConsistencyLevel: 'Session' }
  }
}

resource sqlDb 'Microsoft.DocumentDB/databaseAccounts/sqlDatabases@2024-05-15' = {
  name: '${cosmos.name}/${databaseName}'
  properties: {
    resource: { id: databaseName }
    options: {
      autoscaleSettings: {
        maxThroughput: maxThroughput
      }
    }
  }
}

Why this matters operationally: throughput configuration is part of the architecture, not an afterthought, because it determines cost ceilings and how the system behaves under load.

Limitation: this still does not include multi-region design, backup policy tuning, or workload-specific indexing optimization.

Next, create the two baseline containers with different partition keys.

// Containers for conversation state and knowledge context
param accountName string
param databaseName string = 'ai-platform'

resource stateContainer 'Microsoft.DocumentDB/databaseAccounts/sqlDatabases/containers@2024-05-15' = {
  name: '${accountName}/${databaseName}/conversation-state'
  properties: {
    resource: {
      id: 'conversation-state'
      partitionKey: { paths: ['/sessionId'], kind: 'Hash' }
      indexingPolicy: { indexingMode: 'consistent' }
    }
  }
}

resource knowledgeContainer 'Microsoft.DocumentDB/databaseAccounts/sqlDatabases/containers@2024-05-15' = {
  name: '${accountName}/${databaseName}/knowledge'
  properties: {
    resource: {
      id: 'knowledge'
      partitionKey: { paths: ['/tenantId'], kind: 'Hash' }
      indexingPolicy: { indexingMode: 'consistent' }
    }
  }
}

Why this matters operationally: separate containers let you tune scale, indexing, retention, and governance independently for state versus grounding data.

Limitation: default indexing is convenient, but not always RU-efficient once query patterns stabilize.

Validation checkpoint: deploy the template and verify the database plus both containers exist in the Cosmos DB account.

If you prefer CLI-driven bootstrap, this script provisions the core resources and app settings.

# Provision production-minded Azure resources and baseline app settings
param(
  [string]$ResourceGroup = "rg-ai-platform",
  [string]$Location = "eastus",
  [string]$CosmosAccount = "cosmosaiplatform123",
  [string]$AppName = "func-ai-agent-123"
)

az group create -n $ResourceGroup -l $Location | Out-Null
az cosmosdb create -g $ResourceGroup -n $CosmosAccount --enable-free-tier false --locations regionName=$Location failoverPriority=0 | Out-Null
az cosmosdb sql database create -g $ResourceGroup -a $CosmosAccount -n "ai-platform" --max-throughput 4000 | Out-Null
az cosmosdb sql container create -g $ResourceGroup -a $CosmosAccount -d "ai-platform" -n "conversation-state" --partition-key-path "/sessionId" | Out-Null
az cosmosdb sql container create -g $ResourceGroup -a $CosmosAccount -d "ai-platform" -n "knowledge" --partition-key-path "/tenantId" | Out-Null
az functionapp identity assign -g $ResourceGroup -n $AppName | Out-Null
az functionapp config appsettings set -g $ResourceGroup -n $AppName --settings COSMOS_DATABASE=ai-platform STATE_CONTAINER=conversation-state CONTEXT_CONTAINER=knowledge | Out-Null

Why this matters operationally: even quick bootstrap scripts should make capacity and naming explicit.

Limitation: this is still bootstrap automation, not a full landing-zone or policy-driven deployment model.

Validation checkpoint: confirm the function app has a managed identity and the app settings were written successfully.

Step 5: Wire identity before data access

The cleanest Azure-native pattern is:

your app uses a managed identity
that identity gets data-plane access to Cosmos DB
the same identity gets permission to call Azure OpenAI
secrets are minimized or eliminated

# Bind managed identity permissions and monitoring hooks
param(
  [string]$ResourceGroup = "rg-ai-platform",
  [string]$CosmosAccount = "cosmosaiplatform123",
  [string]$AppName = "func-ai-agent-123",
  [string]$OpenAIName = "aoai-platform-123"
)

$principalId = az functionapp identity show -g $ResourceGroup -n $AppName --query principalId -o tsv
$cosmosId = az cosmosdb show -g $ResourceGroup -n $CosmosAccount --query id -o tsv
$openAiId = az cognitiveservices account show -g $ResourceGroup -n $OpenAIName --query id -o tsv

az role assignment create --assignee-object-id $principalId --assignee-principal-type ServicePrincipal --role "Cosmos DB Built-in Data Contributor" --scope $cosmosId | Out-Null
az role assignment create --assignee-object-id $principalId --assignee-principal-type ServicePrincipal --role "Cognitive Services OpenAI User" --scope $openAiId | Out-Null
az monitor app-insights component create -g $ResourceGroup -a "$AppName-ai" -l eastus --application-type web | Out-Null
$cs = az monitor app-insights component show -g $ResourceGroup -a "$AppName-ai" --query connectionString -o tsv
az functionapp config appsettings set -g $ResourceGroup -n $AppName --settings APPLICATIONINSIGHTS_CONNECTION_STRING=$cs | Out-Null

Why this matters operationally: identity boundaries are one of the fastest ways to reduce risk and improve auditability in agent systems.

Limitation: role assignment alone does not define your full governance model. You still need environment separation, approved access paths, and review gates.

A practical rule: do not let agents talk directly to “whatever data exists.” Give them approved access paths.

That can mean:

direct Cosmos DB container access for tightly bounded runtime logic
Data API builder endpoints for internal service consumption
MCP-oriented access patterns when you need controlled agent integrations (https://learn.microsoft.com/en-us/azure/data-api-builder/)

Validation checkpoint: confirm the managed identity can authenticate and has the expected roles on both Cosmos DB and Azure OpenAI.

Step 6: Create Cosmos DB clients and validate the containers

Now move into the application layer. This sample uses DefaultAzureCredential and creates the database plus two containers if they do not already exist.

# Create Cosmos DB clients with managed identity and ensure containers exist
from azure.identity import DefaultAzureCredential
from azure.cosmos import CosmosClient, PartitionKey
import os

credential = DefaultAzureCredential()
client = CosmosClient(os.environ["COSMOS_ENDPOINT"], credential=credential)

db = client.create_database_if_not_exists(id="ai-platform")
state = db.create_container_if_not_exists(
    id="conversation-state",
    partition_key=PartitionKey(path="/sessionId"),
)
context = db.create_container_if_not_exists(
    id="knowledge",
    partition_key=PartitionKey(path="/tenantId"),
)

print(db.id, state.id, context.id)

Why this matters operationally: using Azure identity instead of connection strings keeps the runtime model aligned with Azure-native security practices.

Limitation: create_*_if_not_exists is useful for validation and development, but many teams will prefer infrastructure-managed creation in production.

Validation checkpoint: run it and confirm it prints ai-platform conversation-state knowledge.

Step 7: Persist and query conversation state correctly

Now write a user turn into the conversation-state container and query all turns for a session ordered by timestamp.

# Store and retrieve conversation state in Cosmos DB
from azure.identity import DefaultAzureCredential
from azure.cosmos import CosmosClient
import os, time

client = CosmosClient(os.environ["COSMOS_ENDPOINT"], credential=DefaultAzureCredential())
container = client.get_database_client("ai-platform").get_container_client("conversation-state")

session_id = "s-1001"
container.upsert_item({
    "id": f"{session_id}-1",
    "sessionId": session_id,
    "role": "user",
    "content": "What changed in our Q2 pipeline?",
    "ts": int(time.time())
})

items = list(container.query_items(
    query="SELECT * FROM c WHERE c.sessionId = @sid ORDER BY c.ts",
    parameters=[{"name": "@sid", "value": session_id}],
    partition_key=session_id,
))
print([{"role": i["role"], "content": i["content"]} for i in items])

Why this matters operationally: the query is aligned to the partition key, which keeps session-history reads efficient and avoids unnecessary fan-out.

Limitation: this pattern is ideal for one-session reads, not for analytics or broad cross-session reporting.

Validation checkpoint: confirm the output includes the user turn you just wrote.

Step 8: Retrieve tenant-scoped grounding metadata

Now fetch a small set of tenant-scoped context records from the knowledge container. This is a simple text filter for support metadata, not vector search.

# Retrieve contextual knowledge from Cosmos DB for grounding
from azure.identity import DefaultAzureCredential
from azure.cosmos import CosmosClient
import os

client = CosmosClient(os.environ["COSMOS_ENDPOINT"], credential=DefaultAzureCredential())
container = client.get_database_client("ai-platform").get_container_client("knowledge")

tenant_id = "contoso"
query = """
SELECT TOP 3 c.title, c.summary
FROM c WHERE c.tenantId = @tenant AND CONTAINS(c.summary, @term, true)
"""
results = list(container.query_items(
    query=query,
    parameters=[{"name": "@tenant", "value": tenant_id}, {"name": "@term", "value": "pipeline"}],
    partition_key=tenant_id,
))
grounding = "\n".join(f"- {r['title']}: {r['summary']}" for r in results)
print(grounding)

Why this matters operationally: grounding the model with explicit tenant-scoped records is more controllable than asking it to “remember” everything.

Limitation: CONTAINS on summary text is a simple metadata filter, not a substitute for vector indexing, ranking, or semantic retrieval.

If you later need vector-enabled retrieval in Cosmos DB, the design changes:

documents include an embedding field
indexing policy and container capabilities must support vector search
the query path becomes similarity-based rather than CONTAINS-based
evaluation shifts from “did I find matching text?” to retrieval quality, latency, and cost

That is a valid next step, but it should be implemented intentionally. This tutorial stays focused on operational state plus simple grounding metadata.

Validation checkpoint: seed a few knowledge records for one tenant and confirm the query returns only that tenant’s matching summaries.

Step 9: Add Azure OpenAI after the data path works

Too many teams start with the model and then scramble to retrofit state, access control, and telemetry.

Do it in this order:

prove the state path
prove the grounding path
prove identity and authorization
then add model inference

This sample calls Azure OpenAI using Azure AD token-based authentication and emits a simple structured telemetry event.

# Call Azure OpenAI with state + Cosmos context and emit structured telemetry
from azure.identity import DefaultAzureCredential, get_bearer_token_provider
from openai import AzureOpenAI
import json, os, time

token_provider = get_bearer_token_provider(DefaultAzureCredential(), "https://cognitiveservices.azure.com/.default")
client = AzureOpenAI(
    azure_endpoint=os.environ["AZURE_OPENAI_ENDPOINT"],
    api_version="2024-02-01",
    azure_ad_token_provider=token_provider,
)

messages = [
    {"role": "system", "content": "Answer using the provided business context."},
    {"role": "user", "content": "Context:\n- Q2 ETL moved to incremental loads\nQuestion: Summarize the impact."},
]
resp = client.chat.completions.create(model=os.environ["AZURE_OPENAI_MODEL"], messages=messages)
event = {"event": "agent_request", "latency_ms": 123, "model": os.environ["AZURE_OPENAI_MODEL"], "ts": int(time.time())}
print(resp.choices[0].message.content)
print(json.dumps(event))

Why this matters operationally: model choice, latency, and event shape are explicit, which is the foundation of operability.

Limitation: this sample emits a simple event shape, not full distributed tracing, prompt lineage, or dependency correlation.

Validation checkpoint: confirm you receive a completion and that the telemetry event includes the expected model name.

Step 10: Assemble the end-to-end request flow

This final application sample reads history from Cosmos DB, retrieves tenant context, calls Azure OpenAI, and persists the assistant reply back into conversation state.

# End-to-end minimal agent request handler using Cosmos DB and Azure OpenAI
from azure.identity import DefaultAzureCredential, get_bearer_token_provider
from azure.cosmos import CosmosClient
from openai import AzureOpenAI
import os, time, json

cred = DefaultAzureCredential()
cosmos = CosmosClient(os.environ["COSMOS_ENDPOINT"], credential=cred)
db = cosmos.get_database_client("ai-platform")
state = db.get_container_client("conversation-state")
knowledge = db.get_container_client("knowledge")
aoai = AzureOpenAI(azure_endpoint=os.environ["AZURE_OPENAI_ENDPOINT"], api_version="2024-02-01",
                   azure_ad_token_provider=get_bearer_token_provider(cred, "https://cognitiveservices.azure.com/.default"))

session_id, tenant_id, prompt = "s-42", "contoso", "How should we explain the new ingestion design?"
ctx = list(knowledge.query_items("SELECT TOP 2 c.summary FROM c WHERE c.tenantId=@t",
                                 parameters=[{"name": "@t", "value": tenant_id}], partition_key=tenant_id))
history = list(state.query_items("SELECT * FROM c WHERE c.sessionId=@s ORDER BY c.ts",
                                 parameters=[{"name": "@s", "value": session_id}], partition_key=session_id))
messages = [{"role": "system", "content": "Use context and keep answers concise."}] + \
           [{"role": h["role"], "content": h["content"]} for h in history] + \
           [{"role": "user", "content": f"Context: {ctx}\nQuestion: {prompt}"}]
reply = aoai.chat.completions.create(model=os.environ["AZURE_OPENAI_MODEL"], messages=messages).choices[0].message.content
state.upsert_item({"id": f"{session_id}-{int(time.time())}", "sessionId": session_id, "role": "assistant", "content": reply, "ts": int(time.time())})
print(json.dumps({"sessionId": session_id, "reply": reply}))

Why this matters operationally: this is the smallest useful slice that has memory, grounding, inference, and persistence in one traceable flow.

Limitation: it still assumes a single request path and does not yet cover retries, fallback models, or tool orchestration.

The sequence diagram below is the mental model operators should be able to follow.

Why this matters operationally: every request should be traceable across app, Cosmos DB, Azure OpenAI, and monitoring.

Limitation: if your traces break at service boundaries, this clean sequence becomes much harder to observe in practice.

Validation checkpoint: run one request end to end and verify that a new assistant turn is written back to conversation-state.

Step 11: Add observability and failure handling

In production, expect these failure domains:

model latency spikes
Cosmos DB 429 throttling
partition hot spots
stale or low-quality memory
tool timeouts
authorization failures

At minimum, capture:

correlation ID
session or workflow ID
model name
prompt metadata, not just raw prompt text
grounding decision summary
Cosmos DB request charge where relevant
latency by dependency
tool outcomes
final response status

Track the things that actually matter:

p95 latency
429 rate
token consumption
failed tool calls
cost growth trends
error rate by dependency

Why this matters operationally: request-level causality is what turns logs into incident response.

Limitation: more telemetry is not automatically better; without naming standards and dashboards, teams still struggle to act on it.

Validation checkpoint: confirm you can answer, for one request, which tenant was queried, which model was called, how long it took, and whether Cosmos DB throttled.

Step 12: Make cost and performance trade-offs explicit

For this architecture, the main cost surfaces are:

Cosmos DB request units
Cosmos DB storage
model tokens
network egress in some integration patterns
observability overhead

Three factors drive Cosmos DB RU efficiency especially hard:

partitioning
indexing policy
document shape

Large noisy documents with fields you never query can waste RU. So can cross-partition queries that ignore the primary access pattern.

When is Cosmos DB retrieval enough for this pattern?

when the grounding corpus is operationally adjacent
when latency matters more than maximum retrieval feature depth
when governance simplicity matters
when you want fewer moving parts

When might you choose a separate retrieval architecture?

when corpus size and retrieval complexity grow independently of operational state
when teams need specialized ranking pipelines
when analytical and retrieval workloads diverge sharply

Why this matters operationally: “supported by Azure” means viable and supportable, not automatically optimal for every workload.

Limitation: the right retrieval boundary still depends on workload shape, governance requirements, and how sophisticated your search needs to be.

Practical controls that pay off quickly:

TTL for ephemeral turns
summarization for long conversations
selective persistence for tool outputs
caching for repeated context
workload-specific throughput planning

Step 13: Keep operational and analytical concerns separate

In this pattern, Cosmos DB remains the operational store for:

low-latency session state
workflow state
tenant-scoped grounding metadata close to the app
auditable interaction records

Bring Fabric in when the question becomes analytics:

downstream analytics
governed BI
broader enterprise data product integration
unified analytics across operational and historical data (https://learn.microsoft.com/en-us/fabric/)

Use Power Platform or Microsoft 365 Copilot extensibility when the requirement is user experience in those ecosystems, not just because those tools are available (https://learn.microsoft.com/en-us/microsoft-365/copilot/).

Why this matters operationally: separating operational and analytical paths keeps latency-sensitive agent flows simpler and easier to govern.

Limitation: some organizations will still need integration patterns across both worlds, so the boundary must be designed rather than assumed.

Step 14: Standardize before you scale

Create reusable templates for:

Cosmos DB account and container setup
managed identity bindings
model configuration
telemetry wiring
naming and tagging
environment separation

Require review for:

partition key changes
new tool onboarding
prompt changes that affect regulated data paths
retention exceptions
cost regressions

Write runbooks early for:

429 throttling response
model fallback behavior
incident triage across app, database, and model dependencies
data retention exceptions
authorization failures

Why this matters operationally: the fastest agent team is usually the one inheriting a disciplined platform, not improvising every deployment.

Limitation: standardization helps only if teams actually adopt the templates and review gates.

Final takeaway

Azure Cosmos DB Agent Kit makes the most sense when you stop asking whether it can power a demo and start asking whether it can anchor a low-latency, governed, observable AI data platform on Azure.

For the support-operations scenario in this tutorial, that means:

model operational state deliberately in Cosmos DB
separate session memory from tenant grounding metadata
use managed identity for Cosmos DB and Azure OpenAI
validate each step before adding more autonomy
add telemetry, retention, and cost guardrails before rollout
keep operational and analytical concerns separate

If you do that, you are not just building an agent.

You are building a platform your teams can actually scale.

What would you standardize first in your environment: partition keys, managed identity, telemetry, or retention?

#CosmosDB #AIAgents #DataArchitecture

Sources & References

Try it yourself

Run this tutorial as a Jupyter notebook: Download runbook.ipynb (29 cells, 29 KB).