add reply_message to LabelConversationReplyTask & LabelConversationReplyPayload

backend/oasst_backend/api/v1/utils.py

@@ -23,19 +23,20 @@ def prepare_message_list(messages: list[Message]) -> list[protocol.Message]:
     return [prepare_message(m) for m in messages]
 
 
+def prepare_conversation_message(message: Message) -> protocol.ConversationMessage:
+    return protocol.ConversationMessage(
+        id=message.id,
+        frontend_message_id=message.frontend_message_id,
+        text=message.text,
+        lang=message.lang,
+        is_assistant=(message.role == "assistant"),
+        emojis=message.emojis or {},
+        user_emojis=message.user_emojis or [],
+    )
+
+
 def prepare_conversation_message_list(messages: list[Message]) -> list[protocol.ConversationMessage]:
-    return [
-        protocol.ConversationMessage(
-            id=message.id,
-            frontend_message_id=message.frontend_message_id,
-            text=message.text,
-            lang=message.lang,
-            is_assistant=(message.role == "assistant"),
-            emojis=message.emojis or {},
-            user_emojis=message.user_emojis or [],
-        )
-        for message in messages
-    ]
+    return [prepare_conversation_message(message) for message in messages]
 
 
 def prepare_conversation(messages: list[Message]) -> protocol.Conversation:

backend/oasst_backend/models/db_payload.py

@@ -117,7 +117,8 @@ class LabelConversationReplyPayload(TaskPayload):
 
     message_id: UUID
     conversation: protocol_schema.Conversation
-    reply: str
+    reply: str  # deprecated
+    reply_message: Optional[protocol_schema.ConversationMessage]
     valid_labels: list[str]
     mandatory_labels: Optional[list[str]]
     mode: Optional[protocol_schema.LabelTaskMode]

backend/oasst_backend/task_repository.py

@@ -101,6 +101,7 @@ class TaskRepository:
                     message_id=task.message_id,
                     conversation=task.conversation,
                     reply=task.reply,
+                    reply_message=task.reply_message,
                     valid_labels=task.valid_labels,
                     mandatory_labels=task.mandatory_labels,
                     mode=task.mode,
@@ -112,6 +113,7 @@ class TaskRepository:
                     message_id=task.message_id,
                     conversation=task.conversation,
                     reply=task.reply,
+                    reply_message=task.reply_message,
                     valid_labels=task.valid_labels,
                     mandatory_labels=task.mandatory_labels,
                     mode=task.mode,

backend/oasst_backend/tree_manager.py

@@ -11,7 +11,11 @@ import numpy as np
 import pydantic
 from fastapi.encoders import jsonable_encoder
 from loguru import logger
-from oasst_backend.api.v1.utils import prepare_conversation, prepare_conversation_message_list
+from oasst_backend.api.v1.utils import (
+    prepare_conversation,
+    prepare_conversation_message,
+    prepare_conversation_message_list,
+)
 from oasst_backend.config import TreeManagerConfiguration, settings
 from oasst_backend.models import Message, MessageReaction, MessageTreeState, Task, TextLabels, User, message_tree_state
 from oasst_backend.prompt_repository import PromptRepository
@@ -379,6 +383,7 @@ class TreeManager:
                             message_id=message.id,
                             conversation=conversation,
                             reply=message.text,
+                            reply_message=prepare_conversation_message(message),
                             valid_labels=list(map(lambda x: x.value, valid_labels)),
                             mandatory_labels=list(map(lambda x: x.value, self.cfg.mandatory_labels_assistant_reply)),
                             mode=label_mode,
@@ -398,6 +403,7 @@ class TreeManager:
                             message_id=message.id,
                             conversation=conversation,
                             reply=message.text,
+                            reply_message=prepare_conversation_message(message),
                             valid_labels=list(map(lambda x: x.value, valid_labels)),
                             mandatory_labels=list(map(lambda x: x.value, self.cfg.mandatory_labels_prompter_reply)),
                             mode=label_mode,

oasst-shared/oasst_shared/schemas/protocol.py

@@ -259,6 +259,7 @@ class LabelConversationReplyTask(AbstractLabelTask):
 
     type: Literal["label_conversation_reply"] = "label_conversation_reply"
     conversation: Conversation  # the conversation so far
+    reply_message: Optional[ConversationMessage]
     reply: str