diff --git a/backend/oasst_backend/api/v1/frontend_messages.py b/backend/oasst_backend/api/v1/frontend_messages.py index f149bebb..dffb2824 100644 --- a/backend/oasst_backend/api/v1/frontend_messages.py +++ b/backend/oasst_backend/api/v1/frontend_messages.py @@ -45,7 +45,7 @@ def get_tree_by_frontend_id( """ pr = PromptRepository(db, api_client) message = pr.fetch_message_by_frontend_message_id(message_id) - tree = pr.fetch_message_tree(message.message_tree_id) + tree = pr.fetch_message_tree(message.message_tree_id, reviewed=False) return utils.prepare_tree(tree, message.message_tree_id) diff --git a/backend/oasst_backend/api/v1/messages.py b/backend/oasst_backend/api/v1/messages.py index fcba59df..409240cb 100644 --- a/backend/oasst_backend/api/v1/messages.py +++ b/backend/oasst_backend/api/v1/messages.py @@ -82,7 +82,7 @@ def get_tree( """ pr = PromptRepository(db, api_client) message = pr.fetch_message(message_id) - tree = pr.fetch_message_tree(message.message_tree_id) + tree = pr.fetch_message_tree(message.message_tree_id, reviewed=False) return utils.prepare_tree(tree, message.message_tree_id) diff --git a/backend/oasst_backend/prompt_repository.py b/backend/oasst_backend/prompt_repository.py index 2dd920e8..35e1ece4 100644 --- a/backend/oasst_backend/prompt_repository.py +++ b/backend/oasst_backend/prompt_repository.py @@ -28,7 +28,7 @@ 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, text, update +from sqlmodel import Session, func, not_, text, update from starlette.status import HTTP_403_FORBIDDEN, HTTP_404_NOT_FOUND @@ -506,10 +506,14 @@ class PromptRepository: messages = self.db.query(Message).filter(Message.parent_id.is_(None)).order_by(func.random()).limit(size).all() return messages - def fetch_message_tree(self, message_tree_id: UUID, reviewed: bool = True) -> list[Message]: + def fetch_message_tree( + self, message_tree_id: UUID, reviewed: bool = True, include_deleted: bool = False + ) -> list[Message]: qry = self.db.query(Message).filter(Message.message_tree_id == message_tree_id) if reviewed: qry = qry.filter(Message.review_result) + if not include_deleted: + qry = qry.filter(not_(Message.deleted)) return qry.all() def fetch_multiple_random_replies(self, max_size: int = 5, message_role: str = None): diff --git a/backend/oasst_backend/tree_manager.py b/backend/oasst_backend/tree_manager.py index 632e76e4..f4338048 100644 --- a/backend/oasst_backend/tree_manager.py +++ b/backend/oasst_backend/tree_manager.py @@ -1075,7 +1075,7 @@ DELETE FROM message WHERE message_tree_id = :message_tree_id; bad_parent_ids = set(m.id for m in replies) logger.debug(f"patching tree {tree_id=}, {bad_parent_ids=}") - tree_messages = self.pr.fetch_message_tree(tree_id) + tree_messages = self.pr.fetch_message_tree(tree_id, reviewed=False, include_deleted=True) logger.debug(f"{tree_id=}, {len(bad_parent_ids)=}, {len(tree_messages)=}") by_id = {m.id: m for m in tree_messages}