Coverage for src/service.py: 47%
205 statements
« prev ^ index » next coverage.py v7.8.0, created at 2025-06-05 23:16 -0700
« prev ^ index » next coverage.py v7.8.0, created at 2025-06-05 23:16 -0700
1"""
2Service layer for Chuck application.
3Provides a facade for all business operations needed by the UI.
4"""
6import json
7import logging
8import jsonschema
9import traceback
10from typing import Dict, Optional, Any, Tuple
12from src.clients.databricks import DatabricksAPIClient
13from src.commands.base import CommandResult
14from src.command_registry import get_command, CommandDefinition
15from src.config import (
16 get_workspace_url,
17 get_databricks_token,
18)
19from src.metrics_collector import get_metrics_collector
22class ChuckService:
23 """Service layer that provides a clean API for the UI to interact with business logic."""
25 def __init__(self, client: Optional[DatabricksAPIClient] = None):
26 """
27 Initialize the service with an optional client.
28 If no client is provided, it attempts to initialize one from config.
29 """
30 self.client = client
31 self.init_error: Optional[str] = None # Store initialization error message
32 if not self.client:
33 try:
34 token = get_databricks_token()
35 workspace_url = get_workspace_url()
36 if token and workspace_url: # Only initialize if both are present
37 self.client = DatabricksAPIClient(workspace_url, token)
38 elif not workspace_url:
39 self.init_error = "Workspace URL not configured."
40 elif not token:
41 self.init_error = "Databricks token not configured."
42 # else: both are missing, client remains None, init_error remains None (or set explicitly)
44 except Exception as e:
45 logging.error(
46 f"Error initializing DatabricksAPIClient in ChuckService: {e}",
47 exc_info=True,
48 )
49 self.client = None
50 self.init_error = f"Client initialization failed: {str(e)}"
52 def _parse_and_validate_tui_args(
53 self,
54 command_def: CommandDefinition,
55 raw_args: Tuple[str, ...],
56 raw_kwargs: Dict[
57 str, Any
58 ], # For potential future use if TUI supports named args
59 ) -> Tuple[Optional[Dict[str, Any]], Optional[CommandResult]]:
60 """
61 Parses raw string arguments from the TUI, converts types based on JSON schema,
62 applies defaults, and validates against the command's schema.
64 Args:
65 command_def: The CommandDefinition for the command.
66 raw_args: A tuple of raw string arguments from the TUI.
67 raw_kwargs: A dictionary of raw keyword arguments (currently unused by TUI).
69 Returns:
70 A tuple: (parsed_args_dict, None) if successful,
71 (None, error_command_result) if parsing/validation fails.
72 """
73 parsed_kwargs: Dict[str, Any] = {}
74 param_definitions = command_def.parameters or {}
75 param_names_ordered = list(
76 param_definitions.keys()
77 ) # Assumes order is defined or consistent
79 # 1. Handle special TUI argument packing (e.g., for agent query, upload content)
80 # This might need more sophisticated logic or command-specific metadata if complex.
81 # For now, assuming direct mapping or simple concatenation for specific commands.
83 if (command_def.name == "agent_query" or command_def.name == "agent") and len(
84 raw_args
85 ) > 0:
86 # Join all raw_args into a single 'query' string
87 # This assumes the command has a parameter named 'query' in its schema.
88 if "query" in param_definitions:
89 query_text = " ".join(raw_args)
90 parsed_kwargs["query"] = query_text
91 # Store the joined text as a "rest" parameter too
92 parsed_kwargs["rest"] = query_text
93 # We don't set raw_args explicitly here, as the command handler will get
94 # the original raw_args directly from the function call
95 raw_args = () # Consumed all raw_args
96 else:
97 return None, CommandResult(
98 False,
99 message=f"Command '{command_def.name}' is misconfigured: expected 'query' parameter.",
100 )
102 elif command_def.name == "bug" and len(raw_args) > 0:
103 # Join all raw_args into a single 'description' string
104 if "description" in param_definitions:
105 description_text = " ".join(raw_args)
106 parsed_kwargs["description"] = description_text
107 # Store the joined text as a "rest" parameter too
108 parsed_kwargs["rest"] = description_text
109 raw_args = () # Consumed all raw_args
110 else:
111 # Fallback to rest parameter
112 parsed_kwargs["rest"] = " ".join(raw_args)
113 raw_args = ()
115 elif command_def.name == "add_stitch_report" and len(raw_args) > 0:
116 # First arg is table_path, rest is notebook name
117 if len(raw_args) >= 1:
118 parsed_kwargs["table_path"] = raw_args[0]
119 if len(raw_args) > 1:
120 # Join remaining arguments as notebook name
121 name_text = " ".join(raw_args[1:])
122 parsed_kwargs["name"] = name_text
123 # Store joined text as rest parameter too
124 parsed_kwargs["rest"] = name_text
125 raw_args = () # Consumed all raw_args
127 elif command_def.name == "upload_file" and len(raw_args) >= 1:
128 # First arg is filename, rest is content (joined)
129 # Assumes parameters "filename" and "content"
130 if "filename" in param_definitions and "content" in param_definitions:
131 parsed_kwargs["filename"] = raw_args[0]
132 if len(raw_args) > 1:
133 parsed_kwargs["content"] = " ".join(raw_args[1:])
134 else:
135 # Content might be optional or handled by schema default/validation
136 # If content is required and not provided, schema validation should catch it.
137 # Or, if it can be truly empty: parsed_kwargs["content"] = ""
138 pass # Let schema validation handle if content is missing but required
139 raw_args = () # Consumed all raw_args
140 else:
141 return None, CommandResult(
142 False,
143 message=f"Command '{command_def.name}' is misconfigured: expected 'filename' and 'content' parameters.",
144 )
146 # 2. Parse arguments - support both positional and flag-style arguments
147 remaining_args = list(raw_args)
149 # First, parse flag-style arguments (--flag value)
150 i = 0
151 while i < len(remaining_args):
152 arg = remaining_args[i]
153 if arg.startswith("--"):
154 # Flag-style argument
155 flag_name = arg[2:] # Remove '--' prefix
156 if flag_name in param_definitions:
157 # Check if we have a value for this flag
158 if i + 1 < len(remaining_args) and not remaining_args[
159 i + 1
160 ].startswith("--"):
161 flag_value = remaining_args[i + 1]
162 parsed_kwargs[flag_name] = flag_value
163 # Remove both flag and value from remaining args
164 remaining_args.pop(i) # Remove flag
165 remaining_args.pop(
166 i
167 ) # Remove value (index shifts after first pop)
168 continue
169 else:
170 # Flag without value - could be boolean flag
171 if (
172 param_definitions.get(flag_name, {}).get("type")
173 == "boolean"
174 ):
175 parsed_kwargs[flag_name] = True
176 remaining_args.pop(i)
177 continue
178 else:
179 usage = (
180 command_def.usage_hint
181 or f"Correct usage for '{command_def.name}' not available."
182 )
183 return None, CommandResult(
184 False,
185 message=f"Flag '--{flag_name}' requires a value. {usage}",
186 )
187 else:
188 usage = (
189 command_def.usage_hint
190 or f"Correct usage for '{command_def.name}' not available."
191 )
192 return None, CommandResult(
193 False,
194 message=f"Unknown flag '--{flag_name}' for command '{command_def.name}'. {usage}",
195 )
196 i += 1
198 # Then, map remaining positional arguments to parameter names
199 for i, arg_val_str in enumerate(remaining_args):
200 if i < len(param_names_ordered):
201 param_name = param_names_ordered[i]
202 # If special handling above or flag parsing already populated this, skip to avoid overwrite
203 if param_name not in parsed_kwargs:
204 parsed_kwargs[param_name] = (
205 arg_val_str # Store as string for now, type conversion next
206 )
207 else:
208 # Too many positional arguments provided
209 usage = (
210 command_def.usage_hint
211 or f"Correct usage for '{command_def.name}' not available."
212 )
213 return None, CommandResult(
214 False,
215 message=f"Too many arguments for command '{command_def.name}'. {usage}",
216 )
218 # 3. Type conversion (string from TUI to schema-defined type) & apply defaults
219 final_args_for_validation: Dict[str, Any] = {}
220 for param_name, schema_prop in param_definitions.items():
221 param_type = schema_prop.get("type")
222 default_value = schema_prop.get("default")
224 if param_name in parsed_kwargs:
225 raw_value = parsed_kwargs[param_name]
226 try:
227 if param_type == "integer":
228 final_args_for_validation[param_name] = int(raw_value)
229 elif param_type == "number": # JSON schema 'number' can be float
230 final_args_for_validation[param_name] = float(raw_value)
231 elif param_type == "boolean":
232 # Handle common boolean strings, jsonschema might do this too
233 if isinstance(raw_value, str):
234 val_lower = raw_value.lower()
235 if val_lower in ["true", "t", "yes", "y", "1"]:
236 final_args_for_validation[param_name] = True
237 elif val_lower in ["false", "f", "no", "n", "0"]:
238 final_args_for_validation[param_name] = False
239 else:
240 # Let jsonschema validation catch this if type is strict
241 final_args_for_validation[param_name] = raw_value
242 else: # Already a bool (e.g. from default or previous step)
243 final_args_for_validation[param_name] = bool(raw_value)
244 elif param_type == "array" or (
245 isinstance(param_type, list) and "array" in param_type
246 ):
247 # Support both JSON array format and comma-separated strings
248 if isinstance(raw_value, str):
249 # Try to parse as JSON first
250 if raw_value.strip().startswith(
251 "["
252 ) and raw_value.strip().endswith("]"):
253 try:
254 final_args_for_validation[param_name] = json.loads(
255 raw_value
256 )
257 except json.JSONDecodeError as je:
258 usage = (
259 command_def.usage_hint
260 or f"Check help for '{command_def.name}'."
261 )
262 return None, CommandResult(
263 False,
264 message=f"Invalid JSON array format for '{param_name}': {str(je)}. {usage}",
265 )
266 else:
267 # Fall back to comma-separated strings
268 final_args_for_validation[param_name] = [
269 s.strip() for s in raw_value.split(",") if s.strip()
270 ]
271 elif isinstance(
272 raw_value, (list, tuple)
273 ): # Already a list or tuple
274 final_args_for_validation[param_name] = list(raw_value)
275 else:
276 return None, CommandResult(
277 False,
278 message=f"Invalid array format for '{param_name}'. Expected JSON array or comma-separated string.",
279 )
280 elif param_type == "string":
281 final_args_for_validation[param_name] = str(raw_value)
282 else: # Unknown type or type not requiring conversion from string (e.g. already processed)
283 final_args_for_validation[param_name] = raw_value
284 except ValueError:
285 usage = (
286 command_def.usage_hint
287 or f"Check help for '{command_def.name}'."
288 )
289 return None, CommandResult(
290 False,
291 message=f"Invalid value for '{param_name}': '{raw_value}'. Expected type '{param_type}'. {usage}",
292 )
293 elif default_value is not None:
294 final_args_for_validation[param_name] = default_value
295 # If not in parsed_kwargs and no default, it's either optional or will be caught by 'required' validation
297 # 4. Incorporate raw_kwargs if TUI ever supports named args directly (e.g. /cmd --option val)
298 # For now, this is a placeholder.
299 # final_args_for_validation.update(raw_kwargs) # If raw_kwargs were typed and validated
301 # 5. Validate against the full JSON schema
302 full_schema_for_validation = {
303 "type": "object",
304 "properties": param_definitions,
305 "required": command_def.required_params or [],
306 }
307 try:
308 jsonschema.validate(
309 instance=final_args_for_validation, schema=full_schema_for_validation
310 )
311 except jsonschema.exceptions.ValidationError as ve:
312 usage = (
313 command_def.usage_hint or f"Use /help {command_def.name} for details."
314 )
315 # More detailed error: ve.message, ve.path, ve.schema_path
316 error_path = " -> ".join(map(str, ve.path)) if ve.path else "argument"
317 return None, CommandResult(
318 False, message=f"Invalid argument '{error_path}': {ve.message}. {usage}"
319 )
321 return final_args_for_validation, None
323 def execute_command(
324 self,
325 command_name_from_ui: str,
326 *raw_args: str,
327 interactive_input: Optional[str] = None,
328 tool_output_callback: Optional[callable] = None,
329 **raw_kwargs: Any, # For future TUI use, e.g. /cmd --named_arg value
330 ) -> CommandResult:
331 """
332 Execute a command looked up from the registry, with argument parsing and validation.
333 """
334 command_def = get_command(
335 command_name_from_ui
336 ) # Handles TUI aliases via registry
338 if not command_def:
339 return CommandResult(
340 False,
341 message=f"Unknown command: '{command_name_from_ui}'. Type /help for list.",
342 )
344 if not command_def.visible_to_user:
345 return CommandResult(
346 False,
347 message=f"Command '{command_name_from_ui}' is not available for direct use.",
348 )
350 # Authentication Check (before argument parsing)
351 effective_client = self.client
352 if command_def.needs_api_client:
353 # Special case for setup_wizard - always allow it to run even without client
354 if command_def.name == "setup_wizard" or command_name_from_ui == "/setup":
355 # Allow setup wizard to run without a client
356 effective_client = None
357 elif not self.client:
358 error_msg = (
359 self.init_error
360 or "Client not initialized. Please set workspace URL and token (e.g. /select-workspace, /set-token, then /status or /connect)."
361 )
362 return CommandResult(False, message=f"Not authenticated. {error_msg}")
364 parsed_args_dict: Dict[str, Any]
365 args_for_handler: Dict[str, Any]
367 # Interactive Mode Handling
368 if command_def.supports_interactive_input:
369 # For interactive commands, the `interactive_input` is the primary payload.
370 # Always include the interactive_input parameter, even if None, for interactive commands
371 # This prevents dropping inputs in complex interactive flows
372 args_for_handler = {"interactive_input": interactive_input}
373 else:
374 # Standard Argument Parsing & Validation
375 parsed_args_dict, error_result = self._parse_and_validate_tui_args(
376 command_def, raw_args, raw_kwargs
377 )
378 if error_result:
379 return error_result
380 if (
381 parsed_args_dict is None
382 ): # Should be caught by error_result, but as a safeguard
383 return CommandResult(
384 False, message="Internal error during argument parsing."
385 )
386 args_for_handler = parsed_args_dict
388 # Pass tool output callback for agent commands
389 if command_def.name == "agent" and tool_output_callback:
390 args_for_handler["tool_output_callback"] = tool_output_callback
392 # Handler Execution
393 try:
394 # All handlers now expect (client, **kwargs)
395 result: CommandResult = command_def.handler(
396 effective_client, **args_for_handler
397 )
399 # Special Command Post-Processing (Example: set-token)
400 if command_def.name == "databricks-login" and result.success:
401 # The handler for set_token now returns data={"reinitialize_client": True} on success
402 if isinstance(result.data, dict) and result.data.get(
403 "reinitialize_client"
404 ):
405 logging.info("Reinitializing client after successful set-token.")
406 self.reinitialize_client()
408 return result
409 except Exception as e_handler:
410 # Handle pagination cancellation specially - let it bubble up
411 from src.exceptions import PaginationCancelled
413 if isinstance(e_handler, PaginationCancelled):
414 raise # Re-raise to bubble up to main TUI loop
416 logging.error(
417 f"Error executing handler for command '{command_def.name}': {e_handler}",
418 exc_info=True,
419 )
421 # Track error event
422 try:
423 metrics_collector = get_metrics_collector()
424 # Create a context string from the command and args
425 command_context = f"command: {command_name_from_ui}"
426 if args_for_handler:
427 # Convert args to a simple string representation for the error report
428 args_str = ", ".join(
429 [
430 f"{k}={v}"
431 for k, v in args_for_handler.items()
432 if k != "interactive_input"
433 ]
434 )
435 if args_str:
436 command_context += f", args: {args_str}"
438 metrics_collector.track_event(
439 prompt=command_context,
440 error=traceback.format_exc(),
441 tools=[{"name": command_def.name, "arguments": args_for_handler}],
442 additional_data={"event_context": "error_report"},
443 )
444 except Exception as metrics_error:
445 # Don't let metrics collection errors affect the primary error handling
446 logging.error(
447 f"Failed to track error metrics: {metrics_error}", exc_info=True
448 )
450 return CommandResult(
451 False,
452 message=f"Error during command execution: {str(e_handler)}",
453 error=e_handler,
454 )
456 def reinitialize_client(self) -> bool:
457 """
458 Reinitialize the API client with current configuration.
459 This should be called after settings like token or workspace URL change.
460 """
461 logging.info("Attempting to reinitialize DatabricksAPIClient...")
462 try:
463 token = get_databricks_token()
464 workspace_url = get_workspace_url()
465 if token and workspace_url:
466 self.client = DatabricksAPIClient(workspace_url, token)
467 self.init_error = None # Clear previous init error
468 logging.info("DatabricksAPIClient reinitialized successfully.")
469 return True
470 else:
471 self.client = None
472 self.init_error = "Cannot reinitialize client: Workspace URL or token missing from config."
473 logging.warning(f"{self.init_error}")
474 return False
475 except Exception as e:
476 self.client = None
477 self.init_error = f"Failed to reinitialize client: {str(e)}"
478 logging.error(self.init_error, exc_info=True)
479 return False