From 448d17dd61db681cfd71c9c34d8dd7f1279891b4 Mon Sep 17 00:00:00 2001 From: Andreas Koepf Date: Fri, 13 Jan 2023 13:57:55 +0000 Subject: [PATCH] Set some message trees into ranking state after seed-data insertion --- backend/main.py | 13 ++++++++++-- backend/oasst_backend/prompt_repository.py | 14 +++++++------ .../realistic/realistic_seed_data.json | 21 ++++++++++++------- 3 files changed, 33 insertions(+), 15 deletions(-) diff --git a/backend/main.py b/backend/main.py index c13e56df..76b0a35b 100644 --- a/backend/main.py +++ b/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: diff --git a/backend/oasst_backend/prompt_repository.py b/backend/oasst_backend/prompt_repository.py index 5bae3b18..cb6e70f7 100644 --- a/backend/oasst_backend/prompt_repository.py +++ b/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 diff --git a/backend/test_data/realistic/realistic_seed_data.json b/backend/test_data/realistic/realistic_seed_data.json index be2a17e9..ce8e7ff3 100644 --- a/backend/test_data/realistic/realistic_seed_data.json +++ b/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",