Feat: Implement performance and citation fixes

app.py

@@ -74,15 +74,31 @@ def process_enhanced_medical_query(message: str, history: List[List[str]]) -> st
 
 def format_enhanced_medical_response(response: EnhancedMedicalResponse) -> str:
     """
-    Format the enhanced medical response for display
+    Format the enhanced medical response for display, ensuring citations are always included.
     """
     formatted_parts = []
     
-    # Main response
-    formatted_parts.append(response.answer)
+    # Main response from the LLM
+    # The new prompt instructs the LLM to include markdown citations like [1], [2]
+    # The final response text is now the primary source of the answer.
+    final_response_text = response.answer
+    formatted_parts.append(final_response_text)
+    
+    # Always add the clinical sources section if sources exist
+    if response.sources:
+        formatted_parts.append("\n\n---\n")
+        formatted_parts.append("### 📋 **Clinical Sources**")
+        # Create a numbered list of sources for clarity
+        for i, source in enumerate(response.sources, 1):
+            # Ensure we don't list more sources than were used for citations
+            if f"[{i}]" in final_response_text:
+                 formatted_parts.append(f"{i}. {source}")
+            else:
+                # If the LLM didn't cite this source, we can choose to omit it or list it as an uncited reference
+                pass # For now, only show cited sources to keep the output clean.
     
     # Enhanced information section
-    formatted_parts.append("\n---\n")
+    formatted_parts.append("\n\n---\n")
     formatted_parts.append("### 📊 **Enhanced Medical Analysis**")
     
     # Safety and verification info
@@ -96,17 +112,11 @@ def format_enhanced_medical_response(response: EnhancedMedicalResponse) -> str:
     formatted_parts.append(f"**🧠 Medical Entities Extracted**: {response.medical_entities_count}")
     formatted_parts.append(f"**🎯 Context Adherence**: {response.context_adherence_score:.1%}")
     formatted_parts.append(f"**📚 Sources Used**: {len(response.sources)}")
-    if hasattr(response, 'processing_time'):
-        formatted_parts.append(f"**⚡ Processing Time**: {response.processing_time:.2f}s")
-    
-    # Sources
-    if response.sources:
-        formatted_parts.append("\n### 📋 **Clinical Sources**")
-        for i, source in enumerate(response.sources[:5], 1):  # Show top 5 sources
-            formatted_parts.append(f"{i}. {source}")
+    if hasattr(response, 'query_time'): # Changed from processing_time to match the object attribute
+        formatted_parts.append(f"**⚡ Processing Time**: {response.query_time:.2f}s")
     
     # Medical disclaimer
-    formatted_parts.append("\n---")
+    formatted_parts.append("\n---\n")
     formatted_parts.append("*This information is for clinical reference based on Sri Lankan guidelines and does not replace professional medical judgment.*")
     
     return "\n".join(formatted_parts)

src/enhanced_groq_medical_rag.py

@@ -363,49 +363,30 @@ class EnhancedGroqMedicalRAG:
             query_analysis = self.analyze_medical_query(query)
             self._stop_timer("Query Analysis")
             
-            # Step 2: Multi-stage comprehensive retrieval
-            all_documents = []
-            seen_content = set()
-            
-            # Stage 2a: Original query retrieval (increased from 15 to 25)
-            stage1_docs = self.vector_store.search(query=query, k=25)
-            for doc in stage1_docs:
-                if doc.content not in seen_content:
-                    all_documents.append(doc)
-                    seen_content.add(doc.content)
-            
-            # Stage 2b: Expanded query retrieval
-            for expanded_query in query_analysis['expanded_queries']:
-                expanded_docs = self.vector_store.search(expanded_query, k=15)
-                for doc in expanded_docs:
-                    if doc.content not in seen_content and len(all_documents) < 50:
-                        all_documents.append(doc)
-                        seen_content.add(doc.content)
-            
-            # Stage 2c: Entity-specific retrieval
-            for entity in query_analysis['medical_entities']:
-                entity_docs = self.vector_store.search(entity, k=10)
-                for doc in entity_docs:
-                    if doc.content not in seen_content and len(all_documents) < 60:
-                        all_documents.append(doc)
-                        seen_content.add(doc.content)
-            
+            # Step 2: Simplified single-stage retrieval
+            self._start_timer("Single Stage Retrieval")
+            NUM_CANDIDATE_DOCS = 40
+            all_documents = self.vector_store.search(query=query_analysis['original_query'], k=NUM_CANDIDATE_DOCS)
+            self._stop_timer("Single Stage Retrieval")
+
             if not all_documents:
                 return self._create_no_results_response(query, self._stop_timer("Total Query Time"))
             
             # Step 3: Advanced multi-criteria re-ranking
+            self._start_timer("Re-ranking")
             reranked_docs = self._advanced_medical_reranking(query_analysis, all_documents)
+            self._stop_timer("Re-ranking")
             
-            # Step 4: Select an initial set of documents, respecting the user's preference for more context.
-            initial_doc_count = 10
-            final_docs = reranked_docs[:initial_doc_count]
+            # Step 4: Select the final documents to be used for context
+            FINAL_DOC_COUNT = 10
+            final_docs = reranked_docs[:FINAL_DOC_COUNT]
             
             # Step 5: Verify coverage and add missing context if needed, up to a hard limit to avoid API errors.
             MAX_FINAL_DOCS = 12
             coverage_score = self._verify_query_coverage(query_analysis, final_docs)
             if coverage_score < 0.7:  # Less than 70% coverage
                 self.logger.info(f"⚠️ Low coverage score ({coverage_score:.1%}). Retrieving additional context...")
-                additional_docs = self._retrieve_missing_context(query_analysis, final_docs, seen_content)
+                additional_docs = self._retrieve_missing_context(query_analysis, final_docs, set()) # Pass an empty set for seen_content
                 remaining_capacity = MAX_FINAL_DOCS - len(final_docs)
                 if remaining_capacity > 0:
                     final_docs.extend(additional_docs[:remaining_capacity])