Create rag_utils.py

rag_utils.py

@@ -0,0 +1,73 @@
+from langchain.embeddings import HuggingFaceEmbeddings
+from langchain.vectorstores import FAISS
+from langchain.text_splitter import RecursiveCharacterTextSplitter
+import streamlit as st
+
+class RAGSystem:
+    def __init__(self):
+        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
+        )
+        
+    def initialize_knowledge_base(self, knowledge_base):
+        """Initialize vector store with construction knowledge"""
+        documents = []
+        
+        for damage_type, cases in knowledge_base.items():
+            for case in cases:
+                doc_text = f"""
+                Damage Type: {damage_type}
+                Severity: {case['severity']}
+                Description: {case['description']}
+                Repair Methods: {', '.join(case['repair_method'])}
+                Cost: {case['estimated_cost']}
+                Timeframe: {case['timeframe']}
+                Location: {case['location']}
+                Required Expertise: {case['required_expertise']}
+                Immediate Action: {case['immediate_action']}
+                Prevention: {case['prevention']}
+                """
+                documents.append(doc_text)
+        
+        splits = self.text_splitter.create_documents(documents)
+        self.vector_store = FAISS.from_documents(splits, self.embeddings)
+        
+    def query_knowledge_base(self, query, k=3):
+        """Query the vector store for relevant information"""
+        if not self.vector_store:
+            raise ValueError("Vector store not initialized")
+            
+        results = self.vector_store.similarity_search(query, k=k)
+        return results
+    
+    def get_enhanced_analysis(self, damage_type, confidence, custom_query=None):
+        """Get enhanced analysis based on damage type and confidence"""
+        if not custom_query:
+            query = f"What are the key considerations and recommendations for {damage_type} with {confidence}% confidence?"
+        else:
+            query = custom_query
+            
+        relevant_docs = self.query_knowledge_base(query)
+        
+        enhanced_info = {
+            "technical_details": [],
+            "safety_considerations": [],
+            "expert_recommendations": []
+        }
+        
+        for doc in relevant_docs:
+            content = doc.page_content
+            
+            if "severity" in content.lower() or "description" in content.lower():
+                enhanced_info["technical_details"].append(content)
+            elif "immediate action" in content.lower() or "prevention" in content.lower():
+                enhanced_info["safety_considerations"].append(content)
+            elif "repair" in content.lower() or "expertise" in content.lower():
+                enhanced_info["expert_recommendations"].append(content)
+                
+        return enhanced_info