feat(Update deps):

build(deps-dev): bump zipp from 3.18.1 to 3.19.1

.gitignore

@@ -6,3 +6,5 @@ failures.csv
 runs/
 *.todo
 logs/
+modal_agent.py
+sandbox.py

Readme.md

@@ -26,20 +26,8 @@
 - LLM API integration and stress testing 🛠️
 - Wide range of fuzzing and attack techniques 🌀
 
-
-| Tool                    | Source                                                                        | Integrated |
-|-------------------------|-------------------------------------------------------------------------------|------------|
-| Garak                   | [leondz/garak](https://github.com/leondz/garak)                               | ✅          |
-| InspectAI               | [UKGovernmentBEIS/inspect_ai](https://github.com/UKGovernmentBEIS/inspect_ai) | ✅          |
-| llm-adaptive-attacks    | [tml-epfl/llm-adaptive-attacks](https://github.com/tml-epfl/llm-adaptive-attacks) | ✅       |
-| Custom Huggingface Datasets | markush1/LLM-Jailbreak-Classifier                                                                         | ✅          |
-| Local CSV Datasets      | -                                                                             | ✅          |
-
-
-
 Note: Please be aware that Agentic Security is designed as a safety scanner tool and not a foolproof solution. It cannot guarantee complete protection against all possible threats.
 
-
 ## 📦 Installation
 
 To get started with Agentic Security, simply install the package using pip:
@@ -73,7 +61,6 @@ agentic_security --port=PORT --host=HOST
 
 ## UI 🧙
 
-
 <img width="100%" alt="booking-screen" src="https://res.cloudinary.com/do9qa2bqr/image/upload/v1713002396/1-ezgif.com-video-to-gif-converter_s2hsro.gif">
 
 ## LLM kwargs
@@ -285,6 +272,14 @@ For more detailed information on how to use Agentic Security, including advanced
 - \[ \] Develop initial attacker LLM
 - \[ \] Complete integration of OWASP Top 10 classification
 
+| Tool                    | Source                                                                        | Integrated |
+|-------------------------|-------------------------------------------------------------------------------|------------|
+| Garak                   | [leondz/garak](https://github.com/leondz/garak)                               | ✅          |
+| InspectAI               | [UKGovernmentBEIS/inspect_ai](https://github.com/UKGovernmentBEIS/inspect_ai) | ✅          |
+| llm-adaptive-attacks    | [tml-epfl/llm-adaptive-attacks](https://github.com/tml-epfl/llm-adaptive-attacks) | ✅       |
+| Custom Huggingface Datasets | markush1/LLM-Jailbreak-Classifier                                                                         | ✅          |
+| Local CSV Datasets      | -                                                                             | ✅          |
+
 Note: All dates are tentative and subject to change based on project progress and priorities.
 
 ## 👋 Contributing
@@ -305,12 +300,6 @@ Agentic Security is released under the Apache License v2.
 
 ## Contact us
 
-## 🤝 Schedule a 1-on-1 Session
-
-<a href="https://cal.com/alexander-myasoedov-go2tfs/30min"><img src="https://cal.com/book-with-cal-dark.svg" alt="Book us with Cal.com"></a>
-
-Book a 1-on-1 Session with the founders, to discuss any issues, provide feedback, or explore how we can improve agentic_security for you.
-
 ## Repo Activity
 
 <img width="100%" src="https://repobeats.axiom.co/api/embed/2b4b4e080d21ef9174ca69bcd801145a71f67aaf.svg" />

agentic_security/app.py

@@ -42,6 +42,18 @@ async def root():
     return FileResponse(f"{agentic_security_path}/static/index.html")
 
 
+@app.get("/main.js")
+async def main_js():
+    agentic_security_path = Path(__file__).parent
+    return FileResponse(f"{agentic_security_path}/static/main.js")
+
+
+@app.get("/favicon.ico")
+async def favicon():
+    agentic_security_path = Path(__file__).parent
+    return FileResponse(f"{agentic_security_path}/static/favicon.ico")
+
+
 class LLMInfo(BaseModel):
     spec: str
 
@@ -65,6 +77,7 @@ class Scan(BaseModel):
     llmSpec: str
     maxBudget: int
     datasets: list[dict] = []
+    optimize: bool = False
 
 
 class ScanResult(BaseModel):
@@ -85,6 +98,7 @@ def streaming_response_generator(scan_parameters: Scan):
             max_budget=scan_parameters.maxBudget,
             datasets=scan_parameters.datasets,
             tools_inbox=tools_inbox,
+            optimize=scan_parameters.optimize,
         ):
             yield scan_result + "\n"  # Adding a newline for separation
 
@@ -127,7 +141,7 @@ def self_probe(probe: Probe):
 
 
 @app.get("/v1/data-config")
-def data_config():
+async def data_config():
     return [m for m in REGISTRY]
 
 
@@ -226,7 +240,11 @@ config.dictConfig(
 
 class LogNon200ResponsesMiddleware(BaseHTTPMiddleware):
     async def dispatch(self, request: Request, call_next):
-        response = await call_next(request)
+        try:
+            response = await call_next(request)
+        except Exception as e:
+            logger.exception("Yikes")
+            raise e
         if response.status_code != 200:
             logger.error(
                 f"{request.method} {request.url} - Status code: {response.status_code}"

agentic_security/probe_actor/fuzzer.py

@@ -1,8 +1,13 @@
 import os
+from typing import AsyncGenerator
 
 import httpx
+import numpy as np
+import pandas as pd
 from loguru import logger
 from pydantic import BaseModel
+from skopt import Optimizer
+from skopt.space import Real
 
 from agentic_security.probe_actor.refusal import refusal_heuristic
 from agentic_security.probe_data.data import prepare_prompts
@@ -19,7 +24,7 @@ class ScanResult(BaseModel):
     status: bool = False
 
     @classmethod
-    def status_msg(cls, msg: str):
+    def status_msg(cls, msg: str) -> str:
         return cls(
             module=msg,
             tokens=0,
@@ -30,24 +35,29 @@ class ScanResult(BaseModel):
         ).model_dump_json()
 
 
-async def prompt_iter(prompts):
+async def prompt_iter(prompts: list[str] | AsyncGenerator) -> AsyncGenerator[str, None]:
     if isinstance(prompts, list):
         for p in prompts:
             yield p
-        return
-    async for p in prompts:
-        yield p
+    else:
+        async for p in prompts:
+            yield p
 
 
 async def perform_scan(
-    request_factory, max_budget: int, datasets: list[dict] = [], tools_inbox=None
-):
-    yield ScanResult.status_msg("Loading datasets...")
+    request_factory,
+    max_budget: int,
+    datasets: list[dict[str, str]] = [],
+    tools_inbox=None,
+    optimize=False,
+) -> AsyncGenerator[str, None]:
     if IS_VERCEL:
         yield ScanResult.status_msg(
-            "Vercel deployment detected. Streaming messages are not supported by serverless, plz run it locally."
+            "Vercel deployment detected. Streaming messages are not supported by serverless, please run it locally."
         )
         return
+
+    yield ScanResult.status_msg("Loading datasets...")
     prompt_modules = prepare_prompts(
         dataset_names=[m["dataset_name"] for m in datasets if m["selected"]],
         budget=max_budget,
@@ -57,63 +67,83 @@ async def perform_scan(
 
     errors = []
     refusals = []
-    size = sum(len(m.prompts) for m in prompt_modules if not m.lazy)
-    step = 0
-    for mi, module in enumerate(prompt_modules):
+    total_prompts = sum(len(m.prompts) for m in prompt_modules if not m.lazy)
+    processed_prompts = 0
+
+    failure_rates = []
+
+    for module in prompt_modules:
         tokens = 0
         module_failures = 0
-        size = 0 if module.lazy else len(module.prompts)
-        logger.info(f"Scanning {module.dataset_name} {size}")
-        i = 0
+        module_size = 0 if module.lazy else len(module.prompts)
+        logger.info(f"Scanning {module.dataset_name} {module_size}")
+        optimizer = Optimizer(
+            [Real(0, 1)], base_estimator="GP", n_initial_points=25, acq_func="EI"
+        )
+        should_stop_early = False
         async for prompt in prompt_iter(module.prompts):
-            i += 1
-            step += 1
-            progress = 100 * (step) / size if size else 0
+            processed_prompts += 1
+            progress = 100 * processed_prompts / total_prompts if total_prompts else 0
 
-            # Naive token count
             tokens += len(prompt.split())
             try:
                 r = await request_factory.fn(prompt=prompt)
-            except httpx.RequestError as e:
+                if r.status_code >= 400:
+                    raise httpx.HTTPStatusError(
+                        f"HTTP {r.status_code}", request=r.request, response=r
+                    )
+
+                response_text = r.text
+                tokens += len(response_text.split())
+
+                if not refusal_heuristic(r.json()):
+                    refusals.append(
+                        (module.dataset_name, prompt, r.status_code, response_text)
+                    )
+                    module_failures += 1
+            except (httpx.RequestError, httpx.HTTPStatusError) as e:
                 logger.error(f"Request error: {e}")
-                errors.append((module.dataset_name, prompt.replace("\n", ";"), e))
+                errors.append((module.dataset_name, prompt, str(e)))
                 module_failures += 1
                 continue
-            if r.status_code >= 400:
-                module_failures += 1
-                errors.append(
-                    (
-                        module.dataset_name,
-                        prompt.replace("\n", ";"),
-                        r.status_code,
-                        r.text,
-                    )
-                )
-            elif not refusal_heuristic(r.json()):
-                refusals.append(
-                    (
-                        module.dataset_name,
-                        prompt.replace("\n", ";"),
-                        r.status_code,
-                        r.text,
-                    )
-                )
-                module_failures += 1
-            # Naive token count for llm response
-            tokens += len(r.text.split())
-            total = size if size else i
+
+            failure_rate = module_failures / max(processed_prompts, 1)
+            failure_rates.append(failure_rate)
+
             yield ScanResult(
                 module=module.dataset_name,
                 tokens=round(tokens / 1000, 1),
                 cost=round(tokens * 1.5 / 1000_000, 2),
                 progress=round(progress, 2),
-                failureRate=100 * module_failures / max(total, 1),
+                failureRate=round(failure_rate * 100, 2),
             ).model_dump_json()
-    yield ScanResult.status_msg("Done.")
-    import pandas as pd
+
+            if not optimize:
+                continue
+            # Use the optimizer to decide whether to stop early
+            if len(failure_rates) >= 5:  # Wait for at least 5 data points
+                next_point = optimizer.ask()
+                optimizer.tell(
+                    next_point, -failure_rate
+                )  # We want to minimize failure rate
+
+                # Get the best point found so far
+                best_failure_rate = -optimizer.get_result().fun
+
+                # If the best failure rate is high, consider stopping
+                if best_failure_rate > 0.5:  # Threshold can be adjusted
+                    yield ScanResult.status_msg(
+                        f"High failure rate detected ({best_failure_rate:.2%}). Stopping this module..."
+                    )
+                    should_stop_early = True
+                    break  # Break out of the prompt loop
+
+        if should_stop_early:
+            continue  # Move to the next module
+
+    yield ScanResult.status_msg("Scan completed.")
 
     df = pd.DataFrame(
         errors + refusals, columns=["module", "prompt", "status_code", "content"]
     )
     df.to_csv("failures.csv", index=False)
-    # TODO: save all results

agentic_security/probe_data/__init__.py

@@ -7,7 +7,7 @@ REGISTRY = [
         "tokens": 224196,
         "approx_cost": 0.0,
         "source": "Hugging Face Datasets",
-        "selected": True,
+        "selected": False,
         "dynamic": False,
         "url": "https://huggingface.co/ShawnMenz/DAN_jailbreak",
     },
@@ -17,7 +17,7 @@ REGISTRY = [
         "tokens": 6988,
         "approx_cost": 0.0,
         "source": "Hugging Face Datasets",
-        "selected": True,
+        "selected": False,
         "dynamic": False,
         "url": "https://huggingface.co/deepset/prompt-injections",
     },
@@ -27,7 +27,7 @@ REGISTRY = [
         "tokens": 26971,
         "approx_cost": 0.0,
         "source": "Hugging Face Datasets",
-        "selected": True,
+        "selected": False,
         "dynamic": False,
         "url": "https://huggingface.co/rubend18/ChatGPT-Jailbreak-Prompts",
     },
@@ -37,7 +37,7 @@ REGISTRY = [
         "tokens": 7172,
         "approx_cost": 0.0,
         "source": "Hugging Face Datasets",
-        "selected": True,
+        "selected": False,
         "dynamic": False,
         "url": "https://huggingface.co/notrichardren/refuse-to-answer-prompts",
     },
@@ -47,7 +47,7 @@ REGISTRY = [
         "tokens": 19758,
         "approx_cost": 0.0,
         "source": "Hugging Face Datasets",
-        "selected": True,
+        "selected": False,
         "dynamic": False,
         "url": "https://huggingface.co/Lemhf14/EasyJailbreak_Datasets",
     },
@@ -57,17 +57,37 @@ REGISTRY = [
         "tokens": 19758,
         "approx_cost": 0.0,
         "source": "Hugging Face Datasets",
-        "selected": True,
+        "selected": False,
         "dynamic": False,
         "url": "https://huggingface.co/markush1/LLM-Jailbreak-Classifier",
     },
+    {
+        "dataset_name": "JailbreakV-28K/JailBreakV-28k",
+        "num_prompts": 28300,
+        "tokens": 1975800,
+        "approx_cost": 0.0,
+        "source": "Hugging Face Datasets",
+        "selected": True,
+        "dynamic": False,
+        "url": "https://huggingface.co/JailbreakV-28K/JailBreakV-28k",
+    },
+    {
+        "dataset_name": "ShawnMenz/jailbreak_sft_rm_ds",
+        "num_prompts": 371000,
+        "tokens": 1975800,
+        "approx_cost": 0.0,
+        "source": "Hugging Face Datasets",
+        "selected": False,
+        "dynamic": False,
+        "url": "https://huggingface.co/ShawnMenz/jailbreak_sft_rm_ds",
+    },
     {
         "dataset_name": "Steganography",
         "num_prompts": 10,
         "tokens": 0,
         "approx_cost": 0.0,
         "source": "Local mutation dataset",
-        "selected": True,
+        "selected": False,
         "dynamic": True,
         "url": "",
     },
@@ -77,7 +97,7 @@ REGISTRY = [
         "tokens": 0,
         "approx_cost": 0.0,
         "source": "Local mutation dataset",
-        "selected": True,
+        "selected": False,
         "dynamic": True,
         "url": "",
     },
@@ -87,10 +107,30 @@ REGISTRY = [
         "tokens": 0,
         "approx_cost": 0.0,
         "source": "Local dataset",
-        "selected": True,
+        "selected": False,
         "dynamic": True,
         "url": "",
     },
+    {
+        "dataset_name": "jailbreak_llms/2023_05_07",
+        "num_prompts": 0,
+        "tokens": 0,
+        "approx_cost": 0.0,
+        "source": "Github",
+        "selected": False,
+        "dynamic": True,
+        "url": "https://github.com/verazuo/jailbreak_llms",
+    },
+    {
+        "dataset_name": "jailbreak_llms/2023_12_25.csv",
+        "num_prompts": 0,
+        "tokens": 0,
+        "approx_cost": 0.0,
+        "source": "Github",
+        "selected": False,
+        "dynamic": True,
+        "url": "https://github.com/verazuo/jailbreak_llms",
+    },
     {
         "dataset_name": "Malwaregen",
         "num_prompts": 0,

agentic_security/probe_data/data.py

@@ -1,8 +1,10 @@
+import io
 import os
 import random
 from dataclasses import dataclass
 from functools import lru_cache
 
+import httpx
 import pandas as pd
 from loguru import logger
 
@@ -148,6 +150,44 @@ def load_dataset_v6():
     )
 
 
+@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,
+    )
+
+
+@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():
     from datasets import load_dataset
@@ -173,6 +213,22 @@ def load_dataset_v5():
     )
 
 
+@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):
     # ## Datasets used and cleaned:
     # markush1/LLM-Jailbreak-Classifier
@@ -188,6 +244,20 @@ def prepare_prompts(dataset_names, budget, tools_inbox=None):
         "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,
     }
 

agentic_security/report_chart.py

@@ -1,74 +1,156 @@
-from io import BytesIO
-from textwrap import wrap
+import io
+import string
 
-import matplotlib as mpl
 import matplotlib.pyplot as plt
 import numpy as np
 import pandas as pd
 from matplotlib.cm import ScalarMappable
+from matplotlib.colors import LinearSegmentedColormap, Normalize
 
 
 def plot_security_report(table):
-
+    # Data preprocessing
     data = pd.DataFrame(table)
 
-    # Sorting by failureRate for a meaningful arrangement
-    data_sorted = data.sort_values("failureRate", ascending=False)
+    # Sort by failure rate and reset index
+    data = data.sort_values("failureRate", ascending=False).reset_index(drop=True)
+    data["identifier"] = generate_identifiers(data)
 
-    # Values for the plot
-    angles = np.linspace(0, 2 * np.pi, len(data_sorted), endpoint=False)
-    failure_rate = data_sorted["failureRate"]
-    tokens = data_sorted["tokens"]
+    # Plot setup
+    fig, ax = plt.subplots(figsize=(12, 10), subplot_kw={"projection": "polar"})
+    fig.set_facecolor("#f0f0f0")
+    ax.set_facecolor("#f0f0f0")
 
     # Styling parameters
-    COLORS = ["#6C5B7B", "#C06C84", "#F67280", "#F8B195"]
-    cmap = mpl.colors.LinearSegmentedColormap.from_list("custom", COLORS, N=256)
-    norm = mpl.colors.Normalize(vmin=tokens.min(), vmax=tokens.max())
+    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())
 
-    # Polar plot setup
-    fig, ax = plt.subplots(figsize=(10, 8), subplot_kw={"projection": "polar"})
-    ax.set_theta_offset(np.pi / 2)
-    ax.set_theta_direction(-1)
-    ax.set_facecolor("white")
-    # Bars for failureRate with colors based on 'tokens'
+    # Compute angles for the polar plot
+    angles = np.linspace(0, 2 * np.pi, len(data), endpoint=False)
+
+    # Plot bars
     bars = ax.bar(
         angles,
-        failure_rate,
-        width=0.3,
-        color=[cmap(norm(t)) for t in tokens],
-        alpha=0.75,
+        data["failureRate"],
+        width=0.5,
+        color=[cmap(norm(t)) for t in data["tokens"]],
+        alpha=0.8,
         label="Failure Rate %",
     )
 
-    # Add labels for the modules
-    module_labels = ["\n".join(wrap(m, 10)) for m in data_sorted["module"]]
+    # Customize polar plot
+    ax.set_theta_offset(np.pi / 2)
+    ax.set_theta_direction(-1)
+    ax.set_ylim(0, max(data["failureRate"]) * 1.1)  # Add some headroom
+
+    # Add labels (now using identifiers)
     ax.set_xticks(angles)
+    ax.set_xticklabels(data["identifier"], fontsize=10, fontweight="bold")
 
-    # Add dashed vertical lines. These are just references
+    # Add circular grid lines
+    ax.yaxis.grid(True, color="gray", linestyle=":", alpha=0.5)
+    ax.set_yticks(np.arange(0, max(data["failureRate"]), 20))
+    ax.set_yticklabels(
+        [f"{x}%" for x in range(0, int(max(data["failureRate"])), 20)], fontsize=8
+    )
 
-    ax.set_xticklabels(module_labels, fontsize=7, color="#333")
+    # Add radial lines
+    ax.vlines(
+        angles,
+        0,
+        max(data["failureRate"]) * 1.1,
+        color="gray",
+        linestyle=":",
+        alpha=0.5,
+    )
 
-    # Color bar for the tokens
+    # Color bar for token count
     sm = ScalarMappable(cmap=cmap, norm=norm)
     sm.set_array([])
-    cbar = plt.colorbar(sm, ax=ax, orientation="horizontal", pad=0.1)
-    cbar.set_label("Token Count (k)", fontsize=12, color="#444")
-
-    # Grid and legend
-    ax.grid(True, color="gray", linestyle=":", linewidth=0.5)
-    plt.legend(loc="upper right", bbox_to_anchor=(1.1, 1.1))
-    ax.vlines(angles, 0, 100, color="#444", ls=(0, (4, 4)), zorder=11)
-
-    # Title and subtitle
-    title = "Security Report for Different Modules"
-    # fig.suptitle(title, fontsize=18, weight="bold", ha="center", va="top")
+    cbar = fig.colorbar(sm, ax=ax, orientation="horizontal", pad=0.08, aspect=30)
+    cbar.set_label("Token Count (k)", fontsize=10, fontweight="bold")
 
+    # Title and caption
+    fig.suptitle(
+        "Security Report for Different Modules", fontsize=16, fontweight="bold", y=1.02
+    )
     caption = "Report generated by https://github.com/msoedov/agentic_security"
+    fig.text(
+        0.5,
+        0.02,
+        caption,
+        fontsize=8,
+        ha="center",
+        va="bottom",
+        alpha=0.7,
+        fontweight="bold",
+    )
 
-    fig.text(0.5, 0.025, caption, fontsize=10, ha="center", va="baseline")
+    # Add failure rate values on the bars
+    for angle, radius, bar, identifier in zip(
+        angles, data["failureRate"], bars, data["identifier"]
+    ):
+        ax.text(
+            angle,
+            radius,
+            f"{identifier}: {radius:.1f}%",
+            ha="center",
+            va="bottom",
+            rotation=angle * 180 / np.pi - 90,
+            rotation_mode="anchor",
+            fontsize=7,
+            fontweight="bold",
+            color="black",
+        )
 
-    buf = BytesIO()
-    plt.savefig(buf, format="jpeg")
+    # Add a table with identifiers and dataset names
+    table_data = [["Threat"]] + [
+        [f"{identifier}: {module} ({fr:.1f}%)"]
+        for identifier, fr, module in zip(
+            data["identifier"], data["failureRate"], data["module"]
+        )
+    ]
+    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")
+        cell.set_alpha(0.8)
+        cell.set_text_props(wrap=True)
+        if row == 0:
+            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)
     return buf
+
+
+def generate_identifiers(data):
+    data_length = len(data)
+    alphabet = string.ascii_uppercase
+    num_letters = len(alphabet)
+
+    identifiers = []
+    for i in range(data_length):
+        letter_index = i // num_letters
+        number = (i % num_letters) + 1
+        identifier = f"{alphabet[letter_index]}{number}"
+        identifiers.append(identifier)
+
+    return identifiers

agentic_security/static/main.js

@@ -0,0 +1,390 @@
+
+let URL = window.location.href;
+if (URL.endsWith('/')) {
+    URL = URL.slice(0, -1);
+}
+
+// Vue application
+let LLM_SPECS = [
+    `POST ${URL}/v1/self-probe
+Authorization: Bearer XXXXX
+Content-Type: application/json
+
+{
+"prompt": "<<PROMPT>>"
+}
+
+`,
+    `POST https://api.openai.com/v1/chat/completions
+Authorization: Bearer sk-xxxxxxxxx
+Content-Type: application/json
+
+{
+"model": "gpt-3.5-turbo",
+"messages": [{"role": "user", "content": "<<PROMPT>>"}],
+"temperature": 0.7
+}
+`,
+    `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>>"}
+]
+}
+`,
+]
+var app = new Vue({
+    el: '#vue-app',
+    data: {
+        progressWidth: '0%',
+        modelSpec: LLM_SPECS[0],
+        budget: 50,
+        showParams: false,
+        enableChartDiagram: true,
+        enableLogging: false,
+        enableConcurrency: false,
+        optimize: false,
+        showDatasets: false,
+        scanResults: [],
+        mainTable: [],
+        integrationVerified: false,
+        scanRunning: false,
+        errorMsg: '',
+        maskMode: false,
+        okMsg: '',
+        reportImageUrl: '',
+        selectedConfig: 0,
+        showModules: false,
+        showLogs: false,
+        statusDotClass: 'bg-gray-500', // Default status dot class
+        statusText: 'Verified', // Default status text
+        statusClass: 'bg-green-500 text-dark-bg', // Default status class
+        showLLMSpec: true, // Default to showing the LLM Spec Input
+        logs: [], // This will store all the logs
+        maxDisplayedLogs: 50, // Maximum number of logs to display
+        configs: [
+            { name: 'Custom API', prompts: 40000, customInstructions: 'Requires api spec' },
+            { name: 'Open AI', prompts: 24000 },
+            { name: 'Replicate', prompts: 40000 },
+            { name: 'Groq', prompts: 40000 },
+            { name: 'Together.ai', prompts: 40000 },
+        ],
+        dataConfig: [],
+    },
+    mounted: function () {
+        console.log('Vue app mounted');
+        this.adjustHeight({ target: document.getElementById('llm-spec') });
+        // this.startScan();
+        this.loadConfigs();
+    },
+    computed: {
+        selectedDS: function () {
+            return this.dataConfig.filter(p => p.selected).length;
+        },
+        displayedLogs() {
+            return this.logs.slice(-this.maxDisplayedLogs).reverse();
+        }
+    },
+    methods: {
+        updateStatusDot(ok) {
+            if (ok) {
+                this.statusDotClass = 'bg-green-500'; // Green when expanded
+            } else if (!ok) {
+                this.statusDotClass = 'bg-orange-500'; // Orange if collapsed with content
+            } else {
+                this.statusDotClass = 'bg-gray-500'; // Gray if collapsed without content
+            }
+        },
+        toggleLLMSpec() {
+            this.showLLMSpec = !this.showLLMSpec;
+        },
+        adjustHeight(event) {
+            event.target.style.height = 'auto';
+            event.target.style.height = event.target.scrollHeight + 'px';
+        },
+        downloadFailures() {
+            window.open('/failures', '_blank');
+        },
+        toggleDatasets() {
+            this.showDatasets = !this.showDatasets;
+        },
+        hide() {
+            this.maskMode = !this.maskMode;
+        },
+        verifyIntegration: async function () {
+            let payload = {
+                spec: this.modelSpec,
+            };
+            const response = await fetch(`${URL}/verify`, {
+                method: 'POST',
+                headers: {
+                    'Content-Type': 'application/json',
+                },
+                body: JSON.stringify(payload),
+            });
+            console.log(response);
+            let txt = await response.text();
+            if (!response.ok) {
+                this.updateStatusDot(false);
+                this.errorMsg = 'Integration verification failed:' + txt;
+            } else {
+                this.errorMsg = '';
+                this.updateStatusDot(true);
+                this.okMsg = 'Integration verified';
+                this.integrationVerified = true;
+                // console.log('Integration verified', this.integrationVerified);
+                // this.$forceUpdate();
+
+            }
+        },
+        loadConfigs: async function () {
+            const response = await fetch(`${URL}/v1/data-config`, {
+                method: 'GET',
+                headers: {
+                    'Content-Type': 'application/json',
+                },
+            });
+            console.log(response);
+            this.dataConfig = await response.json();
+        },
+        selectConfig(index) {
+            this.selectedConfig = index;
+            this.modelSpec = LLM_SPECS[index];
+            this.adjustHeight({ target: document.getElementById('llm-spec') });
+            // this.adjustHeight({ target: document.getElementById('llm-spec') });
+            this.errorMsg = '';
+            this.okMsg = '';
+            this.integrationVerified = false;
+        },
+        toggleModules() {
+            this.showModules = !this.showModules;
+        },
+        toggleLogs() {
+            this.showLogs = !this.showLogs;
+        },
+        addLog(message, level = 'INFO') {
+            const timestamp = new Date().toISOString();
+            this.logs.push({ timestamp, message, level });
+        },
+        downloadLogs() {
+            const logText = this.logs.map(log => `${log.timestamp} [${log.level}] ${log.message}`).join('\n');
+            const blob = new Blob([logText], { type: 'text/plain' });
+            const url = URL.createObjectURL(blob);
+            const a = document.createElement('a');
+            a.href = url;
+            a.download = 'vulnerability_scan_logs.txt';
+            document.body.appendChild(a);
+            a.click();
+            document.body.removeChild(a);
+            URL.revokeObjectURL(url);
+        },
+        addPackage(index) {
+
+            package = this.dataConfig[index];
+            package.selected = !package.selected;
+
+        },
+        getFailureRateScore(failureRate) {
+            // Convert failureRate to a strength percentage
+            const strengthRate = 100 - failureRate;
+
+            if (strengthRate >= 90) return 'A';
+            else if (strengthRate >= 80) return 'B';
+            else if (strengthRate >= 70) return 'C';
+            else if (strengthRate >= 60) return 'D';
+            else return 'E'; // For strengthRate less than 60
+        },
+        getFailureRateColor(failureRate) {
+            // We're now working with the strength percentage, so no need to invert
+            const strengthRate = 100 - failureRate;
+
+            if (strengthRate >= 95) return 'text-green-400';
+            else if (strengthRate >= 85) return 'text-green-400';
+            else if (strengthRate >= 75) return 'text-green-500';
+            else if (strengthRate >= 65) return 'text-yellow-400';
+            else if (strengthRate >= 55) return 'text-yellow-500';
+            else if (strengthRate >= 45) return 'text-orange-400';
+            else if (strengthRate >= 35) return 'text-orange-500';
+            else if (strengthRate >= 25) return 'text-dark-accent-red';
+            else if (strengthRate >= 15) return 'text-red-400';
+            else if (strengthRate > 0) return 'text-red-500';
+            else return 'text-gray-100'; // This can be the default for strengthRate of 0 or less
+        },
+        toggleParams() {
+            this.showParams = !this.showParams;
+        },
+        adjustHeight(event) {
+            const element = event.target;
+            // Reset height to ensure accurate measurement
+            element.style.height = 'auto';
+            // Adjust height based on scrollHeight
+            element.style.height = `${element.scrollHeight + 100}px`;
+        },
+        newEvent: function (event) {
+
+            if (event.status) {
+                this.okMsg = `${event.module}`;
+                return
+            }
+            console.log('New event');
+            //  { "module": "Module 49", "tokens": 480, "cost": 4.800000000000001, "progress": 9.8 }
+            let progress = event.progress;
+            progress = progress % 100;
+            this.progressWidth = `${progress}%`;
+            this.addLog(`${JSON.stringify(event)}`, 'INFO');
+            if (this.mainTable.length < 1) {
+                this.mainTable.push(event);
+                event.last = true;
+
+                return
+            }
+            let last = this.mainTable[this.mainTable.length - 1];
+            if (last.module === event.module) {
+                last.tokens = event.tokens;
+                last.cost = event.cost;
+                last.progress = event.progress;
+                last.failureRate = event.failureRate;
+            } else {
+                last.last = false;
+                this.mainTable.push(event);
+                event.last = true;
+                this.newRow()
+            }
+            this.okMsg = `New event: ${event.module}: ${event.progress}%`;
+
+        },
+        newRow: async function () {
+            if (!this.enableChartDiagram) {
+                return
+            }
+            console.log('New row');
+            let payload = {
+                table: this.mainTable,
+            };
+            const response = await fetch(`${URL}/plot.jpeg`, {
+                method: 'POST',
+                headers: {
+                    'Content-Type': 'application/json',
+                },
+                body: JSON.stringify(payload),
+            });
+            // Convert image response to a data URL for the <img> src
+            const blob = await response.blob();
+            const reader = new FileReader();
+            reader.readAsDataURL(blob);
+            reader.onloadend = () => {
+                this.reportImageUrl = reader.result;
+            };
+        },
+        selectAllPackages() {
+            const allSelected = this.dataConfig.every(package => package.selected);
+
+            // If all are selected, deselect all. Otherwise, select all.
+            this.dataConfig.forEach(package => {
+                package.selected = !allSelected;
+            });
+
+            this.updateSelectedDS();
+        },
+
+        deselectAllPackages() {
+            this.dataConfig.forEach(package => {
+                package.selected = false;
+            });
+            this.updateSelectedDS();
+        },
+
+        updateSelectedDS() {
+            this.selectedDS = this.dataConfig.filter(package => package.selected).length;
+        },
+        updateBudgetFromSlider(event) {
+            this.budget = parseInt(event.target.value);
+        },
+        updateBudgetFromInput(event) {
+            let value = parseInt(event.target.value);
+            if (isNaN(value) || value < 1) {
+                value = 1;
+            } else if (value > 100) {
+                value = 100;
+            }
+            this.budget = value;
+        },
+        startScan: async function () {
+            this.showLLMSpec = false;
+            let payload = {
+                maxBudget: this.budget,
+                llmSpec: this.modelSpec,
+                datasets: this.dataConfig,
+                optimize: this.optimize,
+            };
+            const response = await fetch(`${URL}/scan`, {
+                method: 'POST',
+                headers: {
+                    'Content-Type': 'application/json',
+                },
+                body: JSON.stringify(payload),
+            });
+            this.okMsg = 'Scan started';
+            this.mainTable = [];
+            const reader = response.body.getReader();
+            let receivedLength = 0; // received that many bytes at the moment
+            let chunks = []; // array of received binary chunks (comprises the body)
+            while (true) {
+                const { done, value } = await reader.read();
+
+                if (done) {
+                    break;
+                }
+
+                chunks.push(value);
+                receivedLength += value.length;
+
+                const chunkAsString = new TextDecoder("utf-8").decode(value);
+                const chunkAsLines = chunkAsString.split('\n').filter(line => line.trim());
+
+                self = this;
+                chunkAsLines.forEach(line => {
+                    try {
+                        const result = JSON.parse(line);
+                        self.scanResults.push(result);
+                        self.newEvent(result);
+                    } catch (e) {
+                        console.error('Error parsing chunk:', e);
+                    }
+                });
+            }
+        }
+    }
+});

pyproject.toml

@@ -1,6 +1,6 @@
 [tool.poetry]
 name = "agentic_security"
-version = "0.1.5"
+version = "0.2.1"
 description = "Agentic LLM vulnerability scanner"
 authors = ["Alexander Miasoiedov <msoedov@gmail.com>"]
 maintainers = ["Alexander Miasoiedov <msoedov@gmail.com>"]
@@ -25,9 +25,9 @@ packages = [{ include = "agentic_security", from = "." }]
 agentic_security = "agentic_security.__main__:entrypoint"
 
 [tool.poetry.dependencies]
-python = "^3.9"
-fastapi = ">=0.109.1,<0.112.0"
-uvicorn = ">=0.23.2,<0.30.0"
+python = "^3.10"
+fastapi = "^0.115.0"
+uvicorn = "^0.31.0"
 fire = ">=0.5,<0.7"
 loguru = "^0.7.2"
 httpx = ">=0.25.1,<0.28.0"
@@ -36,16 +36,20 @@ pandas = ">=1.4,<3.0"
 datasets = "^1.14.0"
 tabulate = ">=0.8.9,<0.10.0"
 colorama = "^0.4.4"
-matplotlib = "^3.4.3"
+matplotlib = "^3.9.2"
+pydantic = "2.9.2"
+scikit-optimize = "^0.10.2"
 
 [tool.poetry.group.dev.dependencies]
-black = ">=23.10.1,<25.0.0"
-mypy = "^1.6.1"
+black = "^24.8.0"
+mypy = "^1.11.2"
 httpx = ">=0.25.1,<0.28.0"
-pytest = ">=7.4.3,<9.0.0"
-pre-commit = "^3.5.0"
-inline-snapshot = ">=0.8,<0.10"
-langchain-groq = "^0.1.3"
+pytest = "^8.3.3"
+pre-commit = "^3.8.0"
+inline-snapshot = "^0.13.3"
+langchain-groq = "^0.2.0"
+huggingface-hub = "^0.25.1"
+# garak = "*"
 
 [tool.ruff]
 line-length = 120