Update src/model.py

src/model.py

@@ -119,7 +119,7 @@ def extract_and_link_sources(text, evidence_snippets):
     
     # Process DOI citations
     for doi in doi_matches:
-        for snippet in evidence_snippets:
+            for snippet in evidence_snippets:
             # Check if this is a direct DOI match
             if 'doi' in snippet and snippet['doi'] == doi:
                 source_map[f"DOI:{doi}"] = {
@@ -133,13 +133,13 @@ def extract_and_link_sources(text, evidence_snippets):
             # Also check the ID field which might contain DOI
             elif snippet.get("id") == f"DOI:{doi}":
                 source_map[f"DOI:{doi}"] = {
-                    "id": snippet["id"],
-                    "title": snippet["title"].strip(),
-                    "url": snippet["url"],
+                        "id": snippet["id"],
+                        "title": snippet["title"].strip(),
+                        "url": snippet["url"],
                     "citation": snippet["citation"],
                     "doi": doi
-                }
-                break
+                    }
+                    break
     
     # Process other citation formats for backward compatibility
     for source_id_match in source_matches:
@@ -147,15 +147,15 @@ def extract_and_link_sources(text, evidence_snippets):
             for snippet in evidence_snippets:
                 if source_id_match == snippet["id"]:
                     source_map[source_id_match] = {
-                        "id": snippet["id"],
-                        "title": snippet["title"].strip(),
-                        "url": snippet["url"],
+                    "id": snippet["id"],
+                    "title": snippet["title"].strip(),
+                    "url": snippet["url"],
                         "citation": snippet["citation"],
                         "pmid": snippet.get("pmid", ""),
                         "doi": snippet.get("doi", "")
-                    }
+                }
                     break
-    
+
     # Replace PMID citations with links
     linked_text = text
     for pmid_key in [f"PMID:{pmid}" for pmid in pmid_matches]:
@@ -186,10 +186,10 @@ def extract_and_link_sources(text, evidence_snippets):
     # Replace other citation formats
     for source_id_key, source_data in source_map.items():
         if not (source_id_key.startswith("PMID:") or source_id_key.startswith("DOI:")):
-            safe_id = re.escape(source_id_key)
-            pattern = f"\\[{safe_id}\\]"
-            replacement = f"[{source_data['title']}]({source_data['url']})"
-            linked_text = re.sub(pattern, replacement, linked_text)
+        safe_id = re.escape(source_id_key)
+        pattern = f"\\[{safe_id}\\]"
+        replacement = f"[{source_data['title']}]({source_data['url']})"
+        linked_text = re.sub(pattern, replacement, linked_text)
     
     # Handle generic [source_id] placeholder
     if "source_id" in source_matches:
@@ -206,7 +206,7 @@ def extract_and_link_sources(text, evidence_snippets):
                     "doi": snippet.get("doi", "")
                 }
                 replacement = f"[{snippet['title']}]({snippet['url']})"
-                linked_text = re.sub(r'\[source_id\]', replacement, linked_text)
+        linked_text = re.sub(r'\[source_id\]', replacement, linked_text)
     
     # Final fallback for any remaining placeholders
     linked_text = re.sub(r'\[source_id\]', "[Medical Reference]", linked_text)
@@ -1283,6 +1283,9 @@ def orchestrator_chat(history, query, use_rag, is_follow_up=False):
         system = {"role": "system", "content": FOLLOW_UP_PROMPT}
     else:
         system = {"role": "system", "content": SYSTEM_PROMPT}
+        
+    # Debug - Print prompt type
+    print(f"Using {'FOLLOW_UP_PROMPT' if is_follow_up else 'SYSTEM_PROMPT'} with query: {query}")
 
     msgs = [system] + history
 
@@ -1363,7 +1366,7 @@ def orchestrator_chat(history, query, use_rag, is_follow_up=False):
            • [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:**
         -   **Follow-up Questions**: Specific questions to gather additional information.
         -   **Reasoning**: Bullet points detailing your clinical reasoning.
@@ -1432,11 +1435,14 @@ def orchestrator_chat(history, query, use_rag, is_follow_up=False):
                 follow_up_questions = "\n".join([f"{i+1}. {q}" for i, q in enumerate(questions) if q])
             else:
                 follow_up_questions = questions
+            
+    # Debug: Print follow-up questions
+    print(f"Follow-up questions generated: {follow_up_questions}")
     else:
         # If RAG is disabled, just parse the response without source processing
         parsed_response = parse_doctor_response(response)
         main_response = response
-        
+
         # Extract reasoning
         reasoning = parsed_response.get("reasoning", [])
         if reasoning:
@@ -1453,6 +1459,9 @@ def orchestrator_chat(history, query, use_rag, is_follow_up=False):
                 follow_up_questions = "\n".join([f"{i+1}. {q}" for i, q in enumerate(questions) if q])
             else:
                 follow_up_questions = questions
+            
+        # Debug: Print follow-up questions
+        print(f"Follow-up questions generated: {follow_up_questions}")
     
     # Return four values: main response, explanation, follow-up questions, and evidence
     return main_response, explanation, follow_up_questions, evidence_snippets