Update app.py

app.py

@@ -57,8 +57,7 @@ class AdvancedPdfChatbot:
             
         refinement_chain = LLMChain(
             llm=self.refinement_llm,
-            prompt=self.refinement_prompt,
-            output_key='refined_query'
+            prompt=self.refinement_prompt
         )
         qa_chain = ConversationalRetrievalChain.from_llm(
             self.llm,
@@ -68,20 +67,6 @@ class AdvancedPdfChatbot:
         )
         self.overall_chain = self.CustomChain(refinement_chain=refinement_chain, qa_chain=qa_chain)
 
-        def setup_conversation_chain(self):
-            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 = self.CustomChain(refinement_chain=refinement_chain, qa_chain=qa_chain)
-
     class CustomChain(Chain):
         def __init__(self, refinement_chain, qa_chain):
             super().__init__()
@@ -90,24 +75,28 @@ class AdvancedPdfChatbot:
     
         @property
         def input_keys(self):
+            """Define the input keys that this chain expects."""
             return ["query", "chat_history"]
     
         @property
         def output_keys(self):
+            """Define the output keys that this chain returns."""
             return ["answer"]
     
         def _call(self, inputs):
             query = inputs['query']
             chat_history = inputs.get('chat_history', [])
             
-            refined_query = self.refinement_chain.run({'query': query, 'chat_history': chat_history})
-            response = self.qa_chain({'question': refined_query, 'chat_history': chat_history})
+            # Run the refinement chain to refine the query
+            refinement_inputs = {'query': query, 'chat_history': chat_history}
+            refined_query = self.refinement_chain.run(refinement_inputs)
+            
+            # Run the QA chain using the refined query and the chat history
+            qa_inputs = {"question": refined_query, "chat_history": chat_history}
+            response = self.qa_chain(qa_inputs)
             
             return {"answer": response['answer']}
 
-
-
-
     def chat(self, query):
         if not self.overall_chain:
             return "Please upload a PDF first."