liqiang
b119135836
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
375 lines
18 KiB
Python
375 lines
18 KiB
Python
# coding: utf-8
|
|
# Copyright (c) 2025 inclusionAI.
|
|
import asyncio
|
|
import time
|
|
import traceback
|
|
|
|
from aworld.core.agent.base import BaseAgent
|
|
|
|
from aworld.core.exceptions import AWorldRuntimeException
|
|
|
|
import aworld.trace as trace
|
|
from typing import List, Callable, Any
|
|
|
|
from aworld.core.common import TaskItem, ActionModel
|
|
from aworld.core.context.base import Context
|
|
|
|
from aworld.agents.llm_agent import Agent
|
|
from aworld.core.event.base import Message, Constants, TopicType, ToolMessage, AgentMessage
|
|
from aworld.core.task import Task, TaskResponse
|
|
from aworld.events.manager import EventManager
|
|
from aworld.logs.util import logger
|
|
from aworld.replay_buffer import EventReplayBuffer
|
|
from aworld.runners import HandlerFactory
|
|
from aworld.runners.handler.base import DefaultHandler
|
|
|
|
from aworld.runners.task_runner import TaskRunner
|
|
from aworld.utils.common import override_in_subclass, new_instance
|
|
from aworld.runners.state_manager import EventRuntimeStateManager
|
|
|
|
|
|
class TaskEventRunner(TaskRunner):
|
|
"""Event driven task runner."""
|
|
|
|
def __init__(self, task: Task, *args, **kwargs):
|
|
super().__init__(task, *args, **kwargs)
|
|
self._task_response = None
|
|
self.event_mng = EventManager(self.context)
|
|
self.hooks = {}
|
|
self.handlers = []
|
|
self.init_messages = []
|
|
self.background_tasks = set()
|
|
self.state_manager = EventRuntimeStateManager.instance()
|
|
self.replay_buffer = EventReplayBuffer()
|
|
|
|
async def do_run(self, context: Context = None):
|
|
if self.swarm and not self.swarm.initialized:
|
|
raise AWorldRuntimeException("swarm needs to use `reset` to init first.")
|
|
if not self.init_messages:
|
|
raise AWorldRuntimeException("no question event to solve.")
|
|
|
|
async with trace.task_span(self.init_messages[0].session_id, self.task):
|
|
for msg in self.init_messages:
|
|
await self.event_mng.emit_message(msg)
|
|
await self._do_run()
|
|
await self._save_trajectories()
|
|
return self._response()
|
|
|
|
async def pre_run(self):
|
|
logger.debug(f"[TaskEventRunner] pre_run start {self.task.id}")
|
|
await super().pre_run()
|
|
self.event_mng.context = self.context
|
|
self.context.event_manager = self.event_mng
|
|
|
|
if self.swarm and not self.swarm.max_steps:
|
|
self.swarm.max_steps = self.task.conf.get('max_steps', 10)
|
|
observation = self.observation
|
|
if not observation:
|
|
raise RuntimeError("no observation, check run process")
|
|
|
|
self._build_first_message()
|
|
|
|
if self.swarm:
|
|
logger.debug(f"swarm: {self.swarm}")
|
|
# register agent handler
|
|
for _, agent in self.swarm.agents.items():
|
|
if override_in_subclass('async_policy', agent.__class__, Agent):
|
|
await self.event_mng.register(Constants.AGENT, agent.id(), agent.async_run)
|
|
else:
|
|
await self.event_mng.register(Constants.AGENT, agent.id(), agent.run)
|
|
# register tool handler
|
|
for key, tool in self.tools.items():
|
|
if tool.handler:
|
|
await self.event_mng.register(Constants.TOOL, tool.name(), tool.handler)
|
|
else:
|
|
await self.event_mng.register(Constants.TOOL, tool.name(), tool.step)
|
|
handlers = self.event_mng.event_bus.get_topic_handlers(
|
|
Constants.TOOL, tool.name())
|
|
if not handlers:
|
|
await self.event_mng.register(Constants.TOOL, Constants.TOOL, tool.step)
|
|
|
|
self._stopped = asyncio.Event()
|
|
|
|
# handler of process in framework
|
|
handler_list = self.conf.get("handlers")
|
|
if handler_list:
|
|
# handler class name
|
|
for hand in handler_list:
|
|
self.handlers.append(new_instance(hand, self))
|
|
else:
|
|
for handler in HandlerFactory:
|
|
self.handlers.append(HandlerFactory(handler, runner=self))
|
|
logger.debug(f"[TaskEventRunner] pre_run finish {self.task.id}")
|
|
|
|
def _build_first_message(self):
|
|
# build the first message
|
|
if self.agent_oriented:
|
|
agents = self.swarm.communicate_agent
|
|
if isinstance(agents, BaseAgent):
|
|
agents = [agents]
|
|
|
|
for agent in agents:
|
|
self.init_messages.append(AgentMessage(payload=self.observation,
|
|
sender='runner',
|
|
receiver=agent.id(),
|
|
session_id=self.context.session_id,
|
|
headers={'context': self.context}))
|
|
else:
|
|
actions: List[ActionModel] = self.observation.content
|
|
action_dict = {}
|
|
for action in actions:
|
|
if action.tool_name not in action_dict:
|
|
action_dict[action.tool_name] = []
|
|
action_dict[action.tool_name].append(action)
|
|
|
|
for tool_name, actions in action_dict.items():
|
|
self.init_messages.append(ToolMessage(payload=actions,
|
|
sender='runner',
|
|
receiver=tool_name,
|
|
session_id=self.context.session_id,
|
|
headers={'context': self.context}))
|
|
|
|
async def _common_process(self, message: Message) -> List[Message]:
|
|
logger.debug(
|
|
f"[TaskEventRunner] _common_process start {self.task.id}, message_id = {message.id}")
|
|
event_bus = self.event_mng.event_bus
|
|
|
|
key = message.category
|
|
transformer = self.event_mng.get_transform_handler(key)
|
|
if transformer:
|
|
message = await event_bus.transform(message, handler=transformer)
|
|
|
|
results = []
|
|
handlers = self.event_mng.get_handlers(key)
|
|
async with trace.message_span(message=message):
|
|
logger.debug(
|
|
f"[TaskEventRunner] start_message_node start {self.task.id}, message_id = {message.id}")
|
|
self.state_manager.start_message_node(message)
|
|
logger.debug(
|
|
f"[TaskEventRunner] start_message_node end {self.task.id}, message_id = {message.id}")
|
|
if handlers:
|
|
if message.topic:
|
|
handlers = {message.topic: handlers.get(message.topic, [])}
|
|
elif message.receiver:
|
|
handlers = {message.receiver: handlers.get(
|
|
message.receiver, [])}
|
|
else:
|
|
logger.warning(
|
|
f"{message.id} no receiver and topic, be ignored.")
|
|
handlers.clear()
|
|
|
|
handle_tasks = []
|
|
for topic, handler_list in handlers.items():
|
|
if not handler_list:
|
|
logger.warning(f"{topic} no handler, ignore.")
|
|
continue
|
|
|
|
for handler in handler_list:
|
|
t = asyncio.create_task(
|
|
self._handle_task(message, handler))
|
|
handle_tasks.append(t)
|
|
logger.debug(
|
|
f"[TaskEventRunner] _common_process handle_tasks collect finished {self.task.id}, message_id = {message.id}")
|
|
|
|
# For _handle_task case, end message node asynchronously
|
|
async def async_end_message_node():
|
|
logger.debug(
|
|
f"[TaskEventRunner] async_end_message_node STARTED {self.task.id}, message_id = {message.id}")
|
|
try:
|
|
# Wait for all _handle_task tasks to complete before ending message node
|
|
if handle_tasks:
|
|
logger.debug(
|
|
f"[TaskEventRunner] async_end_message_node {self.task.id} Before gather {len(handle_tasks)} tasks")
|
|
await asyncio.gather(*handle_tasks)
|
|
logger.debug(
|
|
f"[TaskEventRunner] async_end_message_node {self.task.id} After gather tasks completed")
|
|
logger.debug(
|
|
f"[TaskEventRunner] _common_process handle_tasks process end_message_node start {self.task.id}, message_id = {message.id}")
|
|
self.state_manager.end_message_node(message)
|
|
logger.debug(
|
|
f"[TaskEventRunner] _common_process handle_tasks process finished {self.task.id}, message_id = {message.id}")
|
|
except Exception as e:
|
|
logger.error(f"Error in async_end_message_node: {e}")
|
|
raise
|
|
|
|
end_node_task = asyncio.create_task(async_end_message_node())
|
|
self.background_tasks.add(end_node_task)
|
|
end_node_task.add_done_callback(self.background_tasks.discard)
|
|
else:
|
|
# not handler, return raw message
|
|
results.append(message)
|
|
|
|
t = asyncio.create_task(self._raw_task(results))
|
|
self.background_tasks.add(t)
|
|
t.add_done_callback(self.background_tasks.discard)
|
|
# wait until it is complete
|
|
await t
|
|
self.state_manager.end_message_node(message)
|
|
logger.debug(
|
|
f"[TaskEventRunner] _common_process return results {self.task.id}, message_id = {message.id}, ")
|
|
return results
|
|
|
|
async def _handle_task(self, message: Message, handler: Callable[..., Any]):
|
|
con = message
|
|
async with trace.handler_span(message=message, handler=handler):
|
|
try:
|
|
logger.debug(
|
|
f"event_runner _handle_task - self: {self}, swarm: {self.swarm}, event_mng: {self.event_mng}, event_bus: {self.event_mng.event_bus}, message: {message}")
|
|
logger.info(
|
|
f"[TaskEventRunner] {self.task.id} _handle_task start, message: {message.id}")
|
|
if asyncio.iscoroutinefunction(handler):
|
|
con = await handler(con)
|
|
else:
|
|
con = handler(con)
|
|
|
|
logger.info(
|
|
f"[TaskEventRunner] {self.task.id} _handle_task finished message= {message.id}, session_id = {self.task.session_id}")
|
|
if isinstance(con, Message):
|
|
# process in framework
|
|
self.state_manager.save_message_handle_result(name=handler.__name__,
|
|
message=message,
|
|
result=con)
|
|
async for event in self._inner_handler_process(
|
|
results=[con],
|
|
handlers=self.handlers
|
|
):
|
|
await self.event_mng.emit_message(event)
|
|
else:
|
|
self.state_manager.save_message_handle_result(name=handler.__name__,
|
|
message=message)
|
|
except Exception as e:
|
|
logger.warning(
|
|
f"{handler} process fail. {traceback.format_exc()}")
|
|
error_msg = Message(
|
|
category=Constants.TASK,
|
|
payload=TaskItem(msg=str(e), data=message),
|
|
sender=self.name,
|
|
session_id=self.context.session_id,
|
|
topic=TopicType.ERROR,
|
|
headers={"context": self.context}
|
|
)
|
|
self.state_manager.save_message_handle_result(name=handler.__name__,
|
|
message=message,
|
|
result=error_msg)
|
|
await self.event_mng.emit_message(error_msg)
|
|
|
|
async def _raw_task(self, messages: List[Message]):
|
|
# process in framework
|
|
async for event in self._inner_handler_process(
|
|
results=messages,
|
|
handlers=self.handlers
|
|
):
|
|
await self.event_mng.emit_message(event)
|
|
|
|
async def _inner_handler_process(self, results: List[Message], handlers: List[DefaultHandler]):
|
|
# can use runtime backend to parallel
|
|
for handler in handlers:
|
|
for result in results:
|
|
async for event in handler.handle(result):
|
|
yield event
|
|
|
|
async def _do_run(self):
|
|
logger.debug(f"[TaskEventRunner] _do_run start {self.task.id}")
|
|
|
|
"""Task execution process in real."""
|
|
start = time.time()
|
|
msg = None
|
|
answer = None
|
|
message = None
|
|
try:
|
|
while True:
|
|
if self.task.timeout > 0 and time.time() - self.start_time > self.task.timeout:
|
|
logger.warn(
|
|
f"[TaskEventRunner] {self.task.id} task timeout after {time.time() - self.start_time} seconds.")
|
|
self._task_response = TaskResponse(answer='',
|
|
success=False,
|
|
context=message.context,
|
|
id=self.task.id,
|
|
time_cost=(time.time() - self.start_time),
|
|
usage=self.context.token_usage,
|
|
msg='cancellation: task timeout',
|
|
status='cancelled')
|
|
await self.stop()
|
|
if await self.is_stopped():
|
|
logger.debug(
|
|
f"[TaskEventRunner] break snap {self.task.id}")
|
|
await self.event_mng.done()
|
|
logger.info(
|
|
f" [TaskEventRunner] stop task {self.task.id}...")
|
|
if self._task_response is None:
|
|
# send msg to output
|
|
self._task_response = TaskResponse(msg=msg,
|
|
answer=answer,
|
|
context=message.context,
|
|
success=True if not msg else False,
|
|
id=self.task.id,
|
|
time_cost=(
|
|
time.time() - start),
|
|
usage=self.context.token_usage,
|
|
status='success' if not msg else 'failed')
|
|
break
|
|
logger.debug(f"[TaskEventRunner] next snap {self.task.id}")
|
|
# consume message
|
|
message: Message = await self.event_mng.consume()
|
|
logger.debug(
|
|
f"[TaskEventRunner] next consume finished {self.task.id}, event_bus: {self.event_mng.event_bus},: message = {message}")
|
|
# use registered handler to process message
|
|
await self._common_process(message)
|
|
logger.debug(
|
|
f"[TaskEventRunner] _common_process finished {self.task.id}")
|
|
except Exception as e:
|
|
logger.error(f"consume message fail. {traceback.format_exc()}")
|
|
error_msg = Message(
|
|
category=Constants.TASK,
|
|
payload=TaskItem(msg=str(e), data=message),
|
|
sender=self.name,
|
|
session_id=self.context.session_id,
|
|
topic=TopicType.ERROR,
|
|
headers={"context": self.context}
|
|
)
|
|
self.state_manager.save_message_handle_result(name=TaskEventRunner.__name__,
|
|
message=message,
|
|
result=error_msg)
|
|
await self.event_mng.emit_message(error_msg)
|
|
finally:
|
|
logger.debug(
|
|
f"[TaskEventRunner] _do_run finished await_is_stopped {self.task.id}")
|
|
if await self.is_stopped():
|
|
logger.info(
|
|
f"[TaskEventRunner] _do_run finished is_stopped {self.task.id}")
|
|
await self.context.update_task_after_run(self._task_response)
|
|
if not self.task.is_sub_task:
|
|
logger.info(f"FINISHED|TaskEventRunner|outputs|{self.task.id} {self.task.is_sub_task}")
|
|
await self.task.outputs.mark_completed()
|
|
|
|
if self.swarm and self.swarm.agents:
|
|
for agent_name, agent in self.swarm.agents.items():
|
|
try:
|
|
if hasattr(agent, 'sandbox') and agent.sandbox:
|
|
await agent.sandbox.cleanup()
|
|
except Exception as e:
|
|
logger.warning(
|
|
f"event_runner Failed to cleanup sandbox for agent {agent_name}: {e}")
|
|
|
|
async def stop(self):
|
|
self._stopped.set()
|
|
|
|
async def is_stopped(self):
|
|
return self._stopped.is_set()
|
|
|
|
def response(self):
|
|
return self._task_response
|
|
|
|
def _response(self):
|
|
if self.context.get_task().conf and self.context.get_task().conf.resp_carry_context == False:
|
|
self._task_response.context = None
|
|
return self._task_response
|
|
|
|
async def _save_trajectories(self):
|
|
try:
|
|
messages = self.event_mng.messages_by_task_id(self.task.id)
|
|
trajectory = await self.replay_buffer.get_trajectory(messages, self.task.id, self.state_manager)
|
|
self._task_response.trajectory = trajectory
|
|
except Exception as e:
|
|
logger.error(f"Failed to get trajectories: {str(e)}.{traceback.format_exc()}")
|