use StreamlitChatMessageHistory

.gitignore

@@ -1 +1,2 @@
 .ipynb_checkpoints
+.DS_Store

app.py

@@ -9,10 +9,13 @@ SystemMessagePromptTemplate,
 HumanMessagePromptTemplate,
 )
 from langchain.memory import ConversationBufferWindowMemory
+from langchain.memory.chat_message_histories import StreamlitChatMessageHistory
 from langchain.schema import AIMessage, HumanMessage
 import streamlit as st
 
 
+
+
 @st.cache_resource
 def set_api_key(api_key):
     os.environ["OPENAI_API_KEY"] = api_key
@@ -21,7 +24,10 @@ def set_api_key(api_key):
 @st.cache_resource
 def get_chain(model_name, temperature):
     llm = ChatOpenAI(model_name=model_name, temperature=temperature)
-    memory = ConversationBufferWindowMemory(k=3, memory_key="chat_history", return_messages=True)
+    msgs = StreamlitChatMessageHistory(key="basic_chat_app")
+    memory = ConversationBufferWindowMemory(k=3, memory_key="chat_history", 
+                                            chat_memory=msgs,
+                                            return_messages=True)
     prompt = ChatPromptTemplate(
     messages=[
         SystemMessagePromptTemplate.from_template(
@@ -40,26 +46,15 @@ def get_chain(model_name, temperature):
     return conversation
 
 
-def display_messages(chain):
-    """
-    Show the messages in the conversation buffer.
-    """
-    for message in chain.memory.buffer:
-        if isinstance(message, AIMessage):
-            role = "assistant"
-        elif isinstance(message, HumanMessage):
-            role = "user"
-        with st.chat_message(role):
-            st.write(message.content)
-            
 
+            
 if __name__ == "__main__":
     st.header("Basic chatbot")
     with st.expander("How conversation history works"):
         st.write("To keep input lengths down and costs reasonable,"
                  " this bot only 'remembers' the past three turns of conversation.")
-        st.write("Each combination of model type and temperature has its own unique chat history.")
-        st.write("To clear the current model's memory and start fresh, click 'Clear history'" )
+        st.write("To clear all memory and start fresh, click 'Clear history'" )
+        
     API_KEY = st.sidebar.text_input(
     'API Key',
     type='password',
@@ -86,14 +81,16 @@ if __name__ == "__main__":
         chain = get_chain(model_name, temperature)
         if st.button("Clear history"):
             chain.memory.clear()
-        display_messages(chain)
+            st.cache_resource.clear()
+        for message in chain.memory.buffer:
+            st.chat_message(message.type).write(message.content)
         text = st.chat_input()
         if text:
-            if text.lower() == "clear":
-                chain.memory.clear()
-            else:
-                result = chain.predict(input=text)
-                display_messages(chain)
+            with st.chat_message("user"):
+                st.write(text)
+            result = chain.predict(input=text)
+            with st.chat_message("assistant"):
+                st.write(result)