Merge branch 'main' of github.com:LAION-AI/Open-Chat-GPT

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 ###

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

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"])

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)

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)

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)

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()

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)

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)

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

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,

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()))

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),
+        )

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