From e9ccba0395da2c79e83ad93408af415dad641886 Mon Sep 17 00:00:00 2001
From: Kevin Hu <kevinhu.sh@gmail.com>
Date: Fri, 24 Jan 2025 11:07:55 +0800
Subject: [PATCH] Add timestamp to messages (#4624)

### What problem does this PR solve?


### Type of change

- [x] New Feature (non-breaking change which adds functionality)
---
 api/db/services/canvas_service.py       | 5 +++--
 api/db/services/conversation_service.py | 9 +++++----
 2 files changed, 8 insertions(+), 6 deletions(-)

diff --git a/api/db/services/canvas_service.py b/api/db/services/canvas_service.py
index 8514b293c..9e26510a9 100644
--- a/api/db/services/canvas_service.py
+++ b/api/db/services/canvas_service.py
@@ -14,6 +14,7 @@
 #  limitations under the License.
 #
 import json
+import time
 import traceback
 from uuid import uuid4
 from agent.canvas import Canvas
@@ -80,7 +81,7 @@ def completion(tenant_id, agent_id, question, session_id=None, stream=True, **kw
             "id": session_id,
             "dialog_id": cvs.id,
             "user_id": kwargs.get("user_id", "") if isinstance(kwargs, dict) else "",
-            "message": [{"role": "assistant", "content": canvas.get_prologue()}],
+            "message": [{"role": "assistant", "content": canvas.get_prologue(), "created_at": time.time()}],
             "source": "agent",
             "dsl": cvs.dsl
         }
@@ -134,7 +135,7 @@ def completion(tenant_id, agent_id, question, session_id=None, stream=True, **kw
                 yield "data:" + json.dumps({"code": 0, "message": "", "data": ans},
                                            ensure_ascii=False) + "\n\n"
 
-            canvas.messages.append({"role": "assistant", "content": final_ans["content"], "id": message_id})
+            canvas.messages.append({"role": "assistant", "content": final_ans["content"], "created_at": time.time(), "id": message_id})
             canvas.history.append(("assistant", final_ans["content"]))
             if final_ans.get("reference"):
                 canvas.reference.append(final_ans["reference"])
diff --git a/api/db/services/conversation_service.py b/api/db/services/conversation_service.py
index 57d2f9314..6dee5a656 100644
--- a/api/db/services/conversation_service.py
+++ b/api/db/services/conversation_service.py
@@ -13,6 +13,7 @@
 #  See the License for the specific language governing permissions and
 #  limitations under the License.
 #
+import time
 from uuid import uuid4
 from api.db import StatusEnum
 from api.db.db_models import Conversation, DB
@@ -75,9 +76,9 @@ def structure_answer(conv, ans, message_id, session_id):
     if not conv.message:
         conv.message = []
     if not conv.message or conv.message[-1].get("role", "") != "assistant":
-        conv.message.append({"role": "assistant", "content": ans["answer"], "id": message_id})
+        conv.message.append({"role": "assistant", "content": ans["answer"], "created_at": time.time(), "id": message_id})
     else:
-        conv.message[-1] = {"role": "assistant", "content": ans["answer"], "id": message_id}
+        conv.message[-1] = {"role": "assistant", "content": ans["answer"], "created_at": time.time(), "id": message_id}
     if conv.reference:
         conv.reference[-1] = reference
     return ans
@@ -94,7 +95,7 @@ def completion(tenant_id, chat_id, question, name="New session", session_id=None
             "id": session_id,
             "dialog_id": chat_id,
             "name": name,
-            "message": [{"role": "assistant", "content": dia[0].prompt_config.get("prologue")}],
+            "message": [{"role": "assistant", "content": dia[0].prompt_config.get("prologue"), "created_at": time.time()}],
             "user_id": kwargs.get("user_id", "")
         }
         ConversationService.save(**conv)
@@ -166,7 +167,7 @@ def iframe_completion(dialog_id, question, session_id=None, stream=True, **kwarg
             "id": session_id,
             "dialog_id": dialog_id,
             "user_id": kwargs.get("user_id", ""),
-            "message": [{"role": "assistant", "content": dia.prompt_config["prologue"]}]
+            "message": [{"role": "assistant", "content": dia.prompt_config["prologue"], "created_at": time.time()}]
         }
         API4ConversationService.save(**conv)
         yield "data:" + json.dumps({"code": 0, "message": "",