344: Create tasks for text labels (#381)

* Implement label task for initial prompts and replies

* Resolve formatting

* Include missing argument

* Modify text_labels API to match new model, update DB schema accordingly

* Send valid labels as part of label tasks

* Send correctly formatted valid_labels list

* Fix request format

* Fix request details for text-frontend reply label task

* Include message_id in tasks

* Address review comments

* Fix alembic tree

backend/alembic/versions/2023_01_05_1745-20cd871f4ec7_added_user_to_textlabels.py

@@ -0,0 +1,30 @@
+"""Added user to TextLabels
+
+Revision ID: 20cd871f4ec7
+Revises: d4161e384f83
+Create Date: 2023-01-05 17:45:15.696468
+
+"""
+import sqlalchemy as sa
+from alembic import op
+from sqlalchemy.dialects import postgresql
+
+# revision identifiers, used by Alembic.
+revision = "20cd871f4ec7"
+down_revision = "3b0adfadbef9"
+branch_labels = None
+depends_on = None
+
+
+def upgrade() -> None:
+    # ### commands auto generated by Alembic - please adjust! ###
+    op.add_column("text_labels", sa.Column("user_id", postgresql.UUID(as_uuid=True), nullable=False))
+    op.create_foreign_key(None, "text_labels", "user", ["user_id"], ["id"])
+    # ### end Alembic commands ###
+
+
+def downgrade() -> None:
+    # ### commands auto generated by Alembic - please adjust! ###
+    op.drop_constraint(None, "text_labels", type_="foreignkey")
+    op.drop_column("text_labels", "user_id")
+    # ### end Alembic commands ###

backend/oasst_backend/api/v1/tasks.py

@@ -119,6 +119,38 @@ def generate_task(
                 conversation=protocol_schema.Conversation(messages=task_messages),
                 replies=replies,
             )
+
+        case protocol_schema.TaskRequestType.label_initial_prompt:
+            logger.info("Generating a LabelInitialPromptTask.")
+            message = pr.fetch_random_initial_prompts(1)[0]
+            task = protocol_schema.LabelInitialPromptTask(
+                message_id=message.id,
+                prompt=message.payload.payload.text,
+                valid_labels=list(map(lambda x: x.value, protocol_schema.TextLabel)),
+            )
+
+        case protocol_schema.TaskRequestType.label_prompter_reply:
+            logger.info("Generating a LabelPrompterReplyTask.")
+            conversation, messages = pr.fetch_multiple_random_replies(max_size=1, message_role="assistant")
+            message = messages[0].payload.payload.text
+            task = protocol_schema.LabelPrompterReplyTask(
+                message_id=message.id,
+                conversation=conversation,
+                reply=message,
+                valid_labels=list(map(lambda x: x.value, protocol_schema.TextLabel)),
+            )
+
+        case protocol_schema.TaskRequestType.label_assistant_reply:
+            logger.info("Generating a LabelAssistantReplyTask.")
+            conversation, messages = pr.fetch_multiple_random_replies(max_size=1, message_role="prompter")
+            message = messages[0].payload.payload.text
+            task = protocol_schema.LabelAssistantReplyTask(
+                message_id=message.id,
+                conversation=conversation,
+                reply=message,
+                valid_labels=list(map(lambda x: x.value, protocol_schema.TextLabel)),
+            )
+
         case _:
             raise OasstError("Invalid request type", OasstErrorCode.TASK_INVALID_REQUEST_TYPE)
 
@@ -256,6 +288,13 @@ def tasks_interaction(
                 pr.store_ranking(interaction)
                 # here we would store the ranking in the database
                 return protocol_schema.TaskDone()
+            case protocol_schema.TextLabels:
+                logger.info(
+                    f"Frontend reports labels of {interaction.message_id=} with {interaction.labels=} by {interaction.user=}."
+                )
+                # TODO: check if the labels are valid?
+                pr.store_text_labels(interaction)
+                return protocol_schema.TaskDone()
             case _:
                 raise OasstError("Invalid response type.", OasstErrorCode.TASK_INVALID_RESPONSE_TYPE)
     except OasstError:

backend/oasst_backend/api/v1/text_labels.py

@@ -1,4 +1,3 @@
-import pydantic
 from fastapi import APIRouter, Depends, HTTPException
 from fastapi.security.api_key import APIKey
 from loguru import logger
@@ -11,17 +10,12 @@ from starlette.status import HTTP_204_NO_CONTENT, HTTP_400_BAD_REQUEST
 router = APIRouter()
 
 
-class LabelTextRequest(pydantic.BaseModel):
-    text_labels: protocol_schema.TextLabels
-    user: protocol_schema.User
-
-
 @router.post("/", status_code=HTTP_204_NO_CONTENT)
 def label_text(
     *,
     db: Session = Depends(deps.get_db),
     api_key: APIKey = Depends(deps.get_api_key),
-    request: LabelTextRequest,
+    text_labels: protocol_schema.TextLabels,
 ) -> None:
     """
     Label a piece of text.
@@ -29,9 +23,9 @@ def label_text(
     api_client = deps.api_auth(api_key, db)
 
     try:
-        logger.info(f"Labeling text {request=}.")
-        pr = PromptRepository(db, api_client, user=request.user)
-        pr.store_text_labels(request.text_labels)
+        logger.info(f"Labeling text {text_labels=}.")
+        pr = PromptRepository(db, api_client, user=text_labels.user)
+        pr.store_text_labels(text_labels)
 
     except Exception:
         logger.exception("Failed to store label.")

backend/oasst_backend/models/db_payload.py

@@ -1,4 +1,5 @@
 from typing import Literal
+from uuid import UUID
 
 from oasst_backend.models.payload_column_type import payload_type
 from oasst_shared.schemas import protocol as protocol_schema
@@ -91,3 +92,37 @@ class RankAssistantRepliesPayload(RankConversationRepliesPayload):
     """A task to rank a set of assistant replies to a conversation."""
 
     type: Literal["rank_assistant_replies"] = "rank_assistant_replies"
+
+
+@payload_type
+class LabelInitialPromptPayload(TaskPayload):
+    """A task to label an initial prompt."""
+
+    type: Literal["label_initial_prompt"] = "label_initial_prompt"
+    message_id: UUID
+    prompt: str
+    valid_labels: list[str]
+
+
+@payload_type
+class LabelConversationReplyPayload(TaskPayload):
+    """A task to label a conversation reply."""
+
+    message_id: UUID
+    conversation: protocol_schema.Conversation
+    reply: str
+    valid_labels: list[str]
+
+
+@payload_type
+class LabelPrompterReplyPayload(LabelConversationReplyPayload):
+    """A task to label a prompter reply."""
+
+    type: Literal["label_prompter_reply"] = "label_prompter_reply"
+
+
+@payload_type
+class LabelAssistantReplyPayload(LabelConversationReplyPayload):
+    """A task to label an assistant reply."""
+
+    type: Literal["label_assistant_reply"] = "label_assistant_reply"

backend/oasst_backend/models/text_labels.py

@@ -15,6 +15,7 @@ class TextLabels(SQLModel, table=True):
             pg.UUID(as_uuid=True), primary_key=True, default=uuid4, server_default=sa.text("gen_random_uuid()")
         ),
     )
+    user_id: UUID = Field(sa_column=sa.Column(pg.UUID(as_uuid=True), sa.ForeignKey("user.id"), nullable=False))
     created_date: Optional[datetime] = Field(
         sa_column=sa.Column(sa.DateTime(), nullable=False, server_default=sa.func.current_timestamp()),
     )

backend/oasst_backend/prompt_repository.py

@@ -282,16 +282,39 @@ class PromptRepository:
                 payload = db_payload.AssistantReplyPayload(type=task.type, conversation=task.conversation)
 
             case protocol_schema.RankInitialPromptsTask:
-                payload = db_payload.RankInitialPromptsPayload(tpye=task.type, prompts=task.prompts)
+                payload = db_payload.RankInitialPromptsPayload(type=task.type, prompts=task.prompts)
 
             case protocol_schema.RankPrompterRepliesTask:
                 payload = db_payload.RankPrompterRepliesPayload(
-                    tpye=task.type, conversation=task.conversation, replies=task.replies
+                    type=task.type, conversation=task.conversation, replies=task.replies
                 )
 
             case protocol_schema.RankAssistantRepliesTask:
                 payload = db_payload.RankAssistantRepliesPayload(
-                    tpye=task.type, conversation=task.conversation, replies=task.replies
+                    type=task.type, conversation=task.conversation, replies=task.replies
+                )
+
+            case protocol_schema.LabelInitialPromptTask:
+                payload = db_payload.LabelInitialPromptPayload(
+                    type=task.type, message_id=task.message_id, prompt=task.prompt, valid_labels=task.valid_labels
+                )
+
+            case protocol_schema.LabelPrompterReplyTask:
+                payload = db_payload.LabelPrompterReplyPayload(
+                    type=task.type,
+                    message_id=task.message_id,
+                    conversation=task.conversation,
+                    reply=task.reply,
+                    valid_labels=task.valid_labels,
+                )
+
+            case protocol_schema.LabelAssistantReplyTask:
+                payload = db_payload.LabelAssistantReplyPayload(
+                    type=task.type,
+                    message_id=task.message_id,
+                    conversation=task.conversation,
+                    reply=task.reply,
+                    valid_labels=task.valid_labels,
                 )
 
             case _:
@@ -388,12 +411,12 @@ class PromptRepository:
     def store_text_labels(self, text_labels: protocol_schema.TextLabels) -> TextLabels:
         model = TextLabels(
             api_client_id=self.api_client.id,
+            message_id=text_labels.message_id,
+            user_id=self.user_id,
             text=text_labels.text,
             labels=text_labels.labels,
         )
-        if text_labels.has_message_id:
-            self.fetch_message_by_frontend_message_id(text_labels.message_id, fail_if_missing=True)
-            model.message_id = text_labels.message_id
+
         self.db.add(model)
         self.db.commit()
         self.db.refresh(model)