From 161c5ab2ffaddbb9489ed24a1899cdf32f59e0b8 Mon Sep 17 00:00:00 2001 From: ykd007 Date: Fri, 15 May 2026 09:09:59 +0530 Subject: [PATCH] Add type hints to MCP function parameters and return types The mcp/main.py functions had bare dict and list return annotations with no generic parameters. This adds proper typing for IDE support and static analysis tools, and fixes a small typo in the verify_llm docstring (spect -> spec). Closes #198 --- agentic_security/mcp/main.py | 18 ++++++++++-------- 1 file changed, 10 insertions(+), 8 deletions(-) diff --git a/agentic_security/mcp/main.py b/agentic_security/mcp/main.py index 7e19e2c6..e256eb09 100644 --- a/agentic_security/mcp/main.py +++ b/agentic_security/mcp/main.py @@ -1,3 +1,5 @@ +from typing import Any + import httpx from mcp.server.fastmcp import FastMCP @@ -12,14 +14,14 @@ @mcp.tool() -async def verify_llm(spec: str) -> dict: +async def verify_llm(spec: str) -> dict[str, Any]: """ Verify an LLM model specification using the FastAPI server Returns: dict: containing the verification result form the FastAPI server - Args: spect(str): The specification of the LLM model to verify. + Args: spec(str): The specification of the LLM model to verify. """ url = f"{AGENTIC_SECURITY}/verify" @@ -34,7 +36,7 @@ async def start_scan( maxBudget: int, optimize: bool = False, enableMultiStepAttack: bool = False, -) -> dict: +) -> dict[str, Any]: """ Start an LLM security scan via the FastAPI server. Returns: @@ -63,7 +65,7 @@ async def start_scan( @mcp.tool() -async def stop_scan() -> dict: +async def stop_scan() -> dict[str, Any]: """Stop an ongoing scan via the FastAPI server. Returns: @@ -76,12 +78,12 @@ async def stop_scan() -> dict: @mcp.tool() -async def get_data_config() -> list: +async def get_data_config() -> list[Any]: """ Retrieve data configuration from the FastAPI server. Returns: - list: The response from the FastAPI server, confirming the scan has been stopped. + list: The response from the FastAPI server containing the data configuration. """ url = f"{AGENTIC_SECURITY}/v1/data-config" async with httpx.AsyncClient() as client: @@ -90,9 +92,9 @@ async def get_data_config() -> list: @mcp.tool() -async def get_spec_templates() -> list: +async def get_spec_templates() -> list[Any]: """ - Retrieve data configuration from the FastAPI server. + Retrieve LLM specification templates from the FastAPI server. Returns: list: The LLM specification templates from the FastAPI server.