Coverage for aipyapp/plugins/p_style_classic.py: 0%

285 statements  

« prev     ^ index     » next       coverage.py v7.10.3, created at 2025-08-11 12:45 +0200

1#!/usr/bin/env python 

2# -*- coding: utf-8 -*- 

3 

4import re 

5import sys 

6from functools import wraps 

7import json 

8 

9from rich.table import Table 

10from rich.syntax import Syntax 

11from rich.markdown import Markdown 

12from rich.tree import Tree 

13from rich.text import Text 

14from rich.console import Console 

15 

16from aipyapp.display import RichDisplayPlugin 

17from live_display import LiveDisplay 

18from aipyapp import T 

19 

20def restore_output(func): 

21 @wraps(func) 

22 def wrapper(self, *args, **kwargs): 

23 old_stdout, old_stderr = sys.stdout, sys.stderr 

24 sys.stdout, sys.stderr = sys.__stdout__, sys.__stderr__ 

25 

26 try: 

27 return func(self, *args, **kwargs) 

28 finally: 

29 sys.stdout, sys.stderr = old_stdout, old_stderr 

30 return wrapper 

31 

32class DisplayClassic(RichDisplayPlugin): 

33 """Classic display style""" 

34 name = "classic" 

35 version = "1.0.0" 

36 description = "Classic display style" 

37 author = "AiPy Team" 

38 

39 def __init__(self, console: Console, quiet: bool = False): 

40 super().__init__(console, quiet) 

41 self.live_display = None 

42 

43 def _get_title(self, title: str, *args, style: str = "info", prefix: str = "\n"): 

44 text = Text(f"{prefix}{title}".format(*args), style=style) 

45 text.highlight_words(args, style="bold white") 

46 return text 

47 

48 def on_exception(self, event): 

49 """异常事件处理""" 

50 msg = event.data.get('msg', '') 

51 exception = event.data.get('exception') 

52 title = self._get_title(T("Exception occurred"), msg, style="error") 

53 tree = Tree(title) 

54 tree.add(exception) 

55 self.console.print(tree) 

56 

57 def on_task_start(self, event): 

58 """任务开始事件处理""" 

59 data = event.data 

60 instruction = data.get('instruction') 

61 title = data.get('title') 

62 if not title: 

63 title = instruction 

64 tree = Tree(f"🚀 {T('Task processing started')}") 

65 tree.add(title) 

66 self.console.print(tree) 

67 

68 def on_query_start(self, event): 

69 """查询开始事件处理""" 

70 data = event.data 

71 llm = data.get('llm', '') 

72 title = self._get_title(T("Sending message to {}"), llm) 

73 self.console.print(title) 

74 

75 def on_round_start(self, event): 

76 """回合开始事件处理""" 

77 data = event.data 

78 instruction = data.get('instruction') 

79 title = data.get('title') 

80 if not title: 

81 title = instruction 

82 prompt = self._get_title(T("Instruction processing started")) 

83 tree = Tree(prompt) 

84 tree.add(title) 

85 self.console.print(tree) 

86 

87 def on_stream_start(self, event): 

88 """流式开始事件处理""" 

89 if not self.quiet: 

90 self.live_display = LiveDisplay() 

91 self.live_display.__enter__() 

92 title = self._get_title(T("Streaming started"), prefix="") 

93 self.console.print(title) 

94 

95 def on_stream_end(self, event): 

96 """流式结束事件处理""" 

97 if self.live_display: 

98 self.live_display.__exit__(None, None, None) 

99 self.live_display = None 

100 

101 def on_stream(self, event): 

102 """LLM 流式响应事件处理""" 

103 response = event.data 

104 lines = response.get('lines') 

105 reason = response.get('reason', False) 

106 if self.live_display: 

107 self.live_display.update_display(lines, reason=reason) 

108 

109 @staticmethod 

110 def convert_front_matter(md_text: str) -> str: 

111 pattern = r"^---\s*\n(.*?)\n---\s*\n" 

112 #return re.sub(pattern, r"```yaml\n\1\n```\n", md_text, flags=re.DOTALL) 

113 return re.sub(pattern, "", md_text, flags=re.DOTALL) 

114 

115 def on_response_complete(self, event): 

116 """LLM 响应完成事件处理""" 

117 data = event.data 

118 llm = data.get('llm', '') 

119 msg = data.get('msg') 

120 if not msg: 

121 title = self._get_title(T("LLM response is empty"), style="error") 

122 self.console.print(title) 

123 return 

124 

125 if msg.role == 'error': 

126 title = self._get_title(T("Failed to receive message"), style="error") 

127 tree = Tree(title) 

128 tree.add(msg.content) 

129 self.console.print(tree) 

130 return 

131 

132 content = self.convert_front_matter(msg.content) 

133 if msg.reason: 

134 content = f"{msg.reason}\n\n-----\n\n{content}" 

135 title = self._get_title(f"{T('Completed receiving message')} ({llm})", style="success") 

136 tree = Tree(title) 

137 tree.add(Markdown(content)) 

138 self.console.print(tree) 

139 

140 def on_task_status(self, event): 

141 """任务状态事件处理""" 

142 status = event.data.get('status') 

143 completed = status.get('completed', False) 

144 style = "success" if completed else "error" 

145 title = self._get_title(T("Task status"), style=style) 

146 tree = Tree(title, guide_style=style) 

147 if completed: 

148 tree.add(T("Completed")) 

149 tree.add(T("Confidence level: {}", status.get('confidence', 0))) 

150 else: 

151 tree.add(T("Failed")) 

152 tree.add(T("Reason: {}", status.get('reason', ''))) 

153 tree.add(T("Suggestion: {}", status.get('suggestion', ''))) 

154 self.console.print(tree) 

155 

156 def on_parse_reply(self, event): 

157 """消息解析结果事件处理""" 

158 ret = event.data.get('result') 

159 if not ret: 

160 return 

161 

162 title = self._get_title(T("Message parse result")) 

163 tree = Tree(title) 

164 

165 if 'blocks' in ret and ret['blocks']: 

166 block_count = len(ret['blocks']) 

167 tree.add(f"{block_count} {T('code blocks')}") 

168 

169 if 'commands' in ret and ret['commands']: 

170 commands = ret['commands'] 

171 # 分别统计和显示不同类型的指令 

172 exec_commands = [cmd for cmd in commands if cmd['type'] == 'exec'] 

173 edit_commands = [cmd for cmd in commands if cmd['type'] == 'edit'] 

174 

175 if exec_commands: 

176 exec_names = [cmd.get('block_name', 'Unknown') for cmd in exec_commands] 

177 exec_str = ", ".join(exec_names[:3]) 

178 if len(exec_names) > 3: 

179 exec_str += f" (+{len(exec_names)-3} more)" 

180 tree.add(f"{T('Execution')}: {exec_str}") 

181 

182 if edit_commands: 

183 edit_names = [cmd['instruction']['name'] for cmd in edit_commands if 'instruction' in cmd] 

184 edit_str = ", ".join(edit_names[:3]) 

185 if len(edit_names) > 3: 

186 edit_str += f" (+{len(edit_names)-3} more)" 

187 tree.add(f"{T('Edit')}: {edit_str}") 

188 

189 if 'call_tool' in ret: 

190 tree.add(T("MCP tool call")) 

191 

192 if 'errors' in ret and ret['errors']: 

193 error_count = len(ret['errors']) 

194 tree.add(f"{error_count} {T('errors')}") 

195 

196 self.console.print(tree) 

197 

198 def on_exec(self, event): 

199 """代码执行开始事件处理""" 

200 block = event.data.get('block') 

201 title = self._get_title(T("Start executing code block {}"), block.name) 

202 self.console.print(title) 

203 

204 def on_edit_start(self, event): 

205 """代码编辑开始事件处理""" 

206 instruction = event.data.get('instruction', {}) 

207 block_name = instruction.get('name', 'Unknown') 

208 old_str = instruction.get('old', '') 

209 new_str = instruction.get('new', '') 

210 

211 title = self._get_title(T("Start editing code block {}"), block_name, style="warning") 

212 tree = Tree(title) 

213 

214 if old_str: 

215 old_preview = old_str[:50] + '...' if len(old_str) > 50 else old_str 

216 tree.add(f"{T('Replace')}: {repr(old_preview)}") 

217 if new_str: 

218 new_preview = new_str[:50] + '...' if len(new_str) > 50 else new_str 

219 tree.add(f"{T('With')}: {repr(new_preview)}") 

220 

221 self.console.print(tree) 

222 

223 def on_edit_result(self, event): 

224 """代码编辑结果事件处理""" 

225 data = event.data 

226 result = data.get('result', {}) 

227 

228 success = result.get('success', False) 

229 message = result.get('message', '') 

230 block_name = result.get('block_name', 'Unknown') 

231 new_version = result.get('new_version') 

232 

233 if success: 

234 style = "success" 

235 title = self._get_title(T("Edit completed {}"), block_name, style=style) 

236 tree = Tree(title) 

237 

238 if message: 

239 tree.add(message) 

240 if new_version: 

241 tree.add(f"{T('New version')}: v{new_version}") 

242 else: 

243 style = "error" 

244 title = self._get_title(T("Edit failed {}"), block_name, style=style) 

245 tree = Tree(title) 

246 tree.add(message or T("Edit operation failed")) 

247 

248 self.console.print(tree) 

249 

250 @restore_output 

251 def on_call_function(self, event): 

252 """函数调用事件处理""" 

253 data = event.data 

254 funcname = data.get('funcname') 

255 title = self._get_title(T("Start calling function {}"), funcname) 

256 self.console.print(title) 

257 

258 @restore_output 

259 def on_call_function_result(self, event): 

260 """函数调用结果事件处理""" 

261 data = event.data 

262 funcname = data.get('funcname') 

263 success = data.get('success', False) 

264 result = data.get('result') 

265 error = data.get('error') 

266 

267 if success: 

268 style = "success" 

269 title = self._get_title(T("Function call result {}"), funcname, style=style) 

270 tree = Tree(title) 

271 if result is not None: 

272 # 格式化并显示结果 

273 if isinstance(result, (dict, list)): 

274 json_result = json.dumps(result, ensure_ascii=False, indent=2, default=str) 

275 tree.add(Syntax(json_result, "json", word_wrap=True, line_range=(0, 10))) 

276 else: 

277 tree.add(str(result)) 

278 else: 

279 tree.add(T("No return value")) 

280 self.console.print(tree) 

281 else: 

282 style = "error" 

283 title = self._get_title(T("Function call failed {}"), funcname, style=style) 

284 tree = Tree(title) 

285 tree.add(error if error else T("Unknown error")) 

286 self.console.print(tree) 

287 

288 def on_exec_result(self, event): 

289 """代码执行结果事件处理""" 

290 data = event.data 

291 result = data.get('result') 

292 block = data.get('block') 

293 

294 try: 

295 success = result['__state__']['success'] 

296 style = "success" if success else "error" 

297 except: 

298 style = "warning" 

299 

300 # 显示说明信息 

301 block_name = getattr(block, 'name', 'Unknown') if block else 'Unknown' 

302 title = self._get_title(T("Execution result {}"), block_name, style=style) 

303 tree = Tree(title) 

304 

305 # JSON格式化和高亮显示结果 

306 json_result = json.dumps(result, ensure_ascii=False, indent=2, default=str) 

307 tree.add(Syntax(json_result, "json", word_wrap=True)) 

308 self.console.print(tree) 

309 

310 def on_mcp_call(self, event): 

311 """工具调用事件处理""" 

312 title = self._get_title(T("Start calling MCP tool")) 

313 self.console.print(title) 

314 

315 def on_mcp_result(self, event): 

316 """MCP 工具调用结果事件处理""" 

317 data = event.data 

318 result = data.get('result') 

319 block = data.get('block') 

320 title = self._get_title(T("MCP tool call result {}"), block.name) 

321 self.console.print(title) 

322 json_result = json.dumps(result, ensure_ascii=False, indent=2, default=str) 

323 self.console.print_json(json_result, style="dim") 

324 

325 def on_round_end(self, event): 

326 """任务总结事件处理""" 

327 summary = event.data['summary'] 

328 usages = summary.get('usages', []) 

329 if usages: 

330 table = Table(title=T("Task Summary"), show_lines=True) 

331 

332 table.add_column(T("Round"), justify="center", style="bold cyan", no_wrap=True) 

333 table.add_column(T("Time(s)"), justify="right") 

334 table.add_column(T("In Tokens"), justify="right") 

335 table.add_column(T("Out Tokens"), justify="right") 

336 table.add_column(T("Total Tokens"), justify="right", style="bold magenta") 

337 

338 round = 1 

339 for row in usages: 

340 table.add_row( 

341 str(round), 

342 str(row["time"]), 

343 str(row["input_tokens"]), 

344 str(row["output_tokens"]), 

345 str(row["total_tokens"]), 

346 ) 

347 round += 1 

348 self.console.print("\n") 

349 self.console.print(table) 

350 

351 summary = summary.get('summary') 

352 title = self._get_title(T("End processing instruction")) 

353 tree = Tree(title) 

354 tree.add(f"{T('Summary')}: {summary}") 

355 self.console.print(tree) 

356 

357 def on_upload_result(self, event): 

358 """云端上传结果事件处理""" 

359 data = event.data 

360 status_code = data.get('status_code', 0) 

361 url = data.get('url', '') 

362 if url: 

363 self.console.print(f"🟢 {T('Article uploaded successfully, {}', url)}", style="success") 

364 else: 

365 self.console.print(f"🔴 {T('Upload failed (status code: {})', status_code)}", style="error") 

366 

367 def on_task_end(self, event): 

368 """任务结束事件处理""" 

369 path = event.data.get('path', '') 

370 title = self._get_title(T("Task completed")) 

371 tree = Tree(title) 

372 tree.add(path) 

373 self.console.print(tree) 

374 

375 def on_runtime_message(self, event): 

376 """Runtime消息事件处理""" 

377 data = event.data 

378 message = data.get('message', '') 

379 status = data.get('status', 'info') 

380 title = self._get_title(message, style=status) 

381 self.console.print(title) 

382 

383 def on_runtime_input(self, event): 

384 """Runtime输入事件处理""" 

385 # 输入事件通常不需要特殊处理,因为input_prompt已经处理了 

386 pass