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
This commit is contained in:
Andreas Köpf
2023-01-21 00:29:53 +01:00
committed by GitHub
parent 2d21b65ed0
commit 94e5d50537
8 changed files with 287 additions and 52 deletions
+37 -4
View File
@@ -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)
+71 -16
View File
@@ -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)
+95 -4
View File
@@ -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)
+6
View File
@@ -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})\$(.*)$"
)
+56 -28
View File
@@ -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 = """
+2
View File
@@ -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)
@@ -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
@@ -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."""