Update app.py

app.py

@@ -69,31 +69,51 @@ class AdvancedPdfChatbot:
         self.overall_chain = self.CustomChain(refinement_chain=refinement_chain, qa_chain=qa_chain)
 
     class CustomChain(Chain):
-        def __init__(self, refinement_chain, qa_chain):
-            super().__init__()
-            self.refinement_chain = refinement_chain
-            self.qa_chain = qa_chain
-
-        @property
-        def input_keys(self):
-            return ["query", "chat_history"]
-
-        @property
-        def output_keys(self):
-            return ["answer"]
-
-        def _call(self, inputs):
-            query = inputs['query']
-            chat_history = inputs.get('chat_history', [])
-            
-            # Run the refinement chain to refine the query
-            refined_query = self.refinement_chain.run({'query': query, 'chat_history': chat_history})
-            
-            # Run the QA chain using the refined query and the chat history
-            response = self.qa_chain.run({"question": refined_query, "chat_history": chat_history})
+    def __init__(self, refinement_chain, qa_chain):
+        super().__init__()
+        self.refinement_chain = refinement_chain
+        self.qa_chain = qa_chain
+
+    @property
+    def input_keys(self):
+        return ["query", "chat_history"]
+
+    @property
+    def output_keys(self):
+        return ["answer"]
+
+    def _call(self, inputs):
+        query = inputs['query']
+        chat_history = inputs.get('chat_history', [])
+        
+        # Run the refinement chain to refine the query
+        refined_query = self.refinement_chain.run({'query': query, 'chat_history': chat_history})
+        
+        # Run the QA chain using the refined query and the chat history
+        response = self.qa_chain({'question': refined_query, 'chat_history': chat_history})
+        
+        # Return the answer
+        return {"answer": response['answer']}
+
+
+
+    def setup_conversation_chain(self):
+        if not self.db:
+            raise ValueError("Database not initialized. Please upload a PDF first.")
             
-            # Return the answer
-            return {"answer": response}
+        refinement_chain = LLMChain(
+            llm=self.refinement_llm,
+            prompt=self.refinement_prompt,
+            output_key='refined_query'
+        )
+        qa_chain = ConversationalRetrievalChain.from_llm(
+            self.llm,
+            retriever=self.db.as_retriever(),
+            memory=self.memory,
+            combine_docs_chain_kwargs={"prompt": self.prompt}
+        )
+        self.overall_chain = CustomChain(refinement_chain=refinement_chain, qa_chain=qa_chain)
+
 
     def chat(self, query):
         if not self.overall_chain: