mirror of
https://github.com/wassname/Open-Assistant.git
synced 2026-07-01 16:50:12 +08:00
Daniel Hug
935e556cf6
* 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
45 lines
1.3 KiB
Python
45 lines
1.3 KiB
Python
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)
|