Skip to content
Open
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension


Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
91 changes: 90 additions & 1 deletion python/packages/azure-cosmos/README.md
Original file line number Diff line number Diff line change
Expand Up @@ -35,4 +35,93 @@ Container naming behavior:
- Container name is configured on the provider (`container_name` or `AZURE_COSMOS_CONTAINER_NAME`)
- `session_id` is used as the Cosmos partition key for reads/writes

See `samples/cosmos_history_provider.py` for a runnable package-local example.
See `samples/history_provider/cosmos_history_basic.py` for a runnable package-local 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)

```python
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

```python
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

```python
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 Azure `TokenCredential`
- **Account key**: Pass a key string via `credential` parameter
- **Environment variables**: Set `AZURE_COSMOS_ENDPOINT`, `AZURE_COSMOS_DATABASE_NAME`,
`AZURE_COSMOS_CONTAINER_NAME`, and `AZURE_COSMOS_KEY` (key not required when using
Azure credentials)
- **Pre-created client**: Pass an existing `CosmosClient` or `ContainerProxy`

### 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/checkpoint_storage/cosmos_checkpoint_workflow.py` for a standalone example,
`samples/checkpoint_storage/cosmos_checkpoint_foundry.py` for an end-to-end example
with Azure AI Foundry agents, or `samples/cosmos_e2e_foundry.py` for both
history and checkpointing together.
Original file line number Diff line number Diff line change
Expand Up @@ -2,6 +2,7 @@

import importlib.metadata

from ._checkpoint_storage import CosmosCheckpointStorage
from ._history_provider import CosmosHistoryProvider

try:
Expand All @@ -10,6 +11,7 @@
__version__ = "0.0.0" # Fallback for development mode

__all__ = [
"CosmosCheckpointStorage",
"CosmosHistoryProvider",
"__version__",
]
Loading
Loading