677 - Add tree message export (LAION-AI#808)

* Added - Basic functions to export trees for users, export-ready trees and specific tree ids to files

* Added print to logger by default for no file specified

* linting to remove extra imports

* Added cli for exporting trees which are ready to export

Fixed some accidental removal

Updated message lookup to use dict for better perf

* removed unused imports

* changed export flag for including deleted prompts back to include_deleted for better understandability

* Use native collection types list, tuple, dict

* pre-commit fix

Co-authored-by: Andreas Köpf <andreas.koepf@provisio.com>

Author: GraemeHarris

backend/main.py

@@ -273,6 +273,24 @@ def get_openapi_schema():
     return json.dumps(app.openapi())
 
 
+def export_ready_trees(file: Optional[str] = None, use_compression: bool = False):
+    try:
+        with Session(engine) as db:
+            api_client = api_auth(settings.OFFICIAL_WEB_API_KEY, db=db)
+            dummy_user = protocol_schema.User(id="__dummy_user__", display_name="Dummy User", auth_method="local")
+
+            ur = UserRepository(db=db, api_client=api_client)
+            tr = TaskRepository(db=db, api_client=api_client, client_user=dummy_user, user_repository=ur)
+            pr = PromptRepository(
+                db=db, api_client=api_client, client_user=dummy_user, user_repository=ur, task_repository=tr
+            )
+            tm = TreeManager(db, pr)
+
+            tm.export_all_ready_trees(file, use_compression=use_compression)
+    except Exception:
+        logger.exception("Error exporting trees.")
+
+
 def main():
     # Importing here so we don't import packages unnecessarily if we're
     # importing main as a module.
@@ -289,11 +307,21 @@ def main():
     )
     parser.add_argument("--host", help="The host to run the server", default="0.0.0.0")
     parser.add_argument("--port", help="The port to run the server", default=8080)
+    parser.add_argument(
+        "--export", help="Export all trees which are ready for exporting.", action=argparse.BooleanOptionalAction
+    )
+    parser.add_argument(
+        "--export-file",
+        help="Name of file to export trees to. If not provided when exporting, output will be send to STDOUT",
+    )
 
     args = parser.parse_args()
 
     if args.print_openapi_schema:
         print(get_openapi_schema())
+    elif args.export:
+        use_compression: bool = ".gz" in args.export_file
+        export_ready_trees(file=args.export_file, use_compression=use_compression)
     else:
         uvicorn.run(app, host=args.host, port=args.port)
 

backend/oasst_backend/prompt_repository.py

@@ -2,7 +2,7 @@
 from collections import defaultdict
 from datetime import datetime
 from http import HTTPStatus
-from typing import List, Optional, Tuple
+from typing import Optional
 from uuid import UUID, uuid4
 
 import oasst_backend.models.db_payload as db_payload
@@ -255,7 +255,7 @@ def store_rating(self, rating: protocol_schema.MessageRating) -> MessageReaction
         return reaction
 
     @managed_tx_method(CommitMode.COMMIT)
-    def store_ranking(self, ranking: protocol_schema.MessageRanking) -> Tuple[MessageReaction, Task]:
+    def store_ranking(self, ranking: protocol_schema.MessageRanking) -> tuple[MessageReaction, Task]:
         # fetch task
         task = self.task_repository.fetch_task_by_frontend_message_id(ranking.message_id)
         self._validate_task(task, frontend_message_id=ranking.message_id)
@@ -345,13 +345,13 @@ def insert_toxicity(self, message_id: UUID, model: str, score: float, label: str
         return message_toxicity
 
     @managed_tx_method(CommitMode.FLUSH)
-    def insert_message_embedding(self, message_id: UUID, model: str, embedding: List[float]) -> MessageEmbedding:
+    def insert_message_embedding(self, message_id: UUID, model: str, embedding: list[float]) -> MessageEmbedding:
         """Insert the embedding of a new message in the database.
 
         Args:
             message_id (UUID): the identifier of the message we want to save its embedding
             model (str): the model used for creating the embedding
-            embedding (List[float]): the values obtained from the message & model
+            embedding (list[float]): the values obtained from the message & model
 
         Raises:
             OasstError: if misses some of the before params
@@ -383,7 +383,7 @@ def insert_reaction(
         return reaction
 
     @managed_tx_method(CommitMode.FLUSH)
-    def store_text_labels(self, text_labels: protocol_schema.TextLabels) -> Tuple[TextLabels, Task, Message]:
+    def store_text_labels(self, text_labels: protocol_schema.TextLabels) -> tuple[TextLabels, Task, Message]:
 
         valid_labels: Optional[list[str]] = None
         mandatory_labels: Optional[list[str]] = None
@@ -529,6 +529,22 @@ def fetch_message_tree(
             qry = qry.filter(not_(Message.deleted))
         return qry.all()
 
+    def fetch_user_message_trees(
+        self, user_id: Message.user_id, reviewed: bool = True, include_deleted: bool = False
+    ) -> list[Message]:
+        qry = self.db.query(Message).filter(Message.user_id == user_id)
+        if reviewed:
+            qry = qry.filter(Message.review_result)
+        if not include_deleted:
+            qry = qry.filter(not_(Message.deleted))
+        return qry.all()
+
+    def fetch_message_trees_ready_for_export(self) -> list[MessageTreeState]:
+        qry = self.db.query(MessageTreeState).filter(
+            MessageTreeState.state == message_tree_state.State.READY_FOR_EXPORT
+        )
+        return qry.all()
+
     def fetch_multiple_random_replies(self, max_size: int = 5, message_role: str = None):
         """
         Fetch a conversation with multiple possible replies to it.

backend/oasst_backend/tree_manager.py

@@ -1,4 +1,6 @@
+import json
 import random
+import sys
 from datetime import datetime
 from enum import Enum
 from http import HTTPStatus
@@ -7,11 +9,13 @@
 
 import numpy as np
 import pydantic
+from fastapi.encoders import jsonable_encoder
 from loguru import logger
 from oasst_backend.api.v1.utils import prepare_conversation, prepare_conversation_message_list
 from oasst_backend.config import TreeManagerConfiguration, settings
 from oasst_backend.models import Message, MessageReaction, MessageTreeState, Task, TextLabels, User, message_tree_state
 from oasst_backend.prompt_repository import PromptRepository
+from oasst_backend.utils import tree_export
 from oasst_backend.utils.database_utils import CommitMode, async_managed_tx_method, managed_tx_method
 from oasst_backend.utils.hugging_face import HfClassificationModel, HfEmbeddingModel, HfUrl, HuggingFaceAPI
 from oasst_backend.utils.ranking import ranked_pairs
@@ -1184,14 +1188,55 @@ def purge_user(self, user_id: UUID, ban: bool = True) -> None:
         if ban:
             self.db.execute(update(User).filter(User.id == user_id).values(deleted=True, enabled=False))
 
+    def export_trees_to_file(
+        self,
+        message_tree_ids: list[str],
+        file=None,
+        reviewed: bool = True,
+        include_deleted: bool = False,
+        use_compression: bool = False,
+    ) -> None:
+        trees_to_export: List[tree_export.ExportMessageTree] = []
+
+        for message_tree_id in message_tree_ids:
+            messages: List[Message] = self.pr.fetch_message_tree(message_tree_id, reviewed, include_deleted)
+            trees_to_export.append(tree_export.build_export_tree(message_tree_id, messages))
+
+        if file:
+            tree_export.write_trees_to_file(file, trees_to_export, use_compression)
+        else:
+            sys.stdout.write(json.dumps(jsonable_encoder(trees_to_export), indent=2))
+
+    def export_all_ready_trees(
+        self, file: str, reviewed: bool = True, include_deleted: bool = False, use_compression: bool = False
+    ) -> None:
+        message_tree_states: MessageTreeState = self.pr.fetch_message_trees_ready_for_export()
+        message_tree_ids = [ms.message_tree_id for ms in message_tree_states]
+        self.export_trees_to_file(message_tree_ids, file, reviewed, include_deleted, use_compression)
+
+    def export_all_user_trees(
+        self,
+        user_id: str,
+        file: str,
+        reviewed: bool = True,
+        include_deleted: bool = False,
+        use_compression: bool = False,
+    ) -> None:
+        messages = self.pr.fetch_user_message_trees(UUID(user_id))
+        message_tree_ids = [ms.message_tree_id for ms in messages]
+        self.export_trees_to_file(message_tree_ids, file, reviewed, include_deleted, use_compression)
+
 
 if __name__ == "__main__":
     from oasst_backend.api.deps import api_auth
+
+    # from oasst_backend.api.deps import create_api_client
     from oasst_backend.database import engine
     from oasst_backend.prompt_repository import PromptRepository
 
     with Session(engine) as db:
         api_client = api_auth(settings.OFFICIAL_WEB_API_KEY, db=db)
+        # api_client = create_api_client(session=db, description="test", frontend_type="bot")
         dummy_user = protocol_schema.User(id="__dummy_user__", display_name="Dummy User", auth_method="local")
 
         pr = PromptRepository(db=db, api_client=api_client, client_user=dummy_user)
@@ -1200,25 +1245,22 @@ def purge_user(self, user_id: UUID, ban: bool = True) -> None:
         tm = TreeManager(db, pr, cfg)
         tm.ensure_tree_states()
 
-        tm.purge_user_messages(user_id=UUID("2ef9ad21-0dc5-442d-8750-6f7f1790723f"), purge_initial_prompts=False)
+        # tm.purge_user_messages(user_id=UUID("2ef9ad21-0dc5-442d-8750-6f7f1790723f"), purge_initial_prompts=False)
         # tm.purge_user(user_id=UUID("2ef9ad21-0dc5-442d-8750-6f7f1790723f"))
         # db.commit()
 
         # print("query_num_active_trees", tm.query_num_active_trees())
         # print("query_incomplete_rankings", tm.query_incomplete_rankings())
         # print("query_replies_need_review", tm.query_replies_need_review())
+        # print("query_incomplete_reply_reviews", tm.query_replies_need_review())
         # print("query_incomplete_initial_prompt_reviews", tm.query_prompts_need_review())
         # print("query_extendible_trees", tm.query_extendible_trees())
         # print("query_extendible_parents", tm.query_extendible_parents())
-        # print("query_tree_size", tm.query_tree_size(message_tree_id=UUID("bdf434cf-4df5-4b74-949c-a5a157bc3292")))
-
-        # print(
-        #     "query_reviews_for_message",
-        #     tm.query_reviews_for_message(message_id=UUID("6a444493-0d48-4316-a9f1-7e263f5a2473")),
-        # )
 
         # print("next_task:", tm.next_task())
 
         # print(
-        #     "query_tree_ranking_results", tm.query_tree_ranking_results(UUID("6036f58f-41b5-48c4-bdd9-b16f34ab1312"))
+        #     ".query_tree_ranking_results", tm.query_tree_ranking_results(UUID("2ac20d38-6650-43aa-8bb3-f61080c0d921"))
         # )
+
+        print(tm.export_trees_to_file(message_tree_ids=["7e75fb38-e664-4e2b-817c-b9a0b01b0074"], file="lol.jsonl"))

backend/oasst_backend/utils/tree_export.py

@@ -0,0 +1,71 @@
+from __future__ import annotations
+
+import gzip
+import json
+from collections import defaultdict
+from typing import Optional, TextIO
+
+from fastapi.encoders import jsonable_encoder
+from oasst_backend.models import Message
+from pydantic import BaseModel
+
+
+class ExportMessageNode(BaseModel):
+    message_id: str
+    parent_id: Optional[str]
+    text: Optional[str]
+    role: str
+    review_count: Optional[int]
+    rank: Optional[int]
+    replies: Optional[list[ExportMessageNode]]
+
+    @classmethod
+    def prep_message_export(cls, message: Message) -> ExportMessageNode:
+        return cls(
+            message_id=str(message.id),
+            parent_id=str(message.parent_id) if message.parent_id else None,
+            text=str(message.payload.payload.text),
+            role=message.role,
+            review_count=message.review_count,
+            rank=message.rank,
+        )
+
+
+class ExportMessageTree(BaseModel):
+    message_tree_id: str
+    replies: Optional[ExportMessageNode]
+
+
+def build_export_tree(message_tree_id: str, messages: list[Message]) -> ExportMessageTree:
+    export_tree = ExportMessageTree(message_tree_id=str(message_tree_id))
+    export_tree_data = [ExportMessageNode.prep_message_export(m) for m in messages]
+
+    message_parents = defaultdict(list)
+    for message in export_tree_data:
+        message_parents[message.parent_id].append(message)
+
+    def build_tree(tree: dict, parent: Optional[str], messages: list[Message]):
+        children = message_parents[parent]
+        tree.replies = children
+
+        for idx, child in enumerate(tree.replies):
+            build_tree(tree.replies[idx], child.message_id, messages)
+
+    build_tree(export_tree, None, export_tree_data)
+
+    return export_tree
+
+
+def write_trees_to_file(file, trees: list[ExportMessageTree], use_compression: bool = True) -> None:
+
+    out_buff: TextIO
+    if use_compression:
+        out_buff = gzip.open(file, "wt", encoding="UTF-8")
+    else:
+        out_buff = open(file, "wt", encoding="UTF-8")
+
+    with out_buff as f:
+        for tree in trees:
+            file_data = jsonable_encoder(tree)
+            json.dump(file_data, f)
+            f.write("\n")