92 lines
3.0 KiB
Python
92 lines
3.0 KiB
Python
|
from enum import Enum
|
||
|
|
from typing import List, Dict, Any, Optional
|
||
|
|
from dataclasses import dataclass, field
|
||
|
|
|
||
|
|
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]
|