Update rag_utils.py

rag_utils.py

@@ -1,154 +1,201 @@
+# rag_utils.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
+from langchain.docstore.document import Document
 import logging
+from typing import List, Dict, Any
+import numpy as np
+from tqdm import tqdm
+import streamlit as st
 
-# Configure logging
-logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(levelname)s - %(message)s')
+# Set up logging
+logging.basicConfig(level=logging.INFO)
+logger = logging.getLogger(__name__)
 
 class RAGSystem:
     def __init__(self):
+        """Initialize RAG system with custom embeddings and configurations"""
         try:
-            # Initialize embeddings
             self.embeddings = HuggingFaceEmbeddings(
-                model_name="sentence-transformers/all-mpnet-base-v2"
+                model_name="sentence-transformers/all-mpnet-base-v2",
+                model_kwargs={'device': 'cuda' if st.cuda.is_available() else 'cpu'}
             )
             self.vector_store = None
             self.text_splitter = RecursiveCharacterTextSplitter(
                 chunk_size=500,
-                chunk_overlap=50
+                chunk_overlap=50,
+                separators=["\n\n", "\n", ". ", ", ", " ", ""]
             )
-            # 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.")
+            logger.info("RAG system initialized successfully")
         except Exception as e:
-            logging.error(f"Failed to initialize RAG system: {str(e)}")
-            raise e
+            logger.error(f"Error initializing RAG system: {str(e)}")
+            raise
 
-    def initialize_knowledge_base(self, knowledge_base):
-        """Initialize vector store with enhanced construction knowledge"""
+    def _create_documents(self, knowledge_base: Dict) -> List[Document]:
+        """Create documents from knowledge base with structured format"""
+        documents = []
         try:
-            documents = []
-            # Validate knowledge base
-            self._validate_knowledge_base(knowledge_base)
-            
-            # 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 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.")
+                for case in cases:
+                    # Create a detailed document for each case
+                    technical_info = f"""
+                    Technical Analysis for {damage_type}:
+                    Severity Level: {case['severity']}
+                    Detailed Description: {case['description']}
+                    Primary Location: {case['location']}
+                    Required Expertise: {case['required_expertise']}
+                    """
+
+                    repair_info = f"""
+                    Repair and Maintenance Information:
+                    Repair Methods: {' -> '.join(case['repair_method'])}
+                    Estimated Cost Range: {case['estimated_cost']}
+                    Expected Timeframe: {case['timeframe']}
+                    """
+
+                    safety_info = f"""
+                    Safety and Prevention Guidelines:
+                    Immediate Actions Required: {case['immediate_action']}
+                    Preventive Measures: {case['prevention']}
+                    Critical Considerations: Special attention needed for {damage_type} in {case['location']}
+                    """
+
+                    # Combine all information
+                    doc_text = f"{technical_info}\n{repair_info}\n{safety_info}"
+                    
+                    # Create metadata for better retrieval
+                    metadata = {
+                        'damage_type': damage_type,
+                        'severity': case['severity'],
+                        'location': case['location'],
+                        'document_type': 'construction_damage_analysis'
+                    }
+
+                    documents.append(Document(
+                        page_content=doc_text,
+                        metadata=metadata
+                    ))
             
-            if not documents:
-                raise ValueError("No valid documents to process.")
+            logger.info(f"Created {len(documents)} documents from knowledge base")
+            return documents
+        except Exception as e:
+            logger.error(f"Error creating documents: {str(e)}")
+            raise
+
+    def initialize_knowledge_base(self, knowledge_base: Dict):
+        """Initialize vector store with construction knowledge"""
+        try:
+            # Create documents
+            documents = self._create_documents(knowledge_base)
             
-            splits = self.text_splitter.create_documents(documents)
-            self.vector_store = FAISS.from_documents(splits, self.embeddings)
+            # Split documents into chunks
+            splits = self.text_splitter.split_documents(documents)
             
-            # 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()
-                }
+            # Create vector store
+            self.vector_store = FAISS.from_documents(
+                documents=splits,
+                embedding=self.embeddings
             )
-            logging.info("Knowledge base initialized successfully.")
+            
+            logger.info("Knowledge base initialized successfully")
+        except Exception as e:
+            logger.error(f"Error initializing knowledge base: {str(e)}")
+            raise
+
+    def _format_response(self, docs: List[Document], damage_type: str, confidence: float) -> Dict[str, List[str]]:
+        """Format retrieved documents into structured response"""
+        response = {
+            "technical_details": [],
+            "safety_considerations": [],
+            "expert_recommendations": []
+        }
+
+        try:
+            for doc in docs:
+                content = doc.page_content
+                # Parse technical details
+                if "Technical Analysis" in content:
+                    response["technical_details"].append(
+                        f"For {damage_type} (Confidence: {confidence:.1f}%):\n" +
+                        content.split("Technical Analysis")[1].split("Repair")[0].strip()
+                    )
+                
+                # Parse safety considerations
+                if "Safety and Prevention" in content:
+                    response["safety_considerations"].append(
+                        content.split("Safety and Prevention")[1].strip()
+                    )
+                
+                # Parse repair recommendations
+                if "Repair and Maintenance" in content:
+                    response["expert_recommendations"].append(
+                        content.split("Repair and Maintenance")[1].split("Safety")[0].strip()
+                    )
+            
+            return response
         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 idx, case in enumerate(cases):
-                for key in required_keys:
-                    if key not in case:
-                        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):
-        """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 get_enhanced_analysis(self, damage_type, confidence, custom_query=None):
-        """Get enhanced analysis with dynamic content generation"""
+            logger.error(f"Error formatting response: {str(e)}")
+            raise
+
+    def get_enhanced_analysis(
+        self, 
+        damage_type: str, 
+        confidence: float, 
+        custom_query: str = None
+    ) -> Dict[str, List[str]]:
+        """Get enhanced analysis with optional custom query support"""
         try:
             if not self.vector_store:
-                raise ValueError("Vector store is not initialized.")
+                raise ValueError("Vector store not initialized")
 
-            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.
-                """
+            # Prepare query
+            if custom_query:
+                query = f"{custom_query} for {damage_type} damage"
             else:
-                base_query = custom_query
-            
+                query = f"""
+                Provide detailed analysis for {damage_type} damage with {confidence}% confidence level.
+                Include technical assessment, safety considerations, and repair recommendations.
+                """
+
             # 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 = {
-                "technical_details": self._extract_technical_details(results, damage_type),
-                "safety_considerations": self._extract_safety_considerations(results),
-                "expert_recommendations": self._extract_recommendations(results, confidence)
+            docs = self.vector_store.similarity_search(
+                query=query,
+                k=3,  # Get top 3 most relevant documents
+                fetch_k=5  # Fetch top 5 for better diversity
+            )
+
+            # Format and return response
+            return self._format_response(docs, damage_type, confidence)
+
+        except Exception as e:
+            logger.error(f"Error getting enhanced analysis: {str(e)}")
+            return {
+                "technical_details": [f"Error retrieving analysis: {str(e)}"],
+                "safety_considerations": ["Please try again or contact support."],
+                "expert_recommendations": ["System currently unavailable."]
             }
-            return enhanced_info
+
+    def get_similar_cases(self, damage_type: str, confidence: float) -> List[Dict[str, Any]]:
+        """Get similar damage cases for comparison"""
+        try:
+            if not self.vector_store:
+                raise ValueError("Vector store not initialized")
+
+            query = f"Find similar cases of {damage_type} damage"
+            docs = self.vector_store.similarity_search(query, k=3)
+            
+            similar_cases = []
+            for doc in docs:
+                if doc.metadata['damage_type'] != damage_type:  # Avoid same damage type
+                    similar_cases.append({
+                        'damage_type': doc.metadata['damage_type'],
+                        'severity': doc.metadata['severity'],
+                        'location': doc.metadata['location'],
+                        'details': doc.page_content[:200] + '...'  # First 200 chars
+                    })
+            
+            return similar_cases
         except Exception as e:
-            logging.error(f"Failed to generate enhanced analysis: {str(e)}")
-            return {"technical_details": [], "safety_considerations": [], "expert_recommendations": []}
+            logger.error(f"Error getting similar cases: {str(e)}")
+            return []