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

chapter9/gaia-experience/AWorld/aworld/runners/task_runner.py

@@ -0,0 +1,160 @@
+# coding: utf-8
+# Copyright (c) 2025 inclusionAI.
+import abc
+import os
+import time
+import uuid
+from typing import Callable, Any
+
+from pydantic import BaseModel
+
+import aworld.tools
+from aworld.config import ConfigDict
+from aworld.config.conf import ToolConfig
+from aworld.core.agent.swarm import Swarm
+from aworld.core.common import Observation
+from aworld.core.context.base import Context
+from aworld.core.context.session import Session
+from aworld.core.tool.base import Tool, AsyncTool
+from aworld.core.task import Task, TaskResponse, Runner
+from aworld.logs.util import logger
+from aworld import trace, cleanup
+from aworld.utils.common import load_module_by_path
+
+
+class TaskRunner(Runner):
+    """Task based runner api class."""
+    __metaclass__ = abc.ABCMeta
+
+    def __init__(self,
+                 task: Task,
+                 *,
+                 agent_oriented: bool = True,
+                 daemon_target: Callable[..., Any] = None):
+        """Task runner initialize.
+
+        Args:
+            task: Task entity to be executed.
+            agent_oriented: Is it an agent oriented task, default is True.
+        """
+        if task.tools is None:
+            task.tools = []
+        if task.tool_names is None:
+            task.tool_names = []
+
+        if agent_oriented:
+            if not task.agent and not task.swarm:
+                raise ValueError("agent and swarm all is None.")
+            if task.agent and task.swarm:
+                logger.warning("agent and swarm all is not None.")
+                raise ValueError("agent and swarm choose one only.")
+            if task.agent:
+                # uniform agent
+                task.swarm = Swarm(task.agent)
+
+        if task.conf is None:
+            task.conf = dict()
+        if isinstance(task.conf, BaseModel):
+            task.conf = task.conf.model_dump()
+        task.conf = ConfigDict(task.conf)
+        check_input = task.conf.get("check_input", False)
+        if check_input and not task.input:
+            raise ValueError("task no input")
+
+        if not task.is_sub_task:
+            self.context = task.context if task.context else Context()
+            self.context.set_task(task)
+        self.task = task
+        self.agent_oriented = agent_oriented
+        self.daemon_target = daemon_target
+        self._use_demon = False if not task.conf else task.conf.get(
+            'use_demon', False)
+        self._exception = None
+        self.start_time = time.time()
+        self.step_agent_counter = {}
+
+    async def pre_run(self):
+        task = self.task
+        # copy context from parent_task(if exists)
+        if task.is_sub_task:
+            task.context = await task.context.build_sub_context(
+                task.input, task.id,
+                agents=task.swarm.agents if task.swarm and task.swarm.agents else None
+            )
+            self.context = task.context
+            self.context.set_task(task)
+        self.swarm = task.swarm
+        self.input = task.input
+        self.outputs = task.outputs
+        self.name = task.name
+        self.conf = task.conf if task.conf else ConfigDict()
+        self.tools = {
+            tool.name(): tool for tool in task.tools} if task.tools else {}
+        task.tool_names.extend(self.tools.keys())
+        # lazy load
+        self.tool_names = task.tool_names
+        self.tools_conf = task.tools_conf
+        if self.tools_conf is None:
+            self.tools_conf = {}
+        # mcp performs special process, use async only in the runn
+        self.tools_conf['mcp'] = ToolConfig(use_async=True, name='mcp')
+        self.endless_threshold = task.endless_threshold
+
+        # build context
+        if task.session_id:
+            session = Session(session_id=task.session_id)
+        else:
+            session = Session(session_id=uuid.uuid4().hex)
+        trace_id = uuid.uuid1().hex if trace.get_current_span(
+        ) is None else trace.get_current_span().get_trace_id()
+        self.context.task_id = self.task.id
+        self.context.trace_id = trace_id
+        self.context.session = session
+        self.context.swarm = self.swarm
+
+        # init tool state by reset(), and ignore them observation
+        observation = None
+        if self.tools:
+            for _, tool in self.tools.items():
+                # use the observation and info of the last one
+                if isinstance(tool, Tool):
+                    tool.context = self.context
+                    observation, info = tool.reset()
+                elif isinstance(tool, AsyncTool):
+                    observation, info = await tool.reset()
+                else:
+                    logger.warning(f"Unsupported tool type: {tool}, will ignored.")
+
+        if observation:
+            if not observation.content:
+                observation.content = self.input
+        else:
+            observation = Observation(content=self.input)
+
+        self.observation = observation
+        if self.swarm:
+            self.swarm.event_driven = task.event_driven
+            self.swarm.reset(observation.content,
+                             context=self.context, tools=self.tool_names)
+
+        self._load_tool_module()
+        logger.info(f'{"sub task: " if self.task.is_sub_task else "main task: "}{self.task.id} started...')
+
+    def _load_tool_module(self):
+        # used to distributed running local tools
+        try:
+            value = os.environ.get(aworld.tools.LOCAL_TOOLS_ENV_VAR, '')
+            if value:
+                for val in value.split(";"):
+                    load_module_by_path(os.path.basename(val).replace("_action", ""),
+                                        val.replace("_action.py", ".py"))
+                    load_module_by_path(os.path.basename(val), val)
+        except:
+            logger.warning(f"{os.environ.get(aworld.tools.LOCAL_TOOLS_ENV_VAR, '')} tools load fail, can't use them!!")
+
+    async def post_run(self):
+        pass
+
+    @abc.abstractmethod
+    async def do_run(self, context: Context = None) -> TaskResponse:
+        """Task do run."""