Message tree state machine (LAION-AI#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>

Author: andreaskoepf

ansible/dev.yaml

@@ -79,6 +79,7 @@
           REDIS_HOST: oasst-redis
           DEBUG_ALLOW_ANY_API_KEY: "true"
           DEBUG_USE_SEED_DATA: "true"
+          DEBUG_ALLOW_SELF_LABELING: "true"
           MAX_WORKERS: "1"
           RATE_LIMIT: "false"
           DEBUG_SKIP_EMBEDDING_COMPUTATION: "true"

backend/alembic/versions/2023_01_08_2208-92a367bb9f40_restructure_message_tree_state_table.py

@@ -0,0 +1,63 @@
+"""restructure message_tree_state table
+
+Revision ID: 92a367bb9f40
+Revises: ba61fe17fb6e
+Create Date: 2023-01-08 22:08:46.458195
+
+"""
+import sqlalchemy as sa
+import sqlmodel
+from alembic import op
+from sqlalchemy.dialects import postgresql
+
+# revision identifiers, used by Alembic.
+revision = "92a367bb9f40"
+down_revision = "aac6b2f66006"
+branch_labels = None
+depends_on = None
+
+
+def upgrade() -> None:
+    # ### commands auto generated by Alembic - please adjust! ###
+    op.drop_table("message_tree_state")
+    op.create_table(
+        "message_tree_state",
+        sa.Column("message_tree_id", postgresql.UUID(as_uuid=True), nullable=False),
+        sa.Column("goal_tree_size", sa.Integer(), nullable=False),
+        sa.Column("max_depth", sa.Integer(), nullable=False),
+        sa.Column("max_children_count", sa.Integer(), nullable=False),
+        sa.Column("state", sqlmodel.sql.sqltypes.AutoString(length=128), nullable=False),
+        sa.Column("active", sa.Boolean(), nullable=False),
+        sa.Column("accepted_messages", sa.Integer(), nullable=False),
+        sa.ForeignKeyConstraint(
+            ["message_tree_id"],
+            ["message.id"],
+        ),
+        sa.PrimaryKeyConstraint("message_tree_id"),
+    )
+    op.create_index(op.f("ix_message_tree_state_active"), "message_tree_state", ["active"], unique=False)
+    op.create_index(op.f("ix_message_tree_state_state"), "message_tree_state", ["state"], unique=False)
+
+    # ### end Alembic commands ###
+
+
+def downgrade() -> None:
+    # ### commands auto generated by Alembic - please adjust! ###
+    op.drop_index(op.f("ix_message_tree_state_state"), table_name="message_tree_state")
+    op.drop_index(op.f("ix_message_tree_state_active"), table_name="message_tree_state")
+    op.drop_table("message_tree_state")
+    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("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 ###

backend/alembic/versions/2023_01_09_0047-05975b274a81_add_review_count_ranking_count_to_.py

@@ -0,0 +1,31 @@
+"""add review_count & ranking_count to message
+
+Revision ID: 05975b274a81
+Revises: 92a367bb9f40
+Create Date: 2023-01-09 00:47:25.496036
+
+"""
+import sqlalchemy as sa
+from alembic import op
+
+# revision identifiers, used by Alembic.
+revision = "05975b274a81"
+down_revision = "92a367bb9f40"
+branch_labels = None
+depends_on = None
+
+
+def upgrade() -> None:
+    # ### commands auto generated by Alembic - please adjust! ###
+    op.add_column("message", sa.Column("review_count", sa.Integer(), server_default=sa.text("0"), nullable=False))
+    op.add_column("message", sa.Column("review_result", sa.Boolean(), server_default=sa.text("false"), nullable=False))
+    op.add_column("message", sa.Column("ranking_count", sa.Integer(), server_default=sa.text("0"), nullable=False))
+    # ### end Alembic commands ###
+
+
+def downgrade() -> None:
+    # ### commands auto generated by Alembic - please adjust! ###
+    op.drop_column("message", "ranking_count")
+    op.drop_column("message", "review_result")
+    op.drop_column("message", "review_count")
+    # ### end Alembic commands ###

backend/main.py

@@ -12,9 +12,12 @@
 from loguru import logger
 from oasst_backend.api.deps import get_dummy_api_client
 from oasst_backend.api.v1.api import api_router
+from oasst_backend.api.v1.utils import prepare_conversation
 from oasst_backend.config import settings
 from oasst_backend.database import engine
+from oasst_backend.models import message_tree_state
 from oasst_backend.prompt_repository import PromptRepository, TaskRepository, UserRepository
+from oasst_backend.tree_manager import TreeManager, TreeManagerConfiguration
 from oasst_shared.exceptions import OasstError, OasstErrorCode
 from oasst_shared.schemas import protocol as protocol_schema
 from pydantic import BaseModel
@@ -116,6 +119,7 @@ class DummyMessage(BaseModel):
                 pr = PromptRepository(
                     db=db, api_client=api_client, client_user=dummy_user, user_repository=ur, task_repository=tr
                 )
+                tm = TreeManager(db, pr, TreeManagerConfiguration())
 
                 with open(settings.DEBUG_USE_SEED_DATA_PATH) as f:
                     dummy_messages_raw = json.load(f)
@@ -138,24 +142,19 @@ class DummyMessage(BaseModel):
                                 msg.parent_message_id, fail_if_missing=True
                             )
                             conversation_messages = pr.fetch_message_conversation(parent_message)
-                            conversation = protocol_schema.Conversation(
-                                messages=[
-                                    protocol_schema.ConversationMessage(
-                                        text=cmsg.text,
-                                        is_assistant=cmsg.role == "assistant",
-                                        message_id=cmsg.id,
-                                        fronend_message_id=cmsg.frontend_message_id,
-                                    )
-                                    for cmsg in conversation_messages
-                                ]
-                            )
+                            conversation = prepare_conversation(conversation_messages)
                             task = tr.store_task(
                                 protocol_schema.AssistantReplyTask(conversation=conversation),
                                 message_tree_id=parent_message.message_tree_id,
                                 parent_message_id=parent_message.id,
                             )
                         tr.bind_frontend_message_id(task.id, msg.task_message_id)
-                        message = pr.store_text_reply(msg.text, msg.task_message_id, msg.user_message_id)
+                        message = pr.store_text_reply(
+                            msg.text, msg.task_message_id, msg.user_message_id, review_count=5, review_result=True
+                        )
+                        if message.parent_id is None:
+                            tm._insert_default_state(root_message_id=message.id, state=message_tree_state.State.GROWING)
+                            db.commit()
 
                         logger.info(
                             f"Inserted: message_id: {message.id}, payload: {message.payload.payload}, parent_message_id: {message.parent_id}"
@@ -168,14 +167,27 @@ class DummyMessage(BaseModel):
             logger.exception("Seed data insertion failed")
 
 
+@app.on_event("startup")
+def ensure_tree_states():
+    try:
+        logger.info("Startup: TreeManager.ensure_tree_states()")
+        cfg = TreeManagerConfiguration()  # TODO: decide where config is stored, e.g. load form json/yaml file
+        with Session(engine) as db:
+            tm = TreeManager(db, None, configuration=cfg)
+            tm.ensure_tree_states()
+
+    except Exception:
+        logger.exception("TreeManager.ensure_tree_states() failed.")
+
+
 app.include_router(api_router, prefix=settings.API_V1_STR)
 
 
 def get_openapi_schema():
     return json.dumps(app.openapi())
 
 
-if __name__ == "__main__":
+def main():
     # Importing here so we don't import packages unnecessarily if we're
     # importing main as a module.
     import argparse
@@ -198,3 +210,7 @@ def get_openapi_schema():
         print(get_openapi_schema())
     else:
         uvicorn.run(app, host=args.host, port=args.port)
+
+
+if __name__ == "__main__":
+    main()