|
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]
|