ai-flow-python/app/service/engine/workflow_engine.py

import logging
import os
import time
from typing import Any

from app.config import settings
from app.service.engine.execution_context import ExecutionContext
from app.service.engine.parameter_resolver import ParameterResolver
from app.service.engine.workflow_result import WorkflowResult
from app.service.spec.step import Step
from app.service.spec.workflow import Workflow
from app.service.task.task_executor import TaskExecutor
from app.service.task.task_param import TaskParam
from app.service.task.task_result import TaskResult
from app.util.task_util import create_task_workspace, cleanup_workspace

logger = logging.getLogger(__name__)


class WorkflowEngine:
    _BASE64_DECODE_TYPE = "base64_decode"
    _BASE64_MISSING_FILES_ERROR = "No files provided for base64 decoding"
    _MAX_PREVIOUS_STEP_RERUNS = 3

    def __init__(self, task_executor: TaskExecutor):
        self._task_executor = task_executor
        self._resolver = ParameterResolver()

    def execute(
        self, task_id: str, workflow: Workflow, inputs: dict[str, Any]
    ) -> WorkflowResult:
        total_steps = len(workflow.order)
        logger.info("========== Workflow start ==========  version=%s, total_steps=%d",
                     workflow.version, total_steps)
        wf_t0 = time.time()
        workspace: str | None = None

        ctx = ExecutionContext(inputs)

        try:
            workspace = create_task_workspace(task_id, workflow.template_path)
            logger.info("Workspace created: %s", workspace)
        except Exception as e:
            logger.error("Failed to create workspace: %s", e)
            return WorkflowResult.failure(f"failed: {e}")

        try:
            for idx, step_id in enumerate(workflow.order):
                step = workflow.steps.get(step_id)
                if step is None:
                    logger.warning("Step '%s' not found in steps definition, skipping", step_id)
                    continue

                step_no = idx + 1
                logger.info("------ Step %d/%d [%s] type=%s start ------",
                             step_no, total_steps, step.id, step.type)
                logger.info("[%s] raw_inputs=%s", step.id, step.inputs)

                step_t0 = time.time()
                try:
                    step_dir = os.path.join(workspace, step.id)
                    os.makedirs(step_dir, exist_ok=True)
                    result = self._execute_step_with_retry(task_id, workspace, step, ctx)
                    step_elapsed = time.time() - step_t0

                    if result.is_failure():
                        recovered_result = self._recover_base64_by_rerunning_previous_step(
                            task_id=task_id,
                            workspace=workspace,
                            workflow=workflow,
                            current_step_index=idx,
                            current_step=step,
                            ctx=ctx,
                            failed_result=result,
                        )
                        if recovered_result is not None:
                            result = recovered_result
                            step_elapsed = time.time() - step_t0

                    if result.is_failure():
                        logger.error("[%s] FAILED in %.2fs after retries. Error: %s",
                                     step_id, step_elapsed, result.error_message)
                        return WorkflowResult.failure(
                            f"Step {step_id} failed: {result.error_message}"
                        )
                except Exception as e:
                    step_elapsed = time.time() - step_t0
                    logger.error("[%s] EXCEPTION in %.2fs: %s", step_id, step_elapsed, e, exc_info=True)
                    return WorkflowResult.failure(f"Step {step_id} failed: {e}")

                logger.info("[%s] SUCCESS in %.2fs  outputs=%s",
                             step.id, step_elapsed, _shorten(result.outputs))
                ctx.put_step_result(step_id, result.outputs)

            logger.info("All steps completed, resolving final outputs ...")
            outputs: dict[str, Any] = {}
            if workflow.outputs:
                for key, value in workflow.outputs.items():
                    outputs[key] = self._resolver.resolve(value, ctx)
            logger.info("Final output_keys=%s", list(outputs.keys()))

            total_elapsed = time.time() - wf_t0
            logger.info("========== Workflow end ==========  total_elapsed=%.2fs, success=True",
                         total_elapsed)
            return WorkflowResult.ok(outputs)
        finally:
            if workspace:
                cleanup_workspace(workspace, task_id)

    def _execute_step_with_retry(
        self, task_id: str, workspace: str, step: Step, ctx: ExecutionContext
    ) -> TaskResult:
        max_retries = (
            settings.step_retry_count if step.retry_count is None else step.retry_count
        )
        last_result: TaskResult | None = None

        for attempt in range(max_retries + 1):
            step_inputs = self._resolve_inputs(step.inputs, ctx)
            task = self._task_executor.resolve(step.type)
            param = TaskParam(
                workspace=workspace,
                step_id=step.id,
                type=step.type,
                resources=step.resources,
                inputs=step_inputs,
            )

            if attempt > 0:
                logger.info("[%s] Retry attempt %d/%d",
                             step.id, attempt, max_retries)

            logger.info("[%s] Executing task type=%s, resolved_input_keys=%s",
                         step.id, step.type, list(step_inputs.keys()))

            attempt_t0 = time.time()
            last_result = task.execute(param)
            attempt_elapsed = time.time() - attempt_t0

            if last_result.is_success():
                logger.info("[%s] Task returned SUCCESS in %.2fs",
                             step.id, attempt_elapsed)
                return last_result

            if attempt < max_retries:
                logger.warning(
                    "[%s] Task returned FAILURE in %.2fs (attempt %d/%d), will retry. Error: %s",
                    step.id, attempt_elapsed,
                    attempt + 1, max_retries + 1, last_result.error_message,
                )
                time.sleep(1)
            else:
                logger.error(
                    "[%s] Task returned FAILURE in %.2fs (final attempt %d/%d). Error: %s",
                    step.id, attempt_elapsed,
                    attempt + 1, max_retries + 1, last_result.error_message,
                )

        return last_result  # type: ignore[return-value]

    def _resolve_inputs(
        self, inputs: dict[str, Any], ctx: ExecutionContext
    ) -> dict[str, Any]:
        return {k: self._resolver.resolve(v, ctx) for k, v in inputs.items()} if inputs else {}

    def _recover_base64_by_rerunning_previous_step(
        self,
        task_id: str,
        workspace: str,
        workflow: Workflow,
        current_step_index: int,
        current_step: Step,
        ctx: ExecutionContext,
        failed_result: TaskResult,
    ) -> TaskResult | None:
        if current_step.type != self._BASE64_DECODE_TYPE:
            return None
        if self._BASE64_MISSING_FILES_ERROR not in failed_result.error_message:
            return None
        if current_step_index <= 0:
            logger.error(
                "[%s] Missing files in base64_decode but no previous step exists",
                current_step.id,
            )
            return None

        previous_step_id = workflow.order[current_step_index - 1]
        previous_step = workflow.steps.get(previous_step_id)
        if previous_step is None:
            logger.error(
                "[%s] Missing files in base64_decode but previous step '%s' not found",
                current_step.id,
                previous_step_id,
            )
            return None

        current_result = failed_result
        for recover_attempt in range(1, self._MAX_PREVIOUS_STEP_RERUNS + 1):
            logger.warning(
                "[%s] Trigger fallback: rerun previous step [%s] (attempt %d/%d) "
                "because base64_decode has no files",
                current_step.id,
                previous_step.id,
                recover_attempt,
                self._MAX_PREVIOUS_STEP_RERUNS,
            )

            previous_result = self._execute_step_with_retry(task_id, workspace, previous_step, ctx)
            if previous_result.is_failure():
                logger.warning(
                    "[%s] Previous step rerun failed on attempt %d/%d: %s",
                    previous_step.id,
                    recover_attempt,
                    self._MAX_PREVIOUS_STEP_RERUNS,
                    previous_result.error_message,
                )
                continue

            ctx.put_step_result(previous_step.id, previous_result.outputs)
            logger.info(
                "[%s] Previous step rerun succeeded on attempt %d/%d, rerunning current step",
                previous_step.id,
                recover_attempt,
                self._MAX_PREVIOUS_STEP_RERUNS,
            )

            current_result = self._execute_step_with_retry(task_id, workspace, current_step, ctx)
            if current_result.is_success():
                logger.info(
                    "[%s] Recovery succeeded after rerunning previous step [%s]",
                    current_step.id,
                    previous_step.id,
                )
                return current_result

            if self._BASE64_MISSING_FILES_ERROR not in current_result.error_message:
                logger.error(
                    "[%s] Current step rerun failed with non-recoverable error: %s",
                    current_step.id,
                    current_result.error_message,
                )
                return current_result

        logger.error(
            "[%s] Recovery failed after %d attempts of rerunning previous step [%s]",
            current_step.id,
            self._MAX_PREVIOUS_STEP_RERUNS,
            previous_step.id,
        )
        return current_result


def _shorten(d: dict | None, max_len: int = 120) -> str:
    if d is None:
        return "null"
    s = str(d)
    return s if len(s) <= max_len else s[:max_len] + "..."