Spaces:
Paused
Paused
File size: 6,423 Bytes
b9f94e1 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 185 | """Workflow engine implementation."""
from __future__ import annotations
import asyncio
import json
import logging
import time
from datetime import UTC, datetime
from pathlib import Path
from typing import TYPE_CHECKING, Any
from hermes.config.settings import get_settings
from hermes.core.types import TaskStatus
if TYPE_CHECKING:
from collections.abc import Callable, Coroutine
logger = logging.getLogger(__name__)
class WorkflowStep:
"""A step in a workflow."""
def __init__(
self,
name: str,
handler: Callable[..., Coroutine[Any, Any, Any]],
dependencies: list[str] | None = None,
retry_count: int = 3,
timeout: float = 300.0,
) -> None:
self.name = name
self.handler = handler
self.dependencies = dependencies or []
self.retry_count = retry_count
self.timeout = timeout
self.status: TaskStatus = TaskStatus.PENDING
self.result: Any = None
self.error: str | None = None
self.start_time: float | None = None
self.end_time: float | None = None
async def execute(self, context: dict[str, Any]) -> Any:
"""Execute the step."""
self.status = TaskStatus.RUNNING
self.start_time = time.monotonic()
for attempt in range(self.retry_count):
try:
self.result = await asyncio.wait_for(
self.handler(context), timeout=self.timeout
)
self.status = TaskStatus.COMPLETED
self.end_time = time.monotonic()
return self.result
except TimeoutError:
logger.warning(f"Step {self.name} timed out (attempt {attempt + 1})")
if attempt == self.retry_count - 1:
self.status = TaskStatus.FAILED
self.error = f"Timeout after {self.timeout}s"
self.end_time = time.monotonic()
raise
except Exception as e:
logger.warning(f"Step {self.name} failed (attempt {attempt + 1}): {e}")
if attempt == self.retry_count - 1:
self.status = TaskStatus.FAILED
self.error = str(e)
self.end_time = time.monotonic()
raise
@property
def duration_ms(self) -> float:
"""Get step duration in milliseconds."""
if self.start_time and self.end_time:
return (self.end_time - self.start_time) * 1000
return 0.0
class Workflow:
"""A workflow definition."""
def __init__(self, name: str, description: str = "") -> None:
self.name = name
self.description = description
self.steps: list[WorkflowStep] = []
self.context: dict[str, Any] = {}
self.status: TaskStatus = TaskStatus.PENDING
self.created_at = datetime.now(UTC)
def add_step(self, step: WorkflowStep) -> None:
"""Add a step to the workflow."""
self.steps.append(step)
async def execute(self, initial_context: dict[str, Any] | None = None) -> dict[str, Any]:
"""Execute the workflow."""
self.status = TaskStatus.RUNNING
self.context = initial_context or {}
try:
completed: set[str] = set()
while len(completed) < len(self.steps):
ready = [
s
for s in self.steps
if s.name not in completed
and s.status == TaskStatus.PENDING
and all(dep in completed for dep in s.dependencies)
]
if not ready:
if any(s.status == TaskStatus.FAILED for s in self.steps):
self.status = TaskStatus.FAILED
break
await asyncio.gather(
*[step.execute(self.context) for step in ready],
return_exceptions=True,
)
for step in ready:
completed.add(step.name)
if step.result:
self.context[step.name] = step.result
if self.status == TaskStatus.RUNNING:
self.status = TaskStatus.COMPLETED
except Exception as e:
self.status = TaskStatus.FAILED
logger.error(f"Workflow {self.name} failed: {e}")
raise
return self.context
class WorkflowEngine:
"""Workflow execution engine with checkpointing."""
def __init__(self) -> None:
self.settings = get_settings()
self._workflows: dict[str, Workflow] = {}
self._checkpoint_dir = Path("data/checkpoints")
async def register_workflow(self, workflow: Workflow) -> None:
"""Register a workflow."""
self._workflows[workflow.name] = workflow
async def execute_workflow(
self, name: str, context: dict[str, Any] | None = None
) -> dict[str, Any]:
"""Execute a registered workflow."""
workflow = self._workflows.get(name)
if not workflow:
raise ValueError(f"Workflow not found: {name}")
await self._save_checkpoint(workflow, "start")
result = await workflow.execute(context)
await self._save_checkpoint(workflow, "complete")
return result
async def _save_checkpoint(self, workflow: Workflow, stage: str) -> None:
"""Save workflow checkpoint."""
try:
self._checkpoint_dir.mkdir(parents=True, exist_ok=True)
checkpoint = {
"workflow": workflow.name,
"stage": stage,
"status": workflow.status.value,
"timestamp": datetime.now(UTC).isoformat(),
"context": workflow.context,
}
path = self._checkpoint_dir / f"{workflow.name}_{stage}.json"
path.write_text(json.dumps(checkpoint, indent=2, default=str), encoding="utf-8")
except Exception as e:
logger.warning(f"Could not save checkpoint: {e}")
def get_workflow(self, name: str) -> Workflow | None:
"""Get a workflow by name."""
return self._workflows.get(name)
|