From 2786ad613241b985a32d0d3b9c920b4d4c531e22 Mon Sep 17 00:00:00 2001 From: Igor Miagkov Date: Wed, 28 Dec 2022 11:36:38 +0400 Subject: [PATCH] Sampling random posts for tasks --- backend/oasst_backend/api/v1/tasks.py | 139 +++++++++--------- backend/oasst_backend/models/post.py | 2 + backend/oasst_backend/models/post_reaction.py | 4 +- backend/oasst_backend/models/work_package.py | 10 ++ 4 files changed, 83 insertions(+), 72 deletions(-) diff --git a/backend/oasst_backend/api/v1/tasks.py b/backend/oasst_backend/api/v1/tasks.py index 0778fd4c..cf568334 100644 --- a/backend/oasst_backend/api/v1/tasks.py +++ b/backend/oasst_backend/api/v1/tasks.py @@ -1,6 +1,6 @@ # -*- coding: utf-8 -*- import random -from typing import Any +from typing import Any, Optional, Tuple from uuid import UUID from fastapi import APIRouter, Depends, HTTPException @@ -15,13 +15,18 @@ from starlette.status import HTTP_400_BAD_REQUEST router = APIRouter() -def generate_task(request: protocol_schema.TaskRequest) -> protocol_schema.Task: +def generate_task( + request: protocol_schema.TaskRequest, pr: PromptRepository +) -> Tuple[protocol_schema.Task, Optional[UUID], Optional[UUID]]: + thread_id = None + parent_post_id = None + match request.type: case protocol_schema.TaskRequestType.random: logger.info("Frontend requested a random task.") while request.type == protocol_schema.TaskRequestType.random: request.type = random.choice(list(protocol_schema.TaskRequestType)).value - return generate_task(request) + return generate_task(request, pr) case protocol_schema.TaskRequestType.summarize_story: logger.info("Generating a SummarizeStoryTask.") task = protocol_schema.SummarizeStoryTask( @@ -41,77 +46,65 @@ def generate_task(request: protocol_schema.TaskRequest) -> protocol_schema.Task: ) case protocol_schema.TaskRequestType.user_reply: logger.info("Generating a UserReplyTask.") - task = protocol_schema.UserReplyTask( - conversation=protocol_schema.Conversation( - messages=[ - protocol_schema.ConversationMessage( - text="Hey, assistant, what's going on in the world?", - is_assistant=False, - ), - protocol_schema.ConversationMessage( - text="I'm not sure I understood correctly, could you rephrase that?", - is_assistant=True, - ), - ], - ) - ) + posts = pr.fetch_random_conversation("assistant") + messages = [ + protocol_schema.ConversationMessage(text=p.payload.payload.text, is_assistant=(p.role == "assistant")) + for p in posts + ] + + task = protocol_schema.UserReplyTask(conversation=protocol_schema.Conversation(messages=messages)) + thread_id = posts[-1].thread_id + parent_post_id = posts[-1].id case protocol_schema.TaskRequestType.assistant_reply: logger.info("Generating a AssistantReplyTask.") - task = protocol_schema.AssistantReplyTask( - conversation=protocol_schema.Conversation( - messages=[ - protocol_schema.ConversationMessage( - text="Hey, assistant, write me an English essay about water.", - is_assistant=False, - ), - ], - ) - ) + posts = pr.fetch_random_conversation("user") + messages = [ + protocol_schema.ConversationMessage(text=p.payload.payload.text, is_assistant=(p.role == "assistant")) + for p in posts + ] + + task = protocol_schema.AssistantReplyTask(conversation=protocol_schema.Conversation(messages=messages)) + thread_id = posts[-1].thread_id + parent_post_id = posts[-1].id case protocol_schema.TaskRequestType.rank_initial_prompts: logger.info("Generating a RankInitialPromptsTask.") - task = protocol_schema.RankInitialPromptsTask( - prompts=[ - "Please write a story about a time you were happy.", - "Please write a story about a time you were sad.", - ] - ) + + posts = pr.fetch_random_initial_prompts() + task = protocol_schema.RankInitialPromptsTask(prompts=[p.payload.payload.text for p in posts]) case protocol_schema.TaskRequestType.rank_user_replies: logger.info("Generating a RankUserRepliesTask.") + conversation, replies = pr.fetch_multiple_random_replies(post_role="assistant") + + messages = [ + protocol_schema.ConversationMessage( + text=p.payload.payload.text, + is_assistant=(p.role == "assistant"), + ) + for p in conversation + ] + replies = [p.payload.payload.text for p in replies] task = protocol_schema.RankUserRepliesTask( conversation=protocol_schema.Conversation( - messages=[ - protocol_schema.ConversationMessage( - text="Hey, assistant, what's going on in the world?", - is_assistant=False, - ), - protocol_schema.ConversationMessage( - text="I'm not sure I understood correctly, could you rephrase that?", - is_assistant=True, - ), - ], + messages=messages, ), - replies=[ - "Oh come oooooon!", - "What are the news?", - ], + replies=replies, ) case protocol_schema.TaskRequestType.rank_assistant_replies: logger.info("Generating a RankAssistantRepliesTask.") + conversation, replies = pr.fetch_multiple_random_replies(post_role="user") + + messages = [ + protocol_schema.ConversationMessage( + text=p.payload.payload.text, + is_assistant=(p.role == "assistant"), + ) + for p in conversation + ] + replies = [p.payload.payload.text for p in replies] task = protocol_schema.RankAssistantRepliesTask( - conversation=protocol_schema.Conversation( - messages=[ - protocol_schema.ConversationMessage( - text="Hey, assistant, what's going on in the world?", - is_assistant=False, - ), - ], - ), - replies=[ - "I'm not sure I understood correctly, could you rephrase that?", - "The world is fine. All good.", - "Crap is hitting the fan. Start farming.", - ], + conversation=protocol_schema.Conversation(messages=messages), + replies=replies, ) case _: raise HTTPException( @@ -121,7 +114,7 @@ def generate_task(request: protocol_schema.TaskRequest) -> protocol_schema.Task: logger.info(f"Generated {task=}.") - return task + return task, thread_id, parent_post_id @router.post("/", response_model=protocol_schema.AnyTask) # work with Union once more types are added @@ -137,10 +130,9 @@ def request_task( api_client = deps.api_auth(api_key, db) try: - task = generate_task(request) - pr = PromptRepository(db, api_client, request.user) - pr.store_task(task) + task, thread_id, parent_post_id = generate_task(request, pr) + pr.store_task(task, thread_id, parent_post_id) except Exception: logger.exception("Failed to generate task.") @@ -190,11 +182,17 @@ def acknowledge_task_failure( """ The frontend reports failure to implement a task. """ - deps.api_auth(api_key, db) - logger.info(f"Frontend reports failure to implement task {task_id=}, {nack_request=}.") - # here we would store the post id in the database for the task - return {} + try: + 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, user=None) + pr.acknowledge_task_failure(task_id) + except (KeyError, RuntimeError): + logger.exception("Failed to acknowledge task.") + raise HTTPException( + status_code=HTTP_400_BAD_REQUEST, + ) @router.post("/interaction") @@ -219,8 +217,9 @@ def post_interaction( ) # here we store the text reply in the database - # ToDo: role user or agent? - pr.store_text_reply(interaction, role="unknown") + pr.store_text_reply( + text=interaction.text, post_id=interaction.post_id, user_post_id=interaction.user_post_id + ) return protocol_schema.TaskDone() case protocol_schema.PostRating: diff --git a/backend/oasst_backend/models/post.py b/backend/oasst_backend/models/post.py index d1569a67..ffa52240 100644 --- a/backend/oasst_backend/models/post.py +++ b/backend/oasst_backend/models/post.py @@ -31,3 +31,5 @@ class Post(SQLModel, table=True): ) payload_type: str = Field(nullable=False, max_length=200) payload: PayloadContainer = Field(sa_column=sa.Column(payload_column_type(PayloadContainer), nullable=True)) + depth: int = Field(sa_column=sa.Column(sa.Integer, default=0, server_default=sa.text("0"), nullable=False)) + children_count: int = Field(sa_column=sa.Column(sa.Integer, default=0, server_default=sa.text("0"), nullable=False)) diff --git a/backend/oasst_backend/models/post_reaction.py b/backend/oasst_backend/models/post_reaction.py index fb8b0fcb..f6141e56 100644 --- a/backend/oasst_backend/models/post_reaction.py +++ b/backend/oasst_backend/models/post_reaction.py @@ -13,8 +13,8 @@ from .payload_column_type import PayloadContainer, payload_column_type class PostReaction(SQLModel, table=True): __tablename__ = "post_reaction" - post_id: Optional[UUID] = Field( - sa_column=sa.Column(pg.UUID(as_uuid=True), sa.ForeignKey("post.id"), nullable=False, primary_key=True) + work_package_id: Optional[UUID] = Field( + sa_column=sa.Column(pg.UUID(as_uuid=True), sa.ForeignKey("work_package.id"), nullable=False, primary_key=True) ) person_id: UUID = Field( sa_column=sa.Column(pg.UUID(as_uuid=True), sa.ForeignKey("person.id"), nullable=False, primary_key=True) diff --git a/backend/oasst_backend/models/work_package.py b/backend/oasst_backend/models/work_package.py index 0a452a71..a89ed646 100644 --- a/backend/oasst_backend/models/work_package.py +++ b/backend/oasst_backend/models/work_package.py @@ -5,6 +5,7 @@ from uuid import UUID, uuid4 import sqlalchemy as sa import sqlalchemy.dialects.postgresql as pg +from sqlalchemy import false from sqlmodel import Field, SQLModel from .payload_column_type import PayloadContainer, payload_column_type @@ -26,3 +27,12 @@ class WorkPackage(SQLModel, table=True): payload_type: str = Field(nullable=False, max_length=200) payload: PayloadContainer = Field(sa_column=sa.Column(payload_column_type(PayloadContainer), nullable=False)) 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())) + frontend_ref_post_id: Optional[str] = None + thread_id: Optional[UUID] = None + parent_post_id: Optional[UUID] = None + + @property + def expired(self) -> bool: + return self.expiry_date is not None and datetime.utcnow() < self.expiry_date