Coverage for src/chuck_data/chuck_data/agent/tool_executor.py: 0%

72 statements  

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

1""" 

2Agent tool schema provider and execution logic. 

3 

4This module is responsible for: 

51. Providing tool schemas to the LLM agent, sourced from the command_registry. 

62. Executing tools (commands) requested by the LLM agent, including: 

7 - Looking up the command definition from the command_registry. 

8 - Validating arguments provided by the LLM against the command's JSON schema. 

9 - Calling the appropriate command handler. 

10 - Returning the result (or error) in a JSON-serializable dictionary format. 

11""" 

12 

13import logging 

14import jsonschema # Requires jsonschema to be installed 

15 

16from ...command_registry import get_command 

17from ...command_registry import ( 

18 get_agent_tool_schemas as get_command_registry_tool_schemas, 

19) 

20from ...commands.base import ( 

21 CommandResult, 

22) # For type hinting and checking handler result 

23from ...clients.databricks import DatabricksAPIClient # For type hinting api_client 

24from typing import Dict, Any, Optional, List 

25 

26 

27# The display_to_user utility and individual tool implementation functions 

28# (like list_models, set_warehouse, tag_pii_columns, scan_schema_for_pii, etc.) 

29# that were previously in this file have been removed. 

30# Their logic now resides within the modular command handlers in src.commands package. 

31 

32# --- PII Scanning and Stitch Setup Functions --- 

33# The original agent_tools.py contained `tag_pii_columns`, `scan_schema_for_pii`, and `setup_stitch` 

34# which were complex and also used by command handlers. 

35# These functions have now been moved to their respective command modules in the src.commands package. 

36# For example, tag_pii_columns logic is now in src.commands.tag_pii, 

37# scan_schema_for_pii is in src.commands.scan_pii, and setup_stitch is in src.commands.setup_stitch. 

38 

39 

40def get_tool_schemas() -> List[Dict[str, Any]]: 

41 """Get all command schemas in agent tool format from the command registry.""" 

42 return get_command_registry_tool_schemas() 

43 

44 

45def execute_tool( 

46 api_client: Optional[DatabricksAPIClient], 

47 tool_name: str, 

48 tool_args: Dict[str, Any], 

49 output_callback: Optional[callable] = None, 

50) -> Dict[str, Any]: 

51 """Execute a tool (command) by its name with the provided arguments. 

52 

53 Args: 

54 api_client: The Databricks API client, passed to the handler if needed. 

55 tool_name: The name of the tool (command) to execute. 

56 tool_args: A dictionary of arguments for the tool, parsed from LLM JSON. 

57 

58 Returns: 

59 A dictionary containing the result of the tool execution, suitable for 

60 JSON serialization and consumption by the LLM agent. 

61 Includes {"error": "..."} if any issues occur. 

62 """ 

63 logging.debug( 

64 f"Agent attempting to execute tool: {tool_name} with args: {tool_args}" 

65 ) 

66 

67 command_def = get_command(tool_name) 

68 

69 if not command_def: 

70 logging.error(f"Agent tool '{tool_name}' not found in command registry.") 

71 return {"error": f"Tool '{tool_name}' not found."} 

72 

73 if not command_def.visible_to_agent: 

74 logging.warning(f"Agent attempted to call non-agent-visible tool: {tool_name}") 

75 return {"error": f"Tool '{tool_name}' is not available to the agent."} 

76 

77 # --- Argument Validation using JSON Schema --- 

78 # Construct the full schema for validation, including required fields 

79 schema_to_validate = { 

80 "type": "object", 

81 "properties": command_def.parameters or {}, 

82 "required": command_def.required_params or [], 

83 } 

84 

85 try: 

86 jsonschema.validate(instance=tool_args, schema=schema_to_validate) 

87 logging.debug(f"Tool arguments for '{tool_name}' validated successfully.") 

88 except jsonschema.exceptions.ValidationError as ve: 

89 logging.error( 

90 f"Validation error for tool '{tool_name}' args {tool_args}: {ve.message}" 

91 ) 

92 # Provide a more user-friendly error message if possible, or just the validation message 

93 return { 

94 "error": f"Invalid arguments for tool '{tool_name}': {ve.message}. Schema: {ve.schema}" 

95 } 

96 except Exception as e_schema: # Catch other schema-related errors if any 

97 logging.error( 

98 f"Schema validation unexpected error for '{tool_name}': {e_schema}" 

99 ) 

100 return { 

101 "error": f"Internal error during argument validation for '{tool_name}': {str(e_schema)}" 

102 } 

103 

104 # --- Client Provisioning and Handler Execution --- 

105 effective_client = None 

106 if command_def.needs_api_client: 

107 if not api_client: 

108 logging.error( 

109 f"Tool '{tool_name}' requires an API client, but none was provided or initialized." 

110 ) 

111 return { 

112 "error": f"Tool '{tool_name}' cannot be executed: API client not available." 

113 } 

114 effective_client = api_client 

115 

116 try: 

117 # Handlers are now standardized to (client, **kwargs) 

118 # The command_def.handler should point to a function from command_handlers.py 

119 logging.debug( 

120 f"Executing agent tool '{tool_name}' with handler: {command_def.handler.__name__}" 

121 ) 

122 result_obj: CommandResult = command_def.handler(effective_client, **tool_args) 

123 

124 if not isinstance(result_obj, CommandResult): 

125 logging.error( 

126 f"Handler for '{tool_name}' did not return a CommandResult object. Got: {type(result_obj)}" 

127 ) 

128 return { 

129 "error": f"Tool '{tool_name}' did not execute correctly (internal error: unexpected result type)." 

130 } 

131 

132 if result_obj.success: 

133 logging.debug( 

134 f"Tool '{tool_name}' executed successfully. Data: {result_obj.data}" 

135 ) 

136 

137 # Use custom output formatter if available, otherwise return data directly 

138 formatted_output = None 

139 if command_def.output_formatter: 

140 try: 

141 formatted_output = command_def.output_formatter(result_obj) 

142 except Exception as e: 

143 logging.warning( 

144 f"Output formatter failed for '{tool_name}': {e}, falling back to default" 

145 ) 

146 

147 # Call output callback to display results immediately if provided 

148 # Always use original data for TUI display to preserve pagination capability 

149 if output_callback and result_obj.data is not None: 

150 try: 

151 output_callback(tool_name, result_obj.data) 

152 except Exception as e: 

153 # Handle pagination cancellation specially - let it bubble up 

154 from ...exceptions import PaginationCancelled 

155 

156 if isinstance(e, PaginationCancelled): 

157 raise # Re-raise to bubble up to agent manager 

158 

159 logging.warning( 

160 f"Tool output callback failed for '{tool_name}': {e}" 

161 ) 

162 

163 # Return the appropriate output 

164 if formatted_output is not None: 

165 return formatted_output 

166 

167 # Ensure data is JSON serializable. If not, this will be an issue for the agent. 

168 # The plan is for CommandResult.data to be JSON serializable. 

169 if result_obj.data is not None: 

170 return result_obj.data # Return the data directly as per plan 

171 else: 

172 # If data is None but success is true, return a success message or an empty dict 

173 # The LLM might expect some JSON output. 

174 return { 

175 "success": True, 

176 "message": result_obj.message or f"Tool '{tool_name}' completed.", 

177 } 

178 else: 

179 logging.error( 

180 f"Tool '{tool_name}' execution failed. Message: {result_obj.message}, Error: {result_obj.error}" 

181 ) 

182 error_payload = { 

183 "error": result_obj.message or f"Tool '{tool_name}' failed." 

184 } 

185 if result_obj.error and hasattr(result_obj.error, "__str__"): 

186 # Add more specific error details if available and simple 

187 error_payload["details"] = str(result_obj.error) 

188 return error_payload 

189 

190 except Exception as e_exec: 

191 # Handle pagination cancellation specially - let it bubble up 

192 from ...exceptions import PaginationCancelled 

193 

194 if isinstance(e_exec, PaginationCancelled): 

195 raise # Re-raise to bubble up to agent manager 

196 

197 logging.warning( 

198 f"Critical error executing tool '{tool_name}' handler '{command_def.handler.__name__}': {e_exec}", 

199 exc_info=True, 

200 ) 

201 return { 

202 "error": f"Unexpected error during execution of tool '{tool_name}': {str(e_exec)}" 

203 }