Underground-Digital/Workflow-Engine

1import uuid2from datetime import datetime, timezone3from enum import Enum4from typing import Optional5 6from pydantic import BaseModel, Field7 8from core.workflow.entities.node_entities import NodeRunResult9from models.workflow import WorkflowNodeExecutionStatus10 11 12class RouteNodeState(BaseModel):13    class Status(Enum):14        RUNNING = "running"15        SUCCESS = "success"16        FAILED = "failed"17        PAUSED = "paused"18 19    id: str = Field(default_factory=lambda: str(uuid.uuid4()))20    """node state id"""21 22    node_id: str23    """node id"""24 25    node_run_result: Optional[NodeRunResult] = None26    """node run result"""27 28    status: Status = Status.RUNNING29    """node status"""30 31    start_at: datetime32    """start time"""33 34    paused_at: Optional[datetime] = None35    """paused time"""36 37    finished_at: Optional[datetime] = None38    """finished time"""39 40    failed_reason: Optional[str] = None41    """failed reason"""42 43    paused_by: Optional[str] = None44    """paused by"""45 46    index: int = 147 48    def set_finished(self, run_result: NodeRunResult) -> None:49        """50        Node finished51 52        :param run_result: run result53        """54        if self.status in {RouteNodeState.Status.SUCCESS, RouteNodeState.Status.FAILED}:55            raise Exception(f"Route state {self.id} already finished")56 57        if run_result.status == WorkflowNodeExecutionStatus.SUCCEEDED:58            self.status = RouteNodeState.Status.SUCCESS59        elif run_result.status == WorkflowNodeExecutionStatus.FAILED:60            self.status = RouteNodeState.Status.FAILED61            self.failed_reason = run_result.error62        else:63            raise Exception(f"Invalid route status {run_result.status}")64 65        self.node_run_result = run_result66        self.finished_at = datetime.now(timezone.utc).replace(tzinfo=None)67 68 69class RuntimeRouteState(BaseModel):70    routes: dict[str, list[str]] = Field(71        default_factory=dict, description="graph state routes (source_node_state_id: target_node_state_id)"72    )73 74    node_state_mapping: dict[str, RouteNodeState] = Field(75        default_factory=dict, description="node state mapping (route_node_state_id: route_node_state)"76    )77 78    def create_node_state(self, node_id: str) -> RouteNodeState:79        """80        Create node state81 82        :param node_id: node id83        """84        state = RouteNodeState(node_id=node_id, start_at=datetime.now(timezone.utc).replace(tzinfo=None))85        self.node_state_mapping[state.id] = state86        return state87 88    def add_route(self, source_node_state_id: str, target_node_state_id: str) -> None:89        """90        Add route to the graph state91 92        :param source_node_state_id: source node state id93        :param target_node_state_id: target node state id94        """95        if source_node_state_id not in self.routes:96            self.routes[source_node_state_id] = []97 98        self.routes[source_node_state_id].append(target_node_state_id)99 100    def get_routes_with_node_state_by_source_node_state_id(self, source_node_state_id: str) -> list[RouteNodeState]:101        """102        Get routes with node state by source node id103 104        :param source_node_state_id: source node state id105        :return: routes with node state106        """107        return [108            self.node_state_mapping[target_state_id] for target_state_id in self.routes.get(source_node_state_id, [])109        ]110