Coverage for src/command_output.py: 0%
162 statements
« prev ^ index » next coverage.py v7.8.0, created at 2025-06-05 22:56 -0700
« 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.
4This module provides consistent formatting of command results for different consumers,
5such as the user TUI interface or the LLM agent.
6"""
8from typing import Dict, Any
9from rich.console import Console
11from src.ui.table_formatter import display_table
13from src.command_result import CommandResult
14from src.ui.theme import (
15 SUCCESS,
16 WARNING,
17 NEUTRAL,
18 SUCCESS_STYLE,
19 TABLE_TITLE_STYLE,
20)
23class OutputFormatter:
24 """Format command results for different consumers (user TUI, agent, etc.)"""
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")
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 ]
46 # Define styling functions
47 def value_style(row):
48 setting = row.get("setting", "")
49 value = row.get("value", "")
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"
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 )
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})
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"
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 )
114 @staticmethod
115 def format_for_agent(result: CommandResult) -> Dict[str, Any]:
116 """
117 Format a command result for agent consumption.
119 Args:
120 result: Command result to format
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 }
134 # Start with a base response
135 response = {"success": True}
137 # Add the message if available
138 if result.message:
139 response["message"] = result.message
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
147 # Copy all data from the result
148 for key, value in result.data.items():
149 response[key] = value
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
158 return response
160 # --- Helper methods for specific command output formatting ---
161 # These would contain the display logic currently in cli.py
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")
169 if not catalogs:
170 console.print(f"[{WARNING}]No catalogs found.[/{WARNING}]")
171 return
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")
179 style_map = {
180 "name": style_name,
181 }
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 )
195 # Display current catalog if set
196 if current_catalog:
197 console.print(
198 f"\nCurrent catalog: [{SUCCESS_STYLE}]{current_catalog}[/{SUCCESS_STYLE}]"
199 )
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")
208 if not schemas:
209 console.print(
210 f"[{WARNING}]No schemas found in catalog '{catalog_name}'.[/{WARNING}]"
211 )
212 return
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")
220 style_map = {
221 "name": style_name,
222 }
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 )
236 # Display current schema if set
237 if current_schema:
238 console.print(
239 f"\nCurrent schema: [{SUCCESS_STYLE}]{current_schema}[/{SUCCESS_STYLE}]"
240 )
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))
250 if not tables:
251 console.print(
252 f"[{WARNING}]No tables found in {catalog_name}.{schema_name}[/{WARNING}]"
253 )
254 return
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
264 style_map = {
265 "created": format_date,
266 "updated": format_date,
267 }
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 )
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)
288 if not models:
289 console.print(f"[{WARNING}]No models found.[/{WARNING}]")
290 return
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 ]
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
309 style_map = {
310 "name": style_name,
311 }
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 )
338 # Display current model if set
339 if current_model:
340 console.print(
341 f"\nCurrent model: [{SUCCESS_STYLE}]{current_model}[/{SUCCESS_STYLE}]"
342 )
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")
350 if not warehouses:
351 console.print(f"[{WARNING}]No warehouses found.[/{WARNING}]")
352 return
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")
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}]"
369 style_map = {
370 "name": style_name,
371 "state": style_state,
372 }
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 )
386 # Display current warehouse if set
387 if current_warehouse:
388 console.print(
389 f"\nCurrent warehouse: [{SUCCESS_STYLE}]{current_warehouse}[/{SUCCESS_STYLE}]"
390 )
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")
398 if not volumes:
399 console.print(f"[{WARNING}]No volumes found.[/{WARNING}]")
400 return
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")
408 style_map = {
409 "name": style_name,
410 }
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 )
424 # Display current volume if set
425 if current_volume:
426 console.print(
427 f"\nCurrent volume: [{SUCCESS_STYLE}]{current_volume}[/{SUCCESS_STYLE}]"
428 )