Coverage for pydantic_ai_jupyter / views.py: 91%

1"""View classes for rendering pydantic-ai events in Jupyter notebooks.""" 

2 

3from __future__ import annotations 

4 

5import json 

6import traceback 

7from html import escape 

8from typing import Any 

9 

10from pydantic import Field 

11from pydantic_ai.messages import ( 

12 FinalResultEvent, 

13 PartEndEvent, 

14 PartStartEvent, 

15 RetryPromptPart, 

16 TextPart, 

17 ThinkingPart, 

18 ToolCallPart, 

19 ToolReturnPart, 

20) 

21 

22from .models import View 

23 

24 

25class ToolCallView(View): 

26 """Renders a tool call with its arguments.""" 

27 

28 tool_name: str 

29 args: str | dict[str, Any] | None = None 

30 tool_call_id: str | None = None 

31 

32 @classmethod 

33 def from_part(cls, part: ToolCallPart) -> ToolCallView: 

34 return cls( 

35 tool_name=part.tool_name, args=part.args, tool_call_id=part.tool_call_id 

36 ) 

37 

38 def render(self) -> str: 

39 args_str = ( 

40 json.dumps(self.args, indent=2) 

41 if isinstance(self.args, dict) 

42 else str(self.args or "{}") 

43 ) 

44 return f""" 

45 <div style="border-left: 3px solid #3b82f6; padding: 8px 12px; margin: 8px 0; background: #eff6ff; border-radius: 4px;"> 

46 <div style="font-weight: 600; color: #1d4ed8; margin-bottom: 4px;"> 

47 🔧 <code style="background: #dbeafe; padding: 2px 6px; border-radius: 3px;">{escape(self.tool_name)}</code> 

48 </div> 

49 <pre style="margin: 0; font-size: 12px; background: #f8fafc; padding: 8px; border-radius: 3px; overflow-x: auto;">{escape(args_str)}</pre> 

50 </div> 

51 """ 

52 

53 def __repr__(self) -> str: 

54 args_str = ( 

55 json.dumps(self.args) 

56 if isinstance(self.args, dict) 

57 else str(self.args or "{}") 

58 ) 

59 if len(args_str) > 200: 

60 args_str = args_str[:200] + "..." 

61 return f"🔧 {self.tool_name}({args_str})" 

62 

63 

64class ToolResultView(View): 

65 """Renders a tool result (success or retry).""" 

66 

67 tool_name: str 

68 content: Any 

69 tool_call_id: str | None = None 

70 is_retry: bool = False 

71 max_length: int = 500 

72 

73 @classmethod 

74 def from_part(cls, part: ToolReturnPart | RetryPromptPart) -> ToolResultView: 

75 return cls( 

76 tool_name=part.tool_name or "unknown", 

77 content=part.content if hasattr(part, "content") else str(part), 

78 tool_call_id=part.tool_call_id, 

79 is_retry=isinstance(part, RetryPromptPart), 

80 ) 

81 

82 def render(self) -> str: 

83 content_str = ( 

84 self.content 

85 if isinstance(self.content, str) 

86 else json.dumps(self.content, indent=2) 

87 ) 

88 if len(content_str) > self.max_length: 

89 content_str = content_str[: self.max_length] + "..." 

90 

91 if self.is_retry: 

92 return f""" 

93 <div style="border-left: 3px solid #f59e0b; padding: 8px 12px; margin: 8px 0; background: #fffbeb; border-radius: 4px;"> 

94 <div style="font-weight: 600; color: #b45309; margin-bottom: 4px;"> 

95 🔄 Retry: <code style="background: #fef3c7; padding: 2px 6px; border-radius: 3px;">{escape(self.tool_name)}</code> 

96 </div> 

97 <pre style="margin: 0; font-size: 12px; background: #fffdf5; padding: 8px; border-radius: 3px; overflow-x: auto; white-space: pre-wrap;">{escape(content_str)}</pre> 

98 </div> 

99 """ 

100 else: 

101 return f""" 

102 <div style="border-left: 3px solid #10b981; padding: 8px 12px; margin: 8px 0; background: #ecfdf5; border-radius: 4px;"> 

103 <div style="font-weight: 600; color: #047857; margin-bottom: 4px;"> 

104 ✅ Result: <code style="background: #d1fae5; padding: 2px 6px; border-radius: 3px;">{escape(self.tool_name)}</code> 

105 </div> 

106 <pre style="margin: 0; font-size: 12px; background: #f0fdf4; padding: 8px; border-radius: 3px; overflow-x: auto; white-space: pre-wrap;">{escape(content_str)}</pre> 

107 </div> 

108 """ 

109 

110 def __repr__(self) -> str: 

111 content_str = ( 

112 self.content 

113 if isinstance(self.content, str) 

114 else json.dumps(self.content, indent=2) 

115 ) 

116 return f"✅ {self.tool_name} → {content_str}" 

117 

118 

119class ErrorView(View): 

120 """Renders an exception with optional traceback.""" 

121 

122 error_type: str 

123 message: str 

124 details: str | None = None 

125 

126 @classmethod 

127 def from_exception(cls, exc: Exception) -> ErrorView: 

128 return cls( 

129 error_type=type(exc).__name__, 

130 message=str(exc), 

131 details=traceback.format_exc(), 

132 ) 

133 

134 def render(self) -> str: 

135 details_html = "" 

136 if self.details: 

137 details = self.details 

138 if len(details) > 1000: 

139 details = details[-1000:] 

140 escaped = escape(details) 

141 

142 details_html = f"""<details style="margin-top: 8px;"> 

143 <summary style="cursor: pointer; color: #991b1b;">Show traceback</summary> 

144 <pre style="margin: 4px 0 0 0; font-size: 11px; background: #fef2f2; padding: 8px; border-radius: 3px; overflow-x: auto; white-space: pre-wrap;">{escaped}</pre> 

145 </details>""" 

146 

147 return f""" 

148 <div style="border-left: 3px solid #dc2626; padding: 8px 12px; margin: 8px 0; background: #fef2f2; border-radius: 4px;"> 

149 <div style="font-weight: 600; color: #dc2626; margin-bottom: 4px;"> 

150 ❌ <code style="background: #fee2e2; padding: 2px 6px; border-radius: 3px;">{escape(self.error_type)}</code> 

151 </div> 

152 <div style="font-size: 13px; color: #7f1d1d;">{escape(self.message)}</div> 

153 {details_html} 

154 </div> 

155 """ 

156 

157 def __repr__(self) -> str: 

158 return f"❌ {self.error_type}: {self.message}" 

159 

160 

161class ThinkingView(View): 

162 """A live-updating view for model thinking/reasoning content.""" 

163 

164 content: str = "" 

165 

166 def render(self) -> str: 

167 escaped = escape(self.content) 

168 

169 return f""" 

170 <details style="margin: 8px 0;" open> 

171 <summary style="cursor: pointer; font-weight: 600; color: #6b7280; font-size: 12px;"> 

172 💭 Thinking... 

173 </summary> 

174 <div style="border-left: 3px solid #9ca3af; padding: 8px 12px; margin: 4px 0; background: #f9fafb; border-radius: 4px;"> 

175 <pre style="margin: 0; font-size: 12px; color: #4b5563; white-space: pre-wrap; font-family: inherit;">{escaped}</pre> 

176 </div> 

177 </details> 

178 """ 

179 

180 def append(self, text: str) -> None: 

181 """Append text and update the display.""" 

182 self.content += text 

183 self.update() 

184 

185 def __repr__(self) -> str: 

186 preview = ( 

187 self.content[:100] + "..." if len(self.content) > 100 else self.content 

188 ) 

189 return f"💭 Thinking: {preview}" 

190 

191 

192class DebugEventView(View): 

193 """Renders debug/lifecycle events in a muted style.""" 

194 

195 event_type: str 

196 summary: str 

197 details: str | None = None 

198 

199 def __repr__(self) -> str: 

200 return f"⚙️ {self.event_type}: {self.summary}" 

201 

202 @classmethod 

203 def from_event(cls, event: Any) -> DebugEventView: 

204 """Create a debug view from any event.""" 

205 event_type = type(event).__name__ 

206 

207 if isinstance(event, PartEndEvent): 

208 part = event.part 

209 part_type = type(part).__name__ 

210 if isinstance(part, TextPart): 

211 preview = ( 

212 part.content[:50] + "..." 

213 if len(part.content) > 50 

214 else part.content 

215 ) 

216 summary = f"{part_type} completed" 

217 details = f"Content: {preview}" 

218 elif isinstance(part, ThinkingPart): 

219 summary = f"{part_type} completed" 

220 details = None 

221 elif isinstance(part, ToolCallPart): 

222 summary = f"{part_type} completed: {part.tool_name}" 

223 details = None 

224 else: 

225 summary = f"{part_type} completed" 

226 details = None 

227 next_kind = getattr(event, "next_part_kind", None) 

228 if next_kind: 

229 summary += f" → {next_kind}" 

230 elif isinstance(event, PartStartEvent): 

231 part = event.part 

232 part_type = type(part).__name__ 

233 if isinstance(part, ToolCallPart): 

234 summary = f"{part_type}: {part.tool_name}" 

235 details = ( 

236 part.args if isinstance(part.args, str) else json.dumps(part.args) 

237 ) 

238 else: 

239 summary = f"{part_type} starting" 

240 details = None 

241 elif isinstance(event, FinalResultEvent): 

242 summary = "Agent completed" 

243 if event.tool_name: 

244 details = f"via tool: {event.tool_name}" 

245 else: 

246 details = "text response" 

247 else: 

248 summary = str(event)[:100] 

249 details = None 

250 

251 return cls(event_type=event_type, summary=summary, details=details) 

252 

253 def render(self) -> str: 

254 details_html = "" 

255 if self.details: 

256 details = self.details 

257 if len(details) > 100: 

258 details = details[:100] + "..." 

259 escaped = escape(details) 

260 details_html = f'<span style="color: #9ca3af;"> · {escaped}</span>' 

261 

262 return f""" 

263 <div style="padding: 2px 8px; margin: 2px 0; font-size: 11px; color: #6b7280; font-family: monospace;"> 

264 ⚙️ <span style="color: #9ca3af;">{self.event_type}</span>: {escape(self.summary)}{details_html} 

265 </div> 

266 """ 

267 

268 

269class StreamingToolCallView(View): 

270 """A live-updating view for tool call arguments as they stream in.""" 

271 

272 tool_name: str = Field(default="", description="The name of the tool being called.") 

273 args: str = Field(default="", description="The arguments being passed to the tool.") 

274 tool_call_id: str | None = Field( 

275 default=None, description="The ID of the tool call." 

276 ) 

277 

278 def render(self) -> str: 

279 args_escaped = escape(self.args) 

280 cursor = ( 

281 '<span style="animation: blink 1s infinite;">▊</span>' 

282 if not args_escaped.endswith("}") 

283 else "" 

284 ) 

285 return f""" 

286 <style>@keyframes blink {{ 50% {{ opacity: 0; }} }}</style> 

287 <div style="border-left: 3px solid #3b82f6; padding: 8px 12px; margin: 8px 0; background: #eff6ff; border-radius: 4px;"> 

288 <div style="font-weight: 600; color: #1d4ed8; margin-bottom: 4px;"> 

289 🔧 <code style="background: #dbeafe; padding: 2px 6px; border-radius: 3px;">{escape(self.tool_name)}</code> 

290 </div> 

291 <pre style="margin: 0; font-size: 12px; background: #f8fafc; padding: 8px; border-radius: 3px; overflow-x: auto;">{args_escaped}{cursor}</pre> 

292 </div> 

293 """ 

294 

295 def append_args(self, delta: str) -> None: 

296 """Append to args and update the display.""" 

297 self.args += delta 

298 self.update() 

299 

300 def append_tool_name(self, delta: str) -> None: 

301 """Append to tool name and update the display.""" 

302 self.tool_name += delta 

303 self.update() 

304 

305 def __repr__(self) -> str: 

306 return f"🔧 {self.tool_name}({self.args})"