crewai_persistence_cosmosdb 0.1.0__tar.gz
This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
- crewai_persistence_cosmosdb-0.1.0/.gitignore +6 -0
- crewai_persistence_cosmosdb-0.1.0/PKG-INFO +263 -0
- crewai_persistence_cosmosdb-0.1.0/README.md +245 -0
- crewai_persistence_cosmosdb-0.1.0/pyproject.toml +37 -0
- crewai_persistence_cosmosdb-0.1.0/src/crewai_persistence_cosmosdb/__init__.py +3 -0
- crewai_persistence_cosmosdb-0.1.0/src/crewai_persistence_cosmosdb/persistence.py +121 -0
- crewai_persistence_cosmosdb-0.1.0/tests/test_persistence.py +212 -0
- crewai_persistence_cosmosdb-0.1.0/uv.lock +4396 -0
|
@@ -0,0 +1,263 @@
|
|
|
1
|
+
Metadata-Version: 2.4
|
|
2
|
+
Name: crewai_persistence_cosmosdb
|
|
3
|
+
Version: 0.1.0
|
|
4
|
+
Summary: Azure CosmosDB persistence backend for CrewAI Flows with built-in message pruning via agentstate-reducer
|
|
5
|
+
Author-email: Kamal <skamalj@github.com>
|
|
6
|
+
Keywords: agent-state,azure,cosmosdb,crewai,flows,message-reducer,persistence
|
|
7
|
+
Classifier: License :: OSI Approved :: MIT License
|
|
8
|
+
Classifier: Operating System :: OS Independent
|
|
9
|
+
Classifier: Programming Language :: Python :: 3
|
|
10
|
+
Classifier: Topic :: Scientific/Engineering :: Artificial Intelligence
|
|
11
|
+
Requires-Python: <3.14,>=3.10
|
|
12
|
+
Requires-Dist: azure-cosmos
|
|
13
|
+
Requires-Dist: azure-identity
|
|
14
|
+
Requires-Dist: crewai<2.0.0,>=1.0.0
|
|
15
|
+
Provides-Extra: reducer
|
|
16
|
+
Requires-Dist: agentstate-reducer>=0.1.1; extra == 'reducer'
|
|
17
|
+
Description-Content-Type: text/markdown
|
|
18
|
+
|
|
19
|
+
# crewai-persistence-cosmosdb
|
|
20
|
+
|
|
21
|
+
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.
|
|
22
|
+
|
|
23
|
+
**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.
|
|
24
|
+
|
|
25
|
+
## Features
|
|
26
|
+
|
|
27
|
+
- **Full flow state persistence** — save and load CrewAI flow state to/from CosmosDB at any step
|
|
28
|
+
- **Built-in message pruning** — optional `MessageReducer` prunes message history at the persistence layer, keeping state documents lean without touching your flow code
|
|
29
|
+
- **Flexible authentication** — key-based or Azure RBAC (Managed Identity, `az login`, service principal)
|
|
30
|
+
- **Auto-creates database and container** — when using key-based auth
|
|
31
|
+
- **Pydantic model support** — accepts both `BaseModel` instances and plain dicts as state data
|
|
32
|
+
|
|
33
|
+
## Installation
|
|
34
|
+
|
|
35
|
+
```bash
|
|
36
|
+
pip install crewai-persistence-cosmosdb
|
|
37
|
+
```
|
|
38
|
+
|
|
39
|
+
With optional message pruning support:
|
|
40
|
+
|
|
41
|
+
```bash
|
|
42
|
+
pip install "crewai-persistence-cosmosdb[reducer]"
|
|
43
|
+
```
|
|
44
|
+
|
|
45
|
+
**Requires Python 3.10+**
|
|
46
|
+
|
|
47
|
+
## Database and Container Setup
|
|
48
|
+
|
|
49
|
+
| Auth mode | Database | Container | Partition key |
|
|
50
|
+
|---|---|---|---|
|
|
51
|
+
| **Key-based** (`COSMOS_KEY` set) | Created automatically if absent | Created automatically if absent | `/flow_uuid` (set by the backend) |
|
|
52
|
+
| **RBAC / Managed Identity** (no key) | **Must pre-exist** | **Must pre-exist** | `/flow_uuid` (must be pre-configured) |
|
|
53
|
+
|
|
54
|
+
**Key-based** is the easiest way to get started — just point the backend at an existing CosmosDB account and it will provision everything.
|
|
55
|
+
|
|
56
|
+
**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:
|
|
57
|
+
|
|
58
|
+
```bash
|
|
59
|
+
az cosmosdb sql database create --account-name <account> --name <db>
|
|
60
|
+
az cosmosdb sql container create \
|
|
61
|
+
--account-name <account> --database-name <db> --name <container> \
|
|
62
|
+
--partition-key-path "/flow_uuid"
|
|
63
|
+
```
|
|
64
|
+
|
|
65
|
+
> **Important:** The partition key path must be `/flow_uuid` regardless of how the container is created.
|
|
66
|
+
|
|
67
|
+
## Authentication
|
|
68
|
+
|
|
69
|
+
### Key-based (development / admin access)
|
|
70
|
+
|
|
71
|
+
```bash
|
|
72
|
+
export COSMOS_ENDPOINT="https://<account>.documents.azure.com:443/"
|
|
73
|
+
export COSMOS_KEY="<your-key>"
|
|
74
|
+
```
|
|
75
|
+
|
|
76
|
+
### Azure RBAC / Managed Identity (production)
|
|
77
|
+
|
|
78
|
+
Set only the endpoint — no key. The backend uses `DefaultAzureCredential`, which resolves in this order: environment service principal → managed identity → `az login`.
|
|
79
|
+
|
|
80
|
+
```bash
|
|
81
|
+
export COSMOS_ENDPOINT="https://<account>.documents.azure.com:443/"
|
|
82
|
+
# COSMOS_KEY not set → DefaultAzureCredential is used
|
|
83
|
+
```
|
|
84
|
+
|
|
85
|
+
For **user-assigned managed identity**:
|
|
86
|
+
|
|
87
|
+
```bash
|
|
88
|
+
export AZURE_CLIENT_ID="<managed-identity-client-id>"
|
|
89
|
+
```
|
|
90
|
+
|
|
91
|
+
For **service principal**:
|
|
92
|
+
|
|
93
|
+
```bash
|
|
94
|
+
export AZURE_TENANT_ID="<tenant-id>"
|
|
95
|
+
export AZURE_CLIENT_ID="<client-id>"
|
|
96
|
+
export AZURE_CLIENT_SECRET="<client-secret>"
|
|
97
|
+
```
|
|
98
|
+
|
|
99
|
+
## Quick Start
|
|
100
|
+
|
|
101
|
+
### Normal flow (stateless steps)
|
|
102
|
+
|
|
103
|
+
```python
|
|
104
|
+
import os
|
|
105
|
+
from crewai.flow.flow import Flow, start, listen
|
|
106
|
+
from crewai_persistence_cosmosdb import CosmosDBFlowPersistence
|
|
107
|
+
|
|
108
|
+
persistence = CosmosDBFlowPersistence(
|
|
109
|
+
endpoint=os.environ["COSMOS_ENDPOINT"],
|
|
110
|
+
database_name="mydb",
|
|
111
|
+
container_name="flow_states",
|
|
112
|
+
key=os.environ.get("COSMOS_KEY"), # omit for RBAC
|
|
113
|
+
)
|
|
114
|
+
|
|
115
|
+
class MyFlow(Flow):
|
|
116
|
+
@start()
|
|
117
|
+
def first_step(self):
|
|
118
|
+
return {"status": "started", "value": 42}
|
|
119
|
+
|
|
120
|
+
@listen(first_step)
|
|
121
|
+
def second_step(self, data):
|
|
122
|
+
persistence.save_state(
|
|
123
|
+
flow_uuid=self.state["id"],
|
|
124
|
+
method_name="second_step",
|
|
125
|
+
state_data=self.state,
|
|
126
|
+
)
|
|
127
|
+
return data
|
|
128
|
+
|
|
129
|
+
flow = MyFlow()
|
|
130
|
+
flow.kickoff()
|
|
131
|
+
```
|
|
132
|
+
|
|
133
|
+
### Conversational flow (with message history)
|
|
134
|
+
|
|
135
|
+
```python
|
|
136
|
+
import os
|
|
137
|
+
import uuid
|
|
138
|
+
from crewai.flow.flow import Flow, start, listen
|
|
139
|
+
from crewai_persistence_cosmosdb import CosmosDBFlowPersistence
|
|
140
|
+
|
|
141
|
+
persistence = CosmosDBFlowPersistence(
|
|
142
|
+
endpoint=os.environ["COSMOS_ENDPOINT"],
|
|
143
|
+
database_name="mydb",
|
|
144
|
+
container_name="flow_states",
|
|
145
|
+
key=os.environ.get("COSMOS_KEY"),
|
|
146
|
+
)
|
|
147
|
+
|
|
148
|
+
FLOW_UUID = str(uuid.uuid4())
|
|
149
|
+
|
|
150
|
+
class ChatFlow(Flow):
|
|
151
|
+
@start()
|
|
152
|
+
def handle_turn(self):
|
|
153
|
+
# Restore prior conversation state if it exists
|
|
154
|
+
prior = persistence.load_state(FLOW_UUID)
|
|
155
|
+
messages = prior.get("messages", []) if prior else []
|
|
156
|
+
|
|
157
|
+
# Add new user message
|
|
158
|
+
messages.append({"role": "human", "content": "Tell me about Azure CosmosDB."})
|
|
159
|
+
|
|
160
|
+
# ... call your LLM here ...
|
|
161
|
+
messages.append({"role": "ai", "content": "CosmosDB is a globally distributed NoSQL database..."})
|
|
162
|
+
|
|
163
|
+
state = {"messages": messages}
|
|
164
|
+
persistence.save_state(
|
|
165
|
+
flow_uuid=FLOW_UUID,
|
|
166
|
+
method_name="handle_turn",
|
|
167
|
+
state_data=state,
|
|
168
|
+
)
|
|
169
|
+
return state
|
|
170
|
+
|
|
171
|
+
flow = ChatFlow()
|
|
172
|
+
flow.kickoff()
|
|
173
|
+
```
|
|
174
|
+
|
|
175
|
+
## API Reference
|
|
176
|
+
|
|
177
|
+
### `CosmosDBFlowPersistence`
|
|
178
|
+
|
|
179
|
+
```python
|
|
180
|
+
CosmosDBFlowPersistence(
|
|
181
|
+
endpoint,
|
|
182
|
+
database_name,
|
|
183
|
+
container_name,
|
|
184
|
+
key=None,
|
|
185
|
+
reducer=None,
|
|
186
|
+
messages_key="messages",
|
|
187
|
+
)
|
|
188
|
+
```
|
|
189
|
+
|
|
190
|
+
| Parameter | Type | Default | Description |
|
|
191
|
+
|---|---|---|---|
|
|
192
|
+
| `endpoint` | `str` | required | CosmosDB account endpoint URL |
|
|
193
|
+
| `database_name` | `str` | required | CosmosDB database name |
|
|
194
|
+
| `container_name` | `str` | required | CosmosDB container name |
|
|
195
|
+
| `key` | `str \| None` | `None` | Account key; omit to use `DefaultAzureCredential` (RBAC) |
|
|
196
|
+
| `reducer` | `MessageReducer \| None` | `None` | Optional pruner — see [Built-in Message Pruning](#built-in-message-pruning) |
|
|
197
|
+
| `messages_key` | `str` | `"messages"` | State key that holds the message list |
|
|
198
|
+
|
|
199
|
+
### Methods
|
|
200
|
+
|
|
201
|
+
| Method | Description |
|
|
202
|
+
|---|---|
|
|
203
|
+
| `init_db()` | Initialise database/container references (called automatically by `__init__`) |
|
|
204
|
+
| `save_state(flow_uuid, method_name, state_data)` | Persist flow state (upsert by `flow_uuid`) |
|
|
205
|
+
| `load_state(flow_uuid)` | Load the most recently saved state; returns `None` if not found |
|
|
206
|
+
|
|
207
|
+
## Built-in Message Pruning
|
|
208
|
+
|
|
209
|
+
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.
|
|
210
|
+
|
|
211
|
+
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.**
|
|
212
|
+
|
|
213
|
+
### Install with reducer support
|
|
214
|
+
|
|
215
|
+
```bash
|
|
216
|
+
pip install "crewai-persistence-cosmosdb[reducer]"
|
|
217
|
+
```
|
|
218
|
+
|
|
219
|
+
### Usage
|
|
220
|
+
|
|
221
|
+
```python
|
|
222
|
+
from agentstate_reducer import MessageReducer
|
|
223
|
+
from agentstate_reducer.models import ReducerConfig
|
|
224
|
+
from crewai_persistence_cosmosdb import CosmosDBFlowPersistence
|
|
225
|
+
|
|
226
|
+
config = ReducerConfig(min_messages=10, max_messages=20)
|
|
227
|
+
reducer = MessageReducer(config=config)
|
|
228
|
+
|
|
229
|
+
persistence = CosmosDBFlowPersistence(
|
|
230
|
+
endpoint=os.environ["COSMOS_ENDPOINT"],
|
|
231
|
+
database_name="mydb",
|
|
232
|
+
container_name="flow_states",
|
|
233
|
+
key=os.environ.get("COSMOS_KEY"),
|
|
234
|
+
reducer=reducer, # prune before each save
|
|
235
|
+
messages_key="messages", # state key holding the message list (default)
|
|
236
|
+
)
|
|
237
|
+
```
|
|
238
|
+
|
|
239
|
+
When `len(messages) > max_messages`, the oldest `human`/`ai` messages are removed until `min_messages` remain. The following are **never** pruned:
|
|
240
|
+
|
|
241
|
+
- Index 0 (typically the system prompt) — controlled by `preserve_first=True`
|
|
242
|
+
- `system` and `function` messages
|
|
243
|
+
- `tool` messages — unless their parent `ai` message is pruned (cascade behaviour, configurable)
|
|
244
|
+
|
|
245
|
+
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`).
|
|
246
|
+
|
|
247
|
+
## Data Model
|
|
248
|
+
|
|
249
|
+
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`).
|
|
250
|
+
|
|
251
|
+
| Field | Description |
|
|
252
|
+
|---|---|
|
|
253
|
+
| `id` | Same as `flow_uuid` (CosmosDB document id) |
|
|
254
|
+
| `flow_uuid` | Unique identifier for the flow run (partition key) |
|
|
255
|
+
| `_method_name` | Name of the flow method that triggered the save |
|
|
256
|
+
| `_saved_at` | ISO-8601 UTC timestamp of the save |
|
|
257
|
+
| _user fields_ | All fields from the original state dict / Pydantic model |
|
|
258
|
+
|
|
259
|
+
CosmosDB system fields (`_rid`, `_self`, `_etag`, `_attachments`, `_ts`) are stripped before `load_state` returns the document.
|
|
260
|
+
|
|
261
|
+
## License
|
|
262
|
+
|
|
263
|
+
MIT
|
|
@@ -0,0 +1,245 @@
|
|
|
1
|
+
# crewai-persistence-cosmosdb
|
|
2
|
+
|
|
3
|
+
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.
|
|
4
|
+
|
|
5
|
+
**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.
|
|
6
|
+
|
|
7
|
+
## Features
|
|
8
|
+
|
|
9
|
+
- **Full flow state persistence** — save and load CrewAI flow state to/from CosmosDB at any step
|
|
10
|
+
- **Built-in message pruning** — optional `MessageReducer` prunes message history at the persistence layer, keeping state documents lean without touching your flow code
|
|
11
|
+
- **Flexible authentication** — key-based or Azure RBAC (Managed Identity, `az login`, service principal)
|
|
12
|
+
- **Auto-creates database and container** — when using key-based auth
|
|
13
|
+
- **Pydantic model support** — accepts both `BaseModel` instances and plain dicts as state data
|
|
14
|
+
|
|
15
|
+
## Installation
|
|
16
|
+
|
|
17
|
+
```bash
|
|
18
|
+
pip install crewai-persistence-cosmosdb
|
|
19
|
+
```
|
|
20
|
+
|
|
21
|
+
With optional message pruning support:
|
|
22
|
+
|
|
23
|
+
```bash
|
|
24
|
+
pip install "crewai-persistence-cosmosdb[reducer]"
|
|
25
|
+
```
|
|
26
|
+
|
|
27
|
+
**Requires Python 3.10+**
|
|
28
|
+
|
|
29
|
+
## Database and Container Setup
|
|
30
|
+
|
|
31
|
+
| Auth mode | Database | Container | Partition key |
|
|
32
|
+
|---|---|---|---|
|
|
33
|
+
| **Key-based** (`COSMOS_KEY` set) | Created automatically if absent | Created automatically if absent | `/flow_uuid` (set by the backend) |
|
|
34
|
+
| **RBAC / Managed Identity** (no key) | **Must pre-exist** | **Must pre-exist** | `/flow_uuid` (must be pre-configured) |
|
|
35
|
+
|
|
36
|
+
**Key-based** is the easiest way to get started — just point the backend at an existing CosmosDB account and it will provision everything.
|
|
37
|
+
|
|
38
|
+
**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:
|
|
39
|
+
|
|
40
|
+
```bash
|
|
41
|
+
az cosmosdb sql database create --account-name <account> --name <db>
|
|
42
|
+
az cosmosdb sql container create \
|
|
43
|
+
--account-name <account> --database-name <db> --name <container> \
|
|
44
|
+
--partition-key-path "/flow_uuid"
|
|
45
|
+
```
|
|
46
|
+
|
|
47
|
+
> **Important:** The partition key path must be `/flow_uuid` regardless of how the container is created.
|
|
48
|
+
|
|
49
|
+
## Authentication
|
|
50
|
+
|
|
51
|
+
### Key-based (development / admin access)
|
|
52
|
+
|
|
53
|
+
```bash
|
|
54
|
+
export COSMOS_ENDPOINT="https://<account>.documents.azure.com:443/"
|
|
55
|
+
export COSMOS_KEY="<your-key>"
|
|
56
|
+
```
|
|
57
|
+
|
|
58
|
+
### Azure RBAC / Managed Identity (production)
|
|
59
|
+
|
|
60
|
+
Set only the endpoint — no key. The backend uses `DefaultAzureCredential`, which resolves in this order: environment service principal → managed identity → `az login`.
|
|
61
|
+
|
|
62
|
+
```bash
|
|
63
|
+
export COSMOS_ENDPOINT="https://<account>.documents.azure.com:443/"
|
|
64
|
+
# COSMOS_KEY not set → DefaultAzureCredential is used
|
|
65
|
+
```
|
|
66
|
+
|
|
67
|
+
For **user-assigned managed identity**:
|
|
68
|
+
|
|
69
|
+
```bash
|
|
70
|
+
export AZURE_CLIENT_ID="<managed-identity-client-id>"
|
|
71
|
+
```
|
|
72
|
+
|
|
73
|
+
For **service principal**:
|
|
74
|
+
|
|
75
|
+
```bash
|
|
76
|
+
export AZURE_TENANT_ID="<tenant-id>"
|
|
77
|
+
export AZURE_CLIENT_ID="<client-id>"
|
|
78
|
+
export AZURE_CLIENT_SECRET="<client-secret>"
|
|
79
|
+
```
|
|
80
|
+
|
|
81
|
+
## Quick Start
|
|
82
|
+
|
|
83
|
+
### Normal flow (stateless steps)
|
|
84
|
+
|
|
85
|
+
```python
|
|
86
|
+
import os
|
|
87
|
+
from crewai.flow.flow import Flow, start, listen
|
|
88
|
+
from crewai_persistence_cosmosdb import CosmosDBFlowPersistence
|
|
89
|
+
|
|
90
|
+
persistence = CosmosDBFlowPersistence(
|
|
91
|
+
endpoint=os.environ["COSMOS_ENDPOINT"],
|
|
92
|
+
database_name="mydb",
|
|
93
|
+
container_name="flow_states",
|
|
94
|
+
key=os.environ.get("COSMOS_KEY"), # omit for RBAC
|
|
95
|
+
)
|
|
96
|
+
|
|
97
|
+
class MyFlow(Flow):
|
|
98
|
+
@start()
|
|
99
|
+
def first_step(self):
|
|
100
|
+
return {"status": "started", "value": 42}
|
|
101
|
+
|
|
102
|
+
@listen(first_step)
|
|
103
|
+
def second_step(self, data):
|
|
104
|
+
persistence.save_state(
|
|
105
|
+
flow_uuid=self.state["id"],
|
|
106
|
+
method_name="second_step",
|
|
107
|
+
state_data=self.state,
|
|
108
|
+
)
|
|
109
|
+
return data
|
|
110
|
+
|
|
111
|
+
flow = MyFlow()
|
|
112
|
+
flow.kickoff()
|
|
113
|
+
```
|
|
114
|
+
|
|
115
|
+
### Conversational flow (with message history)
|
|
116
|
+
|
|
117
|
+
```python
|
|
118
|
+
import os
|
|
119
|
+
import uuid
|
|
120
|
+
from crewai.flow.flow import Flow, start, listen
|
|
121
|
+
from crewai_persistence_cosmosdb import CosmosDBFlowPersistence
|
|
122
|
+
|
|
123
|
+
persistence = CosmosDBFlowPersistence(
|
|
124
|
+
endpoint=os.environ["COSMOS_ENDPOINT"],
|
|
125
|
+
database_name="mydb",
|
|
126
|
+
container_name="flow_states",
|
|
127
|
+
key=os.environ.get("COSMOS_KEY"),
|
|
128
|
+
)
|
|
129
|
+
|
|
130
|
+
FLOW_UUID = str(uuid.uuid4())
|
|
131
|
+
|
|
132
|
+
class ChatFlow(Flow):
|
|
133
|
+
@start()
|
|
134
|
+
def handle_turn(self):
|
|
135
|
+
# Restore prior conversation state if it exists
|
|
136
|
+
prior = persistence.load_state(FLOW_UUID)
|
|
137
|
+
messages = prior.get("messages", []) if prior else []
|
|
138
|
+
|
|
139
|
+
# Add new user message
|
|
140
|
+
messages.append({"role": "human", "content": "Tell me about Azure CosmosDB."})
|
|
141
|
+
|
|
142
|
+
# ... call your LLM here ...
|
|
143
|
+
messages.append({"role": "ai", "content": "CosmosDB is a globally distributed NoSQL database..."})
|
|
144
|
+
|
|
145
|
+
state = {"messages": messages}
|
|
146
|
+
persistence.save_state(
|
|
147
|
+
flow_uuid=FLOW_UUID,
|
|
148
|
+
method_name="handle_turn",
|
|
149
|
+
state_data=state,
|
|
150
|
+
)
|
|
151
|
+
return state
|
|
152
|
+
|
|
153
|
+
flow = ChatFlow()
|
|
154
|
+
flow.kickoff()
|
|
155
|
+
```
|
|
156
|
+
|
|
157
|
+
## API Reference
|
|
158
|
+
|
|
159
|
+
### `CosmosDBFlowPersistence`
|
|
160
|
+
|
|
161
|
+
```python
|
|
162
|
+
CosmosDBFlowPersistence(
|
|
163
|
+
endpoint,
|
|
164
|
+
database_name,
|
|
165
|
+
container_name,
|
|
166
|
+
key=None,
|
|
167
|
+
reducer=None,
|
|
168
|
+
messages_key="messages",
|
|
169
|
+
)
|
|
170
|
+
```
|
|
171
|
+
|
|
172
|
+
| Parameter | Type | Default | Description |
|
|
173
|
+
|---|---|---|---|
|
|
174
|
+
| `endpoint` | `str` | required | CosmosDB account endpoint URL |
|
|
175
|
+
| `database_name` | `str` | required | CosmosDB database name |
|
|
176
|
+
| `container_name` | `str` | required | CosmosDB container name |
|
|
177
|
+
| `key` | `str \| None` | `None` | Account key; omit to use `DefaultAzureCredential` (RBAC) |
|
|
178
|
+
| `reducer` | `MessageReducer \| None` | `None` | Optional pruner — see [Built-in Message Pruning](#built-in-message-pruning) |
|
|
179
|
+
| `messages_key` | `str` | `"messages"` | State key that holds the message list |
|
|
180
|
+
|
|
181
|
+
### Methods
|
|
182
|
+
|
|
183
|
+
| Method | Description |
|
|
184
|
+
|---|---|
|
|
185
|
+
| `init_db()` | Initialise database/container references (called automatically by `__init__`) |
|
|
186
|
+
| `save_state(flow_uuid, method_name, state_data)` | Persist flow state (upsert by `flow_uuid`) |
|
|
187
|
+
| `load_state(flow_uuid)` | Load the most recently saved state; returns `None` if not found |
|
|
188
|
+
|
|
189
|
+
## Built-in Message Pruning
|
|
190
|
+
|
|
191
|
+
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.
|
|
192
|
+
|
|
193
|
+
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.**
|
|
194
|
+
|
|
195
|
+
### Install with reducer support
|
|
196
|
+
|
|
197
|
+
```bash
|
|
198
|
+
pip install "crewai-persistence-cosmosdb[reducer]"
|
|
199
|
+
```
|
|
200
|
+
|
|
201
|
+
### Usage
|
|
202
|
+
|
|
203
|
+
```python
|
|
204
|
+
from agentstate_reducer import MessageReducer
|
|
205
|
+
from agentstate_reducer.models import ReducerConfig
|
|
206
|
+
from crewai_persistence_cosmosdb import CosmosDBFlowPersistence
|
|
207
|
+
|
|
208
|
+
config = ReducerConfig(min_messages=10, max_messages=20)
|
|
209
|
+
reducer = MessageReducer(config=config)
|
|
210
|
+
|
|
211
|
+
persistence = CosmosDBFlowPersistence(
|
|
212
|
+
endpoint=os.environ["COSMOS_ENDPOINT"],
|
|
213
|
+
database_name="mydb",
|
|
214
|
+
container_name="flow_states",
|
|
215
|
+
key=os.environ.get("COSMOS_KEY"),
|
|
216
|
+
reducer=reducer, # prune before each save
|
|
217
|
+
messages_key="messages", # state key holding the message list (default)
|
|
218
|
+
)
|
|
219
|
+
```
|
|
220
|
+
|
|
221
|
+
When `len(messages) > max_messages`, the oldest `human`/`ai` messages are removed until `min_messages` remain. The following are **never** pruned:
|
|
222
|
+
|
|
223
|
+
- Index 0 (typically the system prompt) — controlled by `preserve_first=True`
|
|
224
|
+
- `system` and `function` messages
|
|
225
|
+
- `tool` messages — unless their parent `ai` message is pruned (cascade behaviour, configurable)
|
|
226
|
+
|
|
227
|
+
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`).
|
|
228
|
+
|
|
229
|
+
## Data Model
|
|
230
|
+
|
|
231
|
+
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`).
|
|
232
|
+
|
|
233
|
+
| Field | Description |
|
|
234
|
+
|---|---|
|
|
235
|
+
| `id` | Same as `flow_uuid` (CosmosDB document id) |
|
|
236
|
+
| `flow_uuid` | Unique identifier for the flow run (partition key) |
|
|
237
|
+
| `_method_name` | Name of the flow method that triggered the save |
|
|
238
|
+
| `_saved_at` | ISO-8601 UTC timestamp of the save |
|
|
239
|
+
| _user fields_ | All fields from the original state dict / Pydantic model |
|
|
240
|
+
|
|
241
|
+
CosmosDB system fields (`_rid`, `_self`, `_etag`, `_attachments`, `_ts`) are stripped before `load_state` returns the document.
|
|
242
|
+
|
|
243
|
+
## License
|
|
244
|
+
|
|
245
|
+
MIT
|
|
@@ -0,0 +1,37 @@
|
|
|
1
|
+
[build-system]
|
|
2
|
+
requires = ["hatchling"]
|
|
3
|
+
build-backend = "hatchling.build"
|
|
4
|
+
|
|
5
|
+
[project]
|
|
6
|
+
name = "crewai_persistence_cosmosdb"
|
|
7
|
+
version = "0.1.0"
|
|
8
|
+
description = "Azure CosmosDB persistence backend for CrewAI Flows with built-in message pruning via agentstate-reducer"
|
|
9
|
+
authors = [{name = "Kamal", email = "skamalj@github.com"}]
|
|
10
|
+
readme = "README.md"
|
|
11
|
+
requires-python = ">=3.10,<3.14"
|
|
12
|
+
dependencies = [
|
|
13
|
+
"crewai>=1.0.0,<2.0.0",
|
|
14
|
+
"azure-cosmos",
|
|
15
|
+
"azure-identity"
|
|
16
|
+
]
|
|
17
|
+
keywords = ["crewai", "flows", "persistence", "cosmosdb", "azure", "agent-state", "message-reducer"]
|
|
18
|
+
classifiers = [
|
|
19
|
+
"Programming Language :: Python :: 3",
|
|
20
|
+
"License :: OSI Approved :: MIT License",
|
|
21
|
+
"Operating System :: OS Independent",
|
|
22
|
+
"Topic :: Scientific/Engineering :: Artificial Intelligence",
|
|
23
|
+
]
|
|
24
|
+
|
|
25
|
+
[project.optional-dependencies]
|
|
26
|
+
reducer = ["agentstate-reducer>=0.1.1"]
|
|
27
|
+
|
|
28
|
+
[dependency-groups]
|
|
29
|
+
dev = [
|
|
30
|
+
"pytest>=7.0",
|
|
31
|
+
"pytest-cov",
|
|
32
|
+
"langchain-openai",
|
|
33
|
+
"crewai_persistence_cosmosdb[reducer]",
|
|
34
|
+
]
|
|
35
|
+
|
|
36
|
+
[tool.hatch.build.targets.wheel]
|
|
37
|
+
packages = ["src/crewai_persistence_cosmosdb"]
|
|
@@ -0,0 +1,121 @@
|
|
|
1
|
+
from __future__ import annotations
|
|
2
|
+
|
|
3
|
+
from datetime import datetime, timezone
|
|
4
|
+
from typing import Any, Dict, Optional, Union
|
|
5
|
+
|
|
6
|
+
from pydantic import BaseModel, ConfigDict, PrivateAttr, Field
|
|
7
|
+
|
|
8
|
+
from azure.cosmos import CosmosClient, PartitionKey
|
|
9
|
+
|
|
10
|
+
try:
|
|
11
|
+
from azure.identity import DefaultAzureCredential
|
|
12
|
+
except ImportError:
|
|
13
|
+
DefaultAzureCredential = None # type: ignore
|
|
14
|
+
|
|
15
|
+
from crewai.flow.persistence.base import FlowPersistence
|
|
16
|
+
|
|
17
|
+
|
|
18
|
+
_COSMOS_SYSTEM_FIELDS = {"_rid", "_self", "_etag", "_attachments", "_ts"}
|
|
19
|
+
|
|
20
|
+
|
|
21
|
+
class CosmosDBFlowPersistence(FlowPersistence):
|
|
22
|
+
"""Azure CosmosDB persistence backend for CrewAI flow state.
|
|
23
|
+
|
|
24
|
+
Parameters
|
|
25
|
+
----------
|
|
26
|
+
endpoint:
|
|
27
|
+
CosmosDB account endpoint URL.
|
|
28
|
+
database_name:
|
|
29
|
+
Name of the database to use.
|
|
30
|
+
container_name:
|
|
31
|
+
Name of the container to use.
|
|
32
|
+
key:
|
|
33
|
+
Account key for key-based authentication. When omitted,
|
|
34
|
+
``DefaultAzureCredential`` is used (RBAC / Managed Identity).
|
|
35
|
+
reducer:
|
|
36
|
+
Optional ``MessageReducer`` (from ``agentstate-reducer``) that prunes
|
|
37
|
+
the message list before each save.
|
|
38
|
+
messages_key:
|
|
39
|
+
Name of the state key that holds the messages list (default ``"messages"``).
|
|
40
|
+
"""
|
|
41
|
+
|
|
42
|
+
model_config = ConfigDict(arbitrary_types_allowed=True)
|
|
43
|
+
|
|
44
|
+
persistence_type: str = Field(default="cosmosdb")
|
|
45
|
+
endpoint: str
|
|
46
|
+
database_name: str
|
|
47
|
+
container_name: str
|
|
48
|
+
key: Optional[str] = None
|
|
49
|
+
reducer: Optional[Any] = None
|
|
50
|
+
messages_key: str = "messages"
|
|
51
|
+
|
|
52
|
+
_client: Any = PrivateAttr()
|
|
53
|
+
_container: Any = PrivateAttr()
|
|
54
|
+
|
|
55
|
+
def model_post_init(self, __context: Any) -> None:
|
|
56
|
+
if self.key:
|
|
57
|
+
self._client = CosmosClient(self.endpoint, credential=self.key)
|
|
58
|
+
else:
|
|
59
|
+
if DefaultAzureCredential is None:
|
|
60
|
+
raise ImportError(
|
|
61
|
+
"azure-identity is required for RBAC authentication. "
|
|
62
|
+
"Install it with: pip install azure-identity"
|
|
63
|
+
)
|
|
64
|
+
self._client = CosmosClient(self.endpoint, credential=DefaultAzureCredential())
|
|
65
|
+
self.init_db()
|
|
66
|
+
|
|
67
|
+
def init_db(self) -> None:
|
|
68
|
+
if self.key:
|
|
69
|
+
db = self._client.create_database_if_not_exists(id=self.database_name)
|
|
70
|
+
self._container = db.create_container_if_not_exists(
|
|
71
|
+
id=self.container_name,
|
|
72
|
+
partition_key=PartitionKey(path="/flow_uuid"),
|
|
73
|
+
)
|
|
74
|
+
else:
|
|
75
|
+
db = self._client.get_database_client(self.database_name)
|
|
76
|
+
self._container = db.get_container_client(self.container_name)
|
|
77
|
+
|
|
78
|
+
def save_state(
|
|
79
|
+
self,
|
|
80
|
+
flow_uuid: str,
|
|
81
|
+
method_name: str,
|
|
82
|
+
state_data: Union[Dict[str, Any], BaseModel],
|
|
83
|
+
) -> None:
|
|
84
|
+
if isinstance(state_data, BaseModel):
|
|
85
|
+
d: Dict[str, Any] = state_data.model_dump()
|
|
86
|
+
else:
|
|
87
|
+
d = dict(state_data)
|
|
88
|
+
|
|
89
|
+
if self.reducer is not None and self.messages_key in d:
|
|
90
|
+
result = self.reducer.reduce(existing=d[self.messages_key], new=[])
|
|
91
|
+
d[self.messages_key] = result.surviving
|
|
92
|
+
|
|
93
|
+
d["id"] = flow_uuid
|
|
94
|
+
d["flow_uuid"] = flow_uuid
|
|
95
|
+
d["_persistence_meta"] = {
|
|
96
|
+
"method_name": method_name,
|
|
97
|
+
"saved_at": datetime.now(timezone.utc).isoformat(),
|
|
98
|
+
}
|
|
99
|
+
|
|
100
|
+
self._container.upsert_item(d)
|
|
101
|
+
|
|
102
|
+
def load_state(self, flow_uuid: str) -> Optional[Dict[str, Any]]:
|
|
103
|
+
query = (
|
|
104
|
+
"SELECT * FROM c "
|
|
105
|
+
"WHERE c.flow_uuid = @flow_uuid "
|
|
106
|
+
"ORDER BY c._saved_at DESC "
|
|
107
|
+
"OFFSET 0 LIMIT 1"
|
|
108
|
+
)
|
|
109
|
+
params = [{"name": "@flow_uuid", "value": flow_uuid}]
|
|
110
|
+
items = list(
|
|
111
|
+
self._container.query_items(
|
|
112
|
+
query=query,
|
|
113
|
+
parameters=params,
|
|
114
|
+
enable_cross_partition_query=False,
|
|
115
|
+
)
|
|
116
|
+
)
|
|
117
|
+
if not items:
|
|
118
|
+
return None
|
|
119
|
+
|
|
120
|
+
_strip = _COSMOS_SYSTEM_FIELDS | {"flow_uuid", "_persistence_meta"}
|
|
121
|
+
return {k: v for k, v in items[0].items() if k not in _strip}
|