diff --git a/agentic_security/mcp/main.py b/agentic_security/mcp/main.py index 7e19e2c6..0859ea01 100644 --- a/agentic_security/mcp/main.py +++ b/agentic_security/mcp/main.py @@ -11,6 +11,63 @@ AGENTIC_SECURITY = "http://0.0.0.0:8718" +# --------------------------------------------------------------------------- +# Prompt templates +# --------------------------------------------------------------------------- + + +@mcp.prompt() +def security_scan_prompt(llm_spec: str, max_budget: int = 1000) -> str: + """Generate a prompt to kick off a full LLM security scan. + + Args: + llm_spec: The LLM specification string identifying the model endpoint. + max_budget: Maximum number of probes to run (defaults to 1000). + """ + return ( + f"Please run a security scan on the following LLM specification:\n\n" + f" Spec: {llm_spec}\n" + f" Max budget: {max_budget} probes\n\n" + f"Use the start_scan tool to initiate the scan, then monitor progress " + f"with get_data_config, and stop it with stop_scan when complete." + ) + + +@mcp.prompt() +def verify_llm_prompt(llm_spec: str) -> str: + """Generate a prompt to verify that an LLM spec is reachable and well-formed. + + Args: + llm_spec: The LLM specification string to verify. + """ + return ( + f"Verify the following LLM specification is valid and reachable:\n\n" + f" Spec: {llm_spec}\n\n" + f"Use the verify_llm tool and report back whether the spec is accepted " + f"by the Agentic Security server." + ) + + +@mcp.prompt() +def adversarial_probe_prompt(llm_spec: str) -> str: + """Generate a prompt for an adversarial probing session with multi-step attacks. + + Args: + llm_spec: The LLM specification string identifying the target model. + """ + return ( + f"Run an adversarial probing session against the LLM described by:\n\n" + f" Spec: {llm_spec}\n\n" + f"Enable multi-step attacks and optimization in the start_scan call. " + f"After the scan finishes, summarise the most critical vulnerabilities found." + ) + + +# --------------------------------------------------------------------------- +# Tools +# --------------------------------------------------------------------------- + + @mcp.tool() async def verify_llm(spec: str) -> dict: """