# coding: utf-8 # Copyright (c) 2025 inclusionAI. import abc import copy import json from typing import AsyncGenerator, List, Dict, Any, Tuple from aworld.agents.llm_agent import Agent from aworld.core.agent.base import is_agent from aworld.core.common import ActionModel, TaskItem, Observation, ActionResult from aworld.core.context.base import Context from aworld.core.event.base import Message, Constants, TopicType, GroupMessage from aworld.logs.util import logger from aworld.output.base import StepOutput from aworld.runners import HandlerFactory from aworld.runners.handler.base import DefaultHandler from aworld.runners.handler.tool import DefaultToolHandler from aworld.runners.state_manager import RuntimeStateManager, RunNodeStatus from aworld.utils.serialized_util import to_serializable from aworld.utils.run_util import exec_agent class GroupHandler(DefaultHandler): __metaclass__ = abc.ABCMeta def __init__(self, runner: 'TaskEventRunner'): super().__init__(runner) self.runner = runner self.swarm = runner.swarm self.endless_threshold = runner.endless_threshold self.task_id = runner.task.id @classmethod def name(cls): return "_group_handler" @HandlerFactory.register(name=f'__{Constants.GROUP}__') class DefaultGroupHandler(GroupHandler): def is_valid_message(self, message: Message): if message.category != Constants.GROUP: return False return True async def _do_handle(self, message: GroupMessage) -> AsyncGenerator[Message, None]: if not self.is_valid_message(message): return self.context = message.context group_id = message.group_id headers = {'context': self.context} state_manager = RuntimeStateManager.instance() if message.topic == TopicType.GROUP_ACTIONS: # message.payload is List[ActionModel] node_ids = [] action_messages = [] agents = [] tools = [] agent_actions_map = {} for action in message.payload: if not isinstance(action, ActionModel): # error message, p2p async for event in self._send_failed_message(message, message.payload, message): yield event return if is_agent(action): agents.append(action) agent_name = action.tool_name if agent_name not in agent_actions_map: agent_actions_map[agent_name] = [] agent_actions_map[agent_name].append(action) else: tools.append(action) # Process each agent's actions agent_messages = {} for agent_name, actions in agent_actions_map.items(): # Get original agent original_agent = self.swarm.agents.get(agent_name) if not original_agent: error_msg = Message( category=Constants.TASK, payload=TaskItem(msg=f"Can not find {agent_name} agent in swarm.", data=actions, stop=True), sender=self.name(), session_id=message.session_id, topic=TopicType.ERROR, headers={'context': self.context} ) yield error_msg return # Create agent copies and execute for each action for action in actions: msg = await self._build_agent_message(action, message) if msg.category != Constants.AGENT: yield msg return self._update_headers(msg, message) agent_copy = self.copy_agent(original_agent) con = action.policy_info if action.params and 'content' in action.params: con = action.params['content'] if agent_name not in agent_messages: agent_messages[agent_name] = [] agent_messages[agent_name].append((con, agent_copy, msg)) agent_node_ids, agent_tasks = await self._parallel_exec_agents_actions(agent_messages, message) node_ids.extend(agent_node_ids) if tools: tool_mapping = {} for action in tools: tool_name = action.tool_name if tool_name not in tool_mapping: tool_mapping[tool_name] = [] tool_mapping[tool_name].append(action) for tool_name, actions in tool_mapping.items(): msg = await self._build_tool_message(actions, message) self._update_headers(msg, message) action_messages.append(msg) node_ids.append(msg.id) # create group group_meta_data = message.headers.copy() group_meta_data["context"] = message.context.deep_copy() group_meta_data["context"].set_task(message.context.get_task()) await state_manager.create_group(group_id, message.session_id, node_ids, message.headers.get('parent_group_id'), group_meta_data) for _, acts in agent_messages.items(): for act in acts: self.runner.state_manager.start_message_node(act[2]) for msg in action_messages: yield msg await self.process_agent_tasks(agent_tasks, message) elif message.topic == TopicType.GROUP_RESULTS: # merge group results action_results = [] group_results = message.payload group_sender = None group_sender_node_id = None agent_context = self.context.deep_copy() agent_context._task = self.context.get_task() receiver_results = {} for node_id, handle_res_list in group_results.items(): if not handle_res_list: logger.warn(f"{self.name()} get group result with empty handle_res.") return node = state_manager._find_node(node_id) tool_call_id = node.metadata.get('root_tool_call_id') is_tool = not tool_call_id and not node.metadata.get('root_agent_id') if not group_sender: group_sender = node.metadata.get('group_sender') if not group_sender_node_id: group_sender_node_id = node.metadata.get('group_sender_node_id') node_results = [] for handle_res in handle_res_list: res_msg = handle_res.result res_status = handle_res.status if res_status == RunNodeStatus.FAILED or not res_msg: logger.warn(f"{self.name()} get group result with failed handle_res: {handle_res}.") return receiver = res_msg.receiver if not receiver: logger.warn(f"{self.name()} get group result with empty receiver: {res_msg}.") continue if receiver != group_sender: if receiver not in receiver_results: receiver_results[receiver] = [] receiver_results[receiver].append(res_msg) else: if is_tool and isinstance(res_msg.payload, Observation): action_results.extend(res_msg.payload.action_result) else: node_results.append(res_msg.payload) self._merge_context(agent_context, res_msg.context) if node_results and tool_call_id: act_res = ActionResult( content=json.dumps(to_serializable(node_results), ensure_ascii=False), tool_call_id=tool_call_id ) action_results.append(act_res) if action_results: group_res_msg = Message( category=Constants.AGENT, payload=Observation(content="", action_result=action_results), caller=message.caller, sender=self.name(), session_id=message.session_id, receiver=group_sender, headers={'context': agent_context} ) receiver_results[group_sender] = [group_res_msg] for receiver, res_msgs in receiver_results.items(): result_message = self._merge_result_messages(res_msgs, message, group_sender_node_id) group_headers = {} group_sender_node = state_manager._find_node(group_sender_node_id) if group_sender_node: group_headers.update(group_sender_node.metadata.copy()) group_headers['level'] = headers.get('level', 0) + 1 group_headers['context'] = result_message.context or self.context result_message.headers = group_headers yield result_message def copy_agent(self, agent: Agent): """Create a copy of the agent Args: agent: Original agent object Returns: Deep copy of the agent """ return agent async def _parallel_exec_agents_actions(self, agent_messages: Dict[str, List[Tuple[str, Agent, Message]]], message: Message): """Execute multiple agent actions in parallel Args: agent_messages: Messages for agent actions """ tasks = {} messages_ids = [] for agent_name, acts in agent_messages.items(): for act in acts: agent_message = act[2] messages_ids.append(agent_message.id) tasks[agent_message.id] = exec_agent(act[0], act[1], self.context, sub_task=True, outputs=self.context.outputs) return messages_ids, tasks async def process_agent_tasks(self, agent_tasks, input_message): """Process agent async tasks Args: agent_tasks: Agent async tasks """ root_agent_set = set() for node_id, task in agent_tasks.items(): res = await task logger.info(f"{node_id} finished task: {res}") state_manager = self.runner.state_manager node = state_manager._find_node(node_id) if not node: logger.warn(f"{self.name()} get group result with empty node.") return root_agent_id = node.metadata.get('root_agent_id') root_agent_set.add(root_agent_id) self.context.merge_sub_context(res.context) msg = Message( category=Constants.AGENT, payload=[ActionModel(policy_info=res.answer, agent_name=root_agent_id)], sender=root_agent_id, session_id=node.session_id, headers={'context': self.context, 'root_agent_id': root_agent_id, 'root_tool_call_id': node.metadata.get('root_tool_call_id')} ) finish_group_messages = [] async for event in self.runner._inner_handler_process( results=[msg], handlers=self.runner.handlers ): # Only AGENT and TASK messages if isinstance(event, Message) and ( event.category == Constants.AGENT or event.category == Constants.TASK): finish_group_messages.append(event) print(f"======== event context: {event.context},.context.task: {event.context.get_task()}") await state_manager.finish_sub_group(node.metadata.get('group_id'), node_id, finish_group_messages) for agent_id in root_agent_set: agent = self.swarm.agents.get(agent_id) if agent: agent._finished = True def _merge_result_messages(self, res_msgs: List[Message], input_message: Message, group_sender_node_id: str): """Merge multiple result messages Args: res_msgs: Result messages """ if len(res_msgs) == 1: return res_msgs[0] input_list = [] new_context = input_message.context.deep_copy() new_context._task = self.context.get_task() for message in res_msgs: map = {} map[message.sender] = message.payload input_list.append(map) new_context.merge_context(message.context) return Message( category=Constants.AGENT, payload=Observation(content=input_list), sender=self.name(), receiver=res_msgs[0].receiver, session_id=res_msgs[0].session_id, headers={ 'context': new_context } ) async def _build_agent_message(self, action: ActionModel, message: Message) -> Message: session_id = message.session_id headers = { "context": message.context, "root_tool_call_id": action.tool_call_id } from_agent = self.swarm.agents.get(action.agent_name) tool_name = action.tool_name if not tool_name: logger.warn(f"{self.name()} get agent action with empty tool_name.") return Message( category=Constants.TASK, payload=TaskItem(msg=f"Empty tool_name in group_action: {action}.", data=action, stop=True), sender=self.name(), session_id=session_id, topic=TopicType.ERROR, headers=headers ) cur_agent = self.swarm.agents.get(tool_name) if not cur_agent: return Message( category=Constants.TASK, payload=TaskItem(msg=f"Can not find {tool_name} agent in swarm.", data=action, stop=True), sender=self.name(), session_id=session_id, topic=TopicType.ERROR, headers=headers ) cur_agent._finished = False con = action.policy_info if action.params and 'content' in action.params: con = action.params['content'] observation = Observation(content=con, observer=from_agent.id(), from_agent_name=from_agent.id()) return Message( category=Constants.AGENT, payload=observation, caller=message.caller, sender=action.agent_name, session_id=session_id, receiver=cur_agent.id(), headers=headers ) async def _build_tool_message(self, actions: List[ActionModel], message: Message): session_id = message.session_id headers = {"context": message.context.deep_copy()} return Message( category=Constants.TOOL, payload=actions, sender=self.name(), session_id=session_id, receiver=DefaultToolHandler.name(), headers=headers ) async def _send_failed_message(self, message, data, result_msg): yield Message( category=Constants.OUTPUT, payload=StepOutput.build_failed_output(name=f"{message.caller or self.name()}", step_num=0, data=result_msg, task_id=self.task_id), sender=self.name(), session_id=self.context.session_id, headers=message.headers ) yield Message( category=Constants.TASK, payload=TaskItem(msg=result_msg, data=data, stop=True), sender=self.name(), session_id=self.context.session_id, topic=TopicType.ERROR, headers=message.headers ) def _update_headers(self, message: Message, parent_message: Message): headers = message.headers.copy() context = message.context.deep_copy() context.set_task(self.context.get_task()) headers['context'] = context headers['group_id'] = parent_message.group_id headers['root_message_id'] = message.id headers['root_agent_id'] = message.receiver if message.category == Constants.AGENT else '' headers['level'] = 0 headers['group_sender'] = parent_message.sender headers['group_sender_node_id'] = parent_message.id headers['parent_group_id'] = parent_message.headers.get('parent_group_id') message.headers = headers def _merge_context(self, context: Context, new_context: Context): if not new_context: return if not context: context = new_context return context.merge_context(new_context)