import time
import re
from typing import Dict, Any, List, Callable, Optional
from concurrent.futures import ThreadPoolExecutor, as_completed
from .workflow_definition import Workflow, WorkflowStep, ExecutionMode
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
Reference in New Issue View Git Blame Copy Permalink