1  # Copyright (c) 2024-2026 Tencent Zhuque Lab. All rights reserved.
  2  #
  3  # Licensed under the Apache License, Version 2.0 (the "License");
  4  # you may not use this file except in compliance with the License.
  5  # You may obtain a copy of the License at
  6  #
  7  #     http://www.apache.org/licenses/LICENSE-2.0
  8  #
  9  # Unless required by applicable law or agreed to in writing, software
 10  # distributed under the License is distributed on an "AS IS" BASIS,
 11  # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 12  # See the License for the specific language governing permissions and
 13  # limitations under the License.
 14  #
 15  # Requirement: Any integration or derivative work must explicitly attribute
 16  # Tencent Zhuque Lab (https://github.com/Tencent/AI-Infra-Guard) in its
 17  # documentation or user interface, as detailed in the NOTICE file.
 18  
 19  """
 20  Task 工具 - 子 Agent 任务执行工具
 21  遵循与 Skill 相同的标准：
 22  - Agent 模板存储在 prompt/agents/ 下
 23  - 支持 YAML Frontmatter 元数据
 24  """
 25  import os
 26  import re
 27  import uuid
 28  import yaml
 29  from typing import Any, Optional, Dict, List
 30  from tools.registry import register_tool
 31  from utils.logging import logger
 32  from utils.tool_context import ToolContext
 33  from utils.config import base_dir
 34  
 35  AGENTS_DIR = os.path.join(base_dir, "prompt", "agents")
 36  
 37  
 38  def parse_agent_file(file_path: str) -> Dict[str, Any]:
 39      """解析带有 YAML Frontmatter 的 Agent 文件"""
 40      if not os.path.exists(file_path):
 41          return {}
 42  
 43      with open(file_path, 'r', encoding='utf-8') as f:
 44          content = f.read()
 45  
 46      meta = {}
 47      body = content
 48  
 49      # 使用正则匹配 YAML Frontmatter
 50      frontmatter_pattern = re.compile(r'^---\s*\n(.*?)\n---\s*\n', re.DOTALL)
 51      match = frontmatter_pattern.match(content)
 52  
 53      if match:
 54          yaml_content = match.group(1)
 55          meta = yaml.safe_load(yaml_content) or {}
 56          body = content[match.end():].strip()
 57  
 58      # 确保 meta 中有基本信息
 59      if 'description' not in meta:
 60          lines = body.split('\n')
 61          desc_lines = []
 62          for line in lines[:5]:
 63              line = line.strip()
 64              if line and not line.startswith('#'):
 65                  desc_lines.append(line)
 66          meta['description'] = ' '.join(desc_lines)[:200]
 67  
 68      return {
 69          "meta": meta,
 70          "content": body,
 71          "raw": content
 72      }
 73  
 74  
 75  def get_all_agents() -> List[Dict[str, Any]]:
 76      """扫描目录获取所有 Agent"""
 77      agents = []
 78  
 79      if not os.path.exists(AGENTS_DIR):
 80          return agents
 81  
 82      for name in os.listdir(AGENTS_DIR):
 83          if name.startswith('.'):
 84              continue
 85  
 86          item_path = os.path.join(AGENTS_DIR, name)
 87  
 88          # 情况1: 直接是 .md 文件
 89          if os.path.isfile(item_path) and name.endswith('.md'):
 90              agent_name = name[:-3]  # 移除 .md
 91              data = parse_agent_file(item_path)
 92              meta = data.get('meta', {})
 93  
 94              agents.append({
 95                  "name": agent_name,
 96                  "title": meta.get('name', agent_name),
 97                  "description": meta.get('description', ''),
 98                  "meta": meta,
 99                  "path": item_path
100              })
101  
102          # 情况2: 是目录，查找 index.md 或同名 .md
103          elif os.path.isdir(item_path):
104              agent_file = None
105              # 优先查找 index.md
106              index_path = os.path.join(item_path, 'index.md')
107              main_path = os.path.join(item_path, f'{name}.md')
108  
109              if os.path.isfile(index_path):
110                  agent_file = index_path
111              elif os.path.isfile(main_path):
112                  agent_file = main_path
113  
114              if agent_file:
115                  data = parse_agent_file(agent_file)
116                  meta = data.get('meta', {})
117  
118                  agents.append({
119                      "name": name,
120                      "title": meta.get('name', name),
121                      "description": meta.get('description', ''),
122                      "path": agent_file
123                  })
124  
125      return sorted(agents, key=lambda x: x['name'])
126  
127  
128  def load_agent_prompt(agent_name: str) -> Optional[Dict[str, Any]]:
129      """
130      加载 Agent 提示词
131      
132      Args:
133          agent_name: Agent 名称
134          
135      Returns:
136          包含 meta 和 content 的字典，如果不存在返回 None
137      """
138      # 直接检查 .md 文件
139      direct_path = os.path.join(AGENTS_DIR, f"{agent_name}.md")
140      if os.path.isfile(direct_path):
141          return parse_agent_file(direct_path)
142  
143      # 检查目录
144      agent_dir = os.path.join(AGENTS_DIR, agent_name)
145      if os.path.isdir(agent_dir):
146          index_path = os.path.join(agent_dir, 'index.md')
147          main_path = os.path.join(agent_dir, f'{agent_name}.md')
148  
149          if os.path.isfile(index_path):
150              return parse_agent_file(index_path)
151          elif os.path.isfile(main_path):
152              return parse_agent_file(main_path)
153  
154      # 尝试模糊匹配
155      all_agents = get_all_agents()
156      for agent in all_agents:
157          if agent['name'] == agent_name:
158              return parse_agent_file(agent['path'])
159  
160      return None
161  
162  
163  @register_tool
164  async def task(
165          prompt: str,
166          subagent_type: str,
167          description: str = "",
168          context: ToolContext = None
169  ) -> dict[str, Any]:
170      """
171      执行子 Agent 任务
172      
173      Args:
174          prompt: 任务提示词
175          subagent_type: 子 Agent 类型
176          description: 任务描述（3-5 个词）
177          context: 工具上下文
178          
179      Returns:
180          包含执行结果的字典
181      """
182      # 加载 Agent 数据
183      agent_instruction = load_agent_prompt(subagent_type)
184  
185      if agent_instruction is None:
186          available = get_all_agents()
187          available_names = [a['name'] for a in available]
188  
189          return {
190              "success": False,
191              "error": f"Unknown agent type: {subagent_type}. Available agents: {', '.join(available_names) if available_names else 'none'}"
192          }
193  
194      # 构建任务提示词
195      task_prompt = f"""
196  Task: {description or 'Execute the following'}
197  
198  {prompt}
199  
200  Please complete this task and provide a summary of your actions and results.
201  """
202  
203      logger.info(f"Executing task with agent '{subagent_type}': {description or prompt[:50]}")
204  
205      result = await context.call_subagent(
206          description, subagent_type, task_prompt, uuid.uuid4().__str__(), context.language, "", {}
207      )
208  
209      return {
210          "success": True,
211          "title": description or f"Task: {subagent_type}",
212          "output": result,
213          "agent": subagent_type,
214      }
215  
216  
217  @register_tool
218  def list_agents(context: ToolContext = None) -> dict[str, Any]:
219      """
220      列出可用的子 Agent
221      
222      Returns:
223          包含 Agent 列表的字典
224      """
225      agents = get_all_agents()
226  
227      if not agents:
228          return {
229              "success": True,
230              "output": "No agents available.",
231              "agents": []
232          }
233  
234      output_lines = ["Available agents:"]
235      for agent in agents:
236          output_lines.append(f"- {agent['name']}: {agent['description']}")
237  
238      return {
239          "success": True,
240          "output": '\n'.join(output_lines),
241          "agents": agents
242      }