app.py

app.py

@@ -1,64 +1,292 @@
+import os
+import time
+import threading
+import shutil
+from pathlib import Path
 import gradio as gr
-from huggingface_hub import InferenceClient
-
-"""
-For more information on `huggingface_hub` Inference API support, please check the docs: https://huggingface.co/docs/huggingface_hub/v0.22.2/en/guides/inference
-"""
-client = InferenceClient("HuggingFaceH4/zephyr-7b-beta")
-
-
-def respond(
-    message,
-    history: list[tuple[str, str]],
-    system_message,
-    max_tokens,
-    temperature,
-    top_p,
-):
-    messages = [{"role": "system", "content": system_message}]
-
-    for val in history:
-        if val[0]:
-            messages.append({"role": "user", "content": val[0]})
-        if val[1]:
-            messages.append({"role": "assistant", "content": val[1]})
-
-    messages.append({"role": "user", "content": message})
-
-    response = ""
-
-    for message in client.chat_completion(
-        messages,
-        max_tokens=max_tokens,
-        stream=True,
-        temperature=temperature,
-        top_p=top_p,
-    ):
-        token = message.choices[0].delta.content
-
-        response += token
-        yield response
-
-
-"""
-For information on how to customize the ChatInterface, peruse the gradio docs: https://www.gradio.app/docs/chatinterface
-"""
-demo = gr.ChatInterface(
-    respond,
-    additional_inputs=[
-        gr.Textbox(value="You are a friendly Chatbot.", label="System message"),
-        gr.Slider(minimum=1, maximum=2048, value=512, step=1, label="Max new tokens"),
-        gr.Slider(minimum=0.1, maximum=4.0, value=0.7, step=0.1, label="Temperature"),
-        gr.Slider(
-            minimum=0.1,
-            maximum=1.0,
-            value=0.95,
-            step=0.05,
-            label="Top-p (nucleus sampling)",
-        ),
-    ],
-)
+from langchain.document_loaders import PyPDFLoader
+from langchain.text_splitter import RecursiveCharacterTextSplitter
+from langchain.embeddings import HuggingFaceEmbeddings
+from langchain.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 = "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"
+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="mixtral-8x7b-32768",
+            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.launch()
+    demo = create_demo()
+    demo.launch(
+        share=True,
+        show_error=True,
+        max_threads=40  # Replace enable_queue with max_threads
+    )
+
+
+
+
+
+