ai-agent-book 精选快照（<2MB 代码与文档，来自 github.com/bojieli/ai-agent-book）

chapter9/gaia-experience/AWorld/aworld/runners/handler/agent.py

@@ -0,0 +1,481 @@
+# coding: utf-8
+# Copyright (c) 2025 inclusionAI.
+import abc
+from typing import AsyncGenerator, Tuple
+
+from aworld.agents.loop_llm_agent import LoopableAgent
+from aworld.core.agent.base import is_agent, AgentFactory
+from aworld.core.agent.swarm import GraphBuildType, AgentGraph
+from aworld.core.common import ActionModel, Observation, TaskItem
+from aworld.core.event.base import Message, Constants, TopicType, AgentMessage
+from aworld.core.exceptions import AWorldRuntimeException
+from aworld.logs.util import logger
+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 RunNode, RunNodeStatus, RunNodeBusiType
+from aworld.runners.utils import endless_detect
+from aworld.output.base import StepOutput
+
+
+class AgentHandler(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
+
+        self.agent_calls = []
+
+    @classmethod
+    def name(cls):
+        return "_agents_handler"
+
+
+@HandlerFactory.register(name=f'__{Constants.AGENT}__')
+class DefaultAgentHandler(AgentHandler):
+    def is_valid_message(self, message: Message):
+        if message.category != Constants.AGENT:
+            if self.swarm and message.sender in self.swarm.agents and message.sender in AgentFactory:
+                if self.agent_calls:
+                    if self.agent_calls[-1] != message.sender:
+                        self.agent_calls.append(message.sender)
+                else:
+                    self.agent_calls.append(message.sender)
+            return False
+        return True
+
+    async def _do_handle(self, message: Message) -> AsyncGenerator[Message, None]:
+        if not self.is_valid_message(message):
+            return
+
+        headers = {"context": message.context}
+        session_id = message.session_id
+        data = message.payload
+        if not data:
+            # error message, p2p
+            yield Message(
+                category=Constants.OUTPUT,
+                payload=StepOutput.build_failed_output(name=f"{message.caller or self.name()}",
+                                                       step_num=0,
+                                                       data="no data to process.",
+                                                       task_id=self.task_id),
+                sender=self.name(),
+                session_id=session_id,
+                headers=headers
+            )
+            yield Message(
+                category=Constants.TASK,
+                payload=TaskItem(msg="no data to process.", data=data, stop=True),
+                sender=self.name(),
+                session_id=session_id,
+                topic=TopicType.ERROR,
+                headers=headers
+            )
+            return
+
+        if isinstance(data, Tuple) and isinstance(data[0], Observation):
+            data = data[0]
+            message.payload = data
+        # data is Observation
+        if isinstance(data, Observation):
+            if not self.swarm:
+                msg = Message(
+                    category=Constants.TASK,
+                    payload=data.content,
+                    sender=data.observer,
+                    session_id=session_id,
+                    topic=TopicType.FINISHED,
+                    headers=headers
+                )
+                logger.info(f"FINISHED|agent handler send finished message: {msg}")
+                yield msg
+                return
+
+            agent = self.swarm.agents.get(message.receiver)
+            # agent + tool completion protocol.
+            if agent and agent.finished and data.info.get('done'):
+                self.swarm.cur_step += 1
+
+                root_agent = self.swarm.communicate_agent
+                if isinstance(root_agent, list):
+                    root_agent = root_agent[0]
+                if agent.id() == root_agent.id():
+                    msg = Message(
+                        category=Constants.TASK,
+                        payload=data.content,
+                        sender=agent.id(),
+                        session_id=session_id,
+                        topic=TopicType.FINISHED,
+                        headers=headers
+                    )
+                    logger.info(f"FINISHED|agent handler send finished message: {msg}")
+                    yield msg
+                else:
+                    msg = Message(
+                        category=Constants.AGENT,
+                        payload=Observation(content=data.content),
+                        sender=agent.id(),
+                        session_id=session_id,
+                        receiver=root_agent.id(),
+                        headers=message.headers
+                    )
+                    logger.info(f"agent handler send agent message: {msg}")
+                    yield msg
+            else:
+                if data.info.get('done'):
+                    agent_name = self.agent_calls[-1]
+                    async for event in self._stop_check(ActionModel(agent_name=agent_name, policy_info=data.content),
+                                                        message):
+                        yield event
+                elif not message.receiver:
+                    agent_name = message.sender
+                    async for event in self._stop_check(ActionModel(agent_name=agent_name, policy_info=data.content),
+                                                        message):
+                        yield event
+                else:
+                    logger.info(f"agent handler send observation message: {message}")
+                    yield message
+            return
+
+        # data is List[ActionModel]
+        for action in data:
+            if not isinstance(action, ActionModel):
+                # error message, p2p
+                yield Message(
+                    category=Constants.OUTPUT,
+                    payload=StepOutput.build_failed_output(name=f"{message.caller or self.name()}",
+                                                           step_num=0,
+                                                           data="action not a ActionModel.",
+                                                           task_id=self.task_id),
+                    sender=self.name(),
+                    session_id=session_id,
+                    headers=headers
+                )
+                msg = Message(
+                    category=Constants.TASK,
+                    payload=TaskItem(msg="action not a ActionModel.", data=data, stop=True),
+                    sender=self.name(),
+                    session_id=session_id,
+                    topic=TopicType.ERROR,
+                    headers=headers
+                )
+                logger.info(f"agent handler send task message: {msg}")
+                yield msg
+                return
+
+        tools = []
+        agents = []
+        for action in data:
+            if is_agent(action):
+                agents.append(action)
+            else:
+                tools.append(action)
+
+        if tools:
+            msg = Message(
+                category=Constants.TOOL,
+                payload=tools,
+                sender=self.name(),
+                session_id=session_id,
+                receiver=DefaultToolHandler.name(),
+                headers=message.headers
+            )
+            logger.info(f"agent handler send tool message: {msg}")
+            yield msg
+        else:
+            yield Message(
+                category=Constants.OUTPUT,
+                payload=StepOutput.build_finished_output(name=f"{message.caller or self.name()}",
+                                                         step_num=0,
+                                                         task_id=self.task_id),
+                sender=self.name(),
+                receiver=agents[0].tool_name,
+                session_id=session_id,
+                headers=headers
+            )
+
+        for agent in agents:
+            async for event in self._agent(agent, message):
+                logger.info(f"agent handler send message: {event}")
+                yield event
+
+    async def _agent(self, action: ActionModel, message: Message):
+        self.agent_calls.append(action.agent_name)
+        agent = self.swarm.agents.get(action.agent_name)
+        # be handoff
+        agent_name = action.tool_name
+        if not agent_name:
+            async for event in self._stop_check(action, message):
+                yield event
+            return
+
+        headers = {"context": message.context}
+        session_id = message.session_id
+        cur_agent = self.swarm.agents.get(agent_name)
+        if not cur_agent or not agent:
+            yield Message(
+                category=Constants.TASK,
+                payload=TaskItem(msg=f"Can not find {agent_name} or {action.agent_name} agent in swarm.",
+                                 data=action,
+                                 stop=True),
+                sender=self.name(),
+                session_id=session_id,
+                topic=TopicType.ERROR,
+                headers=headers
+            )
+            return
+
+        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=agent.id(), from_agent_name=agent.id())
+
+        if agent.handoffs and agent_name not in agent.handoffs:
+            if message.caller:
+                message.receiver = message.caller
+                message.caller = ''
+                yield message
+            else:
+                yield Message(category=Constants.TASK,
+                              payload=TaskItem(msg=f"Can not handoffs {agent_name} agent ", data=observation),
+                              sender=self.name(),
+                              session_id=session_id,
+                              topic=TopicType.RERUN,
+                              headers=headers)
+            return
+
+        headers = message.headers.copy()
+        # headers.update({"agent_as_tool": True})
+        yield Message(
+            category=Constants.AGENT,
+            payload=observation,
+            caller=message.caller,
+            sender=action.agent_name,
+            session_id=session_id,
+            receiver=action.tool_name,
+            headers=headers,
+        )
+
+    async def _stop_check(self, action: ActionModel, message: Message) -> AsyncGenerator[Message, None]:
+        if GraphBuildType.TEAM.value == self.swarm.build_type:
+            async for event in self._team_stop_check(action, message):
+                yield event
+        elif GraphBuildType.HANDOFF.value == self.swarm.build_type:
+            async for event in self._handoff_stop_check(action, message):
+                yield event
+        else:
+            async for event in self._workflow_stop_check(action, message):
+                yield event
+
+    async def _workflow_stop_check(self, action: ActionModel, message: Message) -> AsyncGenerator[Message, None]:
+        # Equivalent to scheduling
+        session_id = message.session_id
+        agent_name = action.agent_name
+        agent = self.swarm.agents.get(agent_name)
+        if not agent:
+            yield Message(
+                category=Constants.TASK,
+                payload=TaskItem(
+                    msg=f"Can not find {action.agent_name} agent in ordered_agents: {self.swarm.ordered_agents}.",
+                    data=action,
+                    stop=True),
+                sender=self.name(),
+                session_id=session_id,
+                topic=TopicType.ERROR,
+                headers=message.headers
+            )
+            return
+
+        receiver = None
+        # loop agent type
+        if isinstance(agent, LoopableAgent):
+            agent.cur_run_times += 1
+            if not agent.finished:
+                receiver = agent.goto
+
+        if receiver:
+            yield Message(
+                category=Constants.AGENT,
+                payload=Observation(content=action.policy_info),
+                sender=agent.id(),
+                session_id=session_id,
+                receiver=receiver,
+                headers=message.headers
+            )
+        else:
+            agent_graph: AgentGraph = self.swarm.agent_graph
+            # next
+            successor = agent_graph.successor.get(agent_name)
+            if not successor:
+                yield Message(
+                    category=Constants.TASK,
+                    payload=action.policy_info,
+                    sender=agent.id(),
+                    session_id=session_id,
+                    topic=TopicType.FINISHED,
+                    headers=message.headers
+                )
+                return
+
+            for k, _ in successor.items():
+                predecessor = agent_graph.predecessor.get(k)
+                if not predecessor:
+                    raise AWorldRuntimeException(f"{k} has no predecessor {agent_name}, may changed during iteration.")
+
+                all_input = {}
+                pre_finished = True
+                for pre_k, _ in predecessor.items():
+                    if pre_k == agent_name:
+                        all_input[agent_name] = action.policy_info
+                        continue
+                    # check all predecessor agent finished
+                    run_node: RunNode = self.runner.state_manager.query_by_task(
+                        task_id=message.context.get_task().id,
+                        busi_typ=RunNodeBusiType.AGENT,
+                        busi_id=pre_k
+                    )
+                    if run_node:
+                        run_node = run_node[0]
+                    else:
+                        raise AWorldRuntimeException(f"{pre_k} can't find in task: {message.context.get_task().id}.")
+                    if run_node.status == RunNodeStatus.RUNNING or run_node.status == RunNodeStatus.INIT:
+                        # mean not finished
+                        pre_finished = False
+                        logger.info(f"{pre_k} not finished, will wait it.")
+                    else:
+                        logger.info(f"{pre_k} finished, result is: {run_node.results}")
+                        payload = run_node.results[-1].result.payload[0]
+                        all_input[pre_k] = payload.policy_info
+
+                if pre_finished:
+                    yield Message(
+                        category=Constants.AGENT,
+                        payload=Observation(content=all_input if len(all_input) > 1 else all_input.get(agent_name)),
+                        sender=agent.id(),
+                        session_id=session_id,
+                        receiver=k,
+                        headers=message.headers
+                    )
+
+    async def _team_stop_check(self, action: ActionModel, message: Message) -> AsyncGenerator[Message, None]:
+        caller = message.caller
+        session_id = message.session_id
+        agent = self.swarm.agents.get(action.agent_name)
+        if ((not caller or caller == self.swarm.communicate_agent.id())
+                and (self.swarm.cur_step >= self.swarm.max_steps or self.swarm.finished or
+                     (agent.id() == self.swarm.agent_graph.root_agent.id() and agent.finished))):
+            logger.info(
+                f"FINISHED|_social_stop_check finished|{self.swarm.cur_step}|{self.swarm.max_steps}|{self.swarm.finished}")
+            yield Message(
+                category=Constants.TASK,
+                payload=action.policy_info,
+                sender=agent.id(),
+                session_id=session_id,
+                topic=TopicType.FINISHED,
+                headers={"context": message.context}
+            )
+        agent = self.swarm.agents.get(action.agent_name)
+        caller = self.swarm.agent_graph.root_agent.id() or message.caller
+        if agent.id() != self.swarm.agent_graph.root_agent.id():
+            logger.info(f"_stop_check Team|{agent.id()} --> {caller}")
+            yield Message(
+                category=Constants.AGENT,
+                payload=Observation(content=action.policy_info),
+                sender=agent.id(),
+                session_id=message.session_id,
+                receiver=caller,
+                headers=message.headers
+            )
+
+    async def _handoff_stop_check(self, action: ActionModel, message: Message) -> AsyncGenerator[Message, None]:
+        headers = {"context": message.context}
+        agent = self.swarm.agents.get(action.agent_name)
+        caller = message.caller
+        session_id = message.session_id
+        if endless_detect(self.agent_calls,
+                          endless_threshold=self.endless_threshold,
+                          root_agent_name=self.swarm.communicate_agent.id()):
+            logger.info(
+                f"FINISHED|_social_stop_check endless_detect|{self.agent_calls}|{self.endless_threshold}|{self.swarm.communicate_agent.id()}")
+            yield Message(
+                category=Constants.TASK,
+                payload=action.policy_info,
+                sender=agent.id(),
+                session_id=session_id,
+                topic=TopicType.FINISHED,
+                headers=headers
+            )
+            return
+
+        if not caller or caller == self.swarm.communicate_agent.id():
+            if self.swarm.cur_step >= self.swarm.max_steps or self.swarm.finished:
+                logger.info(
+                    f"FINISHED|_social_stop_check finished|{self.swarm.cur_step}|{self.swarm.max_steps}|{self.swarm.finished}")
+                yield Message(
+                    category=Constants.TASK,
+                    payload=action.policy_info,
+                    sender=agent.id(),
+                    session_id=session_id,
+                    topic=TopicType.FINISHED,
+                    headers=headers
+                )
+            else:
+                self.swarm.cur_step += 1
+                logger.info(f"_social_stop_check execute loop {self.swarm.cur_step}.")
+                yield Message(
+                    category=Constants.AGENT,
+                    payload=Observation(content=action.policy_info),
+                    sender=agent.id(),
+                    session_id=session_id,
+                    receiver=self.swarm.communicate_agent.id(),
+                    headers=message.headers
+                )
+        else:
+            idx = 0
+            for idx, name in enumerate(self.agent_calls[::-1]):
+                if name == agent.id():
+                    break
+            idx = len(self.agent_calls) - idx - 1
+            if idx:
+                caller = self.agent_calls[idx - 1]
+
+            yield Message(
+                category=Constants.AGENT,
+                payload=Observation(content=action.policy_info),
+                sender=agent.id(),
+                session_id=session_id,
+                receiver=caller,
+                headers=message.headers
+            )
+
+    def is_group_finish(self, event: Message) -> bool:
+        """Determine if an event triggers group completion"""
+        if not isinstance(event, Message) or not event.group_id:
+            return False
+
+        agent_id = event.sender
+        if not agent_id:
+            return False
+
+        agent = self.swarm.agents.get(agent_id)
+        if not agent:
+            return False
+
+        return agent._finished and agent.id() == event.headers.get('root_agent_id', '')
+
+    async def post_handle(self, input: Message, output: Message) -> Message:
+        new_context = output.context.deep_copy()
+        new_context._task = output.context.get_task()
+        output.context = new_context
+        if self.is_group_finish(output):
+            from aworld.runners.state_manager import RuntimeStateManager
+            state_mng = RuntimeStateManager.instance()
+            await state_mng.finish_sub_group(output.group_id, output.headers.get('root_message_id'),
+                                             [output])
+            return None
+        return output