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-devwith Cosmos DBcosmos-tggi2gmkw22w4, databasecosmos-dbtggi2gmkw22w4, containerconversations - App Config
appcs-tggi2gmkw22w4populated withCOSMOS_DB_ENDPOINTandCONVERSATIONS_DATABASE_CONTAINER agent.py,config.py, andtools.pyfrom the previous post- Virtual environment activated with
openai,azure-appconfiguration,azure-identity, andrequestsinstalled
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 oscosmos = """from azure.cosmos import CosmosClientfrom azure.identity import DefaultAzureCredentialimport uuidfrom datetime import datetime, timezonedef 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 documentdef 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 jsonfrom openai import AzureOpenAIfrom config import get_configfrom tools import get_weather, WEATHER_TOOL_DEFINITIONfrom cosmos import save_conversation, get_conversation_historyPRINCIPAL_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 answerif __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.modelFROM cWHERE 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:
| Field | Purpose |
|---|---|
id | Unique run identifier — traceable back to a specific agent invocation |
principal_id | Partition key — enables per-user history queries and RBAC scoping |
timestamp | ISO 8601 UTC — audit trail, correlatable with APIM logs |
question | Original user input — searchable for pattern analysis |
tool_calls | Full tool call log including arguments and results — debugging and audit |
answer | Final agent response — quality review and feedback loops |
model | Model version — tracks which model version answered which questions |
prompt_tokens / completion_tokens | Cumulative across both LLM calls — accurate per-conversation cost |
total_tokens | Sum of both calls — FinOps input per user per conversation |
apim_gateway | Gateway 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
| Pitfall | Fix |
|---|---|
| Cosmos DB firewall blocks local IP | Portal → Networking → All networks for dev, or add specific IP |
| 403 on Cosmos DB write | Assign Cosmos DB Built-in Data Contributor data plane role to your principal |
CosmosResourceNotFoundError | Verify database name (cosmos-dbtggi2gmkw22w4) and container name (conversations) match exactly |
| Partition key mismatch | Container was created with /principal_id — every document must include this field |
DefaultAzureCredential fails locally | Run az login and ensure the correct subscription is selected |
| Never use connection strings | Use 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:
| Store | What it holds | Who writes it |
|---|---|---|
Hub Cosmos DB ai-usage-container | Per-LLM-call usage events (tokens, model, gateway, IP) | APIM gateway automatically |
Spoke Cosmos DB conversations | Per-run conversation documents (question, tools, answer, cumulative tokens) | Agent code explicitly |
App Config appcs-tggi2gmkw22w4 | All configuration keys for the spoke | Spoke 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.
