Update app.py

app.py

@@ -140,19 +140,22 @@ def get_response(system_message, chat_history, user_text, max_new_tokens=256, co
 
 # Chat interface
 chat_interface = st.container()
-with chat_interface:
-    output_container = st.container()
-
-    # Display chat messages
-    with output_container:
-        for idx, message in enumerate(st.session_state.chat_history):
-            if message['role'] == 'system':
-                continue
-            with st.chat_message(message['role'], avatar=st.session_state.avatars[message['role']]):
-                st.markdown(message['content'])
-
-                # If this is the last assistant message, add the "Continue" button
-                # if idx == len(st.session_state.chat_history) - 1 and message['role'] == 'assistant':
+def refresh_chat():
+    with chat_interface:
+        output_container = st.container()
+    
+        # Display chat messages
+        with output_container:
+            for idx, message in enumerate(st.session_state.chat_history):
+                if message['role'] == 'system':
+                    continue
+                with st.chat_message(message['role'], avatar=st.session_state.avatars[message['role']]):
+                    st.markdown(message['content'])
+    
+                    # If this is the last assistant message, add the "Continue" button
+                    # if idx == len(st.session_state.chat_history) - 1 and message['role'] == 'assistant':
+
+refresh_chat()
 
 if st.button("Continue"):
     st.session_state.need_continue = True
@@ -196,6 +199,7 @@ if st.session_state.get('need_continue', False):
                 continue_last=True
             )
             st.markdown(response)
+        refresh_chat()
 
     # Reset the flag
     st.session_state.need_continue = False