diff --git a/agentic_security/mcp/client.py b/agentic_security/mcp/client.py index e8a29cc9..5b2a9a52 100644 --- a/agentic_security/mcp/client.py +++ b/agentic_security/mcp/client.py @@ -1,4 +1,5 @@ import asyncio +from typing import Any from mcp import ClientSession, StdioServerParameters from mcp.client.stdio import stdio_client @@ -13,7 +14,7 @@ ) -async def run() -> None: +async def run() -> tuple[Any, Any, Any] | None: try: logger.info( "Starting stdio client session with server parameters: %s", server_params diff --git a/agentic_security/mcp/main.py b/agentic_security/mcp/main.py index 7e19e2c6..4a374cd4 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,7 +14,7 @@ @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 @@ -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: @@ -48,7 +50,7 @@ async def start_scan( """ url = f"{AGENTIC_SECURITY}/scan" - payload = { + payload: dict[str, Any] = { "llmSpec": llmSpec, "maxBudget": maxBudget, "datasets": [], @@ -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,7 +78,7 @@ async def stop_scan() -> dict: @mcp.tool() -async def get_data_config() -> list: +async def get_data_config() -> list[dict[str, Any]]: """ Retrieve data configuration from the FastAPI server. @@ -90,7 +92,7 @@ async def get_data_config() -> list: @mcp.tool() -async def get_spec_templates() -> list: +async def get_spec_templates() -> list[dict[str, Any]]: """ Retrieve data configuration from the FastAPI server.