diff --git a/backend/alembic/versions/2023_02_01_2146-9e7ec4a9e3f2_add_skip_bool_skip_reason_to_task.py b/backend/alembic/versions/2023_02_01_2146-9e7ec4a9e3f2_add_skip_bool_skip_reason_to_task.py new file mode 100644 index 00000000..92c36adc --- /dev/null +++ b/backend/alembic/versions/2023_02_01_2146-9e7ec4a9e3f2_add_skip_bool_skip_reason_to_task.py @@ -0,0 +1,30 @@ +"""add skip bool & skip_reason to task + +Revision ID: 9e7ec4a9e3f2 +Revises: 7b8f0011e0b0 +Create Date: 2023-02-01 21:46:49.971052 + +""" +import sqlalchemy as sa +import sqlmodel +from alembic import op + +# revision identifiers, used by Alembic. +revision = "9e7ec4a9e3f2" +down_revision = "55361f323d12" +branch_labels = None +depends_on = None + + +def upgrade() -> None: + # ### commands auto generated by Alembic - please adjust! ### + op.add_column("task", sa.Column("skipped", sa.Boolean(), server_default=sa.text("false"), nullable=False)) + op.add_column("task", sa.Column("skip_reason", sqlmodel.sql.sqltypes.AutoString(length=512), nullable=True)) + # ### end Alembic commands ### + + +def downgrade() -> None: + # ### commands auto generated by Alembic - please adjust! ### + op.drop_column("task", "skip_reason") + op.drop_column("task", "skipped") + # ### end Alembic commands ### diff --git a/backend/oasst_backend/api/v1/tasks.py b/backend/oasst_backend/api/v1/tasks.py index 7db42662..c817ba67 100644 --- a/backend/oasst_backend/api/v1/tasks.py +++ b/backend/oasst_backend/api/v1/tasks.py @@ -131,7 +131,7 @@ def tasks_acknowledge_failure( logger.info(f"Frontend reports failure to implement task {task_id=}, {nack_request=}.") api_client = deps.api_auth(api_key, db) pr = PromptRepository(db, api_client, frontend_user=frontend_user) - pr.task_repository.acknowledge_task_failure(task_id) + pr.skip_task(task_id=task_id, reason=nack_request.reason) except (KeyError, RuntimeError): logger.exception("Failed to not acknowledge task.") raise OasstError("Failed to not acknowledge task.", OasstErrorCode.TASK_NACK_FAILED) diff --git a/backend/oasst_backend/models/task.py b/backend/oasst_backend/models/task.py index 7f91b157..ad41c78a 100644 --- a/backend/oasst_backend/models/task.py +++ b/backend/oasst_backend/models/task.py @@ -31,6 +31,8 @@ class Task(SQLModel, table=True): api_client_id: UUID = Field(nullable=False, foreign_key="api_client.id") ack: Optional[bool] = None done: bool = Field(sa_column=sa.Column(sa.Boolean, nullable=False, server_default=false())) + skipped: bool = Field(sa_column=sa.Column(sa.Boolean, nullable=False, server_default=false())) + skip_reason: str = Field(nullable=True, max_length=512) frontend_message_id: Optional[str] = None message_tree_id: Optional[UUID] = None parent_message_id: Optional[UUID] = None diff --git a/backend/oasst_backend/prompt_repository.py b/backend/oasst_backend/prompt_repository.py index 5d46aa2d..b805a4d7 100644 --- a/backend/oasst_backend/prompt_repository.py +++ b/backend/oasst_backend/prompt_repository.py @@ -36,6 +36,21 @@ from sqlalchemy.orm import Query from sqlalchemy.orm.attributes import flag_modified from sqlmodel import JSON, Session, and_, func, literal_column, not_, or_, text, update +_task_type_and_reaction = ( + ( + (db_payload.PrompterReplyPayload, db_payload.AssistantReplyPayload), + protocol_schema.EmojiCode.skip_reply, + ), + ( + (db_payload.LabelInitialPromptPayload, db_payload.LabelConversationReplyPayload), + protocol_schema.EmojiCode.skip_labeling, + ), + ( + (db_payload.RankInitialPromptsPayload, db_payload.RankConversationRepliesPayload), + protocol_schema.EmojiCode.skip_ranking, + ), +) + class PromptRepository: def __init__( @@ -1001,7 +1016,31 @@ WHERE message.id = cc.id; message_trees=result.get(None, 0), ) - def handle_message_emoji(self, message_id: UUID, op: protocol_schema.EmojiOp, emoji: protocol_schema) -> Message: + @managed_tx_method() + def skip_task(self, task_id: UUID, reason: str): + self.ensure_user_is_enabled() + + task = self.task_repository.fetch_task_by_id(task_id) + self._validate_task(task) + + if not task.collective: + task.skipped = True + task.skip_reason = reason + self.db.add(task) + + def handle_cancel_emoji(task_payload: db_payload.TaskPayload) -> Message | None: + for types, emoji in _task_type_and_reaction: + for t in types: + if isinstance(task_payload, t): + return self.handle_message_emoji(task.parent_message_id, protocol_schema.EmojiOp.add, emoji) + return None + + task_payload: db_payload.TaskPayload = task.payload.payload + handle_cancel_emoji(task_payload) + + def handle_message_emoji( + self, message_id: UUID, op: protocol_schema.EmojiOp, emoji: protocol_schema.EmojiCode + ) -> Message: self.ensure_user_is_enabled() message = self.fetch_message(message_id) diff --git a/backend/oasst_backend/task_repository.py b/backend/oasst_backend/task_repository.py index db4ba576..b721caf5 100644 --- a/backend/oasst_backend/task_repository.py +++ b/backend/oasst_backend/task_repository.py @@ -167,21 +167,6 @@ class TaskRepository: task.done = True self.db.add(task) - @managed_tx_method(CommitMode.COMMIT) - def acknowledge_task_failure(self, task_id): - # find task - task: Task = self.db.query(Task).filter(Task.id == task_id, Task.api_client_id == self.api_client.id).first() - if task is None: - raise OasstError(f"Task for {task_id=} not found", OasstErrorCode.TASK_NOT_FOUND, HTTP_404_NOT_FOUND) - if task.expired: - raise OasstError("Task already expired.", OasstErrorCode.TASK_EXPIRED) - if task.done or task.ack is not None: - raise OasstError("Task already updated.", OasstErrorCode.TASK_ALREADY_UPDATED) - - task.ack = False - # ToDo: check race-condition, transaction - self.db.add(task) - @managed_tx_method(CommitMode.COMMIT) def insert_task( self, diff --git a/backend/oasst_backend/tree_manager.py b/backend/oasst_backend/tree_manager.py index 709fc622..672c9081 100644 --- a/backend/oasst_backend/tree_manager.py +++ b/backend/oasst_backend/tree_manager.py @@ -13,7 +13,16 @@ from fastapi.encoders import jsonable_encoder from loguru import logger from oasst_backend.api.v1.utils import prepare_conversation, prepare_conversation_message_list from oasst_backend.config import TreeManagerConfiguration, settings -from oasst_backend.models import Message, MessageReaction, MessageTreeState, Task, TextLabels, User, message_tree_state +from oasst_backend.models import ( + Message, + MessageEmoji, + MessageReaction, + MessageTreeState, + Task, + TextLabels, + User, + message_tree_state, +) from oasst_backend.prompt_repository import PromptRepository from oasst_backend.utils import tree_export from oasst_backend.utils.database_utils import CommitMode, async_managed_tx_method, managed_tx_method @@ -22,7 +31,7 @@ from oasst_backend.utils.ranking import ranked_pairs from oasst_shared.exceptions.oasst_api_error import OasstError, OasstErrorCode from oasst_shared.schemas import protocol as protocol_schema from oasst_shared.utils import utcnow -from sqlmodel import Session, func, not_, or_, text, update +from sqlmodel import Session, and_, func, not_, or_, text, update class TaskType(Enum): @@ -889,6 +898,14 @@ class TreeManager: self.db.query(Message) .select_from(MessageTreeState) .join(Message, MessageTreeState.message_tree_id == Message.message_tree_id) + .outerjoin( + MessageEmoji, + and_( + Message.id == MessageEmoji.message_id, + MessageEmoji.user_id == self.pr.user_id, + MessageEmoji.emoji == protocol_schema.EmojiCode.skip_labeling, + ), + ) .filter( MessageTreeState.active, MessageTreeState.state == state, @@ -896,6 +913,7 @@ class TreeManager: not_(Message.deleted), Message.review_count < required_reviews, Message.lang == lang, + MessageEmoji.message_id.is_(None), ) ) @@ -948,12 +966,18 @@ SELECT m.parent_id, m.role, COUNT(m.id) children_count, MIN(m.ranking_count) chi mts.message_tree_id FROM message_tree_state mts INNER JOIN message m ON mts.message_tree_id = m.message_tree_id + LEFT JOIN message_emoji me on + (m.id = me.message_id + AND :skip_user_id IS NOT NULL + AND me.user_id = :skip_user_id + AND me.emoji = :skip_ranking) WHERE mts.active -- only consider active trees AND mts.state = :ranking_state -- message tree must be in ranking state AND m.review_result -- must be reviewed AND m.lang = :lang -- matches lang AND NOT m.deleted -- not deleted AND m.parent_id IS NOT NULL -- ignore initial prompts + AND me.message_id IS NULL -- no skip ranking emoji for user GROUP BY m.parent_id, m.role, mts.message_tree_id HAVING COUNT(m.id) > 1 and MIN(m.ranking_count) < :num_required_rankings """ @@ -979,6 +1003,8 @@ HAVING(COUNT(mr.message_id) FILTER (WHERE mr.user_id = :user_id) = 0) "ranking_state": message_tree_state.State.RANKING, "lang": lang, "user_id": user_id, + "skip_user_id": self.pr.user_id, + "skip_ranking": protocol_schema.EmojiCode.skip_ranking, }, ) return [IncompleteRankingsRow.from_orm(x) for x in r.all()] @@ -988,6 +1014,11 @@ HAVING(COUNT(mr.message_id) FILTER (WHERE mr.user_id = :user_id) = 0) SELECT m.id as parent_id, m.role as parent_role, m.depth, m.message_tree_id, COUNT(c.id) active_children_count FROM message_tree_state mts INNER JOIN message m ON mts.message_tree_id = m.message_tree_id -- all elements of message tree + LEFT JOIN message_emoji me ON + (m.id = me.message_id + AND :skip_user_id IS NOT NULL + AND me.user_id = :skip_user_id + AND me.emoji = :skip_reply) LEFT JOIN message c ON m.id = c.parent_id -- child nodes WHERE mts.active -- only consider active trees AND mts.state = :growing_state -- message tree must be growing @@ -995,6 +1026,7 @@ WHERE mts.active -- only consider active trees AND m.depth < mts.max_depth -- ignore leaf nodes as parents AND m.review_result -- parent node must have positive review AND m.lang = :lang -- parent matches lang + AND me.message_id IS NULL -- no skip reply emoji for user AND NOT coalesce(c.deleted, FALSE) -- don't count deleted children AND (c.review_result OR coalesce(c.review_count, 0) < :num_reviews_reply) -- don't count children with negative review but count elements under review GROUP BY m.id, m.role, m.depth, m.message_tree_id, mts.max_children_count @@ -1015,6 +1047,8 @@ HAVING COUNT(c.id) < mts.max_children_count -- below maximum number of children "num_prompter_replies": self.cfg.num_prompter_replies, "lang": lang, "user_id": user_id, + "skip_user_id": self.pr.user_id, + "skip_reply": protocol_schema.EmojiCode.skip_reply, }, ) @@ -1053,6 +1087,8 @@ HAVING COUNT(m.id) < mts.goal_tree_size "num_prompter_replies": self.cfg.num_prompter_replies, "lang": lang, "user_id": user_id, + "skip_user_id": self.pr.user_id, + "skip_reply": protocol_schema.EmojiCode.skip_reply, }, ) return [ActiveTreeSizeRow.from_orm(x) for x in r.all()] @@ -1525,7 +1561,8 @@ if __name__ == "__main__": with Session(engine) as db: api_client = api_auth(settings.OFFICIAL_WEB_API_KEY, db=db) # api_client = create_api_client(session=db, description="test", frontend_type="bot") - dummy_user = protocol_schema.User(id="__dummy_user__", display_name="Dummy User", auth_method="local") + # dummy_user = protocol_schema.User(id="__dummy_user__", display_name="Dummy User", auth_method="local") + dummy_user = protocol_schema.User(id="1234", display_name="bulb", auth_method="local") pr = PromptRepository(db=db, api_client=api_client, client_user=dummy_user) cfg = TreeManagerConfiguration() @@ -1540,14 +1577,18 @@ if __name__ == "__main__": # print("query_incomplete_rankings", tm.query_incomplete_rankings()) # print("query_replies_need_review", tm.query_replies_need_review()) # print("query_incomplete_reply_reviews", tm.query_replies_need_review()) - # print("query_incomplete_initial_prompt_reviews", tm.query_prompts_need_review()) + xs = tm.query_prompts_need_review(lang="en") + print("xs", len(xs)) + for x in xs: + print(x.id, x.emojis) + # print("query_incomplete_initial_prompt_reviews", tm.query_prompts_need_review(lang="en")) # print("query_extendible_trees", tm.query_extendible_trees()) # print("query_extendible_parents", tm.query_extendible_parents()) # print("next_task:", tm.next_task()) - print( - ".query_tree_ranking_results", tm.query_tree_ranking_results(UUID("21f9d585-d22c-44ab-a696-baa3d83b5f1b")) - ) + # print( + # ".query_tree_ranking_results", tm.query_tree_ranking_results(UUID("21f9d585-d22c-44ab-a696-baa3d83b5f1b")) + # ) # print(tm.export_trees_to_file(message_tree_ids=["7e75fb38-e664-4e2b-817c-b9a0b01b0074"], file="lol.jsonl")) diff --git a/oasst-shared/oasst_shared/schemas/protocol.py b/oasst-shared/oasst_shared/schemas/protocol.py index 1139286c..6bd95dc5 100644 --- a/oasst-shared/oasst_shared/schemas/protocol.py +++ b/oasst-shared/oasst_shared/schemas/protocol.py @@ -489,6 +489,11 @@ class EmojiCode(str, enum.Enum): poop = "poop" # 💩 skull = "skull" # 💀 + # skip task system uses special emoji codes + skip_reply = "skip_reply" + skip_ranking = "skip_ranking" + skip_labeling = "skip_labeling" + class EmojiOp(str, enum.Enum): togggle = "toggle"