diff --git a/docs/examples/index.md b/docs/examples/index.md index 908a422a2..920e78f88 100644 --- a/docs/examples/index.md +++ b/docs/examples/index.md @@ -11,9 +11,9 @@ Learn by example! Explore working code samples demonstrating Flock's features an Explore the numbered example folders in this repository: -- **01 — The Declarative Way** (`examples/01-the-declarative-way`) — minimal and focused -- **02 — The Blackboard** (`examples/02-the-blackboard`) — architecture overview -- **03 — The Dashboard** (`examples/03-the-dashboard`) — real-time monitoring +- **01 — Getting Started** (`examples/01-getting-started`) — minimal and focused +- **02 — Patterns** (`examples/02-patterns`) — architecture and flow patterns +- **04 — Misc** (`examples/04-misc`) — dashboard and persistence demos ### Component Examples **Learn to build custom components and engines** @@ -44,12 +44,19 @@ These examples show how to extend Flock with custom logic: ### Feature Examples **Focused examples for specific capabilities** -Feature-focused examples are integrated into the folders above (e.g., dashboard edge cases). Additional feature demos may be added over time. +Feature-focused examples are integrated into the folders above (for example dashboard edge cases). Additional feature demos may be added over time. + +### Dapr Examples +**Run Flock with Dapr-backed state stores** + +- **12 — Dapr** (`examples/12-dapr`) — in-memory, Redis encrypted, and PostgreSQL state-store setups + +Start here for setup details: `examples/12-dapr/README.md` ### Dashboard Examples **Interactive dashboard demonstrations** -Check out `examples/03-the-dashboard` to explore: +Check out `examples/04-misc` to explore: - **Declarative Pizza** - Single-agent dashboard demo - **Edge Cases** - Multi-agent cascades and filtering @@ -108,10 +115,10 @@ git clone https://github.com/whiteducksoftware/flock.git cd flock # Run a minimal example -python examples/01-the-declarative-way/01_declarative_pizza.py +python examples/01-getting-started/01_declarative_pizza.py # Run with dashboard -python examples/03-the-dashboard/01_declarative_pizza.py +python examples/04-misc/02-dashboard-edge-cases.py ``` --- @@ -151,7 +158,7 @@ async def main(): asyncio.run(main()) ``` -Run it locally: `python examples/01-the-declarative-way/01_declarative_pizza.py` +Run it locally: `python examples/01-getting-started/01_declarative_pizza.py` --- @@ -159,7 +166,7 @@ Run it locally: `python examples/01-the-declarative-way/01_declarative_pizza.py` **What it demonstrates:** Agent cascades, filtering, and real-time updates -Run: `python examples/03-the-dashboard/02-dashboard-edge-cases.py` +Run: `python examples/04-misc/02-dashboard-edge-cases.py` --- @@ -180,7 +187,7 @@ await orchestrator.serve( ) ``` -Run: `python examples/03-the-dashboard/01_declarative_pizza.py` +Run: `python examples/04-misc/02-dashboard-edge-cases.py` --- diff --git a/docs/getting-started/index.md b/docs/getting-started/index.md index 4f7e779d0..63c0532ba 100644 --- a/docs/getting-started/index.md +++ b/docs/getting-started/index.md @@ -40,6 +40,14 @@ Welcome to Flock! This section will help you get up and running quickly with the [:octicons-arrow-right-24: Server Components Concepts](server-components-concepts.md) +- **🧱 Dapr State Store** + + --- + + Optional distributed persistence with Dapr-backed state stores. + + [:octicons-arrow-right-24: Dapr State Store Guide](../guides/dapr-state-store.md) + --- diff --git a/docs/getting-started/quick-start.md b/docs/getting-started/quick-start.md index 3d961872f..71240052e 100644 --- a/docs/getting-started/quick-start.md +++ b/docs/getting-started/quick-start.md @@ -30,6 +30,8 @@ export DEFAULT_MODEL="openai/gpt-4.1" **That's it.** Flock works with any LiteLLM-supported model (OpenAI, Anthropic, Azure, local models, etc.). +Need distributed blackboard persistence? See [Dapr State Store Integration](../guides/dapr-state-store.md). + --- ## Your First Agent (60 Seconds) diff --git a/docs/guides/blackboard.md b/docs/guides/blackboard.md index e163f4088..742753bc4 100644 --- a/docs/guides/blackboard.md +++ b/docs/guides/blackboard.md @@ -531,20 +531,26 @@ await flock.store.clear() ### Memory Management -**Current limitation:** Blackboard is in-memory only (v0.5.0) +Flock supports multiple blackboard storage backends depending on your deployment needs. ```python -# ⚠️ In-memory only -# After 10,000 artifacts, memory usage grows -# Restart required to clear +# In-memory store (default) +flock = Flock("openai/gpt-4.1") + +# SQLite store (local durable history) +from flock.store import SQLiteBlackboardStore +flock = Flock("openai/gpt-4.1", store=SQLiteBlackboardStore(".flock/history.db")) -# ✅ v1.0 will support: -# - Redis backend (distributed state) -# - PostgreSQL backend (persistent history) -# - Automatic artifact expiration -# - Query by time range +# Dapr store (distributed backend) +from flock.storage import DaprStateBlackboardStore, DaprStateBlackboardConfig +flock = Flock( + "openai/gpt-4.1", + store=DaprStateBlackboardStore(DaprStateBlackboardConfig(store_name="flockstate")), +) ``` +For backend-specific guidance, see the [Persistent Blackboard guide](persistent-blackboard.md) and the [Dapr State Store guide](dapr-state-store.md). + ### Observability **Enable tracing to see artifact flow:** diff --git a/docs/guides/dapr-state-store.md b/docs/guides/dapr-state-store.md new file mode 100644 index 000000000..299029b86 --- /dev/null +++ b/docs/guides/dapr-state-store.md @@ -0,0 +1,184 @@ +--- +title: Dapr State Store Integration +description: Use Dapr-supported state stores as an optional distributed blackboard backend for Flock. +tags: + - blackboard + - persistence + - dapr + - integrations +search: + boost: 1.4 +--- + +# Dapr State Store Integration + +Flock includes Dapr-backed blackboard storage as an optional feature. This lets you keep your existing agent and artifact model while switching the backend from local storage to any Dapr-supported state store. + +Use Dapr storage when you need distributed state, backend-level encryption, or backend-specific TTL and query capabilities. + +--- + +## Install + +For this repository: + +```bash +uv sync --extra dapr +``` + +For external projects using Flock: + +```bash +uv add "flock-core[dapr]" +``` + +If you only need local durable history in a single process, prefer SQLite via [Persistent Blackboard History](persistent-blackboard.md). + +--- + +## Quick Start + +```python +from flock import Flock +from flock.storage import ( + DaprStateBlackboardConfig, + DaprStateBlackboardStore, + DaprStateBlackboardStoreClientConfig, +) + +client_config = DaprStateBlackboardStoreClientConfig( + dapr_grpc_endpoint="localhost:50001", +) + +store_config = DaprStateBlackboardConfig( + store_name="flockstate", # must match Dapr component metadata.name + supports_transactions=True, + supports_etag=True, + consistency="strong", # eventual | strong | unspecified + client_config=client_config, +) + +store = DaprStateBlackboardStore(config=store_config) + +flock = Flock( + model="openai/gpt-4.1", + store=store, +) +``` + +The rest of your agent setup stays unchanged. + +### Backend Capability Notes + +- `supports_transactions=True` is recommended for unencrypted Redis/PostgreSQL backends. +- For encrypted Redis (`encrypted_backend=True`), transactions are automatically disabled by the store implementation because of a Dapr runtime limitation. +- Enable `supports_dapr_query_lang=True` only when your selected state store/component actually supports Dapr query API. + +--- + +## Backend Matrix + +The examples in this repository include three reference Dapr setups: + +| Backend | Example Directory | Encryption | Transactions | Query API | TTL | +| --- | --- | --- | --- | --- | --- | +| In-memory | `examples/12-dapr/inmemory/` | No | No | No | No | +| Redis (encrypted) | `examples/12-dapr/redis_encrypted/` | Yes | Yes | Yes | Yes | +| PostgreSQL 17 | `examples/12-dapr/postgresql_unencrypted/` | No | Yes | No | No | + +--- + +## Configuration Reference + +### DaprStateBlackboardConfig + +| Field | Type | Default | Description | +| --- | --- | --- | --- | +| store_name | str | statestore | Dapr component name (metadata.name) | +| supports_ttl | bool | False | Enable TTL-based expiration if backend supports TTL | +| encrypted_backend | bool | False | Use Dapr encryption metadata path | +| backend_encryption_key | str \| None | None | Optional encryption key when encrypted backend is enabled | +| supports_transactions | bool | False | Use execute_state_transaction for atomic writes | +| supports_dapr_query_lang | bool | False | Enable Dapr query API paths | +| supports_etag | bool | False | Enable optimistic concurrency with ETags | +| etag_max_retries | int | 3 | Retries on ETag conflict | +| consistency | str | unspecified | eventual, strong, or unspecified | +| entries_ttl_seconds | int \| None | None | TTL in seconds when supports_ttl is enabled | +| client_config | DaprStateBlackboardStoreClientConfig \| None | None | Optional Dapr client settings | + +### DaprStateBlackboardStoreClientConfig + +| Field | Type | Default | Description | +| --- | --- | --- | --- | +| dapr_grpc_endpoint | str \| None | None | Dapr gRPC endpoint, for example localhost:50001 | +| headers_callback | Callable \| None | None | Request headers callback | +| interceptors | list \| None | None | gRPC interceptors | +| http_timeout_seconds | int \| None | None | HTTP timeout for Dapr SDK calls | +| max_grpc_message_length | int \| None | None | Max gRPC message size | +| retry_policy | RetryPolicy \| None | None | Dapr SDK retry policy | + +--- + +## Known Limitations + +- Encryption and transactions cannot be combined due to a Dapr runtime transaction serialization issue; Flock automatically falls back to non-transactional writes for encrypted backends. +- Dapr query support depends on the selected state store and component implementation. +- TTL support depends on backend capabilities. +- Distributed writes can still race across multiple Flock instances; use ETags where supported. + +--- + +## Run The Example Stacks + +Only run one backend stack at a time because they share host ports. + +1. In-memory backend: + +```bash +cd examples/12-dapr/inmemory +cp secrets.example.json secrets.json +docker compose up -d +``` + +2. Redis encrypted backend: + +```bash +cd examples/12-dapr/redis_encrypted +cp secrets.example.json secrets.json +docker compose up -d +``` + +3. PostgreSQL backend: + +```bash +cd examples/12-dapr/postgresql_unencrypted +cp secrets.example.json secrets.json +docker compose up -d +``` + +For all three setups, copy `secrets.example.json` to `secrets.json` and fill in required keys before running. + +Then run the matching example script from repository root: + +```bash +export DAPR_GRPC_ENDPOINT="localhost:50001" + +uv run python examples/12-dapr/inmemory/flock_dapr_inmemory.py +# or +uv run python examples/12-dapr/redis_encrypted/flock_dapr_redis.py +# or +uv run python examples/12-dapr/postgresql_unencrypted/flock_dapr_postgresql.py +``` + +--- + +## Choosing SQLite vs Dapr + +- Choose SQLite for single-node durable history and simple local operations. +- Choose Dapr for shared distributed state, pluggable backend infrastructure, and backend-specific capabilities. + +For SQLite-focused persistence and dashboard history workflows, see the [Persistent Blackboard guide](persistent-blackboard.md). + +For full Dapr example details, see [examples/12-dapr/README.md](https://github.com/whiteducksoftware/flock/tree/main/examples/12-dapr/README.md). + +You can also browse runnable setup variants in [Examples Index](../examples/index.md). diff --git a/docs/guides/index.md b/docs/guides/index.md index 29aeeadc1..aed177511 100644 --- a/docs/guides/index.md +++ b/docs/guides/index.md @@ -80,6 +80,14 @@ Comprehensive guides for building production-ready multi-agent systems with Floc
+- **🧱 Dapr State Store Integration** + + --- + + Plug Redis, PostgreSQL, or other Dapr-supported state stores into Flock's blackboard backend. + + [:octicons-arrow-right-24: Dapr State Store Guide](dapr-state-store.md) + - **🦞 OpenClaw Integration** --- diff --git a/docs/guides/persistent-blackboard.md b/docs/guides/persistent-blackboard.md index 0fd4fae5a..779e174f2 100644 --- a/docs/guides/persistent-blackboard.md +++ b/docs/guides/persistent-blackboard.md @@ -21,7 +21,7 @@ The in-memory blackboard is perfect for local prototyping, but production teams - **Replay & Audits** — Inspect any published artifact (including payloads, tags, visibility, version, correlation IDs) even after the orchestrator restarts. - **Lifecycle Analytics** — Summaries and agent history endpoints expose production/consumption counts, visibility breakdowns, and tag trends over arbitrary windows. - **Operator Experience** — Dashboard users can scroll back through previous runs, inspect payload details, and view retention banners that explain how far history goes. -- **Future Backends** — The updated store interface (`FilterConfig`, `ArtifactEnvelope`) paves the way for Postgres, DuckDB, or cloud warehouses with identical semantics. +- **Alternative Backends** — The same store contract also supports distributed backends via Dapr state stores. --- @@ -48,6 +48,23 @@ Key capabilities: --- +## Choosing SQLite vs Dapr + +Use this quick matrix when deciding which store to use: + +| Option | Best For | Persistence Scope | Tradeoffs | +| --- | --- | --- | --- | +| In-memory (default) | Local prototyping and tests | Process lifetime only | Fastest setup, but data is lost on restart | +| SQLiteBlackboardStore | Single-node durable history | Local file persistence | Simple operations, but not a shared distributed backend | +| DaprStateBlackboardStore | Shared/distributed blackboard workloads | Depends on chosen Dapr state store | Flexible backend choice, but requires Dapr sidecar + component setup | + +- Use SQLite when you want local durable history for a single-node deployment. +- Use Dapr-backed storage when you need shared/distributed state or backend-specific capabilities such as store-managed encryption and TTL. + +See [Dapr State Store Integration](dapr-state-store.md) for distributed setup options. + +--- + ## Historical APIs Once the SQLite store is active, the HTTP control plane exposes richer endpoints: diff --git a/docs/index.md b/docs/index.md index 86fa65da3..868ea0c69 100644 --- a/docs/index.md +++ b/docs/index.md @@ -120,6 +120,14 @@ await flock.serve(dashboard=True) [Guides →](guides/index.md){ .md-button } +- :material-database: **Dapr State Store** + + --- + + Use Redis, PostgreSQL, or other Dapr-supported state stores as distributed blackboard backends. + + [Dapr Guide →](guides/dapr-state-store.md){ .md-button } + - :material-api: **API Reference** --- diff --git a/docs/reference/configuration.md b/docs/reference/configuration.md index 9beebd169..a27aaf667 100644 --- a/docs/reference/configuration.md +++ b/docs/reference/configuration.md @@ -128,5 +128,17 @@ Environment variables are resolved in this order: For detailed configuration options, see: - [Installation Guide](../getting-started/installation.md) for environment setup - [.envtemplate](https://github.com/whiteducksoftware/flock/blob/main/.envtemplate) for all available options + +## Storage Backends + +Flock can run with different blackboard backends: + +- In-memory blackboard (default, no extra setup) +- SQLite durable history (single-node persistence) +- Dapr-backed state stores (distributed backend, optional dependency) + +See: +- [Persistent Blackboard History](../guides/persistent-blackboard.md) +- [Dapr State Store Integration](../guides/dapr-state-store.md) - [DSPy Engine Guide](../guides/dspy-engine.md) for `lm_kwargs`, adapters, and Azure auth wiring - [Tracing Configuration](../guides/tracing/) for telemetry settings diff --git a/examples/12-dapr/README.md b/examples/12-dapr/README.md new file mode 100644 index 000000000..0fb179bf7 --- /dev/null +++ b/examples/12-dapr/README.md @@ -0,0 +1,301 @@ +# 🐦 flock-dapr + +[![License: MIT](https://img.shields.io/badge/License-MIT-blue.svg)](LICENSE) +[![Python 3.12+](https://img.shields.io/badge/python-3.12%2B-blue.svg)](https://www.python.org/downloads/) +[![Flock ≥ 0.5](https://img.shields.io/badge/flock--core-%E2%89%A5%200.5-orange.svg)](https://pypi.org/project/flock-core/) +[![Dapr 1.13–1.14](https://img.shields.io/badge/dapr-1.13--1.14-blueviolet.svg)](https://dapr.io/) + +> **Bring your own state store.** +> Plug any Dapr-supported backend — Redis, PostgreSQL, CosmosDB, and more — into +> [Flock](https://whiteducksoftware.github.io/flock/)'s Blackboard. + +> Looking for docs-site version? See the +> [Dapr State Store Integration guide](https://whiteducksoftware.github.io/flock/guides/dapr-state-store/). + +--- + +## Why flock-dapr? + +Flock ships with an **in-memory** store and a **SQLite** store for local +development and single-node persistence. That's great for prototyping, but +production workloads often need: + +- 🔄 **Distributed state** — multiple Flock instances sharing a single + blackboard (Redis, CosmosDB, …) +- 🔒 **Encryption at rest** — leveraging Dapr's built-in + `primaryEncryptionKey` support +- ⏱️ **TTL-based expiration** — automatic cleanup of stale artifacts +- 🏗️ **Operational flexibility** — swap backends without changing + application code + +**flock-dapr** implements Flock's +[`BlackboardStore`](https://whiteducksoftware.github.io/flock/guides/persistent-blackboard/) +contract on top of [Dapr State Management](https://docs.dapr.io/developing-applications/building-blocks/state-management/), +so you can point Flock at *any* Dapr state-store component and keep the rest of +your agent code untouched. + +--- + +## 🚀 Installation + +From this repository root: + +```bash +uv sync --extra dapr +``` + +For external projects: + +```bash +uv add "flock-core[dapr]" +``` + +If you only need single-node durable history, consider SQLite first via the +[Persistent Blackboard guide](https://whiteducksoftware.github.io/flock/guides/persistent-blackboard/). + +--- + +## ⚡ Quick Start + +```python +import asyncio + +from flock import Flock + +from flock.storage import ( + DaprStateBlackboardConfig, + DaprStateBlackboardStore, + DaprStateBlackboardStoreClientConfig, +) + +# 1. Configure the Dapr state store connection +client_config = DaprStateBlackboardStoreClientConfig() + +store_config = DaprStateBlackboardConfig( + store_name="flockstate", # must match your Dapr component name + supports_transactions=True, # Redis, PostgreSQL, CosmosDB + supports_etag=True, # optimistic concurrency control + consistency="strong", # "eventual" | "strong" | "unspecified" +) + +# 2. Create the store and wire it into Flock +dapr_store = DaprStateBlackboardStore(config=store_config) + +flock = Flock( + model="openai/gpt-4.1", + store=dapr_store, # ← swap in the Dapr backend +) + +# 3. Define agents as usual — they know nothing about the backend +# bug_detector = flock.agent("bugs").consumes(Code).publishes(BugReport) + +async def main(): + await flock.serve(dashboard=True) + +asyncio.run(main()) +``` + +That's it — your entire agent swarm now reads and writes through Dapr. No +other code changes required. 🎉 + +--- + +## 🗄️ Supported Backends + +flock-dapr works with **any** [Dapr state-store component](https://docs.dapr.io/reference/components-reference/supported-state-stores/). +The repository ships with ready-to-use example environments for three +backends: + +| Backend | Example Directory | Encryption | Transactions | Query API | TTL | +|---|---|---|---|---|---| +| **In-memory** | `examples/12-dapr/inmemory/` | — | — | — | — | +| **Redis** (redis-stack) | `examples/12-dapr/redis_encrypted/` | ✅ Primary + secondary key | ✅ | ✅ | ✅ | +| **PostgreSQL 17** | `examples/12-dapr/postgresql_unencrypted/` | — | ✅ | — | — | + +> **Tip:** All setups use the same `DaprStateBlackboardStore` class — only the +> Dapr component YAML and config flags change. + +Other state stores known to work with Dapr (CosmosDB, DynamoDB, Cassandra, …) +should work out of the box; just provide the matching Dapr component definition +and adjust the config flags accordingly. + +--- + +## ⚙️ Configuration + +### `DaprStateBlackboardConfig` + +| Field | Type | Default | Description | +|---|---|---|---| +| `store_name` | `str` | `"statestore"` | Dapr component name (`metadata.name` in the component YAML) | +| `supports_ttl` | `bool` | `False` | Enable TTL-based entry expiration (backend must support it) | +| `encrypted_backend` | `bool` | `False` | Indicate the backend uses Dapr encryption (`primaryEncryptionKey`) | +| `backend_encryption_key` | `str \| None` | `None` | Encryption key (only used when `encrypted_backend=True`) | +| `supports_transactions` | `bool` | `False` | Use `execute_state_transaction` for atomic index updates (Redis, PostgreSQL, CosmosDB) | +| `supports_dapr_query_lang` | `bool` | `False` | Use Dapr's query API for `query_artifacts` / `fetch_graph_artifacts` | +| `supports_etag` | `bool` | `False` | Optimistic concurrency control via ETags (first-write-wins with auto-retry) | +| `etag_max_retries` | `int` | `3` | Max retries on ETag conflict (ignored when `supports_etag=False`) | +| `consistency` | `str` | `"unspecified"` | Consistency level: `"eventual"`, `"strong"`, or `"unspecified"` | +| `entries_ttl_seconds` | `int \| None` | `None` | TTL in seconds for state entries (requires `supports_ttl=True`) | +| `client_config` | `...ClientConfig \| None` | `None` | Optional Dapr client settings (see below) | + +### `DaprStateBlackboardStoreClientConfig` + +| Field | Type | Default | Description | +|---|---|---|---| +| `dapr_grpc_endpoint` | `str \| None` | `None` | Dapr runtime gRPC address (e.g. `localhost:50001`) | +| `headers_callback` | `Callable \| None` | `None` | Callable returning `dict[str, str]` headers per request | +| `interceptors` | `list[...] \| None` | `None` | gRPC client interceptors | +| `http_timeout_seconds` | `int \| None` | `None` | HTTP timeout for Dapr connections | +| `max_grpc_message_length` | `int \| None` | `None` | Max gRPC message size in bytes | +| `retry_policy` | `RetryPolicy \| None` | `None` | Dapr SDK retry policy | + +--- + +## 🏗️ Architecture + +``` +┌─────────────────────────────┐ +│ Your Flock App │ +│ agents · types · workflows │ +└────────────┬────────────────┘ + │ store=dapr_store +┌────────────▼────────────────┐ +│ DaprStateBlackboardStore │ +│ serialize · index · r/w │ +└────────────┬────────────────┘ + │ gRPC +┌────────────▼────────────────┐ +│ Dapr Sidecar │ +│ state management building │ +│ block + encryption + TTL │ +└────────────┬────────────────┘ + │ +┌────────────▼────────────────┐ +│ State Store Component │ +│ Redis · PostgreSQL · ... │ +└─────────────────────────────┘ +``` + +### Key Design Decisions + +- **Transactional vs non-transactional paths** — When the backend supports + transactions, artifact writes and index updates are wrapped in + `execute_state_transaction` for atomicity. When encryption is enabled, + transactions are automatically disabled (see [Known Limitations](#-known-limitations)) + and the store falls back to individual `save_state` calls with lazy index + reconciliation on read. + +- **Index-based key management** — Dapr state stores are key-value stores. + flock-dapr maintains secondary indexes manually: + + | Key pattern | Value | + |---|---| + | `artifact:{uuid}` | Serialized `Artifact` JSON | + | `idx:artifacts` | JSON list of all artifact UUID strings | + | `idx:type:{type_name}` | JSON list of UUIDs for that artifact type | + | `consumptions:{artifact_id}` | JSON list of `ConsumptionRecord` dicts | + | `snapshot:{agent_name}` | Serialized `AgentSnapshotRecord` JSON | + | `idx:snapshots` | JSON list of agent name strings | + +- **Custom serialization layer** — Artifacts use Pydantic's built-in + serialization; dataclass-based types (`ConsumptionRecord`, + `AgentSnapshotRecord`) are handled via a dedicated + [serialization module](../../src/flock/storage/dapr/_serialization.py) that deals with + `UUID`, `datetime`, and visibility discriminators. + +--- + +## ⚠️ Known Limitations + +| Limitation | Details | +|---|---| +| **Encryption disables transactions** | Dapr's Go runtime corrupts values in `ExecuteStateTransaction` before encrypting (converts `[]byte` via `fmt.Appendf`). flock-dapr auto-detects this and falls back to non-transactional writes. Index consistency is maintained via lazy reconciliation. | +| **Query API varies by backend** | `supports_dapr_query_lang=True` requires a backend that implements [Dapr's query API](https://docs.dapr.io/developing-applications/building-blocks/state-management/howto-state-query-api/) (e.g. Redis with RediSearch). PostgreSQL v2 does not support it. | +| **TTL depends on the state store** | Set `supports_ttl=True` only if the backing store actually supports TTL; otherwise Dapr will return errors. | +| **No distributed locking** | Concurrent writes from multiple Flock instances to the same index can race. Use `supports_etag=True` for optimistic concurrency control to mitigate this. | + +--- + +## 🛠️ Development Setup + +The `examples/` directory ships three Docker Compose stacks — one for each +reference backend. **They share the same host ports, so run only one at a +time.** + +### 1. Choose a backend + +**In-memory (simplest — no external services):** + +```bash +cd examples/12-dapr/inmemory +cp secrets.example.json secrets.json # fill in your LLM API key, model, etc. +docker compose up -d +``` + +> State lives inside the Dapr sidecar and is lost on restart. Great for quick +> iteration without any database dependencies. + +**Redis (encrypted):** + +```bash +cd examples/12-dapr/redis_encrypted +cp secrets.example.json secrets.json # fill in your LLM API key, model, etc. +docker compose up -d +``` + +**PostgreSQL (unencrypted):** + +```bash +cd examples/12-dapr/postgresql_unencrypted +cp secrets.example.json secrets.json # fill in your LLM API key, model, etc. +docker compose up -d +``` + +### 2. Configure secrets + +Edit `secrets.json` with your values: + +```jsonc +{ + "api_key": "sk-...", // your LLM API key + "base_url": "https://...", // LLM endpoint + "api_version": "2024-12-01-preview", // API version (Azure OpenAI) + "state_store_name": "flockstate", // must match component YAML + "default_model": "openai/gpt-4.1" // model identifier +} +``` + + Note: the Redis example also expects `default_model` in `secrets.json`. + +### 3. Run an example + +```bash +# Make sure the Dapr sidecar is reachable +export DAPR_GRPC_ENDPOINT="localhost:50001" + +# In-memory example +uv run python examples/12-dapr/inmemory/flock_dapr_inmemory.py + +# Redis example +uv run python examples/12-dapr/redis_encrypted/flock_dapr_redis.py + +# PostgreSQL example +uv run python examples/12-dapr/postgresql_unencrypted/flock_dapr_postgresql.py +``` + +### 4. Dapr component files + +Each example stack ships three component definitions under `components/`: + +| File | Purpose | +|---|---| +| `secretstore.yaml` | Local file-based secret store (`secrets.json`) | +| `statestore.yaml` | State store component (in-memory, Redis, or PostgreSQL) | +| `resiliency.yaml` | Retry and circuit-breaker policies | + +--- + +

+ Built with 🦆 by white duck GmbH +

diff --git a/examples/12-dapr/inmemory/components/resiliency.yaml b/examples/12-dapr/inmemory/components/resiliency.yaml new file mode 100644 index 000000000..f53827f02 --- /dev/null +++ b/examples/12-dapr/inmemory/components/resiliency.yaml @@ -0,0 +1,27 @@ +apiVersion: dapr.io/v1alpha1 +kind: Resiliency +metadata: + name: flockdevresiliency +scopes: + - flock-dev + +spec: + policies: + retries: + retryForever: + policy: constant + duration: 5s + maxRetries: -1 + + circuitBreakers: + simpleCB: + maxRequests: 1 + timeout: 5s + trip: consecutiveFailures >= 5 + + targets: + components: + statestore: + outbound: + retry: retryForever + circuitBreaker: simpleCB diff --git a/examples/12-dapr/inmemory/components/secretstore.yaml b/examples/12-dapr/inmemory/components/secretstore.yaml new file mode 100644 index 000000000..da6fe319d --- /dev/null +++ b/examples/12-dapr/inmemory/components/secretstore.yaml @@ -0,0 +1,15 @@ +apiVersion: dapr.io/v1alpha1 +kind: Component +metadata: + name: flock-dev-secretstore + namespace: development +spec: + type: secretstores.local.file + version: v1 + metadata: + - name: secretsFile + value: /secrets/secrets.json + - name: nestedSeparator + value: ":" + - name: multiValued + value: "false" diff --git a/examples/12-dapr/inmemory/components/statestore.yaml b/examples/12-dapr/inmemory/components/statestore.yaml new file mode 100644 index 000000000..bc4bd9f81 --- /dev/null +++ b/examples/12-dapr/inmemory/components/statestore.yaml @@ -0,0 +1,9 @@ +apiVersion: dapr.io/v1alpha1 +kind: Component +metadata: + name: flockstate + namespace: development +spec: + type: state.in-memory + version: v1 + metadata: [] diff --git a/examples/12-dapr/inmemory/docker-compose.yml b/examples/12-dapr/inmemory/docker-compose.yml new file mode 100644 index 000000000..c312f4cea --- /dev/null +++ b/examples/12-dapr/inmemory/docker-compose.yml @@ -0,0 +1,47 @@ +# Lightweight in-memory dev stack. +# No external state-store service — state lives inside the Dapr sidecar. +# Shares the same host ports as the other dev stacks — run only one at a time. + +services: + + # Sidecar for Flock Dev app + flock-dapr: + image: "daprio/daprd:edge" + command: [ + "./daprd", + "--app-id", "flock-dev", + "--placement-host-address", "placement:50006", + "--scheduler-host-address", "scheduler:50007", + "--resources-path", "/components" + ] + volumes: + - "./components/:/components" + - "./secrets.json:/secrets/secrets.json:ro" + ports: + - "50001:50001" # Dapr gRPC port (used by DaprClient on the host) + - "3500:3500" # Dapr HTTP port + networks: + - dapr-net + depends_on: + - placement + - scheduler + + placement: + image: "daprio/placement" + command: ["./placement", "--port", "50006"] + ports: + - "50006:50006" + networks: + - dapr-net + + scheduler: + image: "daprio/scheduler" + command: ["./scheduler", "--port", "50007", "--etcd-data-dir", "/data"] + user: root + volumes: + - "./dapr-etcd-data/:/data" + networks: + - dapr-net + +networks: + dapr-net: null diff --git a/examples/12-dapr/inmemory/flock_dapr_inmemory.py b/examples/12-dapr/inmemory/flock_dapr_inmemory.py new file mode 100644 index 000000000..96dc55b3d --- /dev/null +++ b/examples/12-dapr/inmemory/flock_dapr_inmemory.py @@ -0,0 +1,177 @@ +import asyncio +import os + +from dapr.clients import DaprClient +from pydantic import BaseModel, Field + +from flock import Flock, PublicVisibility, flock_type +from flock.logging.logging import configure_logging, get_logger +from flock.storage.dapr import ( + DaprStateBlackboardConfig, + DaprStateBlackboardStore, + DaprStateBlackboardStoreClientConfig, +) + + +logger = get_logger(__name__) + + +@flock_type +class BandConcept(BaseModel): + genre: str = Field(description="Musical genre (rock, jazz, metal, pop, etc.)") + vibe: str = Field(description="The band's vibe or aesthetic") + target_audience: str = Field(description="Who should love this band?") + + +@flock_type +class BandLineup(BaseModel): + band_name: str = Field(description="Cool band name") + members: list[dict[str, str]] = Field( + description="List of band members with their roles" + ) + origin_story: str = Field(description="How the band formed", min_length=100) + signature_sound: str = Field(description="What makes their sound unique") + + +@flock_type +class Album(BaseModel): + title: str = Field(description="Album title in ALL CAPS") + tracklist: list[dict[str, str]] = Field( + description="Songs with titles and brief descriptions", + min_length=8, + max_length=12, + ) + genre_fusion: str = Field(description="How this album blends genres") + standout_track: str = Field(description="The track that'll be a hit") + production_notes: str = Field(description="Special production techniques") + + +@flock_type +class MarketingCopy(BaseModel): + press_release: str = Field( + description="Professional press release announcing the album", min_length=200 + ) + social_media_hook: str = Field( + description="Catchy social post (280 chars max)", max_length=280 + ) + billboard_tagline: str = Field( + description="10-word tagline for billboards", max_length=100 + ) + target_playlists: list[str] = Field( + description="Spotify/Apple Music playlists to pitch to", + min_length=3, + max_length=5, + ) + + +FLOCK_SECRET_STORE = "flock-dev-secretstore" +FLOCK_BASE_URL_SECRET_KEY = "base_url" +FLOCK_API_VERSION_SECRET_KEY = "api_version" +FLOCK_API_KEY_SECRET_KEY = "api_key" +FLOCK_STATE_STORE_SECRET_KEY = "state_store_name" +FLOCK_DEFAULT_MODEL_SECRET_KEY = "default_model" + + +async def full_blown_flock_test(): + """Test the blackboard with a small team of agents.""" + # Configure global log levels + configure_logging("ERROR", external_level="ERROR") + # Get Dapr secrets + base_url: str | None = None + api_version: str | None = None + state_store_name: str | None = None + api_key: str | None = None + default_model: str | None = None + with DaprClient() as client: + logger.info(f"Retrieving {FLOCK_BASE_URL_SECRET_KEY} secret") + base_url = client.get_secret( + store_name=FLOCK_SECRET_STORE, key=FLOCK_BASE_URL_SECRET_KEY + ).secret.get(FLOCK_BASE_URL_SECRET_KEY) + logger.info(f"Got {FLOCK_BASE_URL_SECRET_KEY}") + logger.info(f"Retrieving {FLOCK_API_VERSION_SECRET_KEY}") + api_version = client.get_secret( + store_name=FLOCK_SECRET_STORE, key=FLOCK_API_VERSION_SECRET_KEY + ).secret.get(FLOCK_API_VERSION_SECRET_KEY) + logger.info(f"Got {FLOCK_API_VERSION_SECRET_KEY}") + logger.info(f"Retrieving {FLOCK_API_KEY_SECRET_KEY}") + api_key = client.get_secret( + store_name=FLOCK_SECRET_STORE, key=FLOCK_API_KEY_SECRET_KEY + ).secret.get(FLOCK_API_KEY_SECRET_KEY) + logger.info(f"Got {FLOCK_API_KEY_SECRET_KEY}") + logger.info(f"Retrieving {FLOCK_STATE_STORE_SECRET_KEY}") + state_store_name = client.get_secret( + store_name=FLOCK_SECRET_STORE, key=FLOCK_STATE_STORE_SECRET_KEY + ).secret.get(FLOCK_STATE_STORE_SECRET_KEY) + logger.info(f"Retrieving {FLOCK_DEFAULT_MODEL_SECRET_KEY}") + default_model = client.get_secret( + store_name=FLOCK_SECRET_STORE, key=FLOCK_DEFAULT_MODEL_SECRET_KEY + ).secret.get(FLOCK_DEFAULT_MODEL_SECRET_KEY) + logger.info(f"Got {FLOCK_DEFAULT_MODEL_SECRET_KEY}") + # Check if all keys have been retrieved + # if one is missing, exit here and throw an exception + if any([ + base_url is None, + api_version is None, + state_store_name is None, + api_key is None, + default_model is None, + ]): + logger.error("UNABLE TO RETRIEVE FULL LIST OF SECRETS!!!") + raise ValueError + # Set required environment-variables before creating Flock-instance + os.environ["AZURE_API_BASE"] = base_url + os.environ["AZURE_API_VERSION"] = api_version + os.environ["AZURE_API_KEY"] = api_key + # Initialize dapr store. + # In-memory: no encryption, no transactions, no TTL, no query API. + client_config = DaprStateBlackboardStoreClientConfig() + store_config = DaprStateBlackboardConfig( + store_name=state_store_name, + supports_ttl=False, + encrypted_backend=False, + backend_encryption_key=None, + supports_transactions=False, + entries_ttl_seconds=None, + client_config=client_config, + supports_dapr_query_lang=False, + supports_etag=False, + consistency="eventual", + ) + dapr_store = DaprStateBlackboardStore(config=store_config) + # Initialize Flock Agent Swarm + flock = Flock( + model=default_model, + max_agent_iterations=100, + no_output=True, + store=dapr_store, # Add Dapr Blackboard Store as backend + ) + _ = ( + flock.agent("talent_scout") + .description("A legendary talent scout who assembles perfect band lineups") + .consumes(BandConcept) + .publishes(BandLineup, visibility=PublicVisibility()) + ) + _ = ( + flock.agent("music_producer") + .description("A visionary music producer who creates debut album concepts") + .consumes(BandLineup) + .publishes(Album, visibility=PublicVisibility()) + ) + _ = ( + flock.agent("marketing_guru") + .description("A marketing genius who writes compelling promotional material") + .consumes(Album) + .publishes(MarketingCopy, visibility=PublicVisibility()) + ) + + await flock.serve(dashboard=True) + + +async def main_test(): + print("Liftoff!") + print("=== Testing Flock with an In-Memory Backend ===") + await full_blown_flock_test() + + +if __name__ == "__main__": + asyncio.run(main_test()) diff --git a/examples/12-dapr/inmemory/secrets.example.json b/examples/12-dapr/inmemory/secrets.example.json new file mode 100644 index 000000000..0f66aafab --- /dev/null +++ b/examples/12-dapr/inmemory/secrets.example.json @@ -0,0 +1,7 @@ +{ + "api_key": "", + "base_url": "", + "api_version": "", + "state_store_name": "", + "default_model": "" +} diff --git a/examples/12-dapr/postgresql_unencrypted/components/resiliency.yaml b/examples/12-dapr/postgresql_unencrypted/components/resiliency.yaml new file mode 100644 index 000000000..fd2036c8b --- /dev/null +++ b/examples/12-dapr/postgresql_unencrypted/components/resiliency.yaml @@ -0,0 +1,27 @@ +apiVersion: dapr.io/v1alpha1 +kind: Resiliency +metadata: + name: flockdevresiliency +scopes: + - flock-dev + +spec: + policies: + retries: + retryForever: + policy: constant + duration: 5s + maxRetries: -1 + + circuitBreakers: + simpleCB: + maxRequests: 1 + timeout: 5s + trip: consecutiveFailures >= 5 + + targets: + components: + statestore: + outbound: + retry: retryForever + circuitBreaker: simpleCB \ No newline at end of file diff --git a/examples/12-dapr/postgresql_unencrypted/components/secretstore.yaml b/examples/12-dapr/postgresql_unencrypted/components/secretstore.yaml new file mode 100644 index 000000000..29370ea1e --- /dev/null +++ b/examples/12-dapr/postgresql_unencrypted/components/secretstore.yaml @@ -0,0 +1,15 @@ +apiVersion: dapr.io/v1alpha1 +kind: Component +metadata: + name: flock-dev-secretstore + namespace: development +spec: + type: secretstores.local.file + version: v1 + metadata: + - name: secretsFile + value: /secrets/secrets.json + - name: nestedSeparator + value: ":" + - name: multiValued + value: "false" \ No newline at end of file diff --git a/examples/12-dapr/postgresql_unencrypted/components/statestore.yaml b/examples/12-dapr/postgresql_unencrypted/components/statestore.yaml new file mode 100644 index 000000000..743aeadf8 --- /dev/null +++ b/examples/12-dapr/postgresql_unencrypted/components/statestore.yaml @@ -0,0 +1,27 @@ +apiVersion: dapr.io/v1alpha1 +kind: Component +metadata: + name: flockstate + namespace: development +spec: + type: state.postgresql + version: v2 + metadata: + - name: host + value: "postgres" + - name: port + value: "5432" + - name: user + value: "postgres" + - name: password + secretKeyRef: + name: postgresql_password + key: postgresql_password + - name: database + value: "flock_dapr" + - name: sslmode + value: "disable" + - name: keyPrefix + value: namespace +auth: + secretStore: flock-dev-secretstore \ No newline at end of file diff --git a/examples/12-dapr/postgresql_unencrypted/docker-compose.yml b/examples/12-dapr/postgresql_unencrypted/docker-compose.yml new file mode 100644 index 000000000..1a91b1951 --- /dev/null +++ b/examples/12-dapr/postgresql_unencrypted/docker-compose.yml @@ -0,0 +1,69 @@ +# NOTE: This stack uses the same host ports as the Redis encrypted stack. +# Do NOT run both stacks simultaneously — they are mutually exclusive. + +services: + + # Sidecar for Flock Dev app + flock-dapr: + image: "daprio/daprd:edge" + command: [ + "./daprd", + "--app-id", "flock-dev", + "--placement-host-address", "placement:50006", + "--scheduler-host-address", "scheduler:50007", + "--resources-path", "/components" + ] + volumes: + - "./components/:/components" + - "./secrets.json:/secrets/secrets.json:ro" + ports: + - "50001:50001" # Dapr gRPC port (used by DaprClient on the host) + - "3500:3500" # Dapr HTTP port + networks: + - dapr-net + depends_on: + placement: + condition: service_started + scheduler: + condition: service_started + postgres: + condition: service_healthy + + placement: + image: "daprio/placement" + command: ["./placement", "--port", "50006"] + ports: + - "50006:50006" + networks: + - dapr-net + + scheduler: + image: "daprio/scheduler" + command: ["./scheduler", "--port", "50007", "--etcd-data-dir", "/data"] + user: root + volumes: + - "./dapr-etcd-data/:/data" + networks: + - dapr-net + + postgres: + image: "postgres:17" + container_name: postgres + ports: + - "5432:5432" + environment: + POSTGRES_PASSWORD: "flock-postgres-dev-2026!" + POSTGRES_DB: "flock_dapr" + networks: + - dapr-net + volumes: + - "./postgres_data:/var/lib/postgresql/data" + healthcheck: + test: ["CMD-SHELL", "pg_isready -U postgres -d flock_dapr"] + interval: 2s + timeout: 5s + retries: 10 + start_period: 5s + +networks: + dapr-net: null \ No newline at end of file diff --git a/examples/12-dapr/postgresql_unencrypted/flock_dapr_postgresql.py b/examples/12-dapr/postgresql_unencrypted/flock_dapr_postgresql.py new file mode 100644 index 000000000..be8909747 --- /dev/null +++ b/examples/12-dapr/postgresql_unencrypted/flock_dapr_postgresql.py @@ -0,0 +1,178 @@ +import asyncio +import os + +from dapr.clients import DaprClient +from pydantic import BaseModel, Field + +from flock import Flock, PublicVisibility, flock_type +from flock.logging.logging import configure_logging, get_logger +from flock.storage import ( + DaprStateBlackboardConfig, + DaprStateBlackboardStore, + DaprStateBlackboardStoreClientConfig, +) + + +logger = get_logger(__name__) + + +@flock_type +class BandConcept(BaseModel): + genre: str = Field(description="Musical genre (rock, jazz, metal, pop, etc.)") + vibe: str = Field(description="The band's vibe or aesthetic") + target_audience: str = Field(description="Who should love this band?") + + +@flock_type +class BandLineup(BaseModel): + band_name: str = Field(description="Cool band name") + members: list[dict[str, str]] = Field( + description="List of band members with their roles" + ) + origin_story: str = Field(description="How the band formed", min_length=100) + signature_sound: str = Field(description="What makes their sound unique") + + +@flock_type +class Album(BaseModel): + title: str = Field(description="Album title in ALL CAPS") + tracklist: list[dict[str, str]] = Field( + description="Songs with titles and brief descriptions", + min_length=8, + max_length=12, + ) + genre_fusion: str = Field(description="How this album blends genres") + standout_track: str = Field(description="The track that'll be a hit") + production_notes: str = Field(description="Special production techniques") + + +@flock_type +class MarketingCopy(BaseModel): + press_release: str = Field( + description="Professional press release announcing the album", min_length=200 + ) + social_media_hook: str = Field( + description="Catchy social post (280 chars max)", max_length=280 + ) + billboard_tagline: str = Field( + description="10-word tagline for billboards", max_length=100 + ) + target_playlists: list[str] = Field( + description="Spotify/Apple Music playlists to pitch to", + min_length=3, + max_length=5, + ) + + +FLOCK_TTL = 300 +FLOCK_SECRET_STORE = "flock-dev-secretstore" +FLOCK_BASE_URL_SECRET_KEY = "base_url" +FLOCK_API_VERSION_SECRET_KEY = "api_version" +FLOCK_API_KEY_SECRET_KEY = "api_key" +FLOCK_STATE_STORE_SECRET_KEY = "state_store_name" +FLOCK_DEFAULT_MODEL_SECRET_KEY = "default_model" + + +async def full_blown_flock_test(): + """Test the blackboard with a small team of agents.""" + # Configure global log levels + configure_logging("ERROR", external_level="ERROR") + # Get Dapr secrets + base_url: str | None = None + api_version: str | None = None + state_store_name: str | None = None + api_key: str | None = None + default_model: str | None = None + with DaprClient() as client: + logger.info(f"Retrieving {FLOCK_BASE_URL_SECRET_KEY} secret") + base_url = client.get_secret( + store_name=FLOCK_SECRET_STORE, key=FLOCK_BASE_URL_SECRET_KEY + ).secret.get(FLOCK_BASE_URL_SECRET_KEY) + logger.info(f"Got {FLOCK_BASE_URL_SECRET_KEY}") + logger.info(f"Retrieving {FLOCK_API_VERSION_SECRET_KEY}") + api_version = client.get_secret( + store_name=FLOCK_SECRET_STORE, key=FLOCK_API_VERSION_SECRET_KEY + ).secret.get(FLOCK_API_VERSION_SECRET_KEY) + logger.info(f"Got {FLOCK_API_VERSION_SECRET_KEY}") + logger.info(f"Retrieving {FLOCK_API_KEY_SECRET_KEY}") + api_key = client.get_secret( + store_name=FLOCK_SECRET_STORE, key=FLOCK_API_KEY_SECRET_KEY + ).secret.get(FLOCK_API_KEY_SECRET_KEY) + logger.info(f"Got {FLOCK_API_KEY_SECRET_KEY}") + logger.info(f"Retrieving {FLOCK_STATE_STORE_SECRET_KEY}") + state_store_name = client.get_secret( + store_name=FLOCK_SECRET_STORE, key=FLOCK_STATE_STORE_SECRET_KEY + ).secret.get(FLOCK_STATE_STORE_SECRET_KEY) + logger.info(f"Retrieving {FLOCK_DEFAULT_MODEL_SECRET_KEY}") + default_model = client.get_secret( + store_name=FLOCK_SECRET_STORE, key=FLOCK_DEFAULT_MODEL_SECRET_KEY + ).secret.get(FLOCK_DEFAULT_MODEL_SECRET_KEY) + logger.info(f"Got {FLOCK_DEFAULT_MODEL_SECRET_KEY}") + # Check if all keys have been retrieved + # if one is missing, exit here and throw an exception + if any([ + base_url is None, + api_version is None, + state_store_name is None, + api_key is None, + default_model is None, + ]): + logger.error("UNABLE TO RETRIEVE FULL LIST OF SECRETS!!!") + raise ValueError + # Set required environment-variables before creating Flock-instance + os.environ["AZURE_API_BASE"] = base_url + os.environ["AZURE_API_VERSION"] = api_version + os.environ["AZURE_API_KEY"] = api_key + # Initialize dapr store. + client_config = DaprStateBlackboardStoreClientConfig() + store_config = DaprStateBlackboardConfig( + store_name=state_store_name, + supports_ttl=False, + encrypted_backend=False, + backend_encryption_key=None, + supports_transactions=True, + entries_ttl_seconds=None, + client_config=client_config, + supports_dapr_query_lang=False, + supports_etag=True, + etag_max_retries=5, + consistency="strong", + ) + dapr_store = DaprStateBlackboardStore(config=store_config) + # Initialize Flock Agent Swarm + flock = Flock( + model=default_model, + max_agent_iterations=100, + no_output=True, + store=dapr_store, # Add Dapr Blackboard Store as backend + ) + _ = ( + flock.agent("talent_scout") + .description("A legendary talent scout who assembles perfect band lineups") + .consumes(BandConcept) + .publishes(BandLineup, visibility=PublicVisibility()) + ) + _ = ( + flock.agent("music_producer") + .description("A visionary music producer who creates debut album concepts") + .consumes(BandLineup) + .publishes(Album, visibility=PublicVisibility()) + ) + _ = ( + flock.agent("marketing_guru") + .description("A marketing genius who writes compelling promotional material") + .consumes(Album) + .publishes(MarketingCopy, visibility=PublicVisibility()) + ) + + await flock.serve(dashboard=True) + + +async def main_test(): + print("Liftoff!") + print("=== Testing Flock with an Unencrypted PostgreSQL Backend ===") + await full_blown_flock_test() + + +if __name__ == "__main__": + asyncio.run(main_test()) diff --git a/examples/12-dapr/postgresql_unencrypted/secrets.example.json b/examples/12-dapr/postgresql_unencrypted/secrets.example.json new file mode 100644 index 000000000..82154ae48 --- /dev/null +++ b/examples/12-dapr/postgresql_unencrypted/secrets.example.json @@ -0,0 +1,8 @@ +{ + "postgresql_password": "", + "api_key": "", + "base_url": "", + "api_version": "", + "state_store_name": "", + "default_model": "" +} \ No newline at end of file diff --git a/examples/12-dapr/redis_encrypted/components/resiliency.yaml b/examples/12-dapr/redis_encrypted/components/resiliency.yaml new file mode 100644 index 000000000..f53827f02 --- /dev/null +++ b/examples/12-dapr/redis_encrypted/components/resiliency.yaml @@ -0,0 +1,27 @@ +apiVersion: dapr.io/v1alpha1 +kind: Resiliency +metadata: + name: flockdevresiliency +scopes: + - flock-dev + +spec: + policies: + retries: + retryForever: + policy: constant + duration: 5s + maxRetries: -1 + + circuitBreakers: + simpleCB: + maxRequests: 1 + timeout: 5s + trip: consecutiveFailures >= 5 + + targets: + components: + statestore: + outbound: + retry: retryForever + circuitBreaker: simpleCB diff --git a/examples/12-dapr/redis_encrypted/components/secretstore.yaml b/examples/12-dapr/redis_encrypted/components/secretstore.yaml new file mode 100644 index 000000000..da6fe319d --- /dev/null +++ b/examples/12-dapr/redis_encrypted/components/secretstore.yaml @@ -0,0 +1,15 @@ +apiVersion: dapr.io/v1alpha1 +kind: Component +metadata: + name: flock-dev-secretstore + namespace: development +spec: + type: secretstores.local.file + version: v1 + metadata: + - name: secretsFile + value: /secrets/secrets.json + - name: nestedSeparator + value: ":" + - name: multiValued + value: "false" diff --git a/examples/12-dapr/redis_encrypted/components/statestore.yaml b/examples/12-dapr/redis_encrypted/components/statestore.yaml new file mode 100644 index 000000000..d818844ae --- /dev/null +++ b/examples/12-dapr/redis_encrypted/components/statestore.yaml @@ -0,0 +1,27 @@ +apiVersion: dapr.io/v1alpha1 +kind: Component +metadata: + name: flockstate + namespace: development +spec: + type: state.redis + version: v1 + metadata: + - name: redisHost + value: "redis:6379" + - name: redisPassword + secretKeyRef: + name: redis_password + key: redis_password + - name: keyPrefix + value: namespace + - name: primaryEncryptionKey + secretKeyRef: + name: encryption_key + key: encryption_key + - name: secondaryEncryptionKey + secretKeyRef: + name: encryption_key_backup + key: encryption_key_backup +auth: + secretStore: flock-dev-secretstore \ No newline at end of file diff --git a/examples/12-dapr/redis_encrypted/docker-compose.yml b/examples/12-dapr/redis_encrypted/docker-compose.yml new file mode 100644 index 000000000..c00019580 --- /dev/null +++ b/examples/12-dapr/redis_encrypted/docker-compose.yml @@ -0,0 +1,56 @@ +services: + + # Sidecar for Flock Dev app + flock-dapr: + image: "daprio/daprd:edge" + command: [ + "./daprd", + "--app-id", "flock-dev", + "--placement-host-address", "placement:50006", + "--scheduler-host-address", "scheduler:50007", + "--resources-path", "/components" + ] + volumes: + - "./components/:/components" + - "./secrets.json:/secrets/secrets.json:ro" + ports: + - "50001:50001" # Dapr gRPC port (used by DaprClient on the host) + - "3500:3500" # Dapr HTTP port + networks: + - dapr-net + depends_on: + - placement + - scheduler + - redis + + placement: + image: "daprio/placement" + command: ["./placement", "--port", "50006"] + ports: + - "50006:50006" + networks: + - dapr-net + + scheduler: + image: "daprio/scheduler" + command: ["./scheduler", "--port", "50007", "--etcd-data-dir", "/data"] + user: root + volumes: + - "./dapr-etcd-data/:/data" + networks: + - dapr-net + + redis: + image: "redis/redis-stack" + container_name: redis + ports: + - "6379:6379" + networks: + - dapr-net + command: ["redis-server", "--loglevel", "warning", "--save", "20", "1", "--requirepass", "flock-redis-dev-2026!", "--loadmodule", "/opt/redis-stack/lib/redisearch.so", "--loadmodule", "/opt/redis-stack/lib/rejson.so"] + volumes: + - "./redis_data:/data" + + +networks: + dapr-net: null diff --git a/examples/12-dapr/redis_encrypted/flock_dapr_redis.py b/examples/12-dapr/redis_encrypted/flock_dapr_redis.py new file mode 100644 index 000000000..6e56a64ce --- /dev/null +++ b/examples/12-dapr/redis_encrypted/flock_dapr_redis.py @@ -0,0 +1,178 @@ +import asyncio +import os + +from dapr.clients import DaprClient +from pydantic import BaseModel, Field + +from flock import Flock, PublicVisibility, flock_type +from flock.logging.logging import configure_logging, get_logger +from flock.storage import ( + DaprStateBlackboardConfig, + DaprStateBlackboardStore, + DaprStateBlackboardStoreClientConfig, +) + + +logger = get_logger(__name__) + + +@flock_type +class BandConcept(BaseModel): + genre: str = Field(description="Musical genre (rock, jazz, metal, pop, etc.)") + vibe: str = Field(description="The band's vibe or aesthetic") + target_audience: str = Field(description="Who should love this band?") + + +@flock_type +class BandLineup(BaseModel): + band_name: str = Field(description="Cool band name") + members: list[dict[str, str]] = Field( + description="List of band members with their roles" + ) + origin_story: str = Field(description="How the band formed", min_length=100) + signature_sound: str = Field(description="What makes their sound unique") + + +@flock_type +class Album(BaseModel): + title: str = Field(description="Album title in ALL CAPS") + tracklist: list[dict[str, str]] = Field( + description="Songs with titles and brief descriptions", + min_length=8, + max_length=12, + ) + genre_fusion: str = Field(description="How this album blends genres") + standout_track: str = Field(description="The track that'll be a hit") + production_notes: str = Field(description="Special production techniques") + + +@flock_type +class MarketingCopy(BaseModel): + press_release: str = Field( + description="Professional press release announcing the album", min_length=200 + ) + social_media_hook: str = Field( + description="Catchy social post (280 chars max)", max_length=280 + ) + billboard_tagline: str = Field( + description="10-word tagline for billboards", max_length=100 + ) + target_playlists: list[str] = Field( + description="Spotify/Apple Music playlists to pitch to", + min_length=3, + max_length=5, + ) + + +FLOCK_TTL = 300 +FLOCK_SECRET_STORE = "flock-dev-secretstore" +FLOCK_BASE_URL_SECRET_KEY = "base_url" +FLOCK_API_VERSION_SECRET_KEY = "api_version" +FLOCK_API_KEY_SECRET_KEY = "api_key" +FLOCK_STATE_STORE_SECRET_KEY = "state_store_name" +FLOCK_DEFAULT_MODEL_SECRET_KEY = "default_model" + + +async def full_blown_flock_test(): + """Test the blackboard with a small team of agents.""" + # Configure global log levels + configure_logging("ERROR", external_level="ERROR") + # Get Dapr secrets + base_url: str | None = None + api_version: str | None = None + state_store_name: str | None = None + api_key: str | None = None + default_model: str | None = None + with DaprClient() as client: + logger.info(f"Retrieving {FLOCK_BASE_URL_SECRET_KEY} secret") + base_url = client.get_secret( + store_name=FLOCK_SECRET_STORE, key=FLOCK_BASE_URL_SECRET_KEY + ).secret.get(FLOCK_BASE_URL_SECRET_KEY) + logger.info(f"Got {FLOCK_BASE_URL_SECRET_KEY}") + logger.info(f"Retrieving {FLOCK_API_VERSION_SECRET_KEY}") + api_version = client.get_secret( + store_name=FLOCK_SECRET_STORE, key=FLOCK_API_VERSION_SECRET_KEY + ).secret.get(FLOCK_API_VERSION_SECRET_KEY) + logger.info(f"Got {FLOCK_API_VERSION_SECRET_KEY}") + logger.info(f"Retrieving {FLOCK_API_KEY_SECRET_KEY}") + api_key = client.get_secret( + store_name=FLOCK_SECRET_STORE, key=FLOCK_API_KEY_SECRET_KEY + ).secret.get(FLOCK_API_KEY_SECRET_KEY) + logger.info(f"Got {FLOCK_API_KEY_SECRET_KEY}") + logger.info(f"Retrieving {FLOCK_STATE_STORE_SECRET_KEY}") + state_store_name = client.get_secret( + store_name=FLOCK_SECRET_STORE, key=FLOCK_STATE_STORE_SECRET_KEY + ).secret.get(FLOCK_STATE_STORE_SECRET_KEY) + logger.info(f"Retrieving {FLOCK_DEFAULT_MODEL_SECRET_KEY}") + default_model = client.get_secret( + store_name=FLOCK_SECRET_STORE, key=FLOCK_DEFAULT_MODEL_SECRET_KEY + ).secret.get(FLOCK_DEFAULT_MODEL_SECRET_KEY) + logger.info(f"Got {FLOCK_DEFAULT_MODEL_SECRET_KEY}") + # Check if all keys have been retrieved + # if one is missing, exit here and throw an exception + if any([ + base_url is None, + api_version is None, + state_store_name is None, + api_key is None, + default_model is None, + ]): + logger.error("UNABLE TO RETRIEVE FULL LIST OF SECRETS!!!") + raise ValueError + # Set required environment-variables before creating Flock-instance + os.environ["AZURE_API_BASE"] = base_url + os.environ["AZURE_API_VERSION"] = api_version + os.environ["AZURE_API_KEY"] = api_key + # Initialize dapr store. + client_config = DaprStateBlackboardStoreClientConfig() + store_config = DaprStateBlackboardConfig( + store_name=state_store_name, + supports_ttl=False, + encrypted_backend=True, + backend_encryption_key=None, + supports_transactions=True, + entries_ttl_seconds=None, + client_config=client_config, + supports_dapr_query_lang=False, + supports_etag=True, + etag_max_retries=5, + consistency="strong", + ) + dapr_store = DaprStateBlackboardStore(config=store_config) + # Initialize Flock Agent Swarm + flock = Flock( + model=default_model, + max_agent_iterations=100, + no_output=True, + store=dapr_store, # Add Dapr Blackboard Store as backend + ) + _ = ( + flock.agent("talent_scout") + .description("A legendary talent scout who assembles perfect band lineups") + .consumes(BandConcept) + .publishes(BandLineup, visibility=PublicVisibility()) + ) + _ = ( + flock.agent("music_producer") + .description("A visionary music producer who creates debut album concepts") + .consumes(BandLineup) + .publishes(Album, visibility=PublicVisibility()) + ) + _ = ( + flock.agent("marketing_guru") + .description("A marketing genius who writes compelling promotional material") + .consumes(Album) + .publishes(MarketingCopy, visibility=PublicVisibility()) + ) + + await flock.serve(dashboard=True) + + +async def main_test(): + print("Liftoff!") + print("=== Testing Flock with an Encrypted Redis Backend ===") + await full_blown_flock_test() + + +if __name__ == "__main__": + asyncio.run(main_test()) diff --git a/examples/12-dapr/redis_encrypted/secrets.example.json b/examples/12-dapr/redis_encrypted/secrets.example.json new file mode 100644 index 000000000..eba12f9c2 --- /dev/null +++ b/examples/12-dapr/redis_encrypted/secrets.example.json @@ -0,0 +1,10 @@ +{ + "redis_password": "", + "encryption_key": "", + "encryption_key_backup": "", + "api_key": "", + "base_url": "", + "api_version": "", + "state_store_name": "", + "default_model": "" +} \ No newline at end of file diff --git a/mkdocs.yml b/mkdocs.yml index 300907490..ec098a46e 100644 --- a/mkdocs.yml +++ b/mkdocs.yml @@ -239,6 +239,7 @@ nav: - Persistent Blackboard: guides/persistent-blackboard.md - MCP Roots: guides/mcp-roots.md - Integrations: + - Dapr State Store: guides/dapr-state-store.md - OpenClaw: guides/openclaw.md - Operations: - Dashboard: guides/dashboard.md @@ -268,6 +269,7 @@ nav: - 04 — Misc: https://github.com/whiteducksoftware/flock/tree/main/examples/04-misc - 05 — Engines: https://github.com/whiteducksoftware/flock/tree/main/examples/05-engines - 09 — Scheduling: https://github.com/whiteducksoftware/flock/tree/main/examples/09-scheduling + - 12 — Dapr: https://github.com/whiteducksoftware/flock/tree/main/examples/12-dapr - ℹ️ About: - Roadmap: about/roadmap.md diff --git a/pyproject.toml b/pyproject.toml index b387d629b..f0908cf59 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -46,6 +46,9 @@ transformers = [ "accelerate>=1.12.0", "bitsandbytes>=0.48.2", ] +dapr = [ + "dapr>=1.13.0,<1.15" +] [[project.authors]] name = "Andre Ratzenberger" diff --git a/src/flock/components/agent/guard.py b/src/flock/components/agent/guard.py index 9ad9eed37..1897da9c6 100644 --- a/src/flock/components/agent/guard.py +++ b/src/flock/components/agent/guard.py @@ -232,7 +232,7 @@ def _handle_verdict( if action == "block": logger.warning(msg) raise GuardBlockedError(verdict) - elif action == "warn": + if action == "warn": logger.warning(msg) elif action == "annotate": logger.info(f"[annotate] {msg}") diff --git a/src/flock/core/orchestrator.py b/src/flock/core/orchestrator.py index 0b38e67fb..adb0689da 100644 --- a/src/flock/core/orchestrator.py +++ b/src/flock/core/orchestrator.py @@ -27,6 +27,7 @@ from flock.core.store import BlackboardStore, ConsumptionRecord from flock.core.subscription import Subscription from flock.core.visibility import PublicVisibility, Visibility +from flock.integrations.openclaw import OpenClawConfig, OpenClawEngine from flock.logging.auto_trace import AutoTracedMeta from flock.logging.logging import get_logger from flock.mcp import ( @@ -34,7 +35,6 @@ FlockMCPConfiguration, ServerParameters, ) -from flock.integrations.openclaw import OpenClawConfig, OpenClawEngine from flock.orchestrator import ( AgentScheduler, ArtifactManager, diff --git a/src/flock/core/store.py b/src/flock/core/store.py index 3cc44f5af..89f3d1d9c 100644 --- a/src/flock/core/store.py +++ b/src/flock/core/store.py @@ -24,7 +24,6 @@ from flock.core.artifacts import Artifact from flock.registry import type_registry -from flock.storage.artifact_aggregator import ArtifactAggregator from flock.utils.type_resolution import TypeResolutionHelper from flock.utils.visibility_utils import deserialize_visibility @@ -186,6 +185,8 @@ class InMemoryBlackboardStore(BlackboardStore): """Simple in-memory implementation suitable for local dev and tests.""" def __init__(self) -> None: + from flock.storage.artifact_aggregator import ArtifactAggregator + self._lock = Lock() self._by_id: dict[UUID, Artifact] = {} self._by_type: dict[str, list[Artifact]] = defaultdict(list) diff --git a/src/flock/integrations/openclaw/streaming.py b/src/flock/integrations/openclaw/streaming.py index 3127fd58c..181225295 100644 --- a/src/flock/integrations/openclaw/streaming.py +++ b/src/flock/integrations/openclaw/streaming.py @@ -7,10 +7,17 @@ from __future__ import annotations import json -from collections.abc import AsyncIterator, Iterable, Iterator, Sequence +from collections.abc import ( + AsyncIterator, + Awaitable, + Callable, + Iterable, + Iterator, + Sequence, +) from dataclasses import dataclass from types import SimpleNamespace -from typing import Any, Awaitable, Callable, Literal, Protocol +from typing import Any, Literal, Protocol import httpx @@ -88,7 +95,7 @@ def flush() -> SSEFrame | None: if not sep: continue - value = raw_value[1:] if raw_value.startswith(" ") else raw_value + value = raw_value.removeprefix(" ") if field == "event": event = value @@ -463,7 +470,7 @@ def flush() -> SSEFrame | None: if not sep: continue - value = raw_value[1:] if raw_value.startswith(" ") else raw_value + value = raw_value.removeprefix(" ") if field == "event": event = value diff --git a/src/flock/storage/__init__.py b/src/flock/storage/__init__.py index faa50a334..0001d4201 100644 --- a/src/flock/storage/__init__.py +++ b/src/flock/storage/__init__.py @@ -1,10 +1,111 @@ """Storage backends for Flock blackboard.""" -from flock.storage.sqlite.query_builder import SQLiteQueryBuilder -from flock.storage.sqlite.schema_manager import SQLiteSchemaManager +from importlib import import_module +from typing import TYPE_CHECKING, Any + + +if TYPE_CHECKING: + from flock.storage.dapr import ( + DaprStateBlackboardConfig, + DaprStateBlackboardStore, + DaprStateBlackboardStoreClientConfig, + create_dapr_client, + deserialize_agent_snapshot, + deserialize_artifact, + deserialize_consumption_records, + deserialize_index, + serialize_agent_snapshot, + serialize_artifact, + serialize_consumption_records, + serialize_index, + ) + from flock.storage.sqlite.query_builder import SQLiteQueryBuilder + from flock.storage.sqlite.schema_manager import SQLiteSchemaManager + + +_EXPORT_MAP: dict[str, tuple[str, str]] = { + "DaprStateBlackboardConfig": ( + "flock.storage.dapr", + "DaprStateBlackboardConfig", + ), + "DaprStateBlackboardStore": ( + "flock.storage.dapr", + "DaprStateBlackboardStore", + ), + "DaprStateBlackboardStoreClientConfig": ( + "flock.storage.dapr", + "DaprStateBlackboardStoreClientConfig", + ), + "create_dapr_client": ( + "flock.storage.dapr", + "create_dapr_client", + ), + "deserialize_agent_snapshot": ( + "flock.storage.dapr", + "deserialize_agent_snapshot", + ), + "deserialize_artifact": ( + "flock.storage.dapr", + "deserialize_artifact", + ), + "deserialize_consumption_records": ( + "flock.storage.dapr", + "deserialize_consumption_records", + ), + "deserialize_index": ( + "flock.storage.dapr", + "deserialize_index", + ), + "serialize_agent_snapshot": ( + "flock.storage.dapr", + "serialize_agent_snapshot", + ), + "serialize_artifact": ( + "flock.storage.dapr", + "serialize_artifact", + ), + "serialize_consumption_records": ( + "flock.storage.dapr", + "serialize_consumption_records", + ), + "serialize_index": ( + "flock.storage.dapr", + "serialize_index", + ), + "SQLiteQueryBuilder": ( + "flock.storage.sqlite.query_builder", + "SQLiteQueryBuilder", + ), + "SQLiteSchemaManager": ( + "flock.storage.sqlite.schema_manager", + "SQLiteSchemaManager", + ), +} + + +def __getattr__(name: str) -> Any: + target = _EXPORT_MAP.get(name) + if target is None: + raise AttributeError(f"module {__name__!r} has no attribute {name!r}") + module_name, symbol_name = target + value = getattr(import_module(module_name), symbol_name) + globals()[name] = value + return value __all__ = [ + "DaprStateBlackboardConfig", + "DaprStateBlackboardStore", + "DaprStateBlackboardStoreClientConfig", "SQLiteQueryBuilder", "SQLiteSchemaManager", + "create_dapr_client", + "deserialize_agent_snapshot", + "deserialize_artifact", + "deserialize_consumption_records", + "deserialize_index", + "serialize_agent_snapshot", + "serialize_artifact", + "serialize_consumption_records", + "serialize_index", ] diff --git a/src/flock/storage/dapr/README.md b/src/flock/storage/dapr/README.md new file mode 100644 index 000000000..f0e671917 --- /dev/null +++ b/src/flock/storage/dapr/README.md @@ -0,0 +1,10 @@ +# Dapr Storage Backend + +This directory contains the Dapr-backed blackboard store implementation used by Flock. + +Canonical documentation lives in: + +- Docs site guide: https://whiteducksoftware.github.io/flock/guides/dapr-state-store/ +- Repository example guide: ../../../../examples/12-dapr/README.md + +Use this README as a pointer only to avoid content drift between implementation docs and user-facing docs. diff --git a/src/flock/storage/dapr/__init__.py b/src/flock/storage/dapr/__init__.py new file mode 100644 index 000000000..5bddb3dd6 --- /dev/null +++ b/src/flock/storage/dapr/__init__.py @@ -0,0 +1,101 @@ +"""Dapr-Backed Blackboard Storage for Flock.""" + +from importlib import import_module +from typing import TYPE_CHECKING, Any + + +if TYPE_CHECKING: + from flock.storage.dapr._client import create_dapr_client + from flock.storage.dapr._serialization import ( + deserialize_agent_snapshot, + deserialize_artifact, + deserialize_consumption_records, + deserialize_index, + serialize_agent_snapshot, + serialize_artifact, + serialize_consumption_records, + serialize_index, + ) + from flock.storage.dapr.dapr_state_blackboard_store import ( + DaprStateBlackboardConfig, + DaprStateBlackboardStore, + DaprStateBlackboardStoreClientConfig, + ) + + +_EXPORT_MAP: dict[str, tuple[str, str]] = { + "DaprStateBlackboardConfig": ( + "flock.storage.dapr.dapr_state_blackboard_store", + "DaprStateBlackboardConfig", + ), + "DaprStateBlackboardStore": ( + "flock.storage.dapr.dapr_state_blackboard_store", + "DaprStateBlackboardStore", + ), + "DaprStateBlackboardStoreClientConfig": ( + "flock.storage.dapr.dapr_state_blackboard_store", + "DaprStateBlackboardStoreClientConfig", + ), + "create_dapr_client": ( + "flock.storage.dapr._client", + "create_dapr_client", + ), + "serialize_agent_snapshot": ( + "flock.storage.dapr._serialization", + "serialize_agent_snapshot", + ), + "serialize_artifact": ( + "flock.storage.dapr._serialization", + "serialize_artifact", + ), + "serialize_consumption_records": ( + "flock.storage.dapr._serialization", + "serialize_consumption_records", + ), + "serialize_index": ( + "flock.storage.dapr._serialization", + "serialize_index", + ), + "deserialize_agent_snapshot": ( + "flock.storage.dapr._serialization", + "deserialize_agent_snapshot", + ), + "deserialize_artifact": ( + "flock.storage.dapr._serialization", + "deserialize_artifact", + ), + "deserialize_consumption_records": ( + "flock.storage.dapr._serialization", + "deserialize_consumption_records", + ), + "deserialize_index": ( + "flock.storage.dapr._serialization", + "deserialize_index", + ), +} + + +def __getattr__(name: str) -> Any: + target = _EXPORT_MAP.get(name) + if target is None: + raise AttributeError(f"module {__name__!r} has no attribute {name!r}") + module_name, symbol_name = target + value = getattr(import_module(module_name), symbol_name) + globals()[name] = value + return value + + +__all__ = [ + "DaprStateBlackboardConfig", + "DaprStateBlackboardStore", + "DaprStateBlackboardStoreClientConfig", + "create_dapr_client", + "deserialize_agent_snapshot", + "deserialize_artifact", + "deserialize_consumption_records", + "deserialize_index", + "serialize_agent_snapshot", + "serialize_artifact", + "serialize_consumption_records", + "serialize_index", +] diff --git a/src/flock/storage/dapr/_client.py b/src/flock/storage/dapr/_client.py new file mode 100644 index 000000000..52105e460 --- /dev/null +++ b/src/flock/storage/dapr/_client.py @@ -0,0 +1,26 @@ +"""Thin wrapper around the Dapr gRPC client. + +Centralises client creation so that address override, retry policies, +and future migration to the async client happen in one place. +""" + +from __future__ import annotations + +from dapr.clients import DaprClient + + +def create_dapr_client(config) -> DaprClient: + """Return a configured :class:`DaprClient`. + + The caller is responsible for closing the client, either + via a context-manager (``with create_dapr_client() as c: ...``) or + by calling ``c.close()`` explicitly. + """ + return DaprClient( + address=config.dapr_grpc_endpoint, + headers_callback=config.header_callback, + interceptors=config.interceptors, + http_timeout_seconds=config.http_timeout_seconds, + max_grpc_message_length=config.max_grpc_message_length, + retry_policy=config.retry_policy, + ) diff --git a/src/flock/storage/dapr/_serialization.py b/src/flock/storage/dapr/_serialization.py new file mode 100644 index 000000000..6b4140c7c --- /dev/null +++ b/src/flock/storage/dapr/_serialization.py @@ -0,0 +1,131 @@ +"""JSON serialization helpers for Dapr state values. + +Artifact uses Pydantic's built-in serialization. The dataclass-based +types (ConsumptionRecord, AgentSnapshotRecord) need manual handling for +UUID and datetime fields. +""" + +from __future__ import annotations + +import json +from dataclasses import asdict +from datetime import datetime +from typing import TYPE_CHECKING, Any +from uuid import UUID + +from flock.core.artifacts import Artifact +from flock.core.visibility import ( + AfterVisibility, + LabelledVisibility, + PrivateVisibility, + PublicVisibility, + TenantVisibility, + Visibility, +) + + +if TYPE_CHECKING: + from flock.core.store import AgentSnapshotRecord, ConsumptionRecord + + +# -- Visibility discriminator map ------------------------------------------- + +_VISIBILITY_MAP: dict[str, type[Visibility]] = { + "Public": PublicVisibility, + "Private": PrivateVisibility, + "Labelled": LabelledVisibility, + "Tenant": TenantVisibility, + "After": AfterVisibility, +} + +# -- JSON encoder for dataclass types ---------------------------------------- + + +def _default(obj: Any) -> Any: + if isinstance(obj, UUID): + return str(obj) + if isinstance(obj, datetime): + return obj.isoformat() + if isinstance(obj, set): + return sorted(obj) + raise TypeError(f"Object of type {type(obj)} is not JSON serializable") + + +# -- Artifact ----------------------------------------------------------------- + + +def serialize_artifact(artifact: Artifact) -> str: + return artifact.model_dump_json() + + +def deserialize_artifact(data: str | bytes) -> Artifact: + artifact = Artifact.model_validate_json(data) + # Pydantic deserializes visibility as the base Visibility class, not the + # correct subclass, because Artifact.visibility is typed as ``Visibility`` + # (not a discriminated union). Re-parse the kind field to get the right + # subclass so that .allows() works. + vis = artifact.visibility + if type(vis) is Visibility: + cls = _VISIBILITY_MAP.get(vis.kind) + if cls is not None: + artifact.visibility = cls.model_validate(vis.model_dump()) + return artifact + + +# -- ConsumptionRecord ------------------------------------------------------- + + +def serialize_consumption_records(records: list[ConsumptionRecord]) -> str: + return json.dumps([asdict(r) for r in records], default=_default) + + +def deserialize_consumption_records(data: str | bytes) -> list[ConsumptionRecord]: + from flock.core.store import ConsumptionRecord + + items = json.loads(data) if data else [] + return [ + ConsumptionRecord( + artifact_id=UUID(item["artifact_id"]), + consumer=item["consumer"], + run_id=item.get("run_id"), + correlation_id=item.get("correlation_id"), + consumed_at=datetime.fromisoformat(item["consumed_at"]), + ) + for item in items + ] + + +# -- AgentSnapshotRecord ----------------------------------------------------- + + +def serialize_agent_snapshot(snapshot: AgentSnapshotRecord) -> str: + return json.dumps(asdict(snapshot), default=_default) + + +def deserialize_agent_snapshot(data: str | bytes) -> AgentSnapshotRecord: + from flock.core.store import AgentSnapshotRecord + + item = json.loads(data) if isinstance(data, (str, bytes)) else data + return AgentSnapshotRecord( + agent_name=item["agent_name"], + description=item["description"], + subscriptions=item["subscriptions"], + output_types=item["output_types"], + labels=item["labels"], + first_seen=datetime.fromisoformat(item["first_seen"]), + last_seen=datetime.fromisoformat(item["last_seen"]), + signature=item["signature"], + ) + + +# -- Index helpers (JSON lists of strings) ------------------------------------ + + +def serialize_index(keys: list[str]) -> str: + return json.dumps(keys) + + +def deserialize_index(data: str | bytes) -> list[str]: + if not data: + return [] + return json.loads(data) diff --git a/src/flock/storage/dapr/dapr_state_blackboard_store.py b/src/flock/storage/dapr/dapr_state_blackboard_store.py new file mode 100644 index 000000000..1f9b8c792 --- /dev/null +++ b/src/flock/storage/dapr/dapr_state_blackboard_store.py @@ -0,0 +1,1217 @@ +"""Dapr-Backed Blackboar Store. + +Utilizes Dapr State-Store Components +as the backend for the Flock Blackboard. +""" + +from __future__ import annotations + +import asyncio +import atexit +import json +from asyncio import Lock +from collections.abc import Callable, Iterable +from typing import TYPE_CHECKING, Any, Literal, TypeVar +from uuid import UUID + +from dapr.clients.grpc._request import ( + TransactionalStateOperation, + TransactionOperationType, +) +from dapr.clients.grpc._state import Concurrency, Consistency, StateItem, StateOptions +from grpc import ( + StatusCode, + StreamStreamClientInterceptor, + StreamUnaryClientInterceptor, + UnaryStreamClientInterceptor, + UnaryUnaryClientInterceptor, +) +from opentelemetry import trace +from pydantic import BaseModel, ConfigDict, Field, model_validator + +from flock.core.artifacts import Artifact +from flock.core.store import ( + AgentSnapshotRecord, + ArtifactEnvelope, + BlackboardStore, + ConsumptionRecord, + FilterConfig, +) +from flock.logging import get_logger +from flock.registry import type_registry +from flock.storage.artifact_aggregator import ArtifactAggregator +from flock.storage.dapr import ( + create_dapr_client, + deserialize_agent_snapshot, + deserialize_artifact, + deserialize_consumption_records, + deserialize_index, + serialize_agent_snapshot, + serialize_artifact, + serialize_consumption_records, + serialize_index, +) +from flock.storage.in_memory.artifact_filter import ArtifactFilter +from flock.storage.in_memory.history_aggregator import HistoryAggregator +from flock.utils.type_resolution import TypeResolutionHelper + + +if TYPE_CHECKING: + from dapr.clients.grpc._response import ( + BulkStatesResponse, + QueryResponse, + ) + from dapr.clients.retry import RetryPolicy + + +T = TypeVar("T") +tracer = trace.get_tracer(__name__) +logger = get_logger(__name__) +registry = type_registry + +# ── Key-naming conventions ─────────────────────────────────────────── +# +# artifact:{uuid} → serialised Artifact (JSON) +# idx:artifacts → JSON list of all artifact UUID strings +# idx:type:{type_name} → JSON list of UUID strings for that type +# consumptions:{artifact_id} → JSON list of ConsumptionRecord dicts +# snapshot:{agent_name} → serialised AgentSnapshotRecord (JSON) +# idx:snapshots → JSON list of agent name strings +# +# Indexes are maintained manually via read-modify-write. For stores +# that support transactions (Redis, PostgreSQL, CosmosDB) the artifact +# write and the index update should be wrapped in +# ``execute_state_transaction`` to guarantee atomicity. +# ───────────────────────────────────────────────────────────────────── + +_IDX_ALL_ARTIFACTS = "idx:artifacts" +_IDX_SNAPSHOTS = "idx:snapshots" + + +def _artifact_key(artifact_id: UUID) -> str: + return f"artifact:{artifact_id}" + + +def _type_index_key(type_name: str) -> str: + return f"idx:type:{type_name}" + + +def _consumptions_key(artifact_id: UUID) -> str: + return f"consumptions:{artifact_id}" + + +def _snapshot_key(agent_name: str) -> str: + return f"snapshot:{agent_name}" + + +def _build_dapr_query(filters: FilterConfig | None) -> str: + """Convert a :class:`FilterConfig` into a Dapr state query JSON string. + + Tags are excluded because the Dapr query language cannot express + set-membership checks against JSON arrays; they must be post-filtered. + """ + conditions: list[dict[str, Any]] = [] + if filters: + if filters.type_names: + names = sorted(filters.type_names) + if len(names) == 1: + conditions.append({"EQ": {"type": names[0]}}) + else: + conditions.append({"IN": {"type": names}}) + if filters.produced_by: + producers = sorted(filters.produced_by) + if len(producers) == 1: + conditions.append({"EQ": {"produced_by": producers[0]}}) + else: + conditions.append({"IN": {"produced_by": producers}}) + if filters.correlation_id: + conditions.append({"EQ": {"correlation_id": filters.correlation_id}}) + if filters.visibility: + kinds = sorted(filters.visibility) + if len(kinds) == 1: + conditions.append({"EQ": {"visibility.kind": kinds[0]}}) + else: + conditions.append({"IN": {"visibility.kind": kinds}}) + if filters.start: + conditions.append({"GTE": {"created_at": filters.start.isoformat()}}) + if filters.end: + conditions.append({"LTE": {"created_at": filters.end.isoformat()}}) + + query: dict[str, Any] = {} + if len(conditions) == 1: + query["filter"] = conditions[0] + elif len(conditions) > 1: + query["filter"] = {"AND": conditions} + query["sort"] = [{"key": "created_at", "order": "ASC"}] + return json.dumps(query) + + +class DaprStateBlackboardStoreClientConfig(BaseModel): + """Optional Config for the underlying Dapr-Client for the Blackboard.""" + + dapr_grpc_endpoint: str | None = Field( + default=None, description="Dapr Runtime gRPC endpoint address. Optional." + ) + headers_callback: Callable[[], dict[str, str]] | None = Field( + default=None, + description="lambda: dict[str, str]. Optional. Generate headers for each request.", + ) + interceptors: ( + list[ + UnaryUnaryClientInterceptor + | UnaryStreamClientInterceptor + | StreamUnaryClientInterceptor + | StreamStreamClientInterceptor + ] + | None + ) = Field(default=None, description="gRPC interceptors") + http_timeout_seconds: int | None = Field( + default=None, + description="Specify a timeout for http-connections to the dapr-backend.", + ) + max_grpc_message_length: int | None = Field( + default=None, + ) + + retry_policy: RetryPolicy | None = Field(default=None, description="Retry-Policy.") + + model_config = ConfigDict(arbitrary_types_allowed=True) + + +class DaprStateBlackboardConfig(BaseModel): + """Configuration for a Dapr-Compatible BlackboardStore.""" + + store_name: str = Field( + default="statestore", description="Name of the state-store for the blackboard." + ) + supports_ttl: bool = Field( + default=False, + description="Whether or not the underlying backend should implement TTL for entries. Only works if the backend really supports it! Otherwise, you will get errors", + ) + encrypted_backend: bool = Field( + default=False, description="Flag to indicate that the backend is encrypted." + ) + backend_encryption_key: str | None = Field( + default=None, + description="Optional. Encryption key for the backend. Only used if `encrypted_backend`==True.", + ) + supports_transactions: bool = Field( + default=False, + description="Whether or not the underlying backend is transactional. (strongly recommended. Examples are: Redis, CosmosDB, PostgresQL).", + ) + + supports_dapr_query_lang: bool = Field( + default=False, + description="If the configured backend supports dapr-queries. Optional.", + ) + + supports_etag: bool = Field( + default=False, + description="Enable optimistic concurrency control via Dapr ETags. " + "When True, all read-modify-write operations pass the ETag received " + "during the preceding read, using first-write-wins semantics. " + "ETag mismatches trigger automatic retries.", + ) + + etag_max_retries: int = Field( + default=3, + description="Maximum number of retries on ETag conflict when " + "supports_etag is True. Ignored when supports_etag is False.", + ge=0, + ) + + consistency: Literal["unspecified", "eventual", "strong"] = Field( + default="unspecified", + description="Consistency level for state operations. " + "Use Consistency.strong for strong consistency or " + "Consistency.eventual for eventual consistency. " + "Defaults to Consistency.unspecified (backend default).", + ) + + entries_ttl_seconds: int | None = Field( + default=None, + description="Optional TTL in Seconds for entries in the underlying state-store.", + ) + + client_config: DaprStateBlackboardStoreClientConfig | None = Field( + default=None, + description="Optional Dapr Client configuration. If not provided, the client will be created using default Dapr-Settings.", + ) + + model_config = ConfigDict(arbitrary_types_allowed=True) + + @model_validator(mode="after") + def _validate_config(self) -> DaprStateBlackboardConfig: + if self.entries_ttl_seconds is not None and not self.supports_ttl: + logger.warning( + "entries_ttl_seconds is set but supports_ttl is False; " + "TTL will not be applied to state entries." + ) + if self.encrypted_backend and self.supports_transactions: + logger.warning( + "Encrypted backend with transactions is not supported due to a " + "bug in the Dapr runtime's ExecuteStateTransaction handler " + "(values are corrupted via Go's fmt.Appendf before encryption). " + "Falling back to non-transactional operations." + ) + self.supports_transactions = False + return self + + +class DaprStateBlackboardStore(BlackboardStore): + """Dapr-backed implementation of :class:`BlackboardStore`. + + Leverages Dapr State Management to persist artifacts, consumption + records, and agent snapshots. Works with any Dapr state-store + component (Redis, PostgreSQL, CosmosDB, …). + + Args: + store_name: The Dapr component name (``metadata.name`` in the + component YAML). Defaults to ``"statestore"``. + + .. warning:: **Encryption and transactions are mutually exclusive.** + + When the Dapr state-store component is configured with + ``primaryEncryptionKey`` (i.e. ``encrypted_backend=True``), the + ``supports_transactions`` flag is automatically forced to + ``False``. This is caused by a bug in the Dapr Go runtime + (`ExecuteStateTransaction handler + `_): + before encrypting, it converts the raw ``[]byte`` value via + ``fmt.Appendf(nil, "%v", req.Value)`` which produces a + space-separated decimal representation (e.g. + ``[91 34 102 ...]``) instead of the original bytes. The + ``SaveState`` / ``SaveBulkState`` handlers do **not** have this + problem — they pass the raw bytes directly to encryption. + + As a consequence, all writes fall back to the non-transactional + path (``save_state`` / ``save_bulk_state``) when encryption is + enabled. Index updates are therefore **not** atomic in this + mode; however, stale index entries are reconciled lazily on read, + so data consistency is maintained eventually. + """ + + def __init__(self, config: DaprStateBlackboardConfig) -> None: + self._store_name = config.store_name + self._client = create_dapr_client(config.client_config) + self._lock = Lock() + self._supports_ttl = config.supports_ttl + self._entries_ttl = config.entries_ttl_seconds + self._encrypted_backend = config.encrypted_backend + self._aggregator = ArtifactAggregator() + self._history_aggregator = HistoryAggregator() + self._supports_transactions = config.supports_transactions + self._supports_dapr_query_lang = config.supports_dapr_query_lang + self._supports_etag = config.supports_etag + self._etag_max_retries = config.etag_max_retries + if config.consistency == "eventual": + self._consistency = Consistency.unspecified + elif config.consistency == "strong": + self._consistency = Consistency.strong + else: + self._consistency = Consistency.unspecified + # Register cleanup + atexit.register(self.close) + logger.info( + f"{__name__} initialized (store={config.store_name}, ttl={config.entries_ttl_seconds}, encrypted={config.encrypted_backend}, transactions={config.supports_transactions}, etag={config.supports_etag}, consistency={config.consistency})" + ) + + def close(self) -> None: + """Release the underlying gRPC channel.""" + logger.info("Closing Dapr-Client...") + self._client.close() + + # ── helpers ────────────────────────────────────────────────────── + def _create_state_metadata_for_save(self, key: str, *, ttl: int | None = None): + """Create the metadata dict for entries.""" + metadata = {} + if self._supports_ttl and ttl and ttl > 0: + metadata["ttlInSeconds"] = str(ttl) + logger.debug(f"Created metadata for entry {key}: {metadata}") + return metadata + + def _build_state_options(self) -> StateOptions | None: + """Build :class:`StateOptions` based on the current configuration. + + Returns ``StateOptions`` with ``concurrency=first_write`` when + ETags are enabled, or with the configured ``consistency`` when + it differs from ``unspecified``. Returns ``None`` when neither + is configured. + """ + concurrency = ( + Concurrency.first_write if self._supports_etag else Concurrency.unspecified + ) + if ( + concurrency == Concurrency.unspecified + and self._consistency == Consistency.unspecified + ): + return None + return StateOptions( + concurrency=concurrency, + consistency=self._consistency, + ) + + @staticmethod + def _is_etag_mismatch(err: Exception) -> bool: + """Return ``True`` if *err* represents a Dapr ETag mismatch.""" + if hasattr(err, "code") and callable(err.code): + code = err.code() + if code == StatusCode.ABORTED: + return True + if code == StatusCode.FAILED_PRECONDITION: + details = ( + err.details() + if hasattr(err, "details") and callable(err.details) + else "" + ) + if details and "etag" in details.lower(): + return True + return False + + async def _retry_on_etag_conflict( + self, + operation: Callable[[], None], + ) -> None: + """Execute *operation* with automatic retry on ETag conflicts. + + Re-raises the original exception after ``_etag_max_retries`` + consecutive failures. + """ + last_err: Exception | None = None + for attempt in range(1, self._etag_max_retries + 1): + try: + operation() + return + except Exception as exc: + if not self._is_etag_mismatch(exc): + raise + last_err = exc + delay = 0.1 * (2 ** (attempt - 1)) # 100ms, 200ms, 400ms … + logger.warning( + f"ETag conflict (attempt {attempt}/{self._etag_max_retries}), " + f"retrying in {delay:.1f}s" + ) + await asyncio.sleep(delay) + raise last_err # type: ignore[misc] + + def _read_index(self, key: str) -> tuple[list[str], str | None]: + """Read a JSON index list from the state store. + + Returns a tuple of ``(items, etag)`` where *etag* is the + state-store ETag when ``supports_etag`` is enabled, or + ``None`` otherwise. + """ + with tracer.start_as_current_span(f"{__name__}._read_index"): + logger.debug(f"reading index for: {key}") + resp = self._client.get_state(self._store_name, key) + etag = resp.etag if self._supports_etag else None + return deserialize_index(resp.text()), etag + + def _write_index( + self, + key: str, + items: list[str], + *, + etag: str | None = None, + ) -> None: + """Overwrite a JSON index list in the state store.""" + with tracer.start_as_current_span(f"{__name__}._write_index"): + logger.debug(f"overwriting json index list for: {key}...") + meta = self._create_state_metadata_for_save(key, ttl=None) + self._client.save_state( + self._store_name, + key, + serialize_index(items), + etag=etag if self._supports_etag else None, + options=self._build_state_options(), + state_metadata=meta, + ) + + def _reconcile_index( + self, + index_key: str, + index_ids: list[str], + live_keys: set[str], + *, + etag: str | None = None, + ) -> list[str]: + """Remove stale entries from an index and persist the cleaned version. + + Compares *index_ids* against *live_keys* (the set of keys whose + bulk-get returned non-empty data). Any id whose corresponding + data key is absent from *live_keys* is considered expired / + deleted and is pruned from the index. + + The cleaned index is written back only when at least one stale + entry is detected, avoiding unnecessary writes. + + Args: + index_key: The state-store key for the index + (e.g. ``idx:artifacts`` or ``idx:type:Foo``). + index_ids: The raw id list read from the index. + live_keys: Set of data-store keys that returned non-empty + data from the bulk-get. + etag: Optional ETag from the preceding index read. + + Returns: + The cleaned list of ids (same order, stale entries removed). + """ + with tracer.start_as_current_span(f"{__name__}._reconcile_index"): + cleaned = [uid for uid in index_ids if uid in live_keys] + stale_count = len(index_ids) - len(cleaned) + if stale_count > 0: + logger.info( + f"Reconciling index '{index_key}': removed {stale_count} stale entry/entries" + ) + self._write_index(index_key, cleaned, etag=etag) + return cleaned + + async def _query_via_dapr_api( + self, filters: FilterConfig | None = None + ) -> list[Artifact]: + """Fetch artifacts using the Dapr state query API (alpha). + + Handles token-based pagination internally so callers always + receive the full result set. Tags are post-filtered in-memory + because the Dapr query language cannot express set-membership + against JSON arrays. + """ + with tracer.start_as_current_span(f"{__name__}._query_via_dapr_api"): + query_json = _build_dapr_query(filters) + logger.debug(f"Dapr query: {query_json}") + artifacts: list[Artifact] = [] + token: str | None = None + while True: + # Inject pagination token for subsequent pages + if token: + query_dict = json.loads(query_json) + query_dict.setdefault("page", {})["token"] = token + current_query = json.dumps(query_dict) + else: + current_query = query_json + response: QueryResponse = self._client.query_state( + store_name=self._store_name, query=current_query + ) + for item in response.results: + if item.error: + logger.error(f"Query item error: {item.error}") + continue + try: + artifact = deserialize_artifact(item.text()) + except Exception: + logger.debug(f"Skipping non-artifact key: {item.key}") + continue + artifacts.append(artifact) + # Continue if more pages available + if response.token: + token = response.token + else: + break + # Post-filter tags (Dapr query cannot handle array membership) + if filters and filters.tags: + artifacts = [a for a in artifacts if filters.tags.issubset(a.tags)] + logger.debug( + f"Dapr query returned {len(artifacts)} artifact(s) after post-filtering" + ) + return artifacts + + def _query_via_index_scan( + self, filters: FilterConfig | None = None + ) -> list[Artifact]: + """Fetch artifacts via index read + bulk get + in-memory filtering. + + Used when the backend does not support the Dapr query API or + when the backend is encrypted (server-side queries are impossible + on encrypted state). + """ + with tracer.start_as_current_span(f"{__name__}._query_via_index_scan"): + filters = filters or FilterConfig() + # Optimisation: narrow the scan when type_names are provided + if filters.type_names: + uid_set: set[str] = set() + for type_name in filters.type_names: + ids, _etag = self._read_index(_type_index_key(type_name)) + uid_set.update(ids) + all_ids = list(uid_set) + else: + all_ids, _etag = self._read_index(_IDX_ALL_ARTIFACTS) + if not all_ids: + return [] + keys = [_artifact_key(UUID(uid)) for uid in all_ids] + items = self._client.get_bulk_state( + self._store_name, keys, parallelism=10 + ).items + # Build live-keys set and deserialize in one pass + live_keys: set[str] = set() + artifacts: list[Artifact] = [] + for item in items: + if item.error or not item.text(): + continue + try: + artifacts.append(deserialize_artifact(item.text())) + live_keys.add(item.key.replace("artifact:", "")) + except Exception: + logger.debug(f"Skipping non-artifact key: {item.key}") + continue + # Reconcile stale index entries + if filters.type_names: + for type_name in filters.type_names: + type_ids, type_etag = self._read_index(_type_index_key(type_name)) + self._reconcile_index( + _type_index_key(type_name), + type_ids, + live_keys, + etag=type_etag, + ) + else: + self._reconcile_index( + _IDX_ALL_ARTIFACTS, all_ids, live_keys, etag=_etag + ) + # Apply full in-memory filtering + artifact_filter = ArtifactFilter(filters) + artifacts = [a for a in artifacts if artifact_filter.matches(a)] + logger.debug( + f"Index scan returned {len(artifacts)} artifact(s) after filtering" + ) + return artifacts + + async def _query_backend( + self, filters: FilterConfig | None = None + ) -> list[Artifact]: + """Return matching artifacts from the backend. + + Dispatches to the Dapr query API when available, otherwise falls + back to a full index scan with in-memory filtering. + """ + with tracer.start_as_current_span(f"{__name__}._query_backend"): + if self._encrypted_backend or not self._supports_dapr_query_lang: + return self._query_via_index_scan(filters) + return await self._query_via_dapr_api(filters) + + async def _get_consumptions_by_artifact_ids( + self, artifact_ids: list[UUID] + ) -> dict[str, list[ConsumptionRecord]]: + """Get a list of consumption-records for a specific artifact by ID.""" + with tracer.start_as_current_span( + f"{__name__}._get_consumptions_by_artifact_ids" + ): + logger.debug(f"Fetching consumptions for {len(artifact_ids)} artifact(s)") + # Retrieve consumptions for each artifact_id + results: dict[str, list[ConsumptionRecord]] = {} + entry_ids: list[str] = [] + for id in artifact_ids: + entry_ids.append(_consumptions_key(id)) # noqa: PERF401 + retrieved_consumption_entries: BulkStatesResponse = ( + self._client.get_bulk_state( + store_name=self._store_name, keys=entry_ids, parallelism=10 + ) + ) + # pre-filter results to exclude errors + # and normalize results + for entry in retrieved_consumption_entries.items: + if entry.error: + logger.error( + f"Error when retrieving consumptions for entry: {entry.key}" + ) + continue + key = entry.key.replace("consumptions:", "") + data = entry.text() + # etag = entry.etag + # deserialize + consumptions_records = deserialize_consumption_records(data) + results[key] = consumptions_records + logger.debug( + f"Retrieved consumptions for {len(results)}/{len(artifact_ids)} artifact(s)" + ) + return results + + async def _publish_transactional(self, artifact: Artifact) -> None: + """Publish an artifact to the blackboard (transactional). + + The artifact data key is saved separately with TTL metadata + (when configured) because ``TransactionalStateOperation`` does + not support per-item metadata. Index updates remain atomic + inside the transaction and are never given a TTL. + """ + with tracer.start_as_current_span(f"{__name__}._publish_transactional"): + async with self._lock: + logger.debug( + f"Publishing artifact {artifact.id} (type={artifact.type}) [transactional]" + ) + key = _artifact_key(artifact.id) + uid = str(artifact.id) + # 1. Read current indexes + all_ids, all_etag = self._read_index(_IDX_ALL_ARTIFACTS) + type_idx, type_etag = self._read_index(_type_index_key(artifact.type)) + if uid not in all_ids: + all_ids.append(uid) + type_idx.append(uid) + # 2. Persist the artifact itself (with TTL when configured) + artifact_meta = self._create_state_metadata_for_save( + key, ttl=self._entries_ttl + ) + self._client.save_state( + self._store_name, + key, + serialize_artifact(artifact), + options=self._build_state_options(), + state_metadata=artifact_meta, + ) + # 3. Atomically update both indexes (no TTL) + artifact_idx_update_op = TransactionalStateOperation( + key=_IDX_ALL_ARTIFACTS, + data=serialize_index(all_ids), + etag=all_etag, + operation_type=TransactionOperationType.upsert, + ) + type_index_update_op = TransactionalStateOperation( + key=_type_index_key(artifact.type), + data=serialize_index(type_idx), + etag=type_etag, + operation_type=TransactionOperationType.upsert, + ) + _ = self._client.execute_state_transaction( + store_name=self._store_name, + operations=[artifact_idx_update_op, type_index_update_op], + ) + logger.info( + f"Published artifact {artifact.id} (type={artifact.type}) [transactional]" + ) + + async def _publish_non_transactional(self, artifact: Artifact) -> None: + """Publish an artifact to the blackboard in a non-transactional manner.""" + with tracer.start_as_current_span(f"{__name__}._publish_non_transactional"): + async with self._lock: + logger.debug( + f"Publishing artifact {artifact.id} (type={artifact.type}) [non-transactional]" + ) + key = _artifact_key(artifact.id) + uid = str(artifact.id) + meta = self._create_state_metadata_for_save(key, ttl=self._entries_ttl) + # 1. Persist the artifact itself (with TTL when configured) + self._client.save_state( + self._store_name, + key, + serialize_artifact(artifact), + options=self._build_state_options(), + state_metadata=meta, + ) + # 2. Append to the global artifact index + all_ids, all_etag = self._read_index(_IDX_ALL_ARTIFACTS) + if uid not in all_ids: + all_ids.append(uid) + self._write_index(_IDX_ALL_ARTIFACTS, all_ids, etag=all_etag) + # 3. Append to the per-type index + type_idx, type_etag = self._read_index(_type_index_key(artifact.type)) + if uid not in type_idx: + type_idx.append(uid) + self._write_index( + _type_index_key(artifact.type), type_idx, etag=type_etag + ) + logger.info( + f"Published artifact {artifact.id} (type={artifact.type}) [non-transactional]" + ) + + async def _record_consumptions_transactional( + self, records: Iterable[ConsumptionRecord] + ) -> None: + """Records the fact that an artifact has been consumed and by whom. (transactional). + + Note: TTL cannot be applied per-item inside a Dapr state + transaction (the Python SDK does not expose per-operation + metadata). Consumption records saved here will not carry TTL. + Orphaned records (whose parent artifact has expired) are + harmless — they are never referenced once the artifact's index + entry is reconciled away. + """ + with tracer.start_as_current_span( + f"{__name__}._record_consumptions_transactional" + ): + async with self._lock: + logger.debug("Recording consumptions [transactional]") + # Group incoming records by artifact-id + by_artifact: dict[UUID, list[ConsumptionRecord]] = {} + transaction_operations: list[TransactionalStateOperation] = [] + for rec in records: + by_artifact.setdefault(rec.artifact_id, []).append(rec) + for artifact_id, new_records in by_artifact.items(): + key = _consumptions_key(artifact_id) + # Read existing consumptions records for this artifact + resp = self._client.get_state(self._store_name, key) + cons_etag = resp.etag if self._supports_etag else None + existing = ( + deserialize_consumption_records(resp.text()) + if resp.text() + else [] + ) + existing.extend(new_records) + transaction_operations.append( + TransactionalStateOperation( + key=key, + operation_type=TransactionOperationType.upsert, + data=serialize_consumption_records(existing), + etag=cons_etag, + ) + ) + # Apply pending operations + _ = self._client.execute_state_transaction( + store_name=self._store_name, + operations=transaction_operations, + ) + logger.info( + f"Recorded consumptions for {len(by_artifact)} artifact(s) [transactional]" + ) + + async def _record_consumptions_non_transactional( + self, + records: Iterable[ConsumptionRecord], + ) -> None: + """Records the fact that an artifact has been consumed and by whom. (non-transactional).""" + with tracer.start_as_current_span( + f"{__name__}._record_consumptions_non_transactional" + ): + async with self._lock: + logger.debug("Recording consumptions [non-transactional]") + # Group incoming records by artifact_id. + by_artifact: dict[UUID, list[ConsumptionRecord]] = {} + state_items: list[StateItem] = [] + for rec in records: + by_artifact.setdefault(rec.artifact_id, []).append(rec) + for artifact_id, new_records in by_artifact.items(): + key = _consumptions_key(artifact_id) + # Read existing consumption records for this artifact. + resp = self._client.get_state(self._store_name, key) + cons_etag = resp.etag if self._supports_etag else None + existing = ( + deserialize_consumption_records(resp.text()) + if resp.text() + else [] + ) + existing.extend(new_records) + consumption_meta = self._create_state_metadata_for_save( + key, ttl=self._entries_ttl + ) + opts = self._build_state_options() + state_items.append( + StateItem( + key=key, + value=serialize_consumption_records(existing), + etag=cons_etag, + options=opts.get_proto() if opts else None, + metadata=consumption_meta, + ) + ) + self._client.save_bulk_state( + store_name=self._store_name, + states=state_items, + ) + logger.info( + f"Recorded consumptions for {len(by_artifact)} artifact(s) [non-transactional]" + ) + + async def _upsert_agent_snapshot_transactional( + self, snapshot: AgentSnapshotRecord + ) -> None: + """Upsert a Snapshot of an agent. (transactional).""" + with tracer.start_as_current_span( + f"{__name__}._upsert_agent_snapshot_transactional" + ): + async with self._lock: + logger.debug( + f"Upserting snapshot for agent '{snapshot.agent_name}' [transactional]" + ) + # 1. Read existing snapshot etag (for OCC) and build operations + snap_resp = self._client.get_state( + self._store_name, _snapshot_key(snapshot.agent_name) + ) + snap_etag = snap_resp.etag if self._supports_etag else None + all_ids, idx_etag = self._read_index(_IDX_SNAPSHOTS) + if snapshot.agent_name not in all_ids: + all_ids.append(snapshot.agent_name) + snapshot_save_op = TransactionalStateOperation( + key=_snapshot_key(snapshot.agent_name), + data=serialize_agent_snapshot(snapshot), + etag=snap_etag, + operation_type=TransactionOperationType.upsert, + ) + append_to_glob_idx_op = TransactionalStateOperation( + key=_IDX_SNAPSHOTS, + data=serialize_index(all_ids), + etag=idx_etag, + operation_type=TransactionOperationType.upsert, + ) + operations = [ + snapshot_save_op, + append_to_glob_idx_op, + ] + _ = self._client.execute_state_transaction( + store_name=self._store_name, operations=operations + ) + logger.info( + f"Upserted snapshot for agent '{snapshot.agent_name}' [transactional]" + ) + + async def _upsert_agent_snapshot_non_transactional( + self, snapshot: AgentSnapshotRecord + ) -> None: + """Upsert a Snapshot of an agent. (non-transactional).""" + with tracer.start_as_current_span( + f"{__name__}._upsert_agent_snapshot_non_transactional" + ): + async with self._lock: + logger.debug( + f"Upserting snapshot for agent '{snapshot.agent_name}' [non-transactional]" + ) + meta = self._create_state_metadata_for_save(snapshot.signature) + # 1. Read existing snapshot to capture etag (for OCC on upsert) + snap_etag: str | None = None + if self._supports_etag: + snap_resp = self._client.get_state( + self._store_name, _snapshot_key(snapshot.agent_name) + ) + snap_etag = snap_resp.etag or None + # 2. Persist the snapshot + self._client.save_state( + self._store_name, + key=_snapshot_key(snapshot.agent_name), + value=serialize_agent_snapshot(snapshot), + etag=snap_etag, + options=self._build_state_options(), + state_metadata=meta, + ) + # 2. Append to the global snapshot index + all_ids, idx_etag = self._read_index(_IDX_SNAPSHOTS) + if snapshot.agent_name not in all_ids: + all_ids.append(snapshot.agent_name) + self._write_index(_IDX_SNAPSHOTS, all_ids, etag=idx_etag) + logger.info( + f"Upserted snapshot for agent '{snapshot.agent_name}' [non-transactional]" + ) + + async def _clear_agent_snapshots_transactional(self) -> None: + """Clear out all agent-snapshots. (transactional).""" + with tracer.start_as_current_span( + f"{__name__}._clear_agent_snapshots_transactional" + ): + async with self._lock: + logger.debug("Clearing all agent snapshots [transactional]") + # 1. Get a list of all snapshot ids. + all_ids, idx_etag = self._read_index(_IDX_SNAPSHOTS) + # 2. Prepare transaction + operations: list[TransactionalStateOperation] = [] + for id in all_ids: # noqa: A001 + operation = TransactionalStateOperation( + operation_type=TransactionOperationType.delete, + key=_snapshot_key(id), + data=str(id), + ) + operations.append(operation) + # 3. Also clear the index itself + operations.append( + TransactionalStateOperation( + operation_type=TransactionOperationType.delete, + key=_IDX_SNAPSHOTS, + etag=idx_etag, + data="", + ) + ) + self._client.execute_state_transaction( + store_name=self._store_name, + operations=operations, + ) + logger.info(f"Cleared {len(all_ids)} agent snapshot(s) [transactional]") + + async def _clear_agent_snapshots_non_transactional(self) -> None: + """Clear out all agent-snapshots. (non-transactional).""" + with tracer.start_as_current_span( + f"{__name__}._clear_agent_snapshots_non_transactional" + ): + async with self._lock: + logger.debug("Clearing all agent snapshots [non-transactional]") + # 1. Get a list of all snapshot ids. + all_ids, idx_etag = self._read_index(_IDX_SNAPSHOTS) + logger.debug("Found %d snapshot(s) to delete", len(all_ids)) + # 2. Perform deletes + for id in all_ids: # noqa: A001 + logger.debug(f"Deleting agent snapshot '{id}'") + _ = self._client.delete_state( + self._store_name, + key=_snapshot_key(id), + options=self._build_state_options(), + ) + # 3. Clear the index + self._write_index(_IDX_SNAPSHOTS, [], etag=idx_etag) + logger.info( + f"Cleared {len(all_ids)} agent snapshot(s) [non-transactional]" + ) + + # ── core artifact operations ───────────────────────────────────── + + async def publish(self, artifact: Artifact) -> None: + """Publish an artifact to the blackboard.""" + with tracer.start_as_current_span(f"{__name__}.publish"): + logger.debug(f"Publishing artifact {artifact.id} (type={artifact.type})") + if self._supports_transactions: + await self._publish_transactional(artifact) + else: + await self._publish_non_transactional(artifact) + + async def get(self, artifact_id: UUID) -> Artifact | None: # type: ignore[override] + with tracer.start_as_current_span(f"{__name__}.get"): + async with self._lock: + logger.debug(f"Getting artifact {artifact_id}") + resp = self._client.get_state( + self._store_name, + _artifact_key(artifact_id), + ) + if not resp.text(): + logger.debug(f"Artifact {artifact_id} not found") + return None + logger.debug(f"Retrieved artifact {artifact_id}") + return deserialize_artifact(resp.text()) + + async def list(self) -> list[Artifact]: # type: ignore[override] + with tracer.start_as_current_span(f"{__name__}.list"): + async with self._lock: + logger.debug("Listing all artifacts") + all_ids, idx_etag = self._read_index(_IDX_ALL_ARTIFACTS) + if not all_ids: + logger.debug("No artifacts found") + return [] + keys = [_artifact_key(UUID(uid)) for uid in all_ids] + items = self._client.get_bulk_state( + self._store_name, + keys, + ).items + # Build a mapping of live keys for reconciliation + live_keys: set[str] = set() + artifacts: list[Artifact] = [] + for item in items: + if not item.text(): + continue + live_keys.add(item.key.replace("artifact:", "")) + artifacts.append(deserialize_artifact(item.text())) + # Reconcile stale index entries + self._reconcile_index( + _IDX_ALL_ARTIFACTS, all_ids, live_keys, etag=idx_etag + ) + logger.debug(f"Listed {len(artifacts)} artifact(s)") + return artifacts + + async def list_by_type(self, type_name: str) -> list[Artifact]: # type: ignore[override] + with tracer.start_as_current_span(f"{__name__}.list_by_type"): + async with self._lock: + logger.debug(f"Listing artifacts by type '{type_name}'") + uids, idx_etag = self._read_index(_type_index_key(type_name)) + if not uids: + logger.debug(f"No artifacts of type '{type_name}' found") + return [] + keys = [_artifact_key(UUID(uid)) for uid in uids] + items = self._client.get_bulk_state( + self._store_name, + keys, + ).items + # Build a mapping of live keys for reconciliation + live_keys: set[str] = set() + artifacts: list[Artifact] = [] + for item in items: + if not item.text(): + continue + live_keys.add(item.key.replace("artifact:", "")) + artifacts.append(deserialize_artifact(item.text())) + # Reconcile stale index entries + self._reconcile_index( + _type_index_key(type_name), uids, live_keys, etag=idx_etag + ) + logger.debug( + f"Listed {len(artifacts)} artifact(s) of type '{type_name}'" + ) + return artifacts + + async def get_by_type( + self, + artifact_type: type[T], + *, + correlation_id: str | None = None, + ) -> list[T]: # type: ignore[override] + with tracer.start_as_current_span(f"{__name__}.get_by_type"): + logger.debug( + f"Getting artifacts by type '{artifact_type.__name__}' (correlation_id={correlation_id})" + ) + artifacts: list[Artifact] = [] + type_name = TypeResolutionHelper.safe_resolve( + registry=registry, type_name=artifact_type.__name__ + ) + artifacts = await self.list_by_type(type_name) + if correlation_id is not None: + before = len(artifacts) + artifacts = [a for a in artifacts if a.correlation_id == correlation_id] + logger.debug( + f"Filtered by correlation_id '{correlation_id}': {before} -> {len(artifacts)} artifact(s)" + ) + logger.debug( + f"Returning {len(artifacts)} artifact(s) of type '{artifact_type.__name__}'" + ) + return [artifact_type(**a.payload) for a in artifacts] # type: ignore[return-value] + + # ── consumption records ────────────────────────────────────────── + + async def record_consumptions( + self, + records: Iterable[ConsumptionRecord], + ) -> None: # type: ignore[override] + """Records the fact that an artifact has been consumed and by whom.""" + with tracer.start_as_current_span(f"{__name__}.record_consumptions"): + logger.debug("Recording consumptions") + if self._supports_transactions: + await self._record_consumptions_transactional(records=records) + else: + await self._record_consumptions_non_transactional(records=records) + + # ── query / aggregation ────────────────────────────────────────── + + async def query_artifacts( + self, + filters: FilterConfig | None = None, + *, + limit: int = 50, + offset: int = 0, + embed_meta: bool = False, + ) -> tuple[list[Artifact | ArtifactEnvelope], int]: + with tracer.start_as_current_span(f"{__name__}.query_artifacts"): + logger.debug( + f"Querying artifacts (filters={filters}, limit={limit}, offset={offset}, embed_meta={embed_meta})" + ) + artifacts = await self._query_backend(filters) + # Sort consistently with InMemoryBlackboardStore + artifacts.sort(key=lambda a: (a.created_at, a.id)) + # Total before pagination + total = len(artifacts) + # Apply offset/limit pagination + offset = max(offset, 0) + if limit <= 0: + page = artifacts[offset:] + else: + page = artifacts[offset : offset + limit] + + if not embed_meta: + logger.debug(f"Query returned {len(page)} artifact(s) (total={total})") + return page, total + + artifact_ids = [a.id for a in page] + consumptions = await self._get_consumptions_by_artifact_ids( + artifact_ids=artifact_ids + ) + envelopes: list[ArtifactEnvelope] = [ + ArtifactEnvelope( + artifact=a, + consumptions=consumptions.get(str(a.id), []), + ) + for a in page + ] + logger.debug(f"Query returned {len(envelopes)} envelope(s) (total={total})") + return envelopes, total + + async def fetch_graph_artifacts( + self, + filters: FilterConfig | None = None, + *, + limit: int = 500, + offset: int = 0, + ) -> tuple[list[ArtifactEnvelope], int]: + """Return artifact envelopes (artifact + consumptions) for graph assembly.""" + logger.debug(f"Fetching graph artifacts (limit={limit}, offset={offset})") + artifacts, total = await self.query_artifacts( + filters=filters, limit=limit, offset=offset, embed_meta=True + ) + envelopes: list[ArtifactEnvelope] = [] + for item in artifacts: + if isinstance(item, ArtifactEnvelope): + envelopes.append(item) + elif isinstance(item, Artifact): + envelopes.append(ArtifactEnvelope(artifact=item)) + logger.debug(f"Fetched {len(envelopes)} graph envelope(s)") + return envelopes, total + + async def summarize_artifacts( + self, + filters: FilterConfig | None = None, + ) -> dict[str, Any]: + """Return aggregate artifact statistics for the given filters.""" + logger.debug(f"Summarizing artifacts (filters={filters})") + filters = filters or FilterConfig() + artifacts, total = await self.query_artifacts( + filters=filters, limit=0, offset=0, embed_meta=False + ) + for artifact in artifacts: + if not isinstance(artifact, Artifact): + raise TypeError("Expected Artifact instance") + # Delegate to aggregator for aggregation logic + is_full_window = filters.start is None and filters.end is None + return self._aggregator.build_summary(artifacts, total, is_full_window) + + async def agent_history_summary( + self, + agent_id: str, + filters: FilterConfig | None = None, + ) -> dict[str, Any]: + """Summarize agent history using history aggregator.""" + logger.debug(f"Summarizing agent history for '{agent_id}' (filters={filters})") + filters = filters or FilterConfig() + envelopes, _ = await self.query_artifacts( + filters=filters, + limit=0, + offset=0, + embed_meta=True, + ) + # Delegate to history aggregator for aggregation logic + return self._history_aggregator.aggregate(envelopes, agent_id) + + # ── agent snapshots ────────────────────────────────────────────── + + async def upsert_agent_snapshot(self, snapshot: AgentSnapshotRecord) -> None: + with tracer.start_as_current_span(f"{__name__}.upsert_agent_snapshot"): + logger.debug(f"Upserting agent snapshot for '{snapshot.agent_name}'") + if self._supports_transactions: + await self._upsert_agent_snapshot_transactional(snapshot) + else: + await self._upsert_agent_snapshot_non_transactional(snapshot) + + async def load_agent_snapshots(self) -> list[AgentSnapshotRecord]: + with tracer.start_as_current_span(f"{__name__}.load_agent_snapshots"): + async with self._lock: + logger.debug("Loading all agent snapshots") + snapshot_records: list[AgentSnapshotRecord] = [] + # 1. Get a list of all snapshot_ids + all_ids, idx_etag = self._read_index(_IDX_SNAPSHOTS) + if not all_ids: + return [] + # 2. Do a bulk read + keys = [_snapshot_key(name) for name in all_ids] + result: BulkStatesResponse = self._client.get_bulk_state( + self._store_name, keys=keys, parallelism=10 + ) + live_keys: set[str] = set() + for item in result.items: + if item.error: + logger.error(f"Error retrieving snapshot: {item.error}") + continue + if not item.text(): + logger.debug(f"Skipping empty snapshot entry: {item.key}") + continue + live_keys.add(item.key.replace("snapshot:", "")) + deserialized = deserialize_agent_snapshot(item.text()) + snapshot_records.append(deserialized) + # Reconcile stale snapshot index entries + self._reconcile_index(_IDX_SNAPSHOTS, all_ids, live_keys, etag=idx_etag) + logger.debug(f"Loaded {len(snapshot_records)} agent snapshot(s)") + return snapshot_records + + async def clear_agent_snapshots(self) -> None: + """Clear out all agent-snapshots.""" + with tracer.start_as_current_span(f"{__name__}.clear_agent_snapshots"): + logger.debug("Clearing all agent snapshots") + if self._supports_transactions: + await self._clear_agent_snapshots_transactional() + else: + await self._clear_agent_snapshots_non_transactional() diff --git a/tests/core/test_conditions_base.py b/tests/core/test_conditions_base.py index 4c4f8f559..506403b8f 100644 --- a/tests/core/test_conditions_base.py +++ b/tests/core/test_conditions_base.py @@ -8,10 +8,13 @@ from dataclasses import dataclass from typing import TYPE_CHECKING -from unittest.mock import AsyncMock, Mock +from unittest.mock import Mock import pytest +from flock.core.conditions import AndCondition + + if TYPE_CHECKING: from flock.core import Flock diff --git a/tests/storage/dapr/test_client.py b/tests/storage/dapr/test_client.py new file mode 100644 index 000000000..6e9c18819 --- /dev/null +++ b/tests/storage/dapr/test_client.py @@ -0,0 +1,58 @@ +from __future__ import annotations + +from types import SimpleNamespace + +import pytest + +from flock.storage.dapr._client import create_dapr_client + + +class _FakeDaprClient: + def __init__(self, **kwargs): + self.kwargs = kwargs + + +def test_create_dapr_client_passes_expected_arguments(monkeypatch) -> None: + captured: dict[str, object] = {} + + def _factory(**kwargs): + captured.update(kwargs) + return _FakeDaprClient(**kwargs) + + monkeypatch.setattr("flock.storage.dapr._client.DaprClient", _factory) + header_callback = lambda: {"x": "1"} + config = SimpleNamespace( + dapr_grpc_endpoint="localhost:50001", + header_callback=header_callback, + interceptors=["i1"], + http_timeout_seconds=10, + max_grpc_message_length=1024, + retry_policy="retry", + ) + + client = create_dapr_client(config) + + assert isinstance(client, _FakeDaprClient) + assert captured == { + "address": "localhost:50001", + "headers_callback": header_callback, + "interceptors": ["i1"], + "http_timeout_seconds": 10, + "max_grpc_message_length": 1024, + "retry_policy": "retry", + } + + +def test_create_dapr_client_requires_header_callback_attribute(monkeypatch) -> None: + monkeypatch.setattr("flock.storage.dapr._client.DaprClient", _FakeDaprClient) + config = SimpleNamespace( + dapr_grpc_endpoint="localhost:50001", + headers_callback=lambda: {"x": "1"}, + interceptors=None, + http_timeout_seconds=None, + max_grpc_message_length=None, + retry_policy=None, + ) + + with pytest.raises(AttributeError): + create_dapr_client(config) diff --git a/tests/storage/dapr/test_dapr_state_blackboard_store_helpers.py b/tests/storage/dapr/test_dapr_state_blackboard_store_helpers.py new file mode 100644 index 000000000..b6cb6139f --- /dev/null +++ b/tests/storage/dapr/test_dapr_state_blackboard_store_helpers.py @@ -0,0 +1,1328 @@ +from __future__ import annotations + +import json +from datetime import UTC, datetime +from types import SimpleNamespace +from uuid import uuid4 + +from dapr.clients.retry import RetryPolicy +from grpc import StatusCode +from pydantic import BaseModel + +from flock.core.artifacts import Artifact +from flock.core.store import ( + AgentSnapshotRecord, + ArtifactEnvelope, + ConsumptionRecord, + FilterConfig, +) +from flock.storage.dapr._serialization import serialize_agent_snapshot +from flock.storage.dapr.dapr_state_blackboard_store import ( + DaprStateBlackboardConfig, + DaprStateBlackboardStoreClientConfig, + DaprStateBlackboardStore, + _artifact_key, + _build_dapr_query, + _consumptions_key, + _snapshot_key, + _type_index_key, +) + + +DaprStateBlackboardStoreClientConfig.model_rebuild( + _types_namespace={"RetryPolicy": RetryPolicy} +) +DaprStateBlackboardConfig.model_rebuild( + _types_namespace={ + "DaprStateBlackboardStoreClientConfig": DaprStateBlackboardStoreClientConfig, + } +) + + +class _DummyClient: + def __init__(self) -> None: + self.closed = False + + def close(self) -> None: + self.closed = True + + +class _FakeError(Exception): + def __init__(self, code: StatusCode, details: str = "") -> None: + self._code = code + self._details = details + + def code(self) -> StatusCode: + return self._code + + def details(self) -> str: + return self._details + + +class _FakeStateResponse: + def __init__(self, data: str = "", etag: str | None = None) -> None: + self._data = data + self.etag = etag + + def text(self) -> str: + return self._data + + +class _FakeBulkItem: + def __init__(self, key: str, data: str = "", error: str | None = None) -> None: + self.key = key + self._data = data + self.error = error + + def text(self) -> str: + return self._data + + +def _make_store( + monkeypatch, **config_overrides +) -> tuple[DaprStateBlackboardStore, _DummyClient]: + dummy_client = _DummyClient() + monkeypatch.setattr( + "flock.storage.dapr.dapr_state_blackboard_store.create_dapr_client", + lambda _cfg: dummy_client, + ) + monkeypatch.setattr( + "flock.storage.dapr.dapr_state_blackboard_store.atexit.register", + lambda *_: None, + ) + + config = DaprStateBlackboardConfig( + store_name="test-store", + **config_overrides, + ) + return DaprStateBlackboardStore(config), dummy_client + + +def test_key_helpers_build_expected_prefixes() -> None: + artifact_id = uuid4() + + assert _artifact_key(artifact_id) == f"artifact:{artifact_id}" + assert _type_index_key("demo.Type") == "idx:type:demo.Type" + assert _consumptions_key(artifact_id) == f"consumptions:{artifact_id}" + assert _snapshot_key("agent-a") == "snapshot:agent-a" + + +def test_build_dapr_query_without_filters_has_only_sort() -> None: + query = json.loads(_build_dapr_query(None)) + + assert query == {"sort": [{"key": "created_at", "order": "ASC"}]} + + +def test_build_dapr_query_with_single_and_multiple_filters() -> None: + single = json.loads(_build_dapr_query(FilterConfig(type_names={"A"}, tags={"x"}))) + assert single["filter"] == {"EQ": {"type": "A"}} + assert single["sort"] == [{"key": "created_at", "order": "ASC"}] + + filters = FilterConfig( + type_names={"B", "A"}, + produced_by={"writer", "reviewer"}, + correlation_id="corr-1", + visibility={"Public", "Private"}, + start=datetime(2026, 1, 1, tzinfo=UTC), + end=datetime(2026, 1, 31, tzinfo=UTC), + tags={"ignored-in-query"}, + ) + multiple = json.loads(_build_dapr_query(filters)) + + assert "AND" in multiple["filter"] + and_conditions = multiple["filter"]["AND"] + assert {"IN": {"type": ["A", "B"]}} in and_conditions + assert {"IN": {"produced_by": ["reviewer", "writer"]}} in and_conditions + assert {"EQ": {"correlation_id": "corr-1"}} in and_conditions + assert {"IN": {"visibility.kind": ["Private", "Public"]}} in and_conditions + assert {"GTE": {"created_at": "2026-01-01T00:00:00+00:00"}} in and_conditions + assert {"LTE": {"created_at": "2026-01-31T00:00:00+00:00"}} in and_conditions + assert "tags" not in json.dumps(multiple) + + +def test_config_validator_warns_for_ttl_without_support(capsys) -> None: + _ = DaprStateBlackboardConfig(supports_ttl=False, entries_ttl_seconds=30) + + assert ( + "entries_ttl_seconds is set but supports_ttl is False" + in capsys.readouterr().out + ) + + +def test_config_validator_disables_transactions_for_encrypted_backend(capsys) -> None: + config = DaprStateBlackboardConfig( + encrypted_backend=True, + supports_transactions=True, + ) + + assert config.supports_transactions is False + assert ( + "Encrypted backend with transactions is not supported" + in capsys.readouterr().out + ) + + +def test_store_metadata_and_close(monkeypatch) -> None: + store, dummy_client = _make_store( + monkeypatch, + supports_ttl=True, + entries_ttl_seconds=60, + ) + + assert store._create_state_metadata_for_save("k", ttl=60) == {"ttlInSeconds": "60"} + assert store._create_state_metadata_for_save("k", ttl=0) == {} + assert store._create_state_metadata_for_save("k", ttl=None) == {} + + store.close() + assert dummy_client.closed + + +def test_build_state_options_respects_etag_and_consistency(monkeypatch) -> None: + etag_store, _ = _make_store(monkeypatch, supports_etag=True) + strong_store, _ = _make_store(monkeypatch, consistency="strong") + default_store, _ = _make_store(monkeypatch) + + etag_options = etag_store._build_state_options() + strong_options = strong_store._build_state_options() + + assert etag_options is not None + assert etag_options.concurrency.name == "first_write" + assert strong_options is not None + assert strong_options.concurrency.name == "unspecified" + assert strong_options.consistency.name == "strong" + assert default_store._build_state_options() is None + + +def test_is_etag_mismatch_handles_expected_grpc_variants() -> None: + assert DaprStateBlackboardStore._is_etag_mismatch(_FakeError(StatusCode.ABORTED)) + assert DaprStateBlackboardStore._is_etag_mismatch( + _FakeError(StatusCode.FAILED_PRECONDITION, details="ETag conflict") + ) + assert not DaprStateBlackboardStore._is_etag_mismatch( + _FakeError(StatusCode.FAILED_PRECONDITION, details="other") + ) + assert not DaprStateBlackboardStore._is_etag_mismatch(RuntimeError("boom")) + + +async def test_retry_on_etag_conflict_retries_then_succeeds(monkeypatch) -> None: + store, _ = _make_store(monkeypatch, supports_etag=True, etag_max_retries=3) + + attempts = {"count": 0} + + def _operation() -> None: + attempts["count"] += 1 + if attempts["count"] < 3: + raise _FakeError(StatusCode.ABORTED) + + sleeps: list[float] = [] + + async def _fake_sleep(delay: float) -> None: + sleeps.append(delay) + + monkeypatch.setattr( + "flock.storage.dapr.dapr_state_blackboard_store.asyncio.sleep", _fake_sleep + ) + + await store._retry_on_etag_conflict(_operation) + + assert attempts["count"] == 3 + assert sleeps == [0.1, 0.2] + + +async def test_retry_on_etag_conflict_raises_after_max_retries(monkeypatch) -> None: + store, _ = _make_store(monkeypatch, supports_etag=True, etag_max_retries=2) + + def _operation() -> None: + raise _FakeError(StatusCode.ABORTED) + + async def _fake_sleep(_delay: float) -> None: + return None + + monkeypatch.setattr( + "flock.storage.dapr.dapr_state_blackboard_store.asyncio.sleep", _fake_sleep + ) + + try: + await store._retry_on_etag_conflict(_operation) + except _FakeError as err: + assert err.code() == StatusCode.ABORTED + else: + raise AssertionError("Expected _FakeError to be raised") + + +def test_exported_package_lazy_getattr_and_missing_symbol() -> None: + import flock.storage.dapr as dapr_pkg + + assert dapr_pkg.serialize_index(["a"]) == '["a"]' + + try: + _ = dapr_pkg.DOES_NOT_EXIST + except AttributeError as err: + assert "has no attribute" in str(err) + else: + raise AssertionError("Expected AttributeError for unknown export") + + +async def test_publish_dispatches_by_transaction_flag(monkeypatch) -> None: + tx_store, _ = _make_store(monkeypatch, supports_transactions=True) + non_tx_store, _ = _make_store(monkeypatch, supports_transactions=False) + + calls: list[str] = [] + + async def _tx(_artifact: Artifact) -> None: + calls.append("tx") + + async def _non_tx(_artifact: Artifact) -> None: + calls.append("non-tx") + + monkeypatch.setattr(tx_store, "_publish_transactional", _tx) + monkeypatch.setattr(tx_store, "_publish_non_transactional", _non_tx) + monkeypatch.setattr(non_tx_store, "_publish_transactional", _tx) + monkeypatch.setattr(non_tx_store, "_publish_non_transactional", _non_tx) + + artifact = Artifact(type="demo.Type", payload={"v": 1}, produced_by="agent") + await tx_store.publish(artifact) + await non_tx_store.publish(artifact) + + assert calls == ["tx", "non-tx"] + + +async def test_query_backend_dispatches_to_expected_strategy(monkeypatch) -> None: + query_store, _ = _make_store(monkeypatch, supports_dapr_query_lang=True) + scan_store, _ = _make_store(monkeypatch, supports_dapr_query_lang=False) + encrypted_store, _ = _make_store( + monkeypatch, + supports_dapr_query_lang=True, + encrypted_backend=True, + ) + + async def _query(_filters): + return [Artifact(type="from-query", payload={}, produced_by="q")] + + def _scan(_filters): + return [Artifact(type="from-scan", payload={}, produced_by="s")] + + monkeypatch.setattr(query_store, "_query_via_dapr_api", _query) + monkeypatch.setattr(query_store, "_query_via_index_scan", _scan) + monkeypatch.setattr(scan_store, "_query_via_dapr_api", _query) + monkeypatch.setattr(scan_store, "_query_via_index_scan", _scan) + monkeypatch.setattr(encrypted_store, "_query_via_dapr_api", _query) + monkeypatch.setattr(encrypted_store, "_query_via_index_scan", _scan) + + assert (await query_store._query_backend())[0].type == "from-query" + assert (await scan_store._query_backend())[0].type == "from-scan" + assert (await encrypted_store._query_backend())[0].type == "from-scan" + + +async def test_query_artifacts_paginates_and_embeds_meta(monkeypatch) -> None: + store, _ = _make_store(monkeypatch) + artifacts = [ + Artifact( + type="T", + payload={"i": 2}, + produced_by="a", + created_at=datetime(2026, 1, 2, tzinfo=UTC), + ), + Artifact( + type="T", + payload={"i": 1}, + produced_by="a", + created_at=datetime(2026, 1, 1, tzinfo=UTC), + ), + ] + first_page_id = artifacts[1].id + rec = ConsumptionRecord( + artifact_id=first_page_id, + consumer="consumer", + consumed_at=datetime(2026, 1, 5, tzinfo=UTC), + ) + + async def _backend(_filters): + return artifacts + + async def _consumptions(*, artifact_ids): + assert artifact_ids + return {str(artifact_ids[0]): [rec]} + + monkeypatch.setattr(store, "_query_backend", _backend) + monkeypatch.setattr(store, "_get_consumptions_by_artifact_ids", _consumptions) + + page, total = await store.query_artifacts(limit=1, offset=0, embed_meta=True) + + assert total == 2 + assert len(page) == 1 + assert isinstance(page[0], ArtifactEnvelope) + assert page[0].artifact.payload["i"] == 1 + assert page[0].consumptions == [rec] + + +async def test_fetch_graph_artifacts_wraps_plain_artifacts(monkeypatch) -> None: + store, _ = _make_store(monkeypatch) + artifact = Artifact(type="T", payload={"x": 1}, produced_by="agent") + + async def _query(**_kwargs): + return ([artifact], 1) + + monkeypatch.setattr(store, "query_artifacts", _query) + + envelopes, total = await store.fetch_graph_artifacts() + + assert total == 1 + assert len(envelopes) == 1 + assert envelopes[0].artifact == artifact + + +async def test_summarize_artifacts_delegates_to_aggregator(monkeypatch) -> None: + store, _ = _make_store(monkeypatch) + artifact = Artifact(type="T", payload={"x": 1}, produced_by="agent") + + async def _query(**_kwargs): + return ([artifact], 1) + + monkeypatch.setattr(store, "query_artifacts", _query) + monkeypatch.setattr( + store._aggregator, + "build_summary", + lambda artifacts, total, full_window: { + "total": total, + "full_window": full_window, + "types": [a.type for a in artifacts], + }, + ) + + result = await store.summarize_artifacts() + + assert result == {"total": 1, "full_window": True, "types": ["T"]} + + +async def test_agent_history_summary_delegates_to_history_aggregator( + monkeypatch, +) -> None: + store, _ = _make_store(monkeypatch) + envelope = ArtifactEnvelope( + artifact=Artifact(type="T", payload={}, produced_by="agent"), + ) + + async def _query(**_kwargs): + return ([envelope], 1) + + monkeypatch.setattr(store, "query_artifacts", _query) + monkeypatch.setattr( + store._history_aggregator, + "aggregate", + lambda envelopes, agent_id: { + "agent_id": agent_id, + "count": len(envelopes), + }, + ) + + result = await store.agent_history_summary("agent-1") + + assert result == {"agent_id": "agent-1", "count": 1} + + +async def test_get_returns_none_for_missing_artifact(monkeypatch) -> None: + store, _ = _make_store(monkeypatch) + missing_id = uuid4() + monkeypatch.setattr( + store._client, + "get_state", + lambda _store, _key: _FakeStateResponse(""), + raising=False, + ) + + result = await store.get(missing_id) + + assert result is None + + +async def test_get_returns_deserialized_artifact(monkeypatch) -> None: + store, _ = _make_store(monkeypatch) + artifact = Artifact(type="T", payload={"a": 1}, produced_by="agent") + monkeypatch.setattr( + store._client, + "get_state", + lambda _store, _key: _FakeStateResponse(artifact.model_dump_json()), + raising=False, + ) + + result = await store.get(artifact.id) + + assert result is not None + assert result.id == artifact.id + assert result.payload == {"a": 1} + + +async def test_list_reads_bulk_and_reconciles_index(monkeypatch) -> None: + store, _ = _make_store(monkeypatch) + artifact = Artifact(type="T", payload={"x": 1}, produced_by="agent") + stale_id = str(uuid4()) + monkeypatch.setattr( + store, "_read_index", lambda _k: ([str(artifact.id), stale_id], "etag-1") + ) + + calls: dict[str, object] = {} + + def _reconcile(index_key, index_ids, live_keys, *, etag=None): + calls["index_key"] = index_key + calls["index_ids"] = index_ids + calls["live_keys"] = live_keys + calls["etag"] = etag + return [str(artifact.id)] + + monkeypatch.setattr(store, "_reconcile_index", _reconcile) + monkeypatch.setattr( + store._client, + "get_bulk_state", + lambda _store, _keys: SimpleNamespace( + items=[ + _FakeBulkItem(f"artifact:{artifact.id}", artifact.model_dump_json()), + _FakeBulkItem(f"artifact:{stale_id}", ""), + ] + ), + raising=False, + ) + + result = await store.list() + + assert [a.id for a in result] == [artifact.id] + assert calls["etag"] == "etag-1" + assert calls["live_keys"] == {str(artifact.id)} + + +async def test_list_by_type_returns_empty_when_index_empty(monkeypatch) -> None: + store, _ = _make_store(monkeypatch) + monkeypatch.setattr(store, "_read_index", lambda _k: ([], "etag-0")) + + result = await store.list_by_type("T") + + assert result == [] + + +async def test_get_by_type_casts_and_filters_correlation(monkeypatch) -> None: + store, _ = _make_store(monkeypatch) + + class DemoModel(BaseModel): + value: int + + matching = Artifact( + type="DemoModel", + payload={"value": 10}, + produced_by="agent", + correlation_id="corr-1", + ) + other = Artifact( + type="DemoModel", + payload={"value": 20}, + produced_by="agent", + correlation_id="corr-2", + ) + + async def _list_by_type(_type): + return [matching, other] + + monkeypatch.setattr(store, "list_by_type", _list_by_type) + + result = await store.get_by_type(DemoModel, correlation_id="corr-1") + + assert len(result) == 1 + assert isinstance(result[0], DemoModel) + assert result[0].value == 10 + + +async def test_record_consumptions_dispatches_by_transaction_flag(monkeypatch) -> None: + tx_store, _ = _make_store(monkeypatch, supports_transactions=True) + non_tx_store, _ = _make_store(monkeypatch, supports_transactions=False) + calls: list[str] = [] + + async def _tx(*, records): + assert records + calls.append("tx") + + async def _non_tx(*, records): + assert records + calls.append("non-tx") + + monkeypatch.setattr(tx_store, "_record_consumptions_transactional", _tx) + monkeypatch.setattr(tx_store, "_record_consumptions_non_transactional", _non_tx) + monkeypatch.setattr(non_tx_store, "_record_consumptions_transactional", _tx) + monkeypatch.setattr(non_tx_store, "_record_consumptions_non_transactional", _non_tx) + + rec = ConsumptionRecord(artifact_id=uuid4(), consumer="agent") + await tx_store.record_consumptions([rec]) + await non_tx_store.record_consumptions([rec]) + + assert calls == ["tx", "non-tx"] + + +async def test_upsert_and_clear_snapshot_dispatches(monkeypatch) -> None: + tx_store, _ = _make_store(monkeypatch, supports_transactions=True) + non_tx_store, _ = _make_store(monkeypatch, supports_transactions=False) + calls: list[str] = [] + + async def _tx_upsert(_snapshot): + calls.append("tx-upsert") + + async def _non_tx_upsert(_snapshot): + calls.append("non-tx-upsert") + + async def _tx_clear(): + calls.append("tx-clear") + + async def _non_tx_clear(): + calls.append("non-tx-clear") + + monkeypatch.setattr(tx_store, "_upsert_agent_snapshot_transactional", _tx_upsert) + monkeypatch.setattr( + tx_store, "_upsert_agent_snapshot_non_transactional", _non_tx_upsert + ) + monkeypatch.setattr( + non_tx_store, "_upsert_agent_snapshot_transactional", _tx_upsert + ) + monkeypatch.setattr( + non_tx_store, "_upsert_agent_snapshot_non_transactional", _non_tx_upsert + ) + + monkeypatch.setattr(tx_store, "_clear_agent_snapshots_transactional", _tx_clear) + monkeypatch.setattr( + tx_store, "_clear_agent_snapshots_non_transactional", _non_tx_clear + ) + monkeypatch.setattr(non_tx_store, "_clear_agent_snapshots_transactional", _tx_clear) + monkeypatch.setattr( + non_tx_store, "_clear_agent_snapshots_non_transactional", _non_tx_clear + ) + + snapshot = AgentSnapshotRecord( + agent_name="agent", + description="desc", + subscriptions=["A"], + output_types=["B"], + labels=["l"], + first_seen=datetime(2026, 1, 1, tzinfo=UTC), + last_seen=datetime(2026, 1, 2, tzinfo=UTC), + signature="sig", + ) + + await tx_store.upsert_agent_snapshot(snapshot) + await non_tx_store.upsert_agent_snapshot(snapshot) + await tx_store.clear_agent_snapshots() + await non_tx_store.clear_agent_snapshots() + + assert calls == ["tx-upsert", "non-tx-upsert", "tx-clear", "non-tx-clear"] + + +async def test_load_agent_snapshots_handles_error_empty_and_valid_entries( + monkeypatch, +) -> None: + store, _ = _make_store(monkeypatch) + snapshot = AgentSnapshotRecord( + agent_name="agent-ok", + description="desc", + subscriptions=["A"], + output_types=["B"], + labels=["l"], + first_seen=datetime(2026, 1, 1, tzinfo=UTC), + last_seen=datetime(2026, 1, 2, tzinfo=UTC), + signature="sig", + ) + + monkeypatch.setattr( + store, "_read_index", lambda _k: (["bad", "empty", "agent-ok"], "etag-snap") + ) + + reconcile_calls: dict[str, object] = {} + + def _reconcile(index_key, index_ids, live_keys, *, etag=None): + reconcile_calls["index_key"] = index_key + reconcile_calls["ids"] = index_ids + reconcile_calls["live"] = live_keys + reconcile_calls["etag"] = etag + return ["agent-ok"] + + monkeypatch.setattr(store, "_reconcile_index", _reconcile) + monkeypatch.setattr( + store._client, + "get_bulk_state", + lambda _store, keys, parallelism=10: SimpleNamespace( + items=[ + _FakeBulkItem("snapshot:bad", "", error="boom"), + _FakeBulkItem("snapshot:empty", ""), + _FakeBulkItem("snapshot:agent-ok", serialize_agent_snapshot(snapshot)), + ] + ), + raising=False, + ) + + result = await store.load_agent_snapshots() + + assert len(result) == 1 + assert result[0].agent_name == "agent-ok" + assert reconcile_calls["live"] == {"agent-ok"} + assert reconcile_calls["etag"] == "etag-snap" + + +def test_read_and_write_index_use_etag_and_state_options(monkeypatch) -> None: + store, _ = _make_store(monkeypatch, supports_etag=True, consistency="strong") + + monkeypatch.setattr( + store._client, + "get_state", + lambda _store, _key: _FakeStateResponse('["a"]', etag="etag-read"), + raising=False, + ) + + save_call: dict[str, object] = {} + + def _save_state(_store, key, value, **kwargs): + save_call["key"] = key + save_call["value"] = value + save_call.update(kwargs) + + monkeypatch.setattr(store._client, "save_state", _save_state, raising=False) + + items, etag = store._read_index("idx:key") + store._write_index("idx:key", ["x"], etag="etag-write") + + assert items == ["a"] + assert etag == "etag-read" + assert save_call["key"] == "idx:key" + assert save_call["value"] == '["x"]' + assert save_call["etag"] == "etag-write" + assert save_call["state_metadata"] == {} + assert save_call["options"] is not None + + +def test_reconcile_index_without_stale_entries_does_not_write(monkeypatch) -> None: + store, _ = _make_store(monkeypatch) + called = {"count": 0} + + def _write(*_args, **_kwargs): + called["count"] += 1 + + monkeypatch.setattr(store, "_write_index", _write) + + result = store._reconcile_index("idx:test", ["a", "b"], {"a", "b"}) + + assert result == ["a", "b"] + assert called["count"] == 0 + + +async def test_query_via_dapr_api_handles_pagination_errors_and_tag_filter( + monkeypatch, +) -> None: + store, _ = _make_store(monkeypatch) + + keep = Artifact( + type="Demo", + payload={"x": 1}, + produced_by="writer", + tags={"keep", "extra"}, + ) + drop = Artifact( + type="Demo", + payload={"x": 2}, + produced_by="writer", + tags={"other"}, + ) + + class _QueryItem: + def __init__(self, key: str, data: str, error: str | None = None) -> None: + self.key = key + self._data = data + self.error = error + + def text(self) -> str: + return self._data + + class _QueryResponse: + def __init__(self, results, token: str | None) -> None: + self.results = results + self.token = token + + calls: list[str] = [] + + def _query_state(*, store_name: str, query: str): + calls.append(query) + if len(calls) == 1: + return _QueryResponse( + [ + _QueryItem("bad", "", error="boom"), + _QueryItem("non-artifact", "not-json"), + _QueryItem("artifact-1", keep.model_dump_json()), + ], + token="next-token", + ) + return _QueryResponse([_QueryItem("artifact-2", drop.model_dump_json())], None) + + monkeypatch.setattr(store._client, "query_state", _query_state, raising=False) + + artifacts = await store._query_via_dapr_api(FilterConfig(tags={"keep"})) + + assert [a.id for a in artifacts] == [keep.id] + assert len(calls) == 2 + assert '"token": "next-token"' in calls[1] + + +def test_query_via_index_scan_type_filters_and_reconcile_paths(monkeypatch) -> None: + store, _ = _make_store(monkeypatch) + artifact = Artifact(type="Demo", payload={"v": 1}, produced_by="writer") + + class _AlwaysMatch: + def __init__(self, _filters) -> None: + pass + + def matches(self, _artifact: Artifact) -> bool: + return True + + monkeypatch.setattr( + "flock.storage.dapr.dapr_state_blackboard_store.ArtifactFilter", + _AlwaysMatch, + ) + + def _read_index(key: str): + if key == _type_index_key("A"): + return ([str(artifact.id)], "etag-a") + if key == _type_index_key("B"): + return ([str(uuid4())], "etag-b") + return ([], None) + + monkeypatch.setattr(store, "_read_index", _read_index) + + monkeypatch.setattr( + store._client, + "get_bulk_state", + lambda _store, _keys, parallelism=10: SimpleNamespace( + items=[ + _FakeBulkItem(f"artifact:{artifact.id}", artifact.model_dump_json()), + _FakeBulkItem(f"artifact:{uuid4()}", "not-json"), + ] + ), + raising=False, + ) + + reconciled: list[tuple[str, list[str], set[str], str | None]] = [] + + def _reconcile(index_key, index_ids, live_keys, *, etag=None): + reconciled.append((index_key, index_ids, live_keys, etag)) + return [str(artifact.id)] + + monkeypatch.setattr(store, "_reconcile_index", _reconcile) + + result = store._query_via_index_scan(FilterConfig(type_names={"A", "B"})) + + assert [a.id for a in result] == [artifact.id] + assert len(reconciled) == 2 + assert {row[0] for row in reconciled} == { + _type_index_key("A"), + _type_index_key("B"), + } + assert all(str(artifact.id) in row[2] for row in reconciled) + + +async def test_get_consumptions_by_artifact_ids_skips_error_entries( + monkeypatch, +) -> None: + store, _ = _make_store(monkeypatch) + artifact_id = uuid4() + + monkeypatch.setattr( + store._client, + "get_bulk_state", + lambda store_name, keys, parallelism=10: SimpleNamespace( + items=[ + _FakeBulkItem(f"consumptions:{artifact_id}", "[]"), + _FakeBulkItem("consumptions:bad", "[]", error="oops"), + ] + ), + raising=False, + ) + + result = await store._get_consumptions_by_artifact_ids([artifact_id]) + + assert result == {str(artifact_id): []} + + +async def test_publish_transactional_writes_artifact_and_index_updates( + monkeypatch, +) -> None: + store, _ = _make_store(monkeypatch, supports_transactions=True, supports_etag=True) + artifact = Artifact(type="Demo", payload={"x": 1}, produced_by="writer") + + def _read_index(key: str): + if key == "idx:artifacts": + return ([], "etag-all") + return ([], "etag-type") + + monkeypatch.setattr(store, "_read_index", _read_index) + + save_calls: list[dict[str, object]] = [] + + def _save_state(_store, key, value, **kwargs): + save_calls.append({"key": key, "value": value, **kwargs}) + + tx_calls: dict[str, object] = {} + + def _execute_state_transaction(*, store_name, operations): + tx_calls["store_name"] = store_name + tx_calls["operations"] = operations + + monkeypatch.setattr(store._client, "save_state", _save_state, raising=False) + monkeypatch.setattr( + store._client, + "execute_state_transaction", + _execute_state_transaction, + raising=False, + ) + + await store._publish_transactional(artifact) + + assert len(save_calls) == 1 + assert save_calls[0]["key"] == f"artifact:{artifact.id}" + operations = tx_calls["operations"] + assert len(operations) == 2 + assert {op.key for op in operations} == {"idx:artifacts", "idx:type:Demo"} + + +async def test_record_consumptions_non_transactional_groups_and_saves_bulk( + monkeypatch, +) -> None: + store, _ = _make_store(monkeypatch, supports_etag=True) + artifact_id = uuid4() + existing = ConsumptionRecord(artifact_id=artifact_id, consumer="existing") + new_record = ConsumptionRecord(artifact_id=artifact_id, consumer="new") + + monkeypatch.setattr( + store._client, + "get_state", + lambda _store, _key: _FakeStateResponse( + json.dumps([ + { + "artifact_id": str(existing.artifact_id), + "consumer": existing.consumer, + "run_id": existing.run_id, + "correlation_id": existing.correlation_id, + "consumed_at": existing.consumed_at.isoformat(), + } + ]), + etag="etag-cons", + ), + raising=False, + ) + + captured: dict[str, object] = {} + + def _save_bulk_state(*, store_name, states): + captured["store_name"] = store_name + captured["states"] = states + + monkeypatch.setattr( + store._client, "save_bulk_state", _save_bulk_state, raising=False + ) + + await store._record_consumptions_non_transactional([new_record]) + + states = captured["states"] + assert len(states) == 1 + payload = json.loads(states[0].value) + assert len(payload) == 2 + assert {entry["consumer"] for entry in payload} == {"existing", "new"} + + +def test_build_dapr_query_single_producer_and_visibility_paths() -> None: + query = json.loads( + _build_dapr_query( + FilterConfig( + produced_by={"solo-producer"}, + visibility={"Public"}, + ) + ) + ) + + and_conditions = query["filter"]["AND"] + assert {"EQ": {"produced_by": "solo-producer"}} in and_conditions + assert {"EQ": {"visibility.kind": "Public"}} in and_conditions + + +def test_store_initializes_eventual_consistency_branch(monkeypatch) -> None: + store, _ = _make_store(monkeypatch, consistency="eventual") + + assert store._consistency.name == "unspecified" + + +async def test_retry_on_etag_conflict_re_raises_non_etag_errors(monkeypatch) -> None: + store, _ = _make_store(monkeypatch, supports_etag=True, etag_max_retries=3) + + async def _sleep_should_not_be_called(_delay: float) -> None: + raise AssertionError("sleep should not be called for non-etag errors") + + monkeypatch.setattr( + "flock.storage.dapr.dapr_state_blackboard_store.asyncio.sleep", + _sleep_should_not_be_called, + ) + + def _operation() -> None: + raise RuntimeError("not an etag conflict") + + try: + await store._retry_on_etag_conflict(_operation) + except RuntimeError as err: + assert "not an etag conflict" in str(err) + else: + raise AssertionError("Expected RuntimeError to be raised") + + +def test_reconcile_index_with_stale_entries_writes_cleaned_index(monkeypatch) -> None: + store, _ = _make_store(monkeypatch) + write_calls: list[tuple[str, list[str], str | None]] = [] + + def _write(index_key: str, items: list[str], *, etag: str | None = None) -> None: + write_calls.append((index_key, items, etag)) + + monkeypatch.setattr(store, "_write_index", _write) + + cleaned = store._reconcile_index( + "idx:artifacts", + ["live-1", "stale-1", "live-2"], + {"live-1", "live-2"}, + etag="etag-reconcile", + ) + + assert cleaned == ["live-1", "live-2"] + assert write_calls == [("idx:artifacts", ["live-1", "live-2"], "etag-reconcile")] + + +async def test_publish_non_transactional_persists_and_updates_indexes( + monkeypatch, +) -> None: + store, _ = _make_store(monkeypatch, supports_etag=True, supports_ttl=True) + artifact = Artifact(type="Demo", payload={"x": 1}, produced_by="writer") + + save_calls: list[dict[str, object]] = [] + index_writes: list[tuple[str, list[str], str | None]] = [] + + def _read_index(key: str): + if key == "idx:artifacts": + return ([], "etag-all") + if key == "idx:type:Demo": + return ([], "etag-type") + raise AssertionError(f"Unexpected index key {key}") + + def _save_state(_store, key, value, **kwargs): + save_calls.append({"key": key, "value": value, **kwargs}) + + def _write_index(key: str, items: list[str], *, etag: str | None = None) -> None: + index_writes.append((key, items, etag)) + + monkeypatch.setattr(store, "_read_index", _read_index) + monkeypatch.setattr(store._client, "save_state", _save_state, raising=False) + monkeypatch.setattr(store, "_write_index", _write_index) + + await store._publish_non_transactional(artifact) + + assert len(save_calls) == 1 + assert save_calls[0]["key"] == f"artifact:{artifact.id}" + assert {w[0] for w in index_writes} == {"idx:artifacts", "idx:type:Demo"} + assert all(str(artifact.id) in w[1] for w in index_writes) + + +async def test_record_consumptions_transactional_groups_and_transacts( + monkeypatch, +) -> None: + store, _ = _make_store(monkeypatch, supports_etag=True) + first_id = uuid4() + second_id = uuid4() + records = [ + ConsumptionRecord(artifact_id=first_id, consumer="a"), + ConsumptionRecord(artifact_id=first_id, consumer="b"), + ConsumptionRecord(artifact_id=second_id, consumer="c"), + ] + + def _get_state(_store, key: str): + if key.endswith(str(first_id)): + return _FakeStateResponse("[]", etag="etag-first") + if key.endswith(str(second_id)): + return _FakeStateResponse("[]", etag="etag-second") + return _FakeStateResponse("[]") + + tx_call: dict[str, object] = {} + + def _execute_state_transaction(*, store_name, operations): + tx_call["store_name"] = store_name + tx_call["operations"] = operations + + monkeypatch.setattr(store._client, "get_state", _get_state, raising=False) + monkeypatch.setattr( + store._client, + "execute_state_transaction", + _execute_state_transaction, + raising=False, + ) + + await store._record_consumptions_transactional(records) + + operations = tx_call["operations"] + assert len(operations) == 2 + assert {op.key for op in operations} == { + f"consumptions:{first_id}", + f"consumptions:{second_id}", + } + + +async def test_upsert_agent_snapshot_transactional_builds_operations( + monkeypatch, +) -> None: + store, _ = _make_store(monkeypatch, supports_etag=True) + snapshot = AgentSnapshotRecord( + agent_name="agent-a", + description="desc", + subscriptions=["Input"], + output_types=["Output"], + labels=["label"], + first_seen=datetime(2026, 1, 1, tzinfo=UTC), + last_seen=datetime(2026, 1, 2, tzinfo=UTC), + signature="sig-a", + ) + + monkeypatch.setattr( + store._client, + "get_state", + lambda _store, _key: _FakeStateResponse("", etag="etag-snap"), + raising=False, + ) + monkeypatch.setattr(store, "_read_index", lambda _k: ([], "etag-idx")) + + tx_call: dict[str, object] = {} + + def _execute_state_transaction(*, store_name, operations): + tx_call["store_name"] = store_name + tx_call["operations"] = operations + + monkeypatch.setattr( + store._client, + "execute_state_transaction", + _execute_state_transaction, + raising=False, + ) + + await store._upsert_agent_snapshot_transactional(snapshot) + + operations = tx_call["operations"] + assert len(operations) == 2 + assert {op.key for op in operations} == { + "snapshot:agent-a", + "idx:snapshots", + } + + +async def test_upsert_agent_snapshot_non_transactional_reads_etag_and_updates_index( + monkeypatch, +) -> None: + store, _ = _make_store(monkeypatch, supports_etag=True) + snapshot = AgentSnapshotRecord( + agent_name="agent-b", + description="desc", + subscriptions=["Input"], + output_types=["Output"], + labels=["label"], + first_seen=datetime(2026, 1, 1, tzinfo=UTC), + last_seen=datetime(2026, 1, 2, tzinfo=UTC), + signature="sig-b", + ) + + monkeypatch.setattr( + store._client, + "get_state", + lambda _store, _key: _FakeStateResponse("", etag="etag-existing"), + raising=False, + ) + monkeypatch.setattr(store, "_read_index", lambda _k: ([], "etag-idx")) + + save_calls: list[dict[str, object]] = [] + index_writes: list[tuple[str, list[str], str | None]] = [] + + def _save_state(_store, key, value, **kwargs): + save_calls.append({"key": key, "value": value, **kwargs}) + + def _write_index(key: str, items: list[str], *, etag: str | None = None) -> None: + index_writes.append((key, items, etag)) + + monkeypatch.setattr(store._client, "save_state", _save_state, raising=False) + monkeypatch.setattr(store, "_write_index", _write_index) + + await store._upsert_agent_snapshot_non_transactional(snapshot) + + assert len(save_calls) == 1 + assert save_calls[0]["etag"] == "etag-existing" + assert index_writes == [("idx:snapshots", ["agent-b"], "etag-idx")] + + +async def test_clear_agent_snapshots_transactional_deletes_all_and_index( + monkeypatch, +) -> None: + store, _ = _make_store(monkeypatch) + monkeypatch.setattr(store, "_read_index", lambda _k: (["a", "b"], "etag-idx")) + + tx_call: dict[str, object] = {} + + def _execute_state_transaction(*, store_name, operations): + tx_call["store_name"] = store_name + tx_call["operations"] = operations + + monkeypatch.setattr( + store._client, + "execute_state_transaction", + _execute_state_transaction, + raising=False, + ) + + await store._clear_agent_snapshots_transactional() + + operations = tx_call["operations"] + assert len(operations) == 3 + assert {op.key for op in operations} == { + "snapshot:a", + "snapshot:b", + "idx:snapshots", + } + + +async def test_clear_agent_snapshots_non_transactional_deletes_and_resets_index( + monkeypatch, +) -> None: + store, _ = _make_store(monkeypatch) + monkeypatch.setattr(store, "_read_index", lambda _k: (["a", "b"], "etag-idx")) + + deleted_keys: list[str] = [] + wrote_index: list[tuple[str, list[str], str | None]] = [] + + def _delete_state(_store, *, key, options=None): + _ = options + deleted_keys.append(key) + + def _write_index(key: str, items: list[str], *, etag: str | None = None) -> None: + wrote_index.append((key, items, etag)) + + monkeypatch.setattr(store._client, "delete_state", _delete_state, raising=False) + monkeypatch.setattr(store, "_write_index", _write_index) + + await store._clear_agent_snapshots_non_transactional() + + assert deleted_keys == ["snapshot:a", "snapshot:b"] + assert wrote_index == [("idx:snapshots", [], "etag-idx")] + + +async def test_list_returns_empty_when_index_empty(monkeypatch) -> None: + store, _ = _make_store(monkeypatch) + monkeypatch.setattr(store, "_read_index", lambda _k: ([], "etag-none")) + + result = await store.list() + + assert result == [] + + +async def test_list_by_type_reads_bulk_and_reconciles_non_empty_index( + monkeypatch, +) -> None: + store, _ = _make_store(monkeypatch) + artifact = Artifact(type="Demo", payload={"v": 1}, produced_by="writer") + stale_id = str(uuid4()) + monkeypatch.setattr( + store, "_read_index", lambda _k: ([str(artifact.id), stale_id], "etag-type") + ) + + reconcile_calls: dict[str, object] = {} + + def _reconcile(index_key, index_ids, live_keys, *, etag=None): + reconcile_calls["index_key"] = index_key + reconcile_calls["index_ids"] = index_ids + reconcile_calls["live_keys"] = live_keys + reconcile_calls["etag"] = etag + return [str(artifact.id)] + + monkeypatch.setattr(store, "_reconcile_index", _reconcile) + monkeypatch.setattr( + store._client, + "get_bulk_state", + lambda _store, _keys: SimpleNamespace( + items=[ + _FakeBulkItem(f"artifact:{artifact.id}", artifact.model_dump_json()), + _FakeBulkItem(f"artifact:{stale_id}", ""), + ] + ), + raising=False, + ) + + result = await store.list_by_type("Demo") + + assert [a.id for a in result] == [artifact.id] + assert reconcile_calls["index_key"] == "idx:type:Demo" + assert reconcile_calls["etag"] == "etag-type" + assert reconcile_calls["live_keys"] == {str(artifact.id)} + + +async def test_query_artifacts_limit_zero_and_without_meta(monkeypatch) -> None: + store, _ = _make_store(monkeypatch) + artifacts = [ + Artifact(type="T", payload={"i": 2}, produced_by="a"), + Artifact(type="T", payload={"i": 1}, produced_by="a"), + ] + + async def _backend(_filters): + return artifacts + + monkeypatch.setattr(store, "_query_backend", _backend) + + page, total = await store.query_artifacts(limit=0, offset=0, embed_meta=False) + + assert total == 2 + assert len(page) == 2 + assert all(isinstance(item, Artifact) for item in page) + + +async def test_fetch_graph_artifacts_keeps_existing_envelopes(monkeypatch) -> None: + store, _ = _make_store(monkeypatch) + envelope = ArtifactEnvelope( + artifact=Artifact(type="T", payload={"x": 1}, produced_by="agent") + ) + + async def _query(**_kwargs): + return ([envelope], 1) + + monkeypatch.setattr(store, "query_artifacts", _query) + + envelopes, total = await store.fetch_graph_artifacts() + + assert total == 1 + assert envelopes == [envelope] + + +async def test_summarize_artifacts_raises_type_error_for_non_artifact_items( + monkeypatch, +) -> None: + store, _ = _make_store(monkeypatch) + envelope = ArtifactEnvelope( + artifact=Artifact(type="T", payload={}, produced_by="agent") + ) + + async def _query(**_kwargs): + return ([envelope], 1) + + monkeypatch.setattr(store, "query_artifacts", _query) + + try: + await store.summarize_artifacts() + except TypeError as err: + assert "Expected Artifact instance" in str(err) + else: + raise AssertionError("Expected TypeError to be raised") + + +async def test_load_agent_snapshots_returns_empty_when_index_is_empty( + monkeypatch, +) -> None: + store, _ = _make_store(monkeypatch) + monkeypatch.setattr(store, "_read_index", lambda _k: ([], "etag-empty")) + + result = await store.load_agent_snapshots() + + assert result == [] diff --git a/tests/storage/dapr/test_serialization.py b/tests/storage/dapr/test_serialization.py new file mode 100644 index 000000000..c7bd86fcb --- /dev/null +++ b/tests/storage/dapr/test_serialization.py @@ -0,0 +1,116 @@ +from __future__ import annotations + +import json +from datetime import UTC, datetime +from uuid import uuid4 + +import pytest + +from flock.core.artifacts import Artifact +from flock.core.store import AgentSnapshotRecord, ConsumptionRecord +from flock.core.visibility import AgentIdentity, PrivateVisibility, Visibility +from flock.storage.dapr._serialization import ( + _VISIBILITY_MAP, + _default, + deserialize_agent_snapshot, + deserialize_artifact, + deserialize_consumption_records, + deserialize_index, + serialize_agent_snapshot, + serialize_artifact, + serialize_consumption_records, + serialize_index, +) + + +def test_default_handles_uuid_datetime_and_set() -> None: + now = datetime(2026, 1, 1, tzinfo=UTC) + value = {"id": uuid4(), "when": now, "tags": {"b", "a"}} + + payload = json.dumps(value, default=_default) + decoded = json.loads(payload) + + assert isinstance(decoded["id"], str) + assert decoded["when"] == now.isoformat() + assert decoded["tags"] == ["a", "b"] + + +def test_default_raises_for_unsupported_type() -> None: + with pytest.raises(TypeError): + _default(object()) + + +def test_artifact_roundtrip_preserves_visibility_subclass() -> None: + artifact = Artifact( + type="demo.Type", + payload={"value": 1}, + produced_by="writer", + visibility=PrivateVisibility(agents={"allowed-agent"}), + ) + + encoded = serialize_artifact(artifact) + decoded = deserialize_artifact(encoded) + + assert isinstance(decoded.visibility, PrivateVisibility) + assert decoded.visibility.kind == "Private" + assert not decoded.visibility.allows(AgentIdentity(name="allowed-agent")) + + +def test_deserialize_artifact_keeps_base_visibility_when_kind_map_missing( + monkeypatch, +) -> None: + artifact = Artifact(type="demo.Type", payload={"value": 1}, produced_by="writer") + monkeypatch.setitem(_VISIBILITY_MAP, "Public", None) # type: ignore[arg-type] + + decoded = deserialize_artifact(serialize_artifact(artifact)) + + assert type(decoded.visibility) is Visibility + + +def test_consumption_records_roundtrip_and_empty_decode() -> None: + record = ConsumptionRecord( + artifact_id=uuid4(), + consumer="agent-a", + run_id="run-1", + correlation_id="corr-1", + consumed_at=datetime(2026, 1, 2, tzinfo=UTC), + ) + + encoded = serialize_consumption_records([record]) + decoded = deserialize_consumption_records(encoded) + + assert decoded == [record] + assert deserialize_consumption_records("") == [] + assert deserialize_consumption_records(encoded.encode("utf-8")) == [record] + + +def test_agent_snapshot_roundtrip_and_dict_input() -> None: + snapshot = AgentSnapshotRecord( + agent_name="agent-a", + description="desc", + subscriptions=["A"], + output_types=["B"], + labels=["l1"], + first_seen=datetime(2026, 1, 3, tzinfo=UTC), + last_seen=datetime(2026, 1, 4, tzinfo=UTC), + signature="sig", + ) + + encoded = serialize_agent_snapshot(snapshot) + decoded_from_json = deserialize_agent_snapshot(encoded) + decoded_from_bytes = deserialize_agent_snapshot(encoded.encode("utf-8")) + decoded_from_dict = deserialize_agent_snapshot(json.loads(encoded)) + + assert decoded_from_json == snapshot + assert decoded_from_bytes == snapshot + assert decoded_from_dict == snapshot + + +def test_index_helpers_roundtrip_and_empty_data() -> None: + keys = ["a", "b"] + + encoded = serialize_index(keys) + + assert deserialize_index(encoded) == keys + assert deserialize_index(encoded.encode("utf-8")) == keys + assert deserialize_index("") == [] diff --git a/tests/storage/test_storage_exports.py b/tests/storage/test_storage_exports.py new file mode 100644 index 000000000..878dfc897 --- /dev/null +++ b/tests/storage/test_storage_exports.py @@ -0,0 +1,41 @@ +from __future__ import annotations + + +def test_storage_package_lazy_getattr_resolves_symbol() -> None: + import flock.storage as storage_pkg + + assert storage_pkg.serialize_index(["a", "b"]) == '["a", "b"]' + + +def test_storage_package_lazy_getattr_missing_symbol() -> None: + import flock.storage as storage_pkg + + try: + _ = storage_pkg.DOES_NOT_EXIST + except AttributeError as err: + assert "has no attribute" in str(err) + else: + raise AssertionError("Expected AttributeError for unknown export") + + +def test_storage_package_exports_include_expected_symbols() -> None: + import flock.storage as storage_pkg + + expected = { + "DaprStateBlackboardConfig", + "DaprStateBlackboardStore", + "DaprStateBlackboardStoreClientConfig", + "SQLiteQueryBuilder", + "SQLiteSchemaManager", + "create_dapr_client", + "deserialize_agent_snapshot", + "deserialize_artifact", + "deserialize_consumption_records", + "deserialize_index", + "serialize_agent_snapshot", + "serialize_artifact", + "serialize_consumption_records", + "serialize_index", + } + + assert expected.issubset(set(storage_pkg.__all__)) diff --git a/tests/test_guard_component.py b/tests/test_guard_component.py index e3caea301..c6e45c96a 100644 --- a/tests/test_guard_component.py +++ b/tests/test_guard_component.py @@ -245,9 +245,7 @@ def mock_ctx(self): @pytest.fixture def inputs(self): - return EvalInputs( - artifacts=[_make_artifact({"question": "Tell me a secret"})] - ) + return EvalInputs(artifacts=[_make_artifact({"question": "Tell me a secret"})]) @pytest.fixture def result(self): @@ -312,13 +310,9 @@ async def test_scan_output_disabled_by_default( assert returned is result @pytest.mark.asyncio - async def test_scan_output_enabled_warn( - self, mock_agent, mock_ctx, inputs, result - ): + async def test_scan_output_enabled_warn(self, mock_agent, mock_ctx, inputs, result): guard = _RejectingGuard( - config=GuardComponentConfig( - scan_output=True, on_output_flagged="warn" - ), + config=GuardComponentConfig(scan_output=True, on_output_flagged="warn"), ) # Should NOT raise (warn mode) returned = await guard.on_post_evaluate(mock_agent, mock_ctx, inputs, result) @@ -329,22 +323,16 @@ async def test_scan_output_enabled_block( self, mock_agent, mock_ctx, inputs, result ): guard = _RejectingGuard( - config=GuardComponentConfig( - scan_output=True, on_output_flagged="block" - ), + config=GuardComponentConfig(scan_output=True, on_output_flagged="block"), ) with pytest.raises(GuardBlockedError): await guard.on_post_evaluate(mock_agent, mock_ctx, inputs, result) @pytest.mark.asyncio - async def test_scan_output_empty_result_skips( - self, mock_agent, mock_ctx, inputs - ): + async def test_scan_output_empty_result_skips(self, mock_agent, mock_ctx, inputs): """When result has no text, output scan is skipped.""" guard = _RejectingGuard( - config=GuardComponentConfig( - scan_output=True, on_output_flagged="block" - ), + config=GuardComponentConfig(scan_output=True, on_output_flagged="block"), ) empty_result = EvalResult(artifacts=[]) returned = await guard.on_post_evaluate( @@ -396,9 +384,7 @@ async def test_extract_prompt_text_tuples(self): @pytest.mark.asyncio async def test_extract_context_documents_tuples(self): """Tuples in context documents are extracted correctly.""" - inputs = EvalInputs( - artifacts=[_make_artifact({"items": ("one", "two")})] - ) + inputs = EvalInputs(artifacts=[_make_artifact({"items": ("one", "two")})]) docs = GuardComponent._extract_context_documents(inputs) assert len(docs) == 1 assert "one" in docs[0] @@ -415,7 +401,9 @@ async def test_extract_result_text_with_lists(self): assert "a" in text @pytest.mark.asyncio - async def test_output_scan_safe_passes_through(self, mock_agent, mock_ctx, inputs, result): + async def test_output_scan_safe_passes_through( + self, mock_agent, mock_ctx, inputs, result + ): """Output scanning with safe verdict returns result unchanged.""" guard = _PassthroughGuard( config=GuardComponentConfig(scan_output=True, on_output_flagged="block"), @@ -535,7 +523,9 @@ async def test_env_fallback_endpoint(self): } with patch.dict( "os.environ", - {"AZURE_CONTENT_SAFETY_ENDPOINT": "https://env.cognitiveservices.azure.com"}, + { + "AZURE_CONTENT_SAFETY_ENDPOINT": "https://env.cognitiveservices.azure.com" + }, ): with patch("httpx.AsyncClient.post", new_callable=AsyncMock) as mock_post: mock_response = MagicMock() @@ -545,7 +535,10 @@ async def test_env_fallback_endpoint(self): await guard._call_shield_api("test", []) call_url = mock_post.call_args[0][0] - assert call_url == "https://env.cognitiveservices.azure.com/contentsafety/text:shieldPrompt" + assert ( + call_url + == "https://env.cognitiveservices.azure.com/contentsafety/text:shieldPrompt" + ) @pytest.mark.asyncio async def test_env_fallback_api_key(self): @@ -554,9 +547,7 @@ async def test_env_fallback_api_key(self): endpoint="https://test.cognitiveservices.azure.com", ), ) - with patch.dict( - "os.environ", {"AZURE_CONTENT_SAFETY_KEY": "env-key-456"} - ): + with patch.dict("os.environ", {"AZURE_CONTENT_SAFETY_KEY": "env-key-456"}): headers = await guard._build_headers() assert headers["Ocp-Apim-Subscription-Key"] == "env-key-456" @@ -570,7 +561,10 @@ async def test_managed_identity_headers(self): ), ) with patch.object( - guard, "_get_managed_identity_token", new_callable=AsyncMock, return_value="mock-token-xyz" + guard, + "_get_managed_identity_token", + new_callable=AsyncMock, + return_value="mock-token-xyz", ): headers = await guard._build_headers() assert headers["Authorization"] == "Bearer mock-token-xyz" @@ -590,7 +584,9 @@ async def test_document_truncation(self, guard): mock_post.return_value = mock_response await guard._call_shield_api("test", ["long document text here"]) - body = mock_post.call_args.kwargs.get("json") or mock_post.call_args[1].get("json") + body = mock_post.call_args.kwargs.get("json") or mock_post.call_args[1].get( + "json" + ) assert body["documents"] == ["long "] @pytest.mark.asyncio @@ -647,9 +643,7 @@ async def test_multiple_guards_compose(self): agent = MagicMock() agent.name = "agent" ctx = MagicMock() - inputs = EvalInputs( - artifacts=[_make_artifact({"q": "hello"})] - ) + inputs = EvalInputs(artifacts=[_make_artifact({"q": "hello"})]) # First guard passes inputs = await pass_guard.on_pre_evaluate(agent, ctx, inputs) diff --git a/uv.lock b/uv.lock index e9c987ed0..37bc59ee7 100644 --- a/uv.lock +++ b/uv.lock @@ -724,6 +724,23 @@ version = "0.9.5" source = { registry = "https://pypi.org/simple" } sdist = { url = "https://files.pythonhosted.org/packages/f1/2a/8c3ac3d8bc94e6de8d7ae270bb5bc437b210bb9d6d9e46630c98f4abd20c/csscompressor-0.9.5.tar.gz", hash = "sha256:afa22badbcf3120a4f392e4d22f9fff485c044a1feda4a950ecc5eba9dd31a05", size = 237808, upload-time = "2017-11-26T21:13:08.238Z" } +[[package]] +name = "dapr" +version = "1.14.1" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "aiohttp" }, + { name = "grpcio" }, + { name = "grpcio-status" }, + { name = "protobuf" }, + { name = "python-dateutil" }, + { name = "typing-extensions" }, +] +sdist = { url = "https://files.pythonhosted.org/packages/d5/6e/5e68994d32e31023b5545e7e1c36e2d835084dc17e6efe8abba170180208/dapr-1.14.1.tar.gz", hash = "sha256:c53923e28d08ee5a0c2e32aec73d15847028bea69674e59d275ef9439da7a3e7", size = 114863, upload-time = "2026-03-25T14:02:52.674Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/68/58/5494f4c7bb5da53f0671600db81501382ea98bba6cb71d46e98adf06d494/dapr-1.14.1-py3-none-any.whl", hash = "sha256:569da8e13baa1065d532f25c892e15365bfedc52cc1f800dd0d5a06a935f0a87", size = 154465, upload-time = "2026-03-25T14:02:51.121Z" }, +] + [[package]] name = "debugpy" version = "1.8.17" @@ -960,7 +977,7 @@ wheels = [ [[package]] name = "flock-core" -version = "0.5.500" +version = "0.5.501" source = { editable = "." } dependencies = [ { name = "aiosqlite" }, @@ -996,6 +1013,9 @@ dependencies = [ azure = [ { name = "azure-identity" }, ] +dapr = [ + { name = "dapr" }, +] semantic = [ { name = "sentence-transformers" }, ] @@ -1047,6 +1067,7 @@ requires-dist = [ { name = "azure-identity", marker = "extra == 'azure'", specifier = ">=1.15.0" }, { name = "bitsandbytes", marker = "extra == 'transformers'", specifier = ">=0.48.2" }, { name = "croniter", specifier = "==6.0.0" }, + { name = "dapr", marker = "extra == 'dapr'", specifier = ">=1.13.0,<1.15" }, { name = "devtools", specifier = "==0.12.2" }, { name = "dspy", specifier = "==3.2.1" }, { name = "duckdb", specifier = "==1.4.1" }, @@ -1075,7 +1096,7 @@ requires-dist = [ { name = "uvicorn", specifier = "==0.38.0" }, { name = "websockets", specifier = "==15.0.1" }, ] -provides-extras = ["azure", "semantic", "transformers"] +provides-extras = ["azure", "semantic", "transformers", "dapr"] [package.metadata.requires-dev] dev = [ @@ -1329,6 +1350,20 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/19/41/0b430b01a2eb38ee887f88c1f07644a1df8e289353b78e82b37ef988fb64/grpcio-1.76.0-cp314-cp314-win_amd64.whl", hash = "sha256:922fa70ba549fce362d2e2871ab542082d66e2aaf0c19480ea453905b01f384e", size = 4834462, upload-time = "2025-10-21T16:22:39.772Z" }, ] +[[package]] +name = "grpcio-status" +version = "1.71.2" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "googleapis-common-protos" }, + { name = "grpcio" }, + { name = "protobuf" }, +] +sdist = { url = "https://files.pythonhosted.org/packages/fd/d1/b6e9877fedae3add1afdeae1f89d1927d296da9cf977eca0eb08fb8a460e/grpcio_status-1.71.2.tar.gz", hash = "sha256:c7a97e176df71cdc2c179cd1847d7fc86cca5832ad12e9798d7fed6b7a1aab50", size = 13677, upload-time = "2025-06-28T04:24:05.426Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/67/58/317b0134129b556a93a3b0afe00ee675b5657f0155509e22fcb853bafe2d/grpcio_status-1.71.2-py3-none-any.whl", hash = "sha256:803c98cb6a8b7dc6dbb785b1111aed739f241ab5e9da0bba96888aa74704cfd3", size = 14424, upload-time = "2025-06-28T04:23:42.136Z" }, +] + [[package]] name = "h11" version = "0.16.0"