feat: revoke msg

YunCheng66 · Mar 22, 2023 · d34eeab · d34eeab
1 parent 701c4a7
commit d34eeab
Show file tree

Hide file tree

Showing 2 changed files with 18 additions and 7 deletions.
diff --git a/modules/context.py b/modules/context.py
@@ -43,10 +43,16 @@ def refresh_last(self) -> None:
             query, output = self.rh[-1]
             self.rh[-1] = (query, parse_codeblock(output))
 
-    def clear(self):
+    def clear(self) -> None:
         self.history = []
         self.rh = []
 
+    def revoke(self) -> List[Tuple[str, str]]:
+        if self.history and self.rh:
+            self.history.pop()
+            self.rh.pop()
+        return self.rh
+
     def limit_round(self):
         hl = len(self.history)
         if hl == 0:

diff --git a/modules/ui.py b/modules/ui.py
@@ -15,12 +15,12 @@ def predict(query, max_length, top_p, temperature, use_stream_chat):
     ctx.limit_round()
     flag = True
     for _, output in infer(
-        query=query,
-        history=ctx.history,
-        max_length=max_length,
-        top_p=top_p,
-        temperature=temperature,
-        use_stream_chat=use_stream_chat
+            query=query,
+            history=ctx.history,
+            max_length=max_length,
+            top_p=top_p,
+            temperature=temperature,
+            use_stream_chat=use_stream_chat
     ):
         if flag:
             ctx.append(query, output)
@@ -89,6 +89,9 @@ def create_ui():
                 with gr.Row():
                     submit = gr.Button("发送", elem_id="c_generate")
 
+                with gr.Row():
+                    revoke_btn = gr.Button("撤回")
+
         submit.click(predict, inputs=[
             input_message,
             max_length,
@@ -100,6 +103,8 @@ def create_ui():
             input_message
         ])
 
+        revoke_btn.click(lambda x: ctx.revoke(), inputs=[input_message], outputs=[chatbot])
+
         clear.click(clear_history, outputs=[chatbot])
         clear_input.click(lambda x: "", inputs=[input_message], outputs=[input_message])