diff --git a/backend/alembic/versions/2022_12_31_0438-8d269bc4fdbd_add_deleted_field_to_post.py b/backend/alembic/versions/2022_12_31_0438-8d269bc4fdbd_add_deleted_field_to_post.py new file mode 100644 index 00000000..786471db --- /dev/null +++ b/backend/alembic/versions/2022_12_31_0438-8d269bc4fdbd_add_deleted_field_to_post.py @@ -0,0 +1,28 @@ +# -*- coding: utf-8 -*- +"""add deleted field to post + +Revision ID: 8d269bc4fdbd +Revises: abb47e9d145a +Create Date: 2022-12-31 04:38:41.799206 + +""" +import sqlalchemy as sa +from alembic import op + +# revision identifiers, used by Alembic. +revision = "8d269bc4fdbd" +down_revision = "abb47e9d145a" +branch_labels = None +depends_on = None + + +def upgrade() -> None: + # ### commands auto generated by Alembic - please adjust! ### + op.add_column("message", sa.Column("deleted", sa.Boolean(), server_default=sa.text("false"), nullable=False)) + # ### end Alembic commands ### + + +def downgrade() -> None: + # ### commands auto generated by Alembic - please adjust! ### + op.drop_column("message", "deleted") + # ### end Alembic commands ### diff --git a/backend/oasst_backend/api/deps.py b/backend/oasst_backend/api/deps.py index 9c4feee2..e0286ba3 100644 --- a/backend/oasst_backend/api/deps.py +++ b/backend/oasst_backend/api/deps.py @@ -4,7 +4,7 @@ from secrets import token_hex from typing import Generator from uuid import UUID -from fastapi import Security +from fastapi import Depends, Security from fastapi.security.api_key import APIKey, APIKeyHeader, APIKeyQuery from loguru import logger from oasst_backend.config import settings @@ -64,3 +64,24 @@ def api_auth( error_code=OasstErrorCode.API_CLIENT_NOT_AUTHORIZED, http_status_code=HTTPStatus.FORBIDDEN, ) + + +def get_api_client( + api_key: APIKey = Depends(get_api_key), + db: Session = Depends(get_db), +): + return api_auth(api_key, db) + + +def get_trusted_api_client( + api_key: APIKey = Depends(get_api_key), + db: Session = Depends(get_db), +): + client = api_auth(api_key, db) + if not client.trusted: + raise OasstError( + "Forbidden", + error_code=OasstErrorCode.API_CLIENT_NOT_AUTHORIZED, + http_status_code=HTTPStatus.FORBIDDEN, + ) + return client diff --git a/backend/oasst_backend/api/v1/api.py b/backend/oasst_backend/api/v1/api.py index b54f3dd0..2286a1ac 100644 --- a/backend/oasst_backend/api/v1/api.py +++ b/backend/oasst_backend/api/v1/api.py @@ -1,7 +1,12 @@ # -*- coding: utf-8 -*- from fastapi import APIRouter -from oasst_backend.api.v1 import tasks, text_labels +from oasst_backend.api.v1 import frontend_messages, frontend_users, messages, stats, tasks, text_labels, users api_router = APIRouter() api_router.include_router(tasks.router, prefix="/tasks", tags=["tasks"]) api_router.include_router(text_labels.router, prefix="/text_labels", tags=["text_labels"]) +api_router.include_router(messages.router, prefix="/messages", tags=["messages"]) +api_router.include_router(frontend_messages.router, prefix="/frontend_messages", tags=["frontend_messages"]) +api_router.include_router(users.router, prefix="/users", tags=["users"]) +api_router.include_router(frontend_users.router, prefix="/frontend_users", tags=["frontend_users"]) +api_router.include_router(stats.router, prefix="/stats", tags=["stats"]) diff --git a/backend/oasst_backend/api/v1/frontend_messages.py b/backend/oasst_backend/api/v1/frontend_messages.py new file mode 100644 index 00000000..85ab74ca --- /dev/null +++ b/backend/oasst_backend/api/v1/frontend_messages.py @@ -0,0 +1,118 @@ +# -*- coding: utf-8 -*- + +from fastapi import APIRouter, Depends +from oasst_backend.api import deps +from oasst_backend.api.v1 import utils +from oasst_backend.exceptions import OasstError, OasstErrorCode +from oasst_backend.models import ApiClient +from oasst_backend.models.db_payload import MessagePayload +from oasst_backend.prompt_repository import PromptRepository +from oasst_shared.schemas import protocol +from sqlmodel import Session + +router = APIRouter() + + +@router.get("/{message_id}") +def get_message_by_frontend_id( + message_id: str, api_client: ApiClient = Depends(deps.get_api_client), db: Session = Depends(deps.get_db) +): + """ + Get a message by its frontend ID. + """ + pr = PromptRepository(db, api_client, user=None) + message = pr.fetch_message_by_frontend_message_id(message_id) + + if not isinstance(message.payload.payload, MessagePayload): + # Unexpected message payload + raise OasstError("Invalid message", OasstErrorCode.INVALID_MESSAGE) + + return protocol.ConversationMessage(text=message.payload.payload.text, is_assistant=(message.role == "assistant")) + + +@router.get("/{message_id}/conversation") +def get_conv_by_frontend_id( + message_id: str, api_client: ApiClient = Depends(deps.get_api_client), db: Session = Depends(deps.get_db) +): + """ + Get a conversation from the tree root and up to the message with given frontend ID. + """ + + pr = PromptRepository(db, api_client, user=None) + message = pr.fetch_message_by_frontend_message_id(message_id) + messages = pr.fetch_message_conversation(message) + return utils.prepare_conversation(messages) + + +@router.get("/{message_id}/tree") +def get_tree_by_frontend_id( + message_id: str, api_client: ApiClient = Depends(deps.get_api_client), db: Session = Depends(deps.get_db) +): + """ + Get all messages belonging to the same message tree. + Message is identified by its frontend ID. + """ + pr = PromptRepository(db, api_client, user=None) + message = pr.fetch_message_by_frontend_message_id(message_id) + tree = pr.fetch_message_tree(message.message_tree_id) + return utils.prepare_tree(tree, message.message_tree_id) + + +@router.get("/{message_id}/children") +def get_children_by_frontend_id( + message_id: str, api_client: ApiClient = Depends(deps.get_api_client), db: Session = Depends(deps.get_db) +): + """ + Get all messages belonging to the same message tree. + """ + pr = PromptRepository(db, api_client, user=None) + message = pr.fetch_message_by_frontend_message_id(message_id) + messages = pr.fetch_message_children(message.id) + return [ + protocol.Message( + id=m.id, parent_id=m.parent_id, text=m.payload.payload.text, is_assistant=(m.role == "assistant") + ) + for m in messages + ] + + +@router.get("/{message_id}/descendants") +def get_descendants_by_frontend_id( + message_id: str, api_client: ApiClient = Depends(deps.get_api_client), db: Session = Depends(deps.get_db) +): + """ + Get a subtree which starts with this message. + The message is identified by its frontend ID. + """ + pr = PromptRepository(db, api_client, user=None) + message = pr.fetch_message_by_frontend_message_id(message_id) + descendants = pr.fetch_message_descendants(message) + return utils.prepare_tree(descendants, message.id) + + +@router.get("/{message_id}/longest_conversation_in_tree") +def get_longest_conv_by_frontend_id( + message_id: str, api_client: ApiClient = Depends(deps.get_api_client), db: Session = Depends(deps.get_db) +): + """ + Get the longest conversation from the tree of the message. + The message is identified by its frontend ID. + """ + pr = PromptRepository(db, api_client, user=None) + message = pr.fetch_message_by_frontend_message_id(message_id) + conv = pr.fetch_longest_conversation(message.message_tree_id) + return utils.prepare_conversation(conv) + + +@router.get("/{message_id}/max_children_in_tree") +def get_max_children_by_frontend_id( + message_id: str, api_client: ApiClient = Depends(deps.get_api_client), db: Session = Depends(deps.get_db) +): + """ + Get message with the most children from the tree of the provided message. + The message is identified by its frontend ID. + """ + pr = PromptRepository(db, api_client, user=None) + message = pr.fetch_message_by_frontend_message_id(message_id) + message, children = pr.fetch_message_with_max_children(message.message_tree_id) + return utils.prepare_tree([message, *children], message.id) diff --git a/backend/oasst_backend/api/v1/frontend_users.py b/backend/oasst_backend/api/v1/frontend_users.py new file mode 100644 index 00000000..35048331 --- /dev/null +++ b/backend/oasst_backend/api/v1/frontend_users.py @@ -0,0 +1,60 @@ +# -*- coding: utf-8 -*- +import datetime +from uuid import UUID + +from fastapi import APIRouter, Depends, Query +from oasst_backend.api import deps +from oasst_backend.models import ApiClient +from oasst_backend.prompt_repository import PromptRepository +from oasst_shared.schemas import protocol +from sqlmodel import Session +from starlette.responses import Response +from starlette.status import HTTP_200_OK + +router = APIRouter() + + +@router.get("/{username}/messages") +def query_frontend_user_messages( + username: str, + api_client_id: UUID = None, + max_count: int = Query(10, gt=0, le=1000), + start_date: datetime.datetime = None, + end_date: datetime.datetime = None, + only_roots: bool = False, + desc: bool = True, + include_deleted: bool = False, + api_client: ApiClient = Depends(deps.get_api_client), + db: Session = Depends(deps.get_db), +): + """ + Query frontend user messages. + """ + pr = PromptRepository(db, api_client, user=None) + messages = pr.query_messages( + username=username, + api_client_id=api_client_id, + desc=desc, + limit=max_count, + start_date=start_date, + end_date=end_date, + only_roots=only_roots, + deleted=None if include_deleted else False, + ) + + return [ + protocol.Message( + id=m.id, parent_id=m.parent_id, text=m.payload.payload.text, is_assistant=(m.role == "assistant") + ) + for m in messages + ] + + +@router.delete("/{username}/messages") +def mark_frontend_user_messages_deleted( + username: str, api_client: ApiClient = Depends(deps.get_trusted_api_client), db: Session = Depends(deps.get_db) +): + pr = PromptRepository(db, api_client, None) + messages = pr.query_messages(username=username, api_client_id=api_client.id) + pr.mark_messages_deleted(messages) + return Response(status_code=HTTP_200_OK) diff --git a/backend/oasst_backend/api/v1/messages.py b/backend/oasst_backend/api/v1/messages.py new file mode 100644 index 00000000..9221afc3 --- /dev/null +++ b/backend/oasst_backend/api/v1/messages.py @@ -0,0 +1,160 @@ +# -*- coding: utf-8 -*- + +import datetime +from uuid import UUID + +from fastapi import APIRouter, Depends, Query, Response +from oasst_backend.api import deps +from oasst_backend.api.v1 import utils +from oasst_backend.exceptions import OasstError, OasstErrorCode +from oasst_backend.models import ApiClient +from oasst_backend.models.db_payload import MessagePayload +from oasst_backend.prompt_repository import PromptRepository +from oasst_shared.schemas import protocol +from sqlmodel import Session +from starlette.status import HTTP_200_OK + +router = APIRouter() + + +@router.get("/") +def query_messages( + username: str = None, + api_client_id: str = None, + max_count: int = Query(10, gt=0, le=1000), + start_date: datetime.datetime = None, + end_date: datetime.datetime = None, + only_roots: bool = False, + desc: bool = True, + allow_deleted: bool = False, + api_client: ApiClient = Depends(deps.get_api_client), + db: Session = Depends(deps.get_db), +): + """ + Query messages. + """ + pr = PromptRepository(db, api_client, user=None) + messages = pr.query_messages( + username=username, + api_client_id=api_client_id, + desc=desc, + limit=max_count, + start_date=start_date, + end_date=end_date, + only_roots=only_roots, + deleted=None if allow_deleted else False, + ) + + return [ + protocol.Message( + id=m.id, parent_id=m.parent_id, text=m.payload.payload.text, is_assistant=(m.role == "assistant") + ) + for m in messages + ] + + +@router.get("/{message_id}") +def get_message( + message_id: UUID, api_client: ApiClient = Depends(deps.get_api_client), db: Session = Depends(deps.get_db) +): + """ + Get a message by its internal ID. + """ + pr = PromptRepository(db, api_client, user=None) + message = pr.fetch_message(message_id) + if not isinstance(message.payload.payload, MessagePayload): + # Unexptcted message payload + raise OasstError("Invalid message", OasstErrorCode.INVALID_MESSAGE) + + return protocol.ConversationMessage(text=message.payload.payload.text, is_assistant=(message.role == "assistant")) + + +@router.get("/{message_id}/conversation") +def get_conv( + message_id: UUID, api_client: ApiClient = Depends(deps.get_api_client), db: Session = Depends(deps.get_db) +): + """ + Get a conversation from the tree root and up to the message with given internal ID. + """ + + pr = PromptRepository(db, api_client, user=None) + messages = pr.fetch_message_conversation(message_id) + return utils.prepare_conversation(messages) + + +@router.get("/{message_id}/tree") +def get_tree( + message_id: UUID, api_client: ApiClient = Depends(deps.get_api_client), db: Session = Depends(deps.get_db) +): + """ + Get all messages belonging to the same message tree. + """ + pr = PromptRepository(db, api_client, user=None) + message = pr.fetch_message(message_id) + tree = pr.fetch_message_tree(message.message_tree_id) + return utils.prepare_tree(tree, message.message_tree_id) + + +@router.get("/{message_id}/children") +def get_children( + message_id: UUID, api_client: ApiClient = Depends(deps.get_api_client), db: Session = Depends(deps.get_db) +): + """ + Get all messages belonging to the same message tree. + """ + pr = PromptRepository(db, api_client, user=None) + messages = pr.fetch_message_children(message_id) + return [ + protocol.Message( + id=m.id, parent_id=m.parent_id, text=m.payload.payload.text, is_assistant=(m.role == "assistant") + ) + for m in messages + ] + + +@router.get("/{message_id}/descendants") +def get_descendants( + message_id: UUID, api_client: ApiClient = Depends(deps.get_api_client), db: Session = Depends(deps.get_db) +): + """ + Get a subtree which starts with this message. + """ + pr = PromptRepository(db, api_client, user=None) + message = pr.fetch_message(message_id) + descendants = pr.fetch_message_descendants(message) + return utils.prepare_tree(descendants, message.id) + + +@router.get("/{message_id}/longest_conversation_in_tree") +def get_longest_conv( + message_id: UUID, api_client: ApiClient = Depends(deps.get_api_client), db: Session = Depends(deps.get_db) +): + """ + Get the longest conversation from the tree of the message. + """ + pr = PromptRepository(db, api_client, user=None) + message = pr.fetch_message(message_id) + conv = pr.fetch_longest_conversation(message.message_tree_id) + return utils.prepare_conversation(conv) + + +@router.get("/{message_id}/max_children_in_tree") +def get_max_children( + message_id: UUID, api_client: ApiClient = Depends(deps.get_api_client), db: Session = Depends(deps.get_db) +): + """ + Get message with the most children from the tree of the provided message. + """ + pr = PromptRepository(db, api_client, user=None) + message = pr.fetch_message(message_id) + message, children = pr.fetch_message_with_max_children(message.message_tree_id) + return utils.prepare_tree([message, *children], message.id) + + +@router.delete("/{message_id}") +def mark_message_deleted( + message_id: UUID, api_client: ApiClient = Depends(deps.get_trusted_api_client), db: Session = Depends(deps.get_db) +): + pr = PromptRepository(db, api_client, None) + pr.mark_messages_deleted(message_id) + return Response(status_code=HTTP_200_OK) diff --git a/backend/oasst_backend/api/v1/stats.py b/backend/oasst_backend/api/v1/stats.py new file mode 100644 index 00000000..831d4df2 --- /dev/null +++ b/backend/oasst_backend/api/v1/stats.py @@ -0,0 +1,17 @@ +# -*- coding: utf-8 -*- +from fastapi import APIRouter, Depends +from oasst_backend.api import deps +from oasst_backend.models import ApiClient +from oasst_backend.prompt_repository import PromptRepository +from sqlmodel import Session + +router = APIRouter() + + +@router.get("/") +def get_message_stats( + db: Session = Depends(deps.get_db), + api_client: ApiClient = Depends(deps.get_trusted_api_client), +): + pr = PromptRepository(db, api_client, None) + return pr.get_stats() diff --git a/backend/oasst_backend/api/v1/users.py b/backend/oasst_backend/api/v1/users.py new file mode 100644 index 00000000..0bac4d6a --- /dev/null +++ b/backend/oasst_backend/api/v1/users.py @@ -0,0 +1,60 @@ +# -*- coding: utf-8 -*- +import datetime +from uuid import UUID + +from fastapi import APIRouter, Depends, Query +from oasst_backend.api import deps +from oasst_backend.models import ApiClient +from oasst_backend.prompt_repository import PromptRepository +from oasst_shared.schemas import protocol +from sqlmodel import Session +from starlette.responses import Response +from starlette.status import HTTP_200_OK + +router = APIRouter() + + +@router.get("/{user_id}/messages") +def query_user_messages( + user_id: UUID, + api_client_id: UUID = None, + max_count: int = Query(10, gt=0, le=1000), + start_date: datetime.datetime = None, + end_date: datetime.datetime = None, + only_roots: bool = False, + desc: bool = True, + include_deleted: bool = False, + api_client: ApiClient = Depends(deps.get_api_client), + db: Session = Depends(deps.get_db), +): + """ + Query user messages. + """ + pr = PromptRepository(db, api_client, user=None) + messages = pr.query_messages( + user_id=user_id, + api_client_id=api_client_id, + desc=desc, + limit=max_count, + start_date=start_date, + end_date=end_date, + only_roots=only_roots, + deleted=None if include_deleted else False, + ) + + return [ + protocol.Message( + id=m.id, parent_id=m.parent_id, text=m.payload.payload.text, is_assistant=(m.role == "assistant") + ) + for m in messages + ] + + +@router.delete("/{user_id}/messages") +def mark_user_messages_deleted( + user_id: UUID, api_client: ApiClient = Depends(deps.get_trusted_api_client), db: Session = Depends(deps.get_db) +): + pr = PromptRepository(db, api_client, None) + messages = pr.query_messages(user_id=user_id) + pr.mark_messages_deleted(messages) + return Response(status_code=HTTP_200_OK) diff --git a/backend/oasst_backend/api/v1/utils.py b/backend/oasst_backend/api/v1/utils.py new file mode 100644 index 00000000..2dac7947 --- /dev/null +++ b/backend/oasst_backend/api/v1/utils.py @@ -0,0 +1,38 @@ +# -*- coding: utf-8 -*- + +from http import HTTPStatus +from uuid import UUID + +from oasst_backend.exceptions import OasstError, OasstErrorCode +from oasst_backend.models import Message +from oasst_backend.models.db_payload import MessagePayload +from oasst_shared.schemas import protocol + + +def prepare_conversation(messages: list[Message]) -> protocol.Conversation: + conv_messages = [] + for message in messages: + if not isinstance(message.payload.payload, MessagePayload): + raise OasstError("Server error", OasstErrorCode.SERVER_ERROR, HTTPStatus.INTERNAL_SERVER_ERROR) + conv_messages.append( + protocol.ConversationMessage(text=message.payload.payload.text, is_assistant=(message.role == "assistant")) + ) + + return protocol.Conversation(messages=conv_messages) + + +def prepare_tree(tree: list[Message], tree_id: UUID) -> protocol.MessageTree: + tree_messages = [] + for message in tree: + if not isinstance(message.payload.payload, MessagePayload): + raise OasstError("Server error", OasstErrorCode.SERVER_ERROR, HTTPStatus.INTERNAL_SERVER_ERROR) + tree_messages.append( + protocol.Message( + id=message.id, + parent_id=message.parent_id, + text=message.payload.payload.text, + is_assistant=(message.role == "assistant"), + ) + ) + + return protocol.MessageTree(id=tree_id, messages=tree_messages) diff --git a/backend/oasst_backend/exceptions.py b/backend/oasst_backend/exceptions.py index ac7366cd..f431b05b 100644 --- a/backend/oasst_backend/exceptions.py +++ b/backend/oasst_backend/exceptions.py @@ -17,6 +17,7 @@ class OasstErrorCode(IntEnum): GENERIC_ERROR = 0 DATABASE_URI_NOT_SET = 1 API_CLIENT_NOT_AUTHORIZED = 2 + SERVER_ERROR = 3 # 1000-2000: tasks endpoint TASK_INVALID_REQUEST_TYPE = 1000 @@ -35,6 +36,7 @@ class OasstErrorCode(IntEnum): USER_NOT_SPECIFIED = 2005 NO_MESSAGE_TREE_FOUND = 2006 NO_REPLIES_FOUND = 2007 + INVALID_MESSAGE = 2008 TASK_NOT_FOUND = 2100 TASK_EXPIRED = 2101 TASK_PAYLOAD_TYPE_MISMATCH = 2102 diff --git a/backend/oasst_backend/journal_writer.py b/backend/oasst_backend/journal_writer.py index 64327366..415d5a47 100644 --- a/backend/oasst_backend/journal_writer.py +++ b/backend/oasst_backend/journal_writer.py @@ -54,7 +54,7 @@ class JournalWriter: self.user = user self.user_id = self.user.id if self.user else None - def log_text_reply(self, task: Task, message_id: UUID, role: str, length: int) -> Journal: + def log_text_reply(self, task: Task, message_id: Optional[UUID], role: str, length: int) -> Journal: return self.log( task_type=task.payload_type, event_type=JournalEventType.text_reply_to_message, @@ -63,7 +63,7 @@ class JournalWriter: message_id=message_id, ) - def log_rating(self, task: Task, message_id: UUID, rating: int) -> Journal: + def log_rating(self, task: Task, message_id: Optional[UUID], rating: int) -> Journal: return self.log( task_type=task.payload_type, event_type=JournalEventType.message_rating, @@ -72,7 +72,7 @@ class JournalWriter: message_id=message_id, ) - def log_ranking(self, task: Task, message_id: UUID, ranking: list[int]) -> Journal: + def log_ranking(self, task: Task, message_id: Optional[UUID], ranking: list[int]) -> Journal: return self.log( task_type=task.payload_type, event_type=JournalEventType.message_ranking, diff --git a/backend/oasst_backend/models/message.py b/backend/oasst_backend/models/message.py index 1425ce98..47512cc7 100644 --- a/backend/oasst_backend/models/message.py +++ b/backend/oasst_backend/models/message.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, Index, SQLModel from .payload_column_type import PayloadContainer, payload_column_type @@ -34,3 +35,4 @@ class Message(SQLModel, table=True): lang: str = Field(nullable=False, max_length=200, default="en-US") 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)) + deleted: bool = Field(sa_column=sa.Column(sa.Boolean, nullable=False, server_default=false())) diff --git a/backend/oasst_backend/prompt_repository.py b/backend/oasst_backend/prompt_repository.py index 15ed3816..b95f9d53 100644 --- a/backend/oasst_backend/prompt_repository.py +++ b/backend/oasst_backend/prompt_repository.py @@ -1,5 +1,8 @@ # -*- coding: utf-8 -*- +import datetime import random +from collections import defaultdict +from http import HTTPStatus from typing import Optional from uuid import UUID, uuid4 @@ -10,7 +13,10 @@ from oasst_backend.journal_writer import JournalWriter from oasst_backend.models import ApiClient, Message, MessageReaction, Task, TextLabels, User from oasst_backend.models.payload_column_type import PayloadContainer from oasst_shared.schemas import protocol as protocol_schema +from oasst_shared.schemas.protocol import SystemStats +from sqlalchemy import update from sqlmodel import Session, func +from starlette.status import HTTP_403_FORBIDDEN class PromptRepository: @@ -21,7 +27,7 @@ class PromptRepository: self.user_id = self.user.id if self.user else None self.journal = JournalWriter(db, api_client, self.user) - def lookup_user(self, client_user: protocol_schema.User) -> User: + def lookup_user(self, client_user: protocol_schema.User) -> Optional[User]: if not client_user: return None user: User = ( @@ -52,6 +58,7 @@ class PromptRepository: return user def validate_frontend_message_id(self, message_id: str) -> None: + # TODO: Should it be replaced with fastapi/pydantic validation? if not isinstance(message_id, str): raise OasstError( f"message_id must be string, not {type(message_id)}", OasstErrorCode.INVALID_FRONTEND_MESSAGE_ID @@ -114,9 +121,7 @@ class PromptRepository: ) return task - def store_text_reply( - self, text: str, frontend_message_id: str, user_frontend_message_id: str, role: str = None - ) -> Message: + def store_text_reply(self, text: str, frontend_message_id: str, user_frontend_message_id: str) -> Message: self.validate_frontend_message_id(frontend_message_id) self.validate_frontend_message_id(user_frontend_message_id) @@ -447,6 +452,13 @@ class PromptRepository: return self.db.query(Message).filter(Message.message_tree_id == message_tree_id).all() def fetch_multiple_random_replies(self, max_size: int = 5, message_role: str = None): + """ + Fetch a conversation with multiple possible replies to it. + + This function finds a random message with >1 replies, + forms a conversation from the corresponding message tree root up to this message + and fetches up to max_size possible replies in continuation to this conversation. + """ parent = self.db.query(Message.id).filter(Message.children_count > 1) if message_role: parent = parent.filter(Message.role == message_role) @@ -473,10 +485,16 @@ class PromptRepository: return conversation, replies - def fetch_message(self, message_id: UUID) -> Optional[Message]: - return self.db.query(Message).filter(Message.id == message_id).one() + def fetch_message(self, message_id: UUID, fail_if_missing: bool = True) -> Optional[Message]: + message = self.db.query(Message).filter(Message.id == message_id).one_or_none() + if fail_if_missing and not message: + raise OasstError("Message not found", OasstErrorCode.MESSAGE_NOT_FOUND) + return message def close_task(self, frontend_message_id: str, allow_personal_tasks: bool = False): + """ + Mark task as done. No further messages will be accepted for this task. + """ self.validate_frontend_message_id(frontend_message_id) task = self.fetch_task_by_frontend_message_id(frontend_message_id) @@ -492,3 +510,195 @@ class PromptRepository: task.done = True self.db.add(task) self.db.commit() + + @staticmethod + def trace_conversation(messages: list[Message] | dict[UUID, Message], last_message: Message) -> list[Message]: + """ + Pick messages from a collection so that the result makes a linear conversation + starting from a message tree root and up to the given message. + Returns an ordered list of messages starting from the message tree root. + """ + if isinstance(messages, list): + messages = {m.id: m for m in messages} + if not isinstance(messages, dict): + # This should not normally happen + raise OasstError("Server error", OasstErrorCode.SERVER_ERROR, HTTPStatus.INTERNAL_SERVER_ERROR) + + conv = [last_message] + while conv[-1].parent_id: + if conv[-1].parent_id not in messages: + # Can't form a continuous conversation + raise OasstError("Server error", OasstErrorCode.SERVER_ERROR, HTTPStatus.INTERNAL_SERVER_ERROR) + + parent_message = messages[conv[-1].parent_id] + conv.append(parent_message) + + return list(reversed(conv)) + + def fetch_message_conversation(self, message: Message | UUID) -> list[Message]: + """ + Fetch a conversation from the tree root and up to this message. + """ + if isinstance(message, UUID): + message = self.fetch_message(message) + + tree_messages = self.fetch_message_tree(message.message_tree_id) + return self.trace_conversation(tree_messages, message) + + def fetch_tree_from_message(self, message: Message | UUID) -> list[Message]: + """ + Fetch message tree this message belongs to. + """ + if isinstance(message, UUID): + message = self.fetch_message(message) + return self.fetch_message_tree(message.message_tree_id) + + def fetch_message_children(self, message: Message | UUID) -> list[Message]: + """ + Get all direct children of this message + """ + if isinstance(message, Message): + message = message.id + + children = self.db.query(Message).filter(Message.parent_id == message).all() + return children + + @staticmethod + def trace_descendants(root: Message, messages: list[Message]) -> list[Message]: + children = defaultdict(list) + for msg in messages: + children[msg.parent_id].append(msg) + + def _traverse_subtree(m: Message): + for child in children[m.id]: + yield child + yield from _traverse_subtree(child) + + return list(_traverse_subtree(root)) + + def fetch_message_descendants(self, message: Message | UUID, max_depth: int = None) -> list[Message]: + """ + Find all descendant messages to this message. + + This function creates a subtree of messages starting from given root message. + """ + if isinstance(message, UUID): + message = self.fetch_message(message) + + desc = self.db.query(Message).filter( + Message.message_tree_id == message.message_tree_id, Message.depth > message.depth + ) + if max_depth is not None: + desc = desc.filter(Message.depth <= max_depth) + + desc = desc.all() + + return self.trace_descendants(message, desc) + + def fetch_longest_conversation(self, message: Message | UUID) -> list[Message]: + tree = self.fetch_tree_from_message(message) + max_message = max(tree, key=lambda m: m.depth) + return self.trace_conversation(tree, max_message) + + def fetch_message_with_max_children(self, message: Message | UUID) -> tuple[Message, list[Message]]: + tree = self.fetch_tree_from_message(message) + max_message = max(tree, key=lambda m: m.children_count) + return max_message, [m for m in tree if m.parent_id == max_message.id] + + def query_messages( + self, + user_id: Optional[UUID] = None, + username: Optional[str] = None, + api_client_id: Optional[UUID] = None, + desc: bool = True, + limit: Optional[int] = 10, + start_date: Optional[datetime.datetime] = None, + end_date: Optional[datetime.datetime] = None, + only_roots: bool = False, + deleted: Optional[bool] = None, + ) -> list[Message]: + if not self.api_client.trusted and not api_client_id: + # Let unprivileged api clients query their own messages without api_client_id being set + api_client_id = self.api_client.id + + if not self.api_client.trusted and api_client_id != self.api_client.id: + # Unprivileged api client asks for foreign messages + raise OasstError("Forbidden", OasstErrorCode.API_CLIENT_NOT_AUTHORIZED, HTTP_403_FORBIDDEN) + + messages = self.db.query(Message) + if user_id: + messages = messages.filter(Message.user_id == user_id) + if username: + messages = messages.join(User) + messages = messages.filter(User.username == username) + if api_client_id: + messages = messages.filter(Message.api_client_id == api_client_id) + + if start_date: + messages = messages.filter(Message.created_date >= start_date) + if end_date: + messages = messages.filter(Message.created_date < end_date) + + if only_roots: + messages = messages.filter(Message.parent_id.is_(None)) + + if deleted is not None: + messages = messages.filter(Message.deleted == deleted) + + if desc: + messages = messages.order_by(Message.created_date.desc()) + else: + messages = messages.order_by(Message.created_date.asc()) + + if limit is not None: + messages = messages.limit(limit) + + # TODO: Pagination could be great at some point + return messages.all() + + def mark_messages_deleted(self, messages: Message | UUID | list[Message | UUID], recursive: bool = True): + """ + Marks deleted messages and all their descendants. + """ + if isinstance(messages, (Message, UUID)): + messages = [messages] + + ids = [] + for message in messages: + if isinstance(message, UUID): + ids.append(message) + elif isinstance(message, Message): + ids.append(message.id) + else: + raise OasstError("Server error", OasstErrorCode.SERVER_ERROR, HTTPStatus.INTERNAL_SERVER_ERROR) + + query = update(Message).where(Message.id.in_(ids)).values(deleted=True) + self.db.execute(query) + + parent_ids = ids + if recursive: + while parent_ids: + query = ( + update(Message).filter(Message.parent_id.in_(parent_ids)).values(deleted=True).returning(Message.id) + ) + + parent_ids = self.db.execute(query).scalars().all() + + self.db.commit() + + def get_stats(self) -> SystemStats: + """ + Get data stats such as number of all messages in the system, + number of deleted and active messages and number of message trees. + """ + deleted = self.db.query(Message.deleted, func.count()).group_by(Message.deleted) + nthreads = self.db.query(None, func.count(Message.id)).filter(Message.parent_id.is_(None)) + query = deleted.union_all(nthreads) + result = {k: v for k, v in query.all()} + + return SystemStats( + all=result.get(True, 0) + result.get(False, 0), + active=result.get(False, 0), + deleted=result.get(True, 0), + message_trees=result.get(None, 0), + ) diff --git a/oasst-shared/oasst_shared/schemas/protocol.py b/oasst-shared/oasst_shared/schemas/protocol.py index 8fe8bdea..ed7dc780 100644 --- a/oasst-shared/oasst_shared/schemas/protocol.py +++ b/oasst-shared/oasst_shared/schemas/protocol.py @@ -38,6 +38,18 @@ class Conversation(BaseModel): messages: list[ConversationMessage] = [] +class Message(ConversationMessage): + id: UUID + parent_id: Optional[UUID] = None + + +class MessageTree(BaseModel): + """All messages belonging to the same message tree.""" + + id: UUID + messages: list[Message] = [] + + class TaskRequest(BaseModel): """The frontend asks the backend for a task.""" @@ -259,3 +271,10 @@ class TextLabels(BaseModel): if not (0 <= value <= 1): raise ValueError(f"Label values must be between 0 and 1, got {value} for {key}.") return v + + +class SystemStats(BaseModel): + all: int = 0 + active: int = 0 + deleted: int = 0 + message_trees: int = 0