app.py

import streamlit as st
import os
from groq import Groq
from transformers import ViTForImageClassification, ViTImageProcessor
from sentence_transformers import SentenceTransformer
from PIL import Image
import torch
import numpy as np
from typing import List, Dict
import faiss
import json

# Initialize sentence transformer for embeddings
@st.cache_resource
def init_embedding_model():
    return SentenceTransformer('all-MiniLM-L6-v2')

# Initialize Groq client
@st.cache_resource
def init_groq_client():
    return Groq(api_key=os.environ.get("GROQ_API_KEY"))
class RAGSystem:
    def __init__(self):
        self.embedding_model = init_embedding_model()
        self.knowledge_base = self.load_knowledge_base()
        self.vector_store = self.create_vector_store()
        
    def load_knowledge_base(self) -> List[Dict]:
        """Load and preprocess knowledge base into a list of documents"""
        # Your existing knowledge base dictionary
        kb = {
            "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",
                        "Surface treatment and waterproofing"
                    ],
                    "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",
                    "testing_required": ["Core testing", "Reinforcement scanning", "Chloride testing"],
                    "common_causes": [
                        "Reinforcement corrosion",
                        "Freeze-thaw cycles",
                        "Poor concrete cover",
                        "Chemical attack"
                    ],
                    "safety_considerations": [
                        "Risk of structural failure",
                        "Falling concrete hazard",
                        "Worker safety during repairs"
                    ]
                },
                {
                    "severity": "Moderate",
                    "description": "Surface spalling without exposed reinforcement",
                    "repair_method": [
                        "Remove loose concrete",
                        "Surface preparation",
                        "Apply repair mortar",
                        "Surface treatment"
                    ],
                    "estimated_cost": "Medium ($5,000-$10,000)",
                    "timeframe": "1-2 weeks",
                    "location": "Non-structural elements",
                    "required_expertise": "Concrete Repair Specialist",
                    "immediate_action": "Remove loose material, protect from water ingress",
                    "prevention": "Surface sealers, proper drainage",
                    "testing_required": ["Surface adhesion testing", "Moisture testing"],
                    "common_causes": [
                        "Surface carbonation",
                        "Impact damage",
                        "Poor curing"
                    ],
                    "safety_considerations": [
                        "Minor falling debris risk",
                        "Dust control during repairs"
                    ]
                }
            ],
            "reinforcement_corrosion": [
                {
                    "severity": "Critical",
                    "description": "Severe corrosion with >30% section loss",
                    "repair_method": [
                        "Structural support installation",
                        "Concrete removal around reinforcement",
                        "Reinforcement replacement",
                        "Corrosion protection application",
                        "High-strength concrete repair",
                        "Cathodic protection installation"
                    ],
                    "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",
                    "testing_required": [
                        "Half-cell potential survey",
                        "Concrete resistivity testing",
                        "Chloride analysis",
                        "Carbonation testing"
                    ],
                    "common_causes": [
                        "Chloride contamination",
                        "Carbonation",
                        "Stray electrical currents",
                        "Poor concrete quality"
                    ],
                    "safety_considerations": [
                        "Structural collapse risk",
                        "Electrical hazards during testing",
                        "Confined space entry"
                    ]
                }
            ],
            "structural_cracks": [
                {
                    "severity": "High",
                    "description": "Active structural cracks >5mm width",
                    "repair_method": [
                        "Structural analysis",
                        "Crack monitoring",
                        "Epoxy injection",
                        "Carbon fiber reinforcement",
                        "Load path modification"
                    ],
                    "estimated_cost": "High ($10,000-$20,000)",
                    "timeframe": "2-4 weeks",
                    "location": "Primary structural elements",
                    "required_expertise": "Structural Engineer",
                    "immediate_action": "Install crack monitors, restrict loading",
                    "prevention": "Proper design, joint maintenance",
                    "testing_required": [
                        "Crack movement monitoring",
                        "Load testing",
                        "Concrete strength testing"
                    ],
                    "common_causes": [
                        "Overloading",
                        "Foundation settlement",
                        "Thermal movements",
                        "Design deficiencies"
                    ],
                    "safety_considerations": [
                        "Structural stability",
                        "Water infiltration",
                        "Working at height"
                    ]
                }
            ],
            "water_damage": [
                {
                    "severity": "Medium",
                    "description": "Active water infiltration with deterioration",
                    "repair_method": [
                        "Water source identification",
                        "Drainage improvement",
                        "Waterproofing membrane installation",
                        "Joint sealing",
                        "Surface treatment"
                    ],
                    "estimated_cost": "Medium ($5,000-$15,000)",
                    "timeframe": "1-3 weeks",
                    "location": "Various locations",
                    "required_expertise": "Waterproofing Specialist",
                    "immediate_action": "Water diversion, dehumidification",
                    "prevention": "Regular maintenance, proper drainage",
                    "testing_required": [
                        "Moisture mapping",
                        "Drainage assessment",
                        "Permeability testing"
                    ],
                    "common_causes": [
                        "Failed waterproofing",
                        "Poor drainage",
                        "Joint failure",
                        "Condensation"
                    ],
                    "safety_considerations": [
                        "Slip hazards",
                        "Electrical safety",
                        "Mold growth"
                    ]
                }
            ],
            "surface_deterioration": [
                {
                    "severity": "Low",
                    "description": "Surface scaling and deterioration",
                    "repair_method": [
                        "Surface cleaning",
                        "Repair material application",
                        "Surface treatment",
                        "Protective coating"
                    ],
                    "estimated_cost": "Low ($2,000-$5,000)",
                    "timeframe": "3-5 days",
                    "location": "Exposed surfaces",
                    "required_expertise": "Concrete Repair Technician",
                    "immediate_action": "Clean and protect surface",
                    "prevention": "Regular maintenance, surface protection",
                    "testing_required": [
                        "Surface strength testing",
                        "Coating adhesion tests"
                    ],
                    "common_causes": [
                        "Freeze-thaw damage",
                        "Chemical exposure",
                        "Poor finishing",
                        "Abrasion"
                    ],
                    "safety_considerations": [
                        "Dust control",
                        "Chemical handling",
                        "PPE requirements"
                    ]
                }
            ],
            "alkali_silica_reaction": [
                {
                    "severity": "High",
                    "description": "Concrete expansion and map cracking due to ASR",
                    "repair_method": [
                        "Expansion monitoring",
                        "Moisture control",
                        "Crack sealing",
                        "Surface treatment",
                        "Structural strengthening"
                    ],
                    "estimated_cost": "High ($15,000-$25,000)",
                    "timeframe": "3-5 weeks",
                    "location": "Concrete elements",
                    "required_expertise": "Materials Engineer + Structural Engineer",
                    "immediate_action": "Monitor expansion, control moisture",
                    "prevention": "Proper aggregate selection, pozzolans",
                    "testing_required": [
                        "Petrographic analysis",
                        "Expansion testing",
                        "Humidity monitoring"
                    ],
                    "common_causes": [
                        "Reactive aggregates",
                        "High alkali cement",
                        "Moisture presence",
                        "Temperature cycles"
                    ],
                    "safety_considerations": [
                        "Progressive deterioration",
                        "Structural integrity",
                        "Long-term monitoring"
                    ]
                }
            ]
        }
        
        # Convert nested knowledge base into flat documents
        documents = []
        for category, items in kb.items():
            for item in items:
                # Create a text representation of the document
                doc_text = f"Category: {category}\n"
                for key, value in item.items():
                    if isinstance(value, list):
                        doc_text += f"{key}: {', '.join(value)}\n"
                    else:
                        doc_text += f"{key}: {value}\n"
                documents.append({
                    "text": doc_text,
                    "metadata": {"category": category}
                })
        
        return documents

    def create_vector_store(self):
        """Create FAISS vector store from knowledge base"""
        # Generate embeddings for all documents
        texts = [doc["text"] for doc in self.knowledge_base]
        embeddings = self.embedding_model.encode(texts)
        
        # Initialize FAISS index
        dimension = embeddings.shape[1]
        index = faiss.IndexFlatL2(dimension)
        index.add(np.array(embeddings).astype('float32'))
        
        return index
    
    def create_vector_store(self):
        """Create FAISS vector store from knowledge base"""
        # Generate embeddings for all documents
        texts = [doc["text"] for doc in self.knowledge_base]
        embeddings = self.embedding_model.encode(texts)
        
        # Initialize FAISS index
        dimension = embeddings.shape[1]
        index = faiss.IndexFlatL2(dimension)
        index.add(np.array(embeddings).astype('float32'))
        
        return index
        
    def get_relevant_context(self, query: str, k: int = 3) -> str:
        """Retrieve relevant context based on query"""
        # Generate query embedding
        query_embedding = self.embedding_model.encode([query])
        
        # Search for similar documents
        D, I = self.vector_store.search(np.array(query_embedding).astype('float32'), k)
        
        # Combine relevant documents into context
        context = "\n\n".join([self.knowledge_base[i]["text"] for i in I[0]])
        return context

    def get_groq_response(query: str, context: str) -> str:
    """Get response from Groq LLM"""
    client = init_groq_client()
    try:
        prompt = f"""Based on the following context about construction defects, please answer the question.
Context:
{context}
Question: {query}
Please provide a detailed and specific answer based on the given context."""

        response = client.chat.completions.create(
            messages=[
                {
                    "role": "system",
                    "content": "You are a construction defect analysis expert. Provide detailed, accurate answers based on the given context."
                },
                {
                    "role": "user",
                    "content": prompt
                }
            ],
            model="llama-3.3-70b-versatile",
        )
        return response.choices[0].message.content
    except Exception as e:
        return f"Error: {str(e)}"

ef main():
    st.set_page_config(
        page_title="Construction Defect RAG Analyzer",
        page_icon="🏗️",
        layout="wide"
    )

    st.title("🏗️ Construction Defect RAG Analyzer")

    # Initialize RAG system
    if 'rag_system' not in st.session_state:
        st.session_state.rag_system = RAGSystem()

    # File upload for image analysis
    uploaded_file = st.file_uploader(
        "Upload a construction image",
        type=['jpg', 'jpeg', 'png']
    )

    # Query input
    user_query = st.text_input("Ask a question about construction defects:")

    if user_query:
        with st.spinner("Processing query..."):
            # Get relevant context using RAG
            context = st.session_state.rag_system.get_relevant_context(user_query)
            
            # Debug view of retrieved context
            if st.checkbox("Show retrieved context"):
                st.subheader("Retrieved Context")
                st.text(context)
            
            # Get response from Groq
            st.subheader("AI Assistant Response")
            response = get_groq_response(user_query, context)
            st.write(response)

    if uploaded_file:
        image = Image.open(uploaded_file)
        st.image(image, caption="Uploaded Image")
        
        # Combine image analysis with RAG
        col1, col2 = st.columns(2)
        
        with col1:
            st.subheader("Image Analysis")
            # Image analysis placeholder
            st.info("Image analysis results would appear here")

        with col2:
            st.subheader("AI Assistant Response")
            if user_query:  # Only show response if there's a query
                # Get relevant context from knowledge base
                context = st.session_state.rag_system.get_relevant_context(user_query)
                
                # Get response from Groq
                with st.spinner("Getting AI response..."):
                    response = get_groq_response(user_query, context)
                    st.write(response)

        # Display knowledge base sections
        if st.checkbox("Show Knowledge Base"):
            st.subheader("Available Knowledge Base")
            kb_data = st.session_state.rag_system.knowledge_base
            for doc in kb_data:
                category = doc["metadata"]["category"]
                with st.expander(category.title()):
                    st.text(doc["text"])

if __name__ == "__main__":
    main()