From 5b69cfe7c343e8672fc350ec35e17f0d046297ca Mon Sep 17 00:00:00 2001
From: Chenlei Hu <chenlei.hu@mail.utoronto.ca>
Date: Sun, 21 Jul 2024 15:29:10 -0400
Subject: [PATCH] Add timestamp to execution messages (#4076)

* Add timestamp to execution messages

* Add execution_end message

* Rename to execution_success
---
 execution.py | 10 +++++++++-
 1 file changed, 9 insertions(+), 1 deletion(-)

diff --git a/execution.py b/execution.py
index 8b3d116c..0a2e62e7 100644
--- a/execution.py
+++ b/execution.py
@@ -3,6 +3,7 @@ import copy
 import logging
 import threading
 import heapq
+import time
 import traceback
 import inspect
 from typing import List, Literal, NamedTuple, Optional
@@ -283,7 +284,11 @@ class PromptExecutor:
         self.success = True
         self.old_prompt = {}
 
-    def add_message(self, event, data, broadcast: bool):
+    def add_message(self, event, data: dict, broadcast: bool):
+        data = {
+            **data,
+            "timestamp": int(time.time() * 1000),
+        }
         self.status_messages.append((event, data))
         if self.server.client_id is not None or broadcast:
             self.server.send_sync(event, data, self.server.client_id)
@@ -394,6 +399,9 @@ class PromptExecutor:
                 if self.success is not True:
                     self.handle_execution_error(prompt_id, prompt, current_outputs, executed, error, ex)
                     break
+            else:
+                # Only execute when the while-loop ends without break
+                self.add_message("execution_success", { "prompt_id": prompt_id }, broadcast=False)
 
             for x in executed:
                 self.old_prompt[x] = copy.deepcopy(prompt[x])