app.py

import os
import time
import threading
import shutil
from pathlib import Path
import gradio as gr
from langchain_community.document_loaders import PyPDFLoader
from langchain.text_splitter import RecursiveCharacterTextSplitter
from langchain.embeddings import HuggingFaceEmbeddings
from langchain_community.vectorstores import FAISS
from langchain.chains import ConversationalRetrievalChain, LLMChain
from langchain.memory import ConversationBufferMemory
from langchain_groq import ChatGroq
from langchain.prompts import PromptTemplate

# Initialize environment variables
GROQ_API_KEY = os.getenv("GROQ_API_KEY")
TEMP_DIR = "temp_uploads"
VECTOR_DIR = "vector_stores"

# Create temporary directories if they don't exist
os.makedirs(TEMP_DIR, exist_ok=True)
os.makedirs(VECTOR_DIR, exist_ok=True)

class DocumentChat:
    def __init__(self):
        self.chain = None
        self.db = None
        self.current_vector_store = None
        self.cleanup_timer = None
        
        # Initialize embedding model
        self.embedding_model = HuggingFaceEmbeddings(
            model_name="sentence-transformers/all-mpnet-base-v2"
        )
        
        # Initialize LLM
        self.llm = ChatGroq(
            api_key=GROQ_API_KEY,
            model_name="deepseek-r1-distill-llama-70b",
            temperature=0.7,
        )
        
        # Initialize memory with output_key
        self.memory = ConversationBufferMemory(
            memory_key="chat_history",
            output_key="answer",
            return_messages=True
        )

        # Add generic question template
        # Add generic question template
        self.generic_template = PromptTemplate(
            input_variables=["question"],
            template="""You are a helpful AI assistant that can:
            1. Read and understand PDF documents that users upload
            2. Answer questions about the contents of uploaded documents
            3. Maintain context through conversation
            4. Process documents and remember their contents for the duration of the chat
            5. Provide accurate and relevant information from the documents

            If the user asks: {question}

            Provide a clear and helpful response about your capabilities in a structured way.
            If the question is about the document and no document is uploaded yet, remind them to upload a document first.

            Remember to be friendly and professional in your response."""
        )
                
        # Create generic chain
        self.generic_chain = LLMChain(
            llm=self.llm,
            prompt=self.generic_template
        )

    def cleanup_files(self, vector_store_path, pdf_path):
        """Clean up files after 10 minutes"""
        time.sleep(600)  # Wait for 10 minutes
        try:
            # Remove vector store
            if os.path.exists(vector_store_path):
                shutil.rmtree(vector_store_path)
            # Remove PDF file
            if os.path.exists(pdf_path):
                os.remove(pdf_path)
            
            # Reset the class variables
            self.chain = None
            self.db = None
            self.current_vector_store = None
            
            print(f"Cleanup completed for: {pdf_path}")
        except Exception as e:
            print(f"Cleanup error: {str(e)}")

    def process_file(self, file_path):
        try:
            if file_path is None:
                return "Please upload a file."
            
            # Generate a unique filename for the vector store
            timestamp = int(time.time())
            vector_store_path = os.path.join(VECTOR_DIR, f"store_{timestamp}")
            
            # Load and process document
            loader = PyPDFLoader(file_path)
            documents = loader.load()
            
            # Split text
            text_splitter = RecursiveCharacterTextSplitter(
                chunk_size=1000,
                chunk_overlap=200
            )
            docs = text_splitter.split_documents(documents)
            
            # Create vector store
            self.db = FAISS.from_documents(docs, self.embedding_model)
            self.db.save_local(vector_store_path)
            self.current_vector_store = vector_store_path
            
            # Create conversation chain
            self.chain = ConversationalRetrievalChain.from_llm(
                llm=self.llm,
                retriever=self.db.as_retriever(
                    search_type="similarity",
                    search_kwargs={"k": 3}
                ),
                memory=self.memory,
                return_source_documents=True,
                combine_docs_chain_kwargs={"prompt": None}
            )
            
            # Start cleanup timer
            self.cleanup_timer = threading.Thread(
                target=self.cleanup_files,
                args=(vector_store_path, file_path)
            )
            self.cleanup_timer.daemon = True
            self.cleanup_timer.start()
            
            return "✅ Document processed successfully! You can now ask questions. Note: The document and its data will be automatically deleted after 10 minutes."
            
        except Exception as e:
            return f"❌ Error processing document: {str(e)}"

    def chat(self, query):
    # List of generic questions to check
        generic_questions = [
            "what can you do?",
            "what are your capabilities?",
            "help",
            "what is this?",
            "how does this work?",
            "what are your functions?",
            "what do you do?",
            "how do i use this?",
            "instructions",
            "guide",
        ]
        
        try:
            # Check if it's a generic question
            if query.lower().strip() in generic_questions:
                # Add document status to response
                status = "I already have a document loaded and ready for questions." if self.chain else "No document is currently loaded. Please upload a PDF document first."
                result = self.generic_chain.run(question=query)
                return f"{result}\n\nCurrent Status: {status}"
            
            # If not generic, process as document question
            if self.chain is None:
                return ("Please upload and process a document first. "
                        "Click the 'Choose Files' button above to upload a PDF document.")
            
            result = self.chain({"question": query})
            return result['answer']
            
        except Exception as e:
            return f"Error processing your question: {str(e)}"

    def reset(self):
        """Reset the chat session"""
        try:
            # Clean up current files if they exist
            if self.current_vector_store and os.path.exists(self.current_vector_store):
                shutil.rmtree(self.current_vector_store)
            
            # Reset all instance variables
            self.chain = None
            self.db = None
            self.current_vector_store = None
            self.memory.clear()
            
            return "Chat session has been reset. You can upload a new document."
        except Exception as e:
            return f"Error resetting chat session: {str(e)}"

def create_demo():
    # Initialize DocumentChat
    doc_chat = DocumentChat()
    
    # Define Gradio interface
    with gr.Blocks(theme=gr.themes.Soft()) as demo:
        gr.Markdown("""
        # 📚 Document Chat Interface
        Upload a PDF document and chat with its contents. Files are automatically deleted after 10 minutes for privacy.
        
        ## Instructions:
        1. Upload a PDF document using the file upload below
        2. Click 'Process Document' and wait for confirmation
        3. Start asking questions about your document
        4. Use 'Reset Chat' to start fresh with a new document
        """)
        
        # Status message
        status_msg = gr.Textbox(label="Status", interactive=False)
        
        with gr.Row():
            with gr.Column(scale=1):
                file_input = gr.File(
                    label="Drop your PDF here",
                    file_types=[".pdf"],
                    type="filepath"
                )
                process_button = gr.Button("📄 Process Document", variant="primary")
                reset_button = gr.Button("🔄 Reset Chat", variant="secondary")
            
            with gr.Column(scale=2):
                chatbot = gr.Chatbot(
                    label="Chat History",
                    height=400,
                    bubble_full_width=False
                )
                msg = gr.Textbox(
                    label="Your Question",
                    placeholder="Ask something about the document or type 'help' for assistance...",
                    lines=2
                )
                send = gr.Button("🚀 Send", variant="primary")
        
        # Event handlers
        def user_message(message, history):
            if not message.strip():
                return "", history
            
            response = doc_chat.chat(message)
            history.append((message, response))
            return "", history

        def reset_chat():
            result = doc_chat.reset()
            return result, None  # Clear chatbot history

        process_button.click(
            fn=doc_chat.process_file,
            inputs=[file_input],
            outputs=[status_msg]
        )
        
        reset_button.click(
            fn=reset_chat,
            inputs=[],
            outputs=[status_msg, chatbot]
        )
        
        # Add both submit methods
        msg.submit(
            fn=user_message,
            inputs=[msg, chatbot],
            outputs=[msg, chatbot]
        )
        
        send.click(
            fn=user_message,
            inputs=[msg, chatbot],
            outputs=[msg, chatbot]
        )
    
    return demo

if __name__ == "__main__":
    demo = create_demo()
    demo.launch(
        share=True,
        show_error=True,
        max_threads=40  # Replace enable_queue with max_threads
    )