Update app.py

app.py

@@ -1,140 +1,107 @@
 import os
 import gradio as gr
-import logging
 from langchain.document_loaders import PyPDFLoader
 from langchain.text_splitter import RecursiveCharacterTextSplitter
 from langchain.embeddings import OpenAIEmbeddings
 from langchain.vectorstores import FAISS
+from langchain.chains import ConversationalRetrievalChain
 from langchain.chat_models import ChatOpenAI
-from langchain.chains import LLMChain
 from langchain.memory import ConversationBufferMemory
-from langchain.prompts import PromptTemplate
-from PyPDF2 import PdfReader
-import spaces
 
+from langchain.prompts import PromptTemplate
 
-class ContextAwareResponseGenerator:
-    def __init__(self, llm):
-        self.llm = llm
-        self.response_prompt = PromptTemplate(
-            input_variables=['context', 'query', 'chat_history'],
-            template="""Given the context, query, and chat history, generate the best response that is clear and helpful. Use structured responses in various formats such as Paragraphs, Headlines and subtexts, bullet points, Sections.
 
-Context: {context}
-Query: {query}
-Chat History: {chat_history}
 
-Choose the most appropriate response structure and generate the response directly, without explicit guidance on which format to use. Your response should be based on the query and context provided."""  
-        )
-        self.response_chain = LLMChain(llm=self.llm, prompt=self.response_prompt)
-    @spaces.GPU
-    def generate_response(self, context, query, chat_history=''):
-        try:
-            # Generate the response content with structure handled by the LLM itself
-            response = self.response_chain.run({
-                'context': context,
-                'query': query,
-                'chat_history': chat_history or "No previous context"
-            })
-            return response.strip()  # LLM decides on the structure
-        except Exception as e:
-            logging.error(f"Response generation error: {e}")
-            return self._default_response(query)
-
-    def _default_response(self, query):
-        return f"I couldn't generate a response for: {query}"
+openai_api_key = os.environ.get("OPENAI_API_KEY")
 
 class AdvancedPdfChatbot:
-    @spaces.GPU
     def __init__(self, openai_api_key):
         os.environ["OPENAI_API_KEY"] = openai_api_key
-        self.llm = ChatOpenAI(temperature=0.2, model_name='gpt-4o')
-        
         self.embeddings = OpenAIEmbeddings()
         self.text_splitter = RecursiveCharacterTextSplitter(chunk_size=1000, chunk_overlap=200)
+        self.llm =  ChatOpenAI(temperature=0,model_name='gpt-4o-mini')
         
         self.memory = ConversationBufferMemory(memory_key="chat_history", return_messages=True)
-        self.response_generator = ContextAwareResponseGenerator(self.llm)
-        
-        self.db = None
-        self.document_context = ""
-        
-    @spaces.GPU
+        self.qa_chain = None
+        self.pdf_path = None
+        self.template = """
+Imagine you are a chat assistant for knowledge retrieval, specializing in providing detailed information with a deep understanding of context.
+Your goal is to generate responses in a structured format that is both informative and engaging.
+
+        """
+        self.prompt = PromptTemplate(template=self.template, input_variables=["context", "question"])
+
     def load_and_process_pdf(self, pdf_path):
-        try:
-            reader = PdfReader(pdf_path)
-            metadata = {
-                "title": reader.metadata.get("/Title", "Untitled"),
-                "author": reader.metadata.get("/Author", "Unknown")
-            }
-            
-            loader = PyPDFLoader(pdf_path)
-            documents = loader.load()
-            texts = self.text_splitter.split_documents(documents)
-            
-            self.db = FAISS.from_documents(texts[:50], self.embeddings)
-            self.document_context = f"Document: {metadata['title']} by {metadata['author']}"
-            
-            return True
-        except Exception as e:
-            logging.error(f"PDF processing error: {e}")
-            return False
-            
-    @spaces.GPU
-    def chat(self, query, is_new_question=False):
-        if not self.db:
-            return "Please upload a PDF first."
-        
-        # Retrieve chat history
-        chat_history = self.memory.load_memory_variables({}).get('chat_history', [])
-        
-        # Reset chat history for new questions
-        if is_new_question:
-            chat_history = []  # For new questions, reset the chat history
-        
-        # Generate context-aware response
-        response = self.response_generator.generate_response(
-            context=self.document_context, 
-            query=query, 
-            chat_history=str(chat_history)
+        loader = PyPDFLoader(pdf_path)
+        documents = loader.load()
+        texts = self.text_splitter.split_documents(documents)
+        self.db = FAISS.from_documents(texts, self.embeddings)
+        self.pdf_path = pdf_path
+        self.setup_conversation_chain()
+
+    def setup_conversation_chain(self):
+        self.qa_chain = ConversationalRetrievalChain.from_llm(
+            self.llm,
+            retriever=self.db.as_retriever(),
+            memory=self.memory,
+            combine_docs_chain_kwargs={"prompt": self.prompt}
         )
-        
-        # Store conversation in memory
-        self.memory.save_context({"input": query}, {"output": response})
-        
-        return response
 
-# Gradio Interface
-pdf_chatbot = AdvancedPdfChatbot(os.environ.get("OPENAI_API_KEY"))
+    def chat(self, query):
+        if not self.qa_chain:
+            return "Please upload a PDF first."
+        result = self.qa_chain({"question": query})
+        return result['answer']
+
+    def get_pdf_path(self):
+        # Return the stored PDF path
+        if self.pdf_path:
+            return self.pdf_path
+        else:
+            return "No PDF uploaded yet."
+
+# Initialize the chatbot
+pdf_chatbot = AdvancedPdfChatbot(openai_api_key)
 
 def upload_pdf(pdf_file):
-    if not pdf_file:
-        return "Upload a PDF file."
-    file_path = pdf_file.name if hasattr(pdf_file, 'name') else pdf_file
-    return "PDF processed successfully" if pdf_chatbot.load_and_process_pdf(file_path) else "Processing failed"
+    if pdf_file is None:
+        return "Please upload a PDF file."
+    file_path = pdf_file.name
+    pdf_chatbot.load_and_process_pdf(file_path)
+    return file_path
 
 def respond(message, history):
-    try:
-        is_new_question = len(history) == 0  # If history is empty, it's a new question
-        bot_message = pdf_chatbot.chat(message, is_new_question)
-        history.append((message, bot_message))
-        return "", history
-    except Exception as e:
-        return f"Error: {e}", history
-
-# Gradio UI
+    bot_message = pdf_chatbot.chat(message)
+    history.append((message, bot_message))
+    return "", history
+
+def clear_chatbot():
+    pdf_chatbot.memory.clear()
+    return []
+
+def get_pdf_path():
+    # Call the method to return the current PDF path
+    return pdf_chatbot.get_pdf_path()
+
+# Create the Gradio interface
 with gr.Blocks() as demo:
-    gr.Markdown("# Advanced PDF Chatbot")
+    gr.Markdown("# PDF Chatbot")
+    
     with gr.Row():
         pdf_upload = gr.File(label="Upload PDF", file_types=[".pdf"])
         upload_button = gr.Button("Process PDF")
 
     upload_status = gr.Textbox(label="Upload Status")
     upload_button.click(upload_pdf, inputs=[pdf_upload], outputs=[upload_status])
-    
+    path_button = gr.Button("Get PDF Path")
+    pdf_path_display = gr.Textbox(label="Current PDF Path")
     chatbot_interface = gr.Chatbot()
-    msg = gr.Textbox(placeholder="Enter your query...")
+    msg = gr.Textbox()
+    clear = gr.Button("Clear")
+
     msg.submit(respond, inputs=[msg, chatbot_interface], outputs=[msg, chatbot_interface])
+    clear.click(clear_chatbot, outputs=[chatbot_interface])
+    path_button.click(get_pdf_path, outputs=[pdf_path_display])
 
 if __name__ == "__main__":
     demo.launch()