Allow to export filtered by tree state and/or only prompts (#1314)

* allow to export filtered by tree state and/or only prompts

* formatting

* add option to filter by lang code

* remove old export code

* export as list in only-deleted case

backend/export.py

@@ -1,36 +1,43 @@
 import argparse
-import json
-import sys
 from pathlib import Path
 from typing import List, Optional
 from uuid import UUID
 
-from fastapi.encoders import jsonable_encoder
 from loguru import logger
 from oasst_backend.database import engine
-from oasst_backend.models import Message, MessageTreeState, message_tree_state
+from oasst_backend.models import Message, MessageTreeState
+from oasst_backend.models.message_tree_state import State as TreeState
 from oasst_backend.utils import tree_export
-from sqlmodel import Session, not_
+from sqlmodel import Session
 
 
-def fetch_tree_ids(db: Session, ready_only: bool = False):
-    tree_qry = db.query(MessageTreeState)
+def fetch_tree_ids(
+    db: Session,
+    state_filter: Optional[TreeState] = None,
+    lang: Optional[str] = None,
+) -> list[tuple[UUID, TreeState]]:
+    tree_qry = (
+        db.query(MessageTreeState)
+        .select_from(MessageTreeState)
+        .join(Message, MessageTreeState.message_tree_id == Message.id)
+    )
 
-    if ready_only:
-        tree_qry = tree_qry.filter(MessageTreeState.state == message_tree_state.State.READY_FOR_EXPORT)
+    if lang is not None:
+        tree_qry = tree_qry.filter(Message.lang == lang)
 
-    message_trees = tree_qry.all()
+    if state_filter:
+        tree_qry = tree_qry.filter(MessageTreeState.state == state_filter)
 
-    message_tree_ids = [tree.message_tree_id for tree in message_trees]
-    return message_tree_ids
+    return [(tree.message_tree_id, tree.state) for tree in tree_qry]
 
 
-def fetch_message_ids(
+def fetch_tree_messages(
     db: Session,
     message_tree_id: Optional[UUID] = None,
     user_id: Optional[UUID] = None,
-    include_deleted: bool = False,
-    deleted_only: bool = False,
+    deleted: bool = None,
+    prompts_only: bool = False,
+    lang: Optional[str] = None,
 ) -> List[Message]:
     qry = db.query(Message)
 
@@ -38,10 +45,12 @@ def fetch_message_ids(
         qry = qry.filter(Message.message_tree_id == message_tree_id)
     if user_id:
         qry = qry.filter(Message.user_id == user_id)
-    if not include_deleted:
-        qry = qry.filter(not_(Message.deleted))
-    if deleted_only:
-        qry = qry.filter(Message.deleted)
+    if deleted is not None:
+        qry = qry.filter(Message.deleted == deleted)
+    if prompts_only:
+        qry = qry.filter(Message.parent_id.is_(None))
+    if lang:
+        qry = qry.filter(Message.lang == lang)
 
     return qry.all()
 
@@ -50,31 +59,49 @@ def export_trees(
     db: Session,
     export_file: Optional[Path] = None,
     use_compression: bool = False,
-    ready_only: bool = False,
-    include_deleted: bool = False,
-    deleted_only: bool = False,
+    deleted: bool = False,
     user_id: Optional[UUID] = None,
-):
+    prompts_only: bool = False,
+    state_filter: Optional[TreeState] = None,
+    lang: Optional[str] = None,
+) -> None:
     trees_to_export: List[tree_export.ExportMessageTree] = []
 
     if user_id:
-        messages = fetch_message_ids(db, user_id=user_id, include_deleted=include_deleted, deleted_only=deleted_only)
-        message_tree_ids = [msg.message_tree_id for msg in messages]
+        messages = fetch_tree_messages(
+            db,
+            user_id=user_id,
+            deleted=deleted,
+            prompts_only=prompts_only,
+            lang=lang,
+        )
+        tree_export.write_messages_to_file(export_file, messages, use_compression)
     else:
-        message_tree_ids = fetch_tree_ids(db, ready_only)
+        message_tree_ids = fetch_tree_ids(db, state_filter, lang=lang)
 
-    message_trees = [
-        fetch_message_ids(db, message_tree_id=tree_id, include_deleted=include_deleted, deleted_only=deleted_only)
-        for tree_id in message_tree_ids
-    ]
+        message_trees = [
+            fetch_tree_messages(
+                db,
+                message_tree_id=tree_id,
+                deleted=deleted,
+                prompts_only=prompts_only,
+                lang=None,
+            )
+            for (tree_id, _) in message_tree_ids
+        ]
 
-    for message_tree_id, message_tree in zip(message_tree_ids, message_trees):
-        trees_to_export.append(tree_export.build_export_tree(message_tree_id, message_tree))
+        # when exporting only-deleted we don't have a porper tree structure, export as list
+        if deleted is True:
+            messages = [m for t in message_trees for m in t]
+            tree_export.write_messages_to_file(export_file, messages, use_compression)
+        else:
+            for (message_tree_id, message_tree_state), message_tree in zip(message_tree_ids, message_trees):
+                t = tree_export.build_export_tree(message_tree_id, message_tree_state, message_tree)
+                if prompts_only:
+                    t.prompt.replies = None
+                trees_to_export.append(t)
 
-    if export_file:
-        tree_export.write_trees_to_file(export_file, trees_to_export, use_compression)
-    else:
-        sys.stdout.write(json.dumps(jsonable_encoder(trees_to_export), indent=2))
+            tree_export.write_trees_to_file(export_file, trees_to_export, use_compression)
 
 
 def validate_args(args):
@@ -83,8 +110,8 @@ def validate_args(args):
 
     args.use_compression = args.export_file is not None and ".gz" in args.export_file
 
-    if args.ready_only and args.user_id is not None:
-        raise ValueError("Cannot use --ready-only when specifying a user ID")
+    if args.state and args.user is not None:
+        raise ValueError("Cannot use --state when specifying a user ID")
 
     if args.export_file is None:
         logger.warning("No export file provided, output will be sent to STDOUT")
@@ -98,11 +125,6 @@ def parse_args():
         type=str,
         help="Name of file to export trees to. If not provided, output will be sent to STDOUT",
     )
-    parser.add_argument(
-        "--ready-only",
-        action="store_true",
-        help="Only export trees which are ready for use",
-    )
     parser.add_argument(
         "--include-deleted",
         action="store_true",
@@ -116,7 +138,22 @@ def parse_args():
     parser.add_argument(
         "--user",
         type=str,
-        help="Only export trees involving the user with the specified ID. Incompatible with --ready-only.",
+        help="Only export trees involving the user with the specified ID. Incompatible with --state.",
+    )
+    parser.add_argument(
+        "--state",
+        type=str,
+        help="all|prompt_lottery_waiting|growing|ready_for_export|aborted_low_grade|halted_by_moderator|backlog_ranking",
+    )
+    parser.add_argument(
+        "--lang",
+        type=str,
+        help="Filter message trees by language code (BCP 47)",
+    )
+    parser.add_argument(
+        "--prompts-only",
+        action="store_true",
+        help="Export a list of initial prompt messages",
     )
 
     args = parser.parse_args()
@@ -127,15 +164,28 @@ def main():
     args = parse_args()
     validate_args(args)
 
+    state_filter: Optional[TreeState] = None
+    if args.state is None:
+        state_filter = TreeState.READY_FOR_EXPORT
+    elif args.state != "all":
+        state_filter = TreeState(args.state)
+
+    deleted: Optional[bool] = False
+    if args.include_deleted:
+        deleted = None
+    if args.deleted_only:
+        deleted = True
+
     with Session(engine) as db:
         export_trees(
             db,
             Path(args.export_file) if args.export_file is not None else None,
             args.use_compression,
-            args.ready_only,
-            args.include_deleted,
-            args.deleted_only,
-            UUID(args.user) if args.user is not None else None,
+            deleted=deleted,
+            user_id=UUID(args.user) if args.user is not None else None,
+            prompts_only=args.prompts_only,
+            state_filter=state_filter,
+            lang=args.lang,
         )
 
 

backend/main.py

@@ -333,24 +333,6 @@ 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 retry_scoring_failed_message_trees():
     try:
         logger.info("TreeManager.retry_scoring_failed_message_trees()")
@@ -382,17 +364,6 @@ 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",
-        default=False,
-        help="Export all trees which are ready for exporting.",
-        action="store_true",
-    )
-    parser.add_argument(
-        "--export-file",
-        type=str,
-        help="Name of file to export trees to. If not provided when exporting, output will be send to STDOUT",
-    )
     parser.add_argument(
         "--retry-scoring",
         default=False,
@@ -405,14 +376,10 @@ def main():
     if args.print_openapi_schema:
         print(get_openapi_schema())
 
-    if args.export:
-        use_compression: bool = ".gz" in args.export_file
-        export_ready_trees(file=args.export_file, use_compression=use_compression)
-
     if args.retry_scoring:
         retry_scoring_failed_message_trees()
 
-    if not (args.export or args.print_openapi_schema or args.retry_scoring):
+    if not (args.print_openapi_schema or args.retry_scoring):
         uvicorn.run(app, host=args.host, port=args.port)
 
 

backend/oasst_backend/prompt_repository.py

@@ -656,12 +656,6 @@ class PromptRepository:
             qry = qry.filter(not_(Message.deleted))
         return self._add_user_emojis_all(qry)
 
-    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,6 +1,4 @@
-import json
 import random
-import sys
 from datetime import datetime, timedelta
 from enum import Enum
 from http import HTTPStatus
@@ -10,7 +8,6 @@ from uuid import UUID
 import numpy as np
 import pydantic
 import sqlalchemy as sa
-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
@@ -25,7 +22,6 @@ from oasst_backend.models import (
     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,
@@ -1673,44 +1669,6 @@ DELETE FROM user_stats WHERE user_id = :user_id;
         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)
-
     @managed_tx_method(CommitMode.COMMIT)
     def retry_scoring_failed_message_trees(self):
         query = self.db.query(MessageTreeState).filter(
@@ -1777,5 +1735,3 @@ if __name__ == "__main__":
         # print(
         #     ".query_tree_ranking_results", tm.query_tree_ranking_results(UUID("21f9d585-d22c-44ab-a696-baa3d83b5f1b"))
         # )
-
-        # print(tm.export_trees_to_file(message_tree_ids=["7e75fb38-e664-4e2b-817c-b9a0b01b0074"], file="lol.jsonl"))

backend/oasst_backend/utils/tree_export.py

@@ -1,12 +1,16 @@
 from __future__ import annotations
 
+import contextlib
 import gzip
 import json
+import sys
 from collections import defaultdict
-from typing import Optional, TextIO
+from typing import Iterable, Optional, TextIO
+from uuid import UUID
 
 from fastapi.encoders import jsonable_encoder
 from oasst_backend.models import Message
+from oasst_backend.models.message_tree_state import State as TreeState
 from pydantic import BaseModel
 
 
@@ -41,10 +45,13 @@ class ExportMessageNode(BaseModel):
 
 class ExportMessageTree(BaseModel):
     message_tree_id: str
+    tree_state: Optional[str]
     prompt: Optional[ExportMessageNode]
 
 
-def build_export_tree(message_tree_id: str, messages: list[Message]) -> ExportMessageTree:
+def build_export_tree(
+    message_tree_id: UUID, message_tree_state: TreeState, messages: list[Message]
+) -> ExportMessageTree:
     export_messages = [ExportMessageNode.prep_message_export(m) for m in messages]
 
     messages_by_parent = defaultdict(list)
@@ -59,19 +66,54 @@ def build_export_tree(message_tree_id: str, messages: list[Message]) -> ExportMe
         return node
 
     prompt = assign_replies(messages_by_parent[None][0])
-    return ExportMessageTree(message_tree_id=str(message_tree_id), prompt=prompt)
+    return ExportMessageTree(message_tree_id=str(message_tree_id), tree_state=message_tree_state, prompt=prompt)
 
 
-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")
+# see https://stackoverflow.com/questions/17602878/how-to-handle-both-with-open-and-sys-stdout-nicely
+@contextlib.contextmanager
+def smart_open(filename: str = None) -> TextIO:
+    if filename and filename != "-":
+        fh = open(filename, "wt", encoding="UTF-8")
     else:
-        out_buff = open(file, "wt", encoding="UTF-8")
+        fh = sys.stdout
+
+    try:
+        yield fh
+    finally:
+        if fh is not sys.stdout:
+            fh.close()
+
+
+def write_trees_to_file(filename: str | None, trees: list[ExportMessageTree], use_compression: bool = True) -> None:
+    out_buff: TextIO
+
+    if use_compression:
+        if not filename:
+            raise RuntimeError("File name must be specified when using compression.")
+        out_buff = gzip.open(filename, "wt", encoding="UTF-8")
+    else:
+        out_buff = smart_open(filename)
 
     with out_buff as f:
         for tree in trees:
-            file_data = jsonable_encoder(tree)
+            file_data = jsonable_encoder(tree, exclude_none=True)
+            json.dump(file_data, f)
+            f.write("\n")
+
+
+def write_messages_to_file(filename: str | None, messages: Iterable[Message], use_compression: bool = True) -> None:
+    out_buff: TextIO
+
+    if use_compression:
+        if not filename:
+            raise RuntimeError("File name must be specified when using compression.")
+        out_buff = gzip.open(filename, "wt", encoding="UTF-8")
+    else:
+        out_buff = smart_open(filename)
+
+    with out_buff as f:
+        for m in messages:
+            export_message = ExportMessageNode.prep_message_export(m)
+            file_data = jsonable_encoder(export_message, exclude_none=True)
             json.dump(file_data, f)
             f.write("\n")