chuk-ai-session-manager 0.7.1__py3-none-any.whl → 0.8__py3-none-any.whl

chuk_ai_session_manager/memory/artifacts_bridge.py

@@ -0,0 +1,469 @@
+# chuk_ai_session_manager/memory/artifacts_bridge.py
+"""
+Artifacts Bridge for AI Virtual Memory.
+
+The ArtifactsBridge integrates with chuk-artifacts for persistent page storage.
+It handles storing/loading pages to/from the L3 and L4 tiers.
+
+When chuk-artifacts is not installed, falls back to in-memory storage.
+
+Design principles:
+- Async-native: All I/O operations are async
+- Optional dependency: Works without chuk-artifacts
+- Pydantic-native: All models are BaseModel subclasses
+- No magic strings: Uses constants for MIME types
+"""
+
+from datetime import datetime
+from typing import Any, Dict, List, Optional, Protocol
+
+from pydantic import BaseModel, Field
+
+from .models import (
+    MEMORY_PAGE_MIME_TYPE,
+    VM_CHECKPOINT_MIME_TYPE,
+    MemoryPage,
+    StorageStats,
+    StorageTier,
+)
+
+# Try to import chuk-artifacts
+try:
+    from chuk_artifacts import ArtifactStore
+    from chuk_artifacts.models import StorageScope
+
+    ARTIFACTS_AVAILABLE = True
+except ImportError:
+    ARTIFACTS_AVAILABLE = False
+    ArtifactStore = None  # type: ignore
+    StorageScope = None  # type: ignore
+
+
+class StorageBackend(Protocol):
+    """Protocol for page storage backends."""
+
+    async def store(
+        self,
+        page: MemoryPage,
+        tier: StorageTier,
+    ) -> str:
+        """Store a page and return its artifact_id."""
+        ...
+
+    async def load(
+        self,
+        artifact_id: str,
+    ) -> Optional[MemoryPage]:
+        """Load a page by artifact_id."""
+        ...
+
+    async def delete(
+        self,
+        artifact_id: str,
+    ) -> bool:
+        """Delete a page by artifact_id."""
+        ...
+
+
+class PageMetadata(BaseModel):
+    """Metadata stored with a page in artifacts."""
+
+    page_id: str
+    modality: str
+    compression_level: int
+    tier: str
+    session_id: Optional[str] = None
+
+
+class CheckpointMetadata(BaseModel):
+    """Metadata stored with a checkpoint."""
+
+    checkpoint_name: str
+    page_count: int
+    session_id: Optional[str] = None
+
+
+class CheckpointEntry(BaseModel):
+    """Entry for a page in a checkpoint manifest."""
+
+    page_id: str
+    artifact_id: str
+    modality: str
+
+
+class CheckpointManifest(BaseModel):
+    """Manifest for a checkpoint."""
+
+    name: str
+    created_at: datetime
+    page_count: int
+    pages: List[CheckpointEntry] = Field(default_factory=list)
+
+
+class InMemoryBackend(BaseModel):
+    """
+    Simple in-memory storage backend for testing/development.
+
+    Not persistent - pages are lost when process exits.
+    """
+
+    pages: Dict[str, bytes] = Field(default_factory=dict)
+    counter: int = Field(default=0)
+
+    model_config = {"arbitrary_types_allowed": True}
+
+    async def store(
+        self,
+        page: MemoryPage,
+        tier: StorageTier,
+    ) -> str:
+        """Store a page in memory."""
+        self.counter += 1
+        artifact_id = f"mem_{tier.value}_{self.counter}"
+
+        # Serialize page to JSON bytes
+        data = page.model_dump_json().encode("utf-8")
+        self.pages[artifact_id] = data
+
+        return artifact_id
+
+    async def load(
+        self,
+        artifact_id: str,
+    ) -> Optional[MemoryPage]:
+        """Load a page from memory."""
+        data = self.pages.get(artifact_id)
+        if not data:
+            return None
+
+        # Deserialize
+        return MemoryPage.model_validate_json(data)
+
+    async def delete(
+        self,
+        artifact_id: str,
+    ) -> bool:
+        """Delete a page from memory."""
+        if artifact_id in self.pages:
+            del self.pages[artifact_id]
+            return True
+        return False
+
+    def clear(self) -> None:
+        """Clear all stored pages."""
+        self.pages.clear()
+        self.counter = 0
+
+
+class ArtifactsBridge(BaseModel):
+    """
+    Bridge between VM and chuk-artifacts storage.
+
+    Handles page serialization/deserialization and tier-appropriate storage.
+
+    Usage:
+        # With chuk-artifacts
+        bridge = ArtifactsBridge()
+        await bridge.configure(artifact_store)
+
+        # Without chuk-artifacts (in-memory fallback)
+        bridge = ArtifactsBridge()
+        await bridge.configure()  # Uses InMemoryBackend
+
+        # Store a page
+        artifact_id = await bridge.store_page(page, StorageTier.L3)
+
+        # Load a page
+        page = await bridge.load_page(artifact_id)
+    """
+
+    # Configuration
+    session_id: Optional[str] = None
+    _backend: Optional[InMemoryBackend] = None
+    _artifact_store: Optional[Any] = None
+    _using_artifacts: bool = False
+
+    model_config = {"arbitrary_types_allowed": True}
+
+    async def configure(
+        self,
+        artifact_store: Optional[Any] = None,
+        session_id: Optional[str] = None,
+    ) -> None:
+        """
+        Configure the bridge with a storage backend.
+
+        Args:
+            artifact_store: Optional ArtifactStore instance
+            session_id: Optional session ID for scoping
+
+        If no artifact_store is provided, uses in-memory backend.
+        """
+        self.session_id = session_id
+
+        if artifact_store and ARTIFACTS_AVAILABLE:
+            self._artifact_store = artifact_store
+            self._using_artifacts = True
+        else:
+            self._backend = InMemoryBackend()
+            self._using_artifacts = False
+
+    @property
+    def is_persistent(self) -> bool:
+        """Check if using persistent storage."""
+        return self._using_artifacts
+
+    async def store_page(
+        self,
+        page: MemoryPage,
+        tier: StorageTier = StorageTier.L3,
+    ) -> str:
+        """
+        Store a page to the specified tier.
+
+        Args:
+            page: MemoryPage to store
+            tier: Target storage tier (L3 or L4)
+
+        Returns:
+            Artifact ID for later retrieval
+        """
+        if self._using_artifacts and self._artifact_store:
+            return await self._store_with_artifacts(page, tier)
+        elif self._backend:
+            return await self._backend.store(page, tier)
+        else:
+            raise RuntimeError("ArtifactsBridge not configured")
+
+    async def _store_with_artifacts(
+        self,
+        page: MemoryPage,
+        tier: StorageTier,
+    ) -> str:
+        """Store using chuk-artifacts."""
+        # Serialize page
+        data = page.model_dump_json().encode("utf-8")
+
+        # Determine scope based on tier
+        if StorageScope:
+            scope = (
+                StorageScope.SESSION if tier == StorageTier.L3 else StorageScope.SANDBOX
+            )
+        else:
+            scope = "session"
+
+        # Build metadata
+        metadata = PageMetadata(
+            page_id=page.page_id,
+            modality=page.modality.value,
+            compression_level=page.compression_level
+            if isinstance(page.compression_level, int)
+            else page.compression_level.value,
+            tier=tier.value,
+            session_id=self.session_id,
+        )
+
+        # Store in artifacts
+        artifact_id = await self._artifact_store.store(
+            data=data,
+            mime=MEMORY_PAGE_MIME_TYPE,
+            scope=scope,
+            summary=f"VM page: {page.page_id}",
+            metadata=metadata.model_dump(),
+        )
+
+        return artifact_id
+
+    async def load_page(
+        self,
+        artifact_id: str,
+    ) -> Optional[MemoryPage]:
+        """
+        Load a page by artifact ID.
+
+        Args:
+            artifact_id: ID returned from store_page
+
+        Returns:
+            MemoryPage or None if not found
+        """
+        if self._using_artifacts and self._artifact_store:
+            return await self._load_with_artifacts(artifact_id)
+        elif self._backend:
+            return await self._backend.load(artifact_id)
+        else:
+            raise RuntimeError("ArtifactsBridge not configured")
+
+    async def _load_with_artifacts(
+        self,
+        artifact_id: str,
+    ) -> Optional[MemoryPage]:
+        """Load using chuk-artifacts."""
+        try:
+            # Get artifact data
+            data = await self._artifact_store.retrieve(artifact_id)
+            if not data:
+                return None
+
+            # Deserialize
+            if isinstance(data, bytes):
+                return MemoryPage.model_validate_json(data)
+            elif isinstance(data, str):
+                return MemoryPage.model_validate_json(data.encode())
+            else:
+                return None
+        except Exception:
+            return None
+
+    async def delete_page(
+        self,
+        artifact_id: str,
+    ) -> bool:
+        """
+        Delete a stored page.
+
+        Args:
+            artifact_id: ID of the artifact to delete
+
+        Returns:
+            True if deleted, False if not found
+        """
+        if self._using_artifacts and self._artifact_store:
+            try:
+                await self._artifact_store.delete(artifact_id)
+                return True
+            except Exception:
+                return False
+        elif self._backend:
+            return await self._backend.delete(artifact_id)
+        else:
+            raise RuntimeError("ArtifactsBridge not configured")
+
+    async def store_checkpoint(
+        self,
+        pages: List[MemoryPage],
+        checkpoint_name: str,
+    ) -> str:
+        """
+        Store a checkpoint of multiple pages.
+
+        Args:
+            pages: List of pages to checkpoint
+            checkpoint_name: Name for the checkpoint
+
+        Returns:
+            Checkpoint artifact ID
+        """
+        # Create checkpoint manifest
+        manifest = CheckpointManifest(
+            name=checkpoint_name,
+            created_at=datetime.utcnow(),
+            page_count=len(pages),
+            pages=[],
+        )
+
+        # Store each page and record artifact IDs
+        for page in pages:
+            artifact_id = await self.store_page(page, StorageTier.L3)
+            manifest.pages.append(
+                CheckpointEntry(
+                    page_id=page.page_id,
+                    artifact_id=artifact_id,
+                    modality=page.modality.value,
+                )
+            )
+
+        # Store manifest itself
+        manifest_data = manifest.model_dump_json().encode("utf-8")
+
+        if self._using_artifacts and self._artifact_store:
+            if StorageScope:
+                scope = StorageScope.SESSION
+            else:
+                scope = "session"
+
+            checkpoint_metadata = CheckpointMetadata(
+                checkpoint_name=checkpoint_name,
+                page_count=len(pages),
+                session_id=self.session_id,
+            )
+
+            checkpoint_id = await self._artifact_store.store(
+                data=manifest_data,
+                mime=VM_CHECKPOINT_MIME_TYPE,
+                scope=scope,
+                summary=f"VM checkpoint: {checkpoint_name}",
+                metadata=checkpoint_metadata.model_dump(),
+            )
+            return checkpoint_id
+        else:
+            # In-memory: store manifest directly
+            if self._backend:
+                self._backend.counter += 1
+                checkpoint_id = f"checkpoint_{self._backend.counter}"
+                self._backend.pages[checkpoint_id] = manifest_data
+                return checkpoint_id
+            raise RuntimeError("ArtifactsBridge not configured")
+
+    async def load_checkpoint(
+        self,
+        checkpoint_id: str,
+    ) -> List[MemoryPage]:
+        """
+        Load all pages from a checkpoint.
+
+        Args:
+            checkpoint_id: ID from store_checkpoint
+
+        Returns:
+            List of MemoryPages
+        """
+        # Load manifest
+        manifest_data: Optional[bytes] = None
+        if self._using_artifacts and self._artifact_store:
+            manifest_data = await self._artifact_store.retrieve(checkpoint_id)
+        elif self._backend:
+            manifest_data = self._backend.pages.get(checkpoint_id)
+        else:
+            raise RuntimeError("ArtifactsBridge not configured")
+
+        if not manifest_data:
+            return []
+
+        if isinstance(manifest_data, bytes):
+            manifest_str = manifest_data.decode("utf-8")
+        else:
+            manifest_str = manifest_data
+
+        manifest = CheckpointManifest.model_validate_json(manifest_str)
+
+        # Load each page
+        pages: List[MemoryPage] = []
+        for entry in manifest.pages:
+            page = await self.load_page(entry.artifact_id)
+            if page:
+                pages.append(page)
+
+        return pages
+
+    def get_stats(self) -> StorageStats:
+        """Get storage statistics."""
+        if self._using_artifacts:
+            return StorageStats(
+                backend="chuk-artifacts",
+                persistent=True,
+                session_id=self.session_id,
+                pages_stored=0,  # Would need to query artifact store
+            )
+        elif self._backend:
+            return StorageStats(
+                backend="in-memory",
+                persistent=False,
+                session_id=self.session_id,
+                pages_stored=len(self._backend.pages),
+            )
+        else:
+            return StorageStats(
+                backend="unconfigured",
+                persistent=False,
+            )