"""Provide modal runner for ChatGPT functionality, including prompt generation and moderation.
It spins up two stubs: ChatGptStub and ConductorStub.
ChatGptStub is responsible for moderating prompts and lyrics, and generating prompts from gpt_description_prompts.
ConductorStub is responsible for distributing queue items to the appropriate model worker.
"""

import json
import logging
import random
import os
import time
from collections import Counter
from collections.abc import Callable
from typing import Any, Dict, Optional

import modal
from openai import OpenAI  # type: ignore[attr-defined]

from suno_utils.worker.settings import bedrock_client
from suno_utils.gpt import chirp_v2
from suno_utils.utils.image_extraction import ImageExtractionError, ImageModerationFailure
from suno_utils.utils.constants import SUNO_AD
from suno_utils.worker.generate_song_lyrics import (
    LyricsLength,
    ModerationError,
    ModerationFailure,
    ModerationReroll,
    ModerationResult,
    ModerationSuccess,
    assemble_twitter_to_song_prompt,
    classify_text_lang,
    does_text_contain_slur,
    get_moderation_result,
    get_stanzas,
    moderate_gpt_description_prompt,
    moderate_user_inputs,
)
from suno_utils.worker.generate_song_lyrics import (
    get_prompt_from_gpt_description_prompt as get_lyrics_prompt_from_gpt_description_prompt,
)
from suno_utils.worker.loader import S3Loader
from suno_utils.worker.lyrics_client import LyricsClient
from suno_utils.worker.modal_base import get_modal_base_image
from suno_utils.worker.modal_model_configuration import DEFAULT_API_TO_MODEL_MAP, EXPERIMENTAL_INFO
from suno_utils.worker.schema import HistoryPrompt, QueueItem
from suno_utils.worker.tracing import distributed_trace, serialize_context, tracer
from suno_utils.worker.vertex_client import make_vertex_client
from suno_utils.worker.copyright_detector import CopyrightDetector
from suno_utils.worker.tag_augmentation import upsample_prompt, downsample_prompt
from suno_utils.worker.modal_model_volume import (
    MODEL_STORE_VOLUME_DIR,
    model_store_volume,
    MODEL_STORE_VOLUME_PREFIX,
    FASTTEXT_CKPT_PATH,
)

CONDUCTOR_DDOG_SERVICE = "conductor-worker"
CHATGPT_DDOG_SERVICE = "chatgpt-worker"

openai_secret = modal.Secret.from_name("openai-secret")
together_secret = modal.Secret.from_name("together-secret")

logger = logging.getLogger(__name__)
logging.basicConfig()
logger.setLevel(logging.DEBUG)

############## CHANGE THESE ##############
DEPLOYMENT_TYPE = "dev"
##########################################
DEPLOYMENT_TYPES = {"dev", "prod"}
assert DEPLOYMENT_TYPE in DEPLOYMENT_TYPES
# orchestrator talks to studio api
APP_NAME = f"orchestrator-{DEPLOYMENT_TYPE}"
assert APP_NAME.endswith(DEPLOYMENT_TYPE)

ResponseDict = dict[str, Any]

if DEPLOYMENT_TYPE == "dev":
    from suno_utils.worker.modal_model_configuration import MODEL_DEV_FNS as MODAL_FNS
elif DEPLOYMENT_TYPE == "prod":
    from suno_utils.worker.modal_model_configuration import MODEL_PROD_FNS as MODAL_FNS
else:
    err_msg = f"DEPLOYMENT_TYPE {DEPLOYMENT_TYPE} unrecognized"
    raise AssertionError(err_msg)

# For dev, we use fixed random seed; for prod, we use random seed based on time
random_seed = 42 if DEPLOYMENT_TYPE == "dev" else int((time.time() * 1000) % 100000000)
print("Random seed set to:", random_seed)
random.seed(random_seed)

# events queue for streaming events
EVENTS_QUEUE_NAME = f"events-queue-{DEPLOYMENT_TYPE}"

DD_SAMPLE_RATE = "1" if DEPLOYMENT_TYPE == "dev" else "0.0001"
# model for unit tests
TEST_MODEL_NAME = "chirp-v3-engine-i"
# models that check for copyright infringement
# we turn on it for all on Jan 23 2025
COPYRIGHT_INFRINGEMENT_MODELS = ["chirp-v3-5-tau", "chirp-v4-tau"]
# modal traffic lookup
MODAL_TRAFFIC_WORKER = "engine-chirpv2_engine_13b_special_8"
# check if the workers exist and fail early
# since these workers are used for worker inits
# if sth is wrong this fails at deploytime
COPYRIGHT_INFRINGEMENT_THRESHOLD = 0.9
BIT_ENABLE_CHECK_FOR_LYRICS_COPYRIGHT_IN_GENERATION = 2
BIT_BYPASS_LYRICS_MODERATION = 4
# looks up the H100 traffic
_ = modal.Cls.from_name(
    f"{MODAL_TRAFFIC_WORKER}_" + f"{'dev' if DEPLOYMENT_TYPE == 'dev' else 'prod'}_240s_H100",
    "ChirpV2Stub",
)().generate
_ = modal.Cls.from_name(
    f"history_encoder-{'dev' if DEPLOYMENT_TYPE == 'dev' else 'prod'}",
    "HistoryEncoderStub",
)().encode_history
_ = modal.Cls.from_name(
    f"sdxl-{'dev' if DEPLOYMENT_TYPE == 'dev' else 'prod'}",
    "StableDiffusion",
)().generate_image_item


def _is_base_v3_model_name(model_name):
    if model_name is None:
        return False
    v3p5_model_prefixes = ["chirp-v3p5", "chirp-v3-5"]
    return "chirp-v3" in model_name and not any(
        model_name.startswith(prefix) for prefix in v3p5_model_prefixes
    )


def _get_lyrics_length_for_model_family(model_name: str | None) -> LyricsLength:
    """Infer the desired number of stanzas for the given model."""

    if model_name and "short" in model_name:  # this is a 30 sec model
        lyrics_length = LyricsLength.SHORT
    elif _is_base_v3_model_name(model_name):
        lyrics_length = LyricsLength.STANDARD
    else:
        lyrics_length = LyricsLength.LONG
    print("for model name:", model_name, "assigning lyrics_length:", lyrics_length)
    return lyrics_length


class ChatGptWorker(S3Loader):
    """Provide nofication functionality for ChatGptStub."""

    def preload(self) -> None:
        """Preload models."""
        start_time = time.time()
        print("MOUNT_PATH:", MODEL_STORE_VOLUME_PREFIX, "FASTTEXT_CKPT_PATH:", FASTTEXT_CKPT_PATH)
        # fast text is a special loading case, it needs the full path
        ckpt_path = os.path.join(MODEL_STORE_VOLUME_PREFIX, FASTTEXT_CKPT_PATH.lstrip("/"))
        print("ckpt_path:", ckpt_path)
        chirp_v2.load_fasttext_model(ckpt_path)

        finish_time = time.time()
        print(f"Preloading took {finish_time - start_time}s")


def _make_success_response_dict(item: QueueItem, prompt_type: str, text_prompt: str) -> ResponseDict:
    return {
        "id": item.id,
        prompt_type: text_prompt,
        "moderation_result": "success",
        "ok": 1,
    }


def _make_failure_response_dict(
    item: QueueItem,
    prompt_type: str,
    text_prompt: str,
    err_msg: str,
) -> ResponseDict:
    return {
        "id": item.id,
        "type": "error",
        "moderation_result": "failure",
        prompt_type: text_prompt,
        "error_message": err_msg,
        "ok": 0,
    }


def _make_reroll_response_dict(
    item: QueueItem,
    prompt_type: str,
    text_prompt: str,
    err_msg: str,
) -> ResponseDict:
    return {
        "id": item.id,
        "type": "error",
        "moderation_result": "reroll",
        prompt_type: text_prompt,
        "error_message": err_msg,
        "ok": 0,
    }


def _make_generic_error_response_dict(
    item: QueueItem,
    prompt_type: str,
    text_prompt: str,
    err_msg: str,
) -> ResponseDict:
    return {
        "id": item.id,
        "type": "error",
        prompt_type: text_prompt,
        "error_message": err_msg,
        "ok": 0,
    }


aws_secret = modal.Secret.from_name("studio-aws")
sample_rule = json.dumps([{"sample_rate": DD_SAMPLE_RATE}])
SECRETS = [
    aws_secret,
    modal.Secret.from_name("openai-secret"),
    modal.Secret.from_name("together-secret"),
    modal.Secret.from_dict(
        {
            "DD_SITE": "datadoghq.com",
            "DD_ENV": DEPLOYMENT_TYPE,
            "DD_SERVICE": "chatgpt-worker",
            "DD_LOGS_ENABLED": "true",
            "DD_TRACE_ENABLED": "true",
            "DD_TRACE_SAMPLING_RULES": sample_rule,
        },
    ),
    modal.Secret.from_name("datadog-metrics"),
    modal.Secret.from_name("api-callback-token"),
    modal.Secret.from_name("google-application-credentials-data"),
    modal.Secret.from_name("lyrics-mod-es-creds"),
]


base_image = (
    get_modal_base_image()
    .pip_install("together==1.4.0")
    .pip_install("elasticsearch>=8.14,<9")
    .pip_install("langdetect>=1,<2")
)


image = base_image.add_local_python_source("suno_utils", copy=False)

app = modal.App(
    APP_NAME,
    image=image,
    secrets=SECRETS,
)


class GenerationModerationSub:
    """General moderation sub class for moderation of arbitrary text."""

    @modal.enter()
    def on_modal_enter(self):
        self.openai_client = OpenAI()

    @tracer.wrap()
    def _moderate_text(
        self,
        queue_item_json: str,
        key: str,
        moderation_function: Callable[[OpenAI, str], ModerationResult],
    ) -> ResponseDict:
        """Moderate arbitrary text, returning a response for ChatGptWorker.notify_finish."""
        from suno_utils.worker.schema import QueueItem

        item = QueueItem.model_validate_json(queue_item_json)
        if (text_prompt := item.metadata.get(key)) is None:
            err_msg = f"item: {item} of type {type(item)} has no key `{key}`, cannot moderate text."
            item.notify_progress(
                _make_generic_error_response_dict(
                    item,
                    key,
                    text_prompt="ERROR: PROMPT NOT FOUND",
                    err_msg=err_msg,
                ),
            )
            raise RuntimeError(err_msg)

        try:
            moderation_result = (moderation_function)(self.openai_client, text_prompt)
        except Exception as e:
            import traceback

            response_dict = {"id": item.id, "ok": 0, "type": "error", "error_message": str(e)}
            item.notify_progress(response_dict)
            traceback.print_exc()
            raise

        match moderation_result:
            case ModerationSuccess():
                response_dict = _make_success_response_dict(item, key, text_prompt)
                item.notify_progress(response_dict)
            case ModerationFailure(err_msg):
                response_dict = _make_failure_response_dict(item, key, text_prompt, err_msg)
                item.notify_progress(response_dict)
            case ModerationReroll(err_msg):
                response_dict = _make_reroll_response_dict(item, key, text_prompt, err_msg)
                item.notify_progress(response_dict)
            case unrecognized:
                err_msg = f"Received unrecognized result from moderation: {unrecognized}; this is a bug."
                response_dict = _make_generic_error_response_dict(item, key, text_prompt, err_msg)
                item.notify_progress(response_dict)
                raise RuntimeError(err_msg)

        return response_dict


@app.cls(
    cpu=1.0,
    memory=1_000,
    min_containers=1 if DEPLOYMENT_TYPE == "dev" else 10,
    timeout=240,
    scaledown_window=240,
    secrets=SECRETS,
)
@modal.concurrent(max_inputs=2)
class TextModerationSub(GenerationModerationSub):
    """Moderate arbitrary text with ChatGPT."""

    # DEPRECATED!  TK

    @modal.enter()
    def on_modal_enter(self):
        self.openai_client = OpenAI()

    @modal.method()
    def moderate_text(self, queue_item_json: str) -> ResponseDict:
        """Moderate text with ChatGPT and report result to results queue."""
        return self._moderate_text(
            queue_item_json,
            key="text",
            moderation_function=get_moderation_result,
        )


@app.cls(
    cpu=1.0,
    memory=1_000,
    min_containers=1 if DEPLOYMENT_TYPE == "dev" else 10,
    timeout=240,
    scaledown_window=240,
    secrets=SECRETS,
)
@modal.concurrent(max_inputs=2)
class TextModerationStub(GenerationModerationSub):
    """Moderate arbitrary text with ChatGPT."""

    @modal.enter()
    def on_modal_enter(self):
        self.openai_client = OpenAI()

    @modal.method()
    def moderate_text(self, queue_item_json: str) -> ResponseDict:
        """Moderate text with ChatGPT and report result to results queue."""
        return self._moderate_text(
            queue_item_json,
            key="text",
            moderation_function=get_moderation_result,
        )


@app.cls(
    cpu=1.0,
    memory=1_000,
    min_containers=1 if DEPLOYMENT_TYPE == "dev" else 10,
    timeout=240,
    scaledown_window=240,
    secrets=SECRETS,
)
@modal.concurrent(max_inputs=2)
class PromptUpsamplerStub:
    """Upsample a prompt for semantic only model (auk+) style descriptions"""

    @modal.enter()
    def on_modal_enter(self):
        self.openai_client = OpenAI()

    @modal.method()
    def upsample(self, prompt: str, is_instrumental: bool = False, check_artist_names=False) -> str:
        """Moderate text with ChatGPT and report result to results queue."""
        print("upsampling:", prompt)
        result = upsample_prompt(
            self.openai_client, prompt, is_instrumental, check_artist_names=check_artist_names
        )
        print("result:", result)
        # TODO(pat): this is rollover cruft, delete later
        if check_artist_names is False:
            return result
        else:
            if isinstance(result, str):
                return {"ok": True, "result": result}
            elif isinstance(result, ModerationFailure):
                return {"ok": False, "error_message": result.err_msg}


def _truncate_downsampled_prompt(prompt: str, max_length=100):
    if len(prompt) < max_length:
        return prompt
    truncated_tags = []
    current_length = 0
    prompt_tags = [tag.strip() for tag in prompt.split(",")]
    for tag in prompt_tags:
        # account for the length of the separator if not the first tag
        new_length = current_length + len(tag) + (0 if not truncated_tags else 2)
        if new_length <= max_length:
            truncated_tags.append(tag)
            current_length = new_length
        else:
            break
    return ", ".join(truncated_tags)


@app.cls(
    cpu=1.0,
    memory=1_000,
    min_containers=1 if DEPLOYMENT_TYPE == "dev" else 100,
    timeout=240,
    scaledown_window=240,
    secrets=SECRETS,
)
@modal.concurrent(max_inputs=2)
class PromptDownsamplerStub:
    """Downsample a prompt to extract genre tags"""

    @modal.enter()
    def on_modal_enter(self):
        self.openai_client = OpenAI()

    @modal.method()
    def downsample(self, prompt: str) -> str | None:
        """Downsample a prompt to extract genre tags."""
        return downsample_prompt(self.openai_client, prompt)

    @modal.method()
    def downsample_tags_for_item(
        self, queue_item_json: str, parent_context: Optional[Dict[str, Any]] = None
    ):
        """Downsample a prompt to extract genre tags."""
        from suno_utils.worker.schema import QueueItem

        item = QueueItem.model_validate_json(queue_item_json)
        logger.debug(queue_item_json)
        prompt = item.metadata.get("tags")
        logger.debug(prompt)
        if prompt is None or len(prompt) == 0:
            logger.debug("returning early due to empty prompt")
            return

        downsampled_prompt = downsample_prompt(self.openai_client, prompt)
        if downsampled_prompt is not None:
            cleaned_downsampled_prompt = _truncate_downsampled_prompt(downsampled_prompt)
            for id in item.all_child_clip_ids:
                payload = {
                    "id": id,
                    "type": "downsampled_tags",
                    "downsampled_tags": cleaned_downsampled_prompt,
                }
                logger.debug(f"sending payload to callback: {payload}")
                item.notify_progress(payload)


@app.cls(
    cpu=1.0,
    memory=1_000,
    min_containers=1 if DEPLOYMENT_TYPE == "dev" else 40,
    timeout=240,
    scaledown_window=240,
    secrets=SECRETS,
    volumes={MODEL_STORE_VOLUME_DIR: model_store_volume},
)
@modal.concurrent(max_inputs=2)
class PromptGenerationStub(GenerationModerationSub):
    """Generate a prompt from a GPT description prompt."""

    @modal.enter()
    def on_modal_enter(self):
        from together import Together

        self.openai_client = OpenAI()
        together_client = Together()
        self.worker = ChatGptWorker()
        self.worker.preload()

        vertex_client = make_vertex_client()

        self.lyrics_client = LyricsClient(
            self.openai_client, together_client, vertex_client, bedrock_client
        )
        self.text_lang_model = chirp_v2.text_lang_model
        if self.text_lang_model is None:
            msg = "chirp_v2.text_lang_model was not loaded successfully."
            raise RuntimeError(msg)
        print("Finished loading lang model.")

        self.events_queue = modal.Queue.from_name(EVENTS_QUEUE_NAME, create_if_missing=True)
        print("Finished init!")

    @modal.method()
    def get_prompt_from_gpt_description_prompt(self, queue_item_json: str) -> ResponseDict:
        """Given gpt_description_prompt in queue_item_json, moderate and return prompt if passing."""
        print("queue_item in get_prompt_from...", queue_item_json)
        item = QueueItem.model_validate_json(queue_item_json)
        item_without_callback = item.copy(exclude={"callback_url"}, deep=True)
        mod_response = self._moderate_text(
            item_without_callback.model_dump_json(),
            key="gpt_description_prompt",
            moderation_function=moderate_gpt_description_prompt,
        )
        if not (mod_response.get("ok") and mod_response.get("moderation_result") == "success"):
            mod_response["type"] = "generate_prompt"
            item.notify_progress(mod_response)
            return mod_response
        # else, gpt_description_prompt passed moderation
        gpt_description_prompt = item.metadata["gpt_description_prompt"]
        make_instrumental = item.metadata.get("make_instrumental", False)
        lyrics_length = _get_lyrics_length_for_model_family(item.model_name)
        lyrics_model = item.metadata.get("lyrics_model", "default")
        use_long_genre_description = False  # Turn this off for now
        print("use_long_genre_description:", use_long_genre_description)
        print("lyrics model in get prompt from gpt description prompt:", lyrics_model)
        need_to_check_generated_lyrics_for_copyright = not lyrics_model.startswith("default")
        are_lyrics_copyrighted_func = (
            CopyrightDetectorStub().are_lyrics_copyrighted.remote
            if need_to_check_generated_lyrics_for_copyright
            else None
        )

        try:
            title, prompt, genre_tags, language = get_lyrics_prompt_from_gpt_description_prompt(
                self.lyrics_client,
                self.text_lang_model,
                gpt_description_prompt,
                are_lyrics_copyrighted_func=are_lyrics_copyrighted_func,
                make_instrumental=make_instrumental,
                lyrics_length=lyrics_length,
                lyrics_model=lyrics_model,
                events_queue=self.events_queue,
                item_id=item.id,
                use_long_genre_description=use_long_genre_description,
                forced_lang=item.metadata.get("gpt_lang", None),
            )
        except ModerationError as e:
            response_dict = _make_generic_error_response_dict(
                item,
                prompt_type="gpt_description_prompt",
                text_prompt=gpt_description_prompt,
                err_msg=str(e),
            )
            response_dict["type"] = "generate_prompt"
            self.events_queue.put(
                {
                    "type": "error",
                    "data": response_dict,
                },
                partition=item.id,
                partition_ttl=60,
            )
            item.notify_progress(response_dict)
            return response_dict
        except Exception as e:
            response_dict = _make_generic_error_response_dict(
                item,
                prompt_type="gpt_description_prompt",
                text_prompt=gpt_description_prompt,
                err_msg="Unexpected Error:" + str(e),
            )
            response_dict["type"] = "generate_prompt"
            self.events_queue.put(
                {
                    "type": "error",
                    "data": response_dict,
                },
                partition=item.id,
                partition_ttl=60,
            )
            item.notify_progress(response_dict)
            return response_dict

        self.events_queue.put(
            {
                "type": "lyrics",
                "data": {
                    "title": title,
                    "prompt": prompt,
                    "genre_tags": genre_tags,
                    "lang": language,
                },
            },
            partition=item.id,
            partition_ttl=60,
        )
        response_dict = {
            "id": item.id,
            "type": "generate_prompt",
            "title": title,
            "prompt": prompt,
            "genre_tags": genre_tags,
            "ok": 1,
        }
        item.notify_progress(response_dict)
        return response_dict


@app.cls(
    cpu=1.0,
    memory=1_000,
    min_containers=1 if DEPLOYMENT_TYPE == "dev" else 100,
    timeout=240,
    scaledown_window=240,
    secrets=SECRETS,
    volumes={MODEL_STORE_VOLUME_DIR: model_store_volume},
)
@modal.concurrent(max_inputs=2)
class ChatGptStub(GenerationModerationSub):
    """Orchestrate lyrics generation and moderation."""

    @modal.enter()
    def on_modal_enter(self):
        """Set up ChatGptStub."""
        print("Start stub init.")

        from together import Together

        self.openai_client = OpenAI()
        together_client = Together()
        vertex_client = make_vertex_client()

        self.lyrics_client = LyricsClient(
            self.openai_client, together_client, vertex_client, bedrock_client
        )
        self.worker = ChatGptWorker()
        self.worker.preload()
        self.text_lang_model = chirp_v2.text_lang_model
        if self.text_lang_model is None:
            msg = "chirp_v2.text_lang_model was not loaded successfully."
            raise RuntimeError(msg)
        print("Finished loading lang model.")

        # Initialize events queue for streaming events
        self.events_queue = modal.Queue.from_name(EVENTS_QUEUE_NAME, create_if_missing=True)

        # save some modal functions -- for the main chirp worker
        self.modal_f_image_extraction = modal.Cls.from_name(
            f"image-extraction-{'dev' if DEPLOYMENT_TYPE == 'dev' else 'prod'}",
            "ImageExtractionStub",
        )().generate_image_description

        self.modal_f_image_extraction_from_url = modal.Cls.from_name(
            f"image-extraction-{'dev' if DEPLOYMENT_TYPE == 'dev' else 'prod'}",
            "ImageExtractionStub",
        )().generate_image_description_from_image_url

        self.modal_f_search_tavily = modal.Cls.from_name(
            f"brave-search-{'dev' if DEPLOYMENT_TYPE == 'dev' else 'prod'}",
            "SearchWorker",
        )().search_with_tavily

        self.modal_compose_search_query = modal.Cls.from_name(
            f"brave-search-{'dev' if DEPLOYMENT_TYPE == 'dev' else 'prod'}",
            "SearchWorker",
        )().compose_search_query

        # Add image generation modal functions
        self.modal_f_image_generator = modal.Cls.from_name(
            f"sdxl-{'dev' if DEPLOYMENT_TYPE == 'dev' else 'prod'}",
            "StableDiffusion",
        )().generate_image_item

        self.modal_f_pro_image_generator = modal.Cls.from_name(
            f"flux-{'dev' if DEPLOYMENT_TYPE == 'dev' else 'prod'}",
            "StableDiffusion",
        )().generate_image_item

        self.modal_f_lyrics_mod = modal.Cls.from_name(
            f"orchestrator-{'dev' if DEPLOYMENT_TYPE == 'dev' else 'prod'}", "CopyrightDetectorStub"
        )().are_lyrics_copyrighted

        print("Finished init!")

    def _predict_lyrics_language(self, lyrics: str) -> str:
        """Moderate lyrics with ChatGPT and report result to results queue."""
        # Could turn it into an API function in the future if needed.
        lyrics_chunks = [chunk.replace("\n", " ").strip() for chunk in lyrics.split("\n\n")]
        # keep only the meaningful ones
        lyrics_chunks = [chunk for chunk in lyrics_chunks if len(chunk) >= 5]
        language_chunks = [classify_text_lang(self.text_lang_model, chunk) for chunk in lyrics_chunks]

        # Count occurrences of each language using Counter
        language_counts = Counter(lang for lang in language_chunks if lang is not None)

        # Find the most common language(s)
        if language_counts:
            most_common = language_counts.most_common(1)
            if most_common:
                most_common_language, count = most_common[0]
                # Check if there are other languages with the same count
                if count / len(language_counts) > 0.5:  # Multiple languages with the same highest count
                    return most_common_language
        return ""  # No languages found is empty string

    def _get_text_from_all_images(
        self, item: QueueItem, image_to_song_s3_ids: list[str]
    ) -> str | ImageModerationFailure | ImageExtractionError:
        """Run image extraction modal runner on all images, save first image as clip cover image, and return descriptions string."""
        print("running image extraction")
        # run image extraction modal runner on each image, and notify failed clip if error
        try:
            images_description = self.modal_f_image_extraction.remote(
                QueueItem(
                    id=item.id,
                    metadata={
                        "image_to_song_s3_ids": image_to_song_s3_ids,
                        "user_context": item.metadata.get("gpt_description_prompt", ""),
                    },
                ).model_dump_json()
            )
        # network errors in image extraction pipeline
        except Exception as e:
            # modal worker tried 3 times
            return ImageExtractionError(e)

        # return text description for ChatGPTStub's lyric generation
        return images_description

    def _get_text_from_image_url(
        self, item: QueueItem
    ) -> str | ImageModerationFailure | ImageExtractionError:
        """Get text from image url."""
        try:
            image_description = self.modal_f_image_extraction_from_url.remote(item.model_dump_json())
        except Exception as e:
            return ImageExtractionError(e)
        return image_description

    def _save_user_cover_image(self, item: QueueItem, s3_id: str):
        """Save the  image to the clip cover image."""

        if item.all_child_clip_ids:
            # TODO: do we need a different image for each clip for stems?
            for clip_id in item.all_child_clip_ids:
                clip_queue_item = QueueItem(id=clip_id, metadata={"image_s3_id": s3_id})

                # save different size copies to s3 for display and video
                self.worker.copy_and_upload_image(clip_queue_item)

                # send progress to finish-clip using callback
                print("sending image progress now")
                item.notify_progress(
                    {
                        "id": clip_id,
                        "type": "image",
                        "image_id": f"image_{clip_id}",  # TODO: should this be s3_id?
                    },
                )

        pass

    def _maybe_process_image_to_song(self, item: QueueItem) -> None:
        """If an image_to_song gen, extract text and generate GPT description prompt for item.  Mutates argument."""
        image_to_song_s3_ids = item.metadata.get("image_to_song_s3_ids")
        if not image_to_song_s3_ids:
            return
        tick = time.time()

        # call OpenAI Vision to get description of image
        description = self._get_text_from_all_images(item, image_to_song_s3_ids)

        # catch image related failures and notify user
        if isinstance(description, ImageModerationFailure):
            err_msg = "We couldn't generate a song for this image. Please try another."
            item.notify_progress(
                {
                    "id": item.id,
                    "type": "error",
                    "error_type": "image_moderation_failure",
                    "error_message": err_msg,
                },
            )
            raise ModerationError(err_msg)
        elif isinstance(description, ImageExtractionError):
            err_msg = f"Image extraction failed: {description}"
            item.notify_progress(
                {
                    "id": item.id,
                    "type": "error",
                    "error_type": "image_extraction_failure",
                    "error_message": err_msg,
                },
            )
            raise ImageExtractionError(err_msg)

        # save first image as display image for each clip
        first_image_s3_id = image_to_song_s3_ids[0]
        self._save_user_cover_image(item, first_image_s3_id)

        # combine image description with user context to make lyric prompt
        item.metadata["gpt_description_prompt"] = (
            (item.metadata.get("gpt_description_prompt", "") or "") + "\n" + description
        )
        tock = time.time()
        print(f"Image to song EXTRACTION took {tock - tick} seconds")

    def _maybe_process_video_to_song(self, item: QueueItem) -> None:
        """If video to song, we pass in the video description as part of the prompt.  Mutates argument."""
        if video_to_song_description := item.metadata.get("video_to_song_description"):
            item.metadata["gpt_description_prompt"] = (
                (item.metadata.get("gpt_description_prompt", "") or "")
                + "\n"
                + video_to_song_description
            )

    def _maybe_process_twitter_to_song(self, item: QueueItem) -> None:
        """If twitter to song, we pass in the twitter content as part of the prompt.  Mutates argument."""
        image_description, search_results = None, None
        if twitter_to_song_content := item.metadata.get("twitter_to_song_content"):
            item.metadata["gpt_description_prompt"] = (
                (item.metadata.get("gpt_description_prompt", "") or "") + "\n" + twitter_to_song_content
            )
            item.metadata["is_twitter_to_song"] = True
            if item.metadata.get("twitter_image_url"):
                image_description = self._get_text_from_image_url(item)
                if isinstance(image_description, ImageModerationFailure):
                    err_msg = "We couldn't generate a song for this image. Please try another."
                    item.notify_progress(
                        {
                            "id": item.id,
                            "type": "error",
                            "error_type": "image_moderation_failure",
                            "error_message": err_msg,
                        },
                    )
                    raise ModerationError(err_msg)
                elif isinstance(image_description, ImageExtractionError):
                    err_msg = f"Image extraction failed: {image_description}"
                    item.notify_progress(
                        {
                            "id": item.id,
                            "type": "error",
                            "error_type": "image_extraction_failure",
                            "error_message": err_msg,
                        },
                    )
                    raise ImageExtractionError(err_msg)
                item.metadata["gpt_description_prompt"] = (
                    (item.metadata.get("gpt_description_prompt", "") or "")
                    + "\n Tweet Image Description: "
                    + image_description
                )
            try:
                search_query = self.modal_compose_search_query.remote(twitter_to_song_content)
                search_results = self.modal_f_search_tavily.remote(search_query, max_results=3, days=7)
            except Exception as e:
                print(f"[Maybe process twitter to song] Error search online: {e}")
                search_results = None
            item.metadata["gpt_description_prompt"] = assemble_twitter_to_song_prompt(
                item.metadata.get("gpt_description_prompt", ""),
                twitter_to_song_content,
                image_description,
                search_results,
            )

    def _start_early_image_generation(self, early_item, title, tags):
        """Helper function to start image generation, handling backlog logic consistently.

        Does not return anything.
        """
        # Only start early image generation if we have BOTH title AND tags
        if not title or not tags:
            return

        # This should only run if it's not an image-to-song (where we already have an image)
        if early_item.is_image_to_song:
            return

        print(f"Starting early image generation with title '{title}' and tags '{tags}'")

        # Mark that we've already started image generation to avoid duplicate runs
        early_item.metadata["image_generation_started"] = True

        # Queue up image generation task
        if isinstance(early_item.model_name, str) and (
            "auk" in early_item.model_name
            or "bluejay" in early_item.model_name
            or "v4" in early_item.model_name
        ):
            try:
                n_backlog = self.modal_f_pro_image_generator.get_current_stats().backlog
                if n_backlog < 25:  # reduce the backlog to 25
                    _ = self.modal_f_pro_image_generator.spawn(
                        early_item.model_dump_json(),
                        parent_context=serialize_context(),
                    )
                else:
                    _ = self.modal_f_image_generator.spawn(
                        early_item.model_dump_json(),
                        parent_context=serialize_context(),
                    )
            except Exception as e:
                print(f"Error starting early image generation: {e}")
        else:
            try:
                _ = self.modal_f_image_generator.spawn(
                    early_item.model_dump_json(), parent_context=serialize_context()
                )
            except Exception as e:
                print(f"Error starting early image generation: {e}")

    @modal.method()
    @distributed_trace("moderate_and_generate", CHATGPT_DDOG_SERVICE, env_name=DEPLOYMENT_TYPE)
    def moderate_and_generate(self, queue_item_json: str, use_openai_mod_for_custom_lyrics=True):
        """Moderate and generate.

        This worker will go through all the moderation steps.
        If successful, will call the orchestrate_and_generate worker.
        """
        item = QueueItem(**json.loads(queue_item_json))
        override_fields = item.metadata.get("override_fields", [])
        print(
            f"Moderator: {item.metadata.get('ids')} is free: {item.is_free_generation}",
        )
        print("in moderate_and_generate with item:", item)
        use_long_genre_description = False
        print(
            "audio model name:",
            item.model_name,
            "use_long_genre_description:",
            use_long_genre_description,
        )
        # generate lyrics and tags
        lyrics_length = _get_lyrics_length_for_model_family(item.model_name)

        try:
            self._maybe_process_image_to_song(item)
            self._maybe_process_video_to_song(item)
            self._maybe_process_twitter_to_song(item)
        except:
            print(f"Got error processing image or video for clip {item.id}. Bailing on generation.")
            return

        using_custom_title = False
        is_twitter_to_song = item.metadata.get("is_twitter_to_song", False)
        if user_prompt := item.metadata.get("gpt_description_prompt"):
            try:
                lyrics_model = item.metadata.get("lyrics_model", "default")
                print("lyrics_model:", lyrics_model)
                are_lyrics_copyrighted_func = (
                    CopyrightDetectorStub().are_lyrics_copyrighted.remote
                    if not lyrics_model.startswith("default")
                    else None
                )

                # Define early title/tags callback function to start image generation early
                title_tags_early_callback = None

                # Only enable early callback if requested in metadata
                if item.metadata.get("extra", {}).get("early_callback", False):

                    def title_tags_early_callback(title: str, tags: str) -> None:
                        if not title:
                            return

                        # Update item with early title/tags
                        early_item = item.model_copy(deep=True)
                        early_item.title = title

                        # Process tags if available
                        if tags:
                            early_item.metadata["tags"] = tags

                        # Notify progress with early title and tags
                        if early_item.all_child_clip_ids:  # Request level item
                            for clip_id in early_item.all_child_clip_ids:
                                notify_data = {
                                    "id": clip_id,
                                    "title": title,
                                    "type": "early_title_tags",
                                }
                                if tags:
                                    notify_data["tags"] = tags

                                early_item.notify_progress(notify_data, blocking=False)
                        else:  # Clip level item
                            notify_data = {
                                "id": early_item.id,
                                "title": title,
                                "type": "early_title_tags",
                            }
                            if tags:
                                notify_data["tags"] = tags

                            early_item.notify_progress(notify_data, blocking=False)

                        # Start image generation if we have both title and tags (using helper)
                        self._start_early_image_generation(early_item, title, tags)

                use_moderation_for_gpt_description_prompt = False if is_twitter_to_song else True
                title, lyrics, genre_tags, language = get_lyrics_prompt_from_gpt_description_prompt(
                    self.lyrics_client,
                    chirp_v2.text_lang_model,
                    user_prompt,
                    use_moderation_for_gpt_description_prompt=use_moderation_for_gpt_description_prompt,
                    check_artist_names=item.metadata.get("check_artist_names", True),
                    make_instrumental=item.metadata.get("make_instrumental", False),
                    lyrics_length=lyrics_length,
                    is_image_or_video_to_song=item.is_image_to_song or item.is_video_to_song,
                    is_twitter_to_song=is_twitter_to_song,
                    lyrics_model=item.metadata.get("lyrics_model", "default"),
                    events_queue=self.events_queue,
                    item_id=item.id,
                    are_lyrics_copyrighted_func=are_lyrics_copyrighted_func,
                    use_long_genre_description=use_long_genre_description,
                    title_tags_callback=title_tags_early_callback,
                    forced_lang=item.metadata.get("gpt_lang", None),
                )

            except ModerationError as mod_error:
                import traceback

                # Write error to request-level event queue
                self.events_queue.put(
                    {
                        "type": "error",
                        "data": {
                            "error_type": "moderation_failure",
                            "error_message": str(mod_error),
                        },
                    },
                    partition=item.id,  # Request-level partition
                    partition_ttl=60,
                )

                # Write error to clip-level event queues
                for clip_id in item.all_child_clip_ids:
                    item.notify_progress(
                        {
                            "id": clip_id,
                            "type": "error",
                            "error_type": "moderation_failure",
                            "error_message": str(mod_error),
                        },
                    )
                    self.events_queue.put(
                        {
                            "type": "error",
                            "data": {
                                "error_type": "moderation_failure",
                                "error_message": str(mod_error),
                            },
                        },
                        partition=clip_id,
                        partition_ttl=60,
                    )
                traceback.print_exc()
                return
            except Exception as e:
                import traceback

                # Write error to request-level event queue
                self.events_queue.put(
                    {
                        "type": "error",
                        "data": {
                            "error_type": "moderation_failure",
                            "error_message": str(e),
                        },
                    },
                    partition=item.id,  # Request-level partition
                    partition_ttl=60,
                )

                # Write error to clip-level event queues
                for clip_id in item.all_child_clip_ids:
                    item.notify_progress(
                        {
                            "id": clip_id,
                            "type": "error",
                            "error_type": "moderation_failure",
                            "error_message": "Generic openAI error",
                        },
                    )
                    self.events_queue.put(
                        {
                            "type": "error",
                            "data": {
                                "error_type": "moderation_failure",
                                "error_message": str(e),
                            },
                        },
                        partition=clip_id,
                        partition_ttl=60,
                    )
                traceback.print_exc()
                return

            if item.title is None or item.title == "":
                # Use generated title if available, otherwise fall back to title_fallback from extra if provided
                if title:
                    item.title = title
                elif item.metadata.get("extra", {}).get("title_fallback"):
                    item.title = item.metadata["extra"]["title_fallback"]
                else:
                    item.title = title  # This will be None or empty string
            else:
                using_custom_title = True
            if item.prompt_text is None or item.prompt_text == "" or "prompt" not in override_fields:
                item.prompt_text = lyrics
            if (
                item.metadata.get("tags") is None
                or item.metadata.get("tags") == ""
                or "tags" not in override_fields
            ):
                item.metadata["tags"] = ", ".join(genre_tags)

            item.metadata["lang"] = language
            logger.info(f"generated {item.metadata['tags']} {lyrics}")

        if (
            user_prompt is None or len(override_fields) > 0 or using_custom_title
        ):  # there's no gpt_description prompt, so we're dealing with custom lyrics.
            # if using override fields or custom title, we need to moderate here
            # IN CASE OF openAI outage, moderation could be down, flag this to False
            # TDDO: backups and fallbacks on this logic
            moderate_lyrics = True

            tags = item.metadata.get("tags") or ""
            language = self._predict_lyrics_language(item.prompt_text or "")
            item.metadata["lang"] = language
            if moderate_lyrics:
                if use_openai_mod_for_custom_lyrics:
                    moderation_result = moderate_user_inputs(
                        self.openai_client, item.prompt_text or "", tags, item.title or ""
                    )
                    logger.debug(moderation_result)
                else:
                    text_to_be_moderated = "\n".join([item.title or "", item.prompt_text or "", tags])

                    moderation_result: ModerationResult = (
                        ModerationSuccess()
                        if not does_text_contain_slur(text_to_be_moderated)
                        else ModerationFailure("manually flagged")
                    )
                # In all cases where lyrics are provided, run lyrics mod v3
                try:
                    print(f"Checking for copyrighted lyrics in Elasticsearch!: {item.prompt_text}")
                    lyrics_are_copyrighted = (
                        self.modal_f_lyrics_mod.spawn(item.prompt_text).get(
                            timeout=30  # normally takes ~400ms but coldstarts on the copyright worker can take 15-30 seconds.
                        )
                        if item.prompt_text
                        else False
                    )
                    if lyrics_are_copyrighted:
                        print("Copyrighted lyrics were found in Elasticsearch.")
                        moderation_result = ModerationFailure("Lyrics contained copyrighted material")
                except Exception as e:
                    print(f"Error checking lyrics copyright: {e}")
                    import traceback

                    traceback.print_exc()
                    lyrics_are_copyrighted = False

                if not isinstance(moderation_result, ModerationSuccess):
                    # Write error to request-level event queue
                    self.events_queue.put(
                        {
                            "type": "error",
                            "data": {
                                "error_type": "moderation_failure",
                                "error_message": moderation_result.err_msg,
                            },
                        },
                        partition=item.id,  # Request-level partition
                        partition_ttl=60,
                    )

                    # Write error to clip-level
                    print("Broadcasting moderation failures")
                    for clip_id in item.all_child_clip_ids:
                        item.notify_progress(
                            {
                                "id": clip_id,
                                "type": "error",
                                "error_type": "moderation_failure",
                                "error_message": moderation_result.err_msg,
                            },
                        )
                        # Also add to clip-level events queue
                        self.events_queue.put(
                            {
                                "type": "error",
                                "data": {
                                    "error_type": "moderation_failure",
                                    "error_message": moderation_result.err_msg,
                                },
                            },
                            partition=clip_id,
                            partition_ttl=60,
                        )
                    return
        if item.metadata.get("make_instrumental", False):
            # make sure we set the langeuge to None if definitely instrumental
            item.metadata["lang"] = None

        # watermarking
        if item.metadata.get("is_copycat") or item.metadata.get("is_audio_watermarked"):
            item.prompt_text = f"{SUNO_AD} {item.prompt_text} {SUNO_AD}"

        # send the lyrics to the latency-sensitiveevents queue
        if item.metadata.get("gpt_description_prompt"):
            event_data = {
                "title": item.title,
                "prompt": item.prompt_text,
                "genre_tags": item.metadata["tags"],
                "lang": item.metadata["lang"],
            }
            self.events_queue.put(
                {
                    "type": "lyrics",
                    "data": event_data,
                },
                partition=item.id,
                partition_ttl=60,
            )
            if user_prompt := item.metadata.get("gpt_description_prompt"):
                if item.all_child_clip_ids:  # this is a request level item
                    for clip_id in item.all_child_clip_ids:
                        self.events_queue.put(
                            {
                                "type": "lyrics",
                                "data": event_data,
                            },
                            partition=clip_id,
                            partition_ttl=60,
                        )

        # at this stage, we can pass this onto orchestrator successfully
        _ = ConductorStub.redirect_and_generate.spawn(
            item.model_dump_json(), parent_context=serialize_context()
        )
        # we notify the update after we spawn to save 1 sec TODO(Chi): revisit if this is necessary when we use TPE for notify_progress
        if user_prompt := item.metadata.get("gpt_description_prompt"):
            # this notification needs to be done at clip level, since this is per clip lyrics?
            # TODO: only for now -- in the future, we should always notify progress on the clip ids
            if item.all_child_clip_ids:  # this is a request level item
                for clip_id in item.all_child_clip_ids:
                    item.notify_progress(
                        {
                            "id": clip_id,
                            "title": item.title,
                            "type": "lyrics",
                            "text": item.prompt_text,
                            "tags": item.metadata["tags"],
                            "lang": item.metadata["lang"],
                        },
                    )
            else:  # this is a clip level item
                item.notify_progress(
                    {
                        "id": item.id,
                        "title": item.title,
                        "type": "lyrics",
                        "text": item.prompt_text,
                        "tags": item.metadata["tags"],
                        "lang": item.metadata["lang"],
                    },
                )
        return item.model_dump_json()

    @modal.method()
    def generate_lyrics_and_genre_from_gpt_instructions(self, queue_item_json: str):
        #  I think this is mostly twitter [?]  PON
        item = QueueItem(**json.loads(queue_item_json))

        # Check that we have exactly one of tweet_context or gpt_description_prompt
        tweet_context = item.metadata.get("tweet_context")
        gpt_description_prompt = item.metadata.get("gpt_description_prompt")

        if bool(tweet_context) == bool(gpt_description_prompt):
            raise ValueError(
                "Must have exactly one of tweet_context or gpt_description_prompt in metadata"
            )

        # Determine if this is a twitter to song request
        is_twitter_to_song = bool(tweet_context)
        user_prompt = tweet_context if is_twitter_to_song else gpt_description_prompt

        if user_prompt:
            try:
                title, lyrics, genre_tags, language = get_lyrics_prompt_from_gpt_description_prompt(
                    self.lyrics_client,
                    chirp_v2.text_lang_model,
                    user_prompt,
                    use_moderation_for_gpt_description_prompt=False,
                    check_artist_names=item.metadata.get("check_artist_names", True),
                    make_instrumental=item.metadata.get("make_instrumental", False),
                    lyrics_length=LyricsLength.STANDARD,
                    is_image_or_video_to_song=item.is_image_to_song or item.is_video_to_song,
                    is_twitter_to_song=is_twitter_to_song,
                    lyrics_model=item.metadata.get("lyrics_model", "default"),
                    events_queue=self.events_queue,
                    item_id=item.id,
                    are_lyrics_copyrighted_func=None,
                    use_long_genre_description=True,
                    forced_lang=item.metadata.get("gpt_lang", None),
                )
                return {
                    "title": title,
                    "lyrics": lyrics,
                    "genre_tags": genre_tags,
                    "language": language,
                }
            except Exception as e:
                return {
                    "error": str(e),
                }
        return {
            "error": "No gpt_description_prompt found",
        }


@app.cls(
    cpu=1.0,
    memory=1_000,
    min_containers=1 if DEPLOYMENT_TYPE == "dev" else 50,
    timeout=240,
    scaledown_window=240,
    secrets=SECRETS,
    volumes={MODEL_STORE_VOLUME_DIR: model_store_volume},
)
@modal.concurrent(max_inputs=4)
class CopyrightDetectorStub:
    """General moderation sub class for moderation of arbitrary text."""

    @modal.enter()
    def on_modal_enter(self):
        self.copyright_detector = CopyrightDetector(
            os.environ["LYRICS_MOD_ES_URL"], os.environ["LYRICS_MOD_ES_KEY"], DEPLOYMENT_TYPE
        )

    @modal.method()
    def are_lyrics_copyrighted(self, lyrics: str, strict: bool = False) -> bool:
        """Determine whether lyrics likely infringe on copyright"""
        print(f"Checking copyright:\n----------\n{lyrics}\n----------")
        result = self.copyright_detector.are_lyrics_copyrighted(lyrics, strict)
        print("copyright result:", result)
        return result


##########################################
# Setup the conductor stub
@app.cls(
    cpu=1.0,
    memory=1_000,
    secrets=SECRETS,
    min_containers=2 if DEPLOYMENT_TYPE == "dev" else 100,
    timeout=240,
    retries=modal.Retries(
        max_retries=2,
        backoff_coefficient=2.0,
        initial_delay=5.0,
    ),
    scaledown_window=240,
)
@modal.concurrent(max_inputs=2)
class ConductorStub(S3Loader):
    """Expose ConductorStub worker as a modal app

    Conductor is responsible for distributing queue items to the appropriate chirp worker.
    It is also responsible for controlling the traffic split for the experimental models.
    Or conducting A/B testing for the hyperparameters.
    Think of the greats: Karajan, Bernstein, Abbado, Haitink, Jansons.
    """

    @modal.enter()
    def on_modal_enter(self):
        """Set up ChatGptStub."""
        # looks up the H100 traffic
        self.modal_f_chirp_main_H100 = modal.Cls.from_name(
            f"{MODAL_TRAFFIC_WORKER}_" + f"{'dev' if DEPLOYMENT_TYPE == 'dev' else 'prod'}_240s_H100",
            "ChirpV2Stub",
        )().generate
        self.modal_f_history_loader = modal.Cls.from_name(
            f"history_encoder-{'dev' if DEPLOYMENT_TYPE == 'dev' else 'prod'}",
            "HistoryEncoderStub",
        )().encode_history
        self.modal_f_image_generator = modal.Cls.from_name(
            f"sdxl-{'dev' if DEPLOYMENT_TYPE == 'dev' else 'prod'}",
            "StableDiffusion",
        )().generate_image_item
        self.modal_f_pro_image_generator = modal.Cls.from_name(
            f"flux-{'dev' if DEPLOYMENT_TYPE == 'dev' else 'prod'}",
            "StableDiffusion",
        )().generate_image_item
        self.modal_f_downsampler = modal.Cls.from_name(
            f"orchestrator-{'dev' if DEPLOYMENT_TYPE == 'dev' else 'prod'}", "PromptDownsamplerStub"
        )().downsample_tags_for_item
        # key is a tuple of worker name and function name
        self.modal_chirp_worker_lookup_cache = {}
        # events queue for streaming events
        self.events_queue = modal.Queue.from_name(EVENTS_QUEUE_NAME, create_if_missing=True)

    def _experimental_model_mixer(self, queue_item: QueueItem) -> QueueItem:
        """Mix the queue request under the hood. For experimental purpose only."""
        temp_queue_item = queue_item.model_copy(deep=True)
        # TODO: probably need some flag to indicate if this user enroll in abtesting or not
        # only a dedicated exp model name will trigger this
        # note that this is a totally fake model, only for traffic split purpose
        # this is experiment split at clip level
        stuidio_api_to_experiment_map = EXPERIMENTAL_INFO["api_model_to_experiment_map"]
        # for bots, do sth special :)
        # now that things are faster we can spare a bit of extra traffic into them
        # but not experiments since they don't label
        if temp_queue_item.is_bot_generation:
            print(f"Warning: bot generation detected. Item ids: {temp_queue_item.all_child_clip_ids}")
            stuidio_api_to_experiment_map = EXPERIMENTAL_INFO["api_model_to_bot_model_map"]
        # AUK infill is a special case -- we will do the mapping here
        if temp_queue_item.model_name == "chirp-auk" and temp_queue_item.is_infill:
            print(f"AUK infill detected. Item ids: {temp_queue_item.all_child_clip_ids}")
            temp_queue_item.model_name = "chirp-auk-infill"
        # Somehow I feel ppl are using the model directly to pass bots is that true?
        if DEPLOYMENT_TYPE == "prod":
            # prod should only have limited models to access
            is_invalid_model_name = (temp_queue_item.model_name not in DEFAULT_API_TO_MODEL_MAP) and (
                temp_queue_item.model_name not in ["chirp-v4-h-api", "chirp-v2-engine-msft-60s"]
            )
            # these are 30b features but the input model name is wrong
            special_30b_models = ["chirp-v3-5-tau", "chirp-v4-tau", "chirp-auk-infill"]
            special_task_models = special_30b_models + [
                "chirp-auk",
                "chirp-auk-o",
                "chirp-bluejay",
                "chirp-bluejay-o",
            ]
            is_unexpected_model_name = (temp_queue_item.model_name not in special_task_models) and (
                temp_queue_item.is_infill
                or temp_queue_item.is_cover_condition
                or temp_queue_item.is_artist_condition
            )
            is_unexpected_model_name_task = (temp_queue_item.model_name not in special_task_models) and (
                temp_queue_item.metadata.get("task") == "infill"
                or temp_queue_item.metadata.get("task") == "cover"
                or temp_queue_item.metadata.get("task") == "artist_consistency"
            )
            is_unexpected_30b_model_in_prod = (
                DEPLOYMENT_TYPE == "prod"
                and (temp_queue_item.model_name in special_30b_models)
                and not (
                    temp_queue_item.is_infill
                    or temp_queue_item.is_cover_condition
                    or temp_queue_item.is_artist_condition
                    # TODO: add extend
                )
            )
            if (
                is_invalid_model_name
                or is_unexpected_model_name
                or is_unexpected_model_name_task
                or is_unexpected_30b_model_in_prod
            ):
                # We will overwrite and give you something from bot model
                if is_invalid_model_name:  # this is more like a bot behavoir
                    print(
                        f"Warning {temp_queue_item.id}: invalid model name in prod: {temp_queue_item.model_name}"
                    )
                    temp_queue_item.model_name = "chirp-v3-5-b"
                elif is_unexpected_model_name:  # this is more like a suno bug
                    print(
                        f"Warning {temp_queue_item.id}: unexpected model name given task: {temp_queue_item.model_name}"
                    )
                    temp_queue_item.model_name = "chirp-v3-5-b"
                elif is_unexpected_model_name_task:
                    print(
                        f"Warning {temp_queue_item.id}: unexpected 30b task in prod: {temp_queue_item.model_name}"
                    )
                    temp_queue_item.model_name = "chirp-v3-5-b"
                elif is_unexpected_30b_model_in_prod:
                    print(
                        f"Warning {temp_queue_item.id}: unexpected 30b model access in prod: {temp_queue_item.model_name}"
                    )
                    # TODO: remove this if once we fix the IOS model name issue
                    if temp_queue_item.metadata.get("source") != "ios":
                        temp_queue_item.model_name = "chirp-v3-5-b"
                # reroute the traffic to bot traffic split
                stuidio_api_to_experiment_map = EXPERIMENTAL_INFO["api_model_to_bot_model_map"]

                # TODO: remove this once we fix the IOS model name issue
                # Cause IOS is bugged rollout on Feb 26, 2025
                if is_unexpected_30b_model_in_prod and temp_queue_item.metadata.get("source") == "ios":
                    if temp_queue_item.model_name == "chirp-v4-tau":
                        temp_queue_item.model_name = "chirp-v4"
                    elif temp_queue_item.model_name == "chirp-v3-5-tau":
                        temp_queue_item.model_name = "chirp-v3-5"
                    stuidio_api_to_experiment_map = EXPERIMENTAL_INFO["api_model_to_experiment_map"]
        # TODO: remove this once we move everything to request level queueItem
        if not temp_queue_item.all_child_clip_ids:
            print(f"Real studio API calls should contain ids!? {queue_item}")
            # This is probably not real API requests but ppl hacking our API
            # In this case, it will default to v3.0 launch :)
            experiment_frac_to_model_name = {
                1: "chirp-v3-engine-i",
            }
            # if API -- won't be in the map, will be passing through
            if temp_queue_item.model_name in stuidio_api_to_experiment_map:
                random_value = random.random()
                for experiment_model_split_frac, model_name in experiment_frac_to_model_name.items():
                    if random_value <= experiment_model_split_frac:
                        temp_queue_item.model_name = model_name
                        logger.debug(f"switched model to {model_name}: {random_value:.3f}")
                        # log this info in metadata
                        temp_queue_item.metadata["experiment"] = model_name
                        temp_queue_item.metadata["experiment_version"] = EXPERIMENTAL_INFO[
                            "experiment_version"
                        ]
                        break

            return temp_queue_item
        # let's do the experiment at request level! This needs to be done a bit more careful
        else:
            # we do not split exp for v2, for instance
            if (
                original_temp_queue_item_model_name := temp_queue_item.model_name
            ) in stuidio_api_to_experiment_map:
                experiment_frac_to_model_names = stuidio_api_to_experiment_map[
                    original_temp_queue_item_model_name
                ]
                # need experiment traffic split
                if len(temp_queue_item.ids or temp_queue_item.multi_ids) != 2:
                    raise ValueError("This is not a valid experiment request")
                random_value = random.random()
                # overwrite 30b experiment temperorily based on task
                # we will leave infill out for auk ab test now
                # if original_temp_queue_item_model_name == "chirp-v4-tau" and (temp_queue_item.is_infill):
                #     random_value = 1.0
                for experiment_model_split_frac, model_names in experiment_frac_to_model_names.items():
                    if random_value <= experiment_model_split_frac:
                        # note that the queue item model name is modified now!
                        temp_queue_item.model_name = model_names[0]  # set a place holder for now
                        # very important that we shuffle the model names here
                        random.shuffle(model_names)
                        temp_queue_item.metadata["experiment"] = model_names
                        temp_queue_item.metadata["experiment_version"] = EXPERIMENTAL_INFO[
                            "experiment_version"
                        ]
                        logger.debug(
                            f"switched models to {temp_queue_item.metadata['experiment']}: {random_value:.3f}"
                        )
                        break
                # For inference experiments
                random_value = random.random()
                inference_experiment = EXPERIMENTAL_INFO["inference_experiment"]
                meets_experimental_condition = (
                    original_temp_queue_item_model_name in inference_experiment
                    and len(temp_queue_item.metadata.get("experiment", [])) == 2  # two model names
                    and len(set(temp_queue_item.metadata.get("experiment", []))) == 1  # but same model
                    and (
                        temp_queue_item.metadata.get("control_sliders", None) is None
                    )  # will not stack slider exps
                )
                experiment_config_split_frac_map = {}
                if meets_experimental_condition:
                    exepriment_model_name = temp_queue_item.metadata["experiment"][0]
                    experiment_config_split_frac_map = inference_experiment[
                        original_temp_queue_item_model_name
                    ].get(exepriment_model_name, {})
                for (
                    experiment_config_split_frac,
                    (param_experiment_name, model_inference_configs),
                ) in experiment_config_split_frac_map.items():
                    if (
                        meets_experimental_condition
                        and experiment_config_split_frac > 0
                        and random_value < experiment_config_split_frac
                    ):
                        # each id will have it's own model configs
                        random.shuffle(model_inference_configs)
                        # task specific experiment filters -- only freedom for diff infill
                        if "freedom" in param_experiment_name and (not temp_queue_item.is_diff_infill):
                            break
                        # rename the exp -- note how we rename it to param_experiment, for logging
                        temp_queue_item.metadata["param_experiment"] = param_experiment_name
                        temp_queue_item.metadata["model_config"] = model_inference_configs
                        logger.debug(
                            f"{temp_queue_item.id}: Switched inference config to {temp_queue_item.metadata['model_config']}: {random_value:.3f}"
                        )
                        break
                meet_mask_control_slider_experimental_condition = (
                    original_temp_queue_item_model_name in inference_experiment
                    and len(temp_queue_item.metadata.get("experiment", [])) == 2  # two model names
                    and len(set(temp_queue_item.metadata.get("experiment", []))) == 1  # but same model
                    and (
                        temp_queue_item.metadata.get("control_sliders", None) is not None
                    )  # will not stack slider exps
                )
                # 10% of control slider mask
                control_slider_experiment_frac = 0.1
                if (
                    meet_mask_control_slider_experimental_condition
                    and random_value < control_slider_experiment_frac
                ):
                    model_inference_configs = [{"masked": True}, {}]
                    random.shuffle(model_inference_configs)
                    temp_queue_item.metadata["model_config"] = model_inference_configs
                    temp_queue_item.metadata["param_experiment"] = "mask_control_slider"
                    logger.debug(
                        f"{temp_queue_item.id}: Switched to mask_control_slider: {random_value:.3f}"
                    )

            return temp_queue_item

    def _notify_queue_status_with_epxerimental_info(self, item: QueueItem) -> None:
        """Send a queue status notification with the experimental info."""
        progress_info = {
            "id": item.id,
            "type": "queued",
        }
        # log the experimental model status
        if experiment_model := item.metadata.get("experiment"):
            progress_info["experiment"] = experiment_model
            progress_info["experiment_version"] = item.metadata.get("experiment_version", "")
            progress_info["model_name"] = item.model_name if item.model_name is not None else ""
        # log the experimental param status
        if epxeriment_request_param := item.metadata.get("param_experiment"):
            progress_info["param_experiment"] = epxeriment_request_param
        if language := item.metadata.get("lang"):
            progress_info["lang"] = language
        item.notify_progress(progress_info)
        # also notify the child clip ids, especially for stems
        for child_id in item.all_child_clip_ids:
            child_item = item.model_copy(deep=True, update={"id": child_id, "ids": None})
            child_item.metadata["multi_ids"] = []
            progress_info["id"] = child_id
            child_item.notify_progress(progress_info)

    def _get_gpt_worker(self, item: QueueItem) -> modal.Function:
        if item.model_name is None:
            raise ValueError("Model name is required to get the GPT worker.")
        stub_name, stub_function = MODAL_FNS[item.model_name]
        # only for the prod 3p5 engine
        if "30b" in stub_name:
            # v4 is on H100 only
            stub_name_worker = f"{stub_name}_H200"
        elif (
            item.model_name.startswith("chirp-v4-up")
            or item.model_name.startswith("chirp-seeds")
            or item.model_name.startswith("chirp-ahi")
        ):
            # upsample model don't have GPU types yet
            stub_name_worker = f"{stub_name}"
        else:
            # we have enough H100s so most other jobs are H100s now
            stub_name_worker = f"{stub_name}_H100"
        # we don't need to lookup it again if we have already done it
        if (stub_name_worker, stub_function) in self.modal_chirp_worker_lookup_cache:
            return self.modal_chirp_worker_lookup_cache[(stub_name_worker, stub_function)]
        else:
            try:
                cls_name = stub_function.split(".")[0]
                func_name = stub_function.split(".")[1]
                chirp_f = getattr(modal.Cls.from_name(stub_name_worker, cls_name)(), func_name)
                assert isinstance(chirp_f, modal.Function)
                self.modal_chirp_worker_lookup_cache[(stub_name_worker, stub_function)] = chirp_f
                return chirp_f
            except Exception as e:
                print(
                    f"Error getting chirp worker for {item.model_name}: {e}. stub_name, stub_function: {stub_name_worker}, {stub_function}"
                )

    def _validate_input_queue_item(self, item: QueueItem) -> bool:
        """Validate the input queue item to check if the request if valid."""
        if (
            item.metadata.get("task") in ["infill", "infill_intro", "infill_outro"]
        ) and not item.is_infill:
            print(f"WTF is happening with {item.id} -- task is infill but not infill arguments.")
            return False
        # add task validations checks -- we should not be passing in clips if task isn't specified
        if not item.is_cover_condition and item.metadata.get("cover_clip_id", None) is not None:
            print(f"WTF is happening with {item.id} -- not cover but without cover id passed in.")
            return False
        if not item.is_artist_condition and item.metadata.get("artist_clip_id", None) is not None:
            print(f"WTF is happening with {item.id} -- not artist but without artist id passed in.")
            return False
        if (
            item.metadata.get("task") == "extend"
            or item.is_cover_extend
            or item.is_cover_infill
            or item.is_artist_extend
            or item.is_artist_infill
            or item.is_artist_cover_extend
        ) and item.prompt_audio is None:
            print(f"WTF is happening with {item.id} -- task is extend but prompt audio id passed in.")
            return False
        if item.metadata.get("task") == "upsample" and not item.is_upsample:
            print(f"WTF is happening with {item.id} -- task is upsample but without upsample arguments.")
            return False
        if item.is_diff_infill and (
            (
                item.metadata.get("infill_start_s", 0) - item.metadata.get("infill_context_start_s", 0)
                > 29
            )
            or (item.metadata.get("infill_context_end_s", 0) - item.metadata.get("infill_end_s", 0) > 29)
        ):
            print(
                f"WTF is happening with {item.id} -- task is diff infill but infill context is too long."
            )
            return False
        return True

    @modal.method()
    @distributed_trace("redirect_and_generate", CONDUCTOR_DDOG_SERVICE, env_name=DEPLOYMENT_TYPE)
    def redirect_and_generate(
        self, queue_item_json: str, history_prompt: HistoryPrompt | None = None
    ) -> None:
        """This worker will distribute the queue item to the appropriate model worker.

        If it is a continue and needs to fetch the history, it will go to the history loader.
            History loader will pass the item back into redirect_and_generate.
        Otherwise, it will go through augmentation and to the GPT worker directly.

        Args:
            queue_item_json: the queue item in json format.
                Note that this queue item can be at request-level or clip-level.
            history_prompt: the history prompt in HistoryPrompt.
        """
        original_item = QueueItem(**json.loads(queue_item_json))
        logger.debug(f"  {original_item}")
        tracer_current_span = tracer.current_span()
        if tracer_current_span:
            tracer_current_span.set_tag("id", original_item.id)
            tracer_current_span.set_tag("model_name", original_item.model_name)
            tracer_current_span.set_tag("title", original_item.title)
        # we need to do some house keeping first
        if "experiment" in original_item.metadata:
            # why would it even get here!? from sloppyness in studio api, carried over
            print(f"Request {original_item.id}: Remove unwanted 'experiment' key")
            # first correct the db logging
            original_item.metadata["experiment"] = []
            self._notify_queue_status_with_epxerimental_info(original_item)
            print(f"Request {original_item.id}: Correct DB info sent")
            original_item.metadata.pop("experiment")
        if "param_experiment" in original_item.metadata:
            # why would it even get here!? from sloppyness in studio api, carried over
            print(f"Request {original_item.id}: Remove unwanted 'param_experiment' key")
            # first correct the db logging
            original_item.metadata["param_experiment"] = ""
            self._notify_queue_status_with_epxerimental_info(original_item)
            print(f"Request {original_item.id}: Correct DB info sent")
            original_item.metadata.pop("param_experiment")
        if "model_config" in original_item.metadata:
            # why would it even get here!? from sloppyness in studio api, carried over
            print(f"Request {original_item.id}: Remove unwanted 'model_config' key")
            original_item.metadata.pop("model_config")
        if "experiment_version" in original_item.metadata:
            # why would it even get here!? from sloppyness in studio api, carried over
            print(f"Request {original_item.id}: Remove unwanted 'experiment_version' key")
            original_item.metadata.pop("experiment_version")

        # validate the queue item -- should meet expected conditions
        if not self._validate_input_queue_item(original_item):
            for index, clip_id in enumerate(original_item.all_child_clip_ids):
                new_item = original_item.model_copy(deep=True, update={"id": clip_id, "ids": None})
                new_item.notify_progress(
                    {
                        "id": new_item.id,
                        "type": "error",
                        "error_type": "invalid_generation_request",
                    },
                )
            return

        if (
            original_item.prompt_audio
            or original_item.is_cover_condition
            or original_item.is_artist_condition
            or original_item.is_infill
            or original_item.is_diff_infill
            or original_item.is_playlist_condition
            or original_item.is_multi_artist_consistency
            or original_item.is_underpainting
            or original_item.is_overpainting
        ) and history_prompt is None:
            # This is a continue job, without history, we need to fetch the history first
            _ = self.modal_f_history_loader.spawn(
                original_item.model_dump_json(), parent_context=serialize_context()
            )
            logger.info(f"Request {original_item.id}: Spawned modal job into history loader")
            # nothing to be done for here, job isn't in queue yet, except upsample / stem
            if (original_item.is_upsample or original_item.is_stem) and original_item.all_child_clip_ids:
                for clip_id in original_item.all_child_clip_ids:
                    curr_item = original_item.model_copy(deep=True, update={"id": clip_id, "ids": None})
                    # empty the multi_ids as well
                    curr_item.metadata["multi_ids"] = []
                    # we need to copy the image -- note that this is independent of the job queue
                    self.copy_and_upload_image(curr_item)
            return
        # ok we have everything, let's proceed!
        if not original_item.is_image_to_song:
            # start the image runner -- this is done at request level, image worker will split to clips
            # note that this is one image gen request, single/multiple clip's image gen
            # Check if image generation was already started from the early callback
            if original_item.metadata.get("image_generation_started", False):
                print(f"Request {original_item.id}: Image generation already started, skipping")
            else:
                if original_item.is_upsample or original_item.is_stem:
                    # we do not want to run the image generator for upsample
                    # it is copied after the request is spawned
                    pass
                elif isinstance(original_item.model_name, str) and (
                    "auk" in original_item.model_name
                    or "bluejay" in original_item.model_name
                    or "v4" in original_item.model_name
                ):
                    n_backlog = self.modal_f_pro_image_generator.get_current_stats().backlog
                    if n_backlog < 25:  # reduce the backlog to 25
                        _ = self.modal_f_pro_image_generator.spawn(
                            original_item.model_dump_json(), parent_context=serialize_context()
                        )
                    else:
                        _ = self.modal_f_image_generator.spawn(
                            original_item.model_dump_json(), parent_context=serialize_context()
                        )
                    if "auk" in original_item.model_name or "bluejay" in original_item.model_name:
                        _ = self.modal_f_downsampler.spawn(
                            original_item.model_dump_json(), parent_context=serialize_context()
                        )
                else:
                    _ = self.modal_f_image_generator.spawn(
                        original_item.model_dump_json(), parent_context=serialize_context()
                    )

        # Do the A/B testing split here
        item = self._experimental_model_mixer(original_item)
        # if this is a clip level job -- once we update all workers these can be removed
        # but if they are here, they are backwards compatible
        if not item.all_child_clip_ids:
            if item.model_name not in MODAL_FNS:
                # this is a bad model name, we should notify at the request level
                item.notify_progress(
                    {
                        "id": item.id,
                        "type": "error",
                        "error_type": "invalid_model_name",
                    },
                )
                return
            # start the audio runner, do a look up first
            gpt_worker = self._get_gpt_worker(item)
            # directly goes to the GPT worker, with empty HistoryPrompt
            _ = gpt_worker.spawn(
                item.model_dump_json(), history_prompt, parent_context=serialize_context()
            )
            logger.info(f"Request {item.id}: Spawned modal job {gpt_worker} item id {item.id}")
            self._notify_queue_status_with_epxerimental_info(item)
        # if this is a request level job:
        else:
            # we split the request back to clips level and pass them on
            items_to_notify = []
            for index, clip_id in enumerate(item.child_primary_clip_ids):
                new_item = item.model_copy(deep=True, update={"id": clip_id, "ids": None})
                if item.multi_ids:  # for stems, select the appropriate ids
                    new_item.metadata["multi_ids"] = item.multi_ids[index]
                if experiment_models := item.metadata.get("experiment"):
                    # reset the new_item's model name
                    new_item.model_name = experiment_models[index]
                    print(
                        f"Request {new_item.id}: Created new item {clip_id} with model name: {new_item.model_name}"
                    )
                if new_item.model_name not in MODAL_FNS:
                    # this is a bad model name, we should notify at the request level
                    print(
                        f"Request {new_item.id}: Incorrect model name {new_item.model_name} for clip {clip_id}"
                    )
                    new_item.notify_progress(
                        {
                            "id": new_item.id,
                            "type": "error",
                            "error_type": "invalid_model_name",
                        },
                    )
                    return
                if experiment_inference_config := item.metadata.get("model_config"):
                    # reset the new_item's model name
                    new_item.metadata["model_config"] = experiment_inference_config[index]
                    if not new_item.metadata["model_config"]:
                        # one of the experiment inference could be empty config
                        # in such cases, we don't bother passing them in
                        new_item.metadata.pop("model_config")
                        # this is what we log into the db. we need this to make db logging work.
                        new_item.metadata.pop("param_experiment")
                    else:
                        print(
                            f"Request {new_item.id}: Created new item {clip_id} with experimental inference config: {new_item.metadata['model_config']}"
                        )
                    # for tag augmentation, the augmented tags are passed from studio api
                    # we will overwrite the tags with augmented_tags
                    # the augmented tags are saved in the metadata
                    # the clips triggered with the tag_aug will have 'tag_aug_vx' in the param_experiment
                    if new_item.metadata.get("model_config", {}).get("tag_aug", False):
                        if new_augmented_tags := new_item.metadata.get("augmented_tags"):
                            print(
                                f"Request {new_item.id}: Created new item {clip_id} with "
                                f"experimental augment tags from: "
                                f"{new_item.metadata['tags']} to "
                                f"{new_augmented_tags}."
                            )
                            new_item.metadata["model_config"] = {}
                            new_item.metadata["tags"] = new_augmented_tags
                        else:
                            # augmented tag not implemented
                            new_item.metadata.pop("model_config")
                            new_item.metadata.pop("param_experiment")

                # start the audio runner, do a look up first
                gpt_worker = self._get_gpt_worker(new_item)
                # for delay experiment -- only sleep for the first request
                if (delay_request_time := new_item.metadata.get("exp_delay_by_sec")) and index == 0:
                    print(
                        f"Request {new_item.id}: is delayed by {delay_request_time}s for user {new_item.metadata.get('user_id')}"
                    )
                    time.sleep(delay_request_time)
                # extra traffic monitoring
                if (
                    DEPLOYMENT_TYPE == "prod"
                    and new_item.model_name
                    and (
                        "v4" in new_item.model_name
                        or "auk" in new_item.model_name
                        or "bluejay" in new_item.model_name
                    )
                ):
                    # check backlog size
                    n_backlog = gpt_worker.get_current_stats().backlog
                    if n_backlog > 200:  # be a little more generous about backlog size
                        print(
                            f"Request {new_item.id}: Current backlog for {new_item.model_name}: {gpt_worker.get_current_stats().backlog}"
                        )
                        new_item.notify_progress(
                            {
                                "id": new_item.id,
                                "type": "error",
                                "error_type": "moderation_failure",
                                "error_message": "Sorry but we are currently experiencing elevated usage. Please try again later. Thanks!",
                            },
                        )
                        # can't pass the info down the line
                        return
                # spin up the gpt worker
                _ = gpt_worker.spawn(
                    new_item.model_dump_json(), history_prompt, parent_context=serialize_context()
                )
                logger.debug(
                    f"Request {new_item.id}: Spawned modal job {gpt_worker}, clip id {new_item.id}"
                )
                items_to_notify.append(new_item)
                self.events_queue.put(
                    {
                        "type": "generate_queued",
                        "data": {
                            "clip_id": new_item.id,
                        },
                    },
                    partition=item.id,
                    partition_ttl=60,
                )
            # cache the items and notify them at the end to not blocking the job spawning
            for curr_item in items_to_notify:
                self._notify_queue_status_with_epxerimental_info(curr_item)
                # note that we do it after the generation request is spawned
                logger.debug(
                    f"Request {new_item.id}: Finished notify clip id {curr_item.id} with model name: {curr_item.model_name}"
                )


def _test_get_prompt_from_gpt_description_prompt(chat_gpt_app: PromptGenerationStub) -> None:
    nice_prompt = (
        QueueItem(
            id="1337",
            model_name=TEST_MODEL_NAME,
            metadata={
                "gpt_description_prompt": "a cumbia song about a cowboy and his cow",
            },
        )
    ).model_dump_json()
    response_dict = chat_gpt_app.get_prompt_from_gpt_description_prompt.remote(nice_prompt)
    print("response_dict:", response_dict)
    partial_expected_response_dict = {
        "id": "1337",
        "ok": 1,
    }
    for k in partial_expected_response_dict:
        assert partial_expected_response_dict[k] == response_dict[k], (
            k,
            partial_expected_response_dict[k],
            response_dict[k],
        )
    assert "[verse" in response_dict["prompt"].lower(), response_dict["prompt"]
    assert isinstance(response_dict["genre_tags"][0], str)


def _test_get_prompt_from_gpt_description_prompt_together_ai(chat_gpt_app: PromptGenerationStub) -> None:
    together_model = (
        "patsuno/Meta-Llama-3.1-8B-Instruct-Reference-remi-test-8B-lr-1en4-ep7-CORRECTED-fdb36cb0"
    )
    nice_prompt = (
        QueueItem(
            id="1337",
            model_name=together_model,
            metadata={
                "gpt_description_prompt": "a sad country song about a cowboy and his cow",
            },
        )
    ).model_dump_json()
    response_dict = chat_gpt_app.get_prompt_from_gpt_description_prompt.remote(nice_prompt)
    print("response_dict:", response_dict)
    partial_expected_response_dict = {
        "id": "1337",
        "ok": 1,
    }
    for k in partial_expected_response_dict:
        assert partial_expected_response_dict[k] == response_dict[k], (
            k,
            partial_expected_response_dict[k],
            response_dict[k],
        )
    assert "[verse" in response_dict["prompt"].lower(), response_dict["prompt"]
    assert isinstance(response_dict["genre_tags"][0], str)


def _test_get_prompt_from_gpt_description_prompt_remi(chat_gpt_app: PromptGenerationStub) -> None:
    # TODO(pat) refactor these tests
    nice_prompt = (
        QueueItem(
            id="1337",
            metadata={
                "gpt_description_prompt": "a sad country song about a cowboy and his cow",
                "lyrics_model": "remi-v1",
            },
        )
    ).model_dump_json()
    response_dict = chat_gpt_app.get_prompt_from_gpt_description_prompt.remote(nice_prompt)
    print("response_dict:", response_dict)
    partial_expected_response_dict = {
        "id": "1337",
        "ok": 1,
    }
    for k in partial_expected_response_dict:
        assert partial_expected_response_dict[k] == response_dict[k], (
            k,
            partial_expected_response_dict[k],
            response_dict[k],
        )
    assert "[verse" in response_dict["prompt"].lower(), response_dict["prompt"]
    assert isinstance(response_dict["genre_tags"][0], str)


def _test_get_prompt_from_gpt_description_prompt_modal_remi(chat_gpt_app: PromptGenerationStub) -> None:
    # TODO(pat) refactor these tests
    modal_remi_model = "my-remi-8b-1"
    nice_prompt = (
        QueueItem(
            id="1337",
            metadata={
                "gpt_description_prompt": "a sad country song about a cowboy and his cow",
                "lyrics_model": modal_remi_model,
            },
        )
    ).model_dump_json()
    response_dict = chat_gpt_app.get_prompt_from_gpt_description_prompt.remote(nice_prompt)
    print("response_dict:", response_dict)
    partial_expected_response_dict = {
        "id": "1337",
        "ok": 1,
    }
    for k in partial_expected_response_dict:
        assert partial_expected_response_dict[k] == response_dict[k], (
            k,
            partial_expected_response_dict[k],
            response_dict[k],
        )
    assert "[verse" in response_dict["prompt"].lower(), response_dict["prompt"]
    assert isinstance(response_dict["genre_tags"][0], str)


def _test_get_prompt_from_gpt_description_prompt_new_chatgpt(chat_gpt_app: PromptGenerationStub) -> None:
    # TODO(pat) refactor these tests
    new_chatgpt = "gpt-4o-2024-11-20"
    nice_prompt = (
        QueueItem(
            id="1337",
            metadata={
                "gpt_description_prompt": "a sad country song about a cowboy and his cow",
                "lyrics_model": new_chatgpt,
            },
        )
    ).model_dump_json()
    response_dict = chat_gpt_app.get_prompt_from_gpt_description_prompt.remote(nice_prompt)
    print("response_dict:", response_dict)
    partial_expected_response_dict = {
        "id": "1337",
        "ok": 1,
    }
    for k in partial_expected_response_dict:
        assert partial_expected_response_dict[k] == response_dict[k], (
            k,
            partial_expected_response_dict[k],
            response_dict[k],
        )
    assert "[verse" in response_dict["prompt"].lower(), response_dict["prompt"]
    assert isinstance(response_dict["genre_tags"][0], str)


def _test_get_prompt_from_gpt_description_prompt_vertex(chat_gpt_app: PromptGenerationStub) -> None:
    # TODO(pat) refactor these tests
    model = "gemini-2.0"
    nice_prompt = (
        QueueItem(
            id="1337",
            metadata={
                "gpt_description_prompt": "a sad country song about a cowboy and his cow",
                "lyrics_model": model,
            },
        )
    ).model_dump_json()
    response_dict = chat_gpt_app.get_prompt_from_gpt_description_prompt.remote(nice_prompt)
    print("response_dict:", response_dict)
    partial_expected_response_dict = {
        "id": "1337",
        "ok": 1,
    }
    for k in partial_expected_response_dict:
        assert partial_expected_response_dict[k] == response_dict[k], (
            k,
            partial_expected_response_dict[k],
            response_dict[k],
        )
    assert "[verse" in response_dict["prompt"].lower(), response_dict["prompt"]
    assert isinstance(response_dict["genre_tags"][0], str)


def _test_get_prompt_from_gpt_description_prompt_remi_hps(chat_gpt_app: PromptGenerationStub) -> None:
    # TODO(pat) refactor these tests
    remi_model = "remi-v1-hps1"
    nice_prompt = (
        QueueItem(
            id="1337",
            metadata={
                "gpt_description_prompt": "a sad country song about a cowboy and his cow",
                "lyrics_model": remi_model,
            },
        )
    ).model_dump_json()
    response_dict = chat_gpt_app.get_prompt_from_gpt_description_prompt.remote(nice_prompt)
    print("response_dict:", response_dict)
    partial_expected_response_dict = {
        "id": "1337",
        "ok": 1,
    }
    for k in partial_expected_response_dict:
        assert partial_expected_response_dict[k] == response_dict[k], (
            k,
            partial_expected_response_dict[k],
            response_dict[k],
        )
    assert "[verse" in response_dict["prompt"].lower(), response_dict["prompt"]
    assert isinstance(response_dict["genre_tags"][0], str)


def _test_moderate_and_generate(chat_gpt_app: ChatGptStub) -> None:
    nice_prompt = (
        QueueItem(
            id="1337",
            model_name=TEST_MODEL_NAME,
            metadata={
                "gpt_description_prompt": "a pop song about a literal banana",
            },
            title="Write me a song",
        )
    ).model_dump_json()
    item = chat_gpt_app.moderate_and_generate.remote(nice_prompt)
    assert item is not None


def _test_moderate_and_generate_stem(chat_gpt_app: ChatGptStub) -> None:
    id_a = "1338a"
    id_b = "1338b"
    a_ids = [id_a + "_stem", id_a + "_complement"]
    b_ids = [id_b + "_stem", id_b + "_complement"]
    nice_prompt = (
        QueueItem(
            id=id_a,
            prompt_audio="a5e2198a-f352-4abb-9a24-7f81b143ded3",  # stone
            model_name="chirp-stem-comp",
            metadata={
                "task": "gen_stem",
                "stem_type_group_name": "Vocals",
                "multi_ids": [a_ids, b_ids],
            },
        )
    ).model_dump_json()
    item = chat_gpt_app.moderate_and_generate.remote(nice_prompt)
    assert item is not None


def _test_moderate_and_generate_stem_multi(chat_gpt_app: ChatGptStub) -> None:
    id_a = "1339"
    a_ids = [f"{id_a}_{i}" for i in range(12)]
    b_ids = [f"{id_a}_{i}" for i in range(12, 24)]
    nice_prompt = (
        QueueItem(
            id=id_a,
            prompt_audio="a5e2198a-f352-4abb-9a24-7f81b143ded3",  # stone
            model_name="chirp-stem",
            metadata={
                "task": "gen_stem",
                "multi_ids": [a_ids, b_ids],
            },
        )
    ).model_dump_json()
    item = chat_gpt_app.moderate_and_generate.remote(nice_prompt)
    assert item is not None


def _test_moderate_and_generate_extract_single_image(chat_gpt_app: ChatGptStub) -> None:
    nice_prompt = (
        QueueItem(
            id="1337",
            model_name=TEST_MODEL_NAME,
            metadata={
                "gpt_description_prompt": "Suno iOS app launch",
                "image_to_song_s3_ids": ["c509b8b5"],
            },
            title="Friday",
        )
    ).model_dump_json()
    item = chat_gpt_app.moderate_and_generate.remote(nice_prompt)
    assert item is not None


def _test_moderate_and_generate_extract_multiple_images(chat_gpt_app: ChatGptStub) -> None:
    nice_prompt = (
        QueueItem(
            id="1337",
            model_name=TEST_MODEL_NAME,
            metadata={
                "gpt_description_prompt": "Suno iOS app launch",
                "image_to_song_s3_ids": [
                    "c509b8b5",
                    "c509b8b6",
                ],  # likely error is rate limiting by HIVE for the same image url!
            },
            title="Friday",
        )
    ).model_dump_json()
    item = chat_gpt_app.moderate_and_generate.remote(nice_prompt)
    assert item is not None


def _test_moderate_and_generate_naughty_gpt_description_prompt(chat_gpt_app: ChatGptStub) -> None:
    naughty_prompt = (
        QueueItem(
            id="1337",
            model_name=TEST_MODEL_NAME,
            metadata={
                # what a world
                "gpt_description_prompt": "What the fuck did you say to me, you little bitch?  You're fucking dead, kiddo.",
            },
            title="Write me a song",
        )
    ).model_dump_json()
    item = chat_gpt_app.moderate_and_generate.remote(naughty_prompt)
    assert item is None


def _test_moderate_and_generate_foreign_languages(chat_gpt_app: ChatGptStub) -> None:
    gpt_prompts = {
        "spanish": "una cancion pop sobre un platano literal",
        "spanish2": "una canción de amor sobre España",
        "chinese": "一首关于真正的香蕉的歌",
        "japanese": "文字通りのバナナについてのポップソング",
        "portuguese": "uma música dubstep sobre uma banana literal",
        "arabic": "أغنية روك عن موزة حرفية",
        "hindi": "एक शाब्दिक केले के बारे में एक रॉक गीत",
        "urdu": "لفظی کیلے کے بارے میں ایک ملکی گانا",
        "russian": "рэп-песня о банане в буквальном смысле слова",
        "dutch": "ik wil een liedje over een bananan vla met stukjes aardbei",
    }
    for language, prompt in gpt_prompts.items():
        nice_prompt = (
            QueueItem(
                id="1337",
                model_name=TEST_MODEL_NAME,
                metadata={
                    "gpt_description_prompt": prompt,
                },
            )
        ).model_dump_json()
        raw_item = chat_gpt_app.moderate_and_generate.remote(nice_prompt)
        assert raw_item is not None
        item = json.loads(raw_item)
        tags = item["metadata"]["tags"]
        title = item["title"]
        assert title
        assert tags
        assert item["title"] != tags
        assert item["prompt_text"]
        # these next two tests might be a little too strict, tags don't
        # necessarily need to come back in English.  feel free to
        # revise if failing
        lowercased_lyrics = item["prompt_text"].lower()
        assert "verse" in lowercased_lyrics or "verso" in lowercased_lyrics

        print(language, prompt, item["prompt_text"])
        assert len(item["prompt_text"]) > 100
        assert len(item["prompt_text"].split("\n")) >= 10


def _test_moderate_generate_ok_title(chat_gpt_app: ChatGptStub) -> None:
    ok_prompt = (
        QueueItem(
            id="1337",
            model_name=TEST_MODEL_NAME,
            metadata={
                "prompt": "please write me a song",
            },
            prompt_text="Please write a song",
            title="I Love You",
        )
    ).model_dump_json()
    item = chat_gpt_app.moderate_and_generate.remote(ok_prompt)
    assert item is not None


def _test_moderate_generate_bad_title(chat_gpt_app: ChatGptStub) -> None:
    bad_prompt = (
        QueueItem(
            id="1337",
            model_name=TEST_MODEL_NAME,
            metadata={
                "prompt": "please write me a song",
            },
            prompt_text="Please write a song",
            title="Fuck you, you are the worst",
        )
    ).model_dump_json()
    item = chat_gpt_app.moderate_and_generate.remote(bad_prompt)
    assert item["moderation_result"] == "success", item


def _test_moderate_text_ok(chat_gpt_app: TextModerationSub) -> None:
    ok_prompt = (
        QueueItem(
            id="1337",
            model_name=TEST_MODEL_NAME,
            metadata={
                "text": "hello how are you",
            },
        )
    ).model_dump_json()
    item = chat_gpt_app.moderate_text.remote(ok_prompt)
    assert item["moderation_result"] == "success", item


def _test_moderate_text_moderately_naughty(chat_gpt_app: TextModerationSub) -> None:
    ok_prompt = (
        QueueItem(
            id="1337",
            model_name=TEST_MODEL_NAME,
            metadata={
                "text": "fuck fuck fuck",
            },
        )
    ).model_dump_json()
    item = chat_gpt_app.moderate_text.remote(ok_prompt)
    assert item["moderation_result"] == "success", item


_MARTYROLOGY_LYRICS = """
[Chorus]
After having done a superhuman good,
Emperor Tal Julian arrived,
He sent him to prison!
After a while he tortured him!
They began by cutting off his hand

[Chorus]
Then he made him swallow molten lead,
Sticking it down his throat by the snout,
On the gridiron frustrated!
From snakes poisoned!
And immersed in the hottest bitumen…

But he only died with a blow!
"""


def _test_moderate_text_martyrology(chat_gpt_app: TextModerationSub) -> None:
    ok_prompt = (
        QueueItem(
            id="1337",
            model_name=TEST_MODEL_NAME,
            metadata={
                "text": _MARTYROLOGY_LYRICS,
            },
        )
    ).model_dump_json()
    item = chat_gpt_app.moderate_text.remote(ok_prompt)
    assert item["moderation_result"] == "success", item


def _test_moderate_and_generate_martyrology(chat_gpt_app: ChatGptStub) -> None:
    prompt = (
        QueueItem(id="1337", model_name=TEST_MODEL_NAME, metadata={}, prompt_text=_MARTYROLOGY_LYRICS)
    ).model_dump_json()
    item = chat_gpt_app.moderate_and_generate.remote(prompt)
    assert item is not None


def _test_moderate_text_very_naughty(chat_gpt_app: TextModerationSub) -> None:
    bad_prompt = (
        QueueItem(
            id="1337",
            model_name=TEST_MODEL_NAME,
            metadata={
                "text": "I'm going to murder you",
            },
        )
    ).model_dump_json()
    item = chat_gpt_app.moderate_text.remote(bad_prompt)
    assert item["moderation_result"] == "failure", item


def _test_disable_artist_name_detection(chat_gpt_app: ChatGptStub) -> None:
    artist_names_prompt = (
        QueueItem(
            id="1337",
            model_name=TEST_MODEL_NAME,
            metadata={
                "gpt_description_prompt": "a song about Ed Sheeran, Michael Jackson and the Beastie Boys",
                "check_artist_names": False,
            },
        )
    ).model_dump_json()
    item = chat_gpt_app.moderate_and_generate.remote(artist_names_prompt)
    assert item is not None


def _test_artist_name_detection_tags(chat_gpt_app: ChatGptStub) -> None:
    artist_names_prompt = (
        QueueItem(
            id="1337",
            model_name=TEST_MODEL_NAME,
            metadata={
                "tags": "taylor swift",
            },
            prompt_text="i am testing",
        )
    ).model_dump_json()
    item = chat_gpt_app.moderate_and_generate.remote(artist_names_prompt)
    assert item is None, item


def _test_artist_name_detection(chat_gpt_app: ChatGptStub) -> None:
    artist_names_prompt = (
        QueueItem(
            id="1337",
            model_name=TEST_MODEL_NAME,
            metadata={"gpt_description_prompt": "a song about ed sheeran"},
        )
    ).model_dump_json()
    item = chat_gpt_app.moderate_and_generate.remote(artist_names_prompt)
    assert item is None


def _test_one_box_instrumentals(chat_gpt_app: ChatGptStub) -> None:
    instrumental_prompt = (
        QueueItem(
            id="1337",
            model_name=TEST_MODEL_NAME,
            metadata={"gpt_description_prompt": "an instrumental funk song", "make_instrumental": True},
        )
    ).model_dump_json()
    item_str = chat_gpt_app.moderate_and_generate.remote(instrumental_prompt)
    item = json.loads(item_str)
    assert item["prompt_text"] == "[Instrumental]", item
    assert not (item["metadata"].get("lang")), item


def _test_no_genres_in_lyrics(chat_gpt_app: ChatGptStub) -> None:
    # This test is a little flaky, feel free to ignore transient errors.
    errors = 0
    for i in range(5):
        prompt = (
            QueueItem(
                id="1337",
                model_name=TEST_MODEL_NAME,
                metadata={
                    "gpt_description_prompt": "a synthpop song about a lost dog",
                },
            )
        ).model_dump_json()
        item_str = chat_gpt_app.moderate_and_generate.remote(prompt)
        item = json.loads(item_str)
        print("synthpop item:", item)
        if "synthpop" in item["prompt_text"]:
            errors += 1
    assert errors <= 2, errors


def _test_lyrics_moderation(chat_gpt_app: ChatGptStub) -> None:
    bad_prompt = "dink"
    queue_item = (
        QueueItem(
            id="1337",
            model_name=TEST_MODEL_NAME,
            metadata={},
            # custom lyrics moderation ignores metadata.prompt and reads prompt_text instead
            prompt_text=bad_prompt,
        )
    ).model_dump_json()
    item = chat_gpt_app.moderate_and_generate.remote(queue_item)
    assert item is None, item


def _test_lyrics_length(chat_gpt_app: ChatGptStub) -> None:
    v3_queue_item = (
        QueueItem(
            id="1337",
            model_name=TEST_MODEL_NAME,
            metadata={
                "gpt_description_prompt": "a song about three scoops of vanilla yogurt",
            },
            prompt_text="",
        )
    ).model_dump_json()
    item_str = chat_gpt_app.moderate_and_generate.remote(v3_queue_item)
    item = json.loads(item_str)
    prompt = item["prompt_text"]
    stanzas = get_stanzas(prompt)
    assert len(stanzas) == 3, (prompt, len(stanzas))

    v2_queue_item = (
        QueueItem(
            id="1337",
            model_name="chirp-v2-xxl-alpha",
            metadata={
                "gpt_description_prompt": "a song about two scoops of vanilla yogurt",
            },
            prompt_text="",
        )
    ).model_dump_json()
    item_str = chat_gpt_app.moderate_and_generate.remote(v2_queue_item)
    item = json.loads(item_str)
    prompt = item["prompt_text"]
    stanzas = get_stanzas(prompt)
    assert len(stanzas) == 6, (prompt, len(stanzas))

    unspecified_queue_item = (
        QueueItem(
            id="1337",
            metadata={
                "gpt_description_prompt": "a song about unspecified scoops of vanilla yogurt",
            },
            prompt_text="",
            model_name="chirp-v3",
        )
    ).model_dump_json()
    item_str = chat_gpt_app.moderate_and_generate.remote(unspecified_queue_item)
    item = json.loads(item_str)
    prompt = item["prompt_text"]
    stanzas = get_stanzas(prompt)
    assert len(stanzas) == 3, (prompt, len(stanzas))


def _test_lyrics_length_model_3p5(chat_gpt_app: ChatGptStub) -> None:
    v3_queue_item = (
        QueueItem(
            id="1337",
            model_name="chirp-v3p5",
            metadata={
                "gpt_description_prompt": "a song about three scoops of vanilla yogurt",
            },
            prompt_text="",
        )
    ).model_dump_json()
    item_str = chat_gpt_app.moderate_and_generate.remote(v3_queue_item)
    assert item_str is not None
    item = json.loads(item_str)
    prompt = item["prompt_text"]
    stanzas = get_stanzas(prompt)
    # we ask for six stanzas, but sometimes chatGPT gets tired, so
    # just make sure we get at least 5 back.
    assert len(stanzas) >= 5, (prompt, len(stanzas))


def _test_image_extraction_screenshot(chat_gpt_app: ChatGptStub) -> None:
    nice_prompt = (
        QueueItem(
            id="1337",
            model_name=TEST_MODEL_NAME,
            metadata={
                "gpt_description_prompt": "festival",
                "image_to_song_s3_ids": ["image_01872ae0-65fe-4549-b514-78bb3ec98299"],
            },
            title="N/A",
        )
    ).model_dump_json()
    item = chat_gpt_app.moderate_and_generate.remote(nice_prompt)
    assert item is not None


def _test_predict_lyrics_language(chat_gpt_app: ChatGptStub) -> None:
    english_prompt = (
        QueueItem(
            id="1337",
            model_name=TEST_MODEL_NAME,
            prompt_text="""
            I am testing this is English

            I am testing this is English
            """,
            metadata={},
            title="N/A",
            ids=[],
        )
    ).model_dump_json()
    item = chat_gpt_app.moderate_and_generate.remote(english_prompt)
    assert item is not None
    item_language = json.loads(item)["metadata"].get("lang")
    assert item_language == "English"
    chinese_prompt = (
        QueueItem(
            id="1337",
            model_name=TEST_MODEL_NAME,
            prompt_text="""
            我在说中文

            我在说中文

            I am testing this is English
            """,
            metadata={},
            title="N/A",
            ids=[],
        )
    ).model_dump_json()
    item = chat_gpt_app.moderate_and_generate.remote(chinese_prompt)
    assert item is not None
    item_language = json.loads(item)["metadata"].get("lang")
    assert item_language == "Chinese"


def _test_copyright_detector(copyright_detector_app):
    good_lyrics = "hi how are you"
    assert not copyright_detector_app.are_lyrics_copyrighted.remote(good_lyrics)
    bad_lyrics = "beat it (beat it), beat it (beat it), no one wants to be defeated"
    assert copyright_detector_app.are_lyrics_copyrighted.remote(bad_lyrics)
    print("copyright detection tests completed")


def _test_prompt_upsampler(prompt_upsampler_app):
    original_prompt = "lo fi japanese citypop"
    upsampled_prompt = prompt_upsampler_app.upsample.remote(original_prompt)
    expected_terms = ["drums", "jazz", "piano", "guitar"]
    found_at_least_one_term = any(term in upsampled_prompt for term in expected_terms)
    assert (
        isinstance(upsampled_prompt, str)
        and len(upsampled_prompt) > len(original_prompt)
        and found_at_least_one_term
    ), upsampled_prompt

    original_prompt = "lo fi japanese citypop"
    upsampled_prompt = prompt_upsampler_app.upsample.remote(original_prompt, is_instrumental=False)
    expected_terms = ["drums", "jazz", "piano", "guitar"]
    found_at_least_one_term = any(term in upsampled_prompt for term in expected_terms)
    assert (
        isinstance(upsampled_prompt, str)
        and len(upsampled_prompt) > len(original_prompt)
        and found_at_least_one_term
    ), upsampled_prompt

    original_prompt = "lo fi japanese citypop"
    upsampled_prompt = prompt_upsampler_app.upsample.remote(original_prompt, is_instrumental=True)
    expected_terms = ["drums", "jazz", "piano", "guitar"]
    found_at_least_one_term = any(term in upsampled_prompt for term in expected_terms)
    assert (
        isinstance(upsampled_prompt, str)
        and len(upsampled_prompt) > len(original_prompt)
        and found_at_least_one_term
    ), upsampled_prompt

    assert "vocals" not in upsampled_prompt, upsampled_prompt

    original_prompt = "taylor swift"
    upsampled_prompt = prompt_upsampler_app.upsample.remote(
        original_prompt, is_instrumental=True, check_artist_names=False
    )
    assert isinstance(upsampled_prompt, str)

    original_prompt = "taylor swift"
    result = prompt_upsampler_app.upsample.remote(
        original_prompt, is_instrumental=True, check_artist_names=True
    )
    assert result["ok"] is False, upsampled_prompt
    assert "taylor" in result["error_message"], upsampled_prompt

    original_prompt = "jazzwave"
    result = prompt_upsampler_app.upsample.remote(
        original_prompt, is_instrumental=True, check_artist_names=True
    )
    assert result["ok"] is True, upsampled_prompt
    assert "jazz" in result["result"]


def _test_prompt_downsampler(prompt_downsampler_app):
    long_prompt = (
        "A vibrant blend of melodic hip hop and pop influences, showcasing an emotional journey "
        "through love and longing. The sound is characterized by rich synthesizers, rhythmic bass, "
        "and danceable beats, all enhancing the energetic vibe of the track. Vocally, it showcases "
        "smooth male vocals that project raw emotion and nostalgia, inviting listeners to connect "
        "with the feelings of lost love. The song thrives on its catchy chorus, engaging the audience "
        "throughout, with a reflective progression that maintains interest and engagement"
    )
    downsampled = prompt_downsampler_app.downsample.remote(long_prompt)
    assert isinstance(downsampled, str)
    assert any(genre in downsampled.lower() for genre in ["hip hop", "pop", "ballad"]), downsampled


def _test_twitter_to_song(chat_gpt_app: ChatGptStub):
    prompt = (
        QueueItem(
            id="1337",
            model_name=TEST_MODEL_NAME,
            metadata={
                "request_log_id": "6a8321c1-579a-4b48-9a89-403d03f476fc",
                "tweet_id": "1912645166415880694",
                "credit_cost": 5,
                "feature_flags": 3,
                "user_handle": "trolltrudy",
                "created_session_id": None,
                "twitter_to_song_log_id": "6a8321c1-579a-4b48-9a89-403d03f476fc",
                "twitter_to_song_content": "@PopCrave Can we send her back",
                "type": "gen",
                "source": "web",
                "prompt": "",
                "tags": "funny, pop, pop music",
                "gpt_prompt": None,
                "gpt_description_prompt": "@PopCrave Can we send her back\n\nTop Comments:\n\n\n@n1k1s_son @Wendys @PopCrave Over 2 years old, scumbag https://t.co/Oj4W05SF6m\n\n@kinematicaI @PopCrave Again. Don't know. Don't care.\n\n@Wendys @PopCrave @trolltrudy  Wendy's wanna send Katy Perry back to Space loooool\n\n@Wendys @PopCrave You are playing on the wrong side Wenzy 🙃\n\n@Powers_HK @Wendys @PopCrave Violence?",
            },
        )
    ).model_dump_json()
    item = chat_gpt_app.moderate_and_generate.remote(prompt)
    assert item is not None


def _test_generate_lyrics_and_genre_from_gpt_instructions(chat_gpt_app: ChatGptStub):
    prompt = (
        QueueItem(
            id="1337",
            model_name=TEST_MODEL_NAME,
            metadata={
                "request_log_id": "6a8321c1-579a-4b48-9a89-403d03f476fc",
                "tweet_id": "1912645166415880694",
                "credit_cost": 5,
                "feature_flags": 3,
                "user_handle": "trolltrudy",
                "created_session_id": None,
                "twitter_to_song_log_id": "6a8321c1-579a-4b48-9a89-403d03f476fc",
                "twitter_to_song_content": "@PopCrave Can we send her back",
                "type": "gen",
                "source": "web",
                "prompt": "",
                "tags": "funny, pop, pop music",
                "gpt_prompt": None,
                "gpt_description_prompt": "@PopCrave Can we send her back\n\nTop Comments:\n\n\n@n1k1s_son @Wendys @PopCrave Over 2 years old, scumbag https://t.co/Oj4W05SF6m\n\n@kinematicaI @PopCrave Again. Don't know. Don't care.\n\n@Wendys @PopCrave @trolltrudy  Wendy's wanna send Katy Perry back to Space loooool\n\n@Wendys @PopCrave You are playing on the wrong side Wenzy 🙃\n\n@Powers_HK @Wendys @PopCrave Violence?",
            },
        )
    ).model_dump_json()
    item = chat_gpt_app.generate_lyrics_and_genre_from_gpt_instructions.remote(prompt)
    assert item is not None


@app.local_entrypoint()
def main() -> None:
    """Test ChatGptStub locally."""
    copyright_detector_app = CopyrightDetectorStub()
    chat_gpt_app = ChatGptStub()
    prompt_generation_app = PromptGenerationStub()
    text_moderation_app = TextModerationSub()
    prompt_upsampler_app = PromptUpsamplerStub()
    prompt_downsampler_app = PromptDownsamplerStub()

    _test_prompt_upsampler(prompt_upsampler_app)
    _test_prompt_downsampler(prompt_downsampler_app)
    _test_copyright_detector(copyright_detector_app)
    _test_get_prompt_from_gpt_description_prompt_new_chatgpt(prompt_generation_app)
    _test_get_prompt_from_gpt_description_prompt_remi_hps(prompt_generation_app)
    _test_get_prompt_from_gpt_description_prompt_modal_remi(prompt_generation_app)

    _test_get_prompt_from_gpt_description_prompt_remi(prompt_generation_app)
    _test_get_prompt_from_gpt_description_prompt_together_ai(prompt_generation_app)
    _test_get_prompt_from_gpt_description_prompt(prompt_generation_app)
    _test_moderate_and_generate(chat_gpt_app)
    _test_moderate_and_generate_naughty_gpt_description_prompt(chat_gpt_app)
    _test_moderate_and_generate_stem(chat_gpt_app)
    _test_moderate_and_generate_stem_multi(chat_gpt_app)

    _test_moderate_and_generate_foreign_languages(chat_gpt_app)

    _test_moderate_generate_ok_title(chat_gpt_app)

    _test_moderate_text_ok(text_moderation_app)
    _test_moderate_text_moderately_naughty(text_moderation_app)
    _test_moderate_text_martyrology(text_moderation_app)
    _test_moderate_and_generate_martyrology(chat_gpt_app)
    _test_moderate_text_very_naughty(text_moderation_app)
    _test_artist_name_detection(chat_gpt_app)
    _test_disable_artist_name_detection(chat_gpt_app)
    _test_artist_name_detection_tags(chat_gpt_app)
    _test_one_box_instrumentals(chat_gpt_app)
    _test_lyrics_moderation(chat_gpt_app)
    _test_lyrics_length(chat_gpt_app)

    _test_no_genres_in_lyrics(chat_gpt_app)

    _test_lyrics_length_model_3p5(chat_gpt_app)
    _test_predict_lyrics_language(chat_gpt_app)
    _test_twitter_to_song(chat_gpt_app)
    _test_generate_lyrics_and_genre_from_gpt_instructions(chat_gpt_app)