Add message tree state model (#414)

* inital commit, created file for MessageTreeState table

* added initial implementation of MessageTreeState SQl model and added it to__init__

* Ran Alembic revision for migration

* reran precommit on

* removed create_data and deleted flag from model

* ran migration revision to remove unwant class variables date_created and deleted

* removed unused imports

* ran pre-commit

* Updated States definiton to be Enums instead fo Constants

* ran pre-commit formatting

* Fixed Enum class

* ran pre-commit

backend/alembic/versions/2023_01_05_1144-d4161e384f83_added_messagetreestate_table.py

@@ -0,0 +1,46 @@
+"""added MessageTreeState table
+
+Revision ID: d4161e384f83
+Revises: 8d269bc4fdbd
+Create Date: 2023-01-05 11:44:02.630633
+
+"""
+import sqlalchemy as sa
+import sqlmodel
+from alembic import op
+from sqlalchemy.dialects import postgresql
+
+# revision identifiers, used by Alembic.
+revision = "d4161e384f83"
+down_revision = "8d269bc4fdbd"
+branch_labels = None
+depends_on = None
+
+
+def upgrade() -> None:
+    # ### commands auto generated by Alembic - please adjust! ###
+    op.create_table(
+        "message_tree_state",
+        sa.Column("id", postgresql.UUID(as_uuid=True), server_default=sa.text("gen_random_uuid()"), nullable=False),
+        sa.Column("created_date", sa.DateTime(), server_default=sa.text("CURRENT_TIMESTAMP"), nullable=False),
+        sa.Column("deleted", sa.Boolean(), server_default=sa.text("false"), nullable=False),
+        sa.Column("message_tree_id", sqlmodel.sql.sqltypes.GUID(), nullable=False),
+        sa.Column("state", sqlmodel.sql.sqltypes.AutoString(length=128), nullable=False),
+        sa.Column("goal_tree_size", sa.Integer(), nullable=False),
+        sa.Column("current_num_non_filtered_messages", sa.Integer(), nullable=False),
+        sa.Column("max_depth", sa.Integer(), nullable=False),
+        sa.PrimaryKeyConstraint("id"),
+    )
+    op.create_index(
+        op.f("ix_message_tree_state_message_tree_id"), "message_tree_state", ["message_tree_id"], unique=False
+    )
+    op.create_index("ix_message_tree_state_tree_id", "message_tree_state", ["message_tree_id"], unique=True)
+    # ### end Alembic commands ###
+
+
+def downgrade() -> None:
+    # ### commands auto generated by Alembic - please adjust! ###
+    op.drop_index("ix_message_tree_state_tree_id", table_name="message_tree_state")
+    op.drop_index(op.f("ix_message_tree_state_message_tree_id"), table_name="message_tree_state")
+    op.drop_table("message_tree_state")
+    # ### end Alembic commands ###

backend/alembic/versions/2023_01_05_1346-3b0adfadbef9_removed_date_created_and_deleted_flag_.py

@@ -0,0 +1,42 @@
+"""removed date_created and deleted flag from message_tree_state
+
+Revision ID: 3b0adfadbef9
+Revises: d4161e384f83
+Create Date: 2023-01-05 13:46:11.338655
+
+"""
+import sqlalchemy as sa
+from alembic import op
+from sqlalchemy.dialects import postgresql
+
+# revision identifiers, used by Alembic.
+revision = "3b0adfadbef9"
+down_revision = "d4161e384f83"
+branch_labels = None
+depends_on = None
+
+
+def upgrade() -> None:
+    # ### commands auto generated by Alembic - please adjust! ###
+    op.drop_column("message_tree_state", "deleted")
+    op.drop_column("message_tree_state", "created_date")
+    # ### end Alembic commands ###
+
+
+def downgrade() -> None:
+    # ### commands auto generated by Alembic - please adjust! ###
+    op.add_column(
+        "message_tree_state",
+        sa.Column(
+            "created_date",
+            postgresql.TIMESTAMP(),
+            server_default=sa.text("CURRENT_TIMESTAMP"),
+            autoincrement=False,
+            nullable=False,
+        ),
+    )
+    op.add_column(
+        "message_tree_state",
+        sa.Column("deleted", sa.BOOLEAN(), server_default=sa.text("false"), autoincrement=False, nullable=False),
+    )
+    # ### end Alembic commands ###

backend/oasst_backend/models/__init__.py

@@ -2,6 +2,7 @@ from .api_client import ApiClient
 from .journal import Journal, JournalIntegration
 from .message import Message
 from .message_reaction import MessageReaction
+from .message_tree_state import MessageTreeState
 from .task import Task
 from .text_labels import TextLabels
 from .user import User
@@ -13,6 +14,7 @@ __all__ = [
     "UserStats",
     "Message",
     "MessageReaction",
+    "MessageTreeState",
     "Task",
     "TextLabels",
     "Journal",

backend/oasst_backend/models/message_tree_state.py

@@ -0,0 +1,44 @@
+from enum import Enum
+from typing import Optional
+from uuid import UUID, uuid4
+
+import sqlalchemy as sa
+import sqlalchemy.dialects.postgresql as pg
+from sqlmodel import Field, Index, SQLModel
+
+# The types of States a message tree can have.
+
+
+class States(Enum):
+    INITIAL = "initial"
+    BREEDING_PHASE = "breeding_phase"
+    RANKING_PHASE = "ranking_phase"
+    READY_FOR_SCORING = "ready_for_scoring"
+    CHILDREN_SCORED = "children_scored"
+    FINAL = "final"
+
+
+VALID_STATES = (
+    States.INITIAL,
+    States.BREEDING_PHASE,
+    States.RANKING_PHASE,
+    States.READY_FOR_SCORING,
+    States.CHILDREN_SCORED,
+    States.FINAL,
+)
+
+
+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(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)