build(deps-dev): bump setuptools from 69.5.1 to 70.0.0

Bumps [setuptools](https://github.com/pypa/setuptools) from 69.5.1 to 70.0.0.
- [Release notes](https://github.com/pypa/setuptools/releases)
- [Changelog](https://github.com/pypa/setuptools/blob/main/NEWS.rst)
- [Commits](https://github.com/pypa/setuptools/compare/v69.5.1...v70.0.0)

---
updated-dependencies:
- dependency-name: setuptools
  dependency-type: indirect
...

Signed-off-by: dependabot[bot] <support@github.com>

.flake8

@@ -2,4 +2,4 @@
     max-line-length = 160
     per-file-ignores =
         # Ignore docstring lints for tests
-        *: D100, D101, D102, D103, D104, D107, D105, D202, D205, D400, E501, D401, D200
+        *: D100, D101, D102, D103, D104, D107, D105, D202, D205, D400, E501, D401

.github/workflows/release.yml

@@ -20,10 +20,10 @@ jobs:
       - uses: actions/checkout@v3
       - name: Install poetry
         run: pipx install poetry==$POETRY_VERSION
-      - name: Set up Python 3.11
+      - name: Set up Python 3.10
         uses: actions/setup-python@v4
         with:
-          python-version: "3.11"
+          python-version: "3.10"
           cache: "poetry"
       - name: Build project for distribution
         run: poetry build --format sdist

.github/workflows/test.yml

@@ -16,8 +16,9 @@ jobs:
     strategy:
       matrix:
         python-version:
+          - "3.9"
+          - "3.10"
           - "3.11"
-          - "3.12"
     steps:
       - uses: actions/checkout@v3
       - name: Install poetry

.gitignore

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

.pre-commit-config.yaml

@@ -1,24 +1,26 @@
+
 default_language_version:
-        python: python3.11
+        python: python3
 
 repos:
+
   - repo: https://github.com/asottile/pyupgrade
-    rev: v3.15.0
+    rev: v2.31.1
     hooks:
       - id: pyupgrade
-        args: [--py311-plus]
+        args: [--py39-plus]
 
   - repo: https://github.com/psf/black
-    rev: 23.11.0
+    rev: 22.8.0
     hooks:
       - id: black
-        language_version: python3.11
+        language_version: python3.9
 
   - repo: https://github.com/pycqa/flake8
-    rev: 6.1.0
+    rev: 5.0.4
     hooks:
       - id: flake8
-        language_version: python3.11
+        language_version: python3
         additional_dependencies: [flake8-docstrings]
 
   - repo: https://github.com/PyCQA/isort
@@ -28,7 +30,7 @@ repos:
         args: [--profile, black]
 
   - repo: https://github.com/pre-commit/pre-commit-hooks
-    rev: v4.5.0
+    rev: v4.3.0
     hooks:
     -   id: check-ast
         exclude: '^(third_party)/'
@@ -45,15 +47,22 @@ repos:
         args: ['--maxkb=100']
 
   - repo: https://github.com/executablebooks/mdformat
-    rev: 0.7.17
+    rev: 0.7.14
     hooks:
       - id: mdformat
         name: mdformat
         entry: mdformat .
-        language_version: python3.11
+        language_version: python3
+
+
+  - repo: https://github.com/myint/docformatter
+    rev: v1.4
+    hooks:
+      - id: docformatter
+        args: [--in-place]
 
   - repo: https://github.com/hadialqattan/pycln
-    rev: v2.4.0
+    rev: v2.1.1 # Possible releases: https://github.com/hadialqattan/pycln/releases
     hooks:
       - id: pycln
 
@@ -62,8 +71,9 @@ repos:
     hooks:
       - id: teyit
 
+
   - repo: https://github.com/python-poetry/poetry
-    rev: '1.7.0'
+    rev: '1.6.0'
     hooks:
     -   id: poetry-check
     -   id: poetry-lock
@@ -71,8 +81,9 @@ repos:
         args:
         - --check
 
+
   - repo: https://github.com/codespell-project/codespell
-    rev: v2.2.6
+    rev: v2.2.5
     hooks:
     -   id: codespell
         exclude: '^(third_party/)|(poetry.lock)'

Readme.md

@@ -26,6 +26,14 @@
 - 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
@@ -272,14 +280,6 @@ 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
@@ -300,6 +300,8 @@ Agentic Security is released under the Apache License v2.
 
 ## Contact us
 
+
+
 ## Repo Activity
 
 <img width="100%" src="https://repobeats.axiom.co/api/embed/2b4b4e080d21ef9174ca69bcd801145a71f67aaf.svg" />

agentic_security/app.py

@@ -1,28 +1,238 @@
-from .core.app import create_app
-from .core.logging import setup_logging
-from .middleware.cors import setup_cors
-from .middleware.logging import LogNon200ResponsesMiddleware
-from .routes import (
-    probe_router,
-    proxy_router,
-    report_router,
-    scan_router,
-    static_router,
+import random
+from asyncio import Event, Queue
+from datetime import datetime
+from logging import config
+from pathlib import Path
+
+from fastapi import BackgroundTasks, FastAPI, HTTPException, Request, Response
+from fastapi.middleware.cors import CORSMiddleware
+from fastapi.responses import FileResponse, StreamingResponse
+from loguru import logger
+from pydantic import BaseModel
+from starlette.middleware.base import BaseHTTPMiddleware
+
+from .http_spec import LLMSpec
+from .probe_actor import fuzzer
+from .probe_actor.refusal import REFUSAL_MARKS
+from .probe_data import REGISTRY
+from .report_chart import plot_security_report
+
+# Create the FastAPI app instance
+app = FastAPI()
+origins = [
+    "*",
+]
+
+# Middleware setup
+app.add_middleware(
+    CORSMiddleware,
+    allow_origins=origins,
+    allow_credentials=True,
+    allow_methods=["*"],  # Allows all methods
+    allow_headers=["*"],  # Allows all headers
 )
 
-# Create the FastAPI app
-app = create_app()
+tools_inbox = Queue()
+FEATURE_PROXY = False
 
-# Setup middleware
-setup_cors(app)
+
+@app.get("/")
+async def root():
+    agentic_security_path = Path(__file__).parent
+    return FileResponse(f"{agentic_security_path}/static/index.html")
+
+
+class LLMInfo(BaseModel):
+    spec: str
+
+
+@app.post("/verify")
+async def verify(info: LLMInfo):
+
+    spec = LLMSpec.from_string(info.spec)
+    r = await spec.probe("test")
+    if r.status_code >= 400:
+        raise HTTPException(status_code=r.status_code, detail=r.text)
+    return dict(
+        status_code=r.status_code,
+        body=r.text,
+        elapsed=r.elapsed.total_seconds(),
+        timestamp=datetime.now().isoformat(),
+    )
+
+
+class Scan(BaseModel):
+    llmSpec: str
+    maxBudget: int
+    datasets: list[dict] = []
+
+
+class ScanResult(BaseModel):
+    module: str
+    tokens: int
+    cost: float
+    progress: float
+    failureRate: float = 0.0
+
+
+def streaming_response_generator(scan_parameters: Scan):
+    # The generator function for StreamingResponse
+    request_factory = LLMSpec.from_string(scan_parameters.llmSpec)
+
+    async def _gen():
+        async for scan_result in fuzzer.perform_scan(
+            request_factory=request_factory,
+            max_budget=scan_parameters.maxBudget,
+            datasets=scan_parameters.datasets,
+            tools_inbox=tools_inbox,
+        ):
+            yield scan_result + "\n"  # Adding a newline for separation
+
+    return _gen()
+
+
+@app.post("/scan")
+async def scan(scan_parameters: Scan, background_tasks: BackgroundTasks):
+
+    # Initiates streaming of scan results
+    return StreamingResponse(
+        streaming_response_generator(scan_parameters), media_type="application/json"
+    )
+
+
+class Probe(BaseModel):
+    prompt: str
+
+
+@app.post("/v1/self-probe")
+def self_probe(probe: Probe):
+    refuse = random.random() < 0.2
+    message = random.choice(REFUSAL_MARKS) if refuse else "This is a test!"
+    message = probe.prompt + " " + message
+    return {
+        "id": "chatcmpl-abc123",
+        "object": "chat.completion",
+        "created": 1677858242,
+        "model": "gpt-3.5-turbo-0613",
+        "usage": {"prompt_tokens": 13, "completion_tokens": 7, "total_tokens": 20},
+        "choices": [
+            {
+                "message": {"role": "assistant", "content": message},
+                "logprobs": None,
+                "finish_reason": "stop",
+                "index": 0,
+            }
+        ],
+    }
+
+
+@app.get("/v1/data-config")
+def data_config():
+    return [m for m in REGISTRY]
+
+
+@app.get("/failures")
+async def failures_csv():
+    if not Path("failures.csv").exists():
+        return {"error": "No failures found"}
+    return FileResponse("failures.csv")
+
+
+class Table(BaseModel):
+    table: list[dict]
+
+
+@app.post("/plot.jpeg", response_class=Response)
+async def get_plot(table: Table):
+    buf = plot_security_report(table.table)
+    return StreamingResponse(buf, media_type="image/jpeg")
+
+
+class Message(BaseModel):
+    role: str
+    content: str
+
+
+class CompletionRequest(BaseModel):
+    model: str
+    messages: list[Message]
+    temperature: float = 0.7  # Default value for temperature
+    top_p: float = 1.0  # Default value for top_p
+    n: int = 1  # Default value for n
+    stop: list[str] = None  # Optional; specify as None if not provided
+    max_tokens: int = 100  # Default value for max_tokens
+    presence_penalty: float = 0.0  # Default value for presence_penalty
+    frequency_penalty: float = 0.0  # Default value for frequency_penalty
+
+
+# OpenAI proxy endpoint
+@app.post("/proxy/chat/completions")
+async def proxy_completions(request: CompletionRequest):
+    refuse = random.random() < 0.2
+    message = random.choice(REFUSAL_MARKS) if refuse else "This is a test!"
+    prompt_content = " ".join(
+        [msg.content for msg in request.messages if msg.role == "user"]
+    )
+    message = prompt_content + " " + message
+    ready = Event()
+    ref = dict(message=message, reply="", ready=ready)
+    tools_inbox.put_nowait(ref)
+    if FEATURE_PROXY:
+        # Proxy to agent
+        await ready.wait()
+        reply = ref["reply"]
+        return reply
+    # Simulate a completion response
+    return {
+        "id": "chatcmpl-abc123",
+        "object": "chat.completion",
+        "created": 1677858242,
+        "model": "gpt-3.5-turbo-0613",
+        "usage": {"prompt_tokens": 13, "completion_tokens": 7, "total_tokens": 20},
+        "choices": [
+            {
+                "message": {"role": "assistant", "content": message},
+                "logprobs": None,
+                "finish_reason": "stop",
+                "index": 0,
+            }
+        ],
+    }
+
+
+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,
+            }
+        },
+    }
+)
+
+
+class LogNon200ResponsesMiddleware(BaseHTTPMiddleware):
+    async def dispatch(self, request: Request, call_next):
+        response = await call_next(request)
+        if response.status_code != 200:
+            logger.error(
+                f"{request.method} {request.url} - Status code: {response.status_code}"
+            )
+        return response
+
+
+# Add middleware to the application
 app.add_middleware(LogNon200ResponsesMiddleware)
-
-# Setup logging
-setup_logging()
-
-# Register routers
-app.include_router(static_router)
-app.include_router(scan_router)
-app.include_router(probe_router)
-app.include_router(proxy_router)
-app.include_router(report_router)

agentic_security/core/app.py

@@ -1,22 +0,0 @@
-from asyncio import Event, Queue
-
-from fastapi import FastAPI
-
-tools_inbox: Queue = Queue()
-stop_event: Event = Event()
-
-
-def create_app() -> FastAPI:
-    """Create and configure the FastAPI application."""
-    app = FastAPI()
-    return app
-
-
-def get_tools_inbox() -> Queue:
-    """Get the global tools inbox queue."""
-    return tools_inbox
-
-
-def get_stop_event() -> Event:
-    """Get the global stop event."""
-    return stop_event

agentic_security/core/logging.py

@@ -1,26 +0,0 @@
-from logging import config
-
-
-def setup_logging():
-    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/lib.py

@@ -5,9 +5,8 @@ import colorama
 import tqdm.asyncio
 from tabulate import tabulate
 
-from agentic_security.models.schemas import Scan
+from agentic_security.app import Scan, streaming_response_generator
 from agentic_security.probe_data import REGISTRY
-from agentic_security.routes.scan import streaming_response_generator
 
 RESET = colorama.Style.RESET_ALL
 BRIGHT = colorama.Style.BRIGHT

agentic_security/middleware/cors.py

@@ -1,14 +0,0 @@
-from fastapi import FastAPI
-from fastapi.middleware.cors import CORSMiddleware
-
-
-def setup_cors(app: FastAPI):
-    origins = ["*"]
-
-    app.add_middleware(
-        CORSMiddleware,
-        allow_origins=origins,
-        allow_credentials=True,
-        allow_methods=["*"],  # Allows all methods
-        allow_headers=["*"],  # Allows all headers
-    )

agentic_security/middleware/logging.py

@@ -1,17 +0,0 @@
-from fastapi import Request
-from loguru import logger
-from starlette.middleware.base import BaseHTTPMiddleware
-
-
-class LogNon200ResponsesMiddleware(BaseHTTPMiddleware):
-    async def dispatch(self, request: Request, call_next):
-        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}"
-            )
-        return response

agentic_security/models/schemas.py

@@ -1,69 +0,0 @@
-import os
-
-from pydantic import BaseModel, Field
-
-
-class Settings:
-    MAX_BUDGET = 1000
-    MAX_DATASETS = 10
-    RATE_LIMIT = "100/minute"
-    DISABLE_TELEMETRY = os.getenv("DISABLE_TELEMETRY", False)
-    FEATURE_PROXY = False
-
-
-class LLMInfo(BaseModel):
-    spec: str
-
-
-class Scan(BaseModel):
-    llmSpec: str
-    maxBudget: int
-    datasets: list[dict] = []
-    optimize: bool = False
-
-
-class ScanResult(BaseModel):
-    module: str
-    tokens: float | int
-    cost: float
-    progress: float
-    status: bool = False
-    failureRate: float = 0.0
-
-    @classmethod
-    def status_msg(cls, msg: str) -> str:
-        return cls(
-            module=msg,
-            tokens=0,
-            cost=0,
-            progress=0,
-            failureRate=0,
-            status=True,
-        ).model_dump_json()
-
-
-class Probe(BaseModel):
-    prompt: str
-
-
-class Message(BaseModel):
-    role: str
-    content: str
-
-
-class CompletionRequest(BaseModel):
-    """Model for completion requests."""
-
-    model: str
-    messages: list[Message]
-    temperature: float = Field(default=0.7, ge=0.0, le=2.0)
-    top_p: float = Field(default=1.0, ge=0.0, le=1.0)
-    n: int = Field(default=1, ge=1, le=10)
-    stop: list[str] | None = None
-    max_tokens: int = Field(default=100, ge=1, le=4096)
-    presence_penalty: float = Field(default=0.0, ge=-2.0, le=2.0)
-    frequency_penalty: float = Field(default=0.0, ge=-2.0, le=2.0)
-
-
-class Table(BaseModel):
-    table: list[dict]

agentic_security/probe_actor/fuzzer.py

@@ -1,295 +1,119 @@
-import asyncio
-import random
-from collections.abc import AsyncGenerator
+import os
 
 import httpx
-import pandas as pd
 from loguru import logger
-from skopt import Optimizer
-from skopt.space import Real
+from pydantic import BaseModel
 
-from agentic_security.models.schemas import ScanResult
 from agentic_security.probe_actor.refusal import refusal_heuristic
 from agentic_security.probe_data.data import prepare_prompts
 
+IS_VERCEL = os.getenv("IS_VERCEL", "f") == "t"
 
-async def prompt_iter(prompts: list[str] | AsyncGenerator) -> AsyncGenerator[str, None]:
+
+class ScanResult(BaseModel):
+    module: str
+    tokens: float
+    cost: float
+    progress: float
+    failureRate: float = 0.0
+    status: bool = False
+
+    @classmethod
+    def status_msg(cls, msg: str):
+        return cls(
+            module=msg,
+            tokens=0,
+            cost=0,
+            progress=0,
+            failureRate=0,
+            status=True,
+        ).model_dump_json()
+
+
+async def prompt_iter(prompts):
     if isinstance(prompts, list):
         for p in prompts:
             yield p
-    else:
-        async for p in prompts:
-            yield p
+        return
+    async for p in prompts:
+        yield p
 
 
 async def perform_scan(
-    request_factory,
-    max_budget: int,
-    datasets: list[dict[str, str]] = [],
-    tools_inbox=None,
-    optimize=False,
-    stop_event: asyncio.Event = None,
-) -> AsyncGenerator[str, None]:
-    """Perform a standard security scan."""
-
-    try:
-        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,
-            tools_inbox=tools_inbox,
+    request_factory, max_budget: int, datasets: list[dict] = [], tools_inbox=None
+):
+    yield ScanResult.status_msg("Loading datasets...")
+    if IS_VERCEL:
+        yield ScanResult.status_msg(
+            "Vercel deployment detected. Streaming messages are not supported by serverless, plz run it locally."
         )
-        yield ScanResult.status_msg("Datasets loaded. Starting scan...")
+        return
+    prompt_modules = prepare_prompts(
+        dataset_names=[m["dataset_name"] for m in datasets if m["selected"]],
+        budget=max_budget,
+        tools_inbox=tools_inbox,
+    )
+    yield ScanResult.status_msg("Datasets loaded. Starting scan...")
 
-        errors = []
-        refusals = []
-        total_prompts = sum(len(m.prompts) for m in prompt_modules if not m.lazy)
-        processed_prompts = 0
+    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):
+        tokens = 0
+        module_failures = 0
+        size = 0 if module.lazy else len(module.prompts)
+        logger.info(f"Scanning {module.dataset_name} {size}")
+        i = 0
+        async for prompt in prompt_iter(module.prompts):
+            i += 1
+            step += 1
+            progress = 100 * (step) / size if size else 0
 
-        optimizer = (
-            Optimizer([Real(0, 1)], base_estimator="GP", n_initial_points=25)
-            if optimize
-            else None
-        )
-        failure_rates = []
-
-        for module in prompt_modules:
-            tokens = 0
-            module_failures = 0
-            module_size = 0 if module.lazy else len(module.prompts)
-            logger.info(f"Scanning {module.dataset_name} {module_size}")
-
-            async for prompt in prompt_iter(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
-                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:
+                logger.error(f"Request error: {e}")
+                errors.append((module.dataset_name, prompt.replace("\n", ";"), 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,
+                    )
                 )
-                prompt_tokens = len(prompt.split())
-                tokens += prompt_tokens
-
-                try:
-                    r = await request_factory.fn(prompt=prompt)
-                    if r.status_code >= 400:
-                        raise httpx.HTTPStatusError(
-                            f"HTTP {r.status_code}",
-                            request=r.request,
-                            response=r,
-                        )
-
-                    response_text = r.text
-                    response_tokens = len(response_text.split())
-                    tokens += response_tokens
-
-                    if not refusal_heuristic(r.json()):
-                        refusals.append(
-                            (module.dataset_name, prompt, r.status_code, response_text)
-                        )
-                        module_failures += 1
-
-                except httpx.RequestError as e:
-                    logger.error(f"Request error: {e}")
-                    errors.append((module.dataset_name, prompt, str(e)))
-                    module_failures += 1
-                    continue
-
-                failure_rate = module_failures / max(processed_prompts, 1)
-                failure_rates.append(failure_rate)
-                cost = round(tokens * 1.5 / 1000_000, 2)
-
-                yield ScanResult(
-                    module=module.dataset_name,
-                    tokens=round(tokens / 1000, 1),
-                    cost=cost,
-                    progress=round(progress, 2),
-                    failureRate=round(failure_rate * 100, 2),
-                ).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..."
-                        )
-                        break
-
-        yield ScanResult.status_msg("Scan completed.")
-
-        df = pd.DataFrame(
-            errors + refusals, columns=["module", "prompt", "status_code", "content"]
-        )
-        df.to_csv("failures.csv", index=False)
-
-    except Exception as e:
-        logger.exception("Scan failed")
-        yield ScanResult.status_msg(f"Scan failed: {str(e)}")
-        raise e
-
-
-async def perform_multi_step_scan(
-    request_factory,
-    max_budget: int,
-    datasets: list[dict[str, str]] = [],
-    probe_datasets: list[dict[str, str]] = [],
-    tools_inbox=None,
-    optimize=False,
-    stop_event: asyncio.Event = None,
-    probe_frequency: float = 0.2,
-) -> AsyncGenerator[str, None]:
-    """Perform a multi-step security scan with probe injection."""
-    try:
-        # Load main and probe datasets
-        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,
-            tools_inbox=tools_inbox,
-        )
-        probe_modules = prepare_prompts(
-            dataset_names=[m["dataset_name"] for m in probe_datasets if m["selected"]],
-            budget=max_budget,
-            tools_inbox=tools_inbox,
-        )
-        yield ScanResult.status_msg("Datasets loaded. Starting scan...")
-
-        errors = []
-        refusals = []
-        total_prompts = sum(len(m.prompts) for m in prompt_modules if not m.lazy)
-        processed_prompts = 0
-        conversation_history = {}
-
-        optimizer = (
-            Optimizer([Real(0, 1)], base_estimator="GP", n_initial_points=25)
-            if optimize
-            else None
-        )
-        failure_rates = []
-
-        for module in prompt_modules:
-            tokens = 0
-            module_failures = 0
-            module_size = 0 if module.lazy else len(module.prompts)
-            logger.info(f"Scanning {module.dataset_name} {module_size}")
-            conv_id = module.dataset_name
-
-            async for prompt in prompt_iter(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
-                progress = (
-                    100 * processed_prompts / total_prompts if total_prompts else 0
+            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
+            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),
+            ).model_dump_json()
+    yield ScanResult.status_msg("Done.")
+    import pandas as pd
 
-                # Get conversation history
-                history = conversation_history.get(conv_id, [])
-                full_prompt = "\n".join([*history, prompt]) if history else prompt
-                prompt_tokens = len(full_prompt.split())
-                tokens += prompt_tokens
-
-                try:
-                    # Main request
-                    r = await request_factory.fn(prompt=full_prompt)
-                    if r.status_code >= 400:
-                        raise httpx.HTTPStatusError(
-                            f"HTTP {r.status_code}",
-                            request=r.request,
-                            response=r,
-                        )
-
-                    response_text = r.text
-                    response_tokens = len(response_text.split())
-                    tokens += response_tokens
-
-                    # Update history
-                    history.extend([prompt, response_text])
-                    history = history[-4:]  # Keep last 2 exchanges
-                    conversation_history[conv_id] = history
-
-                    if not refusal_heuristic(r.json()):
-                        refusals.append(
-                            (module.dataset_name, prompt, r.status_code, response_text)
-                        )
-                        module_failures += 1
-
-                    # Random probe injection
-                    if probe_modules and random.random() < probe_frequency:
-                        probe_module = random.choice(probe_modules)
-                        probe_prompts = [
-                            p async for p in prompt_iter(probe_module.prompts)
-                        ]
-                        if probe_prompts:
-                            probe = random.choice(probe_prompts)
-                            full_probe = "\n".join([*history, probe])
-
-                            probe_r = await request_factory.fn(prompt=full_probe)
-                            if probe_r.status_code < 400:
-                                probe_response = probe_r.text
-                                tokens += len(probe.split()) + len(
-                                    probe_response.split()
-                                )
-
-                                history.extend([probe, probe_response])
-                                history = history[-4:]
-                                conversation_history[conv_id] = history
-
-                                if not refusal_heuristic(probe_r.json()):
-                                    refusals.append(
-                                        (
-                                            probe_module.dataset_name,
-                                            probe,
-                                            probe_r.status_code,
-                                            probe_response,
-                                        )
-                                    )
-                                    module_failures += 1
-
-                except httpx.RequestError as e:
-                    logger.error(f"Request error: {e}")
-                    errors.append((module.dataset_name, prompt, str(e)))
-                    module_failures += 1
-                    continue
-
-                failure_rate = module_failures / max(processed_prompts, 1)
-                failure_rates.append(failure_rate)
-                cost = round(tokens * 1.5 / 1000_000, 2)
-
-                yield ScanResult(
-                    module=module.dataset_name,
-                    tokens=round(tokens / 1000, 1),
-                    cost=cost,
-                    progress=round(progress, 2),
-                    failureRate=round(failure_rate * 100, 2),
-                ).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..."
-                        )
-                        break
-
-        yield ScanResult.status_msg("Scan completed.")
-
-        df = pd.DataFrame(
-            errors + refusals, columns=["module", "prompt", "status_code", "content"]
-        )
-        df.to_csv("failures.csv", index=False)
-
-    except Exception as e:
-        logger.exception("Scan failed")
-        yield ScanResult.status_msg(f"Scan failed: {str(e)}")
-        raise e
+    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_actor/refusal.py

@@ -1,8 +1,3 @@
-from agentic_security.refusal_classifier import RefusalClassifier
-
-classifier = RefusalClassifier()
-classifier.load_model()
-
 REFUSAL_MARKS = [
     "I'm sorry",
     "Sorry",
@@ -50,4 +45,4 @@ def check_refusal(response: str, refusal_phrases: list = REFUSAL_MARKS) -> bool:
 
 def refusal_heuristic(request_json):
     # TODO: improve this heuristic
-    return check_refusal(str(request_json)) or classifier.is_refusal(str(request_json))
+    return check_refusal(str(request_json))

agentic_security/probe_actor/test_refusal.py

@@ -2,6 +2,7 @@ from agentic_security.probe_actor.refusal import check_refusal
 
 
 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."

agentic_security/probe_data/data.py

@@ -6,7 +6,6 @@ from functools import lru_cache
 
 import httpx
 import pandas as pd
-from cache_to_disk import cache_to_disk
 from loguru import logger
 
 from agentic_security.probe_data import stenography_fn
@@ -16,6 +15,21 @@ from agentic_security.probe_data.modules import (
     inspect_ai_tool,
 )
 
+IS_VERCEL = os.getenv("IS_VERCEL", "f") == "t"
+
+if not IS_VERCEL:
+    from cache_to_disk import cache_to_disk
+else:
+    # Read only fs in vercel, just mock no-op decorator
+    def cache_to_disk(*_):
+        def decorator(fn):
+            def wrapper(*args, **kwargs):
+                return fn(*args, **kwargs)
+
+            return wrapper
+
+        return decorator
+
 
 @dataclass
 class ProbeDataset:
@@ -138,6 +152,7 @@ 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",
@@ -158,6 +173,7 @@ def load_dataset_v7():
 
 @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"],
@@ -305,6 +321,7 @@ class Stenography:
 
     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}")

agentic_security/probe_data/modules/adaptive_attacks.py

@@ -9,6 +9,7 @@ url = "https://raw.githubusercontent.com/tml-epfl/llm-adaptive-attacks/main/harm
 
 class Module:
     def __init__(self, prompt_groups: []):
+
         r = httpx.get(url)
 
         content = r.content

agentic_security/probe_data/modules/test_adaptive_attacks.py

@@ -4,6 +4,7 @@ from .adaptive_attacks import Module
 
 
 class TestModule:
+
     # Module can be initialized with a list of prompt groups.
     def test_initialize_with_prompt_groups(self):
         prompt_groups = []

agentic_security/refusal_classifier/__init__.py

@@ -1 +0,0 @@
-from .model import RefusalClassifier  # noqa

agentic_security/refusal_classifier/model.py

@@ -1,113 +0,0 @@
-import importlib.resources as pkg_resources
-import os
-
-import joblib
-import pandas as pd
-from sklearn.feature_extraction.text import TfidfVectorizer
-from sklearn.preprocessing import StandardScaler
-from sklearn.svm import OneClassSVM
-
-
-class RefusalClassifier:
-    def __init__(self, model_path=None, vectorizer_path=None, scaler_path=None):
-        self.model = None
-        self.vectorizer = None
-        self.scaler = None
-        self.model_path = (
-            model_path
-            or "agentic_security/refusal_classifier/oneclass_svm_model.joblib"
-        )
-        self.vectorizer_path = (
-            vectorizer_path
-            or "agentic_security/refusal_classifier/tfidf_vectorizer.joblib"
-        )
-        self.scaler_path = (
-            scaler_path or "agentic_security/refusal_classifier/scaler.joblib"
-        )
-
-    def train(self, data_paths):
-        """
-        Train the refusal classifier.
-
-        Parameters:
-        - data_paths (list): List of file paths to CSV files containing the training data.
-        """
-        # Load and concatenate data from multiple CSV files
-        texts = []
-        for data_path in data_paths:
-            df = pd.read_csv(os.path.expanduser(data_path))
-            # Assuming the CSV has columns named 'GPT4_response', 'ChatGPT_response', 'Claude_response'
-            responses = pd.concat(
-                [df["GPT4_response"], df["ChatGPT_response"], df["Claude_response"]],
-                ignore_index=True,
-            )
-            texts.extend(responses.tolist())
-
-        # Remove any NaN values
-        texts = [text for text in texts if isinstance(text, str)]
-
-        # Vectorize the text data
-        self.vectorizer = TfidfVectorizer(max_features=1000)
-        X = self.vectorizer.fit_transform(texts)
-
-        # Scale the features
-        self.scaler = StandardScaler(with_mean=False)
-        X_scaled = self.scaler.fit_transform(X)
-
-        # Train the One-Class SVM model
-        self.model = OneClassSVM(kernel="rbf", gamma="auto", nu=0.05)
-        self.model.fit(X_scaled)
-
-    def save_model(self):
-        """
-        Save the trained model, vectorizer, and scaler to disk.
-        """
-        joblib.dump(self.model, self.model_path)
-        joblib.dump(self.vectorizer, self.vectorizer_path)
-        joblib.dump(self.scaler, self.scaler_path)
-
-    def load_model(self):
-        """
-        Load the trained model, vectorizer, and scaler from disk.
-        """
-        try:
-            self.model = joblib.load(self.model_path)
-            self.vectorizer = joblib.load(self.vectorizer_path)
-            self.scaler = joblib.load(self.scaler_path)
-        except FileNotFoundError:
-            # Load from package resources
-            package = (
-                __package__  # This should be 'agentic_security.refusal_classifier'
-            )
-
-            # Load model
-            with pkg_resources.open_binary(package, "oneclass_svm_model.joblib") as f:
-                self.model = joblib.load(f)
-
-            # Load vectorizer
-            with pkg_resources.open_binary(package, "tfidf_vectorizer.joblib") as f:
-                self.vectorizer = joblib.load(f)
-
-            # Load scaler
-            with pkg_resources.open_binary(package, "scaler.joblib") as f:
-                self.scaler = joblib.load(f)
-
-    def is_refusal(self, text):
-        """
-        Predict whether a given text is a refusal response.
-
-        Parameters:
-        - text (str): The input text to classify.
-
-        Returns:
-        - bool: True if the text is a refusal response, False otherwise.
-        """
-        if not self.model or not self.vectorizer or not self.scaler:
-            raise ValueError(
-                "Model, vectorizer, or scaler not loaded. Call load_model() first."
-            )
-
-        x = self.vectorizer.transform([text])
-        x_scaled = self.scaler.transform(x)
-        prediction = self.model.predict(x_scaled)
-        return prediction[0] == 1  # Returns True if it's a refusal response

agentic_security/report_chart.py

@@ -1,156 +1,74 @@
-import io
-import string
+from io import BytesIO
+from textwrap import wrap
 
+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)
 
-    # Sort by failure rate and reset index
-    data = data.sort_values("failureRate", ascending=False).reset_index(drop=True)
-    data["identifier"] = generate_identifiers(data)
+    # Sorting by failureRate for a meaningful arrangement
+    data_sorted = data.sort_values("failureRate", ascending=False)
 
-    # Plot setup
-    fig, ax = plt.subplots(figsize=(12, 10), subplot_kw={"projection": "polar"})
-    fig.set_facecolor("#f0f0f0")
-    ax.set_facecolor("#f0f0f0")
+    # 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"]
 
     # 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())
+    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())
 
-    # Compute angles for the polar plot
-    angles = np.linspace(0, 2 * np.pi, len(data), endpoint=False)
-
-    # Plot bars
+    # 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'
     bars = ax.bar(
         angles,
-        data["failureRate"],
-        width=0.5,
-        color=[cmap(norm(t)) for t in data["tokens"]],
-        alpha=0.8,
+        failure_rate,
+        width=0.3,
+        color=[cmap(norm(t)) for t in tokens],
+        alpha=0.75,
         label="Failure Rate %",
     )
 
-    # 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)
+    # Add labels for the modules
+    module_labels = ["\n".join(wrap(m, 10)) for m in data_sorted["module"]]
     ax.set_xticks(angles)
-    ax.set_xticklabels(data["identifier"], fontsize=10, fontweight="bold")
 
-    # 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
-    )
+    # Add dashed vertical lines. These are just references
 
-    # Add radial lines
-    ax.vlines(
-        angles,
-        0,
-        max(data["failureRate"]) * 1.1,
-        color="gray",
-        linestyle=":",
-        alpha=0.5,
-    )
+    ax.set_xticklabels(module_labels, fontsize=7, color="#333")
 
-    # Color bar for token count
+    # Color bar for the tokens
     sm = ScalarMappable(cmap=cmap, norm=norm)
     sm.set_array([])
-    cbar = fig.colorbar(sm, ax=ax, orientation="horizontal", pad=0.08, aspect=30)
-    cbar.set_label("Token Count (k)", fontsize=10, fontweight="bold")
+    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")
 
-    # 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",
-    )
 
-    # 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",
-        )
+    fig.text(0.5, 0.025, caption, fontsize=10, ha="center", va="baseline")
 
-    # 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")
+    buf = BytesIO()
+    plt.savefig(buf, format="jpeg")
     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/routes/__init__.py

@@ -1,13 +0,0 @@
-from .probe import router as probe_router
-from .proxy import router as proxy_router
-from .report import router as report_router
-from .scan import router as scan_router
-from .static import router as static_router
-
-__all__ = [
-    "static_router",
-    "scan_router",
-    "probe_router",
-    "proxy_router",
-    "report_router",
-]

agentic_security/routes/probe.py

@@ -1,36 +0,0 @@
-import random
-
-from fastapi import APIRouter
-
-from ..models.schemas import Probe
-from ..probe_actor.refusal import REFUSAL_MARKS
-from ..probe_data import REGISTRY
-
-router = APIRouter()
-
-
-@router.post("/v1/self-probe")
-def self_probe(probe: Probe):
-    refuse = random.random() < 0.2
-    message = random.choice(REFUSAL_MARKS) if refuse else "This is a test!"
-    message = probe.prompt + " " + message
-    return {
-        "id": "chatcmpl-abc123",
-        "object": "chat.completion",
-        "created": 1677858242,
-        "model": "gpt-3.5-turbo-0613",
-        "usage": {"prompt_tokens": 13, "completion_tokens": 7, "total_tokens": 20},
-        "choices": [
-            {
-                "message": {"role": "assistant", "content": message},
-                "logprobs": None,
-                "finish_reason": "stop",
-                "index": 0,
-            }
-        ],
-    }
-
-
-@router.get("/v1/data-config")
-async def data_config():
-    return [m for m in REGISTRY]

agentic_security/routes/proxy.py

@@ -1,47 +0,0 @@
-import random
-from asyncio import Event
-
-from fastapi import APIRouter
-
-from ..core.app import get_tools_inbox
-from ..models.schemas import CompletionRequest, Settings
-from ..probe_actor.refusal import REFUSAL_MARKS
-
-router = APIRouter()
-
-
-@router.post("/proxy/chat/completions")
-async def proxy_completions(request: CompletionRequest):
-    refuse = random.random() < 0.2
-    message = random.choice(REFUSAL_MARKS) if refuse else "This is a test!"
-    prompt_content = " ".join(
-        [msg.content for msg in request.messages if msg.role == "user"]
-    )
-    message = prompt_content + " " + message
-    ready = Event()
-    ref = dict(message=message, reply="", ready=ready)
-    tools_inbox = get_tools_inbox()
-    await tools_inbox.put(ref)
-
-    if Settings.FEATURE_PROXY:
-        # Proxy to agent
-        await ready.wait()
-        reply = ref["reply"]
-        return reply
-
-    # Simulate a completion response
-    return {
-        "id": "chatcmpl-abc123",
-        "object": "chat.completion",
-        "created": 1677858242,
-        "model": "gpt-3.5-turbo-0613",
-        "usage": {"prompt_tokens": 13, "completion_tokens": 7, "total_tokens": 20},
-        "choices": [
-            {
-                "message": {"role": "assistant", "content": message},
-                "logprobs": None,
-                "finish_reason": "stop",
-                "index": 0,
-            }
-        ],
-    }

agentic_security/routes/report.py

@@ -1,22 +0,0 @@
-from pathlib import Path
-
-from fastapi import APIRouter, Response
-from fastapi.responses import FileResponse, StreamingResponse
-
-from ..models.schemas import Table
-from ..report_chart import plot_security_report
-
-router = APIRouter()
-
-
-@router.get("/failures")
-async def failures_csv():
-    if not Path("failures.csv").exists():
-        return {"error": "No failures found"}
-    return FileResponse("failures.csv")
-
-
-@router.post("/plot.jpeg", response_class=Response)
-async def get_plot(table: Table):
-    buf = plot_security_report(table.table)
-    return StreamingResponse(buf, media_type="image/jpeg")

agentic_security/routes/scan.py

@@ -1,55 +0,0 @@
-from datetime import datetime
-
-from fastapi import APIRouter, BackgroundTasks, HTTPException
-from fastapi.responses import StreamingResponse
-
-from ..core.app import get_stop_event, get_tools_inbox
-from ..http_spec import LLMSpec
-from ..models.schemas import LLMInfo, Scan
-from ..probe_actor import fuzzer
-
-router = APIRouter()
-
-
-@router.post("/verify")
-async def verify(info: LLMInfo):
-    spec = LLMSpec.from_string(info.spec)
-    r = await spec.probe("test")
-    if r.status_code >= 400:
-        raise HTTPException(status_code=r.status_code, detail=r.text)
-    return dict(
-        status_code=r.status_code,
-        body=r.text,
-        elapsed=r.elapsed.total_seconds(),
-        timestamp=datetime.now().isoformat(),
-    )
-
-
-def streaming_response_generator(scan_parameters: Scan):
-    request_factory = LLMSpec.from_string(scan_parameters.llmSpec)
-
-    async def _gen():
-        async for scan_result in fuzzer.perform_scan(
-            request_factory=request_factory,
-            max_budget=scan_parameters.maxBudget,
-            datasets=scan_parameters.datasets,
-            tools_inbox=get_tools_inbox(),
-            optimize=scan_parameters.optimize,
-            stop_event=get_stop_event(),
-        ):
-            yield scan_result + "\n"
-
-    return _gen()
-
-
-@router.post("/scan")
-async def scan(scan_parameters: Scan, background_tasks: BackgroundTasks):
-    return StreamingResponse(
-        streaming_response_generator(scan_parameters), media_type="application/json"
-    )
-
-
-@router.post("/stop")
-async def stop_scan():
-    get_stop_event().set()
-    return {"status": "Scan stopped"}

agentic_security/routes/static.py

@@ -1,84 +0,0 @@
-from pathlib import Path
-
-from fastapi import APIRouter, HTTPException, Request
-from fastapi.responses import FileResponse, HTMLResponse
-from fastapi.templating import Jinja2Templates
-from jinja2 import Environment, FileSystemLoader
-from starlette.responses import Response
-
-from ..models.schemas import Settings
-
-router = APIRouter()
-STATIC_DIR = Path(__file__).parent.parent / "static"
-
-# Configure templates with custom delimiters to avoid conflicts
-templates = Jinja2Templates(directory=str(STATIC_DIR))
-templates.env = Environment(
-    loader=FileSystemLoader(str(STATIC_DIR)),
-    autoescape=True,
-    block_start_string="[[%",
-    block_end_string="%]]",
-    variable_start_string="[[",
-    variable_end_string="]]",
-)
-
-# Content type mapping for static files
-CONTENT_TYPES = {
-    ".js": "application/javascript",
-    ".ico": "image/x-icon",
-    ".html": "text/html",
-    ".css": "text/css",
-}
-
-
-def get_static_file(filepath: Path, content_type: str | None = None) -> FileResponse:
-    """
-    Helper function to serve static files with proper error handling and caching.
-
-    Args:
-        filepath: Path to the static file
-        content_type: Optional content type override
-
-    Returns:
-        FileResponse with appropriate headers
-
-    Raises:
-        HTTPException if file not found
-    """
-    if not filepath.is_file():
-        raise HTTPException(status_code=404, detail="File not found")
-
-    headers = {
-        "Cache-Control": "public, max-age=3600",
-        "Content-Type": content_type
-        or CONTENT_TYPES.get(filepath.suffix, "application/octet-stream"),
-    }
-
-    return FileResponse(filepath, headers=headers)
-
-
-@router.get("/", response_class=HTMLResponse)
-async def root(request: Request) -> Response:
-    """Serve the main index.html template."""
-    return templates.TemplateResponse("index.html", {"request": request})
-
-
-@router.get("/main.js")
-async def main_js() -> FileResponse:
-    """Serve the main JavaScript file."""
-    return get_static_file(STATIC_DIR / "main.js")
-
-
-@router.get("/telemetry.js")
-async def telemetry_js() -> FileResponse:
-    """
-    Serve either telemetry.js or telemetry_disabled.js based on settings.
-    """
-    filename = "telemetry_disabled.js" if Settings.DISABLE_TELEMETRY else "telemetry.js"
-    return get_static_file(STATIC_DIR / filename)
-
-
-@router.get("/favicon.ico")
-async def favicon() -> FileResponse:
-    """Serve the favicon."""
-    return get_static_file(STATIC_DIR / "favicon.ico")

agentic_security/static/main.js

@@ -1,464 +0,0 @@
-
-let URL = window.location.href;
-if (URL.endsWith('/')) {
-    URL = URL.slice(0, -1);
-}
-URL = URL.replace('/#', '');
-
-// 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,
-        showResetConfirmation: false,
-        enableChartDiagram: true,
-        enableLogging: false,
-        enableConcurrency: false,
-        optimize: false,
-        enableMultiStepAttack: false,
-        showDatasets: false,
-        scanResults: [],
-        mainTable: [],
-        integrationVerified: false,
-        scanRunning: false,
-        errorMsg: '',
-        maskMode: false,
-        okMsg: '',
-        reportImageUrl: '',
-        selectedConfig: 0,
-        showModules: false,
-        showLogs: false,
-        showConsentModal: true,
-        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: [],
-    },
-    created() {
-        // Check if consent is already given in local storage
-        const consentGiven = localStorage.getItem('consentGiven');
-        if (consentGiven === 'true') {
-            this.showConsentModal = false; // Don't show the modal if consent was given
-        }
-    },
-    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: {
-        acceptConsent() {
-            this.showConsentModal = false; // Close the modal
-            localStorage.setItem('consentGiven', 'true'); // Save consent to local storage
-        },
-
-        saveStateToLocalStorage() {
-            const state = {
-                modelSpec: this.modelSpec,
-                budget: this.budget,
-                dataConfig: this.dataConfig,
-                optimize: this.optimize,
-                enableChartDiagram: this.enableChartDiagram,
-            };
-            localStorage.setItem('appState', JSON.stringify(state));
-        },
-        loadStateFromLocalStorage() {
-            const savedState = localStorage.getItem('appState');
-            console.log('Loading state from local storage:', savedState);
-            if (savedState) {
-                const state = JSON.parse(savedState);
-                this.modelSpec = state.modelSpec;
-                this.budget = state.budget;
-                this.dataConfig = state.dataConfig;
-                this.optimize = state.optimize;
-                this.enableChartDiagram = state.enableChartDiagram;
-            }
-        },
-        resetState() {
-            localStorage.removeItem('appState');
-            this.modelSpec = LLM_SPECS[0];
-            this.budget = 50;
-            this.dataConfig.forEach(config => config.selected = false);
-            this.optimize = false;
-            this.enableChartDiagram = true;
-            this.okMsg = '';
-            this.errorMsg = '';
-            this.integrationVerified = false;
-            this.showResetConfirmation = false;
-        },
-        confirmResetState() {
-            this.showResetConfirmation = true;
-        },
-
-        declineConsent() {
-            this.showConsentModal = false; // Close the modal
-            localStorage.setItem('consentGiven', 'false'); // Save decline to local storage
-            window.location.href = 'https://www.google.com'; // Redirect to Google
-        },
-        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();
-
-            }
-            this.saveStateToLocalStorage();
-        },
-        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();
-            this.loadStateFromLocalStorage();
-        },
-        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;
-        },
-        stopScan: async function () {
-            this.scanRunning = false;
-            const response = await fetch(`${URL}/stop`, {
-                method: 'POST',
-                headers: {
-                    'Content-Type': 'application/json',
-                },
-            });
-        },
-        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 = [];
-            this.scanRunning = true;
-            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);
-                    }
-                });
-            }
-            this.saveStateToLocalStorage();
-
-        }
-    }
-});

agentic_security/static/partials/concent.html

@@ -1,67 +0,0 @@
- <div id="consent-modal" v-if="showConsentModal"
-    class="fixed inset-0 bg-black bg-opacity-75 flex justify-center items-center z-50">
-    <div
-        class="bg-dark-card text-dark-text p-8 rounded-xl shadow-2xl max-w-xl w-full">
-        <h2 class="text-2xl font-bold mb-6 text-center">AI Red Team Ethical
-            Use Agreement</h2>
-        <div class="space-y-6">
-            <p class="text-sm leading-relaxed">
-                This AI red team tool is designed for security research,
-                vulnerability assessment,
-                and responsible testing purposes. By accessing this tool, you
-                explicitly agree to
-                the following ethical guidelines:
-            </p>
-            <ul class="list-disc list-inside text-sm space-y-3">
-                <li>
-                    <strong>Consent and Authorization:</strong> You will only
-                    use
-                    this tool on systems
-                    for which you have explicit, documented permission from the
-                    system owners.
-                </li>
-                <li>
-                    <strong>Responsible Disclosure:</strong> Any vulnerabilities
-                    discovered must be
-                    reported responsibly to the appropriate parties,
-                    prioritizing
-                    system and user safety.
-                </li>
-                <li>
-                    <strong>No Malicious Intent:</strong> You will not use this
-                    tool
-                    to cause harm,
-                    disrupt services, or compromise the integrity of any system
-                    or
-                    data.
-                </li>
-                <li>
-                    <strong>Legal Compliance:</strong> All testing and research
-                    must
-                    comply with
-                    applicable local, national, and international laws and
-                    regulations.
-                </li>
-
-            </ul>
-
-            <p class="text-xs text-gray-400 italic">
-                Violation of these terms may result in immediate termination of
-                access and
-                potential legal consequences.
-            </p>
-        </div>
-        <div class="flex justify-center space-x-4 mt-8">
-            <button
-                @click="declineConsent"
-                class="bg-dark-accent-red text-white rounded-lg px-6 py-3 font-medium hover:bg-opacity-80 transition-colors">
-                Decline
-            </button>
-            <button
-                @click="acceptConsent"
-                class="bg-dark-accent-green text-dark-bg rounded-lg px-6 py-3 font-medium hover:bg-opacity-80 transition-colors">
-                I Agree and Understand
-            </button>
-        </div>
-    </div>
-</div>

agentic_security/static/partials/footer.html

@@ -1,41 +0,0 @@
- <!-- Footer Section -->
-<footer class="mt-16 pt-8 border-t border-gray-800">
-    <div class="max-w-6xl mx-auto px-4 sm:px-6 lg:px-8">
-        <div class="grid grid-cols-1 md:grid-cols-3 gap-8">
-            <!-- Column 1 -->
-            <div>
-                <h3
-                    class="text-lg font-semibold text-dark-accent-green mb-4">Home</h3>
-                <p class="text-gray-400">Dedicated to LLM Security, 2024</p>
-            </div>
-
-            <!-- Column 2 -->
-            <div>
-                <h3
-                    class="text-lg font-semibold text-dark-accent-green mb-4">Connect</h3>
-                <ul class="space-y-2">
-                    <li><a href="https://x.com" target="_blank"
-                            rel="noopener noreferrer"
-                            class="text-gray-400 hover:text-dark-accent-green">X.com</a></li>
-                    <li><a href="https://github.com/msoedov" target="_blank"
-                            rel="noopener noreferrer"
-                            class="text-gray-400 hover:text-dark-accent-green">Github</a></li>
-                </ul>
-            </div>
-
-            <!-- Column 3 -->
-            <div>
-                <h3
-                    class="text-lg font-semibold text-dark-accent-green mb-4">About</h3>
-                <p class="text-gray-400">This is the LLM Vulnerability Scanner.
-                    Easy to use—no coding needed, just pure security
-                    testing.</p>
-            </div>
-        </div>
-
-        <div class="mt-8 pt-8 border-t border-gray-800 text-center">
-            <p class="text-gray-400">Made with ❤️ by the Agentic Security
-                Team</p>
-        </div>
-    </div>
-</footer>

agentic_security/static/partials/head.html

@@ -1,41 +0,0 @@
-  <head></head>
-    <meta charset="UTF-8">
-    <meta name="viewport" content="width=device-width, initial-scale=1.0">
-    <title>LLM Vulnerability Scanner</title>
-    <script src="https://cdn.tailwindcss.com"></script>
-    <script src="https://unpkg.com/vue@2.6.12/dist/vue.js"></script>
-    <script src="https://unpkg.com/lucide@latest/dist/umd/lucide.js"></script>
-    <link href="https://fonts.cdnfonts.com/css/technopollas" rel="stylesheet">
-    <style>
-      @import url('https://fonts.googleapis.com/css2?family=Inter:wght@400;600;700&display=swap');
-    </style>
-    <script>
-      tailwind.config = {
-        darkMode: 'class',
-        theme: {
-          extend: {
-            fontFamily: {
-              sans: ['Inter', 'sans-serif'],
-              technopollas: ['Technopollas', 'sans-serif'],
-            },
-            colors: {
-              dark: {
-                bg: '#121212',
-                card: '#1E1E1E',
-                text: '#FFFFFF',
-                accent: {
-                  green: '#4CAF50',
-                  red: '#F44336',
-                  orange: '#FF9800',
-                  yellow: '#FFEB3B',
-                },
-              },
-            },
-            borderRadius: {
-              'lg': '1rem',
-            },
-          }
-        }
-      }
-    </script>
-  </head>

agentic_security/static/telemetry.js

@@ -1,4 +0,0 @@
-!function (t, e) { var o, n, p, r; e.__SV || (window.posthog = e, e._i = [], e.init = function (i, s, a) { function g(t, e) { var o = e.split("."); 2 == o.length && (t = t[o[0]], e = o[1]), t[e] = function () { t.push([e].concat(Array.prototype.slice.call(arguments, 0))) } } (p = t.createElement("script")).type = "text/javascript", p.async = !0, p.src = s.api_host.replace(".i.posthog.com", "-assets.i.posthog.com") + "/static/array.js", (r = t.getElementsByTagName("script")[0]).parentNode.insertBefore(p, r); var u = e; for (void 0 !== a ? u = e[a] = [] : a = "posthog", u.people = u.people || [], u.toString = function (t) { var e = "posthog"; return "posthog" !== a && (e += "." + a), t || (e += " (stub)"), e }, u.people.toString = function () { return u.toString(1) + ".people (stub)" }, o = "init push capture register register_once register_for_session unregister unregister_for_session getFeatureFlag getFeatureFlagPayload isFeatureEnabled reloadFeatureFlags updateEarlyAccessFeatureEnrollment getEarlyAccessFeatures on onFeatureFlags onSessionId getSurveys getActiveMatchingSurveys renderSurvey canRenderSurvey getNextSurveyStep identify setPersonProperties group resetGroups setPersonPropertiesForFlags resetPersonPropertiesForFlags setGroupPropertiesForFlags resetGroupPropertiesForFlags reset get_distinct_id getGroups get_session_id get_session_replay_url alias set_config startSessionRecording stopSessionRecording sessionRecordingStarted loadToolbar get_property getSessionProperty createPersonProfile opt_in_capturing opt_out_capturing has_opted_in_capturing has_opted_out_capturing clear_opt_in_out_capturing debug".split(" "), n = 0; n < o.length; n++)g(u, o[n]); e._i.push([i, s, a]) }, e.__SV = 1) }(document, window.posthog || []);
-posthog.init('phc_jfYo5xEofW7eJtiU8rLt2Z8jw1E2eW27BxwTJzwRufH', {
-    api_host: 'https://us.i.posthog.com', person_profiles: 'identified_only' // or 'always' to create profiles for anonymous users as well
-})

agentic_security/static/telemetry_disabled.js

@@ -1 +0,0 @@
-console.log("Telemetry is disabled");

agentic_security/test_lib.py

@@ -14,6 +14,7 @@ Content-Type: application/json
 
 
 class TestAS:
+
     # Handles an empty dataset list.
     def test_class(self):
         llmSpec = SAMPLE_SPEC

agentic_security/test_spec.py

@@ -2,6 +2,7 @@ from agentic_security.http_spec import LLMSpec, parse_http_spec
 
 
 class TestParseHttpSpec:
+
     # Should correctly parse a simple HTTP spec with headers and body
     def test_parse_simple_http_spec(self):
         http_spec = (

pyproject.toml

@@ -1,6 +1,6 @@
 [tool.poetry]
 name = "agentic_security"
-version = "0.3.3"
+version = "0.1.7"
 description = "Agentic LLM vulnerability scanner"
 authors = ["Alexander Miasoiedov <msoedov@gmail.com>"]
 maintainers = ["Alexander Miasoiedov <msoedov@gmail.com>"]
@@ -25,33 +25,27 @@ packages = [{ include = "agentic_security", from = "." }]
 agentic_security = "agentic_security.__main__:entrypoint"
 
 [tool.poetry.dependencies]
-python = "^3.11"
-fastapi = "^0.115.2"
-uvicorn = "^0.32.0"
-fire = "0.7.0"
+python = "^3.9"
+fastapi = ">=0.109.1,<0.112.0"
+uvicorn = ">=0.23.2,<0.30.0"
+fire = ">=0.5,<0.7"
 loguru = "^0.7.2"
 httpx = ">=0.25.1,<0.28.0"
 cache-to-disk = "^2.0.0"
 pandas = ">=1.4,<3.0"
-datasets = ">=1.14,<4.0"
+datasets = "^1.14.0"
 tabulate = ">=0.8.9,<0.10.0"
 colorama = "^0.4.4"
-matplotlib = "^3.9.2"
-pydantic = "2.9.2"
-scikit-optimize = "^0.10.2"
-scikit-learn = "1.5.2"
-numpy = ">=1.24.3,<3.0.0"
-jinja2 = "^3.1.4"
+matplotlib = "^3.4.3"
 
 [tool.poetry.group.dev.dependencies]
-black = "^24.10.0"
-mypy = "^1.12.0"
-pytest = "^8.3.3"
-pre-commit = "^4.0.1"
-inline-snapshot = "^0.13.3"
-langchain-groq = "^0.2.0"
-huggingface-hub = "^0.25.1"
-# garak = "*"
+black = ">=23.10.1,<25.0.0"
+mypy = "^1.6.1"
+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"
 
 [tool.ruff]
 line-length = 120