Adding codebeaver.yml

.github/workflows/release.yml

@@ -9,7 +9,7 @@ on:
       - 0.*
 
 env:
-  POETRY_VERSION: "1.8.5"
+  POETRY_VERSION: "1.7.1"
 
 jobs:
   if_release:

.github/workflows/test.yml

@@ -7,7 +7,7 @@ on:
     branches: [main]
 
 env:
-  POETRY_VERSION: "1.8.5"
+  POETRY_VERSION: "1.7.1"
   OPENAI_API_KEY: "sk-fake"
 
 jobs:

.gitignore

@@ -17,5 +17,4 @@ inv/
 scripts/
 docx/
 agentic_security.toml
-/venv
-*.csv
+/venv

Dockerfile

@@ -1,5 +1,5 @@
 # Build stage
-FROM python:3.11-slim AS builder
+FROM python:3.11-slim as builder
 
 WORKDIR /app
 
@@ -14,9 +14,6 @@ RUN poetry self add "poetry-plugin-export"
 # Copy only dependency files to leverage Docker layer caching
 COPY pyproject.toml poetry.lock ./
 
-# update lock file to avoid failure
-RUN poetry lock
-
 # Install dependencies
 RUN poetry export -f requirements.txt --without-hashes -o requirements.txt
 RUN pip install --no-cache-dir -r requirements.txt

Readme.md

@@ -83,7 +83,8 @@ agentic_security --port=PORT --host=HOST
 
 ## UI 🧙
 
-<img width="100%" alt="booking-screen" src="https://raw.githubusercontent.com/msoedov/agentic_security/refs/heads/main/docs/images/demo.gif">
+<img width="100%" alt="booking-screen" src="https://res.cloudinary.com/dq0w2rtm9/image/upload/v1736433557/z0bsyzhsqlgcr3w4ovwp.gif">
+<img width="100%" alt="booking-screen" src="https://res.cloudinary.com/dq0w2rtm9/image/upload/v1741192668/final_aa9jhb.gif">
 
 ## LLM kwargs
 
@@ -402,16 +403,6 @@ This setup ensures a continuous integration approach towards maintaining securit
 
 The `Module` class is designed to manage prompt processing and interaction with external AI models and tools. It supports fetching, processing, and posting prompts asynchronously for model vulnerabilities. Check out [module.md](https://github.com/msoedov/agentic_security/blob/main/docs/module.md) for details.
 
-
-## MCP server
-
-```shell
-pip install -U mcp
-
-# From cloned directory
-mcp install agentic_security/mcp/main.py
-```
-
 ## Documentation
 
 For more detailed information on how to use Agentic Security, including advanced features and customization options, please refer to the official documentation.
@@ -438,7 +429,6 @@ We’re just getting started! Here’s what’s on the horizon:
 
 Note: All dates are tentative and subject to change based on project progress and priorities.
 
-
 ## 👋 Contributing
 
 Contributions to Agentic Security are welcome! If you'd like to contribute, please follow these steps:

agentic_security/__init__.py

@@ -1,3 +1,3 @@
-from .lib import SecurityScanner
+from .lib import AgenticSecurity
 
-__all__ = ["SecurityScanner"]
+__all__ = ["AgenticSecurity"]

agentic_security/__main__.py

@@ -5,7 +5,7 @@ import fire
 import uvicorn
 
 from agentic_security.app import app
-from agentic_security.lib import SecurityScanner
+from agentic_security.lib import AgenticSecurity
 from agentic_security.misc.banner import init_banner
 
 
@@ -32,14 +32,14 @@ class CLI:
         Run Agentic Security in CI mode.
         """
         sys.path.append(os.path.dirname("."))
-        SecurityScanner().entrypoint()
+        AgenticSecurity().entrypoint()
 
     def init(self, host: str = "0.0.0.0", port: int = 8718):
         """
         Generate the default CI configuration file.
         """
         sys.path.append(os.path.dirname("."))
-        SecurityScanner().generate_default_settings(host, port)
+        AgenticSecurity().generate_default_settings(host, port)
 
     i = init
 
@@ -48,7 +48,7 @@ class CLI:
         List all available security checks.
         """
         sys.path.append(os.path.dirname("."))
-        SecurityScanner().list_checks()
+        AgenticSecurity().list_checks()
 
 
 def main():

agentic_security/agents/operator_crew.py

@@ -1,4 +1,5 @@
 import asyncio
+import logging
 import os
 from typing import Any
 
@@ -9,11 +10,12 @@ from pydantic import BaseModel, ConfigDict, Field
 
 # Assuming LLMSpec is defined elsewhere; placeholder import
 from agentic_security.http_spec import LLMSpec
-from agentic_security.logutils import logger
 
 LLM_SPECS = []  # Populate with LLM spec strings if needed
 
 # Configure logging
+logging.basicConfig(level=logging.INFO)
+logger = logging.getLogger(__name__)
 
 
 # Define AgentSpecification model

agentic_security/agents/operator_pydantic.py

@@ -1,4 +1,5 @@
 import asyncio
+import logging
 from typing import Any
 
 import httpx
@@ -7,10 +8,13 @@ from pydantic_ai import Agent, RunContext, Tool
 
 # Assuming LLMSpec is defined elsewhere; placeholder import
 from agentic_security.http_spec import LLMSpec
-from agentic_security.logutils import logger
 
 LLM_SPECS = []  # Populate this list with LLM spec strings if needed
 
+# Configure logging
+logging.basicConfig(level=logging.INFO)
+logger = logging.getLogger(__name__)
+
 
 # Define AgentSpecification model
 class AgentSpecification(BaseModel):

agentic_security/config.py

@@ -1,10 +1,9 @@
 from functools import lru_cache
 
 import tomli
+from loguru import logger
 
-from agentic_security.logutils import logger
-
-SETTINGS_VERSION = 2
+SETTINGS_VERSION = 1
 
 
 @lru_cache(maxsize=1)
@@ -143,13 +142,6 @@ use_disk_cache = false
 retry = 3
 timeout_connect = 30
 timeout_response = 90
-
-[fuzzer]
-max_prompt_lenght = 2048
-budget_multiplier = 100000000
-initial_optimizer_points = 25
-min_failure_samples = 5
-failure_rate_threshold = 0.5
 """.replace(
                     "$HOST", host
                 )

agentic_security/core/app.py

@@ -4,14 +4,10 @@ from asyncio import Event, Queue
 from fastapi import FastAPI
 from fastapi.responses import ORJSONResponse
 
-from agentic_security.http_spec import LLMSpec
-
 tools_inbox: Queue = Queue()
 stop_event: Event = Event()
 current_run: str = {"spec": "", "id": ""}
-_secrets: dict[str, str] = {}
-
-current_run: dict[str, int | LLMSpec] = {"spec": "", "id": ""}
+_secrets = {}
 
 
 def create_app() -> FastAPI:
@@ -30,29 +26,29 @@ def get_stop_event() -> Event:
     return stop_event
 
 
-def get_current_run() -> dict[str, int | LLMSpec]:
+def get_current_run() -> str:
     """Get the current run id."""
     return current_run
 
 
-def set_current_run(spec: LLMSpec) -> dict[str, int | LLMSpec]:
+def set_current_run(spec):
     """Set the current run id."""
     current_run["id"] = hash(id(spec))
     current_run["spec"] = spec
     return current_run
 
 
-def get_secrets() -> dict[str, str]:
+def get_secrets():
     return _secrets
 
 
-def set_secrets(secrets: dict[str, str]) -> dict[str, str]:
+def set_secrets(secrets):
     _secrets.update(secrets)
     expand_secrets(_secrets)
     return _secrets
 
 
-def expand_secrets(secrets: dict[str, str]) -> None:
+def expand_secrets(secrets):
     for key in secrets:
         val = secrets[key]
         if val.startswith("$"):

agentic_security/core/logging.py

@@ -1,5 +1,26 @@
-from agentic_security.logutils import set_log_level_to_info
+from logging import config
 
 
 def setup_logging():
-    return set_log_level_to_info()
+    config.dictConfig(
+        {
+            "version": 1,
+            "disable_existing_loggers": True,
+            "handlers": {
+                "console": {
+                    "class": "logging.StreamHandler",
+                },
+            },
+            "root": {
+                "handlers": ["console"],
+                "level": "INFO",
+            },
+            "loggers": {
+                "uvicorn.access": {
+                    "level": "ERROR",  # Set higher log level to suppress info logs globally
+                    "handlers": ["console"],
+                    "propagate": False,
+                }
+            },
+        }
+    )

agentic_security/core/test_app.py

@@ -6,12 +6,30 @@ from agentic_security.core.app import expand_secrets
 
 
 @pytest.fixture(autouse=True)
+def reset_globals():
+    """
+    Reset globals (_secrets, current_run, tools_inbox, stop_event) before each test.
+    This ensures tests run in a clean state.
+    """
+    from agentic_security.core.app import _secrets, current_run, get_tools_inbox, get_stop_event
+    _secrets.clear()
+    current_run["spec"] = ""
+    current_run["id"] = ""
+    # Clear tools_inbox queue
+    queue = get_tools_inbox()
+    while not queue.empty():
+        queue.get_nowait()
+    # Reset stop_event if it is set
+    event = get_stop_event()
+    if event.is_set():
+        event.clear()
 def setup_env_vars():
     # Set up environment variables for testing
     os.environ["TEST_ENV_VAR"] = "test_value"
 
 
 def test_expand_secrets_with_env_var():
+    os.environ["TEST_ENV_VAR"] = "test_value"
     secrets = {"secret_key": "$TEST_ENV_VAR"}
     expand_secrets(secrets)
     assert secrets["secret_key"] == "test_value"
@@ -27,3 +45,180 @@ def test_expand_secrets_without_dollar_sign():
     secrets = {"secret_key": "plain_value"}
     expand_secrets(secrets)
     assert secrets["secret_key"] == "plain_value"
+
+import asyncio
+from fastapi import FastAPI
+from fastapi.responses import ORJSONResponse
+from agentic_security.core.app import create_app, get_tools_inbox, get_stop_event, get_current_run, set_current_run, get_secrets, set_secrets, expand_secrets
+
+class DummyLLMSpec:
+    """A dummy LLMSpec for testing purposes."""
+    pass
+
+def test_create_app():
+    """Test that create_app returns a FastAPI app with ORJSONResponse."""
+    app = create_app()
+    assert isinstance(app, FastAPI)
+    assert app.router.default_response_class == ORJSONResponse
+
+def test_get_tools_inbox():
+    """Test that get_tools_inbox returns a Queue instance."""
+    queue = get_tools_inbox()
+    from asyncio import Queue
+    assert isinstance(queue, Queue)
+
+def test_get_stop_event():
+    """Test that get_stop_event returns an Event instance."""
+    event = get_stop_event()
+    from asyncio import Event
+    assert isinstance(event, Event)
+
+def test_get_current_run_initial():
+    """Test that get_current_run returns the initial current run dictionary."""
+    current = get_current_run()
+    # The initial dictionary should have an empty spec and id.
+    assert current["spec"] == ""
+    assert current["id"] == ""
+
+def test_set_current_run():
+    """Test that set_current_run updates the current run with the dummy LLMSpec."""
+    dummy_spec = DummyLLMSpec()
+    updated = set_current_run(dummy_spec)
+    assert updated["spec"] is dummy_spec
+    # Ensure that the id is computed as hash(id(dummy_spec))
+    expected_id = hash(id(dummy_spec))
+    assert updated["id"] == expected_id
+
+def test_get_and_set_secrets():
+    """Test that set_secrets updates the secrets dictionary and get_secrets returns the updated values."""
+    # Clear any previously set secrets
+    secrets_before = get_secrets().copy()
+    os.environ["MY_SECRET"] = "secret_value"
+    new_secrets = {"key1": "$MY_SECRET", "key2": "plain"}
+    updated = set_secrets(new_secrets)
+    assert updated["key1"] == "secret_value"
+    assert updated["key2"] == "plain"
+
+def test_expand_secrets_multiple_keys():
+    """Test expand_secrets with multiple keys, including one with an environment variable,
+    one with a non-existent variable, and one that is plain."""
+    os.environ["TEST_ENV_VAR"] = "test_value"
+    secrets = {"env_key": "$TEST_ENV_VAR", "nonexistent_key": "$NON_EXISTENT", "plain_key": "value"}
+    expand_secrets(secrets)
+    assert secrets["env_key"] == "test_value"
+    # For a non-existent environment variable, os.getenv returns None
+    assert secrets["nonexistent_key"] is None
+    # Plain values should not be changed.
+    assert secrets["plain_key"] == "value"
+def test_expand_secrets_with_space_after_dollar():
+    """Test expand_secrets when the value has a dollar sign followed by a space.
+    Since the value does not start strictly with "$", the secret remains unchanged.
+    Also verifies that the stripping in expand_secrets (via strip("$"))
+    will remove both dollar and any whitespace if the value actually started with '$'.
+    """
+    os.environ["SPACED_VAR"] = "spaced_value"
+    secrets = {"key": "$ SPACED_VAR"}
+    expand_secrets(secrets)
+    # " $ SPACED_VAR" after strip("$") becomes " SPACED_VAR" which is not a valid env key so returns None.
+    assert secrets["key"] is None
+
+def test_set_secrets_update_existing():
+    """Test that set_secrets updates an existing secret and retains previously set keys."""
+    os.environ["VAR1"] = "value1"
+    os.environ["VAR2"] = "value2"
+    result_first = set_secrets({"a": "$VAR1", "b": "b_val"})
+    assert result_first["a"] == "value1"
+    # Change VAR1 in environment and update secret "a", and add secret "c"
+    os.environ["VAR1"] = "new_value1"
+    result_second = set_secrets({"a": "$VAR1", "c": "$VAR2"})
+    assert result_second["a"] == "new_value1"
+    assert result_second["b"] == "b_val"
+    assert result_second["c"] == "value2"
+
+def test_tools_inbox_state():
+    """Test that get_tools_inbox returns the same queue instance 
+    and that the queue state persists across multiple calls.
+    """
+    from asyncio import Queue
+    inbox1 = get_tools_inbox()
+    inbox1.put_nowait("message")
+    inbox2 = get_tools_inbox()
+    # inbox2 should contain the "message" from inbox1
+    msg = inbox2.get_nowait()
+    assert msg == "message"
+
+def test_stop_event_state():
+    """Test that stop_event can be set and cleared, and its state persists."""
+    event = get_stop_event()
+    # Initially the event should not be set
+    assert not event.is_set()
+    event.set()
+    assert event.is_set()
+    event.clear()
+    assert not event.is_set()
+
+def test_set_current_run_returns_global_dict():
+    """Test that set_current_run returns the same global current_run dictionary
+    as returned by get_current_run.
+    """
+    dummy_spec = DummyLLMSpec()
+    updated = set_current_run(dummy_spec)
+    current = get_current_run()
+    assert updated is current
+def test_get_secrets_initial():
+    """Test that get_secrets returns an empty dictionary initially."""
+    assert get_secrets() == {}
+
+def test_set_secrets_empty():
+    """Test that setting an empty secrets dictionary does not modify existing secrets."""
+    # first set initial secrets
+    initial = {"key": "value"}
+    set_secrets(initial)
+    # update with an empty dict – the existing keys remain
+    result = set_secrets({})
+    assert result == initial
+
+def test_update_current_run_twice():
+    """Test updating current run twice with different LLMSpec values."""
+    dummy1 = DummyLLMSpec()
+    dummy2 = DummyLLMSpec()
+    set_current_run(dummy1)
+    first = get_current_run().copy()
+    set_current_run(dummy2)
+    second = get_current_run().copy()
+    # first update should hold dummy1, second should hold dummy2
+    assert first["spec"] is dummy1
+    assert second["spec"] is dummy2
+    # Ensure that id has changed (using hash(id(dummy_spec)))
+    assert first["id"] != second["id"]
+
+def test_expand_secrets_trailing_whitespace():
+    """Test expand_secrets when the secret value has trailing whitespace after the dollar sign.
+    The trailing whitespace remains after stripping only the dollar sign, so the looked-up environment variable key will not match.
+    """
+    os.environ["TRIM_TEST"] = "trimmed"
+    secrets = {"key": "$TRIM_TEST "}
+    expand_secrets(secrets)
+    # Since "TRIM_TEST " (with trailing space) is not set in the environment, the secret should be None.
+    assert secrets["key"] is None
+def test_expand_secrets_empty_dict():
+    """Test expand_secrets with an empty dictionary does nothing."""
+    secrets = {}
+    expand_secrets(secrets)
+    assert secrets == {}
+
+def test_expand_secrets_with_non_string_value():
+    """Test that expand_secrets raises an AttributeError when a secret value is not a string."""
+    secrets = {"key": 123}
+    with pytest.raises(AttributeError):
+        expand_secrets(secrets)
+
+def test_expand_secrets_multiple_dollar_signs():
+    """Test expand_secrets with a value that contains multiple leading dollar signs.
+    The extra dollar signs are removed by the strip method.
+    """
+    os.environ["MULTI_DOLLAR_VAR"] = "multi_value"
+    secrets = {"key": "$$MULTI_DOLLAR_VAR"}
+    expand_secrets(secrets)
+    # After stripping, "$$MULTI_DOLLAR_VAR".strip("$") returns "MULTI_DOLLAR_VAR"
+    assert secrets["key"] == "multi_value"

agentic_security/http_spec.py

@@ -1,6 +1,5 @@
 import base64
 from enum import Enum
-from urllib.parse import urlparse
 
 import httpx
 from pydantic import BaseModel
@@ -160,14 +159,6 @@ def parse_http_spec(http_spec: str) -> LLMSpec:
     # Extract the method and URL from the first line
     method, url = lines[0].split(" ")[0:2]
 
-    # Check url validity
-    valid_url = urlparse(url)
-    # if missing the correct formatting ://, urlparse.netloc will be empty
-    if valid_url.scheme not in ("http", "https") or not valid_url.netloc:
-        raise InvalidHTTPSpecError(
-            f"Invalid URL: {url}. Ensure it starts with 'http://' or 'https://'"
-        )
-
     # Initialize headers and body
     headers = {}
     body = ""

agentic_security/lib.py

@@ -4,12 +4,12 @@ from datetime import datetime
 
 import colorama
 import tqdm.asyncio
+from loguru import logger
 from rich.console import Console
 from rich.table import Table
 from tabulate import tabulate
 
 from agentic_security.config import SettingsMixin  # Importing the configuration mixin
-from agentic_security.logutils import logger
 from agentic_security.primitives import Scan
 from agentic_security.probe_data import REGISTRY
 from agentic_security.routes.scan import streaming_response_generator
@@ -23,7 +23,7 @@ YELLOW = colorama.Fore.YELLOW
 BLUE = colorama.Fore.BLUE
 
 
-class SecurityScanner(SettingsMixin):
+class AgenticSecurity(SettingsMixin):
     @classmethod
     async def async_scan(
         cls,
@@ -202,7 +202,7 @@ class SecurityScanner(SettingsMixin):
         self.load_config(self.default_path)
         logger.info("Configuration loaded successfully.")
         print(self.config)
-        datasets = list(self.get_config_value("modules", {}).values())
+        datasets = list(self.get_config_value("modules").values())
         for d in datasets:
             d["selected"] = True
         self.scan(

agentic_security/logutils.py

@@ -1,146 +0,0 @@
-# import sys
-
-# from loguru import logger
-
-# # Define custom colors
-# BLUE = "#89CFF0"
-# BROWN = "#8B4513"  # Brown for DEBUG
-
-# # Define custom log level colors
-# logger.level("DEBUG", color=f"<fg {BROWN}>")
-# logger.level("INFO", color=f"<fg {BLUE}>")
-
-# # Define custom log format with aligned messages and colored levels
-# LOG_FORMAT = (
-#     "<level>{level:<8}</level> "  # Properly formatted and colored log level
-#     "<level>{message:<100}</level> "  # Left-aligned message for readability
-#     "<cyan>{file.name}</cyan>:<cyan>{line}</cyan>"  # File name and line number in cyan
-# )
-
-# # Remove default handlers and add a new one with custom formatting
-# logger.remove()
-# logger.add(sys.stdout, format=LOG_FORMAT, level="DEBUG", colorize=True)
-import logging
-import logging.config
-import time
-from collections.abc import Callable, Coroutine
-from functools import wraps
-from os import getenv
-from typing import Any, ParamSpec, TypeVar
-
-LOGGER_NAME = None
-
-LOGGING_CONFIG = {
-    "version": 1,
-    "disable_existing_loggers": False,
-    "formatters": {
-        "rich": {"format": "%(message)s", "datefmt": "[%X]"},
-    },
-    "handlers": {
-        "rich": {
-            "class": "rich.logging.RichHandler",
-            "level": "INFO",
-            "formatter": "rich",
-            "show_time": False,
-            "rich_tracebacks": False,
-            "show_path": lambda: True if getenv("API_RUNTIME") == "dev" else False,
-            "tracebacks_show_locals": False,
-        },
-    },
-    "loggers": {
-        "": {  # Root logger configuration
-            "level": "INFO",
-            "handlers": ["rich"],
-            "propagate": True,
-        },
-        "httpx": {  # Disable httpx logging
-            "level": "WARNING",  # Suppress DEBUG and INFO messages from httpx
-            "handlers": [],
-            "propagate": False,
-        },
-        "uvicorn.access": {  # Disable uvicorn.access logging
-            "level": "WARNING",  # Suppress DEBUG and INFO messages from uvicorn.access
-            "handlers": [],
-            "propagate": False,
-        },
-    },
-}
-
-
-def configure_logging():
-    # Apply the dictionary configuration
-    logging.config.dictConfig(LOGGING_CONFIG)
-
-    # Get and return the logger
-    logger = logging.getLogger(LOGGER_NAME)
-    return logger
-
-
-logger: logging.Logger = configure_logging()
-
-
-def set_log_level_to_debug():
-    logger = logging.getLogger(LOGGER_NAME)
-    logger.setLevel(logging.DEBUG)
-    # Update handler level as well
-    for handler in logger.handlers:
-        handler.setLevel(logging.DEBUG)
-
-
-def set_log_level_to_info():
-    logger = logging.getLogger(LOGGER_NAME)
-    logger.setLevel(logging.INFO)
-    # Update handler level as well
-    for handler in logger.handlers:
-        handler.setLevel(logging.INFO)
-
-
-# Set initial log level
-set_log_level_to_info()
-
-
-# Define generic type variables for return type and parameters
-R = TypeVar("R")
-P = ParamSpec("P")
-
-
-def time_execution_sync(
-    additional_text: str = "",
-) -> Callable[[Callable[P, R]], Callable[P, R]]:
-    def decorator(func: Callable[P, R]) -> Callable[P, R]:
-        @wraps(func)
-        def wrapper(*args: P.args, **kwargs: P.kwargs) -> R:
-            start_time = time.time()
-            result = func(*args, **kwargs)
-            execution_time = time.time() - start_time
-            logger.debug(
-                f"{additional_text} Execution time: {execution_time:.2f} seconds"
-            )
-            return result
-
-        return wrapper
-
-    return decorator
-
-
-def time_execution_async(
-    additional_text: str = "",
-) -> Callable[
-    [Callable[P, Coroutine[Any, Any, R]]], Callable[P, Coroutine[Any, Any, R]]
-]:
-    def decorator(
-        func: Callable[P, Coroutine[Any, Any, R]]
-    ) -> Callable[P, Coroutine[Any, Any, R]]:
-        @wraps(func)
-        async def wrapper(*args: P.args, **kwargs: P.kwargs) -> R:
-            start_time = time.time()
-            result = await func(*args, **kwargs)
-            execution_time = time.time() - start_time
-            logger.debug(
-                f"{additional_text} Execution time: {execution_time:.2f} seconds"
-            )
-            return result
-
-        return wrapper
-
-    return decorator

agentic_security/mcp/client.py

@@ -1,53 +0,0 @@
-import asyncio
-
-from mcp import ClientSession, StdioServerParameters
-from mcp.client.stdio import stdio_client
-
-# Create server parameters for stdio connection
-server_params = StdioServerParameters(
-    command="python",  # Executable
-    args=["agentic_security/mcp/main.py"],  # Your server script
-    env=None,  # Optional environment variables
-)
-
-
-async def run():
-    async with stdio_client(server_params) as (read, write):
-        async with ClientSession(read, write) as session:
-            # Initialize the connection
-            await session.initialize()
-
-            # List available prompts, resources, and tools
-            prompts = await session.list_prompts()
-            print(f"Available prompts: {prompts}")
-
-            resources = await session.list_resources()
-            print(f"Available resources: {resources}")
-
-            tools = await session.list_tools()
-            print(f"Available tools: {tools}")
-
-            # Call the echo tool
-            echo_result = await session.call_tool(
-                "echo_tool", arguments={"message": "Hello from client!"}
-            )
-            print(f"Tool result: {echo_result}")
-
-            # # Read the echo resource
-            # echo_content, mime_type = await session.read_resource(
-            #     "echo://Hello_resource"
-            # )
-            # print(f"Resource content: {echo_content}")
-            # print(f"Resource MIME type: {mime_type}")
-
-            # # Get and use the echo prompt
-            # prompt_result = await session.get_prompt(
-            #     "echo_prompt", arguments={"message": "Hello prompt!"}
-            # )
-            # print(f"Prompt result: {prompt_result}")
-
-            # You can perform additional operations here as needed
-
-
-if __name__ == "__main__":
-    asyncio.run(run())

agentic_security/mcp/main.py

@@ -1,76 +0,0 @@
-import httpx
-from mcp.server.fastmcp import FastMCP
-
-# Initialize MCP server
-mcp = FastMCP(
-    name="Agentic Security MCP Server",
-    description="MCP server to interact with LLM scanning test",
-    dependencies=["httpx"],
-)
-
-# FastAPI Server Configuration
-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."""
-    url = f"{AGENTIC_SECURITY}/verify"
-    async with httpx.AsyncClient() as client:
-        response = await client.post(url, json={"spec": spec})
-        return response.json()
-
-
-@mcp.tool()
-async def start_scan(
-    llmSpec: str,
-    maxBudget: int,
-    optimize: bool = False,
-    enableMultiStepAttack: bool = False,
-) -> dict:
-    """Start an LLM security scan via the FastAPI server."""
-    url = f"{AGENTIC_SECURITY}/scan"
-    payload = {
-        "llmSpec": llmSpec,
-        "maxBudget": maxBudget,
-        "datasets": [],
-        "optimize": optimize,
-        "enableMultiStepAttack": enableMultiStepAttack,
-        "probe_datasets": [],
-        "secrets": {},
-    }
-    async with httpx.AsyncClient() as client:
-        response = await client.post(url, json=payload)
-        return response.json()
-
-
-@mcp.tool()
-async def stop_scan() -> dict:
-    """Stop an ongoing scan via the FastAPI server."""
-    url = f"{AGENTIC_SECURITY}/stop"
-    async with httpx.AsyncClient() as client:
-        response = await client.post(url)
-        return response.json()
-
-
-@mcp.tool()
-async def get_data_config() -> list:
-    """Retrieve data configuration from the FastAPI server."""
-    url = f"{AGENTIC_SECURITY}/v1/data-config"
-    async with httpx.AsyncClient() as client:
-        response = await client.get(url)
-        return response.json()
-
-
-@mcp.tool()
-async def get_spec_templates() -> list:
-    """Retrieve data configuration from the FastAPI server."""
-    url = f"{AGENTIC_SECURITY}/v1/llm-specs"
-    async with httpx.AsyncClient() as client:
-        response = await client.get(url)
-        return response.json()
-
-
-# Run the MCP server
-if __name__ == "__main__":
-    mcp.run()

agentic_security/middleware/logging.py

@@ -1,8 +1,7 @@
 from fastapi import Request
+from loguru import logger
 from starlette.middleware.base import BaseHTTPMiddleware
 
-from agentic_security.logutils import logger
-
 
 class LogNon200ResponsesMiddleware(BaseHTTPMiddleware):
     async def dispatch(self, request: Request, call_next):

agentic_security/misc/banner.py

@@ -8,14 +8,14 @@ except ImportError:
 
 
 def generate_banner(
-    title: str = "Agentic Security",
-    font: str = "slant",
-    version: str = "v2.1.0",
-    tagline: str = "Proactive Threat Detection & Automated Security Protocols",
-    author: str = "Developed by: [Security Team]",
-    website: str = "Website: https://github.com/msoedov/agentic_security",
-    warning: str | None = "",  # Using Optional for warning since it might be None
-) -> str:
+    title="Agentic Security",
+    font="slant",
+    version="v2.1.0",
+    tagline="Proactive Threat Detection & Automated Security Protocols",
+    author="Developed by: [Security Team]",
+    website="Website: https://github.com/msoedov/agentic_security",
+    warning="",
+):
     """Generate a visually enhanced banner with dynamic width and borders."""
     # Define the text elements
 
@@ -80,7 +80,6 @@ def generate_banner(
 
 
 def init_banner():
-    return
     ver = version("agentic_security")
     try:
         print(generate_banner(version=ver))

agentic_security/primitives/__init__.py

@@ -1,6 +1,5 @@
-# noqa
-from agentic_security.primitives.models import CompletionRequest  # noqa
 from agentic_security.primitives.models import (  # noqa
+    CompletionRequest,
     FileProbeResponse,
     LLMInfo,
     Message,

agentic_security/probe_actor/fuzzer.py

@@ -3,44 +3,32 @@ import random
 import time
 from collections.abc import AsyncGenerator
 from json import JSONDecodeError
-from typing import Any
 
 import httpx
+import pandas as pd
+from loguru import logger
 from skopt import Optimizer
 from skopt.space import Real
 
-from agentic_security.config import settings_var
 from agentic_security.http_spec import Modality
-from agentic_security.logutils import logger
 from agentic_security.primitives import Scan, ScanResult
 from agentic_security.probe_actor.cost_module import calculate_cost
 from agentic_security.probe_actor.refusal import refusal_heuristic
-from agentic_security.probe_actor.state import FuzzerState
 from agentic_security.probe_data import audio_generator, image_generator, msj_data
 from agentic_security.probe_data.data import prepare_prompts
 
-MAX_PROMPT_LENGTH = settings_var("fuzzer.max_prompt_lenght", 2048)
-BUDGET_MULTIPLIER = settings_var("fuzzer.budget_multiplier", 100000000)
-INITIAL_OPTIMIZER_POINTS = settings_var("fuzzer.initial_optimizer_points", 25)
-MIN_FAILURE_SAMPLES = settings_var("min_failure_samples", 5)
-FAILURE_RATE_THRESHOLD = settings_var("failure_rate_threshold", 0.5)
+# TODO: full log file
+
+MAX_PROMPT_LENGTH = 2048
+BUDGET_MULTIPLIER = 100_000_000
+INITIAL_OPTIMIZER_POINTS = 25
+MIN_FAILURE_SAMPLES = 5
+FAILURE_RATE_THRESHOLD = 0.5
 
 
 async def generate_prompts(
     prompts: list[str] | AsyncGenerator,
 ) -> AsyncGenerator[str, None]:
-    """
-    Asynchronously generates and yields individual prompts.
-
-    If the input is a list of strings, the function sequentially yields each string.
-    If the input is an asynchronous generator, it forwards each generated prompt.
-
-    Args:
-        prompts (list[str] | AsyncGenerator): A list of strings or an asynchronous generator of prompts.
-
-    Yields:
-        str: An individual prompt from the list or the asynchronous generator.
-    """
     if isinstance(prompts, list):
         for prompt in prompts:
             yield prompt
@@ -49,21 +37,7 @@ async def generate_prompts(
             yield prompt
 
 
-def get_modality_adapter(llm_spec):
-    """
-    Returns the appropriate request adapter based on the modality of the LLM specification.
-
-    Depending on the modality of `llm_spec`, the function selects the corresponding request adapter.
-    If the modality is IMAGE or AUDIO, it returns an adapter for handling the respective type.
-    If the modality is TEXT or an unrecognized type, it returns `llm_spec` as is.
-
-    Args:
-        llm_spec: An object containing modality information for the LLM.
-
-    Returns:
-        RequestAdapter | llm_spec: An instance of the appropriate request adapter
-        or the original `llm_spec` if no adaptation is needed.
-    """
+def multi_modality_spec(llm_spec):
     match llm_spec.modality:
         case Modality.IMAGE:
             return image_generator.RequestAdapter(llm_spec)
@@ -76,71 +50,40 @@ def get_modality_adapter(llm_spec):
 
 
 async def process_prompt(
-    request_factory,
-    prompt: str,
-    tokens: int,
-    module_name: str,
-    fuzzer_state: FuzzerState,
+    request_factory, prompt, tokens, module_name, refusals, errors, outputs
 ) -> tuple[int, bool]:
     """
-    Processes a single prompt using the provided request factory and updates tracking lists.
-
-    This function sends the given `prompt` to the `request_factory`, checks for errors, and updates
-    the `tokens`, `refusals`, `errors`, and `outputs` lists accordingly. If the request fails or
-    the response indicates a refusal, the function records the issue and returns the updated token count
-    along with a boolean indicating whether the prompt was refused.
-
-    Args:
-        request_factory: An object with a `fn` method used to send the prompt.
-        prompt (str): The input prompt to be processed.
-        tokens (int): The current token count, which will be updated.
-        module_name (str): The name of the module handling the request.
-        fuzzer_state: State tracking object for the fuzzer
-
-    Returns:
-        tuple[int, bool]: Updated token count and a boolean indicating if the prompt was refused.
+    Process a single prompt and update the token count and failure status.
     """
     try:
         response = await request_factory.fn(prompt=prompt)
-
-        # Handle HTTP errors
         if response.status_code == 422:
             logger.error(f"Invalid prompt: {prompt}, error=422")
-            fuzzer_state.add_error(module_name, prompt, 422, "Invalid prompt")
+            errors.append((module_name, prompt, 422, "Invalid prompt"))
             return tokens, True
 
         if response.status_code >= 400:
             logger.error(f"HTTP {response.status_code} {response.content=}")
-            fuzzer_state.add_error(
-                module_name, prompt, response.status_code, response.text
-            )
+            errors.append((module_name, prompt, response.status_code, response.text))
             return tokens, True
-
-        # Process successful response
         response_text = response.text
         tokens += len(response_text.split())
 
-        # Check if the response indicates a refusal
         refused = refusal_heuristic(response.json())
         if refused:
-            fuzzer_state.add_refusal(
-                module_name, prompt, response.status_code, response_text
-            )
+            refusals.append((module_name, prompt, response.status_code, response_text))
 
-        fuzzer_state.add_output(module_name, prompt, response_text, refused)
+        outputs.append((module_name, prompt, response_text, refused))
         return tokens, refused
 
     except httpx.RequestError as exc:
         logger.error(f"Request error: {exc}")
-        fuzzer_state.add_error(module_name, prompt, "?", str(exc))
+        errors.append((module_name, prompt, "?", str(exc)))
         return tokens, True
     except JSONDecodeError as json_decode_error:
-        logger.error(f"JSON error: {json_decode_error}")
-        fuzzer_state.add_error(module_name, prompt, "?", str(json_decode_error))
+        logger.error(f"Jason error: {json_decode_error}")
+        errors.append((module_name, prompt, "?", str(json_decode_error)))
         return tokens, True
-    except Exception as e:
-        logger.exception(f"Unexpected error: {e}")
-        return tokens, False
 
 
 async def process_prompt_batch(
@@ -148,29 +91,14 @@ async def process_prompt_batch(
     prompts: list[str],
     tokens: int,
     module_name: str,
-    fuzzer_state: FuzzerState,
+    refusals,
+    errors,
+    outputs,
 ) -> tuple[int, int]:
-    """
-    Processes a batch of prompts asynchronously and aggregates the results.
-
-    This function sends multiple prompts concurrently using `process_prompt`,
-    collects the token count and failure status for each prompt, and returns
-    the total number of tokens processed and the number of failed prompts.
-
-    Args:
-        request_factory: An object with a `fn` method used to send the prompts.
-        prompts (list[str]): A list of input prompts to be processed.
-        tokens (int): The initial token count, which will be updated.
-        module_name (str): The name of the module handling the request.
-        fuzzer_state: State tracking object for the fuzzer
-
-    Returns:
-        tuple[int, int]:
-            - Total number of tokens processed.
-            - Number of failed prompts.
-    """
     tasks = [
-        process_prompt(request_factory, p, tokens, module_name, fuzzer_state)
+        process_prompt(
+            request_factory, p, tokens, module_name, refusals, errors, outputs
+        )
         for p in prompts
     ]
     results = await asyncio.gather(*tasks)
@@ -179,154 +107,7 @@ async def process_prompt_batch(
     return total_tokens, failures
 
 
-async def scan_module(
-    request_factory,
-    module,
-    fuzzer_state: FuzzerState,
-    processed_prompts: int = 0,
-    total_prompts: int = 0,
-    max_budget: int = 0,
-    total_tokens: int = 0,
-    optimize: bool = False,
-    stop_event: asyncio.Event | None = None,
-) -> AsyncGenerator[dict[str, Any], None]:
-    """
-    Scan a single module.
-
-    Args:
-        request_factory: The factory for creating requests
-        module: The prompt module to scan
-        fuzzer_state: State tracking object for the fuzzer
-        processed_prompts: Number of prompts processed so far
-        total_prompts: Total number of prompts to process
-        max_budget: Maximum token budget
-        total_tokens: Current token count
-        optimize: Whether to use optimization
-        stop_event: Event to stop scanning
-
-    Yields:
-        ScanResult objects as the scan progresses
-    """
-    tokens = 0
-    module_failures = 0
-    module_prompts = 0
-    failure_rates = []
-    should_stop = False
-
-    # Initialize optimizer if optimization is enabled
-    optimizer = (
-        Optimizer(
-            [Real(0, 1)], base_estimator="GP", n_initial_points=INITIAL_OPTIMIZER_POINTS
-        )
-        if optimize
-        else None
-    )
-
-    module_size = 0 if module.lazy else len(module.prompts)
-    logger.info(f"Scanning {module.dataset_name} {module_size}")
-    yield ScanResult(
-        module=module.dataset_name,
-        tokens=0,
-        cost=0,
-        progress=0,
-        failureRate=0,
-        prompt="",
-        latency=0,
-        model="",
-    ).model_dump_json()
-
-    async for prompt in generate_prompts(module.prompts):
-        if stop_event and stop_event.is_set():
-            stop_event.clear()
-            logger.info("Scan stopped by user.")
-            yield ScanResult.status_msg("Scan stopped by user.")
-            return
-
-        processed_prompts += 1
-        module_prompts += 1
-
-        # Calculate progress based on total processed prompts
-        progress = 100 * processed_prompts / total_prompts if total_prompts else 0
-        progress = progress % 100
-
-        total_tokens -= tokens
-        start = time.time()
-
-        tokens, failed = await process_prompt(
-            request_factory,
-            prompt,
-            tokens,
-            module.dataset_name,
-            fuzzer_state=fuzzer_state,
-        )
-
-        end = time.time()
-        total_tokens += tokens
-
-        if failed:
-            module_failures += 1
-
-        failure_rate = module_failures / max(module_prompts, 1)
-        failure_rates.append(failure_rate)
-        cost = calculate_cost(tokens)
-
-        response_text = fuzzer_state.get_last_output(prompt) or ""
-
-        yield ScanResult(
-            module=module.dataset_name,
-            tokens=round(tokens / 1000, 1),
-            cost=cost,
-            progress=round(progress, 2),
-            failureRate=round(failure_rate * 100, 2),
-            prompt=prompt[:MAX_PROMPT_LENGTH],
-            latency=end - start,
-            model=response_text,
-        ).model_dump_json()
-
-        # Optimization logic
-        if optimize and optimizer and len(failure_rates) >= MIN_FAILURE_SAMPLES:
-            next_point = optimizer.ask()
-            optimizer.tell(next_point, -failure_rate)
-            best_failure_rate = -optimizer.get_result().fun
-            if best_failure_rate > FAILURE_RATE_THRESHOLD:
-                yield ScanResult.status_msg(
-                    f"High failure rate detected ({best_failure_rate:.2%}). Stopping this module..."
-                )
-                should_stop = True
-                break
-
-        # Budget check
-        if total_tokens > max_budget:
-            logger.info(
-                f"Scan ran out of budget and stopped. {total_tokens=} {max_budget=}"
-            )
-            yield ScanResult.status_msg(
-                f"Scan ran out of budget and stopped. {total_tokens=} {max_budget=}"
-            )
-            should_stop = True
-            break
-
-        if should_stop:
-            break
-
-    return
-
-
 async def with_error_handling(agen):
-    """
-    Wraps an asynchronous generator with error handling.
-
-    This function iterates over an asynchronous generator, yielding its values.
-    If an exception occurs, it logs the error and yields a failure message.
-    Finally, it ensures that a completion message is always yielded.
-
-    Args:
-        agen: An asynchronous generator that produces scan results.
-
-    Yields:
-        ScanResult: Either a successful result, an error message if an
-        exception occurs, or a completion message at the end.
-    """
     try:
         async for t in agen:
             yield t
@@ -342,37 +123,14 @@ async def perform_single_shot_scan(
     max_budget: int,
     datasets: list[dict[str, str]] = [],
     tools_inbox=None,
-    optimize: bool = False,
-    stop_event: asyncio.Event | None = None,
+    optimize=False,
+    stop_event: asyncio.Event = None,
     secrets: dict[str, str] = {},
 ) -> AsyncGenerator[str, None]:
-    """
-    Perform a standard security scan using a given request factory.
-
-    This function processes security scan prompts from selected datasets while
-    respecting a predefined token budget. It supports optimization, failure tracking,
-    and early stopping based on budget constraints or user intervention.
-
-    Args:
-        request_factory: A factory function that generates requests for processing prompts.
-        max_budget (int): The maximum token budget for the scan.
-        datasets (list[dict[str, str]], optional): A list of datasets containing security prompts.
-        tools_inbox: Optional additional tools for processing (default: None).
-        optimize (bool, optional): Whether to enable failure rate optimization (default: False).
-        stop_event (asyncio.Event, optional): An event to signal early termination (default: None).
-        secrets (dict[str, str], optional): A dictionary of secrets for authentication (default: {}).
-
-    Yields:
-        str: JSON-encoded scan results or status messages.
-
-    The function iterates over prompts, processes them asynchronously, and updates
-    failure statistics and token usage. If the scan exceeds the budget or failure rate is too high,
-    it stops execution. Results are saved to a CSV file upon completion.
-    """
+    """Perform a standard security scan."""
     max_budget = max_budget * BUDGET_MULTIPLIER
     selected_datasets = [m for m in datasets if m["selected"]]
-    request_factory = get_modality_adapter(request_factory)
-
+    request_factory = multi_modality_spec(request_factory)
     yield ScanResult.status_msg("Loading datasets...")
     prompt_modules = prepare_prompts(
         dataset_names=[m["dataset_name"] for m in selected_datasets],
@@ -382,35 +140,108 @@ async def perform_single_shot_scan(
     )
     yield ScanResult.status_msg("Datasets loaded. Starting scan...")
 
-    fuzzer_state = FuzzerState()
+    errors = []
+    refusals = []
+    outputs = []
     total_prompts = sum(len(m.prompts) for m in prompt_modules if not m.lazy)
     processed_prompts = 0
 
+    optimizer = (
+        Optimizer([Real(0, 1)], base_estimator="GP", n_initial_points=25)
+        if optimize
+        else None
+    )
+    failure_rates = []
+
     total_tokens = 0
+    tokens = 0
+    should_stop = False
     for module in prompt_modules:
-        module_gen = scan_module(
-            request_factory=request_factory,
-            module=module,
-            fuzzer_state=fuzzer_state,
-            processed_prompts=processed_prompts,
-            total_prompts=total_prompts,
-            max_budget=max_budget,
-            total_tokens=total_tokens,
-            optimize=optimize,
-            stop_event=stop_event,
-        )
-        try:
-            async for result in module_gen:
-                yield result
-        except Exception:
-            logger.error("Module exception")
-            continue
-        # Update processed_prompts count
+        if should_stop:
+            break
+        tokens = 0
+        module_failures = 0
         module_size = 0 if module.lazy else len(module.prompts)
-        processed_prompts += module_size
+        logger.info(f"Scanning {module.dataset_name} {module_size}")
+        module_prompts = 0  # Reset for each module
+
+        async for prompt in generate_prompts(module.prompts):
+            if stop_event and stop_event.is_set():
+                stop_event.clear()
+                logger.info("Scan stopped by user.")
+                yield ScanResult.status_msg("Scan stopped by user.")
+                return
+
+            processed_prompts += 1
+            module_prompts += 1  # Fixed increment syntax
+            # Calculate progress based on total processed prompts
+            progress = 100 * processed_prompts / total_prompts if total_prompts else 0
+            progress = progress % 100
+
+            total_tokens -= tokens
+            start = time.time()
+            tokens, failed = await process_prompt(
+                request_factory,
+                prompt,
+                tokens,
+                module.dataset_name,
+                refusals,
+                errors,
+                outputs,
+            )
+            end = time.time()
+            total_tokens += tokens
+
+            if failed:
+                module_failures += 1
+            failure_rate = module_failures / max(module_prompts, 1)
+            failure_rates.append(failure_rate)
+            cost = calculate_cost(tokens)
+
+            last_output = outputs[-1] if outputs else None
+            if last_output and last_output[1] == prompt:
+                response_text = last_output[2]
+            else:
+                response_text = ""
+
+            yield ScanResult(
+                module=module.dataset_name,
+                tokens=round(tokens / 1000, 1),
+                cost=cost,
+                progress=round(progress, 2),
+                failureRate=round(failure_rate * 100, 2),
+                prompt=prompt[:MAX_PROMPT_LENGTH],
+                latency=end - start,
+                model=response_text,
+            ).model_dump_json()
+
+            if optimize and len(failure_rates) >= 5:
+                next_point = optimizer.ask()
+                optimizer.tell(next_point, -failure_rate)
+                best_failure_rate = -optimizer.get_result().fun
+                if best_failure_rate > 0.5:
+                    yield ScanResult.status_msg(
+                        f"High failure rate detected ({best_failure_rate:.2%}). Stopping this module..."
+                    )
+                    should_stop = True
+                    break
+            if total_tokens > max_budget:
+                logger.info(
+                    f"Scan ran out of budget and stopped. {total_tokens=} {max_budget=}"
+                )
+                yield ScanResult.status_msg(
+                    f"Scan ran out of budget and stopped. {total_tokens=} {max_budget=}"
+                )
+                should_stop = True
+                break
 
     yield ScanResult.status_msg("Scan completed.")
-    fuzzer_state.export_failures("failures.csv")
+
+    failure_data = errors + refusals
+    df = pd.DataFrame(
+        failure_data, columns=["module", "prompt", "status_code", "content"]
+    )
+    df.to_csv("failures.csv", index=False)
 
 
 async def perform_many_shot_scan(
@@ -419,39 +250,14 @@ async def perform_many_shot_scan(
     datasets: list[dict[str, str]] = [],
     probe_datasets: list[dict[str, str]] = [],
     tools_inbox=None,
-    optimize: bool = False,
-    stop_event: asyncio.Event | None = None,
+    optimize=False,
+    stop_event: asyncio.Event = None,
     probe_frequency: float = 0.2,
     max_ctx_length: int = 10_000,
     secrets: dict[str, str] = {},
 ) -> AsyncGenerator[str, None]:
-    """
-    Perform a multi-step security scan with probe injection.
-
-    This function executes a security scan while periodically injecting probe datasets
-    to test system robustness. It tracks failures, optimizes scan efficiency,
-    and ensures adherence to a predefined token budget.
-
-    Args:
-        request_factory: A factory function that generates requests for processing prompts.
-        max_budget (int): The maximum token budget for the scan.
-        datasets (list[dict[str, str]], optional): The main datasets for scanning.
-        probe_datasets (list[dict[str, str]], optional): Additional datasets for probe injection.
-        tools_inbox: Optional tools for additional processing (default: None).
-        optimize (bool, optional): Whether to enable failure rate optimization (default: False).
-        stop_event (asyncio.Event, optional): An event to signal early termination (default: None).
-        probe_frequency (float, optional): The probability of probe injection (default: 0.2).
-        max_ctx_length (int, optional): The maximum context length before resetting (default: 10,000 tokens).
-        secrets (dict[str, str], optional): A dictionary of secrets for authentication (default: {}).
-
-    Yields:
-        str: JSON-encoded scan results or status messages.
-
-    This function iterates over prompts, injects probe prompts at random intervals,
-    processes them asynchronously, and tracks failure rates. If failure rates exceed a threshold
-    or budget is exhausted, the scan is stopped early. Results are saved to a CSV file upon completion.
-    """
-    request_factory = get_modality_adapter(request_factory)
+    """Perform a multi-step security scan with probe injection."""
+    request_factory = multi_modality_spec(request_factory)
     # Load main and probe datasets
     yield ScanResult.status_msg("Loading datasets...")
     prompt_modules = prepare_prompts(
@@ -463,10 +269,17 @@ async def perform_many_shot_scan(
     msj_modules = msj_data.prepare_prompts(probe_datasets)
     yield ScanResult.status_msg("Datasets loaded. Starting scan...")
 
-    fuzzer_state = FuzzerState()
+    errors = []
+    refusals = []
+    outputs = []
     total_prompts = sum(len(m.prompts) for m in prompt_modules if not m.lazy)
     processed_prompts = 0
 
+    optimizer = (
+        Optimizer([Real(0, 1)], base_estimator="GP", n_initial_points=25)
+        if optimize
+        else None
+    )
     failure_rates = []
 
     for module in prompt_modules:
@@ -480,7 +293,6 @@ async def perform_many_shot_scan(
                 logger.info("Scan stopped by user.")
                 yield ScanResult.status_msg("Scan stopped by user.")
                 return
-
             tokens = 0
             processed_prompts += 1
             progress = 100 * processed_prompts / total_prompts if total_prompts else 0
@@ -508,7 +320,9 @@ async def perform_many_shot_scan(
                     full_prompt,
                     tokens,
                     module.dataset_name,
-                    fuzzer_state=fuzzer_state,
+                    refusals,
+                    errors,
+                    outputs,
                 )
                 if failed:
                     module_failures += 1
@@ -529,48 +343,30 @@ async def perform_many_shot_scan(
                 prompt=prompt[:MAX_PROMPT_LENGTH],
             ).model_dump_json()
 
-            if optimize and len(failure_rates) >= MIN_FAILURE_SAMPLES:
-                yield ScanResult.status_msg(
-                    f"High failure rate detected ({failure_rate:.2%}). Stopping this module..."
-                )
-                break
+            if optimize and len(failure_rates) >= 5:
+                next_point = optimizer.ask()
+                optimizer.tell(next_point, -failure_rate)
+                best_failure_rate = -optimizer.get_result().fun
+                if best_failure_rate > 0.5:
+                    yield ScanResult.status_msg(
+                        f"High failure rate detected ({best_failure_rate:.2%}). Stopping this module..."
+                    )
+                    break
 
     yield ScanResult.status_msg("Scan completed.")
-    fuzzer_state.export_failures("failures.csv")
+
+    df = pd.DataFrame(
+        errors + refusals, columns=["module", "prompt", "status_code", "content"]
+    )
+    df.to_csv("failures.csv", index=False)
 
 
 def scan_router(
     request_factory,
     scan_parameters: Scan,
     tools_inbox=None,
-    stop_event: asyncio.Event | None = None,
+    stop_event: asyncio.Event = None,
 ):
-    """
-    Route scan requests to the appropriate scanning function.
-
-    This function determines whether to perform a multi-step or single-shot
-    security scan based on the provided scan parameters.
-
-    Args:
-        request_factory: A factory function to generate requests for processing prompts.
-        scan_parameters (Scan): An object containing the parameters for the scan, including:
-            - enableMultiStepAttack (bool): Whether to perform a multi-step scan.
-            - maxBudget (int): The maximum token budget for the scan.
-            - datasets (list[dict[str, str]]): The datasets to scan.
-            - probe_datasets (list[dict[str, str]], optional): Datasets for probe injection (multi-step only).
-            - optimize (bool): Whether to enable optimization.
-            - secrets (dict[str, str], optional): A dictionary of secrets for authentication.
-        tools_inbox: Optional tools for additional processing (default: None).
-        stop_event (asyncio.Event, optional): An event to signal early termination (default: None).
-
-    Returns:
-        A function wrapped with `with_error_handling`, which executes either:
-        - `perform_many_shot_scan` for multi-step scanning.
-        - `perform_single_shot_scan` for single-shot scanning.
-
-    The function ensures that the appropriate scanning method is chosen based on
-    the `enableMultiStepAttack` flag in `scan_parameters`.
-    """
     if scan_parameters.enableMultiStepAttack:
         return with_error_handling(
             perform_many_shot_scan(

agentic_security/probe_actor/operator.py

@@ -1,4 +1,5 @@
 import asyncio
+import logging
 from typing import Any
 
 import httpx
@@ -6,10 +7,13 @@ from pydantic import BaseModel, Field
 from pydantic_ai import Agent, RunContext
 
 from agentic_security.http_spec import LLMSpec
-from agentic_security.logutils import logger
 
 LLM_SPECS = []
 
+# Configure logging
+logging.basicConfig(level=logging.INFO)
+logger = logging.getLogger(__name__)
+
 
 class AgentSpecification(BaseModel):
     name: str | None = Field(None, description="Name of the LLM/agent")

agentic_security/probe_actor/state.py

@@ -1,47 +0,0 @@
-import pandas as pd
-
-
-class FuzzerState:
-    """Container for tracking scan results"""
-
-    def __init__(self):
-        self.errors = []
-        self.refusals = []
-        self.outputs = []
-
-    def add_error(
-        self,
-        module_name: str,
-        prompt: str,
-        status_code: int | str,
-        error_msg: str,
-    ):
-        """Add an error to the state"""
-        self.errors.append((module_name, prompt, status_code, error_msg))
-
-    def add_refusal(
-        self, module_name: str, prompt: str, status_code: int, response_text: str
-    ):
-        """Add a refusal to the state"""
-        self.refusals.append((module_name, prompt, status_code, response_text))
-
-    def add_output(
-        self, module_name: str, prompt: str, response_text: str, refused: bool
-    ):
-        """Add an output to the state"""
-        self.outputs.append((module_name, prompt, response_text, refused))
-
-    def get_last_output(self, prompt: str) -> str | None:
-        """Get the last output for a given prompt"""
-        for output in reversed(self.outputs):
-            if output[1] == prompt:
-                return output[2]
-        return None
-
-    def export_failures(self, filename: str = "failures.csv"):
-        """Export failures to a CSV file"""
-        failure_data = self.errors + self.refusals
-        df = pd.DataFrame(
-            failure_data, columns=["module", "prompt", "status_code", "content"]
-        )
-        df.to_csv(filename, index=False)

agentic_security/probe_actor/test_refusal.py

@@ -0,0 +1,51 @@
+from agentic_security.probe_actor.refusal import DefaultRefusalClassifier
+from agentic_security.probe_actor.refusal import refusal_heuristic, RefusalClassifierManager, RefusalClassifierPlugin
+
+
+class TestCheckRefusal:
+    # The function correctly identifies a refusal phrase in the response.
+    def test_identify_refusal_phrase(self):
+        response = "I'm sorry, but I cannot provide that information."
+        assert DefaultRefusalClassifier().is_refusal(response)
+
+    # The response is an empty string.
+    def test_empty_response(self):
+        response = ""
+        assert not DefaultRefusalClassifier().is_refusal(response)
+    def test_non_refusal_phrase(self):
+        """Test that a response without refusal phrases returns False."""
+        response = "Here is some helpful information."
+        assert not DefaultRefusalClassifier().is_refusal(response)
+
+    def test_refusal_heuristic(self):
+        """Test that the refusal_heuristic function correctly identifies refusal phrases."""
+        response = "Sorry, I cannot assist with your request."
+        # refusal_heuristic will cast the argument to a string and use the plugin manager to check for refusal markers
+        assert refusal_heuristic(response)
+
+    def test_plugin_manager_with_dummy_plugin(self):
+        """Test that RefusalClassifierManager returns True if any registered plugin returns True."""
+        class AlwaysFalsePlugin(RefusalClassifierPlugin):
+            def is_refusal(self, response: str) -> bool:
+                return False
+
+        class AlwaysTruePlugin(RefusalClassifierPlugin):
+            def is_refusal(self, response: str) -> bool:
+                return True
+
+        manager = RefusalClassifierManager()
+        manager.register_plugin("false", AlwaysFalsePlugin())
+        manager.register_plugin("true", AlwaysTruePlugin())
+        response = "Any response text"
+        assert manager.is_refusal(response)
+
+    def test_default_classifier_with_empty_phrases(self):
+        """Test that DefaultRefusalClassifier returns False when provided with an empty refusal phrase list."""
+        classifier = DefaultRefusalClassifier(refusal_phrases=[])
+        response = "I do not have any info."
+        assert not classifier.is_refusal(response)
+
+    def test_case_sensitivity(self):
+        """Test that string matching is case-sensitive."""
+        response = "i'm sorry, but can you help me?"  # lower-case "i'm sorry" does not match "I'm sorry" and no extra refusal phrases are present
+        assert not DefaultRefusalClassifier().is_refusal(response)

agentic_security/probe_data/__init__.py

@@ -1,4 +1,4 @@
-from .data import load_local_csv, load_local_csv_files
+from .data import load_local_csv
 
 REGISTRY_V0 = [
     {
@@ -484,18 +484,3 @@ REGISTRY = REGISTRY_V0 + [
         "modality": "text",
     },
 ]
-
-for ds in load_local_csv_files():
-    REGISTRY.append(
-        {
-            "dataset_name": ds.dataset_name,
-            "num_prompts": len(ds.prompts),
-            "tokens": ds.prompts,
-            "approx_cost": 0.0,
-            "is_active": True,
-            "source": f"Local file dataset: {ds.metadata['src']}",
-            "selected": False,
-            "url": "",
-            "modality": "text",
-        }
-    )

agentic_security/probe_data/audio_generator.py

@@ -1,5 +1,4 @@
 import base64
-import logging
 import os
 import platform
 import subprocess
@@ -8,16 +7,6 @@ import uuid
 import httpx
 from cache_to_disk import cache_to_disk
 
-# Configure logging
-logging.basicConfig(level=logging.INFO)
-logger = logging.getLogger(__name__)
-
-
-class AudioGenerationError(Exception):
-    """Custom exception for errors during audio generation."""
-
-    pass
-
 
 def encode(content: bytes) -> str:
     encoded_content = base64.b64encode(content).decode("utf-8")
@@ -52,25 +41,12 @@ def generate_audio_mac_wav(prompt: str) -> bytes:
         # Read the WAV file into memory
         with open(temp_wav_path, "rb") as f:
             audio_bytes = f.read()
-
-    except subprocess.CalledProcessError as e:
-        logger.error(f"Subprocess error: {e}")
-        raise AudioGenerationError("Failed to generate or convert audio.") from e
-    except FileNotFoundError as e:
-        logger.error(f"File not found: {e}")
-        raise AudioGenerationError("Required file not found.") from e
-    except Exception as e:
-        logger.exception("Unexpected error occurred.")
-        raise AudioGenerationError(
-            "An unexpected error occurred during audio generation."
-        ) from e
     finally:
-        for path in (temp_aiff_path, temp_wav_path):
-            try:
-                if os.path.exists(path):
-                    os.remove(path)
-            except Exception as e:
-                logger.warning(f"Failed to delete temporary file {path}: {e}")
+        # Clean up the temporary files
+        if os.path.exists(temp_aiff_path):
+            os.remove(temp_aiff_path)
+        if os.path.exists(temp_wav_path):
+            os.remove(temp_wav_path)
 
     # Return the audio bytes
     return audio_bytes

agentic_security/probe_data/data.py

@@ -1,16 +1,14 @@
 import io
 import os
 import random
-from collections.abc import Callable, Iterator
-from functools import partial
-from typing import Any, TypeVar
+from functools import lru_cache
 
 import httpx
 import pandas as pd
 from cache_to_disk import cache_to_disk
 from datasets import load_dataset
+from loguru import logger
 
-from agentic_security.logutils import logger
 from agentic_security.probe_data import stenography_fn
 from agentic_security.probe_data.models import ProbeDataset
 from agentic_security.probe_data.modules import (
@@ -21,293 +19,281 @@ from agentic_security.probe_data.modules import (
     rl_model,
 )
 
-# Type aliases for clarity
-T = TypeVar("T")
-FilterFn = Callable[[pd.Series], bool]
-ColumnMappings = dict[str, str]
-DatasetLoader = Callable[[], ProbeDataset]
-TransformFn = Callable[[str], str]
 
+@cache_to_disk()
+def load_dataset_general(
+    dataset_name,
+    dataset_split="train",
+    column_mappings=None,
+    filter_fn=None,
+    custom_url=None,
+    additional_metadata=None,
+):
+    """
+    Generalized function to load datasets with flexible configurations.
 
-# Core data loading utilities
-def fetch_csv_content(url: str) -> str:
-    """Fetch CSV content from a URL."""
-    response = httpx.get(url)
-    response.raise_for_status()  # Raise exception for bad responses
-    return response.content.decode("utf-8")
+    :param dataset_name: Name of the dataset or URL for custom CSVs
+    :param dataset_split: Split to load from the dataset (e.g., "train")
+    :param column_mappings: Dictionary mapping dataset columns to expected keys, e.g., {'prompt': 'query'}
+    :param filter_fn: A filtering function that takes a row and returns True/False
+    :param custom_url: URL for custom CSV datasets
+    :param additional_metadata: Additional metadata to include in the ProbeDataset
+    :return: A ProbeDataset object with the processed data
+    """
+    if custom_url:
+        logger.info(f"Loading custom CSV dataset from {custom_url}")
+        r = httpx.get(custom_url)
+        content = r.content
+        df = pd.read_csv(io.StringIO(content.decode("utf-8")))
+    else:
+        logger.info(f"Loading dataset {dataset_name} from Hugging Face datasets")
+        dataset = load_dataset(dataset_name)
+        df = pd.DataFrame(dataset[dataset_split])
 
+    # Apply column mappings if provided
+    if column_mappings:
+        df.rename(columns=column_mappings, inplace=True)
 
-def load_df_from_source(source: str, is_url: bool = False) -> pd.DataFrame:
-    """Load DataFrame from either URL or Hugging Face dataset."""
-    if is_url:
-        logger.info(f"Loading custom CSV from {source}")
-        return pd.read_csv(io.StringIO(fetch_csv_content(source)))
-    logger.info(f"Loading dataset {source} from Hugging Face")
-    return pd.DataFrame(load_dataset(source)["train"])
-
-
-def transform_df(
-    df: pd.DataFrame,
-    mappings: ColumnMappings | None = None,
-    filter_fn: FilterFn | None = None,
-) -> pd.DataFrame:
-    """Apply transformations to DataFrame."""
-    result = df.copy()
-    if mappings:
-        result.rename(columns=mappings, inplace=True)
+    # Filter rows if filter_fn is provided
     if filter_fn:
-        result = result[result.apply(filter_fn, axis=1)]
-    return result
+        df = df[df.apply(filter_fn, axis=1)]
 
+    # Extract prompts
+    prompts = df[column_mappings.get("prompt", "prompt")].tolist()
 
-def create_probe_dataset(
-    name: str, prompts: list[str], metadata: dict[str, Any] | None = None
-) -> ProbeDataset:
-    """Create a ProbeDataset from prompts."""
-    metadata = metadata or {}
     return ProbeDataset(
-        dataset_name=name,
-        metadata=metadata,
+        dataset_name=dataset_name,
+        metadata=additional_metadata or {},
         prompts=prompts,
-        tokens=sum(len(str(s).split()) for s in prompts),
+        tokens=count_words_in_list(prompts),
         approx_cost=0.0,
     )
 
 
-# Generalized dataset loader
+def count_words_in_list(str_list):
+    """Calculate the total number of words in a given list of strings.
+
+    :param str_list: List of strings
+    :return: Total number of words across all strings in the list
+    """
+    total_words = sum(len(str(s).split()) for s in str_list)
+    return total_words
+
+
 @cache_to_disk()
-def load_dataset_generic(
-    name: str,
-    mappings: ColumnMappings | None = None,
-    filter_fn: FilterFn | None = None,
-    url: str | None = None,
-    metadata: dict[str, Any] | None = None,
-) -> ProbeDataset:
-    """Load and process a dataset with flexible configuration."""
-    try:
-        df = load_df_from_source(url or name, is_url=bool(url))
-        transformed_df = transform_df(df, mappings, filter_fn)
-
-        # Determine which column to use as the prompt source
-        prompt_col = None
-        if mappings and "prompt" in mappings:
-            prompt_col = mappings["prompt"]
-        elif "prompt" in transformed_df.columns:
-            prompt_col = "prompt"
-        else:
-            # Try to find a suitable text column
-            text_columns = [
-                col
-                for col in transformed_df.columns
-                if any(
-                    keyword in col.lower()
-                    for keyword in ["prompt", "text", "query", "question"]
-                )
-            ]
-            if text_columns:
-                prompt_col = text_columns[0]
-                logger.info(f"Using column '{prompt_col}' as prompt source")
-            else:
-                logger.error(f"No suitable prompt column found in dataset {name}")
-                return create_probe_dataset(name, [], metadata)
-
-        # Extract prompts and filter out empty ones
-        prompts = [
-            p
-            for p in transformed_df[prompt_col].tolist()
-            if p and isinstance(p, (str, int, float))
-        ]
-        return create_probe_dataset(name, prompts, metadata)
-    except Exception as e:
-        logger.error(f"Error loading dataset {name}: {e}")
-        return create_probe_dataset(name, [], {"error": str(e)})
-
-
-# Dataset-specific configurations
-DATASET_CONFIGS = {
-    "ShawnMenz/DAN_jailbreak": {
-        "mappings": {"prompt": "prompt"},
-        "filter_fn": lambda row: row["jailbreak"],
-    },
-    "deepset/prompt-injections": {
-        "mappings": {"prompt": "text"},
-        "filter_fn": lambda row: row["label"],
-    },
-    "notrichardren/refuse-to-answer-prompts": {
-        "mappings": {"prompt": "claim"},
-        "filter_fn": lambda row: row["label"],
-    },
-    "rubend18/ChatGPT-Jailbreak-Prompts": {
-        "mappings": {"prompt": "Prompt"},
-        "filter_fn": lambda row: row["label"],
-    },
-    "markush1/LLM-Jailbreak-Classifier": {
-        "mappings": {"prompt": "prompt"},
-        "filter_fn": lambda row: row["classification"] == "jailbreak",
-    },
-    "ShawnMenz/jailbreak_sft_rm_ds": {
-        "url": "hf://datasets/ShawnMenz/jailbreak_sft_rm_ds/jailbreak_sft_rm_ds.csv",
-        "mappings": {"prompt": "prompt"},
-        "filter_fn": lambda row: row["jailbreak"] == "jailbreak",
-    },
-    "verazuo/jailbreak_llms/2023_05_07": {
-        "url": "https://raw.githubusercontent.com/verazuo/jailbreak_llms/main/data/prompts/jailbreak_prompts_2023_05_07.csv",
-        "mappings": {"prompt": "prompt"},
-        "filter_fn": lambda row: bool(row["jailbreak"]),
-    },
-    "verazuo/jailbreak_llms/2023_12_25": {
-        "url": "https://raw.githubusercontent.com/verazuo/jailbreak_llms/main/data/prompts/jailbreak_prompts_2023_12_25.csv",
-        "mappings": {"prompt": "prompt"},
-        "filter_fn": lambda row: bool(row["jailbreak"]),
-    },
-}
-
-# Additional generic dataset configurations
-DATASET_CONFIGS_GENERICS = {
-    "simonycl/aya-23-8B_advbench_jailbreak": {"mappings": {"prompt": "prompt"}},
-    "acmc/jailbreaks_dataset_with_perplexity_bigcode_starcoder2-3b_bigcode_starcoder2-7b": {},
-    "karanxa/dolphin-jailbreak-finetuning-dataset": {"mappings": {"prompt": "text"}},
-    "karanxa/llama-2-jailbreak-dataset": {"mappings": {"prompt": "text"}},
-    "karanxa/llama2-uncensored-jailbreak-dataset-finetuning": {
-        "mappings": {"prompt": "text"}
-    },
-    "liuyanchen1015/Llama-3.2-1B_jailbreak_responses": {
-        "mappings": {"prompt": "jailbreak_prompt_text"}
-    },
-    "liuyanchen1015/Llama-3.2-1B-Instruct_jailbreak_responses": {
-        "mappings": {"prompt": "jailbreak_prompt_text"}
-    },
-    "liuyanchen1015/Llama-3.2-1B-Instruct_jailbreak_responses_with_judgment": {
-        "mappings": {"prompt": "jailbreak_prompt_text"}
-    },
-    "jackhhao/jailbreak-classification": {"mappings": {"prompt": "prompt"}},
-    "walledai/JailbreakBench": {"mappings": {"prompt": "prompt"}},
-    "walledai/JailbreakHub": {"mappings": {"prompt": "prompt"}},
-    "Granther/evil-jailbreak": {"mappings": {"prompt": "text"}},
-    "sevdeawesome/jailbreak_success": {"mappings": {"prompt": "jailbreak_prompt_text"}},
-    "IDA-SERICS/Disaster-tweet-jailbreaking": {"mappings": {"prompt": "prompt_attack"}},
-    "GeorgeDaDude/Jailbreak_Complete_DS_labeled": {"mappings": {"prompt": "question"}},
-    "dayone3nder/jailbreak_prompt_JBB_sft_trainset": {"mappings": {"prompt": "prompt"}},
-    "dayone3nder/general_safe_mix_jailbreak_prompt_JBB_trainset": {
-        "mappings": {"prompt": "prompt"}
-    },
-}
-
-
-# Dataset factory
-def create_dataset_loader(name: str, config: dict[str, Any]) -> DatasetLoader:
-    """Create a dataset loader from configuration."""
-    return partial(
-        load_dataset_generic,
-        name=name,
-        mappings=config.get("mappings"),
-        filter_fn=config.get("filter_fn"),
-        url=config.get("url"),
-        metadata={"source": name, "config": str(config)},
-    )
-
-
-# Specialized loaders
-@cache_to_disk()
-def load_multi_dataset(name: str, sub_datasets: list[str]) -> ProbeDataset:
-    """Load and combine multiple sub-datasets."""
-    prompts = []
-    for sub in sub_datasets:
-        try:
-            dataset = load_dataset(name, sub)
-            if "query" in dataset["train"].features:
-                prompts.extend(dataset["train"]["query"])
-            else:
-                logger.warning(f"No 'query' column in {name}/{sub}")
-        except Exception as e:
-            logger.error(f"Error loading {name}/{sub}: {e}")
-
-    return create_probe_dataset(
-        f"{name}_combined", prompts, {"source": name, "sub_datasets": sub_datasets}
+def load_dataset_v1():
+    return load_dataset_general(
+        dataset_name="ShawnMenz/DAN_jailbreak",
+        column_mappings={"prompt": "prompt", "jailbreak": "jailbreak"},
+        filter_fn=lambda row: row["jailbreak"],
     )
 
 
 @cache_to_disk()
-def load_jailbreak_v28k() -> ProbeDataset:
-    """Load JailBreakV-28K dataset."""
-    try:
-        df = pd.read_csv(
-            "hf://datasets/JailbreakV-28K/JailBreakV-28k/JailBreakV_28K.csv"
-        )
-        prompts = df["jailbreak_query"].tolist()
-        return create_probe_dataset(
-            "JailbreakV-28K/JailBreakV-28k",
-            prompts,
-            {"source": "JailbreakV-28K/JailBreakV-28k"},
-        )
-    except Exception as e:
-        logger.error(f"Error loading JailbreakV-28K: {e}")
-        return create_probe_dataset("JailbreakV-28K/JailBreakV-28k", [])
+def load_dataset_v2():
+    return load_dataset_general(
+        dataset_name="deepset/prompt-injections",
+        column_mappings={"prompt": "text", "jailbreak": "label"},
+        filter_fn=lambda row: row["label"],
+    )
 
 
 @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)})
+def load_dataset_v4():
+    return load_dataset_general(
+        dataset_name="notrichardren/refuse-to-answer-prompts",
+        column_mappings={"prompt": "claim", "jailbreak": "label"},
+        filter_fn=lambda row: row["label"],
+    )
 
 
-@cache_to_disk(1)
-def load_csv(file: str) -> ProbeDataset:
-    """Load prompts from local CSV files."""
-    prompts = []
-    try:
-        df = pd.read_csv(file)
-        prompts = df["prompt"].tolist()
-        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(f"fs://{file}", prompts, {"src": str(file)})
+@cache_to_disk()
+def load_dataset_v3():
+    return load_dataset_general(
+        dataset_name="rubend18/ChatGPT-Jailbreak-Prompts",
+        column_mappings={"prompt": "Prompt"},
+        filter_fn=lambda row: row["label"],
+    )
 
 
-@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")]
-    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}")
-
-    return datasets
+@cache_to_disk()
+def load_dataset_v6():
+    dataset = load_dataset("markush1/LLM-Jailbreak-Classifier")
+    bad_prompts = [
+        p
+        for p, c in zip(dataset["train"]["prompt"], dataset["train"]["classification"])
+        if c == "jailbreak"
+    ]
+    return ProbeDataset(
+        dataset_name="markush1/LLM-Jailbreak-Classifier",
+        metadata={},
+        prompts=bad_prompts,
+        tokens=count_words_in_list(bad_prompts),
+        approx_cost=0.0,
+    )
 
 
-# Stenography transformer
-class StenographyTransformer:
-    """Apply stenography transformations to datasets."""
+@cache_to_disk()
+def load_dataset_v7():
+    splits = {
+        "mini_JailBreakV_28K": "JailBreakV_28K/mini_JailBreakV_28K.csv",
+        "JailBreakV_28K": "JailBreakV_28K/JailBreakV_28K.csv",
+    }
+    df = pd.read_csv(
+        "hf://datasets/JailbreakV-28K/JailBreakV-28k/" + splits["JailBreakV_28K"]
+    )
+    bad_prompts = df["jailbreak_query"].tolist()
+    print(df.shape)
+    return ProbeDataset(
+        dataset_name="JailbreakV-28K/JailBreakV-28k",
+        metadata={},
+        prompts=bad_prompts,
+        tokens=count_words_in_list(bad_prompts),
+        approx_cost=0.0,
+    )
 
-    TRANSFORMATIONS = {
+
+@cache_to_disk()
+def load_dataset_v8():
+    df = pd.read_csv(
+        "hf://datasets/ShawnMenz/jailbreak_sft_rm_ds/jailbreak_sft_rm_ds.csv",
+        names=["jailbreak", "prompt"],
+    )
+    filtered = df[df["jailbreak"] == "jailbreak"]["prompt"].tolist()
+    return ProbeDataset(
+        dataset_name="JailbreakV-28K/JailBreakV-28k",
+        metadata={},
+        prompts=filtered,
+        tokens=count_words_in_list(filtered),
+        approx_cost=0.0,
+    )
+
+
+@cache_to_disk()
+def load_dataset_v5():
+    ds = []
+    for c in [
+        "AdvBench",
+        "ForbiddenQuestion",
+        "MJP",
+        "MaliciousInstruct",
+        "QuestionList",
+    ]:
+        dataset = load_dataset("Lemhf14/EasyJailbreak_Datasets", c)
+        bad_prompts = dataset["train"]["query"]
+        ds.extend(bad_prompts)
+
+    return ProbeDataset(
+        dataset_name="Lemhf14/EasyJailbreak_Datasets",
+        metadata={},
+        prompts=ds,
+        tokens=count_words_in_list(ds),
+        approx_cost=0.0,
+    )
+
+
+@cache_to_disk()
+def load_generic_csv(url, name, column="prompt", predicator=None):
+    r = httpx.get(url)
+    content = r.content
+    df = pd.read_csv(io.StringIO(content.decode("utf-8")))
+    logger.info(f"Loaded {len(df)} prompts from {url}")
+    filtered_prompts = df[df.apply(predicator, axis=1)][column].tolist()
+    return ProbeDataset(
+        dataset_name=name,
+        metadata={},
+        prompts=filtered_prompts,
+        tokens=count_words_in_list(filtered_prompts),
+        approx_cost=0.0,
+    )
+
+
+def prepare_prompts(dataset_names, budget, tools_inbox=None, options=[]):
+    # ## Datasets used and cleaned:
+    # markush1/LLM-Jailbreak-Classifier
+    # 1. Open-Orca/OpenOrca
+    # 2. ShawnMenz/DAN_jailbreak
+    # 3. EddyLuo/JailBreakV_28K
+    # 4. https://raw.githubusercontent.com/verazuo/jailbreak_llms/main/data/jailbreak_prompts.csv
+
+    dataset_map = {
+        "ShawnMenz/DAN_jailbreak": load_dataset_v1,
+        "deepset/prompt-injections": load_dataset_v2,
+        "notrichardren/refuse-to-answer-prompts": load_dataset_v4,
+        "rubend18/ChatGPT-Jailbreak-Prompts": load_dataset_v3,
+        "Lemhf14/EasyJailbreak_Datasets": load_dataset_v5,
+        "markush1/LLM-Jailbreak-Classifier": load_dataset_v6,
+        "JailbreakV-28K/JailBreakV-28k": load_dataset_v7,
+        "ShawnMenz/jailbreak_sft_rm_ds": load_dataset_v8,
+        "verazuo/jailbreak_llms/2023_05_07": lambda: load_generic_csv(
+            url="https://raw.githubusercontent.com/verazuo/jailbreak_llms/main/data/prompts/jailbreak_prompts_2023_05_07.csv",
+            name="verazuo/jailbreak_llms/2023_05_07",
+            column="prompt",
+            predicator=lambda x: bool(x["jailbreak"]),
+        ),
+        "verazuo/jailbreak_llms/2023_12_25.csv": lambda: load_generic_csv(
+            url="https://raw.githubusercontent.com/verazuo/jailbreak_llms/main/data/prompts/jailbreak_prompts_2023_12_25.csv.csv",
+            name="verazuo/jailbreak_llms/2023_12_25.csv",
+            column="prompt",
+            predicator=lambda x: bool(x["jailbreak"]),
+        ),
+        "Custom CSV": load_local_csv,
+    }
+    dataset_map.update(dataset_map_generics)
+
+    group = []
+    for dataset_name in dataset_names:
+        if dataset_name in dataset_map:
+            logger.info(f"Loading {dataset_name}")
+            try:
+                group.append(dataset_map[dataset_name]())
+            except Exception as e:
+                logger.error(f"Error loading {dataset_name}: {e}")
+
+    dynamic_datasets = {
+        "AgenticBackend": lambda opts: dataset_from_iterator(
+            "AgenticBackend",
+            fine_tuned.Module(group, tools_inbox=tools_inbox, opts=opts).apply(),
+            lazy=True,
+        ),
+        "Steganography": lambda opts: Stenography(group),
+        "llm-adaptive-attacks": lambda opts: dataset_from_iterator(
+            "llm-adaptive-attacks",
+            adaptive_attacks.Module(group, tools_inbox=tools_inbox, opts=opts).apply(),
+        ),
+        "Garak": lambda opts: dataset_from_iterator(
+            "Garak",
+            garak_tool.Module(group, tools_inbox=tools_inbox, opts=opts).apply(),
+            lazy=True,
+        ),
+        "Reinforcement Learning Optimization": lambda opts: dataset_from_iterator(
+            "Reinforcement Learning Optimization",
+            rl_model.Module(group, tools_inbox=tools_inbox, opts=opts).apply(),
+            lazy=True,
+        ),
+        "InspectAI": lambda opts: dataset_from_iterator(
+            "InspectAI",
+            inspect_ai_tool.Module(group, tools_inbox=tools_inbox).apply(),
+            lazy=True,
+        ),
+        "GPT fuzzer": lambda opts: [],
+    }
+
+    dynamic_groups = []
+    options = options or [{} for _ in dataset_names]
+    for dataset_name, opts in zip(dataset_names, options):
+        if dataset_name in dynamic_datasets:
+            logger.info(f"Loading {dataset_name}")
+
+            ds = dynamic_datasets[dataset_name](opts)
+
+            for g in ds:
+                dynamic_groups.append(g)
+    return group + dynamic_groups
+
+
+class Stenography:
+    fn_library = {
         "rot5": stenography_fn.rot5,
         "rot13": stenography_fn.rot13,
         "base64": stenography_fn.base64_encode,
@@ -320,164 +306,151 @@ class StenographyTransformer:
         "zigzag_obfuscation": stenography_fn.zigzag_obfuscation,
     }
 
-    def __init__(self, datasets: list[ProbeDataset]):
-        self.datasets = datasets
+    def __init__(self, prompt_groups: [ProbeDataset]):
+        self.prompt_groups = prompt_groups
 
-    def transform(self) -> Iterator[ProbeDataset]:
-        """Generate transformed datasets."""
-        if not self.datasets:
-            logger.warning("No datasets provided to StenographyTransformer")
-            return
+    def __iter__(self):
+        return self.apply()
 
-        for dataset in self.datasets:
-            if not dataset.prompts:
-                logger.warning(
-                    f"Dataset {dataset.dataset_name} has no prompts to transform"
+    def apply(self):
+        for prompt_group in self.prompt_groups:
+            size = len(prompt_group.prompts)
+            for name, fn in self.fn_library.items():
+                logger.info(f"Applying {name} to {prompt_group.dataset_name}")
+                selected_prompts = random.sample(prompt_group.prompts, int(size * 0.25))
+                prompts = [
+                    f"Follow the instructions in {name}: {fn(p)}"
+                    for p in selected_prompts
+                ]
+                yield ProbeDataset(
+                    dataset_name=f"stenography.{name}({prompt_group.dataset_name})",
+                    metadata={},
+                    prompts=prompts,
+                    tokens=count_words_in_list(prompts),
+                    approx_cost=0.0,
                 )
-                continue
-
-            size = len(dataset.prompts)
-            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))
-                try:
-                    transformed = [
-                        f"Follow the instructions in {name}: {fn(str(p))}"
-                        for p in sampled
-                    ]
-                    transformed_dataset = create_probe_dataset(
-                        f"stenography.{name}({dataset.dataset_name})",
-                        transformed,
-                        {"source": dataset.dataset_name, "transformation": name},
-                    )
-                    yield transformed_dataset
-                except Exception as e:
-                    logger.error(
-                        f"Error applying {name} to {dataset.dataset_name}: {e}"
-                    )
 
 
-def dataset_from_iterator(
-    name: str, iterator: Iterator[str], lazy: bool = False
-) -> list[ProbeDataset]:
-    """Convert an iterator into a list of ProbeDataset objects."""
-    try:
-        prompts = list(iterator) if not lazy else iterator
-        tokens = sum(len(str(s).split()) for s in prompts) if not lazy else 0
-        dataset = ProbeDataset(
-            dataset_name=name,
-            metadata={"source": name, "lazy": lazy},
-            prompts=prompts,
-            tokens=tokens,
-            approx_cost=0.0,
-            lazy=lazy,
-        )
-        return [dataset]
-    except Exception as e:
-        logger.error(f"Error creating dataset from iterator {name}: {e}")
-        return [create_probe_dataset(name, [], {"error": str(e)})]
+@lru_cache
+def load_local_csv() -> ProbeDataset:
+    csv_files = [f for f in os.listdir(".") if f.endswith(".csv")]
+    logger.info(f"Found {len(csv_files)} CSV files")
+    logger.info(f"CSV files: {csv_files}")
+    prompt_list = []
 
-
-# Main dataset preparation
-def prepare_prompts(
-    dataset_names: list[str],
-    budget: float,
-    tools_inbox=None,
-    options: list[dict] = None,
-) -> list[ProbeDataset]:
-    """Prepare datasets based on names and options."""
-    # Base dataset loaders
-    logger.info(f"Preparing datasets: {dataset_names}")
-    dataset_loaders = {
-        **{k: create_dataset_loader(k, v) for k, v in DATASET_CONFIGS.items()},
-        **{k: create_dataset_loader(k, v) for k, v in DATASET_CONFIGS_GENERICS.items()},
-        "Lemhf14/EasyJailbreak_Datasets": partial(
-            load_multi_dataset,
-            name="Lemhf14/EasyJailbreak_Datasets",
-            sub_datasets=[
-                "AdvBench",
-                "ForbiddenQuestion",
-                "MJP",
-                "MaliciousInstruct",
-                "QuestionList",
-            ],
-        ),
-        "JailbreakV-28K/JailBreakV-28k": load_jailbreak_v28k,
-        "Local CSV": load_local_csv,
-        "Custom CSV": load_local_csv,
-    }
-
-    # Dynamic dataset loaders
-    dynamic_loaders = {
-        "AgenticBackend": lambda opts: dataset_from_iterator(
-            "AgenticBackend",
-            fine_tuned.Module(
-                opts["datasets"], tools_inbox=tools_inbox, opts=opts
-            ).apply(),
-            lazy=True,
-        ),
-        "Steganography": lambda opts: list(
-            StenographyTransformer(opts["datasets"]).transform()
-        ),
-        "llm-adaptive-attacks": lambda opts: dataset_from_iterator(
-            "llm-adaptive-attacks",
-            adaptive_attacks.Module(
-                opts["datasets"], tools_inbox=tools_inbox, opts=opts
-            ).apply(),
-        ),
-        "Garak": lambda opts: dataset_from_iterator(
-            "Garak",
-            garak_tool.Module(
-                opts["datasets"], tools_inbox=tools_inbox, opts=opts
-            ).apply(),
-            lazy=True,
-        ),
-        "Reinforcement Learning Optimization": lambda opts: dataset_from_iterator(
-            "Reinforcement Learning Optimization",
-            rl_model.Module(
-                opts["datasets"], tools_inbox=tools_inbox, opts=opts
-            ).apply(),
-            lazy=True,
-        ),
-        "InspectAI": lambda opts: dataset_from_iterator(
-            "InspectAI",
-            inspect_ai_tool.Module([], tools_inbox=tools_inbox).apply(),
-            lazy=True,
-        ),
-        "GPT fuzzer": lambda opts: [],
-    }
-
-    datasets = []
-    options = options or [dict(datasets=datasets) for _ in dataset_names]
-
-    # Load base datasets
-    for name, opts in zip(dataset_names, options):
-        if name not in dataset_loaders:
-            continue
+    for file in csv_files:
         try:
-            datasets.append(dataset_loaders[name]())
+            df = pd.read_csv(file)
         except Exception as e:
-            logger.error(f"Error loading {name}: {e}")
-
-    # Load dynamic datasets and apply transformations
-    for name, opts in zip(dataset_names, options):
-        if name not in dynamic_loaders:
+            logger.error(f"Error reading {file}: {e}")
             continue
-        logger.info(f"Loading dynamic dataset {name} {opts}")
-        opts["datasets"] = datasets
-        try:
-            dynamic_result = dynamic_loaders[name](opts)
-            datasets.extend(dynamic_result)
-        except Exception as e:
-            logger.exception(f"Error loading dynamic {name}: {e}")
+        # Check if 'prompt' column exists
+        if "prompt" in df.columns:
+            prompt_list.extend(df["prompt"].tolist())
+        else:
+            logger.warning(f"File {file} does not contain a 'prompt' column")
+    return ProbeDataset(
+        dataset_name="Local CSV",
+        metadata={"src": str(csv_files)},
+        prompts=prompt_list,
+        tokens=count_words_in_list(prompt_list),
+        approx_cost=0.0,
+    )
 
-    # Load csv datasets and apply transformations
-    for name, opts in zip(dataset_names, options):
-        if not name.endswith(".csv"):
-            continue
-        logger.info(f"Loading csv dataset {name} {opts}")
-        datasets.append(load_csv(name))
 
-    return datasets
+def dataset_from_iterator(name: str, iterator, lazy=False) -> list:
+    """Convert an iterator into a list of prompts and create a ProbeDataset
+    object.
+
+    Args:
+        name (str): The name of the dataset.
+        iterator (iterator): An iterator that generates prompts.
+
+    Returns:
+        list: A list containing a single ProbeDataset object.
+    """
+    prompts = list(iterator) if not lazy else iterator
+    tokens = count_words_in_list(prompts) if not lazy else 0
+    dataset = ProbeDataset(
+        dataset_name=name,
+        metadata={},
+        prompts=prompts,
+        tokens=tokens,
+        approx_cost=0.0,
+        lazy=lazy,
+    )
+    return [dataset]
+
+
+# TODO: refactor this abstraction
+
+dataset_map_generics = {
+    "simonycl/aya-23-8B_advbench_jailbreak": lambda: load_dataset_general(
+        dataset_name="simonycl/aya-23-8B_advbench_jailbreak",
+        column_mappings={"prompt": "prompt"},
+    ),
+    "acmc/jailbreaks_dataset_with_perplexity_bigcode_starcoder2-3b_bigcode_starcoder2-7b": lambda: load_dataset_general(
+        dataset_name="acmc/jailbreaks_dataset_with_perplexity_bigcode_starcoder2-3b_bigcode_starcoder2-7b"
+    ),
+    "karanxa/dolphin-jailbreak-finetuning-dataset": lambda: load_dataset_general(
+        dataset_name="karanxa/dolphin-jailbreak-finetuning-dataset",
+        column_mappings={"prompt": "text"},
+    ),
+    "karanxa/llama-2-jailbreak-dataset": lambda: load_dataset_general(
+        dataset_name="karanxa/llama-2-jailbreak-dataset",
+        column_mappings={"prompt": "text"},
+    ),
+    "karanxa/llama2-uncensored-jailbreak-dataset-finetuning": lambda: load_dataset_general(
+        dataset_name="karanxa/llama2-uncensored-jailbreak-dataset-finetuning",
+        column_mappings={"prompt": "text"},
+    ),
+    "liuyanchen1015/Llama-3.2-1B_jailbreak_responses": lambda: load_dataset_general(
+        dataset_name="liuyanchen1015/Llama-3.2-1B_jailbreak_responses",
+        column_mappings={"prompt": "jailbreak_prompt_text"},
+    ),
+    "liuyanchen1015/Llama-3.2-1B-Instruct_jailbreak_responses": lambda: load_dataset_general(
+        dataset_name="liuyanchen1015/Llama-3.2-1B-Instruct_jailbreak_responses",
+        column_mappings={"prompt": "jailbreak_prompt_text"},
+    ),
+    "liuyanchen1015/Llama-3.2-1B-Instruct_jailbreak_responses_with_judgment": lambda: load_dataset_general(
+        dataset_name="liuyanchen1015/Llama-3.2-1B-Instruct_jailbreak_responses_with_judgment",
+        column_mappings={"prompt": "jailbreak_prompt_text"},
+    ),
+    "jackhhao/jailbreak-classification": lambda: load_dataset_general(
+        dataset_name="jackhhao/jailbreak-classification",
+        column_mappings={"prompt": "prompt"},
+    ),
+    "markush1/LLM-Jailbreak-Classifier": lambda: load_dataset_general(
+        dataset_name="markush1/LLM-Jailbreak-Classifier",
+        column_mappings={"prompt": "prompt"},
+    ),
+    "walledai/JailbreakBench": lambda: load_dataset_general(
+        dataset_name="walledai/JailbreakBench", column_mappings={"prompt": "prompt"}
+    ),
+    "walledai/JailbreakHub": lambda: load_dataset_general(
+        dataset_name="walledai/JailbreakHub", column_mappings={"prompt": "prompt"}
+    ),
+    "Granther/evil-jailbreak": lambda: load_dataset_general(
+        dataset_name="Granther/evil-jailbreak", column_mappings={"prompt": "text"}
+    ),
+    "sevdeawesome/jailbreak_success": lambda: load_dataset_general(
+        dataset_name="sevdeawesome/jailbreak_success",
+        column_mappings={"prompt": "jailbreak_prompt_text"},
+    ),
+    "IDA-SERICS/Disaster-tweet-jailbreaking": lambda: load_dataset_general(
+        dataset_name="IDA-SERICS/Disaster-tweet-jailbreaking",
+        column_mappings={"prompt": "prompt_attack"},
+    ),
+    "GeorgeDaDude/Jailbreak_Complete_DS_labeled": lambda: load_dataset_general(
+        dataset_name="GeorgeDaDude/Jailbreak_Complete_DS_labeled",
+        column_mappings={"prompt": "question"},
+    ),
+    "dayone3nder/jailbreak_prompt_JBB_sft_trainset": lambda: load_dataset_general(
+        dataset_name="dayone3nder/jailbreak_prompt_JBB_sft_trainset",
+        column_mappings={"prompt": "prompt"},
+    ),
+    "dayone3nder/general_safe_mix_jailbreak_prompt_JBB_trainset": lambda: load_dataset_general(
+        dataset_name="dayone3nder/general_safe_mix_jailbreak_prompt_JBB_trainset",
+        column_mappings={"prompt": "prompt"},
+    ),
+}

agentic_security/probe_data/modules/adaptive_attacks.py

@@ -3,8 +3,7 @@ import io
 
 import httpx
 import pandas as pd
-
-from agentic_security.logutils import logger
+from loguru import logger
 
 url = "https://raw.githubusercontent.com/tml-epfl/llm-adaptive-attacks/main/harmful_behaviors/harmful_behaviors_pair.csv"
 

agentic_security/probe_data/modules/fine_tuned.py

@@ -3,8 +3,7 @@ import os
 import uuid as U
 
 import httpx
-
-from agentic_security.logutils import logger
+from loguru import logger
 
 AUTH_TOKEN: str = os.getenv("AS_TOKEN", "gh0-5f4a8ed2-37c6-4bd7-a0cf-7070eae8115b")
 
@@ -66,7 +65,7 @@ class Module:
                 return {}
 
     async def fetch_prompts(self) -> list[str]:
-        api_url = "https://mcp.metaheuristic.co/infer"
+        api_url = "https://edge.metaheuristic.co/infer"
         headers = {
             "Authorization": f"Bearer {AUTH_TOKEN}",
             "Content-Type": "application/json",

agentic_security/probe_data/modules/garak_tool.py

@@ -4,7 +4,7 @@ import json
 import os
 import subprocess
 
-from agentic_security.logutils import logger
+from loguru import logger
 
 # TODO: add probes modules
 

agentic_security/probe_data/modules/inspect_ai_tool.py

@@ -2,7 +2,7 @@ import asyncio
 import importlib.util
 import os
 
-from agentic_security.logutils import logger
+from loguru import logger
 
 inspect_ai_task = (
     __file__.replace("inspect_ai_tool.py", "inspect_ai_task.py")

agentic_security/probe_data/modules/rl_model.py

@@ -8,8 +8,7 @@ from typing import Deque
 
 import numpy as np
 import requests
-
-from agentic_security.logutils import logger
+from loguru import logger
 
 AUTH_TOKEN: str = os.getenv("AS_TOKEN", "gh0-5f4a8ed2-37c6-4bd7-a0cf-7070eae8115b")
 
@@ -216,7 +215,7 @@ class Module:
         self.run_id = U.uuid4().hex
         self.batch_size = self.opts.get("batch_size", 500)
         self.rl_model = CloudRLPromptSelector(
-            prompt_groups, "https://mcp.metaheuristic.co", run_id=self.run_id
+            prompt_groups, "https://edge.metaheuristic.co", run_id=self.run_id
         )
 
     async def apply(self):

agentic_security/probe_data/modules/test_rl_model.py

@@ -36,7 +36,7 @@ def mock_requests() -> Mock:
 def mock_rl_selector() -> Mock:
     return CloudRLPromptSelector(
         dataset_prompts,
-        api_url="https://mcp.metaheuristic.co",
+        api_url="https://edge.metaheuristic.co",
     )
 
 
@@ -94,7 +94,7 @@ class TestCloudRLPromptSelector:
     def test_select_next_prompt_success_service(self, dataset_prompts):
         selector = CloudRLPromptSelector(
             dataset_prompts,
-            api_url="https://mcp.metaheuristic.co",
+            api_url="https://edge.metaheuristic.co",
         )
         next_prompt = selector.select_next_prompt(
             "How does RL work?", passed_guard=True

agentic_security/probe_data/msj_data.py

@@ -1,6 +1,6 @@
 from dataclasses import dataclass
 
-from cache_to_disk import cache_to_disk  # noqa
+from cache_to_disk import cache_to_disk
 
 
 # TODO: refactor this class to use from .data
@@ -22,7 +22,7 @@ class ProbeDataset:
         }
 
 
-# @cache_to_disk(n_days_to_cache=1)
+@cache_to_disk()
 def load_dataset_generic(name, getter=lambda x: x["train"]["prompt"]):
     from datasets import load_dataset
 

agentic_security/report_chart.py

@@ -7,61 +7,9 @@ import pandas as pd
 from matplotlib.cm import ScalarMappable
 from matplotlib.colors import LinearSegmentedColormap, Normalize
 
-from agentic_security.logutils import logger
 
-from .primitives import Table
-
-
-def plot_security_report(table: Table) -> io.BytesIO:
-    """
-    Generates a polar plot representing the security report based on the given data.
-
-    Args:
-        table (Table): The input data table containing security metrics.
-
-    Returns:
-        io.BytesIO: A buffer containing the generated plot image in PNG format.
-                    Returns an empty buffer in case of an error.
-    """
-    try:
-        return _plot_security_report(table=table)
-    except (TypeError, ValueError, OverflowError, IndexError, Exception) as e:
-        logger.error(f"Error in generating the security report: {e} {table}")
-    return io.BytesIO()
-
-
-def generate_identifiers(data: pd.DataFrame) -> list[str]:
-    """
-    Generates unique identifiers for the given dataset.
-
-    Args:
-        data (pd.DataFrame): A pandas DataFrame containing security-related data.
-
-    Returns:
-        list[str]: A list of generated identifiers. Returns a list with an empty string in case of an error.
-    """
-    return _generate_identifiers(data=data)
-
-
-def _plot_security_report(table: Table) -> io.BytesIO:
-    """
-    Generates a polar plot-based security report visualizing the failure rates for different modules.
-
-    This function processes the input data, sorts it by failure rate, and generates a polar plot
-    where each bar represents the failure rate for a specific module. The plot includes identifiers,
-    color-coding based on token count, failure rate values on the bars, and a table listing the modules
-    and their corresponding failure rates.
-
-    Args:
-        table (Table): A table-like structure (e.g., pandas DataFrame) containing security report data
-                        with columns for failure rate, tokens, and modules.
-
-    Returns:
-        io.BytesIO: A buffer containing the generated plot image in PNG format.
-    """
+def plot_security_report(table):
     # Data preprocessing
-    logger.info("Data preprocessing started.")
-
     data = pd.DataFrame(table)
 
     # Sort by failure rate and reset index
@@ -72,10 +20,10 @@ def _plot_security_report(table: Table) -> io.BytesIO:
     fig, ax = plt.subplots(figsize=(12, 10), subplot_kw={"projection": "polar"})
     fig.set_facecolor("#f0f0f0")
     ax.set_facecolor("#f0f0f0")
-    logger.info("Plot setup complete.")
 
     # Styling parameters
     colors = ["#6C5B7B", "#C06C84", "#F67280", "#F8B195"][::-1]  # Pastel palette
+    # colors = ["#440154", "#3b528b", "#21908c", "#5dc863"]  # Viridis-inspired palette
     cmap = LinearSegmentedColormap.from_list("custom", colors, N=256)
     norm = Normalize(vmin=data["tokens"].min(), vmax=data["tokens"].max())
 
@@ -126,10 +74,7 @@ def _plot_security_report(table: Table) -> io.BytesIO:
 
     # Title and caption
     fig.suptitle(
-        "Security Report for Different Modules",
-        fontsize=16,
-        fontweight="bold",
-        y=1.02,
+        "Security Report for Different Modules", fontsize=16, fontweight="bold", y=1.02
     )
     caption = "Report generated by https://github.com/msoedov/agentic_security"
     fig.text(
@@ -167,12 +112,17 @@ def _plot_security_report(table: Table) -> io.BytesIO:
             data["identifier"], data["failureRate"], data["module"]
         )
     ]
-    table = ax.table(cellText=table_data, loc="right", cellLoc="left")
+    table = ax.table(
+        cellText=table_data,
+        loc="right",
+        cellLoc="left",
+    )
     table.auto_set_font_size(False)
     table.set_fontsize(8)
 
     # Adjust table style
     table.scale(1, 0.7)
+
     for (row, col), cell in table.get_celld().items():
         cell.set_edgecolor("none")
         cell.set_facecolor("#f0f0f0" if row % 2 == 0 else "#e0e0e0")
@@ -182,33 +132,17 @@ def _plot_security_report(table: Table) -> io.BytesIO:
             cell.set_text_props(fontweight="bold")
 
     # Adjust layout and save
+
     plt.tight_layout()
     buf = io.BytesIO()
     plt.savefig(buf, format="png", dpi=300, bbox_inches="tight")
     plt.close(fig)
     buf.seek(0)
-    logger.info("Report successfully generated and saved to buffer.")
     return buf
 
 
-def _generate_identifiers(data: pd.DataFrame) -> list[str]:
-    """
-    Generates a list of unique identifiers for each row in the given DataFrame.
-
-    The identifiers are based on the English alphabet, with each identifier consisting
-    of a letter followed by a number. The letter represents the "group" of identifiers
-    (using a letter from A to Z) and the number is a counter within that group. For example:
-    - A1, A2, ..., A26, B1, B2, ..., Z1, Z2, ...
-
-    Args:
-        data (pd.DataFrame): The input DataFrame containing data for which identifiers
-                             are to be generated.
-
-    Returns:
-        list[str]: A list of unique identifiers as strings, one for each row in the DataFrame.
-    """
+def generate_identifiers(data):
     data_length = len(data)
-
     alphabet = string.ascii_uppercase
     num_letters = len(alphabet)
 

agentic_security/routes/_specs.py

@@ -1,397 +0,0 @@
-_SPECS = [
-    """POST ${SELF_URL}/v1/self-probe
-Authorization: Bearer XXXXX
-Content-Type: application/json
-
-{
-"prompt": "<<PROMPT>>"
-}
-
-""",
-    """POST https://api.openai.com/v1/chat/completions
-Authorization: Bearer $OPENAI_API_KEY
-Content-Type: application/json
-
-{
-"model": "gpt-3.5-turbo",
-"messages": [{"role": "user", "content": "<<PROMPT>>"}],
-"temperature": 0.7
-}
-""",
-    """
-POST https://api.deepseek.com/chat/completions
-Authorization: Bearer $DEEPSEEK_API_KEY
-Content-Type: application/json
-
-{
-  "model": "deepseek-chat",
-  "messages": [
-    {"role": "system", "content": "You are a helpful assistant."},
-    {"role": "user", "content": "<<PROMPT>>"}
-  ],
-  "stream": false
-}
-""",
-    """POST https://api.replicate.com/v1/models/mistralai/mixtral-8x7b-instruct-v0.1/predictions
-Authorization: Bearer $APIKEY
-Content-Type: application/json
-
-{
-"input": {
-"top_k": 50,
-"top_p": 0.9,
-"prompt": "Write a bedtime story about neural networks I can read to my toddler",
-"temperature": 0.6,
-"max_new_tokens": 1024,
-"prompt_template": "<s>[INST] <<PROMPT>> [/INST] ",
-"presence_penalty": 0,
-"frequency_penalty": 0
-}
-}
-""",
-    """POST https://api.groq.com/v1/request_manager/text_completion
-Authorization: Bearer $APIKEY
-Content-Type: application/json
-
-{
-"model_id": "codellama-34b",
-"system_prompt": "You are helpful and concise coding assistant",
-"user_prompt": "<<PROMPT>>"
-}
-""",
-    """POST https://api.together.xyz/v1/chat/completions
-Authorization: Bearer $TOGETHER_API_KEY
-Content-Type: application/json
-
-{
-"model": "mistralai/Mixtral-8x7B-Instruct-v0.1",
-"messages": [
-{"role": "system", "content": "You are an expert travel guide"},
-{"role": "user", "content": "<<PROMPT>>"}
-]
-}
-""",
-    """POST ${SELF_URL}/v1/self-probe-image
-Authorization: Bearer XXXXX
-Content-Type: application/json
-
-[
-    {
-      "role": "user",
-      "content": [
-        {
-          "type": "text",
-          "text": "What is in this image?",
-        },
-        {
-          "type": "image_url",
-          "image_url": {
-            "url":  f"data:image/jpeg;base64,{<<BASE64_IMAGE>>}"
-          },
-        },
-      ],
-    }
-]
-""",
-    """POST ${SELF_URL}/v1/self-probe-file
-Authorization: Bearer $API_KEY
-Content-Type: multipart/form-data
-
-{
-  "file": "@./sample_audio.m4a",
-  "model": "whisper-large-v3"
-}
-""",
-    """POST https://api.gemini.com/v1/generate
-Authorization: Bearer $GEMINI_API_KEY
-Content-Type: application/json
-
-{
-  "model": "gemini-latest",
-  "prompt": "<<PROMPT>>",
-  "temperature": 0.8,
-  "max_tokens": 150,
-  "top_p": 1.0,
-  "frequency_penalty": 0,
-  "presence_penalty": 0
-}
-""",
-    """POST https://api.anthropic.com/v1/complete
-Authorization: Bearer $ANTHROPIC_API_KEY
-Content-Type: application/json
-
-{
-  "model": "claude-v1.3",
-  "prompt": "<<PROMPT>>",
-  "temperature": 0.7,
-  "max_tokens_to_sample": 256,
-  "stop_sequences": ["\n\nHuman:"]
-}
-""",
-    """POST https://api.cohere.ai/generate
-Authorization: Bearer $COHERE_API_KEY
-Content-Type: application/json
-
-{
-  "model": "command-xlarge-nightly",
-  "prompt": "<<PROMPT>>",
-  "max_tokens": 300,
-  "temperature": 0.75,
-  "k": 0,
-  "p": 0.75
-}
-""",
-    """POST https://<<RESOURCE_NAME>>.openai.azure.com/openai/deployments/<<DEPLOYMENT_NAME>>/completions?api-version=2023-06-01-preview
-Authorization: Bearer $AZURE_API_KEY
-Content-Type: application/json
-
-{
-  "prompt": "<<PROMPT>>",
-  "max_tokens": 150,
-  "temperature": 0.7,
-  "top_p": 0.9,
-  "frequency_penalty": 0,
-  "presence_penalty": 0
-}
-""",
-    """POST https://api.assemblyai.com/v2/transcript
-Authorization: Bearer $ASSEMBLY_API_KEY
-Content-Type: application/json
-
-{
-  "audio_url": "<<AUDIO_FILE_URL>>"
-}
-""",
-    """POST https://api.openrouter.ai/v1/chat/completions
-Authorization: Bearer $OPENROUTER_API_KEY
-Content-Type: application/json
-
-{
-  "model": "openrouter-latest",
-  "prompt": "<<PROMPT>>",
-  "temperature": 0.7,
-  "max_tokens": 150,
-  "top_p": 0.9,
-  "frequency_penalty": 0,
-  "presence_penalty": 0
-}
-""",
-]
-
-
-LLM_SPECS = [
-    """POST ${SELF_URL}/v1/self-probe
-Authorization: Bearer XXXXX
-Content-Type: application/json
-
-{
-"prompt": "<<PROMPT>>"
-}
-
-""",
-    """POST https://api.openai.com/v1/chat/completions
-Authorization: Bearer $OPENAI_API_KEY
-Content-Type: application/json
-
-{
-"model": "gpt-3.5-turbo",
-"messages": [{"role": "user", "content": "<<PROMPT>>"}],
-"temperature": 0.7
-}
-""",
-    """
-POST https://api.deepseek.com/chat/completions
-Authorization: Bearer $DEEPSEEK_API_KEY
-Content-Type: application/json
-
-{
-  "model": "deepseek-chat",
-  "messages": [
-    {"role": "system", "content": "You are a helpful assistant."},
-    {"role": "user", "content": "<<PROMPT>>"}
-  ],
-  "stream": false
-}
-""",
-    """POST https://api.replicate.com/v1/models/mistralai/mixtral-8x7b-instruct-v0.1/predictions
-Authorization: Bearer $APIKEY
-Content-Type: application/json
-
-{
-"input": {
-"top_k": 50,
-"top_p": 0.9,
-"prompt": "Write a bedtime story about neural networks I can read to my toddler",
-"temperature": 0.6,
-"max_new_tokens": 1024,
-"prompt_template": "<s>[INST] <<PROMPT>> [/INST] ",
-"presence_penalty": 0,
-"frequency_penalty": 0
-}
-}
-""",
-    """POST https://api.groq.com/v1/request_manager/text_completion
-Authorization: Bearer $APIKEY
-Content-Type: application/json
-
-{
-"model_id": "codellama-34b",
-"system_prompt": "You are helpful and concise coding assistant",
-"user_prompt": "<<PROMPT>>"
-}
-""",
-    """POST https://api.together.xyz/v1/chat/completions
-Authorization: Bearer $TOGETHER_API_KEY
-Content-Type: application/json
-
-{
-"model": "mistralai/Mixtral-8x7B-Instruct-v0.1",
-"messages": [
-{"role": "system", "content": "You are an expert travel guide"},
-{"role": "user", "content": "<<PROMPT>>"}
-]
-}
-""",
-    """POST ${SELF_URL}/v1/self-probe-image
-Authorization: Bearer XXXXX
-Content-Type: application/json
-
-[
-    {
-      "role": "user",
-      "content": [
-        {
-          "type": "text",
-          "text": "What is in this image?",
-        },
-        {
-          "type": "image_url",
-          "image_url": {
-            "url":  f"data:image/jpeg;base64,{<<BASE64_IMAGE>>}"
-          },
-        },
-      ],
-    }
-]
-""",
-    """POST ${SELF_URL}/v1/self-probe-file
-Authorization: Bearer $API_KEY
-Content-Type: multipart/form-data
-
-{
-  "file": "@./sample_audio.m4a",
-  "model": "whisper-large-v3"
-}
-""",
-    """POST https://api.gemini.com/v1/generate
-Authorization: Bearer $GEMINI_API_KEY
-Content-Type: application/json
-
-{
-  "model": "gemini-latest",
-  "prompt": "<<PROMPT>>",
-  "temperature": 0.8,
-  "max_tokens": 150,
-  "top_p": 1.0,
-  "frequency_penalty": 0,
-  "presence_penalty": 0
-}
-""",
-    """POST https://api.anthropic.com/v1/complete
-Authorization: Bearer $ANTHROPIC_API_KEY
-Content-Type: application/json
-
-{
-  "model": "claude-v1.3",
-  "prompt": "<<PROMPT>>",
-  "temperature": 0.7,
-  "max_tokens_to_sample": 256,
-  "stop_sequences": ["\n\nHuman:"]
-}
-""",
-    """POST https://api.cohere.ai/generate
-Authorization: Bearer $COHERE_API_KEY
-Content-Type: application/json
-
-{
-  "model": "command-xlarge-nightly",
-  "prompt": "<<PROMPT>>",
-  "max_tokens": 300,
-  "temperature": 0.75,
-  "k": 0,
-  "p": 0.75
-}
-""",
-    """POST https://<<RESOURCE_NAME>>.openai.azure.com/openai/deployments/<<DEPLOYMENT_NAME>>/completions?api-version=2023-06-01-preview
-Authorization: Bearer $AZURE_API_KEY
-Content-Type: application/json
-
-{
-  "prompt": "<<PROMPT>>",
-  "max_tokens": 150,
-  "temperature": 0.7,
-  "top_p": 0.9,
-  "frequency_penalty": 0,
-  "presence_penalty": 0
-}
-""",
-    """POST https://api.assemblyai.com/v2/transcript
-Authorization: Bearer $ASSEMBLY_API_KEY
-Content-Type: application/json
-
-{
-  "audio_url": "<<AUDIO_FILE_URL>>"
-}
-""",
-    """POST https://api.openrouter.ai/v1/chat/completions
-Authorization: Bearer $OPENROUTER_API_KEY
-Content-Type: application/json
-
-{
-  "model": "openrouter-latest",
-  "prompt": "<<PROMPT>>",
-  "temperature": 0.7,
-  "max_tokens": 150,
-  "top_p": 0.9,
-  "frequency_penalty": 0,
-  "presence_penalty": 0
-}
-""",
-]
-
-
-LLM_CONFIGS = [
-    {
-        "name": "Custom API",
-        "prompts": 40000,
-        "customInstructions": "Requires api spec",
-        "logo": "/icons/myshell.png",
-    },
-    {"name": "Open AI", "prompts": 24000, "logo": "/icons/openai.png"},
-    {"name": "Deepseek v1", "prompts": 24000, "logo": "/icons/deepseek.png"},
-    {"name": "Replicate", "prompts": 40000, "logo": "/icons/replicate.png"},
-    {"name": "Groq", "prompts": 40000, "logo": "/icons/groq.png"},
-    {"name": "Together.ai", "prompts": 40000, "logo": "/icons/together.png"},
-    {
-        "name": "Custom API Image",
-        "prompts": 40000,
-        "customInstructions": "Requires api spec",
-        "modality": "Image",
-        "logo": "/icons/myshell.png",
-    },
-    {
-        "name": "Custom API Files",
-        "prompts": 40000,
-        "customInstructions": "Requires api spec",
-        "modality": "Files",
-        "logo": "/icons/myshell.png",
-    },
-    {"name": "Gemini", "prompts": 40000, "logo": "/icons/gemini.png"},
-    {"name": "Claude", "prompts": 40000, "logo": "/icons/claude.png"},
-    {"name": "Cohere", "prompts": 40000, "logo": "/icons/cohere.png"},
-    {"name": "Azure OpenAI", "prompts": 40000, "logo": "/icons/azureai.png"},
-    {"name": "assemblyai", "prompts": 40000, "logo": "/icons/myshell.png"},
-    {"name": "OpenRouter.ai", "prompts": 40000, "logo": "/icons/openrouter.png"},
-]
-
-LLM_SPECS = [dict(spec=spec, **d) for spec, d in zip(_SPECS, LLM_CONFIGS)]

agentic_security/routes/probe.py

@@ -6,7 +6,6 @@ from fastapi.responses import JSONResponse
 from ..primitives import FileProbeResponse, Probe
 from ..probe_actor.refusal import REFUSAL_MARKS
 from ..probe_data import REGISTRY
-from ._specs import LLM_SPECS
 
 router = APIRouter()
 
@@ -74,12 +73,6 @@ async def data_config():
     return [m for m in REGISTRY]
 
 
-@router.get("/v1/llm-specs", response_model=list)
-def get_llm_specs():
-    """Returns the LLM API specifications."""
-    return LLM_SPECS
-
-
 @router.get("/health")
 async def health_check():
     """Health check endpoint."""

agentic_security/routes/proxy.py

@@ -2,8 +2,7 @@ import random
 from asyncio import Event
 
 from fastapi import APIRouter
-
-from agentic_security.logutils import logger
+from loguru import logger
 
 from ..core.app import get_current_run, get_tools_inbox
 from ..primitives import CompletionRequest, Settings

agentic_security/routes/scan.py

@@ -1,6 +1,4 @@
-from collections.abc import Generator
 from datetime import datetime
-from typing import Any
 
 from fastapi import (
     APIRouter,
@@ -12,12 +10,11 @@ from fastapi import (
     UploadFile,
 )
 from fastapi.responses import StreamingResponse
-
-from agentic_security.logutils import logger
+from loguru import logger
 
 from ..core.app import get_stop_event, get_tools_inbox, set_current_run
 from ..dependencies import InMemorySecrets, get_in_memory_secrets
-from ..http_spec import InvalidHTTPSpecError, LLMSpec
+from ..http_spec import LLMSpec
 from ..primitives import LLMInfo, Scan
 from ..probe_actor import fuzzer
 
@@ -27,12 +24,10 @@ router = APIRouter()
 @router.post("/verify")
 async def verify(
     info: LLMInfo, secrets: InMemorySecrets = Depends(get_in_memory_secrets)
-) -> dict[str, int | str | float]:
+):
     spec = LLMSpec.from_string(info.spec)
     try:
         r = await spec.verify()
-    except InvalidHTTPSpecError as e:
-        raise HTTPException(status_code=400, detail=str(e))
     except Exception as e:
         logger.exception(e)
         raise HTTPException(status_code=400, detail=str(e))
@@ -47,7 +42,7 @@ async def verify(
     )
 
 
-def streaming_response_generator(scan_parameters: Scan) -> Generator[str, Any, None]:
+def streaming_response_generator(scan_parameters: Scan):
     request_factory = LLMSpec.from_string(scan_parameters.llmSpec)
     set_current_run(request_factory)
 
@@ -68,7 +63,7 @@ async def scan(
     scan_parameters: Scan,
     background_tasks: BackgroundTasks,
     secrets: InMemorySecrets = Depends(get_in_memory_secrets),
-) -> StreamingResponse:
+):
     scan_parameters.with_secrets(secrets)
     return StreamingResponse(
         streaming_response_generator(scan_parameters), media_type="application/json"
@@ -76,7 +71,7 @@ async def scan(
 
 
 @router.post("/stop")
-async def stop_scan() -> dict[str, str]:
+async def stop_scan():
     get_stop_event().set()
     return {"status": "Scan stopped"}
 
@@ -90,7 +85,7 @@ async def scan_csv(
     maxBudget: int = Query(10_000),
     enableMultiStepAttack: bool = Query(False),
     secrets: InMemorySecrets = Depends(get_in_memory_secrets),
-) -> StreamingResponse:
+):
     # TODO: content dataset to fuzzer
     content = await file.read()  # noqa
     llm_spec = await llmSpec.read()

agentic_security/routes/telemetry.py

@@ -1,8 +1,7 @@
 import sentry_sdk
+from loguru import logger
 from sentry_sdk.integrations.logging import ignore_logger
 
-from agentic_security.logutils import logger
-
 from ..primitives import Settings
 
 

agentic_security/static/base.js

@@ -101,7 +101,7 @@ Content-Type: application/json
 ]
 `,
   `POST ${SELF_URL}/v1/self-probe-file
-Authorization: Bearer $API_KEY
+Authorization: Bearer $GROQ_API_KEY
 Content-Type: multipart/form-data
 
 {
@@ -204,9 +204,8 @@ let LLM_CONFIGS = [
   { name: 'Claude', prompts: 40000, logo: '/icons/claude.png' },
   { name: 'Cohere', prompts: 40000, logo: '/icons/cohere.png' },
   { name: 'Azure OpenAI', prompts: 40000, logo: '/icons/azureai.png' },
-  { name: 'assemblyai', prompts: 40000, logo: fallbackIcon },
   { name: 'OpenRouter.ai', prompts: 40000, logo: '/icons/openrouter.png' },
-
+  { name: 'assemblyai', prompts: 40000, logo: fallbackIcon },
 ];
 function has_image(spec) {
   return spec.includes('<<BASE64_IMAGE>>');

agentic_security/static/index.html

@@ -63,32 +63,6 @@
           </div>
         </section>
       </main>
-      <div class="fixed top-6 right-6 z-50 space-y-3">
-        <transition-group name="toast">
-            <div
-                v-for="toast in toasts"
-                :key="toast.id"
-                class="flex items-center p-3 rounded-xl shadow-xl text-white max-w-md animate-toast-in border border-opacity-30"
-                :class="{
-                    'bg-success-toast border-accent-green': toast.type === 'success',
-                    'bg-error-toast border-accent-red': toast.type === 'error',
-                    'bg-info-toast border-accent-orange': toast.type === 'info'
-                }"
-            >
-                <span class="flex-1 font-medium tracking-wide text-sm">{{ toast.message }}</span>
-                <button
-                    @click="removeToast(toast.id)"
-                    class="ml-3 focus:outline-none hover:opacity-80 transition-opacity"
-                >
-                    <svg class="w-4 h-4" fill="none" stroke="currentColor" viewBox="0 0 24 24">
-                        <path stroke-linecap="round" stroke-linejoin="round" stroke-width="2" d="M6 18L18 6M6 6l12 12" />
-                    </svg>
-                </button>
-            </div>
-        </transition-group>
-    </div>
-
-
       <main class="max-w-6xl mx-auto space-y-8">
         <section class="bg-dark-card rounded-lg p-6 shadow-lg" v-show="false">
           <h2 class="text-2xl font-bold mb-4">Select a Config</h2>

agentic_security/static/main.js

@@ -25,8 +25,6 @@ var app = new Vue({
         showModules: false,
         showLogs: false,
         showConsentModal: true,
-        toasts: [], // Array to store toast notifications
-        toastTimeout: 3000, // Duration in milliseconds (3 seconds)
         statusDotClass: 'bg-gray-500', // Default status dot class
         statusText: 'Verified', // Default status text
         statusClass: 'bg-green-500 text-dark-bg', // Default status class
@@ -95,36 +93,21 @@ var app = new Vue({
 
     },
     methods: {
-        showToast(message, type = 'success') {
-            const id = Date.now(); // Unique ID for each toast
-            this.toasts.push({ id, message, type });
-
-            // Automatically remove toast after timeout
-            setTimeout(() => {
-                this.removeToast(id);
-            }, this.toastTimeout);
-        },
-
-        removeToast(id) {
-            this.toasts = this.toasts.filter(toast => toast.id !== id);
-        },
         focusTextarea() {
             this.isFocused = true;
-            // Remove 'self' assignment if not used elsewhere
+            self = this.$refs;
             this.$nextTick(() => {
-                this.$refs.textarea.focus();
-                this.adjustHeight({ target: this.$refs.textarea });
+                // Focus the textarea after rendering
+                self.textarea.focus();
+                this.adjustHeight({ target: self.textarea });
             });
-            // Correct the event listener to use handleOutsideClick
-            document.addEventListener("mousedown", this.handleOutsideClick);
+            document.addEventListener("mousedown", this.handleClickOutside);
+
         },
         handleOutsideClick(event) {
-            if (!this.$refs.textarea) {
-                return
-            }
-            if (!this.$refs.textarea.contains(event.target)) {
+            if (!this.$refs.container.contains(event.target)) {
                 this.isFocused = false;
-                document.removeEventListener("mousedown", this.handleOutsideClick);
+                document.removeEventListener("mousedown", this.handleClickOutside);
             }
         },
         unfocusTextarea() {
@@ -132,12 +115,7 @@ var app = new Vue({
         },
         acceptConsent() {
             this.showConsentModal = false; // Close the modal
-
-            try {
-                localStorage.setItem('consentGiven', 'true'); // Save consent to local storage
-            } catch (e) {
-                this.showToast('Failed to save consent', 'error'); // Show error if saving fails
-            }
+            localStorage.setItem('consentGiven', 'true'); // Save consent to local storage
         },
 
         saveStateToLocalStorage() {
@@ -178,7 +156,6 @@ var app = new Vue({
             this.integrationVerified = false;
             this.showResetConfirmation = false;
             this.enableMultiStepAttack = false;
-            this.showToast('All settings have been reset to default', 'info');
         },
         confirmResetState() {
             this.showResetConfirmation = true;
@@ -217,39 +194,31 @@ var app = new Vue({
                 spec: this.modelSpec,
             };
             let startTime = performance.now(); // Capture start time
-
-            try {
-                const response = await fetch(`${SELF_URL}/verify`, {
-                    method: 'POST',
-                    headers: {
-                        'Content-Type': 'application/json',
-                    },
-                    body: JSON.stringify(payload),
-                });
-
-                let r = await response.json();
-
-                let endTime = performance.now(); // Capture end time
-                let latency = ((endTime - startTime) / 1000).toFixed(3); // Calculate latency in milliseconds
-                this.latency = latency;
-
-                if (!response.ok) {
-                    this.updateStatusDot(false);
-                    this.errorMsg = 'Integration verification failed:' + JSON.stringify(r);
-                    this.showToast('Integration verification failed', 'error');
-                } else {
-                    this.errorMsg = '';
-                    this.updateStatusDot(true);
-                    this.okMsg = 'Integration verified';
-                    this.showToast('Integration verified successfully', 'success');
-                    this.integrationVerified = true;
-                }
-            } catch (error) {
+            const response = await fetch(`${SELF_URL}/verify`, {
+                method: 'POST',
+                headers: {
+                    'Content-Type': 'application/json',
+                },
+                body: JSON.stringify(payload),
+            });
+            console.log(response);
+            let r = await response.json();
+            let endTime = performance.now(); // Capture end time
+            let latency = endTime - startTime; // Calculate latency in milliseconds
+            latency = latency.toFixed(3) / 1000; // Round to 2 decimal places
+            this.latency = latency;
+            if (!response.ok) {
+                this.updateStatusDot(false);
+                this.errorMsg = 'Integration verification failed:' + JSON.stringify(r);
+            } else {
+                this.errorMsg = '';
                 this.updateStatusDot(true);
-                this.errorMsg = 'Server unreachable';
-                this.showToast('Network error', 'error');
-            }
+                this.okMsg = 'Integration verified';
+                this.integrationVerified = true;
+                // console.log('Integration verified', this.integrationVerified);
+                // this.$forceUpdate();
 
+            }
             this.saveStateToLocalStorage();
         },
         loadConfigs: async function () {
@@ -271,7 +240,6 @@ var app = new Vue({
             this.errorMsg = '';
             this.okMsg = '';
             this.integrationVerified = false;
-            this.showToast(`Config ${index + 1} selected`, 'info');
         },
         toggleModules() {
             this.showModules = !this.showModules;
@@ -359,7 +327,6 @@ var app = new Vue({
                 return
             }
             console.log('New row');
-            this.showToast('New module', 'success');
             let payload = {
                 table: this.mainTable,
             };
@@ -470,8 +437,6 @@ var app = new Vue({
                     }
                 });
             }
-            this.scanRunning = false;
-            this.showToast('Scan finished successfully', 'success');
             this.saveStateToLocalStorage();
 
         }

agentic_security/static/partials/head.html

@@ -19,7 +19,7 @@
               technopollas: ['Technopollas', 'sans-serif'],
             },
             colors: {
-                dark: {
+                t1: {
                 bg: '#0D0D0D', // Jet Black
                 card: '#1A1A1A', // Dark Carbon Fiber
                 text: '#FFFFFF',
@@ -30,7 +30,7 @@
                   yellow: '#2E4053', // Jet Black
                 },
               },
-              dark1: {
+              dark: {
                 bg: '#121212',
                 card: '#1E1E1E',
                 text: '#FFFFFF',
@@ -86,7 +86,6 @@
         }
       }
     </script>
-
     <style>
 .scrollbar-hide::-webkit-scrollbar {
   display: none;
@@ -97,55 +96,4 @@
 }
 
     </style>
-<style>
-/* Toast-specific colors */
-.bg-success-toast {
-    background: #1C3F74
-}
-
-.bg-error-toast {
-    background: #85144B
-}
-
-.bg-info-toast {
-    background: #FFC300
-}
-
-.border-accent-green, .border-accent-red, .border-accent-orange {
-    border-color: rgba(255, 255, 255, 0.1); /* Subtle white border for depth */
-}
-
-/* Animation */
-.animate-toast-in {
-    animation: tSlideIn 0.4s cubic-bezier(0.25, 0.1, 0.25, 1);
-}
-
-@keyframes tSlideIn {
-    from {
-        transform: translateX(120%) scale(0.95);
-        opacity: 0;
-    }
-    to {
-        transform: translateX(0) scale(1);
-        opacity: 1;
-    }
-}
-
-.toast-enter-active,
-.toast-leave-active {
-    transition: all 0.4s cubic-bezier(0.25, 0.1, 0.25, 1);
-}
-
-.toast-enter-from,
-.toast-leave-to {
-    transform: translateX(120%) scale(0.95);
-    opacity: 0;
-}
-
-/* Hover effect */
-[toast-type]:hover {
-    transform: translateY(-2px);
-    box-shadow: 0 6px 20px rgba(0, 0, 0, 0.6);
-}
-</style>
   </head>

agentic_security/test_lib.py

@@ -8,7 +8,7 @@ import time
 import pytest
 
 import agentic_security.test_spec_assets as test_spec_assets
-from agentic_security.lib import SecurityScanner
+from agentic_security.lib import AgenticSecurity
 
 
 def has_module(module_name):
@@ -57,7 +57,7 @@ class TestLibraryLevel:
         maxBudget = 1000000
         max_th = 0.3
         datasets = make_test_registry()
-        result = SecurityScanner.scan(llmSpec, maxBudget, datasets, max_th)
+        result = AgenticSecurity.scan(llmSpec, maxBudget, datasets, max_th)
         assert isinstance(result, dict)
         print(result)
         assert len(result) in [0, 1]
@@ -68,7 +68,7 @@ class TestLibraryLevel:
         maxBudget = 1000
         max_th = 0.3
         datasets = make_test_registry()
-        result = SecurityScanner.scan(
+        result = AgenticSecurity.scan(
             llmSpec, maxBudget, datasets, max_th, enableMultiStepAttack=True
         )
         assert isinstance(result, dict)
@@ -93,7 +93,7 @@ class TestLibraryLevel:
                 "opts": {"port": 9094},
             },
         ]
-        result = SecurityScanner.scan(llmSpec, maxBudget, datasets, max_th)
+        result = AgenticSecurity.scan(llmSpec, maxBudget, datasets, max_th)
         assert isinstance(result, dict)
         print(result)
         assert len(result) in [0, 1]
@@ -120,7 +120,7 @@ class TestLibraryLevel:
                 "modality": "text",
             },
         ]
-        result = SecurityScanner.scan(llmSpec, maxBudget, datasets, max_th)
+        result = AgenticSecurity.scan(llmSpec, maxBudget, datasets, max_th)
         assert isinstance(result, dict)
         print(result)
         assert len(result) in [0, 1]
@@ -148,7 +148,7 @@ class TestLibraryLevel:
                 "modality": "text",
             },
         ]
-        result = SecurityScanner.scan(llmSpec, maxBudget, datasets, max_th)
+        result = AgenticSecurity.scan(llmSpec, maxBudget, datasets, max_th)
         assert isinstance(result, dict)
         print(result)
         assert len(result) in [0, 1]
@@ -164,10 +164,10 @@ class TestEntrypointCI:
             temp_path = os.path.join(tmpdir, "custom_agesec.toml")
 
             # Override default_path to the temporary path
-            SecurityScanner.default_path = temp_path
+            AgenticSecurity.default_path = temp_path
 
             # Generate the default configuration
-            security = SecurityScanner()
+            security = AgenticSecurity()
             security.generate_default_settings()
 
             # Check that the config file was created at the temporary path
@@ -189,17 +189,17 @@ class TestEntrypointCI:
             temp_path = os.path.join(tmpdir, "custom_agesec.toml")
 
             # Override default_path to the temporary path
-            SecurityScanner.default_path = temp_path
+            AgenticSecurity.default_path = temp_path
 
             # Generate the default configuration
-            security = SecurityScanner()
+            security = AgenticSecurity()
             security.generate_default_settings()
 
             # Load the generated configuration
-            SecurityScanner.load_config(temp_path)
+            AgenticSecurity.load_config(temp_path)
 
             # Validate loaded configuration
-            config = SecurityScanner.config
+            config = AgenticSecurity.config
             assert (
                 config["general"]["maxBudget"] == 1000000
             ), "maxBudget should be 1000000"

codebeaver.yml

@@ -0,0 +1,2 @@
+from: python-pytest-poetry
+# This file was generated automatically by CodeBeaver based on your repository. Learn how to customize it here: https://docs.codebeaver.ai/configuration/

docs/http_spec.md

@@ -33,7 +33,7 @@ The `LLMSpec` class is the core of the HTTP specification. It provides the follo
 ### Methods
 
 - **`from_string(http_spec: str) -> LLMSpec`**: Parses an HTTP specification string into an `LLMSpec` object.
-- **`validate(prompt: str, encoded_image: str, encoded_audio: str, files: dict) -> null`**: Validates the request parameters based on the specified modality.
+- **`validate(prompt: str, encoded_image: str, encoded_audio: str, files: dict) -> None`**: Validates the request parameters based on the specified modality.
 - **`probe(prompt: str, encoded_image: str = "", encoded_audio: str = "", files: dict = {}) -> httpx.Response`**: Sends an HTTP request using the specified parameters.
 - **`verify() -> httpx.Response`**: Verifies the HTTP specification by sending a test request.
 
@@ -52,11 +52,12 @@ Authorization: Bearer sk-xxxxxxxxx
 Content-Type: application/json
 
 {
-  "model": "gpt-3.5-turbo",
-  "messages": [{"role": "user", "content": "<<PROMPT>>"}],
-  "temperature": 0.7
+    "model": "gpt-3.5-turbo",
+    "messages": [{"role": "user", "content": "<<PROMPT>>"}],
+    "temperature": 0.7
 }
 """
+
 spec = LLMSpec.from_string(http_spec)
 response = await spec.probe("What is the capital of France?")
 ```
@@ -70,11 +71,12 @@ Authorization: Bearer sk-xxxxxxxxx
 Content-Type: application/json
 
 {
-  "model": "gpt-4-vision-preview",
-  "messages": [{"role": "user", "content": "What is in this image? <<BASE64_IMAGE>>"}],
-  "temperature": 0.7
+    "model": "gpt-4-vision-preview",
+    "messages": [{"role": "user", "content": "What is in this image? <<BASE64_IMAGE>>"}],
+    "temperature": 0.7
 }
 """
+
 spec = LLMSpec.from_string(http_spec)
 encoded_image = encode_image_base64_by_url("https://example.com/image.jpg")
 response = await spec.probe("What is in this image?", encoded_image=encoded_image)
@@ -89,11 +91,12 @@ Authorization: Bearer sk-xxxxxxxxx
 Content-Type: application/json
 
 {
-  "model": "whisper-large-v3",
-  "messages": [{"role": "user", "content": "Transcribe this audio: <<BASE64_AUDIO>>"}],
-  "temperature": 0.7
+    "model": "whisper-large-v3",
+    "messages": [{"role": "user", "content": "Transcribe this audio: <<BASE64_AUDIO>>"}],
+    "temperature": 0.7
 }
 """
+
 spec = LLMSpec.from_string(http_spec)
 encoded_audio = encode_audio_base64_by_url("https://example.com/audio.mp3")
 response = await spec.probe("Transcribe this audio:", encoded_audio=encoded_audio)
@@ -108,11 +111,12 @@ Authorization: Bearer sk-xxxxxxxxx
 Content-Type: multipart/form-data
 
 {
-  "model": "gpt-3.5-turbo",
-  "messages": [{"role": "user", "content": "Process this file: <<FILE>>"}],
-  "temperature": 0.7
+    "model": "gpt-3.5-turbo",
+    "messages": [{"role": "user", "content": "Process this file: <<FILE>>"}],
+    "temperature": 0.7
 }
 """
+
 spec = LLMSpec.from_string(http_spec)
 files = {"file": ("document.txt", open("document.txt", "rb"))}
 response = await spec.probe("Process this file:", files=files)

docs/probe_data.md

@@ -54,15 +54,20 @@ The `probe_data` module is a core component of the Agentic Security project, res
 
 - **Classes:**
   - `PromptSelectionInterface`: Abstract base class for prompt selection strategies.
+
     - Methods:
       - `select_next_prompt(current_prompt: str, passed_guard: bool) -> str`: Selects next prompt
       - `select_next_prompts(current_prompt: str, passed_guard: bool) -> list[str]`: Selects multiple prompts
-      - `update_rewards(previous_prompt: str, current_prompt: str, reward: float, passed_guard: bool) -> null`: Updates rewards
+      - `update_rewards(previous_prompt: str, current_prompt: str, reward: float, passed_guard: bool) -> None`: Updates rewards
+
   - `RandomPromptSelector`: Basic random selection with history tracking.
+
     - Parameters:
       - `prompts: list[str]`: List of available prompts
       - `history_size: int = 3`: Size of history to prevent cycles
+
   - `CloudRLPromptSelector`: Cloud-based RL implementation with fallback.
+
     - Parameters:
       - `prompts: list[str]`: List of available prompts
       - `api_url: str`: URL of RL service
@@ -70,7 +75,9 @@ The `probe_data` module is a core component of the Agentic Security project, res
       - `history_size: int = 300`: Size of history
       - `timeout: int = 5`: Request timeout
       - `run_id: str = ""`: Unique run identifier
+
   - `QLearningPromptSelector`: Local Q-learning implementation.
+
     - Parameters:
       - `prompts: list[str]`: List of available prompts
       - `learning_rate: float = 0.1`: Learning rate
@@ -79,11 +86,13 @@ The `probe_data` module is a core component of the Agentic Security project, res
       - `exploration_decay: float = 0.995`: Exploration decay rate
       - `min_exploration: float = 0.01`: Minimum exploration rate
       - `history_size: int = 300`: Size of history
-- **Module**: Main class that uses CloudRLPromptSelector.
-  - Parameters:
-    - `prompt_groups: list[str]`: Groups of prompts
-    - `tools_inbox: asyncio.Queue`: Queue for tool communication
-    - `opts: dict = {}`: Configuration options
+
+  - `Module`: Main class that uses CloudRLPromptSelector.
+
+    - Parameters:
+      - `prompt_groups: list[str]`: Groups of prompts
+      - `tools_inbox: asyncio.Queue`: Queue for tool communication
+      - `opts: dict = {}`: Configuration options
 
 ## Usage Examples
 
@@ -110,9 +119,10 @@ from agentic_security.probe_data.modules.rl_model import QLearningPromptSelector
 
 prompts = ["What is AI?", "Explain machine learning"]
 selector = QLearningPromptSelector(prompts)
+
 current_prompt = "What is AI?"
-next_prompt = selector.select_next_prompt(current_prompt, passed_guard=true)
-selector.update_rewards(current_prompt, next_prompt, reward=1.0, passed_guard=true)
+next_prompt = selector.select_next_prompt(current_prompt, passed_guard=True)
+selector.update_rewards(current_prompt, next_prompt, reward=1.0, passed_guard=True)
 ```
 
 ## Conclusion

poetry.lock

@@ -1,4 +1,4 @@
-# This file is automatically @generated by Poetry 1.8.5 and should not be changed by hand.
+# This file is automatically @generated by Poetry 1.8.2 and should not be changed by hand.
 
 [[package]]
 name = "aiohappyeyeballs"
@@ -663,13 +663,13 @@ tests = ["pytest", "pytest-cov", "pytest-xdist"]
 
 [[package]]
 name = "datasets"
-version = "3.4.0"
+version = "3.3.0"
 description = "HuggingFace community-driven open-source library of datasets"
 optional = false
 python-versions = ">=3.9.0"
 files = [
-    {file = "datasets-3.4.0-py3-none-any.whl", hash = "sha256:35ef5182bddd38f7aa774d9f33c3e8b8e9c9c7ea41b4b7969fde431919cb556b"},
-    {file = "datasets-3.4.0.tar.gz", hash = "sha256:f3defae5d9c79ff586db3b17389fdde01704ffea015293a050d7e8ab6816bad8"},
+    {file = "datasets-3.3.0-py3-none-any.whl", hash = "sha256:22312d09626f8fc3aa0a237b0c164997f5903bddd4c4c9e27dbaf563754c681b"},
+    {file = "datasets-3.3.0.tar.gz", hash = "sha256:54c607b06f6eaa1572e21e200d2870d89d50e3bcc622dc2021a53a6ce4f684c2"},
 ]
 
 [package.dependencies]
@@ -691,15 +691,15 @@ xxhash = "*"
 [package.extras]
 audio = ["librosa", "soundfile (>=0.12.1)", "soxr (>=0.4.0)"]
 benchmarks = ["tensorflow (==2.12.0)", "torch (==2.0.1)", "transformers (==4.30.1)"]
-dev = ["Pillow (>=9.4.0)", "absl-py", "decorator", "elasticsearch (>=7.17.12,<8.0.0)", "faiss-cpu (>=1.8.0.post1)", "jax (>=0.3.14)", "jaxlib (>=0.3.14)", "joblib (<1.3.0)", "joblibspark", "librosa", "lz4", "moto[server]", "polars[timezone] (>=0.20.0)", "protobuf (<4.0.0)", "py7zr", "pyav", "pyspark (>=3.4)", "pytest", "pytest-datadir", "pytest-xdist", "rarfile (>=4.0)", "ruff (>=0.3.0)", "s3fs", "s3fs (>=2021.11.1)", "soundfile (>=0.12.1)", "soundfile (>=0.12.1)", "soxr (>=0.4.0)", "sqlalchemy", "tensorflow (>=2.16.0)", "tensorflow (>=2.6.0)", "tensorflow (>=2.6.0)", "tiktoken", "torch", "torch (>=2.0.0)", "torchdata", "torchvision", "transformers", "transformers (>=4.42.0)", "zstandard"]
+dev = ["Pillow (>=9.4.0)", "absl-py", "decorator", "decord (==0.6.0)", "elasticsearch (>=7.17.12,<8.0.0)", "faiss-cpu (>=1.8.0.post1)", "jax (>=0.3.14)", "jaxlib (>=0.3.14)", "joblib (<1.3.0)", "joblibspark", "librosa", "lz4", "moto[server]", "polars[timezone] (>=0.20.0)", "protobuf (<4.0.0)", "py7zr", "pyspark (>=3.4)", "pytest", "pytest-datadir", "pytest-xdist", "rarfile (>=4.0)", "ruff (>=0.3.0)", "s3fs", "s3fs (>=2021.11.1)", "soundfile (>=0.12.1)", "soundfile (>=0.12.1)", "soxr (>=0.4.0)", "sqlalchemy", "tensorflow (>=2.16.0)", "tensorflow (>=2.6.0)", "tensorflow (>=2.6.0)", "tiktoken", "torch", "torch (>=2.0.0)", "torchdata", "transformers", "transformers (>=4.42.0)", "zstandard"]
 docs = ["s3fs", "tensorflow (>=2.6.0)", "torch", "transformers"]
 jax = ["jax (>=0.3.14)", "jaxlib (>=0.3.14)"]
 quality = ["ruff (>=0.3.0)"]
 s3 = ["s3fs"]
 tensorflow = ["tensorflow (>=2.6.0)"]
 tensorflow-gpu = ["tensorflow (>=2.6.0)"]
-tests = ["Pillow (>=9.4.0)", "absl-py", "decorator", "elasticsearch (>=7.17.12,<8.0.0)", "faiss-cpu (>=1.8.0.post1)", "jax (>=0.3.14)", "jaxlib (>=0.3.14)", "joblib (<1.3.0)", "joblibspark", "librosa", "lz4", "moto[server]", "polars[timezone] (>=0.20.0)", "protobuf (<4.0.0)", "py7zr", "pyav", "pyspark (>=3.4)", "pytest", "pytest-datadir", "pytest-xdist", "rarfile (>=4.0)", "s3fs (>=2021.11.1)", "soundfile (>=0.12.1)", "soundfile (>=0.12.1)", "soxr (>=0.4.0)", "sqlalchemy", "tensorflow (>=2.16.0)", "tensorflow (>=2.6.0)", "tiktoken", "torch (>=2.0.0)", "torchdata", "torchvision", "transformers (>=4.42.0)", "zstandard"]
-tests-numpy2 = ["Pillow (>=9.4.0)", "absl-py", "decorator", "elasticsearch (>=7.17.12,<8.0.0)", "jax (>=0.3.14)", "jaxlib (>=0.3.14)", "joblib (<1.3.0)", "joblibspark", "lz4", "moto[server]", "polars[timezone] (>=0.20.0)", "protobuf (<4.0.0)", "py7zr", "pyav", "pyspark (>=3.4)", "pytest", "pytest-datadir", "pytest-xdist", "rarfile (>=4.0)", "s3fs (>=2021.11.1)", "soundfile (>=0.12.1)", "soundfile (>=0.12.1)", "soxr (>=0.4.0)", "sqlalchemy", "tiktoken", "torch (>=2.0.0)", "torchdata", "torchvision", "transformers (>=4.42.0)", "zstandard"]
+tests = ["Pillow (>=9.4.0)", "absl-py", "decorator", "decord (==0.6.0)", "elasticsearch (>=7.17.12,<8.0.0)", "faiss-cpu (>=1.8.0.post1)", "jax (>=0.3.14)", "jaxlib (>=0.3.14)", "joblib (<1.3.0)", "joblibspark", "librosa", "lz4", "moto[server]", "polars[timezone] (>=0.20.0)", "protobuf (<4.0.0)", "py7zr", "pyspark (>=3.4)", "pytest", "pytest-datadir", "pytest-xdist", "rarfile (>=4.0)", "s3fs (>=2021.11.1)", "soundfile (>=0.12.1)", "soundfile (>=0.12.1)", "soxr (>=0.4.0)", "sqlalchemy", "tensorflow (>=2.16.0)", "tensorflow (>=2.6.0)", "tiktoken", "torch (>=2.0.0)", "torchdata", "transformers (>=4.42.0)", "zstandard"]
+tests-numpy2 = ["Pillow (>=9.4.0)", "absl-py", "decorator", "decord (==0.6.0)", "elasticsearch (>=7.17.12,<8.0.0)", "jax (>=0.3.14)", "jaxlib (>=0.3.14)", "joblib (<1.3.0)", "joblibspark", "lz4", "moto[server]", "polars[timezone] (>=0.20.0)", "protobuf (<4.0.0)", "py7zr", "pyspark (>=3.4)", "pytest", "pytest-datadir", "pytest-xdist", "rarfile (>=4.0)", "s3fs (>=2021.11.1)", "soundfile (>=0.12.1)", "soundfile (>=0.12.1)", "soxr (>=0.4.0)", "sqlalchemy", "tiktoken", "torch (>=2.0.0)", "torchdata", "transformers (>=4.42.0)", "zstandard"]
 torch = ["torch"]
 vision = ["Pillow (>=9.4.0)"]
 
@@ -1164,13 +1164,13 @@ zstd = ["zstandard (>=0.18.0)"]
 
 [[package]]
 name = "huggingface-hub"
-version = "0.29.3"
+version = "0.28.1"
 description = "Client library to download and publish models, datasets and other repos on the huggingface.co hub"
 optional = false
 python-versions = ">=3.8.0"
 files = [
-    {file = "huggingface_hub-0.29.3-py3-none-any.whl", hash = "sha256:0b25710932ac649c08cdbefa6c6ccb8e88eef82927cacdb048efb726429453aa"},
-    {file = "huggingface_hub-0.29.3.tar.gz", hash = "sha256:64519a25716e0ba382ba2d3fb3ca082e7c7eb4a2fc634d200e8380006e0760e5"},
+    {file = "huggingface_hub-0.28.1-py3-none-any.whl", hash = "sha256:aa6b9a3ffdae939b72c464dbb0d7f99f56e649b55c3d52406f49e0a5a620c0a7"},
+    {file = "huggingface_hub-0.28.1.tar.gz", hash = "sha256:893471090c98e3b6efbdfdacafe4052b20b84d59866fb6f54c33d9af18c303ae"},
 ]
 
 [package.dependencies]
@@ -1234,13 +1234,13 @@ files = [
 
 [[package]]
 name = "inline-snapshot"
-version = "0.20.9"
+version = "0.20.3"
 description = "golden master/snapshot/approval testing library which puts the values right into your source code"
 optional = false
 python-versions = ">=3.8"
 files = [
-    {file = "inline_snapshot-0.20.9-py3-none-any.whl", hash = "sha256:e4bffcfb39d79bba0320a138e47217601823bbf6948760d75fa175d62a7b893d"},
-    {file = "inline_snapshot-0.20.9.tar.gz", hash = "sha256:ce41505690f0f6ca96466776bc3e015ccc7412ecb65b74f4e7e90f1a834cbac4"},
+    {file = "inline_snapshot-0.20.3-py3-none-any.whl", hash = "sha256:1ea999fbf38dd11cc72d0e1a0b9303c63d496b77bdc406a394fe2424ae842f70"},
+    {file = "inline_snapshot-0.20.3.tar.gz", hash = "sha256:7a353170b7e42aa89086c7ba790a973c9645523acf985532648dabd7ee2d71f2"},
 ]
 
 [package.dependencies]
@@ -1931,13 +1931,13 @@ pygments = ">2.12.0"
 
 [[package]]
 name = "mkdocs-material"
-version = "9.6.10"
+version = "9.6.7"
 description = "Documentation that simply works"
 optional = false
 python-versions = ">=3.8"
 files = [
-    {file = "mkdocs_material-9.6.10-py3-none-any.whl", hash = "sha256:36168548df4e2ddeb9a334ddae4ab9c388ccfea4dd50ffee657d22b93dcb1c3e"},
-    {file = "mkdocs_material-9.6.10.tar.gz", hash = "sha256:25a453c1f24f34fcf1f53680c03d2c1421b52ce5247f4468153c87a70cd5f1fc"},
+    {file = "mkdocs_material-9.6.7-py3-none-any.whl", hash = "sha256:8a159e45e80fcaadd9fbeef62cbf928569b93df954d4dc5ba76d46820caf7b47"},
+    {file = "mkdocs_material-9.6.7.tar.gz", hash = "sha256:3e2c1fceb9410056c2d91f334a00cdea3215c28750e00c691c1e46b2a33309b4"},
 ]
 
 [package.dependencies]
@@ -1971,26 +1971,27 @@ files = [
 
 [[package]]
 name = "mkdocstrings"
-version = "0.29.0"
+version = "0.28.2"
 description = "Automatic documentation from sources, for MkDocs."
 optional = false
 python-versions = ">=3.9"
 files = [
-    {file = "mkdocstrings-0.29.0-py3-none-any.whl", hash = "sha256:8ea98358d2006f60befa940fdebbbc88a26b37ecbcded10be726ba359284f73d"},
-    {file = "mkdocstrings-0.29.0.tar.gz", hash = "sha256:3657be1384543ce0ee82112c3e521bbf48e41303aa0c229b9ffcccba057d922e"},
+    {file = "mkdocstrings-0.28.2-py3-none-any.whl", hash = "sha256:57f79c557e2718d217d6f6a81bf75a0de097f10e922e7e5e00f085c3f0ff6895"},
+    {file = "mkdocstrings-0.28.2.tar.gz", hash = "sha256:9b847266d7a588ea76a8385eaebe1538278b4361c0d1ce48ed005be59f053569"},
 ]
 
 [package.dependencies]
 Jinja2 = ">=2.11.1"
 Markdown = ">=3.6"
 MarkupSafe = ">=1.1"
-mkdocs = ">=1.6"
+mkdocs = ">=1.4"
 mkdocs-autorefs = ">=1.4"
+mkdocs-get-deps = ">=0.2"
 pymdown-extensions = ">=6.3"
 
 [package.extras]
 crystal = ["mkdocstrings-crystal (>=0.3.4)"]
-python = ["mkdocstrings-python (>=1.16.2)"]
+python = ["mkdocstrings-python (>=0.5.2)"]
 python-legacy = ["mkdocstrings-python-legacy (>=0.2.1)"]
 
 [[package]]
@@ -2346,79 +2347,90 @@ files = [
 
 [[package]]
 name = "orjson"
-version = "3.10.16"
+version = "3.10.15"
 description = "Fast, correct Python JSON library supporting dataclasses, datetimes, and numpy"
 optional = false
-python-versions = ">=3.9"
+python-versions = ">=3.8"
 files = [
-    {file = "orjson-3.10.16-cp310-cp310-macosx_10_15_x86_64.macosx_11_0_arm64.macosx_10_15_universal2.whl", hash = "sha256:4cb473b8e79154fa778fb56d2d73763d977be3dcc140587e07dbc545bbfc38f8"},
-    {file = "orjson-3.10.16-cp310-cp310-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:622a8e85eeec1948690409a19ca1c7d9fd8ff116f4861d261e6ae2094fe59a00"},
-    {file = "orjson-3.10.16-cp310-cp310-manylinux_2_17_armv7l.manylinux2014_armv7l.whl", hash = "sha256:c682d852d0ce77613993dc967e90e151899fe2d8e71c20e9be164080f468e370"},
-    {file = "orjson-3.10.16-cp310-cp310-manylinux_2_17_i686.manylinux2014_i686.whl", hash = "sha256:8c520ae736acd2e32df193bcff73491e64c936f3e44a2916b548da048a48b46b"},
-    {file = "orjson-3.10.16-cp310-cp310-manylinux_2_17_ppc64le.manylinux2014_ppc64le.whl", hash = "sha256:134f87c76bfae00f2094d85cfab261b289b76d78c6da8a7a3b3c09d362fd1e06"},
-    {file = "orjson-3.10.16-cp310-cp310-manylinux_2_17_s390x.manylinux2014_s390x.whl", hash = "sha256:b59afde79563e2cf37cfe62ee3b71c063fd5546c8e662d7fcfc2a3d5031a5c4c"},
-    {file = "orjson-3.10.16-cp310-cp310-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:113602f8241daaff05d6fad25bd481d54c42d8d72ef4c831bb3ab682a54d9e15"},
-    {file = "orjson-3.10.16-cp310-cp310-musllinux_1_2_aarch64.whl", hash = "sha256:4fc0077d101f8fab4031e6554fc17b4c2ad8fdbc56ee64a727f3c95b379e31da"},
-    {file = "orjson-3.10.16-cp310-cp310-musllinux_1_2_armv7l.whl", hash = "sha256:9c6bf6ff180cd69e93f3f50380224218cfab79953a868ea3908430bcfaf9cb5e"},
-    {file = "orjson-3.10.16-cp310-cp310-musllinux_1_2_i686.whl", hash = "sha256:5673eadfa952f95a7cd76418ff189df11b0a9c34b1995dff43a6fdbce5d63bf4"},
-    {file = "orjson-3.10.16-cp310-cp310-musllinux_1_2_x86_64.whl", hash = "sha256:5fe638a423d852b0ae1e1a79895851696cb0d9fa0946fdbfd5da5072d9bb9551"},
-    {file = "orjson-3.10.16-cp310-cp310-win32.whl", hash = "sha256:33af58f479b3c6435ab8f8b57999874b4b40c804c7a36b5cc6b54d8f28e1d3dd"},
-    {file = "orjson-3.10.16-cp310-cp310-win_amd64.whl", hash = "sha256:0338356b3f56d71293c583350af26f053017071836b07e064e92819ecf1aa055"},
-    {file = "orjson-3.10.16-cp311-cp311-macosx_10_15_x86_64.macosx_11_0_arm64.macosx_10_15_universal2.whl", hash = "sha256:44fcbe1a1884f8bc9e2e863168b0f84230c3d634afe41c678637d2728ea8e739"},
-    {file = "orjson-3.10.16-cp311-cp311-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:78177bf0a9d0192e0b34c3d78bcff7fe21d1b5d84aeb5ebdfe0dbe637b885225"},
-    {file = "orjson-3.10.16-cp311-cp311-manylinux_2_17_armv7l.manylinux2014_armv7l.whl", hash = "sha256:12824073a010a754bb27330cad21d6e9b98374f497f391b8707752b96f72e741"},
-    {file = "orjson-3.10.16-cp311-cp311-manylinux_2_17_i686.manylinux2014_i686.whl", hash = "sha256:ddd41007e56284e9867864aa2f29f3136bb1dd19a49ca43c0b4eda22a579cf53"},
-    {file = "orjson-3.10.16-cp311-cp311-manylinux_2_17_ppc64le.manylinux2014_ppc64le.whl", hash = "sha256:0877c4d35de639645de83666458ca1f12560d9fa7aa9b25d8bb8f52f61627d14"},
-    {file = "orjson-3.10.16-cp311-cp311-manylinux_2_17_s390x.manylinux2014_s390x.whl", hash = "sha256:9a09a539e9cc3beead3e7107093b4ac176d015bec64f811afb5965fce077a03c"},
-    {file = "orjson-3.10.16-cp311-cp311-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:31b98bc9b40610fec971d9a4d67bb2ed02eec0a8ae35f8ccd2086320c28526ca"},
-    {file = "orjson-3.10.16-cp311-cp311-musllinux_1_2_aarch64.whl", hash = "sha256:0ce243f5a8739f3a18830bc62dc2e05b69a7545bafd3e3249f86668b2bcd8e50"},
-    {file = "orjson-3.10.16-cp311-cp311-musllinux_1_2_armv7l.whl", hash = "sha256:64792c0025bae049b3074c6abe0cf06f23c8e9f5a445f4bab31dc5ca23dbf9e1"},
-    {file = "orjson-3.10.16-cp311-cp311-musllinux_1_2_i686.whl", hash = "sha256:ea53f7e68eec718b8e17e942f7ca56c6bd43562eb19db3f22d90d75e13f0431d"},
-    {file = "orjson-3.10.16-cp311-cp311-musllinux_1_2_x86_64.whl", hash = "sha256:a741ba1a9488c92227711bde8c8c2b63d7d3816883268c808fbeada00400c164"},
-    {file = "orjson-3.10.16-cp311-cp311-win32.whl", hash = "sha256:c7ed2c61bb8226384c3fdf1fb01c51b47b03e3f4536c985078cccc2fd19f1619"},
-    {file = "orjson-3.10.16-cp311-cp311-win_amd64.whl", hash = "sha256:cd67d8b3e0e56222a2e7b7f7da9031e30ecd1fe251c023340b9f12caca85ab60"},
-    {file = "orjson-3.10.16-cp312-cp312-macosx_10_15_x86_64.macosx_11_0_arm64.macosx_10_15_universal2.whl", hash = "sha256:6d3444abbfa71ba21bb042caa4b062535b122248259fdb9deea567969140abca"},
-    {file = "orjson-3.10.16-cp312-cp312-macosx_15_0_arm64.whl", hash = "sha256:30245c08d818fdcaa48b7d5b81499b8cae09acabb216fe61ca619876b128e184"},
-    {file = "orjson-3.10.16-cp312-cp312-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:a0ba1d0baa71bf7579a4ccdcf503e6f3098ef9542106a0eca82395898c8a500a"},
-    {file = "orjson-3.10.16-cp312-cp312-manylinux_2_17_armv7l.manylinux2014_armv7l.whl", hash = "sha256:eb0beefa5ef3af8845f3a69ff2a4aa62529b5acec1cfe5f8a6b4141033fd46ef"},
-    {file = "orjson-3.10.16-cp312-cp312-manylinux_2_17_i686.manylinux2014_i686.whl", hash = "sha256:6daa0e1c9bf2e030e93c98394de94506f2a4d12e1e9dadd7c53d5e44d0f9628e"},
-    {file = "orjson-3.10.16-cp312-cp312-manylinux_2_17_ppc64le.manylinux2014_ppc64le.whl", hash = "sha256:9da9019afb21e02410ef600e56666652b73eb3e4d213a0ec919ff391a7dd52aa"},
-    {file = "orjson-3.10.16-cp312-cp312-manylinux_2_17_s390x.manylinux2014_s390x.whl", hash = "sha256:daeb3a1ee17b69981d3aae30c3b4e786b0f8c9e6c71f2b48f1aef934f63f38f4"},
-    {file = "orjson-3.10.16-cp312-cp312-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:80fed80eaf0e20a31942ae5d0728849862446512769692474be5e6b73123a23b"},
-    {file = "orjson-3.10.16-cp312-cp312-musllinux_1_2_aarch64.whl", hash = "sha256:73390ed838f03764540a7bdc4071fe0123914c2cc02fb6abf35182d5fd1b7a42"},
-    {file = "orjson-3.10.16-cp312-cp312-musllinux_1_2_armv7l.whl", hash = "sha256:a22bba012a0c94ec02a7768953020ab0d3e2b884760f859176343a36c01adf87"},
-    {file = "orjson-3.10.16-cp312-cp312-musllinux_1_2_i686.whl", hash = "sha256:5385bbfdbc90ff5b2635b7e6bebf259652db00a92b5e3c45b616df75b9058e88"},
-    {file = "orjson-3.10.16-cp312-cp312-musllinux_1_2_x86_64.whl", hash = "sha256:02c6279016346e774dd92625d46c6c40db687b8a0d685aadb91e26e46cc33e1e"},
-    {file = "orjson-3.10.16-cp312-cp312-win32.whl", hash = "sha256:7ca55097a11426db80f79378e873a8c51f4dde9ffc22de44850f9696b7eb0e8c"},
-    {file = "orjson-3.10.16-cp312-cp312-win_amd64.whl", hash = "sha256:86d127efdd3f9bf5f04809b70faca1e6836556ea3cc46e662b44dab3fe71f3d6"},
-    {file = "orjson-3.10.16-cp313-cp313-macosx_10_15_x86_64.macosx_11_0_arm64.macosx_10_15_universal2.whl", hash = "sha256:148a97f7de811ba14bc6dbc4a433e0341ffd2cc285065199fb5f6a98013744bd"},
-    {file = "orjson-3.10.16-cp313-cp313-macosx_15_0_arm64.whl", hash = "sha256:1d960c1bf0e734ea36d0adc880076de3846aaec45ffad29b78c7f1b7962516b8"},
-    {file = "orjson-3.10.16-cp313-cp313-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:a318cd184d1269f68634464b12871386808dc8b7c27de8565234d25975a7a137"},
-    {file = "orjson-3.10.16-cp313-cp313-manylinux_2_17_armv7l.manylinux2014_armv7l.whl", hash = "sha256:df23f8df3ef9223d1d6748bea63fca55aae7da30a875700809c500a05975522b"},
-    {file = "orjson-3.10.16-cp313-cp313-manylinux_2_17_i686.manylinux2014_i686.whl", hash = "sha256:b94dda8dd6d1378f1037d7f3f6b21db769ef911c4567cbaa962bb6dc5021cf90"},
-    {file = "orjson-3.10.16-cp313-cp313-manylinux_2_17_ppc64le.manylinux2014_ppc64le.whl", hash = "sha256:f12970a26666a8775346003fd94347d03ccb98ab8aa063036818381acf5f523e"},
-    {file = "orjson-3.10.16-cp313-cp313-manylinux_2_17_s390x.manylinux2014_s390x.whl", hash = "sha256:15a1431a245d856bd56e4d29ea0023eb4d2c8f71efe914beb3dee8ab3f0cd7fb"},
-    {file = "orjson-3.10.16-cp313-cp313-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:c83655cfc247f399a222567d146524674a7b217af7ef8289c0ff53cfe8db09f0"},
-    {file = "orjson-3.10.16-cp313-cp313-musllinux_1_2_aarch64.whl", hash = "sha256:fa59ae64cb6ddde8f09bdbf7baf933c4cd05734ad84dcf4e43b887eb24e37652"},
-    {file = "orjson-3.10.16-cp313-cp313-musllinux_1_2_armv7l.whl", hash = "sha256:ca5426e5aacc2e9507d341bc169d8af9c3cbe88f4cd4c1cf2f87e8564730eb56"},
-    {file = "orjson-3.10.16-cp313-cp313-musllinux_1_2_i686.whl", hash = "sha256:6fd5da4edf98a400946cd3a195680de56f1e7575109b9acb9493331047157430"},
-    {file = "orjson-3.10.16-cp313-cp313-musllinux_1_2_x86_64.whl", hash = "sha256:980ecc7a53e567169282a5e0ff078393bac78320d44238da4e246d71a4e0e8f5"},
-    {file = "orjson-3.10.16-cp313-cp313-win32.whl", hash = "sha256:28f79944dd006ac540a6465ebd5f8f45dfdf0948ff998eac7a908275b4c1add6"},
-    {file = "orjson-3.10.16-cp313-cp313-win_amd64.whl", hash = "sha256:fe0a145e96d51971407cb8ba947e63ead2aa915db59d6631a355f5f2150b56b7"},
-    {file = "orjson-3.10.16-cp39-cp39-macosx_10_15_x86_64.macosx_11_0_arm64.macosx_10_15_universal2.whl", hash = "sha256:c35b5c1fb5a5d6d2fea825dec5d3d16bea3c06ac744708a8e1ff41d4ba10cdf1"},
-    {file = "orjson-3.10.16-cp39-cp39-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:c9aac7ecc86218b4b3048c768f227a9452287001d7548500150bb75ee21bf55d"},
-    {file = "orjson-3.10.16-cp39-cp39-manylinux_2_17_armv7l.manylinux2014_armv7l.whl", hash = "sha256:6e19f5102fff36f923b6dfdb3236ec710b649da975ed57c29833cb910c5a73ab"},
-    {file = "orjson-3.10.16-cp39-cp39-manylinux_2_17_i686.manylinux2014_i686.whl", hash = "sha256:17210490408eb62755a334a6f20ed17c39f27b4f45d89a38cd144cd458eba80b"},
-    {file = "orjson-3.10.16-cp39-cp39-manylinux_2_17_ppc64le.manylinux2014_ppc64le.whl", hash = "sha256:fbbe04451db85916e52a9f720bd89bf41f803cf63b038595674691680cbebd1b"},
-    {file = "orjson-3.10.16-cp39-cp39-manylinux_2_17_s390x.manylinux2014_s390x.whl", hash = "sha256:6a966eba501a3a1f309f5a6af32ed9eb8f316fa19d9947bac3e6350dc63a6f0a"},
-    {file = "orjson-3.10.16-cp39-cp39-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:01e0d22f06c81e6c435723343e1eefc710e0510a35d897856766d475f2a15687"},
-    {file = "orjson-3.10.16-cp39-cp39-musllinux_1_2_aarch64.whl", hash = "sha256:7c1e602d028ee285dbd300fb9820b342b937df64d5a3336e1618b354e95a2569"},
-    {file = "orjson-3.10.16-cp39-cp39-musllinux_1_2_armv7l.whl", hash = "sha256:d230e5020666a6725629df81e210dc11c3eae7d52fe909a7157b3875238484f3"},
-    {file = "orjson-3.10.16-cp39-cp39-musllinux_1_2_i686.whl", hash = "sha256:0f8baac07d4555f57d44746a7d80fbe6b2c4fe2ed68136b4abb51cfec512a5e9"},
-    {file = "orjson-3.10.16-cp39-cp39-musllinux_1_2_x86_64.whl", hash = "sha256:524e48420b90fc66953e91b660b3d05faaf921277d6707e328fde1c218b31250"},
-    {file = "orjson-3.10.16-cp39-cp39-win32.whl", hash = "sha256:a9f614e31423d7292dbca966a53b2d775c64528c7d91424ab2747d8ab8ce5c72"},
-    {file = "orjson-3.10.16-cp39-cp39-win_amd64.whl", hash = "sha256:c338dc2296d1ed0d5c5c27dfb22d00b330555cb706c2e0be1e1c3940a0895905"},
-    {file = "orjson-3.10.16.tar.gz", hash = "sha256:d2aaa5c495e11d17b9b93205f5fa196737ee3202f000aaebf028dc9a73750f10"},
+    {file = "orjson-3.10.15-cp310-cp310-macosx_10_15_x86_64.macosx_11_0_arm64.macosx_10_15_universal2.whl", hash = "sha256:552c883d03ad185f720d0c09583ebde257e41b9521b74ff40e08b7dec4559c04"},
+    {file = "orjson-3.10.15-cp310-cp310-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:616e3e8d438d02e4854f70bfdc03a6bcdb697358dbaa6bcd19cbe24d24ece1f8"},
+    {file = "orjson-3.10.15-cp310-cp310-manylinux_2_17_armv7l.manylinux2014_armv7l.whl", hash = "sha256:7c2c79fa308e6edb0ffab0a31fd75a7841bf2a79a20ef08a3c6e3b26814c8ca8"},
+    {file = "orjson-3.10.15-cp310-cp310-manylinux_2_17_ppc64le.manylinux2014_ppc64le.whl", hash = "sha256:73cb85490aa6bf98abd20607ab5c8324c0acb48d6da7863a51be48505646c814"},
+    {file = "orjson-3.10.15-cp310-cp310-manylinux_2_17_s390x.manylinux2014_s390x.whl", hash = "sha256:763dadac05e4e9d2bc14938a45a2d0560549561287d41c465d3c58aec818b164"},
+    {file = "orjson-3.10.15-cp310-cp310-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:a330b9b4734f09a623f74a7490db713695e13b67c959713b78369f26b3dee6bf"},
+    {file = "orjson-3.10.15-cp310-cp310-manylinux_2_5_i686.manylinux1_i686.whl", hash = "sha256:a61a4622b7ff861f019974f73d8165be1bd9a0855e1cad18ee167acacabeb061"},
+    {file = "orjson-3.10.15-cp310-cp310-musllinux_1_2_aarch64.whl", hash = "sha256:acd271247691574416b3228db667b84775c497b245fa275c6ab90dc1ffbbd2b3"},
+    {file = "orjson-3.10.15-cp310-cp310-musllinux_1_2_armv7l.whl", hash = "sha256:e4759b109c37f635aa5c5cc93a1b26927bfde24b254bcc0e1149a9fada253d2d"},
+    {file = "orjson-3.10.15-cp310-cp310-musllinux_1_2_i686.whl", hash = "sha256:9e992fd5cfb8b9f00bfad2fd7a05a4299db2bbe92e6440d9dd2fab27655b3182"},
+    {file = "orjson-3.10.15-cp310-cp310-musllinux_1_2_x86_64.whl", hash = "sha256:f95fb363d79366af56c3f26b71df40b9a583b07bbaaf5b317407c4d58497852e"},
+    {file = "orjson-3.10.15-cp310-cp310-win32.whl", hash = "sha256:f9875f5fea7492da8ec2444839dcc439b0ef298978f311103d0b7dfd775898ab"},
+    {file = "orjson-3.10.15-cp310-cp310-win_amd64.whl", hash = "sha256:17085a6aa91e1cd70ca8533989a18b5433e15d29c574582f76f821737c8d5806"},
+    {file = "orjson-3.10.15-cp311-cp311-macosx_10_15_x86_64.macosx_11_0_arm64.macosx_10_15_universal2.whl", hash = "sha256:c4cc83960ab79a4031f3119cc4b1a1c627a3dc09df125b27c4201dff2af7eaa6"},
+    {file = "orjson-3.10.15-cp311-cp311-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:ddbeef2481d895ab8be5185f2432c334d6dec1f5d1933a9c83014d188e102cef"},
+    {file = "orjson-3.10.15-cp311-cp311-manylinux_2_17_armv7l.manylinux2014_armv7l.whl", hash = "sha256:9e590a0477b23ecd5b0ac865b1b907b01b3c5535f5e8a8f6ab0e503efb896334"},
+    {file = "orjson-3.10.15-cp311-cp311-manylinux_2_17_ppc64le.manylinux2014_ppc64le.whl", hash = "sha256:a6be38bd103d2fd9bdfa31c2720b23b5d47c6796bcb1d1b598e3924441b4298d"},
+    {file = "orjson-3.10.15-cp311-cp311-manylinux_2_17_s390x.manylinux2014_s390x.whl", hash = "sha256:ff4f6edb1578960ed628a3b998fa54d78d9bb3e2eb2cfc5c2a09732431c678d0"},
+    {file = "orjson-3.10.15-cp311-cp311-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:b0482b21d0462eddd67e7fce10b89e0b6ac56570424662b685a0d6fccf581e13"},
+    {file = "orjson-3.10.15-cp311-cp311-manylinux_2_5_i686.manylinux1_i686.whl", hash = "sha256:bb5cc3527036ae3d98b65e37b7986a918955f85332c1ee07f9d3f82f3a6899b5"},
+    {file = "orjson-3.10.15-cp311-cp311-musllinux_1_2_aarch64.whl", hash = "sha256:d569c1c462912acdd119ccbf719cf7102ea2c67dd03b99edcb1a3048651ac96b"},
+    {file = "orjson-3.10.15-cp311-cp311-musllinux_1_2_armv7l.whl", hash = "sha256:1e6d33efab6b71d67f22bf2962895d3dc6f82a6273a965fab762e64fa90dc399"},
+    {file = "orjson-3.10.15-cp311-cp311-musllinux_1_2_i686.whl", hash = "sha256:c33be3795e299f565681d69852ac8c1bc5c84863c0b0030b2b3468843be90388"},
+    {file = "orjson-3.10.15-cp311-cp311-musllinux_1_2_x86_64.whl", hash = "sha256:eea80037b9fae5339b214f59308ef0589fc06dc870578b7cce6d71eb2096764c"},
+    {file = "orjson-3.10.15-cp311-cp311-win32.whl", hash = "sha256:d5ac11b659fd798228a7adba3e37c010e0152b78b1982897020a8e019a94882e"},
+    {file = "orjson-3.10.15-cp311-cp311-win_amd64.whl", hash = "sha256:cf45e0214c593660339ef63e875f32ddd5aa3b4adc15e662cdb80dc49e194f8e"},
+    {file = "orjson-3.10.15-cp312-cp312-macosx_10_15_x86_64.macosx_11_0_arm64.macosx_10_15_universal2.whl", hash = "sha256:9d11c0714fc85bfcf36ada1179400862da3288fc785c30e8297844c867d7505a"},
+    {file = "orjson-3.10.15-cp312-cp312-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:dba5a1e85d554e3897fa9fe6fbcff2ed32d55008973ec9a2b992bd9a65d2352d"},
+    {file = "orjson-3.10.15-cp312-cp312-manylinux_2_17_armv7l.manylinux2014_armv7l.whl", hash = "sha256:7723ad949a0ea502df656948ddd8b392780a5beaa4c3b5f97e525191b102fff0"},
+    {file = "orjson-3.10.15-cp312-cp312-manylinux_2_17_ppc64le.manylinux2014_ppc64le.whl", hash = "sha256:6fd9bc64421e9fe9bd88039e7ce8e58d4fead67ca88e3a4014b143cec7684fd4"},
+    {file = "orjson-3.10.15-cp312-cp312-manylinux_2_17_s390x.manylinux2014_s390x.whl", hash = "sha256:dadba0e7b6594216c214ef7894c4bd5f08d7c0135f4dd0145600be4fbcc16767"},
+    {file = "orjson-3.10.15-cp312-cp312-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:b48f59114fe318f33bbaee8ebeda696d8ccc94c9e90bc27dbe72153094e26f41"},
+    {file = "orjson-3.10.15-cp312-cp312-manylinux_2_5_i686.manylinux1_i686.whl", hash = "sha256:035fb83585e0f15e076759b6fedaf0abb460d1765b6a36f48018a52858443514"},
+    {file = "orjson-3.10.15-cp312-cp312-musllinux_1_2_aarch64.whl", hash = "sha256:d13b7fe322d75bf84464b075eafd8e7dd9eae05649aa2a5354cfa32f43c59f17"},
+    {file = "orjson-3.10.15-cp312-cp312-musllinux_1_2_armv7l.whl", hash = "sha256:7066b74f9f259849629e0d04db6609db4cf5b973248f455ba5d3bd58a4daaa5b"},
+    {file = "orjson-3.10.15-cp312-cp312-musllinux_1_2_i686.whl", hash = "sha256:88dc3f65a026bd3175eb157fea994fca6ac7c4c8579fc5a86fc2114ad05705b7"},
+    {file = "orjson-3.10.15-cp312-cp312-musllinux_1_2_x86_64.whl", hash = "sha256:b342567e5465bd99faa559507fe45e33fc76b9fb868a63f1642c6bc0735ad02a"},
+    {file = "orjson-3.10.15-cp312-cp312-win32.whl", hash = "sha256:0a4f27ea5617828e6b58922fdbec67b0aa4bb844e2d363b9244c47fa2180e665"},
+    {file = "orjson-3.10.15-cp312-cp312-win_amd64.whl", hash = "sha256:ef5b87e7aa9545ddadd2309efe6824bd3dd64ac101c15dae0f2f597911d46eaa"},
+    {file = "orjson-3.10.15-cp313-cp313-macosx_10_15_x86_64.macosx_11_0_arm64.macosx_10_15_universal2.whl", hash = "sha256:bae0e6ec2b7ba6895198cd981b7cca95d1487d0147c8ed751e5632ad16f031a6"},
+    {file = "orjson-3.10.15-cp313-cp313-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:f93ce145b2db1252dd86af37d4165b6faa83072b46e3995ecc95d4b2301b725a"},
+    {file = "orjson-3.10.15-cp313-cp313-manylinux_2_17_armv7l.manylinux2014_armv7l.whl", hash = "sha256:7c203f6f969210128af3acae0ef9ea6aab9782939f45f6fe02d05958fe761ef9"},
+    {file = "orjson-3.10.15-cp313-cp313-manylinux_2_17_ppc64le.manylinux2014_ppc64le.whl", hash = "sha256:8918719572d662e18b8af66aef699d8c21072e54b6c82a3f8f6404c1f5ccd5e0"},
+    {file = "orjson-3.10.15-cp313-cp313-manylinux_2_17_s390x.manylinux2014_s390x.whl", hash = "sha256:f71eae9651465dff70aa80db92586ad5b92df46a9373ee55252109bb6b703307"},
+    {file = "orjson-3.10.15-cp313-cp313-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:e117eb299a35f2634e25ed120c37c641398826c2f5a3d3cc39f5993b96171b9e"},
+    {file = "orjson-3.10.15-cp313-cp313-manylinux_2_5_i686.manylinux1_i686.whl", hash = "sha256:13242f12d295e83c2955756a574ddd6741c81e5b99f2bef8ed8d53e47a01e4b7"},
+    {file = "orjson-3.10.15-cp313-cp313-musllinux_1_2_aarch64.whl", hash = "sha256:7946922ada8f3e0b7b958cc3eb22cfcf6c0df83d1fe5521b4a100103e3fa84c8"},
+    {file = "orjson-3.10.15-cp313-cp313-musllinux_1_2_armv7l.whl", hash = "sha256:b7155eb1623347f0f22c38c9abdd738b287e39b9982e1da227503387b81b34ca"},
+    {file = "orjson-3.10.15-cp313-cp313-musllinux_1_2_i686.whl", hash = "sha256:208beedfa807c922da4e81061dafa9c8489c6328934ca2a562efa707e049e561"},
+    {file = "orjson-3.10.15-cp313-cp313-musllinux_1_2_x86_64.whl", hash = "sha256:eca81f83b1b8c07449e1d6ff7074e82e3fd6777e588f1a6632127f286a968825"},
+    {file = "orjson-3.10.15-cp313-cp313-win32.whl", hash = "sha256:c03cd6eea1bd3b949d0d007c8d57049aa2b39bd49f58b4b2af571a5d3833d890"},
+    {file = "orjson-3.10.15-cp313-cp313-win_amd64.whl", hash = "sha256:fd56a26a04f6ba5fb2045b0acc487a63162a958ed837648c5781e1fe3316cfbf"},
+    {file = "orjson-3.10.15-cp38-cp38-macosx_10_15_x86_64.macosx_11_0_arm64.macosx_10_15_universal2.whl", hash = "sha256:5e8afd6200e12771467a1a44e5ad780614b86abb4b11862ec54861a82d677746"},
+    {file = "orjson-3.10.15-cp38-cp38-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:da9a18c500f19273e9e104cca8c1f0b40a6470bcccfc33afcc088045d0bf5ea6"},
+    {file = "orjson-3.10.15-cp38-cp38-manylinux_2_17_armv7l.manylinux2014_armv7l.whl", hash = "sha256:bb00b7bfbdf5d34a13180e4805d76b4567025da19a197645ca746fc2fb536586"},
+    {file = "orjson-3.10.15-cp38-cp38-manylinux_2_17_ppc64le.manylinux2014_ppc64le.whl", hash = "sha256:33aedc3d903378e257047fee506f11e0833146ca3e57a1a1fb0ddb789876c1e1"},
+    {file = "orjson-3.10.15-cp38-cp38-manylinux_2_17_s390x.manylinux2014_s390x.whl", hash = "sha256:dd0099ae6aed5eb1fc84c9eb72b95505a3df4267e6962eb93cdd5af03be71c98"},
+    {file = "orjson-3.10.15-cp38-cp38-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:7c864a80a2d467d7786274fce0e4f93ef2a7ca4ff31f7fc5634225aaa4e9e98c"},
+    {file = "orjson-3.10.15-cp38-cp38-manylinux_2_5_i686.manylinux1_i686.whl", hash = "sha256:c25774c9e88a3e0013d7d1a6c8056926b607a61edd423b50eb5c88fd7f2823ae"},
+    {file = "orjson-3.10.15-cp38-cp38-musllinux_1_2_aarch64.whl", hash = "sha256:e78c211d0074e783d824ce7bb85bf459f93a233eb67a5b5003498232ddfb0e8a"},
+    {file = "orjson-3.10.15-cp38-cp38-musllinux_1_2_armv7l.whl", hash = "sha256:43e17289ffdbbac8f39243916c893d2ae41a2ea1a9cbb060a56a4d75286351ae"},
+    {file = "orjson-3.10.15-cp38-cp38-musllinux_1_2_i686.whl", hash = "sha256:781d54657063f361e89714293c095f506c533582ee40a426cb6489c48a637b81"},
+    {file = "orjson-3.10.15-cp38-cp38-musllinux_1_2_x86_64.whl", hash = "sha256:6875210307d36c94873f553786a808af2788e362bd0cf4c8e66d976791e7b528"},
+    {file = "orjson-3.10.15-cp38-cp38-win32.whl", hash = "sha256:305b38b2b8f8083cc3d618927d7f424349afce5975b316d33075ef0f73576b60"},
+    {file = "orjson-3.10.15-cp38-cp38-win_amd64.whl", hash = "sha256:5dd9ef1639878cc3efffed349543cbf9372bdbd79f478615a1c633fe4e4180d1"},
+    {file = "orjson-3.10.15-cp39-cp39-macosx_10_15_x86_64.macosx_11_0_arm64.macosx_10_15_universal2.whl", hash = "sha256:ffe19f3e8d68111e8644d4f4e267a069ca427926855582ff01fc012496d19969"},
+    {file = "orjson-3.10.15-cp39-cp39-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:d433bf32a363823863a96561a555227c18a522a8217a6f9400f00ddc70139ae2"},
+    {file = "orjson-3.10.15-cp39-cp39-manylinux_2_17_armv7l.manylinux2014_armv7l.whl", hash = "sha256:da03392674f59a95d03fa5fb9fe3a160b0511ad84b7a3914699ea5a1b3a38da2"},
+    {file = "orjson-3.10.15-cp39-cp39-manylinux_2_17_ppc64le.manylinux2014_ppc64le.whl", hash = "sha256:3a63bb41559b05360ded9132032239e47983a39b151af1201f07ec9370715c82"},
+    {file = "orjson-3.10.15-cp39-cp39-manylinux_2_17_s390x.manylinux2014_s390x.whl", hash = "sha256:3766ac4702f8f795ff3fa067968e806b4344af257011858cc3d6d8721588b53f"},
+    {file = "orjson-3.10.15-cp39-cp39-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:7a1c73dcc8fadbd7c55802d9aa093b36878d34a3b3222c41052ce6b0fc65f8e8"},
+    {file = "orjson-3.10.15-cp39-cp39-manylinux_2_5_i686.manylinux1_i686.whl", hash = "sha256:b299383825eafe642cbab34be762ccff9fd3408d72726a6b2a4506d410a71ab3"},
+    {file = "orjson-3.10.15-cp39-cp39-musllinux_1_2_aarch64.whl", hash = "sha256:abc7abecdbf67a173ef1316036ebbf54ce400ef2300b4e26a7b843bd446c2480"},
+    {file = "orjson-3.10.15-cp39-cp39-musllinux_1_2_armv7l.whl", hash = "sha256:3614ea508d522a621384c1d6639016a5a2e4f027f3e4a1c93a51867615d28829"},
+    {file = "orjson-3.10.15-cp39-cp39-musllinux_1_2_i686.whl", hash = "sha256:295c70f9dc154307777ba30fe29ff15c1bcc9dfc5c48632f37d20a607e9ba85a"},
+    {file = "orjson-3.10.15-cp39-cp39-musllinux_1_2_x86_64.whl", hash = "sha256:63309e3ff924c62404923c80b9e2048c1f74ba4b615e7584584389ada50ed428"},
+    {file = "orjson-3.10.15-cp39-cp39-win32.whl", hash = "sha256:a2f708c62d026fb5340788ba94a55c23df4e1869fec74be455e0b2f5363b8507"},
+    {file = "orjson-3.10.15-cp39-cp39-win_amd64.whl", hash = "sha256:efcf6c735c3d22ef60c4aa27a5238f1a477df85e9b15f2142f9d669beb2d13fd"},
+    {file = "orjson-3.10.15.tar.gz", hash = "sha256:05ca7fe452a2e9d8d9d706a2984c95b9c2ebc5db417ce0b7a49b91d50642a23e"},
 ]
 
 [[package]]
@@ -2713,13 +2725,13 @@ testing = ["pytest", "pytest-benchmark"]
 
 [[package]]
 name = "pre-commit"
-version = "4.2.0"
+version = "4.1.0"
 description = "A framework for managing and maintaining multi-language pre-commit hooks."
 optional = false
 python-versions = ">=3.9"
 files = [
-    {file = "pre_commit-4.2.0-py2.py3-none-any.whl", hash = "sha256:a009ca7205f1eb497d10b845e52c838a98b6cdd2102a6c8e4540e94ee75c58bd"},
-    {file = "pre_commit-4.2.0.tar.gz", hash = "sha256:601283b9757afd87d40c4c4a9b2b5de9637a8ea02eaff7adc2d0fb4e04841146"},
+    {file = "pre_commit-4.1.0-py2.py3-none-any.whl", hash = "sha256:d29e7cb346295bcc1cc75fc3e92e343495e3ea0196c9ec6ba53f49f10ab6ae7b"},
+    {file = "pre_commit-4.1.0.tar.gz", hash = "sha256:ae3f018575a588e30dfddfab9a05448bfbd6b73d78709617b5a2b853549716d4"},
 ]
 
 [package.dependencies]
@@ -4469,4 +4481,4 @@ propcache = ">=0.2.0"
 [metadata]
 lock-version = "2.0"
 python-versions = "^3.11"
-content-hash = "04c8dc31a9c1a1b2bbb607041a31fa291ed6cdf0b060d30161786da7e8fbab4e"
+content-hash = "35e03dba41d30cf6129a4a4f3107eca560f779205b21d3ffb2871eeffc5d5a64"

pyproject.toml

@@ -1,6 +1,6 @@
 [tool.poetry]
 name = "agentic_security"
-version = "0.7.1"
+version = "0.6.0"
 description = "Agentic LLM vulnerability scanner"
 authors = ["Alexander Miasoiedov <msoedov@gmail.com>"]
 maintainers = ["Alexander Miasoiedov <msoedov@gmail.com>"]
@@ -69,7 +69,7 @@ pytest-mock = "^3.14.0"
 black = ">=24.10,<26.0"
 mypy = "^1.12.0"
 pre-commit = "^4.0.1"
-huggingface-hub = ">=0.25.1,<0.30.0"
+huggingface-hub = ">=0.25.1,<0.29.0"
 
 # Docs
 mkdocs = ">=1.4.2"
@@ -87,13 +87,7 @@ build-backend = "poetry.core.masonry.api"
 
 
 [tool.pytest.ini_options]
-addopts = "--durations=5 -m 'not slow' -n 3"
+addopts = "--durations=5 -m 'not slow' -n auto"
 asyncio_mode = "auto"
 asyncio_default_fixture_loop_scope = "function"
 markers = "slow: marks tests as slow"
-
-[project]
-# MCP requires the following fields to be present in the pyproject.toml file
-name = "agentic_security"
-version = "1.0.0"
-requires-python = ">=3.11"

tests/conftest.py

@@ -1,17 +1,8 @@
 import os
 
 import pytest
-from cache_to_disk import delete_old_disk_caches
-
-from agentic_security.logutils import logger
 
 
 def pytest_runtest_setup(item):
     if "slow" in item.keywords and not os.getenv("RUN_SLOW_TESTS"):
         pytest.skip("Skipping slow test")
-
-
-@pytest.fixture(autouse=True, scope="session")
-def setup_delete_old_disk_caches():
-    logger.info("delete_old_disk_caches")
-    delete_old_disk_caches()

tests/probe_actor/test_fuzzer.py

@@ -7,7 +7,6 @@ import pytest
 
 from agentic_security.primitives import Scan
 from agentic_security.probe_actor.fuzzer import (
-    FuzzerState,
     generate_prompts,
     perform_many_shot_scan,
     perform_single_shot_scan,
@@ -208,7 +207,9 @@ class TestProcessPrompt(unittest.IsolatedAsyncioTestCase):
             prompt="test prompt",
             tokens=0,
             module_name="module_a",
-            fuzzer_state=FuzzerState(),
+            refusals=[],
+            errors=[],
+            outputs=[],
         )
 
         self.assertEqual(tokens, 3)  # Tokens from "Valid response text"
@@ -225,17 +226,20 @@ class TestProcessPrompt(unittest.IsolatedAsyncioTestCase):
             )
         )
 
-        fuzzer_state = FuzzerState()
+        refusals = []
+        outputs = []
         tokens, refusal = await process_prompt(
             request_factory=mock_request_factory,
             prompt="test prompt",
             tokens=0,
             module_name="module_a",
-            fuzzer_state=fuzzer_state,
+            refusals=refusals,
+            errors=[],
+            outputs=outputs,
         )
 
         self.assertEqual(tokens, 3)  # Tokens from "Response indicating refusal"
-        # self.assertFalse(fuzzer_state.refusals)
+        self.assertFalse(refusal)
 
     async def test_http_error_response(self):
         mock_request_factory = Mock()
@@ -248,13 +252,15 @@ class TestProcessPrompt(unittest.IsolatedAsyncioTestCase):
             )
         )
 
-        fuzzer_state = FuzzerState()
+        refusals = []
         await process_prompt(
             request_factory=mock_request_factory,
             prompt="test prompt",
             tokens=0,
             module_name="module_a",
-            fuzzer_state=fuzzer_state,
+            refusals=refusals,
+            errors=[],
+            outputs=[],
         )
 
     async def test_request_error(self):
@@ -263,14 +269,18 @@ class TestProcessPrompt(unittest.IsolatedAsyncioTestCase):
             side_effect=httpx.RequestError("Connection error")
         )
 
-        fuzzer_state = FuzzerState()
+        errors = []
         tokens, refusal = await process_prompt(
             request_factory=mock_request_factory,
             prompt="test prompt",
             tokens=0,
             module_name="module_a",
-            fuzzer_state=fuzzer_state,
+            refusals=[],
+            errors=errors,
+            outputs=[],
         )
 
         self.assertEqual(tokens, 0)
         self.assertTrue(refusal)
+        self.assertEqual(len(errors), 1)
+        self.assertIn("Connection error", errors[0][3])

tests/test_report_chart.py

@@ -0,0 +1,136 @@
+import io
+import string
+import pytest
+import pandas as pd
+import numpy as np
+from agentic_security.report_chart import plot_security_report, generate_identifiers
+
+class TestReportChart:
+    """Test suite for agentic_security.report_chart module."""
+
+    def test_generate_identifiers_short(self):
+        """Test generate_identifiers with a small dataset."""
+        df = pd.DataFrame([{'dummy': i} for i in range(5)])
+        identifiers = generate_identifiers(df)
+        expected = ['A1', 'A2', 'A3', 'A4', 'A5']
+        assert identifiers == expected
+
+    def test_generate_identifiers_edge(self):
+        """Test generate_identifiers with more than 26 items to cover cycling over the alphabet."""
+        n = 30
+        df = pd.DataFrame([{'dummy': i} for i in range(n)])
+        identifiers = generate_identifiers(df)
+        # For i=25, identifier should be A26, and for i=26, identifier should be B1
+        assert identifiers[25] == 'A26'
+        assert identifiers[26] == 'B1'
+        assert len(identifiers) == n
+
+    def test_generate_identifiers_empty(self):
+        """Test generate_identifiers with an empty dataframe."""
+        df = pd.DataFrame([])
+        identifiers = generate_identifiers(df)
+        assert identifiers == []
+
+    def test_plot_security_report_png_output(self):
+        """Test plot_security_report returns valid PNG output."""
+        # Create a sample table with required columns
+        table = [
+            {"failureRate": 10, "tokens": 100, "module": "Module1"},
+            {"failureRate": 30, "tokens": 200, "module": "Module2"},
+            {"failureRate": 20, "tokens": 150, "module": "Module3"},
+        ]
+        buf = plot_security_report(table)
+        # Check that buf is a BytesIO object and starts with PNG header bytes
+        assert isinstance(buf, io.BytesIO)
+        buf.seek(0)
+        header = buf.read(8)
+        assert header.startswith(b'\x89PNG')
+
+    def test_plot_security_report_ordering(self, monkeypatch):
+        """Test that the table embedded in the plot contains correctly sorted order by descending failure rate."""
+        table = [
+            {"failureRate": 15, "tokens": 110, "module": "ModuleA"},
+            {"failureRate": 25, "tokens": 210, "module": "ModuleB"},
+            {"failureRate": 5,  "tokens": 90,  "module": "ModuleC"},
+        ]
+        result_holder = {}
+        from matplotlib.axes import Axes
+        original_table = Axes.table
+        def fake_table(self, *args, **kwargs):
+            result_holder['cellText'] = kwargs.get('cellText')
+            return original_table(self, *args, **kwargs)
+        monkeypatch.setattr(Axes, "table", fake_table)
+        plot_security_report(table)
+        cell_text = result_holder.get('cellText')
+        assert cell_text is not None
+        # Verify header row in the table
+        assert cell_text[0] == ["Threat"]
+        # Since the data are sorted (highest failure rate first), ModuleB (25.0%) should appear in one of the rows.
+        found = any("ModuleB (25.0%)" in row[0] for row in cell_text[1:])
+        assert found
+
+    def test_plot_security_report_one_entry(self):
+        """Test plot_security_report with a single entry."""
+        table = [{"failureRate": 50, "tokens": 300, "module": "OnlyModule"}]
+        buf = plot_security_report(table)
+        assert isinstance(buf, io.BytesIO)
+        buf.seek(0)
+        content = buf.read()
+        assert content.startswith(b'\x89PNG')
+    def test_generate_identifiers_many(self):
+        """Test generate_identifiers with 52 items to verify identifier sequence."""
+        n = 52
+        df = pd.DataFrame([{'dummy': i} for i in range(n)])
+        identifiers = generate_identifiers(df)
+        assert identifiers[0] == "A1"
+        assert identifiers[25] == "A26"
+        assert identifiers[26] == "B1"
+        assert identifiers[51] == "B26"
+
+    def test_plot_security_report_missing_failureRate(self):
+        """Test plot_security_report raises KeyError when 'failureRate' column is missing."""
+        table = [{"tokens": 100, "module": "Mod1"}]  # Missing 'failureRate'
+        with pytest.raises(KeyError):
+            plot_security_report(table)
+
+    def test_plot_security_report_missing_tokens(self):
+        """Test plot_security_report raises KeyError when 'tokens' column is missing."""
+        table = [{"failureRate": 10, "module": "Mod1"}]  # Missing 'tokens'
+        with pytest.raises(KeyError):
+            plot_security_report(table)
+
+    def test_plot_security_report_empty_table(self):
+        """Test plot_security_report raises KeyError when the table is empty."""
+        table = []
+        with pytest.raises(KeyError):
+            plot_security_report(table)
+    def test_plot_security_report_missing_module(self):
+        """Test plot_security_report raises KeyError when 'module' column is missing."""
+        table = [{"failureRate": 10, "tokens": 100}]  # Missing 'module'
+        with pytest.raises(KeyError):
+            plot_security_report(table)
+
+    def test_plot_security_report_failure_rate_labels(self, monkeypatch):
+        """Test that plot_security_report calls ax.text for each failure rate bar label."""
+        table = [
+            {"failureRate": 10, "tokens": 100, "module": "Mod1"},
+            {"failureRate": 20, "tokens": 150, "module": "Mod2"},
+            {"failureRate": 30, "tokens": 200, "module": "Mod3"},
+        ]
+        # Count the number of times ax.text is called for drawing failure rate labels.
+        call_count = [0]
+        from matplotlib.axes import Axes
+        original_text = Axes.text
+        def fake_text(self, *args, **kwargs):
+            call_count[0] += 1
+            return original_text(self, *args, **kwargs)
+        monkeypatch.setattr(Axes, "text", fake_text)
+        plot_security_report(table)
+        # The loop inside plot_security_report calls ax.text once for each data point.
+        assert call_count[0] == len(table)
+
+    def test_plot_security_report_non_numeric_failureRate(self):
+        """Test that plot_security_report raises an exception when failureRate is non-numeric."""
+        table = [{"failureRate": "invalid", "tokens": 100, "module": "ModX"}]
+        with pytest.raises(Exception):
+            plot_security_report(table)

tests/test_scan.py

@@ -0,0 +1,126 @@
+import io
+import asyncio
+import json
+from datetime import datetime, timedelta
+from threading import Event
+import pytest
+from fastapi import FastAPI
+from fastapi.testclient import TestClient
+from agentic_security.routes import scan
+
+# Dummy LLMSpec for success tests
+class DummyLLMSpec:
+    def __init__(self, spec_string):
+        self.spec_string = spec_string
+    async def verify(self):
+        class DummyResponse:
+            status_code = 200
+            text = "verification succeeded"
+            elapsed = timedelta(seconds=0.5)
+        return DummyResponse()
+    @classmethod
+    def from_string(cls, spec_string):
+        return DummyLLMSpec(spec_string)
+
+# Dummy scan_router generator to simulate streaming responses
+async def dummy_scan_router(request_factory, scan_parameters, tools_inbox, stop_event):
+    for i in range(2):
+        yield f"result {i}"
+
+# Define a dummy Secrets class for testing purposes.
+class DummySecrets:
+    def __init__(self):
+        self.secrets = {}
+
+# Create FastAPI app for testing and include the scan router.
+@pytest.fixture
+def app():
+    app = FastAPI()
+    app.include_router(scan.router)
+    return app
+
+@pytest.fixture
+def client(app):
+    return TestClient(app)
+
+@pytest.fixture(autouse=True)
+def patch_dependencies(monkeypatch):
+    # Patch LLMSpec used in the routes with our dummy implementation.
+    monkeypatch.setattr(scan, "LLMSpec", DummyLLMSpec)
+    # Patch fuzzer.scan_router to use our dummy scanning generator.
+    monkeypatch.setattr(scan.fuzzer, "scan_router", dummy_scan_router)
+    # Patch get_stop_event to return a dummy Event.
+    dummy_event = Event()
+    monkeypatch.setattr(scan, "get_stop_event", lambda: dummy_event)
+    # Patch get_tools_inbox to return None.
+    monkeypatch.setattr(scan, "get_tools_inbox", lambda: None)
+    # Patch set_current_run to be a no-op.
+    monkeypatch.setattr(scan, "set_current_run", lambda x: None)
+    # Patch get_in_memory_secrets to return a DummySecrets instance.
+    monkeypatch.setattr(scan, "get_in_memory_secrets", lambda: DummySecrets())
+    # Ensure Scan.with_secrets is a no-op if not already implemented.
+    if not hasattr(scan.Scan, "with_secrets"):
+        monkeypatch.setattr(scan.Scan, "with_secrets", lambda self, secrets: None)
+
+def test_verify_success(client):
+    """Test /verify endpoint for a successful verification."""
+    data = {"spec": "dummy"}
+    response = client.post("/verify", json=data)
+    res_json = response.json()
+    assert response.status_code == 200
+    assert res_json["status_code"] == 200
+    assert res_json["body"] == "verification succeeded"
+    assert "elapsed" in res_json
+    assert "timestamp" in res_json
+
+def test_verify_failure(client, monkeypatch):
+    """Test /verify endpoint when verification fails."""
+    class DummyLLMSpecFailure:
+        def __init__(self, spec_string):
+            self.spec_string = spec_string
+        async def verify(self):
+            raise Exception("verification error")
+        @classmethod
+        def from_string(cls, spec_string):
+            return DummyLLMSpecFailure(spec_string)
+    monkeypatch.setattr(scan, "LLMSpec", DummyLLMSpecFailure)
+    data = {"spec": "bad"}
+    response = client.post("/verify", json=data)
+    assert response.status_code == 400
+    assert "verification error" in response.text
+
+def test_scan(client):
+    """Test /scan endpoint to ensure streaming response works."""
+    data = {"llmSpec": "dummy", "optimize": False, "maxBudget": 10, "enableMultiStepAttack": False}
+    response = client.post("/scan", json=data)
+    assert response.status_code == 200
+    content = list(response.iter_lines())
+    expected = ["result 0", "result 1"]
+    assert content == expected
+
+def test_stop_scan(client):
+    """Test /stop endpoint to ensure scan stopping functionality."""
+    dummy_event = scan.get_stop_event()
+    dummy_event.clear()
+    response = client.post("/stop")
+    assert response.status_code == 200
+    assert response.json() == {"status": "Scan stopped"}
+    assert dummy_event.is_set()
+
+def test_scan_csv(client):
+    """Test /scan-csv endpoint with CSV file and llmSpec upload."""
+    csv_content = b"col1,col2\nvalue1,value2"
+    llm_spec_content = b"dummy"
+    files = {
+        "file": ("dummy.csv", csv_content, "text/csv"),
+        "llmSpec": ("spec.txt", llm_spec_content, "text/plain"),
+    }
+    response = client.post(
+        "/scan-csv",
+        files=files,
+        data={"optimize": "false", "maxBudget": "10", "enableMultiStepAttack": "false"},
+    )
+    assert response.status_code == 200
+    content = list(response.iter_lines())
+    expected = ["result 0", "result 1"]
+    assert content == expected