src/streamlit_app.py

import streamlit as st
import os
import tempfile

# Updated Imports for Modern LangChain (v0.2+)
from langchain_community.document_loaders import TextLoader
from langchain_text_splitters import RecursiveCharacterTextSplitter
from langchain_openai import OpenAIEmbeddings, ChatOpenAI
from langchain_chroma import Chroma
from langchain_core.prompts import ChatPromptTemplate

# These are the specific imports that often cause issues in newer versions
# If these fail, ensure you have `pip install langchain` (the main package)
from langchain_classic.chains import create_retrieval_chain
from langchain_classic.chains.combine_documents import create_stuff_documents_chain

# --- Page Config ---
st.set_page_config(page_title="ACC Co-Pilot", page_icon="🏗️", layout="wide")

st.title("🏗️ Automated Compliance Checking (ACC) Co-Pilot")
st.markdown("""
**Objective:** Assist engineers in verifying design parameters against building codes using AI.
This tool interprets regulatory text to provide `COMPLIANT` or `NON-COMPLIANT` verdicts.
""")

# --- Sidebar: Configuration & Data ---
with st.sidebar:
    st.header("1. Configuration")
    
    # Try to load API Key from Secrets (HF Spaces) or User Input
    api_key = os.environ.get("OPENAI_API_KEY")
    if not api_key:
        api_key = st.text_input("Enter OpenAI API Key", type="password")
    
    if not api_key:
        st.warning("Please provide an OpenAI API Key to proceed.")
        st.stop()
        
    os.environ["OPENAI_API_KEY"] = api_key

    st.header("2. Regulatory Knowledge Base")
    # Default mock data
    default_text = """SECTION 402: BUILDING ENVELOPE REQUIREMENTS
402.1 General. Building thermal envelope assemblies shall comply with the insulation requirements of this section.
402.2 Specific Insulation Requirements (Zone 4):
    a) Roofs: Insulation entirely above deck shall have a minimum R-value of R-30 continuous insulation (c.i.).
    b) Walls (Above Grade): Metal framed walls shall have a minimum R-value of R-13 + R-7.5 c.i.
    c) Floors: Mass floors shall have a minimum R-10 c.i.
SECTION 503: HVAC SYSTEMS
503.1 Fan Control. Each cooling system listed in Table 503.1 shall be designed to vary the indoor fan airflow as a function of load.
503.2 Temperature Reset. Supply air temperature reset controls shall be provided for variable air volume (VAV) systems.
"""
    
    # Allow user to edit the "Regulation" to test different rules
    reg_text = st.text_area("Current Regulation Text:", value=default_text, height=300)

# --- Logic: Vector Database (Cached) ---
@st.cache_resource(show_spinner=False)
def process_regulatory_text(text):
    """
    Takes the raw text, chunks it, and creates a Chroma VectorStore.
    Cached so we don't re-embed on every interaction.
    """
    # 1. Save to temp file (Loader requirement)
    with tempfile.NamedTemporaryFile(mode="w", delete=False, suffix=".txt") as tmp_file:
        tmp_file.write(text)
        tmp_path = tmp_file.name

    # 2. Load and Split
    loader = TextLoader(tmp_path)
    docs = loader.load()
    text_splitter = RecursiveCharacterTextSplitter(chunk_size=500, chunk_overlap=50)
    splits = text_splitter.split_documents(docs)

    # 3. Embed and Store
    embeddings = OpenAIEmbeddings()
    # Initialize Chroma properly
    vectorstore = Chroma.from_documents(documents=splits, embedding=embeddings)
    
    # Cleanup temp file
    try:
        os.remove(tmp_path)
    except OSError:
        pass
    
    return vectorstore

# Initialize retriever as None to handle cases where indexing fails or hasn't run
retriever = None

if reg_text:
    with st.spinner("Indexing Regulatory Documents..."):
        try:
            vectorstore = process_regulatory_text(reg_text)
            retriever = vectorstore.as_retriever(search_kwargs={"k": 2})
            st.sidebar.success("✅ Knowledge Base Indexed")
        except Exception as e:
            st.error(f"Error indexing documents: {e}")
            st.stop()

# --- Logic: The Chain ---
# Only build the chain if the retriever is ready
if retriever:
    # Define the System Prompt
    system_prompt = (
        "You are an expert Automated Compliance Checking (ACC) Co-Pilot. "
        "Your goal is to assist building engineers in verifying design parameters "
        "against the provided building codes.\n\n"
        "Rules:\n"
        "1. Use ONLY the context provided below to answer the question.\n"
        "2. If the context does not contain the answer, say 'I cannot find a specific regulation'.\n"
        "3. Always cite the specific Section Number (e.g., Section 402.1) in your answer.\n"
        "4. If the user provides a design value, compare it to the requirement and explicitly state 'COMPLIANT' or 'NON-COMPLIANT'.\n\n"
        "Context:\n"
        "{context}"
    )

    prompt_template = ChatPromptTemplate.from_messages(
        [
            ("system", system_prompt),
            ("human", "{input}"),
        ]
    )

    llm = ChatOpenAI(model_name="gpt-4", temperature=0)
    question_answer_chain = create_stuff_documents_chain(llm, prompt_template)
    rag_chain = create_retrieval_chain(retriever, question_answer_chain)

    # --- Main Interface ---

    st.header("🔍 Compliance Verification")

    # Pre-filled examples for non-technical demo flow
    examples = [
        "I am designing a metal framed wall above grade in Zone 4. My design uses R-10 insulation + R-5 continuous insulation. Is this compliant?",
        "What are the requirements for fan controls in cooling systems?",
        "Does a Mass Floor require continuous insulation?"
    ]

    selected_example = st.selectbox("Select a scenario or type your own below:", [""] + examples)

    if selected_example:
        user_input = selected_example
    else:
        user_input = ""

    query = st.text_input("Enter your design parameter or question:", value=user_input)

    if st.button("Check Compliance"):
        if not query:
            st.warning("Please enter a query.")
        else:
            with st.spinner("Analyzing against regulations..."):
                # Invoke the chain
                response = rag_chain.invoke({"input": query})
                
                # Display Answer
                st.subheader("💡 Analysis Report")
                st.write(response["answer"])
                
                # Display Citations (Expandable)
                with st.expander("📂 Source Documentation (Citations)"):
                    for i, doc in enumerate(response["context"]):
                        st.markdown(f"**Source {i+1}**")
                        st.info(doc.page_content)
else:
    st.info("Please provide an OpenAI API Key and Regulation Text to initialize the system.")

# --- Footer ---
st.markdown("---")
st.caption("ACC Co-Pilot Demo | Powered by LangChain, OpenAI & Streamlit")