Build latest book artifacts / build (push) Canceled after 0s
dependency resolution / resolve (3.11) (push) Canceled after 0s
dependency resolution / resolve (3.13) (push) Canceled after 0s
deploy-pages / build (push) Canceled after 0s
deploy-pages / deploy (push) Canceled after 0s
i18n consistency check / check (push) Canceled after 0s
provider adoption tests / test (chapter2/context-compression) (push) Canceled after 0s
provider adoption tests / test (chapter2/prompt-injection) (push) Canceled after 0s
provider adoption tests / test (chapter2/system-hint) (push) Canceled after 0s
provider adoption tests / test (chapter3/log-sanitization) (push) Canceled after 0s
web-search-agent tests / test (push) Canceled after 0s
web-search-agent tests / agentbook (push) Canceled after 0s
825 lines
28 KiB
Python
825 lines
28 KiB
Python
import abc
|
|
import time
|
|
import asyncio
|
|
from pydantic import BaseModel
|
|
from typing import Optional, List
|
|
from aworld.core.event.base import Message
|
|
from enum import Enum
|
|
from abc import ABC, abstractmethod, ABCMeta
|
|
from aworld.core.agent.base import is_agent_by_name
|
|
from aworld.core.tool.tool_desc import is_tool_by_name
|
|
from aworld.core.singleton import InheritanceSingleton, SingletonMeta
|
|
from aworld.core.event.base import Constants
|
|
from aworld.logs.util import logger
|
|
from aworld.events.util import send_message
|
|
|
|
|
|
class RunNodeBusiType(Enum):
|
|
AGENT = 'AGENT'
|
|
TOOL = 'TOOL'
|
|
TASK = 'TASK'
|
|
TOOL_CALLBACK = 'TOOL_CALLBACK'
|
|
HUMAN = 'HUMAN'
|
|
|
|
@staticmethod
|
|
def from_message_category(category: str) -> 'RunNodeBusiType':
|
|
if category == Constants.AGENT:
|
|
return RunNodeBusiType.AGENT
|
|
if category == Constants.TOOL:
|
|
return RunNodeBusiType.TOOL
|
|
if category == Constants.TASK:
|
|
return RunNodeBusiType.TASK
|
|
if category == Constants.TOOL_CALLBACK:
|
|
return RunNodeBusiType.TOOL_CALLBACK
|
|
if category == Constants.HUMAN:
|
|
return RunNodeBusiType.HUMAN
|
|
return None
|
|
|
|
|
|
class RunNodeStatus(Enum):
|
|
INIT = 'INIT'
|
|
RUNNING = 'RUNNING'
|
|
BREAKED = 'BREAKED'
|
|
SUCCESS = 'SUCCESS'
|
|
FAILED = 'FAILED'
|
|
TIMEOUT = 'TIMEOUT'
|
|
|
|
|
|
class HandleResult(BaseModel):
|
|
name: str = None
|
|
status: RunNodeStatus = None
|
|
result_msg: Optional[str] = None
|
|
result: Optional[Message] = None
|
|
|
|
|
|
class RunNode(BaseModel):
|
|
# {busi_id}_{busi_type}
|
|
node_id: Optional[str] = None
|
|
task_id: Optional[str] = None
|
|
busi_type: str = None
|
|
busi_id: Optional[str] = None
|
|
session_id: Optional[str] = None
|
|
msg_id: Optional[str] = None # input message id
|
|
# busi_id of node that send the input message
|
|
msg_from: Optional[str] = None
|
|
parent_node_id: Optional[str] = None
|
|
status: RunNodeStatus = None
|
|
result_msg: Optional[str] = None
|
|
results: Optional[List[HandleResult]] = None
|
|
create_time: Optional[float] = None
|
|
execute_time: Optional[float] = None
|
|
end_time: Optional[float] = None
|
|
group_id: Optional[str] = None
|
|
# sub_group_root_id required when group_id is not None
|
|
sub_group_root_id: Optional[str] = None
|
|
# metadata is used to store the context of the sub task when group_id is not None
|
|
metadata: Optional[dict] = None
|
|
|
|
def has_finished(self):
|
|
return self.status in [RunNodeStatus.SUCCESS, RunNodeStatus.FAILED, RunNodeStatus.TIMEOUT]
|
|
|
|
|
|
class SubGroup(BaseModel):
|
|
'''
|
|
SubGroup represents an execution chain pointing to the root node
|
|
'''
|
|
root_node_id: Optional[str] = None
|
|
session_id: Optional[str] = None
|
|
group_id: Optional[str] = None
|
|
create_time: Optional[float] = None
|
|
execute_time: Optional[float] = None
|
|
end_time: Optional[float] = None
|
|
status: RunNodeStatus = None
|
|
result_msg: Optional[str] = None
|
|
results: Optional[List[HandleResult]] = None
|
|
metadata: Optional[dict] = None
|
|
|
|
def has_finished(self):
|
|
return self.status in [RunNodeStatus.SUCCESS, RunNodeStatus.FAILED, RunNodeStatus.TIMEOUT]
|
|
|
|
|
|
class NodeGroup(BaseModel):
|
|
'''
|
|
Node group, used to manage sub group
|
|
'''
|
|
group_id: str = None
|
|
session_id: str = None
|
|
# subtask root node id list
|
|
root_node_ids: List[str] = None
|
|
finished: Optional[bool] = False
|
|
finish_notified: Optional[bool] = False
|
|
create_time: Optional[float] = None
|
|
execute_time: Optional[float] = None
|
|
end_time: Optional[float] = None
|
|
status: RunNodeStatus = None
|
|
# failed subtask root node id list
|
|
failed_root_node_ids: Optional[List[str]] = None
|
|
parent_group_id: Optional[str] = None
|
|
metadata: Optional[dict] = None
|
|
|
|
def has_finished(self):
|
|
return self.status in [RunNodeStatus.SUCCESS, RunNodeStatus.FAILED, RunNodeStatus.TIMEOUT]
|
|
|
|
|
|
class NodeGroupDetail(NodeGroup):
|
|
sub_groups: Optional[List[SubGroup]] = None
|
|
|
|
|
|
class StateStorage:
|
|
__metaclass__ = abc.ABCMeta
|
|
|
|
@abstractmethod
|
|
def get(self, node_id: str) -> RunNode:
|
|
pass
|
|
|
|
@abstractmethod
|
|
def insert(self, node: RunNode):
|
|
pass
|
|
|
|
@abstractmethod
|
|
def update(self, node: RunNode):
|
|
pass
|
|
|
|
@abstractmethod
|
|
def query(self, session_id: str) -> List[RunNode]:
|
|
pass
|
|
|
|
@abstractmethod
|
|
def query_by_task_id(self, task_id: str) -> List[RunNode]:
|
|
pass
|
|
|
|
|
|
class NodeGroupStorage:
|
|
__metaclass__ = abc.ABCMeta
|
|
|
|
@abstractmethod
|
|
def get(self, group_id: str) -> NodeGroup:
|
|
pass
|
|
|
|
@abstractmethod
|
|
def insert(self, node_group: NodeGroup):
|
|
pass
|
|
|
|
@abstractmethod
|
|
def update(self, node_group: NodeGroup):
|
|
pass
|
|
|
|
|
|
class SubGroupStorage:
|
|
__metaclass__ = abc.ABCMeta
|
|
|
|
@abstractmethod
|
|
def get(self, node_id: str) -> SubGroup:
|
|
pass
|
|
|
|
@abstractmethod
|
|
def insert(self, sub_group: SubGroup):
|
|
pass
|
|
|
|
@abstractmethod
|
|
def update(self, sub_group: SubGroup):
|
|
pass
|
|
|
|
|
|
class StateStorageMeta(SingletonMeta, ABCMeta):
|
|
pass
|
|
|
|
|
|
class InMemoryStateStorage(StateStorage, InheritanceSingleton, metaclass=StateStorageMeta):
|
|
'''
|
|
In memory state storage
|
|
'''
|
|
|
|
def __init__(self, max_session=1000):
|
|
self._max_session = max_session
|
|
self._nodes = {} # {node_id: RunNode}
|
|
self._ordered_session_ids = []
|
|
self._session_nodes = {} # {session_id: [RunNode, RunNode]}
|
|
|
|
def get(self, node_id: str) -> RunNode:
|
|
return self._nodes.get(node_id)
|
|
|
|
def insert(self, node: RunNode):
|
|
if node.session_id not in self._ordered_session_ids:
|
|
self._ordered_session_ids.append(node.session_id)
|
|
self._session_nodes.update({node.session_id: []})
|
|
if node.node_id not in self._nodes:
|
|
self._nodes.update({node.node_id: node})
|
|
self._session_nodes[node.session_id].append(node)
|
|
|
|
if len(self._ordered_session_ids) > self._max_session:
|
|
oldest_session_id = self._ordered_session_ids.pop(0)
|
|
session_nodes = self._session_nodes.pop(oldest_session_id)
|
|
for node in session_nodes:
|
|
self._nodes.pop(node.node_id)
|
|
# logger.info(f"storage nodes: {self._nodes}")
|
|
|
|
def update(self, node: RunNode):
|
|
self._nodes[node.node_id] = node
|
|
|
|
def query(self, session_id: str, msg_id: str = None) -> List[RunNode]:
|
|
session_nodes = self._session_nodes.get(session_id, [])
|
|
if msg_id:
|
|
return [node for node in session_nodes if node.msg_id == msg_id]
|
|
return session_nodes
|
|
|
|
def query_by_task_id(self, task_id: str) -> List[RunNode]:
|
|
return [node for node in self._nodes.values() if node.task_id == task_id]
|
|
|
|
|
|
class InMemoryNodeGroupStorage(NodeGroupStorage, InheritanceSingleton, metaclass=StateStorageMeta):
|
|
'''
|
|
In memory node group storage
|
|
'''
|
|
|
|
def __init__(self):
|
|
self.node_groups = {}
|
|
|
|
def get(self, group_id: str) -> NodeGroup:
|
|
return self.node_groups.get(group_id)
|
|
|
|
def insert(self, node_group: NodeGroup):
|
|
self.node_groups[node_group.group_id] = node_group
|
|
|
|
def update(self, node_group: NodeGroup):
|
|
self.node_groups[node_group.group_id] = node_group
|
|
|
|
|
|
class InMemorySubGroupStorage(SubGroupStorage, InheritanceSingleton, metaclass=StateStorageMeta):
|
|
'''
|
|
In memory sub task storage
|
|
'''
|
|
|
|
def __init__(self):
|
|
self.sub_groups = {}
|
|
|
|
def get(self, node_id: str) -> SubGroup:
|
|
return self.sub_groups.get(node_id)
|
|
|
|
def insert(self, sub_group: SubGroup):
|
|
self.sub_groups[sub_group.root_node_id] = sub_group
|
|
|
|
def update(self, sub_group: SubGroup):
|
|
self.sub_groups[sub_group.root_node_id] = sub_group
|
|
|
|
|
|
class RuntimeStateManager(InheritanceSingleton):
|
|
'''
|
|
Runtime state manager
|
|
'''
|
|
|
|
def __init__(self,
|
|
storage: StateStorage = InMemoryStateStorage.instance()):
|
|
self.storage = storage
|
|
self._node_group_manager = None
|
|
|
|
@property
|
|
def node_group_manager(self):
|
|
if not self._node_group_manager:
|
|
self._node_group_manager = NodeGroupManager(node_state_manager=self)
|
|
return self._node_group_manager
|
|
|
|
def create_node(self,
|
|
busi_type: RunNodeBusiType,
|
|
busi_id: str,
|
|
session_id: str,
|
|
node_id: str = None,
|
|
task_id: str = None,
|
|
parent_node_id: str = None,
|
|
msg_id: str = None,
|
|
msg_from: str = None,
|
|
group_id: str = None,
|
|
sub_group_root_id: str = None,
|
|
metadata: Optional[dict] = None) -> RunNode:
|
|
'''
|
|
create node and insert to storage
|
|
'''
|
|
node_id = node_id or msg_id
|
|
node = self._find_node(node_id)
|
|
if node:
|
|
# raise Exception(f"node already exist, node_id: {node_id}")
|
|
return
|
|
if parent_node_id:
|
|
parent_node = self._find_node(parent_node_id)
|
|
if not parent_node:
|
|
logger.warning(
|
|
f"parent node not exist, parent_node_id: {parent_node_id}")
|
|
node = RunNode(node_id=node_id,
|
|
busi_type=busi_type.name,
|
|
busi_id=busi_id,
|
|
session_id=session_id,
|
|
task_id=task_id,
|
|
msg_id=msg_id,
|
|
msg_from=msg_from,
|
|
parent_node_id=parent_node_id,
|
|
status=RunNodeStatus.INIT,
|
|
create_time=time.time(),
|
|
group_id=group_id,
|
|
sub_group_root_id=sub_group_root_id,
|
|
metadata=metadata)
|
|
self.storage.insert(node)
|
|
# create sub group if node is the root node of sub group
|
|
if group_id and sub_group_root_id and node_id == sub_group_root_id:
|
|
self.node_group_manager.create_sub_group(group_id, session_id, sub_group_root_id, metadata)
|
|
return node
|
|
|
|
def run_node(self, node_id: str):
|
|
'''
|
|
set node status to RUNNING and update to storage
|
|
'''
|
|
logger.debug(f"====== set node {node_id} running =======")
|
|
node = self._node_exist(node_id)
|
|
node.status = RunNodeStatus.RUNNING
|
|
node.execute_time = time.time()
|
|
self.storage.update(node)
|
|
# update sub group status if node is the root node
|
|
if node.group_id and node.sub_group_root_id and node.node_id == node.sub_group_root_id:
|
|
self.node_group_manager.run_sub_group(node_id)
|
|
|
|
def save_result(self,
|
|
node_id: str,
|
|
result: HandleResult):
|
|
'''
|
|
save node execute result and update to storage
|
|
'''
|
|
node = self._node_exist(node_id)
|
|
if not node.results:
|
|
node.results = []
|
|
node.results.append(result)
|
|
self.storage.update(node)
|
|
|
|
def break_node(self, node_id):
|
|
'''
|
|
set node status to BREAKED and update to storage
|
|
'''
|
|
node = self._node_exist(node_id)
|
|
node.status = RunNodeStatus.BREAKED
|
|
self.storage.update(node)
|
|
|
|
def run_succeed(self,
|
|
node_id,
|
|
result_msg=None,
|
|
results: List[HandleResult] = None):
|
|
'''
|
|
set node status to SUCCESS and update to storage
|
|
'''
|
|
node = self._node_exist(node_id)
|
|
node.status = RunNodeStatus.SUCCESS
|
|
node.result_msg = result_msg
|
|
node.end_time = time.time()
|
|
if results:
|
|
if not node.results:
|
|
node.results = []
|
|
node.results.extend(results)
|
|
logger.debug(f"====== run_succeed set node {node_id} succeed: {node} =======")
|
|
|
|
self.storage.update(node)
|
|
|
|
def run_failed(self,
|
|
node_id,
|
|
result_msg=None,
|
|
results: List[HandleResult] = None):
|
|
'''
|
|
set node status to FAILED and update to storage
|
|
'''
|
|
node = self._node_exist(node_id)
|
|
node.status = RunNodeStatus.FAILED
|
|
node.result_msg = result_msg
|
|
node.end_time = time.time()
|
|
if results:
|
|
if not node.results:
|
|
node.results = []
|
|
node.results.extend(results)
|
|
self.storage.update(node)
|
|
|
|
def run_timeout(self,
|
|
node_id,
|
|
result_msg=None):
|
|
'''
|
|
set node status to TIMEOUT and update to storage
|
|
'''
|
|
node = self._node_exist(node_id)
|
|
node.status = RunNodeStatus.TIMEOUT
|
|
node.result_msg = result_msg
|
|
self.storage.update(node)
|
|
|
|
def finish_sub_task(self, node_id: str):
|
|
'''
|
|
finish sub task with node_id as the root node
|
|
'''
|
|
node = self._node_exist(node_id)
|
|
node.sub_task_finished = True
|
|
self.storage.update(node)
|
|
|
|
def get_node(self, node_id: str) -> RunNode:
|
|
'''
|
|
get node from storage
|
|
'''
|
|
return self._find_node(node_id)
|
|
|
|
def get_nodes(self, session_id: str) -> List[RunNode]:
|
|
'''
|
|
get nodes from storage
|
|
'''
|
|
return self.storage.query(session_id)
|
|
|
|
def _node_exist(self, node_id: str):
|
|
node = self._find_node(node_id)
|
|
if not node:
|
|
raise Exception(f"node not found, node_id: {node_id}")
|
|
return node
|
|
|
|
def _find_node(self, node_id: str):
|
|
return self.storage.get(node_id)
|
|
|
|
def _judge_msg_from_busi_type(self, msg_from: str) -> RunNodeBusiType:
|
|
'''
|
|
judge msg_from busi_type
|
|
'''
|
|
if is_agent_by_name(msg_from):
|
|
return RunNodeBusiType.AGENT
|
|
if is_tool_by_name(msg_from):
|
|
return RunNodeBusiType.TOOL
|
|
return RunNodeBusiType.TASK
|
|
|
|
async def wait_for_node_completion(self, node_id: str, timeout: float = 600.0, interval: float = 1.0) -> RunNode:
|
|
'''Poll for node status until completion or timeout.
|
|
|
|
Args:
|
|
node_id: Node ID
|
|
timeout: Timeout threshold in seconds
|
|
interval: Polling interval in seconds
|
|
|
|
Returns:
|
|
RunNode: Node object
|
|
|
|
Raises:
|
|
Exception: If node does not exist
|
|
TimeoutError: If waiting times out
|
|
'''
|
|
start_time = time.time()
|
|
log_start_time = time.strftime("%Y-%m-%d %H:%M:%S", time.localtime())
|
|
logger.info(f"wait for node completion: {node_id}, start_time:{log_start_time}")
|
|
|
|
while True:
|
|
node = self._find_node(node_id)
|
|
if not node:
|
|
raise Exception(f"Node not found, node_id: {node_id}")
|
|
|
|
# Check if node has completed
|
|
if node.status in [RunNodeStatus.SUCCESS, RunNodeStatus.FAILED, RunNodeStatus.BREAKED,
|
|
RunNodeStatus.TIMEOUT]:
|
|
return node
|
|
|
|
# Check if timed out
|
|
if time.time() - start_time > timeout:
|
|
self.run_timeout(node_id, result_msg=f"Waiting for node completion timed out after {timeout} seconds")
|
|
node = self._find_node(node_id)
|
|
return node
|
|
|
|
# Wait for the specified interval before polling again
|
|
await asyncio.sleep(interval)
|
|
|
|
async def create_group(self, group_id: str,
|
|
session_id: str,
|
|
root_node_ids: List[str] = None,
|
|
parent_group_id: Optional[str] = None,
|
|
metadata: Optional[dict] = None) -> NodeGroup:
|
|
'''
|
|
create node group
|
|
'''
|
|
return await self.node_group_manager.create_group(group_id, session_id, root_node_ids, parent_group_id, metadata)
|
|
|
|
async def finish_sub_group(self,
|
|
group_id: str,
|
|
root_node_id: str,
|
|
results: List[Message] = None,
|
|
result_msg: str = None):
|
|
'''
|
|
finish sub group
|
|
'''
|
|
handle_results = []
|
|
for msg in results:
|
|
handle_result = HandleResult(
|
|
status=RunNodeStatus.FAILED if msg.is_error() else RunNodeStatus.SUCCESS,
|
|
result=msg,
|
|
name=msg.sender
|
|
)
|
|
handle_results.append(handle_result)
|
|
await self.node_group_manager.finish_sub_group(group_id, root_node_id, handle_results, result_msg)
|
|
|
|
def get_group(self, group_id: str) -> NodeGroup:
|
|
'''
|
|
get group basic info
|
|
'''
|
|
return self.node_group_manager.get_group(group_id)
|
|
|
|
def query_group_detail(self, group_id: str) -> NodeGroupDetail:
|
|
'''
|
|
query group detail info with all sub group info
|
|
'''
|
|
return self.node_group_manager.query_group_detail(group_id)
|
|
|
|
def query_by_task(self, task_id: str, busi_typ: RunNodeBusiType = None, busi_id: str = None) -> List[RunNode]:
|
|
all_task_nodes = self.storage.query_by_task_id(task_id)
|
|
if (not busi_typ and busi_id) or (busi_typ and not busi_id):
|
|
raise Exception("busi_typ and busi_id must be both None or not None")
|
|
if busi_typ and busi_id:
|
|
result_nodes = [node for node in all_task_nodes if node.busi_type == busi_typ.name and node.busi_id == busi_id]
|
|
else:
|
|
result_nodes = all_task_nodes
|
|
result_nodes.sort(key=lambda x: x.create_time if x.create_time else 0, reverse=True)
|
|
return result_nodes
|
|
|
|
|
|
class NodeGroupManager(InheritanceSingleton):
|
|
'''
|
|
Node group manager, used to manage node group
|
|
'''
|
|
|
|
def __init__(self,
|
|
sub_group_storage: SubGroupStorage = InMemorySubGroupStorage.instance(),
|
|
node_group_storage: NodeGroupStorage = InMemoryNodeGroupStorage.instance(),
|
|
node_state_manager: RuntimeStateManager = None):
|
|
self.sub_group_storage = sub_group_storage
|
|
self.node_group_storage = node_group_storage
|
|
self.node_state_manager = node_state_manager
|
|
|
|
async def create_group(self, group_id: str,
|
|
session_id: str,
|
|
root_node_ids: List[str] = None,
|
|
parent_group_id: Optional[str] = None,
|
|
metadata: Optional[dict] = None) -> NodeGroup:
|
|
'''
|
|
create node group
|
|
'''
|
|
group = self._find_group(group_id)
|
|
if group:
|
|
raise Exception(f"group already exist, group_id: {group_id}")
|
|
node_group = NodeGroup(
|
|
session_id=session_id,
|
|
group_id=group_id,
|
|
root_node_ids=root_node_ids,
|
|
parent_group_id=parent_group_id,
|
|
metadata=metadata,
|
|
create_time=time.time(),
|
|
update_time=time.time(),
|
|
status=RunNodeStatus.INIT,
|
|
)
|
|
self.node_group_storage.insert(node_group)
|
|
await self._check_subgroup_status(group_id, root_node_ids)
|
|
|
|
def create_sub_group(self,
|
|
group_id: str,
|
|
session_id: str,
|
|
root_node_id: str,
|
|
metadata: Optional[dict] = None) -> SubGroup:
|
|
'''
|
|
create sub group
|
|
'''
|
|
subgroup = self._find_subgroup(root_node_id)
|
|
if subgroup:
|
|
raise Exception(f"subgroup already exist, group_id: {group_id}, root_node_id: {root_node_id}")
|
|
run_node = self.node_state_manager.get_node(root_node_id)
|
|
if not run_node:
|
|
raise Exception(f"run node not found, root_node_id: {root_node_id}")
|
|
|
|
sub_group = SubGroup(
|
|
session_id=session_id,
|
|
group_id=group_id,
|
|
root_node_id=root_node_id,
|
|
metadata=metadata,
|
|
create_time=time.time(),
|
|
update_time=time.time(),
|
|
status=RunNodeStatus.INIT,
|
|
)
|
|
self.sub_group_storage.insert(sub_group)
|
|
return sub_group
|
|
|
|
def run_sub_group(self,
|
|
root_node_id: str):
|
|
'''
|
|
run sub group
|
|
'''
|
|
subgroup = self._subgroup_exist(root_node_id)
|
|
if not subgroup:
|
|
raise Exception(f"subgroup not found, root_node_id: {root_node_id}")
|
|
|
|
subgroup.execute_time = time.time()
|
|
subgroup.status = RunNodeStatus.RUNNING
|
|
self.sub_group_storage.update(subgroup)
|
|
self.run_group(subgroup.group_id)
|
|
|
|
def run_group(self, group_id):
|
|
group = self.node_group_storage.get(group_id)
|
|
if group.status == RunNodeStatus.INIT:
|
|
group.status = RunNodeStatus.RUNNING
|
|
group.execute_time = time.time()
|
|
self.node_group_storage.update(group)
|
|
|
|
async def finish_sub_group(self,
|
|
group_id: str,
|
|
root_node_id: str,
|
|
results: List[HandleResult] = None,
|
|
result_msg: str = None):
|
|
'''
|
|
finish sub task with node_id as the root node
|
|
'''
|
|
subgroup = self.sub_group_storage.get(root_node_id)
|
|
if not subgroup:
|
|
raise Exception(f"subgroup not found, group_id: {group_id}, root_node_id: {root_node_id}")
|
|
if subgroup.group_id != group_id:
|
|
raise Exception(f"subgroup group_id not match, group_id: {group_id}, root_node_id: {root_node_id}")
|
|
|
|
group = self._group_exist(group_id)
|
|
subgroup.end_time = time.time()
|
|
subgroup.results = results
|
|
subgroup.result_msg = result_msg
|
|
subgroup.status = RunNodeStatus.SUCCESS
|
|
for result in results:
|
|
if result.status == RunNodeStatus.FAILED:
|
|
subgroup.status = RunNodeStatus.FAILED
|
|
self.sub_group_storage.update(subgroup)
|
|
# check all subgroup status and update group status
|
|
await self._check_subgroup_status(group_id, group.root_node_ids)
|
|
|
|
async def _check_subgroup_status(self, group_id, root_node_ids: List[str]):
|
|
'''
|
|
check subgroups status and update group status, if group finished, send group finish message
|
|
'''
|
|
all_subgroups_finished = True
|
|
failed_subgroups = []
|
|
for root_node_id in root_node_ids:
|
|
subgroup = self.sub_group_storage.get(root_node_id)
|
|
if not subgroup or not subgroup.has_finished():
|
|
all_subgroups_finished = False
|
|
break
|
|
if subgroup.status == RunNodeStatus.FAILED or subgroup.status == RunNodeStatus.TIMEOUT:
|
|
failed_subgroups.append(subgroup)
|
|
|
|
if all_subgroups_finished:
|
|
group = self._group_exist(group_id)
|
|
if failed_subgroups:
|
|
group.status = RunNodeStatus.FAILED
|
|
group.failed_root_node_ids = [subgroup.root_node_id for subgroup in failed_subgroups]
|
|
else:
|
|
group.status = RunNodeStatus.SUCCESS
|
|
group.end_time = time.time()
|
|
self.node_group_storage.update(group)
|
|
await self._send_group_finish_message(group_id)
|
|
|
|
async def _send_group_finish_message(self, group_id: str):
|
|
'''
|
|
Currently, for simple implementation, concurrency control needs to be considered in a distributed environment
|
|
'''
|
|
group = self._group_exist(group_id)
|
|
if group.finish_notified:
|
|
logger.warning(f"group finish message already sent, group_id: {group_id}")
|
|
return
|
|
group_results = {}
|
|
metadata = group.metadata
|
|
for root_node_id in group.root_node_ids:
|
|
subgroup = self.sub_group_storage.get(root_node_id)
|
|
group_results[root_node_id] = subgroup.results
|
|
if not metadata:
|
|
metadata = subgroup.metadata
|
|
|
|
metadata = metadata or {}
|
|
if group.parent_group_id:
|
|
metadata.update({
|
|
"parent_group_id": group.parent_group_id
|
|
})
|
|
message = Message(
|
|
category="group",
|
|
payload=group_results,
|
|
sender="node_group_manager",
|
|
session_id=group.session_id,
|
|
topic="__group_results",
|
|
headers=metadata
|
|
)
|
|
await send_message(message)
|
|
group.finish_notified = True
|
|
self.node_group_storage.update(group)
|
|
|
|
def get_group(self, group_id: str) -> NodeGroup:
|
|
'''
|
|
get group basic info
|
|
'''
|
|
return self._find_group(group_id)
|
|
|
|
def query_group_detail(self, group_id: str) -> NodeGroupDetail:
|
|
'''
|
|
query group detail info with all sub group info
|
|
'''
|
|
group = self._find_group(group_id)
|
|
if not group:
|
|
return None
|
|
sub_groups = []
|
|
for root_node_id in group.root_node_ids:
|
|
subgroup = self._find_subgroup(root_node_id)
|
|
if subgroup:
|
|
sub_groups.append(subgroup)
|
|
return NodeGroupDetail(
|
|
group_id=group.group_id,
|
|
root_node_ids=group.root_node_ids,
|
|
parent_group_id=group.parent_group_id,
|
|
metadata=group.metadata,
|
|
create_time=group.create_time,
|
|
execute_time=group.execute_time,
|
|
end_time=group.end_time,
|
|
status=group.status,
|
|
failed_root_node_ids=group.failed_root_node_ids,
|
|
sub_groups=sub_groups
|
|
)
|
|
|
|
def _find_subgroup(self, root_node_id: str) -> SubGroup:
|
|
return self.sub_group_storage.get(root_node_id)
|
|
|
|
def _subgroup_exist(self, root_node_id: str) -> SubGroup:
|
|
subgroup = self._find_subgroup(root_node_id)
|
|
if not subgroup:
|
|
raise Exception(f"subgroup not found, root_node_id: {root_node_id}")
|
|
return subgroup
|
|
|
|
def _find_group(self, group_id: str) -> NodeGroup:
|
|
return self.node_group_storage.get(group_id)
|
|
|
|
def _group_exist(self, group_id: str) -> NodeGroup:
|
|
group = self._find_group(group_id)
|
|
if not group:
|
|
raise Exception(f"group not found, group_id: {group_id}")
|
|
return group
|
|
|
|
|
|
class EventRuntimeStateManager(RuntimeStateManager):
|
|
|
|
def __init__(self, storage: StateStorage = InMemoryStateStorage.instance()):
|
|
super().__init__(storage)
|
|
|
|
def start_message_node(self, message: Message):
|
|
'''
|
|
create and start node while message handle started.
|
|
'''
|
|
metadata = message.headers
|
|
run_node_busi_type = RunNodeBusiType.from_message_category(
|
|
message.category)
|
|
logger.debug(
|
|
f"start message node: {message.receiver}, busi_type={run_node_busi_type}, node_id={message.id}")
|
|
if run_node_busi_type:
|
|
self.create_node(
|
|
node_id=message.id,
|
|
busi_type=run_node_busi_type,
|
|
busi_id=message.receiver or "",
|
|
session_id=message.session_id,
|
|
task_id=message.task_id,
|
|
msg_id=message.id,
|
|
msg_from=message.sender,
|
|
group_id=metadata.get("group_id") if metadata else None,
|
|
sub_group_root_id=metadata.get("root_message_id") if metadata else None,
|
|
metadata=metadata)
|
|
self.run_node(message.id)
|
|
|
|
def save_message_handle_result(self, name: str, message: Message, result: Message = None):
|
|
'''
|
|
save message handle result
|
|
'''
|
|
run_node_busi_type = RunNodeBusiType.from_message_category(
|
|
message.category)
|
|
if run_node_busi_type:
|
|
if result and result.is_error():
|
|
handle_result = HandleResult(
|
|
name=name,
|
|
status=RunNodeStatus.FAILED,
|
|
result=result)
|
|
else:
|
|
handle_result = HandleResult(
|
|
name=name,
|
|
status=self.get_node(message.id).status if self.get_node(message.id) else RunNodeStatus.FAILED,
|
|
result=result)
|
|
self.save_result(node_id=message.id, result=handle_result)
|
|
|
|
def end_message_node(self, message: Message):
|
|
'''
|
|
end node while message handle finished.
|
|
'''
|
|
run_node_busi_type = RunNodeBusiType.from_message_category(
|
|
message.category)
|
|
if run_node_busi_type:
|
|
node = self._node_exist(node_id=message.id)
|
|
status = RunNodeStatus.SUCCESS
|
|
if node.results:
|
|
for result in node.results:
|
|
if result.status == RunNodeStatus.FAILED:
|
|
status = RunNodeStatus.FAILED
|
|
break
|
|
if status == RunNodeStatus.FAILED:
|
|
self.run_failed(node_id=message.id)
|
|
else:
|
|
self.run_succeed(node_id=message.id)
|
|
|
|
def get_message_node_status(self, message: Message) -> RunNodeStatus:
|
|
node = self.get_node(node_id=message.id)
|
|
if not node:
|
|
return RunNodeStatus.INIT
|
|
return node.status
|