experimental_files/app.py

"""
Streamlit App for RAG Pipeline with PDF Upload
This app allows users to upload PDF files, process them, and query them using RAG.
"""

import streamlit as st
import os
import tempfile
from pathlib import Path
from typing import List

# Import RAG pipeline components
from src.rag_pipeline import (
    process_pdfs_in_directory,
    documents_chunking,
    EmbeddingModel,
    VectorStore,
    RagRetriever,
    create_groq_llm,
    rag_pipeline,
    rag_pipeline_with_memory,
    summarize_answer,
)

# Page configuration
st.set_page_config(
    page_title="RAG Pipeline - PDF Query System",
    page_icon="📚",
    layout="wide"
)

# Initialize session state
if 'vectorstore' not in st.session_state:
    st.session_state.vectorstore = None
if 'retriever' not in st.session_state:
    st.session_state.retriever = None
if 'llm' not in st.session_state:
    st.session_state.llm = None
if 'embedding_manager' not in st.session_state:
    st.session_state.embedding_manager = None
if 'documents_processed' not in st.session_state:
    st.session_state.documents_processed = False
if 'chunked_documents' not in st.session_state:
    st.session_state.chunked_documents = None
if 'chat_history' not in st.session_state:
    st.session_state.chat_history = []


def initialize_components():
    """Initialize RAG components if not already initialized."""
    if st.session_state.embedding_manager is None:
        with st.spinner("Loading embedding model..."):
            st.session_state.embedding_manager = EmbeddingModel()
    
    if st.session_state.llm is None:
        try:
            with st.spinner("Initializing Groq LLM..."):
                st.session_state.llm = create_groq_llm()
        except ValueError as e:
            st.error(f"Error initializing LLM: {e}")
            st.info("Please make sure GROQ_API_KEY is set in your .env file.")
            return False
    
    return True


def process_uploaded_pdfs(uploaded_files, chunk_size, chunk_overlap):
    """Process uploaded PDF files."""
    if not uploaded_files:
        st.warning("Please upload at least one PDF file.")
        return False
    
    # Create temporary directory for uploaded files
    with tempfile.TemporaryDirectory() as temp_dir:
        # Save uploaded files to temporary directory
        for uploaded_file in uploaded_files:
            file_path = os.path.join(temp_dir, uploaded_file.name)
            with open(file_path, "wb") as f:
                f.write(uploaded_file.getbuffer())
        
        # Process PDFs
        with st.spinner("Loading PDF documents..."):
            documents = process_pdfs_in_directory(temp_dir)
        
        if not documents:
            st.error("No documents were loaded. Please check your PDF files.")
            return False
        
        st.success(f"Loaded {len(documents)} document(s)")
        
        # Chunk documents
        with st.spinner(f"Chunking documents (chunk_size={chunk_size}, overlap={chunk_overlap})..."):
            chunked_documents = documents_chunking(
                documents,
                chunk_size=chunk_size,
                chunk_overlap=chunk_overlap
            )
        
        st.session_state.chunked_documents = chunked_documents
        
        # Initialize components
        if not initialize_components():
            return False
        
        # Generate embeddings
        with st.spinner("Generating embeddings..."):
            texts = [doc.page_content for doc in chunked_documents]
            embeddings = st.session_state.embedding_manager.generate_embedding(texts)
        
        # Initialize or get vector store
        if st.session_state.vectorstore is None:
            st.session_state.vectorstore = VectorStore()
        
        # Add documents to vector store
        with st.spinner("Adding documents to vector store..."):
            st.session_state.vectorstore.add_documents(
                documents=chunked_documents,
                embeddings=embeddings
            )
        
        # Initialize or recreate retriever (to ensure it has latest method signature)
        st.session_state.retriever = RagRetriever(
            vector_store=st.session_state.vectorstore,
            embedding_manager=st.session_state.embedding_manager
        )
        
        st.session_state.documents_processed = True
        return True


def main():
    """Main Streamlit app."""
    st.title("📚 RAG Pipeline - PDF Query System")
    st.markdown("Upload PDF files and query them using Retrieval-Augmented Generation (RAG)")
    
    # Sidebar for configuration
    with st.sidebar:
        st.header("⚙️ Configuration")
        
        st.subheader("Chunking Parameters")
        chunk_size = st.slider(
            "Chunk Size",
            min_value=200,
            max_value=2000,
            value=800,
            step=100,
            help="Size of each text chunk in characters"
        )
        chunk_overlap = st.slider(
            "Chunk Overlap",
            min_value=0,
            max_value=500,
            value=200,
            step=50,
            help="Number of overlapping characters between chunks"
        )
        
        st.subheader("Query Parameters")
        top_k = st.slider(
            "Top K Results",
            min_value=1,
            max_value=10,
            value=5,
            help="Number of document chunks to retrieve"
        )
        
        st.subheader("🔍 Metadata Filters")
        st.caption("Filter documents by metadata for faster, more accurate retrieval")
        use_filters = st.checkbox("Enable Metadata Filtering", value=False)
        
        metadata_filters = {}
        if use_filters and st.session_state.vectorstore:
            try:
                # Get sample metadata to show available fields
                sample_results = st.session_state.vectorstore.collection.get(limit=1)
                if sample_results.get('metadatas') and len(sample_results['metadatas']) > 0:
                    sample_meta = sample_results['metadatas'][0]
                    available_fields = list(sample_meta.keys())
                    
                    # Filter by source file
                    if 'source' in available_fields or 'source_file' in available_fields:
                        source_field = 'source' if 'source' in available_fields else 'source_file'
                        # Get unique sources
                        all_results = st.session_state.vectorstore.collection.get()
                        unique_sources = set()
                        for meta in all_results.get('metadatas', []):
                            source = meta.get(source_field) or meta.get('source_file')
                            if source:
                                unique_sources.add(source)
                        
                        if unique_sources:
                            selected_sources = st.multiselect(
                                "Filter by Source File",
                                options=sorted(unique_sources),
                                help="Select one or more source files to search in"
                            )
                            if selected_sources:
                                # Always store as list for consistent handling
                                metadata_filters[source_field] = selected_sources if isinstance(selected_sources, list) else [selected_sources]
                    
                    # Filter by page number
                    if 'page' in available_fields:
                        page_filter = st.text_input(
                            "Filter by Page Number (optional)",
                            placeholder="e.g., 1, 2, 3 or leave empty",
                            help="Enter page numbers separated by commas"
                        )
                        if page_filter:
                            try:
                                pages = [int(p.strip()) for p in page_filter.split(',')]
                                metadata_filters['page'] = pages if len(pages) > 1 else pages[0]
                            except:
                                st.warning("Invalid page number format")
                    
                    # Custom metadata filter
                    with st.expander("Advanced: Custom Metadata Filter"):
                        filter_key = st.text_input("Metadata Key", placeholder="e.g., author, title")
                        filter_value = st.text_input("Metadata Value", placeholder="e.g., John Doe")
                        if filter_key and filter_value:
                            metadata_filters[filter_key] = filter_value
            except Exception as e:
                st.warning(f"Could not load metadata filters: {str(e)}")
        
        st.divider()
        
        st.subheader("System Status")
        if st.session_state.documents_processed:
            st.success("✅ Documents Processed")
            if st.session_state.chunked_documents:
                st.info(f"📄 {len(st.session_state.chunked_documents)} chunks available")
            if st.session_state.vectorstore:
                try:
                    count = st.session_state.vectorstore.collection.count()
                    st.info(f"🗄️ {count} documents in vector store")
                except:
                    st.warning("⚠️ Could not check vector store count")
        else:
            st.info("⏳ No documents processed")
        
        if st.button("🔄 Reset System", type="secondary"):
            st.session_state.vectorstore = None
            st.session_state.retriever = None
            st.session_state.llm = None
            st.session_state.embedding_manager = None
            st.session_state.documents_processed = False
            st.session_state.chunked_documents = None
            st.session_state.chat_history = []
            st.rerun()
        
        # Reinitialize retriever if it exists but doesn't have the new method signature
        if st.session_state.retriever and st.session_state.vectorstore and st.session_state.embedding_manager:
            import inspect
            sig = inspect.signature(st.session_state.retriever.retrieve)
            if 'metadata_filters' not in sig.parameters:
                st.session_state.retriever = RagRetriever(
                    vector_store=st.session_state.vectorstore,
                    embedding_manager=st.session_state.embedding_manager
                )
        
        st.divider()
        
        if st.button("🗑️ Clear Chat History", type="secondary"):
            st.session_state.chat_history = []
            st.rerun()
    
    # Main content area
    tab1, tab2 = st.tabs(["📤 Upload & Process", "💬 Chat"])
    
    with tab1:
        st.header("Upload PDF Files")
        st.markdown("Upload one or more PDF files to process and add to the knowledge base.")
        
        uploaded_files = st.file_uploader(
            "Choose PDF files",
            type=['pdf'],
            accept_multiple_files=True,
            help="You can upload multiple PDF files at once"
        )
        
        if uploaded_files:
            st.info(f"📎 {len(uploaded_files)} file(s) selected")
            
            # Display file names
            with st.expander("View uploaded files"):
                for file in uploaded_files:
                    st.write(f"- {file.name} ({file.size:,} bytes)")
            
            if st.button("🚀 Process Documents", type="primary"):
                success = process_uploaded_pdfs(uploaded_files, chunk_size, chunk_overlap)
                if success:
                    st.balloons()
                    st.success("✅ Documents processed successfully! You can now chat with them in the Chat tab.")
                else:
                    st.error("❌ Failed to process documents. Please check the error messages above.")
    
    with tab2:
        st.header("💬 Chat with Documents")
        st.markdown("Ask questions about the uploaded PDF documents. The chat remembers previous conversations.")
        
        if not st.session_state.documents_processed:
            st.warning("⚠️ Please upload and process documents first in the 'Upload & Process' tab.")
        else:
            # Display chat history
            chat_container = st.container()
            with chat_container:
                if st.session_state.chat_history:
                    for message in st.session_state.chat_history:
                        role = message.get("role", "user")
                        content = message.get("content", "")
                        
                        if role == "user":
                            with st.chat_message("user"):
                                st.write(content)
                        elif role == "assistant":
                            with st.chat_message("assistant"):
                                st.write(content)
                                # Show sources if available
                                if "sources" in message:
                                    with st.expander("📄 Sources"):
                                        for i, source in enumerate(message["sources"], 1):
                                            st.markdown(f"**Source {i}** (Score: {source.get('score', 0):.4f})")
                                            st.caption(f"Preview: {source.get('preview', '')[:200]}...")
                else:
                    st.info("👋 Start a conversation by asking a question below!")
            
            # Chat input
            query = st.chat_input(
                "Ask a question about the documents...",
                key="chat_input"
            )
            
            # Handle query
            if query:
                if st.session_state.retriever and st.session_state.llm:
                    # Add user message to chat history
                    st.session_state.chat_history.append({
                        "role": "user",
                        "content": query
                    })
                    
                    with st.spinner("Thinking..."):
                        try:
                            # First, retrieve documents to check if we have results
                            # Use score_threshold=0 to get all results, even with low similarity scores
                            # Handle both old and new method signatures
                            try:
                                results = st.session_state.retriever.retrieve(
                                    query=query,
                                    top_k=top_k,
                                    score_threshold=0,  # Get all results, even with low scores
                                    metadata_filters=metadata_filters if use_filters else None
                                )
                            except TypeError:
                                # Fallback for old method signature (without metadata_filters)
                                results = st.session_state.retriever.retrieve(
                                    query=query,
                                    top_k=top_k,
                                    score_threshold=0
                                )
                            
                            # Debug: Show retrieval info
                            if not results:
                                st.warning(f"⚠️ No documents retrieved. Vector store has {st.session_state.vectorstore.collection.count()} documents.")
                            
                            # Prepare sources for display
                            sources = [{
                                "score": r.get("score", 0),
                                "preview": r.get("document", "")[:300] + "..."
                            } for r in results] if results else []
                            
                            # Get answer using RAG pipeline with memory
                            answer = rag_pipeline_with_memory(
                                query=query,
                                retriever=st.session_state.retriever,
                                llm=st.session_state.llm,
                                conversation_history=st.session_state.chat_history[:-1],  # Exclude current query
                                top_k=top_k,
                                metadata_filters=metadata_filters if use_filters else None
                            )
                            
                            # Create concise summary for memory (store full answer in message, summary in history)
                            concise_answer = summarize_answer(answer, st.session_state.llm, max_length=150)
                            
                            # Add assistant response to chat history
                            # Store full answer for display, but concise version for memory
                            st.session_state.chat_history.append({
                                "role": "assistant",
                                "content": answer,  # Full answer for display
                                "concise": concise_answer,  # Concise version for memory
                                "sources": sources
                            })
                            
                            st.rerun()
                            
                        except Exception as e:
                            st.error(f"Error processing query: {str(e)}")
                            st.exception(e)
                            # Remove the user message if there was an error
                            if st.session_state.chat_history and st.session_state.chat_history[-1]["role"] == "user":
                                st.session_state.chat_history.pop()
                else:
                    st.error("System not properly initialized. Please process documents first.")
            
            # Example queries
            if not st.session_state.chat_history:
                st.divider()
                st.subheader("💡 Example Queries")
                example_queries = [
                    "What is the main topic of the document?",
                    "Summarize the key points",
                    "What are the main findings?",
                ]
                
                cols = st.columns(len(example_queries))
                for i, example in enumerate(example_queries):
                    with cols[i]:
                        if st.button(f"📝 {example[:30]}...", key=f"example_{i}"):
                            # Simulate chat input
                            st.session_state.chat_history.append({
                                "role": "user",
                                "content": example
                            })
                            st.rerun()


if __name__ == "__main__":
    main()