feat: support restart webui

SipherAGI#243
bluelovers · Jun 4, 2024 · 689195c · 689195c
1 parent fcab048
commit 689195c
Show file tree

Hide file tree

Showing 2 changed files with 22 additions and 1 deletion.
diff --git a/agent_scheduler/task_runner.py b/agent_scheduler/task_runner.py
@@ -620,6 +620,27 @@ def __on_completed(self):
         elif action == "Stop webui":
             log.info("[AgentScheduler] Stopping webui...")
             _exit(0)
+        elif action == "Restart webui":
+            log.info("[AgentScheduler] Restarting webui...")
+            with gr.Blocks() as demo:
+                gr.HTML("""
+    <script type="text/javascript">
+    console.info(`Restart webui...after 5sec`);
+    
+    var requestPing = function() {
+        requestGet("./internal/ping", {}, function(data) {
+            location.reload();
+        }, function() {
+            console.info(`Waiting webui start...`);
+            setTimeout(requestPing, 500);
+        });
+    };
+    
+    setTimeout(requestPing, 5000);
+    </script>
+                """)
+            shared.state.request_restart()
+            demo.launch()
 
         if command:
             subprocess.Popen(command)

diff --git a/scripts/task_scheduler.py b/scripts/task_scheduler.py
@@ -41,7 +41,7 @@
 placement_under_generate = "Under Generate button"
 placement_between_prompt_and_generate = "Between Prompt and Generate button"
 
-completion_action_choices = ["Do nothing", "Shut down", "Restart", "Sleep", "Hibernate", "Stop webui"]
+completion_action_choices = ["Do nothing", "Shut down", "Restart", "Sleep", "Hibernate", "Stop webui", "Restart webui"]
 
 task_filter_choices = ["All", "Bookmarked", "Done", "Failed", "Interrupted"]