memory.py

"""
============================================================================
 MEMORY MODULE — Sliding Window Array with Index Locking
============================================================================
 Manages conversational state and token budgeting.
 
 Architecture (the Freight Train):
   Index [0]  = System Prompt (the Engine — NEVER detached)
   Index [1:] = Chat History  (cargo cars — oldest pruned when overweight)
   
 The memory manager counts tokens and automatically slices out the oldest
 user/assistant message pairs when approaching the context window limit.
============================================================================
"""

import os
import tiktoken
from pathlib import Path


class MemoryManager:
    """
    Sliding-window conversation memory with locked system prompt.
    
    The system prompt at index [0] is immutable and survives all pruning.
    Chat history grows at the tail and is pruned from the head (oldest first).
    """

    def __init__(self, system_prompt_path: str = "system_protocol.txt", max_tokens: int = None):
        """
        Initialize the memory manager.

        Args:
            system_prompt_path: Path to the system protocol file.
            max_tokens: Maximum token budget. Defaults to env var or 6000.
        """
        # Load the system prompt — the immovable Engine
        self._system_prompt = self._load_system_prompt(system_prompt_path)

        # Token budget — how heavy the train can get before we drop cargo
        self._max_tokens = max_tokens or int(os.getenv("MAX_TOKEN_BUDGET", "6000"))

        # Reserve ~20% of budget for the model's response generation
        self._response_reserve = int(self._max_tokens * 0.20)

        # The conversation payload — Index [0] is permanently locked
        self._messages: list[dict] = [
            {"role": "system", "content": self._system_prompt}
        ]

        # Token counter — using cl100k_base (GPT-4/3.5 tokenizer) as a
        # reasonable cross-model approximation. Not perfect for Llama/Gemma
        # but close enough for budget management.
        try:
            self._encoder = tiktoken.get_encoding("cl100k_base")
        except Exception:
            # Fallback: rough 4-chars-per-token estimation if tiktoken fails
            self._encoder = None

        # Cache the system prompt token count (it never changes)
        self._system_token_count = self._count_tokens(self._system_prompt)

    # ------------------------------------------------------------------
    # SWARM CONTEXT INJECTION
    # ------------------------------------------------------------------

    def inject_swarm_context(self, context_str: str) -> None:
        """Inject swarm metadata or file content into the system prompt temporarily."""
        new_prompt = f"{self._system_prompt}\n\n[SWARM CONTEXT]\n{context_str}\n"
        self._messages[0] = {"role": "system", "content": new_prompt}
        self._enforce_budget()

    def reset_system_prompt(self) -> None:
        """Restore the original system prompt without injected context."""
        self._messages[0] = {"role": "system", "content": self._system_prompt}
        self._enforce_budget()

    # ------------------------------------------------------------------
    # PUBLIC API
    # ------------------------------------------------------------------

    def add_user_message(self, content: str) -> None:
        """Append a user message to the conversation history."""
        self._messages.append({"role": "user", "content": content})
        self._enforce_budget()

    def add_assistant_message(self, content: str) -> None:
        """Append an assistant response to the conversation history."""
        self._messages.append({"role": "assistant", "content": content})
        self._enforce_budget()

    def get_payload(self) -> list[dict]:
        """
        Return the full message array ready to send to an API.
        
        Returns:
            List of message dicts with system prompt at [0] and
            chat history at [1:].
        """
        return list(self._messages)  # Return a copy to prevent mutation

    def get_history_length(self) -> int:
        """Return the number of chat messages (excluding system prompt)."""
        return len(self._messages) - 1

    def get_total_tokens(self) -> int:
        """Return the estimated total token count of the current payload."""
        return sum(self._count_tokens(m["content"]) for m in self._messages)

    def get_budget_info(self) -> dict:
        """Return a snapshot of the current memory budget state."""
        total = self.get_total_tokens()
        usable = self._max_tokens - self._response_reserve
        return {
            "total_tokens": total,
            "max_tokens": self._max_tokens,
            "usable_budget": usable,
            "remaining": usable - total,
            "system_tokens": self._system_token_count,
            "history_messages": self.get_history_length(),
            "utilization_pct": round((total / usable) * 100, 1) if usable > 0 else 0,
        }

    def clear_history(self) -> None:
        """Wipe all chat history but preserve the system prompt (Index [0])."""
        self._messages = [self._messages[0]]

    def update_system_prompt(self, new_prompt: str) -> None:
        """Hot-swap the system prompt without losing chat history."""
        self._messages[0] = {"role": "system", "content": new_prompt}
        self._system_prompt = new_prompt
        self._system_token_count = self._count_tokens(new_prompt)
        self._enforce_budget()

    # ------------------------------------------------------------------
    # PRIVATE INTERNALS
    # ------------------------------------------------------------------

    def _load_system_prompt(self, path: str) -> str:
        """Load the system protocol file from disk."""
        filepath = Path(path)
        if not filepath.exists():
            # Search relative to this module's directory as fallback
            filepath = Path(__file__).parent / path

        if filepath.exists():
            return filepath.read_text(encoding="utf-8").strip()

        # If neither path works, return a minimal default
        return (
            "You are Ted, an advanced AI developer assistant. "
            "Be direct, technical, and precise."
        )

    def _count_tokens(self, text: str) -> int:
        """
        Estimate the token count of a string.
        
        Uses tiktoken's cl100k_base encoder when available,
        falls back to a ~4 chars/token heuristic otherwise.
        """
        if self._encoder:
            return len(self._encoder.encode(text))
        # Rough heuristic fallback
        return len(text) // 4

    def _enforce_budget(self) -> None:
        """
        Prune the oldest chat messages if total tokens exceed the budget.
        
        This is the core "unhook the oldest cargo cars" operation.
        Index [0] (the Engine / system prompt) is NEVER touched.
        Messages are removed in pairs (user + assistant) from index [1]
        to maintain conversational coherence.
        """
        usable_budget = self._max_tokens - self._response_reserve

        while self.get_total_tokens() > usable_budget and len(self._messages) > 1:
            # Remove the oldest non-system message (Index [1])
            removed = self._messages.pop(1)
            # If we removed a user message and the next one is an assistant
            # response to it, remove that too to keep pairs intact
            if (
                removed["role"] == "user"
                and len(self._messages) > 1
                and self._messages[1]["role"] == "assistant"
            ):
                self._messages.pop(1)

    # ------------------------------------------------------------------
    # DUNDER METHODS
    # ------------------------------------------------------------------

    def __len__(self) -> int:
        return len(self._messages)

    def __repr__(self) -> str:
        info = self.get_budget_info()
        return (
            f"MemoryManager("
            f"messages={info['history_messages']}, "
            f"tokens={info['total_tokens']}/{info['usable_budget']}, "
            f"utilization={info['utilization_pct']}%)"
        )