diff --git a/plugins/sub_agent.py b/plugins/sub_agent.py new file mode 100644 index 0000000..a0c4a9c --- /dev/null +++ b/plugins/sub_agent.py @@ -0,0 +1,511 @@ +""" +子 Agent 插件 - 多 Agent 协作系统 +主 AI 可以创建子 Agent、下达指令、获取结果 +""" +import os +import json +import uuid +import time +import threading +from datetime import datetime +from typing import Dict, List, Optional + + +class SubAgent: + """单个子 Agent 实例""" + + def __init__(self, agent_id: str, name: str, task: str, system_prompt: str): + self.id = agent_id + self.name = name + self.task = task + self.status = "idle" # idle / running / done / failed + self.system_prompt = system_prompt + self.conversation = [] # 独立对话历史 + self.result = None + self.error = None + self.created_at = datetime.now().strftime("%H:%M:%S") + self.finished_at = None + self.steps = 0 # 已执行步数 + self.max_steps = 20 + + def to_dict(self): + return { + "id": self.id, + "name": self.name, + "task": self.task, + "status": self.status, + "steps": self.steps, + "created_at": self.created_at, + "finished_at": self.finished_at, + "has_result": self.result is not None, + "has_error": self.error is not None, + } + + +class Plugin: + """子 Agent 系统 - 创建、管理、协作""" + + def __init__(self): + self.usage = """子 Agent 系统 - 多 Agent 协作 + +管理: + spawn <名称> <任务描述> - 创建子 Agent 并分配任务 + spawn <名称> <任务> --system <提示词> - 创建并指定系统提示词 + list - 列出所有子 Agent + status - 查看子 Agent 状态 + result - 获取子 Agent 执行结果 + send <消息> - 向子 Agent 发送消息 + kill - 终止子 Agent + clean - 清理已完成的 Agent + stats - 统计信息 + +消息格式 (AI 调用): + {"action": "use_tool", "tool": "sub_agent", "args": "spawn coder 请帮我重构 main.py"} + {"action": "use_tool", "tool": "sub_agent", "args": "result agent_xxx"} + {"action": "use_tool", "tool": "sub_agent", "args": "send agent_xxx 用户要求加上类型提示"} + +说明: + - 子 Agent 独立运行,有自己的对话历史 + - 子 Agent 可以使用 code_editor、git_tools 等工具 + - 主 AI 可以同时管理多个子 Agent + - 子 Agent 完成后通过 result 获取结果 +""" + self.cli = None + self.agents: Dict[str, SubAgent] = {} + self._lock = threading.Lock() + + def set_cli(self, cli): + self.cli = cli + + def get_tool_info(self): + return { + "name": "sub_agent", + "description": "子 Agent 系统 - 创建子 Agent 执行独立任务,支持并行多任务、结果回收、消息传递", + "keywords": ["agent", "子agent", "并行", "任务", "协程", "多agent", "spawn", "sub"], + "usage": self.usage + } + + def get_mcp_definition(self): + return { + "name": "sub_agent", + "description": "子 Agent 多任务协作系统", + "inputSchema": { + "type": "object", + "properties": { + "operation": {"type": "string", "description": "操作: spawn/list/status/result/send/kill/clean/stats"}, + "name": {"type": "string", "description": "Agent 名称"}, + "task": {"type": "string", "description": "任务描述"}, + "agent_id": {"type": "string", "description": "Agent ID"}, + "message": {"type": "string", "description": "发送的消息"} + }, + "required": ["operation"] + } + } + + def convert_mcp_args(self, arguments): + op = arguments.get("operation", "") + name = arguments.get("name", "") + task = arguments.get("task", "") + return f"{op} {name} {task}".strip() + + # ── 路由 ── + + def handle(self, args: str) -> str: + try: + parts = args.strip().split(maxsplit=1) + if not parts: + return "错误:请提供操作类型" + + op = parts[0].lower() + rest = parts[1] if len(parts) > 1 else "" + + handlers = { + "spawn": self._spawn, + "list": self._list, + "status": self._status, + "result": self._result, + "send": self._send, + "kill": self._kill, + "clean": self._clean, + "stats": self._stats, + } + + handler = handlers.get(op) + if not handler: + return f"错误:不支持的操作 '{op}'" + return handler(rest) + + except Exception as e: + return f"子 Agent 错误: {str(e)}" + + # ── 操作 ── + + def _spawn(self, args: str) -> str: + """创建子 Agent""" + # 解析 --system 参数 + system_prompt = "" + if " --system " in args: + args, system_prompt = args.split(" --system ", 1) + system_prompt = system_prompt.strip() + + parts = args.strip().split(maxsplit=1) + if len(parts) < 2: + return "错误:格式: spawn <名称> <任务描述>" + + name = parts[0].strip() + task = parts[1].strip() + + # 生成 ID + agent_id = f"agent_{name}_{uuid.uuid4().hex[:6]}" + + # 默认系统提示词 + if not system_prompt: + system_prompt = ( + f"你是子 Agent「{name}」,负责执行以下任务:\n" + f"{task}\n\n" + "你可以使用 code_editor、git_tools、cmd_executor、file_manager 等工具。\n" + "完成任务后,用以下格式报告结果:\n" + '{"status": "done", "summary": "简要总结", "details": "详细结果"}\n\n' + "请开始执行任务。" + ) + + agent = SubAgent(agent_id, name, task, system_prompt) + with self._lock: + self.agents[agent_id] = agent + + # 在后台线程执行 + thread = threading.Thread(target=self._run_agent, args=(agent,), daemon=True) + thread.start() + + return (f"✅ 已创建子 Agent\n" + f" ID: {agent_id}\n" + f" 名称: {name}\n" + f" 任务: {task}\n" + f" 状态: running") + + def _run_agent(self, agent: SubAgent): + """在后台执行子 Agent""" + agent.status = "running" + + try: + engine = None + if self.cli and hasattr(self.cli, 'current_engine'): + engine = self.cli.current_engine + + if not engine: + agent.status = "failed" + agent.error = "无可用的 AI 引擎" + agent.finished_at = datetime.now().strftime("%H:%M:%S") + return + + # 第一轮:让子 Agent 开始执行任务 + current_input = agent.task + max_steps = agent.max_steps + + while agent.steps < max_steps and agent.status == "running": + agent.steps += 1 + + # 调用引擎 + response = engine.generate_response( + current_input, + system_prompt=agent.system_prompt + ) + + if not response: + agent.status = "failed" + agent.error = "引擎返回空响应" + break + + # 记录对话 + agent.conversation.append({"role": "user", "content": current_input}) + agent.conversation.append({"role": "assistant", "content": response}) + + # 检查是否包含完成信号 + if self._is_task_done(response): + agent.status = "done" + agent.result = response + break + + # 检查是否需要工具调用 + tool_result = self._try_execute_tool(response) + if tool_result: + current_input = f"工具执行结果:\n{tool_result}\n\n请继续执行任务。" + else: + # 没有工具调用,视为完成 + agent.status = "done" + agent.result = response + break + + if agent.status == "running": + agent.status = "done" + agent.result = "已达到最大步数限制,任务暂停。" + + except Exception as e: + agent.status = "failed" + agent.error = str(e) + + agent.finished_at = datetime.now().strftime("%H:%M:%S") + + def _is_task_done(self, response: str) -> bool: + """检查响应中是否包含完成信号""" + done_markers = ['"status": "done"', '"status":"done"', + '任务完成', '已完成', '执行完毕'] + return any(m in response for m in done_markers) + + def _try_execute_tool(self, response: str) -> Optional[str]: + """尝试从响应中提取并执行工具调用""" + if not self.cli: + return None + + # 查找 JSON 工具调用 + import re + + # 模式: {"action": "use_tool", "tool": "...", "args": "..."} + patterns = [ + r'\{"action"\s*:\s*"use_tool"\s*,\s*"tool"\s*:\s*"([^"]+)"\s*,\s*"args"\s*:\s*"([^"]*(?:\\.[^"]*)*)"\s*\}', + r'\{"action"\s*:\s*"use_tool"\s*,\s*"tool"\s*:\s*"([^"]+)"\s*,\s*"args"\s*:\s*(\{[^}]+\})\s*\}', + ] + + for pattern in patterns: + match = re.search(pattern, response, re.DOTALL) + if match: + tool_name = match.group(1) + tool_args = match.group(2) + + # 反转义 + tool_args = tool_args.replace('\\"', '"').replace('\\n', '\n') + + # 查找并执行工具 + if hasattr(self.cli, 'liugin_manager'): + for tool in self.cli.liugin_manager.tools: + if tool.get('name') == tool_name: + handler = tool.get('handler') + if handler: + try: + result = handler(tool_args) + return result + except Exception as e: + return f"工具执行错误: {e}" + break + + # 模式2: 多个 JSON 指令 (数组) + array_match = re.search(r'\[\s*\{.*?\}\s*\]', response, re.DOTALL) + if array_match: + try: + commands = json.loads(array_match.group()) + if isinstance(commands, list): + results = [] + for cmd in commands: + if isinstance(cmd, dict) and cmd.get('action') == 'use_tool': + tool_name = cmd.get('tool', '') + tool_args = cmd.get('args', '') + if isinstance(tool_args, dict): + tool_args = json.dumps(tool_args, ensure_ascii=False) + + for tool in self.cli.liugin_manager.tools: + if tool.get('name') == tool_name: + handler = tool.get('handler') + if handler: + try: + r = handler(tool_args) + results.append(f"[{tool_name}] {r}") + except Exception as e: + results.append(f"[{tool_name}] 错误: {e}") + break + if results: + return '\n'.join(results) + except json.JSONDecodeError: + pass + + return None + + def _list(self, args: str) -> str: + """列出所有子 Agent""" + with self._lock: + if not self.agents: + return "没有子 Agent" + + result = [f"子 Agent 列表 ({len(self.agents)} 个):\n"] + + # 按状态分组 + running = [a for a in self.agents.values() if a.status == "running"] + done = [a for a in self.agents.values() if a.status == "done"] + failed = [a for a in self.agents.values() if a.status == "failed"] + + if running: + result.append(f"🟢 运行中 ({len(running)}):") + for a in running: + result.append(f" {a.id} [{a.name}] 步骤:{a.steps} 任务:{a.task[:40]}") + + if done: + result.append(f"✅ 已完成 ({len(done)}):") + for a in done: + result.append(f" {a.id} [{a.name}] 任务:{a.task[:40]}") + + if failed: + result.append(f"❌ 失败 ({len(failed)}):") + for a in failed: + result.append(f" {a.id} [{a.name}] 错误:{a.error[:40] if a.error else '?'}") + + return '\n'.join(result) + + def _status(self, args: str) -> str: + """查看子 Agent 状态""" + agent_id = args.strip() + if not agent_id: + return "错误:请提供 Agent ID" + + agent = self.agents.get(agent_id) + if not agent: + # 模糊匹配 + matches = [a for a in self.agents.values() if agent_id in a.id or agent_id == a.name] + if matches: + agent = matches[0] + else: + return f"错误:未找到 Agent '{agent_id}'" + + lines = [ + f"Agent: {agent.id}", + f" 名称: {agent.name}", + f" 任务: {agent.task}", + f" 状态: {agent.status}", + f" 步骤: {agent.steps}/{agent.max_steps}", + f" 创建: {agent.created_at}", + ] + if agent.finished_at: + lines.append(f" 完成: {agent.finished_at}") + if agent.error: + lines.append(f" 错误: {agent.error}") + lines.append(f" 对话轮次: {len(agent.conversation)}") + + return '\n'.join(lines) + + def _result(self, args: str) -> str: + """获取子 Agent 结果""" + agent_id = args.strip() + if not agent_id: + return "错误:请提供 Agent ID" + + agent = self.agents.get(agent_id) + if not agent: + matches = [a for a in self.agents.values() if agent_id in a.id or agent_id == a.name] + if matches: + agent = matches[0] + else: + return f"错误:未找到 Agent '{agent_id}'" + + if agent.status == "running": + return f"Agent {agent.id} 仍在运行中 (步骤 {agent.steps}/{agent.max_steps})" + + if agent.status == "failed": + return f"Agent {agent.id} 执行失败: {agent.error}" + + if agent.result: + result = agent.result + if len(result) > 3000: + result = result[:3000] + "\n... (截断)" + return f"Agent {agent.id} 结果:\n{result}" + + return f"Agent {agent.id} 暂无结果" + + def _send(self, args: str) -> str: + """向子 Agent 发送消息""" + parts = args.strip().split(maxsplit=1) + if len(parts) < 2: + return "错误:格式: send <消息>" + + agent_id = parts[0] + message = parts[1] + + agent = self.agents.get(agent_id) + if not agent: + matches = [a for a in self.agents.values() if agent_id in a.id or agent_id == a.name] + if matches: + agent = matches[0] + else: + return f"错误:未找到 Agent '{agent_id}'" + + if agent.status != "running": + return f"Agent {agent.id} 未在运行 (状态: {agent.status})" + + # 在后台线程处理消息 + thread = threading.Thread(target=self._process_message, args=(agent, message), daemon=True) + thread.start() + + return f"✅ 消息已发送给 {agent.id}" + + def _process_message(self, agent: SubAgent, message: str): + """处理发送给子 Agent 的消息""" + try: + engine = None + if self.cli and hasattr(self.cli, 'current_engine'): + engine = self.cli.current_engine + + if not engine: + return + + agent.conversation.append({"role": "user", "content": message}) + agent.steps += 1 + + response = engine.generate_response( + message, + system_prompt=agent.system_prompt + ) + + if response: + agent.conversation.append({"role": "assistant", "content": response}) + + if self._is_task_done(response): + agent.status = "done" + agent.result = response + agent.finished_at = datetime.now().strftime("%H:%M:%S") + else: + tool_result = self._try_execute_tool(response) + if tool_result: + # 继续执行 + self._process_message(agent, f"工具结果:\n{tool_result}\n请继续。") + + except Exception as e: + agent.error = str(e) + + def _kill(self, args: str) -> str: + """终止子 Agent""" + agent_id = args.strip() + agent = self.agents.get(agent_id) + if not agent: + matches = [a for a in self.agents.values() if agent_id in a.id or agent_id == a.name] + if matches: + agent = matches[0] + else: + return f"错误:未找到 Agent '{agent_id}'" + + agent.status = "failed" + agent.error = "被主 Agent 终止" + agent.finished_at = datetime.now().strftime("%H:%M:%S") + return f"✅ 已终止 {agent.id}" + + def _clean(self, args: str) -> str: + """清理已完成的 Agent""" + with self._lock: + to_remove = [aid for aid, a in self.agents.items() + if a.status in ("done", "failed")] + for aid in to_remove: + del self.agents[aid] + + return f"✅ 已清理 {len(to_remove)} 个已完成的 Agent" + + def _stats(self, args: str) -> str: + """统计信息""" + with self._lock: + total = len(self.agents) + running = sum(1 for a in self.agents.values() if a.status == "running") + done = sum(1 for a in self.agents.values() if a.status == "done") + failed = sum(1 for a in self.agents.values() if a.status == "failed") + total_steps = sum(a.steps for a in self.agents.values()) + + return (f"📊 子 Agent 统计\n" + f" 总数: {total}\n" + f" 运行中: {running}\n" + f" 已完成: {done}\n" + f" 失败: {failed}\n" + f" 总步数: {total_steps}") diff --git a/tests/test_plugins.py b/tests/test_plugins.py index 3d7eef3..585a241 100644 --- a/tests/test_plugins.py +++ b/tests/test_plugins.py @@ -353,6 +353,65 @@ class TestBrowserAuto: assert len(info["keywords"]) > 0 +# ══════════════════════════════════════ +# 引擎加载测试 +# ══════════════════════════════════════ + +class TestSubAgent: + """sub_agent 功能测试""" + + @pytest.fixture + def sa(self): + mod = load_module("sub_agent", os.path.join(PLUGINS_DIR, "sub_agent.py")) + return mod.Plugin() + + def test_list_empty(self, sa): + result = sa.handle("list") + assert "没有" in result + + def test_stats_empty(self, sa): + result = sa.handle("stats") + assert "总数: 0" in result + + def test_spawn(self, sa): + result = sa.handle("spawn tester 测试任务") + assert "已创建" in result + assert "agent_tester" in result + + def test_spawn_and_list(self, sa): + sa.handle("spawn a 任务A") + sa.handle("spawn b 任务B") + import time; time.sleep(0.1) + result = sa.handle("list") + assert "agent_a" in result + assert "agent_b" in result + + def test_status(self, sa): + sa.handle("spawn x 某任务") + import time; time.sleep(0.1) + # 获取 agent ID + agent_id = list(sa.agents.keys())[0] + result = sa.handle(f"status {agent_id}") + assert "名称" in result + + def test_kill(self, sa): + sa.handle("spawn y 要终止的任务") + import time; time.sleep(0.1) + agent_id = list(sa.agents.keys())[0] + result = sa.handle(f"kill {agent_id}") + assert "已终止" in result + + def test_clean(self, sa): + sa.handle("spawn z 清理测试") + import time; time.sleep(0.1) + result = sa.handle("clean") + assert "已清理" in result + + def test_invalid_agent(self, sa): + result = sa.handle("status nonexistent") + assert "未找到" in result + + # ══════════════════════════════════════ # 引擎加载测试 # ══════════════════════════════════════