src/agents.py

"""
src/agents.py — PharmGPTAgent orchestration layer.

Pipeline per user message
─────────────────────────
  1. SafetyTriageTool   → emergency / self-harm → return escalation message
  2. MedicalScopeTool   → out-of-scope          → return redirect message
                          greeting              → return welcome message
  3. OllamaChatClient   → build message list with system prompt + history
                          call chat/completions endpoint
  4. ResponseGuardrailTool → append disclaimer, strip stray image markdown
  5. Return (response_text, status_tag)

status_tag values:
  "ok"                — normal LLM response returned
  "greeting"          — deterministic greeting response
  "blocked_scope"     — out-of-scope query
  "blocked_emergency" — safety escalation
  "error"             — LLM / network error
"""
import logging
from typing import Dict, Iterator, List, Optional, Tuple

from src.llm_client import LLMError, OllamaChatClient
from src.tools import MedicalScopeTool, ResponseGuardrailTool, SafetyTriageTool

logger = logging.getLogger(__name__)

# ─── System prompt ────────────────────────────────────────────────────────────
_SYSTEM_PROMPT = """You are PharmGPT, an expert AI assistant specialising in medicine, pharmaceuticals, and healthcare.

Your responsibilities:
- Answer questions about medications, drug interactions, dosages, contraindications, and side effects.
- Explain diseases, symptoms, diagnostic criteria, and evidence-based treatment options.
- Clarify pharmacological mechanisms and therapeutic classes.
- Support patients, caregivers, students, and healthcare professionals with accurate, up-to-date information.
- Always remind users to consult a licensed healthcare provider for personal medical decisions.

Constraints (strictly enforced):
- Respond only with text. Do NOT generate or reference images.
- Do NOT claim to diagnose, prescribe, or replace professional medical judgment.
- Do NOT speculate beyond established medical knowledge.
- If unsure, say so clearly rather than guessing.
- Keep answers clear, empathetic, and concise unless depth is explicitly requested."""


class PharmGPTAgent:
    """
    Orchestrates the tool pipeline and delegates to the Ollama LLM for
    medical responses.  One instance can be reused across many turns.
    """

    def __init__(self) -> None:
        self._safety = SafetyTriageTool()
        self._scope = MedicalScopeTool()
        self._guardrail = ResponseGuardrailTool()
        self._client: Optional[OllamaChatClient] = None
        self._last_status: str = "ok"

    @property
    def last_status(self) -> str:
        """Status tag of the most recent stream() or run() call."""
        return self._last_status

    def _get_client(self) -> OllamaChatClient:
        """Lazy-initialise the LLM client so config errors surface at runtime."""
        if self._client is None:
            self._client = OllamaChatClient()
        return self._client

    def run(
        self,
        user_message: str,
        conversation_history: Optional[List[Dict[str, str]]] = None,
    ) -> Tuple[str, str]:
        """
        Process one user turn and return (response_text, status_tag).

        Args:
            user_message:          The raw text input from the user.
            conversation_history:  All previous {"role", "content"} turns
                                   (system messages excluded).

        Returns:
            A tuple of (response_text: str, status_tag: str).
        """
        history: List[Dict[str, str]] = list(conversation_history or [])

        # ── 1. Emergency / safety check ───────────────────────────────────────
        safety = self._safety.run(user_message)
        if not safety.allowed:
            logger.warning(
                "SAFETY_TRIAGE blocked | query_snippet=%.60s", user_message
            )
            return safety.message, "blocked_emergency"

        # ── 2. Scope / greeting check ─────────────────────────────────────────
        scope = self._scope.run(user_message)
        if not scope.allowed:
            logger.info(
                "SCOPE_BLOCKED | tag=%s | query_snippet=%.60s",
                scope.tag,
                user_message,
            )
            return scope.message, "blocked_scope"

        if scope.tag == "greeting":
            return scope.message, "greeting"

        # ── 3. LLM call ───────────────────────────────────────────────────────
        messages = history + [{"role": "user", "content": user_message}]
        try:
            client = self._get_client()
            raw = client.chat(messages, system_prompt=_SYSTEM_PROMPT)
        except LLMError as exc:
            logger.error("LLM_ERROR | %s", exc)
            return (
                "⚠️ I'm having trouble reaching the language model right now.\n\n"
                f"**Technical detail:** {exc}\n\n"
                "Please check the model configuration in the Space secrets "
                "and ensure the Ollama endpoint is reachable.",
                "error",
            )

        # ── 4. Guardrail ──────────────────────────────────────────────────────
        final = self._guardrail.run(raw)
        logger.info(
            "LLM_OK | chars=%d | history_turns=%d", len(final), len(history)
        )
        return final, "ok"

    def stream(
        self,
        user_message: str,
        conversation_history: Optional[List[Dict[str, str]]] = None,
    ) -> Iterator[str]:
        """
        Generator version of run().  Yields text chunks for st.write_stream().
        After exhausting the generator, read .last_status for the status tag.
        """
        history: List[Dict[str, str]] = list(conversation_history or [])

        # ── 1. Safety ─────────────────────────────────────────────────────────
        safety = self._safety.run(user_message)
        if not safety.allowed:
            logger.warning("SAFETY_TRIAGE blocked | query_snippet=%.60s", user_message)
            self._last_status = "blocked_emergency"
            yield safety.message
            return

        # ── 2. Scope / greeting ───────────────────────────────────────────────
        scope = self._scope.run(user_message)
        if not scope.allowed:
            logger.info("SCOPE_BLOCKED | tag=%s | query_snippet=%.60s", scope.tag, user_message)
            self._last_status = "blocked_scope"
            yield scope.message
            return

        if scope.tag == "greeting":
            self._last_status = "greeting"
            yield scope.message
            return

        # ── 3. LLM stream ─────────────────────────────────────────────────────
        messages = history + [{"role": "user", "content": user_message}]
        self._last_status = "ok"
        try:
            client = self._get_client()
            for chunk in client.stream_chat(messages, system_prompt=_SYSTEM_PROMPT):
                yield chunk
            # ── 4. Guardrail: append disclaimer after stream ends ─────────────
            yield (
                "\n\n---\n"
                "> **\u2695\ufe0f Medical Disclaimer:** This information is for general "
                "educational purposes only and is **not** a substitute for professional "
                "medical advice, diagnosis, or treatment. Always consult a qualified "
                "healthcare provider for any medical questions or concerns."
            )
            logger.info("LLM_STREAM_OK | history_turns=%d", len(history))
        except LLMError as exc:
            logger.error("LLM_STREAM_ERROR | %s", exc)
            self._last_status = "error"
            yield (
                "\u26a0\ufe0f I'm having trouble reaching the language model right now.\n\n"
                f"**Technical detail:** {exc}\n\n"
                "Please check the model configuration in the Space secrets "
                "and ensure the Ollama endpoint is reachable."
            )