Update rag_utils.py

rag_utils.py

@@ -13,6 +13,7 @@ logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(levelname)s - %(
 class RAGSystem:
     def __init__(self):
         try:
+            # Initialize embeddings
             self.embeddings = HuggingFaceEmbeddings(
                 model_name="sentence-transformers/all-mpnet-base-v2"
             )
@@ -39,34 +40,38 @@ class RAGSystem:
             # Validate knowledge base
             self._validate_knowledge_base(knowledge_base)
             
-            # Add expert insights and case studies
+            # Generate insights and case studies
             expert_insights = self._generate_expert_insights(knowledge_base)
             case_studies = self._generate_case_studies()
             
             for damage_type, cases in knowledge_base.items():
-                for case in cases:
-                    # Combine basic info with expert insights
-                    relevant_insight = expert_insights.get(damage_type, "")
-                    relevant_cases = case_studies.get(damage_type, "")
-                    
-                    doc_text = f"""
-                    Damage Type: {damage_type}
-                    Severity: {case['severity']}
-                    Description: {case['description']}
-                    Technical Details: {case['description']}
-                    Expert Insight: {relevant_insight}
-                    Case Studies: {relevant_cases}
-                    Repair Methods: {', '.join(case['repair_method'])}
-                    Cost Considerations: {case['estimated_cost']}
-                    Implementation Timeline: {case['timeframe']}
-                    Location Specifics: {case['location']}
-                    Required Expertise Level: {case['required_expertise']}
-                    Emergency Protocol: {case['immediate_action']}
-                    Preventive Measures: {case['prevention']}
-                    Long-term Implications: Analysis of long-term structural integrity impact
-                    Environmental Factors: Consideration of environmental conditions
-                    """
-                    documents.append(doc_text)
+                for idx, case in enumerate(cases):
+                    try:
+                        # Combine insights into document text
+                        relevant_insight = expert_insights.get(damage_type, "")
+                        relevant_cases = case_studies.get(damage_type, "")
+                        
+                        doc_text = f"""
+                        Damage Type: {damage_type}
+                        Severity: {case['severity']}
+                        Description: {case['description']}
+                        Technical Details: {case['description']}
+                        Expert Insight: {relevant_insight}
+                        Case Studies: {relevant_cases}
+                        Repair Methods: {', '.join(case['repair_method'])}
+                        Cost Considerations: {case['estimated_cost']}
+                        Implementation Timeline: {case['timeframe']}
+                        Location Specifics: {case['location']}
+                        Required Expertise Level: {case['required_expertise']}
+                        Emergency Protocol: {case['immediate_action']}
+                        Preventive Measures: {case['prevention']}
+                        """
+                        documents.append(doc_text)
+                    except KeyError as e:
+                        logging.warning(f"Missing key {str(e)} in {damage_type}, case {idx + 1}. Skipping.")
+            
+            if not documents:
+                raise ValueError("No valid documents to process.")
             
             splits = self.text_splitter.create_documents(documents)
             self.vector_store = FAISS.from_documents(splits, self.embeddings)
@@ -89,10 +94,11 @@ class RAGSystem:
         """Validate the structure of the knowledge base."""
         required_keys = ['severity', 'description', 'repair_method', 'estimated_cost', 'timeframe', 'location', 'required_expertise', 'immediate_action', 'prevention']
         for damage_type, cases in knowledge_base.items():
-            for case in cases:
+            for idx, case in enumerate(cases):
                 for key in required_keys:
                     if key not in case:
-                        raise ValueError(f"Missing required field '{key}' in {damage_type}")
+                        logging.error(f"Missing required field '{key}' in {damage_type}, case {idx + 1}")
+                        raise ValueError(f"Missing required field '{key}' in {damage_type}, case {idx + 1}")
         logging.info("Knowledge base validation passed.")
 
     def _get_qa_prompt(self):
@@ -116,26 +122,12 @@ class RAGSystem:
             input_variables=["context", "question"]
         )
 
-    def _generate_expert_insights(self, knowledge_base):
-        """Generate expert insights for each damage type"""
-        insights = {}
-        for damage_type in knowledge_base.keys():
-            insights[damage_type] = f"Expert analysis for {damage_type} including latest research findings and industry best practices."
-        return insights
-
-    def _generate_case_studies(self):
-        """Generate relevant case studies for each damage type"""
-        return {
-            "spalling": "Case studies of successful spalling repairs in similar structures",
-            "reinforcement_corrosion": "Examples of corrosion mitigation in harsh environments",
-            "structural_crack": "Analysis of crack progression and successful interventions",
-            "dampness": "Case studies of effective moisture control solutions",
-            "no_damage": "Preventive maintenance success stories"
-        }
-
     def get_enhanced_analysis(self, damage_type, confidence, custom_query=None):
         """Get enhanced analysis with dynamic content generation"""
         try:
+            if not self.vector_store:
+                raise ValueError("Vector store is not initialized.")
+
             if not custom_query:
                 base_query = f"""
                 Provide a comprehensive analysis for {damage_type} damage with {confidence}% confidence level.
@@ -146,6 +138,9 @@ class RAGSystem:
             
             # Get relevant documents
             results = self.qa_chain.run(base_query)
+            if not results:
+                logging.warning("No results returned for query.")
+                return {"technical_details": [], "safety_considerations": [], "expert_recommendations": []}
             
             # Process and categorize the response
             enhanced_info = {
@@ -156,13 +151,4 @@ class RAGSystem:
             return enhanced_info
         except Exception as e:
             logging.error(f"Failed to generate enhanced analysis: {str(e)}")
-            return None
-
-    def _extract_technical_details(self, results, damage_type):
-        return [f"Detailed technical analysis for {damage_type}", results]
-
-    def _extract_safety_considerations(self, results):
-        return [f"Safety analysis based on current conditions", results]
-
-    def _extract_recommendations(self, results, confidence):
-        return [f"Prioritized recommendations based on {confidence}% confidence", results]
+            return {"technical_details": [], "safety_considerations": [], "expert_recommendations": []}