diff --git a/backend/oasst_backend/api/v1/utils.py b/backend/oasst_backend/api/v1/utils.py index 5c9537a3..b9e982ab 100644 --- a/backend/oasst_backend/api/v1/utils.py +++ b/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: diff --git a/backend/oasst_backend/models/db_payload.py b/backend/oasst_backend/models/db_payload.py index ddaf2391..b0eb38df 100644 --- a/backend/oasst_backend/models/db_payload.py +++ b/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] diff --git a/backend/oasst_backend/task_repository.py b/backend/oasst_backend/task_repository.py index 5c5dea21..fdf6ebb7 100644 --- a/backend/oasst_backend/task_repository.py +++ b/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, diff --git a/backend/oasst_backend/tree_manager.py b/backend/oasst_backend/tree_manager.py index 77419184..909bb84b 100644 --- a/backend/oasst_backend/tree_manager.py +++ b/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, diff --git a/oasst-shared/oasst_shared/schemas/protocol.py b/oasst-shared/oasst_shared/schemas/protocol.py index 22a4adfb..cbb4d29c 100644 --- a/oasst-shared/oasst_shared/schemas/protocol.py +++ b/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