support pagination for chat list (#2773)

close #2584. ~~Doesn't support return null cursor when the chats in
response contain first or last chat. So the frontend need to do 1 more
request~~
Frontend not support this yet

inference/server/oasst_inference_server/routes/chats.py

@@ -1,8 +1,9 @@
 import asyncio
+import base64
 
 import fastapi
 import pydantic
-from fastapi import Depends
+from fastapi import Depends, Query
 from loguru import logger
 from oasst_inference_server import auth, chat_utils, deps, models, queueing
 from oasst_inference_server.schemas import chat as chat_schema
@@ -21,12 +22,47 @@
 async def list_chats(
     include_hidden: bool = False,
     ucr: UserChatRepository = Depends(deps.create_user_chat_repository),
+    limit: int | None = Query(10, gt=0, le=100),
+    after: str | None = None,
+    before: str | None = None,
 ) -> chat_schema.ListChatsResponse:
     """Lists all chats."""
     logger.info("Listing all chats.")
-    chats = await ucr.get_chats(include_hidden=include_hidden)
+
+    def encode_cursor(chat: models.DbChat):
+        return base64.b64encode(chat.id.encode()).decode()
+
+    def decode_cursor(cursor: str | None):
+        if cursor is None:
+            return None
+        return base64.b64decode(cursor.encode()).decode()
+
+    chats = await ucr.get_chats(
+        include_hidden=include_hidden, limit=limit + 1, after=decode_cursor(after), before=decode_cursor(before)
+    )
+
+    num_rows = len(chats)
+    chats = chats if num_rows <= limit else chats[:-1]  # remove extra item
+    chats = chats if before is None else chats[::-1]  # reverse if query in backward direction
+
+    def get_cursors():
+        prev, next = None, None
+        if num_rows > 0:
+            if (num_rows > limit and before) or after:
+                prev = encode_cursor(chats[0])
+            if num_rows > limit or before:
+                next = encode_cursor(chats[-1])
+        else:
+            if after:
+                prev = after
+            if before:
+                next = before
+        return prev, next
+
+    prev, next = get_cursors()
+
     chats_list = [chat.to_list_read() for chat in chats]
-    return chat_schema.ListChatsResponse(chats=chats_list)
+    return chat_schema.ListChatsResponse(chats=chats_list, next=next, prev=prev)
 
 
 @router.post("")

inference/server/oasst_inference_server/schemas/chat.py

@@ -70,6 +70,8 @@ class ChatRead(ChatListRead):
 
 class ListChatsResponse(pydantic.BaseModel):
     chats: list[ChatListRead]
+    next: str | None = None
+    prev: str | None = None
 
 
 class MessageCancelledException(Exception):

inference/server/oasst_inference_server/user_chat_repository.py

@@ -14,12 +14,30 @@ class UserChatRepository(pydantic.BaseModel):
     class Config:
         arbitrary_types_allowed = True
 
-    async def get_chats(self, include_hidden: bool = False) -> list[models.DbChat]:
+    async def get_chats(
+        self,
+        include_hidden: bool = False,
+        limit: int | None = None,
+        before: str | None = None,
+        after: str | None = None,
+    ) -> list[models.DbChat]:
+        if after is not None and before is not None:
+            raise fastapi.HTTPException(status_code=400, detail="Cannot specify both after and before.")
+
         query = sqlmodel.select(models.DbChat)
         query = query.where(models.DbChat.user_id == self.user_id)
+
         if not include_hidden:
             query = query.where(models.DbChat.hidden.is_(False))
-        query = query.order_by(models.DbChat.created_at.desc())
+        if limit is not None:
+            query = query.limit(limit)
+        if before is not None:
+            query = query.where(models.DbChat.id > before)
+        if after is not None:
+            query = query.where(models.DbChat.id < after)
+
+        query = query.order_by(models.DbChat.created_at.desc() if before is None else models.DbChat.created_at)
+
         return (await self.session.exec(query)).all()
 
     async def get_chat_by_id(self, chat_id: str, include_messages: bool = True) -> models.DbChat: