crewai-persistence-cosmosdb

Azure CosmosDB persistence backend for CrewAI Flows. Persists flow state between steps so your flows can resume from any saved checkpoint.

What makes this different: it has message history pruning built in. Pass a MessageReducer and the persistence layer automatically caps your message list before writing to CosmosDB — no changes to your flow code or state model required. This is the only CrewAI CosmosDB persistence backend with this capability.

Features

• Full flow state persistence — save and load CrewAI flow state to/from CosmosDB at any step
• Built-in message pruning — optional MessageReducer prunes message history at the persistence layer, keeping state documents lean without touching your flow code
• Flexible authentication — key-based or Azure RBAC (Managed Identity, az login, service principal)
• Auto-creates database and container — when using key-based auth
• Pydantic model support — accepts both BaseModel instances and plain dicts as state data

Installation

pip install crewai-persistence-cosmosdb

With optional message pruning support:

pip install "crewai-persistence-cosmosdb[reducer]"

Requires Python 3.10+

Database and Container Setup

Auth mode	Database	Container	Partition key
Key-based (COSMOS_KEY set)	Created automatically if absent	Created automatically if absent	/flow_uuid (set by the backend)
RBAC / Managed Identity (no key)	Must pre-exist	Must pre-exist	/flow_uuid (must be pre-configured)

Key-based is the easiest way to get started — just point the backend at an existing CosmosDB account and it will provision everything.

RBAC is recommended for production. Because the backend only calls get_database_client / get_container_client (no write permissions needed at setup time), the database and container must already be provisioned before the backend is initialised. Create them via the Azure portal, Terraform, Bicep, or the Azure CLI:

az cosmosdb sql database create --account-name <account> --name <db>
az cosmosdb sql container create \
  --account-name <account> --database-name <db> --name <container> \
  --partition-key-path "/flow_uuid"

Important: The partition key path must be /flow_uuid regardless of how the container is created.

Authentication

Key-based (development / admin access)

export COSMOS_ENDPOINT="https://<account>.documents.azure.com:443/"
export COSMOS_KEY="<your-key>"

Azure RBAC / Managed Identity (production)

Set only the endpoint — no key. The backend uses DefaultAzureCredential, which resolves in this order: environment service principal → managed identity → az login.

export COSMOS_ENDPOINT="https://<account>.documents.azure.com:443/"
# COSMOS_KEY not set → DefaultAzureCredential is used

For user-assigned managed identity:

export AZURE_CLIENT_ID="<managed-identity-client-id>"

For service principal:

export AZURE_TENANT_ID="<tenant-id>"
export AZURE_CLIENT_ID="<client-id>"
export AZURE_CLIENT_SECRET="<client-secret>"

Quick Start

Normal flow (stateless steps)

import os
from crewai.flow.flow import Flow, start, listen
from crewai_persistence_cosmosdb import CosmosDBFlowPersistence

persistence = CosmosDBFlowPersistence(
    endpoint=os.environ["COSMOS_ENDPOINT"],
    database_name="mydb",
    container_name="flow_states",
    key=os.environ.get("COSMOS_KEY"),   # omit for RBAC
)

class MyFlow(Flow):
    @start()
    def first_step(self):
        return {"status": "started", "value": 42}

    @listen(first_step)
    def second_step(self, data):
        persistence.save_state(
            flow_uuid=self.state["id"],
            method_name="second_step",
            state_data=self.state,
        )
        return data

flow = MyFlow()
flow.kickoff()

Conversational flow (with message history)

import os
import uuid
from crewai.flow.flow import Flow, start, listen
from crewai_persistence_cosmosdb import CosmosDBFlowPersistence

persistence = CosmosDBFlowPersistence(
    endpoint=os.environ["COSMOS_ENDPOINT"],
    database_name="mydb",
    container_name="flow_states",
    key=os.environ.get("COSMOS_KEY"),
)

FLOW_UUID = str(uuid.uuid4())

class ChatFlow(Flow):
    @start()
    def handle_turn(self):
        # Restore prior conversation state if it exists
        prior = persistence.load_state(FLOW_UUID)
        messages = prior.get("messages", []) if prior else []

        # Add new user message
        messages.append({"role": "human", "content": "Tell me about Azure CosmosDB."})

        # ... call your LLM here ...
        messages.append({"role": "ai", "content": "CosmosDB is a globally distributed NoSQL database..."})

        state = {"messages": messages}
        persistence.save_state(
            flow_uuid=FLOW_UUID,
            method_name="handle_turn",
            state_data=state,
        )
        return state

flow = ChatFlow()
flow.kickoff()

API Reference

CosmosDBFlowPersistence

CosmosDBFlowPersistence(
    endpoint,
    database_name,
    container_name,
    key=None,
    reducer=None,
    messages_key="messages",
)

Parameter	Type	Default	Description
endpoint	str	required	CosmosDB account endpoint URL
database_name	str	required	CosmosDB database name
container_name	str	required	CosmosDB container name
key	str | None	None	Account key; omit to use DefaultAzureCredential (RBAC)
reducer	MessageReducer | None	None	Optional pruner — see Built-in Message Pruning
messages_key	str	"messages"	State key that holds the message list

Methods

Method	Description
init_db()	Initialise database/container references (called automatically by __init__)
save_state(flow_uuid, method_name, state_data)	Persist flow state (upsert by flow_uuid)
load_state(flow_uuid)	Load the most recently saved state; returns None if not found

Built-in Message Pruning

Long-running conversational flows accumulate message history with every turn. Left unchecked this inflates document size, increases CosmosDB storage costs, and eventually blows past LLM context limits.

This backend solves that at the persistence layer: pass a MessageReducer and it automatically prunes the message list inside save_state() before the document is written to CosmosDB. Your flow code and state model stay untouched.

Install with reducer support

pip install "crewai-persistence-cosmosdb[reducer]"

Usage

from agentstate_reducer import MessageReducer
from agentstate_reducer.models import ReducerConfig
from crewai_persistence_cosmosdb import CosmosDBFlowPersistence

config = ReducerConfig(min_messages=10, max_messages=20)
reducer = MessageReducer(config=config)

persistence = CosmosDBFlowPersistence(
    endpoint=os.environ["COSMOS_ENDPOINT"],
    database_name="mydb",
    container_name="flow_states",
    key=os.environ.get("COSMOS_KEY"),
    reducer=reducer,         # prune before each save
    messages_key="messages", # state key holding the message list (default)
)

When len(messages) > max_messages, the oldest human/ai messages are removed until min_messages remain. The following are never pruned:

• Index 0 (typically the system prompt) — controlled by preserve_first=True
• system and function messages
• tool messages — unless their parent ai message is pruned (cascade behaviour, configurable)

See agentstate-reducer on PyPI for full configuration options: preserve_first, cascade_tool_messages, summarize_fn, and role alias support (user/assistant/agent).

Data Model

Each call to save_state upserts a single CosmosDB document partitioned by flow_uuid. Only the latest state for each flow run is stored (upsert overwrites on id = flow_uuid).

Field	Description
id	Same as flow_uuid (CosmosDB document id)
flow_uuid	Unique identifier for the flow run (partition key)
_method_name	Name of the flow method that triggered the save
_saved_at	ISO-8601 UTC timestamp of the save
user fields	All fields from the original state dict / Pydantic model

CosmosDB system fields (_rid, _self, _etag, _attachments, _ts) are stripped before load_state returns the document.

License

MIT

Long-term memory via on_prune (agentstate-reducer >= 0.4.0)

Messages pruned from the flow state are exactly the ones leaving the model's view. The persistence layer forwards a memory namespace to the reducer, and any on_prune hook receives (pruned_messages, namespace) — so pruned turns can flow straight into CrewAI's unified Memory (or any store), with no package coupling:

from agentstate_reducer import MessageReducer, ReducerConfig, Background
from crewai.memory import Memory

memory = Memory(storage=...)                       # e.g. crewai-memory-dynamodb

def remember(pruned, namespace):
    text = "\n".join(m["content"] for m in pruned)
    memory.remember_many(memory.extract_memories(text), scope=namespace)   # LLM calls -> run off-path

reducer = MessageReducer(config=ReducerConfig(max_messages=20, on_prune=[Background(remember)]))

class SupportState(BaseModel):
    id: str = ""
    memory_namespace: str = "/user/kamal"          # long-term scope: the USER, not the flow
    messages: list = []

@persist(CosmosDBFlowPersistence(..., reducer=reducer))
class SupportFlow(Flow[SupportState]):
    ...

The persistence layer reads memory_namespace (or whatever ReducerConfig.namespace_key names) from the flow state on every save_state and passes it through untouched. If the state never sets it, the namespace falls back to "/flow/<flow_uuid>". Each pruned message reaches the hooks once.