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

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

@@ -0,0 +1,95 @@
+# aworld/runners/handler/output.py
+import json
+from typing import AsyncGenerator
+from aworld.core.task import TaskResponse
+from aworld.models.model_response import ModelResponse
+from aworld.runners import HandlerFactory
+from aworld.runners.handler.base import DefaultHandler
+from aworld.output.base import StepOutput, MessageOutput, Output
+from aworld.core.common import TaskItem
+from aworld.core.event.base import Message, Constants, TopicType
+from aworld.logs.util import logger
+from aworld.runners.hook.hook_factory import HookFactory
+from aworld.runners.hook.hooks import HookPoint
+
+
+@HandlerFactory.register(name=f'__{Constants.OUTPUT}__')
+class DefaultOutputHandler(DefaultHandler):
+    def __init__(self, runner):
+        super().__init__(runner)
+        self.runner = runner
+        self.hooks = {}
+        if runner.task.hooks:
+            for k, vals in runner.task.hooks.items():
+                self.hooks[k] = []
+                for v in vals:
+                    cls = HookFactory.get_class(v)
+                    if cls:
+                        self.hooks[k].append(cls)
+
+    def is_valid_message(self, message: Message):
+        if message.category != Constants.OUTPUT:
+            return False
+        return True
+
+    async def _do_handle(self, message):
+        if not self.is_valid_message(message):
+            return
+        # 1. get outputs
+        outputs = self.runner.task.outputs
+        if not outputs:
+            yield Message(
+                category=Constants.TASK,
+                payload=TaskItem(msg="Cannot get outputs.",
+                                 data=message, stop=True),
+                sender=self.name(),
+                session_id=self.runner.context.session_id,
+                topic=TopicType.ERROR,
+                headers={"context": message.context}
+            )
+            return
+
+        # 2. Call OUTPUT_PROCESS hooks to process data in the message
+        async for event in self.run_hooks(message, HookPoint.OUTPUT_PROCESS):
+            # If hook returns a processed message, use the processed message
+            if event and isinstance(event, Message) and event.payload:
+                message.payload = event.payload
+
+        # 3. build Output
+        payload = message.payload
+        mark_complete = False
+        output = None
+        try:
+            if isinstance(payload, Output):
+                output = payload
+                output.task_id = self.runner.task.id
+            elif isinstance(payload, TaskResponse):
+                logger.info(
+                    f"FINISHED|output get task_response with usage: {json.dumps(payload.usage)}")
+                if message.topic == TopicType.FINISHED or message.topic == TopicType.ERROR:
+                    mark_complete = True
+            elif isinstance(payload, ModelResponse) or isinstance(payload, AsyncGenerator):
+                output = MessageOutput(source=payload, task_id=self.runner.task.id)
+        except Exception as e:
+            logger.warning(f"Failed to parse output: {e}")
+            yield Message(
+                category=Constants.TASK,
+                payload=TaskItem(msg="Failed to parse output.",
+                                 data=payload, stop=True),
+                sender=self.name(),
+                session_id=self.runner.context.session_id,
+                topic=TopicType.ERROR,
+                headers={"context": message.context}
+            )
+        finally:
+            if output:
+                if not output.metadata:
+                    output.metadata = {}
+                output.metadata['sender'] = message.sender
+                output.metadata['receiver'] = message.receiver
+                await outputs.add_output(output)
+            if mark_complete:
+                logger.info(f"FINISHED|output mark_completed|{self.runner.task.id}")
+                await outputs.mark_completed()
+
+        return