Spaces:

Vishnu-bot
/

pdf-chatbot

Sleeping

Update app.py

e4f7344 verified 4 months ago

9.48 kB

import streamlit as st
import os
import time
import hashlib
from io import BytesIO
from langchain_groq import ChatGroq
from langchain_huggingface import HuggingFaceEmbeddings
from langchain_text_splitters import RecursiveCharacterTextSplitter
from langchain.chains import create_retrieval_chain, create_history_aware_retriever
from langchain.chains.combine_documents import create_stuff_documents_chain
from langchain_community.vectorstores import FAISS
from langchain_community.document_loaders import PyPDFLoader
from langchain_core.prompts import ChatPromptTemplate, MessagesPlaceholder
from langchain_core.messages import HumanMessage, AIMessage
from langchain_core.documents import Document

# ==========================================
# 10. SYSTEM CONFIGURATIONS & CONSTANTS
# ==========================================
EMBEDDING_MODEL = "sentence-transformers/all-MiniLM-L6-v2"
GROQ_MODEL = "llama-3.3-70b-versatile"
CHUNK_SIZE = 700
CHUNK_OVERLAP = 100
MAX_HISTORY_PAIRS = 10  # Keeps up to 20 messages total (10 user, 10 assistant)

# Retriever Configuration (MMR)
SEARCH_TYPE = "mmr"
K_DOCS = 5
FETCH_K = 20

# ==========================================
# 11. SECURE API KEY MANAGEMENT
# ==========================================
# Hugging Face Spaces automatically exposes repo secrets to os.environ
groq_api_key = os.environ.get("GROQ_API_KEY")

if not groq_api_key:
    # Fallback for local testing if not deployed to HF Spaces yet
    from dotenv import load_dotenv
    load_dotenv()
    groq_api_key = os.getenv("GROQ_API_KEY")

# Initialize LLM
if groq_api_key:
    llm = ChatGroq(groq_api_key=groq_api_key, model_name=GROQ_MODEL)
else:
    llm = None

st.set_page_config(page_title="Production RAG Chatbot", layout="wide")
st.title("📚 Advanced RAG Document Q&A with Chat History")

# ==========================================
# 2. CACHED EMBEDDING MODEL LOADING
# ==========================================
@st.cache_resource
def load_embedding_model():
    """Loads and caches the embedding model globally across sessions."""
    return HuggingFaceEmbeddings(
        model_name=EMBEDDING_MODEL,
        model_kwargs={"device": "cpu"},
        encode_kwargs={"normalize_embeddings": True}
    )

embeddings = load_embedding_model()

# ==========================================
# 1. FILE HASHING & PROCESSING
# ==========================================
def calculate_files_hash(uploaded_files):
    """Generates a unique SHA-256 hash combining all uploaded files' contents."""
    hasher = hashlib.sha256()
    for uploaded_file in sorted(uploaded_files, key=lambda x: x.name):
        # Read file chunks to handle large files efficiently without memory spikes
        bytes_data = uploaded_file.getvalue()
        hasher.update(bytes_data)
    return hasher.hexdigest()

def process_uploaded_files(uploaded_files):
    """Loads, splits, and embeds all pages from uploaded PDFs."""
    all_docs = []
    
    for uploaded_file in uploaded_files:
        # Streamlit files exist in memory; write temporarily to parse via PyPDFLoader
        temp_filename = f"temp_{uploaded_file.name}"
        with open(temp_filename, "wb") as f:
            f.write(uploaded_file.getvalue())
        
        try:
            loader = PyPDFLoader(temp_filename)
            # 4. Process ALL documents (removed [:50] restriction)
            for page in loader.load():
                # Add human-readable metadata for tracking source info
                page.metadata["source"] = uploaded_file.name
                all_docs.append(page)
        finally:
            if os.path.exists(temp_filename):
                os.remove(temp_filename)
                
    text_splitter = RecursiveCharacterTextSplitter(
        chunk_size=CHUNK_SIZE, 
        chunk_overlap=CHUNK_OVERLAP
    )
    final_documents = text_splitter.split_documents(all_docs)
    
    # Create the FAISS Vector Database
    vector_store = FAISS.from_documents(final_documents, embeddings)
    return vector_store

# ==========================================
# 3. CHAIN INITIALIZATION (BUILD ONCE)
# ==========================================
def build_retrieval_chain(vector_store):
    """Builds and caches the entire LangChain execution pipeline."""
    # 5. Configure Retriever with Maximal Marginal Relevance (MMR)
    retriever = vector_store.as_retriever(
        search_type=SEARCH_TYPE,
        search_kwargs={"k": K_DOCS, "fetch_k": FETCH_K}
    )
    
    contextualize_q_prompt = ChatPromptTemplate.from_messages([
        ("system", "Given a chat history and the latest user question which might reference context in the chat history, formulate a standalone question which can be understood without the chat history. Do NOT answer the question, just reformulate it if needed and otherwise return it as is."),
        MessagesPlaceholder(variable_name="chat_history"),
        ("human", "{input}")
    ])
    
    qa_prompt = ChatPromptTemplate.from_messages([
        ("system", "Answer the user's questions based ONLY on the provided context. Provide the most accurate response.\n\nContext:\n{context}"),
        MessagesPlaceholder(variable_name="chat_history"),
        ("human", "{input}")
    ])
    
    history_aware_retriever = create_history_aware_retriever(llm, retriever, contextualize_q_prompt)
    document_chain = create_stuff_documents_chain(llm, qa_prompt)
    
    return create_retrieval_chain(history_aware_retriever, document_chain)

# ==========================================
# STREAMLIT STATE & INTERFACE MANAGEMENT
# ==========================================
if "chat_history" not in st.session_state:
    st.session_state.chat_history = []
if "current_file_hash" not in st.session_state:
    st.session_state.current_file_hash = None
if "retreival_chain" not in st.session_state:
    st.session_state.retreival_chain = None

# Sidebar for file upload
with st.sidebar:
    st.header("Document Management")
    uploaded_files = st.file_uploader(
        "Upload your Research Papers (PDF)", 
        type=["pdf"], 
        accept_multiple_files=True
    )
    
    if not groq_api_key:
        st.error("GROQ_API_KEY is missing. Please configure it in your Space Secrets environment.")

# Handle Vector Database compilation conditionally
if uploaded_files and llm:
    new_hash = calculate_files_hash(uploaded_files)
    
    # Trigger processing only if files have changed or chain hasn't been built yet
    if st.session_state.current_file_hash != new_hash:
        with st.sidebar:
            with st.spinner("Processing new documents..."):
                try:
                    vs = process_uploaded_files(uploaded_files)
                    # Compile and store the chain once
                    st.session_state.retreival_chain = build_retrieval_chain(vs)
                    st.session_state.current_file_hash = new_hash
                    st.success("Vector Store & Chains Ready!")
                except Exception as e:
                    st.error(f"Error parsing documents: {e}")

# Display historic chat items
for message in st.session_state.chat_history:
    role = "user" if isinstance(message, HumanMessage) else "assistant"
    with st.chat_message(role):
        st.write(message.content)

# Chat execution block
if user_prompt := st.chat_input("Ask something about your uploaded documents..."):
    if not uploaded_files:
        st.warning("Please upload at least one PDF document first.")
    elif not st.session_state.retreival_chain:
        st.error("Retrieval pipeline is not initialized.")
    else:
        # User input visual representation
        with st.chat_message("user"):
            st.write(user_prompt)
            
        with st.chat_message("assistant"):
            with st.spinner("Analyzing context..."):
                # 7. Use perf_counter to accurately capture real elapsed turnaround time
                start_time = time.perf_counter()
                
                response = st.session_state.retreival_chain.invoke({
                    'input': user_prompt,
                    'chat_history': st.session_state.chat_history
                })
                
                elapsed_time = time.perf_counter() - start_time
                
                st.write(response['answer'])
                st.caption(f"Response compiled in {elapsed_time:.2f} seconds")
                
                # 9. Improved Reference Tracker with metadata source & page tracking
                with st.expander("View Reference Passages"):
                    for i, doc in enumerate(response['context']):
                        source_file = doc.metadata.get("source", "Unknown Document")
                        # Page index typically starts at 0 internally; add 1 for user reading clarity
                        page_num = doc.metadata.get("page", 0) + 1
                        
                        st.markdown(f"**Source {i+1}:** {source_file} (Page {page_num})")
                        st.write(doc.page_content)
                        st.divider()
                        
        # 8. Slide window memory to prevent unbounded chat history growth
        st.session_state.chat_history.append(HumanMessage(content=user_prompt))
        st.session_state.chat_history.append(AIMessage(content=response['answer']))
        
        if len(st.session_state.chat_history) > (MAX_HISTORY_PAIRS * 2):
            st.session_state.chat_history = st.session_state.chat_history[-(MAX_HISTORY_PAIRS * 2):]

	import streamlit as st
	import os
	import time
	import hashlib
	from io import BytesIO
	from langchain_groq import ChatGroq
	from langchain_huggingface import HuggingFaceEmbeddings
	from langchain_text_splitters import RecursiveCharacterTextSplitter
	from langchain.chains import create_retrieval_chain, create_history_aware_retriever
	from langchain.chains.combine_documents import create_stuff_documents_chain
	from langchain_community.vectorstores import FAISS
	from langchain_community.document_loaders import PyPDFLoader
	from langchain_core.prompts import ChatPromptTemplate, MessagesPlaceholder
	from langchain_core.messages import HumanMessage, AIMessage
	from langchain_core.documents import Document

	# ==========================================
	# 10. SYSTEM CONFIGURATIONS & CONSTANTS
	# ==========================================
	EMBEDDING_MODEL = "sentence-transformers/all-MiniLM-L6-v2"
	GROQ_MODEL = "llama-3.3-70b-versatile"
	CHUNK_SIZE = 700
	CHUNK_OVERLAP = 100
	MAX_HISTORY_PAIRS = 10 # Keeps up to 20 messages total (10 user, 10 assistant)

	# Retriever Configuration (MMR)
	SEARCH_TYPE = "mmr"
	K_DOCS = 5
	FETCH_K = 20

	# ==========================================
	# 11. SECURE API KEY MANAGEMENT
	# ==========================================
	# Hugging Face Spaces automatically exposes repo secrets to os.environ
	groq_api_key = os.environ.get("GROQ_API_KEY")

	if not groq_api_key:
	# Fallback for local testing if not deployed to HF Spaces yet
	from dotenv import load_dotenv
	load_dotenv()
	groq_api_key = os.getenv("GROQ_API_KEY")

	# Initialize LLM
	if groq_api_key:
	llm = ChatGroq(groq_api_key=groq_api_key, model_name=GROQ_MODEL)
	else:
	llm = None

	st.set_page_config(page_title="Production RAG Chatbot", layout="wide")
	st.title("📚 Advanced RAG Document Q&A with Chat History")

	# ==========================================
	# 2. CACHED EMBEDDING MODEL LOADING
	# ==========================================
	@st.cache_resource
	def load_embedding_model():
	"""Loads and caches the embedding model globally across sessions."""
	return HuggingFaceEmbeddings(
	model_name=EMBEDDING_MODEL,
	model_kwargs={"device": "cpu"},
	encode_kwargs={"normalize_embeddings": True}
	)

	embeddings = load_embedding_model()

	# ==========================================
	# 1. FILE HASHING & PROCESSING
	# ==========================================
	def calculate_files_hash(uploaded_files):
	"""Generates a unique SHA-256 hash combining all uploaded files' contents."""
	hasher = hashlib.sha256()
	for uploaded_file in sorted(uploaded_files, key=lambda x: x.name):
	# Read file chunks to handle large files efficiently without memory spikes
	bytes_data = uploaded_file.getvalue()
	hasher.update(bytes_data)
	return hasher.hexdigest()

	def process_uploaded_files(uploaded_files):
	"""Loads, splits, and embeds all pages from uploaded PDFs."""
	all_docs = []

	for uploaded_file in uploaded_files:
	# Streamlit files exist in memory; write temporarily to parse via PyPDFLoader
	temp_filename = f"temp_{uploaded_file.name}"
	with open(temp_filename, "wb") as f:
	f.write(uploaded_file.getvalue())

	try:
	loader = PyPDFLoader(temp_filename)
	# 4. Process ALL documents (removed [:50] restriction)
	for page in loader.load():
	# Add human-readable metadata for tracking source info
	page.metadata["source"] = uploaded_file.name
	all_docs.append(page)
	finally:
	if os.path.exists(temp_filename):
	os.remove(temp_filename)

	text_splitter = RecursiveCharacterTextSplitter(
	chunk_size=CHUNK_SIZE,
	chunk_overlap=CHUNK_OVERLAP
	)
	final_documents = text_splitter.split_documents(all_docs)

	# Create the FAISS Vector Database
	vector_store = FAISS.from_documents(final_documents, embeddings)
	return vector_store

	# ==========================================
	# 3. CHAIN INITIALIZATION (BUILD ONCE)
	# ==========================================
	def build_retrieval_chain(vector_store):
	"""Builds and caches the entire LangChain execution pipeline."""
	# 5. Configure Retriever with Maximal Marginal Relevance (MMR)
	retriever = vector_store.as_retriever(
	search_type=SEARCH_TYPE,
	search_kwargs={"k": K_DOCS, "fetch_k": FETCH_K}
	)

	contextualize_q_prompt = ChatPromptTemplate.from_messages([
	("system", "Given a chat history and the latest user question which might reference context in the chat history, formulate a standalone question which can be understood without the chat history. Do NOT answer the question, just reformulate it if needed and otherwise return it as is."),
	MessagesPlaceholder(variable_name="chat_history"),
	("human", "{input}")
	])

	qa_prompt = ChatPromptTemplate.from_messages([
	("system", "Answer the user's questions based ONLY on the provided context. Provide the most accurate response.\n\nContext:\n{context}"),
	MessagesPlaceholder(variable_name="chat_history"),
	("human", "{input}")
	])

	history_aware_retriever = create_history_aware_retriever(llm, retriever, contextualize_q_prompt)
	document_chain = create_stuff_documents_chain(llm, qa_prompt)

	return create_retrieval_chain(history_aware_retriever, document_chain)

	# ==========================================
	# STREAMLIT STATE & INTERFACE MANAGEMENT
	# ==========================================
	if "chat_history" not in st.session_state:
	st.session_state.chat_history = []
	if "current_file_hash" not in st.session_state:
	st.session_state.current_file_hash = None
	if "retreival_chain" not in st.session_state:
	st.session_state.retreival_chain = None

	# Sidebar for file upload
	with st.sidebar:
	st.header("Document Management")
	uploaded_files = st.file_uploader(
	"Upload your Research Papers (PDF)",
	type=["pdf"],
	accept_multiple_files=True
	)

	if not groq_api_key:
	st.error("GROQ_API_KEY is missing. Please configure it in your Space Secrets environment.")

	# Handle Vector Database compilation conditionally
	if uploaded_files and llm:
	new_hash = calculate_files_hash(uploaded_files)

	# Trigger processing only if files have changed or chain hasn't been built yet
	if st.session_state.current_file_hash != new_hash:
	with st.sidebar:
	with st.spinner("Processing new documents..."):
	try:
	vs = process_uploaded_files(uploaded_files)
	# Compile and store the chain once
	st.session_state.retreival_chain = build_retrieval_chain(vs)
	st.session_state.current_file_hash = new_hash
	st.success("Vector Store & Chains Ready!")
	except Exception as e:
	st.error(f"Error parsing documents: {e}")

	# Display historic chat items
	for message in st.session_state.chat_history:
	role = "user" if isinstance(message, HumanMessage) else "assistant"
	with st.chat_message(role):
	st.write(message.content)

	# Chat execution block
	if user_prompt := st.chat_input("Ask something about your uploaded documents..."):
	if not uploaded_files:
	st.warning("Please upload at least one PDF document first.")
	elif not st.session_state.retreival_chain:
	st.error("Retrieval pipeline is not initialized.")
	else:
	# User input visual representation
	with st.chat_message("user"):
	st.write(user_prompt)

	with st.chat_message("assistant"):
	with st.spinner("Analyzing context..."):
	# 7. Use perf_counter to accurately capture real elapsed turnaround time
	start_time = time.perf_counter()

	response = st.session_state.retreival_chain.invoke({
	'input': user_prompt,
	'chat_history': st.session_state.chat_history
	})

	elapsed_time = time.perf_counter() - start_time

	st.write(response['answer'])
	st.caption(f"Response compiled in {elapsed_time:.2f} seconds")

	# 9. Improved Reference Tracker with metadata source & page tracking
	with st.expander("View Reference Passages"):
	for i, doc in enumerate(response['context']):
	source_file = doc.metadata.get("source", "Unknown Document")
	# Page index typically starts at 0 internally; add 1 for user reading clarity
	page_num = doc.metadata.get("page", 0) + 1

	st.markdown(f"Source {i+1}: {source_file} (Page {page_num})")
	st.write(doc.page_content)
	st.divider()

	# 8. Slide window memory to prevent unbounded chat history growth
	st.session_state.chat_history.append(HumanMessage(content=user_prompt))
	st.session_state.chat_history.append(AIMessage(content=response['answer']))

	if len(st.session_state.chat_history) > (MAX_HISTORY_PAIRS * 2):
	st.session_state.chat_history = st.session_state.chat_history[-(MAX_HISTORY_PAIRS * 2):]