chatbot_final.py

import streamlit as st
import os
import shutil
from langchain_community.document_loaders import PyPDFLoader
from langchain_text_splitters import RecursiveCharacterTextSplitter
from langchain_huggingface import HuggingFaceEmbeddings
from langchain_community.vectorstores import FAISS
from langchain_classic.retrievers import EnsembleRetriever
from langchain_core.prompts import ChatPromptTemplate
from langchain_core.output_parsers import StrOutputParser
from langchain_core.runnables import RunnablePassthrough
import pandas as pd
from datetime import datetime
 
# if "page" not in st.session_state:
#     st.session_state.page = "reader"
 
# if st.session_state.page == "reader":
#     document_reader_page()
# elif st.session_state.page == "loader":
#     pdf_loader_page()
# elif st.session_state.page == "iso":
#     chatbot_page("Chatbuddy Application (ISO26262)", r"C:\Users\RVX2KOR\Desktop\LLMs\faissdb", "messages_iso")
# elif st.session_state.page == "vda360":
#     chatbot_page("Chatbuddy Application (VDA360)", r"C:\Users\RVX2KOR\Desktop\Olllama\vda-360", "messages_vda360")
# elif st.session_state.page == "vda305":
#     chatbot_page("Chatbuddy Application (VDA305)", r"C:\Users\RVX2KOR\Desktop\Olllama\vda_305", "messages_vda305")
# elif st.session_state.page == "obd":
#     chatbot_page("Chatbuddy Application (OBD_J1979-2_202104)", r"C:\Users\RVX2KOR\Desktop\Olllama\obd_store", "messages_obd")
# elif st.session_state.page == "standards":
#     standards_page()
 
# -------------------------
# Configuration / paths (all overridable via environment variables)
try:
    from dotenv import load_dotenv
    load_dotenv()
except ImportError:
    pass

PDF_STORE_DIR = os.environ.get("PDF_STORE_DIR", "/data/pdfs")
EMBEDDINGS_STORE_DIR = os.environ.get("EMBEDDINGS_STORE_DIR", "/data/embeddings")
USER_XL_PATH = os.environ.get("USER_XL_PATH", "/data/users/users.xlsx")

# LLM backend: "ollama" or "huggingface"
LLM_BACKEND = os.environ.get("LLM_BACKEND", "huggingface")
OLLAMA_BASE_URL = os.environ.get("OLLAMA_BASE_URL", "http://localhost:11434")
OLLAMA_MODEL = os.environ.get("OLLAMA_MODEL", "mistral")
HF_MODEL_ID = os.environ.get("HF_MODEL_ID", "mistralai/Mistral-7B-Instruct-v0.3")
HF_TOKEN = os.environ.get("HF_TOKEN", "")


def get_llm():
    """Return the configured LLM instance."""
    if LLM_BACKEND == "ollama":
        from langchain_community.llms import Ollama
        return Ollama(model=OLLAMA_MODEL, base_url=OLLAMA_BASE_URL)
    else:
        from langchain_community.llms import HuggingFaceEndpoint
        return HuggingFaceEndpoint(
            repo_id=HF_MODEL_ID,
            huggingfacehub_api_token=HF_TOKEN,
            temperature=0.3,
            max_new_tokens=512,
        )


def run_rag(retriever, query):
    """Run a RAG chain using the modern LCEL approach."""
    prompt = ChatPromptTemplate.from_template(
        "Answer the question based only on the context below.\n\n"
        "Context:\n{context}\n\nQuestion: {question}"
    )

    def format_docs(docs):
        return "\n\n".join(doc.page_content for doc in docs)

    chain = (
        {"context": retriever | format_docs, "question": RunnablePassthrough()}
        | prompt
        | get_llm()
        | StrOutputParser()
    )
    return chain.invoke(query)

# Pre-built FAISS index paths (default to subdirs inside EMBEDDINGS_STORE_DIR)
ISO_FAISS_PATH = os.environ.get("ISO_FAISS_PATH", os.path.join(EMBEDDINGS_STORE_DIR, "iso26262"))
VDA360_FAISS_PATH = os.environ.get("VDA360_FAISS_PATH", os.path.join(EMBEDDINGS_STORE_DIR, "vda360"))
VDA305_FAISS_PATH = os.environ.get("VDA305_FAISS_PATH", os.path.join(EMBEDDINGS_STORE_DIR, "vda305"))
OBD_FAISS_PATH = os.environ.get("OBD_FAISS_PATH", os.path.join(EMBEDDINGS_STORE_DIR, "obd_store"))

STANDARDS_EMBEDDINGS_FOLDER = "standards_all"
STANDARDS_EMBEDDINGS_PATH = os.path.join(EMBEDDINGS_STORE_DIR, STANDARDS_EMBEDDINGS_FOLDER)

os.makedirs(os.path.dirname(USER_XL_PATH), exist_ok=True)

if not os.path.exists(USER_XL_PATH):
    df = pd.DataFrame(columns=["name", "ntid", "login_time"])
    df.to_excel(USER_XL_PATH, index=False)

os.makedirs(PDF_STORE_DIR, exist_ok=True)
os.makedirs(EMBEDDINGS_STORE_DIR, exist_ok=True)
 
# Set page config once
st.set_page_config(page_title="Document Reader System", layout="wide")
 
# -------------------------
# Helper: create a combined FAISS from all PDFs in PDF_STORE_DIR
def build_standards_combined(save_path=STANDARDS_EMBEDDINGS_PATH):
    try:
        pdf_files = [os.path.join(PDF_STORE_DIR, f) for f in os.listdir(PDF_STORE_DIR) if f.lower().endswith(".pdf")]
        if not pdf_files:
            st.warning("No PDFs found in PDF store folder. Place PDFs in PDF_STORE_DIR and try again.")
            return False
 
        all_docs = []
        splitter = RecursiveCharacterTextSplitter(chunk_size=10000, chunk_overlap=100)
 
        for pdf in pdf_files:
            st.write(f"Loading and splitting: {os.path.basename(pdf)}")
            loader = PyPDFLoader(pdf)
            pages = loader.load()
            docs = splitter.split_documents(pages)
            all_docs.extend(docs)
 
        if not all_docs:
            st.error("No text extracted from PDFs (empty documents).")
            return False
 
        st.write("Creating embeddings (this may take a while for many/large PDFs)...")
        embeddings = HuggingFaceEmbeddings(model_name="sentence-transformers/all-MiniLM-L6-v2")
        combined_vs = FAISS.from_documents(all_docs, embedding=embeddings)
 
        # remove old combined folder if present to avoid conflicts
        if os.path.exists(save_path):
            try:
                shutil.rmtree(save_path)
            except Exception as e:
                st.warning(f"Couldn't remove old standards folder: {e}")
 
        os.makedirs(save_path, exist_ok=True)
        combined_vs.save_local(save_path)
 
        st.session_state["standards_vs"] = combined_vs
        st.success(f"Combined standards embeddings created and saved to:\n{save_path}")
        return True
    except Exception as e:
        st.error(f"Error while building standards embeddings: {e}")
        return False
 
 
# -------------------------
# Process a single PDF
def process_pdf(file_path, save_folder):
    try:
        st.write("Step 1: Initializing PyPDFLoader...")
        loader = PyPDFLoader(file_path)
        st.write("Step 2: Loading data...")
        data = loader.load()
        st.write("Step 3: Splitting text...")
        text_splitter = RecursiveCharacterTextSplitter(chunk_size=1000, chunk_overlap=100)
        text_chunks = text_splitter.split_documents(data)
        st.write("Step 4: Initializing embeddings...")
        embeddings = HuggingFaceEmbeddings(model_name="sentence-transformers/all-MiniLM-L6-v2")
        st.write("Step 5: Creating FAISS vector store...")
        vector_store = FAISS.from_documents(text_chunks, embedding=embeddings)
 
        save_path = os.path.join(EMBEDDINGS_STORE_DIR, save_folder)
        os.makedirs(save_path, exist_ok=True)
 
        st.write("Step 6: Saving vector store locally...")
        vector_store.save_local(save_path)
 
        st.success(f"Process completed successfully! Stored in {save_path}")
        return True
    except Exception as e:
        st.error(f"An error occurred: {e}")
        return False
def login_page():
    st.title("🔐 Login")
 
    st.write("Please enter your details to continue")
 
    name = st.text_input("Name")
    ntid = st.text_input("NTID")
 
    if st.button("Login"):
        if name.strip() == "" or ntid.strip() == "":
            st.warning("Both Name and NTID are required")
            return
 
        df = pd.read_excel(USER_XL_PATH)
 
        existing = df[
            (df["name"].str.lower() == name.lower()) &
            (df["ntid"].str.lower() == ntid.lower())
        ]
 
        if existing.empty:
            df = pd.concat(
                [df, pd.DataFrame([{
                    "name": name,
                    "ntid": ntid,
                    "login_time": datetime.now()
                }])],
                ignore_index=True
            )
            df.to_excel(USER_XL_PATH, index=False)
 
        st.session_state.logged_in = True
        st.session_state.user_name = name
        st.session_state.user_ntid = ntid
 
        st.success(f"Welcome {name} 👋")
        st.session_state.page = "reader"
        st.rerun()
 
 
# -------------------------
# Page 1: Document Reader (home)
def document_reader_page():
    st.title("GENAI for EBS")
    st.write("Click a button below to proceed.")
 
    if st.button("Document Loader"):
        st.session_state.page = "loader"
        st.rerun()
 
    if st.button("ISO26262"):
        st.session_state.page = "iso"
        st.rerun()
 
    if st.button("VDA360"):
        st.session_state.page = "vda360"
        st.rerun()
 
    if st.button("VDA305"):
        st.session_state.page = "vda305"
        st.rerun()
 
    if st.button("OBD_J1979-2_202104"):
        st.session_state.page = "obd"
        st.rerun()
 
    if st.button("Standards"):
        st.session_state.page = "standards"
        st.rerun()
 
 
# -------------------------
# Page 2: PDF Loader + QA Chat
def pdf_loader_page():
    st.title("PDF Loader")
    col1, col2 = st.columns([3, 1])
 
    # List existing projects
    with col2:
        st.markdown("### 📂 Uploaded PDFs")
        existing_folders = [d for d in os.listdir(EMBEDDINGS_STORE_DIR)
                            if os.path.isdir(os.path.join(EMBEDDINGS_STORE_DIR, d))]
        if existing_folders:
            for folder in existing_folders:
                st.markdown(f"- {folder}")
        else:
            st.info("No PDFs found")
 
    # Upload or load
    with col1:
        st.write("Answer the question below to proceed.")
        if "prev_choice" not in st.session_state:
            st.session_state.prev_choice = None
 
        choice = st.radio("Is PDF already uploaded?", ["Select an option", "Yes", "No"])
        if choice != st.session_state.prev_choice:
            st.session_state.pdf_processed = False
            st.session_state.messages_loader = []
            st.session_state.current_project = None
        st.session_state.prev_choice = choice
 
        if choice == "Yes":
            if existing_folders:
                selected_folder = st.selectbox("Select an existing PDF/project:", existing_folders)
                if st.button("Load Selected PDF"):
                    st.session_state.current_project = selected_folder
                    st.session_state.pdf_processed = True
                    st.success(f"Loaded embeddings from '{selected_folder}' successfully!")
            else:
                st.info("No uploaded PDFs found. Please select 'No' to upload a new one.")
 
        elif choice == "No":
            st.write("Upload a PDF file to create embeddings and ask questions.")
            project_name = st.text_input("Enter a name for this PDF/project:")
            uploaded_file = st.file_uploader("Browse PDF", type="pdf")
 
            if uploaded_file is not None and project_name.strip() != "":
                file_path = os.path.join(PDF_STORE_DIR, uploaded_file.name)
                with open(file_path, "wb") as f:
                    f.write(uploaded_file.getbuffer())
 
                st.success("File uploaded successfully! Click 'Load' to process it.")
                if st.button("Load") or st.session_state.pdf_processed is False:
                    with st.spinner("Processing PDF and creating embeddings..."):
                        if process_pdf(file_path, project_name):
                            st.session_state.pdf_processed = True
                            st.session_state.current_project = project_name
                        else:
                            st.error("Failed to process the PDF.")
                            return
 
            elif uploaded_file is not None and project_name.strip() == "":
                st.warning("Please enter a project name before loading the PDF.")
 
        else:
            st.info("Please select Yes or No to continue.")
 
        # QA interface
        if st.session_state.get("pdf_processed") and st.session_state.get("current_project"):
            st.markdown("---")
            st.subheader("Ask questions about your PDF")
            query = st.text_input("Type your question here:", key="qa_query")
            if query:
                embeddings = HuggingFaceEmbeddings(model_name="sentence-transformers/all-MiniLM-L6-v2")
                vector_store_path = os.path.join(EMBEDDINGS_STORE_DIR, st.session_state.current_project)
                try:
                    vector_store = FAISS.load_local(vector_store_path, embeddings, allow_dangerous_deserialization=True)
                    retriever = vector_store.as_retriever()
                    response = run_rag(retriever, query)
                    if 'messages_loader' not in st.session_state:
                        st.session_state.messages_loader = []
                    st.session_state.messages_loader.append({"query": query, "response": response})
                except Exception as e:
                    st.error(f"Failed to load embeddings: {e}")
 
            if 'messages_loader' in st.session_state:
                for msg in st.session_state.messages_loader:
                    with st.chat_message("user"):
                        st.markdown(f"**You:** {msg['query']}")
                    with st.chat_message("assistant"):
                        st.markdown(f"**Bot:** {msg['response']}")
 
        st.markdown("---")
        if st.button("Back"):
            st.session_state.page = "reader"
            st.session_state.pdf_processed = False
            st.session_state.messages_loader = []
            st.session_state.current_project = None
            st.rerun()
 
 
# -------------------------
# Generic Chatbot Page Template
def chatbot_page(title, faiss_path, msg_key):
    st.title(title)
 
    embeddings = HuggingFaceEmbeddings(model_name="sentence-transformers/all-MiniLM-L6-v2")
    try:
        faiss_index = FAISS.load_local(faiss_path, embeddings, allow_dangerous_deserialization=True)
        retriever = faiss_index.as_retriever()
    except Exception as e:
        st.error(f"Failed to load FAISS index from {faiss_path}: {e}")
        retriever = None

    if retriever is None:
        st.error("Retriever not available.")
        if st.button("Back"):
            st.session_state.page = "reader"
            st.rerun()
        return

    if msg_key not in st.session_state:
        st.session_state[msg_key] = []

    query = st.text_input("Ask a question about your documents:", "")
    if query:
        response = run_rag(retriever, query)
        st.session_state[msg_key].append({"query": query, "response": response})
 
    for msg in st.session_state[msg_key]:
        with st.chat_message("user"):
            st.markdown(f"**Hello Robot**\n{msg['query']}")
        with st.chat_message("assistant"):
            st.markdown(f"**Hello Human**\n{msg['response']}")
 
    st.markdown("---")
    if st.button("Back"):
        st.session_state.page = "reader"
        st.rerun()
 
# Page 4: Standards (creates/loads combined embeddings from all PDFs)
def standards_page():
    st.title("Standards Chatbot")
    if not st.session_state.get("standards_loaded", False):
        st.session_state.messages_standards = []
        st.session_state.selected_standards = []
    # RIGHT PANEL
    col1, col2 = st.columns([3, 1])
    with col2:
        st.markdown("### 📂 Available Embeddings (Projects)")
 
        existing_folders = [d for d in os.listdir(EMBEDDINGS_STORE_DIR)
                            if os.path.isdir(os.path.join(EMBEDDINGS_STORE_DIR, d))]
 
        selected_folders = st.multiselect(
            "Select one or more projects:",
            existing_folders
        )
 
        if st.button("Load Selected Embeddings"):
            if selected_folders:
                st.session_state.standards_loaded = True
                st.session_state.selected_standards = selected_folders
                st.success(f"Loaded embeddings for: {', '.join(selected_folders)}")
            else:
                st.warning("Please select at least one project.")
 
    # MAIN QA INTERFACE
    with col1:
        if st.session_state.get("standards_loaded", False):
            query = st.text_input("Ask a question about selected PDFs:")
 
            if query:
                # Load all selected FAISS indexes and merge
                embeddings = HuggingFaceEmbeddings(model_name="sentence-transformers/all-MiniLM-L6-v2")
 
                retrievers = []
                for folder in st.session_state.selected_standards:
                    path = os.path.join(EMBEDDINGS_STORE_DIR, folder)
                    try:
                        vs = FAISS.load_local(path, embeddings, allow_dangerous_deserialization=True)
                        retrievers.append(vs.as_retriever())
                    except Exception as e:
                        st.error(f"Failed to load {folder}: {e}")
 
                if retrievers:
                    # Merge retrievers into one
                   
                    combined_retriever = EnsembleRetriever(retrievers=retrievers, weights=[1]*len(retrievers))
 
                    response = run_rag(combined_retriever, query)
 
                    # Store chat history
                    if "messages_standards" not in st.session_state:
                        st.session_state.messages_standards = []
                    st.session_state.messages_standards.append({"query": query, "response": response})
 
            # Show chat history
            if "messages_standards" in st.session_state:
                for msg in st.session_state.messages_standards:
                    with st.chat_message("user"):
                        st.markdown(f"**You:** {msg['query']}")
                    with st.chat_message("assistant"):
                        st.markdown(f"**Bot:** {msg['response']}")
 
    st.markdown("---")
    if st.button("Back"):
        st.session_state.page = "reader"
        st.session_state.standards_loaded = False
        st.session_state.messages_standards = []
        st.session_state.selected_standards = []
        st.rerun()
       
# -------------------------
# Page routing
# if "page" not in st.session_state:
#     st.session_state.page = "reader"
# Auto-login via environment variables (for containerised / dev deployments)
_AUTO_LOGIN_NAME = os.environ.get("AUTO_LOGIN_NAME", "").strip()
_AUTO_LOGIN_NTID = os.environ.get("AUTO_LOGIN_NTID", "").strip()

if "logged_in" not in st.session_state:
    st.session_state.logged_in = False

if not st.session_state.logged_in and _AUTO_LOGIN_NAME and _AUTO_LOGIN_NTID:
    st.session_state.logged_in = True
    st.session_state.user_name = _AUTO_LOGIN_NAME
    st.session_state.user_ntid = _AUTO_LOGIN_NTID

if not st.session_state.logged_in:
    login_page()
    st.stop()
if "page" not in st.session_state:
    st.session_state.page = "reader"
if st.session_state.page == "reader":
    document_reader_page()
elif st.session_state.page == "loader":
    pdf_loader_page()
elif st.session_state.page == "iso":
    chatbot_page("Chatbuddy Application (ISO26262)", ISO_FAISS_PATH, "messages_iso")
elif st.session_state.page == "vda360":
    chatbot_page("Chatbuddy Application (VDA360)", VDA360_FAISS_PATH, "messages_vda360")
elif st.session_state.page == "vda305":
    chatbot_page("Chatbuddy Application (VDA305)", VDA305_FAISS_PATH, "messages_vda305")
elif st.session_state.page == "obd":
    chatbot_page("Chatbuddy Application (OBD_J1979-2_202104)", OBD_FAISS_PATH, "messages_obd")
elif st.session_state.page == "standards":
    standards_page()