app.py

import streamlit as st
import os
import tempfile
from dotenv import load_dotenv
import PyPDF2
import faiss
import numpy as np
from sentence_transformers import SentenceTransformer
from langchain.text_splitter import RecursiveCharacterTextSplitter
from langchain.embeddings import HuggingFaceEmbeddings
from langchain.llms import HuggingFacePipeline
from langchain.chains import RetrievalQA
from langchain.vectorstores import FAISS
from langchain.document_loaders import PyPDFLoader
from transformers import pipeline, AutoTokenizer, AutoModelForCausalLM
import torch

# Load environment variables
load_dotenv()

# Configuration from environment variables
EMBEDDING_MODEL = os.getenv("EMBEDDING_MODEL", "sentence-transformers/all-MiniLM-L6-v2")
LLM_MODEL = os.getenv("LLM_MODEL", "microsoft/DialoGPT-medium")
CHUNK_SIZE = int(os.getenv("CHUNK_SIZE", "1000"))
CHUNK_OVERLAP = int(os.getenv("CHUNK_OVERLAP", "200"))
MAX_TOKENS = int(os.getenv("MAX_TOKENS", "256"))
TEMPERATURE = float(os.getenv("TEMPERATURE", "0.7"))
STREAMLIT_PORT = int(os.getenv("STREAMLIT_SERVER_PORT", "8501"))

def validate_environment():
    """Validate environment variables and return status"""
    status = {
        "valid": True,
        "warnings": [],
        "errors": []
    }
    
    # Check if .env file exists
    if not os.path.exists(".env"):
        status["warnings"].append("No .env file found. Using default values.")
    
    # Validate numeric values
    if CHUNK_SIZE <= 0:
        status["errors"].append("CHUNK_SIZE must be positive")
        status["valid"] = False
    
    if CHUNK_OVERLAP < 0:
        status["errors"].append("CHUNK_OVERLAP must be non-negative")
        status["valid"] = False
    
    if MAX_TOKENS <= 0:
        status["errors"].append("MAX_TOKENS must be positive")
        status["valid"] = False
    
    if not (0 <= TEMPERATURE <= 2):
        status["warnings"].append("TEMPERATURE should be between 0 and 2")
    
    return status

# Configure Streamlit page
st.set_page_config(
    page_title="RAG PDF Chatbot",
    page_icon="📚",
    layout="wide",
    initial_sidebar_state="expanded"
)

# Custom CSS for better UI
st.markdown("""
<style>
    .main-header {
        font-size: 3rem;
        color: #1f77b4;
        text-align: center;
        margin-bottom: 2rem;
    }
    .chat-message {
        padding: 1rem;
        border-radius: 0.5rem;
        margin: 1rem 0;
        color: #333333;
        font-size: 16px;
        line-height: 1.5;
    }
    .user-message {
        background-color: #f0f8ff;
        border-left: 4px solid #2196f3;
        color: #1a1a1a;
    }
    .bot-message {
        background-color: #f5f5f5;
        border-left: 4px solid #4caf50;
        color: #1a1a1a;
    }
    .sidebar-content {
        background-color: #f8f9fa;
        padding: 1rem;
        border-radius: 0.5rem;
    }
    .chat-message strong {
        color: #2c3e50;
        font-weight: 600;
    }
</style>
""", unsafe_allow_html=True)

class RAGChatbot:
    def __init__(self):
        self.embeddings = None
        self.vectorstore = None
        self.qa_chain = None
        self.llm = None
        self.documents = []
        
    def initialize_models(self):
        """Initialize the embedding model and LLM"""
        try:
            # Initialize embeddings using sentence transformers (no API token needed)
            self.embeddings = HuggingFaceEmbeddings(
                model_name=EMBEDDING_MODEL,
                model_kwargs={'device': 'cpu'},
                encode_kwargs={'normalize_embeddings': True}
            )
            
            # Initialize LLM using local transformers (no API token needed)
            tokenizer = AutoTokenizer.from_pretrained(LLM_MODEL)
            model = AutoModelForCausalLM.from_pretrained(
                LLM_MODEL,
                torch_dtype=torch.float16 if torch.cuda.is_available() else torch.float32,
                device_map="auto" if torch.cuda.is_available() else None
            )
            
            # Create pipeline
            pipe = pipeline(
                "text-generation",
                model=model,
                tokenizer=tokenizer,
                max_new_tokens=MAX_TOKENS,
                temperature=TEMPERATURE,
                do_sample=True,
                pad_token_id=tokenizer.eos_token_id,
                truncation=True
            )
            
            self.llm = HuggingFacePipeline(pipeline=pipe)
            
            return True
        except Exception as e:
            st.error(f"Error initializing models: {str(e)}")
            return False
    
    def process_pdf(self, pdf_file):
        """Process uploaded PDF file and create vector store"""
        try:
            # Save uploaded file temporarily
            with tempfile.NamedTemporaryFile(delete=False, suffix=".pdf") as tmp_file:
                tmp_file.write(pdf_file.read())
                tmp_file_path = tmp_file.name
            
            # Load PDF using PyPDFLoader
            loader = PyPDFLoader(tmp_file_path)
            self.documents = loader.load()
            
            # Split documents into chunks
            text_splitter = RecursiveCharacterTextSplitter(
                chunk_size=CHUNK_SIZE,
                chunk_overlap=CHUNK_OVERLAP,
                length_function=len,
            )
            
            texts = text_splitter.split_documents(self.documents)
            
            # Create vector store
            self.vectorstore = FAISS.from_documents(texts, self.embeddings)
            
            # Create custom prompt template for better understanding
            from langchain.prompts import PromptTemplate
            
            custom_prompt = PromptTemplate(
                template="""Context: {context}

Question: {question}

Answer:""",
                input_variables=["context", "question"]
            )
            
            # Create QA chain with enhanced retrieval
            self.qa_chain = RetrievalQA.from_chain_type(
                llm=self.llm,
                chain_type="stuff",
                retriever=self.vectorstore.as_retriever(
                    search_type="similarity_score_threshold",
                    search_kwargs={
                        "k": 5,  # Increased to get more relevant documents
                        "score_threshold": 0.3  # Lower threshold for more flexible matching
                    }
                ),
                return_source_documents=True,
                chain_type_kwargs={"prompt": custom_prompt}
            )
            
            # Clean up temporary file
            os.unlink(tmp_file_path)
            
            return True, len(texts)
            
        except Exception as e:
            st.error(f"Error processing PDF: {str(e)}")
            return False, 0
    
    def ask_question(self, question):
        """Ask a question and get answer from the RAG system"""
        try:
            if self.qa_chain is None:
                return "Please upload a PDF file first.", []
            
            # Truncate question if it's too long
            if len(question) > 500:
                question = question[:500] + "..."
            
            # Use simple similarity search for better results
            simple_retriever = self.vectorstore.as_retriever(search_kwargs={"k": 5})
            docs = simple_retriever.get_relevant_documents(question)
            
            if not docs:
                return "I couldn't find relevant information in the document to answer your question. Please try rephrasing or asking about a different topic.", []
            
            # Use the LLM directly with retrieved documents
            context = "\n\n".join([doc.page_content for doc in docs])
            
            # Try multiple prompt approaches for better results
            prompts_to_try = [
                f"""Here's some information: {context}

Question: {question}

Answer:""",
                f"""Based on this text: {context}

{question}

Response:""",
                f"""Context: {context}

Q: {question}
A:""",
                f"""Information: {context}

{question}

Answer based on the information above:"""
            ]
            
            answer = ""
            for prompt in prompts_to_try:
                try:
                    response = self.llm(prompt)
                    answer = response.strip()
                    
                    # Check if we got a good answer
                    if (len(answer) > 20 and 
                        answer.lower().strip() != question.lower().strip() and
                        not answer.lower().startswith(question.lower())):
                        break
                except:
                    continue
            
            # If still no good answer, use the last attempt
            if not answer or len(answer) < 10:
                answer = response.strip() if 'response' in locals() else ""
            source_docs = docs
            
            # Clean up the answer more aggressively
            # Remove everything before the first meaningful content
            lines = answer.split('\n')
            cleaned_lines = []
            found_content = False
            
            for line in lines:
                line = line.strip()
                
                # Skip empty lines at the beginning
                if not line and not found_content:
                    continue
                
                # Skip template markers
                if line.lower() in ['context:', 'question:', 'answer:']:
                    continue
                
                # Skip lines that are just template text
                if any(template_text in line.lower() for template_text in [
                    'use the following pieces',
                    'answer the following question',
                    'based on the provided context',
                    'according to the context',
                    'the context shows that',
                    'from the context'
                ]):
                    continue
                
                # If we find actual content, start collecting
                if line and not any(template_word in line.lower() for template_word in ['context:', 'question:', 'answer:']):
                    found_content = True
                    cleaned_lines.append(line)
            
            answer = '\n'.join(cleaned_lines).strip()
            
            # If answer is empty or too short, provide a fallback
            if len(answer) < 10 or answer.lower().strip() == question.lower().strip():
                # Provide a summary of the context as fallback
                context_summary = context[:500] + "..." if len(context) > 500 else context
                answer = f"Based on the document content, here's what I found: {context_summary}"
                
                # If still no good answer, provide a generic response
                if len(answer) < 20:
                    answer = "I found relevant information in the document, but I'm having trouble formatting the response. Please try rephrasing your question."
            
            return answer, source_docs
            
        except Exception as e:
            error_msg = str(e)
            if "max_length" in error_msg or "max_new_tokens" in error_msg:
                return "The question or context is too long. Please try asking a shorter, more specific question.", []
            else:
                return f"Error getting answer: {error_msg}", []

def main():
    # Header
    st.markdown('<h1 class="main-header">PDF PARSER</h1>', unsafe_allow_html=True)
    st.markdown("Upload a PDF file and ask questions about its content!")
    
    # Validate environment
    env_status = validate_environment()
    if not env_status["valid"]:
        st.error("❌ Environment configuration errors:")
        for error in env_status["errors"]:
            st.error(f"• {error}")
        return
    
    if env_status["warnings"]:
        for warning in env_status["warnings"]:
            st.warning(f"⚠️ {warning}")
    
    # Initialize session state
    if "chatbot" not in st.session_state:
        st.session_state.chatbot = RAGChatbot()
        st.session_state.chat_history = []
        st.session_state.pdf_processed = False
    
    # Sidebar for file upload and settings
    with st.sidebar:
        st.markdown('<div class="sidebar-content">', unsafe_allow_html=True)
        st.header("📁 Upload PDF")
        
        uploaded_file = st.file_uploader(
            "Choose a PDF file",
            type="pdf",
            help="Upload a PDF file to start chatting about its content"
        )
        
        if uploaded_file is not None:
            if st.button("Process PDF", type="primary"):
                with st.spinner("Processing PDF and initializing models..."):
                    # Initialize models if not already done
                    if st.session_state.chatbot.embeddings is None:
                        if not st.session_state.chatbot.initialize_models():
                            st.error("Failed to initialize models. Please check your configuration.")
                            return
                    
                    # Process PDF
                    success, num_chunks = st.session_state.chatbot.process_pdf(uploaded_file)
                    
                    if success:
                        st.session_state.pdf_processed = True
                        st.success(f"✅ PDF processed successfully! Created {num_chunks} text chunks.")
                        st.session_state.chat_history = []  # Clear chat history
                    else:
                        st.error("❌ Failed to process PDF. Please try again.")
        
        st.markdown("</div>", unsafe_allow_html=True)
        
    
    # Main chat interface
    if not st.session_state.pdf_processed:
        st.info("👈 Please upload and process a PDF file using the sidebar to start chatting!")
        
        # Show example questions
        st.markdown("### 💡 Example Questions You Can Ask:")
        example_questions = [
            "What is the main topic of this document?",
            "Can you summarize the key points?",
            "What are the important findings or conclusions?",
            "Are there any specific recommendations mentioned?",
            "What methodology was used in this study?",
            "Tell me about the results",
            "What does this document say about...?",
            "Explain the main concepts",
            "What are the key takeaways?",
            "How does this relate to...?"
        ]
        
        for i, question in enumerate(example_questions, 1):
            st.markdown(f"{i}. {question}")
        
        st.info("💡 **Tip**: Ask questions in any way you like - the bot understands context and relevance!")
    
    else:
        # Chat interface
        st.markdown("### 💬 Chat with your PDF")
        
        # Display chat history
        for message in st.session_state.chat_history:
            if message["role"] == "user":
                st.markdown(f"""
                <div class="chat-message user-message">
                    <strong>You:</strong> {message["content"]}
                </div>
                """, unsafe_allow_html=True)
            else:
                st.markdown(f"""
                <div class="chat-message bot-message">
                    <strong>🤖 Assistant:</strong> {message["content"]}
                </div>
                """, unsafe_allow_html=True)
                
        
        # Chat input
        user_question = st.chat_input("Ask a question about the PDF content...")
        
        if user_question:
            # Add user message to chat history
            st.session_state.chat_history.append({
                "role": "user",
                "content": user_question
            })
            
            # Get answer from chatbot
            with st.spinner("Thinking..."):
                answer, sources = st.session_state.chatbot.ask_question(user_question)
            
            # Add bot response to chat history
            st.session_state.chat_history.append({
                "role": "assistant",
                "content": answer,
                "sources": sources
            })
            
            # Rerun to display new messages
            st.rerun()
        
        # Clear chat button
        if st.button("🗑️ Clear Chat History"):
            st.session_state.chat_history = []
            st.rerun()

if __name__ == "__main__":
    main()