rp/pr/workflows/workflow_engine.py

import re
import time
from concurrent.futures import ThreadPoolExecutor, as_completed
from typing import Any, Callable, Dict, List, Optional

from .workflow_definition import ExecutionMode, Workflow, WorkflowStep


class WorkflowExecutionContext:
    def __init__(self):
        self.variables: Dict[str, Any] = {}
        self.step_results: Dict[str, Any] = {}
        self.execution_log: List[Dict[str, Any]] = []

    def set_variable(self, name: str, value: Any):
        self.variables[name] = value

    def get_variable(self, name: str, default: Any = None) -> Any:
        return self.variables.get(name, default)

    def set_step_result(self, step_id: str, result: Any):
        self.step_results[step_id] = result

    def get_step_result(self, step_id: str) -> Any:
        return self.step_results.get(step_id)

    def log_event(self, event_type: str, step_id: str, details: Dict[str, Any]):
        self.execution_log.append(
            {
                "timestamp": time.time(),
                "event_type": event_type,
                "step_id": step_id,
                "details": details,
            }
        )


class WorkflowEngine:
    def __init__(self, tool_executor: Callable, max_workers: int = 5):
        self.tool_executor = tool_executor
        self.max_workers = max_workers

    def _evaluate_condition(
        self, condition: str, context: WorkflowExecutionContext
    ) -> bool:
        if not condition:
            return True

        try:
            safe_locals = {
                "variables": context.variables,
                "results": context.step_results,
            }
            return eval(condition, {"__builtins__": {}}, safe_locals)
        except Exception:
            return False

    def _substitute_variables(
        self, arguments: Dict[str, Any], context: WorkflowExecutionContext
    ) -> Dict[str, Any]:
        substituted = {}
        for key, value in arguments.items():
            if isinstance(value, str):
                pattern = r"\$\{([^}]+)\}"
                matches = re.findall(pattern, value)
                for match in matches:
                    if match.startswith("step."):
                        step_id = match.split(".", 1)[1]
                        replacement = context.get_step_result(step_id)
                        if replacement is not None:
                            value = value.replace(f"${{{match}}}", str(replacement))
                    elif match.startswith("var."):
                        var_name = match.split(".", 1)[1]
                        replacement = context.get_variable(var_name)
                        if replacement is not None:
                            value = value.replace(f"${{{match}}}", str(replacement))
                substituted[key] = value
            else:
                substituted[key] = value
        return substituted

    def _execute_step(
        self, step: WorkflowStep, context: WorkflowExecutionContext
    ) -> Dict[str, Any]:
        if not self._evaluate_condition(step.condition, context):
            context.log_event("skipped", step.step_id, {"reason": "condition_not_met"})
            return {"status": "skipped", "step_id": step.step_id}

        arguments = self._substitute_variables(step.arguments, context)

        start_time = time.time()
        retry_attempts = 0
        last_error = None

        while retry_attempts <= step.retry_count:
            try:
                context.log_event(
                    "executing",
                    step.step_id,
                    {
                        "tool": step.tool_name,
                        "arguments": arguments,
                        "attempt": retry_attempts + 1,
                    },
                )

                result = self.tool_executor(step.tool_name, arguments)

                execution_time = time.time() - start_time
                context.set_step_result(step.step_id, result)
                context.log_event(
                    "completed",
                    step.step_id,
                    {
                        "execution_time": execution_time,
                        "result_size": len(str(result)) if result else 0,
                    },
                )

                return {
                    "status": "success",
                    "step_id": step.step_id,
                    "result": result,
                    "execution_time": execution_time,
                }

            except Exception as e:
                last_error = str(e)
                retry_attempts += 1
                if retry_attempts <= step.retry_count:
                    time.sleep(1 * retry_attempts)

        context.log_event("failed", step.step_id, {"error": last_error})
        return {
            "status": "failed",
            "step_id": step.step_id,
            "error": last_error,
            "execution_time": time.time() - start_time,
        }

    def _get_next_steps(
        self, completed_step: WorkflowStep, result: Dict[str, Any], workflow: Workflow
    ) -> List[WorkflowStep]:
        next_steps = []

        if result["status"] == "success" and completed_step.on_success:
            for step_id in completed_step.on_success:
                step = workflow.get_step(step_id)
                if step:
                    next_steps.append(step)

        elif result["status"] == "failed" and completed_step.on_failure:
            for step_id in completed_step.on_failure:
                step = workflow.get_step(step_id)
                if step:
                    next_steps.append(step)

        elif workflow.execution_mode == ExecutionMode.SEQUENTIAL:
            current_index = workflow.steps.index(completed_step)
            if current_index + 1 < len(workflow.steps):
                next_steps.append(workflow.steps[current_index + 1])

        return next_steps

    def execute_workflow(
        self, workflow: Workflow, initial_variables: Optional[Dict[str, Any]] = None
    ) -> WorkflowExecutionContext:
        context = WorkflowExecutionContext()

        if initial_variables:
            context.variables.update(initial_variables)

        if workflow.variables:
            context.variables.update(workflow.variables)

        context.log_event("workflow_started", "workflow", {"name": workflow.name})

        if workflow.execution_mode == ExecutionMode.PARALLEL:
            with ThreadPoolExecutor(max_workers=self.max_workers) as executor:
                futures = {
                    executor.submit(self._execute_step, step, context): step
                    for step in workflow.steps
                }

                for future in as_completed(futures):
                    step = futures[future]
                    try:
                        result = future.result()
                        context.log_event("step_completed", step.step_id, result)
                    except Exception as e:
                        context.log_event(
                            "step_failed", step.step_id, {"error": str(e)}
                        )

        else:
            pending_steps = workflow.get_initial_steps()
            executed_step_ids = set()

            while pending_steps:
                step = pending_steps.pop(0)

                if step.step_id in executed_step_ids:
                    continue

                result = self._execute_step(step, context)
                executed_step_ids.add(step.step_id)

                next_steps = self._get_next_steps(step, result, workflow)
                pending_steps.extend(next_steps)

        context.log_event(
            "workflow_completed",
            "workflow",
            {
                "total_steps": len(context.step_results),
                "executed_steps": list(context.step_results.keys()),
            },
        )

        return context