Files
2026-07-13 13:06:23 +08:00

147 lines
4.4 KiB
Python

from copy import copy
from enum import Enum
from dataclasses import dataclass, field
from datetime import datetime
from typing import Any, Awaitable, Optional, Union, Callable, TypedDict, Literal
from .utils import (
string_to_md5_hash,
generate_uuid,
function_or_method_to_string,
function_or_method_to_repr,
)
class ReturnBehavior(Enum):
DISPATCH = "dispatch"
GOTO = "goto"
ABORT = "abort"
INPUT = "input"
class TaskStatus(Enum):
RUNNING = "running"
SUCCESS = "success"
FAILURE = "failure"
PENDING = "pending"
class InvokeInterCache(TypedDict):
result: Any
already_sent_to_event_group: set[str]
GroupEventReturns = dict[str, Any]
@dataclass
class EventGroupInput:
group_name: str
results: GroupEventReturns
behavior: ReturnBehavior = ReturnBehavior.DISPATCH
@dataclass
class EventInput(EventGroupInput):
task_id: str = field(default_factory=generate_uuid)
@classmethod
def from_input(cls: "EventInput", input_data: dict[str, Any]) -> "EventInput":
return cls(
group_name="user_input", results=input_data, behavior=ReturnBehavior.INPUT
)
@dataclass
class _SpecialEventReturn:
behavior: ReturnBehavior
returns: Any
def __post_init__(self):
if not isinstance(self.behavior, ReturnBehavior):
raise TypeError(
f"behavior must be a ReturnBehavior, not {type(self.behavior)}"
)
# (group_event_results, global ctx set by user) -> result
EventFunction = Callable[
[Optional[EventInput], Optional[Any]], Awaitable[Union[Any, _SpecialEventReturn]]
]
@dataclass
class EventGroup:
name: str
events_hash: str
events: dict[str, "BaseEvent"]
retrigger_type: Literal["all", "any"] = "all"
def hash(self) -> str:
return self.events_hash
class BaseEvent:
parent_groups: dict[str, EventGroup]
func_inst: EventFunction
id: str
repr_name: str
def __init__(
self,
func_inst: EventFunction,
parent_groups: Optional[dict[str, EventGroup]] = None,
):
self.parent_groups = parent_groups or {}
self.func_inst = func_inst
self.id = string_to_md5_hash(function_or_method_to_string(self.func_inst))
self.repr_name = function_or_method_to_repr(self.func_inst)
self.meta = {"func_body": function_or_method_to_string(self.func_inst)}
def debug_string(self, exclude_events: Optional[set[str]] = None) -> str:
exclude_events = exclude_events or set([self.id])
parents_str = format_parents(self.parent_groups, exclude_events=exclude_events)
return f"{self.repr_name}\n{parents_str}"
def __repr__(self) -> str:
return f"Node(source={self.repr_name})"
async def solo_run(
self, event_input: EventInput, global_ctx: Any = None
) -> Awaitable[Any]:
return await self.func_inst(event_input, global_ctx)
@dataclass
class Task:
task_id: str
status: TaskStatus = TaskStatus.PENDING
created_at: datetime = field(default_factory=datetime.now)
upated_at: datetime = field(default_factory=datetime.now)
def format_parents(parents: dict[str, EventGroup], exclude_events: set[str], indent=""):
# Below code is ugly
# But it works and only for debug display
result = []
for i, parent_group in enumerate(parents.values()):
is_last_group = i == len(parents) - 1
group_prefix = "└─ " if is_last_group else "├─ "
result.append(indent + group_prefix + f"<{parent_group.name}>")
for j, parent in enumerate(parent_group.events.values()):
root_events = copy(exclude_events)
is_last = j == len(parent_group.events) - 1
child_indent = indent + (" " if is_last_group else "│ ")
inter_indent = " " if is_last else "│ "
prefix = "└─ " if is_last else "├─ "
if parent.id in root_events:
result.append(f"{child_indent}{prefix}{parent.repr_name} <loop>")
continue
root_events.add(parent.id)
parent_debug = parent.debug_string(exclude_events=root_events).split("\n")
parent_debug = [p for p in parent_debug if p.strip()]
result.append(f"{child_indent}{prefix}{parent.repr_name}")
for line in parent_debug[1:]:
result.append(f"{child_indent}{inter_indent}{line}")
return "\n".join(result)