client/components/chatUI.py

from datetime import date
import streamlit as st
from utils.api import ask_question

_ACCESSED = date.today().strftime("%d %B %Y")

_CITATIONS = {
    "PubMedQA Dataset": {
        "authors": "Jin, Q., Dhingra, B., Liu, Z., Cohen, W. and Lu, X.",
        "year": "2019",
        "title": "PubMedQA: A Biomedical Research Question Answering Dataset",
        "venue": "Proceedings of EMNLP 2019",
        "url": "https://arxiv.org/abs/1909.06146",
        "hf_url": "https://huggingface.co/datasets/qiaojin/PubMedQA",
    },
    "Mental Health Counseling Dataset": {
        "authors": "Amod",
        "year": "2023",
        "title": "Mental Health Counseling Conversations",
        "venue": "HuggingFace Datasets",
        "url": "https://huggingface.co/datasets/Amod/mental_health_counseling_conversations",
        "hf_url": "https://huggingface.co/datasets/Amod/mental_health_counseling_conversations",
    },
    "Medical MediQA Dataset": {
        "authors": "Han, X. et al. (MedAlpaca)",
        "year": "2023",
        "title": "MedAlpaca — Medical Meadow MediQA",
        "venue": "HuggingFace Datasets",
        "url": "https://huggingface.co/datasets/medalpaca/medical_meadow_mediqa",
        "hf_url": "https://huggingface.co/datasets/medalpaca/medical_meadow_mediqa",
    },
    "MedQA-USMLE Dataset": {
        "authors": "Jin, D., Pan, E., Oufattole, N., Weng, W., Fang, H. and Szolovits, P.",
        "year": "2021",
        "title": "What Disease does this Patient Have? A Large-scale Open Domain Question Answering Dataset from Medical Exams",
        "venue": "Applied Sciences, 11(14), 6421",
        "url": "https://arxiv.org/abs/2009.13081",
        "hf_url": "https://huggingface.co/datasets/GBaker/MedQA-USMLE-4-options-hf",
    },
}

_SUGGESTIONS = [
    "Symptoms of Type 2 diabetes?",
    "How does hypertension affect the kidneys?",
    "Mechanism of action of ibuprofen?",
    "Difference between Type 1 and Type 2 diabetes?",
    "Common treatments for depression?",
]

_FRIENDLY_ERRORS = {
    "Connection refused": "The assistant is warming up. Please wait a moment and try again.",
    "NewConnectionError": "The assistant is warming up. Please wait a moment and try again.",
    "localhost":          "The assistant is warming up. Please wait a moment and try again.",
    "500":                "The assistant encountered an issue. Please try rephrasing your question.",
    "504":                "The request timed out. Please try a shorter question.",
    "404":                "Knowledge base is being set up. Load a dataset from the sidebar first.",
}


def _friendly_error(exc: Exception) -> str:
    msg = str(exc)
    for keyword, friendly in _FRIENDLY_ERRORS.items():
        if keyword in msg:
            return friendly
    return "Something went wrong. Please try again."


def render_chat() -> None:
    if "messages" not in st.session_state:
        st.session_state.messages = []
    if "pending_question" not in st.session_state:
        st.session_state.pending_question = None

    if not st.session_state.messages and st.session_state.pending_question is None:
        st.markdown(
            '<div class="cm-welcome">'
            "<p>Ask about symptoms, medications, or clinical conditions.<br>"
            "Upload your own documents in the sidebar to extend the knowledge base.</p>"
            "</div>",
            unsafe_allow_html=True,
        )
        row1 = st.columns(3)
        row2 = st.columns(2)
        for col, suggestion in zip(row1 + row2, _SUGGESTIONS):
            with col:
                if st.button(suggestion, key=f"sug_{suggestion[:15]}", use_container_width=True):
                    st.session_state.messages.append({"role": "user", "content": suggestion})
                    st.session_state.pending_question = suggestion
                    st.rerun()

    for msg in st.session_state.messages:
        with st.chat_message(msg["role"]):
            st.markdown(msg["content"])
            if msg.get("sources"):
                _render_references(msg["sources"])

    if st.session_state.pending_question:
        question = st.session_state.pending_question
        st.session_state.pending_question = None
        _fetch_and_store(question)
        st.rerun()

    if prompt := st.chat_input("Ask a medical question…"):
        st.session_state.messages.append({"role": "user", "content": prompt})
        st.session_state.pending_question = prompt
        st.rerun()


def _fetch_and_store(question: str) -> None:
    with st.chat_message("assistant"):
        with st.spinner("Searching knowledge base…"):
            try:
                data = ask_question(question)

                if "error" in data:
                    friendly = _friendly_error(Exception(data["error"]))
                    st.warning(friendly)
                    st.session_state.messages.append({"role": "assistant", "content": friendly})
                    return

                answer = data["response"]
                sources = [s for s in data.get("sources", []) if s and s != "unknown"]
                st.markdown(answer)
                if sources:
                    _render_references(sources)
                st.session_state.messages.append(
                    {"role": "assistant", "content": answer, "sources": sources}
                )
            except Exception as exc:
                friendly = _friendly_error(exc)
                st.warning(friendly)
                st.session_state.messages.append({"role": "assistant", "content": friendly})


def _render_references(sources: list[str]) -> None:
    with st.expander(f"References ({len(sources)})", expanded=False):
        for i, src in enumerate(sources, 1):
            if src in _CITATIONS:
                c = _CITATIONS[src]
                st.markdown(
                    f"**[{i}]** {c['authors']} ({c['year']}) "
                    f"*{c['title']}*, {c['venue']}. "
                    f"Available at: [{c['url']}]({c['url']}) "
                    f"(Accessed: {_ACCESSED})."
                )
            else:
                # Uploaded PDF
                st.markdown(f"**[{i}]** 📄 Uploaded document: *{src}*")