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

import logging
import os
import time
from typing import Any

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:
    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("[%s] ========== Workflow start ==========  version=%s, total_steps=%d",
                     task_id, workflow.version, total_steps)
        wf_t0 = time.time()

        ctx = ExecutionContext(inputs)

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

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

            logger.info("[%s] ------ Step %d/%d [%s] type=%s start ------",
                         task_id, idx, total_steps, step.id, step.type)
            logger.info("[%s] [%s] raw_inputs=%s", task_id, 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():
                    logger.error("[%s] [%s] FAILED in %.2fs after retries. Error: %s",
                                  task_id, 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] [%s] EXCEPTION in %.2fs: %s", task_id, step_id, step_elapsed, e, exc_info=True)
                return WorkflowResult.failure(f"Step {step_id} failed: {e}")

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

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

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

    def _execute_step_with_retry(
        self, task_id: str, workspace: str, step: Step, ctx: ExecutionContext
    ) -> TaskResult:
        max_retries = 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] [%s] Retry attempt %d/%d",
                             task_id, step.id, attempt, max_retries)

            logger.info("[%s] [%s] Executing task type=%s, resolved_input_keys=%s",
                         task_id, 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] [%s] Task returned SUCCESS in %.2fs",
                             task_id, step.id, attempt_elapsed)
                return last_result

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