""" 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()