Tag Archives: Azure Cosmos DB

Finding the Right Memory: Vector, Full-Text, and Hybrid Search in Cosmos DB

Posted on by
1

Post 3 of 6 on Cosmos DB agent memory search, because storing memory well doesn’t guarantee you’ll retrieve the right piece.

Post 2 ended with the schema settled and retrieval still open. This post closes that gap: the practical mechanics of Cosmos DB agent memory search, one container, four query patterns. Run the same question four different ways against that container, and it comes back with four different answers, because “find the right memory” isn’t one query pattern; it’s at least three, and knowing which one to reach for is most of the job.

Vector Indexing for Cosmos DB Agent Memory Search

Cosmos DB supports two vector index types, and the right one depends almost entirely on how many vectors you’re searching, not on anything specific to agents.

Diagram showing quantizedFlat versus DiskANN as the vector index choice by dataset size for Cosmos DB agent memory search
quantizedFlat suits smaller workloads; DiskANN takes over once vector counts reach the hundreds of thousands.

quantizedFlat compresses each vector and scans the compressed space exactly. It suits smaller workloads (tens of thousands of vectors) and trades a small amount of accuracy for lower RU cost and faster scans. For a single tenant’s short-term memory, this is often enough on its own.

DiskANN, on the other hand, indexes vectors for approximate nearest-neighbor search and scales to hundreds of thousands or billions of embeddings, with dynamic updates and strong recall even at that size. Post 1 already leaned on DiskANN as part of the case for Cosmos DB as a unified store; this is the mechanism behind that claim.

Sharding the Vector Index for Multitenant Isolation

DiskANN doesn’t have to search across every vector in the container. A vectorIndexShardKey partitions the index itself by a property you choose: session, user, or tenant, so a query only searches candidates within that shard instead of the whole container.

That maps directly onto the partition key work from post 2: set the vectorIndexShardKey to tenantId, or to the same [tenantId, threadId] pair you already use as the partition key, and semantic search for one tenant never touches another tenant’s vectors. A global, unsharded index still works and makes searching everything at once simpler, but it’sonly appropriate for a single-tenant app or a genuinely shared knowledge base where cross-tenant recall is the point rather than a leak.

Full-Text Search: When Precision Beats Semantics

Vector search finds what’s semantically similar. Sometimes semantically similar isn’t what you want — a customer asking about “the refund policy” needs the actual refund policy language, not five conceptually related passages about returns in general.

Full-text search on Cosmos DB handles that case through BM25, a statistical ranking function that scores by term frequency and document length. Cosmos DB applies linguistic processing automatically: tokenization, stemming, case normalization, so “running” still matches “run” or “ran.” It’s the right tool whenever exact terms or phrases carry meaning that a vector embedding would blur.

Hybrid Search: Combining Both with RRF

Most agent memory queries don’t need to choose between semantic and lexical relevance; they need a blend of both. That’s what Reciprocal Rank Fusion (RRF) does: it takes the vector-similarity ranking and the BM25 ranking for the same result set and merges them into one combined rank, instead of forcing a pick between the two.

Diagram showing vector rank and BM25 rank merging through Reciprocal Rank Fusion for Cosmos DB agent memory search
RRF blends vector similarity and keyword relevance into one ranked result, rather than picking one over the other.

In practice, this shows up as a single ORDER BY RANK RRF(...) clause, which the next section demonstrates directly.

Four Ways to Ask the Same Question

Take the turn-based schema from post 2 — tenantId, threadId, turnIndex, messages, embedding, content and run the same underlying question against it four ways. (content is a flat, denormalized copy of the turn’s text, added specifically because Cosmos DB doesn’t support wildcard array paths like /messages/*/content in a full-text policy or index the full-text and hybrid queries below point at c.content rather than c.messages for exactly that reason.)

Most recent, by recency:

SELECT TOP 5 c.messages, c.turnIndex
FROM c
WHERE c.tenantId = @tenantId AND c.threadId = @threadId
ORDER BY c.turnIndex DESC

Semantic, by vector similarity:

SELECT TOP 5 c.messages, VectorDistance(c.embedding, @queryVector) AS score
FROM c
WHERE c.tenantId = @tenantId AND c.threadId = @threadId
ORDER BY VectorDistance(c.embedding, @queryVector)

Hybrid, blending both with RRF:

SELECT TOP 5 c.messages, VectorDistance(c.embedding, @queryVector) AS score
FROM c
WHERE c.tenantId = @tenantId AND c.threadId = @threadId
ORDER BY VectorDistance(c.embedding, @queryVector)

Keyword, by exact phrase:

SELECT TOP 5 c.messages, c.turnIndex
FROM c
WHERE c.tenantId = @tenantId AND c.threadId = @threadId
  AND FULLTEXTCONTAINS(c.content, @phrase)
ORDER BY c.turnIndex DESC

Run all four against a thread where a customer asked about refunds three times, in different words, across twenty turns, and the differences stop being theoretical fast: recency surfaces whichever turn happened most recently, even if it’s off-topic; semantic search pulls in every conceptually related turn, including the ones that used different words entirely; hybrid balances the two; keyword search returns only the turns that used the customer’s actual phrase, and ranks them by recency underneath that filter.

Running These Queries in Data Explorer

The four queries above use parameterized SQL, the same form search.py, from the companion repo behind this series, sends through the Python SDK, which binds @tenantId, @queryVector, and @phrase properly before the query runs. Paste them as-is into the Azure Portal’s Data Explorer query pane instead, and two things break, neither of which is a schema or code bug:

Data Explorer’s query box doesn’t bind named parameters. A query that leaves @tenantId unresolved either matches nothing and returns “No results” silently, or for VectorDistance() inside ORDER BY and FullTextScore() fails to compile outright, because both functions require their arguments to resolve to literal values at query-compile time rather than at execution time.

Swap every @parameter for a literal value and all four run cleanly. Against the seeded sample data (tenantId = "contoso", threadId = "thread-1234", searching for "refund"):

Recency, with literals:

SELECT TOP 5 c.messages, c.turnIndex 
  FROM c WHERE c.tenantId = "contoso" AND c.threadId = "thread-1234" 
  ORDER BY c.turnIndex DESC

Semantic, with literals:

SELECT TOP 5 c.messages, VectorDistance(c.embedding, [0.8196, 0.6392, -0.2471, 0.1608, -0.8667, 0.4902, -0.2549, 0.2235]) AS score
FROM c
WHERE c.tenantId = "contoso" AND c.threadId = "thread-1234"
ORDER BY VectorDistance(c.embedding, [0.8196, 0.6392, -0.2471, 0.1608, -0.8667, 0.4902, -0.2549, 0.2235])

Hybrid, with literals:

SELECT TOP 5 c.messages, c.turnIndex
FROM c
WHERE c.tenantId = "contoso" AND c.threadId = "thread-1234"
ORDER BY RANK RRF(
    VectorDistance(c.embedding, [0.8196, 0.6392, -0.2471, 0.1608, -0.8667, 0.4902, -0.2549, 0.2235]),
    FullTextScore(c.content, "refund")
)

Keyword, with literals:

SELECT TOP 5 c.messages, c.turnIndex
FROM c
WHERE c.tenantId = "contoso" AND c.threadId = "thread-1234"
  AND FULLTEXTCONTAINS(c.content, "refund")
ORDER BY c.turnIndex DESC

Pitfalls

Reaching for DiskANN on a small dataset. DiskANN’s approximate search and sharding options solve a scale problem. Below roughly ten thousand vectors, quantizedFlat gets equivalent recall for less operational complexity and lower RU cost. Default to DiskANN because it sounds like the “serious” choice, and you’ve added index-shard decisions to a workload that never needed them.

A global vector index in a multitenant app. Skip the vectorIndexShardKey, and a semantic query searches every candidate in the entire container, tenant boundaries or not. Nothing stops the query from surfacing another tenant’s conceptually similar memory in the result set unless a WHERE clause happens to filter it back out after the fact, and relying on a filter to catch what the index itself should have scoped is the kind of gap that shows up in an audit, not in testing.

Forgetting WHERE filters still apply. Vector and hybrid queries look like they replace normal filtering, but ORDER BY VectorDistance(...) or ORDER BY RANK RRF(...) still runs inside a WHERE-scoped query, same as any other. Leave the WHERE c.tenantId = @tenantId AND c.threadId = @threadId clause off a semantic query, and it searches everything the container holds, not just the thread the agent is currently in.

Next: Coordinating Multiple Agents

That settles Cosmos DB agent memory search for a single agent working alone. Coordinating what several agents know about the same conversation is a different problem, and it’s where change feed, a mechanism post 1 already covered as a callback to the 2023 retail monitoring work, comes back to tie multi-agent state together. That’s post 4.

Sources

Designing a Cosmos DB Agent Memory Schema

Posted on by
1

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.

Diagram comparing GUID, threadId, and tenantId-plus-threadId partition key strategies for a Cosmos DB agent memory schema
GUID maximizes write distribution, threadId keeps a conversation together, and tenantId+threadId adds tenant isolation on top.

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.
  • threadId as 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.

Diagram comparing one-document-per-turn, per-response, and per-thread shapes for a Cosmos DB agent memory schema
One document per turn is the recommended default one document per thread risks unbounded item growth.

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 CosmosClient
client = 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 Foundry Citadel Platform Azure: Conversation Persistence with Cosmos DB

Posted on by
1

In the previous post, we connected a real tool-calling agent to the Microsoft Foundry Citadel Platform on Azure, routing every LLM call through the APIM governance hub in Sweden Central. The agent answered weather questions; the hub captured usage events in Cosmos DB; and Application Insights confirmed that both LLM calls were governed. The agent worked, but it had no memory. Every run started fresh, with no record of what was asked or answered.

This post adds conversation persistence to the Microsoft Foundry Citadel Platform on Azure. Every agent run now produces a structured document in the spoke’s Cosmos DB conversations container: the user’s question, the tool call made, the tool result, the agent’s answer, token counts, model version, and timestamp. The agent gains a memory layer, marking the transition as the spoke’s data tier becomes active.

What We Build

Each agent run writes one document to the spoke Cosmos DB:

{
  "id": "run-20260625-143022-stockholm",
  "principal_id": "steefjan@msn.com",
  "timestamp": "2026-06-25T14:30:22.441Z",
  "question": "What is the weather like in Stockholm right now?",
  "tool_calls": [
    {
      "name": "get_weather",
      "arguments": {"location": "Stockholm"},
      "result": {
        "location": "Stockholm, Sweden",
        "temperature_celsius": 22.5,
        "wind_speed_kmh": 7.2,
        "condition": "Overcast"
      }
    }
  ],
  "answer": "The weather in Stockholm is overcast with a temperature of 22.5°C...",
  "model": "gpt-4o-2024-11-20",
  "prompt_tokens": 234,
  "completion_tokens": 67,
  "total_tokens": 301,
  "apim_gateway": "apim-wpvlimv4ngkns.azure-api.net"
}

The partition key is /principal_id matching the container definition deployed by the spoke Bicep template. In addition, this arrangement ensures that all conversations for a given user are grouped into the same logical partition, making per-user history queries efficient.

The agent writes the document after completing the run, so a failed or incomplete run leaves no record.Moreover, it’s clean, simple, and auditable.

Prerequisites

From the previous two posts you should have:

  • Hub deployed in rg-ai-hub-gateway-dev
  • Spoke deployed in rg-ai-spoke-dev with Cosmos DB cosmos-tggi2gmkw22w4, database cosmos-dbtggi2gmkw22w4, container conversations
  • App Config appcs-tggi2gmkw22w4 populated with COSMOS_DB_ENDPOINT and CONVERSATIONS_DATABASE_CONTAINER
  • agent.py, config.py, and tools.py from the previous post
  • Virtual environment activated with openai, azure-appconfiguration, azure-identity, and requests installed

Step 1 — Install the Cosmos DB SDK

With your virtual environment activated:

pip install azure-cosmos

Pitfall: Cosmos DB Public Network Access

If your Cosmos DB has firewall rules enabled (which the spoke Bicep template sets by default), your local IP needs to be in the allowed list or public access needs to be set to All networks for dev. Check via the portal: cosmos-tggi2gmkw22w4 (your instance) → Networking → Public access → All networks → Save. In production this would be networkIsolation=true with private endpoints only.

Step 2 — Extend Config to Read Cosmos DB Settings

The spoke App Config already contains COSMOS_DB_ENDPOINT and CONVERSATIONS_DATABASE_CONTAINER — populated automatically during deployment. Extend config.py to pull these:

$lines = @(
    "from azure.appconfiguration import AzureAppConfigurationClient",
    "from azure.identity import DefaultAzureCredential",
    "",
    "APP_CONFIG_ENDPOINT = 'https://appcs-tggi2gmkw22w4.azconfig.io'",
    "LABEL = 'ai-lz'",
    "",
    "def get_config() -> dict:",
    "    credential = DefaultAzureCredential()",
    "    client = AzureAppConfigurationClient(",
    "        base_url=APP_CONFIG_ENDPOINT,",
    "        credential=credential",
    "    )",
    "    keys = [",
    "        'AI_FOUNDRY_PROJECT_ENDPOINT',",
    "        'CHAT_DEPLOYMENT_NAME',",
    "        'APIM_GATEWAY_URL',",
    "        'APIM_SUBSCRIPTION_KEY',",
    "        'COSMOS_DB_ENDPOINT',",
    "        'CONVERSATIONS_DATABASE_CONTAINER',",
    "        'DATABASE_NAME',",
    "    ]",
    "    config = {}",
    "    for key in keys:",
    "        setting = client.get_configuration_setting(key=key, label=LABEL)",
    "        config[key] = setting.value",
    "    return config",
    "",
    "if __name__ == '__main__':",
    "    cfg = get_config()",
    "    for k, v in cfg.items():",
    "        print(f'{k}: {v[:40]}...')"
)
[System.IO.File]::WriteAllLines("$PWD\config.py", $lines, [System.Text.UTF8Encoding]::new($false))

Test it:

python config.py

You should now see seven keys including COSMOS_DB_ENDPOINT pointing to https://cosmos-tggi2gmkw22w4.documents.azure.com:443/ and CONVERSATIONS_DATABASE_CONTAINER set to conversations.

Step 3 — Create cosmos.py

Create a dedicated Cosmos DB module:

$code = 'import os
cosmos = """from azure.cosmos import CosmosClient
from azure.identity import DefaultAzureCredential
import uuid
from datetime import datetime, timezone
def get_cosmos_client(endpoint):
    return CosmosClient(url=endpoint, credential=DefaultAzureCredential())
def save_conversation(endpoint, database_name, container_name, principal_id, question, tool_calls, answer, model, prompt_tokens, completion_tokens, total_tokens, apim_gateway):
    container = get_cosmos_client(endpoint).get_database_client(database_name).get_container_client(container_name)
    now = datetime.now(timezone.utc)
    fmt = \"%Y%m%d-%H%M%S\"
    doc_id = f\"run-{now.strftime(fmt)}-{str(uuid.uuid4())[:8]}\"
    document = {\"id\": doc_id, \"principal_id\": principal_id, \"timestamp\": now.isoformat(), \"question\": question, \"tool_calls\": tool_calls, \"answer\": answer, \"model\": model, \"prompt_tokens\": prompt_tokens, \"completion_tokens\": completion_tokens, \"total_tokens\": total_tokens, \"apim_gateway\": apim_gateway}
    container.create_item(body=document)
    print(f\"Saved conversation: {doc_id}\")
    return document
def get_conversation_history(endpoint, database_name, container_name, principal_id, limit=5):
    container = get_cosmos_client(endpoint).get_database_client(database_name).get_container_client(container_name)
    query = f\"SELECT TOP {limit} c.id, c.timestamp, c.question, c.answer, c.total_tokens FROM c WHERE c.principal_id = @principal_id ORDER BY c._ts DESC\"
    return list(container.query_items(query=query, parameters=[{\"name\": \"@principal_id\", \"value\": principal_id}], partition_key=principal_id))
"""
agent = """import json
from openai import AzureOpenAI
from config import get_config
from tools import get_weather, WEATHER_TOOL_DEFINITION
from cosmos import save_conversation, get_conversation_history
PRINCIPAL_ID = \"steefjan@msn.com\"
def run_agent_with_memory(user_question):
    cfg = get_config()
    apim_base = cfg[\"APIM_GATEWAY_URL\"].rstrip(\"/\").replace(\"/openai\", \"\")
    client = AzureOpenAI(azure_endpoint=apim_base, api_key=cfg[\"APIM_SUBSCRIPTION_KEY\"], api_version=\"2024-02-01\")
    messages = [{\"role\": \"user\", \"content\": user_question}]
    print(f\"Sending request via APIM: {apim_base}\")
    response = client.chat.completions.create(model=cfg[\"CHAT_DEPLOYMENT_NAME\"], messages=messages, tools=[WEATHER_TOOL_DEFINITION], tool_choice=\"auto\")
    msg = response.choices[0].message
    messages.append(msg)
    tool_calls_log = []
    answer = msg.content or \"\"
    total_prompt = response.usage.prompt_tokens
    total_completion = response.usage.completion_tokens
    if msg.tool_calls:
        for tool_call in msg.tool_calls:
            args = json.loads(tool_call.function.arguments)
            print(f\"  -> Tool call: get_weather({args})\")
            result_str = get_weather(**args)
            result_json = json.loads(result_str)
            print(f\"  -> Tool result: {result_str}\")
            tool_calls_log.append({\"name\": tool_call.function.name, \"arguments\": args, \"result\": result_json})
            messages.append({\"role\": \"tool\", \"tool_call_id\": tool_call.id, \"content\": result_str})
        response2 = client.chat.completions.create(model=cfg[\"CHAT_DEPLOYMENT_NAME\"], messages=messages)
        answer = response2.choices[0].message.content
        total_prompt += response2.usage.prompt_tokens
        total_completion += response2.usage.completion_tokens
    save_conversation(endpoint=cfg[\"COSMOS_DB_ENDPOINT\"], database_name=cfg[\"DATABASE_NAME\"], container_name=cfg[\"CONVERSATIONS_DATABASE_CONTAINER\"], principal_id=PRINCIPAL_ID, question=user_question, tool_calls=tool_calls_log, answer=answer, model=cfg[\"CHAT_DEPLOYMENT_NAME\"], prompt_tokens=total_prompt, completion_tokens=total_completion, total_tokens=total_prompt+total_completion, apim_gateway=apim_base.replace(\"https://\", \"\"))
    return answer
if __name__ == \"__main__\":
    cfg = get_config()
    print(\"=== Recent conversation history ===\")
    history = get_conversation_history(cfg[\"COSMOS_DB_ENDPOINT\"], cfg[\"DATABASE_NAME\"], cfg[\"CONVERSATIONS_DATABASE_CONTAINER\"], PRINCIPAL_ID, 3)
    if history:
        for h in history:
            print(f\"  [{h[\"timestamp\"]}] Q: {h[\"question\"][:60]}...\")
    else:
        print(\"  No previous conversations found.\")
    print()
    question = \"What is the weather like in Amsterdam right now?\"
    print(f\"Question: {question}\")
    answer = run_agent_with_memory(question)
    print(f\"Answer: {answer}\")
"""
with open("cosmos.py", "w", encoding="utf-8") as f:
    f.write(cosmos)
with open("agent_with_memory.py", "w", encoding="utf-8") as f:
    f.write(agent)
print("Done")
'
[System.IO.File]::WriteAllText("$PWD\write_files.py", $code, [System.Text.UTF8Encoding]::new($false))
python write_files.py

Pitfall: Managed Identity RBAC for Cosmos DB

The CosmosClient with DefaultAzureCredential uses your Azure CLI identity locally. That identity needs the Cosmos DB Built-in Data Contributor role on the Cosmos DB account — not a standard Azure RBAC role, but a Cosmos DB data plane role. The spoke deployment should have assigned this automatically via the assignCosmosDBCosmosDbBuiltInDataContributorExecutor deployment. If you get a 403, verify:

az cosmosdb sql role assignment list `
  --account-name cosmos-tggi2gmkw22w4 `
  --resource-group rg-ai-spoke-dev `
  --output table

Your principal ID (8e856fa1-f4c4-4a02-91a5-a6ccc6afc6b3) should appear with role definition ID ending in 00000000-0000-0000-0000-000000000002 (Built-in Data Contributor). If not, assign it:

az cosmosdb sql role assignment create `
  --account-name cosmos-tggi2gmkw22w4 `
  --resource-group rg-ai-spoke-dev `
  --role-definition-id /subscriptions/dc0f4d72-3734-4b03-8884-ccfb9c2c4cc7/resourceGroups/rg-ai-spoke-dev/providers/Microsoft.DocumentDB/databaseAccounts/cosmos-tggi2gmkw22w4/sqlRoleDefinitions/00000000-0000-0000-0000-000000000002 `
  --principal-id 8e856fa1-f4c4-4a02-91a5-a6ccc6afc6b3 `
  --scope /subscriptions/dc0f4d72-3734-4b03-8884-ccfb9c2c4cc7/resourceGroups/rg-ai-spoke-dev/providers/Microsoft.DocumentDB/databaseAccounts/cosmos-tggi2gmkw22w4

Pitfall: No Connection Strings

Never use Cosmos DB connection strings or account keys in the agent code. The pattern here uses DefaultAzureCredential throughout — locally it picks up your az login identity, in production it uses the spoke’s Managed Identity. This is the NEN 7510 and cVGZ security baseline compliant approach.

Step 4 — Create agent_with_memory.py

$lines = @(
    "import json",
    "from openai import AzureOpenAI",
    "from config import get_config",
    "from tools import get_weather, WEATHER_TOOL_DEFINITION",
    "from cosmos import save_conversation, get_conversation_history",
    "",
    "PRINCIPAL_ID = 'steefjan@msn.com'",
    "",
    "def run_agent_with_memory(user_question: str) -> str:",
    "    cfg = get_config()",
    "    apim_base = cfg['APIM_GATEWAY_URL'].rstrip('/').replace('/openai', '')",
    "",
    "    client = AzureOpenAI(",
    "        azure_endpoint=apim_base,",
    "        api_key=cfg['APIM_SUBSCRIPTION_KEY'],",
    "        api_version='2024-02-01',",
    "    )",
    "",
    "    messages = [{'role': 'user', 'content': user_question}]",
    "    print(f'Sending request via APIM: {apim_base}')",
    "",
    "    # First LLM call - tool decision",
    "    response = client.chat.completions.create(",
    "        model=cfg['CHAT_DEPLOYMENT_NAME'],",
    "        messages=messages,",
    "        tools=[WEATHER_TOOL_DEFINITION],",
    "        tool_choice='auto',",
    "    )",
    "",
    "    msg = response.choices[0].message",
    "    messages.append(msg)",
    "    first_usage = response.usage",
    "",
    "    tool_calls_log = []",
    "    answer = msg.content or ''",
    "    total_prompt_tokens = first_usage.prompt_tokens",
    "    total_completion_tokens = first_usage.completion_tokens",
    "",
    "    # Handle tool calls",
    "    if msg.tool_calls:",
    "        for tool_call in msg.tool_calls:",
    "            args = json.loads(tool_call.function.arguments)",
    "            print(f'  -> Tool call: get_weather({args})')",
    "            result_str = get_weather(**args)",
    "            result_json = json.loads(result_str)",
    "            print(f'  -> Tool result: {result_str}')",
    "",
    "            tool_calls_log.append({",
    "                'name': tool_call.function.name,",
    "                'arguments': args,",
    "                'result': result_json,",
    "            })",
    "",
    "            messages.append({",
    "                'role': 'tool',",
    "                'tool_call_id': tool_call.id,",
    "                'content': result_str,",
    "            })",
    "",
    "        # Second LLM call - synthesis",
    "        response2 = client.chat.completions.create(",
    "            model=cfg['CHAT_DEPLOYMENT_NAME'],",
    "            messages=messages,",
    "        )",
    "        answer = response2.choices[0].message.content",
    "        total_prompt_tokens += response2.usage.prompt_tokens",
    "        total_completion_tokens += response2.usage.completion_tokens",
    "",
    "    total_tokens = total_prompt_tokens + total_completion_tokens",
    "",
    "    # Save to Cosmos DB conversations container in the spoke",
    "    save_conversation(",
    "        endpoint=cfg['COSMOS_DB_ENDPOINT'],",
    "        database_name=cfg['DATABASE_NAME'],",
    "        container_name=cfg['CONVERSATIONS_DATABASE_CONTAINER'],",
    "        principal_id=PRINCIPAL_ID,",
    "        question=user_question,",
    "        tool_calls=tool_calls_log,",
    "        answer=answer,",
    "        model=cfg['CHAT_DEPLOYMENT_NAME'],",
    "        prompt_tokens=total_prompt_tokens,",
    "        completion_tokens=total_completion_tokens,",
    "        total_tokens=total_tokens,",
    "        apim_gateway=apim_base.replace('https://', ''),",
    "    )",
    "",
    "    return answer",
    "",
    "if __name__ == '__main__':",
    "    # Show last 3 conversations before running",
    "    from config import get_config",
    "    cfg = get_config()",
    "    print('=== Recent conversation history ===')",
    "    history = get_conversation_history(",
    "        endpoint=cfg['COSMOS_DB_ENDPOINT'],",
    "        database_name=cfg['DATABASE_NAME'],",
    "        container_name=cfg['CONVERSATIONS_DATABASE_CONTAINER'],",
    "        principal_id=PRINCIPAL_ID,",
    "        limit=3,",
    "    )",
    "    if history:",
    "        for h in history:",
    "            print(f'  [{h[chr(116)+chr(105)+chr(109)+chr(101)+chr(115)+chr(116)+chr(97)+chr(109)+chr(112)]}] Q: {h[chr(113)+chr(117)+chr(101)+chr(115)+chr(116)+chr(105)+chr(111)+chr(110))[:60]}...')",
    "    else:",
    "        print('  No previous conversations found.')",
    "    print()",
    "",
    "    question = 'What is the weather like in Amsterdam right now?'",
    "    print(f'Question: {question}')",
    "    answer = run_agent_with_memory(question)",
    "    print(f'Answer: {answer}')"
)
[System.IO.File]::WriteAllLines("$PWD\agent_with_memory.py", $lines, [System.Text.UTF8Encoding]::new($false))

Run it:

python agent_with_memory.py

A successful run looks like this:

Run it a second time and the history section will show the previous exchange:

Step 5 — Validate in Cosmos DB Data Explorer

Go to the Azure Portal → cosmos-tggi2gmkw22w4 → Data Explorer → cosmos-dbtggi2gmkw22w4 → conversations → Items.

You should see your conversation document with all fields populated. The partition key /principal_id should match steefjan@msn.com.

To query all conversations for a user:

SELECT * FROM c WHERE c.principal_id = 'steefjan@msn.com' ORDER BY c._ts DESC

To get a summary of all runs with token totals:

SELECT c.id, c.timestamp, c.question, c.total_tokens, c.model
FROM c
WHERE c.principal_id = 'steefjan@msn.com'
ORDER BY c._ts DESC

Step 6 — What’s in the Document

Looking at a stored conversation document, every field serves a purpose:

FieldPurpose
idUnique run identifier — traceable back to a specific agent invocation
principal_idPartition key — enables per-user history queries and RBAC scoping
timestampISO 8601 UTC — audit trail, correlatable with APIM logs
questionOriginal user input — searchable for pattern analysis
tool_callsFull tool call log including arguments and results — debugging and audit
answerFinal agent response — quality review and feedback loops
modelModel version — tracks which model version answered which questions
prompt_tokens / completion_tokensCumulative across both LLM calls — accurate per-conversation cost
total_tokensSum of both calls — FinOps input per user per conversation
apim_gatewayGateway used — identifies which hub instance served the request

The token counts here are cumulative across both LLM calls (tool decision and synthesis), yielding a true per-conversation costrather than a per-call figure. This is more useful for FinOps reporting you care about the cost of answering a question, not the cost of individual API calls within that answer.

Pitfalls Summary

PitfallFix
Cosmos DB firewall blocks local IPPortal → Networking → All networks for dev, or add specific IP
403 on Cosmos DB writeAssign Cosmos DB Built-in Data Contributor data plane role to your principal
CosmosResourceNotFoundErrorVerify database name (cosmos-dbtggi2gmkw22w4) and container name (conversations) match exactly
Partition key mismatchContainer was created with /principal_id — every document must include this field
DefaultAzureCredential fails locallyRun az login and ensure the correct subscription is selected
Never use connection stringsUse DefaultAzureCredential throughout — locally via az login, in production via Managed Identity

What the Full Citadel Data Layer Now Looks Like

After this post, the spoke’s data tier is fully active:

StoreWhat it holdsWho writes it
Hub Cosmos DB ai-usage-containerPer-LLM-call usage events (tokens, model, gateway, IP)APIM gateway automatically
Spoke Cosmos DB conversationsPer-run conversation documents (question, tools, answer, cumulative tokens)Agent code explicitly
App Config appcs-tggi2gmkw22w4All configuration keys for the spokeSpoke deployment automatically

The hub’s ai-usage-container captures the infrastructure view of every API call, governed and logged. The spoke’s conversations container captures the application view of every user interaction, structured and queryable. Together, they give you both compliance evidence and application telemetry from a single agent run.

What’s Next

The next post in this series showcases the Citadel Kill Switch and explains how it stops a governed agent when necessary. It details how the five-layer containment system in APIM effectively shuts down the process without affecting the spoke or agent code. The conversation history you’ve created illustrates the clear before-and-after contrast: requests flow to Cosmos DB and then abruptly halt at the gateway layer.