Update app.py

app.py

@@ -26,8 +26,6 @@ def add_bg_from_local(image_file):
     background-size: cover}}</style>""", unsafe_allow_html=True)
 add_bg_from_local(image_bg)
 
-#st.header("Welcome")
-
 st.markdown("""
     <svg width="600" height="100">
         <text x="50%" y="50%" font-family="San serif" font-size="42px" fill="Black" text-anchor="middle" stroke="white"
@@ -36,7 +34,6 @@ st.markdown("""
     </svg>
 """, unsafe_allow_html=True)
 
-
 def get_answer(query, chain):
     try:
         response = chain.invoke(query)
@@ -45,8 +42,6 @@ def get_answer(query, chain):
         st.error(f"Error in get_answer: {e}")
         return None
 
-#st.title("MULTIMODAL DOC QA")
-
 uploaded_file = st.file_uploader("File upload", type="pdf")
 if uploaded_file is not None:
     temp_file_path = os.path.join("temp", uploaded_file.name)
@@ -56,7 +51,6 @@ if uploaded_file is not None:
 
     path = os.path.abspath(temp_file_path)
     st.write(f"File saved to: {path}")
-
     st.write("Document uploaded successfully!")
 
 if st.button("Start Processing"):
@@ -82,7 +76,6 @@ st.markdown("""
     </style>
     """, unsafe_allow_html=True)
 
-
 if user_input := st.chat_input("User Input"):
     if 'chain' in st.session_state and 'image_vdb' in st.session_state:
         chain = st.session_state['chain']
@@ -90,15 +83,24 @@ 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})
+                
+                # Save context in memory
+                memory.save_context(
+                    {"role": "user", "content": user_input},
+                    {"role": "assistant", "content": response}
+                )
+
+                # Append messages to session state for display
+                st.session_state.messages.append({"role": "user", "content": user_input})
+                st.session_state.messages.append({"role": "assistant", "content": response})
+
                 try:
                     query_and_print_results(image_vdb, user_input)
                 except Exception as e: