Fixed issues 191 195

agentic_security/mcp/client.py

@@ -11,13 +11,13 @@ server_params = StdioServerParameters(
 )
 
 
-async def run():
+async def run()-> None:  
     async with stdio_client(server_params) as (read, write):
         async with ClientSession(read, write) as session:
-            # Initialize the connection
+            # Initialize the connection --> connection doesnt work 
             await session.initialize()
 
-            # List available prompts, resources, and tools
+            # List available prompts, resources, and tools --> no avalialbe tools 
             prompts = await session.list_prompts()
             print(f"Available prompts: {prompts}")
 
@@ -27,7 +27,7 @@ async def run():
             tools = await session.list_tools()
             print(f"Available tools: {tools}")
 
-            # Call the echo tool
+            # Call the echo tool --> echo tool iisue 
             echo_result = await session.call_tool(
                 "echo_tool", arguments={"message": "Hello from client!"}
             )

agentic_security/mcp/main.py

@@ -14,7 +14,15 @@ AGENTIC_SECURITY = "http://0.0.0.0:8718"
 
 @mcp.tool()
 async def verify_llm(spec: str) -> dict:
-    """Verify an LLM model specification using the FastAPI server."""
+    """
+    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.
+    
+    """
     url = f"{AGENTIC_SECURITY}/verify"
     async with httpx.AsyncClient() as client:
         response = await client.post(url, json={"spec": spec})
@@ -28,7 +36,18 @@ async def start_scan(
     optimize: bool = False,
     enableMultiStepAttack: bool = False,
 ) -> dict:
-    """Start an LLM security scan via the FastAPI server."""
+    """
+    Start an LLM security scan via the FastAPI server.
+    Returns:
+        dict: The scan initiation result from the FastAPI server.
+    
+    Args: 
+        llmSpec (str): The specification of the LLM model.
+        maxBudget (int): The maximum budget for the scan.
+        optimize (bool, optional): Whether to enable optimization during scanning. Defaults to False.
+        enableMultiStepAttack (bool, optional): Whether to enable multi-step attack 
+    
+    """
     url = f"{AGENTIC_SECURITY}/scan"
     payload = {
         "llmSpec": llmSpec,
@@ -46,7 +65,11 @@ async def start_scan(
 
 @mcp.tool()
 async def stop_scan() -> dict:
-    """Stop an ongoing scan via the FastAPI server."""
+    """Stop an ongoing scan via the FastAPI server.
+
+    Returns:
+        dict: The confirmation from the FastAPI server that the scan has been stopped.
+    """
     url = f"{AGENTIC_SECURITY}/stop"
     async with httpx.AsyncClient() as client:
         response = await client.post(url)
@@ -55,7 +78,12 @@ async def stop_scan() -> dict:
 
 @mcp.tool()
 async def get_data_config() -> list:
-    """Retrieve data configuration from the FastAPI server."""
+    """
+    Retrieve data configuration from the FastAPI server.
+
+    Returns:
+        list: The response from the FastAPI server, confirming the scan has been stopped.
+    """
     url = f"{AGENTIC_SECURITY}/v1/data-config"
     async with httpx.AsyncClient() as client:
         response = await client.get(url)
@@ -64,7 +92,12 @@ async def get_data_config() -> list:
 
 @mcp.tool()
 async def get_spec_templates() -> list:
-    """Retrieve data configuration from the FastAPI server."""
+    """
+    Retrieve data configuration from the FastAPI server.
+
+    Returns:
+        list: The LLM specification templates from the FastAPI server.
+    """
     url = f"{AGENTIC_SECURITY}/v1/llm-specs"
     async with httpx.AsyncClient() as client:
         response = await client.get(url)