Update pages/2_🧠_context_aware_chatbot.py

pages/2_🧠_context_aware_chatbot.py

@@ -6,6 +6,8 @@ from langchain_openai import ChatOpenAI
 from langchain.chains.conversation.base import ConversationChain
 from langchain.memory.buffer import ConversationBufferMemory
 from langchain.memory import ConversationSummaryMemory
+from langchain.chains.conversational_retrieval.base import ConversationalRetrievalChain
+from langchain.prompts import PromptTemplate
 
 st.set_page_config(page_title="Context aware chatbot", page_icon="🧠")
 st.header('Context aware chatbot')
@@ -18,11 +20,33 @@ class ContextChatbot:
 
     @st.cache_resource
     def setup_chain(_self):
+        prompt_template = """
+        Use the following pieces of context to answer the user's question about sickle cell.
+        If you don't know the answer, just say that you don't know, don't try to make up an answer.
+        ----------------
+        {context}
+
+        Question: {question}
+        """
+
+        PROMPT = PromptTemplate(
+            template=prompt_template, input_variables=["context", "question"]
+        )
+        chain_type_kwargs = {"prompt": PROMPT}
         summarizer = ChatOpenAI(model_name= "gpt-4o", temperature=0, streaming=True)
-        memory = ConversationSummaryMemory(llm = summarizer)
         llm = ChatOpenAI(model_name=_self.openai_model, temperature=0, streaming=True)
-        chain = ConversationChain(llm=llm, memory=memory, verbose=True)
-        return chain
+
+        chat_history = []
+
+        qa = ConversationalRetrievalChain.from_llm(
+            llm = llm,
+            chain_type = "stuff",
+            memory = ConversationSummaryMemory(llm = summarizer, memory_key='chat_history', input_key='question', output_key= 'answer', return_messages=True),
+            retriever = vectorstore.as_retriever(k = 3, search_type="mmr"),
+            return_source_documents=True,
+            combine_docs_chain_kwargs=chain_type_kwargs
+        )
+        return qa
 
     @utils.enable_chat_history
     def main(self):