ai-flow-python/app/service/task/task_executor.py

import logging
from typing import Callable

from app.service.task.base_task import BaseTask
from app.service.task.task_param import TaskParam
from app.service.task.task_result import TaskResult

logger = logging.getLogger(__name__)


class _NoopTask(BaseTask):
    def get_type(self) -> str:
        return "noop"

    def execute(self, param: TaskParam) -> TaskResult:
        logger.warning("[noop] No task implementation found for type=%s, returning empty success", param.type)
        return TaskResult.success({})


class TaskExecutor:
    def __init__(self):
        self._registry: dict[str, BaseTask] = {}

    def register(self, type_name: str, task: BaseTask) -> "TaskExecutor":
        self._registry[type_name] = task
        return self

    def resolve(self, type_name: str) -> BaseTask:
        task = self._registry.get(type_name)
        if task is None:
            logger.warning("No registered task for type '%s', using NoopTask", type_name)
            return _NoopTask()
        return task

    def registered_types(self) -> set[str]:
        return set(self._registry.keys())

    @classmethod
    def create_default(cls) -> "TaskExecutor":
        from app.service.task.impl.kontext_pro_task import KontextProTask
        from app.service.task.impl.kontext_max_task import KontextMaxTask
        from app.service.task.impl.keling_image2video_task import KelingImage2VideoTask
        from app.service.task.impl.keling_video2effect_task import KelingVideo2EffectTask
        from app.service.task.impl.smart_beauty_task import SmartBeautyTask
        from app.service.task.impl.video_editor_task import VideoEditorTask
        from app.service.task.impl.image_editor_task import ImageEditorTask
        from app.service.task.impl.file_download_task import FileDownloadTask
        from app.service.task.impl.file_upload_task import FileUploadTask
        from app.service.task.impl.base64_decode_task import Base64DecodeTask
        from app.service.task.impl.base64_encode_task import Base64EncodeTask
        from app.service.task.impl.nano_banana_task import NanoBananaTask
        from app.service.task.impl.google_veo3_task import GoogleVeo3Task
        from app.service.task.impl.seedance_video_task import SeedanceVideoTask
        from app.service.task.impl.seedream_image_task import SeedreamImageTask

        executor = cls()
        tasks: list[BaseTask] = [
            KontextProTask(),
            KontextMaxTask(),
            KelingImage2VideoTask(),
            KelingVideo2EffectTask(),
            SmartBeautyTask(),
            VideoEditorTask(),
            ImageEditorTask(),
            FileDownloadTask(),
            FileUploadTask(),
            Base64DecodeTask(),
            Base64EncodeTask(),
            NanoBananaTask(),
            GoogleVeo3Task(),
            SeedanceVideoTask(),
            SeedreamImageTask(),
        ]
        for t in tasks:
            executor.register(t.get_type(), t)
        return executor