324 lines
13 KiB
Python
324 lines
13 KiB
Python
# tool_utils/communicator.py
|
|
|
|
"""
|
|
Utility module for Python tools designed to integrate with ProjectUtility.
|
|
|
|
Provides standardized functions for argument parsing based on a tool's
|
|
configuration and for sending structured messages (progress, status, results)
|
|
via standard output using JSON Lines, as well as handling standard error
|
|
and exit codes according to the ProjectUtility Tool Interface specification.
|
|
"""
|
|
|
|
import argparse
|
|
import json
|
|
import sys
|
|
import os
|
|
import logging
|
|
from typing import Dict, Any, List, Optional, Tuple, Callable
|
|
|
|
# Set up a basic logger for the utility module itself
|
|
# Note: Tools using this should ideally configure their own application-level logging.
|
|
util_logger = logging.getLogger(__name__)
|
|
# Add a basic handler if no other logging is configured by the tool
|
|
if not util_logger.hasHandlers():
|
|
util_logger.addHandler(logging.StreamHandler(sys.stderr))
|
|
util_logger.setLevel(logging.WARNING) # Default to WARNING to avoid being too noisy
|
|
|
|
# --- Constants for message types ---
|
|
MSG_TYPE_PROGRESS = "progress"
|
|
MSG_TYPE_STATUS = "status"
|
|
MSG_TYPE_LOG = "log"
|
|
MSG_TYPE_RESULT = "result"
|
|
|
|
# --- Argument Parsing ---
|
|
|
|
def parse_arguments(tool_parameters: List[Dict[str, Any]]) -> argparse.Namespace:
|
|
"""
|
|
Parses command-line arguments based on the tool parameter definitions.
|
|
|
|
Follows the convention:
|
|
- String, Int, Float, File, Folder: --name <value>
|
|
- Boolean: --name (flag is present if True)
|
|
|
|
Args:
|
|
tool_parameters: A list of dictionaries, where each dictionary defines
|
|
a parameter (matching the structure in tool_config.json,
|
|
requiring at least 'name', 'type', 'required', 'description').
|
|
|
|
Returns:
|
|
An argparse.Namespace object containing the parsed arguments.
|
|
The script will exit if '--help' is invoked.
|
|
"""
|
|
parser = argparse.ArgumentParser(
|
|
description="Tool script adhering to ProjectUtility interface.",
|
|
formatter_class=argparse.ArgumentDefaultsHelpFormatter # Show defaults in help
|
|
)
|
|
|
|
for param in tool_parameters:
|
|
name = param.get("name")
|
|
ptype = str(param.get("type", "string")).lower() # Default to string if missing
|
|
required = param.get("required", False)
|
|
default = param.get("default")
|
|
description = param.get("description", "") # Help text
|
|
|
|
if not name:
|
|
util_logger.warning(f"Skipping parameter definition missing 'name': {param}")
|
|
continue
|
|
|
|
arg_name = f"--{name}"
|
|
|
|
try:
|
|
if ptype == "boolean":
|
|
# For booleans, store_true means the arg name implies True
|
|
# If required=True, it doesn't make much sense for a boolean flag,
|
|
# but argparse handles 'required' flags correctly if needed.
|
|
# Default behavior: if flag present -> True, else False.
|
|
# We don't use the 'default' from config directly here as store_true handles it.
|
|
parser.add_argument(
|
|
arg_name,
|
|
action="store_true", # Automatically sets default=False
|
|
help=f"(Flag) {description}",
|
|
required=required # Can be required, though unusual for flags
|
|
)
|
|
# Note: If a boolean needs to default to True and be toggled off,
|
|
# use action="store_false" and adjust logic/help text.
|
|
else:
|
|
# For other types, specify the type for basic conversion by argparse
|
|
# Note: 'file' and 'folder' are treated as strings here; validation
|
|
# would happen within the tool's logic.
|
|
arg_type = str
|
|
if ptype == "integer":
|
|
arg_type = int
|
|
elif ptype == "float":
|
|
arg_type = float
|
|
|
|
parser.add_argument(
|
|
arg_name,
|
|
type=arg_type,
|
|
required=required,
|
|
default=default,
|
|
help=description,
|
|
metavar=f'<{ptype.upper()}>' # Provides hint in help message, e.g. <STRING>
|
|
)
|
|
except Exception as e:
|
|
util_logger.error(f"Failed to add argument for parameter '{name}': {e}", exc_info=True)
|
|
|
|
|
|
# Parse known args to allow flexibility if gestore sends extra args
|
|
args, unknown_args = parser.parse_known_args()
|
|
|
|
if unknown_args:
|
|
util_logger.warning(f"Received unrecognized arguments: {unknown_args}")
|
|
|
|
# Post-parsing validation (optional but recommended)
|
|
# Example: Check if required arguments that are *not* booleans actually have a value
|
|
# (argparse usually handles this with 'required=True', but defensive check)
|
|
for param in tool_parameters:
|
|
name = param.get("name")
|
|
required = param.get("required", False)
|
|
ptype = str(param.get("type", "string")).lower()
|
|
if required and name and ptype != "boolean":
|
|
if getattr(args, name.replace('-', '_'), None) is None: # Check if value is None
|
|
# Argparse should have already exited if required=True and missing,
|
|
# but this handles edge cases or if required=False but tool logic needs it.
|
|
print_error(f"Required parameter '--{name}' is missing or has no value.", exit_code=2)
|
|
# print_error exits the script
|
|
|
|
|
|
util_logger.debug(f"Arguments parsed successfully: {args}")
|
|
return args
|
|
|
|
# --- Standard Output Communication (JSON Lines) ---
|
|
|
|
def _send_json_message(msg_type: str, data: Dict[str, Any]) -> None:
|
|
"""Formats and sends a JSON message to stdout."""
|
|
message = {"type": msg_type, **data} # Combine type and data payload
|
|
try:
|
|
json_str = json.dumps(message, ensure_ascii=False)
|
|
print(json_str, file=sys.stdout, flush=True) # Print and FLUSH!
|
|
util_logger.debug(f"Sent message: {json_str}")
|
|
except (TypeError, ValueError) as e:
|
|
util_logger.error(f"Failed to serialize message to JSON: {message}. Error: {e}", exc_info=True)
|
|
# Also report the serialization error via stderr
|
|
print(f"TOOL_UTILS_ERROR: Failed to send JSON message. Type: {msg_type}. Error: {e}", file=sys.stderr, flush=True)
|
|
|
|
|
|
def send_progress(value: float, message: Optional[str] = None) -> None:
|
|
"""Sends a progress update message to stdout.
|
|
|
|
Args:
|
|
value: Progress value (ideally between 0.0 and 1.0).
|
|
message: Optional descriptive message about the progress.
|
|
"""
|
|
payload = {"value": max(0.0, min(1.0, value))} # Clamp value between 0 and 1
|
|
if message is not None:
|
|
payload["message"] = str(message)
|
|
_send_json_message(MSG_TYPE_PROGRESS, payload)
|
|
|
|
|
|
def send_status(message: str) -> None:
|
|
"""Sends a status message (general information) to stdout.
|
|
|
|
Args:
|
|
message: The status message string.
|
|
"""
|
|
if not message:
|
|
util_logger.warning("Attempted to send an empty status message.")
|
|
return
|
|
_send_json_message(MSG_TYPE_STATUS, {"message": str(message)})
|
|
|
|
|
|
def send_log(level: str, message: str) -> None:
|
|
"""Sends a log message to stdout (for potential capture by the manager).
|
|
Note: This is distinct from the tool's internal logging to stderr/file.
|
|
|
|
Args:
|
|
level: Log level (e.g., "info", "warning", "debug", "error").
|
|
message: The log message string.
|
|
"""
|
|
if not message:
|
|
util_logger.warning("Attempted to send an empty log message.")
|
|
return
|
|
_send_json_message(MSG_TYPE_LOG, {"level": str(level).lower(), "message": str(message)})
|
|
|
|
|
|
def send_result(result_data: Dict[str, Any]) -> None:
|
|
"""Sends a result data message to stdout.
|
|
|
|
Args:
|
|
result_data: A dictionary containing the results to be sent.
|
|
Must be JSON serializable.
|
|
"""
|
|
if not isinstance(result_data, dict):
|
|
util_logger.error(f"Result data must be a dictionary, got {type(result_data)}. Cannot send result.")
|
|
print(f"TOOL_UTILS_ERROR: Result data must be a dictionary. Data: {result_data}", file=sys.stderr, flush=True)
|
|
return
|
|
_send_json_message(MSG_TYPE_RESULT, {"data": result_data})
|
|
|
|
|
|
# --- Standard Error and Exit Codes ---
|
|
|
|
def print_error(message: str, exit_code: Optional[int] = 1) -> None:
|
|
"""
|
|
Prints an error message to stderr and optionally exits the script.
|
|
|
|
Args:
|
|
message: The error message string.
|
|
exit_code: The exit code to use (e.g., 1 for general error).
|
|
If None, the script will not exit after printing.
|
|
Defaults to 1.
|
|
"""
|
|
print(f"ERROR: {message}", file=sys.stderr, flush=True) # Print and FLUSH!
|
|
util_logger.error(message) # Also log it via logging if configured
|
|
|
|
if exit_code is not None:
|
|
util_logger.info(f"Exiting with code {exit_code} due to error.")
|
|
sys.exit(exit_code)
|
|
|
|
|
|
def exit_success(message: Optional[str] = None) -> None:
|
|
"""
|
|
Exits the script with exit code 0 (success).
|
|
Optionally prints a final success message to stdout (as status).
|
|
|
|
Args:
|
|
message: Optional final success message to send as status.
|
|
"""
|
|
if message:
|
|
send_status(f"Success: {message}") # Send final message via standard channel
|
|
util_logger.info("Exiting with code 0 (Success).")
|
|
sys.exit(0)
|
|
|
|
|
|
# --- Helper to Load Tool Config (Optional but useful for tools) ---
|
|
|
|
def load_tool_config(config_filename: str = "tool_config.json") -> Optional[Dict[str, Any]]:
|
|
"""
|
|
Loads the tool's configuration file (typically 'tool_config.json')
|
|
from the directory containing the currently running script.
|
|
|
|
Args:
|
|
config_filename: The name of the configuration file.
|
|
|
|
Returns:
|
|
A dictionary with the loaded configuration data, or None if the
|
|
file is not found or cannot be parsed.
|
|
"""
|
|
script_dir = os.path.dirname(os.path.abspath(sys.argv[0])) # Directory of the main script
|
|
config_path = os.path.join(script_dir, config_filename)
|
|
|
|
if not os.path.isfile(config_path):
|
|
util_logger.error(f"Tool configuration file not found: {config_path}")
|
|
print(f"TOOL_UTILS_ERROR: Configuration file '{config_filename}' not found in script directory.", file=sys.stderr, flush=True)
|
|
return None
|
|
|
|
try:
|
|
with open(config_path, 'r', encoding='utf-8') as f:
|
|
config_data = json.load(f)
|
|
if not isinstance(config_data, dict):
|
|
util_logger.error(f"Invalid format in configuration file: {config_path}. Expected a JSON object.")
|
|
print(f"TOOL_UTILS_ERROR: Configuration file '{config_filename}' does not contain a valid JSON object.", file=sys.stderr, flush=True)
|
|
return None
|
|
util_logger.debug(f"Loaded tool configuration from: {config_path}")
|
|
return config_data
|
|
except json.JSONDecodeError as e:
|
|
util_logger.error(f"Invalid JSON in configuration file: {config_path}. Error: {e}")
|
|
print(f"TOOL_UTILS_ERROR: Invalid JSON in '{config_filename}'. Error: {e}", file=sys.stderr, flush=True)
|
|
return None
|
|
except Exception as e:
|
|
util_logger.exception(f"Failed to load or parse tool configuration file: {config_path}")
|
|
print(f"TOOL_UTILS_ERROR: Could not load '{config_filename}'. Error: {e}", file=sys.stderr, flush=True)
|
|
return None
|
|
|
|
|
|
# --- Example Usage (within a tool script) ---
|
|
#
|
|
# if __name__ == "__main__":
|
|
# # 1. Load configuration to get parameter definitions
|
|
# config = load_tool_config()
|
|
# if config is None:
|
|
# print_error("Failed to load tool configuration.", exit_code=3)
|
|
#
|
|
# tool_params_def = config.get("parameters", [])
|
|
#
|
|
# # 2. Parse arguments based on the definition
|
|
# try:
|
|
# args = parse_arguments(tool_params_def)
|
|
# except Exception as e:
|
|
# print_error(f"Error parsing arguments: {e}", exit_code=2)
|
|
#
|
|
# # 3. Use the communicator functions during execution
|
|
# send_status("Tool execution started.")
|
|
# try:
|
|
# # --- Your Tool's Logic ---
|
|
# user_input = args.user_text # Access parsed args
|
|
# print(f"Tool received text: {user_input}") # Normal diagnostic print to stdout (optional)
|
|
# sys.stdout.flush()
|
|
#
|
|
# total_steps = 5
|
|
# for i in range(total_steps):
|
|
# # Simulate work
|
|
# import time
|
|
# time.sleep(0.5)
|
|
# progress_val = (i + 1) / total_steps
|
|
# send_progress(progress_val, f"Step {i+1} completed.")
|
|
# if i == 2:
|
|
# send_log("warning", "Potential issue detected during step 3.")
|
|
#
|
|
# # Simulate an error condition
|
|
# # if args.user_text == "error":
|
|
# # print_error("User triggered a simulated error!", exit_code=5)
|
|
#
|
|
# # --- End of Tool's Logic ---
|
|
#
|
|
# # 4. Send results (if any)
|
|
# final_result = {"processed_text": user_input.upper(), "steps_done": total_steps}
|
|
# send_result(final_result)
|
|
#
|
|
# # 5. Exit successfully
|
|
# exit_success("Tool finished processing.")
|
|
#
|
|
# except Exception as e:
|
|
# # Catch unexpected errors in the tool's logic
|
|
# print_error(f"An unexpected error occurred in the tool: {e}", exit_code=1)
|
|
# |