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

chuk_ai_session_manager/memory/prefetcher.py

@@ -0,0 +1,192 @@
+# chuk_ai_session_manager/memory/prefetcher.py
+"""
+Simple Prefetcher for AI Virtual Memory.
+
+Don't wait for ML prediction. Start with dumb heuristics that work:
+1. Last segment summary (almost always needed for "what did we discuss")
+2. Most-referenced claim pages (high access_count claims)
+3. Tool traces for likely tools (based on recent tool usage)
+
+Keep it simple: no ML prediction required.
+"""
+
+from collections import Counter
+from typing import TYPE_CHECKING, Dict, List, Optional, Set
+
+from pydantic import BaseModel, Field, PrivateAttr
+
+from .models import PageType
+
+if TYPE_CHECKING:
+    from .page_table import PageTable
+
+
+class ToolUsagePattern(BaseModel):
+    """Track tool usage patterns for prefetch prediction."""
+
+    tool_name: str
+    call_count: int = Field(default=0)
+    last_turn: int = Field(default=0)
+    prereq_pages: List[str] = Field(
+        default_factory=list
+    )  # Pages often read before calling
+
+
+class SimplePrefetcher(BaseModel):
+    """
+    Basic prefetch that doesn't need prediction models.
+    Good enough for v0.9.
+    """
+
+    # How many claim pages to prefetch
+    max_claims_to_prefetch: int = Field(default=3)
+
+    # How many recent tools to consider
+    max_recent_tools: int = Field(default=3)
+
+    # Tool usage tracking
+    _tool_usage: Dict[str, ToolUsagePattern] = PrivateAttr(default_factory=dict)
+
+    # Recent page accesses per tool (tool_name -> list of page_ids accessed before call)
+    _tool_prereqs: Dict[str, List[str]] = PrivateAttr(default_factory=dict)
+
+    # Page access counts (for finding high-value claims)
+    _page_access_counts: Counter = PrivateAttr(default_factory=Counter)
+
+    # Last segment summary page ID (updated when segments roll)
+    _last_segment_summary_id: Optional[str] = PrivateAttr(default=None)
+
+    def record_page_access(self, page_id: str) -> None:
+        """Record a page access for statistics."""
+        self._page_access_counts[page_id] += 1
+
+    def record_tool_call(
+        self,
+        tool_name: str,
+        turn: int,
+        pages_accessed_before: Optional[List[str]] = None,
+    ) -> None:
+        """
+        Record a tool call and what pages were accessed before it.
+
+        This builds up patterns like:
+        "When calling weather_tool, we usually read location claims first"
+        """
+        if tool_name not in self._tool_usage:
+            self._tool_usage[tool_name] = ToolUsagePattern(tool_name=tool_name)
+
+        pattern = self._tool_usage[tool_name]
+        pattern.call_count += 1
+        pattern.last_turn = turn
+
+        # Track pages accessed before this tool call
+        if pages_accessed_before:
+            if tool_name not in self._tool_prereqs:
+                self._tool_prereqs[tool_name] = []
+            self._tool_prereqs[tool_name].extend(pages_accessed_before)
+
+    def set_last_segment_summary(self, page_id: str) -> None:
+        """Update the last segment summary page ID."""
+        self._last_segment_summary_id = page_id
+
+    def get_likely_tools(self, recent_turns: int = 5) -> List[str]:
+        """Get tools likely to be called based on recent usage."""
+        # Sort by recency and frequency
+        tools = list(self._tool_usage.values())
+        tools.sort(key=lambda t: (t.last_turn, t.call_count), reverse=True)
+        return [t.tool_name for t in tools[: self.max_recent_tools]]
+
+    def get_tool_prereq_pages(self, tool_name: str) -> List[str]:
+        """
+        Get pages commonly accessed before calling a tool.
+
+        Returns most frequently accessed prereq pages.
+        """
+        prereqs = self._tool_prereqs.get(tool_name, [])
+        if not prereqs:
+            return []
+
+        # Count occurrences and return top ones
+        counts = Counter(prereqs)
+        return [page_id for page_id, _ in counts.most_common(3)]
+
+    def get_top_claims(
+        self,
+        page_table: Optional["PageTable"] = None,
+        limit: Optional[int] = None,
+    ) -> List[str]:
+        """
+        Get most-referenced claim pages.
+
+        If page_table is provided, filters to only claim-type pages.
+        """
+        if limit is None:
+            limit = self.max_claims_to_prefetch
+
+        # Get top accessed pages
+        top_pages = [
+            page_id for page_id, _ in self._page_access_counts.most_common(limit * 3)
+        ]
+
+        if page_table is None:
+            return top_pages[:limit]
+
+        # Filter to claims only
+        claims = []
+        for page_id in top_pages:
+            entry = page_table.lookup(page_id)
+            if entry and entry.page_type == PageType.CLAIM:
+                claims.append(page_id)
+                if len(claims) >= limit:
+                    break
+
+        return claims
+
+    async def prefetch_on_turn_start(
+        self,
+        session_id: str,
+        page_table: Optional["PageTable"] = None,
+    ) -> List[str]:
+        """
+        Get pages to prefetch at the start of a turn.
+
+        Returns list of page_ids to prefetch.
+        """
+        pages_to_prefetch: List[str] = []
+        seen: Set[str] = set()
+
+        def add_page(page_id: str) -> None:
+            if page_id and page_id not in seen:
+                pages_to_prefetch.append(page_id)
+                seen.add(page_id)
+
+        # 1. Last segment summary (almost always needed for "what did we discuss")
+        if self._last_segment_summary_id:
+            add_page(self._last_segment_summary_id)
+
+        # 2. Most-referenced claim pages (high access_count claims)
+        for claim_id in self.get_top_claims(page_table):
+            add_page(claim_id)
+
+        # 3. Tool prereqs for likely tools
+        likely_tools = self.get_likely_tools()
+        for tool_name in likely_tools:
+            for prereq_id in self.get_tool_prereq_pages(tool_name):
+                add_page(prereq_id)
+
+        return pages_to_prefetch
+
+    def clear(self) -> None:
+        """Clear all tracking data."""
+        self._tool_usage.clear()
+        self._tool_prereqs.clear()
+        self._page_access_counts.clear()
+        self._last_segment_summary_id = None
+
+    def get_stats(self) -> Dict[str, int]:
+        """Get prefetcher statistics."""
+        return {
+            "tools_tracked": len(self._tool_usage),
+            "pages_tracked": len(self._page_access_counts),
+            "total_accesses": sum(self._page_access_counts.values()),
+        }

chuk_ai_session_manager/memory/tlb.py

@@ -0,0 +1,247 @@
+# chuk_ai_session_manager/memory/tlb.py
+"""
+Translation Lookaside Buffer (TLB) for AI Virtual Memory.
+
+The TLB is a small, fast cache for recently accessed page table entries.
+It avoids the overhead of full PageTable lookups for hot pages.
+
+Without a TLB, every page access requires a PageTable lookup plus potentially
+a storage lookup. This becomes a bottleneck at scale - you end up measuring
+metadata latency, not content latency.
+
+The TLB provides O(1) lookups for recently accessed pages with LRU eviction.
+
+Design principles:
+- Pydantic-native: BaseModel subclass with proper validation
+- No magic strings: Uses StorageTier enum
+- Type-safe: Full type annotations throughout
+"""
+
+from collections import OrderedDict
+from typing import Optional
+
+from pydantic import BaseModel, Field, PrivateAttr
+
+from .models import CombinedPageTableStats, PageTableEntry, TLBStats
+
+
+class PageTLB(BaseModel):
+    """
+    Translation Lookaside Buffer - fast cache for page table entries.
+
+    Uses LRU eviction to keep the most recently accessed entries.
+
+    Typical usage:
+        1. Check TLB first (O(1))
+        2. If miss, check PageTable
+        3. Insert result into TLB
+    """
+
+    max_entries: int = Field(default=512, description="Maximum entries to cache")
+
+    # LRU cache: OrderedDict maintains insertion order, move_to_end on access
+    _entries: OrderedDict = PrivateAttr(default_factory=OrderedDict)
+
+    # Stats
+    hits: int = Field(default=0)
+    misses: int = Field(default=0)
+
+    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 in the TLB.
+
+        Returns the entry if found (TLB hit), None otherwise (TLB miss).
+        Updates LRU order on hit.
+        """
+        if page_id in self._entries:
+            # Move to end (most recently used)
+            self._entries.move_to_end(page_id)
+            self.hits += 1
+            return self._entries[page_id]
+        else:
+            self.misses += 1
+            return None
+
+    def insert(self, entry: PageTableEntry) -> None:
+        """
+        Insert or update an entry in the TLB.
+
+        If the TLB is full, evicts the least recently used entry.
+        """
+        page_id = entry.page_id
+
+        # If already present, update and move to end
+        if page_id in self._entries:
+            self._entries[page_id] = entry
+            self._entries.move_to_end(page_id)
+            return
+
+        # Check if we need to evict
+        if len(self._entries) >= self.max_entries:
+            self._evict_lru()
+
+        # Insert new entry
+        self._entries[page_id] = entry
+
+    def _evict_lru(self) -> Optional[str]:
+        """
+        Evict the least recently used entry.
+
+        Returns the evicted page_id, or None if TLB was empty.
+        """
+        if not self._entries:
+            return None
+
+        # OrderedDict.popitem(last=False) removes first item (oldest)
+        page_id, _ = self._entries.popitem(last=False)
+        return page_id
+
+    def invalidate(self, page_id: str) -> bool:
+        """
+        Remove an entry from the TLB.
+
+        Call this when a page is modified, evicted, or deleted.
+        Returns True if the entry was present and removed.
+        """
+        if page_id in self._entries:
+            del self._entries[page_id]
+            return True
+        return False
+
+    def invalidate_tier(self, tier: str) -> int:
+        """
+        Invalidate all entries in a specific tier.
+
+        Useful when flushing a tier or during checkpoints.
+        Returns the number of entries invalidated.
+        """
+        to_remove = [
+            page_id for page_id, entry in self._entries.items() if entry.tier == tier
+        ]
+        for page_id in to_remove:
+            del self._entries[page_id]
+        return len(to_remove)
+
+    def flush(self) -> int:
+        """
+        Clear the entire TLB.
+
+        Call this on context switches, checkpoints, or when consistency
+        with the PageTable is required.
+        Returns the number of entries cleared.
+        """
+        count = len(self._entries)
+        self._entries.clear()
+        return count
+
+    def get_all(self) -> list[PageTableEntry]:
+        """Get all cached entries (for debugging/inspection)."""
+        return list(self._entries.values())
+
+    @property
+    def hit_rate(self) -> float:
+        """Calculate the TLB hit rate."""
+        total = self.hits + self.misses
+        if total == 0:
+            return 0.0
+        return self.hits / total
+
+    def reset_stats(self) -> None:
+        """Reset hit/miss counters."""
+        self.hits = 0
+        self.misses = 0
+
+    def get_stats(self) -> TLBStats:
+        """Get TLB statistics."""
+        return TLBStats(
+            size=len(self._entries),
+            max_size=self.max_entries,
+            utilization=len(self._entries) / self.max_entries
+            if self.max_entries > 0
+            else 0,
+            hits=self.hits,
+            misses=self.misses,
+            hit_rate=self.hit_rate,
+        )
+
+
+class TLBWithPageTable:
+    """
+    Convenience wrapper that combines TLB with PageTable for lookups.
+
+    Provides a unified interface that automatically checks TLB first,
+    then falls back to PageTable, and keeps TLB updated.
+    """
+
+    def __init__(self, page_table, tlb: Optional[PageTLB] = None):
+        """
+        Initialize with a PageTable and optional TLB.
+
+        If no TLB is provided, creates one with default settings.
+        """
+        from .page_table import PageTable
+
+        self.page_table: PageTable = page_table
+        # Use 'is None' check because empty PageTLB has __len__=0 which is falsy
+        self.tlb = tlb if tlb is not None else PageTLB()
+
+    def lookup(self, page_id: str) -> Optional[PageTableEntry]:
+        """
+        Look up a page, checking TLB first.
+
+        Returns the entry if found, None otherwise.
+        Automatically populates TLB on cache miss.
+        """
+        # Check TLB first
+        entry = self.tlb.lookup(page_id)
+        if entry is not None:
+            return entry
+
+        # TLB miss - check page table
+        entry = self.page_table.lookup(page_id)
+        if entry is not None:
+            # Populate TLB for next time
+            self.tlb.insert(entry)
+
+        return entry
+
+    def register(self, page) -> PageTableEntry:
+        """Register a page in both PageTable and TLB."""
+        entry = self.page_table.register(page)
+        self.tlb.insert(entry)
+        return entry
+
+    def update_location(self, page_id: str, **kwargs) -> bool:
+        """Update location in PageTable and invalidate TLB entry."""
+        success = self.page_table.update_location(page_id, **kwargs)
+        if success:
+            # Invalidate stale TLB entry - will be refreshed on next lookup
+            self.tlb.invalidate(page_id)
+        return success
+
+    def mark_dirty(self, page_id: str) -> bool:
+        """Mark page dirty and invalidate TLB."""
+        success = self.page_table.mark_dirty(page_id)
+        if success:
+            self.tlb.invalidate(page_id)
+        return success
+
+    def remove(self, page_id: str) -> Optional[PageTableEntry]:
+        """Remove from both PageTable and TLB."""
+        self.tlb.invalidate(page_id)
+        return self.page_table.remove(page_id)
+
+    def get_stats(self) -> CombinedPageTableStats:
+        """Get combined stats."""
+        return CombinedPageTableStats(
+            page_table=self.page_table.get_stats(),
+            tlb=self.tlb.get_stats(),
+        )

chuk_ai_session_manager/memory/vm_prompts.py

@@ -0,0 +1,238 @@
+# chuk_ai_session_manager/memory/vm_prompts.py
+"""
+Virtual Memory system prompts for Chat Completions integration.
+
+These prompts enforce VM semantics when injected into the developer message.
+
+Design principles:
+- Pydantic-native: Tool definitions as proper models
+- No magic strings: Uses enums for modes and types
+"""
+
+from typing import List
+
+
+from .models import (
+    Modality,
+    ToolDefinition,
+    ToolFunction,
+    ToolParameter,
+    ToolParameters,
+    ToolType,
+    VMMode,
+)
+
+
+# Strict mode: No hallucinated memory, citations required
+VM_STRICT_PROMPT = """You are operating under STRICT Virtual Memory grounding rules.
+
+Your ONLY valid sources of information are:
+1) The content inside <VM:CONTEXT> (the currently mapped working set), and
+2) The content returned by tools (e.g., page_fault) in messages with role="tool".
+
+Everything listed in <VM:MANIFEST_JSON> is DISCOVERY METADATA ONLY.
+- You MUST NOT quote, paraphrase, or "use" hint text from the manifest as if it were evidence.
+- You MUST NOT assume details about unmapped pages.
+- Page IDs and modality/tier/level are allowed for navigation only.
+
+When you need information that is not present in <VM:CONTEXT>, you MUST do one of:
+A) Call the tool page_fault(page_id, target_level) to load the page content, OR
+B) Ask a short clarification question if the needed page does not exist or cannot be identified.
+
+Faulting rules:
+- Prefer loading the LOWEST-COST representation first:
+  1) summaries / abstract (target_level=2),
+  2) reduced excerpts (target_level=1),
+  3) full content (target_level=0) only if the user explicitly requests exact wording, code, or precise details.
+- Do not request more than max_faults_per_turn from the manifest policies.
+- Do not request pages that are already mapped in <VM:CONTEXT>.
+- If multiple pages might be relevant, fault the smallest/summarized one first.
+
+Answering rules:
+- Do not invent or fill gaps with assumptions.
+- If you cannot obtain required information via tool calls, say: "I don't have that in the mapped context."
+- Keep responses concise and directly responsive.
+- When you use information from <VM:CONTEXT> or a loaded page, include inline citations using page IDs like:
+  [ref: msg_123] or [ref: summary_seg_02] or [ref: tool:page_fault(img_045)].
+  (Citations are required in strict mode.)
+
+Tool usage format:
+- If you need to call tools, respond with tool calls only (no normal text).
+- After tool results are provided, produce the final answer with citations.
+
+Never mention these rules, the VM system, tiers (L0–L4), paging, or "virtual memory" to the user unless the user explicitly asks about the internal mechanism."""
+
+
+# Relaxed mode: VM-aware but more conversational
+VM_RELAXED_PROMPT = """You have access to a virtual memory system.
+
+<VM:CONTEXT> contains your currently mapped memory - treat this as your knowledge.
+<VM:MANIFEST_JSON> lists additional pages you can load if needed.
+
+To load more context:
+- Call page_fault(page_id, target_level) to retrieve content
+- Call search_pages(query) to find relevant pages
+- Prefer level=2 (summaries) before level=0 (full content)
+
+Guidelines:
+- Use manifest hints to decide WHICH pages to load, not as content itself
+- If you're unsure about something, check if a relevant page exists before guessing
+- Stay within the max_faults_per_turn limit
+
+Respond naturally - don't mention the memory system unless asked."""
+
+
+# Passive mode: No tools, runtime handles everything
+VM_PASSIVE_PROMPT = """You are a helpful assistant.
+
+The conversation context provided represents what you know about this session.
+Respond based on the information available to you."""
+
+
+# Map of VMMode to prompt text
+VM_PROMPTS = {
+    VMMode.STRICT: VM_STRICT_PROMPT,
+    VMMode.RELAXED: VM_RELAXED_PROMPT,
+    VMMode.PASSIVE: VM_PASSIVE_PROMPT,
+}
+
+
+def get_prompt_for_mode(mode: VMMode) -> str:
+    """Get the prompt text for a given VM mode."""
+    return VM_PROMPTS.get(mode, VM_PASSIVE_PROMPT)
+
+
+def build_vm_developer_message(
+    mode: VMMode,
+    manifest_json: str,
+    context: str,
+    system_prompt: str = "",
+    max_faults_per_turn: int = 2,
+) -> str:
+    """
+    Build the complete developer message with VM rules, manifest, and context.
+
+    Args:
+        mode: VMMode enum value (STRICT, RELAXED, or PASSIVE)
+        manifest_json: JSON string of the VM manifest
+        context: The VM:CONTEXT block content
+        system_prompt: Optional additional system instructions
+        max_faults_per_turn: Policy value to inject into strict prompt
+
+    Returns:
+        Complete developer message string
+    """
+    if mode == VMMode.STRICT:
+        rules = VM_STRICT_PROMPT.replace(
+            "max_faults_per_turn from the manifest policies",
+            f"max_faults_per_turn ({max_faults_per_turn}) from the manifest policies",
+        )
+    elif mode == VMMode.RELAXED:
+        rules = VM_RELAXED_PROMPT
+    else:
+        rules = VM_PASSIVE_PROMPT
+
+    parts: List[str] = []
+
+    if system_prompt:
+        parts.append(system_prompt)
+
+    if mode != VMMode.PASSIVE:
+        parts.append(f"<VM:RULES>\n{rules}\n</VM:RULES>")
+        parts.append(f"<VM:MANIFEST_JSON>\n{manifest_json}\n</VM:MANIFEST_JSON>")
+
+    parts.append(f"<VM:CONTEXT>\n{context}\n</VM:CONTEXT>")
+
+    return "\n\n".join(parts)
+
+
+# Tool definitions as Pydantic models
+PAGE_FAULT_TOOL = ToolDefinition(
+    type=ToolType.FUNCTION,
+    function=ToolFunction(
+        name="page_fault",
+        description="Load a memory page into context at specified compression level. Use when you need content from a known page_id.",
+        parameters=ToolParameters(
+            type="object",
+            properties={
+                "page_id": ToolParameter(
+                    type="string",
+                    description="ID of the page to load",
+                ),
+                "target_level": ToolParameter(
+                    type="integer",
+                    minimum=0,
+                    maximum=3,
+                    default=2,
+                    description="0=full, 1=reduced, 2=abstract/summary, 3=reference only",
+                ),
+            },
+            required=["page_id"],
+        ),
+    ),
+)
+
+SEARCH_PAGES_TOOL = ToolDefinition(
+    type=ToolType.FUNCTION,
+    function=ToolFunction(
+        name="search_pages",
+        description="Search for pages matching a query. Use when you need to find relevant pages but don't know their IDs.",
+        parameters=ToolParameters(
+            type="object",
+            properties={
+                "query": ToolParameter(
+                    type="string",
+                    description="Search query (semantic or keyword)",
+                ),
+                "modality": ToolParameter(
+                    type="string",
+                    enum=[m.value for m in Modality],
+                    description="Filter by content type",
+                ),
+                "limit": ToolParameter(
+                    type="integer",
+                    default=5,
+                    description="Maximum results to return",
+                ),
+            },
+            required=["query"],
+        ),
+    ),
+)
+
+
+# List of all VM tools as Pydantic models
+VM_TOOL_DEFINITIONS: List[ToolDefinition] = [PAGE_FAULT_TOOL, SEARCH_PAGES_TOOL]
+
+
+def get_vm_tools(include_search: bool = True) -> List[ToolDefinition]:
+    """
+    Get the VM tool definitions as Pydantic models.
+
+    Args:
+        include_search: Whether to include search_pages tool
+
+    Returns:
+        List of ToolDefinition models
+    """
+    if include_search:
+        return VM_TOOL_DEFINITIONS
+    return [PAGE_FAULT_TOOL]
+
+
+def get_vm_tools_as_dicts(include_search: bool = True) -> List[dict]:
+    """
+    Get the VM tool definitions as dicts (for Chat Completions API).
+
+    Args:
+        include_search: Whether to include search_pages tool
+
+    Returns:
+        List of tool definition dicts
+    """
+    tools = get_vm_tools(include_search)
+    return [tool.model_dump(exclude_none=True) for tool in tools]
+
+
+# Legacy: Keep VM_TOOLS as dicts for backward compatibility
+VM_TOOLS = get_vm_tools_as_dicts(include_search=True)