| 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: |
| |
| st.markdown(f"**[{i}]** 📄 Uploaded document: *{src}*") |
|
|
