* Add CosmosCheckpointStorage for Python workflow checkpointing Add native Cosmos DB NoSQL support for workflow checkpoint storage in the Python agent-framework-azure-cosmos package, achieving parity with the existing .NET CosmosCheckpointStore. New files: - _checkpoint_storage.py: CosmosCheckpointStorage implementing the CheckpointStorage protocol with 6 methods (save, load, list_checkpoints, delete, get_latest, list_checkpoint_ids) - test_cosmos_checkpoint_storage.py: Unit and integration tests - workflow_checkpointing.py: Sample demonstrating Cosmos DB-backed workflow checkpoint/resume Auth support: - Managed identity / RBAC via Azure credential objects (DefaultAzureCredential, ManagedIdentityCredential, etc.) - Key-based auth via account key string or AZURE_COSMOS_KEY env var - Pre-created CosmosClient or ContainerProxy Key design decisions: - Partition key: /workflow_name for efficient per-workflow queries - Serialization: Reuses encode/decode_checkpoint_value for full Python object fidelity (hybrid JSON + pickle approach) - Container auto-creation via create_container_if_not_exists Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> * Adding cosmos checkpointer * Resolving comments * Fixing builds * Adding sample for history provider and checkpoint storage * Resolving comments * fixing builds * Resolving comments --------- Co-authored-by: Aayush Kataria <aayushkataria@Aayushs-MacBook-Pro-2.local> Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Co-authored-by: Evan Mattson <35585003+moonbox3@users.noreply.github.com>
4.2 KiB
Get Started with Microsoft Agent Framework Azure Cosmos DB
Please install this package via pip:
pip install agent-framework-azure-cosmos --pre
Azure Cosmos DB History Provider
The Azure Cosmos DB integration provides CosmosHistoryProvider for persistent conversation history storage.
Basic Usage Example
from azure.identity.aio import DefaultAzureCredential
from agent_framework_azure_cosmos import CosmosHistoryProvider
provider = CosmosHistoryProvider(
endpoint="https://<account>.documents.azure.com:443/",
credential=DefaultAzureCredential(),
database_name="agent-framework",
container_name="chat-history",
)
Credentials follow the same pattern used by other Azure connectors in the repository:
- Pass a credential object (for example
DefaultAzureCredential) - Or pass a key string directly
- Or set
AZURE_COSMOS_KEYin the environment
Container naming behavior:
- Container name is configured on the provider (
container_nameorAZURE_COSMOS_CONTAINER_NAME) session_idis used as the Cosmos partition key for reads/writes
See samples/02-agents/conversations/cosmos_history_provider.py for a runnable example.
Cosmos DB Workflow Checkpoint Storage
CosmosCheckpointStorage implements the CheckpointStorage protocol, enabling
durable workflow checkpointing backed by Azure Cosmos DB NoSQL. Workflows can be
paused and resumed across process restarts by persisting checkpoint state in Cosmos DB.
Basic Usage
Managed Identity / RBAC (recommended for production)
from azure.identity.aio import DefaultAzureCredential
from agent_framework import WorkflowBuilder
from agent_framework_azure_cosmos import CosmosCheckpointStorage
checkpoint_storage = CosmosCheckpointStorage(
endpoint="https://<account>.documents.azure.com:443/",
credential=DefaultAzureCredential(),
database_name="agent-framework",
container_name="workflow-checkpoints",
)
Account Key
from agent_framework_azure_cosmos import CosmosCheckpointStorage
checkpoint_storage = CosmosCheckpointStorage(
endpoint="https://<account>.documents.azure.com:443/",
credential="<your-account-key>",
database_name="agent-framework",
container_name="workflow-checkpoints",
)
Then use with a workflow
from agent_framework import WorkflowBuilder
# Build a workflow with checkpointing enabled
workflow = WorkflowBuilder(
start_executor=start,
checkpoint_storage=checkpoint_storage,
).build()
# Run the workflow — checkpoints are automatically saved after each superstep
result = await workflow.run(message="input data")
# Resume from a checkpoint
latest = await checkpoint_storage.get_latest(workflow_name=workflow.name)
if latest:
resumed = await workflow.run(checkpoint_id=latest.checkpoint_id)
Authentication Options
CosmosCheckpointStorage supports the same authentication modes as CosmosHistoryProvider:
- Managed identity / RBAC (recommended): Pass
DefaultAzureCredential(),ManagedIdentityCredential(), or any AzureTokenCredential - Account key: Pass a key string via
credentialparameter - Environment variables: Set
AZURE_COSMOS_ENDPOINT,AZURE_COSMOS_DATABASE_NAME,AZURE_COSMOS_CONTAINER_NAME, andAZURE_COSMOS_KEY(key not required when using Azure credentials) - Pre-created client: Pass an existing
CosmosClientorContainerProxy
Database and Container Setup
The database and container are created automatically on first use (via
create_database_if_not_exists and create_container_if_not_exists). The container
uses /workflow_name as the partition key. You can also pre-create them in the Azure
portal with this partition key configuration.
Environment Variables
| Variable | Description |
|---|---|
AZURE_COSMOS_ENDPOINT |
Cosmos DB account endpoint |
AZURE_COSMOS_DATABASE_NAME |
Database name |
AZURE_COSMOS_CONTAINER_NAME |
Container name |
AZURE_COSMOS_KEY |
Account key (optional if using Azure credentials) |
See samples/03-workflows/checkpoint/cosmos_workflow_checkpointing.py for a standalone example,
or samples/03-workflows/checkpoint/cosmos_workflow_checkpointing_foundry.py for an end-to-end
example with Azure AI Foundry agents.