Enhance Q&A breakdown agent

pipeline.py

@@ -75,6 +75,21 @@ def _format_intake_question(question: dict, round_idx: int, max_rounds: int, tar
     return prompt_text
 
 
+def _format_qa_transcript(qa_pairs: list) -> str:
+    if not qa_pairs:
+        return ""
+    lines = []
+    for idx, qa in enumerate(qa_pairs, 1):
+        question = qa.get("question", "").strip()
+        answer = qa.get("answer", "").strip()
+        if question:
+            lines.append(f"Q{idx}: {question}")
+        if answer:
+            lines.append(f"A{idx}: {answer}")
+        lines.append("")
+    return "\n".join(lines).strip()
+
+
 def _format_insights_block(insights: dict) -> str:
     if not insights:
         return ""
@@ -160,13 +175,15 @@ def _handle_clinical_answer(session_id: str, answer_text: str):
         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)
+        transcript = _format_qa_transcript(state["answers"])
         _clear_clinical_intake_state(session_id)
         return {
             "type": "insights",
             "insights": insights,
             "insights_block": insights_block,
             "refined_query": refined_query,
-            "qa_pairs": state["answers"]
+            "qa_pairs": state["answers"],
+            "qa_transcript": transcript
         }
     state["pending_question_index"] = next_index
     state["current_round"] = len(state["answers"]) + 1
@@ -240,7 +257,11 @@ def stream_chat(
             "activated": False,
             "rounds": 0,
             "reason": "",
-            "insights": []
+            "insights": [],
+            "plan": [],
+            "qa_pairs": [],
+            "transcript": "",
+            "insights_block": ""
         }
     }
     def record_stage(stage_name: str, start_time: float):
@@ -262,6 +283,8 @@ def stream_chat(
         {"role": "assistant", "content": ""}
     ]
     
+    clinical_intake_context_block = ""
+    
     if not enable_clinical_intake:
         _clear_clinical_intake_state(user_id)
     else:
@@ -281,11 +304,24 @@ def stream_chat(
                 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)
+                pipeline_diagnostics["clinical_intake"]["qa_pairs"] = intake_result.get("qa_pairs", [])
+                pipeline_diagnostics["clinical_intake"]["transcript"] = intake_result.get("qa_transcript", "")
+                pipeline_diagnostics["clinical_intake"]["insights_block"] = intake_result.get("insights_block", "")
+                base_refined = intake_result.get("refined_query", message)
+                summary_section = ""
+                transcript_section = ""
+                if intake_result.get("insights_block"):
+                    summary_section = f"Clinical intake summary:\n{intake_result['insights_block']}"
+                if intake_result.get("qa_transcript"):
+                    transcript_section = f"Clinical intake Q&A transcript:\n{intake_result['qa_transcript']}"
+                sections = [base_refined, summary_section, transcript_section]
+                message = "\n\n---\n\n".join([section for section in sections if section])
+                clinical_intake_context_block = "\n\n".join([seg for seg in [summary_section, transcript_section] if seg])
         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", "")
+            pipeline_diagnostics["clinical_intake"]["plan"] = triage_plan.get("questions", [])
             needs_intake = triage_plan.get("needs_additional_info") and triage_plan.get("questions")
             if needs_intake:
                 first_prompt = _start_clinical_intake_session(
@@ -460,13 +496,14 @@ def stream_chat(
     
     base_system_prompt = system_prompt if system_prompt else "As a medical specialist, provide clinical and concise answers. Use Markdown format with bullet points. Do not use tables."
     
-    combined_context = ""
+    context_sections = []
+    if clinical_intake_context_block:
+        context_sections.append("Clinical Intake Context:\n" + clinical_intake_context_block)
     if rag_contexts:
-        combined_context += "Document Context:\n" + "\n\n".join(rag_contexts[:4])
+        context_sections.append("Document Context:\n" + "\n\n".join(rag_contexts[:4]))
     if search_contexts:
-        if combined_context:
-            combined_context += "\n\n"
-        combined_context += "Web Search Context:\n" + "\n\n".join(search_contexts)
+        context_sections.append("Web Search Context:\n" + "\n\n".join(search_contexts))
+    combined_context = "\n\n".join(context_sections)
     
     logger.info(f"[MEDSWIN] Executing {len(breakdown.get('sub_topics', []))} tasks sequentially...")
     medswin_answers = []

supervisor.py

@@ -160,6 +160,38 @@ Keep strategies focused and avoid overlap."""
         }
 
 
+def _prepare_clinical_question_plan(plan: dict, safe_rounds: int) -> dict:
+    """Normalize Gemini question plan to 1-5 sequential prompts."""
+    if not isinstance(plan, dict):
+        return {"questions": []}
+    questions = plan.get("questions", [])
+    if not isinstance(questions, list):
+        questions = []
+    cleaned = []
+    for idx, raw in enumerate(questions):
+        if not isinstance(raw, dict):
+            continue
+        question_text = (raw.get("question") or "").strip()
+        if not question_text:
+            continue
+        entry = dict(raw)
+        entry["question"] = question_text
+        entry["order"] = entry.get("order") or raw.get("id") or (idx + 1)
+        cleaned.append(entry)
+    cleaned.sort(key=lambda item: item.get("order", 0))
+    cleaned = cleaned[:max(1, min(5, safe_rounds))]
+    for idx, item in enumerate(cleaned, 1):
+        item["order"] = idx
+    plan["questions"] = cleaned
+    if cleaned:
+        plan["max_rounds"] = min(len(cleaned), safe_rounds)
+        plan["needs_additional_info"] = bool(plan.get("needs_additional_info", True))
+    else:
+        plan["needs_additional_info"] = False
+        plan["max_rounds"] = 0
+    return plan
+
+
 async def gemini_supervisor_rag_brainstorm_async(query: str, retrieved_docs: str, time_elapsed: float) -> dict:
     """Gemini Supervisor: In RAG mode, brainstorm retrieved documents into 1-4 short contexts"""
     max_doc_length = 3000
@@ -274,6 +306,7 @@ Guidelines:
         json_end = response.rfind('}') + 1
         if json_start >= 0 and json_end > json_start:
             plan = json.loads(response[json_start:json_end])
+            plan = _prepare_clinical_question_plan(plan, safe_rounds)
             return plan
         raise ValueError("Clinical intake JSON not found")
     except Exception as exc: