Update app.py

app.py

@@ -37,7 +37,7 @@ if "OPENAI_API_KEY" in os.environ:
         loader2 = PyPDFLoader("Tbank resources.pdf")
         data1 = loader1.load()
         data2 = loader2.load()
-        data= data1 + data2
+        data = data1 + data2
         text_splitter = RecursiveCharacterTextSplitter(chunk_size=1000, chunk_overlap=200)
         all_splits = text_splitter.split_documents(data)
         vectorstore = Chroma.from_documents(documents=all_splits, embedding=OpenAIEmbeddings())
@@ -80,14 +80,17 @@ if "OPENAI_API_KEY" in os.environ:
             ]
         )
 
-        memory = ConversationBufferMemory(memory_key="chat_history", return_messages=True)
         document_chain = create_stuff_documents_chain(chat, question_answering_prompt)
 
-        return retriever, document_chain, memory
+        return retriever, document_chain
 
     # Load components
     with st.spinner("Initializing Tbank Assistant..."):
-        retriever, document_chain, memory = initialize_components()
+        retriever, document_chain = initialize_components()
+
+    # Initialize memory for each session
+    if "memory" not in st.session_state:
+        st.session_state.memory = ConversationBufferMemory(memory_key="chat_history", return_messages=True)
 
     # Chat interface
     st.subheader("Chat with Tbank Assistant")
@@ -118,7 +121,7 @@ if "OPENAI_API_KEY" in os.environ:
             response = document_chain.invoke(
                 {
                     "context": docs,
-                    "chat_history": memory.load_memory_variables({})["chat_history"],
+                    "chat_history": st.session_state.memory.load_memory_variables({})["chat_history"],
                     "messages": [
                         HumanMessage(content=prompt)
                     ],
@@ -133,7 +136,7 @@ if "OPENAI_API_KEY" in os.environ:
         st.session_state.messages.append({"role": "assistant", "content": full_response})
 
         # Update memory
-        memory.save_context({"input": prompt}, {"output": full_response})
+        st.session_state.memory.save_context({"input": prompt}, {"output": full_response})
 
 else:
     st.warning("Please enter your OpenAI API Key in the sidebar to start the chatbot.")