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

chuk_ai_session_manager/memory/mutation_log.py

@@ -0,0 +1,186 @@
+# chuk_ai_session_manager/memory/mutation_log.py
+"""
+Lite Mutation Log for AI Virtual Memory.
+
+Append-only log of page operations for:
+- Debugging: "What was in context for turn T?"
+- Replay: Reconstruct state for testing
+- Grounding story: Prove what the model saw
+
+This is a lite version focused on basic tracking.
+Full event-sourcing with time-travel comes in v0.15.
+"""
+
+import uuid
+from collections import defaultdict
+from datetime import datetime
+from typing import Dict, List, Optional
+
+from pydantic import BaseModel, Field, PrivateAttr
+
+from .models import (
+    Actor,
+    MutationType,
+    PageMutation,
+    StorageTier,
+)
+
+
+class ContextSnapshot(BaseModel):
+    """Snapshot of what pages were in context at a given turn."""
+
+    turn: int
+    page_ids: List[str] = Field(default_factory=list)
+    timestamp: datetime = Field(default_factory=datetime.utcnow)
+
+
+class MutationLogLite(BaseModel):
+    """
+    Append-only log of page operations.
+
+    Not full event-sourcing, but enough for:
+    - Debugging: "what was in context for turn T?"
+    - Replay: Reconstruct state for testing
+    - Grounding story: Prove what the model saw
+    """
+
+    session_id: str = Field(default="")
+
+    # Append-only mutation list
+    _mutations: List[PageMutation] = PrivateAttr(default_factory=list)
+
+    # Index: page_id -> list of mutations
+    _by_page: Dict[str, List[PageMutation]] = PrivateAttr(
+        default_factory=lambda: defaultdict(list)
+    )
+
+    # Context snapshots per turn
+    _context_snapshots: Dict[int, ContextSnapshot] = PrivateAttr(default_factory=dict)
+
+    def append(self, mutation: PageMutation) -> None:
+        """Append a mutation to the log."""
+        self._mutations.append(mutation)
+        self._by_page[mutation.page_id].append(mutation)
+
+    def record_mutation(
+        self,
+        page_id: str,
+        mutation_type: MutationType,
+        tier_after: StorageTier,
+        tier_before: Optional[StorageTier] = None,
+        actor: Actor = Actor.SYSTEM,
+        cause: Optional[str] = None,
+        turn: int = 0,
+    ) -> PageMutation:
+        """Create and append a mutation record."""
+        mutation = PageMutation(
+            mutation_id=str(uuid.uuid4())[:8],
+            page_id=page_id,
+            turn=turn,
+            mutation_type=mutation_type,
+            tier_before=tier_before,
+            tier_after=tier_after,
+            actor=actor,
+            cause=cause,
+        )
+        self.append(mutation)
+        return mutation
+
+    def record_context_at_turn(self, turn: int, page_ids: List[str]) -> None:
+        """
+        Snapshot the context (L0 pages) at a turn.
+
+        Call this at the start or end of each turn to enable
+        "what was in context for turn T?" queries.
+        """
+        self._context_snapshots[turn] = ContextSnapshot(
+            turn=turn,
+            page_ids=list(page_ids),
+        )
+
+    def get_context_at_turn(self, turn: int) -> List[str]:
+        """
+        Replay: what page_ids were in L0 at turn T?
+
+        Returns empty list if no snapshot for that turn.
+        """
+        snapshot = self._context_snapshots.get(turn)
+        if snapshot:
+            return snapshot.page_ids
+        return []
+
+    def get_history(self, page_id: str) -> List[PageMutation]:
+        """All mutations for a page."""
+        return list(self._by_page.get(page_id, []))
+
+    def get_mutations_by_actor(self, actor: Actor) -> List[PageMutation]:
+        """All mutations by a specific actor."""
+        return [m for m in self._mutations if m.actor == actor]
+
+    def get_mutations_by_type(self, mutation_type: MutationType) -> List[PageMutation]:
+        """All mutations of a specific type."""
+        return [m for m in self._mutations if m.mutation_type == mutation_type]
+
+    def get_mutations_in_turn(self, turn: int) -> List[PageMutation]:
+        """All mutations in a specific turn."""
+        return [m for m in self._mutations if m.turn == turn]
+
+    def get_pages_created_in_turn(self, turn: int) -> List[str]:
+        """Get page IDs created in a specific turn."""
+        return [
+            m.page_id
+            for m in self._mutations
+            if m.turn == turn and m.mutation_type == MutationType.CREATE
+        ]
+
+    def get_pages_faulted_in_turn(self, turn: int) -> List[str]:
+        """Get page IDs faulted in during a specific turn."""
+        return [
+            m.page_id
+            for m in self._mutations
+            if m.turn == turn and m.mutation_type == MutationType.FAULT_IN
+        ]
+
+    def get_pages_evicted_in_turn(self, turn: int) -> List[str]:
+        """Get page IDs evicted in a specific turn."""
+        return [
+            m.page_id
+            for m in self._mutations
+            if m.turn == turn and m.mutation_type == MutationType.EVICT
+        ]
+
+    def get_all_mutations(self) -> List[PageMutation]:
+        """Get all mutations in chronological order."""
+        return list(self._mutations)
+
+    def mutation_count(self) -> int:
+        """Total number of mutations logged."""
+        return len(self._mutations)
+
+    def page_count(self) -> int:
+        """Number of unique pages with mutations."""
+        return len(self._by_page)
+
+    def clear(self) -> None:
+        """Clear all mutations (for testing)."""
+        self._mutations.clear()
+        self._by_page.clear()
+        self._context_snapshots.clear()
+
+    def get_summary(self) -> Dict[str, int]:
+        """Get summary statistics."""
+        by_type: Dict[MutationType, int] = defaultdict(int)
+        for m in self._mutations:
+            by_type[m.mutation_type] += 1
+
+        return {
+            "total_mutations": len(self._mutations),
+            "unique_pages": len(self._by_page),
+            "context_snapshots": len(self._context_snapshots),
+            "creates": by_type.get(MutationType.CREATE, 0),
+            "faults": by_type.get(MutationType.FAULT_IN, 0),
+            "evictions": by_type.get(MutationType.EVICT, 0),
+            "compressions": by_type.get(MutationType.COMPRESS, 0),
+            "pins": by_type.get(MutationType.PIN, 0),
+            "unpins": by_type.get(MutationType.UNPIN, 0),
+        }

chuk_ai_session_manager/memory/pack_cache.py

@@ -0,0 +1,206 @@
+# chuk_ai_session_manager/memory/pack_cache.py
+"""
+Context Pack Cache - caches packed context to avoid re-packing on small incremental turns.
+
+The TLB caches address translations. But in LLM terms, the expensive part is often:
+- Selecting pages
+- Choosing compression levels
+- Rendering into VM:CONTEXT
+
+This cache stores the packed context output to avoid recomputing it when:
+- User says something short
+- Assistant responds
+- No pages changed
+
+Cache is keyed by: (session_id, model_id, token_budget, working_set_hash)
+"""
+
+import hashlib
+from collections import OrderedDict
+from datetime import datetime
+from typing import Dict, List, Optional
+
+from pydantic import BaseModel, Field
+
+
+class PackedContext(BaseModel):
+    """A cached packed context result."""
+
+    # The packed content
+    vm_context: str = Field(..., description="The VM:CONTEXT formatted content")
+    vm_manifest_json: str = Field(..., description="The VM:MANIFEST_JSON content")
+
+    # Metadata
+    page_ids: List[str] = Field(default_factory=list, description="Pages included")
+    tokens_used: int = Field(default=0)
+    created_at: datetime = Field(default_factory=datetime.utcnow)
+
+    # Cache key components (for debugging)
+    session_id: str = Field(default="")
+    model_id: str = Field(default="")
+    token_budget: int = Field(default=0)
+    working_set_hash: str = Field(default="")
+
+
+class CacheEntry(BaseModel):
+    """Internal cache entry with LRU tracking."""
+
+    packed: PackedContext
+    last_accessed: datetime = Field(default_factory=datetime.utcnow)
+    access_count: int = Field(default=0)
+
+
+class ContextPackCache:
+    """
+    Cache packed context to avoid re-packing on small incremental turns.
+    Invalidate on working set changes.
+
+    This drastically reduces overhead for "small incremental turns" -
+    user says something short, assistant responds, no pages changed, reuse the pack.
+    """
+
+    def __init__(self, max_entries: int = 32):
+        self.max_entries = max_entries
+        self._cache: OrderedDict[str, CacheEntry] = OrderedDict()
+        self._stats = {
+            "hits": 0,
+            "misses": 0,
+            "evictions": 0,
+            "invalidations": 0,
+        }
+
+    def _make_key(
+        self,
+        session_id: str,
+        model_id: str,
+        token_budget: int,
+        working_set_hash: str,
+    ) -> str:
+        """Create cache key from components."""
+        key_str = f"{session_id}:{model_id}:{token_budget}:{working_set_hash}"
+        return hashlib.sha256(key_str.encode()).hexdigest()[:16]
+
+    @staticmethod
+    def compute_working_set_hash(
+        page_ids: List[str], versions: Optional[Dict[str, int]] = None
+    ) -> str:
+        """
+        Compute a hash of the working set for cache key.
+
+        Args:
+            page_ids: List of page IDs in working set
+            versions: Optional dict of page_id -> version number
+
+        Returns:
+            Hash string representing the working set state
+        """
+        # Sort for consistency
+        sorted_ids = sorted(page_ids)
+
+        if versions:
+            # Include versions in hash
+            parts = [f"{pid}:{versions.get(pid, 0)}" for pid in sorted_ids]
+        else:
+            parts = sorted_ids
+
+        content = "|".join(parts)
+        return hashlib.sha256(content.encode()).hexdigest()[:16]
+
+    def get(
+        self,
+        session_id: str,
+        model_id: str,
+        token_budget: int,
+        working_set_hash: str,
+    ) -> Optional[PackedContext]:
+        """
+        O(1) lookup for cached pack.
+
+        Returns None if not found or expired.
+        """
+        key = self._make_key(session_id, model_id, token_budget, working_set_hash)
+
+        if key not in self._cache:
+            self._stats["misses"] += 1
+            return None
+
+        # Move to end (most recently used)
+        self._cache.move_to_end(key)
+
+        entry = self._cache[key]
+        entry.last_accessed = datetime.utcnow()
+        entry.access_count += 1
+
+        self._stats["hits"] += 1
+        return entry.packed
+
+    def put(
+        self,
+        session_id: str,
+        model_id: str,
+        token_budget: int,
+        working_set_hash: str,
+        packed: PackedContext,
+    ) -> None:
+        """Store packed context, evict LRU if full."""
+        key = self._make_key(session_id, model_id, token_budget, working_set_hash)
+
+        # Update packed metadata
+        packed.session_id = session_id
+        packed.model_id = model_id
+        packed.token_budget = token_budget
+        packed.working_set_hash = working_set_hash
+
+        # Evict if at capacity
+        while len(self._cache) >= self.max_entries:
+            # Remove oldest (first item)
+            self._cache.popitem(last=False)
+            self._stats["evictions"] += 1
+
+        self._cache[key] = CacheEntry(packed=packed)
+
+    def invalidate_session(self, session_id: str) -> int:
+        """
+        Invalidate all cached packs for session.
+
+        Called when working set changes.
+        Returns number of entries invalidated.
+        """
+        keys_to_remove = [
+            key
+            for key, entry in self._cache.items()
+            if entry.packed.session_id == session_id
+        ]
+
+        for key in keys_to_remove:
+            del self._cache[key]
+            self._stats["invalidations"] += 1
+
+        return len(keys_to_remove)
+
+    def invalidate_all(self) -> None:
+        """Clear entire cache."""
+        count = len(self._cache)
+        self._cache.clear()
+        self._stats["invalidations"] += count
+
+    @property
+    def size(self) -> int:
+        """Current number of cached entries."""
+        return len(self._cache)
+
+    @property
+    def hit_rate(self) -> float:
+        """Cache hit rate."""
+        total = self._stats["hits"] + self._stats["misses"]
+        if total == 0:
+            return 0.0
+        return self._stats["hits"] / total
+
+    def get_stats(self) -> Dict[str, int]:
+        """Get cache statistics."""
+        return {
+            **self._stats,
+            "size": len(self._cache),
+            "max_size": self.max_entries,
+        }

chuk_ai_session_manager/memory/page_table.py

@@ -0,0 +1,275 @@
+# chuk_ai_session_manager/memory/page_table.py
+"""
+Page Table implementation for AI Virtual Memory.
+
+The PageTable maps page IDs to their metadata (location, state, access patterns).
+It's the core data structure for tracking all pages in the system.
+
+Like an OS page table, it doesn't hold content - just metadata about where
+content lives and its current state.
+
+Design principles:
+- Pydantic-native: BaseModel subclass with proper validation
+- No magic strings: Uses StorageTier and Modality enums
+- Type-safe: Full type annotations throughout
+"""
+
+from datetime import datetime
+from typing import Dict, List, Optional, Set
+
+from pydantic import BaseModel, Field, PrivateAttr
+
+from .models import (
+    CompressionLevel,
+    MemoryPage,
+    Modality,
+    PageTableEntry,
+    PageTableStats,
+    StorageTier,
+)
+
+
+class PageTable(BaseModel):
+    """
+    Maps page IDs to their current location and state.
+
+    The page table is the source of truth for:
+    - Where each page lives (which tier)
+    - Whether it's dirty (needs flush)
+    - Access patterns (for eviction decisions)
+    """
+
+    # Core mapping
+    entries: Dict[str, PageTableEntry] = Field(default_factory=dict)
+
+    # Indexes for fast lookup - use PrivateAttr for internal state
+    _by_tier: Dict[StorageTier, Set[str]] = PrivateAttr(
+        default_factory=lambda: {t: set() for t in StorageTier}
+    )
+    _by_modality: Dict[Modality, Set[str]] = PrivateAttr(
+        default_factory=lambda: {m: set() for m in Modality}
+    )
+    _dirty_pages: Set[str] = PrivateAttr(default_factory=set)
+
+    model_config = {"arbitrary_types_allowed": True}
+
+    def __len__(self) -> int:
+        return len(self.entries)
+
+    def __contains__(self, page_id: str) -> bool:
+        return page_id in self.entries
+
+    def lookup(self, page_id: str) -> Optional[PageTableEntry]:
+        """
+        Look up a page entry by ID.
+
+        Returns None if page doesn't exist.
+        Does NOT update access tracking (use mark_accessed for that).
+        """
+        return self.entries.get(page_id)
+
+    def register(self, page: MemoryPage) -> PageTableEntry:
+        """
+        Register a new page in the table.
+
+        Creates a PageTableEntry from the MemoryPage and adds it to indexes.
+        """
+        entry = PageTableEntry(
+            page_id=page.page_id,
+            tier=page.storage_tier,
+            artifact_id=page.artifact_id,
+            compression_level=page.compression_level,
+            dirty=page.dirty,
+            last_accessed=page.last_accessed,
+            access_count=page.access_count,
+            size_tokens=page.size_tokens or page.estimate_tokens(),
+            modality=page.modality,
+            # v0.8 fields
+            page_type=page.page_type,
+            provenance=page.provenance,
+            pinned=page.pinned,
+        )
+
+        self._add_entry(entry)
+        return entry
+
+    def register_entry(self, entry: PageTableEntry) -> None:
+        """Register an existing entry directly."""
+        self._add_entry(entry)
+
+    def _add_entry(self, entry: PageTableEntry) -> None:
+        """Internal: add entry to table and indexes."""
+        page_id = entry.page_id
+
+        # Remove from old indexes if updating
+        if page_id in self.entries:
+            old_entry = self.entries[page_id]
+            self._by_tier.get(old_entry.tier, set()).discard(page_id)
+            self._by_modality.get(old_entry.modality, set()).discard(page_id)
+            if old_entry.dirty:
+                self._dirty_pages.discard(page_id)
+
+        # Add to table
+        self.entries[page_id] = entry
+
+        # Update indexes
+        if entry.tier not in self._by_tier:
+            self._by_tier[entry.tier] = set()
+        self._by_tier[entry.tier].add(page_id)
+
+        if entry.modality not in self._by_modality:
+            self._by_modality[entry.modality] = set()
+        self._by_modality[entry.modality].add(page_id)
+
+        if entry.dirty:
+            self._dirty_pages.add(page_id)
+
+    def remove(self, page_id: str) -> Optional[PageTableEntry]:
+        """
+        Remove a page from the table.
+
+        Returns the removed entry, or None if not found.
+        """
+        entry = self.entries.pop(page_id, None)
+        if entry:
+            self._by_tier.get(entry.tier, set()).discard(page_id)
+            self._by_modality.get(entry.modality, set()).discard(page_id)
+            self._dirty_pages.discard(page_id)
+        return entry
+
+    def update_location(
+        self,
+        page_id: str,
+        tier: StorageTier,
+        artifact_id: Optional[str] = None,
+        compression_level: Optional[CompressionLevel] = None,
+    ) -> bool:
+        """
+        Update a page's location (tier movement).
+
+        Returns True if successful, False if page not found.
+        """
+        entry = self.entries.get(page_id)
+        if not entry:
+            return False
+
+        # Update tier index
+        self._by_tier.get(entry.tier, set()).discard(page_id)
+        entry.tier = tier
+        if tier not in self._by_tier:
+            self._by_tier[tier] = set()
+        self._by_tier[tier].add(page_id)
+
+        # Update other fields if provided
+        if artifact_id is not None:
+            entry.artifact_id = artifact_id
+        if compression_level is not None:
+            entry.compression_level = compression_level
+
+        return True
+
+    def mark_accessed(self, page_id: str) -> bool:
+        """
+        Mark a page as accessed (updates LRU tracking).
+
+        Returns True if successful, False if page not found.
+        """
+        entry = self.entries.get(page_id)
+        if not entry:
+            return False
+
+        entry.mark_accessed()
+        return True
+
+    def mark_dirty(self, page_id: str) -> bool:
+        """
+        Mark a page as dirty (modified).
+
+        Returns True if successful, False if page not found.
+        """
+        entry = self.entries.get(page_id)
+        if not entry:
+            return False
+
+        entry.dirty = True
+        self._dirty_pages.add(page_id)
+        return True
+
+    def mark_clean(self, page_id: str) -> bool:
+        """
+        Mark a page as clean (flushed).
+
+        Returns True if successful, False if page not found.
+        """
+        entry = self.entries.get(page_id)
+        if not entry:
+            return False
+
+        entry.dirty = False
+        entry.last_flushed = datetime.utcnow()
+        self._dirty_pages.discard(page_id)
+        return True
+
+    def get_by_tier(self, tier: StorageTier) -> List[PageTableEntry]:
+        """Get all entries in a specific tier."""
+        page_ids = self._by_tier.get(tier, set())
+        return [self.entries[pid] for pid in page_ids if pid in self.entries]
+
+    def get_by_modality(self, modality: Modality) -> List[PageTableEntry]:
+        """Get all entries of a specific modality."""
+        page_ids = self._by_modality.get(modality, set())
+        return [self.entries[pid] for pid in page_ids if pid in self.entries]
+
+    def get_dirty_pages(self) -> List[PageTableEntry]:
+        """Get all dirty (modified) pages."""
+        return [self.entries[pid] for pid in self._dirty_pages if pid in self.entries]
+
+    def get_working_set(self) -> List[PageTableEntry]:
+        """Get all pages in L0 and L1 (the working set)."""
+        l0 = self.get_by_tier(StorageTier.L0)
+        l1 = self.get_by_tier(StorageTier.L1)
+        return l0 + l1
+
+    def get_eviction_candidates(
+        self,
+        tier: StorageTier = StorageTier.L1,
+        limit: int = 10,
+    ) -> List[PageTableEntry]:
+        """
+        Get pages that are candidates for eviction, sorted by LRU.
+
+        Returns oldest-accessed pages first.
+        """
+        entries = self.get_by_tier(tier)
+        # Sort by last_accessed (oldest first)
+        entries.sort(key=lambda e: e.last_accessed)
+        return entries[:limit]
+
+    def get_stats(self) -> "PageTableStats":
+        """Get page table statistics."""
+        return PageTableStats(
+            total_pages=len(self.entries),
+            dirty_pages=len(self._dirty_pages),
+            pages_by_tier={t: len(self._by_tier.get(t, set())) for t in StorageTier},
+            pages_by_modality={
+                m: len(self._by_modality.get(m, set())) for m in Modality
+            },
+        )
+
+    def get_total_tokens(self, tiers: Optional[List[StorageTier]] = None) -> int:
+        """
+        Get total estimated tokens across specified tiers.
+
+        If tiers is None, counts all pages.
+        """
+        total = 0
+        entries_list = list(self.entries.values())
+
+        if tiers:
+            entries_list = [e for e in entries_list if e.tier in tiers]
+
+        for entry in entries_list:
+            if entry.size_tokens:
+                total += entry.size_tokens
+
+        return total