Coverage for src/chuck_data/command_output.py: 0%

162 statements  

« prev     ^ index     » next       coverage.py v7.8.0, created at 2025-06-05 22:56 -0700

1""" 

2Unified command output handling for both user and agent interfaces. 

3 

4This module provides consistent formatting of command results for different consumers, 

5such as the user TUI interface or the LLM agent. 

6""" 

7 

8from typing import Dict, Any 

9from rich.console import Console 

10 

11from .ui.table_formatter import display_table 

12 

13from .command_result import CommandResult 

14from .ui.theme import ( 

15 SUCCESS, 

16 WARNING, 

17 NEUTRAL, 

18 SUCCESS_STYLE, 

19 TABLE_TITLE_STYLE, 

20) 

21 

22 

23class OutputFormatter: 

24 """Format command results for different consumers (user TUI, agent, etc.)""" 

25 

26 @staticmethod 

27 def _display_status(data: Dict[str, Any], console: Console) -> None: 

28 """Display current status information including connection status and permissions.""" 

29 workspace_url = data.get("workspace_url", "Not set") 

30 active_catalog = data.get("active_catalog", "Not set") 

31 active_schema = data.get("active_schema", "Not set") 

32 active_model = data.get("active_model", "Not set") 

33 warehouse_id = data.get("warehouse_id", "Not set") 

34 connection_status = data.get("connection_status", "Unknown") 

35 

36 # Prepare settings for display 

37 status_items = [ 

38 {"setting": "Workspace URL", "value": workspace_url}, 

39 {"setting": "Active Catalog", "value": active_catalog}, 

40 {"setting": "Active Schema", "value": active_schema}, 

41 {"setting": "Active Model", "value": active_model}, 

42 {"setting": "Active Warehouse", "value": warehouse_id}, 

43 {"setting": "Connection Status", "value": connection_status}, 

44 ] 

45 

46 # Define styling functions 

47 def value_style(row): 

48 setting = row.get("setting", "") 

49 value = row.get("value", "") 

50 

51 # Special handling for connection status 

52 if setting == "Connection Status": 

53 if ( 

54 value == "Connected - token is valid" 

55 or value == "Connected (client present)." 

56 ): 

57 return "green" 

58 elif ( 

59 "Invalid" in value 

60 or "Not connected" in value 

61 or "error" in value.lower() 

62 ): 

63 return "red" 

64 else: 

65 return "yellow" 

66 # General styling for values 

67 elif value != "Not set": 

68 return "green" 

69 else: 

70 return "yellow" 

71 

72 # Display the status table 

73 display_table( 

74 console=console, 

75 data=status_items, 

76 columns=["setting", "value"], 

77 headers=["Setting", "Value"], 

78 title="Current Configuration", 

79 style_map={"value": value_style}, 

80 title_style=TABLE_TITLE_STYLE, 

81 show_lines=False, 

82 ) 

83 

84 # If permissions data is available, display it 

85 permissions_data = data.get("permissions") 

86 if permissions_data: 

87 # Display permissions information 

88 permissions_items = [] 

89 for key, value in permissions_data.items(): 

90 permissions_items.append({"permission": key, "status": value}) 

91 

92 # Define styling function for permission status 

93 def status_style(row): 

94 status = row.get("status", "") 

95 if status == "OK": 

96 return "green" 

97 elif status == "ERROR": 

98 return "red" 

99 else: 

100 return "yellow" 

101 

102 # Display the permissions table 

103 display_table( 

104 console=console, 

105 data=permissions_items, 

106 columns=["permission", "status"], 

107 headers=["Permission", "Status"], 

108 title="API Permissions", 

109 style_map={"status": status_style}, 

110 title_style=TABLE_TITLE_STYLE, 

111 show_lines=False, 

112 ) 

113 

114 @staticmethod 

115 def format_for_agent(result: CommandResult) -> Dict[str, Any]: 

116 """ 

117 Format a command result for agent consumption. 

118 

119 Args: 

120 result: Command result to format 

121 

122 Returns: 

123 Dictionary formatted for agent consumption 

124 """ 

125 if not result.success: 

126 return { 

127 "error": ( 

128 str(result.error) 

129 if result.error 

130 else result.message or "Unknown error" 

131 ) 

132 } 

133 

134 # Start with a base response 

135 response = {"success": True} 

136 

137 # Add the message if available 

138 if result.message: 

139 response["message"] = result.message 

140 

141 # Add displayed_to_user flag for tools that show output directly 

142 if result.data and isinstance(result.data, dict): 

143 # For commands that display directly to user, indicate this to the agent 

144 # This is a temporary approach - eventually the registry would indicate 

145 # whether a command displays directly to the user 

146 

147 # Copy all data from the result 

148 for key, value in result.data.items(): 

149 response[key] = value 

150 

151 # Add displayed_to_user flag for appropriate commands 

152 # Ideally this would come from the command registry 

153 response["displayed_to_user"] = True 

154 else: 

155 # If data isn't a dict, include it as-is 

156 response["data"] = result.data 

157 

158 return response 

159 

160 # --- Helper methods for specific command output formatting --- 

161 # These would contain the display logic currently in cli.py 

162 

163 @staticmethod 

164 def _display_catalogs(data: Dict[str, Any], console: Console) -> None: 

165 """Display catalogs in a nicely formatted way.""" 

166 catalogs = data.get("catalogs", []) 

167 current_catalog = data.get("current_catalog") 

168 

169 if not catalogs: 

170 console.print(f"[{WARNING}]No catalogs found.[/{WARNING}]") 

171 return 

172 

173 # Define a style map for conditional formatting 

174 def style_name(row): 

175 if row.get("name") == current_catalog: 

176 return f"[{SUCCESS_STYLE}]{row.get('name')}[/{SUCCESS_STYLE}]" 

177 return row.get("name") 

178 

179 style_map = { 

180 "name": style_name, 

181 } 

182 

183 # Display the catalogs table 

184 display_table( 

185 console=console, 

186 data=catalogs, 

187 columns=["name", "type", "comment"], 

188 headers=["Name", "Type", "Comment"], 

189 title="Available Catalogs", 

190 style_map=style_map, 

191 title_style=TABLE_TITLE_STYLE, 

192 show_lines=False, 

193 ) 

194 

195 # Display current catalog if set 

196 if current_catalog: 

197 console.print( 

198 f"\nCurrent catalog: [{SUCCESS_STYLE}]{current_catalog}[/{SUCCESS_STYLE}]" 

199 ) 

200 

201 @staticmethod 

202 def _display_schemas(data: Dict[str, Any], console: Console) -> None: 

203 """Display schemas in a nicely formatted way.""" 

204 schemas = data.get("schemas", []) 

205 catalog_name = data.get("catalog_name", "") 

206 current_schema = data.get("current_schema") 

207 

208 if not schemas: 

209 console.print( 

210 f"[{WARNING}]No schemas found in catalog '{catalog_name}'.[/{WARNING}]" 

211 ) 

212 return 

213 

214 # Define a style map for conditional formatting 

215 def style_name(row): 

216 if row.get("name") == current_schema: 

217 return f"[{SUCCESS_STYLE}]{row.get('name')}[/{SUCCESS_STYLE}]" 

218 return row.get("name") 

219 

220 style_map = { 

221 "name": style_name, 

222 } 

223 

224 # Display the schemas table 

225 display_table( 

226 console=console, 

227 data=schemas, 

228 columns=["name", "comment"], 

229 headers=["Name", "Comment"], 

230 title=f"Schemas in catalog '{catalog_name}'", 

231 style_map=style_map, 

232 title_style=TABLE_TITLE_STYLE, 

233 show_lines=False, 

234 ) 

235 

236 # Display current schema if set 

237 if current_schema: 

238 console.print( 

239 f"\nCurrent schema: [{SUCCESS_STYLE}]{current_schema}[/{SUCCESS_STYLE}]" 

240 ) 

241 

242 @staticmethod 

243 def _display_tables(data: Dict[str, Any], console: Console) -> None: 

244 """Display tables in a nicely formatted way.""" 

245 tables = data.get("tables", []) 

246 catalog_name = data.get("catalog_name", "") 

247 schema_name = data.get("schema_name", "") 

248 total_count = data.get("total_count", len(tables)) 

249 

250 if not tables: 

251 console.print( 

252 f"[{WARNING}]No tables found in {catalog_name}.{schema_name}[/{WARNING}]" 

253 ) 

254 return 

255 

256 # Define a style map for conditional formatting 

257 def format_date(row, col_name): 

258 date_str = row.get(col_name, "") 

259 if not date_str: 

260 return "" 

261 # Format could be improved based on actual date format in the data 

262 return date_str 

263 

264 style_map = { 

265 "created": format_date, 

266 "updated": format_date, 

267 } 

268 

269 # Display the tables 

270 display_table( 

271 console=console, 

272 data=tables, 

273 columns=["name", "table_type", "column_count", "created", "updated"], 

274 headers=["Table Name", "Type", "# Cols", "Created", "Last Updated"], 

275 title=f"Tables in {catalog_name}.{schema_name} ({total_count} total)", 

276 style_map=style_map, 

277 title_style=TABLE_TITLE_STYLE, 

278 show_lines=True, 

279 ) 

280 

281 @staticmethod 

282 def _display_models(data: Dict[str, Any], console: Console) -> None: 

283 """Display models in a nicely formatted way.""" 

284 models = data.get("models", []) 

285 current_model = data.get("current_model") 

286 is_detailed = data.get("is_detailed", False) 

287 

288 if not models: 

289 console.print(f"[{WARNING}]No models found.[/{WARNING}]") 

290 return 

291 

292 # Define a style map for conditional formatting 

293 def style_name(row): 

294 model_name = row.get("name") 

295 # Check if this is a recommended model 

296 is_recommended = model_name in [ 

297 "databricks-meta-llama-3-3-70b-instruct", 

298 "databricks-claude-3-7-sonnet", 

299 ] 

300 

301 if model_name == current_model: 

302 if is_recommended: 

303 return f"[bold green]{model_name} (recommended)[/bold green]" 

304 return f"[bold green]{model_name}[/bold green]" 

305 elif is_recommended: 

306 return f"{model_name} [green](recommended)[/green]" 

307 return model_name 

308 

309 style_map = { 

310 "name": style_name, 

311 } 

312 

313 if is_detailed: 

314 # Display detailed models with more columns 

315 display_table( 

316 console=console, 

317 data=models, 

318 columns=["name", "creator", "created", "status", "description"], 

319 headers=["Model Name", "Created By", "Date", "Status", "Description"], 

320 title="Available Models", 

321 style_map=style_map, 

322 title_style=TABLE_TITLE_STYLE, 

323 show_lines=True, 

324 ) 

325 else: 

326 # Display simple models list 

327 display_table( 

328 console=console, 

329 data=models, 

330 columns=["name", "created", "status"], 

331 headers=["Model Name", "Created", "Status"], 

332 title="Available Models", 

333 style_map=style_map, 

334 title_style=TABLE_TITLE_STYLE, 

335 show_lines=False, 

336 ) 

337 

338 # Display current model if set 

339 if current_model: 

340 console.print( 

341 f"\nCurrent model: [{SUCCESS_STYLE}]{current_model}[/{SUCCESS_STYLE}]" 

342 ) 

343 

344 @staticmethod 

345 def _display_warehouses(data: Dict[str, Any], console: Console) -> None: 

346 """Display warehouses in a nicely formatted way.""" 

347 warehouses = data.get("warehouses", []) 

348 current_warehouse = data.get("current_warehouse") 

349 

350 if not warehouses: 

351 console.print(f"[{WARNING}]No warehouses found.[/{WARNING}]") 

352 return 

353 

354 # Define a style map for conditional formatting 

355 def style_name(row): 

356 if row.get("name") == current_warehouse: 

357 return f"[{SUCCESS_STYLE}]{row.get('name')}[/{SUCCESS_STYLE}]" 

358 return row.get("name") 

359 

360 def style_state(row): 

361 state = row.get("state", "").lower() 

362 if state == "running": 

363 return f"[{SUCCESS}]{state}[/{SUCCESS}]" 

364 elif state == "stopped": 

365 return f"[{NEUTRAL}]{state}[/{NEUTRAL}]" 

366 else: 

367 return f"[{WARNING}]{state}[/{WARNING}]" 

368 

369 style_map = { 

370 "name": style_name, 

371 "state": style_state, 

372 } 

373 

374 # Display the warehouses 

375 display_table( 

376 console=console, 

377 data=warehouses, 

378 columns=["name", "size", "state", "auto_stop_mins", "created_by"], 

379 headers=["Warehouse Name", "Size", "State", "Auto Stop", "Created By"], 

380 title="Available Warehouses", 

381 style_map=style_map, 

382 title_style=TABLE_TITLE_STYLE, 

383 show_lines=False, 

384 ) 

385 

386 # Display current warehouse if set 

387 if current_warehouse: 

388 console.print( 

389 f"\nCurrent warehouse: [{SUCCESS_STYLE}]{current_warehouse}[/{SUCCESS_STYLE}]" 

390 ) 

391 

392 @staticmethod 

393 def _display_volumes(data: Dict[str, Any], console: Console) -> None: 

394 """Display volumes in a nicely formatted way.""" 

395 volumes = data.get("volumes", []) 

396 current_volume = data.get("current_volume") 

397 

398 if not volumes: 

399 console.print(f"[{WARNING}]No volumes found.[/{WARNING}]") 

400 return 

401 

402 # Define a style map for conditional formatting 

403 def style_name(row): 

404 if row.get("name") == current_volume: 

405 return f"[{SUCCESS_STYLE}]{row.get('name')}[/{SUCCESS_STYLE}]" 

406 return row.get("name") 

407 

408 style_map = { 

409 "name": style_name, 

410 } 

411 

412 # Display the volumes 

413 display_table( 

414 console=console, 

415 data=volumes, 

416 columns=["name", "type", "catalog", "schema", "owner", "created"], 

417 headers=["Volume Name", "Type", "Catalog", "Schema", "Owner", "Created"], 

418 title="Available Volumes", 

419 style_map=style_map, 

420 title_style=TABLE_TITLE_STYLE, 

421 show_lines=False, 

422 ) 

423 

424 # Display current volume if set 

425 if current_volume: 

426 console.print( 

427 f"\nCurrent volume: [{SUCCESS_STYLE}]{current_volume}[/{SUCCESS_STYLE}]" 

428 )