fix(csv ds loading):

Add MseeP.ai badge

Dockerfile

@@ -19,6 +19,10 @@ RUN poetry lock
 
 # Install dependencies
 RUN poetry export -f requirements.txt --without-hashes -o requirements.txt
+
+# Install wheel (required to build packages like fire)
+RUN pip install --upgrade pip setuptools wheel
+
 RUN pip install --no-cache-dir -r requirements.txt
 
 # Runtime stage

Readme.md

@@ -21,9 +21,7 @@
   <a href="https://pypi.org/project/agentic-security/">
     <img alt="PyPI Version" src="https://img.shields.io/pypi/v/agentic-security?style=for-the-badge&logo=pypi&labelColor=000000&color=00CCFF" />
   </a>
-  <a href="https://discord.gg/stw3DfZQ">
-    <img alt="Join Discord" src="https://img.shields.io/badge/Discord-Join%20Us-black?style=for-the-badge&logo=discord&labelColor=000000&color=DD55FF" />
-  </a>
+
 </p>
 
 

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 does not 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!"}
             )
@@ -47,6 +47,7 @@ async def run():
             # print(f"Prompt result: {prompt_result}")
 
             # You can perform additional operations here as needed
+            return prompts, resources, tools
 
 
 if __name__ == "__main__":

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)

agentic_security/probe_data/data.py

@@ -245,61 +245,47 @@ def load_jailbreak_v28k() -> ProbeDataset:
         return create_probe_dataset("JailbreakV-28K/JailBreakV-28k", [])
 
 
-@cache_to_disk()
-def load_local_csv() -> ProbeDataset:
-    """Load prompts from local CSV files."""
-    csv_files = [f for f in os.listdir(".") if f.endswith(".csv")]
-    logger.info(f"Found {len(csv_files)} CSV files: {csv_files}")
-
-    prompts = []
-    for file in csv_files:
-        try:
-            df = pd.read_csv(file)
-            if "prompt" in df.columns:
-                prompts.extend(df["prompt"].tolist())
-            else:
-                logger.warning(f"File {file} lacks a suitable prompt column")
-        except Exception as e:
-            logger.error(f"Error reading {file}: {e}")
-
-    return create_probe_dataset("Local CSV", prompts, {"src": str(csv_files)})
-
-
 @cache_to_disk(1)
-def load_csv(file: str) -> ProbeDataset:
-    """Load prompts from local CSV files."""
+def file_dataset(file) -> list[str]:
     prompts = []
     try:
-        df = pd.read_csv(file)
-        prompts = df["prompt"].tolist()
+        df = pd.read_csv(os.path.join("./datasets", file), encoding_errors="ignore")
         if "prompt" in df.columns:
-            prompts.extend(df["prompt"].tolist())
+            prompts = df["prompt"].tolist()
         else:
             logger.warning(f"File {file} lacks a suitable prompt column")
     except Exception as e:
         logger.error(f"Error reading {file}: {e}")
+    return prompts
+
+
+def load_local_csv() -> ProbeDataset:
+    """Load prompts from local CSV files."""
+    os.makedirs("./datasets", exist_ok=True)
+    csv_files = [f for f in os.listdir("./datasets") if f.endswith(".csv")]
+    logger.info(f"Found {len(csv_files)} CSV files: {csv_files}")
+
+    prompts = []
+    for file in csv_files:
+        prompts.extend(file_dataset(file))
+    return create_probe_dataset("Local CSV", prompts, {"src": str(csv_files)})
+
+
+def load_csv(file: str) -> ProbeDataset:
+    """Load prompts from local CSV files."""
+    prompts = file_dataset(file)
     return create_probe_dataset(f"fs://{file}", prompts, {"src": str(file)})
 
 
-@cache_to_disk(1)
 def load_local_csv_files() -> list[ProbeDataset]:
     """Load prompts from local CSV files and return a list of ProbeDataset objects."""
-    csv_files = [f for f in os.listdir(".") if f.endswith(".csv")]
+    csv_files = [f for f in os.listdir("./datasets") if f.endswith(".csv")]
     logger.info(f"Found {len(csv_files)} CSV files: {csv_files}")
 
     datasets = []
 
     for file in csv_files:
-        try:
-            df = pd.read_csv(file)
-            if "prompt" in df.columns:
-                prompts = df["prompt"].tolist()
-                datasets.append(create_probe_dataset(file, prompts, {"src": file}))
-            else:
-                logger.warning(f"File {file} lacks a suitable prompt column")
-        except Exception as e:
-            logger.error(f"Error reading {file}: {e}")
-
+        datasets.append(create_probe_dataset(file, file_dataset(file), {"src": file}))
     return datasets
 
 
@@ -336,12 +322,20 @@ class StenographyTransformer:
                 )
                 continue
 
-            size = len(dataset.prompts)
+            # Convert async generator (if any) to a list
+            prompts = (
+                list(dataset.prompts) if hasattr(dataset.prompts, "__iter__") else []
+            )
+            size = len(prompts)
+
+            if size == 0:
+                continue
+
             sample_size = max(1, int(size * 0.25))  # At least 1 sample
 
             for name, fn in self.TRANSFORMATIONS.items():
                 logger.info(f"Applying {name} to {dataset.dataset_name}")
-                sampled = random.sample(dataset.prompts, min(sample_size, size))
+                sampled = random.sample(prompts, min(sample_size, size))
                 try:
                     transformed = [
                         f"Follow the instructions in {name}: {fn(str(p))}"

poetry.lock

@@ -149,23 +149,24 @@ files = [
 
 [[package]]
 name = "anyio"
-version = "4.4.0"
+version = "4.9.0"
 description = "High level compatibility layer for multiple asynchronous event loop implementations"
 optional = false
-python-versions = ">=3.8"
+python-versions = ">=3.9"
 files = [
-    {file = "anyio-4.4.0-py3-none-any.whl", hash = "sha256:c1b2d8f46a8a812513012e1107cb0e68c17159a7a594208005a57dc776e1bdc7"},
-    {file = "anyio-4.4.0.tar.gz", hash = "sha256:5aadc6a1bbb7cdb0bede386cac5e2940f5e2ff3aa20277e991cf028e0585ce94"},
+    {file = "anyio-4.9.0-py3-none-any.whl", hash = "sha256:9f76d541cad6e36af7beb62e978876f3b41e3e04f2c1fbf0884604c0a9c4d93c"},
+    {file = "anyio-4.9.0.tar.gz", hash = "sha256:673c0c244e15788651a4ff38710fea9675823028a6f08a5eda409e0c9840a028"},
 ]
 
 [package.dependencies]
 idna = ">=2.8"
 sniffio = ">=1.1"
+typing_extensions = {version = ">=4.5", markers = "python_version < \"3.13\""}
 
 [package.extras]
-doc = ["Sphinx (>=7)", "packaging", "sphinx-autodoc-typehints (>=1.2.0)", "sphinx-rtd-theme"]
-test = ["anyio[trio]", "coverage[toml] (>=7)", "exceptiongroup (>=1.2.0)", "hypothesis (>=4.0)", "psutil (>=5.9)", "pytest (>=7.0)", "pytest-mock (>=3.6.1)", "trustme", "uvloop (>=0.17)"]
-trio = ["trio (>=0.23)"]
+doc = ["Sphinx (>=8.2,<9.0)", "packaging", "sphinx-autodoc-typehints (>=1.2.0)", "sphinx_rtd_theme"]
+test = ["anyio[trio]", "blockbuster (>=1.5.23)", "coverage[toml] (>=7)", "exceptiongroup (>=1.2.0)", "hypothesis (>=4.0)", "psutil (>=5.9)", "pytest (>=7.0)", "trustme", "truststore (>=0.9.1)", "uvloop (>=0.21)"]
+trio = ["trio (>=0.26.1)"]
 
 [[package]]
 name = "appnope"
@@ -1162,6 +1163,17 @@ http2 = ["h2 (>=3,<5)"]
 socks = ["socksio (==1.*)"]
 zstd = ["zstandard (>=0.18.0)"]
 
+[[package]]
+name = "httpx-sse"
+version = "0.4.0"
+description = "Consume Server-Sent Event (SSE) messages with HTTPX."
+optional = false
+python-versions = ">=3.8"
+files = [
+    {file = "httpx-sse-0.4.0.tar.gz", hash = "sha256:1e81a3a3070ce322add1d3529ed42eb5f70817f45ed6ec915ab753f961139721"},
+    {file = "httpx_sse-0.4.0-py3-none-any.whl", hash = "sha256:f329af6eae57eaa2bdfd962b42524764af68075ea87370a2de920af5341e318f"},
+]
+
 [[package]]
 name = "huggingface-hub"
 version = "0.29.3"
@@ -1796,6 +1808,32 @@ files = [
 [package.dependencies]
 traitlets = "*"
 
+[[package]]
+name = "mcp"
+version = "1.6.0"
+description = "Model Context Protocol SDK"
+optional = false
+python-versions = ">=3.10"
+files = [
+    {file = "mcp-1.6.0-py3-none-any.whl", hash = "sha256:7bd24c6ea042dbec44c754f100984d186620d8b841ec30f1b19eda9b93a634d0"},
+    {file = "mcp-1.6.0.tar.gz", hash = "sha256:d9324876de2c5637369f43161cd71eebfd803df5a95e46225cab8d280e366723"},
+]
+
+[package.dependencies]
+anyio = ">=4.5"
+httpx = ">=0.27"
+httpx-sse = ">=0.4"
+pydantic = ">=2.7.2,<3.0.0"
+pydantic-settings = ">=2.5.2"
+sse-starlette = ">=1.6.1"
+starlette = ">=0.27"
+uvicorn = ">=0.23.1"
+
+[package.extras]
+cli = ["python-dotenv (>=1.0.0)", "typer (>=0.12.4)"]
+rich = ["rich (>=13.9.4)"]
+ws = ["websockets (>=15.0.1)"]
+
 [[package]]
 name = "mdit-py-plugins"
 version = "0.4.2"
@@ -3116,6 +3154,26 @@ files = [
 [package.dependencies]
 typing-extensions = ">=4.6.0,<4.7.0 || >4.7.0"
 
+[[package]]
+name = "pydantic-settings"
+version = "2.8.1"
+description = "Settings management using Pydantic"
+optional = false
+python-versions = ">=3.8"
+files = [
+    {file = "pydantic_settings-2.8.1-py3-none-any.whl", hash = "sha256:81942d5ac3d905f7f3ee1a70df5dfb62d5569c12f51a5a647defc1c3d9ee2e9c"},
+    {file = "pydantic_settings-2.8.1.tar.gz", hash = "sha256:d5c663dfbe9db9d5e1c646b2e161da12f0d734d422ee56f567d0ea2cee4e8585"},
+]
+
+[package.dependencies]
+pydantic = ">=2.7.0"
+python-dotenv = ">=0.21.0"
+
+[package.extras]
+azure-key-vault = ["azure-identity (>=1.16.0)", "azure-keyvault-secrets (>=4.8.0)"]
+toml = ["tomli (>=2.0.1)"]
+yaml = ["pyyaml (>=6.0.1)"]
+
 [[package]]
 name = "pyfiglet"
 version = "1.0.2"
@@ -3280,6 +3338,20 @@ files = [
 [package.dependencies]
 six = ">=1.5"
 
+[[package]]
+name = "python-dotenv"
+version = "1.1.0"
+description = "Read key-value pairs from a .env file and set them as environment variables"
+optional = false
+python-versions = ">=3.9"
+files = [
+    {file = "python_dotenv-1.1.0-py3-none-any.whl", hash = "sha256:d7c01d9e2293916c18baf562d95698754b0dbbb5e74d457c45d4f6561fb9d55d"},
+    {file = "python_dotenv-1.1.0.tar.gz", hash = "sha256:41f90bc6f5f177fb41f53e87666db362025010eb28f60a01c9143bfa33a2b2d5"},
+]
+
+[package.extras]
+cli = ["click (>=5.0)"]
+
 [[package]]
 name = "python-multipart"
 version = "0.0.20"
@@ -3909,6 +3981,25 @@ files = [
     {file = "soupsieve-2.6.tar.gz", hash = "sha256:e2e68417777af359ec65daac1057404a3c8a5455bb8abc36f1a9866ab1a51abb"},
 ]
 
+[[package]]
+name = "sse-starlette"
+version = "2.1.3"
+description = "SSE plugin for Starlette"
+optional = false
+python-versions = ">=3.8"
+files = [
+    {file = "sse_starlette-2.1.3-py3-none-any.whl", hash = "sha256:8ec846438b4665b9e8c560fcdea6bc8081a3abf7942faa95e5a744999d219772"},
+    {file = "sse_starlette-2.1.3.tar.gz", hash = "sha256:9cd27eb35319e1414e3d2558ee7414487f9529ce3b3cf9b21434fd110e017169"},
+]
+
+[package.dependencies]
+anyio = "*"
+starlette = "*"
+uvicorn = "*"
+
+[package.extras]
+examples = ["fastapi"]
+
 [[package]]
 name = "stack-data"
 version = "0.6.3"
@@ -4469,4 +4560,4 @@ propcache = ">=0.2.0"
 [metadata]
 lock-version = "2.0"
 python-versions = "^3.11"
-content-hash = "04c8dc31a9c1a1b2bbb607041a31fa291ed6cdf0b060d30161786da7e8fbab4e"
+content-hash = "f8777299520e96d3c5d14f5a216b36ed869a6429434f1366761cf1cdfe07ef52"

pyproject.toml

@@ -1,6 +1,6 @@
 [tool.poetry]
 name = "agentic_security"
-version = "0.7.1"
+version = "0.7.3"
 description = "Agentic LLM vulnerability scanner"
 authors = ["Alexander Miasoiedov <msoedov@gmail.com>"]
 maintainers = ["Alexander Miasoiedov <msoedov@gmail.com>"]
@@ -52,6 +52,7 @@ sentry_sdk = "^2.22.0"
 orjson = "^3.10"
 pyfiglet = "^1.0.2"
 termcolor = "^2.4.0"
+mcp = "^1.4.1"
 
 # garak = { version = "*", optional = true }
 pytest-xdist = "3.6.1"

tests/test_mcp.py

@@ -0,0 +1,12 @@
+import pytest
+
+from agentic_security.mcp.client import run
+
+
+@pytest.mark.asyncio
+async def test_mcp_echo_tool():
+    """Test the echo tool functionality"""
+    prompts, resources, tools = await run()
+    assert prompts
+    assert resources
+    assert tools

ui/package-lock.json

@@ -6891,9 +6891,9 @@
       }
     },
     "node_modules/http-proxy-middleware": {
-      "version": "2.0.7",
-      "resolved": "https://registry.npmjs.org/http-proxy-middleware/-/http-proxy-middleware-2.0.7.tgz",
-      "integrity": "sha512-fgVY8AV7qU7z/MmXJ/rxwbrtQH4jBQ9m7kp3llF0liB7glmFeVZFBepQb32T3y8n8k2+AEYuMPCpinYW+/CuRA==",
+      "version": "2.0.9",
+      "resolved": "https://registry.npmjs.org/http-proxy-middleware/-/http-proxy-middleware-2.0.9.tgz",
+      "integrity": "sha512-c1IyJYLYppU574+YI7R4QyX2ystMtVXZwIdzazUIPIJsHuWNd+mho2j+bKoHftndicGj9yh+xjd+l0yj7VeT1Q==",
       "dev": true,
       "license": "MIT",
       "dependencies": {