From f32ec78988210c9d482883cc46540bc8da3a0c5c Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Andreas=20K=C3=B6pf?= Date: Wed, 11 Jan 2023 14:46:29 +0100 Subject: [PATCH] Returning from /api/v1/frontend_users/{auth_method}/{username} fixes #615 --- .../oasst_backend/api/v1/frontend_users.py | 19 +++++++++++++------ backend/oasst_backend/api/v1/utils.py | 3 ++- backend/oasst_backend/prompt_repository.py | 14 ++++++++------ .../exceptions/oasst_api_error.py | 1 + oasst-shared/oasst_shared/schemas/protocol.py | 9 ++++++--- 5 files changed, 30 insertions(+), 16 deletions(-) diff --git a/backend/oasst_backend/api/v1/frontend_users.py b/backend/oasst_backend/api/v1/frontend_users.py index 31f14c64..1caf1f46 100644 --- a/backend/oasst_backend/api/v1/frontend_users.py +++ b/backend/oasst_backend/api/v1/frontend_users.py @@ -15,7 +15,7 @@ from starlette.status import HTTP_204_NO_CONTENT router = APIRouter() -@router.get("/{auth_method}/{username}", response_model=protocol.User) +@router.get("/{auth_method}/{username}", response_model=protocol.FrontEndUser) def query_frontend_user( auth_method: str, username: str, @@ -28,11 +28,14 @@ def query_frontend_user( """ ur = UserRepository(db, api_client) user = ur.query_frontend_user(auth_method, username, api_client_id) - return protocol.User(id=user.username, display_name=user.display_name, auth_method=user.auth_method) + return protocol.FrontEndUser( + id=user.username, display_name=user.display_name, auth_method=user.auth_method, user_id=user.id + ) -@router.get("/{username}/messages", response_model=list[protocol.Message]) +@router.get("/{auth_method}/{username}/messages", response_model=list[protocol.Message]) def query_frontend_user_messages( + auth_method: str, username: str, api_client_id: UUID = None, max_count: int = Query(10, gt=0, le=1000), @@ -49,6 +52,7 @@ def query_frontend_user_messages( """ pr = PromptRepository(db, api_client) messages = pr.query_messages( + auth_method=auth_method, username=username, api_client_id=api_client_id, desc=desc, @@ -61,10 +65,13 @@ def query_frontend_user_messages( return utils.prepare_message_list(messages) -@router.delete("/{username}/messages", status_code=HTTP_204_NO_CONTENT) +@router.delete("/{auth_method}/{username}/messages", status_code=HTTP_204_NO_CONTENT) def mark_frontend_user_messages_deleted( - username: str, api_client: ApiClient = Depends(deps.get_trusted_api_client), db: Session = Depends(deps.get_db) + auth_method: str, + username: str, + api_client: ApiClient = Depends(deps.get_trusted_api_client), + db: Session = Depends(deps.get_db), ): pr = PromptRepository(db, api_client) - messages = pr.query_messages(username=username, api_client_id=api_client.id) + messages = pr.query_messages(auth_method=auth_method, username=username, api_client_id=api_client.id) pr.mark_messages_deleted(messages) diff --git a/backend/oasst_backend/api/v1/utils.py b/backend/oasst_backend/api/v1/utils.py index a8b54483..035f6d19 100644 --- a/backend/oasst_backend/api/v1/utils.py +++ b/backend/oasst_backend/api/v1/utils.py @@ -7,6 +7,7 @@ from oasst_shared.schemas import protocol def prepare_message(m: Message) -> protocol.Message: return protocol.Message( id=m.id, + frontend_message_id=m.frontend_message_id, parent_id=m.parent_id, text=m.text, is_assistant=(m.role == "assistant"), @@ -23,7 +24,7 @@ def prepare_conversation_message_list(messages: list[Message]) -> list[protocol. protocol.ConversationMessage( text=message.text, is_assistant=(message.role == "assistant"), - message_id=message.id, + id=message.id, frontend_message_id=message.frontend_message_id, ) for message in messages diff --git a/backend/oasst_backend/prompt_repository.py b/backend/oasst_backend/prompt_repository.py index 75b9b6dc..5bae3b18 100644 --- a/backend/oasst_backend/prompt_repository.py +++ b/backend/oasst_backend/prompt_repository.py @@ -245,10 +245,10 @@ class PromptRepository: ) last_conv_message = task_payload.conversation.messages[-1] - parent_msg = self.fetch_message(last_conv_message.message_id) + parent_msg = self.fetch_message(last_conv_message.id) # store reaction to message - ranked_message_ids = [task_payload.reply_messages[i].message_id for i in ranking.ranking] + ranked_message_ids = [task_payload.reply_messages[i].id for i in ranking.ranking] for mid in ranked_message_ids: message = self.fetch_message(mid) if message.parent_id != parent_msg.id: @@ -273,7 +273,7 @@ class PromptRepository: ) # store reaction to message - ranked_message_ids = [task_payload.prompt_messages[i].message_id for i in ranking.ranking] + ranked_message_ids = [task_payload.prompt_messages[i].id for i in ranking.ranking] reaction_payload = db_payload.RankingReactionPayload( ranking=ranking.ranking, ranked_message_ids=ranked_message_ids ) @@ -617,6 +617,7 @@ class PromptRepository: def query_messages( self, user_id: Optional[UUID] = None, + auth_method: Optional[str] = None, username: Optional[str] = None, api_client_id: Optional[UUID] = None, desc: bool = True, @@ -637,9 +638,11 @@ class PromptRepository: messages = self.db.query(Message) if user_id: messages = messages.filter(Message.user_id == user_id) - if username: + 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) + messages = messages.filter(User.username == username, User.auth_method == auth_method) if api_client_id: messages = messages.filter(Message.api_client_id == api_client_id) @@ -662,7 +665,6 @@ class PromptRepository: 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): diff --git a/oasst-shared/oasst_shared/exceptions/oasst_api_error.py b/oasst-shared/oasst_shared/exceptions/oasst_api_error.py index c365eed2..e6eba233 100644 --- a/oasst-shared/oasst_shared/exceptions/oasst_api_error.py +++ b/oasst-shared/oasst_shared/exceptions/oasst_api_error.py @@ -44,6 +44,7 @@ class OasstErrorCode(IntEnum): BROKEN_CONVERSATION = 2009 TREE_NOT_IN_GROWING_STATE = 2010 CORRUPT_RANKING_RESULT = 2011 + AUTH_AND_USERNAME_REQUIRED = 2012 TEXT_LABELS_WRONG_MESSAGE_ID = 2050 TEXT_LABELS_INVALID_LABEL = 2051 diff --git a/oasst-shared/oasst_shared/schemas/protocol.py b/oasst-shared/oasst_shared/schemas/protocol.py index 374e8d26..63c9ded0 100644 --- a/oasst-shared/oasst_shared/schemas/protocol.py +++ b/oasst-shared/oasst_shared/schemas/protocol.py @@ -29,13 +29,17 @@ class User(BaseModel): auth_method: Literal["discord", "local"] +class FrontEndUser(User): + user_id: UUID + + class ConversationMessage(BaseModel): """Represents a message in a conversation between the user and the assistant.""" + id: Optional[UUID] = None + frontend_message_id: Optional[str] = None text: str is_assistant: bool - message_id: Optional[UUID] = None - frontend_message_id: Optional[str] = None class Conversation(BaseModel): @@ -45,7 +49,6 @@ class Conversation(BaseModel): class Message(ConversationMessage): - id: UUID parent_id: Optional[UUID] = None created_date: Optional[datetime] = None