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

Code before:

            if num_red_flag is not None and num_red_flag >= self.cfg.auto_mod_red_flags:
                if m.parent_id is None:
                    logger.warning(
                        f"[AUTO MOD] Halting tree {m.message_tree_id}, initial prompt got too many red flags ({m.emojis})."
                    )
                    self.enter_low_grade_state(m.message_tree_id)
                else:
                    logger.warning(f"[AUTO MOD] Deleting message {m.id=}, it received too many red flags ({m.emojis}).")
                    self.pr.mark_messages_deleted(m.id, recursive=True)

            num_skip_reply = m.emojis.get(protocol_schema.EmojiCode.skip_reply)
            if num_skip_reply is not None and num_skip_reply >= self.cfg.auto_mod_max_skip_reply:
                logger.warning(
                    f"[AUTO MOD] Halting tree {m.message_tree_id} due to high skip-reply count of message {m.id=} ({m.emojis})."
                )
                self.halt_tree(m.id, halt=True)

    def determine_task_availability(self, lang: str) -> dict[protocol_schema.TaskRequestType, int]:
        self.pr.ensure_user_is_enabled()

        if not lang:
            lang = "en"
            logger.warning("Task availability request without lang tag received, assuming lang='en'.")

        init_prompt_disabled_langs = os.getenv("INIT_PROMPT_DISABLED_LANGS", "none")
        init_prompt_disabled_langs = init_prompt_disabled_langs.split(",")
        if lang in init_prompt_disabled_langs:
            num_missing_prompts = 0
        else:
            num_missing_prompts = self._prompt_lottery(lang=lang, max_activate=1)

        self._auto_moderation(lang=lang)
        extendible_parents, _ = self.query_extendible_parents(lang=lang)
        prompts_need_review = self.query_prompts_need_review(lang=lang)
        replies_need_review = self.query_replies_need_review(lang=lang)
        incomplete_rankings = self.query_incomplete_rankings(lang=lang)

        return self._determine_task_availability_internal(
            num_missing_prompts=num_missing_prompts,
            extendible_parents=extendible_parents,
            prompts_need_review=prompts_need_review,
            replies_need_review=replies_need_review,
            incomplete_rankings=incomplete_rankings,
        )

    @staticmethod
    def _get_label_descriptions(valid_labels: list[TextLabels]) -> list[protocol_schema.LabelDescription]:
        return [
            protocol_schema.LabelDescription(
                name=l.value, widget=l.widget.value, display_text=l.display_text, help_text=l.help_text

Review comment (refactor):
Could you move this to the [TreeManagerConfiguration](https://github.com/LAION-AI/Open-Assistant/blob/main/backend/oasst_backend/config.py)? The setting then still can be set via env-variable and this setting would be handled in the same way all other TreeManager settings are handled.

Code after:

            if num_red_flag is not None and num_red_flag >= self.cfg.auto_mod_red_flags:
                if m.parent_id is None:
                    logger.warning(
                        f"[AUTO MOD] Halting tree {m.message_tree_id}, initial prompt got too many red flags ({m.emojis})."
                    )
                    self.enter_low_grade_state(m.message_tree_id)
                else:
                    logger.warning(f"[AUTO MOD] Deleting message {m.id=}, it received too many red flags ({m.emojis}).")
                    self.pr.mark_messages_deleted(m.id, recursive=True)

            num_skip_reply = m.emojis.get(protocol_schema.EmojiCode.skip_reply)
            if num_skip_reply is not None and num_skip_reply >= self.cfg.auto_mod_max_skip_reply:
                logger.warning(
                    f"[AUTO MOD] Halting tree {m.message_tree_id} due to high skip-reply count of message {m.id=} ({m.emojis})."
                )
                self.halt_tree(m.id, halt=True)

    def determine_task_availability(self, lang: str) -> dict[protocol_schema.TaskRequestType, int]:
        self.pr.ensure_user_is_enabled()

        if not lang:
            lang = "en"
            logger.warning("Task availability request without lang tag received, assuming lang='en'.")

        if lang in self.cfg.init_prompt_disabled_langs_list:
            num_missing_prompts = 0
        else:
            num_missing_prompts = self._prompt_lottery(lang=lang, max_activate=1)

        self._auto_moderation(lang=lang)
        extendible_parents, _ = self.query_extendible_parents(lang=lang)
        prompts_need_review = self.query_prompts_need_review(lang=lang)
        replies_need_review = self.query_replies_need_review(lang=lang)
        incomplete_rankings = self.query_incomplete_rankings(lang=lang)

        return self._determine_task_availability_internal(
            num_missing_prompts=num_missing_prompts,
            extendible_parents=extendible_parents,
            prompts_need_review=prompts_need_review,
            replies_need_review=replies_need_review,
            incomplete_rankings=incomplete_rankings,
        )

    @staticmethod
    def _get_label_descriptions(valid_labels: list[TextLabels]) -> list[protocol_schema.LabelDescription]:
        return [
            protocol_schema.LabelDescription(
                name=l.value, widget=l.widget.value, display_text=l.display_text, help_text=l.help_text
            )
            for l in valid_labels