models/base_assistant.py

"""
base_assistant.py
~~~~~~~~~~~~~~~~~
Abstract base class defining the unified interface for all AI assistants
in this project.

Inheritance hierarchy
---------------------
    BaseAssistant  (this file)
    ├── OSSAssistant   — local HuggingFace model (Qwen2.5-0.5B-Instruct)
    └── GroqAssistant  — cloud API (Llama 3 via Groq)

Design principles
-----------------
1. **Single interface** — callers (Streamlit app, evaluation suite, scripts)
   depend only on BaseAssistant. Swapping backends requires no call-site changes.

2. **Concrete shared code lives here** — reset(), history, turn_count, memory,
   get_info(), chat_batch() are identical for every backend; defining them once
   avoids drift.

3. **Minimal abstract surface** — subclasses must implement only three things:
       chat(msg)        → str
       stream(msg)      → Iterator[str]
       backend_name     → str   (abstract property)

4. **Safety built-in** — a SafetyGuard is attached at this layer. Callers
   should prefer safe_chat() / safe_stream() which run input + output checks
   automatically. The underlying chat() / stream() remain available for
   pipelines that handle safety separately.

5. **Evaluation-ready** — get_info() returns a structured dict; chat_batch()
   runs a list of prompts with safety checks and fresh memory per prompt.

Example usage
-------------
    # Swap backends with zero call-site change
    from models.oss_assistant  import OSSAssistant
    from models.groq_assistant import GroqAssistant
    from models.base_assistant import BaseAssistant

    def run(bot: BaseAssistant, prompt: str) -> str:
        return bot.chat(prompt)

    oss_reply  = run(OSSAssistant(),  "What is entropy?")
    groq_reply = run(GroqAssistant(), "What is entropy?")

    # Evaluation batch (fresh session, no carry-over)
    results = bot.chat_batch(["Q1", "Q2", "Q3"])
"""

from __future__ import annotations

from abc import ABC, abstractmethod
from typing import Dict, Iterator, List, Optional

from models.memory_manager import ConversationMemory
from models.safety_guard import SafetyGuard, SafetyConfig, SafetyResult
from models.logger_config import logger


# ─────────────────────────────────────────────────────────────────────────────
# Base class
# ─────────────────────────────────────────────────────────────────────────────

class BaseAssistant(ABC):
    """
    Unified interface for all conversational AI backends.

    Subclasses MUST:
    ----------------
    1. Set ``self._memory`` (a ConversationMemory instance) in ``__init__``
       before calling any base-class method.
    2. Implement ``chat(user_message) -> str``.
    3. Implement ``stream(user_message) -> Iterator[str]``.
    4. Implement the ``backend_name`` abstract property.

    Subclasses MUST NOT:
    --------------------
    - Redefine ``reset()``, ``history``, ``turn_count``, or ``memory`` —
      these are provided here and delegate to ``self._memory``.
    - Redefine ``get_info()`` or ``chat_batch()`` unless specialised behaviour
      is needed.
    """

    # Every subclass __init__ must assign this before any method is called.
    _memory: ConversationMemory

    # SafetyGuard instance — created with default config if not supplied.
    # Override by passing safety_config= to the subclass constructor,
    # or replace self._guard at any time.
    _guard: SafetyGuard

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

    @abstractmethod
    def chat(self, user_message: str) -> str:
        """
        Send a message and return the complete assistant reply as a string.

        Must:
        - Validate the input (return a safe string on empty input, never raise).
        - Add the user message to ``self._memory`` before calling the backend.
        - Save the assistant reply to ``self._memory`` on success.
        - Call ``self._memory.rollback_last_user_message()`` on failure so
          memory stays consistent.

        Args:
            user_message: Raw text from the user.

        Returns:
            The assistant's reply as a plain string. Error messages are
            returned as strings (prefixed with ``[Error: …]``), never raised.
        """

    @abstractmethod
    def stream(self, user_message: str) -> Iterator[str]:
        """
        Send a message and yield the reply incrementally as text chunks.

        Designed for ``st.write_stream()`` in Streamlit:
            full_reply = st.write_stream(bot.stream(user_message))

        Contract:
        - Must yield at least one string (possibly the full reply as one chunk
          if the backend does not support true streaming).
        - Must update ``self._memory`` atomically: only after the full reply
          has been accumulated, or rollback on error.
        - Must never raise — yield an error string instead.

        Args:
            user_message: Raw text from the user.

        Yields:
            str: Incremental text chunks of the assistant's reply.
        """

    @property
    @abstractmethod
    def backend_name(self) -> str:
        """
        Short human-readable identifier for this backend.

        Examples: ``"Qwen2.5-0.5B-Instruct"``, ``"llama-3.3-70b-versatile"``.
        Used in ``get_info()`` dict, evaluation reports, and UI labels.
        """

    # ── Concrete shared methods ───────────────────────────────────────────────
    # These work identically for every backend because they all delegate to
    # self._memory, which every subclass is required to set.

    def reset(self) -> None:
        """
        Clear all conversation history.
        The backend (model / API client) stays initialised.
        """
        self._memory.clear()
        logger.info(f"{self.__class__.__name__}: conversation reset.")

    @property
    def guard(self) -> SafetyGuard:
        """
        Access the attached SafetyGuard.
        Replace with a custom instance to change rules at runtime:
            bot.guard = SafetyGuard(SafetyConfig(min_block_severity="high"))
        """
        # Lazily create a default guard if the subclass didn't set one.
        if not hasattr(self, "_guard") or self._guard is None:
            self._guard = SafetyGuard()
        return self._guard

    @guard.setter
    def guard(self, value: SafetyGuard) -> None:
        self._guard = value
        logger.info(f"{self.__class__.__name__}: SafetyGuard replaced | {repr(value)}")

    # ── Safety-wrapped public methods ─────────────────────────────────────────
    # Prefer these over chat() / stream() in production.

    def safe_chat(self, user_message: str) -> str:
        """
        Run input safety check, call chat(), then run output safety check.

        If the input is blocked, returns the guard's safe fallback immediately
        (model is never called). If the output is blocked, returns the output
        guard's fallback instead of the raw reply.

        Args:
            user_message: Raw text from the user.

        Returns:
            Safe assistant reply (or a refusal message).
        """
        # ── Input check ───────────────────────────────────────────────────────
        in_result = self.guard.check_input(user_message)
        if in_result.blocked:
            return in_result.safe_response

        # ── Model call ────────────────────────────────────────────────────────
        reply = self.chat(user_message)

        # ── Output check ──────────────────────────────────────────────────────
        out_result = self.guard.check_output(reply)
        if out_result.blocked:
            # Roll back memory so the blocked exchange isn't retained
            self._memory.rollback_last_user_message()
            return out_result.safe_response

        return reply

    def safe_stream(self, user_message: str) -> Iterator[str]:
        """
        Input-check, then yield chunks from stream(), then output-check the
        accumulated reply. If the final reply is blocked, the already-yielded
        chunks cannot be un-sent — instead a [RESPONSE REMOVED] notice is
        appended so the caller knows the content was moderated.

        Args:
            user_message: Raw text from the user.

        Yields:
            str: Text chunks, followed by a moderation notice if blocked.
        """
        # ── Input check ───────────────────────────────────────────────────────
        in_result = self.guard.check_input(user_message)
        if in_result.blocked:
            yield in_result.safe_response
            return

        # ── Stream from model ─────────────────────────────────────────────────
        accumulated: List[str] = []
        for chunk in self.stream(user_message):
            accumulated.append(chunk)
            yield chunk

        # ── Output check on assembled reply ───────────────────────────────────
        full_reply = "".join(accumulated)
        out_result = self.guard.check_output(full_reply)
        if out_result.blocked:
            self._memory.rollback_last_user_message()
            yield (
                f"\n\n---\n"
                f"*[Content moderated — {out_result.rule_name}. "
                f"{out_result.safe_response}]*"
            )

    @property
    def history(self) -> List[Dict[str, str]]:
        """
        Read-only view of the current conversation history.

        Returns a list of ``{"role": ..., "content": ...}`` dicts compatible
        with the OpenAI / HuggingFace chat-template format.
        """
        return self._memory.as_message_list()

    @property
    def turn_count(self) -> int:
        """Number of *complete* human/AI turns in the current session."""
        return self._memory.complete_turn_count

    @property
    def memory(self) -> ConversationMemory:
        """Direct (read-intended) access to the ConversationMemory object."""
        return self._memory

    # ── Evaluation helpers ────────────────────────────────────────────────────

    def get_info(self) -> Dict[str, object]:
        """
        Return a structured snapshot of this assistant's current state.
        """
        g = self.guard
        return {
            "backend":           self.__class__.__name__,
            "backend_name":      self.backend_name,
            "turn_count":        self.turn_count,
            "memory_turns":      self._memory.turn_count,
            "system_prompt":     self._memory.system_prompt[:80],
            "safety_enabled":    True,
            "safety_rules":      sum(len(v) for v in g._compiled.values()),
            "safety_min_sev":    g.config.min_block_severity,
        }

    def chat_batch(
        self,
        inputs: List[str],
        *,
        reset_between: bool = True,
    ) -> List[Dict[str, str]]:
        """
        Run a list of prompts and return structured results.

        Designed for evaluation pipelines (DeepEval, custom scorers) where
        you need input/output pairs without Streamlit overhead.

        Parameters
        ----------
        inputs : list of str
            Prompts to send, in order.
        reset_between : bool, default True
            If True, conversation history is cleared before each prompt so
            each input is evaluated in isolation (standard for benchmarking).
            If False, all prompts share a single growing conversation.

        Returns
        -------
        list of dicts, one per input:
            {
                "input":  str,   — the original prompt
                "output": str,   — the assistant's reply
                "turn":   int,   — turn number (1-indexed)
                "error":  bool,  — True if output starts with "[Error"
            }
        """
        results: List[Dict[str, str]] = []

        for i, prompt in enumerate(inputs, 1):
            if reset_between:
                self.reset()

            logger.debug(
                f"chat_batch [{i}/{len(inputs)}] | "
                f"reset_between={reset_between} | "
                f"prompt={prompt[:60]!r}"
            )

            try:
                # Use safe_chat so every batch call is also safety-checked
                output = self.safe_chat(prompt)
            except Exception as exc:
                output = f"[Error: {exc}]"
                logger.error(f"chat_batch: unexpected exception on turn {i}: {exc}")

            is_blocked = output == self.guard.config.refusal_harmful \
                or output == self.guard.config.refusal_jailbreak \
                or output == self.guard.config.refusal_injection

            results.append({
                "input":   prompt,
                "output":  output,
                "turn":    i,
                "error":   output.startswith("[Error") or output.startswith("[Rate"),
                "blocked": is_blocked,
            })

        return results

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

    def __repr__(self) -> str:
        return (
            f"{self.__class__.__name__}("
            f"backend={self.backend_name!r}, "
            f"turns={self.turn_count})"
        )