Update app.py

app.py

@@ -5,10 +5,44 @@ from langchain.text_splitter import RecursiveCharacterTextSplitter
 from langchain.embeddings import OpenAIEmbeddings
 from langchain.vectorstores import FAISS
 from langchain.chat_models import ChatOpenAI
-from langchain.chains import ConversationalRetrievalChain
+from langchain.chains import ConversationalRetrievalChain, LLMChain
 from langchain.memory import ConversationBufferMemory
 from langchain.prompts import PromptTemplate
 
+class QueryRefiner:
+    def __init__(self):
+        self.refinement_llm = ChatOpenAI(temperature=0.2, model_name='gpt-3.5-turbo')
+        self.refinement_prompt = PromptTemplate(
+            input_variables=['query', 'context'],
+            template="""Refine and enhance the following query for maximum clarity and precision:
+
+Original Query: {query}
+Document Context: {context}
+
+Enhanced Query Requirements:
+- Clarify any ambiguous terms
+- Add specific context-driven details
+- Ensure precise information retrieval
+- Restructure for optimal comprehension
+
+Refined Query:"""
+        )
+        self.refinement_chain = LLMChain(
+            llm=self.refinement_llm, 
+            prompt=self.refinement_prompt
+        )
+
+    def refine_query(self, original_query, context_hints=''):
+        try:
+            refined_query = self.refinement_chain.run({
+                'query': original_query, 
+                'context': context_hints or "General academic document"
+            })
+            return refined_query.strip()
+        except Exception as e:
+            print(f"Query refinement error: {e}")
+            return original_query
+
 class AdvancedPdfChatbot:
     def __init__(self, openai_api_key):
         os.environ["OPENAI_API_KEY"] = openai_api_key
@@ -17,21 +51,18 @@ class AdvancedPdfChatbot:
         self.llm = ChatOpenAI(temperature=0, model_name='gpt-4')
         
         self.memory = ConversationBufferMemory(memory_key="chat_history", return_messages=True)
+        self.query_refiner = QueryRefiner()
         self.db = None
         self.chain = None
         
-        self.template = """
-        You are a study partner assistant helping students analyze PDF documents.
-        
-        Answer the question based only on the most recent provided resources.
-        Provide the most relevant and concise answer possible. Give a structured response in parts when needed according to complexity of the question and details needed, have headlines or bullet points only when necessary
-        
-        Context: {context}
-        Question: {question}
-        Answer:
-        """
         self.qa_prompt = PromptTemplate(
-            template=self.template, 
+            template="""You are an expert academic assistant analyzing a document.
+
+Context: {context}
+Question: {question}
+
+Provide a comprehensive, precise answer based strictly on the document's content.
+If the answer isn't directly available, explain why.""",
             input_variables=["context", "question"]
         )
 
@@ -43,7 +74,7 @@ class AdvancedPdfChatbot:
         
         self.chain = ConversationalRetrievalChain.from_llm(
             llm=self.llm,
-            retriever=self.db.as_retriever(),
+            retriever=self.db.as_retriever(search_kwargs={"k": 3}),
             memory=self.memory,
             combine_docs_chain_kwargs={"prompt": self.qa_prompt}
         )
@@ -52,13 +83,26 @@ class AdvancedPdfChatbot:
         if not self.chain:
             return "Please upload a PDF first."
         
-        result = self.chain({"question": query})
+        context_hints = self._extract_document_type()
+        refined_query = self.query_refiner.refine_query(query, context_hints)
+        
+        result = self.chain({"question": refined_query})
         return result['answer']
 
+    def _extract_document_type(self):
+        """Extract basic document characteristics"""
+        if not self.db:
+            return ""
+        try:
+            first_doc = list(self.db.docstore._dict.values())[0].page_content[:500]
+            return f"Document appears to cover: {first_doc[:100]}..."
+        except:
+            return "Academic/technical document"
+
     def clear_memory(self):
         self.memory.clear()
 
-# Gradio interface setup remains mostly the same
+# Gradio Interface
 pdf_chatbot = AdvancedPdfChatbot(os.environ.get("OPENAI_API_KEY"))
 
 def upload_pdf(pdf_file):
@@ -85,9 +129,9 @@ def clear_chatbot():
     pdf_chatbot.clear_memory()
     return []
 
-# Gradio interface
+# Gradio UI
 with gr.Blocks() as demo:
-    gr.Markdown("# PDF Chatbot")
+    gr.Markdown("# Advanced PDF Chatbot")
     
     with gr.Row():
         pdf_upload = gr.File(label="Upload PDF", file_types=[".pdf"])
@@ -95,9 +139,13 @@ with gr.Blocks() as demo:
 
     upload_status = gr.Textbox(label="Upload Status")
     upload_button.click(upload_pdf, inputs=[pdf_upload], outputs=[upload_status])
+    
     chatbot_interface = gr.Chatbot()
-    msg = gr.Textbox()
+    msg = gr.Textbox(placeholder="Enter your query...")
     msg.submit(respond, inputs=[msg, chatbot_interface], outputs=[msg, chatbot_interface])
+    
+    clear_button = gr.Button("Clear Conversation")
+    clear_button.click(clear_chatbot, outputs=[chatbot_interface])
 
 if __name__ == "__main__":
     demo.launch()