Azure Cosmos DB Agent Memory and Agentic Retrieval Toolkits

By Law Wen Feng, Principal Solution Architect | wenfeng.my

If you have been building AI agents in the Azure ecosystem over the past year, you have probably hit the same wall I have: agents are stateless by default. The moment you need one to remember context across sessions, pull from a knowledge base, or ground its answers in enterprise data, you are duct-taping a vector store to a chat model and hoping it scales.

Most teams I work with in Southeast Asia — Kuala Lumpur, Singapore, Jakarta — start their agent journey the same way: an LLM, a prompt, maybe a tool call or two. It works beautifully for demos. But the moment you ask the agent to recall what happened three conversations ago, or retrieve the right clause from a 50,000-page compliance library, the architecture crumbles.

This article is about using Azure Cosmos DB as the database layer for AI agents — not as a generic "use Cosmos for everything" pitch, but as a practical, field-tested approach to agents that remember, retrieve, and act on enterprise data. And it covers the two purpose-built toolkits Microsoft shipped for exactly this problem: the Agent Memory Toolkit and the Agentic Retrieval toolkit.

The Agent Memory Problem

An AI agent has three core capabilities: reasoning (the LLM), acting (tool calls), and — critically — remembering. Without memory, your agent perpetually starts from zero. It cannot learn from past interactions, maintain user preferences, or recall which documents it already retrieved.

Enterprise agents really need three kinds of memory:

Short-term memory — the conversation context window. The LLM handles this natively, but it is bounded and ephemeral.
Long-term memory — persistent facts, user preferences, past decisions, interaction history. This needs a durable store.
Episodic memory — the ability to retrieve relevant past episodes or knowledge chunks when facing a new task. This is where retrieval-augmented generation (RAG) comes in.

The usual response is to bolt on three separate systems: a cache for conversation history, a relational database for structured state, and a vector database for semantic search. That is three operational burdens, three integration points, three things to monitor, secure, and scale.

Cosmos DB can collapse all three into a single, globally distributed database with native vector indexing.

Why Cosmos DB for Agent Memory

I have spent enough time with Pinecone, Weaviate, and Milvus to respect what they do. They are purpose-built vector databases, and they do vector search extremely well. But here is what I have learned building production agents for enterprise clients: you rarely need just vector search.

You need vector search plus structured queries. Semantic retrieval plus transactional consistency. Global distribution plus single-digit-millisecond latency. And you need all of it in a service your security team already trusts.

Cosmos DB for NoSQL gives you:

Native vector indexing with DiskANN-based algorithms (the same research behind Microsoft's large-scale search). Three index types: flat for brute-force exact search, quantizedFlat for compressed fast search, and diskANN for high-accuracy approximate search at scale.
Hybrid queries that combine vector similarity with SQL-style filtering — tenant, document type, recency, permissions — in a single query.
Full-text and hybrid search alongside vector search in the same container.
Native TTL to expire ephemeral memories with zero cleanup jobs.
Change feed for event-driven reactions when new memories are written.
Global distribution with multi-region writes, which matters when your agents serve users across ASEAN from multiple Azure regions.

Setting Up a Cosmos DB Container for Agent Memory

First, provision the infrastructure. Two details here trip people up, so I will call them out explicitly: the account-create command uses --locations (plural), and the vector embedding policy cannot be modified in place after container creation — you can only add new vector paths, or drop a path and re-add it with new settings, so define it up front.

# Account (note: --locations, not --location) az cosmosdb create \ --name agent-memory-prod \ --resource-group rg-ai-agents \ --locations regionName=southeastasia failoverPriority=0 \ --default-consistency-levelSession

# Database az cosmosdb sql database create \ --account-name agent-memory-prod \ --resource-group rg-ai-agents \ --name AgentMemoryDB \ --throughput 400 # Container: indexing policy (with vectorIndexes) and vector embedding policy # are passed as two separate parameters az cosmosdb sql container create \ --account-name agent-memory-prod \ --resource-group rg-ai-agents \ --database-name AgentMemoryDB \ --name memories \ --partition-key-path "/tenantId" \ --throughput 400 \ --idx @index-policy.json \ --vector-embeddings@embedding-policy.json

The two policy files:

// index-policy.json { "indexingMode": "consistent", "automatic": true, "includedPaths": [{ "path": "/*" }], "excludedPaths": [{ "path": "/vector/*" }], "vectorIndexes": [{ "path": "/vector", "type": "quantizedFlat" }] } // embedding-policy.json { "vectorEmbeddings": [{ "path": "/vector", "dataType": "float32", "dimensions": 1536, "distanceFunction": "cosine" }] }

The equivalent Bicep resource (Microsoft.DocumentDB/databaseAccounts/sqlDatabases/containers, API version 2025-10-15 or later) takes the same two structures: properties.indexingPolicy.vectorIndexes and properties.vectorEmbeddingPolicy.vectorEmbeddings.

One sizing note: quantizedFlat and diskANN indexes require at least 1,000 vectors before the index kicks in — below that threshold Cosmos DB falls back to a brute-force scan. For agent memory stores that grow over time this is fine; for small tests, expect full-scan behaviour.

Storing and Retrieving Agent Memories

A single Cosmos DB document can hold structured metadata, raw text, and the vector embedding together. No separate collections needed.

from azure.cosmos import CosmosClient from openai import AzureOpenAI from datetime import datetime, timezone cosmos = CosmosClient( "https://agent-memory-prod.documents.azure.com:443/", credential="<managed-identity-or-key>" ) container = cosmos.get_database_client("AgentMemoryDB") \
.get_container_client("memories") oai = AzureOpenAI( azure_endpoint="https://your-openai.openai.azure.com/", api_version="2025-06-01" ) def get_embedding(text: str) -> list[float]: resp = oai.embeddings.create( model="text-embedding-3-small", input=text) return resp.data[0].embedding def store_memory(tenant_id, user_id, session_id, memory_type, content, metadata=None): """Store an agent memory with its vector embedding.""" return container.upsert_item({ "id": f"{session_id}-{datetime.now(timezone.utc).timestamp()}", "tenantId": tenant_id, "userId": user_id, "sessionId": session_id, "memoryType": memory_type, # conversation | fact | preference | document "content": content, "vector": get_embedding(content), "createdAt": datetime.now(timezone.utc).isoformat(), "ttl": 2592000, # 30 days; requires TTL enabled at container level "metadata": metadata or {} })

Notice the partition strategy: tenantId as the partition key keeps all of a tenant's memories co-located. For multi-tenant agent deployments you get isolation by design, and tenant-scoped queries hit a single logical partition.

Retrieval in an agentic context is not "search and return top-k". The agent reasons about what to retrieve, when, and how to use it. Here is a hybrid query that combines semantic similarity with structured filters in a single round-trip:

def agentic_retrieve(tenant_id, query_text, memory_type=None, top_k=10): """Vector similarity + structured filters in one query.""" qv = get_embedding(query_text) where = ["c.tenantId = @tid"] params = [{"name": "@tid", "value": tenant_id}] if memory_type: where.append("c.memoryType = @mt") params.append({"name": "@mt", "value": memory_type}) params.append({"name": "@topk", "value": int(top_k)}) # TOP accepts a parameter in parameterized queries query = f""" SELECT TOP @topk c.id, c.content, c.memoryType, VectorDistance(c.vector, @embedding) AS similarity FROM c WHERE {" AND ".join(where)} ORDER BY VectorDistance(c.vector, @embedding) """ params.append({"name": "@embedding", "value": qv}) return list(container.query_items( query=query, parameters=params, partition_key=tenant_id))

Two things that bite teams here. First, always include a TOP clause: without it, the vector search tries to score far more candidates than you need, burning RUs and adding latency. Second, pass TOP as a parameter in parameterized queries (as shown above) rather than f-string interpolation — Cosmos DB supports TOP @param, which keeps the query cacheable and avoids building strings around user-influenced values.

The Agent Memory Toolkit: Stop Building This Yourself

Everything above is the do-it-yourself path. The reason this topic deserves fresh attention in 2026 is that Microsoft shipped two purpose-built toolkits on top of Cosmos DB for exactly this workload.

The Agent Memory Toolkit (AzureCosmosDB/AgentMemoryToolkit on GitHub, public preview) is a Python SDK — pip install azure-cosmos-agent-memory — that gives your agent both raw conversation history and derived memory: thread summaries, extracted facts, and cross-thread user profiles, all searchable semantically. The distillation pipeline can run in-process (zero extra infrastructure) or in a sibling Azure Durable Functions app that watches the Cosmos DB change feed. Sync and async clients mirror each other.

from azure.cosmos.agent_memory import CosmosMemoryClient memory = CosmosMemoryClient( cosmos_endpoint=COSMOS_DB_ENDPOINT, cosmos_database="ai_memory", cosmos_container="memories", ai_foundry_endpoint=AI_FOUNDRY_ENDPOINT, embedding_deployment_name="text-embedding-3-large", chat_deployment_name="gpt-4o-mini", use_default_credential=True, # Entra ID, no API keys ) memory.connect_cosmos() # auto-creates database + containers memory.upsert_memory(user_id="user-001", thread_id=THREAD, role="user", content="I love Cosmos DB.") memory.process_now(user_id="user-001", thread_id=THREAD) # -> thread summary + fact extraction + user profile update hits = memory.search_cosmos(user_id="user-001", query_text="Cosmos DB preferences", top=5)

If you are on the Microsoft Agent Framework, it gets even shorter: the agent-framework-azure-cosmos-memory package (preview) ships a CosmosMemoryContextProvider you attach to your agent once. Its before_run hook searches the user's memory and injects relevant context; its after_run hook stores new turns and lets the toolkit extract facts and summaries in the background. Memory follows the user across threads and sessions via a stable user_id, with no orchestration code on your side.

The companion piece is the Agentic Retrieval accelerator (AzureCosmosDB/AgenticRetrieval): a self-correcting, multi-stage RAG pipeline for questions that need multi-hop reasoning. Instead of one search-and-answer pass, it drafts a preliminary answer, identifies what is missing, decomposes the gap into focused sub-questions, retrieves targeted evidence per sub-question across multiple Cosmos DB containers, and synthesises a grounded final answer. It can optionally use the Cosmos DB semantic reranker to reorder results before synthesis. For high-stakes domains — legal, financial, medical — that completeness loop is exactly what a naive single-shot RAG pipeline cannot deliver.

And for agents that speak MCP, the Cosmos DB MCP Toolkit (AzureCosmosDB/MCPToolKit) is a production-ready MCP server with Entra ID authentication, document operations, vector search, hybrid search with reciprocal rank fusion, and schema discovery — deployable to Azure Container Apps in minutes.

Where You Still Build Your Own

The toolkits are preview, Python-first, and opinionated. You still want the manual pattern from the previous section when:

Your stack is .NET or Java today (the SDKs will catch up; the DIY path works now).
You need a partition key or retention policy the toolkit's defaults do not match.
Your compliance posture requires you to own every line of data-access code.

In those cases the reference architecture is straightforward: documents arrive from SharePoint, Blob Storage, or databases; an ingestion pipeline (Azure Functions works well) chunks and embeds them into Cosmos DB with provenance metadata; the agent's tool-calling loop invokes retrieval only when reasoning requires it; and the change feed fans new memories out to anything else that needs them.

Pitfalls I See Repeatedly

1. Embedding everything without a filtering strategy. Teams dump every conversation turn into the vector store. After a few weeks, retrieval quality drops because the index is polluted with trivial exchanges. Classify memories at write time. Store only meaningful interactions and assign TTLs per type — conversation turns might expire in 7 days while extracted facts persist. And remember: per-item ttl only takes effect when TTL is enabled at the container level.

2. Over-relying on vector similarity. A user asking "what did the finance team approve last quarter?" needs temporal filtering, not just semantic matching. Always combine vector search with structured filters — that hybrid query is Cosmos DB's actual superpower over standalone vector databases.

3. Getting the partition key wrong. Partition by document ID instead of tenant ID and your cross-tenant queries become expensive fan-outs, and you lose isolation guarantees. Partition by the highest-level isolation boundary.

4. Ignoring the 1,000-vector index threshold. Below it, quantizedFlat and diskANN silently fall back to brute force. It is fine for small stores, but if your latency numbers look odd at low volume, that is why.

5. Forgetting cost mechanics. Cosmos DB bills per RU. Vector queries cost more than point reads, so always use TOP, filter by partition key where possible, watch RU consumption in Azure Monitor, and consider serverless capacity mode for agent workloads with spiky traffic — which is most of them.

6. Not using the change feed. When a new memory is written, the change feed can trigger downstream updates — search indexes, caches, monitoring. The Agent Memory Toolkit's Durable Functions processor is built exactly this way.

Key Takeaways

Cosmos DB collapses the agent memory stack. One container handles conversation state, long-term facts, and RAG retrieval — no separate vector and relational databases to stitch together.
Hybrid queries are the differentiator. Vector similarity combined with SQL-style filters (tenant, user, time range, type) in a single query is what makes it viable for enterprise agent workloads.
Use the toolkits before building your own. The Agent Memory Toolkit, Agentic Retrieval accelerator, and MCP Toolkit are in public preview and encode patterns most teams otherwise rebuild badly.
Get the data model right. Partition by tenant, define the vector embedding policy at container creation (it cannot be modified in place), enable TTL at container level, and always include a TOP clause in vector queries.
Watch the RU meter. Vector search at scale is real money. Classify memories, expire what does not matter, and alert on cost anomalies early.

The enterprise agents that succeed in production will be the ones with a solid data foundation underneath. Cosmos DB now gives you that foundation — with the agent-specific tooling on top — without the operational overhead of stitching together multiple specialised databases.

This article is part of my ongoing series on building production AI systems on Azure. If you are working on agent architectures in the Southeast Asian enterprise space, feel free to reach out — I am always happy to compare notes.

Follow me on LinkedIn or read more at wenfeng.my.