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
 
 
 
@@ -19,20 +19,38 @@ from langchain.memory import ConversationBufferMemory#, PostgresChatMessageHisto
 API_TOKEN = os.getenv('HF_INFER_API')
 POSTGRE_URL = os.environ['POSTGRE_URL']
 
+@st.cache_resource
+def get_llm_chain():
+    return 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
+    )
+
+
+@st.cache_resource
+def get_memory():
+    return PostgresChatMessageHistory(connection_string=POSTGRE_URL, session_id=str(datetime.timestamp(datetime.now())))
 
 
 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.sessio_state.memory = get_memory()
     st.session_state.memory.chat_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
-    )
+    # 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
+    # )
+
+    st.session_state['chain'] = get_llm_chain()
+
+
 
 st.title("Chat With Me")
 st.subheader("by Jonathan Jordan")
@@ -58,10 +76,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})