Update app.py

app.py

@@ -6,11 +6,16 @@ from src.utils.ingest_image import extract_and_store_images
 from src.utils.text_qa import qa_bot
 from src.utils.image_qa import query_and_print_results
 import nest_asyncio
+from langchain.memory import ConversationBufferWindowMemory
+from langchain_community.chat_message_histories import StreamlitChatMessageHistory
+from dotenv import load_dotenv
 nest_asyncio.apply()
 
-from dotenv import load_dotenv
 load_dotenv()
 
+memory_storage = StreamlitChatMessageHistory(key="chat_messages")
+memory = ConversationBufferWindowMemory(memory_key="chat_history", human_prefix="User", chat_memory=memory_storage, k=3)
+
 def get_answer(query, chain):
     try:
         response = chain.invoke(query)
@@ -23,17 +28,13 @@ st.title("MULTIMODAL DOC QA")
 
 uploaded_file = st.file_uploader("File upload", type="pdf")
 if uploaded_file is not None:
-    # Save the uploaded file to a temporary location
     temp_file_path = os.path.join("temp", uploaded_file.name)
-    os.makedirs("temp", exist_ok=True)  # Ensure the temp directory exists
+    os.makedirs("temp", exist_ok=True)
     with open(temp_file_path, "wb") as f:
         f.write(uploaded_file.getbuffer())
 
-    # Get the absolute path of the saved file
-    #temp_dir = tempfile.mkdtemp()
     path = os.path.abspath(temp_file_path)
     st.write(f"File saved to: {path}")
-    print(path)
 
     st.write("Document uploaded successfully!")
 
@@ -44,8 +45,8 @@ if st.button("Start Processing"):
                 client = create_vector_database(path)
                 image_vdb = extract_and_store_images(path)
                 chain = qa_bot(client)
-                st.session_state['chain'] = chain  # Store chain in session state
-                st.session_state['image_vdb'] = image_vdb  # Store image_vdb in session state
+                st.session_state['chain'] = chain
+                st.session_state['image_vdb'] = image_vdb
                 st.success("Processing complete.")
             except Exception as e:
                 st.error(f"Error during processing: {e}")
@@ -59,11 +60,15 @@ if user_input := st.chat_input("User Input"):
 
         with st.chat_message("user"):
             st.markdown(user_input)
+        memory.save_context({"role": "user", "content": user_input})
 
         with st.spinner("Generating Response..."):
             response = get_answer(user_input, chain)
             if response:
                 st.markdown(response)
+                with st.chat_message("assistant"):
+                    st.markdown(response)
+                memory.save_context({"role": "assistant", "content": response})
                 try:
                     query_and_print_results(image_vdb, user_input)
                 except Exception as e:
@@ -72,3 +77,14 @@ if user_input := st.chat_input("User Input"):
                 st.error("Failed to generate response.")
     else:
         st.error("Please start processing before entering user input.")
+
+if "messages" not in st.session_state:
+    st.session_state.messages = []
+
+for message in st.session_state.messages:
+    with st.chat_message(message["role"]):
+        st.write(message["content"])
+
+for i, msg in enumerate(memory_storage.messages):
+    name = "user" if i % 2 == 0 else "assistant"
+    st.chat_message(name).markdown(msg.content)