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/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.")