Update app.py

app.py

@@ -69,16 +69,19 @@ st.sidebar.markdown("*Generated content may be inaccurate or false.*")
 def edit_message(position):
     returnText = st.text_area("message:", value = st.session_state.messages[position-1]["content"])
     
-    if st.button("Submit"):
+    if st.button("Save"):
         st.session_state.messages[position-1]["content"] = returnText
         st.rerun()
-    
+    if st.session_state.messages[position-1]["role"] == "user"
+        if st.button("Save & Retry"):
+            st.session_state.messages[position-1]["content"] = returnText
+            del session_state.messages[position-1:]
+            st.session_state.instant_request = True
 
 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]
 
-
 st.subheader(f'{selected_model}')
 
 # Initialize chat history
@@ -96,6 +99,29 @@ for message in st.session_state.messages:
         if message["role"] == "assistant":
             col2.button("", icon = ":material/delete:", key="button_remove_message_"+str(pos), args=[pos], on_click=remove_message)
 
+if "instant_request" not in st.session_state:
+    st.session_state.instant_request = False
+
+
+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
+    
+
 # Accept user input
 if prompt := st.chat_input(f"Hi I'm {selected_model}, ask me a question"):
     # Display user message in chat message container and Add user message to chat history