add to protocol.Message

backend/oasst_backend/api/v1/frontend_messages.py

@@ -7,7 +7,6 @@ from oasst_backend.exceptions import OasstError, OasstErrorCode
 from oasst_backend.models import ApiClient
 from oasst_backend.models.db_payload import MessagePayload
 from oasst_backend.prompt_repository import PromptRepository
-from oasst_shared.schemas import protocol
 from sqlmodel import Session
 
 router = APIRouter()
@@ -27,7 +26,7 @@ def get_message_by_frontend_id(
         # Unexpected message payload
         raise OasstError("Invalid message", OasstErrorCode.INVALID_MESSAGE)
 
-    return protocol.ConversationMessage(text=message.payload.payload.text, is_assistant=(message.role == "assistant"))
+    return utils.prepare_message(message)
 
 
 @router.get("/{message_id}/conversation")
@@ -68,12 +67,7 @@ def get_children_by_frontend_id(
     pr = PromptRepository(db, api_client, user=None)
     message = pr.fetch_message_by_frontend_message_id(message_id)
     messages = pr.fetch_message_children(message.id)
-    return [
-        protocol.Message(
-            id=m.id, parent_id=m.parent_id, text=m.payload.payload.text, is_assistant=(m.role == "assistant")
-        )
-        for m in messages
-    ]
+    return utils.prepare_message_list(messages)
 
 
 @router.get("/{message_id}/descendants")

backend/oasst_backend/api/v1/frontend_users.py

@@ -4,9 +4,9 @@ from uuid import UUID
 
 from fastapi import APIRouter, Depends, Query
 from oasst_backend.api import deps
+from oasst_backend.api.v1 import utils
 from oasst_backend.models import ApiClient
 from oasst_backend.prompt_repository import PromptRepository
-from oasst_shared.schemas import protocol
 from sqlmodel import Session
 from starlette.responses import Response
 from starlette.status import HTTP_200_OK
@@ -41,13 +41,7 @@ def query_frontend_user_messages(
         only_roots=only_roots,
         deleted=None if include_deleted else False,
     )
-
-    return [
-        protocol.Message(
-            id=m.id, parent_id=m.parent_id, text=m.payload.payload.text, is_assistant=(m.role == "assistant")
-        )
-        for m in messages
-    ]
+    return utils.prepare_message_list(messages)
 
 
 @router.delete("/{username}/messages")

backend/oasst_backend/api/v1/messages.py

@@ -10,7 +10,6 @@ from oasst_backend.exceptions import OasstError, OasstErrorCode
 from oasst_backend.models import ApiClient
 from oasst_backend.models.db_payload import MessagePayload
 from oasst_backend.prompt_repository import PromptRepository
-from oasst_shared.schemas import protocol
 from sqlmodel import Session
 from starlette.status import HTTP_200_OK
 
@@ -45,12 +44,7 @@ def query_messages(
         deleted=None if allow_deleted else False,
     )
 
-    return [
-        protocol.Message(
-            id=m.id, parent_id=m.parent_id, text=m.payload.payload.text, is_assistant=(m.role == "assistant")
-        )
-        for m in messages
-    ]
+    return utils.prepare_message_list(messages)
 
 
 @router.get("/{message_id}")
@@ -66,7 +60,7 @@ def get_message(
         # Unexptcted message payload
         raise OasstError("Invalid message", OasstErrorCode.INVALID_MESSAGE)
 
-    return protocol.ConversationMessage(text=message.payload.payload.text, is_assistant=(message.role == "assistant"))
+    return utils.prepare_message(message)
 
 
 @router.get("/{message_id}/conversation")
@@ -104,12 +98,7 @@ def get_children(
     """
     pr = PromptRepository(db, api_client, user=None)
     messages = pr.fetch_message_children(message_id)
-    return [
-        protocol.Message(
-            id=m.id, parent_id=m.parent_id, text=m.payload.payload.text, is_assistant=(m.role == "assistant")
-        )
-        for m in messages
-    ]
+    return utils.prepare_message_list(messages)
 
 
 @router.get("/{message_id}/descendants")

backend/oasst_backend/api/v1/utils.py

@@ -9,6 +9,22 @@ from oasst_backend.models.db_payload import MessagePayload
 from oasst_shared.schemas import protocol
 
 
+def prepare_message(m: Message) -> protocol.Message:
+    if not isinstance(m.payload.payload, MessagePayload):
+        raise OasstError("Server error", OasstErrorCode.SERVER_ERROR, HTTPStatus.INTERNAL_SERVER_ERROR)
+    return protocol.Message(
+        id=m.id,
+        parent_id=m.parent_id,
+        text=m.payload.payload.text,
+        is_assistant=(m.role == "assistant"),
+        created_date=m.created_date,
+    )
+
+
+def prepare_message_list(messages: list[Message]) -> list[protocol.Message]:
+    return [prepare_message(m) for m in messages]
+
+
 def prepare_conversation(messages: list[Message]) -> protocol.Conversation:
     conv_messages = []
     for message in messages:
@@ -26,13 +42,6 @@ def prepare_tree(tree: list[Message], tree_id: UUID) -> protocol.MessageTree:
     for message in tree:
         if not isinstance(message.payload.payload, MessagePayload):
             raise OasstError("Server error", OasstErrorCode.SERVER_ERROR, HTTPStatus.INTERNAL_SERVER_ERROR)
-        tree_messages.append(
-            protocol.Message(
-                id=message.id,
-                parent_id=message.parent_id,
-                text=message.payload.payload.text,
-                is_assistant=(message.role == "assistant"),
-            )
-        )
+        tree_messages.append(prepare_message(message))
 
     return protocol.MessageTree(id=tree_id, messages=tree_messages)