Update app.py

app.py

@@ -11,7 +11,7 @@ from typing import Optional
 
 from langchain.prompts import ChatPromptTemplate, MessagesPlaceholder
 from langchain_core.chat_history import BaseChatMessageHistory
-from langchain.memory import ConversationBufferMemory, PostgresChatMessageHistory
+from langchain.memory import ConversationBufferMemory#, PostgresChatMessageHistory
 
 
 
@@ -22,16 +22,16 @@ POSTGRE_URL = os.environ['POSTGRE_URL']
 
 
 if 'memory' not in st.session_state:
-    # st.session_state['memory'] = ConversationBufferMemory(return_messages=True)
+    st.session_state['memory'] = ConversationBufferMemory(return_messages=True)
     
-    st.session_state.memory = PostgresChatMessageHistory(connection_string=POSTGRE_URL, session_id=str(datetime.timestamp(datetime.now())))
-    st.session_state.memory.add_ai_message("Hello, My name is Jonathan Jordan. You can call me Jojo. How can I help you today?")
+    # st.session_state.memory = PostgresChatMessageHistory(connection_string=POSTGRE_URL, session_id=str(datetime.timestamp(datetime.now())))
+    # st.session_state.memory.add_ai_message("Hello, My name is Jonathan Jordan. You can call me Jojo. How can I help you today?")
 
 if 'chain' not in st.session_state:
     st.session_state['chain'] = custom_chain_with_history(
         llm=CustomLLM(repo_id="mistralai/Mixtral-8x7B-Instruct-v0.1", model_type='text-generation', api_token=API_TOKEN, stop=["\n<|","<|"], temperature=0.001), 
-        # memory=st.session_state.memory.chat_memory,
-        memory=st.session_state.memory
+        memory=st.session_state.memory.chat_memory,
+        # memory=st.session_state.memory
     )
 
 st.title("Chat With Me")
@@ -58,10 +58,10 @@ if prompt := st.chat_input("Ask me anything.."):
     # Display assistant response in chat message container
     with st.chat_message("assistant"):
         st.markdown(response)
-    st.session_state.memory.add_user_message(prompt)
-    st.session_state.memory.add_ai_message(response)
-    # st.session_state.memory.save_context({"question":prompt}, {"output":response})
-    # st.session_state.memory.chat_memory.messages = st.session_state.memory.chat_memory.messages[-15:]
+    # st.session_state.memory.add_user_message(prompt)
+    # st.session_state.memory.add_ai_message(response)
+    st.session_state.memory.save_context({"question":prompt}, {"output":response})
+    st.session_state.memory.chat_memory.messages = st.session_state.memory.chat_memory.messages[-15:]
     # Add assistant response to chat history
     st.session_state.messages.append({"role": "assistant", "content": response})