adjust memory

app.py

@@ -22,12 +22,17 @@ def set_api_key(api_key):
 
 
 @st.cache_resource
-def get_chain(model_name, temperature):
-    llm = ChatOpenAI(model_name=model_name, temperature=temperature)
+def setup_memory():
     msgs = StreamlitChatMessageHistory(key="basic_chat_app")
     memory = ConversationBufferWindowMemory(k=3, memory_key="chat_history", 
                                             chat_memory=msgs,
                                             return_messages=True)
+    return memory
+
+
+@st.cache_resource
+def get_chain(model_name, _memory, temperature):
+    llm = ChatOpenAI(model_name=model_name, temperature=temperature)
     prompt = ChatPromptTemplate(
     messages=[
         SystemMessagePromptTemplate.from_template(
@@ -46,7 +51,6 @@ def get_chain(model_name, temperature):
     return conversation
 
 
-
             
 if __name__ == "__main__":
     st.header("Basic chatbot")
@@ -56,10 +60,10 @@ if __name__ == "__main__":
         st.write("To clear all memory and start fresh, click 'Clear history'" )
         
     API_KEY = st.sidebar.text_input(
-    'API Key',
-    type='password',
-    help="Enter your OpenAI API key to use this app",
-    value=None)
+        'API Key',
+        type='password',
+        help="Enter your OpenAI API key to use this app",
+        value=None)
 
     model_name = st.sidebar.selectbox(
         label = "Choose a model",
@@ -78,7 +82,8 @@ if __name__ == "__main__":
     
     if API_KEY is not None:
         set_api_key(API_KEY)
-        chain = get_chain(model_name, temperature)
+        memory = setup_memory()
+        chain = get_chain(model_name, memory, temperature)
         if st.button("Clear history"):
             chain.memory.clear()
             st.cache_resource.clear()