Update rag_engine.py

rag_engine.py

@@ -1,53 +1,41 @@
 """
 RAG Query Engine for Lab Report Decoder
-Uses Hugging Face models for embeddings and generation
+Uses Hugging Face models - OPTIMIZED for speed
 """
 
 from sentence_transformers import SentenceTransformer
-from transformers import pipeline, AutoTokenizer, AutoModelForCausalLM
+from transformers import pipeline
 import chromadb
-from chromadb.config import Settings
 from typing import List, Dict
 from pdf_extractor import LabResult
-import torch
+import os
 
 class LabReportRAG:
-    """RAG system for explaining lab results using Hugging Face models"""
+    """RAG system for explaining lab results - Fast and efficient"""
     
     def __init__(self, db_path: str = "./chroma_db"):
-        """Initialize the RAG system with Hugging Face models"""
+        """Initialize the RAG system with fast models"""
         
-        print("🔄 Loading Hugging Face models...")
+        print("🔄 Loading models (optimized for speed)...")
         
-        # Use smaller, faster models for embeddings
+        # Fast embedding model
         self.embedding_model = SentenceTransformer('all-MiniLM-L6-v2')
+        print("✅ Embeddings loaded")
         
-        # Use a medical-focused or general LLM
-        # Options: 
-        # - "microsoft/Phi-3-mini-4k-instruct" (good balance)
-        # - "google/flan-t5-base" (lighter)
-        # - "meta-llama/Llama-2-7b-chat-hf" (requires auth)
-        
-        model_name = "microsoft/Phi-3-mini-4k-instruct"
-        
+        # Use FAST text generation model
+        print("🔄 Loading text generation model...")
         try:
-            self.tokenizer = AutoTokenizer.from_pretrained(model_name, trust_remote_code=True)
-            self.llm = AutoModelForCausalLM.from_pretrained(
-                model_name,
-                trust_remote_code=True,
-                torch_dtype=torch.float16 if torch.cuda.is_available() else torch.float32,
-                device_map="auto" if torch.cuda.is_available() else None
-            )
-            print(f"✅ Loaded model: {model_name}")
-        except Exception as e:
-            print(f"⚠️ Could not load {model_name}, falling back to simpler model")
-            # Fallback to lighter model
+            # Use Flan-T5 - much faster than Phi-3
             self.text_generator = pipeline(
-                "text-generation",
-                model="google/flan-t5-base",
-                max_length=512
+                "text2text-generation",
+                model="google/flan-t5-small",  # Even smaller/faster
+                max_length=256,
+                device=-1  # Force CPU (HF Spaces default)
             )
-            self.llm = None
+            print("✅ Text generation model loaded (Flan-T5-small)")
+        except Exception as e:
+            print(f"⚠️ Model loading error: {e}")
+            self.text_generator = None
         
         # Load vector store
         try:
@@ -55,49 +43,13 @@ class LabReportRAG:
             self.collection = self.client.get_collection("lab_reports")
             print("✅ Vector database loaded")
         except Exception as e:
-            print(f"⚠️ No vector database found. Please run build_vector_db.py first.")
+            print(f"⚠️ Vector database not found: {e}")
             self.collection = None
     
-    def _generate_with_phi(self, prompt: str, max_tokens: int = 512) -> str:
-        """Generate text using Phi-3 model"""
-        inputs = self.tokenizer(prompt, return_tensors="pt", truncation=True, max_length=2048)
-        
-        if torch.cuda.is_available():
-            inputs = {k: v.to('cuda') for k, v in inputs.items()}
-        
-        outputs = self.llm.generate(
-            **inputs,
-            max_new_tokens=max_tokens,
-            temperature=0.7,
-            do_sample=True,
-            top_p=0.9
-        )
-        
-        response = self.tokenizer.decode(outputs[0], skip_special_tokens=True)
-        # Remove the prompt from response
-        response = response.replace(prompt, "").strip()
-        return response
-    
-    def _generate_with_fallback(self, prompt: str) -> str:
-        """Generate text using fallback pipeline"""
-        result = self.text_generator(prompt, max_length=512, num_return_sequences=1)
-        return result[0]['generated_text']
-    
-    def _generate_text(self, prompt: str) -> str:
-        """Generate text using available model"""
-        try:
-            if self.llm is not None:
-                return self._generate_with_phi(prompt)
-            else:
-                return self._generate_with_fallback(prompt)
-        except Exception as e:
-            print(f"Generation error: {e}")
-            return "Sorry, I encountered an error generating the explanation."
-    
-    def _retrieve_context(self, query: str, k: int = 3) -> str:
+    def _retrieve_context(self, query: str, k: int = 2) -> str:
         """Retrieve relevant context from vector database"""
         if self.collection is None:
-            return "No medical reference data available."
+            return "Limited medical information available."
         
         try:
             # Create query embedding
@@ -110,144 +62,233 @@ class LabReportRAG:
             )
             
             # Combine documents
-            if results and results['documents']:
-                context = "\n\n".join(results['documents'][0])
-                return context
+            if results and results['documents'] and len(results['documents'][0]) > 0:
+                context = "\n".join(results['documents'][0])
+                # Limit context length for speed
+                return context[:1000]
             else:
-                return "No relevant information found."
+                return "No specific information found."
         except Exception as e:
             print(f"Retrieval error: {e}")
-            return "Error retrieving medical information."
+            return "Error retrieving information."
+    
+    def _generate_text(self, prompt: str) -> str:
+        """Generate text - with fallback to template-based"""
+        if self.text_generator is None:
+            return "AI model not available. Using basic explanation."
+        
+        try:
+            # Generate with timeout protection
+            result = self.text_generator(
+                prompt,
+                max_length=256,
+                do_sample=True,
+                temperature=0.7,
+                num_return_sequences=1
+            )
+            return result[0]['generated_text'].strip()
+        except Exception as e:
+            print(f"Generation error: {e}")
+            return "Unable to generate detailed explanation."
     
     def explain_result(self, result: LabResult) -> str:
         """Generate explanation for a single lab result"""
         
-        # Retrieve relevant context
-        query = f"{result.test_name} {result.status} meaning causes treatment"
-        context = self._retrieve_context(query, k=3)
+        print(f"  Explaining: {result.test_name} ({result.status})...")
         
-        # Create prompt
-        prompt = f"""You are a helpful medical assistant. Explain this lab result in simple terms.
-
-Medical Information:
-{context}
-
-Lab Test: {result.test_name}
-Value: {result.value} {result.unit}
-Reference Range: {result.reference_range}
-Status: {result.status}
+        # Quick template-based explanation for speed
+        if result.status == 'normal':
+            return self._explain_normal(result)
+        elif result.status == 'high':
+            return self._explain_high(result)
+        elif result.status == 'low':
+            return self._explain_low(result)
+        else:
+            return self._explain_unknown(result)
+    
+    def _explain_normal(self, result: LabResult) -> str:
+        """Fast template for normal results"""
+        context = self._retrieve_context(f"{result.test_name} normal meaning", k=1)
+        
+        explanation = f"""✅ Your {result.test_name} level of {result.value} {result.unit} is within the normal range ({result.reference_range}).
 
-Please explain:
-1. What this test measures
-2. What this result means
-3. Possible causes if abnormal
-4. Dietary recommendations if applicable
+This indicates healthy levels. """
+        
+        if context and len(context) > 20:
+            # Add context if available
+            explanation += f"\n\n{context[:300]}"
+        
+        return explanation
+    
+    def _explain_high(self, result: LabResult) -> str:
+        """Fast template for high results"""
+        context = self._retrieve_context(f"{result.test_name} high causes treatment", k=2)
+        
+        explanation = f"""⚠️ Your {result.test_name} level of {result.value} {result.unit} is ABOVE the normal range ({result.reference_range}).
 
-Keep it simple and clear. Answer:"""
+"""
+        
+        if context and len(context) > 20:
+            explanation += f"{context[:400]}\n\n"
+        
+        explanation += "💡 Recommendation: Discuss these results with your healthcare provider for personalized advice."
+        
+        return explanation
+    
+    def _explain_low(self, result: LabResult) -> str:
+        """Fast template for low results"""
+        context = self._retrieve_context(f"{result.test_name} low causes treatment", k=2)
+        
+        explanation = f"""⚠️ Your {result.test_name} level of {result.value} {result.unit} is BELOW the normal range ({result.reference_range}).
 
-        # Generate explanation
-        explanation = self._generate_text(prompt)
+"""
+        
+        if context and len(context) > 20:
+            explanation += f"{context[:400]}\n\n"
+        
+        explanation += "💡 Recommendation: Consult with your healthcare provider about these results."
         
         return explanation
     
+    def _explain_unknown(self, result: LabResult) -> str:
+        """Template for unknown status"""
+        return f"""Your {result.test_name} result is {result.value} {result.unit}.
+
+Reference range: {result.reference_range}
+
+We couldn't automatically determine if this is within normal range. Please consult your healthcare provider to interpret this result."""
+    
     def explain_all_results(self, results: List[LabResult]) -> Dict[str, str]:
-        """Generate explanations for all lab results"""
+        """Generate explanations for all lab results - FAST"""
         explanations = {}
         
-        for result in results:
-            print(f"Explaining {result.test_name}...")
-            explanation = self.explain_result(result)
-            explanations[result.test_name] = explanation
+        print(f"🧠 Generating explanations for {len(results)} results...")
+        
+        for i, result in enumerate(results, 1):
+            print(f"  [{i}/{len(results)}] {result.test_name}...")
+            try:
+                explanation = self.explain_result(result)
+                explanations[result.test_name] = explanation
+            except Exception as e:
+                print(f"    Error: {e}")
+                explanations[result.test_name] = f"Unable to generate explanation for {result.test_name}."
         
+        print("✅ All explanations generated")
         return explanations
     
     def answer_followup_question(self, question: str, lab_results: List[LabResult]) -> str:
-        """Answer follow-up questions about lab results"""
+        """Answer follow-up questions - FAST"""
+        
+        print(f"💬 Processing question: {question[:50]}...")
         
         # Create context from lab results
-        results_context = "\n".join([
-            f"{r.test_name}: {r.value} {r.unit} (Status: {r.status}, Range: {r.reference_range})"
-            for r in lab_results
-        ])
+        results_summary = []
+        for r in lab_results[:10]:  # Limit to first 10 for speed
+            results_summary.append(
+                f"{r.test_name}: {r.value} {r.unit} ({r.status})"
+            )
+        results_context = "\n".join(results_summary)
         
-        # Retrieve relevant medical information
-        medical_context = self._retrieve_context(question, k=3)
+        # Get relevant medical info
+        medical_context = self._retrieve_context(question, k=2)
         
-        # Create prompt
-        prompt = f"""You are a medical assistant. Answer this question based on the patient's lab results and medical information.
-
-Patient's Lab Results:
-{results_context}
-
-Medical Information:
-{medical_context}
-
-Question: {question}
-
-Provide a clear, helpful answer. Answer:"""
+        # Simple template-based response for speed
+        if "food" in question.lower() or "eat" in question.lower() or "diet" in question.lower():
+            answer = f"""Based on your lab results:\n\n{results_context}\n\n"""
+            if medical_context and len(medical_context) > 20:
+                answer += f"{medical_context[:500]}"
+            else:
+                answer += "For dietary recommendations specific to your results, please consult with a healthcare provider or nutritionist."
         
-        # Generate answer
-        answer = self._generate_text(prompt)
+        elif "why" in question.lower() or "cause" in question.lower():
+            answer = f"""Regarding your question about your results:\n\n"""
+            if medical_context and len(medical_context) > 20:
+                answer += f"{medical_context[:500]}"
+            else:
+                answer += "There can be various causes for abnormal lab results. Your healthcare provider can help identify the specific cause in your case."
+        
+        else:
+            # General question
+            if medical_context and len(medical_context) > 20:
+                answer = medical_context[:500]
+            else:
+                answer = f"""Based on your results:\n{results_context}\n\nFor specific medical advice about your results, please consult with your healthcare provider."""
         
+        print("✅ Answer generated")
         return answer
     
     def generate_summary(self, results: List[LabResult]) -> str:
-        """Generate overall summary of lab results"""
+        """Generate overall summary - FAST"""
+        
+        print("📊 Generating summary...")
         
         abnormal = [r for r in results if r.status in ['high', 'low']]
         normal = [r for r in results if r.status == 'normal']
         
         if not abnormal:
-            return "✅ Great news! All your lab results are within normal ranges. Keep up the good work with your health!"
-        
-        # Get context about abnormal results
-        queries = [f"{r.test_name} {r.status}" for r in abnormal]
-        combined_query = " ".join(queries)
-        context = self._retrieve_context(combined_query, k=4)
-        
-        # Create summary prompt
-        abnormal_list = "\n".join([
-            f"- {r.test_name}: {r.value} {r.unit} ({r.status})"
-            for r in abnormal
-        ])
-        
-        prompt = f"""Provide a brief summary of these lab results.
+            return """✅ Excellent news! All your lab results are within normal ranges. 
 
-Normal Results: {len(normal)} tests
-Abnormal Results: {len(abnormal)} tests
-
-Abnormal Tests:
-{abnormal_list}
+This suggests that the tested parameters are functioning well. Continue maintaining your current health habits, and follow your healthcare provider's recommendations for routine monitoring."""
+        
+        # Build summary
+        summary = f"""📊 Lab Results Summary
 
-Medical Context:
-{context}
+Total Tests: {len(results)}
+✅ Normal: {len(normal)}
+⚠️ Abnormal: {len(abnormal)}
 
-Write a 2-3 paragraph summary explaining what these results mean overall and general recommendations. Be reassuring but honest. Summary:"""
+"""
+        
+        if abnormal:
+            summary += "**Tests Outside Normal Range:**\n"
+            for r in abnormal[:5]:  # Limit to first 5
+                status_emoji = "↑" if r.status == "high" else "↓"
+                summary += f"{status_emoji} {r.test_name}: {r.value} {r.unit} ({r.status})\n"
+            
+            if len(abnormal) > 5:
+                summary += f"... and {len(abnormal) - 5} more\n"
+            
+            summary += "\n"
         
-        # Generate summary
-        summary = self._generate_text(prompt)
+        # Get context for abnormal results
+        if abnormal:
+            abnormal_names = ", ".join([r.test_name for r in abnormal[:3]])
+            context = self._retrieve_context(f"{abnormal_names} interpretation", k=2)
+            
+            if context and len(context) > 20:
+                summary += f"**Key Information:**\n{context[:400]}\n\n"
+        
+        summary += """**Next Steps:**
+1. Review these results with your healthcare provider
+2. Discuss any concerns or symptoms you're experiencing
+3. Follow recommended treatment or monitoring plans
+
+Remember: These results are for educational purposes. Always consult your doctor for medical advice."""
         
+        print("✅ Summary generated")
         return summary
 
 
-# Example usage
+# Test if ran directly
 if __name__ == "__main__":
-    from pdf_extractor import LabResult
-    
-    # Initialize RAG system
-    print("Initializing RAG system...")
-    rag = LabReportRAG()
-    
-    # Example result
-    test_result = LabResult(
-        test_name="Hemoglobin",
-        value="10.5",
-        unit="g/dL",
-        reference_range="12.0-15.5",
-        status="low"
-    )
+    print("Testing RAG system...")
     
-    # Generate explanation
-    print("\nGenerating explanation...")
-    explanation = rag.explain_result(test_result)
-    print(f"\n{explanation}")
+    try:
+        rag = LabReportRAG()
+        print("\n✅ RAG system initialized successfully!")
+        
+        # Test with example
+        from pdf_extractor import LabResult
+        test_result = LabResult(
+            test_name="Hemoglobin",
+            value="10.5",
+            unit="g/dL",
+            reference_range="12.0-15.5",
+            status="low"
+        )
+        
+        explanation = rag.explain_result(test_result)
+        print(f"\nTest Explanation:\n{explanation}")
+        
+    except Exception as e:
+        print(f"\n❌ Error: {e}")