models/memory_manager.py

"""
memory_manager.py
~~~~~~~~~~~~~~~~~
Production-ready conversational memory for the OSS AI assistant.

Design goals
------------
1. Retain recent history   — stores every human/AI exchange as typed dicts.
2. Configurable window     — `max_turns` caps how many pairs are kept.
3. Token-overflow guard    — a pre-inference check counts real tokens and
                             drops the oldest turn(s) until the prompt fits
                             within `max_context_tokens`.
4. Context preservation    — the system prompt is never counted against the
                             history budget; it is always included.
5. Reset support           — `clear()` wipes history; model stays loaded.

Memory layout (internal)
------------------------
Each item in `_turns` is a pair:
    {"human": str, "ai": str | None}

`ai` is None between the moment the user sends a message and the moment
the assistant replies — this lets the caller inspect in-flight state.

Public interface mirrors the OpenAI / HF chat-template message format:
    [{"role": "user", "content": "..."}, {"role": "assistant", "content": "..."}]
"""

from __future__ import annotations

from dataclasses import dataclass, field
from typing import Dict, List, Optional, TYPE_CHECKING

from models.logger_config import logger

if TYPE_CHECKING:
    # Avoid a hard import of transformers here; it is optional for callers
    # that only use memory without a model.
    from transformers import PreTrainedTokenizerBase


# ─────────────────────────────────────────────────────────────────────────────
# Internal storage type
# ─────────────────────────────────────────────────────────────────────────────

@dataclass
class Turn:
    """One complete human/AI exchange."""
    human: str
    ai: Optional[str] = None          # None until the assistant has replied

    @property
    def is_complete(self) -> bool:
        return self.ai is not None


# ─────────────────────────────────────────────────────────────────────────────
# Memory class
# ─────────────────────────────────────────────────────────────────────────────

class ConversationMemory:
    """
    Sliding-window conversational memory with token-budget enforcement.

    Parameters
    ----------
    max_turns : int
        Hard cap on the number of turns kept in memory. Oldest turns are
        evicted first when the cap is reached. Default: 10.
    max_context_tokens : int
        Soft token budget for the entire prompt (system + history).
        If a tokenizer is provided, oldest turns are dropped until the
        rendered prompt fits within this budget. Default: 1 800 (leaves
        ~200 tokens headroom for a 2 048-token context model like Qwen-0.5B).
    system_prompt : str
        Prepended to every prompt. Never counted against the history window.
    """

    def __init__(
        self,
        max_turns: int = 10,
        max_context_tokens: int = 1_800,
        system_prompt: str = (
            "You are a helpful, harmless, and honest AI assistant. "
            "Answer concisely and accurately."
        ),
    ) -> None:
        if max_turns < 1:
            raise ValueError("max_turns must be >= 1")
        if max_context_tokens < 64:
            raise ValueError("max_context_tokens must be >= 64")

        self.max_turns = max_turns
        self.max_context_tokens = max_context_tokens
        self.system_prompt = system_prompt

        self._turns: List[Turn] = []

        logger.debug(
            f"ConversationMemory created | max_turns={max_turns} "
            f"| max_context_tokens={max_context_tokens}"
        )

    # ── Write API ─────────────────────────────────────────────────────────────

    def add_user_message(self, text: str) -> None:
        """
        Record the user's message.  Opens a new Turn; `ai` is None until
        `add_assistant_reply()` is called.

        Immediately applies the turn-count window so the upcoming prompt
        is already trimmed before inference starts.
        """
        text = text.strip()
        if not text:
            raise ValueError("User message must not be empty.")

        self._turns.append(Turn(human=text))
        self._apply_turn_window()
        logger.debug(f"Memory: user message added | turns={len(self._turns)}")

    def add_assistant_reply(self, text: str) -> None:
        """
        Attach the assistant's reply to the most-recent open Turn.
        Raises RuntimeError if there is no open turn (i.e., add_user_message
        was not called first, or the turn is already complete).
        """
        text = text.strip()
        if not self._turns or self._turns[-1].is_complete:
            raise RuntimeError(
                "Cannot add assistant reply: no open turn found. "
                "Call add_user_message() first."
            )
        self._turns[-1].ai = text
        logger.debug(f"Memory: assistant reply saved | turns={len(self._turns)}")

    def rollback_last_user_message(self) -> None:
        """
        Remove the most-recent incomplete turn (user message without a reply).
        Called by OSSAssistant when inference fails, keeping memory consistent.
        """
        if self._turns and not self._turns[-1].is_complete:
            self._turns.pop()
            logger.warning("Memory: rolled back incomplete turn after inference error.")

    # ── Read API ──────────────────────────────────────────────────────────────

    def as_message_list(self) -> List[Dict[str, str]]:
        """
        Return history as a flat list of chat-template dicts:
            [{"role": "user", "content": "..."}, {"role": "assistant", "content": "..."}, ...]

        Incomplete turns (no AI reply yet) are included as user-only entries
        so that the model still sees the current user message.
        """
        messages: List[Dict[str, str]] = []
        for turn in self._turns:
            messages.append({"role": "user", "content": turn.human})
            if turn.ai is not None:
                messages.append({"role": "assistant", "content": turn.ai})
        return messages

    def as_prompt_messages(self) -> List[Dict[str, str]]:
        """
        Return the full prompt list: [system] + history.
        This is what you pass to `tokenizer.apply_chat_template()`.
        """
        return [
            {"role": "system", "content": self.system_prompt},
            *self.as_message_list(),
        ]

    def as_plain_text(self) -> str:
        """Human-readable transcript — useful for debugging or UI display."""
        lines: List[str] = []
        for i, turn in enumerate(self._turns, 1):
            lines.append(f"[{i}] User: {turn.human}")
            if turn.ai is not None:
                lines.append(f"[{i}]   AI: {turn.ai}")
            else:
                lines.append(f"[{i}]   AI: (awaiting reply)")
        return "\n".join(lines) if lines else "(no history)"

    # ── Token-budget enforcement ───────────────────────────────────────────────

    def enforce_token_budget(
        self, tokenizer: "PreTrainedTokenizerBase"
    ) -> int:
        """
        Drop oldest complete turns until the rendered prompt fits within
        `max_context_tokens`.  Returns the final token count.

        Only complete turns are dropped — the most-recent (in-flight) turn
        is always kept so the model sees the current question.

        Call this AFTER `add_user_message()` and BEFORE running inference.
        """
        while True:
            token_count = self._count_tokens(tokenizer)
            if token_count <= self.max_context_tokens:
                break

            # Find the oldest *complete* turn to evict
            evict_idx = next(
                (i for i, t in enumerate(self._turns) if t.is_complete), None
            )
            if evict_idx is None:
                # Nothing left to drop — the single current message is too long.
                # Log a warning; the tokenizer's own truncation will handle it.
                logger.warning(
                    f"Token budget exceeded ({token_count} > {self.max_context_tokens}) "
                    "but no complete turns to evict. Tokenizer truncation will apply."
                )
                break

            evicted = self._turns.pop(evict_idx)
            logger.debug(
                f"Token budget: evicted turn (tokens={token_count} > "
                f"{self.max_context_tokens}) | "
                f"human={evicted.human[:40]!r}"
            )

        final_count = self._count_tokens(tokenizer)
        logger.debug(f"Token budget enforced | final_tokens={final_count}")
        return final_count

    def _count_tokens(self, tokenizer: "PreTrainedTokenizerBase") -> int:
        """
        Render the current prompt to a string and count its tokens.
        Uses the model's native chat template for accuracy.
        """
        try:
            prompt_text = tokenizer.apply_chat_template(
                self.as_prompt_messages(),
                tokenize=False,
                add_generation_prompt=True,
            )
            return len(tokenizer.encode(prompt_text))
        except Exception as exc:
            logger.warning(f"Token counting failed ({exc}); skipping budget check.")
            return 0

    # ── Window management ─────────────────────────────────────────────────────

    def _apply_turn_window(self) -> None:
        """
        Enforce `max_turns`: drop oldest complete turns until the count fits.
        The in-flight turn (incomplete) is never evicted here.
        """
        complete_turns = [t for t in self._turns if t.is_complete]
        overflow = len(complete_turns) - self.max_turns

        if overflow <= 0:
            return

        # Remove the `overflow` oldest complete turns
        removed = 0
        i = 0
        while removed < overflow and i < len(self._turns):
            if self._turns[i].is_complete:
                self._turns.pop(i)
                removed += 1
            else:
                i += 1

        logger.debug(f"Turn window: evicted {removed} old turn(s) | remaining={len(self._turns)}")

    # ── Reset ─────────────────────────────────────────────────────────────────

    def clear(self) -> None:
        """Wipe all conversation history. Model stays loaded."""
        count = len(self._turns)
        self._turns.clear()
        logger.info(f"ConversationMemory cleared ({count} turn(s) removed).")

    # ── Introspection ─────────────────────────────────────────────────────────

    @property
    def turn_count(self) -> int:
        """Total turns stored (complete + in-flight)."""
        return len(self._turns)

    @property
    def complete_turn_count(self) -> int:
        """Turns where the assistant has already replied."""
        return sum(1 for t in self._turns if t.is_complete)

    @property
    def is_empty(self) -> bool:
        return len(self._turns) == 0

    def __repr__(self) -> str:
        return (
            f"ConversationMemory("
            f"turns={self.turn_count}, "
            f"max_turns={self.max_turns}, "
            f"max_context_tokens={self.max_context_tokens})"
        )


# ─────────────────────────────────────────────────────────────────────────────
# Backwards-compatible alias (keeps old import from previous memory_manager.py)
# ─────────────────────────────────────────────────────────────────────────────

class ConversationMemoryManager(ConversationMemory):
    """
    Legacy alias kept for backwards compatibility.
    Prefer ConversationMemory in new code.
    """

    def add_turn(self, human_text: str, ai_text: str) -> None:
        """Old-style API: add a complete turn in one call."""
        self.add_user_message(human_text)
        self.add_assistant_reply(ai_text)

    def get_history(self) -> List[Dict[str, str]]:
        return self.as_message_list()

    def get_history_as_text(self) -> str:
        return self.as_plain_text()