medibot.py

import os
import tempfile
import streamlit as st
from dotenv import load_dotenv

# Document Loaders (Verified standard LangChain Community imports)
from langchain_community.document_loaders import (
    PyMuPDFLoader, 
    TextLoader, 
    CSVLoader, 
    Docx2txtLoader
)
from langchain_text_splitters import RecursiveCharacterTextSplitter

# Embeddings & Vector Store
from langchain_huggingface import HuggingFaceEmbeddings
from langchain_community.vectorstores import FAISS

# LLM & Chains
from langchain_core.prompts import PromptTemplate
from langchain_groq import ChatGroq
from langchain.chains import create_retrieval_chain
from langchain.chains.combine_documents import create_stuff_documents_chain

# Load environment variables (Local development backup)
load_dotenv()

# Configuration
DB_FAISS_PATH = "vectorstore/db_faiss"
os.makedirs("vectorstore", exist_ok=True)

# --- Helper Functions ---

def get_document_loader(file_path):
    """Select the correct loader based on file extension."""
    ext = file_path.split('.')[-1].lower()
    if ext == 'pdf':
        return PyMuPDFLoader(file_path)  # High-speed C-based parser
    elif ext == 'docx':
        return Docx2txtLoader(file_path)
    elif ext == 'csv':
        return CSVLoader(file_path)
    elif ext == 'txt':
        return TextLoader(file_path, encoding='utf-8')
    else:
        raise ValueError(f"Unsupported file type: {ext}")

@st.cache_resource
def get_embedding_model():
    return HuggingFaceEmbeddings(model_name='sentence-transformers/all-MiniLM-L6-v2')

def get_vectorstore():
    """Loads the existing FAISS database into session state for instant access."""
    if 'vectorstore' in st.session_state and st.session_state.vectorstore is not None:
        return st.session_state.vectorstore
        
    if os.path.exists(DB_FAISS_PATH) and os.listdir(DB_FAISS_PATH):
        try:
            # FIXED: Explicitly named keyword argument for embeddings
            db = FAISS.load_local(
                folder_path=DB_FAISS_PATH, 
                embeddings=get_embedding_model(), 
                allow_dangerous_deserialization=True
            )
            st.session_state.vectorstore = db
            return db
        except Exception:
            return None
    return None

# --- Main Application ---

def main():
    st.set_page_config(page_title="RAG Document Assistant", layout="wide")
    
    # Initialize Session State
    if 'messages' not in st.session_state:
        st.session_state.messages = []
    if 'processed_files' not in st.session_state:
        st.session_state.processed_files = []
    if 'total_chunks' not in st.session_state:
        st.session_state.total_chunks = 0
    if 'vectorstore' not in st.session_state:
        st.session_state.vectorstore = None

    # ==========================================
    # SIDEBAR: File Upload & Processing
    # ==========================================
    with st.sidebar:
        st.title("📚 RAG Assistant Panel")
        
        uploaded_files = st.file_uploader(
            "Upload Documents", 
            type=["pdf", "docx", "txt", "csv"], 
            accept_multiple_files=True
        )
        
        if st.button("Process & Index Documents", type="primary"):
            if uploaded_files:
                all_documents = []
                newly_processed = []
                
                with st.spinner("Reading and indexing files incredibly fast..."):
                    with tempfile.TemporaryDirectory() as temp_dir:
                        for uploaded_file in uploaded_files:
                            temp_file_path = os.path.join(temp_dir, uploaded_file.name)
                            
                            with open(temp_file_path, "wb") as f:
                                f.write(uploaded_file.getbuffer())
                            
                            try:
                                loader = get_document_loader(temp_file_path)
                                docs = loader.load()
                                all_documents.extend(docs)
                                newly_processed.append(uploaded_file.name)
                            except Exception as e:
                                st.error(f"Failed to process {uploaded_file.name}: {e}")
                    
                    if all_documents:
                        text_splitter = RecursiveCharacterTextSplitter(chunk_size=1000, chunk_overlap=100)
                        text_chunks = text_splitter.split_documents(all_documents)
                        
                        embeddings = get_embedding_model()
                        existing_db = get_vectorstore()
                        
                        if existing_db is not None:
                            new_db = FAISS.from_documents(text_chunks, embeddings)
                            existing_db.merge_from(new_db)
                            existing_db.save_local(DB_FAISS_PATH)
                            st.session_state.vectorstore = existing_db
                        else:
                            db = FAISS.from_documents(text_chunks, embeddings)
                            db.save_local(DB_FAISS_PATH)
                            st.session_state.vectorstore = db
                        
                        st.session_state.total_chunks += len(text_chunks)
                        for file_name in newly_processed:
                            if file_name not in st.session_state.processed_files:
                                st.session_state.processed_files.append(file_name)
                                
                        st.success("Database successfully updated!")
            else:
                st.warning("Please upload files first.")
                
        st.markdown("---")
        
        st.subheader("📊 Database Stats")
        st.metric(label="Total Chunks Stored", value=st.session_state.total_chunks)
        
        st.markdown("**Processed Files:**")
        if not st.session_state.processed_files:
            st.caption("No files processed yet in this session.")
        for file_name in st.session_state.processed_files:
            st.caption(f"✅ {file_name}")

    # ==========================================
    # MAIN WINDOW: Chat Interface
    # ==========================================
    st.title("📖 RAG Document Assistant")
    st.caption("Upload documents in the sidebar and ask questions based on their context.")

    for message in st.session_state.messages:
        st.chat_message(message['role']).markdown(message['content'])

    user_prompt = st.chat_input("Ask Questions About Your Documents")

    if user_prompt:
        st.chat_message('user').markdown(user_prompt)
        st.session_state.messages.append({'role': 'user', 'content': user_prompt})
        
        vectorstore = get_vectorstore()
        
        if vectorstore is None:
            warning_msg = "Please upload and process a document in the sidebar first!"
            st.chat_message('assistant').markdown(warning_msg)
            st.session_state.messages.append({'role': 'assistant', 'content': warning_msg})
        else:
            try:
                # FIXED: Seamless compatibility handling for Streamlit Secrets or local environment variables
                GROQ_API_KEY = st.secrets.get("GROQ_API_KEY", os.environ.get("GROQ_API_KEY"))
                
                llm = ChatGroq(
                    model="llama-3.1-8b-instant",
                    temperature=0.5,
                    max_tokens=512,
                    api_key=GROQ_API_KEY,
                )

                CUSTOM_PROMPT_TEMPLATE = """
                Use the pieces of information provided in the context to answer the user's question.
                If you don't know the answer, just say that you don't know, don't try to make up an answer. 
                Don't provide anything out of the given context.

                Context: {context}
                Question: {input}

                Start the answer directly. No small talk please.
                """
                prompt_template = PromptTemplate(
                    template=CUSTOM_PROMPT_TEMPLATE, 
                    input_variables=["context", "input"]
                )

                combine_docs_chain = create_stuff_documents_chain(llm, prompt_template)
                retriever = vectorstore.as_retriever(search_kwargs={'k': 3})
                rag_chain = create_retrieval_chain(retriever, combine_docs_chain)

                with st.spinner("Searching documents..."):
                    response = rag_chain.invoke({'input': user_prompt})

                result = response["answer"]
                source_documents = response["context"]
                
                source_metadata = [f"- Source: `{doc.metadata.get('source', 'Unknown')}` | Page: `{doc.metadata.get('page', 'N/A')}`" for doc in source_documents]
                result_to_show = f"{result}\n\n**Source Docs:**\n" + "\n".join(set(source_metadata))
                
                st.chat_message('assistant').markdown(result_to_show)
                st.session_state.messages.append({'role': 'assistant', 'content': result_to_show})

            except Exception as e:
                st.error(f"Error during retrieval: {str(e)}")

if __name__ == "__main__":
    main()