Update src/model.py

src/model.py

@@ -1196,7 +1196,8 @@ def parse_doctor_response(response_text):
         "diagnosis": "",
         "treatment": "",
         "reasoning": [],
-        "sources": []
+        "sources": [],
+        "follow_up_questions": []
     }
 
     # Try to extract diagnosis
@@ -1210,7 +1211,7 @@ def parse_doctor_response(response_text):
         parsed["treatment"] = treatment_match.group(2).strip()
 
     # Try to extract reasoning if present
-    reasoning_match = re.search(r'(?i)reasoning:?\s*(.*?)(?:\n\n\Z|\Z)', response_text, re.DOTALL)
+    reasoning_match = re.search(r'(?i)(?:#{1,3}\s*)?reasoning:?\s*(.*?)(?:(?:#{1,3}\s*(?:follow[ -]?up questions|sources))|(?:\n\n\Z)|\Z)', response_text, re.DOTALL)
     if reasoning_match:
         reasoning_text = reasoning_match.group(1).strip()
         # Split into bullet points if present
@@ -1221,11 +1222,28 @@ def parse_doctor_response(response_text):
                 parsed["reasoning"][0] = parsed["reasoning"][0].lstrip('- ')
         else:
             parsed["reasoning"] = [reasoning_text]
+    
+    # Extract follow-up questions
+    follow_up_match = re.search(r'(?i)(?:#{1,3}\s*)?follow[ -]?up questions:?\s*(.*?)(?:(?:#{1,3}\s*(?:reasoning|sources))|(?:\n\n\Z)|\Z)', response_text, re.DOTALL)
+    if follow_up_match:
+        follow_up_text = follow_up_match.group(1).strip()
+        # Extract numbered questions
+        questions = re.findall(r'(?:^|\n)\s*\d+\.\s*(.*?)(?=\n\s*\d+\.|\Z)', follow_up_text, re.DOTALL)
+        if questions:
+            parsed["follow_up_questions"] = [q.strip() for q in questions if q.strip()]
+        else:
+            # Try to extract bullet points if no numbered questions
+            questions = re.findall(r'(?:^|\n)\s*[-•*]\s*(.*?)(?=\n\s*[-•*]|\Z)', follow_up_text, re.DOTALL)
+            if questions:
+                parsed["follow_up_questions"] = [q.strip() for q in questions if q.strip()]
+            elif follow_up_text.strip():
+                # If no formatting but there's text, use the whole text
+                parsed["follow_up_questions"] = [follow_up_text.strip()]
 
     # Extract sources/references
-    sources_match = re.search(r'(?i)(sources|references):?\s*(.*?)(?:\n\n\Z|\Z)', response_text, re.DOTALL)
+    sources_match = re.search(r'(?i)(?:#{1,3}\s*)?(?:sources|references):?\s*(.*?)(?:\n\n\Z|\Z)', response_text, re.DOTALL)
     if sources_match:
-        sources_text = sources_match.group(2).strip()
+        sources_text = sources_match.group(1).strip()
         # Split into individual sources
         if '\n' in sources_text:
             parsed["sources"] = [item.strip() for item in sources_text.split('\n') if item.strip()]
@@ -1323,41 +1341,53 @@ def orchestrator_chat(history, query, use_rag, is_follow_up=False):
     # Add instructions for structured output
     if use_rag:
         output_instructions = """
-        Please structure your response clearly.
-        **Priority 1: Ask Clarifying Questions**
-        If the user's query lacks detail for a proper assessment (e.g., age, specific symptoms, medical history, duration, severity), your HIGHEST priority is to ask these questions first. Do not provide a diagnosis or plan until sufficient information is gathered.
-        **Priority 2: Main Response (After Clarification)**
-        Once sufficient information is available (either initially or after asking questions), provide:
-        1. A direct answer to the patient's concerns.
-        2. If appropriate, a clear diagnosis or differential diagnosis with likelihood assessments.
-        3. Recommendations for a treatment plan or next steps.
-        4. IMPORTANT: You MUST cite between 2-3 different medical evidence sources using either:
+        Please structure your response as follows:
+
+        **Main Response**
+        Always provide a direct answer to the doctor's query first, even if you don't have all the information you'd like. Include:
+        1. Your best answer based on the information already provided
+        2. If appropriate, a differential diagnosis with likelihood assessments
+        3. Recommendations for treatment or next steps
+        4. Cite 2-3 different medical evidence sources using:
            • [PMID:123456] format for PubMed articles
            • [DOI:10.xxxx/yyyy] format for Europe PMC articles without PMID
-           
-           Use no more than 3 sources and no fewer than 2 sources.
-        **After your main response, ALWAYS include these sections:**
-        -   **Reasoning**: Bullet points detailing your clinical reasoning.
-        -   **Sources**: A list of all references cited in your main response (2-3 sources), formatted as:
-             - PMID: 12345678 - Author et al. (Year). Title. Journal.
-               URL: https://pubmed.ncbi.nlm.nih.gov/12345678/
-             - DOI: 10.xxxx/yyyy - Author et al. (Year). Title. Journal.
-               URL: https://doi.org/10.xxxx/yyyy
-        IMPORTANT: Only cite sources that were provided in the evidence. Do not fabricate references, PMIDs, or DOIs.
+
+        **Follow-up Questions**
+        IMPORTANT: Place all clarifying questions in this separate section, not in your main response.
+        List 2-4 specific questions that would help you provide a more precise assessment.
+        Format these as a numbered list.
+
+        **Reasoning**
+        Provide bullet points explaining your clinical reasoning.
+
+        **Sources**
+        List all references cited (2-3 sources), formatted as:
+        - PMID: 12345678 - Author et al. (Year). Title. Journal.
+          URL: https://pubmed.ncbi.nlm.nih.gov/12345678/
+        - DOI: 10.xxxx/yyyy - Author et al. (Year). Title. Journal.
+          URL: https://doi.org/10.xxxx/yyyy
+
+        IMPORTANT: Only cite sources provided in the evidence. Do not fabricate references.
         """
     else:
         # Different instructions when RAG is disabled - no mention of sources or citations
         output_instructions = """
-        Please structure your response clearly.
-        **Priority 1: Ask Clarifying Questions**
-        If the user's query lacks detail for a proper assessment (e.g., age, specific symptoms, medical history, duration, severity), your HIGHEST priority is to ask these questions first. Do not provide a diagnosis or plan until sufficient information is gathered.
-        **Priority 2: Main Response (After Clarification)**
-        Once sufficient information is available (either initially or after asking questions), provide:
-        1.  A direct answer to the patient's concerns.
-        2.  If appropriate, a clear diagnosis or differential diagnosis.
-        3.  Recommendations for a treatment plan or next steps.
-        **After your main response, ALWAYS include this section:**
-        -   **Reasoning**: Bullet points detailing your clinical reasoning.
+        Please structure your response as follows:
+
+        **Main Response**
+        Always provide a direct answer to the doctor's query first, even if you don't have all the information you'd like. Include:
+        1. Your best answer based on the information already provided
+        2. If appropriate, a differential diagnosis
+        3. Recommendations for treatment or next steps
+
+        **Follow-up Questions**
+        IMPORTANT: Place all clarifying questions in this separate section, not in your main response.
+        List 2-4 specific questions that would help you provide a more precise assessment.
+        Format these as a numbered list.
+
+        **Reasoning**
+        Provide bullet points explaining your clinical reasoning.
+
         IMPORTANT: Since database search is disabled, do not include citations or sources in your response.
         """
 
@@ -1370,6 +1400,7 @@ def orchestrator_chat(history, query, use_rag, is_follow_up=False):
     # Extract and process sources
     explanation = None
     evidence = None
+    follow_up_questions = []
     
     if use_rag:
         # Process the response to replace source placeholders with actual links
@@ -1388,6 +1419,9 @@ def orchestrator_chat(history, query, use_rag, is_follow_up=False):
                 explanation = "\n".join([f"- {r}" for r in reasoning])
             else:
                 explanation = reasoning
+                
+        # Extract follow-up questions
+        follow_up_questions = parsed_response.get("follow_up_questions", [])
     else:
         # If RAG is disabled, just parse the response without source processing
         parsed_response = parse_doctor_response(response)
@@ -1400,9 +1434,12 @@ def orchestrator_chat(history, query, use_rag, is_follow_up=False):
                 explanation = "\n".join([f"- {r}" for r in reasoning])
             else:
                 explanation = reasoning
+                
+        # Extract follow-up questions
+        follow_up_questions = parsed_response.get("follow_up_questions", [])
     
     # Return the response (with links if RAG is enabled), explanation, and evidence
-    return main_response, explanation, evidence_snippets
+    return main_response, explanation, follow_up_questions, evidence_snippets
 
 # Enhanced interactive loop with better handling of consultations
 def run_consultation(use_rag=True):
@@ -1428,7 +1465,7 @@ def run_consultation(use_rag=True):
             print("\nSearching medical databases...")
             
         # Process query
-        reply, explanation, evidence = orchestrator_chat(history, query, use_rag, is_follow_up)
+        reply, explanation, follow_up_questions, evidence = orchestrator_chat(history, query, use_rag, is_follow_up)
 
         # Display the AI response
         print("\n" + "=" * 30)