Post 2 of 6 on Cosmos DB agent memory schema design before you write a line of agent code.
Post 1 made the case for one database instead of three. This post is about the decisions that determine whether that one database actually holds up: partition key and item shape, both of which you choose before an agent ever writes a turn. Get these wrong, and no amount of DiskANN or 99.999% SLA saves you from a hot partition or a rewrite-the-whole-item cost curve later.
Partition Keys: The Foundation of a Cosmos DB Agent Memory Schema
Because Cosmos DB automatically partitions data, the partition key is the single most consequential choice in the schema. It decides how writes distribute, how queries scope, and eventually how much a bad decision costs to unwind. Three strategies cover most agent memory scenarios, and each trades distribution for locality differently.
Partition Keys
- GUID as the partition key. Every item lands in its own logical partition. Writes distribute as evenly as possible, so this works well for high-volume, write-heavy logging where you rarely need to reassemble a conversation later think raw telemetry more than chat history. The cost shows up on read: reconstructing a thread means a cross-partition query.
threadIdas the partition key. All turns in a conversation share one partition key, so “give me the last 10 turns” or “vector search within this thread” both stay inside a single partition. This is the default for conversational agents and RAG apps, provided threads are numerous and varied enough to avoid concentrating writes on a few hot values.[tenantId, threadId]as a hierarchical partition key. This is where the series comes full circle: I covered hierarchical partition keys in Azure Cosmos DB’s Latest Performance Features back in 2023, for the same tenant-then-item pattern. Nothing about the mechanism changed for agent workloads: threads still colocate under their tenant, tenant-level queries still avoid scanning every partition, and Cosmos DB still sub-partitions past the 20 GB logical-partition ceiling the same way it always did. Only the data moving through it is new.
There’s no universally correct answer here, so match the strategy to the query pattern that matters most: threadId if “give me this conversation” dominates, GUID if raw write throughput dominates, hierarchical if tenant isolation is a governance requirement and not just a nice-to-have.
Three Ways to Shape the Memory Item
Partition key decides where data lives. Item shape decides what it costs to read and write once it’s there, and that’s the other half of a Cosmos DB agent memory schema.
One document per turn (recommended default). Each item holds a complete exchange: a user prompt, the agent’s reply, and any tool call in between, and carries threadId, turnIndex, and an embedding alongside it. Most single- and multi-agent apps default to this shape, because it balances a small, cheap-to-write item against enough context to be useful on its own. “Latest N turns” is a simple ORDER BY turnIndex query, and ttl can expire old turns individually instead of touching the whole thread.
One document per response. Every user message, agent reply, and tool result gets its own item, all sharing a threadId. This is the most granular option, useful when you need to embed and search every single utterance independently, but it multiplies item count and RU cost on read. It loses the natural “question and answer together” unit that a semantic cache wants.
One document per thread. The whole conversation lives in one item that grows with every append. Reading the full history is a single read, which sounds appealing until a long-running thread turns every new turn into a full-item rewrite. Treat this as an anti-pattern unless the thread is short and you can bound its length by design: a five-turn onboarding flow, maybe; an open-ended assistant conversation, no.
TTL as Memory Lifecycle Management
Short-term memory should disappear on its own, and time-to-live is how Cosmos DB does that without a cleanup job. Set a default ttl on the container, and every item expires after that many seconds unless it overrides the value itself; set the container default to -1 instead, and Cosmos DB turns on ttl without expiring anything, unless an item sets its own positive ttl field. That second mode is the more useful one for agent memory, since it lets long-term memories (ttl: -1 on the item, meaning never expire) sit in the same container as short-term turns (ttl: 3600, gone in an hour) without a separate container or a background job doing the deleting.
A Worked Schema
Here’s the turn-based model from the previous section as an actual container and item.
Create the container with a hierarchical partition key, and turn on ttl at the container level:
az cosmosdb sql container create \ --account-name my-cosmos-account \ --database-name agentmemory \ --name turns \ --resource-group my-rg \ --partition-key-path "/tenantId" "/threadId" \ --ttl -1
A single turn, ready to insert:
{ "id": "b9c5b6ce-2d9a-4a2b-9d76-0f5f9b2a9a91", "tenantId": "contoso", "threadId": "thread-1234", "turnIndex": 7, "messages": [ { "role": "user", "content": "What's our refund policy for accessories?" }, { "role": "agent", "content": "Refund policy is 30 days for unopened items." } ], "embedding": [0.013, -0.092, 0.551], "ttl": 3600}
And writing it with the Python SDK:
from azure. cosmos import CosmosClientclient = CosmosClient(url, credential)container = client.get_database_client("agentmemory").get_container_client("turns")container.upsert_item({ "id": "b9c5b6ce-2d9a-4a2b-9d76-0f5f9b2a9a91", "tenantId": "contoso", "threadId": "thread-1234", "turnIndex": 7, "messages": [ {"role": "user", "content": "What's our refund policy for accessories?"}, {"role": "agent", "content": "Refund policy is 30 days for unopened items."}, ], "embedding": [0.013, -0.092, 0.551], "ttl": 3600,})
Swap ttl: 3600 for ttl: -1 on any item you want to keep past the container default as summarized long-term memory; for instance, it survives while the rest of the thread ages out on schedule.
Seeing It Live
The rest of this section shows the same schema running against a real Cosmos DB account, not just described on paper.
The provision.py creating the turns container with the [/tenantId, /threadId] hierarchical partition key and container-level ttl.
The same container in Data Explorer’s Scale & Settings pane partition key and ttl exactly as provision.py set them.
A seeded turn in Data Explorer’s Items view tenantId, threadId, turnIndex, messages, and embedding, matching the JSON above field for field.
Pitfalls
Hot partitions from low-cardinality tenant keys. A hierarchical [tenantId, threadId] key only distributes well if tenants themselves are numerous and reasonably balanced in volume. One enterprise customer generating 80% of total traffic under a single tenantId value creates a hot partition no amount of threadId variety underneath it fixes. Check tenant volume distribution before committing to this key, not after.
Unbounded thread-per-item growth. It’s tempting to reach for one-document-per-thread because “just read the whole conversation” feels simpler in application code. In practice, an item that grows by one append per turn racks up RU cost on every single write as the item gets larger, and Cosmos DB item size limits eventually cap how long a thread can run at all. If a thread’s length isn’t predictable and short, don’t model it this way.
Forgetting embeddings need to be top-level fields. If you nest an embedding array inside a messages object, it won’t be indexable for vector search — it has to sit at the top level of the item, alongside threadId and turnIndex, for the container’s vector policy to pick it up. This one is easy to miss because the item still writes successfully; it just never shows up in a vector query, and that failure is silent until someone notices recall is worse than expected.
Next: Finding the Right Memory
The schema in this post gets data into Cosmos DB efficiently. It doesn’t yet get the right memory back out at the moment an agent needs it — that’s the job of vector, full-text, and hybrid search, which is where post 3 picks up. That settles the Cosmos DB agent memory schema; retrieval is the next problem worth solving properly.
Sources
- Microsoft Learn — Agent memories in Azure Cosmos DB for NoSQL
- Microsoft Learn — Time to live (TTL) in Azure Cosmos DB.
- Microsoft Learn — Hierarchical partition keys
- Your own archive — Azure Cosmos DB’s Latest Performance Features
- GitHub — steefjan1/cosmos-agent-memory-lab — the runnable sample behind this schema, and behind posts 3 and 4
