Update app.py

app.py

@@ -6,6 +6,7 @@ from dotenv import load_dotenv
 import time
 from langchain.vectorstores import Chroma
 from langchain.embeddings import HuggingFaceEmbeddings
+from langchain.memory import ConversationBufferMemory
 from langchain_core.prompts import ChatPromptTemplate, PromptTemplate
 from langchain_groq import ChatGroq
 from langchain.chains import RetrievalQA, LLMChain
@@ -248,19 +249,22 @@ Remember to maintain a professional, clear tone and ensure all medicine recommen
 Answer:"""
 
 # Create the QA chain with correct variables
+memory = ConversationBufferMemory(
+    memory_key="chat_history",
+    input_key="question",
+    output_key="answer"
+)
+
 qa = RetrievalQA.from_chain_type(
     llm=llm,
     chain_type="stuff",
     retriever=retriever,
-    return_source_documents=True,
+    memory=memory,
     chain_type_kwargs={
         "prompt": PromptTemplate(
             template=prompt_template,
             input_variables=["context", "question"]
-        ),
-        "memory_key": "chat_history",  # Specify the memory key
-        "input_key": "question",  # Specify which key to use for memory
-        "output_key": "answer"  # Specify the output key
+        )
     }
 )