Improved RAG/Memory. Still needs work.

.gitignore

@@ -8,3 +8,10 @@ wheels/
 
 # Virtual environments
 .venv
+data/chroma_db/chroma.sqlite3
+data/chroma_db/83aaf16d-d56b-4094-b418-2d729ad113ea/data_level0.bin
+MiraiAssist.rar
+local_digest.txt
+data/chroma_db/83aaf16d-d56b-4094-b418-2d729ad113ea/header.bin
+data/chroma_db/83aaf16d-d56b-4094-b418-2d729ad113ea/length.bin
+data/chroma_db/83aaf16d-d56b-4094-b418-2d729ad113ea/link_lists.bin

LICENSE

@@ -0,0 +1,21 @@
+MIT License
+
+Copyright (c) 2025 Nighthawk42
+
+Permission is hereby granted, free of charge, to any person obtaining a copy
+of this software and associated documentation files (the "Software"), to deal
+in the Software without restriction, including without limitation the rights
+to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+copies of the Software, and to permit persons to whom the Software is
+furnished to do so, subject to the following conditions:
+
+The above copyright notice and this permission notice shall be included in all
+copies or substantial portions of the Software.
+
+THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+SOFTWARE.

README.md

@@ -1,7 +1,18 @@
- ## MiraiAssist v0.3.1-RAG
+ ## MiraiAssist v0.0.1-RAG
 
  A modular Python voice/text assistant framework featuring RAG (Retrieval-Augmented Generation) for long-term conversation memory, real-time STT/TTS, and a customizable UI.
 
+## Project Status
+This project is very much in the alpha state. It works. But it still has quirks and improvement. Especially in the GUI department. I have little to no experience designing nice GUIs.
+
+## TODO
+
+Add a hot-word like "Hey Mirai".
+
+Design a better GUI. Maybe use a TUI instead?
+
+Move everything to an embedded Python.
+
  ## Features
 
  *   **Voice & Text Input:** Interact via Push-to-Talk, a Record button, or a text input box.
@@ -125,20 +136,11 @@
  │   ├── system_manager.py
  │   ├── tts_manager.py
  │   └── ui_manager.py
- ├── placeholder.png          Replace with your screenshot
  ├── config.yaml
  ├── main.py
  └── README.md
  ```
 
- ## Contributing
-
- *(Add contribution guidelines here)*
-
- ## License
-
- *(Specify your license here, e.g., MIT License)*
-
  ## Acknowledgements
 
  *   CustomTkinter
@@ -150,4 +152,4 @@
  *   Tiktoken (OpenAI)
  *   Rich (Textualize)
  *   OpenAI Python Client
- *   Hugging Face Hub & Transformers
+ *   Hugging Face Hub & Transformers

config.yaml

@@ -4,13 +4,12 @@
 
 # --- LLM (Large Language Model) Settings ---
 llm:
-  api_base_url: "http://localhost:1234/v1" # e.g., LM Studio, Ollama, OpenAI
-  api_key_env_var: "NONE"                  # Environment variable for API key. "NONE" or null if not needed.
-  model_name: "lmstudio-community/Meta-Llama-3-8B-Instruct-GGUF" # Model identifier for the API
-  # !!! SET THIS ACCURATELY FOR YOUR MODEL !!! (e.g., Llama3 8k=8192, GPT-4 128k=131072)
-  # Used for Tiktoken length checks if available. Set to 0 to disable checks.
+  api_base_url: "http://localhost:1234/v1"
+  api_key_env_var: "NONE"
+  model_name: "lmstudio-community/gemma-3-4b-it-qat-q4_0-gguf" # For the API call
+  tokenizer_source_for_estimation: "google/gemma-2b-it" # Or appropriate Gemma base for tokenizer
   model_context_window: 8192
-  system_prompt: "You are Mirei, a helpful and concise AI assistant integrated into a desktop application. Respond clearly and directly. You can use markdown formatting."
+  system_prompt: "You are Mirai, a helpful and concise AI assistant integrated into a desktop application. Respond clearly and directly. You can use markdown formatting."
   temperature: 0.7      # 0=deterministic, >1 more creative
   max_tokens: 1536      # Max tokens for the LLM's *response* (ensure less than model_context_window)
   timeout_seconds: 120.0
@@ -18,13 +17,17 @@ llm:
 
 # --- Context Manager (RAG) Settings ---
 context_manager:
-  storage_path: "data/conversation_state.json" # Stores full raw history
-  vector_db_path: "data/chroma_db"             # Stores vector index
-  # Embedding Model: https://www.sbert.net/docs/pretrained_models.html
-  embedding_model_name: "all-MiniLM-L6-v2"     # Model for text embeddings ("multi-qa-MiniLM-L6-cos-v1" is another option)
-  collection_name: "mirei_chat_history"        # ChromaDB collection name
-  retrieval_results: 3                         # How many relevant history chunks to retrieve
-  include_recent_messages: 2                   # How many recent messages to *always* include with RAG results
+  storage_path: "data/conversation_state.json" # Full raw conversation history
+  vector_db_path: "data/chroma_db" # ChromaDB persistence path
+  embedding_model_name: "all-MiniLM-L6-v2" # Sentence-transformers model for embeddings
+  collection_name: "mirei_chat_history" # ChromaDB collection name
+  retrieval_results: 3 # How many relevant history chunks to retrieve
+  include_recent_messages: 2 # How many recent messages to *always* include
+
+# --- Memory Manager Settings ---
+memory_manager:
+  short_term_window_turns: 2    # Number of recent conversational turns (user+assistant)
+  long_term_retrieval_count: 3  # Number of relevant history messages from RAG
 
 # --- STT (Speech-to-Text) Settings ---
 stt:

data/conversation_state.json

@@ -0,0 +1 @@
+[]

main.py

@@ -1,14 +1,13 @@
-# ================================================
-# FILE: main.py (RAG + UI Text Input Integration - Corrected)
-# ================================================
+# C:\Users\Nighthawk\Desktop\MiraiAssist\main.py
 
 from __future__ import annotations
 
+import os
 import sys
 import signal
 import queue
 import threading
-import tkinter
+import tkinter # Keep for basic error dialogs
 import traceback
 import logging
 from pathlib import Path
@@ -17,24 +16,24 @@ from typing import Optional, Dict, Any
 # Use tkinter for basic error dialogs if GUI fails early
 from tkinter import Tk, messagebox
 
-# Import customtkinter only if GUI is intended
-try:
-    import customtkinter
-    CUSTOMTKINTER_AVAILABLE = True
-except ImportError:
-    CUSTOMTKINTER_AVAILABLE = False
-
 # --- PATH & LOGGING SETUP ---
 project_root = Path(__file__).resolve().parent
 if str(project_root) not in sys.path:
      sys.path.insert(0, str(project_root))
 
+# --- Import customtkinter only if GUI is intended ---
+try:
+    import customtkinter
+    CUSTOMTKINTER_AVAILABLE = True
+except ImportError:
+    CUSTOMTKINTER_AVAILABLE = False
+    # We will check this later before initializing the UI part
+
 # Basic logging setup FIRST (refined later by SystemManager)
-for h in logging.root.handlers[:]:
-    logging.root.removeHandler(h)
 logging.basicConfig(
     level=logging.INFO,
-    format="%(asctime)s [%(levelname)s] %(name)s: %(message)s"
+    format="%(asctime)s [%(levelname)s] %(name)s: %(message)s",
+    handlers=[logging.StreamHandler(sys.stderr)] # Log to stderr initially
 )
 logger = logging.getLogger(__name__) # Logger for main.py
 
@@ -44,36 +43,31 @@ try:
     from modules.system_manager import SystemManager, LoggingSetupError, RequirementError
     from modules.audio_manager import AudioManager, AudioManagerError
     from modules.stt_manager import STTManager, STTManagerError
-    from modules.context_manager import ContextManager, ContextManagerError # RAG version
-    from modules.llm_manager import LLMManager, LLMManagerError       # RAG version
+    from modules.context_manager import ContextManager, ContextManagerError
+    from modules.memory_manager import MemoryManager, MemoryManagerError # IMPORT MemoryManager
+    from modules.llm_manager import LLMManager, LLMManagerError
     from modules.tts_manager import TTSManager, TTSManagerError
-    from modules.ui_manager import UIManager                          # UI with text input
+    from modules.ui_manager import UIManager
 except ImportError as exc:
     print(f"FATAL ERROR: Failed to import core modules: {exc}", file=sys.stderr)
     traceback.print_exc()
-    try:
-        root = Tk(); root.withdraw()
-        messagebox.showerror("Startup Error", f"Failed to import core modules:\n\n{exc}\n\nPlease check dependencies (e.g., run 'uv sync') and logs.")
+    try: root = Tk(); root.withdraw(); messagebox.showerror("Startup Error", f"Failed to import core modules:\n\n{exc}\n\nPlease check dependencies (e.g., run 'uv sync') and logs.")
     except Exception: pass
     sys.exit(1)
 except Exception as e:
     print(f"FATAL ERROR: Unexpected error during initial imports: {e}", file=sys.stderr)
     traceback.print_exc()
-    try:
-         root = Tk(); root.withdraw()
-         messagebox.showerror("Startup Error", f"An unexpected error occurred during startup:\n\n{e}\n\nPlease check logs.")
+    try: root = Tk(); root.withdraw(); messagebox.showerror("Startup Error", f"An unexpected error occurred during startup:\n\n{e}\n\nPlease check logs.")
     except Exception: pass
     sys.exit(1)
 
-# Check if customtkinter loaded if we intend to use it
 if not CUSTOMTKINTER_AVAILABLE:
-     print("FATAL ERROR: CustomTkinter library is required but not found. Please install it.", file=sys.stderr)
-     sys.exit(1)
+    print("FATAL ERROR: CustomTkinter library is required but not found. Please install it (`uv add customtkinter`).", file=sys.stderr)
+    sys.exit(1)
+
 
-# --- SHARED QUEUE ---
 gui_queue: queue.Queue[Dict[str, Any]] = queue.Queue()
 
-# --- APPLICATION CONTROLLER ---
 class MiraiAppController:
     APP_NAME    = UIManager.APP_NAME
     APP_VERSION = UIManager.APP_VERSION
@@ -81,10 +75,10 @@ class MiraiAppController:
     def __init__(self) -> None:
         logger.info(f"Initializing {self.APP_NAME} Controller v{self.APP_VERSION}")
 
-        # Manager references initialization
         self.cfg: Optional[ConfigManager] = None
         self.sysman: Optional[SystemManager] = None
         self.ctx: Optional[ContextManager] = None
+        self.memman: Optional[MemoryManager] = None # ADDED MemoryManager attribute
         self.audio: Optional[AudioManager] = None
         self.stt: Optional[STTManager] = None
         self.llm: Optional[LLMManager] = None
@@ -94,10 +88,12 @@ class MiraiAppController:
         self._backend_ready = False
         self._shutting_down = False
 
-        # --- Initialize Backend ---
         try:
             self._initialize_backend()
+            self._backend_ready = True
+            logger.info("Backend initialization successful.")
         except (ConfigError, LoggingSetupError, RequirementError, ContextManagerError,
+                  MemoryManagerError, # ADDED MemoryManagerError to exception list
                   AudioManagerError, STTManagerError, LLMManagerError, TTSManagerError) as e:
             logger.critical(f"CRITICAL Backend initialization failed: {type(e).__name__}: {e}", exc_info=True)
             self._show_startup_error_dialog(f"Backend Initialization Failed:\n\n{type(e).__name__}: {e}\n\nPlease check configuration and logs.")
@@ -107,203 +103,185 @@ class MiraiAppController:
             self._show_startup_error_dialog(f"Unexpected Critical Error during Startup:\n\n{e}\n\nPlease check logs.")
             sys.exit(1)
 
-        if not self._backend_ready:
-            logger.critical("Backend initialization process completed but backend is not marked as ready.")
-            self._show_startup_error_dialog("Backend initialization did not complete successfully.\nCannot start UI. Check logs for details.")
-            sys.exit(1)
-
-        # --- Initialize UI ---
         try:
             self.ui = UIManager(gui_queue, self.cfg)
-            # ---> Make sure the callback method exists before calling build_window <---
-            if not hasattr(self, '_handle_text_input_submit'):
-                 # This is a sanity check, the AttributeError should prevent getting here usually
-                 raise AttributeError("Internal Error: _handle_text_input_submit method is missing in MiraiAppController.")
-
             self.ui.build_window(
-                on_close_callback      = self._handle_close_request,
-                record_callback        = self._handle_record_button_click,
-                theme_change_callback  = self._handle_theme_change,
-                clear_history_callback = self._handle_clear_history,
-                about_callback         = self._handle_about,
-                text_submit_callback   = self._handle_text_input_submit # Correctly pass the method reference
+                on_close_callback=self._handle_close_request,
+                record_callback=self._handle_record_button_click,
+                theme_change_callback=self._handle_theme_change,
+                clear_history_callback=self._handle_clear_history,
+                about_callback=self._handle_about,
+                text_submit_callback=self._handle_text_input_submit
             )
             self.ui.update_status("Ready")
             self.ui.set_record_button_state("Record", enabled=True)
             self._setup_ptt_bindings_via_ui()
-
-            # --- Populate UI with loaded history ---
             self._populate_initial_ui_history()
-
+            logger.info("UI initialized successfully.")
         except Exception as e:
             logger.critical(f"Failed to initialize UI Manager: {e}", exc_info=True)
-            self._stop_backend_managers() # Attempt cleanup
+            self._stop_backend_managers()
             self._show_startup_error_dialog(f"FATAL ERROR: UI Initialization failed: {e}\n\nPlease check logs.")
             sys.exit(1)
 
-        # --- Start GUI Queue Polling ---
-        if self.ui and self.ui.root:
+        if self.ui and self.ui.root and self.ui.root.winfo_exists():
             self.ui.root.after(100, self._poll_gui_queue)
         else:
-             logger.critical("UI Root window not available after initialization. Cannot poll queue.")
-             self._handle_close_request()
-             sys.exit(1)
-
+            logger.critical("UI Root window not available after initialization. Cannot poll queue.")
+            self._handle_close_request()
+            sys.exit(1)
 
     def _show_startup_error_dialog(self, message: str):
-        """Helper to show critical startup errors via Tkinter."""
         try:
             root = Tk(); root.withdraw()
             messagebox.showerror(f"{self.APP_NAME} - Critical Startup Error", message)
         except Exception as tk_error:
             print(f"\nCRITICAL STARTUP ERROR (Tkinter Dialog Failed: {tk_error}):\n{message}\n", file=sys.stderr)
 
-
     def _initialize_backend(self) -> None:
-        """Initializes backend managers sequentially. Raises exceptions on failure."""
-        logger.info("Starting backend initialization...")
-        self.cfg = ConfigManager()
-        self.cfg.load()
-        logger.info("Configuration loaded.")
+        logger.info("Starting backend initialization sequence...")
+        self.cfg = ConfigManager(); self.cfg.load(); logger.info("Configuration loaded.")
 
-        self.sysman = SystemManager(self.cfg)
-        self.sysman.setup_logging()
-        self.sysman.log_system_info()
+        self.sysman = SystemManager(self.cfg); self.sysman.setup_logging(); logger.info("Logging configured.")
+        self.sysman.log_system_info(); self.sysman.verify_requirements(); logger.info("Requirement checks passed.")
 
-        self.sysman.verify_requirements()
-        logger.info("Requirement checks passed.")
+        # Initialize ContextManager first (data layer)
+        self.ctx = ContextManager(self.cfg); logger.info("Context manager (RAG data layer) initialized.")
 
-        self.ctx = ContextManager(self.cfg)
-        logger.info("Context manager (RAG) ready.")
+        # Initialize MemoryManager (strategy layer, depends on ContextManager)
+        if not self.ctx: # Should ideally be caught by earlier checks if ctx fails
+            raise MemoryManagerError("ContextManager failed to initialize before MemoryManager.")
+        self.memman = MemoryManager(self.cfg, self.ctx); logger.info("Memory manager (context strategy) initialized.")
 
-        self.audio = AudioManager(self.cfg, gui_queue)
-        logger.info("Audio manager ready.")
+        self.audio = AudioManager(self.cfg, gui_queue); logger.info("Audio manager initialized.")
+        self.stt = STTManager(self.cfg); logger.info("STT manager initialized.")
 
-        self.stt = STTManager(self.cfg)
-        logger.info("STT manager ready.")
-
-        if not self.ctx: raise LLMManagerError("ContextManager failed to initialize before LLMManager.")
+        # Initialize LLMManager (depends on MemoryManager)
         self.llm = LLMManager(self.cfg, gui_queue)
-        self.llm.set_context_manager(self.ctx)
-        logger.info("LLM manager ready.")
+        if not self.memman: # Should be caught if memman fails
+            raise LLMManagerError("MemoryManager failed to initialize before LLMManager.")
+        self.llm.set_memory_manager(self.memman) # Link MemoryManager to LLMManager
+        logger.info("LLM manager initialized and linked with MemoryManager.")
 
         if not self.audio: raise TTSManagerError("AudioManager failed to initialize before TTSManager.")
-        self.tts = TTSManager(self.cfg, gui_queue, self.audio)
-        logger.info("TTS manager ready.")
-
-        self._backend_ready = True
-        logger.info("Backend initialization successful.")
-
+        self.tts = TTSManager(self.cfg, gui_queue, self.audio); logger.info("TTS manager initialized.")
 
     def _populate_initial_ui_history(self):
-        """Loads history from ContextManager into the UI."""
-        if self.ctx and self.ui:
-            initial_history = self.ctx.history
+        if self.memman and self.ui and self.ui.root: # Check and use memman
+            initial_history = self.memman.get_full_history() # Get history via memman
             if initial_history:
-                logger.info(f"Populating UI with {len(initial_history)} loaded messages from history file.")
-                def populate_ui():
+                logger.info(f"Populating UI with {len(initial_history)} messages from loaded history.")
+                def populate_task():
                     if not self.ui or not self.ui.history_textbox or not self.ui.history_textbox.winfo_exists():
                         logger.warning("Cannot populate history, UI textbox not available.")
                         return
+                    self.ui.clear_history_display()
                     for message in initial_history:
                         role = message.get("role")
                         content = message.get("content")
-                        if role and content and role in ["user", "assistant"]:
+                        if role and content and role in ["user", "assistant", "system"]:
                             self.ui.append_history(role, content)
                     self.ui.update_status("Ready (History Loaded)")
-                self.ui.schedule_task(populate_ui)
+                self.ui.schedule_task(populate_task)
             else:
-                 logger.info("No previous conversation history found to display.")
-
+                logger.info("No previous conversation history found or loaded.")
+        else:
+             logger.warning("Cannot populate initial UI history: MemoryManager or UIManager not ready.") # Updated log message
 
     def _poll_gui_queue(self) -> None:
-        """Polls the shared queue for messages from backend threads."""
         if self._shutting_down: return
         try:
             while not gui_queue.empty():
                 msg = gui_queue.get_nowait()
                 self._handle_message(msg)
         except queue.Empty: pass
-        except Exception as e: logger.error(f"Error processing GUI queue: {e}", exc_info=True)
+        except Exception as e:
+            logger.error(f"Error processing GUI queue: {e}", exc_info=True)
+            if self.ui: self.ui.log(f"Queue processing error: {e}", tag="critical")
         finally:
             if not self._shutting_down and self.ui and self.ui.root:
                 try:
                     if self.ui.root.winfo_exists():
                         self.ui.root.after(100, self._poll_gui_queue)
-                    else: logger.warning("UI root window destroyed, stopping GUI queue polling.")
-                except tkinter.TclError: logger.warning("TclError checking UI window, stopping GUI queue polling (likely closing).")
-
+                    else:
+                        logger.warning("UI root window destroyed, stopping GUI queue polling.")
+                except tkinter.TclError:
+                    logger.warning("TclError checking UI window, stopping GUI queue polling (likely closing).")
 
     def _handle_message(self, message: Dict[str, Any]) -> None:
-        """Routes messages from the queue to appropriate UI actions."""
         if self._shutting_down or not self.ui: return
-
         msg_type = message.get("type")
         payload = message.get("payload")
         tag = message.get("tag", "info")
-
-        if msg_type == "log": self.ui.log(str(payload), tag=tag)
-        elif msg_type == "status": self.ui.update_status(str(payload))
-        elif msg_type == "audio_ready": self._process_audio_ready(payload)
-        elif msg_type == "stt_result": self._process_stt_result(payload)
-        elif msg_type == "llm_chunk": self._process_llm_chunk(payload)
-        elif msg_type == "llm_result": self._process_llm_result(payload)
-        elif msg_type == "shutdown_request": self._handle_close_request()
-        else:
-            logger.warning(f"Received unknown message type in GUI queue: {msg_type}")
-            self.ui.log(f"Unknown message type: {msg_type}", tag="warning")
-
+        try:
+            if msg_type == "log": self.ui.log(str(payload), tag=tag)
+            elif msg_type == "status": self.ui.update_status(str(payload))
+            elif msg_type == "audio_ready": self._process_audio_ready(payload)
+            elif msg_type == "stt_result": self._process_stt_result(payload)
+            elif msg_type == "llm_chunk": self._process_llm_chunk(payload)
+            elif msg_type == "llm_result": self._process_llm_result(payload)
+            elif msg_type == "shutdown_request": self._handle_close_request()
+            else:
+                logger.warning(f"Received unknown message type in GUI queue: {msg_type}")
+                self.ui.log(f"Unknown message type received: {msg_type}", tag="warning")
+        except AttributeError as e:
+            if "winfo_exists" in str(e):
+                 logger.warning(f"UI widget likely destroyed during message handling ({msg_type}). Error: {e}")
+            else:
+                 logger.error(f"AttributeError handling message type '{msg_type}': {e}", exc_info=True)
+        except Exception as e:
+            logger.error(f"Error handling message type '{msg_type}': {e}", exc_info=True)
+            if self.ui: self.ui.log(f"Error handling message: {e}", tag="error")
 
     def _process_audio_ready(self, payload: Any):
-        """Handles the 'audio_ready' message."""
         if isinstance(payload, dict) and "filepath" in payload:
             path = str(payload["filepath"])
             duration = payload.get("duration", 0.0)
-            self.ui.update_status("Transcribing…")
-            self.ui.log(f"Transcribing recorded audio ({duration:.2f}s)…", tag="info")
+            self.ui.update_status("Transcribing...")
+            self.ui.log(f"Transcribing recorded audio ({duration:.2f}s)...", tag="info")
             self._run_stt_in_background(path)
         else:
-            self.ui.log("Error: Invalid audio data received.", tag="error")
-            self.ui.update_status("ERROR: Invalid Audio")
+            if self.ui:
+                self.ui.log("Error: Invalid audio data received.", tag="error")
+                self.ui.update_status("ERROR: Invalid Audio")
 
     def _process_stt_result(self, payload: Any):
-        """Handles the 'stt_result' message and triggers the LLM."""
         if isinstance(payload, str):
             text = payload.strip()
             if text:
-                self.ui.log("Transcription complete.", tag="info")
-                # Display the user's transcribed message *before* submitting
-                self.ui.append_history("user", text)
-                # Use common submission logic
+                if self.ui:
+                    self.ui.log("Transcription complete.", tag="info")
+                    self.ui.append_history("user", text)
                 self._submit_user_input(text)
             else:
-                self.ui.log("No speech detected in audio.", tag="warning")
-                self.ui.update_status("Ready (No speech detected)")
-                self.ui.append_history_event("(No speech detected)")
-        elif payload is None: # Explicit None indicates error
-            self.ui.log("Speech transcription failed. Check logs.", tag="error")
-            self.ui.update_status("ERROR: Transcription Failed")
-            self.ui.append_history_event("(Transcription Error)")
+                if self.ui:
+                    self.ui.log("No speech detected in audio.", tag="warning")
+                    self.ui.update_status("Ready (No speech)")
+                    self.ui.append_history_event("(No speech detected)")
+        elif payload is None:
+            if self.ui:
+                self.ui.log("Speech transcription failed. Check logs.", tag="error")
+                self.ui.update_status("ERROR: Transcription Failed")
+                self.ui.append_history_event("(Transcription Error)")
         else:
-            self.ui.log(f"Invalid STT payload type: {type(payload)}", tag="error")
-            self.ui.update_status("ERROR: Invalid STT Data")
+            if self.ui:
+                self.ui.log(f"Invalid STT payload type: {type(payload)}", tag="error")
+                self.ui.update_status("ERROR: Invalid STT Data")
 
     def _process_llm_chunk(self, payload: Any):
-        """Handles incoming LLM stream chunks."""
         if isinstance(payload, dict) and "delta" in payload:
-            if not self.ui._assistant_streaming:
-                self.ui.start_assistant_stream()
-            self.ui.append_stream_chunk(str(payload["delta"]))
+            if self.ui:
+                if not self.ui._assistant_streaming: # Assuming _assistant_streaming is an attribute in UIManager
+                    self.ui.start_assistant_stream()
+                self.ui.append_stream_chunk(str(payload["delta"]))
         else:
             logger.error(f"Invalid llm_chunk payload: {payload}")
 
     def _process_llm_result(self, payload: Any):
-        """Handles the final LLM result and triggers TTS."""
-        self.ui.finish_assistant_stream()
+        if self.ui: self.ui.finish_assistant_stream()
 
         if not isinstance(payload, dict):
-            self.ui.log(f"Invalid llm_result payload type: {type(payload)}", tag="error")
-            self.ui.update_status("ERROR: Invalid LLM Data")
+            if self.ui: self.ui.log(f"Invalid llm_result payload type: {type(payload)}", tag="error"); self.ui.update_status("ERROR: Invalid LLM Data")
             return
 
         text = payload.get("text")
@@ -311,87 +289,81 @@ class MiraiAppController:
         err_msg = payload.get("error_message", "")
 
         if error:
-            self.ui.log(f"LLM Error: {err_msg}", tag="error")
-            self.ui.update_status("ERROR: LLM Failed")
-            self.ui.append_history_event(f"(LLM Error: {err_msg[:60]}...)")
+            if self.ui: self.ui.log(f"LLM Error: {err_msg}", tag="error"); self.ui.update_status("ERROR: LLM Failed"); self.ui.append_history_event(f"(LLM Error: {err_msg[:60]}...)")
         elif text:
             if self.tts and self.audio:
-                self.ui.update_status("Speaking…")
-                self.ui.log("Sending response to TTS…", tag="info")
+                if self.ui: self.ui.update_status("Speaking..."); self.ui.log("Sending response to TTS...", tag="info")
                 if not self.audio.is_playing and not self.audio.is_recording:
                     self.tts.speak_text(text)
                 else:
-                    logger.warning("Audio manager busy. Cannot speak TTS response.")
-                    self.ui.log("Audio busy; cannot speak.", tag="warning")
-                    self.ui.update_status("Ready (Audio busy)")
+                    logger.warning("Audio manager busy. Cannot speak TTS response now.")
+                    if self.ui: self.ui.log("Audio busy; cannot speak.", tag="warning"); self.ui.update_status("Ready (Audio busy)")
             else:
                 logger.error("TTS or Audio manager not available to speak.")
-                self.ui.update_status("ERROR: TTS Not Ready")
-        else: # No error, but no text
-            self.ui.log("Assistant provided no response.", tag="warning")
-            self.ui.update_status("Ready (No response)")
-            self.ui.append_history_event("(Assistant gave no response)")
-
+                if self.ui: self.ui.update_status("ERROR: TTS Not Ready"); self.ui.update_status("Ready") # Reset status
+        else:
+            if self.ui: self.ui.log("Assistant provided no speakable response.", tag="warning"); self.ui.update_status("Ready (No response)"); self.ui.append_history_event("(Assistant gave no response)")
 
     def _run_stt_in_background(self, filepath: str) -> None:
-        """Runs STT transcription in a separate thread."""
         if not self.stt:
-            logger.error("STT Manager not available for background task.")
+            logger.error("STT Manager not available.")
             gui_queue.put({"type": "stt_result", "payload": None})
             return
-
         def worker():
             result: Optional[str] = None
-            try: result = self.stt.transcribe(filepath)
-            except Exception as e: logger.error(f"STT background worker exception: {e}", exc_info=True)
+            try:
+                result = self.stt.transcribe(filepath)
+            except Exception as e:
+                logger.error(f"STT background thread exception: {e}", exc_info=True)
             finally:
                 gui_queue.put({"type": "stt_result", "payload": result})
                 try:
                     p = Path(filepath)
-                    if p.exists() and "temp_input" in p.name: p.unlink(); logger.debug(f"Deleted temporary STT input file: {filepath}")
-                except Exception as e: logger.warning(f"Error deleting temp file {filepath}: {e}")
-
+                    if p.exists() and "temp_input" in p.name and p.suffix == ".wav":
+                         p.unlink()
+                         logger.debug(f"Deleted temporary STT input file: {filepath}")
+                except Exception as e:
+                    logger.warning(f"Error deleting temp STT file {filepath}: {e}")
         threading.Thread(target=worker, name="STTWorker", daemon=True).start()
 
-
     def _submit_user_input(self, text: str):
-        """Handles user text input, updating context and triggering LLM."""
         if not text: return
-        if not self.ctx or not self.llm or not self.ui:
-             logger.error("Cannot process user input: Core components missing.")
-             self.ui.update_status("ERROR: Core component missing")
-             self.ui.log("Internal error: Cannot process user input.", tag="error")
+        if not self.memman or not self.llm or not self.ui: # Check memman
+             logger.error("Cannot process user input: Core components missing (memman, llm, or ui).")
+             if self.ui:
+                self.ui.update_status("ERROR: Core component error")
+                self.ui.log("Internal error processing input.", tag="error")
              return
-
-        # Note: UI display (append_history) is now handled by the *caller*
-        # (_process_stt_result or _handle_text_input_submit) before calling this.
-
         try:
-            self.ctx.add_message("user", text)
-            logger.debug(f"User input added to RAG context: '{text[:50]}...'")
+            self.memman.add_message("user", text) # Use memman to add the user's message
+            logger.debug(f"User input added to memory via MemoryManager: '{text[:50]}...'")
         except Exception as e:
-            logger.error(f"Error adding user message to context: {e}", exc_info=True)
-            self.ui.log(f"Error saving context: {e}", tag="error")
-            # Decide if we should stop here if context fails
-            # return # Optionally stop if context add fails
+            logger.error(f"Error adding user message to memory: {e}", exc_info=True)
+            if self.ui: self.ui.log(f"Error saving to memory: {e}", tag="error")
+            # Decide if we should stop if context add fails. For now, let LLM proceed.
 
-        # Run LLM
-        self.ui.update_status("Thinking...")
-        self.ui.log("Sending text to LLM (with RAG)...", tag="info")
+        if self.ui:
+            self.ui.update_status("Thinking...")
+            self.ui.log("Sending request to LLM (with MemoryManager)...", tag="info") # Updated log
         self.llm.run_llm_in_background(text)
 
-
     def _handle_record_button_click(self) -> None:
-        if self._shutting_down or not self._backend_ready or not self.audio or not self.ui: return
+        if self._shutting_down or not self._backend_ready or not self.audio or not self.ui:
+            logger.debug("Ignoring record button click (shutting down or backend not ready).")
+            return
         if self.audio.is_playing:
             logger.info("Record button clicked while playing; stopping playback first.")
             self.audio.stop_playback()
             self.ui.schedule_task(self._toggle_recording)
-        else: self._toggle_recording()
+        else:
+            self._toggle_recording()
 
     def _toggle_recording(self) -> None:
         if not self.audio or not self.ui: return
-        default_color = customtkinter.ThemeManager.theme["CTkButton"]["fg_color"]
+        default_color = "grey"
+        try: default_color = customtkinter.ThemeManager.theme["CTkButton"]["fg_color"]
+        except (KeyError, TypeError): logger.warning("Could not get default button theme color.")
+
         if self.audio.is_recording:
             self.audio.stop_recording()
             self.ui.set_record_button_state("Record", color=default_color, enabled=True)
@@ -401,22 +373,42 @@ class MiraiAppController:
 
     def _setup_ptt_bindings_via_ui(self) -> None:
         if not self.cfg or not self.audio or not self.ui: return
-        combo = self.cfg.get("activation", "push_to_talk_key", "<Control-space>")
-        if not combo: logger.info("Push-to-Talk key not configured. PTT disabled."); return
+        combo = self.cfg.get("activation", "push_to_talk_key")
+        if not combo or not isinstance(combo, str):
+            logger.info("Push-to-Talk key not configured or invalid. PTT disabled.")
+            return
         try:
-            parts = combo.strip("<>").split("-"); key = parts[-1] if parts else ""
-            if not key: logger.error(f"Invalid PTT key format: '{combo}'"); return
-            success = self.ui.bind_ptt(combo, f"<KeyRelease-{key}>", self._handle_ptt_start, self._handle_ptt_stop)
-            if success: self.ui.log(f"Push-to-Talk enabled ({combo})", tag="info")
-            else: self.ui.log(f"Failed to bind PTT key ({combo}). See logs.", tag="error")
-        except Exception as e: logger.error(f"Error setting up PTT bindings for '{combo}': {e}", exc_info=True)
+            if not (combo.startswith("<") and combo.endswith(">") and "-" in combo):
+                logger.error(f"Invalid PTT key format in config: '{combo}'. Expected like '<Control-space>'.")
+                return
+            parts = combo.strip("<>").split('-')
+            key = parts[-1] if parts else ""
+            if not key:
+                 logger.error(f"Could not extract key from PTT combo: '{combo}'")
+                 return
+            release_event = f"<KeyRelease-{key}>"
+            success = self.ui.bind_ptt(
+                press_event=combo,
+                release_event=release_event,
+                ptt_start_callback=self._handle_ptt_start,
+                ptt_stop_callback=self._handle_ptt_stop
+            )
+            if success:
+                self.ui.log(f"Push-to-Talk enabled ({combo})", tag="info")
+                self.ui.set_record_button_state(f"Record (Hold {combo})", enabled=True)
+            else:
+                if self.ui: self.ui.log(f"Failed to bind PTT key ({combo}). See logs.", tag="error")
+        except Exception as e:
+            logger.error(f"Error setting up PTT bindings for '{combo}': {e}", exc_info=True)
+            if self.ui: self.ui.log(f"Error binding PTT key ({combo}).", tag="error")
 
     def _handle_ptt_start(self) -> None:
         if self._shutting_down or not self._backend_ready or not self.audio or not self.ui: return
         if self.audio.is_playing:
             logger.info("PTT pressed while playing; stopping playback first.")
             self.audio.stop_playback()
-            self.ui.schedule_task(self._handle_ptt_start)
+            if self.ui and self.ui.root: # Check if UI root exists before scheduling
+                self.ui.schedule_task(lambda: self.ui.root.after(50, self._handle_ptt_start))
             return
         if not self.audio.is_recording:
             logger.debug("PTT Start: Triggering recording toggle.")
@@ -428,10 +420,8 @@ class MiraiAppController:
             logger.debug("PTT Stop: Triggering recording toggle.")
             self._toggle_recording()
 
-
-    # ────────────────── UI MENU & TEXT INPUT CALLBACKS ────────────────────────
     def _handle_theme_change(self, mode: str) -> None:
-        if self.cfg:
+        if self.cfg and self.ui:
             logger.info(f"Saving theme preference: {mode}")
             self.cfg.update_value("gui", "theme_preference", mode.lower())
             self.cfg.save()
@@ -439,93 +429,121 @@ class MiraiAppController:
 
     def _handle_clear_history(self) -> None:
         logger.info("Clear History requested by user.")
-        if self.ctx and self.ui:
+        if self.memman and self.ui: # Use memman
             try:
-                self.ctx.clear_context()
+                self.memman.clear_memory() # Use memman to clear context
                 self.ui.clear_history_display()
-                self.ui.log("Conversation history and vector index cleared.", tag="info")
+                self.ui.log("Conversation history, memory, and vector index cleared.", tag="info")
                 self.ui.update_status("Ready (History Cleared)")
             except Exception as e:
-                logger.error(f"Error clearing context: {e}", exc_info=True)
-                self.ui.log(f"Error clearing context: {e}", tag="error")
-        else: logger.warning("ContextManager or UIManager not available for clear history.")
+                logger.error(f"Error clearing memory: {e}", exc_info=True)
+                if self.ui: self.ui.log(f"Error clearing memory: {e}", tag="error")
+        else:
+            logger.warning("MemoryManager or UIManager not available for clear history.")
 
     def _handle_about(self) -> None:
-        if self.ui: self.ui.show_about_dialog(self.APP_NAME, self.APP_VERSION)
+        if self.ui:
+            self.ui.show_about_dialog(self.APP_NAME, self.APP_VERSION)
 
-    # <<< METHOD DEFINITION IS NOW CORRECTLY HERE >>>
-    def _handle_text_input_submit(self, text: str):
-        """Callback triggered by UIManager when text is submitted."""
-        logger.info(f"Text input received: '{text[:60]}...'")
+    def _handle_text_input_submit(self, text: str) -> None:
+        logger.info(f"Text input submitted: '{text[:60]}...'")
         if self._shutting_down or not self._backend_ready:
             logger.warning("Ignoring text input during shutdown or if backend not ready.")
             return
         if not text:
             logger.warning("Received empty text input from UI callback.")
             return
-
-        # Display the user's typed message FIRST
         if self.ui:
             self.ui.append_history("user", text)
         else:
-             logger.error("UI not available to display submitted text.")
-             return
-
-        # Use the common submission logic
+            logger.error("UI not available to display submitted text.")
+            return
         self._submit_user_input(text)
 
-
-    # ────────────────── SHUTDOWN ─────────────────────────────────────────────
     def _handle_close_request(self) -> None:
-        if self._shutting_down: return
+        if self._shutting_down:
+            logger.debug("Shutdown already in progress.")
+            return
         self._shutting_down = True
         logger.info("Shutdown requested. Initiating graceful shutdown...")
-        if self.ui:
-            try: self.ui.update_status("Shutting down…")
+
+        if self.ui and self.ui.root and self.ui.root.winfo_exists():
+            try: self.ui.update_status("Shutting down...")
             except Exception: pass
+
+        # Save raw history via ContextManager (which MemoryManager uses indirectly for data persistence)
         if self.ctx:
             try:
-                logger.info("Saving full conversation history before shutdown...")
+                logger.info("Saving final conversation history (via ContextManager)...")
                 self.ctx.save_context()
                 logger.info("Conversation history saved.")
+            except Exception as e:
+                logger.error(f"Failed to save context during shutdown: {e}", exc_info=True)
+            # Shutdown ContextManager (e.g., release Chroma resources)
+            try:
                 self.ctx.shutdown()
-            except Exception as e: logger.error(f"Failed to save/shutdown context: {e}", exc_info=True)
+            except Exception as e:
+                 logger.error(f"Error shutting down ContextManager: {e}", exc_info=True)
+        
+        # Shutdown MemoryManager (if it has specific shutdown tasks in the future)
+        if self.memman:
+            try:
+                self.memman.shutdown()
+            except Exception as e:
+                logger.error(f"Error shutting down MemoryManager: {e}", exc_info=True)
+
         self._stop_backend_managers()
         if self.ui:
-            try: logger.info("Destroying UI window..."); self.ui.destroy(); logger.info("UI window destroyed.")
-            except Exception as e: logger.error(f"Error destroying UI window: {e}", exc_info=True)
-        logger.info("MiraiAssist shutdown complete.")
-
+            try:
+                logger.info("Destroying UI window...")
+                self.ui.destroy()
+                logger.info("UI window destroyed.")
+            except Exception as e:
+                logger.error(f"Error destroying UI window during shutdown: {e}", exc_info=True)
+        logger.info(f"{self.APP_NAME} shutdown complete.")
 
     def _stop_backend_managers(self) -> None:
         logger.info("Stopping backend managers...")
         if self.audio:
             try: logger.debug("Stopping AudioManager..."); self.audio.stop(); logger.debug("AudioManager stopped.")
             except Exception as e: logger.warning(f"Error stopping AudioManager: {e}", exc_info=True)
+        # Add other explicit stop calls here if managers require them
+        logger.info("Backend managers stopped.")
 
-
-    # ────────────────── MAIN LOOP ─────────────────────────────────────────────
     def run(self) -> None:
-        if not self._backend_ready: logger.critical("Backend not ready. Cannot start UI main loop."); sys.exit(1)
-        if self.ui:
-            try: logger.info("Entering UI main loop."); self.ui.run(); logger.info("UI main loop exited normally.")
-            except Exception as e: logger.critical(f"UI main loop crashed: {e}", exc_info=True); self._handle_close_request(); sys.exit(1)
-        else: logger.critical("Cannot run: UI Manager not initialized."); sys.exit(1)
+        if not self._backend_ready:
+            logger.critical("Backend not ready. Cannot start UI main loop.")
+            sys.exit(1)
+        if not self.ui:
+            logger.critical("Cannot run: UI Manager not initialized.")
+            sys.exit(1)
+        try:
+            logger.info("Entering UI main loop.")
+            self.ui.run()
+            logger.info("UI main loop exited normally.")
+        except Exception as e:
+            logger.critical(f"UI main loop crashed: {e}", exc_info=True)
+            self._handle_close_request()
+            sys.exit(1)
 
-# ──────────────────────────────────────────────────────────────────────────────
-# SIGNAL HANDLER & ENTRY POINT
-# ──────────────────────────────────────────────────────────────────────────────
 _controller_instance: Optional[MiraiAppController] = None
-
 def _signal_handler(sig, frame) -> None:
-    logger.warning(f"Received signal {sig}; requesting graceful shutdown.")
+    logger.warning(f"Received signal {signal.Signals(sig).name} ({sig}). Requesting graceful shutdown.")
     global _controller_instance
     if _controller_instance and not _controller_instance._shutting_down:
-        try: gui_queue.put_nowait({"type": "shutdown_request"})
-        except queue.Full: logger.error("GUI queue full during signal handling. Forcing exit."); sys.exit(1)
-        except Exception as e: logger.error(f"Error putting shutdown request on queue: {e}. Forcing exit."); sys.exit(1)
-    elif _controller_instance and _controller_instance._shutting_down: logger.warning("Shutdown already in progress. Signal ignored.")
-    else: logger.info("Controller instance not found during signal handling. Exiting."); sys.exit(0)
+        try:
+            gui_queue.put_nowait({"type": "shutdown_request"})
+        except queue.Full:
+            logger.error("GUI queue full during signal handling. Forcing exit.")
+            os._exit(1) # Force exit if queue is unresponsive
+        except Exception as e:
+            logger.error(f"Error putting shutdown request on queue: {e}. Forcing exit.")
+            os._exit(1)
+    elif _controller_instance and _controller_instance._shutting_down:
+        logger.warning("Shutdown already in progress. Signal ignored.")
+    else:
+        logger.info("Controller instance not found during signal handling. Exiting.")
+        sys.exit(0)
 
 def main() -> None:
     global _controller_instance
@@ -536,17 +554,24 @@ def main() -> None:
         controller = MiraiAppController()
         _controller_instance = controller
         controller.run()
-    except SystemExit: raise # Allow clean exits
+    except SystemExit:
+        logger.info("SystemExit caught. Exiting application.")
+    except KeyboardInterrupt:
+         logger.info("KeyboardInterrupt caught. Initiating shutdown...")
+         if _controller_instance and not _controller_instance._shutting_down:
+              _controller_instance._handle_close_request()
     except Exception as e:
-        logger.critical(f"Unhandled top-level exception: {e}", exc_info=True)
+        logger.critical(f"Unhandled top-level exception in main: {e}", exc_info=True)
         try:
             root = Tk(); root.withdraw()
             messagebox.showerror("Critical Error", f"A critical unhandled error occurred:\n\n{e}\n\nPlease check logs.")
-        except Exception: print(f"\nCRITICAL UNHANDLED ERROR (Dialog Failed): {e}\n", file=sys.stderr); traceback.print_exc()
+        except Exception:
+            print(f"\nCRITICAL UNHANDLED ERROR (Dialog Failed): {e}\n", file=sys.stderr)
+            traceback.print_exc()
         finally:
-            if controller and hasattr(controller, '_handle_close_request'):
+            if _controller_instance and hasattr(_controller_instance, '_handle_close_request') and not _controller_instance._shutting_down:
                  logger.info("Attempting cleanup after top-level error...")
-                 try: controller._handle_close_request()
+                 try: _controller_instance._handle_close_request()
                  except Exception as cleanup_e: logger.error(f"Error during final cleanup attempt: {cleanup_e}")
             sys.exit(1)
 

modules/__init__.py

@@ -1,4 +1,26 @@
 # modules/__init__.py
 
 # This file marks the 'modules' directory as a Python package.
-# It can remain empty or be used for package-level initializations if needed later.
+
+# For easier imports if needed, e.g., from modules import ConfigManager
+from .audio_manager import AudioManager, AudioManagerError
+from .config_manager import ConfigManager, ConfigError
+from .context_manager import ContextManager, ContextManagerError
+from .llm_manager import LLMManager, LLMManagerError
+from .memory_manager import MemoryManager, MemoryManagerError # ADDED MemoryManager
+from .stt_manager import STTManager, STTManagerError
+from .system_manager import SystemManager, LoggingSetupError, RequirementError
+from .tts_manager import TTSManager, TTSManagerError
+from .ui_manager import UIManager
+
+__all__ = [
+    "AudioManager", "AudioManagerError",
+    "ConfigManager", "ConfigError",
+    "ContextManager", "ContextManagerError",
+    "LLMManager", "LLMManagerError",
+    "MemoryManager", "MemoryManagerError", # ADDED MemoryManager
+    "STTManager", "STTManagerError",
+    "SystemManager", "LoggingSetupError", "RequirementError",
+    "TTSManager", "TTSManagerError",
+    "UIManager",
+]

modules/context_manager.py

@@ -1,37 +1,44 @@
-# ================================================
-# FILE: modules/context_manager.py
-# ================================================
+# modules/context_manager.py
 
+import datetime
 import json
 import logging
 from pathlib import Path
 import time
 import shutil
-from typing import List, Dict, Any, Optional
+from typing import List, Dict, Any, Optional, Tuple
 
+# --- Dependency Imports with Checks ---
 try:
     from sentence_transformers import SentenceTransformer
     SENTENCE_TRANSFORMERS_AVAILABLE = True
 except ImportError:
     SENTENCE_TRANSFORMERS_AVAILABLE = False
-    # Dummy class for type hinting if needed, error raised in init
-    class SentenceTransformer: pass
+    class SentenceTransformer: # Dummy for type hints
+        def __init__(self, *args, **kwargs): pass
+        def encode(self, *args, **kwargs): return []
 
 try:
     import chromadb
-    from chromadb.config import Settings as ChromaSettings # Use specific Settings import
+    from chromadb.config import Settings as ChromaSettings
     CHROMA_AVAILABLE = True
 except ImportError:
     CHROMA_AVAILABLE = False
-    # Dummy classes/module
-    class chromadb:
+    class chromadb: # Dummy for type hints
          @staticmethod
-         def PersistentClient(*args, **kwargs): pass
-         class Collection: pass
+         def PersistentClient(*args, **kwargs): return ChromaClientDummy()
+    class ChromaClientDummy:
+        def get_or_create_collection(self, *args, **kwargs): return ChromaCollectionDummy()
+    class ChromaCollectionDummy:
+        def count(self): return 0
+        def get(self, *args, **kwargs): return {"ids": []}
+        def add(self, *args, **kwargs): pass
+        def query(self, *args, **kwargs): return {"ids": [[]], "documents": [[]], "metadatas": [[]], "distances": [[]]}
+        def delete(self, *args, **kwargs): pass
     class ChromaSettings: pass
 
 
-# Use relative import for ConfigManager
+# Local Imports
 from .config_manager import ConfigManager
 
 logger = logging.getLogger(__name__)
@@ -42,343 +49,331 @@ class ContextManagerError(Exception):
 
 class ContextManager:
     """
-    Manages conversation history using RAG.
-
-    - Stores full history in JSON.
-    - Indexes messages into a ChromaDB vector store.
-    - Retrieves relevant past messages based on user queries.
+    Manages conversation history using Retrieval-Augmented Generation (RAG).
+    - Stores the full, raw conversation history chronologically in a JSON file.
+    - Indexes each message into a persistent ChromaDB vector store.
+    - Provides methods to add messages and retrieve relevant past messages.
     """
     DEFAULT_STORAGE_PATH = "data/conversation_state.json"
     DEFAULT_VECTOR_DB_PATH = "data/chroma_db"
     DEFAULT_EMBEDDING_MODEL = "all-MiniLM-L6-v2"
     DEFAULT_COLLECTION_NAME = "mirei_chat_history"
-    DEFAULT_RETRIEVAL_RESULTS = 3
-    DEFAULT_INCLUDE_RECENT = 2
+    DEFAULT_RAG_N_RESULTS_FALLBACK = 3 # Fallback if n_results not provided to retrieve_relevant_context
 
     def __init__(self, config: ConfigManager):
-        """Initializes RAG components and loads history."""
         logger.info("Initializing ContextManager (RAG)...")
+
         if not SENTENCE_TRANSFORMERS_AVAILABLE:
-            raise ContextManagerError("Required library 'sentence-transformers' not installed. Run: uv add sentence-transformers")
+            raise ContextManagerError(
+                "Required library 'sentence-transformers' not installed. Run: uv add sentence-transformers"
+            )
         if not CHROMA_AVAILABLE:
-             raise ContextManagerError("Required library 'chromadb' not installed. Run: uv add chromadb")
+            raise ContextManagerError("Required library 'chromadb' not installed. Run: uv add chromadb")
 
         cfg_section = config.get("context_manager", default={})
-
         self.storage_path: Path = Path(cfg_section.get("storage_path", self.DEFAULT_STORAGE_PATH)).resolve()
         self.vector_db_path: Path = Path(cfg_section.get("vector_db_path", self.DEFAULT_VECTOR_DB_PATH)).resolve()
         self.embedding_model_name: str = cfg_section.get("embedding_model_name", self.DEFAULT_EMBEDDING_MODEL)
         self.collection_name: str = cfg_section.get("collection_name", self.DEFAULT_COLLECTION_NAME)
-        self.n_retrieval_results: int = int(cfg_section.get("retrieval_results", self.DEFAULT_RETRIEVAL_RESULTS))
-        self.n_include_recent: int = int(cfg_section.get("include_recent_messages", self.DEFAULT_INCLUDE_RECENT))
+        # n_retrieval_results and n_include_recent from context_manager config are now primarily
+        # used as fallbacks if methods are called without specific counts, or for direct use if any.
+        # MemoryManager will use its own config for prompt construction.
+        self.n_retrieval_results_fallback: int = int(cfg_section.get("retrieval_results", self.DEFAULT_RAG_N_RESULTS_FALLBACK))
 
-        self.messages: List[Dict[str, str]] = []
+
+        self.messages: List[Dict[str, Any]] = []
         self.embedding_model: Optional[SentenceTransformer] = None
-        self.chroma_client: Optional[chromadb.ClientAPI] = None # Use ClientAPI type hint
-        self.collection: Optional[chromadb.Collection] = None
+        self.chroma_client: Optional[chromadb.API] = None
+        self.collection: Optional[chromadb.Collection] = None # Type hint for Chroma's Collection
 
-        # 1. Load Embedding Model
-        try:
-            logger.info(f"Loading embedding model: {self.embedding_model_name}")
-            self.embedding_model = SentenceTransformer(self.embedding_model_name)
-            logger.info("Embedding model loaded successfully.")
-        except Exception as e:
-            logger.critical(f"Failed to load SentenceTransformer model '{self.embedding_model_name}': {e}", exc_info=True)
-            raise ContextManagerError(f"Failed to load embedding model: {e}") from e
-
-        # 2. Initialize ChromaDB
-        try:
-            logger.info(f"Initializing ChromaDB client at: {self.vector_db_path}")
-            # Ensure directory exists for persistent client
-            self.vector_db_path.mkdir(parents=True, exist_ok=True)
-            self.chroma_client = chromadb.PersistentClient(
-                path=str(self.vector_db_path),
-                settings=ChromaSettings(anonymized_telemetry=False) # Disable telemetry
-            )
-            # Get or create the collection
-            logger.info(f"Getting or creating Chroma collection: {self.collection_name}")
-            self.collection = self.chroma_client.get_or_create_collection(
-                name=self.collection_name,
-                # Optionally specify embedding function if not using default OpenAI
-                # metadata={"hnsw:space": "cosine"} # Default is L2, cosine often better for ST
-            )
-            logger.info(f"ChromaDB collection '{self.collection_name}' ready. Item count: {self.collection.count()}")
-
-        except Exception as e:
-             logger.critical(f"Failed to initialize ChromaDB client or collection: {e}", exc_info=True)
-             raise ContextManagerError(f"ChromaDB initialization failed: {e}") from e
-
-        # 3. Load full history from JSON
-        self._load_full_history()
-
-        # 4. Index loaded history (can be slow for large histories on first run)
-        self._initial_index()
+        self._load_embedding_model()
+        self._initialize_vector_db()
+        self._load_full_history() # This now ensures 'original_index' is present
+        self._synchronize_index()
 
         logger.info("ContextManager (RAG) initialized successfully.")
 
-    def _load_full_history(self):
-        """Loads the complete conversation history from JSON, backing up corrupted files."""
-        # This function remains largely the same as the improved version from before
-        # Just ensures self.messages holds the full history.
+    def _load_embedding_model(self) -> None:
+        logger.info(f"Loading embedding model: '{self.embedding_model_name}'...")
+        start_time = time.time()
+        try:
+            self.embedding_model = SentenceTransformer(self.embedding_model_name)
+            _ = self.embedding_model.encode(["test warm-up"], show_progress_bar=False) # Warm-up/check
+            load_time = time.time() - start_time
+            logger.info(f"Embedding model '{self.embedding_model_name}' loaded in {load_time:.2f}s.")
+        except Exception as e:
+            logger.critical(f"Failed to load SentenceTransformer model '{self.embedding_model_name}': {e}", exc_info=True)
+            raise ContextManagerError(f"Embedding model load failed: {e}") from e
+
+    def _initialize_vector_db(self) -> None:
+        logger.info(f"Initializing ChromaDB client at: {self.vector_db_path}")
+        try:
+            self.vector_db_path.mkdir(parents=True, exist_ok=True)
+            self.chroma_client = chromadb.PersistentClient(
+                path=str(self.vector_db_path),
+                settings=ChromaSettings(anonymized_telemetry=False)
+            )
+            self.collection = self.chroma_client.get_or_create_collection(
+                name=self.collection_name,
+                # metadata={"hnsw:space": "cosine"} # Optional: Explicitly set distance metric if needed
+            )
+            logger.info(f"ChromaDB collection '{self.collection_name}' ready. Initial item count: {self.collection.count()}")
+        except Exception as e:
+            logger.critical(f"Failed to initialize ChromaDB: {e}", exc_info=True)
+            raise ContextManagerError(f"ChromaDB initialization failed: {e}") from e
+
+    def _load_full_history(self) -> None:
         if self.storage_path.exists() and self.storage_path.is_file():
             try:
                 logger.info(f"Loading full conversation history from {self.storage_path}")
                 with self.storage_path.open("r", encoding="utf-8") as f:
                     content = f.read()
-                    if not content.strip():
-                        logger.warning(f"History file {self.storage_path} is empty.")
-                        self.messages = []
-                        return
-                    loaded_data = json.loads(content)
+                if not content.strip():
+                    logger.warning(f"History file '{self.storage_path}' is empty.")
+                    self.messages = []; return
 
+                loaded_data = json.loads(content)
                 if isinstance(loaded_data, list):
-                    self.messages = [
-                        msg for msg in loaded_data
-                        if isinstance(msg, dict) and "role" in msg and "content" in msg
-                    ]
-                    logger.info(f"Loaded {len(self.messages)} messages from history file.")
-                    if len(self.messages) != len(loaded_data):
-                         logger.warning("Some invalid message formats found in history file were skipped.")
+                    valid_messages = []
+                    for i, msg_dict in enumerate(loaded_data):
+                        if (isinstance(msg_dict, dict) and
+                                "role" in msg_dict and isinstance(msg_dict["role"], str) and
+                                "content" in msg_dict and # Allow empty content for system messages potentially
+                                msg_dict["role"] in ["user", "assistant", "system"]): # Allow system role
+                            # Ensure 'original_index' is present and correct
+                            msg_copy = msg_dict.copy()
+                            msg_copy['original_index'] = i # The index in the loaded list is its original_index
+                            valid_messages.append(msg_copy)
+                        else:
+                            logger.warning(f"Skipping invalid message format at index {i} in history: {msg_dict}")
+                    self.messages = valid_messages
+                    logger.info(f"Loaded {len(self.messages)} valid messages from history.")
                 else:
-                    logger.warning(f"History file {self.storage_path} does not contain a list. Starting fresh.")
+                    logger.warning(f"History file '{self.storage_path}' not a list. Starting fresh.")
                     self.messages = []
-
-            except (json.JSONDecodeError, IOError, Exception) as e:
-                error_type = type(e).__name__
-                logger.error(f"Failed to load/parse history file {self.storage_path} ({error_type}): {e}. Backing up and starting fresh.", exc_info=True)
-                try:
-                    backup_path = self.storage_path.with_name(
-                        f"{self.storage_path.stem}_corrupted_{int(time.time())}{self.storage_path.suffix}"
-                    )
-                    shutil.move(str(self.storage_path), str(backup_path))
-                    logger.info(f"Backed up corrupted history file to: {backup_path}")
-                except Exception as backup_e:
-                    logger.error(f"Failed to back up corrupted history file {self.storage_path}: {backup_e}", exc_info=True)
-                self.messages = []
+            except (json.JSONDecodeError, IOError) as e:
+                logger.error(f"Failed to load/parse history '{self.storage_path}' ({type(e).__name__}): {e}. Backing up.", exc_info=False)
+                self._backup_corrupted_file(self.storage_path); self.messages = []
+            except Exception as e:
+                logger.error(f"Unexpected error loading history '{self.storage_path}': {e}. Backing up.", exc_info=True)
+                self._backup_corrupted_file(self.storage_path); self.messages = []
         else:
-            logger.info(f"History file not found at {self.storage_path}. Starting with empty history.")
+            logger.info(f"History file not found at '{self.storage_path}'. Starting empty history.")
             self.messages = []
 
-    def _initial_index(self):
-        """Indexes messages from the loaded history if they aren't already in ChromaDB."""
-        if not self.collection or not self.embedding_model:
-            logger.error("Cannot perform initial index: Chroma collection or embedding model not available.")
-            return
+    def _backup_corrupted_file(self, file_path: Path) -> None:
+        try:
+            backup_dir = file_path.parent / "backups"; backup_dir.mkdir(exist_ok=True)
+            timestamp = datetime.datetime.now().strftime("%Y%m%d_%H%M%S")
+            backup_path = backup_dir / f"{file_path.stem}_corrupted_{timestamp}{file_path.suffix}"
+            shutil.move(str(file_path), str(backup_path))
+            logger.info(f"Backed up corrupted file to: {backup_path}")
+        except Exception as backup_e:
+            logger.error(f"Could not back up file '{file_path}': {backup_e}", exc_info=True)
 
-        logger.info("Performing initial check/indexing of loaded history...")
+    def _synchronize_index(self) -> None:
+        if not self.collection or not self.embedding_model:
+            logger.error("Cannot sync index: DB or model not ready."); return
+        logger.info("Synchronizing vector index with loaded history...")
         start_time = time.time()
-        added_count = 0
-        existing_ids = set(self.collection.get(include=[])['ids']) # Efficient way to get all IDs
+        messages_to_index: List[Dict[str, Any]] = [] # List of message dicts
+        try:
+            existing_db_ids_result = self.collection.get(include=[])
+            existing_db_ids = set(existing_db_ids_result.get('ids', []))
+            logger.debug(f"Found {len(existing_db_ids)} existing IDs in Chroma.")
 
-        ids_to_add = []
-        embeddings_to_add = []
-        documents_to_add = []
-        metadatas_to_add = []
+            for msg in self.messages:
+                # 'original_index' should have been set during _load_full_history or add_message
+                msg_original_idx = msg.get('original_index')
+                if msg_original_idx is None:
+                    logger.error(f"Message found without original_index during sync: {str(msg)[:100]}. Skipping.")
+                    continue
+                msg_id = f"msg_{msg_original_idx}"
+                if msg_id not in existing_db_ids:
+                    messages_to_index.append(msg) # msg already contains 'original_index'
 
-        for i, msg in enumerate(self.messages):
-            msg_id = f"msg_{i}" # Simple index-based ID
-            if msg_id not in existing_ids:
-                content = msg.get("content", "")
-                role = msg.get("role", "unknown")
+            if not messages_to_index: logger.info("Vector index is synchronized."); return
+            logger.info(f"Found {len(messages_to_index)} messages from history to index...")
+
+            ids_to_add, docs_to_add, metas_to_add = [], [], []
+            for message in messages_to_index:
+                content = message.get("content", "")
+                role = message.get("role", "unknown")
+                msg_original_idx = message['original_index'] # Should exist
                 if content: # Only index messages with content
+                    msg_id = f"msg_{msg_original_idx}"
                     ids_to_add.append(msg_id)
-                    # Embedding happens in batch later
-                    documents_to_add.append(content)
-                    metadatas_to_add.append({"role": role, "index": i})
-                    added_count += 1
+                    docs_to_add.append(content)
+                    metas_to_add.append({"role": role, "original_index": msg_original_idx})
 
-        # Batch embedding and adding
-        if ids_to_add:
-            logger.info(f"Found {added_count} messages from history to index...")
-            try:
-                # Calculate embeddings in batch
-                embeddings_to_add = self.embedding_model.encode(documents_to_add, show_progress_bar=False).tolist()
+            if ids_to_add:
+                logger.debug(f"Encoding {len(docs_to_add)} documents for batch indexing...")
+                embeddings = self.embedding_model.encode(docs_to_add, show_progress_bar=False).tolist()
+                logger.debug(f"Adding {len(ids_to_add)} items to ChromaDB...")
+                self.collection.add(ids=ids_to_add, embeddings=embeddings, documents=docs_to_add, metadatas=metas_to_add)
+                logger.info(f"Successfully indexed {len(ids_to_add)} messages from history.")
+            else:
+                 logger.info("No valid messages found to index after filtering.")
+        except Exception as e:
+            logger.error(f"Error during index synchronization: {e}", exc_info=True)
+        finally:
+             sync_time = time.time() - start_time
+             logger.info(f"Index sync check completed in {sync_time:.2f}s.")
 
-                # Add to ChromaDB in batch
-                self.collection.add(
-                    ids=ids_to_add,
-                    embeddings=embeddings_to_add,
-                    documents=documents_to_add,
-                    metadatas=metadatas_to_add
-                )
-                logger.info(f"Successfully indexed {added_count} messages.")
-            except Exception as e:
-                logger.error(f"Error during batch indexing: {e}", exc_info=True)
-                # Potential issue: partial add? Chroma handles batches transactionally usually.
-        else:
-            logger.info("No new messages from loaded history needed indexing.")
-
-        end_time = time.time()
-        logger.info(f"Initial indexing check completed in {end_time - start_time:.2f} seconds.")
-
-    def _index_message(self, msg_index: int, message: Dict[str, str]):
-        """Adds a single message to the vector store."""
+    def _index_message(self, msg_original_index: int, message: Dict[str, Any]) -> None:
         if not self.collection or not self.embedding_model:
-            logger.error("Cannot index message: Chroma collection or embedding model not available.")
-            return
-
-        msg_id = f"msg_{msg_index}"
+            logger.error("Cannot index: DB or model not ready."); return
+        msg_id = f"msg_{msg_original_index}"
         content = message.get("content", "")
         role = message.get("role", "unknown")
-
-        if not content:
-             logger.debug(f"Skipping indexing for message {msg_id} (no content).")
-             return
-
+        if not content: logger.debug(f"Skipping indexing for msg_{msg_original_index} (no content)."); return
         try:
-            logger.debug(f"Indexing message: {msg_id} (Role: {role})")
             embedding = self.embedding_model.encode([content], show_progress_bar=False)[0].tolist()
             self.collection.add(
-                ids=[msg_id],
-                embeddings=[embedding],
-                documents=[content],
-                metadatas=[{"role": role, "index": msg_index}]
+                ids=[msg_id], embeddings=[embedding], documents=[content],
+                metadatas=[{"role": role, "original_index": msg_original_index}]
             )
+            logger.debug(f"Indexed message: {msg_id} (Role: {role}, OrigIdx: {msg_original_index})")
         except Exception as e:
-            logger.error(f"Failed to index message {msg_id}: {e}", exc_info=True)
+            logger.error(f"Failed to index msg {msg_id}: {e}", exc_info=True)
 
-
-    def add_message(self, role: str, content: str):
-        """
-        Adds a message to the in-memory history and indexes it in the vector store.
-        Does NOT save the JSON file automatically.
-        """
-        if role not in ("user", "assistant"):
-            raise ValueError(f"Invalid message role: '{role}'. Must be 'user' or 'assistant'.")
+    def add_message(self, role: str, content: str) -> None:
+        if role not in ("user", "assistant", "system"):
+            raise ValueError(f"Invalid role: '{role}'. Must be 'user', 'assistant', or 'system'.")
         if not isinstance(content, str):
-             logger.warning(f"Message content is not a string (type: {type(content)}). Converting to string.")
-             content = str(content)
+            content = str(content)
 
-        logger.debug(f"Adding message to memory - Role: {role}, Content: '{content[:50]}...'")
         new_message = {"role": role, "content": content}
+        # Assign 'original_index' based on its future position in self.messages
+        new_message['original_index'] = len(self.messages)
         self.messages.append(new_message)
-        new_message_index = len(self.messages) - 1
+        logger.debug(
+            f"Added message to memory (OrigIdx: {new_message['original_index']}): Role={role}, Content='{content[:50]}...'"
+        )
+        self._index_message(new_message['original_index'], new_message)
 
-        # Index the new message immediately
-        self._index_message(new_message_index, new_message)
-
-        # NOTE: No condensation/truncation happens here anymore
-        # NOTE: No automatic JSON save happens here anymore
-
-    def retrieve_relevant_context(self, query: str) -> List[Dict[str, str]]:
-        """Retrieves messages from history relevant to the query."""
+    def retrieve_relevant_context(self, query: str, n_results: Optional[int] = None) -> List[Dict[str, Any]]:
         if not self.collection or not self.embedding_model:
-             logger.error("Cannot retrieve context: Chroma collection or embedding model not available.")
-             return []
-        if not query:
-             logger.warning("Cannot retrieve context for empty query.")
-             return []
-        if self.collection.count() == 0:
-             logger.debug("Skipping retrieval: Vector store is empty.")
-             return []
+            logger.error("Cannot retrieve: DB or model not ready."); return []
+        if not query: logger.warning("Cannot retrieve context for empty query."); return []
+        
+        collection_count = self.collection.count()
+        if collection_count == 0: logger.debug("Skipping retrieval: Vector store empty."); return []
 
+        num_to_retrieve = n_results if n_results is not None else self.n_retrieval_results_fallback
+        num_to_retrieve = min(num_to_retrieve, collection_count)
+        if num_to_retrieve <= 0: logger.debug("No results to retrieve."); return []
 
         try:
-            logger.debug(f"Retrieving {self.n_retrieval_results} relevant messages for query: '{query[:60]}...'")
+            logger.info(f"Retrieving up to {num_to_retrieve} messages for query: '{query[:60]}...'")
             start_time = time.time()
-
             query_embedding = self.embedding_model.encode([query], show_progress_bar=False)[0].tolist()
-
             results = self.collection.query(
-                query_embeddings=[query_embedding],
-                n_results=min(self.n_retrieval_results, self.collection.count()), # Don't request more than available
-                include=["documents", "metadatas", "distances"] # Include distance for potential filtering/logging
+                query_embeddings=[query_embedding], n_results=num_to_retrieve,
+                include=["documents", "metadatas", "distances"]
             )
+            retrieval_time = time.time() - start_time
+            logger.debug(f"ChromaDB query in {retrieval_time:.3f}s.")
 
-            end_time = time.time()
-            logger.debug(f"Retrieval query finished in {end_time - start_time:.3f} seconds.")
-
-            # Process results
             retrieved_messages = []
-            if results and results.get("ids") and results["ids"][0]: # Chroma returns lists within lists
-                retrieved_ids = results["ids"][0]
-                documents = results["documents"][0]
-                metadatas = results["metadatas"][0]
-                distances = results["distances"][0]
-
-                for i, doc_id in enumerate(retrieved_ids):
-                     role = metadatas[i].get("role", "unknown")
-                     content = documents[i]
-                     distance = distances[i]
-                     logger.debug(f"  Retrieved: ID={doc_id}, Role={role}, Distance={distance:.4f}, Content='{content[:50]}...'")
-                     retrieved_messages.append({"role": role, "content": content})
-            else:
-                 logger.debug("No relevant messages found by Chroma query.")
-
-            # Sort by original index? Chroma doesn't guarantee order, but similarity search is the primary goal.
-            # If original order is desired *among retrieved items*, we'd need to sort by metadata['index'] here.
-            # For now, return in similarity order as Chroma gives them.
-
+            if results and results.get("ids") and results["ids"][0]:
+                for i, doc_id in enumerate(results["ids"][0]):
+                    metadata = results["metadatas"][0][i] if results["metadatas"] and results["metadatas"][0] else {}
+                    original_idx_val = metadata.get("original_index", -1)
+                    try: original_idx = int(original_idx_val)
+                    except (ValueError, TypeError): original_idx = -1; logger.warning(f"Invalid original_index {original_idx_val}")
+                    
+                    retrieved_messages.append({
+                        "role": metadata.get("role", "unknown"),
+                        "content": results["documents"][0][i] if results["documents"] and results["documents"][0] else "",
+                        "metadata": {"original_index": original_idx, "distance": results["distances"][0][i] if results["distances"] and results["distances"][0] else float('inf')}
+                    })
+            logger.info(f"Retrieved {len(retrieved_messages)} relevant messages.")
             return retrieved_messages
-
         except Exception as e:
-             logger.error(f"Error during context retrieval: {e}", exc_info=True)
-             return [] # Return empty list on error
+            logger.error(f"Error during context retrieval: {e}", exc_info=True); return []
 
-    def get_recent_messages(self, num_turns: int) -> List[Dict[str, str]]:
-         """Gets the last N turns (user+assistant pairs) from history."""
-         if num_turns <= 0:
-             return []
-         # A turn is typically user + assistant, so num_messages = num_turns * 2
-         num_messages = num_turns * 2
-         return self.messages[-num_messages:] # Slice the end of the list
+    def get_recent_messages(self, num_messages_to_fetch: int) -> List[Dict[str, Any]]:
+        """Gets the last N messages. Ensures 'original_index' is present."""
+        if num_messages_to_fetch <= 0: return []
+        
+        start_idx_slice = max(0, len(self.messages) - num_messages_to_fetch)
+        recent_slice = self.messages[start_idx_slice:]
 
-    def save_context(self):
-        """Saves the current full conversation history atomically to the JSON storage file."""
-        # This function remains the same as the improved atomic save version
+        # Ensure all messages in the slice have 'original_index'.
+        # This primarily safeguards against older data formats if any were loaded
+        # without 'original_index' (though _load_full_history attempts to add it).
+        processed_recent: List[Dict[str, Any]] = []
+        for i, msg_dict in enumerate(recent_slice):
+            msg_copy = msg_dict.copy() # Work with a copy
+            if 'original_index' not in msg_copy or not isinstance(msg_copy['original_index'], int):
+                # Fallback: if somehow original_index is missing or invalid from the loaded message
+                calculated_original_idx = start_idx_slice + i
+                msg_copy['original_index'] = calculated_original_idx
+                logger.warning(
+                    f"ContextManager.get_recent_messages: Re-calculated missing/invalid 'original_index' "
+                    f"({calculated_original_idx}) for recent message: {str(msg_dict.get('content',''))[:30]}..."
+                )
+            processed_recent.append(msg_copy)
+        
+        logger.debug(f"Retrieved {len(processed_recent)} messages ({num_messages_to_fetch} requested).")
+        return processed_recent
+
+    def save_context(self) -> None:
         temp_path = self.storage_path.with_suffix(f"{self.storage_path.suffix}.tmp")
-        final_path = self.storage_path
-
         try:
-            final_path.parent.mkdir(parents=True, exist_ok=True)
-            logger.info(f"Saving full history ({len(self.messages)} messages) atomically to {final_path}")
+            self.storage_path.parent.mkdir(parents=True, exist_ok=True)
+            logger.info(f"Saving full history ({len(self.messages)} messages) to {self.storage_path}")
+            # Save role, content, and original_index to allow reconstruction
+            messages_to_save = [
+                {"role": msg["role"], "content": msg["content"], "original_index": msg.get("original_index", idx)}
+                for idx, msg in enumerate(self.messages)
+            ]
             with temp_path.open("w", encoding="utf-8") as f:
-                json.dump(self.messages, f, ensure_ascii=False, indent=2)
-            shutil.move(str(temp_path), str(final_path))
-            logger.info(f"Full history saved successfully to {final_path}")
-        except (IOError, OSError) as e:
-            logger.error(f"Failed to write history file to {final_path} (or temp file {temp_path}): {e}", exc_info=True)
-            if temp_path.exists():
-                try: temp_path.unlink()
-                except OSError: pass
+                json.dump(messages_to_save, f, ensure_ascii=False, indent=2)
+            shutil.move(str(temp_path), str(self.storage_path))
+            logger.info("Full history saved successfully.")
         except Exception as e:
-            logger.error(f"Unexpected error saving history to {final_path}: {e}", exc_info=True)
-            if temp_path.exists():
-                try: temp_path.unlink()
-                except OSError: pass
+            logger.error(f"Failed to save history '{self.storage_path}': {e}", exc_info=True)
+            if temp_path.exists(): temp_path.unlink(missing_ok=True)
 
-    def clear_context(self):
-        """Clears history in memory, clears the vector store, and saves the empty state."""
-        logger.info("Clearing conversation context (memory, vector store, and file)...")
+    def clear_context(self) -> None:
+        logger.warning("Clearing conversation context (Memory, Vector Store, File)...")
         self.messages = []
-
-        # Clear the Chroma collection
         if self.collection:
             try:
-                logger.warning(f"Deleting all items from Chroma collection: {self.collection_name}")
-                existing_ids = self.collection.get(include=[])['ids']
-                if existing_ids:
-                    self.collection.delete(ids=existing_ids)
-                logger.info("Chroma collection cleared.")
+                count = self.collection.count()
+                if count > 0:
+                    logger.info(f"Deleting {count} items from Chroma collection '{self.collection_name}'...")
+                    # Efficient way to clear a Chroma collection (if API supports `delete_collection`)
+                    # Or, if not, delete all items by IDs.
+                    # For current chromadb versions, deleting by IDs is standard.
+                    # If the collection can be deleted and recreated:
+                    # self.chroma_client.delete_collection(name=self.collection_name)
+                    # self.collection = self.chroma_client.get_or_create_collection(name=self.collection_name)
+                    # logger.info(f"Chroma collection '{self.collection_name}' deleted and recreated.")
+                    # --- OR ---
+                    all_ids_result = self.collection.get(include=[]) # Only need IDs
+                    all_ids = all_ids_result.get('ids', [])
+                    if all_ids:
+                        self.collection.delete(ids=all_ids)
+                        logger.info(f"Deleted {len(all_ids)} items from Chroma collection.")
+                    else:
+                         logger.info("Chroma collection was already empty (no IDs to delete).")
+                else:
+                    logger.info("Chroma collection already empty.")
             except Exception as e:
-                logger.error(f"Failed to clear Chroma collection '{self.collection_name}': {e}", exc_info=True)
-                # Continue with clearing memory and file even if DB clear fails
-
-        # Save the empty context to file
-        self.save_context()
+                logger.error(f"Failed to clear ChromaDB collection: {e}", exc_info=True)
+        self.save_context() # Save the empty state
+        logger.warning("Conversation context cleared.")
 
     @property
-    def history(self) -> List[Dict[str, str]]:
-         """Provides read-only access to the full message history."""
-         return list(self.messages) # Return a copy
+    def history(self) -> List[Dict[str, Any]]:
+        return list(self.messages) # Return a shallow copy
 
-    def shutdown(self):
-         """Cleanly shuts down components (if necessary)."""
-         # ChromaDB PersistentClient doesn't explicitly require shutdown usually,
-         # but can be good practice if there were explicit connections.
-         logger.info("ContextManager shutting down...")
-         # Unload embedding model? Not strictly necessary unless memory is critical.
-         self.embedding_model = None
-         self.chroma_client = None # Clear references
-         self.collection = None
-         logger.info("ContextManager resources released.")
+    def shutdown(self) -> None:
+        logger.info("ContextManager shutting down...")
+        self.embedding_model = None
+        self.chroma_client = None # Chroma client usually handles its own persistence
+        self.collection = None
+        logger.info("ContextManager shutdown complete.")

modules/llm_manager.py

@@ -1,6 +1,4 @@
-# ================================================
-# FILE: modules/llm_manager.py
-# ================================================
+# modules/llm_manager.py
 
 import logging
 import os
@@ -8,9 +6,10 @@ import queue
 import threading
 import asyncio
 import re
-from typing import Optional, List, Dict, Any
+import time
+from typing import Optional, List, Dict, Any, Union
 
-# Import OpenAI library
+# --- Dependency Imports with Checks ---
 try:
     from openai import (
         AsyncOpenAI, APIError, APIConnectionError, APITimeoutError,
@@ -19,128 +18,235 @@ try:
     OPENAI_AVAILABLE = True
 except ImportError:
     OPENAI_AVAILABLE = False
-    # Dummy classes for OpenAI API errors
+    # Dummy classes for type hints if library not installed
     class AsyncOpenAI:
-        pass
-    class APIError(Exception):
-        pass
-    class APIConnectionError(APIError):
-        pass
-    class APITimeoutError(APIConnectionError):
-        pass
-    class RateLimitError(APIError):
-        status_code = 429
-        message = "Rate limit exceeded."
-    class InternalServerError(APIError):
-        status_code = 500
-        message = "Internal server error."
-    class AuthenticationError(APIError):
-        status_code = 401
-        message = "Authentication error."
-    class BadRequestError(APIError):
-        status_code = 400
-        message = "Bad request."
+        def __init__(self, *args, **kwargs): pass
+        class chat:
+            class completions:
+                @staticmethod
+                async def create(*args, **kwargs):
+                    if False: yield # Make it an async generator type
+                    return
+    class APIError(Exception): status_code: Optional[int] = None; message: str = "OpenAI API Error"
+    class APIConnectionError(APIError): pass
+    class APITimeoutError(APIConnectionError): pass
+    class RateLimitError(APIError): status_code = 429; message = "Rate limit exceeded."
+    class InternalServerError(APIError): status_code = 500; message = "Internal server error."
+    class AuthenticationError(APIError): status_code = 401; message = "Authentication error."
+    class BadRequestError(APIError): status_code = 400; message = "Bad request."
 
-# Import Tiktoken
 try:
     import tiktoken
     TIKTOKEN_AVAILABLE = True
+    # from tiktoken import Encoding as TiktokenEncoding # More specific type
+    TiktokenEncoding = Any
 except ImportError:
     TIKTOKEN_AVAILABLE = False
-    class Tiktoken: pass # Dummy
+    class TiktokenEncoding: pass
+    class tiktoken:
+        @staticmethod
+        def encoding_for_model(model: str) -> Optional[TiktokenEncoding]: return None
+        @staticmethod
+        def get_encoding(encoding: str) -> Optional[TiktokenEncoding]: return None
 
+try:
+    from transformers import AutoTokenizer, PreTrainedTokenizerBase, AutoConfig
+    TRANSFORMERS_AVAILABLE = True
+except ImportError:
+    TRANSFORMERS_AVAILABLE = False
+    class PreTrainedTokenizerBase: pass # Dummy for type hints
+    class AutoTokenizer:
+        @staticmethod
+        def from_pretrained(model_name: str, **kwargs) -> Optional[PreTrainedTokenizerBase]: return None
+    class AutoConfig:
+         @staticmethod
+         def from_pretrained(model_name: str, **kwargs) -> Any: return None
+
+
+# Local Imports
 from .config_manager import ConfigManager
-from .context_manager import ContextManager
+from .memory_manager import MemoryManager # Changed from ContextManager
 
 logger = logging.getLogger(__name__)
 
-class LLMManagerError(Exception): pass
+class LLMManagerError(Exception):
+    """Custom exception for LLMManager operational errors."""
+    pass
 
 class LLMManager:
-    DEFAULT_SYSTEM_PROMPT = "You are a helpful AI assistant."
+    """
+    Manages asynchronous communication with an OpenAI-compatible LLM API,
+    integrating context from MemoryManager and appropriate tokenization.
+    """
+    DEFAULT_SYSTEM_PROMPT = "You are Mirei, a helpful and concise AI assistant. Respond clearly and directly using markdown. Use provided context when relevant."
     DEFAULT_TEMPERATURE = 0.7
     DEFAULT_MAX_TOKENS = 1536
     DEFAULT_TIMEOUT = 120.0
     DEFAULT_RETRIES = 1
+    DEFAULT_CONTEXT_WINDOW = 0 # 0 means disable token checks / truncation
+    DEFAULT_TOKENIZER_PREFERENCE = "auto" # "auto", "tiktoken", "transformers", "heuristic"
+    DEFAULT_CHARS_PER_TOKEN = 4 # For heuristic estimation
+    PROMPT_TRUNCATION_BUFFER = 100 # Tokens reserved (max_tokens for response + buffer)
 
     def __init__(self, config: ConfigManager, gui_queue: queue.Queue):
-        logger.info("Initializing LLMManager...")
-        if not OPENAI_AVAILABLE: raise LLMManagerError("Required library 'openai' not installed.")
-        # Warn if tiktoken missing but don't block unless context window configured
-        if not TIKTOKEN_AVAILABLE: logger.warning("Tiktoken library not found. Token length checking will be unavailable.")
+        logger.info("Initializing LLMManager (with MemoryManager & Tokenizer Logic)...")
+
+        if not OPENAI_AVAILABLE:
+            raise LLMManagerError("Required 'openai' library not installed. Run: uv add openai")
 
         self.config = config
         self.gui_queue = gui_queue
         self.llm_config = config.get_llm_config()
-        self.context_manager: Optional[ContextManager] = None
+        self.memory_manager: Optional[MemoryManager] = None # Will be linked via set_memory_manager
 
-        self.api_base_url: Optional[str] = self.llm_config.get("api_base_url")
-        self.api_key_env_var: Optional[str] = self.llm_config.get("api_key_env_var")
-        self.model_name: Optional[str] = self.llm_config.get("model_name")
-        self.system_prompt: str = self.llm_config.get("system_prompt", self.DEFAULT_SYSTEM_PROMPT)
-        self.temperature: float = float(self.llm_config.get("temperature", self.DEFAULT_TEMPERATURE))
-        self.max_tokens: int = int(self.llm_config.get("max_tokens", self.DEFAULT_MAX_TOKENS))
-        self.timeout: float = float(self.llm_config.get("timeout_seconds", self.DEFAULT_TIMEOUT))
-        self.max_retries: int = int(self.llm_config.get("max_retries", self.DEFAULT_RETRIES))
+        self._load_config_values()
+        self._validate_config()
 
-        # API Key handling... (same as before)
-        self.api_key: Optional[str] = None
-        if self.api_key_env_var and self.api_key_env_var.upper() != "NONE":
-            self.api_key = os.environ.get(self.api_key_env_var)
-            if not self.api_key: logger.warning(f"LLM API key env var '{self.api_key_env_var}' set but not found.")
-        else: logger.info("No LLM API key environment variable specified.")
-
-        if not self.api_base_url: raise LLMManagerError("LLM 'api_base_url' is missing.")
-        if not self.model_name: raise LLMManagerError("LLM 'model_name' is missing.")
-
-        # --- Tiktoken Initialization ---
-        self.encoder = None
-        self.model_context_window = int(self.llm_config.get("model_context_window", 0))
-
-        if self.model_context_window <= 0:
-            logger.warning("LLM 'model_context_window' not configured or invalid in config. Token length checking disabled.")
-        elif not TIKTOKEN_AVAILABLE:
-            logger.warning("LLM 'model_context_window' configured, but Tiktoken library not found. Token length checking disabled.")
-            self.model_context_window = 0 # Disable checking if lib missing
-        else:
-            # Try getting encoder for the specific model name
-            model_name_for_encoder = self.llm_config.get("model_name") # Use the model name from config
-            try:
-                self.encoder = tiktoken.encoding_for_model(model_name_for_encoder)
-                logger.info(f"Initialized tiktoken encoder for model: {model_name_for_encoder}")
-            except KeyError:
-                logger.warning(f"Tiktoken encoder not found for model '{model_name_for_encoder}'. Falling back to 'cl100k_base'.")
-                try:
-                    self.encoder = tiktoken.get_encoding("cl100k_base") # Common base
-                except Exception as enc_e:
-                     logger.error(f"Failed to load fallback tiktoken encoder 'cl100k_base': {enc_e}")
-                     self.encoder = None; self.model_context_window = 0 # Disable if fallback fails
-            except Exception as e:
-                logger.error(f"Failed to initialize tiktoken encoder: {e}", exc_info=True)
-                self.encoder = None; self.model_context_window = 0 # Disable on other errors
-        # --- END Tiktoken Initialization ---
-
-        # Initialize OpenAI Client
-        try:
-            client_api_key = self.api_key if self.api_key else "DUMMY_KEY"
-            self.client = AsyncOpenAI(
-                base_url=self.api_base_url, api_key=client_api_key,
-                timeout=self.timeout, max_retries=self.max_retries
-            )
-            logger.info(f"AsyncOpenAI client initialized. Base URL: {self.api_base_url}, Model: {self.model_name}")
-        except Exception as e:
-            logger.critical(f"Failed to initialize AsyncOpenAI client: {e}", exc_info=True)
-            raise LLMManagerError(f"OpenAI client initialization failed: {e}") from e
+        self.api_key = self._load_api_key()
+        self.tokenizer: Optional[Union[TiktokenEncoding, PreTrainedTokenizerBase]] = None
+        self.tokenizer_type: Optional[str] = None
+        self._initialize_tokenizer() # Now uses refined logic
+        self._log_tokenizer_status()
+        self._initialize_openai_client()
 
         self._is_processing_lock = threading.Lock()
         self._is_processing = False
         logger.info("LLMManager initialized successfully.")
 
-    def set_context_manager(self, context_manager: ContextManager) -> None:
-        if not isinstance(context_manager, ContextManager):
-            raise TypeError(f"context_manager must be instance of ContextManager, got {type(context_manager)}")
-        self.context_manager = context_manager
-        logger.info("RAG ContextManager linked to LLMManager.")
+    def _load_config_values(self) -> None:
+        self.api_base_url = self.llm_config.get("api_base_url")
+        self.api_key_env_var = self.llm_config.get("api_key_env_var")
+        self.model_name = self.llm_config.get("model_name") # Used for API call
+        # For tokenizer loading, we might use a different identifier if specified
+        self.tokenizer_source_identifier = self.llm_config.get("tokenizer_source_for_estimation", self.model_name)
+
+        self.system_prompt = self.llm_config.get("system_prompt", self.DEFAULT_SYSTEM_PROMPT)
+        self.temperature = float(self.llm_config.get("temperature", self.DEFAULT_TEMPERATURE))
+        self.max_tokens = int(self.llm_config.get("max_tokens", self.DEFAULT_MAX_TOKENS))
+        self.timeout = float(self.llm_config.get("timeout_seconds", self.DEFAULT_TIMEOUT))
+        self.max_retries = int(self.llm_config.get("max_retries", self.DEFAULT_RETRIES))
+        self.model_context_window = int(self.llm_config.get("model_context_window", self.DEFAULT_CONTEXT_WINDOW))
+        self.tokenizer_preference = self.llm_config.get("tokenizer_preference", self.DEFAULT_TOKENIZER_PREFERENCE).lower()
+        self.chars_per_token_estimate = int(self.llm_config.get("chars_per_token_estimate", self.DEFAULT_CHARS_PER_TOKEN))
+        if self.chars_per_token_estimate <= 0: self.chars_per_token_estimate = self.DEFAULT_CHARS_PER_TOKEN
+
+    def _validate_config(self) -> None:
+        if not self.api_base_url: raise LLMManagerError("LLM 'api_base_url' missing.")
+        if not self.model_name: raise LLMManagerError("LLM 'model_name' missing.")
+        if self.model_context_window > 0 and self.max_tokens >= self.model_context_window:
+            logger.warning(
+                f"Configured 'max_tokens' ({self.max_tokens}) is >= 'model_context_window' ({self.model_context_window}). "
+                "This leaves no room for the prompt. LLM calls may fail. Adjust config."
+            )
+
+    def _load_api_key(self) -> Optional[str]:
+        key, env_var = None, self.api_key_env_var
+        if env_var and env_var.upper() != "NONE":
+            key = os.environ.get(env_var)
+            if not key: logger.warning(f"LLM API key env var '{env_var}' set but not found.")
+            else: logger.debug("LLM API key loaded from environment.")
+        else: logger.info("No LLM API key env var (or set to NONE).")
+        return key
+
+    def _initialize_tokenizer(self) -> None:
+        if self.model_context_window <= 0:
+            logger.warning("model_context_window <= 0. Token checking/specific tokenizer loading disabled. Using heuristic.")
+            self.tokenizer_type = 'heuristic'; self.tokenizer = None; return
+
+        pref = self.tokenizer_preference
+        # Use tokenizer_source_identifier for loading the tokenizer
+        identifier_for_tokenizer = self.tokenizer_source_identifier
+        logger.info(f"Initializing tokenizer (Preference: '{pref}', Source for Tokenizer: '{identifier_for_tokenizer}')...")
+
+        load_successful = False
+        if pref == "tiktoken": load_successful = self._try_load_tiktoken(identifier_for_tokenizer)
+        elif pref == "transformers": load_successful = self._try_load_transformers(identifier_for_tokenizer)
+        elif pref == "auto": load_successful = self._try_auto_load_tokenizer(identifier_for_tokenizer)
+        elif pref == "heuristic": self.tokenizer_type = 'heuristic'; load_successful = True
+        else:
+            logger.error(f"Invalid 'tokenizer_preference': '{pref}'. Defaulting to heuristic.")
+            self.tokenizer_type = 'heuristic'; load_successful = True
+
+        if not load_successful:
+            logger.warning(f"Tokenizer init failed for '{identifier_for_tokenizer}' (pref: '{pref}'). Falling back to heuristic.")
+            self.tokenizer_type = 'heuristic'; self.tokenizer = None
+
+    def _try_auto_load_tokenizer(self, model_identifier: str) -> bool:
+        logger.debug(f"Auto-detecting tokenizer for: '{model_identifier}'")
+        model_lower = model_identifier.lower()
+        if model_lower.startswith("gpt-") or "ada" in model_lower or "babbage" in model_lower or "curie" in model_lower or "davinci" in model_lower or "text-embedding-" in model_lower :
+            logger.debug(f"Auto-detect: '{model_identifier}' suggests OpenAI model. Trying Tiktoken first.")
+            if self._try_load_tiktoken(model_identifier): return True
+            logger.debug(f"Tiktoken failed for '{model_identifier}'. Trying Transformers as broader attempt.")
+            if self._try_load_transformers(model_identifier): return True
+            return False
+        logger.debug(f"Auto-detect: '{model_identifier}' not an explicit OpenAI pattern. Trying Transformers first.")
+        if self._try_load_transformers(model_identifier): return True
+        logger.debug(f"Transformers failed for '{model_identifier}'. Trying Tiktoken as general fallback.")
+        if self._try_load_tiktoken(model_identifier): return True # Tiktoken tries cl100k_base
+        logger.warning(f"Auto-detection failed for '{model_identifier}'. No suitable tokenizer found by auto logic."); return False
+
+    def _try_load_tiktoken(self, model_identifier: str) -> bool:
+        if not TIKTOKEN_AVAILABLE: logger.warning("Tiktoken library not available."); return False
+        try:
+            self.tokenizer = tiktoken.encoding_for_model(model_identifier)
+            self.tokenizer_type = 'tiktoken'
+            logger.info(f"Successfully loaded Tiktoken for model: '{model_identifier}'")
+            return True
+        except KeyError:
+            logger.debug(f"Tiktoken: No direct encoding for '{model_identifier}'. Trying 'cl100k_base'.")
+            try:
+                self.tokenizer = tiktoken.get_encoding("cl100k_base")
+                self.tokenizer_type = 'tiktoken'
+                logger.info("Successfully loaded Tiktoken with 'cl100k_base' fallback.")
+                return True
+            except Exception as e_fallback: logger.warning(f"Tiktoken fallback load failed: {e_fallback}"); return False
+        except Exception as e: logger.error(f"Tiktoken init error for '{model_identifier}': {e}", exc_info=True); return False
+
+    def _try_load_transformers(self, model_identifier: str) -> bool:
+        if not TRANSFORMERS_AVAILABLE: logger.warning("Transformers library not available."); return False
+        try:
+            # AutoConfig.from_pretrained(model_identifier, trust_remote_code=True) # Optional pre-check
+            self.tokenizer = AutoTokenizer.from_pretrained(model_identifier, trust_remote_code=True, use_fast=True)
+            self.tokenizer_type = 'transformers'
+            logger.info(f"Successfully loaded Transformers tokenizer for: '{model_identifier}'")
+            return True
+        except OSError as e:
+            logger.warning(f"Transformers: Failed to load tokenizer for '{model_identifier}'. If local, ensure tokenizer files (tokenizer.model, etc.) are present or provide Hub ID. Error: {e}")
+            return False
+        except Exception as e:
+            logger.error(f"Transformers: Unexpected error for '{model_identifier}': {e}", exc_info=True)
+            return False
+
+    def _log_tokenizer_status(self) -> None:
+        if self.model_context_window <= 0:
+             logger.warning("LLM 'model_context_window' <= 0. Token length checking disabled.")
+             self.tokenizer_type = 'heuristic' # Ensure type reflects disabled checks
+        elif self.tokenizer_type == 'heuristic':
+             logger.warning(f"Using heuristic token counting (1 token ≈ {self.chars_per_token_estimate} chars).")
+        elif self.tokenizer:
+             logger.info(f"Initialized '{self.tokenizer_type}' tokenizer for '{self.tokenizer_source_identifier}'.")
+        else: # Should be covered by heuristic fallback, but as a safeguard
+             logger.error("Tokenizer initialization failed. Using heuristic token counting.")
+             self.tokenizer_type = 'heuristic'
+
+    def _initialize_openai_client(self) -> None:
+        try:
+            client_api_key = self.api_key if self.api_key else "placeholder_if_not_needed"
+            self.client = AsyncOpenAI(
+                base_url=self.api_base_url, api_key=client_api_key,
+                timeout=self.timeout, max_retries=self.max_retries
+            )
+            logger.info(f"AsyncOpenAI client initialized. Target: {self.api_base_url}")
+        except Exception as e:
+            logger.critical(f"Failed to initialize AsyncOpenAI client: {e}", exc_info=True)
+            raise LLMManagerError(f"OpenAI client initialization failed: {e}") from e
+
+    def set_memory_manager(self, memory_manager: MemoryManager) -> None: # Changed type hint
+        if not isinstance(memory_manager, MemoryManager):
+            raise TypeError("Invalid MemoryManager provided to LLMManager.")
+        self.memory_manager = memory_manager
+        logger.info("MemoryManager linked successfully to LLMManager.")
 
     def _get_current_system_prompt(self) -> Dict[str, str]:
         return {"role": "system", "content": self.system_prompt}
@@ -148,172 +254,267 @@ class LLMManager:
     def _filter_think_tags(self, text: str) -> str:
         if not text or "<think>" not in text: return text
         try:
-             think_pattern = r"<think>.*?</think>"
-             filtered_text = re.sub(think_pattern, "", text, flags=re.DOTALL).strip()
-             if len(text) != len(filtered_text):
-                 logger.debug("Filtered <think> blocks.")
-                 if not filtered_text and text: logger.warning("LLM response was only <think> blocks.")
-             return filtered_text
-        except Exception as e:
-             logger.error(f"Error filtering <think> tags: {e}", exc_info=True)
-             return text
+            filtered_text = re.sub(r"<think>.*?</think>", "", text, flags=re.DOTALL).strip()
+            if len(text) != len(filtered_text): logger.debug("Filtered <think> blocks.")
+            if not filtered_text and text: logger.warning("LLM response was only <think> blocks.")
+            return filtered_text
+        except Exception as e: logger.error(f"Error filtering <think> tags: {e}"); return text
 
-    def _estimate_prompt_tokens(self, messages: List[Dict[str, str]]) -> int:
-        """Estimates token count for messages using tiktoken, including overhead."""
-        if not self.encoder: return 0 # Cannot estimate without encoder
+    def _estimate_prompt_tokens(self, messages: List[Dict[str, Any]]) -> int:
+        if self.tokenizer_type == 'heuristic' or self.tokenizer is None or self.model_context_window <= 0:
+            char_count = sum(len(str(msg.get("content", ""))) for msg in messages)
+            estimated = char_count // self.chars_per_token_estimate
+            overhead = len(messages) * 4 # Rough overhead per message
+            final_estimate = estimated + overhead
+            logger.debug(f"Token estimation (heuristic): ~{final_estimate} tokens for {len(messages)} messages.")
+            return final_estimate
 
         num_tokens = 0
         try:
-            for message in messages:
-                num_tokens += 4  # Approximation for message overhead (role, separators)
-                for key, value in message.items():
-                    if value:
-                        # Ensure value is a string before encoding
-                        value_str = str(value)
-                        num_tokens += len(self.encoder.encode(value_str))
-            num_tokens += 3  # Approximation for priming assistant response
+            if self.tokenizer_type == 'tiktoken':
+                for message in messages:
+                    num_tokens += 4
+                    for key, value in message.items():
+                        if value: num_tokens += len(self.tokenizer.encode(str(value)))
+                    if message.get("role") == "assistant": num_tokens += 1
+                num_tokens += 3
+            elif self.tokenizer_type == 'transformers' and isinstance(self.tokenizer, PreTrainedTokenizerBase):
+                # Try to apply chat template for more accuracy if available, else sum parts.
+                try:
+                    # This is the ideal way IF the tokenizer has a well-defined chat template
+                    # and the messages are in the format it expects.
+                    # We might need to convert our messages list to what tokenizer.apply_chat_template expects
+                    # For now, using a simpler sum as a robust estimation.
+                    # chat_prompt = self.tokenizer.apply_chat_template(messages, tokenize=False, add_generation_prompt=False)
+                    # num_tokens = len(self.tokenizer.encode(chat_prompt))
+
+                    # Simpler sum-of-parts approach (fallback if apply_chat_template is tricky)
+                    current_tokens = 0
+                    for message in messages:
+                        role_tokens = len(self.tokenizer.encode(str(message.get("role", "")), add_special_tokens=False))
+                        content_tokens = len(self.tokenizer.encode(str(message.get("content", "")), add_special_tokens=False))
+                        current_tokens += role_tokens + content_tokens + 4 # Rough overhead
+                    if self.tokenizer.bos_token_id is not None: current_tokens +=1
+                    if self.tokenizer.eos_token_id is not None: current_tokens +=1
+                    num_tokens = current_tokens
+
+                except Exception as template_e:
+                    logger.warning(f"Failed to apply chat template for Transformers token estimation: {template_e}. Summing parts.")
+                    current_tokens = 0
+                    for message in messages:
+                        role_tokens = len(self.tokenizer.encode(str(message.get("role", "")), add_special_tokens=False))
+                        content_tokens = len(self.tokenizer.encode(str(message.get("content", "")), add_special_tokens=False))
+                        current_tokens += role_tokens + content_tokens + 4
+                    if self.tokenizer.bos_token_id is not None: current_tokens +=1
+                    if self.tokenizer.eos_token_id is not None: current_tokens +=1
+                    num_tokens = current_tokens
+            else:
+                logger.error(f"Unknown tokenizer type '{self.tokenizer_type}' during estimation. Using heuristic.")
+                return self._fallback_to_heuristic_estimation(messages)
+
+            logger.debug(f"Token estimation ({self.tokenizer_type}): {num_tokens} tokens for {len(messages)} messages.")
             return num_tokens
         except Exception as e:
-            logger.error(f"Error during token estimation: {e}", exc_info=True)
-            return 999999 # Return large number to likely trigger truncation on error
+            logger.error(f"{self.tokenizer_type} estimation error: {e}. Falling back to heuristic.", exc_info=True)
+            return self._fallback_to_heuristic_estimation(messages)
 
-    async def _stream_llm_response(self, user_input: str):
-        if not self.client: logger.error("LLM client not initialized."); self._signal_error("LLM Client Not Ready"); return
-        if not self.context_manager: logger.error("ContextManager not set."); self._signal_error("Context Manager Not Set"); return
+    def _fallback_to_heuristic_estimation(self, messages: List[Dict[str, Any]]) -> int:
+        """Utility to force heuristic estimation and log it."""
+        original_tokenizer_type = self.tokenizer_type
+        self.tokenizer_type = 'heuristic' # Force heuristic for this call
+        self.tokenizer = None # Clear potentially problematic tokenizer for safety
+        estimated = self._estimate_prompt_tokens(messages) # Recursive call hits heuristic branch
+        # Don't restore tokenizer_type here; if we fell back, we stay heuristic until next re-init or successful load.
+        logger.debug(f"Fell back to heuristic, estimated ~{estimated} tokens for {len(messages)} messages.")
+        return estimated
 
-        logger.info("Starting LLM RAG request.")
-        self.gui_queue.put({"type": "status", "payload": "Retrieving Context..."})
+    def _truncate_prompt(self, messages: List[Dict[str, Any]], max_prompt_tokens: int) -> List[Dict[str, Any]]:
+        estimated_tokens = self._estimate_prompt_tokens(messages)
+        logger.debug(f"Truncating prompt. Current estimated tokens: {estimated_tokens}, Target: <= {max_prompt_tokens}")
 
-        full_response_text = ""
-        error_occurred = False; error_message = ""; status_code = None
-        retrieved_context: List[Dict[str, str]] = []; recent_messages: List[Dict[str, str]] = []
+        if len(messages) <= 1: # Should at least have system or user
+            if estimated_tokens > max_prompt_tokens:
+                raise LLMManagerError(f"Cannot truncate: Single message prompt ({estimated_tokens} tokens) exceeds limit ({max_prompt_tokens}).")
+            return messages
+
+        # Identify system prompt (if any) and last user message to preserve them
+        system_prompt_msg: Optional[Dict[str, Any]] = None
+        last_user_msg_idx = -1
+
+        if messages[0].get("role") == "system":
+            system_prompt_msg = messages[0]
+            core_messages_start_idx = 1
+        else:
+            core_messages_start_idx = 0
+        
+        # Find the last user message
+        for i in range(len(messages) - 1, core_messages_start_idx -1, -1):
+            if messages[i].get("role") == "user":
+                last_user_msg_idx = i
+                break
+        
+        if last_user_msg_idx == -1 and messages[-1].get("role") != "system": # No user message, but not just system
+            last_user_msg_idx = len(messages) -1 # Treat the last message as immutable if no explicit user message
+        
+        # Messages that can be removed (between system prompt and last user message, or all but last if no system/user)
+        mutable_history: List[Dict[str, Any]] = []
+        final_fixed_messages: List[Dict[str, Any]] = []
+
+        if system_prompt_msg:
+            mutable_history = messages[core_messages_start_idx : last_user_msg_idx if last_user_msg_idx != -1 else len(messages)]
+            final_fixed_messages.append(system_prompt_msg)
+        else:
+            mutable_history = messages[core_messages_start_idx : last_user_msg_idx if last_user_msg_idx != -1 else len(messages)]
+
+        if last_user_msg_idx != -1 and last_user_msg_idx < len(messages): # Ensure last_user_msg_idx is valid
+             # Add messages before last user message to mutable
+             if last_user_msg_idx > core_messages_start_idx :
+                  mutable_history = messages[core_messages_start_idx : last_user_msg_idx]
+             else: # last_user_msg_idx is the first core message or doesn't exist
+                  mutable_history = [] # No history to remove before last user message
+
+             if last_user_msg_idx < len(messages): # If there IS a last user message (not just system)
+                final_fixed_messages.append(messages[last_user_msg_idx])
+        elif not final_fixed_messages: # No system and no identified last user, keep last message
+             if messages:
+                  mutable_history = messages[:-1]
+                  final_fixed_messages.append(messages[-1])
+
+
+        current_messages_for_truncation = (([system_prompt_msg] if system_prompt_msg else []) +
+                                          mutable_history +
+                                          ([messages[last_user_msg_idx]] if last_user_msg_idx != -1 and last_user_msg_idx < len(messages) else []))
+
+
+        while self._estimate_prompt_tokens(current_messages_for_truncation) > max_prompt_tokens and mutable_history:
+            removed = mutable_history.pop(0) # Remove oldest from the mutable part
+            logger.debug(f"Truncating message: Role={removed.get('role')}, Content='{str(removed.get('content'))[:30]}...'")
+            current_messages_for_truncation = (([system_prompt_msg] if system_prompt_msg else []) +
+                                              mutable_history +
+                                              ([messages[last_user_msg_idx]] if last_user_msg_idx != -1 and last_user_msg_idx < len(messages) else []))
+            logger.debug(f" > New estimated tokens: {self._estimate_prompt_tokens(current_messages_for_truncation)}")
+
+        final_prompt_construct = current_messages_for_truncation
+        final_tokens = self._estimate_prompt_tokens(final_prompt_construct)
+
+        if final_tokens > max_prompt_tokens:
+            raise LLMManagerError(f"Prompt too long ({final_tokens} > {max_prompt_tokens}) after trying to truncate. Critical messages might be too large.")
+
+        logger.info(f"Prompt truncated to approx {final_tokens} tokens.")
+        return final_prompt_construct
+
+    async def _stream_llm_response(self, user_input: str) -> None:
+        if not self.client: self._signal_error("LLM Client Error"); return
+        if not self.memory_manager: self._signal_error("MemoryManager Link Error"); return # Use MemoryManager
+
+        logger.info("Starting LLM stream (using MemoryManager).")
+        self.gui_queue.put({"type": "status", "payload": "Constructing Context..."})
+
+        full_response_text = ""; error_occurred = False; error_message = ""; status_code = None
+        final_messages_sent: List[Dict[str, Any]] = []
 
         try:
-            # 1. Retrieve Context
-            try: retrieved_context = self.context_manager.retrieve_relevant_context(user_input)
-            except Exception as rag_e: logger.error(f"Failed RAG retrieval: {rag_e}", exc_info=True); self.gui_queue.put({"type": "log", "payload": "Warn: History retrieval failed.", "tag": "warning"})
-
-            # 2. Get Recent Messages
-            if self.context_manager.n_include_recent > 0:
-                 try: recent_messages = self.context_manager.get_recent_messages(self.context_manager.n_include_recent)
-                 except Exception as recent_e: logger.error(f"Failed get recent msgs: {recent_e}", exc_info=True)
-
-            # 3. Construct Initial Message List
-            messages: List[Dict[str, Any]] = []
-            messages.append(self._get_current_system_prompt())
-            if retrieved_context:
-                try: retrieved_context.sort(key=lambda x: x.get('metadata', {}).get('index', float('inf')))
-                except Exception: pass
-                context_str = "\n".join([f"- {msg['role']}: {msg['content']}" for msg in retrieved_context])
-                messages.append({"role": "system", "content": f"## Relevant Context:\n{context_str}\n## End Context"})
-                logger.debug(f"Added {len(retrieved_context)} retrieved messages.")
-            if recent_messages:
-                 messages.extend(recent_messages); logger.debug(f"Added {len(recent_messages)} recent messages.")
-            messages.append({"role": "user", "content": user_input})
-
-            # 4. Check and Truncate Prompt if Necessary
-            if self.encoder and self.model_context_window > 0:
-                # Reserve space for response and a small buffer
-                buffer = 50 # Add a small buffer for safety
-                max_prompt_tokens = self.model_context_window - self.max_tokens - buffer
-                if max_prompt_tokens <= 0:
-                     logger.warning("Configured model_context_window too small for max_tokens. Check config.")
-                     max_prompt_tokens = self.model_context_window // 2 # Use half window as fallback limit
+            context_messages = self.memory_manager.construct_prompt_context(user_input)
+            messages = [self._get_current_system_prompt()] + context_messages + [{"role": "user", "content": user_input}]
+            final_messages_sent = messages
 
+            if self.model_context_window > 0:
+                max_prompt_tokens = self.model_context_window - self.max_tokens - self.PROMPT_TRUNCATION_BUFFER
+                if max_prompt_tokens <= 50: max_prompt_tokens = 50 # Min sensible limit
                 estimated_tokens = self._estimate_prompt_tokens(messages)
-                logger.debug(f"Est. prompt tokens: {estimated_tokens}. Max allowed: {max_prompt_tokens}")
-
-                # Truncation Loop
-                while estimated_tokens > max_prompt_tokens and len(messages) > 2: # Keep SysPrompt+UserQuery minimum
-                    logger.warning(f"Prompt too long ({estimated_tokens} > {max_prompt_tokens}). Reducing...")
-                    # Remove the oldest context message (index 1, after system prompt)
-                    removed_message = messages.pop(1)
-                    logger.debug(f"Removed oldest context: Role={removed_message.get('role')}, Content='{str(removed_message.get('content'))[:30]}...'")
-                    estimated_tokens = self._estimate_prompt_tokens(messages) # Recalculate
-
-                # Final check after truncation attempts
                 if estimated_tokens > max_prompt_tokens:
-                    logger.error(f"Cannot reduce prompt ({estimated_tokens}) below limit ({max_prompt_tokens}). Aborting request.")
-                    self._signal_error("Prompt Too Long", f"Cannot fit prompt within model limit ({self.model_context_window}).")
-                    return # Stop processing
+                    logger.warning(f"Prompt tokens ({estimated_tokens}) > limit ({max_prompt_tokens}). Truncating...")
+                    final_messages_sent = self._truncate_prompt(messages, max_prompt_tokens)
+            # ... (Rest of API call and streaming logic remains the same as your provided `_stream_llm_response`) ...
+            # Ensure to use `final_messages_sent` for the API call.
+            # When assistant response is received, add it using `self.memory_manager.add_message("assistant", ...)`.
 
-            # 5. Make API Call
-            logger.debug(f"Sending {len(messages)} final messages to LLM.")
+            # (Beginning of existing API call block)
+            logger.info(f"Sending {len(final_messages_sent)} messages to LLM API.")
             self.gui_queue.put({"type": "status", "payload": "Thinking..."})
             stream = await self.client.chat.completions.create(
-                model=self.model_name, messages=messages,
-                temperature=self.temperature, max_tokens=self.max_tokens, stream=True
+                model=self.model_name,
+                messages=final_messages_sent, # Use the potentially truncated list
+                temperature=self.temperature,
+                max_tokens=self.max_tokens,
+                stream=True
             )
 
-            # 6. Process Stream Safely
-            logger.debug("LLM stream opened. Reading chunks...")
+            logger.debug("LLM response stream opened.")
             async for chunk in stream:
-                choice = None; delta_content = None; finish_reason = None
-                if chunk.choices:
-                    choice = chunk.choices[0]
-                    if choice.delta: delta_content = choice.delta.content
-                    finish_reason = choice.finish_reason
+                delta_content = chunk.choices[0].delta.content if chunk.choices and chunk.choices[0].delta else None
+                finish_reason = chunk.choices[0].finish_reason if chunk.choices and chunk.choices[0].finish_reason else None
                 if delta_content:
                     full_response_text += delta_content
                     self.gui_queue.put({"type": "llm_chunk", "payload": {"delta": delta_content}})
                 if finish_reason:
-                     logger.info(f"LLM stream finished. Reason: {finish_reason}")
-                     if finish_reason == "length":
-                         logger.warning("LLM response may be truncated (max_tokens).")
-                         self.gui_queue.put({"type": "log", "payload": "Assistant response might be cut short.", "tag": "warning"})
-                     elif finish_reason not in ["stop", None]:
-                         logger.warning(f"LLM stopped unexpectedly: {finish_reason}")
+                    logger.info(f"LLM stream ended. Finish reason: '{finish_reason}'")
+                    if finish_reason == "length":
+                        logger.warning("LLM response potentially truncated due to 'max_tokens' limit.")
+                        self.gui_queue.put({"type": "log", "payload": "Assistant response may be incomplete (max tokens).", "tag": "warning"})
+                    break 
 
-            if not full_response_text.strip() and not error_occurred:
-                logger.warning("LLM returned empty response after stream.")
+            if not full_response_text.strip():
+                logger.warning("LLM stream completed but yielded empty text content.")
+            # (End of existing API call block - ensure error handling follows)
 
-        # 7. Handle API Errors
-        except AuthenticationError as e: error_occurred=True; status_code=getattr(e,'status_code',401); error_message=f"Auth error ({status_code}): Check API key/perms. {getattr(e,'message',str(e))}"; logger.error(error_message,exc_info=True)
-        except BadRequestError as e: error_occurred=True; status_code=getattr(e,'status_code',400); error_message=f"Bad request ({status_code}): Prompt issue? {getattr(e,'message',str(e))}"; logger.error(error_message,exc_info=True)
-        except APIConnectionError as e: error_occurred=True; error_message=f"Network error: {self.api_base_url}. {e}"; logger.error(error_message,exc_info=True)
-        except APITimeoutError as e: error_occurred=True; error_message=f"LLM request timed out ({self.timeout}s). {e}"; logger.error(error_message)
-        except RateLimitError as e: error_occurred=True; status_code=getattr(e,'status_code',429); error_message=f"Rate limit ({status_code}). {getattr(e,'message',str(e))}"; logger.error(error_message)
-        except InternalServerError as e: error_occurred=True; status_code=getattr(e,'status_code',500); error_message=f"LLM server error ({status_code}): {getattr(e,'message',str(e))}"; logger.error(error_message,exc_info=True)
-        except APIError as e: error_occurred=True; status_code=getattr(e,'status_code','N/A'); error_message=f"LLM API error ({status_code}): {getattr(e,'message',str(e))}"; logger.error(error_message,exc_info=True)
-        except Exception as e: error_occurred=True; error_message=f"Unexpected LLM comm error: {type(e).__name__}: {e}"; logger.error(error_message,exc_info=True)
+        except LLMManagerError as prep_err:
+            error_occurred = True; error_message = str(prep_err); logger.error(error_message, exc_info=False)
+        except AuthenticationError as e: error_occurred=True; status_code=getattr(e,'status_code',401); error_message=f"Auth Error ({status_code}): Check API key. {getattr(e,'message',str(e))}"
+        except BadRequestError as e: error_occurred=True; status_code=getattr(e,'status_code',400); error_message=f"Bad Request ({status_code}): Invalid prompt/model? {getattr(e,'message',str(e))}"
+        except APIConnectionError as e: error_occurred=True; error_message=f"Network Error connecting to {self.api_base_url}. {e}"
+        except APITimeoutError as e: error_occurred=True; error_message=f"LLM request timed out ({self.timeout}s). {e}"
+        except RateLimitError as e: error_occurred=True; status_code=getattr(e,'status_code',429); error_message=f"Rate Limit Error ({status_code}). {getattr(e,'message',str(e))}"
+        except InternalServerError as e: error_occurred=True; status_code=getattr(e,'status_code',500); error_message=f"LLM Server Error ({status_code}): {getattr(e,'message',str(e))}"
+        except APIError as e: error_occurred=True; status_code=getattr(e,'status_code','N/A'); error_message=f"Generic LLM API Error ({status_code}): {getattr(e,'message',str(e))}"
+        except Exception as e: # Catch other prep or API errors
+            error_occurred = True; error_message = f"Unexpected error during LLM processing: {e}"; logger.error(error_message, exc_info=True)
 
-        # 8. Finalize
-        finally:
-            filtered_response_text = self._filter_think_tags(full_response_text)
-            if not error_occurred and filtered_response_text.strip():
-                try:
-                    if self.context_manager:
-                        self.context_manager.add_message("assistant", filtered_response_text)
-                        logger.debug("Assistant response added/indexed.")
-                    else: error_occurred=True; error_message="Internal error: Context manager lost."
-                except Exception as e: logger.error(f"Could not save/index assist reply: {e}",exc_info=True); error_occurred=True; error_message=f"Failed to save context: {e}"
+        filtered_response_text = self._filter_think_tags(full_response_text)
 
-            final_payload = {"text": filtered_response_text if not error_occurred else None, "error": error_occurred, "error_message": error_message if error_occurred else None, "status_code": status_code if error_occurred else None}
-            self.gui_queue.put({"type": "llm_result", "payload": final_payload})
-            if error_occurred: self._signal_error(f"LLM Failed: {error_message.split('.')[0]}", log_message=error_message)
-            logger.info("LLM RAG stream processing finished.")
+        if not error_occurred and filtered_response_text.strip():
+            try:
+                if self.memory_manager: # Use MemoryManager
+                    self.memory_manager.add_message("assistant", filtered_response_text)
+                    logger.debug("Assistant response added via MemoryManager.")
+                else: logger.error("MemoryManager reference lost post-processing.")
+            except Exception as ctx_e: logger.error(f"Failed to add assistant response via MemoryManager: {ctx_e}", exc_info=True)
 
-    def run_llm_in_background(self, user_input: str):
+        final_payload = {
+            "text": filtered_response_text if not error_occurred else None,
+            "error": error_occurred, "error_message": error_message if error_occurred else None,
+            "status_code": status_code if error_occurred and status_code else (200 if not error_occurred else None)
+        }
+        self.gui_queue.put({"type": "llm_result", "payload": final_payload})
+
+        if error_occurred: self._signal_error(f"LLM Failed ({status_code or 'N/A'})", error_message)
+        logger.info("LLM stream processing method finished.")
+
+
+    def run_llm_in_background(self, user_input: str) -> None:
         if not self._is_processing_lock.acquire(blocking=False):
-            logger.warning("LLM busy. Request ignored."); self.gui_queue.put({"type": "log", "payload": "Assistant is busy.", "tag": "warning"}); return
+            logger.warning("LLM busy. Request ignored.")
+            self.gui_queue.put({"type": "log", "payload": "Assistant is busy.", "tag": "warning"})
+            return
         self._is_processing = True
-        logger.info("Starting LLM background thread.")
+        logger.info(f"Starting LLM background thread for: '{user_input[:50]}...'")
         thread = threading.Thread(target=self._run_llm_thread_target, args=(user_input,), daemon=True, name="LLMStreamThread")
         thread.start()
 
-    def _run_llm_thread_target(self, user_input: str):
+    def _run_llm_thread_target(self, user_input: str) -> None:
         loop = None
         try:
             try: loop = asyncio.get_running_loop()
             except RuntimeError: loop = asyncio.new_event_loop(); asyncio.set_event_loop(loop)
             loop.run_until_complete(self._stream_llm_response(user_input))
-        except Exception as e: logger.error(f"Fatal error in LLM thread: {e}", exc_info=True); self._signal_error("LLM Task Failed", str(e))
+        except Exception as e:
+            logger.error(f"Critical error in LLM background thread: {e}", exc_info=True)
+            self._signal_error("LLM Task Failed Unexpectedly", f"Error: {e}")
         finally:
             self._is_processing = False
             self._is_processing_lock.release()
             logger.debug("LLM processing lock released.")
 
-    def _signal_error(self, status_message: str, log_message: Optional[str] = None):
+    def _signal_error(self, status_message: str, log_message: Optional[str] = None) -> None:
         self.gui_queue.put({"type": "status", "payload": f"ERROR: {status_message}"})
-        log_msg = log_message if log_message else status_message
-        self.gui_queue.put({"type": "log", "payload": log_msg, "tag": "error"})
+        self.gui_queue.put({"type": "log", "payload": log_message or status_message, "tag": "error"})

modules/memory_manager.py

@@ -0,0 +1,227 @@
+# modules/memory_manager.py
+
+import logging
+from typing import List, Dict, Any, Optional
+
+# Local Imports
+from .config_manager import ConfigManager
+from .context_manager import ContextManager
+
+logger = logging.getLogger(__name__)
+
+class MemoryManagerError(Exception):
+    """Custom exception for MemoryManager specific errors."""
+    pass
+
+class MemoryManager:
+    """
+    Manages the construction of conversation context for the LLM,
+    orchestrating short-term (recent) and long-term (RAG) memory
+    retrieval from the ContextManager.
+    """
+    DEFAULT_STM_WINDOW_TURNS = 2
+    DEFAULT_LTM_RETRIEVAL_COUNT = 3
+
+    def __init__(self, config: ConfigManager, context_manager: ContextManager):
+        logger.info("Initializing MemoryManager...")
+        if not isinstance(config, ConfigManager):
+            raise MemoryManagerError("Invalid ConfigManager instance provided to MemoryManager.")
+        if not isinstance(context_manager, ContextManager):
+            raise MemoryManagerError("Invalid ContextManager instance provided to MemoryManager.")
+
+        self.config = config
+        self.context_manager = context_manager
+
+        mem_cfg = self.config.get("memory_manager", default={})
+        self.short_term_window_turns: int = int(mem_cfg.get(
+            "short_term_window_turns", self.DEFAULT_STM_WINDOW_TURNS
+        ))
+        self.long_term_retrieval_count: int = int(mem_cfg.get(
+            "long_term_retrieval_count", self.DEFAULT_LTM_RETRIEVAL_COUNT
+        ))
+
+        if self.short_term_window_turns < 0:
+            logger.warning(f"MemoryManager 'short_term_window_turns' ({self.short_term_window_turns}) cannot be negative. Setting to 0.")
+            self.short_term_window_turns = 0
+        if self.long_term_retrieval_count < 0:
+            logger.warning(f"MemoryManager 'long_term_retrieval_count' ({self.long_term_retrieval_count}) cannot be negative. Setting to 0.")
+            self.long_term_retrieval_count = 0
+
+        logger.info(
+            f"MemoryManager configured: STM Turns={self.short_term_window_turns}, "
+            f"LTM Count={self.long_term_retrieval_count}"
+        )
+        logger.info("MemoryManager initialized successfully.")
+
+    def add_message(self, role: str, content: str) -> None:
+        """
+        Adds a message to the underlying ContextManager, which handles
+        in-memory history and vector indexing.
+        """
+        try:
+            self.context_manager.add_message(role, content)
+            logger.debug(f"MemoryManager: Message (Role: {role}) passed to ContextManager.")
+        except Exception as e:
+            logger.error(f"MemoryManager: Error adding message via ContextManager: {e}", exc_info=True)
+
+    def construct_prompt_context(self, current_query: str) -> List[Dict[str, str]]:
+        """
+        Constructs a list of context messages for the LLM prompt.
+        The returned context should be *prior to* the current_query, as
+        LLMManager will append the current_query.
+        It attempts to enforce alternating user/assistant roles.
+        """
+        logger.debug(f"Constructing prompt context leading up to query: '{current_query[:50]}...'")
+
+        # 1. Retrieve Long-Term Memory (RAG)
+        retrieved_ltm: List[Dict[str, Any]] = []
+        if self.long_term_retrieval_count > 0:
+            try:
+                retrieved_ltm = self.context_manager.retrieve_relevant_context(
+                    query=current_query,
+                    n_results=self.long_term_retrieval_count
+                )
+                logger.debug(f"Retrieved {len(retrieved_ltm)} LTM messages via RAG.")
+            except Exception as e:
+                logger.error(f"Error retrieving LTM from ContextManager: {e}", exc_info=True)
+        else:
+            logger.debug("LTM retrieval skipped (long_term_retrieval_count is 0).")
+
+        # 2. Retrieve Short-Term Memory (Recent Messages)
+        num_recent_messages_to_fetch = self.short_term_window_turns * 2
+        recent_stm: List[Dict[str, Any]] = []
+        if num_recent_messages_to_fetch > 0:
+            try:
+                # This fetches messages from history which now includes the current_query
+                recent_stm = self.context_manager.get_recent_messages(num_recent_messages_to_fetch)
+                logger.debug(f"Retrieved {len(recent_stm)} candidate STM messages (recent).")
+            except Exception as e:
+                logger.error(f"Error retrieving STM from ContextManager: {e}", exc_info=True)
+        else:
+            logger.debug("STM retrieval skipped (short_term_window_turns is 0).")
+
+        # 3. Combine and initially sort all candidate messages by original_index
+        combined_candidates_dict: Dict[int, Dict[str, str]] = {}
+        current_query_original_index: Optional[int] = None
+
+        # Determine original_index of current_query (it's the last in context_manager.messages)
+        if self.context_manager.messages:
+            last_message_in_full_history = self.context_manager.messages[-1]
+            if last_message_in_full_history.get("role") == "user" and \
+               last_message_in_full_history.get("content") == current_query:
+                current_query_original_index = last_message_in_full_history.get("original_index")
+                logger.debug(f"Identified current_query to exclude with original_index: {current_query_original_index}")
+
+        for msg_source_name, msg_source_list in [("LTM", retrieved_ltm), ("STM", recent_stm)]:
+            for msg in msg_source_list:
+                original_index = msg.get("metadata", {}).get("original_index") if msg_source_name == "LTM" else msg.get("original_index")
+
+                if original_index is not None and isinstance(original_index, int):
+                    # Exclude the current_query itself from the context being built
+                    if current_query_original_index is not None and original_index == current_query_original_index:
+                        logger.debug(f"Skipping current_query (OrigIdx: {original_index}) from {msg_source_name} during initial assembly.")
+                        continue
+                    
+                    content = msg.get("content", "").strip()
+                    role = msg.get("role", "unknown")
+                    if content and role in ["user", "assistant"]: # Only consider valid roles and non-empty content
+                        combined_candidates_dict[original_index] = {"role": role, "content": content}
+                else:
+                    logger.warning(f"{msg_source_name} message missing valid 'original_index': {msg.get('content', '')[:30]}...")
+        
+        sorted_indices = sorted(combined_candidates_dict.keys())
+        chronological_context: List[Dict[str, str]] = [
+            combined_candidates_dict[idx] for idx in sorted_indices
+        ]
+        logger.debug(f"Assembled {len(chronological_context)} chronological context candidates (pre-alternation).")
+
+        # 4. Enforce alternating roles to build final_context_messages
+        final_context_messages: List[Dict[str, str]] = []
+        last_added_role: Optional[str] = None
+
+        for msg in chronological_context:
+            current_role = msg["role"] # role should be "user" or "assistant" at this point
+            current_content = msg["content"] # content should be non-empty and stripped
+
+            if not final_context_messages: # First message to add to context
+                # The very first message in history (after system prompt, handled by LLMManager)
+                # ideally should be a 'user' message for most models.
+                # However, if RAG pulls an 'assistant' message as the oldest relevant,
+                # and there's no preceding 'user' message in `chronological_context`,
+                # we might have an issue.
+                # For now, let's just add the first valid message.
+                final_context_messages.append(msg)
+                last_added_role = current_role
+                logger.debug(f"Alternation: Adding first message to context: Role={current_role}, Content='{current_content[:30]}...'")
+            elif current_role != last_added_role:
+                final_context_messages.append(msg)
+                last_added_role = current_role
+                logger.debug(f"Alternation: Adding message (role changed): Role={current_role}, Content='{current_content[:30]}...'")
+            else: # Roles are the same as the last added message
+                if current_role == "user":
+                    # Merge with the previous user message
+                    final_context_messages[-1]["content"] = (final_context_messages[-1]["content"] + "\n" + current_content).strip()
+                    logger.debug(f"Alternation: Merged user message. New combined content starts: '{final_context_messages[-1]['content'][:30]}...'")
+                elif current_role == "assistant":
+                    # Replace the previous assistant message with this (presumably more relevant or later chronological) one
+                    logger.debug(f"Alternation: Replacing previous assistant message ('{final_context_messages[-1]['content'][:30]}...') with new one ('{current_content[:30]}...').")
+                    final_context_messages[-1] = msg
+                    # last_added_role remains "assistant"
+        
+        # Final check: The context being returned to LLMManager should not cause an
+        # [System, Assistant, User (current_query)] sequence if the history is short
+        # and only an assistant message was selected for context.
+        # If the very first message of our context is "assistant", and there's nothing before it,
+        # it means the LLM prompt will be System, Assistant, User(current). This is often bad.
+        # So, if `final_context_messages` has only one message and it's an assistant, we might clear it,
+        # or if it starts with assistant and the *overall true history* implies a user should have come before it.
+        # This specific edge case (first message in context being assistant) is what was causing the issue.
+        # LLMManager adds System then current User. Context goes in between.
+        # Prompt: System, [Context Messages], User (current)
+        # If Context Messages = [Assistant, User, Assistant]
+        # Result: System, Assistant, User, Assistant, User ( PROBLEM: S, A)
+        if final_context_messages and final_context_messages[0].get("role") == "assistant":
+            # If the very first message in our constructed context is 'assistant',
+            # it will directly follow the 'system' prompt if no other 'user' message
+            # precedes it from an earlier part of history not included in this RAG/STM window.
+            # This is a common cause for the alternation error.
+            # We remove this leading assistant message to allow the subsequent 'user' (current_query)
+            # to follow the system prompt, or to allow a 'user' message later in final_context_messages
+            # to be the first non-system message.
+            logger.warning(
+                f"Alternation: First message in constructed context is 'assistant' ('{final_context_messages[0]['content'][:30]}...'). "
+                "Removing it to prevent System-Assistant start for the LLM."
+            )
+            final_context_messages.pop(0)
+            
+            # After removing, if the new first message is same role as next, re-evaluate (simple fix)
+            if len(final_context_messages) >= 2 and final_context_messages[0].get("role") == final_context_messages[1].get("role"):
+                logger.debug("Post-pop alternation check: Consecutive roles found after removing leading assistant.")
+                if final_context_messages[0].get("role") == "user": # Two users
+                    merged_user_content = (final_context_messages[0]["content"] + "\n" + final_context_messages[1]["content"]).strip()
+                    final_context_messages[0]["content"] = merged_user_content
+                    final_context_messages.pop(1)
+                    logger.debug("Merged consecutive users after pop.")
+                # Not typically expecting two assistants after pop, but could be added if needed.
+
+
+        logger.info(f"MemoryManager returning final alternating prompt context with {len(final_context_messages)} messages.")
+        return final_context_messages
+
+    def clear_memory(self) -> None:
+        logger.info("MemoryManager: Clearing all memory via ContextManager.")
+        try:
+            self.context_manager.clear_context()
+        except Exception as e:
+            logger.error(f"MemoryManager: Error clearing memory via ContextManager: {e}", exc_info=True)
+
+    def get_full_history(self) -> List[Dict[str, Any]]:
+        try:
+            return self.context_manager.history
+        except Exception as e:
+            logger.error(f"MemoryManager: Error retrieving full history from ContextManager: {e}", exc_info=True)
+            return []
+
+    def shutdown(self) -> None:
+        logger.info("MemoryManager shutting down...")
+        logger.info("MemoryManager shutdown complete.")

modules/system_manager.py

@@ -1,14 +1,5 @@
-# ================================================
-# FILE: modules/system_manager.py (Corrected Checks)
-# ================================================
-"""
-SystemManager for MiraiAssist.
+# modules/system_manager.py
 
-Handles:
-- Application-wide logging setup (console + file), optionally using Rich for console.
-- System information logging (OS, Python, hardware).
-- Critical runtime requirement validation.
-"""
 from __future__ import annotations
 
 import sys
@@ -23,30 +14,26 @@ import shutil
 from datetime import datetime
 from pathlib import Path
 from types import ModuleType
-from typing import Any, Dict, List, Optional, Tuple, Union
+from typing import Any, Dict, List, Optional, Tuple, Union, Callable
 from urllib.parse import urlparse
 
-# Use relative import for ConfigManager and other local modules
 from .config_manager import ConfigManager
-# Corrected import path if context_manager is directly under modules
-from .context_manager import ContextManager
+from .context_manager import ContextManager # For default paths
 
-# --- Rich Integration ---
 try:
     from rich.logging import RichHandler
     RICH_AVAILABLE = True
 except ImportError:
     RICH_AVAILABLE = False
-    class RichHandler: pass # Dummy class
-# ------------------------
+    class RichHandler: # type: ignore
+        def __init__(self, *args, **kwargs): pass # Dummy for type hints
 
-# Optional PyTorch import
 try:
     import torch
     TORCH_AVAILABLE = True
 except ImportError:
     TORCH_AVAILABLE = False
-    class torch: # Dummy class
+    class torch: # type: ignore # Dummy class
         @staticmethod
         def cuda_is_available(): return False
         @staticmethod
@@ -54,7 +41,7 @@ except ImportError:
         @staticmethod
         def cuda_get_device_name(i): return ""
         __version__ = "Not Installed"
-        class version: cuda = "N/A"
+        class version: cuda = "N/A" # type: ignore
 
 logger = logging.getLogger(__name__)
 
@@ -68,14 +55,7 @@ class RequirementError(Exception):
 
 
 class SystemManager:
-    """
-    Manages system-level tasks like logging setup, information gathering,
-    and requirement verification for the MiraiAssist application.
-    """
-
-    MIN_PYTHON_VERSION: Tuple[int, int] = (3, 9) # Minimum required Python version
-
-    # Default logging settings (used if config is missing keys)
+    MIN_PYTHON_VERSION: Tuple[int, int] = (3, 9)
     DEFAULT_LOG_FORMAT = '%(asctime)s [%(levelname)-8s] %(name)-25s %(message)s'
     DEFAULT_DATE_FORMAT = '%Y-%m-%d %H:%M:%S'
     DEFAULT_CONSOLE_LEVEL = "INFO"
@@ -84,13 +64,10 @@ class SystemManager:
     DEFAULT_LOG_DIR = "logs"
     DEFAULT_APP_LOG_FILE = "mirai_assist.log"
     DEFAULT_ERROR_LOG_FILE = "errors.log"
-    DEFAULT_LOG_MAX_BYTES = 5 * 1024 * 1024 # 5 MB
+    DEFAULT_LOG_MAX_BYTES = 5 * 1024 * 1024
     DEFAULT_LOG_BACKUP_COUNT = 3
 
     def __init__(self, cfg: ConfigManager) -> None:
-        """
-        Initializes the SystemManager.
-        """
         if not cfg.is_loaded:
             raise ValueError("ConfigManager must be loaded before initializing SystemManager.")
         self.cfg = cfg
@@ -98,12 +75,7 @@ class SystemManager:
         self.console_handler: Optional[logging.Handler] = None
         logger.debug("SystemManager initialized.")
 
-    # --------------------------------------------------------------------- #
-    # Logging Setup
-    # --------------------------------------------------------------------- #
-
     def _rotate_log_file(self, file_path: Path) -> None:
-        """Rotates a single log file if it exists and is non-empty."""
         try:
             if file_path.exists() and file_path.stat().st_size > 0:
                 timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
@@ -116,29 +88,20 @@ class SystemManager:
             logger.error(f"Unexpected error rotating log file '{file_path}': {e}", exc_info=True)
 
     def setup_logging(self) -> None:
-        """
-        Configures application-wide logging based on settings in ConfigManager.
-        Sets up console (optionally with Rich) and rotating file handlers.
-        """
         logger.info("Setting up application logging...")
         try:
             log_cfg = self.cfg.get("logging", default={})
-
-            # Basic Logger Setup
             root_logger = logging.getLogger()
             if root_logger.hasHandlers():
                  logger.debug("Clearing existing logging handlers.")
                  for handler in root_logger.handlers[:]:
-                      root_logger.removeHandler(handler)
-                      handler.close()
+                      root_logger.removeHandler(handler); handler.close()
             root_logger.setLevel(logging.DEBUG)
 
-            # Standard Formatter (for file logs primarily)
             log_format = log_cfg.get("format", self.DEFAULT_LOG_FORMAT)
             date_format = log_cfg.get("date_format", self.DEFAULT_DATE_FORMAT)
             standard_formatter = logging.Formatter(log_format, datefmt=date_format)
 
-            # Console Handler (Rich or Standard)
             console_level_str = log_cfg.get("console_log_level", self.DEFAULT_CONSOLE_LEVEL).upper()
             console_level = getattr(logging, console_level_str, logging.INFO)
             use_rich = log_cfg.get("rich_console_logging", False)
@@ -147,37 +110,24 @@ class SystemManager:
                 logger.info("Configuring Rich console handler.")
                 rich_keywords = log_cfg.get("rich_keywords", [])
                 if not isinstance(rich_keywords, list) or not rich_keywords: rich_keywords = None
-
                 rich_handler = RichHandler(
-                    level=console_level,
-                    show_time=log_cfg.get("rich_show_time", True),
-                    show_level=log_cfg.get("rich_show_level", True),
-                    show_path=log_cfg.get("rich_show_path", False),
-                    markup=log_cfg.get("rich_markup", True),
-                    rich_tracebacks=log_cfg.get("rich_tracebacks", True),
-                    tracebacks_show_locals=log_cfg.get("rich_tracebacks_show_locals", False),
-                    keywords=rich_keywords,
+                    level=console_level, show_time=log_cfg.get("rich_show_time", True),
+                    show_level=log_cfg.get("rich_show_level", True), show_path=log_cfg.get("rich_show_path", False),
+                    markup=log_cfg.get("rich_markup", True), rich_tracebacks=log_cfg.get("rich_tracebacks", True),
+                    tracebacks_show_locals=log_cfg.get("rich_tracebacks_show_locals", False), keywords=rich_keywords,
                 )
-                root_logger.addHandler(rich_handler)
-                self.console_handler = rich_handler
+                root_logger.addHandler(rich_handler); self.console_handler = rich_handler
                 logger.info(f"Rich console logging enabled at level: {console_level_str}")
-
             elif use_rich and not RICH_AVAILABLE:
-                logger.warning("Rich console logging enabled in config, but 'rich' library not found. Falling back to standard handler.")
-                # Fall through to standard handler setup
-                use_rich = False # Ensure we proceed with standard handler
-
-            if not use_rich: # Standard StreamHandler
+                logger.warning("Rich console logging enabled, but 'rich' not found. Falling back.")
+                use_rich = False
+            if not use_rich:
                 logger.info("Configuring standard console handler.")
                 stream_handler = logging.StreamHandler(sys.stdout)
-                stream_handler.setFormatter(standard_formatter)
-                stream_handler.setLevel(console_level)
-                root_logger.addHandler(stream_handler)
-                self.console_handler = stream_handler
+                stream_handler.setFormatter(standard_formatter); stream_handler.setLevel(console_level)
+                root_logger.addHandler(stream_handler); self.console_handler = stream_handler
                 logger.info(f"Standard console logging enabled at level: {console_level_str}")
 
-
-            # File Logging
             if log_cfg.get("file_logging_enabled", True):
                 log_dir_path_str = log_cfg.get("log_directory", self.DEFAULT_LOG_DIR)
                 log_directory = Path(log_dir_path_str).resolve()
@@ -188,7 +138,6 @@ class SystemManager:
                     logger.error(f"Failed to create log directory '{log_directory}': {e}. File logging disabled.", exc_info=True)
                     return
 
-                # Application Log File Handler
                 app_log_filename = log_cfg.get("application_log_file", self.DEFAULT_APP_LOG_FILE)
                 app_log_path = log_directory / app_log_filename
                 app_log_level_str = log_cfg.get("application_log_level", self.DEFAULT_APP_LOG_LEVEL).upper()
@@ -200,13 +149,11 @@ class SystemManager:
                     app_file_handler = logging.handlers.RotatingFileHandler(
                         app_log_path, maxBytes=max_bytes, backupCount=backup_count, encoding="utf-8"
                     )
-                    app_file_handler.setFormatter(standard_formatter)
-                    app_file_handler.setLevel(app_log_level)
+                    app_file_handler.setFormatter(standard_formatter); app_file_handler.setLevel(app_log_level)
                     root_logger.addHandler(app_file_handler)
-                    logger.info(f"Application file logging configured: '{app_log_path.name}' at level {app_log_level_str}")
-                except Exception as e: logger.error(f"Failed to setup application file logger '{app_log_path}': {e}", exc_info=True)
+                    logger.info(f"Application file logging: '{app_log_path.name}' at level {app_log_level_str}")
+                except Exception as e: logger.error(f"Failed to setup app file logger '{app_log_path}': {e}", exc_info=True)
 
-                # Error Log File Handler
                 error_log_filename = log_cfg.get("error_log_file", self.DEFAULT_ERROR_LOG_FILE)
                 error_log_path = log_directory / error_log_filename
                 error_log_level_str = log_cfg.get("error_log_level", self.DEFAULT_ERROR_LOG_LEVEL).upper()
@@ -216,296 +163,195 @@ class SystemManager:
                     error_file_handler = logging.handlers.RotatingFileHandler(
                         error_log_path, maxBytes=max_bytes, backupCount=backup_count, encoding="utf-8"
                     )
-                    error_file_handler.setFormatter(standard_formatter)
-                    error_file_handler.setLevel(error_log_level)
+                    error_file_handler.setFormatter(standard_formatter); error_file_handler.setLevel(error_log_level)
                     root_logger.addHandler(error_file_handler)
-                    logger.info(f"Error file logging configured: '{error_log_path.name}' at level {error_log_level_str}")
+                    logger.info(f"Error file logging: '{error_log_path.name}' at level {error_log_level_str}")
                 except Exception as e: logger.error(f"Failed to setup error file logger '{error_log_path}': {e}", exc_info=True)
             else:
                 logger.info("File logging is disabled via configuration.")
-
             logger.info("Logging setup complete.")
-
         except Exception as e:
-            logger.exception("An critical error occurred during logging setup.")
+            logger.exception("Critical error during logging setup.") # Use logger.exception for auto exc_info
             raise LoggingSetupError(f"Failed to configure logging: {e}") from e
 
-    # --------------------------------------------------------------------- #
-    # System Information Logging
-    # --------------------------------------------------------------------- #
-
     def log_system_info(self) -> None:
-        """Logs key system and environment details."""
-        if self._sysinfo_logged:
-            logger.debug("System info already logged. Skipping.")
-            return
-
+        if self._sysinfo_logged: logger.debug("System info already logged."); return
         logger.info("----- System Information -----")
         try:
             logger.info(f"OS            : {platform.system()} {platform.release()} ({platform.machine()})")
             logger.info(f"Python        : {platform.python_version()} ({platform.python_implementation()})")
             logger.info(f"Python Path   : {sys.executable}")
-            # Use project_root derived in main.py if available via cfg perhaps, or re-derive
-            project_root_path = Path(__file__).resolve().parents[1] # Assuming modules/system_manager.py
-            logger.info(f"Project Root  : {project_root_path}") # Adjust if structure differs
+            project_root_path = Path(__file__).resolve().parents[1]
+            logger.info(f"Project Root  : {project_root_path}")
             logger.info(f"Config File   : {self.cfg.config_path}")
-
-            # uv Version (Best Effort)
             try:
                 uv_path = shutil.which("uv")
                 if uv_path:
                     result = subprocess.run([uv_path, "--version"], capture_output=True, text=True, timeout=3, check=False, encoding='utf-8')
                     uv_version = result.stdout.strip() if result.returncode == 0 else f"Error ({result.returncode})"
-                else: uv_version = "<'uv' command not found in PATH>"
+                else: uv_version = "<'uv' command not found>"
                 logger.info(f"uv Version    : {uv_version}")
             except Exception as e: logger.info(f"uv Version    : <Error checking: {e}>")
 
-            # PyTorch & CUDA
             logger.info(f"PyTorch       : Version {torch.__version__}" if TORCH_AVAILABLE else "PyTorch       : Not Installed")
             if TORCH_AVAILABLE:
                 try:
                     cuda_available = torch.cuda.is_available()
-                    logger.info(f"  CUDA Status : {'Available' if cuda_available else 'Not Available or Not Setup'}")
+                    logger.info(f"  CUDA Status : {'Available' if cuda_available else 'Not Available'}")
                     if cuda_available:
-                        cuda_version = getattr(torch.version, "cuda", "Unknown")
+                        cuda_version = getattr(torch.version, "cuda", "Unknown") # type: ignore
                         logger.info(f"  CUDA Version: {cuda_version}")
                         device_count = torch.cuda.device_count()
                         logger.info(f"  GPU Count   : {device_count}")
                         for i in range(device_count):
                             try: gpu_name = torch.cuda.get_device_name(i); logger.info(f"  GPU {i}       : {gpu_name}")
-                            except Exception as e: logger.warning(f"  GPU {i}       : <Error getting name: {e}>")
+                            except Exception as e_gpu: logger.warning(f"  GPU {i}       : <Error name: {e_gpu}>")
                     else:
-                         # Check config vs reality
-                         stt_device = self.cfg.get("stt", "device", "cpu").lower()
-                         tts_device = self.cfg.get("tts", "device", "cpu").lower()
-                         if 'cuda' in [stt_device, tts_device]:
-                              logger.warning("  Configuration requests CUDA, but torch.cuda.is_available() is False.")
-                except Exception as e: logger.error(f"  Error checking PyTorch/CUDA details: {e}", exc_info=True)
+                         stt_dev = self.cfg.get("stt", "device", "cpu").lower()
+                         tts_dev = self.cfg.get("tts", "device", "cpu").lower() # Placeholder if TTS had GPU option
+                         if 'cuda' in [stt_dev, tts_dev]: logger.warning("  Config requests CUDA, but torch.cuda.is_available() is False.")
+                except Exception as e_torch: logger.error(f"  Error PyTorch/CUDA details: {e_torch}", exc_info=True)
 
-            # Context Store Path (RAG version)
             context_cfg = self.cfg.get("context_manager", default={})
             storage_path = context_cfg.get("storage_path", ContextManager.DEFAULT_STORAGE_PATH)
             vector_db_path = context_cfg.get("vector_db_path", ContextManager.DEFAULT_VECTOR_DB_PATH)
             logger.info(f"Context Store : History='{storage_path}', Vector DB='{vector_db_path}'")
             logger.info("------------------------------")
             self._sysinfo_logged = True
-
         except Exception as e:
             logger.error(f"Error gathering system information: {e}", exc_info=True)
             logger.info("----- System Information End (incomplete) -----")
 
-
-    # --------------------------------------------------------------------- #
-    # Requirement Verification
-    # --------------------------------------------------------------------- #
-
     def _check_python_version(self) -> bool:
-        """Checks if the current Python version meets the minimum requirement."""
         if sys.version_info < self.MIN_PYTHON_VERSION:
-            logger.critical(f"CRITICAL: Python version {self.MIN_PYTHON_VERSION[0]}.{self.MIN_PYTHON_VERSION[1]}+ required. Found: {platform.python_version()}")
+            logger.critical(f"CRITICAL: Python {self.MIN_PYTHON_VERSION[0]}.{self.MIN_PYTHON_VERSION[1]}+ required. Found: {platform.python_version()}")
             return False
         logger.info(f"✓ Python version check passed ({platform.python_version()})")
-        return True # <<< Added return True
+        return True
 
     def _check_import(self, module_name: str, package_name: Optional[str] = None, purpose: str = "") -> bool:
-        """Checks if a module can be imported."""
         install_name = package_name if package_name else module_name
         purpose_str = f" ({purpose})" if purpose else ""
         try:
             importlib.import_module(module_name)
             logger.info(f"✓ Dependency check passed: {module_name}{purpose_str}")
-            return True # <<< Added return True
+            return True
         except ImportError:
-            logger.critical(f"CRITICAL: Missing required module '{module_name}'{purpose_str}. Install with: uv add {install_name}")
+            logger.critical(f"CRITICAL: Missing module '{module_name}'{purpose_str}. Install: uv add {install_name}")
             return False
         except Exception as e:
-             logger.critical(f"CRITICAL: Error importing module '{module_name}'{purpose_str}: {e}", exc_info=True)
+             logger.critical(f"CRITICAL: Error importing '{module_name}'{purpose_str}: {e}", exc_info=True)
              return False
 
     def _check_cuda_availability(self) -> bool:
-        """Checks if CUDA is available via PyTorch, if configured to be used."""
         stt_device = self.cfg.get("stt", "device", "cpu").lower()
-        tts_device = self.cfg.get("tts", "device", "cpu").lower() # Check TTS too if it might use GPU
-        needs_cuda = 'cuda' in [stt_device, tts_device]
+        # Example: if embedding model could also use GPU via context_manager config
+        embedding_device_cfg = self.cfg.get("context_manager", "embedding_device", "cpu").lower()
+        needs_cuda = 'cuda' in [stt_device, embedding_device_cfg]
 
         if not needs_cuda:
-            logger.info("✓ CUDA check skipped (not configured for use in STT/TTS).")
-            return True # <<< Added return True
-
+            logger.info("✓ CUDA check skipped (not configured for STT/Embeddings).")
+            return True
         if not TORCH_AVAILABLE:
-            logger.critical("CRITICAL: CUDA requested (device='cuda'), but PyTorch is not installed.")
+            logger.critical("CRITICAL: CUDA requested, but PyTorch is not installed.")
             return False
-
         if not torch.cuda.is_available():
-            logger.critical("CRITICAL: CUDA requested (device='cuda'), but torch.cuda.is_available() returned False. Check drivers and PyTorch CUDA build.")
+            logger.critical("CRITICAL: CUDA requested, but torch.cuda.is_available() is False. Check drivers/PyTorch CUDA build.")
             return False
-
         logger.info("✓ CUDA availability check passed.")
-        return True # <<< Added return True
+        return True
 
     def _check_llm_endpoint(self) -> bool:
-        """Checks basic network connectivity to the configured LLM API endpoint."""
         base_url: Optional[str] = self.cfg.get("llm", "api_base_url")
         if not base_url:
-            logger.critical("CRITICAL: LLM endpoint URL (llm.api_base_url) is not configured.")
+            logger.critical("CRITICAL: LLM endpoint URL (llm.api_base_url) not configured.")
             return False
-
         try:
             parsed_url = urlparse(base_url)
-            hostname = parsed_url.hostname
-            port = parsed_url.port
-
-            if not hostname:
-                 logger.critical(f"CRITICAL: Could not parse hostname from LLM endpoint URL: {base_url}")
-                 return False
-
+            hostname, port = parsed_url.hostname, parsed_url.port
+            if not hostname: logger.critical(f"CRITICAL: Cannot parse hostname from LLM URL: {base_url}"); return False
             if port is None: port = 443 if parsed_url.scheme == "https" else 80
-
-            logger.info(f"Checking LLM endpoint connectivity: {hostname}:{port} (from {base_url})")
+            logger.info(f"Checking LLM endpoint: {hostname}:{port} (from {base_url})")
             with socket.create_connection((hostname, port), timeout=5.0):
-                logger.info(f"✓ LLM endpoint check passed: Successfully connected to {hostname}:{port}")
-            return True # <<< Added return True (after successful connection)
-
-        except socket.timeout:
-            logger.critical(f"CRITICAL: Cannot reach LLM endpoint: Connection to {hostname}:{port} timed out.")
-            return False
-        except socket.gaierror as e:
-             logger.critical(f"CRITICAL: Cannot reach LLM endpoint: DNS resolution failed for {hostname}. Error: {e}")
-             return False
-        except OSError as e:
-            logger.critical(f"CRITICAL: Cannot reach LLM endpoint {hostname}:{port}. Error: {e}")
-            return False
-        except Exception as e:
-             logger.critical(f"CRITICAL: Unexpected error checking LLM endpoint {base_url}: {e}", exc_info=True)
-             return False
+                logger.info(f"✓ LLM endpoint check passed: Connected to {hostname}:{port}")
+            return True
+        except socket.timeout: logger.critical(f"CRITICAL: LLM endpoint timeout: {hostname}:{port}."); return False
+        except socket.gaierror as e: logger.critical(f"CRITICAL: LLM endpoint DNS fail: {hostname}. Error: {e}"); return False
+        except OSError as e: logger.critical(f"CRITICAL: LLM endpoint OSError {hostname}:{port}. Error: {e}"); return False
+        except Exception as e: logger.critical(f"CRITICAL: LLM endpoint error {base_url}: {e}", exc_info=True); return False
 
     def _check_context_store_writability(self) -> bool:
-        """Checks if the context storage file path and vector DB path are writable."""
         context_cfg = self.cfg.get("context_manager", default={})
-        storage_path_str = context_cfg.get("storage_path", ContextManager.DEFAULT_STORAGE_PATH)
-        vector_db_path_str = context_cfg.get("vector_db_path", ContextManager.DEFAULT_VECTOR_DB_PATH)
-        store_path = Path(storage_path_str).resolve()
-        vector_path = Path(vector_db_path_str).resolve()
-
-        paths_to_check = {
-            "History": store_path.parent, # Check parent dir for JSON file
-            "Vector DB": vector_path      # Check ChromaDB dir itself
-        }
+        store_path = Path(context_cfg.get("storage_path", ContextManager.DEFAULT_STORAGE_PATH)).resolve()
+        vector_path = Path(context_cfg.get("vector_db_path", ContextManager.DEFAULT_VECTOR_DB_PATH)).resolve()
+        paths_to_check = {"History JSON Dir": store_path.parent, "Vector DB Dir": vector_path}
         all_writable = True
-
         for name, dir_path in paths_to_check.items():
             try:
-                # Check if directory exists, try to create if not
-                if not dir_path.exists():
-                    logger.info(f"Context store ({name}) directory does not exist, attempting to create: {dir_path}")
-                    dir_path.mkdir(parents=True, exist_ok=True)
-                elif not dir_path.is_dir():
-                     logger.critical(f"CRITICAL: Context store ({name}) path exists but is not a directory: {dir_path}")
-                     all_writable = False
-                     continue # Stop checking this path, move to next
-
-                # Check directory write permissions
-                temp_file_path = dir_path / f".writetest_{os.getpid()}_{datetime.now().timestamp()}"
+                if not dir_path.exists(): dir_path.mkdir(parents=True, exist_ok=True); logger.info(f"Created context dir ({name}): {dir_path}")
+                elif not dir_path.is_dir(): logger.critical(f"CRITICAL: Context path ({name}) not a dir: {dir_path}"); all_writable = False; continue
+                temp_file = dir_path / f".writetest_{os.getpid()}_{datetime.now().timestamp()}"
                 try:
-                    with open(temp_file_path, 'w') as f: f.write('test')
-                    temp_file_path.unlink() # Clean up the temporary file
-                    logger.info(f"✓ Context store ({name}) writability check passed (directory: {dir_path})")
-                    # Don't return True here yet, need to check all paths
-                except OSError as e:
-                     logger.critical(f"CRITICAL: Cannot write to context store ({name}) directory '{dir_path}'. Check permissions. Error: {e}")
-                     all_writable = False
-                finally:
-                     temp_file_path.unlink(missing_ok=True) # Ensure cleanup
-
-            except OSError as e:
-                logger.critical(f"CRITICAL: Error accessing or creating context store ({name}) directory '{dir_path}': {e}")
-                all_writable = False
-            except Exception as e:
-                 logger.critical(f"CRITICAL: Unexpected error checking context store ({name}) writability ({dir_path}): {e}", exc_info=True)
-                 all_writable = False
-
-        # Return the overall result after checking all paths
+                    with open(temp_file, 'w') as f: f.write('test')
+                    logger.info(f"✓ Context store ({name}) writability passed (directory: {dir_path})")
+                except OSError as e: logger.critical(f"CRITICAL: Cannot write to context dir ({name}) '{dir_path}': {e}"); all_writable = False
+                finally: temp_file.unlink(missing_ok=True)
+            except Exception as e: logger.critical(f"CRITICAL: Error checking context dir ({name}) '{dir_path}': {e}", exc_info=True); all_writable = False
         return all_writable
 
-    # Helper for Rich dependency check called by verify_requirements
     def _check_rich_dependency(self) -> bool:
-         """Checks Rich install only if enabled in config. Returns True if disabled or installed."""
          if not self.cfg.get("logging", "rich_console_logging", False):
-              logger.info("✓ Dependency check skipped: Rich (Console Logging disabled in config)")
-              return True # Pass if not enabled
-         # If enabled, check import
+              logger.info("✓ Dependency check skipped: Rich (Console Logging disabled)")
+              return True
          if not RICH_AVAILABLE:
-             logger.critical("CRITICAL: Missing optional module 'rich' required for rich_console_logging. Install with: uv add rich")
+             logger.critical("CRITICAL: Missing 'rich' for rich_console_logging. Install: uv add rich")
              return False
          logger.info("✓ Dependency check passed: Rich (for Console Logging)")
-         return True # <<< Added return True
-
-
-    # --- Main Verification Method ---
+         return True
 
     def verify_requirements(self) -> None:
-        """
-        Runs all critical pre-flight checks for the application.
-        """
         logger.info("Running critical requirement checks...")
         failures: List[str] = []
-
-        # Define Checks
-        checks_to_run = [
+        checks_to_run: List[Tuple[Callable[[], bool], str]] = [
             (self._check_python_version, "Python Version"),
-            # Core Dependencies
             (lambda: self._check_import("yaml", "pyyaml", "Config parsing"), "PyYAML"),
             (lambda: self._check_import("customtkinter", purpose="GUI Toolkit"), "CustomTkinter"),
             (lambda: self._check_import("pyaudio", purpose="Audio I/O"), "PyAudio"),
             (lambda: self._check_import("numpy", purpose="Audio/Numeric Processing"), "NumPy"),
             (lambda: self._check_import("soundfile", purpose="Audio File I/O"), "SoundFile"),
-             # Rich (optional, check if enabled via _check_rich_dependency)
             (self._check_rich_dependency, "Rich (Console Logging)"),
-            # AI Modules
             (lambda: self._check_import("faster_whisper", package_name="faster-whisper", purpose="STT"), "Faster Whisper"),
             (lambda: self._check_import("openai", purpose="LLM Client"), "OpenAI Client"),
-            # (lambda: self._check_import("tiktoken", purpose="LLM Tokenizer"), "TikToken"), # Still optional
             (lambda: self._check_import("kokoro", purpose="TTS"), "Kokoro TTS"),
-            # RAG Dependencies
             (lambda: self._check_import("sentence_transformers", purpose="RAG Embeddings"), "Sentence Transformers"),
             (lambda: self._check_import("chromadb", purpose="RAG Vector Store"), "ChromaDB"),
-            # Hardware/Connectivity
+            # Added Transformers and Tokenizers library checks
+            (lambda: self._check_import("transformers", purpose="Universal Tokenizer"), "Transformers Library"),
+            (lambda: self._check_import("tokenizers", purpose="Core Tokenizer for Transformers"), "Tokenizers (HF)"),
             (self._check_cuda_availability, "CUDA Availability"),
             (self._check_llm_endpoint, "LLM Endpoint Connectivity"),
             (self._check_context_store_writability, "Context Store Writability"),
         ]
-
-        # Run Checks
         all_passed = True
         for check_func, check_name in checks_to_run:
             try:
-                # Execute the check function and store the boolean result
                 result = check_func()
-                # Explicitly check if the result is False (covers None or other non-True values implicitly)
-                if result is False: # Check specifically for False
-                    # Avoid adding Rich failure here if it passed because it was disabled
-                    if not (check_name == "Rich (Console Logging)" and self.cfg.get("logging", "rich_console_logging", False) is False):
+                if result is False:
+                    if not (check_name == "Rich (Console Logging)" and not self.cfg.get("logging", "rich_console_logging", False)):
                          failures.append(check_name)
-                         all_passed = False
+                    all_passed = False
                 elif result is not True:
-                     # Log if a check didn't return True or False explicitly (potential bug)
-                     logger.warning(f"Requirement check '{check_name}' returned non-boolean value: {result}. Treating as failure.")
-                     failures.append(f"{check_name} (Bad Return)")
-                     all_passed = False
-
+                     logger.warning(f"Req check '{check_name}' returned non-boolean: {result}. Treating as fail.")
+                     failures.append(f"{check_name} (Bad Return)"); all_passed = False
             except Exception as e:
-                 # Catch unexpected errors within the check function itself
-                 logger.critical(f"CRITICAL: Unexpected error during requirement check '{check_name}': {e}", exc_info=True)
-                 failures.append(f"{check_name} (Error)")
-                 all_passed = False
+                 logger.critical(f"CRITICAL: Error during req check '{check_name}': {e}", exc_info=True)
+                 failures.append(f"{check_name} (Error)"); all_passed = False
 
-
-        # Report Results
         if not all_passed:
             failure_summary = ", ".join(failures)
             logger.critical("*** APPLICATION STARTUP BLOCKED ***")
             logger.critical(f"Failed requirement checks: {failure_summary}")
-            logger.critical("Please resolve the issues listed above and restart the application.")
+            logger.critical("Please resolve the issues and restart.")
             raise RequirementError(f"Failed checks: {failure_summary}")
         else:
             logger.info("✓ All critical system requirement checks passed.")

modules/ui_manager.py

@@ -1,5 +1,5 @@
 # ================================================
-# FILE: modules/ui_manager.py (Textbox Input)
+# FILE: modules/ui_manager.py
 # ================================================
 from __future__ import annotations