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.
@@ -0,0 +1,6 @@
1
+ __pycache__/
2
+ *.py[cod]
3
+ .venv/
4
+ dist/
5
+ *.egg-info/
6
+ .pytest_cache/
@@ -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,3 @@
1
+ from crewai_persistence_cosmosdb.persistence import CosmosDBFlowPersistence
2
+
3
+ __all__ = ["CosmosDBFlowPersistence"]
@@ -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}