Coverage for agent.py: 17%

1 import json

2 from rich.console import Console

3 from rich.panel import Panel

4 from rich.markdown import Markdown

5 from qrclaw.config import MAX_ITERATIONS, COMPRESS_THRESHOLD

6 from qrclaw.providers import provider

7 from qrclaw.providers.base import LLMResponse

8 from qrclaw.tools.registry import get_schemas, execute, need_confirm

9 from qrclaw.memory.session import Session

10 from qrclaw.memory import compressor, LongTermMemory

11 from qrclaw.prompt import build_system_prompt

12 from qrclaw.cli.display import show_plan_progress

13 from qrclaw.workspace import Workspace

14 from qrclaw.logger import get_logger

15

16 logger = get_logger("qrclaw.agent")

17

18

19# 用 threading.local 隔离每个线程的 session/workspace

20# 多个子 agent 并行时，各自的上下文互不干扰

21 import threading

22 _thread_local = threading.local()

23

24 def set_session(session: Session):

25 """注入当前线程的 session"""

26 _thread_local.session = session

27

28 def get_session() -> Session | None:

29 """获取当前线程的 session"""

30 return getattr(_thread_local, "session", None)

31

32 def set_workspace(workspace: Workspace):

33 """注入当前线程的 workspace"""

34 _thread_local.workspace = workspace

35

36 def get_workspace() -> Workspace | None:

37 """获取当前线程的 workspace"""

38 return getattr(_thread_local, "workspace", None)

39

40 def set_agent_depth(depth: int):

41 """设置当前 agent 的深度（0=顶层 agent）"""

42 _thread_local.agent_depth = depth

43

44 def get_agent_depth() -> int:

45 """获取当前 agent 的深度，默认为 0（顶层）"""

46 return getattr(_thread_local, "agent_depth", 0)

47

48 def is_sub_agent() -> bool:

49 """判断当前是否是子 agent"""

50 return get_agent_depth() > 0

51

52

53 def _dump_assistant_msg(response: LLMResponse) -> dict:

54 """把 LLMResponse 转成可存入 session 的 assistant 消息 dict"""

55 msg: dict = {"role": "assistant", "content": response.content or ""}

56 if response.tool_calls:

57 tc_list = []

58 for tc in response.tool_calls:

59 entry = {

60 "id": tc.id,

61 "type": "function",

62 "function": {"name": tc.name, "arguments": tc.arguments},

63 }

64 # 保留 thought_signature，回传给 Vertex AI 时需要

65 if tc.thought_signature:

66 entry["__thought_signature__"] = tc.thought_signature

67 tc_list.append(entry)

68 msg["tool_calls"] = tc_list

69 return msg

70

71

72 def run(user_input: str, session: Session, console: Console, workspace: Workspace, auto_confirm: bool = False):

73 logger.info(f"收到用户输入: {user_input[:100]}...")

74

75 set_session(session)

76 set_workspace(workspace)

77 session.add({"role": "user", "content": user_input})

78

79 # system prompt 每次实时构建，不存进 session

80 # 这样工作目录、工具列表永远是最新的

81 from qrclaw.skills.registry import SkillRegistry

82 tool_names = [s["function"]["name"] for s in get_schemas()]

83 memory = LongTermMemory(workspace.memory_file)

84 skill_registry = SkillRegistry()

85 skill_registry.load_from_dir(workspace.skills_dir)

86 system_prompt = {

87 "role": "system",

88 "content": build_system_prompt(

89 tool_names,

90 memory,

91 skill_registry,

92 active_plan=session.active_plan,

93 heartbeat_file=workspace.heartbeat_file,

94 is_sub_agent=is_sub_agent(),

95 )

96 }

97 logger.debug(f"System prompt 已构建，可用工具: {', '.join(tool_names)}")

98

99 # 权限拒绝计数器

100 permission_denied_count = 0

101 MAX_PERMISSION_DENIED = 2

102

103 for iteration in range(MAX_ITERATIONS):

104 logger.debug(f"开始第 {iteration + 1} 轮推理")

105

106 # 每次调 LLM 时把 system prompt 拼到最前面

107 messages = [system_prompt, *session.messages]

108

109 # spinner 只包住 LLM 请求这一步，拿到响应立即退出

110 with console.status("[bold yellow]思考中...[/bold yellow]", spinner="dots"):

111 logger.debug(f"调用 LLM，消息数: {len(messages)}")

112 try:

113 response = provider.chat(messages, tools=get_schemas())

114 session.update_tokens(

115 prompt_tokens=response.prompt_tokens,

116 completion_tokens=response.completion_tokens,

117 total_tokens=response.total_tokens,

118 )

119 logger.info(f"LLM 响应成功，使用 {response.total_tokens} tokens")

120 except Exception as e:

121 logger.error(f"LLM 调用失败: {e}", exc_info=True)

122 raise

123

124 finish_reason = response.finish_reason

125

126 # 检查是否需要压缩

127 if response.prompt_tokens > COMPRESS_THRESHOLD:

128 logger.warning(f"Prompt tokens ({response.prompt_tokens}) 超过阈值 ({COMPRESS_THRESHOLD})，触发压缩")

129 compressor.summarize(session)

130 else:

131 logger.debug(f"Prompt tokens: {response.prompt_tokens}, 阈值: {COMPRESS_THRESHOLD}")

132

133 if finish_reason == "stop":

134 logger.info(f"推理完成，最终答案长度: {len(response.content)} 字符")

135 session.add({"role": "assistant", "content": response.content})

136 console.print() # 添加空行

137 console.print(Panel(

138 Markdown(response.content, code_theme="ansi_dark"),

139 title="[bold green]Agent[/bold green]",

140 border_style="green",

141 expand=True,

142 ))

143 console.print() # 添加空行

144 return response.content

145

146 if finish_reason == "length":

147 logger.warning("LLM 响应被截断 (finish_reason=length)")

148 console.print("[bold red]警告：回复被截断，请尝试简化任务[/bold red]\n")

149 return "错误：回复被截断"

150

151 # assistant 消息只存一次，在工具循环之前

152 assistant_msg_saved = False

153

154 for tc in response.tool_calls:

155 name = tc.name

156 arguments = tc.arguments

157 logger.info(f"工具调用: {name}, 参数: {arguments[:200]}...")

158

159 # 把 JSON 字符串格式化后高亮显示

160 try:

161 args_formatted = json.dumps(json.loads(arguments), ensure_ascii=False, indent=2)

162 except Exception:

163 args_formatted = arguments

164 console.print() # 添加空行

165 console.print(Panel(

166 f"[bold cyan]{name}[/bold cyan]\n" + args_formatted,

167 title="[bold yellow]▶ 调用工具[/bold yellow]",

168 border_style="yellow",

169 expand=False,

170 ))

171 console.print() # 添加空行

172

173 # 高风险工具执行前，询问用户确认（子 agent 自动跳过）

174 if need_confirm(name) and not auto_confirm:

175 logger.debug(f"工具 {name} 需要用户确认")

176 console.print(f"[bold red]⚠ 需要确认[/bold red] 是否允许执行？(y/n) ", end="")

177 choice = input().strip().lower()

178 if choice != "y":

179 result = "用户拒绝执行此操作"

180 logger.warning(f"用户拒绝执行工具: {name}")

181 console.print(Panel(result, title="[bold red]已拒绝[/bold red]", border_style="red", expand=False))

182 if not assistant_msg_saved:

183 session.add(_dump_assistant_msg(response))

184 assistant_msg_saved = True

185 session.add({"role": "tool", "tool_call_id": tc.id, "content": result})

186 break # 退出工具循环，回到外层让LLM重新推理

187

188 if not assistant_msg_saved:

189 session.add(_dump_assistant_msg(response))

190 assistant_msg_saved = True

191

192 try:

193 result = execute(name, arguments)

194 logger.info(f"工具执行成功: {name}, 结果长度: {len(result)} 字符")

195 # 工具执行成功，重置权限拒绝计数器

196 permission_denied_count = 0

197 except PermissionError as e:

198 # 权限拒绝

199 permission_denied_count += 1

200 result = str(e)

201 logger.warning(f"权限拒绝 ({permission_denied_count}/{MAX_PERMISSION_DENIED}): {name}, 原因: {e}")

202

203 # 检查是否达到上限

204 if permission_denied_count >= MAX_PERMISSION_DENIED:

205 console.print(Panel(

206 f"[bold red]您没有权限执行这个操作！[/bold red]\n\n"

207 f"连续 {MAX_PERMISSION_DENIED} 次权限拒绝，任务终止。\n\n"

208 f"最后一次拒绝原因:\n{e}",

209 title="[bold red]⛔ 权限不足[/bold red]",

210 border_style="red",

211 expand=False,

212 ))

213 console.print()

214 return f"错误：连续 {MAX_PERMISSION_DENIED} 次权限拒绝，任务终止"

215 except Exception as e:

216 result = f"工具执行失败: {str(e)}"

217 logger.error(f"工具执行失败: {name}, 错误: {e}", exc_info=True)

218

219 preview = result[:200] + "..." if len(result) > 200 else result

220 console.print(Panel(

221 preview,

222 title="[bold blue]◀ 工具结果[/bold blue]",

223 border_style="blue",

224 expand=False,

225 ))

226 console.print() # 添加空行

227

228 session.add({

229 "role": "tool",

230 "tool_call_id": tc.id,

231 "content": result,

232 })

233

234 # 如果已经达到权限拒绝上限，直接退出

235 if permission_denied_count >= MAX_PERMISSION_DENIED:

236 break

237

238 if name in ("create_plan", "complete_step"):

239 show_plan_progress(console, session)

240

241 logger.warning("达到最大迭代次数，强制退出")

242 return "错误：达到最大迭代次数"

243

244

245 def run_sub_agent(task: str, sub_workspace: Workspace) -> str:

246 """

247 以静默模式运行子 agent，返回结果字符串。

248 子 agent 不打印到用户终端，结果直接返回给调用方（主 agent）。

249 任务完成后自动清理工作空间（保留 logs，删除 sessions/skills/MEMORY.md）。

250

251 重要：子 agent 不允许再派生子 agent，防止无限嵌套。

252

253 Args:

254 task: 子 agent 要执行的任务描述

255 sub_workspace: 子 agent 的工作空间

256 Returns:

257 str: 子 agent 的最终回复

258 """

259 import shutil

260 from io import StringIO

261 from rich.console import Console as RichConsole

262 from qrclaw.memory.session import Session

263

264 logger.info(f"启动子 agent: {sub_workspace.agent_id}, 任务: {task[:100]}...")

265

266 # 设置子 agent 深度（当前深度 + 1）

267 current_depth = get_agent_depth()

268 set_agent_depth(current_depth + 1)

269 logger.info(f"子 agent 深度: {current_depth + 1}")

270

271 buffer = StringIO()

272 sub_console = RichConsole(file=buffer, highlight=False)

273 sub_session = Session(sessions_dir=sub_workspace.sessions_dir, resume=False)

274

275 try:

276 result = run(task, sub_session, sub_console, sub_workspace, auto_confirm=True)

277 result = result or "子 agent 未返回结果"

278 logger.info(f"子 agent {sub_workspace.agent_id} 执行完毕，结果长度: {len(result)} 字符")

279 finally:

280 # 恢复深度

281 set_agent_depth(current_depth)

282

283 # 清理工作空间：保留 logs，删除 sessions/skills/MEMORY.md

284 try:

285 if sub_workspace.sessions_dir.exists():

286 shutil.rmtree(sub_workspace.sessions_dir)

287 if sub_workspace.skills_dir.exists():

288 shutil.rmtree(sub_workspace.skills_dir)

289 if sub_workspace.memory_file.exists():

290 sub_workspace.memory_file.unlink()

291 logger.info(f"子 agent {sub_workspace.agent_id} 工作空间已清理（logs 保留）")

292 except Exception as e:

293 logger.warning(f"清理子 agent 工作空间失败: {e}")

294

295 return result