Message tree state machine (#555)

* add query_incomplete_rankings()

* Add SQL queries for TreeManager task selection

* first working version of TreeManager.next_task()

* remove old generate_task(), add mandatory_labels to text_labels task

* Add ConversationMessage list to Ranking tasks

* add more sophisticated sql queries to find extendible trees

* add TreeManager.query_extendible_parents()

* fix task validation, seed data insertion (reviewed)

* provide user for task selection in text-frontend

* enter 'growing' state

* enter 'aborted_low_grade' state

* enter 'ranking' state

* check tree 'growing' state upon relpy insertion

* exclude user from labeling their own messages (added DEBUG_ALLOW_SELF_LABELING setting)

* add DEBUG_ALLOW_SELF_LABELING to docker-compose.yaml

* fix ranking submission

* add query_tree_ranking_results()

* add ranked_message_ids to RankingReactionPayload

* fix reply_messages instead of prompt_messages

* incorment 'ranking_count' of ranked replies

* added logic to check_condition_for_scoring_state

* changes to msg_tree_state_machine

* pre-commit changes

* enter 'ready_for_scoring' state

* re-add HF embedding call (lost during merge)

* use prepare_conversation() helper for seed-data creation

* Partially add user specified task selection

Co-authored-by: Daniel Hug <danielpatrickhug@gmail.com>

backend/oasst_backend/models/db_payload.py

@@ -1,4 +1,4 @@
-from typing import Literal
+from typing import Literal, Optional
 from uuid import UUID
 
 from oasst_backend.models.payload_column_type import payload_type
@@ -28,7 +28,7 @@ class RateSummaryPayload(TaskPayload):
 @payload_type
 class InitialPromptPayload(TaskPayload):
     type: Literal["initial_prompt"] = "initial_prompt"
-    hint: str
+    hint: str | None
 
 
 @payload_type
@@ -64,12 +64,13 @@ class RatingReactionPayload(ReactionPayload):
 class RankingReactionPayload(ReactionPayload):
     type: Literal["message_ranking"] = "message_ranking"
     ranking: list[int]
+    ranked_message_ids: list[UUID]
 
 
 @payload_type
 class RankConversationRepliesPayload(TaskPayload):
     conversation: protocol_schema.Conversation  # the conversation so far
-    replies: list[str]
+    reply_messages: list[protocol_schema.ConversationMessage]
 
 
 @payload_type
@@ -77,7 +78,7 @@ class RankInitialPromptsPayload(TaskPayload):
     """A task to rank a set of initial prompts."""
 
     type: Literal["rank_initial_prompts"] = "rank_initial_prompts"
-    prompts: list[str]
+    prompt_messages: list[protocol_schema.ConversationMessage]
 
 
 @payload_type
@@ -102,6 +103,7 @@ class LabelInitialPromptPayload(TaskPayload):
     message_id: UUID
     prompt: str
     valid_labels: list[str]
+    mandatory_labels: Optional[list[str]]
 
 
 @payload_type
@@ -112,6 +114,7 @@ class LabelConversationReplyPayload(TaskPayload):
     conversation: protocol_schema.Conversation
     reply: str
     valid_labels: list[str]
+    mandatory_labels: Optional[list[str]]
 
 
 @payload_type

backend/oasst_backend/models/message.py

@@ -41,6 +41,10 @@ class Message(SQLModel, table=True):
     children_count: int = Field(sa_column=sa.Column(sa.Integer, default=0, server_default=sa.text("0"), nullable=False))
     deleted: bool = Field(sa_column=sa.Column(sa.Boolean, nullable=False, server_default=false()))
 
+    review_count: int = Field(sa_column=sa.Column(sa.Integer, default=0, server_default=sa.text("0"), nullable=False))
+    review_result: bool = Field(sa_column=sa.Column(sa.Boolean, default=False, server_default=false(), nullable=False))
+    ranking_count: int = Field(sa_column=sa.Column(sa.Integer, default=0, server_default=sa.text("0"), nullable=False))
+
     def ensure_is_message(self) -> None:
         if not self.payload or not isinstance(self.payload.payload, MessagePayload):
             raise OasstError("Invalid message", OasstErrorCode.INVALID_MESSAGE, HTTPStatus.INTERNAL_SERVER_ERROR)

backend/oasst_backend/models/message_tree_state.py

@@ -1,29 +1,28 @@
 from enum import Enum
-from typing import Optional
-from uuid import UUID, uuid4
+from uuid import UUID
 
 import sqlalchemy as sa
 import sqlalchemy.dialects.postgresql as pg
-from sqlmodel import Field, Index, SQLModel
+from sqlmodel import Field, SQLModel
 
 
-class States(str, Enum):
+class State(str, Enum):
     """States of the Open-Assistant message tree state machine."""
 
     INITIAL_PROMPT_REVIEW = "initial_prompt_review"
     """In this state the message tree consists only of a single inital prompt root node.
-    Initial prompt labeling tasks will determine if the tree goes into `breeding_phase` or
-    `aborted_low_grade`."""
+    Initial prompt labeling tasks will determine if the tree goes into `growing` or
+    `aborted_low_grade` state."""
 
-    BREEDING_PHASE = "breeding_phase"
+    GROWING = "growing"
     """Assistant & prompter human demonstrations are collected. Concurrently labeling tasks
     are handed out to check if the quality of the replies surpasses the minimum acceptable
     quality.
     When the required number of messages passing the initial labelling-quality check has been
-    collected the tree will enter `ranking_phase`. If too many poor-quality labelling responses
+    collected the tree will enter `ranking`. If too many poor-quality labelling responses
     are received the tree can also enter the `aborted_low_grade` state."""
 
-    RANKING_PHASE = "ranking_phase"
+    RANKING = "ranking"
     """The tree has been successfully populated with the desired number of messages. Ranking
     tasks are now handed out for all nodes with more than one child."""
 
@@ -46,28 +45,26 @@ class States(str, Enum):
 
 
 VALID_STATES = (
-    States.INITIAL_PROMPT_REVIEW,
-    States.BREEDING_PHASE,
-    States.RANKING_PHASE,
-    States.READY_FOR_SCORING,
-    States.READY_FOR_EXPORT,
-    States.ABORTED_LOW_GRADE,
+    State.INITIAL_PROMPT_REVIEW,
+    State.GROWING,
+    State.RANKING,
+    State.READY_FOR_SCORING,
+    State.READY_FOR_EXPORT,
+    State.ABORTED_LOW_GRADE,
 )
 
-TERMINAL_STATES = (States.READY_FOR_EXPORT, States.ABORTED_LOW_GRADE, States.SCORING_FAILED, States.HALTED_BY_MODERATOR)
+TERMINAL_STATES = (State.READY_FOR_EXPORT, State.ABORTED_LOW_GRADE, State.SCORING_FAILED, State.HALTED_BY_MODERATOR)
 
 
 class MessageTreeState(SQLModel, table=True):
     __tablename__ = "message_tree_state"
-    __table_args__ = (Index("ix_message_tree_state_tree_id", "message_tree_id", unique=True),)
 
-    id: Optional[UUID] = Field(
-        sa_column=sa.Column(
-            pg.UUID(as_uuid=True), primary_key=True, default=uuid4, server_default=sa.text("gen_random_uuid()")
-        ),
+    message_tree_id: UUID = Field(
+        sa_column=sa.Column(pg.UUID(as_uuid=True), sa.ForeignKey("message.id"), primary_key=True)
     )
-    message_tree_id: UUID = Field(nullable=False, index=True)
-    state: str = Field(nullable=False, max_length=128)
     goal_tree_size: int = Field(nullable=False)
-    current_num_non_filtered_messages: int = Field(nullable=False)
     max_depth: int = Field(nullable=False)
+    max_children_count: int = Field(nullable=False)
+    state: str = Field(nullable=False, max_length=128, index=True)
+    active: bool = Field(nullable=False, index=True)
+    accepted_messages: int = Field(nullable=False, default=0)