server/modules/llm.py

import re
from langchain_core.prompts import PromptTemplate
from langchain_core.runnables import RunnablePassthrough, RunnableLambda
from langchain_core.output_parsers import StrOutputParser
from langchain_groq import ChatGroq
from core.settings import get_settings

settings = get_settings()

_REWRITE_PROMPT = PromptTemplate.from_template(
    "Rewrite the following into a clear, standalone medical question.\n\n"
    "Question: {question}\n\n"
    "Rewritten:"
)

_ANSWER_PROMPT = PromptTemplate.from_template(
    "You are CureMind, an AI medical information assistant.\n\n"
    "Instructions:\n"
    "1. If the context below is relevant to the question, use it and cite sources with [number] notation.\n"
    "2. If the context is not relevant or is empty, answer from your general medical knowledge "
    "and begin your answer with: 'Based on general medical knowledge:'\n"
    "3. Never provide personal diagnoses or treatment prescriptions.\n"
    "4. Be clear, accurate, and concise.\n\n"
    "Context:\n{context}\n\n"
    "Question:\n{question}\n\n"
    "Answer:"
)

_SOURCE_LABELS: dict[str, str] = {
    "pubmedqa": "PubMedQA Dataset",
    "mental_health_counseling": "Mental Health Counseling Dataset",
    "medical_meadow_mediqa": "Medical MediQA Dataset",
    "medqa_usmle": "MedQA-USMLE Dataset",
}


def _format_source(metadata: dict) -> str:
    src = metadata.get("source", "unknown")
    if src in _SOURCE_LABELS:
        return _SOURCE_LABELS[src]
    filename = src.replace("\\", "/").split("/")[-1]
    page = metadata.get("page")
    return f"{filename} (p.{int(page) + 1})" if page is not None else filename


_REASONING_TRIGGERS = frozenset(
    {"why", "how", "explain", "cause", "mechanism", "pathophysiology", "difference", "compare"}
)


def _build_models() -> dict[str, ChatGroq]:
    kwargs = {"groq_api_key": settings.groq_api_key}
    return {
        "fast": ChatGroq(model_name="qwen/qwen3-32b", temperature=0.3, **kwargs),
        "reasoning": ChatGroq(model_name="deepseek-r1-distill-qwen-32b", temperature=0.2, **kwargs),
        "fallback": ChatGroq(model_name="llama-3.3-70b-versatile", temperature=0.3, **kwargs),
    }


def _format_docs(docs: list) -> str:
    return "\n\n".join(
        f"[{i + 1}] {doc.page_content}\n(Source: {doc.metadata.get('source', 'unknown')})"
        for i, doc in enumerate(docs)
    )


def _build_output(x: dict) -> dict:
    answer = re.sub(r"<think>.*?</think>", "", x["answer"], flags=re.DOTALL).strip()
    used_general_knowledge = answer.lower().startswith("based on general medical knowledge")
    sources = (
        []
        if used_general_knowledge or not x["docs"]
        else list(dict.fromkeys(_format_source(doc.metadata) for doc in x["docs"]))
    )
    return {
        "question": x["question"],
        "rewritten_question": x["rewritten_question"],
        "answer": answer,
        "sources": sources,
    }


def get_llm_chain(retriever):
    models = _build_models()

    def _select_model(x: dict) -> ChatGroq:
        q = x["question"].lower()
        if any(w in q for w in _REASONING_TRIGGERS) or len(q) > 120:
            return models["reasoning"]
        return models["fast"]

    def _generate_answer(x: dict) -> str:
        try:
            return (_ANSWER_PROMPT | _select_model(x) | StrOutputParser()).invoke(x)
        except Exception:
            return (_ANSWER_PROMPT | models["fallback"] | StrOutputParser()).invoke(x)

    return (
        {"question": RunnablePassthrough()}
        | RunnablePassthrough.assign(
            rewritten_question=_REWRITE_PROMPT | models["fast"] | StrOutputParser()
        )
        | RunnablePassthrough.assign(
            docs=lambda x: retriever.invoke(x["rewritten_question"])
        )
        | RunnablePassthrough.assign(
            context=lambda x: _format_docs(x["docs"])
        )
        | RunnablePassthrough.assign(
            answer=RunnableLambda(_generate_answer)
        )
        | RunnableLambda(lambda x: _build_output(x))
    )