rp/pr/workflows/workflow_definition.py

from dataclasses import dataclass, field
from enum import Enum
from typing import Any, Dict, List, Optional


class ExecutionMode(Enum):
    SEQUENTIAL = "sequential"
    PARALLEL = "parallel"
    CONDITIONAL = "conditional"


@dataclass
class WorkflowStep:
    tool_name: str
    arguments: Dict[str, Any]
    step_id: str
    condition: Optional[str] = None
    on_success: Optional[List[str]] = None
    on_failure: Optional[List[str]] = None
    retry_count: int = 0
    timeout_seconds: int = 300

    def to_dict(self) -> Dict[str, Any]:
        return {
            "tool_name": self.tool_name,
            "arguments": self.arguments,
            "step_id": self.step_id,
            "condition": self.condition,
            "on_success": self.on_success,
            "on_failure": self.on_failure,
            "retry_count": self.retry_count,
            "timeout_seconds": self.timeout_seconds,
        }

    @staticmethod
    def from_dict(data: Dict[str, Any]) -> "WorkflowStep":
        return WorkflowStep(
            tool_name=data["tool_name"],
            arguments=data["arguments"],
            step_id=data["step_id"],
            condition=data.get("condition"),
            on_success=data.get("on_success"),
            on_failure=data.get("on_failure"),
            retry_count=data.get("retry_count", 0),
            timeout_seconds=data.get("timeout_seconds", 300),
        )


@dataclass
class Workflow:
    name: str
    description: str
    steps: List[WorkflowStep]
    execution_mode: ExecutionMode = ExecutionMode.SEQUENTIAL
    variables: Dict[str, Any] = field(default_factory=dict)
    tags: List[str] = field(default_factory=list)

    def to_dict(self) -> Dict[str, Any]:
        return {
            "name": self.name,
            "description": self.description,
            "steps": [step.to_dict() for step in self.steps],
            "execution_mode": self.execution_mode.value,
            "variables": self.variables,
            "tags": self.tags,
        }

    @staticmethod
    def from_dict(data: Dict[str, Any]) -> "Workflow":
        return Workflow(
            name=data["name"],
            description=data["description"],
            steps=[WorkflowStep.from_dict(step) for step in data["steps"]],
            execution_mode=ExecutionMode(data.get("execution_mode", "sequential")),
            variables=data.get("variables", {}),
            tags=data.get("tags", []),
        )

    def add_step(self, step: WorkflowStep):
        self.steps.append(step)

    def get_step(self, step_id: str) -> Optional[WorkflowStep]:
        for step in self.steps:
            if step.step_id == step_id:
                return step
        return None

    def get_initial_steps(self) -> List[WorkflowStep]:
        if self.execution_mode == ExecutionMode.SEQUENTIAL:
            return [self.steps[0]] if self.steps else []
        elif self.execution_mode == ExecutionMode.PARALLEL:
            return self.steps
        else:
            return [step for step in self.steps if not step.condition]