Files
Open-Assistant/backend/oasst_backend/models/db_payload.py
T
Andreas Köpf 14fa08e2e7 Message tree state machine (#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>
2023-01-11 10:54:03 +01:00

132 lines
3.3 KiB
Python

from typing import Literal, Optional
from uuid import UUID
from oasst_backend.models.payload_column_type import payload_type
from oasst_shared.schemas import protocol as protocol_schema
from pydantic import BaseModel
@payload_type
class TaskPayload(BaseModel):
type: str
@payload_type
class SummarizationStoryPayload(TaskPayload):
type: Literal["summarize_story"] = "summarize_story"
story: str
@payload_type
class RateSummaryPayload(TaskPayload):
type: Literal["rate_summary"] = "rate_summary"
full_text: str
summary: str
scale: protocol_schema.RatingScale
@payload_type
class InitialPromptPayload(TaskPayload):
type: Literal["initial_prompt"] = "initial_prompt"
hint: str | None
@payload_type
class PrompterReplyPayload(TaskPayload):
type: Literal["prompter_reply"] = "prompter_reply"
conversation: protocol_schema.Conversation
hint: str | None
@payload_type
class AssistantReplyPayload(TaskPayload):
type: Literal["assistant_reply"] = "assistant_reply"
conversation: protocol_schema.Conversation
@payload_type
class MessagePayload(BaseModel):
text: str
@payload_type
class ReactionPayload(BaseModel):
type: str
@payload_type
class RatingReactionPayload(ReactionPayload):
type: Literal["message_rating"] = "message_rating"
rating: str
@payload_type
class RankingReactionPayload(ReactionPayload):
type: Literal["message_ranking"] = "message_ranking"
ranking: list[int]
ranked_message_ids: list[UUID]
@payload_type
class RankConversationRepliesPayload(TaskPayload):
conversation: protocol_schema.Conversation # the conversation so far
reply_messages: list[protocol_schema.ConversationMessage]
@payload_type
class RankInitialPromptsPayload(TaskPayload):
"""A task to rank a set of initial prompts."""
type: Literal["rank_initial_prompts"] = "rank_initial_prompts"
prompt_messages: list[protocol_schema.ConversationMessage]
@payload_type
class RankPrompterRepliesPayload(RankConversationRepliesPayload):
"""A task to rank a set of prompter replies to a conversation."""
type: Literal["rank_prompter_replies"] = "rank_prompter_replies"
@payload_type
class RankAssistantRepliesPayload(RankConversationRepliesPayload):
"""A task to rank a set of assistant replies to a conversation."""
type: Literal["rank_assistant_replies"] = "rank_assistant_replies"
@payload_type
class LabelInitialPromptPayload(TaskPayload):
"""A task to label an initial prompt."""
type: Literal["label_initial_prompt"] = "label_initial_prompt"
message_id: UUID
prompt: str
valid_labels: list[str]
mandatory_labels: Optional[list[str]]
@payload_type
class LabelConversationReplyPayload(TaskPayload):
"""A task to label a conversation reply."""
message_id: UUID
conversation: protocol_schema.Conversation
reply: str
valid_labels: list[str]
mandatory_labels: Optional[list[str]]
@payload_type
class LabelPrompterReplyPayload(LabelConversationReplyPayload):
"""A task to label a prompter reply."""
type: Literal["label_prompter_reply"] = "label_prompter_reply"
@payload_type
class LabelAssistantReplyPayload(LabelConversationReplyPayload):
"""A task to label an assistant reply."""
type: Literal["label_assistant_reply"] = "label_assistant_reply"