backend/app/agents/planner.py

"""
DeepMed-AI — agents/planner.py
PlannerAgent: quyết định dùng RAG retriever hay LLM trực tiếp.
Nhận diện câu hỏi y tế, tên thuốc, và câu hỏi follow-up.
"""

import re

from app.core.logging_config import logger
from app.core.state import AgentState

# ── Từ khoá y tế — Tiếng Việt + Tiếng Anh ─────────────────────────────────────
MEDICAL_KEYWORDS_VI = [
    # Triệu chứng
    "sốt", "đau", "đau đầu", "buồn nôn", "nôn", "tiêu chảy", "ho",
    "mụn", "da", "phát ban", "ngứa", "cảm", "cúm", "khó thở",
    "đau ngực", "đau bụng", "đau lưng", "đau khớp", "đau cơ",
    "mệt mỏi", "yếu", "chóng mặt", "mất trí nhớ", "co giật",
    "tê", "sưng", "chảy máu", "bầm tím", "giảm cân", "tăng cân",
    "chán ăn", "mất ngủ", "khó ngủ", "hoa mắt",
    # Bệnh lý
    "ung thư", "tiểu đường", "huyết áp", "tim mạch", "đột quỵ", "hen suyễn",
    "viêm phổi", "viêm phế quản", "covid", "corona", "nhiễm trùng",
    "viêm", "vi khuẩn", "vi rút", "nấm", "viêm khớp", "loãng xương",
    "tuyến giáp", "thận", "gan", "viêm gan", "trầm cảm",
    "lo âu", "tâm thần", "alzheimer", "parkinson", "động kinh",
    "viêm ruột", "táo bón", "xuất huyết", "nhiễm khuẩn",
    # Điều trị / Thuốc / Tra cứu
    "điều trị", "thuốc", "phác đồ", "liều lượng", "tác dụng phụ",
    "chẩn đoán", "tiên lượng", "phẫu thuật", "thủ thuật",
    "xét nghiệm", "kết quả máu", "siêu âm", "x-quang", "mri",
    "ct scan", "sinh thiết", "tầm soát", "phòng ngừa", "vaccine",
    "tiêm phòng", "hồi phục", "mãn tính", "cấp tính",
    "hội chứng", "rối loạn", "triệu chứng", "bác sĩ", "bệnh viện",
    "y tế", "sức khỏe", "bệnh nhân", "khám", "toa thuốc",
    "kháng sinh", "giảm đau", "hạ sốt", "chống viêm",
    # *** Tra cứu thông tin thuốc (giá, danh mục, hãng sản xuất...) ***
    "giá", "giá thuốc", "đơn giá", "bao nhiêu tiền", "chi phí",
    "danh mục", "danh mục thuốc", "hãng sản xuất", "nhà sản xuất",
    "hoạt chất", "thành phần", "hàm lượng", "nồng độ", "quy cách",
    "chỉ định", "chống chỉ định", "liều dùng", "cách dùng",
    "tương tác", "bảo quản", "hạn dùng", "đường dùng",
    # Bộ phận cơ thể
    "tim", "phổi", "thận", "gan", "não", "dạ dày", "ruột",
    "máu", "xương", "cơ", "thần kinh", "mắt", "tai", "họng",
    "cổ", "cột sống", "khớp", "đầu", "ngực", "bụng", "chân", "tay",
    "răng", "miệng", "lưỡi", "mũi",
]

MEDICAL_KEYWORDS_EN = [
    "fever", "pain", "headache", "nausea", "vomiting", "diarrhea", "cough",
    "acne", "skin", "rash", "itch", "cold", "flu", "shortness of breath",
    "chest pain", "abdominal pain", "back pain", "joint pain", "muscle pain",
    "fatigue", "weakness", "dizziness", "confusion", "memory loss", "seizure",
    "numbness", "tingling", "swelling", "bleeding", "bruising",
    "cancer", "diabetes", "hypertension", "heart disease", "stroke", "asthma",
    "copd", "pneumonia", "bronchitis", "covid", "coronavirus",
    "infection", "virus", "bacteria", "fungal", "arthritis", "osteoporosis",
    "thyroid", "kidney disease", "liver disease", "hepatitis", "depression",
    "anxiety", "bipolar", "schizophrenia", "alzheimer", "parkinson", "epilepsy",
    "treatment", "therapy", "medication", "medicine", "prescription", "dosage",
    "side effects", "diagnosis", "prognosis", "surgery", "operation",
    "procedure", "test", "lab results", "blood test", "x-ray", "mri",
    "ct scan", "ultrasound", "biopsy", "screening", "prevention", "vaccine",
    "immunization", "rehabilitation", "recovery", "chronic", "acute",
    "syndrome", "disorder", "symptom", "cure", "remedy", "doctor", "hospital",
    "price", "cost", "how much",  # drug price lookups
    "heart", "lung", "kidney", "liver", "brain", "stomach", "intestine",
    "blood", "bone", "muscle", "nerve", "eye", "ear", "throat",
    "neck", "spine", "joint", "head", "chest", "abdomen", "leg", "arm",
]

ALL_MEDICAL_KEYWORDS = set(MEDICAL_KEYWORDS_VI + MEDICAL_KEYWORDS_EN)

# Drug suffix patterns for Title-case detection (Ceftriaxone, Amoxicillin, etc.)
_DRUG_SUFFIXES = re.compile(
    r'\b[A-Z][a-z]{3,}(?:in|ol|am|on|id|il|en|ax|im|an|ex|yl|ne|ate|ide|ine|cin|min|lin|one|fen|lol|tan|pin|pam|ban|mab|nib)\b'
)


def _has_drug_name_in_text(text: str) -> bool:
    """Detect drug names in text: ALLCAPS (MIDANTIN) or Title-case with drug suffix (Ceftriaxone)."""
    # ALL-CAPS words ≥ 3 chars: MIDANTIN, ACC, BFS
    if re.search(r'\b[A-Z][A-Z0-9\-]{2,}\b', text):
        return True
    # Title-case with pharmaceutical suffix: Ceftriaxone, Amoxicillin, Metoprolol
    if _DRUG_SUFFIXES.search(text):
        return True
    return False


def _conversation_mentions_drug(state: AgentState) -> bool:
    """Check if recent conversation history mentions a drug name.
    
    This handles follow-up questions like:
    User: "Ceftriaxone 2000 là gì?"
    User: "giá bao nhiêu?"  ← should still route to retriever!
    """
    for item in state.get("conversation_history", [])[-6:]:
        content = item.get("content", "")
        if _has_drug_name_in_text(content):
            return True
    return False


def PlannerAgent(state: AgentState) -> AgentState:
    """Quyết định dùng RAG retriever hay LLM trực tiếp.

    Logic (ưu tiên từ trên xuống):
    1. Câu hỏi chứa tên thuốc (ALLCAPS hoặc Title-case) → retriever
    2. Câu hỏi chứa từ khoá y tế/dược (giá, liều, thuốc...) → retriever
    3. Conversation history gần đây nhắc đến thuốc (follow-up) → retriever
    4. Ngược lại → LLM trực tiếp
    """
    question = state["question"]
    question_lower = question.lower()

    # Check 1: Drug name in current question
    has_drug = _has_drug_name_in_text(question)

    # Check 2: Medical/pharmaceutical keywords
    has_medical_kw = any(kw in question_lower for kw in ALL_MEDICAL_KEYWORDS)

    # Check 3: Follow-up — recent conversation mentioned a drug
    is_drug_followup = (not has_drug and not has_medical_kw
                        and _conversation_mentions_drug(state))

    use_retriever = has_drug or has_medical_kw or is_drug_followup
    state["current_tool"] = "retriever" if use_retriever else "llm_agent"
    state["retry_count"] = 0

    if is_drug_followup:
        logger.info("Planner: Follow-up detected — recent history has drug name → retriever")
    elif has_drug and not has_medical_kw:
        logger.info("Planner: Drug name '%s' detected → retriever", question[:50])

    return state