app.py

"""
MyeAI (Myeloma AI)
===========================================================
Takes a participant's answers to the 16 Shared Decision Making questions
and runs an adaptive follow-up conversation using Google's Gemini API.

GROUNDING:
General informational questions from the user are answered STRICTLY from the
documents in the uploaded ZIP corpus (Archive_2.zip) using retrieval-augmented
generation. At startup the app extracts text from every PDF/PNG in the corpus,
embeds it with Gemini embeddings, and builds a searchable index. Each user
question retrieves the most relevant passages, and Gemini is instructed to
answer only from those passages (or say the answer isn't in the documents).

Deploy on Hugging Face Spaces (SDK: Gradio).
- Upload Archive_2.zip to the Space repo root (or set KB_ZIP / KB_DIR).
- Set your Gemini key as a Space Secret named GEMINI_API_KEY.
- Recommended packages.txt: poppler-utils, tesseract-ocr  (see notes at bottom).
"""

import os
import json

# ----------------------------------------------------------------------
# ZeroGPU compatibility shim (harmless on CPU basic; skipped if not present).
# Satisfies "No @spaces.GPU function detected during startup" if the Space is
# on ZeroGPU. This app is API-only and does not use a GPU — CPU basic is fine.
# ----------------------------------------------------------------------
try:
    import spaces

    @spaces.GPU
    def _zerogpu_warmup():
        return None
except Exception:
    pass

from google import genai
from google.genai import types
import gradio as gr

from knowledge_base import KnowledgeBase

# ----------------------------------------------------------------------
# The 16 profile questions (verbatim from the intake form)
# ----------------------------------------------------------------------
QUESTIONS = [
    "I prefer my healthcare team (hematologist, nurse practitioner, or physician assistant), and I collaborate in deciding which treatment for relapsed or refractory myeloma is best for me",
    "It is important for me to understand my treatment options for relapsed or refractory myeloma",
    "I trust my healthcare team very much, that's why I leave it up to them to recommend therapy that is right for me",
    "I live alone with no caregiver",
    "I have strong social support and a network that can help with my treatment appointments",
    "I will do whatever it takes to so I can be cured, cancer-free, kill all the myeloma, or at least be in remission and live a long life",
    "I would like to take an aggressive approach to treat my myeloma",
    "I am willing to endure as many side effects as possible to control my myeloma",
    "I prefer to receive treatment in an outpatient setting",
    "I prefer to take medications at home",
    "I prefer to take the least possible amount of pills to control my cancer",
    "Quality of life is more important to me than quantity of life",
    "Clinical drug trial participation is of interest to me",
    "My out-of-pocket cost of treatment is important to me",
    "I prefer to continue an active lifestyle during my myeloma treatment",
    "I worry about how my treatment will affect future treatment options",
]

# Short tags used for the summary view
QUESTION_TAGS = [
    "Shared decision-making",
    "Wants to understand options",
    "Defers to care team",
    "Lives alone / no caregiver",
    "Strong social support",
    "Goal: cure / long life",
    "Wants aggressive approach",
    "Tolerant of side effects",
    "Prefers outpatient",
    "Prefers meds at home",
    "Prefers fewer pills",
    "Quality over quantity",
    "Interested in clinical trials",
    "Out-of-pocket cost matters",
    "Wants active lifestyle",
    "Worries about future options",
]

MODEL = "gemini-2.5-flash"
MAX_FOLLOWUPS = 8  # cap the adaptive conversation

# The conversation always opens with this fixed question from the patient.
FIRST_QUESTION = "Using my profile, what are the best treatment options for my myeloma?"

# ----------------------------------------------------------------------
# Knowledge base — built once at startup.
# ----------------------------------------------------------------------
KB = KnowledgeBase()
_KB_BUILD_ATTEMPTED = False


def ensure_kb(client):
    """Build the KB on first successful client. Returns (ready, message)."""
    global _KB_BUILD_ATTEMPTED
    if KB.ready:
        return True, ""
    if _KB_BUILD_ATTEMPTED and KB.error:
        return False, KB.error
    _KB_BUILD_ATTEMPTED = True
    KB.build(client)
    if KB.ready:
        return True, ""
    return False, (KB.error or "Knowledge base could not be initialized.")


# ----------------------------------------------------------------------
# Prompt building
# ----------------------------------------------------------------------
def build_profile_text(answers):
    """answers: list of 'Yes'/'No' aligned to QUESTIONS."""
    lines = []
    for q, tag, a in zip(QUESTIONS, QUESTION_TAGS, answers):
        lines.append(f"- [{a.upper()}] ({tag}) {q}")
    return "\n".join(lines)


def detect_tensions(answers):
    """Flag notable patterns/contradictions worth probing. Pure logic, no API."""
    a = {i: (answers[i].strip().lower() == "yes") for i in range(len(answers))}
    flags = []

    if (a.get(5) or a.get(6) or a.get(7)) and a.get(11):
        flags.append("Wants an aggressive/curative approach but also values quality of life over quantity. Worth clarifying how they weigh these when they conflict.")
    if a.get(7) and (a.get(14) or a.get(11)):
        flags.append("Willing to endure many side effects, yet wants to stay active / prioritizes quality of life. Probe acceptable side-effect threshold.")
    if a.get(3) and not a.get(4):
        flags.append("Lives alone with no caregiver and limited social support. Treatment logistics and safety monitoring need exploration.")
    if a.get(3) and a.get(9):
        flags.append("Lives alone but prefers taking medications at home. Explore support for safe self-administration and side-effect monitoring.")
    if a.get(2) and (a.get(0) or a.get(1)):
        flags.append("Says they leave decisions to the care team, yet also wants to collaborate / understand options. Clarify how involved they actually want to be.")
    if a.get(13) and a.get(12):
        flags.append("Cost matters and they're open to clinical trials — trials may reduce drug cost; worth surfacing.")
    if a.get(10) and (a.get(6) or a.get(7)):
        flags.append("Prefers the fewest pills possible but wants an aggressive approach. Explore tolerance for treatment intensity vs convenience.")
    if a.get(8) and a.get(3):
        flags.append("Prefers outpatient treatment but lives alone — explore transport and post-visit support.")
    return flags


def system_instruction():
    return (
        "You are a warm, plain-spoken health navigator helping a multiple myeloma patient "
        "(relapsed or refractory) prepare for a shared decision-making conversation with their "
        "care team. You are NOT a doctor and you never give medical advice, diagnoses, dosing, or "
        "treatment recommendations. Your job is to (a) ask thoughtful follow-up questions that help "
        "the patient clarify their own values, priorities, constraints, and concerns, and (b) answer "
        "the patient's general informational questions using ONLY the reference documents provided to "
        "you in the RETRIEVED CONTEXT for that turn.\n\n"
        "STRICT GROUNDING RULES (most important):\n"
        "A. For any general/informational question the patient asks (e.g. 'what is CAR T-cell "
        "therapy?', 'what are the side effects of stem cell transplant?', 'what does relapsed mean?'), "
        "you MUST base your answer solely on the text inside the RETRIEVED CONTEXT block for that turn. "
        "Do NOT use outside knowledge, and do NOT add facts that are not present in that context.\n"
        "B. If the RETRIEVED CONTEXT does not contain enough information to answer, say so plainly: "
        "\"I couldn't find that in the reference documents provided for this tool. Please ask your care "
        "team.\" Do not guess or fill gaps from general knowledge.\n"
        "C. When you use information from the context, attribute it naturally to the document "
        "TITLE(S) shown in each passage's [Document: ...] label — use that title, not a number, and "
        "never say \"Source 1/2/3\" (e.g. \"According to the NCCN Guideline ...\" or \"The CAR T cell "
        "Therapy — International Myeloma Foundation page notes ...\"). Use the title exactly as shown; "
        "do not invent or embellish document names.\n"
        "D. Never invent drug names, doses, statistics, or study results that are not in the context.\n\n"
        "CONVERSATION RULES:\n"
        "1. When you ask a follow-up question, ask exactly ONE question per turn, 1-3 sentences, "
        "conversational and jargon-free.\n"
        "2. Base follow-up questions on the patient's profile and previous answers. Prioritize the "
        "FLAGGED TENSIONS provided — gently explore apparent contradictions without judgment.\n"
        "3. Do not repeat questions already asked. Build on what they say.\n"
        "4. Never recommend or rank treatments. If the patient asks for medical advice or 'what should "
        "I do', kindly redirect them to their care team, then (if appropriate) continue with a follow-up "
        "question.\n"
        "5. Tone: empathetic, respectful, never alarming.\n"
        "6. TREATMENT-OPTIONS QUESTION: When the patient asks what the best treatment options are for "
        "their myeloma (e.g. 'Using my profile, what are the best treatment options for my myeloma?'), "
        "combine TWO sources: their 16 yes/no profile answers AND the RETRIEVED CONTEXT documents. Do "
        "this:\n"
        "   - Open with one warm sentence.\n"
        "   - Walk through how THEIR specific stated priorities map to the kinds of treatment approaches "
        "described in the RETRIEVED CONTEXT for relapsed/refractory myeloma. Tie each point explicitly "
        "back to their answers (for example: a preference for medications at home and fewest pills points "
        "toward asking about convenient outpatient or oral regimens; openness to clinical trials means "
        "trials are worth raising; willingness to endure side effects and wanting an aggressive approach "
        "points toward asking about more intensive options; prioritizing quality of life points toward "
        "regimens that protect daily functioning; living alone with limited support points toward "
        "logistics and monitoring).\n"
        "   - Only describe treatment approaches that actually appear in the RETRIEVED CONTEXT. Do NOT "
        "invent clinical details, drug names, dosing, or anything not supported by the context or implied "
        "by their answers. Do NOT rank, prescribe, or state which option is medically best. Frame "
        "everything as 'based on what you told us and these reference materials, here are options and "
        "questions to raise with your care team.'\n"
        "   - End by reminding them their care team must confirm what is medically appropriate, then ask "
        "your first single follow-up question drawn from their profile and the flagged tensions.\n"
        "7. When you judge that you have enough to summarize their priorities (or after several "
        "exchanges), instead of asking another question, output a final summary. Begin that final "
        "message with the exact token <<SUMMARY>> on its own line, then 4-7 short bullet points "
        "capturing the patient's key priorities, constraints, and questions to raise with their care team."
    )


# ----------------------------------------------------------------------
# Gemini client helpers
# ----------------------------------------------------------------------
def get_client(user_key):
    key = (user_key or "").strip() or os.environ.get("GEMINI_API_KEY", "").strip()
    if not key:
        return None, "No API key found. Set GEMINI_API_KEY as a Space secret."
    try:
        client = genai.Client(api_key=key)
        return client, None
    except Exception as e:
        return None, f"Could not initialize Gemini client: {e}"


def to_gemini_contents(history):
    """history: list of {'role': 'user'|'assistant', 'content': str}
    Returns Gemini Content list (assistant -> 'model')."""
    contents = []
    for m in history:
        role = "model" if m["role"] == "assistant" else "user"
        contents.append(types.Content(role=role, parts=[types.Part(text=m["content"])]))
    return contents


def _latest_user_query(convo):
    for m in reversed(convo):
        if m["role"] == "user":
            return m["content"]
    return ""


def gemini_reply(client, convo, n_followups):
    """convo: internal history list. Retrieves KB context for the latest user
    turn and injects it, then asks Gemini to answer grounded in that context."""
    # Retrieve relevant passages for the most recent user message.
    context_block = ""
    query = _latest_user_query(convo)
    if KB.ready and query.strip():
        try:
            retrieved = KB.retrieve(client, query)
            if retrieved:
                context_block = KB.context_block(retrieved)
        except Exception:
            context_block = ""

    contents = to_gemini_contents(convo)

    # Build the per-turn system instruction with the retrieved context appended.
    sys = system_instruction()
    if context_block:
        sys += (
            "\n\n==================== RETRIEVED CONTEXT (reference documents) ====================\n"
            "The following passages were retrieved from the uploaded reference documents for THIS "
            "turn. Answer general/informational questions using ONLY these passages. If they do not "
            "contain the answer, say you couldn't find it in the reference documents.\n\n"
            f"{context_block}\n"
            "================================================================================\n"
        )
    else:
        sys += (
            "\n\n[No reference passages were retrieved for this turn. If the patient asked a general "
            "informational question, tell them you couldn't find it in the reference documents and "
            "suggest they ask their care team. You may still ask a values-clarifying follow-up "
            "question or work with their profile.]\n"
        )

    if n_followups >= MAX_FOLLOWUPS - 1:
        sys += "\n\nYou have asked enough questions. Provide the final <<SUMMARY>> now."

    cfg = types.GenerateContentConfig(
        system_instruction=sys,
        temperature=0.3,          # lower temp -> stays closer to the sources
        # gemini-2.5-flash has "thinking" on by default, and those thinking
        # tokens are drawn from max_output_tokens. That could exhaust the budget
        # and truncate the visible answer mid-sentence (e.g. the long treatment-
        # options response). Disable thinking (budget=0) so the whole budget goes
        # to the answer, and raise the ceiling for extra headroom.
        max_output_tokens=8192,
        thinking_config=types.ThinkingConfig(thinking_budget=0),
    )
    resp = client.models.generate_content(model=MODEL, contents=contents, config=cfg)
    return (resp.text or "").strip()


# ----------------------------------------------------------------------
# Gradio app
# ----------------------------------------------------------------------
def start_chat(api_key, *radio_values):
    answers = [v if v in ("Yes", "No") else None for v in radio_values]
    if any(v is None for v in answers):
        missing = [i + 1 for i, v in enumerate(answers) if v is None]
        gr.Warning(f"Please answer all 16 questions. Missing: {missing}")
        return (gr.update(), [], [], 0, gr.update(visible=True), gr.update(visible=False))

    client, err = get_client(api_key)
    if err:
        gr.Warning(err)
        return (gr.update(), [], [], 0, gr.update(visible=True), gr.update(visible=False))

    # Build the knowledge base on first use (may take a bit on cold start).
    ready, msg = ensure_kb(client)
    if not ready:
        gr.Warning(f"Reference documents unavailable: {msg}")
        # We still allow the conversation, but answers will note missing docs.

    profile = build_profile_text(answers)
    flags = detect_tensions(answers)
    flag_text = "\n".join(f"- {f}" for f in flags) if flags else "- (No obvious contradictions; explore their highest-stakes priorities.)"

    seed = (
        "Here is the patient's completed profile (16 yes/no answers):\n\n"
        f"{profile}\n\n"
        "FLAGGED TENSIONS / PRIORITIES TO EXPLORE:\n"
        f"{flag_text}\n\n"
        "Acknowledge in one short sentence that you have their profile. Do NOT ask a question yet "
        "and do NOT list treatment options yet. Simply invite them to ask their question."
    )

    convo = [{"role": "user", "content": seed}]
    try:
        reply = gemini_reply(client, convo, 0)
    except Exception as e:
        gr.Warning(f"Gemini error: {e}")
        return (gr.update(), [], [], 0, gr.update(visible=True), gr.update(visible=False))

    convo.append({"role": "assistant", "content": reply})
    chat_display = [{"role": "assistant", "content": reply}]
    return (
        chat_display,
        chat_display,
        convo,
        0,
        gr.update(visible=False),
        gr.update(visible=True),
    )


def respond(user_msg, chat_display, convo, n_followups, api_key):
    user_msg = (user_msg or "").strip()
    if not user_msg:
        return chat_display, chat_display, convo, n_followups, ""

    client, err = get_client(api_key)
    if err:
        gr.Warning(err)
        return chat_display, chat_display, convo, n_followups, user_msg

    ensure_kb(client)  # no-op if already built

    convo = convo + [{"role": "user", "content": user_msg}]
    chat_display = chat_display + [{"role": "user", "content": user_msg}]

    try:
        reply = gemini_reply(client, convo, n_followups)
    except Exception as e:
        gr.Warning(f"Gemini error: {e}")
        return chat_display, chat_display, convo, n_followups, ""

    convo = convo + [{"role": "assistant", "content": reply}]

    display_reply = reply
    if "<<SUMMARY>>" in reply:
        display_reply = reply.replace("<<SUMMARY>>", "").strip()
        display_reply = "**Your Priorities Summary**\n\n" + display_reply + (
            "\n\n_This summary reflects what you shared. Please bring it to your care team. "
            "It is not medical advice._"
        )

    chat_display = chat_display + [{"role": "assistant", "content": display_reply}]
    return chat_display, chat_display, convo, n_followups + 1, ""


def reset_all():
    radio_resets = [gr.update(value=None) for _ in QUESTIONS]
    return (
        [],
        [],
        [],
        0,
        gr.update(visible=True),
        gr.update(visible=False),
        *radio_resets,
    )


CSS = """
@import url('https://fonts.googleapis.com/css2?family=Inter:wght@400;500;600;700&display=swap');
.gradio-container, .gradio-container * {
    font-family: 'Inter', -apple-system, BlinkMacSystemFont, 'Segoe UI', Roboto, Helvetica, Arial, sans-serif !important;
}
.gradio-container {
    background: #ffffff !important;
    color: #111111 !important;
    max-width: 880px !important;
    margin: 0 auto !important;
}
body, .gradio-container .prose, .gradio-container p,
.gradio-container span, .gradio-container label {
    color: #111111 !important;
}
#app-header {
    border-bottom: 1px solid #e6e6e6;
    padding: 4px 0 18px 0;
    margin-bottom: 6px;
}
#app-title {
    font-size: 1.7rem;
    font-weight: 700;
    color: #111111;
    letter-spacing: -0.01em;
    margin: 0;
}
.q-card {
    border: 1px solid #e2e2e2 !important;
    border-radius: 8px !important;
    padding: 16px 18px !important;
    margin-bottom: 12px !important;
    background: #ffffff !important;
    box-shadow: 0 1px 2px rgba(0,0,0,0.04) !important;
}
.q-card label, .q-card span {
    color: #111111 !important;
    font-weight: 500 !important;
}
.gradio-container input[type="radio"] + span,
.gradio-container .wrap label {
    color: #111111 !important;
}
button.primary, .gradio-container button.primary {
    background: #7a0c2e !important;
    color: #ffffff !important;
    border: none !important;
    border-radius: 6px !important;
    font-weight: 600 !important;
}
button.primary:hover, .gradio-container button.primary:hover {
    background: #5f0a24 !important;
}
button.secondary, .gradio-container button.secondary {
    background: #ffffff !important;
    color: #111111 !important;
    border: 1px solid #cfcfcf !important;
    border-radius: 6px !important;
    font-weight: 600 !important;
}
.gradio-container .chatbot, .gradio-container [class*="chatbot"] {
    background: #ffffff !important;
    border: 1px solid #e2e2e2 !important;
    border-radius: 8px !important;
}
.gradio-container .message.bot, .gradio-container .message.user {
    color: #111111 !important;
}
.gradio-container textarea, .gradio-container input[type="text"],
.gradio-container input[type="password"] {
    background: #ffffff !important;
    color: #111111 !important;
    border: 1px solid #cfcfcf !important;
    border-radius: 6px !important;
}
.gradio-container .label-wrap, .gradio-container details summary {
    color: #111111 !important;
}
.sample-q {
    background: #f7f7f7 !important;
    border: 1px solid #e2e2e2 !important;
    border-left: 3px solid #7a0c2e !important;
    border-radius: 6px !important;
    padding: 10px 14px !important;
    margin: 12px 0 4px 0 !important;
}
.sample-q-label {
    display: block;
    font-size: 0.85rem;
    font-weight: 600;
    color: #5a5a5a !important;
    margin-bottom: 4px;
}
.sample-q-text {
    display: block;
    font-size: 0.95rem;
    color: #111111 !important;
    background: #ffffff !important;
    border: 1px solid #e2e2e2 !important;
    border-radius: 4px !important;
    padding: 6px 10px !important;
    font-family: 'Inter', sans-serif !important;
}
footer {visibility: hidden;}
"""

with gr.Blocks(title="MyeAI (Myeloma AI)", css=CSS,
               theme=gr.themes.Base(
                   primary_hue=gr.themes.colors.red,
                   neutral_hue=gr.themes.colors.gray,
                   font=[gr.themes.GoogleFont("Inter"), "system-ui", "sans-serif"],
               )) as demo:
    with gr.Column(elem_id="app-header"):
        gr.HTML(
            "<h1 id='app-title'>MyeAI (Myeloma AI)</h1>"
        )

    api_key = gr.State("")
    convo_state = gr.State([])
    display_state = gr.State([])
    followups = gr.State(0)

    with gr.Group(visible=True) as intake_panel:
        gr.Markdown("**Answer all 16 questions, then click Start Follow-up.**")
        radios = []
        for i, q in enumerate(QUESTIONS):
            with gr.Row(elem_classes="q-card"):
                r = gr.Radio(["Yes", "No"], label=f"{i+1}. {q}")
                radios.append(r)
        start_btn = gr.Button("Start Follow-up", variant="primary")

    with gr.Group(visible=False) as chat_panel:
        chatbot = gr.Chatbot(label="Follow-up Conversation", type="messages", height=460)
        gr.HTML(
            "<div class='sample-q'>"
            "<span class='sample-q-label'>Sample question &mdash; copy and paste it into the box below to begin:</span>"
            f"<code class='sample-q-text'>{FIRST_QUESTION}</code>"
            "</div>"
        )
        with gr.Row():
            msg = gr.Textbox(placeholder="Type or paste your question here...", show_label=False, scale=8)
            send_btn = gr.Button("Send", variant="primary", scale=1)
        restart_btn = gr.Button("Start Over")

    start_btn.click(
        start_chat,
        inputs=[api_key] + radios,
        outputs=[chatbot, display_state, convo_state, followups, intake_panel, chat_panel],
    )
    send_btn.click(
        respond,
        inputs=[msg, display_state, convo_state, followups, api_key],
        outputs=[chatbot, display_state, convo_state, followups, msg],
    )
    msg.submit(
        respond,
        inputs=[msg, display_state, convo_state, followups, api_key],
        outputs=[chatbot, display_state, convo_state, followups, msg],
    )
    restart_btn.click(
        reset_all,
        inputs=None,
        outputs=[chatbot, display_state, convo_state, followups, intake_panel, chat_panel] + radios,
    )

if __name__ == "__main__":
    demo.launch()

# ----------------------------------------------------------------------
# HUGGING FACE SPACE SETUP NOTES
# ----------------------------------------------------------------------
# 1) Files in the Space repo:
#      app.py
#      knowledge_base.py
#      Archive_2.zip          <- the uploaded corpus (repo root)
#      requirements.txt       <- google-genai, gradio, numpy, pypdf, pdf2image, pytesseract, pillow
#      packages.txt           <- poppler-utils, tesseract-ocr   (system deps for PDF text + OCR)
# 2) Space Secret:
#      GEMINI_API_KEY = <your key>
# 3) Hardware: CPU basic is sufficient (no GPU needed).
# 4) On first launch the app extracts + embeds the corpus once and caches the
#    index to _kb_index.npz, so subsequent restarts are fast.