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

1#!/usr/bin/env python

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

3

4 import re

5 import sys

6 from functools import wraps

7 import json

8

9 from rich.table import Table

10 from rich.syntax import Syntax

11 from rich.markdown import Markdown

12 from rich.tree import Tree

13 from rich.text import Text

14 from rich.console import Console

15

16 from aipyapp.display import RichDisplayPlugin

17 from live_display import LiveDisplay

18 from aipyapp import T

19

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

32 class 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