2227 Workflow State Machine Afd19371

Workflow State Machine

Workflow State Machine

from enum import Enum
from dataclasses import dataclass, field
from typing import Dict, Optional, Callable, Any
from datetime import datetime
import json

class WorkflowState(Enum):
    CREATED = "created"
    VALIDATED = "validated"
    QUEUED = "queued"
    GENERATING = "generating"
    POST_PROCESSING = "post_processing"
    COMPLETED = "completed"
    FAILED = "failed"
    CANCELLED = "cancelled"

@dataclass
class WorkflowJob:
    id: str
    prompt: str
    state: WorkflowState = WorkflowState.CREATED
    params: Dict = field(default_factory=dict)
    klingai_job_id: Optional[str] = None
    video_url: Optional[str] = None
    processed_url: Optional[str] = None
    error: Optional[str] = None
    metadata: Dict = field(default_factory=dict)
    created_at: datetime = field(default_factory=datetime.utcnow)
    updated_at: datetime = field(default_factory=datetime.utcnow)
    history: list = field(default_factory=list)

class WorkflowEngine:
    """State machine for video generation workflows."""

    # Valid state transitions
    TRANSITIONS = {
        WorkflowState.CREATED: [WorkflowState.VALIDATED, WorkflowState.FAILED],
        WorkflowState.VALIDATED: [WorkflowState.QUEUED, WorkflowState.FAILED],
        WorkflowState.QUEUED: [WorkflowState.GENERATING, WorkflowState.CANCELLED],
        WorkflowState.GENERATING: [WorkflowState.POST_PROCESSING, WorkflowState.FAILED],
        WorkflowState.POST_PROCESSING: [WorkflowState.COMPLETED, WorkflowState.FAILED],
        WorkflowState.COMPLETED: [],  # Terminal state
        WorkflowState.FAILED: [],     # Terminal state
        WorkflowState.CANCELLED: [],  # Terminal state
    }

    def __init__(self):
        self.handlers: Dict[WorkflowState, Callable] = {}
        self.on_transition: Optional[Callable] = None

    def register_handler(self, state: WorkflowState, handler: Callable):
        """Register a handler for a workflow state."""
        self.handlers[state] = handler

    def transition(self, job: WorkflowJob, new_state: WorkflowState, **kwargs):
        """Transition job to new state."""
        if new_state not in self.TRANSITIONS.get(job.state, []):
            raise ValueError(
                f"Invalid transition: {job.state.value} -> {new_state.value}"
            )

        old_state = job.state
        job.state = new_state
        job.updated_at = datetime.utcnow()
        job.history.append({
            "from": old_state.value,
            "to": new_state.value,
            "timestamp": job.updated_at.isoformat(),
            "data": kwargs
        })

        # Update job with any additional data
        for key, value in kwargs.items():
            if hasattr(job, key):
                setattr(job, key, value)

        if self.on_transition:
            self.on_transition(job, old_state, new_state)

    async def process(self, job: WorkflowJob) -> WorkflowJob:
        """Process job through workflow."""
        while job.state not in [
            WorkflowState.COMPLETED,
            WorkflowState.FAILED,
            WorkflowState.CANCELLED
        ]:
            handler = self.handlers.get(job.state)
            if not handler:
                raise RuntimeError(f"No handler for state: {job.state.value}")

            try:
                await handler(self, job)
            except Exception as e:
                self.transition(job, WorkflowState.FAILED, error=str(e))

        return job