Update app.py

app.py

@@ -10,7 +10,7 @@ from langchain.llms import HuggingFacePipeline
 from langchain.chains import ConversationChain
 from langchain.memory import ConversationBufferMemory
 from langchain.llms import HuggingFaceHub
-from langchain.memory import ConversationTokenBufferMemory
+from langchain.memory import ConversationBufferWindowMemory
 
 from pathlib import Path
 import chromadb
@@ -128,9 +128,8 @@ def initialize_llmchain(temperature, max_tokens, top_k, vector_db, progress=gr.P
                                                           "load_in_8bit": True})
     
     progress(0.75, desc="Defining buffer memory...")
-    memory = ConversationBufferMemory(memory_key="chat_history",output_key='answer',return_messages=True)
-    #memory = ConversationTokenBufferMemory(llm = llm, max_token_limit=100)
-    # retriever=vector_db.as_retriever(search_type="similarity", search_kwargs={'k': 3})
+    #memory = ConversationBufferMemory(memory_key="chat_history",output_key='answer',return_messages=True)
+    memory = ConversationBufferWindowMemory(memory_key = 'history', k=3)
     retriever=vector_db.as_retriever()
     progress(0.8, desc="Defining retrieval chain...")
     qa_chain = ConversationalRetrievalChain.from_llm(llm,retriever=retriever,chain_type="stuff", 
@@ -177,9 +176,9 @@ def initialize_LLM(llm_temperature, max_tokens, top_k, vector_db, progress=gr.Pr
     return qa_chain, "Complete!"
 
 
-def format_chat_history(message, chat_history):
+def format_chat_history(message, history):
     formatted_chat_history = []
-    for user_message, bot_message in chat_history:
+    for user_message, bot_message in history:
         formatted_chat_history.append(f"User: {user_message}")
         formatted_chat_history.append(f"Assistant: {bot_message}")
     return formatted_chat_history