fetch whole message tree for purge (including non-reviewed & deleted)

LAION-AI · Jan 19, 2023 · 7a1a8c8 · 7a1a8c8
1 parent 1a93c21
commit 7a1a8c8
Show file tree

Hide file tree

Showing 4 changed files with 9 additions and 5 deletions.
diff --git 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
@@ -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
@@ -28,7 +28,7 @@
 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 @@ def fetch_random_initial_prompts(self, size: int = 5):
         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
@@ -1075,7 +1075,7 @@ def purge_user_messages(
             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}