Skip to content

Commit

Permalink
fix(fmt):
Browse files Browse the repository at this point in the history
  • Loading branch information
msoedov committed Nov 29, 2024
1 parent 107181f commit f57f3e9
Show file tree
Hide file tree
Showing 21 changed files with 41 additions and 116 deletions.
31 changes: 13 additions & 18 deletions .pre-commit-config.yaml
Original file line number Diff line number Diff line change
@@ -1,26 +1,24 @@

default_language_version:
python: python3
python: python3.11

repos:

- repo: https://github.com/asottile/pyupgrade
rev: v2.31.1
rev: v3.15.0
hooks:
- id: pyupgrade
args: [--py39-plus]
args: [--py311-plus]

- repo: https://github.com/psf/black
rev: 22.8.0
rev: 23.11.0
hooks:
- id: black
language_version: python3.9
language_version: python3.11

- repo: https://github.com/pycqa/flake8
rev: 5.0.4
rev: 6.1.0
hooks:
- id: flake8
language_version: python3
language_version: python3.11
additional_dependencies: [flake8-docstrings]

- repo: https://github.com/PyCQA/isort
Expand All @@ -30,7 +28,7 @@ repos:
args: [--profile, black]

- repo: https://github.com/pre-commit/pre-commit-hooks
rev: v4.3.0
rev: v4.5.0
hooks:
- id: check-ast
exclude: '^(third_party)/'
Expand All @@ -47,16 +45,15 @@ repos:
args: ['--maxkb=100']

- repo: https://github.com/executablebooks/mdformat
rev: 0.7.14
rev: 0.7.17
hooks:
- id: mdformat
name: mdformat
entry: mdformat .
language_version: python3

language_version: python3.11

- repo: https://github.com/hadialqattan/pycln
rev: v2.1.1 # Possible releases: https://github.com/hadialqattan/pycln/releases
rev: v2.4.0
hooks:
- id: pycln

Expand All @@ -65,19 +62,17 @@ repos:
hooks:
- id: teyit


- repo: https://github.com/python-poetry/poetry
rev: '1.6.0'
rev: '1.7.0'
hooks:
- id: poetry-check
- id: poetry-lock
name: validate poetry lock
args:
- --check


- repo: https://github.com/codespell-project/codespell
rev: v2.2.5
rev: v2.2.6
hooks:
- id: codespell
exclude: '^(third_party/)|(poetry.lock)'
Expand Down
4 changes: 2 additions & 2 deletions agentic_security/app.py
Original file line number Diff line number Diff line change
Expand Up @@ -3,11 +3,11 @@
from .middleware.cors import setup_cors
from .middleware.logging import LogNon200ResponsesMiddleware
from .routes import (
static_router,
scan_router,
probe_router,
proxy_router,
report_router,
scan_router,
static_router,
)

# Create the FastAPI app
Expand Down
1 change: 1 addition & 0 deletions agentic_security/core/app.py
Original file line number Diff line number Diff line change
@@ -1,4 +1,5 @@
from asyncio import Event, Queue

from fastapi import FastAPI

tools_inbox: Queue = Queue()
Expand Down
2 changes: 1 addition & 1 deletion agentic_security/lib.py
Original file line number Diff line number Diff line change
Expand Up @@ -6,8 +6,8 @@
from tabulate import tabulate

from agentic_security.models.schemas import Scan
from agentic_security.routes.scan import 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
Expand Down
2 changes: 1 addition & 1 deletion agentic_security/middleware/logging.py
Original file line number Diff line number Diff line change
@@ -1,6 +1,6 @@
from fastapi import Request
from starlette.middleware.base import BaseHTTPMiddleware
from loguru import logger
from starlette.middleware.base import BaseHTTPMiddleware


class LogNon200ResponsesMiddleware(BaseHTTPMiddleware):
Expand Down
1 change: 1 addition & 0 deletions agentic_security/models/schemas.py
Original file line number Diff line number Diff line change
@@ -1,4 +1,5 @@
import os

from pydantic import BaseModel, Field


Expand Down
5 changes: 3 additions & 2 deletions agentic_security/probe_actor/fuzzer.py
Original file line number Diff line number Diff line change
@@ -1,16 +1,17 @@
import asyncio
import os
import random
from typing import AsyncGenerator
from collections.abc import AsyncGenerator

import httpx
import pandas as pd
from loguru import logger
from skopt import Optimizer
from skopt.space import Real

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
from agentic_security.models.schemas import ScanResult

IS_VERCEL = os.getenv("IS_VERCEL", "f") == "t"

Expand Down
1 change: 0 additions & 1 deletion agentic_security/probe_actor/test_refusal.py
Original file line number Diff line number Diff line change
Expand Up @@ -2,7 +2,6 @@


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."
Expand Down
3 changes: 0 additions & 3 deletions agentic_security/probe_data/data.py
Original file line number Diff line number Diff line change
Expand Up @@ -152,7 +152,6 @@ 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",
Expand All @@ -173,7 +172,6 @@ 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"],
Expand Down Expand Up @@ -321,7 +319,6 @@ def __iter__(self):

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}")
Expand Down
1 change: 0 additions & 1 deletion agentic_security/probe_data/modules/adaptive_attacks.py
Original file line number Diff line number Diff line change
Expand Up @@ -9,7 +9,6 @@

class Module:
def __init__(self, prompt_groups: []):

r = httpx.get(url)

content = r.content
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -4,7 +4,6 @@


class TestModule:

# Module can be initialized with a list of prompt groups.
def test_initialize_with_prompt_groups(self):
prompt_groups = []
Expand Down
4 changes: 2 additions & 2 deletions agentic_security/routes/__init__.py
Original file line number Diff line number Diff line change
@@ -1,8 +1,8 @@
from .static import router as static_router
from .scan import router as scan_router
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",
Expand Down
2 changes: 2 additions & 0 deletions agentic_security/routes/probe.py
Original file line number Diff line number Diff line change
@@ -1,5 +1,7 @@
import random

from fastapi import APIRouter

from ..models.schemas import Probe
from ..probe_actor.refusal import REFUSAL_MARKS
from ..probe_data import REGISTRY
Expand Down
4 changes: 3 additions & 1 deletion agentic_security/routes/proxy.py
Original file line number Diff line number Diff line change
@@ -1,9 +1,11 @@
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
from ..core.app import get_tools_inbox

router = APIRouter()

Expand Down
2 changes: 2 additions & 0 deletions agentic_security/routes/report.py
Original file line number Diff line number Diff line change
@@ -1,6 +1,8 @@
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

Expand Down
6 changes: 4 additions & 2 deletions agentic_security/routes/scan.py
Original file line number Diff line number Diff line change
@@ -1,10 +1,12 @@
from datetime import datetime

from fastapi import APIRouter, BackgroundTasks, HTTPException
from fastapi.responses import StreamingResponse
from ..models.schemas import LLMInfo, Scan

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
from ..core.app import get_tools_inbox, get_stop_event

router = APIRouter()

Expand Down
2 changes: 2 additions & 0 deletions agentic_security/routes/static.py
Original file line number Diff line number Diff line change
@@ -1,6 +1,8 @@
from pathlib import Path

from fastapi import APIRouter
from fastapi.responses import FileResponse

from ..models.schemas import Settings

router = APIRouter()
Expand Down
1 change: 0 additions & 1 deletion agentic_security/test_lib.py
Original file line number Diff line number Diff line change
Expand Up @@ -14,7 +14,6 @@


class TestAS:

# Handles an empty dataset list.
def test_class(self):
llmSpec = SAMPLE_SPEC
Expand Down
1 change: 0 additions & 1 deletion agentic_security/test_spec.py
Original file line number Diff line number Diff line change
Expand Up @@ -2,7 +2,6 @@


class TestParseHttpSpec:

# Should correctly parse a simple HTTP spec with headers and body
def test_parse_simple_http_spec(self):
http_spec = (
Expand Down
Loading

0 comments on commit f57f3e9

Please sign in to comment.