pydantic-ai-cosmosdb-persistence

Azure CosmosDB backend for pydantic-ai-persistence — a PydanticAI StepStore (durable events / snapshots / tool-effect ledger) and message-history store, on CosmosDB.

Install

pip install pydantic-ai-cosmosdb-persistence
# or:  pip install "pydantic-ai-persistence[cosmosdb]"

Usage

from pydantic_ai import Agent
from pydantic_ai_harness.step_persistence import StepPersistence
from pydantic_ai_cosmosdb_persistence import CosmosDBStepStore

store = CosmosDBStepStore(
    endpoint="https://<acct>.documents.azure.com:443/",
    key="<key>",
    database_name="pai",
    container_name="steps",
    max_snapshots_per_run=20,
)
agent = Agent("openai:gpt-4o", capabilities=[StepPersistence(store=store)])

CosmosDBHistoryStore(endpoint=…, key=…, database_name=…, container_name=…) provides save/load message history.

Data model

A container partitioned by /PK; each item is {id, PK, SK, data} (the id is a URL-encoded SK since Cosmos ids can't contain # / \ ?; the raw SK is a field for range queries). Database and container are auto-created (key-based auth).

Note: PydanticAI's StepStore is a beta/experimental harness feature.

License

MIT