Metadata-Version: 2.4
Name: crewai_persistence_cosmosdb
Version: 0.1.0
Summary: Azure CosmosDB persistence backend for CrewAI Flows with built-in message pruning via agentstate-reducer
Author-email: Kamal <skamalj@github.com>
Keywords: agent-state,azure,cosmosdb,crewai,flows,message-reducer,persistence
Classifier: License :: OSI Approved :: MIT License
Classifier: Operating System :: OS Independent
Classifier: Programming Language :: Python :: 3
Classifier: Topic :: Scientific/Engineering :: Artificial Intelligence
Requires-Python: <3.14,>=3.10
Requires-Dist: azure-cosmos
Requires-Dist: azure-identity
Requires-Dist: crewai<2.0.0,>=1.0.0
Provides-Extra: reducer
Requires-Dist: agentstate-reducer>=0.1.1; extra == 'reducer'
Description-Content-Type: text/markdown

# crewai-persistence-cosmosdb

Azure CosmosDB persistence backend for [CrewAI Flows](https://docs.crewai.com/concepts/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

```bash
pip install crewai-persistence-cosmosdb
```

With optional message pruning support:

```bash
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:

```bash
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)

```bash
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`.

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

For **user-assigned managed identity**:

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

For **service principal**:

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

## Quick Start

### Normal flow (stateless steps)

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

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

```python
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](#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

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

### Usage

```python
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](https://pypi.org/project/agentstate-reducer/) 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
