AAP MCP Server

#!/usr/bin/env python3 """ AAP Controller Template & Survey Management Tool """ from typing import Any, Dict, Optional, Union from fastmcp import FastMCP from pydantic import Field from connectors.aap_connector import get_aap_connector def register_template_tools(mcp: FastMCP): """Register template management tools with the MCP server""" @mcp.tool() def template_survey_management( action: str = Field(description="Action: list_job_templates, create_job_template, update_job_template, delete_job_template, launch_job_template, list_workflow_templates, create_workflow_template, get_survey, update_survey, get_workflow_survey, update_workflow_survey, list_system_templates"), template_id: Optional[Union[int, float]] = Field(None, description="Template ID"), template_data: Optional[Dict[str, Any]] = Field(None, description="Template data"), survey_data: Optional[Dict[str, Any]] = Field(None, description="Survey specification data"), launch_data: Optional[Dict[str, Any]] = Field(None, description="Launch data with extra vars"), filters: Optional[Dict[str, Any]] = Field(None, description="Filters for listing") ) -> Dict[str, Any]: """ Template and survey management tool. Handles job templates, workflow templates, and survey specifications. """ try: # Job Template Operations if action == "list_job_templates": params = filters or {} return get_aap_connector().get("job_templates/", params) elif action == "create_job_template": if not template_data: return {"error": "template_data is required"} return get_aap_connector().post("job_templates/", template_data) elif action == "update_job_template": if not template_id or not template_data: return {"error": "template_id and template_data are required"} return get_aap_connector().patch(f"job_templates/{template_id}/", template_data) elif action == "delete_job_template": if not template_id: return {"error": "template_id is required"} return get_aap_connector().delete(f"job_templates/{template_id}/") elif action == "launch_job_template": if not template_id: return {"error": "template_id is required"} data = launch_data or {} return get_aap_connector().post(f"job_templates/{template_id}/launch/", data) # Workflow Template Operations elif action == "list_workflow_templates": params = filters or {} return get_aap_connector().get("workflow_job_templates/", params) elif action == "create_workflow_template": if not template_data: return {"error": "template_data is required"} return get_aap_connector().post("workflow_job_templates/", template_data) # Job Template Survey Operations elif action == "get_survey": if not template_id: return {"error": "template_id is required"} return get_aap_connector().get(f"job_templates/{template_id}/survey_spec/") elif action == "update_survey": if not template_id or not survey_data: return {"error": "template_id and survey_data are required"} return get_aap_connector().post(f"job_templates/{template_id}/survey_spec/", survey_data) # Workflow Template Survey Operations elif action == "get_workflow_survey": if not template_id: return {"error": "template_id is required"} return get_aap_connector().get(f"workflow_job_templates/{template_id}/survey_spec/") elif action == "update_workflow_survey": if not template_id or not survey_data: return {"error": "template_id and survey_data are required"} try: # First enable survey on the workflow template if not already enabled get_aap_connector().patch(f"workflow_job_templates/{template_id}/", {"survey_enabled": True}) # Validate survey_data format if "spec" not in survey_data: return {"error": "survey_data must contain 'spec' field"} # Ensure proper format for all questions for i, question in enumerate(survey_data["spec"]): if "question_name" not in question or "variable" not in question or "type" not in question: return {"error": f"Question {i+1} missing required fields: question_name, variable, or type"} # Set default values if missing if "new_question" not in question: question["new_question"] = True if "required" not in question: question["required"] = True if "default" not in question: question["default"] = "" if "question_description" not in question: question["question_description"] = "" # Handle type-specific fields if question["type"] == "multiplechoice": if "choices" not in question: question["choices"] = [] question.setdefault("min", 0) question.setdefault("max", 1024) elif question["type"] in ["text", "textarea"]: question.setdefault("choices", "") question.setdefault("min", 0) question.setdefault("max", 1024) elif question["type"] == "integer": question.setdefault("choices", "") question.setdefault("min", 0) question.setdefault("max", 9999) # Ensure default is integer for integer type if "default" in question and question["default"] != "": try: question["default"] = int(question["default"]) except (ValueError, TypeError): pass elif question["type"] == "float": question.setdefault("choices", "") question.setdefault("min", 0.0) question.setdefault("max", 9999.0) # Ensure default is float for float type if "default" in question and question["default"] != "": try: question["default"] = float(question["default"]) except (ValueError, TypeError): pass else: question.setdefault("choices", "") question.setdefault("min", 0) question.setdefault("max", 1024) # Ensure min/max are always integers for non-float types if question["type"] != "float": if isinstance(question.get("min"), str): try: question["min"] = int(question["min"]) if question["min"] else 0 except ValueError: question["min"] = 0 if isinstance(question.get("max"), str): try: question["max"] = int(question["max"]) if question["max"] else 1024 except ValueError: question["max"] = 1024 # Make the API call response = get_aap_connector().post(f"workflow_job_templates/{template_id}/survey_spec/", survey_data) return { "success": True, "survey_created": True, "workflow_id": template_id, "survey_data": survey_data, "response": response } except Exception as e: return { "error": f"Failed to create workflow survey: {str(e)}", "workflow_id": template_id, "survey_data_provided": survey_data, "debug_info": "Check survey spec format and workflow template status" } # System Templates elif action == "list_system_templates": params = filters or {} return get_aap_connector().get("system_job_templates/", params) else: return {"error": f"Unknown action: {action}"} except Exception as e: return {"error": f"Template/Survey management failed: {str(e)}"}