From 94e5d50537dad94bac6f22502ae25cf54edbc0b1 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Andreas=20K=C3=B6pf?= Date: Sat, 21 Jan 2023 00:29:53 +0100 Subject: [PATCH] add users/cursor keyset_pagination endpoint (#866) * add users/cursor endpoint * add messages/cursor endpoint * add user/{user_id}/messages/cursor, frontend_user/{auth_method}/{username}/messages/cursor * user regex to parse cursor value --- .../oasst_backend/api/v1/frontend_users.py | 41 +++++++- backend/oasst_backend/api/v1/messages.py | 87 +++++++++++++--- backend/oasst_backend/api/v1/users.py | 99 ++++++++++++++++++- backend/oasst_backend/api/v1/utils.py | 6 ++ backend/oasst_backend/prompt_repository.py | 84 ++++++++++------ backend/oasst_backend/user_repository.py | 2 + .../exceptions/oasst_api_error.py | 4 + oasst-shared/oasst_shared/schemas/protocol.py | 16 +++ 8 files changed, 287 insertions(+), 52 deletions(-) diff --git a/backend/oasst_backend/api/v1/frontend_users.py b/backend/oasst_backend/api/v1/frontend_users.py index f2fc3181..5d96aa90 100644 --- a/backend/oasst_backend/api/v1/frontend_users.py +++ b/backend/oasst_backend/api/v1/frontend_users.py @@ -5,6 +5,7 @@ from uuid import UUID from fastapi import APIRouter, Depends, Query from oasst_backend.api import deps from oasst_backend.api.v1 import utils +from oasst_backend.api.v1.messages import get_messages_cursor from oasst_backend.models import ApiClient from oasst_backend.prompt_repository import PromptRepository from oasst_backend.user_repository import UserRepository @@ -76,20 +77,47 @@ def query_frontend_user_messages( Query frontend user messages. """ pr = PromptRepository(db, api_client) - messages = pr.query_messages( + messages = pr.query_messages_ordered_by_created_date( auth_method=auth_method, username=username, api_client_id=api_client_id, desc=desc, limit=max_count, - start_date=start_date, - end_date=end_date, + gte_created_date=start_date, + lte_created_date=end_date, only_roots=only_roots, deleted=None if include_deleted else False, ) return utils.prepare_message_list(messages) +@router.get("/{auth_method}/{username}/messages/cursor", response_model=protocol.MessagePage) +def query_frontend_user_messages_cursor( + auth_method: str, + username: str, + lt: Optional[str] = None, + gt: Optional[str] = None, + only_roots: Optional[bool] = False, + include_deleted: Optional[bool] = False, + max_count: Optional[int] = Query(10, gt=0, le=1000), + desc: Optional[bool] = False, + api_client: ApiClient = Depends(deps.get_api_client), + db: Session = Depends(deps.get_db), +): + return get_messages_cursor( + lt=lt, + gt=gt, + auth_method=auth_method, + username=username, + only_roots=only_roots, + include_deleted=include_deleted, + max_count=max_count, + desc=desc, + api_client=api_client, + db=db, + ) + + @router.delete("/{auth_method}/{username}/messages", status_code=HTTP_204_NO_CONTENT) def mark_frontend_user_messages_deleted( auth_method: str, @@ -98,5 +126,10 @@ def mark_frontend_user_messages_deleted( db: Session = Depends(deps.get_db), ): pr = PromptRepository(db, api_client) - messages = pr.query_messages(auth_method=auth_method, username=username, api_client_id=api_client.id) + messages = pr.query_messages_ordered_by_created_date( + auth_method=auth_method, + username=username, + api_client_id=api_client.id, + limit=None, + ) pr.mark_messages_deleted(messages) diff --git a/backend/oasst_backend/api/v1/messages.py b/backend/oasst_backend/api/v1/messages.py index 409240cb..2dcca64b 100644 --- a/backend/oasst_backend/api/v1/messages.py +++ b/backend/oasst_backend/api/v1/messages.py @@ -1,4 +1,5 @@ -import datetime +from datetime import datetime +from typing import Optional from uuid import UUID from fastapi import APIRouter, Depends, Query @@ -6,8 +7,8 @@ from oasst_backend.api import deps from oasst_backend.api.v1 import utils from oasst_backend.models import ApiClient from oasst_backend.prompt_repository import PromptRepository +from oasst_shared.exceptions.oasst_api_error import OasstError, OasstErrorCode from oasst_shared.schemas import protocol -from oasst_shared.utils import unaware_to_utc from sqlmodel import Session from starlette.status import HTTP_204_NO_CONTENT @@ -16,31 +17,30 @@ router = APIRouter() @router.get("/", response_model=list[protocol.Message]) 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, + auth_method: Optional[str] = None, + username: Optional[str] = None, + api_client_id: Optional[str] = None, + max_count: Optional[int] = Query(10, gt=0, le=1000), + start_date: Optional[datetime] = None, + end_date: Optional[datetime] = None, + only_roots: Optional[bool] = False, + desc: Optional[bool] = True, + allow_deleted: Optional[bool] = False, api_client: ApiClient = Depends(deps.get_api_client), db: Session = Depends(deps.get_db), ): """ Query messages. """ - start_date = unaware_to_utc(start_date) - end_date = unaware_to_utc(end_date) - pr = PromptRepository(db, api_client) - messages = pr.query_messages( + messages = pr.query_messages_ordered_by_created_date( + auth_method=auth_method, username=username, api_client_id=api_client_id, desc=desc, limit=max_count, - start_date=start_date, - end_date=end_date, + gte_created_date=start_date, + lte_created_date=end_date, only_roots=only_roots, deleted=None if allow_deleted else False, ) @@ -48,6 +48,61 @@ def query_messages( return utils.prepare_message_list(messages) +@router.get("/cursor", response_model=protocol.MessagePage) +def get_messages_cursor( + lt: Optional[str] = None, + gt: Optional[str] = None, + user_id: Optional[UUID] = None, + auth_method: Optional[str] = None, + username: Optional[str] = None, + api_client_id: Optional[str] = None, + only_roots: Optional[bool] = False, + include_deleted: Optional[bool] = False, + max_count: Optional[int] = Query(10, gt=0, le=1000), + desc: Optional[bool] = False, + api_client: ApiClient = Depends(deps.get_api_client), + db: Session = Depends(deps.get_db), +): + def split_cursor(x: str | None) -> tuple[datetime, UUID]: + if not x: + return None, None + try: + m = utils.split_uuid_pattern.match(x) + if m: + return datetime.fromisoformat(m[2]), UUID(m[1]) + return datetime.fromisoformat(x), None + except ValueError: + raise OasstError("Invalid cursor value", OasstErrorCode.INVALID_CURSOR_VALUE) + + lte_created_date, lt_id = split_cursor(lt) + gte_created_date, gt_id = split_cursor(gt) + + pr = PromptRepository(db, api_client) + messages = pr.query_messages_ordered_by_created_date( + user_id=user_id, + auth_method=auth_method, + username=username, + api_client_id=api_client_id, + gte_created_date=gte_created_date, + gt_id=gt_id, + lte_created_date=lte_created_date, + lt_id=lt_id, + only_roots=only_roots, + deleted=None if include_deleted else False, + desc=desc, + limit=max_count, + ) + + items = utils.prepare_message_list(messages) + n, p = None, None + if len(items) > 0: + p = str(items[0].id) + "$" + items[0].created_date.isoformat() + n = str(items[-1].id) + "$" + items[-1].created_date.isoformat() + + order = "desc" if desc else "asc" + return protocol.MessagePage(prev=p, next=n, sort_key="created_date", order=order, items=items) + + @router.get("/{message_id}", response_model=protocol.Message) def get_message( message_id: UUID, api_client: ApiClient = Depends(deps.get_api_client), db: Session = Depends(deps.get_db) diff --git a/backend/oasst_backend/api/v1/users.py b/backend/oasst_backend/api/v1/users.py index 0b31495a..63a55691 100644 --- a/backend/oasst_backend/api/v1/users.py +++ b/backend/oasst_backend/api/v1/users.py @@ -5,10 +5,12 @@ from uuid import UUID from fastapi import APIRouter, Depends, Query from oasst_backend.api import deps from oasst_backend.api.v1 import utils +from oasst_backend.api.v1.messages import get_messages_cursor from oasst_backend.models import ApiClient, User from oasst_backend.prompt_repository import PromptRepository from oasst_backend.user_repository import UserRepository from oasst_backend.user_stats_repository import UserStatsRepository, UserStatsTimeFrame +from oasst_shared.exceptions.oasst_api_error import OasstError, OasstErrorCode from oasst_shared.schemas import protocol from sqlmodel import Session from starlette.status import HTTP_204_NO_CONTENT @@ -70,6 +72,70 @@ def get_users_ordered_by_display_name( return [u.to_protocol_frontend_user() for u in users] +@router.get("/cursor", response_model=protocol.FrontEndUserPage) +def get_users_cursor( + lt: Optional[str] = None, + gt: Optional[str] = None, + sort_key: Optional[str] = Query("username", max_length=32), + max_count: Optional[int] = Query(100, gt=0, le=10000), + api_client_id: Optional[UUID] = None, + search_text: Optional[str] = None, + auth_method: Optional[str] = None, + api_client: ApiClient = Depends(deps.get_api_client), + db: Session = Depends(deps.get_db), +): + def split_cursor(x: str | None) -> tuple[str, UUID]: + if not x: + return None, None + m = utils.split_uuid_pattern.match(x) + if m: + return m[2], UUID(m[1]) + return x, None + + items: list[protocol.FrontEndUser] + n, p = None, None + if sort_key == "username": + lte_username, lt_id = split_cursor(lt) + gte_username, gt_id = split_cursor(gt) + items = get_users_ordered_by_username( + api_client_id=api_client_id, + gte_username=gte_username, + gt_id=gt_id, + lte_username=lte_username, + lt_id=lt_id, + auth_method=auth_method, + search_text=search_text, + max_count=max_count, + api_client=api_client, + db=db, + ) + if len(items) > 0: + p = str(items[0].user_id) + "$" + items[0].id + n = str(items[-1].user_id) + "$" + items[-1].id + elif sort_key == "display_name": + lte_display_name, lt_id = split_cursor(lt) + gte_display_name, gt_id = split_cursor(gt) + items = get_users_ordered_by_display_name( + api_client_id=api_client_id, + gte_display_name=gte_display_name, + gt_id=gt_id, + lte_display_name=lte_display_name, + lt_id=lt_id, + auth_method=auth_method, + search_text=search_text, + max_count=max_count, + api_client=api_client, + db=db, + ) + if len(items) > 0: + p = str(items[0].user_id) + "$" + items[0].display_name + n = str(items[-1].user_id) + "$" + items[-1].display_name + else: + raise OasstError(f"Unsupported sort key: '{sort_key}'", OasstErrorCode.SORT_KEY_UNSUPPORTED) + + return protocol.FrontEndUserPage(prev=p, next=n, sort_key=sort_key, order="asc", items=items) + + @router.get("/{user_id}", response_model=protocol.FrontEndUser) def get_user( user_id: UUID, @@ -130,13 +196,13 @@ def query_user_messages( Query user messages. """ pr = PromptRepository(db, api_client) - messages = pr.query_messages( + messages = pr.query_messages_ordered_by_created_date( user_id=user_id, api_client_id=api_client_id, desc=desc, limit=max_count, - start_date=start_date, - end_date=end_date, + gte_created_date=start_date, + lte_created_date=end_date, only_roots=only_roots, deleted=None if include_deleted else False, ) @@ -144,12 +210,37 @@ def query_user_messages( return utils.prepare_message_list(messages) +@router.get("/{user_id}/messages/cursor", response_model=protocol.MessagePage) +def query_user_messages_cursor( + user_id: Optional[UUID], + lt: Optional[str] = None, + gt: Optional[str] = None, + only_roots: Optional[bool] = False, + include_deleted: Optional[bool] = False, + max_count: Optional[int] = Query(10, gt=0, le=1000), + desc: Optional[bool] = False, + api_client: ApiClient = Depends(deps.get_api_client), + db: Session = Depends(deps.get_db), +): + return get_messages_cursor( + lt=lt, + gt=gt, + user_id=user_id, + only_roots=only_roots, + include_deleted=include_deleted, + max_count=max_count, + desc=desc, + api_client=api_client, + db=db, + ) + + @router.delete("/{user_id}/messages", status_code=HTTP_204_NO_CONTENT) 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) - messages = pr.query_messages(user_id=user_id) + messages = pr.query_messages_ordered_by_created_date(user_id=user_id, limit=None) pr.mark_messages_deleted(messages) diff --git a/backend/oasst_backend/api/v1/utils.py b/backend/oasst_backend/api/v1/utils.py index 49cbbd29..99161e32 100644 --- a/backend/oasst_backend/api/v1/utils.py +++ b/backend/oasst_backend/api/v1/utils.py @@ -1,3 +1,4 @@ +import re from uuid import UUID from oasst_backend.models import Message @@ -43,3 +44,8 @@ def prepare_tree(tree: list[Message], tree_id: UUID) -> protocol.MessageTree: tree_messages.append(prepare_message(message)) return protocol.MessageTree(id=tree_id, messages=tree_messages) + + +split_uuid_pattern = re.compile( + r"^([0-9a-fA-F]{8}-[0-9a-fA-F]{4}-[0-9a-fA-F]{4}-[0-9a-fA-F]{4}-[0-9a-fA-F]{12})\$(.*)$" +) diff --git a/backend/oasst_backend/prompt_repository.py b/backend/oasst_backend/prompt_repository.py index 5e928d8b..abf5b721 100644 --- a/backend/oasst_backend/prompt_repository.py +++ b/backend/oasst_backend/prompt_repository.py @@ -1,6 +1,6 @@ -import datetime import random from collections import defaultdict +from datetime import datetime from http import HTTPStatus from typing import List, Optional, Tuple from uuid import UUID, uuid4 @@ -28,7 +28,8 @@ from oasst_backend.utils.database_utils import CommitMode, managed_tx_method from oasst_shared.exceptions import OasstError, OasstErrorCode from oasst_shared.schemas import protocol as protocol_schema from oasst_shared.schemas.protocol import SystemStats -from sqlmodel import Session, func, not_, text, update +from oasst_shared.utils import unaware_to_utc +from sqlmodel import Session, and_, func, not_, or_, text, update from starlette.status import HTTP_403_FORBIDDEN, HTTP_404_NOT_FOUND @@ -664,58 +665,85 @@ class PromptRepository: 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( + def query_messages_ordered_by_created_date( self, user_id: Optional[UUID] = None, auth_method: Optional[str] = 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, + gte_created_date: Optional[datetime] = None, + gt_id: Optional[UUID] = None, + lte_created_date: Optional[datetime] = None, + lt_id: Optional[UUID] = None, only_roots: bool = False, deleted: Optional[bool] = None, + desc: bool = False, + limit: Optional[int] = 100, ) -> 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: + if 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) + if 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) + qry = self.db.query(Message) if user_id: - messages = messages.filter(Message.user_id == user_id) + qry = qry.filter(Message.user_id == user_id) if username or auth_method: if not username and auth_method: raise OasstError("Auth method or username missing.", OasstErrorCode.AUTH_AND_USERNAME_REQUIRED) - messages = messages.join(User) - messages = messages.filter(User.username == username, User.auth_method == auth_method) + qry = qry.join(User) + qry = qry.filter(User.username == username, User.auth_method == auth_method) if api_client_id: - messages = messages.filter(Message.api_client_id == api_client_id) + qry = qry.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) + gte_created_date = unaware_to_utc(gte_created_date) + lte_created_date = unaware_to_utc(lte_created_date) + + if gte_created_date is not None: + if gt_id: + qry = qry.filter( + or_( + Message.created_date > gte_created_date, + and_(Message.created_date == gte_created_date, Message.id > gt_id), + ) + ) + else: + qry = qry.filter(Message.created_date >= gte_created_date) + elif gt_id: + raise OasstError("Need id and date for keyset pagination", OasstErrorCode.GENERIC_ERROR) + + if lte_created_date is not None: + if lt_id: + qry = qry.filter( + or_( + Message.created_date < lte_created_date, + and_(Message.created_date == lte_created_date, Message.id < lt_id), + ) + ) + else: + qry = qry.filter(Message.created_date <= lte_created_date) + elif lt_id: + raise OasstError("Need id and date for keyset pagination", OasstErrorCode.GENERIC_ERROR) if only_roots: - messages = messages.filter(Message.parent_id.is_(None)) + qry = qry.filter(Message.parent_id.is_(None)) if deleted is not None: - messages = messages.filter(Message.deleted == deleted) + qry = qry.filter(Message.deleted == deleted) if desc: - messages = messages.order_by(Message.created_date.desc()) + qry = qry.order_by(Message.created_date.desc(), Message.id.desc()) else: - messages = messages.order_by(Message.created_date.asc()) + qry = qry.order_by(Message.created_date.asc(), Message.id.asc()) if limit is not None: - messages = messages.limit(limit) + qry = qry.limit(limit) - return messages.all() + return qry.all() def update_children_counts(self, message_tree_id: UUID): sql_update_children_count = """ diff --git a/backend/oasst_backend/user_repository.py b/backend/oasst_backend/user_repository.py index c0c2a88d..1e9ac78f 100644 --- a/backend/oasst_backend/user_repository.py +++ b/backend/oasst_backend/user_repository.py @@ -202,9 +202,11 @@ class UserRepository: ) -> list[User]: if not self.api_client.trusted: if not api_client_id: + # Let unprivileged api clients query their own users without api_client_id being set api_client_id = self.api_client.id if api_client_id != self.api_client.id: + # Unprivileged api client asks for foreign users raise OasstError("Forbidden", OasstErrorCode.API_CLIENT_NOT_AUTHORIZED, HTTP_403_FORBIDDEN) qry = self.db.query(User).order_by(User.display_name, User.id) diff --git a/oasst-shared/oasst_shared/exceptions/oasst_api_error.py b/oasst-shared/oasst_shared/exceptions/oasst_api_error.py index e8cd2359..bed6f942 100644 --- a/oasst-shared/oasst_shared/exceptions/oasst_api_error.py +++ b/oasst-shared/oasst_shared/exceptions/oasst_api_error.py @@ -19,6 +19,10 @@ class OasstErrorCode(IntEnum): API_CLIENT_NOT_AUTHORIZED = 2 ROOT_TOKEN_NOT_AUTHORIZED = 3 DATABASE_MAX_RETRIES_EXHAUSTED = 4 + + SORT_KEY_UNSUPPORTED = 100 + INVALID_CURSOR_VALUE = 101 + TOO_MANY_REQUESTS = 429 SERVER_ERROR0 = 500 diff --git a/oasst-shared/oasst_shared/schemas/protocol.py b/oasst-shared/oasst_shared/schemas/protocol.py index d31f1c79..f2164e8f 100644 --- a/oasst-shared/oasst_shared/schemas/protocol.py +++ b/oasst-shared/oasst_shared/schemas/protocol.py @@ -37,6 +37,18 @@ class FrontEndUser(User): created_date: Optional[datetime] = None +class PageResult(BaseModel): + prev: str | None + next: str | None + sort_key: str + items: list + order: Literal["asc", "desc"] + + +class FrontEndUserPage(PageResult): + items: list[FrontEndUser] + + class ConversationMessage(BaseModel): """Represents a message in a conversation between the user and the assistant.""" @@ -58,6 +70,10 @@ class Message(ConversationMessage): created_date: Optional[datetime] = None +class MessagePage(PageResult): + items: list[Message] + + class MessageTree(BaseModel): """All messages belonging to the same message tree."""