Update app.py

app.py

@@ -83,6 +83,26 @@ def remove_message(position):
     st.toast("try to remove message no: " + str(position-1) + " and "+ str(position))
     del st.session_state.messages[position-2:position]
 
+def ask_assistant_write_stream():
+    # Display assistant response in chat message container
+    assistant = ask_assistant_stream(model_links[selected_model], st.session_state.messages, temp_values, max_token_value)
+    pos = len(st.session_state.messages)+1
+    if "stream" in assistant: 
+        with st.chat_message("assistant"):
+            col1, col2 = st.columns([9,1])
+            response = col1.write_stream(assistant["stream"])
+            col2.button("", icon = ":material/edit:", key="button_edit_message_"+str(pos), args=[pos], on_click=edit_message)
+            col2.button("", icon = ":material/delete:", key="button_remove_message_"+str(pos), args=[pos], on_click=remove_message)
+    else:
+        with st.chat_message("assistant"):
+            col1, col2 = st.columns([9,1])
+            response = col1.write("Failure!")
+            col2.button("", icon = ":material/delete:", key="button_remove_message_"+str(pos), args=[pos], on_click=remove_message)
+
+    st.session_state.messages.append({"role": "assistant", "content": response})
+    
+
+
 st.subheader(f'{selected_model}')
 
 # Initialize chat history
@@ -105,23 +125,7 @@ if "instant_request" not in st.session_state:
 
 
 if st.session_state.instant_request:
-    assistant = ask_assistant_stream(model_links[selected_model], st.session_state.messages, temp_values, max_token_value)
-    pos = len(st.session_state.messages)+1
-    if "stream" in assistant: 
-        with st.chat_message("assistant"):
-            col1, col2 = st.columns([9,1])
-            response = col1.write_stream(assistant["stream"])
-            col2.button("", icon = ":material/edit:", key="button_edit_message_"+str(pos), args=[pos], on_click=edit_message)
-            col2.button("", icon = ":material/delete:", key="button_remove_message_"+str(pos), args=[pos], on_click=remove_message)
-    else:
-        with st.chat_message("assistant"):
-            col1, col2 = st.columns([9,1])
-            response = col1.write("Failure!")
-            col2.button("", icon = ":material/delete:", key="button_remove_message_"+str(pos), args=[pos], on_click=remove_message)
-
-    st.session_state.messages.append({"role": "assistant", "content": response})
-    st.session_state.instant_request = False
-    
+    ask_assistant_write_stream()
 
 # Accept user input
 if prompt := st.chat_input(f"Hi I'm {selected_model}, ask me a question"):
@@ -134,18 +138,4 @@ if prompt := st.chat_input(f"Hi I'm {selected_model}, ask me a question"):
     st.session_state.messages.append({"role": "user", "content": prompt})
     
     # Display assistant response in chat message container
-    assistant = ask_assistant_stream(model_links[selected_model], st.session_state.messages, temp_values, max_token_value)
-    pos = len(st.session_state.messages)+1
-    if "stream" in assistant: 
-        with st.chat_message("assistant"):
-            col1, col2 = st.columns([9,1])
-            response = col1.write_stream(assistant["stream"])
-            col2.button("", icon = ":material/edit:", key="button_edit_message_"+str(pos), args=[pos], on_click=edit_message)
-            col2.button("", icon = ":material/delete:", key="button_remove_message_"+str(pos), args=[pos], on_click=remove_message)
-    else:
-        with st.chat_message("assistant"):
-            col1, col2 = st.columns([9,1])
-            response = col1.write("Failure!")
-            col2.button("", icon = ":material/delete:", key="button_remove_message_"+str(pos), args=[pos], on_click=remove_message)
-
-    st.session_state.messages.append({"role": "assistant", "content": response})
+    ask_assistant_write_stream()