Coverage for src/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
« prev ^ index » next coverage.py v7.8.0, created at 2025-06-05 22:56 -0700
1"""
2Agent tool schema provider and execution logic.
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"""
13import logging
14import jsonschema # Requires jsonschema to be installed
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
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.
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.
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()
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.
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.
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 )
67 command_def = get_command(tool_name)
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."}
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."}
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 }
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 }
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
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)
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 }
132 if result_obj.success:
133 logging.debug(
134 f"Tool '{tool_name}' executed successfully. Data: {result_obj.data}"
135 )
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 )
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
156 if isinstance(e, PaginationCancelled):
157 raise # Re-raise to bubble up to agent manager
159 logging.warning(
160 f"Tool output callback failed for '{tool_name}': {e}"
161 )
163 # Return the appropriate output
164 if formatted_output is not None:
165 return formatted_output
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
190 except Exception as e_exec:
191 # Handle pagination cancellation specially - let it bubble up
192 from ..exceptions import PaginationCancelled
194 if isinstance(e_exec, PaginationCancelled):
195 raise # Re-raise to bubble up to agent manager
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 }