Add timestamp to execution messages (#4076)

* Add timestamp to execution messages

* Add execution_end message

* Rename to execution_success
This commit is contained in:
Chenlei Hu 2024-07-21 15:29:10 -04:00 committed by GitHub
parent 95fa9545f1
commit 5b69cfe7c3
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
1 changed files with 9 additions and 1 deletions

View File

@ -3,6 +3,7 @@ import copy
import logging import logging
import threading import threading
import heapq import heapq
import time
import traceback import traceback
import inspect import inspect
from typing import List, Literal, NamedTuple, Optional from typing import List, Literal, NamedTuple, Optional
@ -283,7 +284,11 @@ class PromptExecutor:
self.success = True self.success = True
self.old_prompt = {} 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)) self.status_messages.append((event, data))
if self.server.client_id is not None or broadcast: if self.server.client_id is not None or broadcast:
self.server.send_sync(event, data, self.server.client_id) self.server.send_sync(event, data, self.server.client_id)
@ -394,6 +399,9 @@ class PromptExecutor:
if self.success is not True: if self.success is not True:
self.handle_execution_error(prompt_id, prompt, current_outputs, executed, error, ex) self.handle_execution_error(prompt_id, prompt, current_outputs, executed, error, ex)
break 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: for x in executed:
self.old_prompt[x] = copy.deepcopy(prompt[x]) self.old_prompt[x] = copy.deepcopy(prompt[x])