ak-dev-new-multimodal-storage

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)
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/<backend>.py.

2. Implement the AttachmentStore

# ak-py/src/agentkernel/core/multimodal/storage/<backend>.py
import logging
from typing import Optional

from .base import AttachmentStore

logger = logging.getLogger("ak.core.multimodal.storage.<backend>")


class <Backend>AttachmentStore(AttachmentStore):
    """<Backend> 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"<Backend> 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)
        - name: str (filename)
        - mime_type: str
        - description: str (LLM-generated)
        - timestamp: float (epoch)

        :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 <backend>")
        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 <backend>")

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:

class _MultimodalStorage<Backend>Config(BaseModel):
    """Configuration for <backend> multimodal attachment storage."""
    # Add backend-specific fields, e.g.:
    endpoint: str = Field(default="...", description="<Backend> 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",
        pattern="^(session_cache|in_memory|redis|dynamodb|<backend>)$",  # ADD <backend>
        description="Storage backend for multimodal attachments.",
    )
    # ... existing backends ...
    <backend>: Optional[_MultimodalStorage<Backend>Config] = None       # ADD THIS

4. Register with the Storage Manager Factory

Update AttachmentStorageManager._build_driver() in ak-py/src/agentkernel/core/multimodal/storage/storage_manager.py:

@staticmethod
def _build_driver(session_id: str) -> AttachmentStore:
    config = AKConfig.get().multimodal
    storage_type = config.storage_type

    # ... existing backends ...

    elif storage_type == "<backend>":
        from .<backend> import <Backend>AttachmentStore

        backend_config = config.<backend>
        if backend_config is None:
            raise ValueError(
                "Multimodal storage_type is '<backend>' but no '<backend>' configuration "
                "is provided under 'multimodal'. Please set AK_MULTIMODAL__<BACKEND>__ENDPOINT etc."
            )
        return <Backend>AttachmentStore(
            session_id=session_id,
            endpoint=backend_config.endpoint,
            ttl=backend_config.ttl,
        )

    else:
        # Default: in_memory
        from .in_memory import InMemoryAttachmentStore
        return InMemoryAttachmentStore(session_id)

5. Add Optional Dependencies

In ak-py/pyproject.toml:

[project.optional-dependencies]
<backend> = [
    "backend-sdk>=x.y.z",
]

6. Add Tests

Create ak-py/tests/test_multimodal_storage_<backend>.py:

import pytest
from agentkernel.core.multimodal.storage.<backend> import <Backend>AttachmentStore


class TestBasicOperations:
    """Test save/get/delete operations."""

    def test_save_and_get(self):
        store = <Backend>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 = <Backend>AttachmentStore(session_id="test-session", ...)
        assert store.get("nonexistent") is None

    def test_delete(self):
        store = <Backend>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 = <Backend>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 = <Backend>AttachmentStore(session_id="session-a", ...)
        store_b = <Backend>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:

multimodal:
  enabled: true
  storage_type: <backend>
  max_attachments: 20
  description_model: gpt-4o
  analysis_model: gpt-4o
  <backend>:
    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)

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/<backend>.py implementing AttachmentStore
• Backend-specific config class in config.py (e.g., _MultimodalStorage<Backend>Config)
• storage_type pattern regex updated in _MultimodalConfig
• 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