Update app.py

app.py

@@ -4,8 +4,9 @@ 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, Chain
-from langchain.chat_models import ChatOpenAI
+from langchain.chains.base import Chain
+from langchain.llms import ChatOpenAI
+from langchain.chains import LLMChain, ConversationalRetrievalChain
 from langchain.memory import ConversationBufferMemory
 from langchain.prompts import PromptTemplate
 
@@ -51,50 +52,39 @@ class AdvancedPdfChatbot:
         self.setup_conversation_chain()
 
     def setup_conversation_chain(self):
-        class CustomChain(Chain):
-            refinement_chain: Chain
-            qa_chain: Chain
-
-            @classmethod
-            def from_llms(cls, refinement_llm, qa_llm, retriever, memory, prompt):
-                refinement_chain = Chain(
-                    llm_chain=LLMChain(
-                        llm=refinement_llm,
-                        prompt=self.refinement_prompt,
-                        output_key='refined_query'
-                    )
-                )
-                qa_chain = ConversationalRetrievalChain.from_llm(
-                    qa_llm,
-                    retriever=retriever,
-                    memory=memory,
-                    combine_docs_chain_kwargs={"prompt": prompt}
-                )
-                return cls(refinement_chain=refinement_chain, qa_chain=qa_chain)
-
-            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})
-                self.qa_chain.memory.save_context({"input": query}, {"output": response['answer']})
-                return {"answer": response['answer']}
-
-            @property
-            def input_keys(self):
-                return ['query', 'chat_history']
-
-            @property
-            def output_keys(self):
-                return ['answer']
-
-        self.overall_chain = CustomChain.from_llms(
-            refinement_llm=self.refinement_llm,
-            qa_llm=self.llm,
+        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,
-            prompt=self.prompt
+            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__()
+            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', [])
+            refined_query = self.refinement_chain.run(query=query, chat_history=chat_history)
+            response = self.qa_chain({"question": refined_query, "chat_history": chat_history})
+            return {"answer": response['answer']}
 
     def chat(self, query):
         if not self.overall_chain: