Coverage for src/service.py: 47%

205 statements  

« 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""" 

5 

6import json 

7import logging 

8import jsonschema 

9import traceback 

10from typing import Dict, Optional, Any, Tuple 

11 

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 

20 

21 

22class ChuckService: 

23 """Service layer that provides a clean API for the UI to interact with business logic.""" 

24 

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) 

43 

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)}" 

51 

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. 

63 

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). 

68 

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 

78 

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. 

82 

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 ) 

101 

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 = () 

114 

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 

126 

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 ) 

145 

146 # 2. Parse arguments - support both positional and flag-style arguments 

147 remaining_args = list(raw_args) 

148 

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 

197 

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 ) 

217 

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") 

223 

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 

296 

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 

300 

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 ) 

320 

321 return final_args_for_validation, None 

322 

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 

337 

338 if not command_def: 

339 return CommandResult( 

340 False, 

341 message=f"Unknown command: '{command_name_from_ui}'. Type /help for list.", 

342 ) 

343 

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 ) 

349 

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}") 

363 

364 parsed_args_dict: Dict[str, Any] 

365 args_for_handler: Dict[str, Any] 

366 

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 

387 

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 

391 

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 ) 

398 

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() 

407 

408 return result 

409 except Exception as e_handler: 

410 # Handle pagination cancellation specially - let it bubble up 

411 from src.exceptions import PaginationCancelled 

412 

413 if isinstance(e_handler, PaginationCancelled): 

414 raise # Re-raise to bubble up to main TUI loop 

415 

416 logging.error( 

417 f"Error executing handler for command '{command_def.name}': {e_handler}", 

418 exc_info=True, 

419 ) 

420 

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}" 

437 

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 ) 

449 

450 return CommandResult( 

451 False, 

452 message=f"Error during command execution: {str(e_handler)}", 

453 error=e_handler, 

454 ) 

455 

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