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
103 lines
3.3 KiB
Python
103 lines
3.3 KiB
Python
# coding: utf-8
|
|
# Copyright (c) 2025 inclusionAI.
|
|
import abc
|
|
import time
|
|
|
|
from typing import TypeVar, Generic, AsyncGenerator
|
|
|
|
from aworld.events.util import send_message
|
|
|
|
from aworld.core.common import TaskItem
|
|
|
|
from aworld.core.event.base import Message, Constants, TopicType, CancelMessage
|
|
from aworld.logs.util import logger
|
|
|
|
IN = TypeVar('IN')
|
|
OUT = TypeVar('OUT')
|
|
|
|
|
|
class Handler(Generic[IN, OUT]):
|
|
__metaclass__ = abc.ABCMeta
|
|
|
|
@abc.abstractmethod
|
|
async def handle(self, data: IN) -> AsyncGenerator[OUT, None]:
|
|
"""Process the data as the expected result.
|
|
|
|
Args:
|
|
data: Data generated while running the task.
|
|
"""
|
|
|
|
@classmethod
|
|
def name(cls):
|
|
"""Handler name."""
|
|
return cls.__name__
|
|
|
|
|
|
class DefaultHandler(Handler[Message, AsyncGenerator[Message, None]]):
|
|
"""Default handler."""
|
|
|
|
def __init__(self, runner: 'TaskEventRunner'):
|
|
self.runner = runner
|
|
self.hooks = None
|
|
|
|
def get_registered_name(self):
|
|
"""Get the registered name of the handler.
|
|
|
|
If the class has a REGISTERED_NAME attribute, return the value of the attribute;
|
|
otherwise return None.
|
|
"""
|
|
return getattr(self.__class__, "REGISTERED_NAME", None)
|
|
|
|
def is_valid_message(self, message: Message):
|
|
"""Validate if the message is valid for this handler.
|
|
|
|
If the class has a REGISTERED_NAME attribute, check if the message's category matches the registered name;
|
|
otherwise return True.
|
|
"""
|
|
registered_name = self.get_registered_name()
|
|
if registered_name is not None:
|
|
return message.category == registered_name
|
|
return True
|
|
|
|
async def handle(self, message: Message) -> AsyncGenerator[Message, None]:
|
|
if not self.is_valid_message(message):
|
|
return
|
|
timeout = message.context.get_task().timeout
|
|
time_cost = time.time() - self.runner.start_time
|
|
if message.topic != TopicType.CANCEL and timeout > 0 and time_cost > timeout:
|
|
logger.warn(
|
|
f"[{self.name()}] {message.context.get_task().id} task timeout after {time_cost} seconds.")
|
|
yield CancelMessage(
|
|
payload=TaskItem(msg="task timeout.", data=message, stop=True),
|
|
sender=self.name(),
|
|
session_id=self.runner.context.session_id,
|
|
headers={"context": message.context}
|
|
)
|
|
return
|
|
async for event in self._do_handle(message):
|
|
msg = await self.post_handle(input=message, output=event)
|
|
if msg:
|
|
yield msg
|
|
|
|
async def _do_handle(self, message: Message) -> AsyncGenerator[Message, None]:
|
|
yield message
|
|
|
|
async def post_handle(self, input:Message, output: Message) -> Message:
|
|
"""Post handle the message.
|
|
Args:
|
|
message: Message generated while running the task.
|
|
"""
|
|
return output
|
|
|
|
async def run_hooks(self, message: Message, hook_point: str) -> AsyncGenerator[Message, None]:
|
|
if not self.hooks:
|
|
return
|
|
hooks = self.hooks.get(hook_point, [])
|
|
for hook in hooks:
|
|
try:
|
|
msg = await hook.exec(message)
|
|
if msg:
|
|
yield msg
|
|
except:
|
|
logger.warning(f"{self.name()}|{hook.point()} {hook.name()} execute fail.")
|