src/tools.py

"""
src/tools.py — Lightweight tool registry for the PharmGPT agent pipeline.

Tools are deterministic (no LLM calls) and act as guard layers around the
main chat completion call.  Each tool's .run() method returns a ToolResult.

Pipeline order enforced by PharmGPTAgent:
  1. SafetyTriageTool   — detect emergencies and self-harm intent first
  2. MedicalScopeTool   — confirm query is health / pharma / greeting related
  3. (LLM call)
  4. ResponseGuardrailTool — append disclaimer, strip stray image markdown
"""
import re
from dataclasses import dataclass, field


# ─── Result container ─────────────────────────────────────────────────────────

@dataclass
class ToolResult:
    allowed: bool
    message: str = ""
    tag: str = ""       # "greeting" | "medical" | "emergency" | "out_of_scope"


# ─── Keyword sets ─────────────────────────────────────────────────────────────

_GREETING_PHRASES: frozenset[str] = frozenset({
    "hello", "hi", "hey", "good morning", "good evening", "good afternoon",
    "how are you", "how do you do", "greetings", "what's up", "howdy",
    "good day", "nice to meet you",
})

_MEDICAL_KEYWORDS: frozenset[str] = frozenset({
    # Substances
    "medicine", "medication", "drug", "pill", "tablet", "capsule", "dose",
    "dosage", "prescription", "pharmacy", "pharmacist", "pharmaceutical",
    "therapeutic", "antibiotic", "vaccine", "supplement", "vitamin",
    "ibuprofen", "paracetamol", "aspirin", "insulin", "penicillin",
    "antihistamine", "analgesic", "antifungal", "anticoagulant",
    "antihypertensive", "antidepressant", "antiviral", "chemotherapy",
    "immunotherapy",
    # Conditions / processes
    "disease", "disorder", "syndrome", "symptom", "diagnosis", "treatment",
    "therapy", "infection", "inflammation", "surgery", "overdose",
    "side effect", "contraindication", "interaction", "clinical",
    "virus", "bacteria", "cancer", "tumor", "diabetes", "hypertension",
    "anxiety", "depression", "pain", "fever", "allergy",
    # Body / systems
    "blood", "heart", "lung", "liver", "kidney", "brain", "anatomy",
    "physiology", "pathology", "pharmacology", "toxicology",
    # People / places
    "doctor", "physician", "nurse", "surgeon", "patient", "caregiver",
    "hospital", "clinic", "health", "wellness", "medical",
    # Specialties
    "cardiology", "neurology", "oncology", "orthopedics", "pediatrics",
    "gynecology", "dermatology", "psychiatry", "radiology", "immunology",
    "endocrinology", "mental health",
    # Actions
    "prescribe", "diagnose", "treat", "vaccinate", "medicate", "operate",
})

_EMERGENCY_KEYWORDS: frozenset[str] = frozenset({
    "suicide", "suicidal", "kill myself", "end my life", "want to die",
    "take my own life", "self harm", "self-harm", "cut myself",
    "overdose", "chest pain", "heart attack", "stroke",
    "can't breathe", "cannot breathe", "difficulty breathing",
    "not breathing", "unconscious", "severe bleeding", "poisoning",
    "emergency", "dying right now",
})

_MEDICAL_DISCLAIMER: str = (
    "\n\n---\n"
    "> **⚕️ 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."
)

_GREETING_RESPONSE: str = (
    "Hello! 👋 I'm **PharmGPT**, your AI-powered pharmaceutical and medical assistant.\n\n"
    "I can help you with:\n"
    "- 💊 Medications, dosages, and drug interactions\n"
    "- 🏥 Diseases, symptoms, and treatment options\n"
    "- 🧬 Pharmacology and drug mechanisms\n"
    "- ❤️ General health and wellness information\n\n"
    "What can I help you with today?"
)

_EMERGENCY_RESPONSE: str = (
    "🚨 **This sounds like an emergency.**\n\n"
    "If you or someone nearby is in immediate danger or experiencing a medical "
    "emergency, please **call emergency services right away**:\n\n"
    "| Region | Number |\n"
    "|--------|--------|\n"
    "| India  | 112 / 102 (ambulance) |\n"
    "| USA    | 911 |\n"
    "| Europe | 112 |\n"
    "| UK     | 999 |\n\n"
    "**Mental health crisis lines (India):**\n"
    "- iCall: **9152987821**\n"
    "- NIMHANS: **080-46110007**\n"
    "- Vandrevala Foundation: **1860-2662-345** (24×7)\n\n"
    "You are not alone. Please reach out — help is available."
)

_OUT_OF_SCOPE_RESPONSE: str = (
    "I specialise in **health, medicine, and pharmaceutical** topics.\n\n"
    "Your question doesn't appear to fall within that scope. Feel free to ask me "
    "about medications, symptoms, diseases, pharmacology, drug interactions, or "
    "general health and wellness!"
)


# ─── Tool classes ─────────────────────────────────────────────────────────────

class SafetyTriageTool:
    """
    First-pass safety check.
    Detects emergency keywords and self-harm intent before any LLM call.
    """

    def run(self, query: str) -> ToolResult:
        q = query.lower()
        if any(kw in q for kw in _EMERGENCY_KEYWORDS):
            return ToolResult(
                allowed=False,
                message=_EMERGENCY_RESPONSE,
                tag="emergency",
            )
        return ToolResult(allowed=True)


class MedicalScopeTool:
    """
    Determines whether a query is greeting / medical / out-of-scope.
    Returns a ToolResult with tag set accordingly.
    """

    def run(self, query: str) -> ToolResult:
        q = query.lower().strip()

        # Exact greeting match
        if q in _GREETING_PHRASES or any(q.startswith(g) for g in _GREETING_PHRASES):
            return ToolResult(allowed=True, message=_GREETING_RESPONSE, tag="greeting")

        # Medical keyword presence
        if any(kw in q for kw in _MEDICAL_KEYWORDS):
            return ToolResult(allowed=True, tag="medical")

        return ToolResult(
            allowed=False,
            message=_OUT_OF_SCOPE_RESPONSE,
            tag="out_of_scope",
        )


class ResponseGuardrailTool:
    """
    Post-processing guardrail applied to LLM output.
    - Strips accidental image markdown (e.g. ![...](...))
    - Appends a medical disclaimer
    """

    def run(self, response: str) -> str:
        # Remove any image markdown that slipped through
        cleaned = re.sub(r"!\[.*?\]\(.*?\)", "", response)
        return cleaned.strip() + _MEDICAL_DISCLAIMER