Update app.py

app.py

@@ -62,6 +62,30 @@ def fill_up_placeholders(txt):
         "" if len(placeholders) >= 1 else txt
     )
 
+def rollback_last(idx, ld, state):
+    res = [state["ppmanager_type"].from_json(json.dumps(ppm_str)) for ppm_str in ld]
+    last_user_message = res[idx].pingpongs[-1].ping
+    res[idx].pingpongs = res[idx].pingpongs[:-1]
+    
+    return (
+        last_user_message,
+        res[idx].build_uis(),
+        str(res),
+        gr.update(interactive=False)
+    )
+
+def reset_chat(idx, ld, state):
+    res = [state["ppmanager_type"].from_json(json.dumps(ppm_str)) for ppm_str in ld]
+    res[idx].pingpongs = []
+        
+    return (
+        "",
+        [],
+        str(res),
+        gr.update(visible=True),
+        gr.update(interactive=False),
+    )
+
 async def chat_stream(
     idx, local_data, instruction_txtbox, chat_state,
     global_context, res_temp, res_topk, res_rpen, res_mnts, res_sample, ctx_num_lconv
@@ -243,13 +267,39 @@ with gr.Blocks(css=MODEL_SELECTION_CSS, theme='gradio/soft') as demo:
                     with gr.Row():
                         ctx_num_lconv = gr.Slider(2, 10, 3, step=1, label="number of recent talks to keep", interactive=True)
 
-    instruction_txtbox.submit(
+    send_event = instruction_txtbox.submit(
         chat_stream,
         [idx, local_data, instruction_txtbox, chat_state,
          global_context, res_temp, res_topk, res_rpen, res_mnts, res_sample, ctx_num_lconv],
         [context_inspector, chatbot, local_data]
     )
 
+    regen_event1 = regenerate.click(
+        rollback_last,
+        [idx, local_data, chat_state],
+        [instruction_txtbox, chatbot, local_data, regenerate]
+    )
+    regen_event2 = regen_event1.then(
+        chat_stream,
+        [idx, local_data, instruction_txtbox, chat_state,
+         global_context, res_temp, res_topk, res_rpen, res_mnts, res_sample, ctx_num_lconv],
+        [context_inspector, chatbot, local_data]
+    )
+    regen_event3 = regen_event2.then(
+        lambda: gr.update(interactive=True),
+        None,
+        regenerate
+    )
+    regen_event4 = regen_event3.then(
+        None, local_data, None, 
+        _js="(v)=>{ setStorage('local_data',v) }"  
+    )
+    
+    stop.click(
+        None, None, None,
+        cancels=[send_event, regen_event1, regen_event2, regen_event3, regen_event4]
+    )
+
     for btn in channel_btns:
         btn.click(
             set_chatbot,
@@ -266,6 +316,16 @@ with gr.Blocks(css=MODEL_SELECTION_CSS, theme='gradio/soft') as demo:
             [btn],
             [instruction_txtbox, example_block]  
         )
+
+    clean.click(
+        reset_chat,
+        [idx, local_data, chat_state],
+        [instruction_txtbox, chatbot, local_data, example_block, regenerate]
+    ).then(
+        None, local_data, None, 
+        _js="(v)=>{ setStorage('local_data',v) }"
+    )
+
     
     placeholder_txt1.change(
         inputs=[template_txt, placeholder_txt1, placeholder_txt2, placeholder_txt3],
@@ -308,7 +368,7 @@ with gr.Blocks(css=MODEL_SELECTION_CSS, theme='gradio/soft') as demo:
         outputs=[instruction_txtbox, placeholder_txt1, placeholder_txt2, placeholder_txt3],
         fn=get_final_template
     )    
-    
+
     demo.load(
         None,
         inputs=None,