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