Coverage for agent.py: 17%

1import json 

2from rich.console import Console 

3from rich.panel import Panel 

4from rich.markdown import Markdown 

5from qrclaw.config import MAX_ITERATIONS, COMPRESS_THRESHOLD 

6from qrclaw.providers import provider 

7from qrclaw.providers.base import LLMResponse 

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

9from qrclaw.memory.session import Session 

10from qrclaw.memory import compressor, LongTermMemory 

11from qrclaw.prompt import build_system_prompt 

12from qrclaw.cli.display import show_plan_progress 

13from qrclaw.workspace import Workspace 

14from qrclaw.logger import get_logger 

15 

16logger = get_logger("qrclaw.agent") 

17 

18 

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

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

21import threading 

22_thread_local = threading.local() 

23 

24def set_session(session: Session): 

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

26 _thread_local.session = session 

27 

28def get_session() -> Session | None: 

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

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

31 

32def set_workspace(workspace: Workspace): 

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

34 _thread_local.workspace = workspace 

35 

36def get_workspace() -> Workspace | None: 

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

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

39 

40def set_agent_depth(depth: int): 

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

42 _thread_local.agent_depth = depth 

43 

44def get_agent_depth() -> int: 

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

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

47 

48def is_sub_agent() -> bool: 

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

50 return get_agent_depth() > 0 

51 

52 

53def _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 

72def 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 

245def 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