From 13d01b5a2fcdfcdc2e7777e232593d7e4285c048 Mon Sep 17 00:00:00 2001 From: Igor Miagkov Date: Fri, 30 Dec 2022 06:03:39 +0400 Subject: [PATCH 1/6] management api --- ...-6cb49da61b74_add_deleted_field_to_post.py | 28 ++ backend/oasst_backend/api/deps.py | 23 +- backend/oasst_backend/api/v1/api.py | 3 +- backend/oasst_backend/api/v1/management.py | 291 ++++++++++++++++++ backend/oasst_backend/exceptions.py | 1 + backend/oasst_backend/models/message.py | 2 + backend/oasst_backend/prompt_repository.py | 149 +++++++++ oasst-shared/oasst_shared/schemas/protocol.py | 12 + 8 files changed, 507 insertions(+), 2 deletions(-) create mode 100644 backend/alembic/versions/2022_12_30_0654-6cb49da61b74_add_deleted_field_to_post.py create mode 100644 backend/oasst_backend/api/v1/management.py diff --git a/backend/alembic/versions/2022_12_30_0654-6cb49da61b74_add_deleted_field_to_post.py b/backend/alembic/versions/2022_12_30_0654-6cb49da61b74_add_deleted_field_to_post.py new file mode 100644 index 00000000..a8f228bd --- /dev/null +++ b/backend/alembic/versions/2022_12_30_0654-6cb49da61b74_add_deleted_field_to_post.py @@ -0,0 +1,28 @@ +# -*- coding: utf-8 -*- +"""add deleted field to post + +Revision ID: 6cb49da61b74 +Revises: 73ce3675c1f5 +Create Date: 2022-12-30 06:54:47.110204 + +""" +import sqlalchemy as sa +from alembic import op + +# revision identifiers, used by Alembic. +revision = "6cb49da61b74" +down_revision = "73ce3675c1f5" +branch_labels = None +depends_on = None + + +def upgrade() -> None: + # ### commands auto generated by Alembic - please adjust! ### + op.add_column("post", 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("post", "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..7c9eb493 100644 --- a/backend/oasst_backend/api/v1/api.py +++ b/backend/oasst_backend/api/v1/api.py @@ -1,7 +1,8 @@ # -*- coding: utf-8 -*- from fastapi import APIRouter -from oasst_backend.api.v1 import tasks, text_labels +from oasst_backend.api.v1 import management, tasks, text_labels 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(management.router, prefix="/management", tags=["management"]) diff --git a/backend/oasst_backend/api/v1/management.py b/backend/oasst_backend/api/v1/management.py new file mode 100644 index 00000000..0e3be690 --- /dev/null +++ b/backend/oasst_backend/api/v1/management.py @@ -0,0 +1,291 @@ +# -*- coding: utf-8 -*- +import datetime +from http import HTTPStatus +from uuid import UUID + +from fastapi import APIRouter, Depends, Query, Response +from oasst_backend.api import deps +from oasst_backend.exceptions import OasstError, OasstErrorCode +from oasst_backend.models import ApiClient, Post +from oasst_backend.models.db_payload import PostPayload +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() + + +def _prepare_conversation(messages: list[Post]) -> protocol.Conversation: + conv_messages = [] + for message in messages: + if not isinstance(message.payload.payload, PostPayload): + 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[Post], tree_id: UUID) -> protocol.MessageTree: + tree_messages = [] + for message in tree: + if not isinstance(message.payload.payload, PostPayload): + 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) + + +@router.get("/message") +def query_messages( + username: str = None, + api_client_id: str = None, + max_count: int = Query(10, gt=0, le=25), + start_date: datetime.datetime = None, + end_date: datetime.datetime = None, + only_roots: bool = False, + desc: bool = True, + api_client: ApiClient = Depends(deps.get_api_client), + db: Session = Depends(deps.get_db), +): + """ + Query messages. + """ + if not api_client.trusted and (api_client_id != api_client.id): + # Unprivileged api client asks for foreign messages + return [] + + pr = PromptRepository(db, api_client, user=None) + messages = pr.query_messages( + username=username, + api_client_id=api_client_id, + desc=desc, + max_count=max_count, + start_date=start_date, + end_date=end_date, + only_roots=only_roots, + ) + + 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/{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) + post = pr.fetch_post(message_id) + if not isinstance(post.payload.payload, PostPayload): + raise OasstError("Invalid message id", OasstErrorCode.INVALID_POST_ID) + + return protocol.ConversationMessage(text=post.payload.payload.text, is_assistant=(post.role == "assistant")) + + +@router.get("/frontend_message/{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_post_by_frontend_post_id(message_id, fail_if_missing=True) + + if not isinstance(message.payload.payload, PostPayload): + raise OasstError("Invalid message id", OasstErrorCode.INVALID_POST_ID) + + return protocol.ConversationMessage(text=message.payload.payload.text, is_assistant=(message.role == "assistant")) + + +@router.get("/message/{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 _prepare_conversation(messages) + + +@router.get("/frontend_message/{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_post_by_frontend_post_id(message_id) + messages = pr.fetch_message_conversation(message) + return _prepare_conversation(messages) + + +@router.get("/message/{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_post(message_id) + tree = pr.fetch_message_tree(message) + return _prepare_tree(tree, message.thread_id) + + +@router.get("/frontend_message/{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_post_by_frontend_post_id(message_id) + tree = pr.fetch_message_tree(message) + return _prepare_tree(tree, message.thread_id) + + +@router.get("/message/{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) + return pr.fetch_message_children(message_id) + + +@router.get("/frontend_message/{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_post_by_frontend_post_id(message_id) + return pr.fetch_message_children(message) + + +@router.get("/message/{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_post(message_id) + descendants = pr.fetch_post_descendants(message) + return _prepare_tree(descendants, message.id) + + +@router.get("/frontend_message/{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_post_by_frontend_post_id(message_id) + descendants = pr.fetch_post_descendants(message) + return _prepare_tree(descendants, message.id) + + +@router.get("/message/{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_post(message_id) + conv = pr.fetch_longest_conversation(message.thread_id) + return _prepare_conversation(conv) + + +@router.get("/frontend_message/{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_post_by_frontend_post_id(message_id) + conv = pr.fetch_longest_conversation(message.thread_id) + return _prepare_conversation(conv) + + +@router.get("/message/{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_post(message_id) + message, children = pr.fetch_message_with_max_children(message.thread_id) + return _prepare_tree([message, *children], message.id) + + +@router.get("/frontend_message/{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_post_by_frontend_post_id(message_id) + message, children = pr.fetch_message_with_max_children(message.thread_id) + return _prepare_tree([message, *children], message.id) + + +@router.delete("/message/{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) + + +@router.delete("/user/{username}/message") +def mark_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/exceptions.py b/backend/oasst_backend/exceptions.py index ac7366cd..7f88caed 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 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..6db35019 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,6 +13,7 @@ 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 sqlalchemy import update from sqlmodel import Session, func @@ -492,3 +496,148 @@ class PromptRepository: task.done = True self.db.add(task) self.db.commit() + + @staticmethod + def trace_conversation(messages: list[Post] | dict[UUID, Post], last_message: Post) -> list[Post]: + """ + 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: Post | UUID) -> list[Post]: + """ + Fetch a conversation from the tree root and up to this message. + """ + if isinstance(message, UUID): + message = self.fetch_post(message) + + tree_messages = self.fetch_thread(message.thread_id) + return self.trace_conversation(tree_messages, message) + + def fetch_message_tree(self, message: Post | UUID) -> list[Post]: + """ + Fetch message tree this message belongs to. + """ + if isinstance(message, UUID): + message = self.fetch_post(message) + return self.fetch_thread(message.thread_id) + + def fetch_message_children(self, message: Post | UUID) -> list[Post]: + """ + Get all direct children of this message + """ + if isinstance(message, Post): + message = message.id + + children = self.db.query(Post).filter(Post.parent_id == message).all() + return children + + @staticmethod + def trace_descendants(root: Post, messages: list[Post]) -> list[Post]: + children = defaultdict(list) + for msg in messages: + children[msg.parent_id].append(msg) + + def _traverse_subtree(m: Post): + for child in children[m.id]: + yield child + yield from _traverse_subtree(child) + + return list(_traverse_subtree(root)) + + def fetch_post_descendants(self, message: Post | UUID, max_depth: int = None) -> list[Post]: + if isinstance(message, UUID): + message = self.fetch_post(message) + + desc = self.db.query(Post).filter(Post.thread_id == message.thread_id, Post.depth > message.depth) + if max_depth is not None: + desc = desc.filter(Post.depth <= max_depth) + + desc = desc.all() + + return self.trace_descendants(message, desc) + + def fetch_longest_conversation(self, message: Post | UUID) -> list[Post]: + tree = self.fetch_message_tree(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: Post | UUID) -> tuple[Post, list[Post]]: + tree = self.fetch_message_tree(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, + username: str = None, + api_client_id: str = None, + desc: bool = True, + max_count: int = 10, + start_date: datetime.datetime = None, + end_date: datetime.datetime = None, + only_roots: bool = False, + ) -> list[Post]: + messages = self.db.query(Post) + if username: + messages = messages.join(Person) + messages = messages.filter(Person.username == username) + if api_client_id: + messages = messages.filter(Post.api_client_id == api_client_id) + + if start_date: + messages = messages.filter(Post.created_date >= start_date) + if end_date: + messages = messages.filter(Post.created_date < end_date) + + if only_roots: + messages = messages.filter(Post.parent_id.is_(None)) + + if desc: + messages = messages.order_by(Post.created_date.desc()) + else: + messages = messages.order_by(Post.created_date.asc()) + + messages = messages.limit(max_count).all() + return messages + + def mark_messages_deleted(self, messages: Post | UUID | list[Post | UUID], recursive: bool = True): + if isinstance(messages, (Post, UUID)): + messages = [messages] + + ids = [] + for message in messages: + if isinstance(message, UUID): + ids.append(message) + elif isinstance(message, Post): + ids.append(message.id) + else: + raise OasstError("Server error", OasstErrorCode.SERVER_ERROR, HTTPStatus.INTERNAL_SERVER_ERROR) + + query = update(Post).where(Post.id.in_(ids)).values(deleted=True) + self.db.execute(query) + + parent_ids = ids + if recursive: + while parent_ids: + query = update(Post).filter(Post.parent_id.in_(parent_ids)).values(deleted=True).returning(Post.id) + + parent_ids = self.db.execute(query).scalars().all() + + self.db.commit() diff --git a/oasst-shared/oasst_shared/schemas/protocol.py b/oasst-shared/oasst_shared/schemas/protocol.py index 8fe8bdea..028154c9 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.""" From 8e1d80956acebbe994a5edd5e1871de6197ecbb6 Mon Sep 17 00:00:00 2001 From: Igor Miagkov Date: Fri, 30 Dec 2022 20:02:23 +0400 Subject: [PATCH 2/6] split message api endpoints --- backend/oasst_backend/api/v1/api.py | 7 +- .../oasst_backend/api/v1/frontend_messages.py | 117 +++++++ .../oasst_backend/api/v1/frontend_users.py | 60 ++++ backend/oasst_backend/api/v1/management.py | 291 ------------------ backend/oasst_backend/api/v1/messages.py | 159 ++++++++++ backend/oasst_backend/api/v1/users.py | 60 ++++ backend/oasst_backend/api/v1/utils.py | 38 +++ backend/oasst_backend/prompt_repository.py | 40 ++- 8 files changed, 470 insertions(+), 302 deletions(-) create mode 100644 backend/oasst_backend/api/v1/frontend_messages.py create mode 100644 backend/oasst_backend/api/v1/frontend_users.py delete mode 100644 backend/oasst_backend/api/v1/management.py create mode 100644 backend/oasst_backend/api/v1/messages.py create mode 100644 backend/oasst_backend/api/v1/users.py create mode 100644 backend/oasst_backend/api/v1/utils.py diff --git a/backend/oasst_backend/api/v1/api.py b/backend/oasst_backend/api/v1/api.py index 7c9eb493..b39d8d59 100644 --- a/backend/oasst_backend/api/v1/api.py +++ b/backend/oasst_backend/api/v1/api.py @@ -1,8 +1,11 @@ # -*- coding: utf-8 -*- from fastapi import APIRouter -from oasst_backend.api.v1 import management, tasks, text_labels +from oasst_backend.api.v1 import frontend_messages, frontend_users, messages, 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(management.router, prefix="/management", tags=["management"]) +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"]) 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..acdd388a --- /dev/null +++ b/backend/oasst_backend/api/v1/frontend_messages.py @@ -0,0 +1,117 @@ +# -*- 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 PostPayload +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_post_by_frontend_post_id(message_id, fail_if_missing=True) + + if not isinstance(message.payload.payload, PostPayload): + raise OasstError("Invalid message id", OasstErrorCode.INVALID_POST_ID) + + 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_post_by_frontend_post_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_post_by_frontend_post_id(message_id) + tree = pr.fetch_message_tree(message) + return utils.prepare_tree(tree, message.thread_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_post_by_frontend_post_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_post_by_frontend_post_id(message_id) + descendants = pr.fetch_post_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_post_by_frontend_post_id(message_id) + conv = pr.fetch_longest_conversation(message.thread_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_post_by_frontend_post_id(message_id) + message, children = pr.fetch_message_with_max_children(message.thread_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..f5df2b24 --- /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=25), + 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/management.py b/backend/oasst_backend/api/v1/management.py deleted file mode 100644 index 0e3be690..00000000 --- a/backend/oasst_backend/api/v1/management.py +++ /dev/null @@ -1,291 +0,0 @@ -# -*- coding: utf-8 -*- -import datetime -from http import HTTPStatus -from uuid import UUID - -from fastapi import APIRouter, Depends, Query, Response -from oasst_backend.api import deps -from oasst_backend.exceptions import OasstError, OasstErrorCode -from oasst_backend.models import ApiClient, Post -from oasst_backend.models.db_payload import PostPayload -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() - - -def _prepare_conversation(messages: list[Post]) -> protocol.Conversation: - conv_messages = [] - for message in messages: - if not isinstance(message.payload.payload, PostPayload): - 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[Post], tree_id: UUID) -> protocol.MessageTree: - tree_messages = [] - for message in tree: - if not isinstance(message.payload.payload, PostPayload): - 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) - - -@router.get("/message") -def query_messages( - username: str = None, - api_client_id: str = None, - max_count: int = Query(10, gt=0, le=25), - start_date: datetime.datetime = None, - end_date: datetime.datetime = None, - only_roots: bool = False, - desc: bool = True, - api_client: ApiClient = Depends(deps.get_api_client), - db: Session = Depends(deps.get_db), -): - """ - Query messages. - """ - if not api_client.trusted and (api_client_id != api_client.id): - # Unprivileged api client asks for foreign messages - return [] - - pr = PromptRepository(db, api_client, user=None) - messages = pr.query_messages( - username=username, - api_client_id=api_client_id, - desc=desc, - max_count=max_count, - start_date=start_date, - end_date=end_date, - only_roots=only_roots, - ) - - 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/{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) - post = pr.fetch_post(message_id) - if not isinstance(post.payload.payload, PostPayload): - raise OasstError("Invalid message id", OasstErrorCode.INVALID_POST_ID) - - return protocol.ConversationMessage(text=post.payload.payload.text, is_assistant=(post.role == "assistant")) - - -@router.get("/frontend_message/{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_post_by_frontend_post_id(message_id, fail_if_missing=True) - - if not isinstance(message.payload.payload, PostPayload): - raise OasstError("Invalid message id", OasstErrorCode.INVALID_POST_ID) - - return protocol.ConversationMessage(text=message.payload.payload.text, is_assistant=(message.role == "assistant")) - - -@router.get("/message/{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 _prepare_conversation(messages) - - -@router.get("/frontend_message/{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_post_by_frontend_post_id(message_id) - messages = pr.fetch_message_conversation(message) - return _prepare_conversation(messages) - - -@router.get("/message/{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_post(message_id) - tree = pr.fetch_message_tree(message) - return _prepare_tree(tree, message.thread_id) - - -@router.get("/frontend_message/{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_post_by_frontend_post_id(message_id) - tree = pr.fetch_message_tree(message) - return _prepare_tree(tree, message.thread_id) - - -@router.get("/message/{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) - return pr.fetch_message_children(message_id) - - -@router.get("/frontend_message/{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_post_by_frontend_post_id(message_id) - return pr.fetch_message_children(message) - - -@router.get("/message/{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_post(message_id) - descendants = pr.fetch_post_descendants(message) - return _prepare_tree(descendants, message.id) - - -@router.get("/frontend_message/{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_post_by_frontend_post_id(message_id) - descendants = pr.fetch_post_descendants(message) - return _prepare_tree(descendants, message.id) - - -@router.get("/message/{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_post(message_id) - conv = pr.fetch_longest_conversation(message.thread_id) - return _prepare_conversation(conv) - - -@router.get("/frontend_message/{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_post_by_frontend_post_id(message_id) - conv = pr.fetch_longest_conversation(message.thread_id) - return _prepare_conversation(conv) - - -@router.get("/message/{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_post(message_id) - message, children = pr.fetch_message_with_max_children(message.thread_id) - return _prepare_tree([message, *children], message.id) - - -@router.get("/frontend_message/{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_post_by_frontend_post_id(message_id) - message, children = pr.fetch_message_with_max_children(message.thread_id) - return _prepare_tree([message, *children], message.id) - - -@router.delete("/message/{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) - - -@router.delete("/user/{username}/message") -def mark_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..bd8e7f12 --- /dev/null +++ b/backend/oasst_backend/api/v1/messages.py @@ -0,0 +1,159 @@ +# -*- 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 PostPayload +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=25), + 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_post(message_id) + if not isinstance(message.payload.payload, PostPayload): + raise OasstError("Invalid message id", OasstErrorCode.INVALID_POST_ID) + + 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_post(message_id) + tree = pr.fetch_message_tree(message) + return utils.prepare_tree(tree, message.thread_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_post(message_id) + descendants = pr.fetch_post_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_post(message_id) + conv = pr.fetch_longest_conversation(message.thread_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_post(message_id) + message, children = pr.fetch_message_with_max_children(message.thread_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/users.py b/backend/oasst_backend/api/v1/users.py new file mode 100644 index 00000000..000b0970 --- /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=25), + 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..48f11038 --- /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 Post +from oasst_backend.models.db_payload import PostPayload +from oasst_shared.schemas import protocol + + +def prepare_conversation(messages: list[Post]) -> protocol.Conversation: + conv_messages = [] + for message in messages: + if not isinstance(message.payload.payload, PostPayload): + 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[Post], tree_id: UUID) -> protocol.MessageTree: + tree_messages = [] + for message in tree: + if not isinstance(message.payload.payload, PostPayload): + 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/prompt_repository.py b/backend/oasst_backend/prompt_repository.py index 6db35019..05902b5a 100644 --- a/backend/oasst_backend/prompt_repository.py +++ b/backend/oasst_backend/prompt_repository.py @@ -15,6 +15,7 @@ from oasst_backend.models.payload_column_type import PayloadContainer from oasst_shared.schemas import protocol as protocol_schema from sqlalchemy import update from sqlmodel import Session, func +from starlette.status import HTTP_403_FORBIDDEN class PromptRepository: @@ -477,8 +478,11 @@ 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.POST_NOT_FOUND) + return message def close_task(self, frontend_message_id: str, allow_personal_tasks: bool = False): self.validate_frontend_message_id(frontend_message_id) @@ -586,15 +590,27 @@ class PromptRepository: def query_messages( self, - username: str = None, - api_client_id: str = None, + user_id: Optional[UUID] = None, + username: Optional[str] = None, + api_client_id: Optional[UUID] = None, desc: bool = True, - max_count: int = 10, - start_date: datetime.datetime = None, - end_date: datetime.datetime = None, + 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[Post]: + 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(Post) + if user_id: + messages = messages.filter(Post.person_id == user_id) if username: messages = messages.join(Person) messages = messages.filter(Person.username == username) @@ -609,13 +625,19 @@ class PromptRepository: if only_roots: messages = messages.filter(Post.parent_id.is_(None)) + if deleted is not None: + messages = messages.filter(Post.deleted == deleted) + if desc: messages = messages.order_by(Post.created_date.desc()) else: messages = messages.order_by(Post.created_date.asc()) - messages = messages.limit(max_count).all() - return messages + 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: Post | UUID | list[Post | UUID], recursive: bool = True): if isinstance(messages, (Post, UUID)): From 475f48b1951c0f5231eba616c2e8d31aa961db9d Mon Sep 17 00:00:00 2001 From: Igor Miagkov Date: Sat, 31 Dec 2022 00:34:00 +0400 Subject: [PATCH 3/6] stats api method --- backend/oasst_backend/api/v1/api.py | 3 ++- backend/oasst_backend/api/v1/stats.py | 18 ++++++++++++++++++ backend/oasst_backend/prompt_repository.py | 12 ++++++++++++ oasst-shared/oasst_shared/schemas/protocol.py | 7 +++++++ 4 files changed, 39 insertions(+), 1 deletion(-) create mode 100644 backend/oasst_backend/api/v1/stats.py diff --git a/backend/oasst_backend/api/v1/api.py b/backend/oasst_backend/api/v1/api.py index b39d8d59..2286a1ac 100644 --- a/backend/oasst_backend/api/v1/api.py +++ b/backend/oasst_backend/api/v1/api.py @@ -1,6 +1,6 @@ # -*- coding: utf-8 -*- from fastapi import APIRouter -from oasst_backend.api.v1 import frontend_messages, frontend_users, messages, tasks, text_labels, users +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"]) @@ -9,3 +9,4 @@ 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/stats.py b/backend/oasst_backend/api/v1/stats.py new file mode 100644 index 00000000..789e53f9 --- /dev/null +++ b/backend/oasst_backend/api/v1/stats.py @@ -0,0 +1,18 @@ +# -*- 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 oasst_shared.schemas import protocol +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 protocol.SystemStats(**pr.get_stats()) diff --git a/backend/oasst_backend/prompt_repository.py b/backend/oasst_backend/prompt_repository.py index 05902b5a..8df17347 100644 --- a/backend/oasst_backend/prompt_repository.py +++ b/backend/oasst_backend/prompt_repository.py @@ -663,3 +663,15 @@ class PromptRepository: parent_ids = self.db.execute(query).scalars().all() self.db.commit() + + def get_stats(self): + deleted = self.db.query(Post.deleted, func.count()).group_by(Post.deleted) + nthreads = self.db.query(None, func.count(Post.id)).filter(Post.parent_id.is_(None)) + query = deleted.union_all(nthreads) + result = {k: v for k, v in query.all()} + return { + "all": result.get(True, 0) + result.get(False, 0), + "active": result.get(False, 0), + "deleted": result.get(True, 0), + "threads": result.get(None, 0), + } diff --git a/oasst-shared/oasst_shared/schemas/protocol.py b/oasst-shared/oasst_shared/schemas/protocol.py index 028154c9..a2a12ddc 100644 --- a/oasst-shared/oasst_shared/schemas/protocol.py +++ b/oasst-shared/oasst_shared/schemas/protocol.py @@ -271,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 + threads: int = 0 From f126b21bb3fbea6017fa51d9ea0fc03244e16b04 Mon Sep 17 00:00:00 2001 From: Igor Miagkov Date: Sat, 31 Dec 2022 04:35:33 +0400 Subject: [PATCH 4/6] adjust names and types to new naming --- .../oasst_backend/api/v1/frontend_messages.py | 30 +++---- backend/oasst_backend/api/v1/messages.py | 26 +++--- backend/oasst_backend/api/v1/utils.py | 12 +-- backend/oasst_backend/journal_writer.py | 6 +- backend/oasst_backend/prompt_repository.py | 90 ++++++++++--------- 5 files changed, 83 insertions(+), 81 deletions(-) diff --git a/backend/oasst_backend/api/v1/frontend_messages.py b/backend/oasst_backend/api/v1/frontend_messages.py index acdd388a..c53f6bae 100644 --- a/backend/oasst_backend/api/v1/frontend_messages.py +++ b/backend/oasst_backend/api/v1/frontend_messages.py @@ -5,7 +5,7 @@ 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 PostPayload +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 @@ -21,10 +21,10 @@ def get_message_by_frontend_id( Get a message by its frontend ID. """ pr = PromptRepository(db, api_client, user=None) - message = pr.fetch_post_by_frontend_post_id(message_id, fail_if_missing=True) + message = pr.fetch_message_by_frontend_message_id(message_id) - if not isinstance(message.payload.payload, PostPayload): - raise OasstError("Invalid message id", OasstErrorCode.INVALID_POST_ID) + if not isinstance(message.payload.payload, MessagePayload): + raise OasstError("Invalid message id", OasstErrorCode.INVALID_FRONTEND_MESSAGE_ID) return protocol.ConversationMessage(text=message.payload.payload.text, is_assistant=(message.role == "assistant")) @@ -38,7 +38,7 @@ def get_conv_by_frontend_id( """ pr = PromptRepository(db, api_client, user=None) - message = pr.fetch_post_by_frontend_post_id(message_id) + message = pr.fetch_message_by_frontend_message_id(message_id) messages = pr.fetch_message_conversation(message) return utils.prepare_conversation(messages) @@ -52,9 +52,9 @@ def get_tree_by_frontend_id( Message is identified by its frontend ID. """ pr = PromptRepository(db, api_client, user=None) - message = pr.fetch_post_by_frontend_post_id(message_id) - tree = pr.fetch_message_tree(message) - return utils.prepare_tree(tree, message.thread_id) + 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") @@ -65,7 +65,7 @@ def get_children_by_frontend_id( Get all messages belonging to the same message tree. """ pr = PromptRepository(db, api_client, user=None) - message = pr.fetch_post_by_frontend_post_id(message_id) + message = pr.fetch_message_by_frontend_message_id(message_id) messages = pr.fetch_message_children(message.id) return [ protocol.Message( @@ -84,8 +84,8 @@ def get_descendants_by_frontend_id( The message is identified by its frontend ID. """ pr = PromptRepository(db, api_client, user=None) - message = pr.fetch_post_by_frontend_post_id(message_id) - descendants = pr.fetch_post_descendants(message) + message = pr.fetch_message_by_frontend_message_id(message_id) + descendants = pr.fetch_message_descendants(message) return utils.prepare_tree(descendants, message.id) @@ -98,8 +98,8 @@ def get_longest_conv_by_frontend_id( The message is identified by its frontend ID. """ pr = PromptRepository(db, api_client, user=None) - message = pr.fetch_post_by_frontend_post_id(message_id) - conv = pr.fetch_longest_conversation(message.thread_id) + message = pr.fetch_message_by_frontend_message_id(message_id) + conv = pr.fetch_longest_conversation(message.message_tree_id) return utils.prepare_conversation(conv) @@ -112,6 +112,6 @@ def get_max_children_by_frontend_id( The message is identified by its frontend ID. """ pr = PromptRepository(db, api_client, user=None) - message = pr.fetch_post_by_frontend_post_id(message_id) - message, children = pr.fetch_message_with_max_children(message.thread_id) + 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/messages.py b/backend/oasst_backend/api/v1/messages.py index bd8e7f12..a812ef61 100644 --- a/backend/oasst_backend/api/v1/messages.py +++ b/backend/oasst_backend/api/v1/messages.py @@ -8,7 +8,7 @@ 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 PostPayload +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 @@ -61,9 +61,9 @@ def get_message( Get a message by its internal ID. """ pr = PromptRepository(db, api_client, user=None) - message = pr.fetch_post(message_id) - if not isinstance(message.payload.payload, PostPayload): - raise OasstError("Invalid message id", OasstErrorCode.INVALID_POST_ID) + message = pr.fetch_message(message_id) + if not isinstance(message.payload.payload, MessagePayload): + raise OasstError("Invalid message id", OasstErrorCode.INVALID_FRONTEND_MESSAGE_ID) return protocol.ConversationMessage(text=message.payload.payload.text, is_assistant=(message.role == "assistant")) @@ -89,9 +89,9 @@ def get_tree( Get all messages belonging to the same message tree. """ pr = PromptRepository(db, api_client, user=None) - message = pr.fetch_post(message_id) - tree = pr.fetch_message_tree(message) - return utils.prepare_tree(tree, message.thread_id) + 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") @@ -119,8 +119,8 @@ def get_descendants( Get a subtree which starts with this message. """ pr = PromptRepository(db, api_client, user=None) - message = pr.fetch_post(message_id) - descendants = pr.fetch_post_descendants(message) + message = pr.fetch_message(message_id) + descendants = pr.fetch_message_descendants(message) return utils.prepare_tree(descendants, message.id) @@ -132,8 +132,8 @@ def get_longest_conv( Get the longest conversation from the tree of the message. """ pr = PromptRepository(db, api_client, user=None) - message = pr.fetch_post(message_id) - conv = pr.fetch_longest_conversation(message.thread_id) + message = pr.fetch_message(message_id) + conv = pr.fetch_longest_conversation(message.message_tree_id) return utils.prepare_conversation(conv) @@ -145,8 +145,8 @@ def get_max_children( Get message with the most children from the tree of the provided message. """ pr = PromptRepository(db, api_client, user=None) - message = pr.fetch_post(message_id) - message, children = pr.fetch_message_with_max_children(message.thread_id) + 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) diff --git a/backend/oasst_backend/api/v1/utils.py b/backend/oasst_backend/api/v1/utils.py index 48f11038..2dac7947 100644 --- a/backend/oasst_backend/api/v1/utils.py +++ b/backend/oasst_backend/api/v1/utils.py @@ -4,15 +4,15 @@ from http import HTTPStatus from uuid import UUID from oasst_backend.exceptions import OasstError, OasstErrorCode -from oasst_backend.models import Post -from oasst_backend.models.db_payload import PostPayload +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[Post]) -> protocol.Conversation: +def prepare_conversation(messages: list[Message]) -> protocol.Conversation: conv_messages = [] for message in messages: - if not isinstance(message.payload.payload, PostPayload): + 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")) @@ -21,10 +21,10 @@ def prepare_conversation(messages: list[Post]) -> protocol.Conversation: return protocol.Conversation(messages=conv_messages) -def prepare_tree(tree: list[Post], tree_id: UUID) -> protocol.MessageTree: +def prepare_tree(tree: list[Message], tree_id: UUID) -> protocol.MessageTree: tree_messages = [] for message in tree: - if not isinstance(message.payload.payload, PostPayload): + if not isinstance(message.payload.payload, MessagePayload): raise OasstError("Server error", OasstErrorCode.SERVER_ERROR, HTTPStatus.INTERNAL_SERVER_ERROR) tree_messages.append( protocol.Message( 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/prompt_repository.py b/backend/oasst_backend/prompt_repository.py index 8df17347..155a6d2a 100644 --- a/backend/oasst_backend/prompt_repository.py +++ b/backend/oasst_backend/prompt_repository.py @@ -26,7 +26,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 = ( @@ -119,9 +119,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) @@ -481,7 +479,7 @@ class PromptRepository: 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.POST_NOT_FOUND) + raise OasstError("Message not found", OasstErrorCode.MESSAGE_NOT_FOUND) return message def close_task(self, frontend_message_id: str, allow_personal_tasks: bool = False): @@ -502,7 +500,7 @@ class PromptRepository: self.db.commit() @staticmethod - def trace_conversation(messages: list[Post] | dict[UUID, Post], last_message: Post) -> list[Post]: + 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. @@ -525,66 +523,68 @@ class PromptRepository: return list(reversed(conv)) - def fetch_message_conversation(self, message: Post | UUID) -> list[Post]: + 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_post(message) + message = self.fetch_message(message) - tree_messages = self.fetch_thread(message.thread_id) + tree_messages = self.fetch_message_tree(message.message_tree_id) return self.trace_conversation(tree_messages, message) - def fetch_message_tree(self, message: Post | UUID) -> list[Post]: + 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_post(message) - return self.fetch_thread(message.thread_id) + message = self.fetch_message(message) + return self.fetch_message_tree(message.message_tree_id) - def fetch_message_children(self, message: Post | UUID) -> list[Post]: + def fetch_message_children(self, message: Message | UUID) -> list[Message]: """ Get all direct children of this message """ - if isinstance(message, Post): + if isinstance(message, Message): message = message.id - children = self.db.query(Post).filter(Post.parent_id == message).all() + children = self.db.query(Message).filter(Message.parent_id == message).all() return children @staticmethod - def trace_descendants(root: Post, messages: list[Post]) -> list[Post]: + 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: Post): + 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_post_descendants(self, message: Post | UUID, max_depth: int = None) -> list[Post]: + def fetch_message_descendants(self, message: Message | UUID, max_depth: int = None) -> list[Message]: if isinstance(message, UUID): - message = self.fetch_post(message) + message = self.fetch_message(message) - desc = self.db.query(Post).filter(Post.thread_id == message.thread_id, Post.depth > message.depth) + 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(Post.depth <= max_depth) + desc = desc.filter(Message.depth <= max_depth) desc = desc.all() return self.trace_descendants(message, desc) - def fetch_longest_conversation(self, message: Post | UUID) -> list[Post]: - tree = self.fetch_message_tree(message) + 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: Post | UUID) -> tuple[Post, list[Post]]: - tree = self.fetch_message_tree(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] @@ -599,7 +599,7 @@ class PromptRepository: end_date: Optional[datetime.datetime] = None, only_roots: bool = False, deleted: Optional[bool] = None, - ) -> list[Post]: + ) -> 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 @@ -608,30 +608,30 @@ class PromptRepository: # Unprivileged api client asks for foreign messages raise OasstError("Forbidden", OasstErrorCode.API_CLIENT_NOT_AUTHORIZED, HTTP_403_FORBIDDEN) - messages = self.db.query(Post) + messages = self.db.query(Message) if user_id: - messages = messages.filter(Post.person_id == user_id) + messages = messages.filter(Message.user_id == user_id) if username: - messages = messages.join(Person) - messages = messages.filter(Person.username == username) + messages = messages.join(User) + messages = messages.filter(User.username == username) if api_client_id: - messages = messages.filter(Post.api_client_id == api_client_id) + messages = messages.filter(Message.api_client_id == api_client_id) if start_date: - messages = messages.filter(Post.created_date >= start_date) + messages = messages.filter(Message.created_date >= start_date) if end_date: - messages = messages.filter(Post.created_date < end_date) + messages = messages.filter(Message.created_date < end_date) if only_roots: - messages = messages.filter(Post.parent_id.is_(None)) + messages = messages.filter(Message.parent_id.is_(None)) if deleted is not None: - messages = messages.filter(Post.deleted == deleted) + messages = messages.filter(Message.deleted == deleted) if desc: - messages = messages.order_by(Post.created_date.desc()) + messages = messages.order_by(Message.created_date.desc()) else: - messages = messages.order_by(Post.created_date.asc()) + messages = messages.order_by(Message.created_date.asc()) if limit is not None: messages = messages.limit(limit) @@ -639,34 +639,36 @@ class PromptRepository: # TODO: Pagination could be great at some point return messages.all() - def mark_messages_deleted(self, messages: Post | UUID | list[Post | UUID], recursive: bool = True): - if isinstance(messages, (Post, UUID)): + def mark_messages_deleted(self, messages: Message | UUID | list[Message | UUID], recursive: bool = True): + if isinstance(messages, (Message, UUID)): messages = [messages] ids = [] for message in messages: if isinstance(message, UUID): ids.append(message) - elif isinstance(message, Post): + elif isinstance(message, Message): ids.append(message.id) else: raise OasstError("Server error", OasstErrorCode.SERVER_ERROR, HTTPStatus.INTERNAL_SERVER_ERROR) - query = update(Post).where(Post.id.in_(ids)).values(deleted=True) + 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(Post).filter(Post.parent_id.in_(parent_ids)).values(deleted=True).returning(Post.id) + 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): - deleted = self.db.query(Post.deleted, func.count()).group_by(Post.deleted) - nthreads = self.db.query(None, func.count(Post.id)).filter(Post.parent_id.is_(None)) + 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 { From 6d98ba1f7553ed08f9bce7da503463d3d8f673dc Mon Sep 17 00:00:00 2001 From: Igor Miagkov Date: Sat, 31 Dec 2022 04:41:27 +0400 Subject: [PATCH 5/6] fix alembic migration --- ...0438-8d269bc4fdbd_add_deleted_field_to_post.py} | 14 +++++++------- 1 file changed, 7 insertions(+), 7 deletions(-) rename backend/alembic/versions/{2022_12_30_0654-6cb49da61b74_add_deleted_field_to_post.py => 2022_12_31_0438-8d269bc4fdbd_add_deleted_field_to_post.py} (59%) diff --git a/backend/alembic/versions/2022_12_30_0654-6cb49da61b74_add_deleted_field_to_post.py b/backend/alembic/versions/2022_12_31_0438-8d269bc4fdbd_add_deleted_field_to_post.py similarity index 59% rename from backend/alembic/versions/2022_12_30_0654-6cb49da61b74_add_deleted_field_to_post.py rename to backend/alembic/versions/2022_12_31_0438-8d269bc4fdbd_add_deleted_field_to_post.py index a8f228bd..786471db 100644 --- a/backend/alembic/versions/2022_12_30_0654-6cb49da61b74_add_deleted_field_to_post.py +++ b/backend/alembic/versions/2022_12_31_0438-8d269bc4fdbd_add_deleted_field_to_post.py @@ -1,28 +1,28 @@ # -*- coding: utf-8 -*- """add deleted field to post -Revision ID: 6cb49da61b74 -Revises: 73ce3675c1f5 -Create Date: 2022-12-30 06:54:47.110204 +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 = "6cb49da61b74" -down_revision = "73ce3675c1f5" +revision = "8d269bc4fdbd" +down_revision = "abb47e9d145a" branch_labels = None depends_on = None def upgrade() -> None: # ### commands auto generated by Alembic - please adjust! ### - op.add_column("post", sa.Column("deleted", sa.Boolean(), server_default=sa.text("false"), nullable=False)) + 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("post", "deleted") + op.drop_column("message", "deleted") # ### end Alembic commands ### From ef3a35ff9c9da7aab45e6f40103f32558f861e9b Mon Sep 17 00:00:00 2001 From: Igor Miagkov Date: Sat, 31 Dec 2022 05:14:21 +0400 Subject: [PATCH 6/6] fixes --- .../oasst_backend/api/v1/frontend_messages.py | 3 +- .../oasst_backend/api/v1/frontend_users.py | 2 +- backend/oasst_backend/api/v1/messages.py | 5 ++- backend/oasst_backend/api/v1/stats.py | 3 +- backend/oasst_backend/api/v1/users.py | 2 +- backend/oasst_backend/exceptions.py | 1 + backend/oasst_backend/prompt_repository.py | 39 +++++++++++++++---- oasst-shared/oasst_shared/schemas/protocol.py | 2 +- 8 files changed, 42 insertions(+), 15 deletions(-) diff --git a/backend/oasst_backend/api/v1/frontend_messages.py b/backend/oasst_backend/api/v1/frontend_messages.py index c53f6bae..85ab74ca 100644 --- a/backend/oasst_backend/api/v1/frontend_messages.py +++ b/backend/oasst_backend/api/v1/frontend_messages.py @@ -24,7 +24,8 @@ def get_message_by_frontend_id( message = pr.fetch_message_by_frontend_message_id(message_id) if not isinstance(message.payload.payload, MessagePayload): - raise OasstError("Invalid message id", OasstErrorCode.INVALID_FRONTEND_MESSAGE_ID) + # Unexpected message payload + raise OasstError("Invalid message", OasstErrorCode.INVALID_MESSAGE) return protocol.ConversationMessage(text=message.payload.payload.text, is_assistant=(message.role == "assistant")) diff --git a/backend/oasst_backend/api/v1/frontend_users.py b/backend/oasst_backend/api/v1/frontend_users.py index f5df2b24..35048331 100644 --- a/backend/oasst_backend/api/v1/frontend_users.py +++ b/backend/oasst_backend/api/v1/frontend_users.py @@ -18,7 +18,7 @@ router = APIRouter() def query_frontend_user_messages( username: str, api_client_id: UUID = None, - max_count: int = Query(10, gt=0, le=25), + max_count: int = Query(10, gt=0, le=1000), start_date: datetime.datetime = None, end_date: datetime.datetime = None, only_roots: bool = False, diff --git a/backend/oasst_backend/api/v1/messages.py b/backend/oasst_backend/api/v1/messages.py index a812ef61..9221afc3 100644 --- a/backend/oasst_backend/api/v1/messages.py +++ b/backend/oasst_backend/api/v1/messages.py @@ -21,7 +21,7 @@ router = APIRouter() def query_messages( username: str = None, api_client_id: str = None, - max_count: int = Query(10, gt=0, le=25), + max_count: int = Query(10, gt=0, le=1000), start_date: datetime.datetime = None, end_date: datetime.datetime = None, only_roots: bool = False, @@ -63,7 +63,8 @@ def get_message( pr = PromptRepository(db, api_client, user=None) message = pr.fetch_message(message_id) if not isinstance(message.payload.payload, MessagePayload): - raise OasstError("Invalid message id", OasstErrorCode.INVALID_FRONTEND_MESSAGE_ID) + # Unexptcted message payload + raise OasstError("Invalid message", OasstErrorCode.INVALID_MESSAGE) return protocol.ConversationMessage(text=message.payload.payload.text, is_assistant=(message.role == "assistant")) diff --git a/backend/oasst_backend/api/v1/stats.py b/backend/oasst_backend/api/v1/stats.py index 789e53f9..831d4df2 100644 --- a/backend/oasst_backend/api/v1/stats.py +++ b/backend/oasst_backend/api/v1/stats.py @@ -3,7 +3,6 @@ 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 oasst_shared.schemas import protocol from sqlmodel import Session router = APIRouter() @@ -15,4 +14,4 @@ def get_message_stats( api_client: ApiClient = Depends(deps.get_trusted_api_client), ): pr = PromptRepository(db, api_client, None) - return protocol.SystemStats(**pr.get_stats()) + return pr.get_stats() diff --git a/backend/oasst_backend/api/v1/users.py b/backend/oasst_backend/api/v1/users.py index 000b0970..0bac4d6a 100644 --- a/backend/oasst_backend/api/v1/users.py +++ b/backend/oasst_backend/api/v1/users.py @@ -18,7 +18,7 @@ router = APIRouter() def query_user_messages( user_id: UUID, api_client_id: UUID = None, - max_count: int = Query(10, gt=0, le=25), + max_count: int = Query(10, gt=0, le=1000), start_date: datetime.datetime = None, end_date: datetime.datetime = None, only_roots: bool = False, diff --git a/backend/oasst_backend/exceptions.py b/backend/oasst_backend/exceptions.py index 7f88caed..f431b05b 100644 --- a/backend/oasst_backend/exceptions.py +++ b/backend/oasst_backend/exceptions.py @@ -36,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/prompt_repository.py b/backend/oasst_backend/prompt_repository.py index 155a6d2a..b95f9d53 100644 --- a/backend/oasst_backend/prompt_repository.py +++ b/backend/oasst_backend/prompt_repository.py @@ -13,6 +13,7 @@ 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 @@ -57,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 @@ -450,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) @@ -483,6 +492,9 @@ class PromptRepository: 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) @@ -565,6 +577,11 @@ class PromptRepository: 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) @@ -640,6 +657,9 @@ class PromptRepository: 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] @@ -666,14 +686,19 @@ class PromptRepository: self.db.commit() - def get_stats(self): + 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 { - "all": result.get(True, 0) + result.get(False, 0), - "active": result.get(False, 0), - "deleted": result.get(True, 0), - "threads": result.get(None, 0), - } + + 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 a2a12ddc..ed7dc780 100644 --- a/oasst-shared/oasst_shared/schemas/protocol.py +++ b/oasst-shared/oasst_shared/schemas/protocol.py @@ -277,4 +277,4 @@ class SystemStats(BaseModel): all: int = 0 active: int = 0 deleted: int = 0 - threads: int = 0 + message_trees: int = 0