Repository: LAION-AI/Open-Assistant
File: backend/export.py (Python)

Code before:
) -> List[Message]:
    qry = db.query(Message)

    if message_tree_id:
        qry = qry.filter(Message.message_tree_id == message_tree_id)
    if user_id:
        qry = qry.filter(Message.user_id == user_id)
    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)
    if review_result is False:
        qry = qry.filter(not_(Message.review_result), Message.review_count > 2)
    elif review_result is True:
        qry = qry.filter(Message.review_result)
    if limit is not None and limit > 0:
        qry = qry.limit(limit)

    return qry.all()


def get_events_for_messages(db: Session, message_ids: list[UUID]) -> dict[UUID, ExportMessageEvent]:
    events = {}
    emojis = db.query(MessageEmoji).filter(MessageEmoji.message_id.in_(message_ids)).all()
    for emoji in emojis:
        event = ExportMessageEventEmoji(user_id=str(emoji.user_id), emoji=emoji.emoji)
        events.setdefault(emoji.message_id, {}).setdefault("emoji", []).append(event)
    reactions: list[MessageReaction] = (
        db.query(MessageReaction).filter(MessageReaction.message_id.in_(message_ids)).all()
    )
    for reaction in reactions:
        match reaction.payload_type:
            case "RatingReactionPayload":
                key = "rating"
                payload: db_payload.RatingReactionPayload = reaction.payload.payload
                event = ExportMessageEventRating(user_id=str(reaction.user_id), rating=payload.rating)
            case "RankingReactionPayload":
                key = "ranking"
                payload: db_payload.RankingReactionPayload = reaction.payload.payload
                event = ExportMessageEventRanking(
                    user_id=str(reaction.user_id),
                    ranking=payload.ranking,
                    ranked_message_ids=[str(id) for id in payload.ranked_message_ids],
                    ranking_parent_id=str(payload.ranking_parent_id) if payload.ranking_parent_id else None,
                    message_tree_id=str(payload.message_tree_id) if payload.message_tree_id else None,
                    not_rankable=payload.not_rankable if payload.not_rankable else None,
                )
            case _:
                raise ValueError(f"Unknown payload type {reaction.payload_type}")

Review comment (suggestion):
(We have to test this with the the production db size, I had similar code before for labels which I re-wrote after the query took > 0.5h to execute for 100k messages .. i.e. `fetch_tree_messages_and_avg_labels` function does the join now directly ...)

Code after:
) -> List[Message]:
    qry = db.query(Message)

    if message_tree_id:
        qry = qry.filter(Message.message_tree_id == message_tree_id)
    if user_id:
        qry = qry.filter(Message.user_id == user_id)
    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)
    if review_result is False:
        qry = qry.filter(not_(Message.review_result), Message.review_count > 2)
    elif review_result is True:
        qry = qry.filter(Message.review_result)
    if limit is not None:
        qry = qry.limit(limit)

    return qry.all()


def get_events_for_messages(db: Session, message_ids: list[UUID]) -> dict[UUID, ExportMessageEvent]:
    events = {}
    emojis = db.query(MessageEmoji).filter(MessageEmoji.message_id.in_(message_ids)).all()
    for emoji in emojis:
        event = ExportMessageEventEmoji(user_id=str(emoji.user_id), emoji=emoji.emoji)
        events.setdefault(emoji.message_id, {}).setdefault("emoji", []).append(event)
    reactions: list[MessageReaction] = (
        db.query(MessageReaction).filter(MessageReaction.message_id.in_(message_ids)).all()
    )
    for reaction in reactions:
        match reaction.payload_type:
            case "RatingReactionPayload":
                key = "rating"
                payload: db_payload.RatingReactionPayload = reaction.payload.payload
                event = ExportMessageEventRating(user_id=str(reaction.user_id), rating=payload.rating)
            case "RankingReactionPayload":
                key = "ranking"
                payload: db_payload.RankingReactionPayload = reaction.payload.payload
                event = ExportMessageEventRanking(
                    user_id=str(reaction.user_id),
                    ranking=payload.ranking,
                    ranked_message_ids=[str(id) for id in payload.ranked_message_ids],
                    ranking_parent_id=str(payload.ranking_parent_id) if payload.ranking_parent_id else None,
                    message_tree_id=str(payload.message_tree_id) if payload.message_tree_id else None,
                    not_rankable=payload.not_rankable if payload.not_rankable else None,
                )
            case _:
                raise ValueError(f"Unknown payload type {reaction.payload_type}")