models/oss_assistant.py

"""
oss_assistant.py
~~~~~~~~~~~~~~~~
Production-ready OSS AI assistant backed by Hugging Face Transformers.

Model  : Qwen/Qwen2.5-0.5B-Instruct  (CPU-compatible, ~1 GB download)
Features:
  - Multi-turn conversation via ConversationMemory (sliding-window + token budget)
  - Token-overflow guard: oldest turns evicted before every inference call
  - Configurable temperature, top-p, repetition penalty, max new tokens
  - Lazy model loading (model loaded once, reused across calls)
  - Thread-safe inference via a module-level lock
  - Streamlit-compatible (returns plain strings, no blocking I/O)
  - Full error handling and loguru logging

Usage (standalone):
    from models.oss_assistant import OSSAssistant
    bot = OSSAssistant()
    print(bot.chat("Hello! What can you do?"))

Usage (Streamlit):
    if "oss_bot" not in st.session_state:
        st.session_state.oss_bot = OSSAssistant()
    response = st.session_state.oss_bot.chat(user_input)
"""

from __future__ import annotations

import threading
import time
from dataclasses import dataclass, field
from typing import List, Dict, Optional

import torch
from transformers import (
    AutoTokenizer,
    AutoModelForCausalLM,
    GenerationConfig,
    TextIteratorStreamer,
)

from models.base_assistant import BaseAssistant
from models.memory_manager import ConversationMemory
from dotenv import load_dotenv

from models.logger_config import logger

load_dotenv()

# ── Module-level lock — prevents concurrent inference on a shared model ───────
_INFERENCE_LOCK = threading.Lock()


# ─────────────────────────────────────────────────────────────────────────────
# Configuration dataclass
# ─────────────────────────────────────────────────────────────────────────────

@dataclass
class AssistantConfig:
    """
    All tunable parameters for the OSS assistant.
    Changing any field on an already-instantiated OSSAssistant takes effect
    on the very next call to chat().
    """

    # Model identity
    model_id: str = "Qwen/Qwen2.5-0.5B-Instruct"

    # Generation knobs
    max_new_tokens: int = 512        # tokens the model may produce per turn
    temperature: float = 0.7         # higher → more creative; lower → more factual
    top_p: float = 0.9               # nucleus sampling threshold
    repetition_penalty: float = 1.1  # penalise repeated tokens

    # Context / memory management
    max_history_turns: int = 10      # sliding window: keep last N human/AI pairs
    max_context_tokens: int = 1_800  # token budget for system + history before inference
    system_prompt: str = (
        "You are a helpful, harmless, and honest AI assistant. "
        "Answer concisely and accurately."
    )

    # Runtime
    device: Optional[str] = None     # None → auto-detect (cuda > mps > cpu)
    torch_dtype: Optional[str] = None  # None → float32 on CPU, float16 on GPU


# ─────────────────────────────────────────────────────────────────────────────
# Main assistant class
# ─────────────────────────────────────────────────────────────────────────────

class OSSAssistant(BaseAssistant):
    """
    Multi-turn conversational assistant using a locally-run HuggingFace model.
    Implements the BaseAssistant interface — interchangeable with GroqAssistant.

    The class is responsible for:
      1. Loading the tokenizer and model (once, lazily).
      2. Delegating memory/history management to BaseAssistant / ConversationMemory.
      3. Running inference with configurable sampling parameters.
      4. Returning the plain assistant reply as a string.
    """

    def __init__(self, config: Optional[AssistantConfig] = None) -> None:
        self.config = config or AssistantConfig()

        # ConversationMemory owns all history, windowing, and token-budget logic.
        self._memory = ConversationMemory(
            max_turns=self.config.max_history_turns,
            max_context_tokens=self.config.max_context_tokens,
            system_prompt=self.config.system_prompt,
        )

        # Lazy-loaded — set by _ensure_model_loaded()
        self._tokenizer: Optional[AutoTokenizer] = None
        self._model: Optional[AutoModelForCausalLM] = None
        self._device: Optional[str] = None

        logger.info(f"OSSAssistant initialised | model={self.config.model_id}")

    # ── Abstract interface implementation ─────────────────────────────────────

    @property
    def backend_name(self) -> str:
        """Model identifier used in get_info() and UI labels."""
        return self.config.model_id

    def chat(self, user_message: str) -> str:
        """
        Send a user message and return the assistant's reply.

        Flow
        ----
        1. Validate input.
        2. Load model (once, lazily).
        3. Register user message in memory (opens a new Turn).
        4. _run_inference() enforces token budget before tokenising.
        5. On success: save assistant reply; on failure: rollback the open Turn.
        """
        if not user_message or not user_message.strip():
            return "Please enter a message."

        self._sync_memory_config()
        self._ensure_model_loaded()
        self._memory.add_user_message(user_message)

        try:
            reply = self._run_inference()
        except Exception as exc:
            logger.error(f"Inference failed: {exc}")
            self._memory.rollback_last_user_message()
            return f"[Error during inference: {exc}]"

        self._memory.add_assistant_reply(reply)
        logger.debug(f"Turn complete | user={user_message[:60]!r} | reply={reply[:60]!r}")
        return reply

    def stream(self, user_message: str) -> Iterator[str]:
        """
        Yield the assistant's reply token-by-token using TextIteratorStreamer.

        A background daemon thread runs model.generate() and feeds tokens into
        a TextIteratorStreamer queue. The main thread yields from the queue,
        making this compatible with Streamlit's st.write_stream().

        Memory is updated after the full reply is assembled (atomic save).
        On error the open turn is rolled back.
        """
        import threading

        user_message = (user_message or "").strip()
        if not user_message:
            yield "Please enter a message."
            return

        self._sync_memory_config()
        self._ensure_model_loaded()
        self._memory.add_user_message(user_message)

        try:
            with _INFERENCE_LOCK:
                # Enforce token budget before building the prompt
                self._memory.enforce_token_budget(self._tokenizer)

                messages = self._memory.as_prompt_messages()
                prompt_text = self._tokenizer.apply_chat_template(
                    messages, tokenize=False, add_generation_prompt=True
                )
                inputs = self._tokenizer(
                    prompt_text,
                    return_tensors="pt",
                    truncation=True,
                    max_length=self.config.max_context_tokens,
                ).to(self._model.device)

                input_len = inputs["input_ids"].shape[1]

                # TextIteratorStreamer lets us yield tokens from the main thread
                # while generate() runs in a background thread.
                streamer = TextIteratorStreamer(
                    self._tokenizer,
                    skip_prompt=True,        # don't echo the prompt back
                    skip_special_tokens=True,
                )

                gen_config = GenerationConfig(
                    max_new_tokens=self.config.max_new_tokens,
                    do_sample=self.config.temperature > 0,
                    temperature=self.config.temperature if self.config.temperature > 0 else None,
                    top_p=self.config.top_p,
                    repetition_penalty=self.config.repetition_penalty,
                    pad_token_id=self._tokenizer.pad_token_id,
                    eos_token_id=self._tokenizer.eos_token_id,
                )

                # Launch generation in a daemon thread so it doesn't block
                gen_thread = threading.Thread(
                    target=self._model.generate,
                    kwargs={**inputs, "generation_config": gen_config, "streamer": streamer},
                    daemon=True,
                )
                gen_thread.start()

                # Collect and yield tokens as they arrive
                reply_parts: list[str] = []
                for token_text in streamer:
                    reply_parts.append(token_text)
                    yield token_text

                gen_thread.join()

        except Exception as exc:
            logger.error(f"Stream inference failed: {exc}")
            self._memory.rollback_last_user_message()
            yield f"[Error during inference: {exc}]"
            return

        full_reply = "".join(reply_parts).strip()
        if full_reply:
            self._memory.add_assistant_reply(full_reply)
        else:
            self._memory.rollback_last_user_message()
            logger.warning("OSS model produced an empty stream reply; turn rolled back.")

    @property
    def is_loaded(self) -> bool:
        """True once the model has been loaded into memory."""
        return self._model is not None

    # ── Model loading ─────────────────────────────────────────────────────────

    def _ensure_model_loaded(self) -> None:
        """Load tokenizer + model exactly once. Thread-safe via the module lock."""
        if self._model is not None:
            return  # already loaded

        with _INFERENCE_LOCK:
            # Double-check inside the lock (another thread may have loaded first)
            if self._model is not None:
                return

            logger.info(f"Loading model: {self.config.model_id}")
            t0 = time.perf_counter()

            # ── Resolve device ─────────────────────────────────────────────
            if self.config.device:
                self._device = self.config.device
            elif torch.cuda.is_available():
                self._device = "cuda"
            elif hasattr(torch.backends, "mps") and torch.backends.mps.is_available():
                self._device = "mps"
            else:
                self._device = "cpu"

            logger.info(f"Inference device: {self._device.upper()}")

            # ── Resolve dtype ──────────────────────────────────────────────
            if self.config.torch_dtype:
                dtype = getattr(torch, self.config.torch_dtype)
            elif self._device == "cpu":
                dtype = torch.float32   # float16 is slower on CPU
            else:
                dtype = torch.float16   # saves VRAM on GPU

            # ── Load tokenizer ─────────────────────────────────────────────
            self._tokenizer = AutoTokenizer.from_pretrained(
                self.config.model_id,
                trust_remote_code=True,  # required for Qwen models
            )
            # Ensure pad token is set (Qwen may not set it by default)
            if self._tokenizer.pad_token is None:
                self._tokenizer.pad_token = self._tokenizer.eos_token

            # ── Load model ─────────────────────────────────────────────────
            self._model = AutoModelForCausalLM.from_pretrained(
                self.config.model_id,
                torch_dtype=dtype,
                device_map="auto",      # auto-places layers on GPU/CPU
                low_cpu_mem_usage=True, # stream weights instead of double-buffering
                trust_remote_code=True,
            )
            self._model.eval()  # disable dropout for deterministic inference

            elapsed = time.perf_counter() - t0
            logger.success(f"Model loaded in {elapsed:.1f}s on {self._device.upper()}")

    # ── Config sync ───────────────────────────────────────────────────────────

    def _sync_memory_config(self) -> None:
        """
        Push live AssistantConfig values into the ConversationMemory object.
        Called at the top of chat() so Streamlit sidebar changes take effect
        immediately without requiring a session restart.
        """
        self._memory.max_turns = self.config.max_history_turns
        self._memory.max_context_tokens = self.config.max_context_tokens
        self._memory.system_prompt = self.config.system_prompt

    # ── Inference ─────────────────────────────────────────────────────────────

    def _run_inference(self) -> str:
        """
        Enforce token budget, tokenize, run generation, decode reply.

        Steps
        -----
        1. enforce_token_budget() — evicts oldest complete turns until the
           rendered prompt fits within max_context_tokens.
        2. apply_chat_template() — renders history to a single string using
           the model's native template (handles all special tokens).
        3. Tokenize with explicit truncation as a last-resort safety net.
        4. Generate with per-call GenerationConfig (picks up sidebar changes).
        5. Slice prompt tokens off the output; decode only new tokens.

        Uses the module-level lock — only one thread runs inference at a time.
        """
        with _INFERENCE_LOCK:
            # ── Step 1: token-budget enforcement ─────────────────────────────
            # Evicts oldest complete turns so the prompt fits in context.
            final_token_count = self._memory.enforce_token_budget(self._tokenizer)
            logger.debug(f"Prompt tokens after budget enforcement: {final_token_count}")

            # ── Step 2: render prompt via chat template ───────────────────────
            messages = self._memory.as_prompt_messages()
            prompt_text = self._tokenizer.apply_chat_template(
                messages,
                tokenize=False,
                add_generation_prompt=True,  # appends the "<|im_start|>assistant" trigger
            )

            # ── Step 3: tokenize (truncation = last-resort safety net) ────────
            inputs = self._tokenizer(
                prompt_text,
                return_tensors="pt",
                padding=True,
                truncation=True,
                max_length=self.config.max_context_tokens,
            ).to(self._model.device)

            input_len = inputs["input_ids"].shape[1]
            logger.debug(f"Tokenized prompt length: {input_len} tokens")

            # ── Step 4: build GenerationConfig per-call ───────────────────────
            # Rebuilding each call means Streamlit slider changes take effect
            # immediately without restarting the session.
            gen_config = GenerationConfig(
                max_new_tokens=self.config.max_new_tokens,
                do_sample=self.config.temperature > 0,
                temperature=self.config.temperature if self.config.temperature > 0 else None,
                top_p=self.config.top_p,
                repetition_penalty=self.config.repetition_penalty,
                pad_token_id=self._tokenizer.pad_token_id,
                eos_token_id=self._tokenizer.eos_token_id,
            )

            # ── Step 5: generate + decode ─────────────────────────────────────
            with torch.inference_mode():  # disables grad tracking — faster + less RAM
                output_ids = self._model.generate(
                    **inputs,
                    generation_config=gen_config,
                )

            # Keep only the newly generated token ids (strip the prompt prefix)
            new_token_ids = output_ids[0][input_len:]
            reply = self._tokenizer.decode(new_token_ids, skip_special_tokens=True)
            logger.debug(f"Generated {len(new_token_ids)} new token(s)")

        return reply.strip()

    # ── Dunder helpers ────────────────────────────────────────────────────────

    def __repr__(self) -> str:
        status = "loaded" if self.is_loaded else "not loaded"
        return (
            f"OSSAssistant(model={self.config.model_id!r}, "
            f"device={self._device!r}, status={status}, turns={self.turn_count})"
        )