Ak Dev New Multimodal Storage

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()

# 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; 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 <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>")

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",
        description="Storage backend for multimodal attachments: a built-in short name "
                     "(session_cache, in_memory, redis, dynamodb, <backend>) or a dotted "
                     "path to an AttachmentStore subclass",  # ADD <backend> to the short-name list
    )
    # ... existing backends ...
    <backend>: Optional[_MultimodalStorage<Backend>Config] = None       # ADD THIS

_BUILTIN_ATTACHMENT_STORES = ["session_cache", "in_memory", "redis", "dynamodb", "<backend>"]  # ADD <backend>

@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 == "<backend>":                                                # ADD THIS
        with require_extra("<backend>", "multimodal.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,
        )

    # ... 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)

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

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

multimodal:
  enabled: true
  storage_type: <backend>
  max_attachments: 20
  description_model: gpt-4o
  analysis_model: gpt-4o
  <backend>:
    endpoint: "https://..."
    ttl: 604800