agent-Specialization/modules/sub_agent/task.py
JOJO 9ed956518c refactor(sub_agent): 子智能体从 Node.js 子进程改为主进程内 Python 协程
- 重写子智能体执行核心,不再启动 easyagent Node.js 子进程
- 新增 modules/sub_agent/ 包集中管理子智能体逻辑
- 工具调用复用主进程 WebTerminal.handle_tool_call,自然经过沙箱/容器链路
- 子智能体模型独立读取 ~/.agents/<mode>/config/sub_agent_models.json
- 支持 8 个工具:read_file/write_file/edit_file(replacements+replace_all)/run_command/web_search/extract_webpage/search_workspace/read_mediafile
- 修复子智能体进度弹窗:标题颜色、write_file 显示、过滤非 progress 条目、统一滚动条样式
- 更新 AGENTS.md / CLAUDE.md 子智能体描述
- 新增 test/test_sub_agent_regression.py 回归测试
2026-06-20 00:26:45 +08:00

359 lines
14 KiB
Python

from __future__ import annotations
import asyncio
import base64
import json
import mimetypes
import time
import uuid
from datetime import datetime
from pathlib import Path
from typing import Any, Dict, List, Optional, TYPE_CHECKING
from modules.sub_agent.toolkit import (
SUB_AGENT_TOOLS,
FINISH_TOOL,
_format_tool_result,
_build_sub_agent_profile,
)
from utils.api_client import DeepSeekClient
from utils.logger import setup_logger
if TYPE_CHECKING:
from modules.sub_agent.manager import SubAgentManager
logger = setup_logger(__name__)
class SubAgentTask:
"""单个后台子智能体任务。"""
def __init__(
self,
manager: "SubAgentManager",
task_record: Dict[str, Any],
task_message: str,
system_prompt: str,
model_key: Optional[str],
thinking_mode: Optional[str],
):
self.manager = manager
self.task_record = task_record
self.task_message = task_message
self.system_prompt = system_prompt
self.model_key = model_key
self.thinking_mode = thinking_mode or "fast"
self.task_id = task_record["task_id"]
self.agent_id = task_record["agent_id"]
self.timeout_seconds = int(task_record.get("timeout_seconds") or 180)
self.deliverables_dir = Path(task_record["deliverables_dir"])
self.output_file = Path(task_record["output_file"])
self.stats_file = Path(task_record["stats_file"])
self.progress_file = Path(task_record["progress_file"])
self.conversation_file = Path(task_record["conversation_file"])
self.messages: List[Dict[str, Any]] = [
{"role": "system", "content": system_prompt},
{"role": "user", "content": task_message},
]
self.stats = {
"runtime_start": time.time() * 1000,
"runtime_seconds": 0,
"files_read": 0,
"write_files": 0,
"edit_files": 0,
"searches": 0,
"web_pages": 0,
"commands": 0,
"api_calls": 0,
"token_usage": {"prompt": 0, "completion": 0, "total": 0},
}
self._stdout_lines: List[str] = []
self._cancelled = False
self._task: Optional[asyncio.Task] = None
def emit(self, type_: str, data: Dict[str, Any]) -> None:
"""输出一行 JSONL 到 progress 文件并缓存。"""
line = json.dumps({"type": type_, **data}, ensure_ascii=False)
self._stdout_lines.append(line)
try:
self.progress_file.parent.mkdir(parents=True, exist_ok=True)
with open(self.progress_file, "a", encoding="utf-8") as f:
f.write(line + "\n")
except Exception:
pass
async def run(self) -> None:
"""主 LLM 循环。"""
try:
await self._run_loop()
except asyncio.CancelledError:
self._cancelled = True
logger.debug(f"[SubAgent] task={self.task_id} 被取消")
# shield 避免取消信号中断最终状态落盘
await asyncio.shield(self._write_failure("子智能体被手动终止"))
raise
except Exception as exc:
logger.exception(f"[SubAgent] task={self.task_id} 执行异常")
await self._write_failure(f"执行异常: {exc}")
async def _run_loop(self) -> None:
client, model_key = self._build_client()
tools = list(SUB_AGENT_TOOLS)
tools.append(FINISH_TOOL)
start_time = time.time()
max_turns = 50
for turn in range(1, max_turns + 1):
if self._cancelled:
break
elapsed = time.time() - start_time
if elapsed > self.timeout_seconds:
await self._write_timeout(elapsed)
return
self.stats["api_calls"] += 1
self.stats["turn_count"] = turn
self.stats["runtime_seconds"] = int(elapsed)
self.emit("stats", {**self.stats, "turn_count": turn})
assistant_message, reasoning, tool_calls, usage = await self._call_model(client, model_key, tools)
if usage:
self._apply_usage(usage)
final_message: Dict[str, Any] = {"role": "assistant", "content": assistant_message}
if reasoning:
final_message["reasoning_content"] = reasoning
if tool_calls:
final_message["tool_calls"] = tool_calls
self.messages.append(final_message)
if not tool_calls:
self.messages.append({
"role": "user",
"content": "如果你已经完成了任务,请调用 finish_task 工具提交完成报告。如果还没有完成,请继续执行任务。",
})
continue
for tool_call in tool_calls:
if self._cancelled:
break
name = tool_call.get("function", {}).get("name", "")
args = self._parse_args(tool_call)
progress_id = tool_call.get("id") or f"tool_{int(time.time() * 1000)}_{uuid.uuid4().hex[:6]}"
if name == "finish_task":
await self._write_finish(args, elapsed)
return
self.emit("progress", {"id": progress_id, "tool": name, "status": "running", "args": args, "ts": int(time.time() * 1000)})
result = await self._execute_tool(name, args)
self.emit("progress", {"id": progress_id, "tool": name, "status": "completed" if result.get("success") else "failed", "args": args, "ts": int(time.time() * 1000)})
self._update_stats(name)
content = _format_tool_result(name, result)
if name == "read_mediafile" and result.get("success"):
content = self._build_media_tool_content(result) or content
self.messages.append({
"role": "tool",
"tool_call_id": tool_call.get("id", progress_id),
"content": content,
})
await self._write_failure("任务执行超过最大轮次限制", max_turns_exceeded=True)
def _build_client(self) -> tuple:
"""加载模型配置并初始化 DeepSeekClient。"""
config_path = self.manager.models_config_file
models: List[Dict[str, Any]] = []
default_key = ""
if Path(config_path).exists():
try:
raw = json.loads(Path(config_path).read_text(encoding="utf-8"))
models = raw.get("models", []) if isinstance(raw, dict) else (raw if isinstance(raw, list) else [])
default_key = str(raw.get("default_model", "")) if isinstance(raw, dict) else ""
except Exception as exc:
logger.error(f"[SubAgent] 加载模型配置失败: {exc}")
model_map = {}
valid_models = []
for item in models:
profile = _build_sub_agent_profile(item)
if profile:
key = profile["name"]
model_map[key] = profile
valid_models.append(key)
chosen_key = self.model_key or default_key
if chosen_key not in model_map and valid_models:
chosen_key = valid_models[0]
if chosen_key not in model_map:
raise RuntimeError(f"未找到可用子智能体模型配置: {config_path}")
client = DeepSeekClient(thinking_mode=(self.thinking_mode == "thinking"), web_mode=True)
client.model_key = chosen_key
client.project_path = str(self.manager.project_path)
if self.thinking_mode == "thinking":
# 子智能体的 thinking 模式应全程使用思考模型
client.deep_thinking_session = True
client.apply_profile(model_map[chosen_key])
return client, chosen_key
async def _call_model(
self,
client: DeepSeekClient,
model_key: str,
tools: List[Dict[str, Any]],
) -> tuple:
"""调用模型并解析 assistant 消息。"""
assistant_message = ""
reasoning = ""
tool_calls: List[Dict[str, Any]] = []
usage = None
async for chunk in client.chat(self.messages, tools=tools, stream=True):
if self._cancelled:
break
if chunk.get("error"):
raise RuntimeError(f"API 调用失败: {chunk.get('error')}")
choice = (chunk.get("choices") or [{}])[0]
delta = choice.get("delta") or {}
if delta.get("content"):
assistant_message += delta["content"]
if delta.get("reasoning_content"):
reasoning += delta["reasoning_content"]
elif delta.get("reasoning_details"):
rd = delta["reasoning_details"]
if isinstance(rd, list):
reasoning += "".join(str(d.get("text") or "") for d in rd)
elif isinstance(rd, str):
reasoning += rd
elif isinstance(rd, dict):
reasoning += str(rd.get("text") or "")
for tc in delta.get("tool_calls") or []:
idx = tc.get("index")
if idx is None:
continue
while len(tool_calls) <= idx:
tool_calls.append({"id": "", "type": "function", "function": {"name": "", "arguments": ""}})
existing = tool_calls[idx]
if tc.get("id"):
existing["id"] = tc["id"]
fn = tc.get("function") or {}
if fn.get("name"):
existing["function"]["name"] += fn["name"]
if fn.get("arguments"):
existing["function"]["arguments"] += fn["arguments"]
if chunk.get("usage"):
usage = chunk["usage"]
return assistant_message, reasoning, tool_calls, usage
def _parse_args(self, tool_call: Dict[str, Any]) -> Dict[str, Any]:
raw = tool_call.get("function", {}).get("arguments") or "{}"
try:
return json.loads(raw)
except Exception:
return {"_raw": raw}
async def _execute_tool(self, name: str, args: Dict[str, Any]) -> Dict[str, Any]:
"""通过 manager 调用主进程执行工具。"""
return await self.manager.execute_tool_for_sub_agent(name, args)
def _update_stats(self, name: str) -> None:
if name == "read_file":
self.stats["files_read"] += 1
elif name == "write_file":
self.stats["write_files"] += 1
elif name == "edit_file":
self.stats["edit_files"] += 1
elif name == "search_workspace":
self.stats["searches"] += 1
elif name in ("web_search", "extract_webpage"):
self.stats["web_pages"] += 1
elif name == "run_command":
self.stats["commands"] += 1
def _apply_usage(self, usage: Any) -> None:
try:
if isinstance(usage, dict):
prompt = usage.get("prompt_tokens") or usage.get("prompt") or 0
completion = usage.get("completion_tokens") or usage.get("completion") or 0
total = usage.get("total_tokens") or usage.get("total") or (prompt + completion)
self.stats["token_usage"]["prompt"] += int(prompt)
self.stats["token_usage"]["completion"] += int(completion)
self.stats["token_usage"]["total"] += int(total)
except Exception:
pass
def _build_media_tool_content(self, result: Dict[str, Any]) -> Any:
"""把 read_mediafile 结果转成 OpenAI 多模态 content。"""
b64 = result.get("b64")
mime = result.get("mime")
file_type = result.get("type")
if not b64 or not mime:
return None
if file_type == "image":
return [
{"type": "text", "text": f"已附加图片: {result.get('path')}"},
{"type": "image_url", "image_url": {"url": f"data:{mime};base64,{b64}"}},
]
if file_type == "video":
return [
{"type": "text", "text": f"已附加视频: {result.get('path')}"},
{"type": "video_url", "video_url": {"url": f"data:{mime};base64,{b64}"}},
]
return None
async def _write_finish(self, args: Dict[str, Any], elapsed: float) -> None:
success = bool(args.get("success", False))
summary = str(args.get("summary") or "").strip()
self._finalize_task(success, summary, elapsed)
async def _write_timeout(self, elapsed: float) -> None:
self._finalize_task(False, "任务超时未完成", elapsed, timeout=True)
async def _write_failure(self, message: str, *, max_turns_exceeded: bool = False, timeout: bool = False) -> None:
elapsed = time.time() - (self.stats["runtime_start"] / 1000)
self._finalize_task(False, message, elapsed, max_turns_exceeded=max_turns_exceeded, timeout=timeout)
def _finalize_task(self, success: bool, summary: str, elapsed: float, *, max_turns_exceeded: bool = False, timeout: bool = False) -> None:
runtime_seconds = int(elapsed)
output_data = {
"success": success,
"summary": summary,
"timeout": timeout,
"max_turns_exceeded": max_turns_exceeded,
"stats": {**self.stats, "runtime_seconds": runtime_seconds, "turn_count": self.stats.get("turn_count", 0)},
}
conversation_data = {
"agent_id": self.agent_id,
"created_at": datetime.fromtimestamp(self.stats["runtime_start"] / 1000).isoformat(),
"completed_at": datetime.now().isoformat(),
"success": success,
"summary": summary,
"messages": self.messages,
"stats": output_data["stats"],
}
stats_data = {**self.stats, "runtime_seconds": runtime_seconds, "turn_count": self.stats.get("turn_count", 0)}
self.output_file.parent.mkdir(parents=True, exist_ok=True)
self.output_file.write_text(json.dumps(output_data, ensure_ascii=False), encoding="utf-8")
self.stats_file.write_text(json.dumps(stats_data, ensure_ascii=False), encoding="utf-8")
self.conversation_file.write_text(json.dumps(conversation_data, ensure_ascii=False), encoding="utf-8")
self.emit("output", output_data)
self.emit("conversation", conversation_data)
self.manager._mark_task_done(self.task_id, success, summary, runtime_seconds)
def cancel(self) -> None:
self._cancelled = True
if self._task and not self._task.done():
self._task.cancel()