Use Q&A breakdown agent

pipeline.py

@@ -3,6 +3,7 @@ import os
 import json
 import time
 import logging
+import threading
 import concurrent.futures
 import gradio as gr
 import spaces
@@ -18,9 +19,168 @@ from supervisor import (
     gemini_supervisor_breakdown, gemini_supervisor_search_strategies,
     gemini_supervisor_rag_brainstorm, execute_medswin_task,
     gemini_supervisor_synthesize, gemini_supervisor_challenge,
-    gemini_supervisor_enhance_answer, gemini_supervisor_check_clarity
+    gemini_supervisor_enhance_answer, gemini_supervisor_check_clarity,
+    gemini_clinical_intake_triage, gemini_summarize_clinical_insights
 )
 
+MAX_CLINICAL_QA_ROUNDS = 5
+_clinical_intake_sessions = {}
+_clinical_intake_lock = threading.Lock()
+
+
+def _get_clinical_intake_state(session_id: str):
+    with _clinical_intake_lock:
+        return _clinical_intake_sessions.get(session_id)
+
+
+def _set_clinical_intake_state(session_id: str, state: dict):
+    with _clinical_intake_lock:
+        _clinical_intake_sessions[session_id] = state
+
+
+def _clear_clinical_intake_state(session_id: str):
+    with _clinical_intake_lock:
+        _clinical_intake_sessions.pop(session_id, None)
+
+
+def _history_to_text(history: list, limit: int = 6) -> str:
+    if not history:
+        return "No prior conversation."
+    recent = history[-limit:]
+    lines = []
+    for turn in recent:
+        role = turn.get("role", "user")
+        content = turn.get("content", "")
+        lines.append(f"{role}: {content}")
+    return "\n".join(lines)
+
+
+def _format_intake_question(question: dict, round_idx: int, max_rounds: int, target_lang: str) -> str:
+    header = f"🩺 Clinical intake question {round_idx}/{max_rounds}"
+    body = question.get("question") or "Could you share a bit more detail so I can give an accurate answer?"
+    focus = question.get("clinical_focus")
+    why = question.get("why_it_matters")
+    prompt_parts = [header, body]
+    if focus:
+        prompt_parts.append(f"Focus: {focus}")
+    if why:
+        prompt_parts.append(f"Why it matters: {why}")
+    prompt_parts.append("Please answer in 1-2 sentences so I can continue.")
+    prompt_text = "\n\n".join(prompt_parts)
+    if target_lang and target_lang != "en":
+        try:
+            prompt_text = translate_text(prompt_text, target_lang=target_lang, source_lang="en")
+        except Exception as exc:
+            logger.warning(f"[INTAKE] Question translation failed: {exc}")
+    return prompt_text
+
+
+def _format_insights_block(insights: dict) -> str:
+    if not insights:
+        return ""
+    lines = []
+    profile = insights.get("patient_profile")
+    if profile:
+        lines.append(f"- Patient profile: {profile}")
+    for finding in insights.get("key_findings", []):
+        title = finding.get("title", "Insight")
+        detail = finding.get("detail", "")
+        implication = finding.get("clinical_implication", "")
+        line = f"- {title}: {detail}"
+        if implication:
+            line += f" (Clinical note: {implication})"
+        lines.append(line)
+    return "\n".join(lines)
+
+
+def _build_refined_query(base_query: str, insights: dict, insights_block: str) -> str:
+    sections = [base_query.strip()] if base_query else []
+    if insights_block:
+        sections.append(f"Clinical intake summary:\n{insights_block}")
+    refined = insights.get("refined_problem_statement")
+    if refined:
+        sections.append(f"Refined problem statement:\n{refined}")
+    handoff = insights.get("handoff_note")
+    if handoff:
+        sections.append(f"Handoff note:\n{handoff}")
+    return "\n\n".join([section for section in sections if section])
+
+
+def _start_clinical_intake_session(session_id: str, plan: dict, base_query: str, original_language: str):
+    questions = plan.get("questions", []) or []
+    if not questions:
+        return None
+    max_rounds = plan.get("max_rounds") or len(questions)
+    max_rounds = max(1, min(MAX_CLINICAL_QA_ROUNDS, max_rounds, len(questions)))
+    state = {
+        "base_query": base_query,
+        "original_language": original_language or "en",
+        "questions": questions,
+        "max_rounds": max_rounds,
+        "current_round": 1,
+        "pending_question_index": 0,
+        "awaiting_answer": True,
+        "answers": [],
+        "decision_reason": plan.get("decision_reason", ""),
+        "initial_hypotheses": plan.get("initial_hypotheses", []),
+        "started_at": time.time()
+    }
+    _set_clinical_intake_state(session_id, state)
+    first_prompt = _format_intake_question(
+        questions[0],
+        round_idx=1,
+        max_rounds=max_rounds,
+        target_lang=state["original_language"]
+    )
+    return first_prompt
+
+
+def _handle_clinical_answer(session_id: str, answer_text: str):
+    state = _get_clinical_intake_state(session_id)
+    if not state:
+        return {"type": "error"}
+    questions = state.get("questions", [])
+    idx = state.get("pending_question_index", 0)
+    if idx >= len(questions):
+        logger.warning("[INTAKE] Pending question index out of range, ending intake session")
+        _clear_clinical_intake_state(session_id)
+        return {"type": "error"}
+    question_meta = questions[idx] or {}
+    qa_entry = {
+        "question": question_meta.get("question", ""),
+        "focus": question_meta.get("clinical_focus"),
+        "why_it_matters": question_meta.get("why_it_matters"),
+        "round": state.get("current_round", len(state.get("answers", [])) + 1),
+        "answer": answer_text.strip()
+    }
+    state["answers"].append(qa_entry)
+    next_index = idx + 1
+    reached_round_limit = len(state["answers"]) >= state["max_rounds"]
+    if reached_round_limit or next_index >= len(questions):
+        insights = gemini_summarize_clinical_insights(state["base_query"], state["answers"])
+        insights_block = _format_insights_block(insights)
+        refined_query = _build_refined_query(state["base_query"], insights, insights_block)
+        _clear_clinical_intake_state(session_id)
+        return {
+            "type": "insights",
+            "insights": insights,
+            "insights_block": insights_block,
+            "refined_query": refined_query,
+            "qa_pairs": state["answers"]
+        }
+    state["pending_question_index"] = next_index
+    state["current_round"] = len(state["answers"]) + 1
+    state["awaiting_answer"] = True
+    _set_clinical_intake_state(session_id, state)
+    next_question = questions[next_index]
+    prompt = _format_intake_question(
+        next_question,
+        round_idx=state["current_round"],
+        max_rounds=state["max_rounds"],
+        target_lang=state["original_language"]
+    )
+    return {"type": "question", "prompt": prompt}
+
 
 @spaces.GPU(max_duration=120)
 def stream_chat(
@@ -37,6 +197,7 @@ def stream_chat(
     use_rag: bool,
     medical_model: str,
     use_web_search: bool,
+    enable_clinical_intake: bool,
     disable_agentic_reasoning: bool,
     show_thoughts: bool,
     request: gr.Request
@@ -73,9 +234,15 @@ def stream_chat(
         "plan": None,
         "strategy_decisions": [],
         "stage_metrics": {},
-        "search": {"strategies": [], "total_results": 0}
+        "search": {"strategies": [], "total_results": 0},
+        "clinical_intake": {
+            "enabled": enable_clinical_intake,
+            "activated": False,
+            "rounds": 0,
+            "reason": "",
+            "insights": []
+        }
     }
-    
     def record_stage(stage_name: str, start_time: float):
         pipeline_diagnostics["stage_metrics"][stage_name] = round(time.time() - start_time, 3)
     
@@ -95,6 +262,47 @@ def stream_chat(
         {"role": "assistant", "content": ""}
     ]
     
+    if not enable_clinical_intake:
+        _clear_clinical_intake_state(user_id)
+    else:
+        intake_state = _get_clinical_intake_state(user_id)
+        if intake_state and intake_state.get("awaiting_answer"):
+            logger.info("[INTAKE] Awaiting patient response - processing answer")
+            intake_result = _handle_clinical_answer(user_id, message)
+            if intake_result.get("type") == "question":
+                logger.info("[INTAKE] Requesting additional follow-up")
+                updated_history[-1]["content"] = intake_result["prompt"]
+                thoughts_text = thought_handler.get_thoughts() if (show_thoughts and thought_handler) else ""
+                yield updated_history, thoughts_text
+                if thought_handler:
+                    logger.removeHandler(thought_handler)
+                return
+            if intake_result.get("type") == "insights":
+                pipeline_diagnostics["clinical_intake"]["activated"] = True
+                pipeline_diagnostics["clinical_intake"]["rounds"] = len(intake_result.get("qa_pairs", []))
+                pipeline_diagnostics["clinical_intake"]["insights"] = intake_result.get("insights", {}).get("key_findings", [])
+                message = intake_result.get("refined_query", message)
+        else:
+            history_context = _history_to_text(history)
+            triage_plan = gemini_clinical_intake_triage(message, history_context, MAX_CLINICAL_QA_ROUNDS)
+            pipeline_diagnostics["clinical_intake"]["reason"] = triage_plan.get("decision_reason", "")
+            needs_intake = triage_plan.get("needs_additional_info") and triage_plan.get("questions")
+            if needs_intake:
+                first_prompt = _start_clinical_intake_session(
+                    user_id,
+                    triage_plan,
+                    message,
+                    original_lang
+                )
+                if first_prompt:
+                    pipeline_diagnostics["clinical_intake"]["activated"] = True
+                    updated_history[-1]["content"] = first_prompt
+                    thoughts_text = thought_handler.get_thoughts() if (show_thoughts and thought_handler) else ""
+                    yield updated_history, thoughts_text
+                    if thought_handler:
+                        logger.removeHandler(thought_handler)
+                    return
+    
     plan = None
     if not disable_agentic_reasoning:
         reasoning_stage_start = time.time()

supervisor.py

@@ -217,6 +217,210 @@ Keep contexts brief and factual. Avoid redundancy."""
         }
 
 
+async def gemini_clinical_intake_triage_async(
+    query: str,
+    history_context: str,
+    max_rounds: int = 5
+) -> dict:
+    """Gemini Intake Agent: Decide if additional clinical intake is needed and plan questions"""
+    history_block = history_context if history_context else "No prior conversation."
+    safe_rounds = max(1, min(5, max_rounds))
+    prompt = f"""You are a clinical intake coordinator helping a medical AI system.
+Your job is to review the patient's latest request and decide if more clinical details are required before analysis.
+
+Patient query:
+"{query}"
+
+Recent conversation (if any):
+{history_block}
+
+Return ONLY valid JSON (no markdown):
+{{
+  "needs_additional_info": true | false,
+  "decision_reason": "brief justification",
+  "max_rounds": {safe_rounds},
+  "questions": [
+    {{
+      "order": 1,
+      "question": "single follow-up question to ask the patient",
+      "clinical_focus": "what aspect it clarifies (e.g., onset, severity, meds)",
+      "why_it_matters": "concise clinical rationale",
+      "optional": false
+    }},
+    ...
+  ],
+  "initial_hypotheses": [
+    "optional bullet on potential etiologies or next steps"
+  ]
+}}
+
+Guidelines:
+- Ask at most {safe_rounds} questions. Use fewer if the query is already specific.
+- Order questions to maximize clinical value.
+- Only mark needs_additional_info true when the current data is insufficient for safe reasoning.
+- Keep wording patient-friendly and concise."""
+
+    system_prompt = "You are a triage clinician. Decide if more intake questions are required and outline them as structured JSON."
+
+    response = await call_agent(
+        user_prompt=prompt,
+        system_prompt=system_prompt,
+        model=GEMINI_MODEL_LITE,
+        temperature=0.15
+    )
+
+    try:
+        json_start = response.find('{')
+        json_end = response.rfind('}') + 1
+        if json_start >= 0 and json_end > json_start:
+            plan = json.loads(response[json_start:json_end])
+            return plan
+        raise ValueError("Clinical intake JSON not found")
+    except Exception as exc:
+        logger.error(f"[GEMINI INTAKE] Triage parsing failed: {exc}")
+        return {
+            "needs_additional_info": False,
+            "decision_reason": "Fallback: proceeding without intake",
+            "max_rounds": safe_rounds,
+            "questions": [],
+            "initial_hypotheses": []
+        }
+
+
+def gemini_clinical_intake_triage(
+    query: str,
+    history_context: str,
+    max_rounds: int = 5
+) -> dict:
+    """Wrapper for synchronous clinical intake triage"""
+    if not MCP_AVAILABLE:
+        logger.warning("[GEMINI INTAKE] MCP unavailable, skipping clinical intake triage")
+        return {
+            "needs_additional_info": False,
+            "decision_reason": "MCP unavailable",
+            "max_rounds": max_rounds,
+            "questions": [],
+            "initial_hypotheses": []
+        }
+
+    try:
+        loop = asyncio.get_event_loop()
+        if loop.is_running():
+            if nest_asyncio:
+                return nest_asyncio.run(
+                    gemini_clinical_intake_triage_async(query, history_context, max_rounds)
+                )
+            raise RuntimeError("nest_asyncio not available")
+        return loop.run_until_complete(
+            gemini_clinical_intake_triage_async(query, history_context, max_rounds)
+        )
+    except Exception as exc:
+        logger.error(f"[GEMINI INTAKE] Triage request failed: {exc}")
+        return {
+            "needs_additional_info": False,
+            "decision_reason": "Triage agent error",
+            "max_rounds": max_rounds,
+            "questions": [],
+            "initial_hypotheses": []
+        }
+
+
+async def gemini_summarize_clinical_insights_async(
+    query: str,
+    qa_pairs: list
+) -> dict:
+    """Gemini Intake Agent: Convert answered intake questions into key clinical insights"""
+    qa_json = json.dumps(qa_pairs[:8])  # guard against very long history
+    prompt = f"""You are a clinical documentation expert.
+Summarize the following intake Q&A into key insights for a supervising medical agent.
+
+Original patient query:
+"{query}"
+
+Collected intake Q&A (JSON):
+{qa_json}
+
+Return ONLY valid JSON:
+{{
+  "patient_profile": "1-2 sentence overview combining key demographics/symptoms",
+  "refined_problem_statement": "what problem the supervisor should solve now",
+  "key_findings": [
+    {{
+      "title": "short label",
+      "detail": "what the patient reported",
+      "clinical_implication": "why it matters"
+    }}
+  ],
+  "handoff_note": "action-oriented instruction for the supervisor (<=2 sentences)"
+}}
+
+Guidelines:
+- Highlight red flags, chronic meds, relevant history, and symptom trajectory.
+- Only include facts explicitly stated in the Q&A."""
+
+    system_prompt = "You transform clinical intake dialogs into structured insights for downstream medical reasoning."
+
+    response = await call_agent(
+        user_prompt=prompt,
+        system_prompt=system_prompt,
+        model=GEMINI_MODEL_LITE,
+        temperature=0.2
+    )
+
+    try:
+        json_start = response.find('{')
+        json_end = response.rfind('}') + 1
+        if json_start >= 0 and json_end > json_start:
+            return json.loads(response[json_start:json_end])
+        raise ValueError("Clinical insight JSON not found")
+    except Exception as exc:
+        logger.error(f"[GEMINI INTAKE] Insight summarization failed: {exc}")
+        return {
+            "patient_profile": "",
+            "refined_problem_statement": query,
+            "key_findings": [
+                {"title": "Patient concern", "detail": query, "clinical_implication": "Requires standard evaluation"}
+            ],
+            "handoff_note": "Proceed with regular workflow."
+        }
+
+
+def gemini_summarize_clinical_insights(query: str, qa_pairs: list) -> dict:
+    """Wrapper for synchronous clinical insight summarization"""
+    if not MCP_AVAILABLE:
+        logger.warning("[GEMINI INTAKE] MCP unavailable, using fallback intake summary")
+        return {
+            "patient_profile": "",
+            "refined_problem_statement": query,
+            "key_findings": [
+                {"title": "Patient concern", "detail": query, "clinical_implication": "Requires standard evaluation"}
+            ],
+            "handoff_note": "Proceed with regular workflow."
+        }
+
+    try:
+        loop = asyncio.get_event_loop()
+        if loop.is_running():
+            if nest_asyncio:
+                return nest_asyncio.run(
+                    gemini_summarize_clinical_insights_async(query, qa_pairs)
+                )
+            raise RuntimeError("nest_asyncio not available")
+        return loop.run_until_complete(
+            gemini_summarize_clinical_insights_async(query, qa_pairs)
+        )
+    except Exception as exc:
+        logger.error(f"[GEMINI INTAKE] Insight summarization request failed: {exc}")
+        return {
+            "patient_profile": "",
+            "refined_problem_statement": query,
+            "key_findings": [
+                {"title": "Patient concern", "detail": query, "clinical_implication": "Requires standard evaluation"}
+            ],
+            "handoff_note": "Proceed with regular workflow."
+        }
+
+
 def gemini_supervisor_breakdown(query: str, use_rag: bool, use_web_search: bool, time_elapsed: float, max_duration: int = 120) -> dict:
     """Wrapper to obtain supervisor breakdown synchronously"""
     if not MCP_AVAILABLE:

ui.py

@@ -144,6 +144,11 @@ def create_demo():
                             "Show agentic thought",
                             size="sm"
                         )
+                    enable_clinical_intake = gr.Checkbox(
+                        value=True,
+                        label="Enable clinical intake (max 5 Q&A)",
+                        info="Ask focused follow-up questions before breaking down the case"
+                    )
                     agentic_thoughts_box = gr.Textbox(
                         label="Agentic Thoughts",
                         placeholder="Internal thoughts from MedSwin and supervisor will appear here...",
@@ -261,6 +266,7 @@ def create_demo():
                         use_rag,
                         medical_model,
                         use_web_search,
+                        enable_clinical_intake,
                         disable_agentic_reasoning,
                         show_thoughts_state
                     ],
@@ -283,6 +289,7 @@ def create_demo():
                         use_rag,
                         medical_model,
                         use_web_search,
+                        enable_clinical_intake,
                         disable_agentic_reasoning,
                         show_thoughts_state
                     ],