Set some message trees into ranking state after seed-data insertion

backend/main.py

@@ -107,6 +107,7 @@ if settings.DEBUG_USE_SEED_DATA:
             parent_message_id: Optional[str]
             text: str
             role: str
+            tree_state: Optional[message_tree_state.State]
 
         try:
             logger.info("Seed data check began")
@@ -157,10 +158,17 @@ if settings.DEBUG_USE_SEED_DATA:
                                 )
                         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, review_count=5, review_result=True
+                            msg.text,
+                            msg.task_message_id,
+                            msg.user_message_id,
+                            review_count=5,
+                            review_result=True,
+                            check_tree_state=False,
                         )
                         if message.parent_id is None:
-                            tm._insert_default_state(root_message_id=message.id, state=message_tree_state.State.GROWING)
+                            tm._insert_default_state(
+                                root_message_id=message.id, state=msg.tree_state or message_tree_state.State.GROWING
+                            )
                             db.commit()
 
                         logger.info(
@@ -168,6 +176,7 @@ if settings.DEBUG_USE_SEED_DATA:
                         )
                     else:
                         logger.debug(f"seed data task found: {task.id}")
+
                 logger.info("Seed data check completed")
 
         except Exception:

backend/oasst_backend/prompt_repository.py

@@ -140,6 +140,7 @@ class PromptRepository:
         user_frontend_message_id: str,
         review_count: int = 0,
         review_result: bool = False,
+        check_tree_state: bool = True,
     ) -> Message:
         validate_frontend_message_id(frontend_message_id)
         validate_frontend_message_id(user_frontend_message_id)
@@ -155,12 +156,13 @@ class PromptRepository:
             parent_message = self.fetch_message(task.parent_message_id)
 
             # check tree state
-            ts = self.fetch_tree_state(parent_message.message_tree_id)
-            if not ts.active or ts.state != message_tree_state.State.GROWING:
-                raise OasstError(
-                    "Message insertion failed. Message tree is no longer in 'growing' state.",
-                    OasstErrorCode.TREE_NOT_IN_GROWING_STATE,
-                )
+            if check_tree_state:
+                ts = self.fetch_tree_state(parent_message.message_tree_id)
+                if not ts.active or ts.state != message_tree_state.State.GROWING:
+                    raise OasstError(
+                        "Message insertion failed. Message tree is no longer in 'growing' state.",
+                        OasstErrorCode.TREE_NOT_IN_GROWING_STATE,
+                    )
 
             parent_message.message_tree_id
             parent_message.children_count += 1

backend/test_data/realistic/realistic_seed_data.json

@@ -228,7 +228,8 @@
     "user_message_id": "rrbopvrb",
     "parent_message_id": null,
     "text": " How can I feed squirrels and make  more come to my home?",
-    "role": "prompter"
+    "role": "prompter",
+    "tree_state": "ranking"
   },
   {
     "task_message_id": "92w6z5iw",
@@ -1054,7 +1055,8 @@
     "user_message_id": "3pvmby5u",
     "parent_message_id": null,
     "text": " My 35 year old son is having pain from his jaw, neck and shoulder blade",
-    "role": "prompter"
+    "role": "prompter",
+    "tree_state": "ranking"
   },
   {
     "task_message_id": "v891mplm",
@@ -1418,7 +1420,8 @@
     "user_message_id": "4j9k8rqq",
     "parent_message_id": null,
     "text": " I have been gaining weight and want to start exercising.  How should I begin?",
-    "role": "prompter"
+    "role": "prompter",
+    "tree_state": "ranking"
   },
   {
     "task_message_id": "kfg4s8bg",
@@ -1600,7 +1603,8 @@
     "user_message_id": "vgtyn75k",
     "parent_message_id": null,
     "text": " What is macular degeneration?",
-    "role": "prompter"
+    "role": "prompter",
+    "tree_state": "ranking"
   },
   {
     "task_message_id": "zusgcbfj",
@@ -2041,7 +2045,8 @@
     "user_message_id": "vtaes87s",
     "parent_message_id": null,
     "text": " I keep biting the inside of my cheek when I'm sleeping. What can I do to prevent this in the future?",
-    "role": "prompter"
+    "role": "prompter",
+    "tree_state": "ranking"
   },
   {
     "task_message_id": "zxumt49x",
@@ -2636,7 +2641,8 @@
     "user_message_id": "v14u082o",
     "parent_message_id": null,
     "text": " Any tips on how to revise something you wrote?",
-    "role": "prompter"
+    "role": "prompter",
+    "tree_state": "ranking"
   },
   {
     "task_message_id": "bcw4led5",
@@ -2986,7 +2992,8 @@
     "user_message_id": "di54e1nb",
     "parent_message_id": null,
     "text": " I am going to Singapore on business.  Can you tell me what you know about Singapore?",
-    "role": "prompter"
+    "role": "prompter",
+    "tree_state": "ranking"
   },
   {
     "task_message_id": "z79vpkaj",