Update functions.py

functions.py

@@ -32,6 +32,7 @@ from langchain.chat_models import ChatOpenAI
 from langchain.callbacks.base import CallbackManager
 from langchain.callbacks.streaming_stdout import StreamingStdOutCallbackHandler
 from langchain.chains import ConversationalRetrievalChain, QAGenerationChain
+from langchain.memory import ConversationBufferMemory
 
 from langchain.prompts.chat import (
     ChatPromptTemplate,
@@ -57,6 +58,8 @@ time_str = time.strftime("%d%m%Y-%H%M%S")
 HTML_WRAPPER = """<div style="overflow-x: auto; border: 1px solid #e6e9ef; border-radius: 0.25rem; padding: 1rem; 
 margin-bottom: 2.5rem">{}</div> """
 
+memory = ConversationBufferMemory(memory_key='chat_history', return_messages=True)
+
 #Stuff Chain Type Prompt template
 
 @st.cache_resource
@@ -230,9 +233,10 @@ def embed_text(query,embedding_model,_docsearch):
     chain = ConversationalRetrievalChain.from_llm(chat_llm, 
                                                   retriever= _docsearch.as_retriever(),
                                                   qa_prompt = load_prompt(),
+                                                  memory = memory,
                                                   return_source_documents=True)
     
-    answer = chain({"question": query, "chat_history": chat_history})
+    answer = chain({"question": query})
 
     return answer