--- name: ak-dev-new-multimodal-storage description: > Step-by-step guide for adding a new multimodal attachment storage backend to Agent Kernel. Use this skill when you need to integrate a new storage service (beyond in-memory, Redis, and DynamoDB) for persisting image and file attachments. Covers implementing the AttachmentStore interface, factory registration, configuration, and testing. license: Apache-2.0 metadata: author: yaalalabs category: developer --- # Adding a New Multimodal Storage Backend This guide walks through adding a new attachment storage backend to Agent Kernel's multimodal subsystem. Use the existing Redis (`ak-py/src/agentkernel/core/multimodal/storage/redis.py`) and DynamoDB (`ak-py/src/agentkernel/core/multimodal/storage/dynamodb.py`) implementations as reference. ## Existing Backends | Backend | Config value | Features | Extras | |---|---|---|---| | In-memory | `in_memory` | Ephemeral `ClassVar` dict, zero setup, single-process only | None | | Redis | `redis` | Persistent, TTL, shared `RedisDriver` (lazy connect, retry, ping/reconnect), distributed | `agentkernel[multimodal,redis]` | | DynamoDB | `dynamodb` | Serverless/AWS, TTL via `expiry_time`, fully managed | `agentkernel[multimodal,aws]` | | Session cache | `session_cache` | Legacy — stores in session `nv_cache` (causes bloat, not recommended) | None | > To run multimodal end-to-end (hooks/tools), you typically need `agentkernel[multimodal]` > plus the backend-specific extra shown above (for example: `agentkernel[multimodal,redis]` > or `agentkernel[multimodal,aws]`). ## Architecture Overview The multimodal storage system uses a simple pluggable pattern: 1. **`AttachmentStore`** (`storage/base.py`) — abstract base class with `save()`, `get()`, `delete()` 2. **`AttachmentData`** (`storage/base.py`) — dataclass representing a stored attachment (id, type, data, name, mime_type, description, timestamp, url) 3. **`AttachmentStorageManager`** (`storage/storage_manager.py`) — high-level API that reads configuration, instantiates the correct `AttachmentStore` via `_build_driver()`, and provides `save_attachment()` / `get_attachment_data()` methods 4. **Callers**: `MultimodalPreHook` saves attachments; `AnalyzeAttachmentsTool` retrieves them ``` MultimodalPreHook / AnalyzeAttachmentsTool → AttachmentStorageManager(session_id) → _build_driver(session_id) # reads config.multimodal.storage_type → InMemoryAttachmentStore # or Redis, DynamoDB, etc. → save_attachment() / get_attachment_data() → AttachmentStore.save() / .get() ``` ## Step-by-Step ### 1. Create the Storage Backend File Create `ak-py/src/agentkernel/core/multimodal/storage/.py`. ### 2. Implement the AttachmentStore ```python # ak-py/src/agentkernel/core/multimodal/storage/.py import logging from typing import Optional from .base import AttachmentStore logger = logging.getLogger("ak.core.multimodal.storage.") class AttachmentStore(AttachmentStore): """ storage backend for multimodal attachments.""" def __init__(self, session_id: str, **kwargs): """ Initialize the store for a specific session. :param session_id: Session identifier for key namespacing. :param kwargs: Backend-specific connection parameters from config. """ self._session_id = session_id # Initialize your client/connection # e.g., self._client = BackendClient(endpoint=kwargs.get("endpoint")) logger.info(f" attachment store initialized for session {session_id}") def save(self, attachment: dict, max_attachments: int) -> str: """ Save an attachment dict and return its ID. The attachment dict contains: - id: str (UUID, already generated by AttachmentStorageManager) - type: str ("image" or "file") - data: str (base64 encoded binary; empty string when url is set) - name: str (filename) - mime_type: str - description: str (LLM-generated) - timestamp: float (epoch) - url: str | None (a remote reference recorded without its bytes) Persist the dict whole rather than mapping these keys one by one. The set grows — `url` was added for remote attachments — and an implementation that enumerates them drops a new key silently: the record still loads, carrying neither bytes nor an address. :param attachment: Attachment data dictionary. :param max_attachments: Maximum attachments per session (prune oldest if exceeded). :return: The attachment ID. """ attachment_id = attachment["id"] key = f"{self._session_id}:{attachment_id}" # 1. Serialize and store the attachment # e.g., self._client.put(key, json.dumps(attachment), ttl=self._ttl) # 2. Enforce max_attachments: prune oldest entries if limit exceeded # (track per-session count via an index or query) logger.debug(f"Saved attachment {attachment_id} to ") return attachment_id def get(self, attachment_id: str) -> Optional[dict]: """ Retrieve a full attachment dict by ID. :param attachment_id: The attachment UUID. :return: Attachment dict or None if not found. """ key = f"{self._session_id}:{attachment_id}" # e.g., raw = self._client.get(key) # return json.loads(raw) if raw else None return None def delete(self, attachment_id: str) -> None: """ Delete an attachment by ID. :param attachment_id: The attachment UUID. """ key = f"{self._session_id}:{attachment_id}" # e.g., self._client.delete(key) # IMPORTANT: also remove the ID from the session index so the count # stays accurate and pruning logic works correctly on future saves. # e.g., fetch the index, remove attachment_id, and write it back logger.debug(f"Deleted attachment {attachment_id} from ") ``` ### Key Implementation Notes - **Key format**: Use `{session_id}:{attachment_id}` for isolation between sessions - **Max attachments pruning**: When `save()` is called and the session exceeds `max_attachments`, delete the oldest entry. Track order via timestamps or a per-session index - **Index consistency on delete**: When `delete()` is called, remove the attachment ID from the per-session index in addition to deleting the data entry. Skipping this step causes the index count to drift — the backend will think more attachments exist than actually do, breaking `max_attachments` enforcement on subsequent saves - **TTL**: If the backend supports time-based expiry (like Redis TTL or DynamoDB `expiry_time`), use it for automatic cleanup. Read the TTL value from the backend-specific config - **Connection management**: Reuse the shared connection drivers in `agentkernel/core/util/driver/` (`RedisDriver`, `DynamoDBDriver`, etc.) — they provide lazy connect, 3-retry back-off, and (for Redis-like backends) ping/reconnect. The store may be instantiated per-request (inside `AttachmentStorageManager.__init__`), so keep the driver construction cheap (no eager connect) - **Serialization**: Attachment dicts must be JSON-serializable. The `data` field contains base64-encoded binary, so all values are strings, floats, or ints ### 3. Add Backend-Specific Configuration Update `ak-py/src/agentkernel/core/config.py`: ```python class _MultimodalStorageConfig(BaseModel): """Configuration for multimodal attachment storage.""" # Add backend-specific fields, e.g.: endpoint: str = Field(default="...", description=" endpoint URL") ttl: int = Field(default=604800, description="Attachment TTL in seconds") # table_name, bucket, prefix, etc. class _MultimodalConfig(BaseModel): # ... existing fields ... storage_type: str = Field( default="in_memory", description="Storage backend for multimodal attachments: a built-in short name " "(session_cache, in_memory, redis, dynamodb, ) or a dotted " "path to an AttachmentStore subclass", # ADD to the short-name list ) # ... existing backends ... : Optional[_MultimodalStorageConfig] = None # ADD THIS ``` `storage_type` is a free-form string (no regex `pattern=`) so that a dotted path resolves as bring-your-own — do not add a `pattern=` constraint back. ### 4. Register with the Storage Manager Factory Update `AttachmentStorageManager._build_driver()` in `ak-py/src/agentkernel/core/multimodal/storage/storage_manager.py`. It shares the house pluggable-backend shape from `core/util/factory.py` (`resolve_dotted`, `require_extra`, `AKConfigError` — the same pattern used by the guardrail, trace, session/thread store, and sandbox provider factories): matching is case-insensitive on `storage_type.lower()`, each built-in's lazy import for an optional-dependency backend is wrapped in `require_extra`, and anything left over is treated as a dotted path to an `AttachmentStore` subclass (bring-your-own): ```python _BUILTIN_ATTACHMENT_STORES = ["session_cache", "in_memory", "redis", "dynamodb", ""] # ADD @staticmethod def _build_driver(session_id: str) -> AttachmentStore: config = AKConfig.get().multimodal storage_type = config.storage_type key = storage_type.lower() # ... existing backends (session_cache, in_memory) ... if key == "": # ADD THIS with require_extra("", "multimodal.storage_type: "): from . import AttachmentStore backend_config = config. if backend_config is None: raise ValueError( "Multimodal storage_type is '' but no '' configuration " "is provided under 'multimodal'. Please set AK_MULTIMODAL____ENDPOINT etc." ) return AttachmentStore( session_id=session_id, endpoint=backend_config.endpoint, ttl=backend_config.ttl, ) # ... existing backends (redis, dynamodb) ... # Bring-your-own: a dotted path to an AttachmentStore subclass (session-scoped). if "." not in storage_type: raise AKConfigError( f"unknown multimodal storage_type '{storage_type}'; expected one of {_BUILTIN_ATTACHMENT_STORES} or a dotted path to an AttachmentStore subclass" ) return resolve_dotted(storage_type, base=AttachmentStore)(session_id) ``` A dotted `storage_type` (e.g. `myorg.storage.CustomAttachmentStore`) resolves via `resolve_dotted` without any factory edit at all — only add an `if` branch here for a first-party, in-repo backend you want addressable by a short name. There is no `else` fallback to `in_memory` anymore: an unrecognized non-dotted value now fails loudly via `AKConfigError`. ### 5. Add Optional Dependencies In `ak-py/pyproject.toml`: ```toml [project.optional-dependencies] = [ "backend-sdk>=x.y.z", ] ``` ### 6. Add Tests Create `ak-py/tests/test_multimodal_storage_.py`: ```python import pytest from agentkernel.core.multimodal.storage. import AttachmentStore class TestBasicOperations: """Test save/get/delete operations.""" def test_save_and_get(self): store = AttachmentStore(session_id="test-session", ...) attachment = { "id": "att-1", "type": "image", "data": "base64data...", "name": "test.jpg", "mime_type": "image/jpeg", "description": "A test image", "timestamp": 1234567890.0, } result_id = store.save(attachment, max_attachments=10) assert result_id == "att-1" retrieved = store.get("att-1") assert retrieved is not None assert retrieved["id"] == "att-1" assert retrieved["data"] == "base64data..." def test_get_nonexistent(self): store = AttachmentStore(session_id="test-session", ...) assert store.get("nonexistent") is None def test_delete(self): store = AttachmentStore(session_id="test-session", ...) attachment = { "id": "att-2", "type": "file", "data": "...", "name": "doc.pdf", "mime_type": "application/pdf", "description": "A PDF", "timestamp": 1234567890.0, } store.save(attachment, max_attachments=10) store.delete("att-2") assert store.get("att-2") is None class TestMaxAttachments: """Test that oldest attachments are pruned when limit is exceeded.""" def test_prune_oldest(self): store = AttachmentStore(session_id="test-session", ...) for i in range(5): store.save( {"id": f"att-{i}", "type": "image", "data": "...", "name": f"img{i}.jpg", "mime_type": "image/jpeg", "description": f"Image {i}", "timestamp": float(i)}, max_attachments=3, ) # Oldest (att-0, att-1) should be pruned assert store.get("att-0") is None assert store.get("att-1") is None assert store.get("att-4") is not None class TestSessionIsolation: """Test that attachments from different sessions are isolated.""" def test_isolation(self): store_a = AttachmentStore(session_id="session-a", ...) store_b = AttachmentStore(session_id="session-b", ...) store_a.save( {"id": "att-1", "type": "image", "data": "a-data", "name": "a.jpg", "mime_type": "image/jpeg", "description": "", "timestamp": 1.0}, max_attachments=10, ) assert store_a.get("att-1") is not None assert store_b.get("att-1") is None # different session ``` ### 7. Add Configuration Example Show the config in `config.yaml`: ```yaml multimodal: enabled: true storage_type: max_attachments: 20 description_model: gpt-4o analysis_model: gpt-4o : endpoint: "https://..." ttl: 604800 ``` ### 8. Add Documentation Add or update `docs/docs/advanced/multimodal.md` with: - Backend description and when to use it - Configuration reference - Required environment variables or credentials - Any infrastructure setup steps (e.g., creating tables, buckets) Then update the landing page inventories in `docs/src/components/*/data.tsx`: add the backend to the **Attachment Storage** card's `tags` under the Remember tab in `FeatureExplorer/data.tsx`; if the vendor is new to the site, add a tile to the **Memory, knowledge & data** row in `IntegrationsMarquee/data.tsx` (role `Memory`, every store it backs listed in `title`), and if it already has a tile for a session or thread store, append the attachment role to that tile's `title` instead. Logo sourcing and the build check are in `ak-dev-sync-docs-from-branch`, *Docs-Site Landing and Features Pages*. Grep `docs/src/pages/*.tsx` for the existing backend names (for example "DynamoDB") in case a features page card enumerates storage backends; the Smart Memory Management card in `docs/src/pages/features.tsx` lists session backends and is the usual place such a roll call appears. ## Reference: Existing Implementations ### Redis (`storage/redis.py`) - JSON serialization per attachment - Key format: `{prefix}{session_id}:{attachment_id}` - Connection pooling with lazy `_ensure_connection()` - TTL support via `client.set(key, json, ex=ttl)` - Per-session index for max attachment enforcement - Connection retry: up to 3 attempts ### DynamoDB (`storage/dynamodb.py`) - Partition key: `session_id`, sort key: `attachment_id` - TTL attribute: `expiry_time` (Unix epoch) - Boto3 client with lazy initialization - JSON serialization of attachment data - Per-session index item (`attachment_id = "_index"`) tracking ordered IDs for max attachment enforcement — same index-list pattern as Redis, read via `get_item` (no queries) ### In-Memory (`storage/in_memory.py`) - `ClassVar` dict shared across all instances - Key format: `{session_id}:{attachment_id}` - Order tracked via a per-session `_index` list in insertion order (stored timestamps are not consulted) - No persistence — lost on process restart ## Checklist - [ ] `ak-py/src/agentkernel/core/multimodal/storage/.py` implementing `AttachmentStore` - [ ] Backend-specific config class in `config.py` (e.g., `_MultimodalStorageConfig`) - [ ] `storage_type` description updated in `_MultimodalConfig` to list the new short name (no `pattern=` regex — the field stays free-form for bring-your-own) - [ ] Registration in `AttachmentStorageManager._build_driver()` factory - [ ] Optional dependencies in `pyproject.toml` - [ ] Unit tests for save/get/delete, max attachments pruning, session isolation - [ ] Configuration example in documentation - [ ] Documentation in `docs/docs/advanced/multimodal.md` - [ ] Landing page inventories: Attachment Storage card tags (`FeatureExplorer/data.tsx`), marquee tile or role (`IntegrationsMarquee/data.tsx`)