Update rag_utils.py

rag_utils.py

@@ -1,136 +1,168 @@
-# knowledge_base.py
-
+from langchain.embeddings import HuggingFaceEmbeddings
+from langchain.vectorstores import FAISS
+from langchain.text_splitter import RecursiveCharacterTextSplitter
+from langchain.chains import RetrievalQA
+from langchain.prompts import PromptTemplate
+from langchain.llms import HuggingFaceHub
+import os
 import logging
 
 # Configure logging
 logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(levelname)s - %(message)s')
 
-# Knowledge base containing structural damage information
-KNOWLEDGE_BASE = {
-    "spalling": [
-        {
-            "severity": "Critical",
-            "description": "Severe concrete spalling with exposed reinforcement and section loss",
-            "repair_method": [
-                "Install temporary support",
-                "Remove deteriorated concrete",
-                "Clean and treat reinforcement",
-                "Apply corrosion inhibitor",
-                "Apply bonding agent",
-                "High-strength repair mortar"
-            ],
-            "estimated_cost": "Very High ($15,000+)",
-            "timeframe": "3-4 weeks",
-            "location": "Primary structural elements",
-            "required_expertise": "Structural Engineer + Specialist Contractor",
-            "immediate_action": "Evacuate area, install temporary support, prevent access",
-            "prevention": "Regular inspections, waterproofing, chloride protection"
-        },
-        {
-            "severity": "High",
-            "description": "Surface spalling with visible reinforcement",
-            "repair_method": [
-                "Remove damaged concrete",
-                "Treat reinforcement",
-                "Apply repair mortar",
-                "Surface treatment"
-            ],
-            "estimated_cost": "High ($8,000-$15,000)",
-            "timeframe": "2-3 weeks",
-            "location": "Structural elements",
-            "required_expertise": "Structural Engineer",
-            "immediate_action": "Area isolation, temporary support assessment",
-            "prevention": "Protective coatings, drainage improvement"
-        }
-    ],
-    "reinforcement_corrosion": [
-        {
-            "severity": "Critical",
-            "description": "Severe corrosion with >30% section loss",
-            "repair_method": [
-                "Structural support",
-                "Remove concrete",
-                "Replace reinforcement",
-                "Corrosion protection",
-                "Concrete repair"
-            ],
-            "estimated_cost": "Critical ($20,000+)",
-            "timeframe": "4-6 weeks",
-            "location": "Load-bearing elements",
-            "required_expertise": "Senior Structural Engineer",
-            "immediate_action": "Immediate evacuation, emergency shoring",
-            "prevention": "Waterproofing, cathodic protection"
-        }
-    ],
-    "structural_crack": [
-        {
-            "severity": "High",
-            "description": "Cracks >5mm in structural elements",
-            "repair_method": [
-                "Structural analysis",
-                "Epoxy injection",
-                "Carbon fiber reinforcement",
-                "Crack monitoring"
-            ],
-            "estimated_cost": "High ($10,000-$20,000)",
-            "timeframe": "2-4 weeks",
-            "location": "Primary structure",
-            "required_expertise": "Structural Engineer",
-            "immediate_action": "Install crack monitors, load restriction",
-            "prevention": "Load management, joint maintenance"
-        }
-    ],
-    "dampness": [
-        {
-            "severity": "Medium",
-            "description": "Active water penetration with efflorescence",
-            "repair_method": [
-                "Water source identification",
-                "Drainage improvement",
-                "Waterproof membrane",
-                "Ventilation"
-            ],
-            "estimated_cost": "Medium ($5,000-$10,000)",
-            "timeframe": "1-2 weeks",
-            "location": "Various",
-            "required_expertise": "Waterproofing Specialist",
-            "immediate_action": "Dehumidification, efflorescence cleaning",
-            "prevention": "Proper drainage, vapor barriers"
-        }
-    ],
-    "no_damage": [
-        {
-            "severity": "Low",
-            "description": "No significant structural issues",
-            "repair_method": [
-                "Regular inspection",
-                "Preventive maintenance"
-            ],
-            "estimated_cost": "Low ($500-$2,000)",
-            "timeframe": "1-2 days",
-            "location": "General",
-            "required_expertise": "Building Inspector",
-            "immediate_action": "Continue monitoring",
-            "prevention": "Regular maintenance schedule"
+class RAGSystem:
+    def __init__(self):
+        try:
+            self.embeddings = HuggingFaceEmbeddings(
+                model_name="sentence-transformers/all-mpnet-base-v2"
+            )
+            self.vector_store = None
+            self.text_splitter = RecursiveCharacterTextSplitter(
+                chunk_size=500,
+                chunk_overlap=50
+            )
+            # Initialize HuggingFace model for text generation
+            self.llm = HuggingFaceHub(
+                repo_id="google/flan-t5-large",
+                task="text-generation",
+                model_kwargs={"temperature": 0.7, "max_length": 512}
+            )
+            logging.info("RAG system initialized successfully.")
+        except Exception as e:
+            logging.error(f"Failed to initialize RAG system: {str(e)}")
+            raise e
+
+    def initialize_knowledge_base(self, knowledge_base):
+        """Initialize vector store with enhanced construction knowledge"""
+        try:
+            documents = []
+            # Validate knowledge base
+            self._validate_knowledge_base(knowledge_base)
+            
+            # Add expert 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)
+            
+            splits = self.text_splitter.create_documents(documents)
+            self.vector_store = FAISS.from_documents(splits, self.embeddings)
+            
+            # Initialize QA chain
+            self.qa_chain = RetrievalQA.from_chain_type(
+                llm=self.llm,
+                chain_type="stuff",
+                retriever=self.vector_store.as_retriever(),
+                chain_type_kwargs={
+                    "prompt": self._get_qa_prompt()
+                }
+            )
+            logging.info("Knowledge base initialized successfully.")
+        except Exception as e:
+            logging.error(f"Failed to initialize knowledge base: {str(e)}")
+            raise e
+
+    def _validate_knowledge_base(self, knowledge_base):
+        """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 key in required_keys:
+                    if key not in case:
+                        raise ValueError(f"Missing required field '{key}' in {damage_type}")
+        logging.info("Knowledge base validation passed.")
+
+    def _get_qa_prompt(self):
+        """Create a custom prompt template for the QA chain"""
+        template = """
+        Context: {context}
+        
+        Question: {question}
+        
+        Provide a detailed analysis considering:
+        1. Technical aspects
+        2. Safety implications
+        3. Cost-benefit analysis
+        4. Long-term considerations
+        5. Best practices and recommendations
+        
+        Answer:
+        """
+        return PromptTemplate(
+            template=template,
+            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"
         }
-    ]
-}
 
-# Damage type metadata
-DAMAGE_TYPES = {
-    0: {'name': 'spalling', 'risk': 'High', 'color': '#ff4d4d'},
-    1: {'name': 'reinforcement_corrosion', 'risk': 'Critical', 'color': '#800000'},
-    2: {'name': 'structural_crack', 'risk': 'High', 'color': '#ff6b6b'},
-    3: {'name': 'dampness', 'risk': 'Medium', 'color': '#4dabf7'},
-    4: {'name': 'no_damage', 'risk': 'Low', 'color': '#40c057'}
-}
+    def get_enhanced_analysis(self, damage_type, confidence, custom_query=None):
+        """Get enhanced analysis with dynamic content generation"""
+        try:
+            if not custom_query:
+                base_query = f"""
+                Provide a comprehensive analysis for {damage_type} damage with {confidence}% confidence level.
+                Include technical assessment, safety implications, and expert recommendations.
+                """
+            else:
+                base_query = custom_query
+            
+            # Get relevant documents
+            results = self.qa_chain.run(base_query)
+            
+            # Process and categorize the response
+            enhanced_info = {
+                "technical_details": self._extract_technical_details(results, damage_type),
+                "safety_considerations": self._extract_safety_considerations(results),
+                "expert_recommendations": self._extract_recommendations(results, confidence)
+            }
+            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]
 
-# Validation for 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 key in required_keys:
-            if key not in case:
-                logging.error(f"Missing required field '{key}' in {damage_type}")
-                raise ValueError(f"Missing required field '{key}' in {damage_type}")
-logging.info("Knowledge base validation passed.")
+    def _extract_recommendations(self, results, confidence):
+        return [f"Prioritized recommendations based on {confidence}% confidence", results]