add include_user params in /messages/cursor` endpoint (LAION-AI#2021)

Needed to display user information in message table (admin interface).

Author: notmd

backend/oasst_backend/api/v1/messages.py

@@ -70,6 +70,7 @@ def get_messages_cursor(
     max_count: Optional[int] = Query(10, gt=0, le=1000),
     desc: Optional[bool] = False,
     lang: Optional[str] = None,
+    include_user: Optional[bool] = None,
     frontend_user: deps.FrontendUserId = Depends(deps.get_frontend_user_id),
     api_client: ApiClient = Depends(deps.get_api_client),
     db: Session = Depends(deps.get_db),
@@ -115,6 +116,7 @@ def split_cursor(x: str | None) -> tuple[datetime, UUID]:
         desc=query_desc,
         limit=qry_max_count,
         lang=lang,
+        include_user=include_user,
     )
 
     num_rows = len(items)

backend/oasst_backend/api/v1/utils.py

@@ -26,6 +26,7 @@ def prepare_message(m: Message) -> protocol.Message:
         model_name=m.model_name,
         message_tree_id=m.message_tree_id,
         rank=m.rank,
+        user=m.user.to_protocol_frontend_user() if m.user else None,
     )
 
 

backend/oasst_backend/models/message.py

@@ -6,6 +6,7 @@
 import sqlalchemy as sa
 import sqlalchemy.dialects.postgresql as pg
 from oasst_backend.models.db_payload import MessagePayload
+from oasst_backend.models.user import User
 from oasst_shared.exceptions.oasst_api_error import OasstError, OasstErrorCode
 from pydantic import PrivateAttr
 from sqlalchemy import false
@@ -65,6 +66,7 @@ def __new__(cls, *args: Any, **kwargs: Any):
     emojis: Optional[dict[str, int]] = Field(default=None, sa_column=sa.Column(pg.JSONB), nullable=False)
     _user_emojis: Optional[list[str]] = PrivateAttr(default=None)
     _user_is_author: Optional[bool] = PrivateAttr(default=None)
+    _user: Optional[bool] = PrivateAttr(default=None)
 
     def ensure_is_message(self) -> None:
         if not self.payload or not isinstance(self.payload.payload, MessagePayload):
@@ -88,3 +90,7 @@ def user_emojis(self) -> str:
     @property
     def user_is_author(self) -> str:
         return self._user_is_author
+
+    @property
+    def user(self) -> User:
+        return self._user

backend/oasst_backend/prompt_repository.py

@@ -888,14 +888,27 @@ def fetch_message_with_max_children(self, message: Message | UUID) -> tuple[Mess
         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 _add_user_emojis_all(self, qry: Query) -> list[Message]:
+    def _add_user_emojis_all(self, qry: Query, include_user: bool = False) -> list[Message]:
         if self.user_id is None:
-            return qry.all()
+            if not include_user:
+                return qry.all()
+
+            messages: list[Message] = []
+
+            for element in qry:
+                message = element["Message"]
+                user = element["User"]
+                message._user = user
+                messages.append(message)
+            return messages
 
         order_by_clauses = qry._order_by_clauses
         sq = qry.subquery("m")
+        select_entities = [Message, func.string_agg(MessageEmoji.emoji, literal_column("','")).label("user_emojis")]
+        if include_user:
+            select_entities.append(User)
         qry = (
-            self.db.query(Message, func.string_agg(MessageEmoji.emoji, literal_column("','")).label("user_emojis"))
+            self.db.query(*select_entities)
             .select_entity_from(sq)
             .outerjoin(
                 MessageEmoji,
@@ -915,7 +928,10 @@ def _add_user_emojis_all(self, qry: Query) -> list[Message]:
             if user_emojis:
                 m._user_emojis = user_emojis.split(",")
             m._user_is_author = self.user_id and self.user_id == m.user_id
+            if include_user:
+                m._user = x["User"]
             messages.append(m)
+
         return messages
 
     def query_messages_ordered_by_created_date(
@@ -934,6 +950,7 @@ def query_messages_ordered_by_created_date(
         desc: bool = False,
         limit: Optional[int] = 100,
         lang: Optional[str] = None,
+        include_user: Optional[bool] = None,
     ) -> list[Message]:
         if not self.api_client.trusted:
             if not api_client_id:
@@ -945,12 +962,15 @@ def query_messages_ordered_by_created_date(
                 raise OasstError("Forbidden", OasstErrorCode.API_CLIENT_NOT_AUTHORIZED, HTTPStatus.FORBIDDEN)
 
         qry = self.db.query(Message)
+        if include_user:
+            qry = self.db.query(Message, User)
         if user_id:
             qry = qry.filter(Message.user_id == user_id)
+        if username or auth_method or include_user:
+            qry = qry.join(User)
         if username or auth_method:
             if not (username and auth_method):
                 raise OasstError("Auth method or username missing.", OasstErrorCode.AUTH_AND_USERNAME_REQUIRED)
-            qry = qry.join(User)
             qry = qry.filter(User.username == username, User.auth_method == auth_method)
         if api_client_id:
             qry = qry.filter(Message.api_client_id == api_client_id)
@@ -1004,7 +1024,7 @@ def query_messages_ordered_by_created_date(
         if limit is not None:
             qry = qry.limit(limit)
 
-        return self._add_user_emojis_all(qry)
+        return self._add_user_emojis_all(qry, include_user=include_user)
 
     def update_children_counts(self, message_tree_id: UUID):
         sql_update_children_count = """

oasst-shared/oasst_shared/schemas/protocol.py

@@ -108,6 +108,7 @@ class Message(ConversationMessage):
     message_tree_id: Optional[UUID]
     ranking_count: Optional[int]
     rank: Optional[int]
+    user: Optional[FrontEndUser]
 
 
 class MessagePage(PageResult):