Microsoft Foundry Citadel Platform Azure: Conversation Persistence with Cosmos DB

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) → NetworkingPublic accessAll 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-tggi2gmkw22w4Data Explorercosmos-dbtggi2gmkw22w4conversationsItems.

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.