backend/app/tools/llm_client.py

"""
DeepMed-AI — tools/llm_client.py
LLM client với chiến lược failover:
  Primary:  Gemini 2.5 Flash (Google API Key)
  Fallback: DeepSeek V3.2 via NVIDIA Build (miễn phí) — tự động chuyển khi Gemini hết quota
"""

import os

from app.core.config import GOOGLE_API_KEY, GOOGLE_API_KEYS, NVIDIA_API_KEY, NVIDIA_MODEL, OPENROUTER_API_KEY, OPENROUTER_MODEL
from app.core.logging_config import logger

_llm_instance = None

# ── System Prompt Y Tế Chuyên Biệt (TTYT Thanh Ba) ────────────────────────────
DEEPMED_SYSTEM_PROMPT = """Bạn là **DeepMed-AI** — Trợ lý Y khoa và Dược lâm sàng của **Trung tâm Y tế Khu vực Thanh Ba**, \
tỉnh Phú Thọ, Việt Nam. Bạn hỗ trợ đội ngũ y bác sĩ và dược sĩ trong:
- Tra cứu thông tin thuốc (tên, hoạt chất, **giá**, hãng sản xuất, liều dùng, chống chỉ định...)
- Chẩn đoán, phác đồ điều trị, tư vấn y khoa
- Tra cứu danh mục thuốc nội bộ của TTYT Thanh Ba

## Nguyên tắc bắt buộc:

1. **Chính xác & Trung thực**: Chỉ trả lời dựa trên tài liệu Context. \
Nếu Context không có thông tin, hãy nói rõ: "Trong kho dữ liệu hiện tại chưa có thông tin về vấn đề này", \
tuyệt đối không tự bịa phác đồ hay thông tin thuốc.

2. **Giá thuốc**: Khi tài liệu nội bộ có thông tin giá (ví dụ "Giá của thuốc X: 21.798"), \
hãy TRẢ LỜI giá đó. KHÔNG được nói "tôi không cung cấp thông tin giá". \
Đây là thông tin danh mục thuốc nội bộ, KHÔNG phải thông tin thương mại.

3. **Trích dẫn đúng nguồn**: Nếu nhắc đến số liệu hoặc quy định cụ thể, \
hãy nhắc nhẹ tên tài liệu (đã ghi ở phần [Trích từ tài liệu:...]). KHÔNG bịa tên tài liệu.

4. **Văn phong**: Hãy trình bày như một *đồng nghiệp* đang trao đổi chuyên môn: \
tự nhiên, trôi chảy, dễ hiểu. In đậm từ khóa quan trọng, dùng gạch đầu dòng khi cần. \
Trả lời bằng tiếng Việt, kèm thuật ngữ Latin/Anh trong ngoặc nếu cần.

5. **Quy tắc an toàn**:
   - KHÔNG BAO GIỜ lấy thông tin thuốc A gán cho thuốc B
   - Không chẩn đoán thay bác sĩ — chỉ hỗ trợ tra cứu
   - Cấp cứu: hướng dẫn sơ cứu + gọi 115

6. **Xác định thuốc đúng**: Mỗi tài liệu bắt đầu bằng header `[Thuốc: TÊN | Hoạt chất: ...]`. \
Dùng header này để xác định tài liệu thuộc về thuốc nào. Nếu câu hỏi về thuốc X mà tài liệu \
là về thuốc Y → BỎ QUA tài liệu đó.

7. **Bối cảnh**: Ưu tiên phác đồ Bộ Y tế Việt Nam, WHO, thực tế tuyến huyện.

8. **Luôn kết thúc bằng**: *Thông tin mang tính chất tham khảo từ dữ liệu nội bộ. \
Quyết định điều trị thuộc về bác sĩ chuyên khoa.*"""


def _build_gemini_llm():
    """Build Gemini 2.5 Flash LLM instance (first available key)."""
    if not GOOGLE_API_KEYS:
        return None
    try:
        from langchain_google_genai import ChatGoogleGenerativeAI
        llm = ChatGoogleGenerativeAI(
            model="gemini-2.5-flash",
            google_api_key=GOOGLE_API_KEYS[0],
            temperature=0.2,
            max_output_tokens=4096,
            thinking={"thinking_budget": 1024},
        )
        logger.info("✅ Gemini 2.5 Flash initialized (key 1/%d)", len(GOOGLE_API_KEYS))
        return llm
    except Exception as e:
        logger.error("❌ Gemini init failed: %s", e)
        return None


def _build_gemini_rotation() -> list:
    """Build a list of Gemini LLM instances — one per API key for rotation."""
    if len(GOOGLE_API_KEYS) <= 1:
        return []
    from langchain_google_genai import ChatGoogleGenerativeAI
    extras = []
    for i, key in enumerate(GOOGLE_API_KEYS[1:], start=2):
        try:
            llm = ChatGoogleGenerativeAI(
                model="gemini-2.5-flash",
                google_api_key=key,
                temperature=0.2,
                max_output_tokens=4096,
                thinking={"thinking_budget": 1024},
            )
            extras.append(llm)
            logger.info("✅ Gemini 2.5 Flash key %d/%d initialized", i, len(GOOGLE_API_KEYS))
        except Exception as e:
            logger.warning("⚠️ Gemini key %d init failed: %s", i, e)
    return extras


def _build_nvidia_llm():
    """Build DeepSeek V3.2 LLM via NVIDIA Build (OpenAI-compatible API, free)."""
    if not NVIDIA_API_KEY:
        return None
    try:
        from langchain_openai import ChatOpenAI
        llm = ChatOpenAI(
            model=NVIDIA_MODEL,
            openai_api_key=NVIDIA_API_KEY,
            openai_api_base="https://integrate.api.nvidia.com/v1",
            temperature=0.2,
            max_tokens=4096,
        )
        logger.info("✅ %s via NVIDIA Build initialized (fallback)", NVIDIA_MODEL)
        return llm
    except Exception as e:
        logger.error("❌ NVIDIA Build init failed: %s", e)
        return None


def _build_openrouter_llm():
    """Build LLM via OpenRouter (legacy fallback)."""
    if not OPENROUTER_API_KEY:
        return None
    try:
        from langchain_openai import ChatOpenAI
        llm = ChatOpenAI(
            model=OPENROUTER_MODEL,
            openai_api_key=OPENROUTER_API_KEY,
            openai_api_base="https://openrouter.ai/api/v1",
            temperature=0.2,
            max_tokens=4096,
            default_headers={
                "HTTP-Referer": "https://deepmed-ai.hf.space",
                "X-Title": "DeepMed-AI",
            },
        )
        logger.info("✅ %s via OpenRouter initialized (fallback)", OPENROUTER_MODEL)
        return llm
    except Exception as e:
        logger.error("❌ OpenRouter init failed: %s", e)
        return None


def get_llm():
    """Return LLM with failover: Gemini → NVIDIA Build → OpenRouter.

    Uses LangChain's .with_fallbacks() so when Gemini hits quota/rate-limit
    errors, requests automatically route to fallback — no manual switching needed.
    """
    global _llm_instance
    if _llm_instance is not None:
        return _llm_instance

    gemini = _build_gemini_llm()
    gemini_extras = _build_gemini_rotation()
    nvidia = _build_nvidia_llm()
    openrouter = _build_openrouter_llm()

    # Build fallback chain: Gemini key1 → key2 → key3 → NVIDIA → OpenRouter
    fallbacks = gemini_extras + [fb for fb in [nvidia, openrouter] if fb is not None]

    if gemini and fallbacks:
        _llm_instance = gemini.with_fallbacks(fallbacks)
        n_keys = len(GOOGLE_API_KEYS)
        ext_fb = _fallback_name()
        if n_keys > 1:
            logger.info("🔄 LLM failover: Gemini ×%d keys → %s", n_keys, ext_fb)
        else:
            logger.info("🔄 LLM failover: Gemini 2.5 Flash → %s", ext_fb)
    elif gemini:
        if gemini_extras:
            _llm_instance = gemini.with_fallbacks(gemini_extras)
            logger.info("🟢 LLM: Gemini 2.5 Flash ×%d keys (no external fallback)", len(GOOGLE_API_KEYS))
        else:
            _llm_instance = gemini
            logger.info("🟢 LLM: Gemini 2.5 Flash only (no fallback configured)")
    elif fallbacks:
        _llm_instance = fallbacks[0] if len(fallbacks) == 1 else fallbacks[0].with_fallbacks(fallbacks[1:])
        logger.info("🟡 LLM: %s only (no Gemini key)", _fallback_name())
    else:
        logger.error("❌ Không có LLM nào khả dụng! Cần GOOGLE_API_KEY hoặc NVIDIA_API_KEY")
        return None

    return _llm_instance


def _fallback_name() -> str:
    """Return human-readable fallback name."""
    parts = []
    if NVIDIA_API_KEY:
        parts.append(f"{NVIDIA_MODEL} (NVIDIA)")
    if OPENROUTER_API_KEY:
        parts.append(f"{OPENROUTER_MODEL} (OpenRouter)")
    return " → ".join(parts) or "none"


def get_llm_name() -> str:
    """Return human-readable name of the active LLM configuration."""
    fb = _fallback_name()
    n_keys = len(GOOGLE_API_KEYS)
    gemini_label = f"Gemini 2.5 Flash ×{n_keys}" if n_keys > 1 else "Gemini 2.5 Flash"
    if GOOGLE_API_KEYS and fb:
        return f"{gemini_label} → {fb}"
    elif GOOGLE_API_KEYS:
        return gemini_label
    elif fb:
        return fb
    return "Không có LLM"