Add ruff

.gitignore

@@ -161,5 +161,6 @@ cython_debug/
 
 .env
 .DS_Store
+.ruff_cache
 TODO.md
 !aisploit/target

aisploit/classifiers/huggingface/pipeline_prompt_injection_identifier.py

@@ -32,17 +32,9 @@ def __init__(
     def score(self, input: str) -> Score[float]:
         result = self._model(input)
 
-        score = (
-            result[0]["score"]
-            if result[0]["label"] == self._injection_label
-            else 1 - result[0]["score"]
-        )
+        score = result[0]["score"] if result[0]["label"] == self._injection_label else 1 - result[0]["score"]
 
-        explanation = (
-            "Prompt injection attack detected"
-            if score > self._threshold
-            else "No prompt injection"
-        )
+        explanation = "Prompt injection attack detected" if score > self._threshold else "No prompt injection"
 
         return Score[float](
             flagged=score > self._threshold,

aisploit/classifiers/openai/moderation.py

@@ -1,5 +1,6 @@
-from typing import Optional
 import os
+from typing import Optional
+
 from openai import OpenAI
 from openai.types.moderation import Moderation
 

aisploit/classifiers/text.py

@@ -1,4 +1,5 @@
 import re
+
 from ..core import BaseTextClassifier, Score
 
 
@@ -51,9 +52,5 @@ def __init__(self, *, substring: str, ignore_case=True, flag_matches=True) -> No
             ignore_case (bool, optional): Flag indicating whether to ignore case when matching substrings. Defaults to True.
             flag_matches (bool, optional): Flag indicating whether matches should be flagged. Defaults to True.
         """
-        compiled_pattern = (
-            re.compile(substring, re.IGNORECASE)
-            if ignore_case
-            else re.compile(substring)
-        )
+        compiled_pattern = re.compile(substring, re.IGNORECASE) if ignore_case else re.compile(substring)
         super().__init__(pattern=compiled_pattern, flag_matches=flag_matches)

aisploit/converters/__init__.py

@@ -1,11 +1,11 @@
 from .base64 import Base64Converter
-from .case import LowercaseConverter, UppercaseConverter, TitlecaseConverter
+from .case import LowercaseConverter, TitlecaseConverter, UppercaseConverter
 from .gender import GenderConverter
 from .join import JoinConverter
 from .keyboard_typo import (
-    KeyboardTypoConverter,
     KEYBOARD_NEIGHBORS_QWERTY,
     KEYBOARD_NEIGHBORS_QWERTZ,
+    KeyboardTypoConverter,
 )
 from .no_op import NoOpConverter
 from .remove_punctuation import RemovePunctuationConverter

aisploit/converters/gender.py

@@ -1,7 +1,9 @@
 import textwrap
-from langchain_core.prompts import ChatPromptTemplate
+
 from langchain_core.output_parsers import StrOutputParser
-from ..core import BaseChatModelConverter, BaseChatModel
+from langchain_core.prompts import ChatPromptTemplate
+
+from ..core import BaseChatModel, BaseChatModelConverter
 
 _template = ChatPromptTemplate.from_template(
     textwrap.dedent(

aisploit/converters/keyboard_typo.py

@@ -76,10 +76,7 @@ def __init__(
     def _convert(self, prompt: str) -> str:
         typoPrompt = ""
         for char in prompt:
-            if (
-                random.random() < self._typo_probability
-                and char.lower() in self._keyboard_neighbors
-            ):
+            if random.random() < self._typo_probability and char.lower() in self._keyboard_neighbors:
                 # Replace the character with a random neighboring key
                 neighbor_keys = self._keyboard_neighbors[char.lower()]
                 typo_char = random.choice(neighbor_keys)

aisploit/converters/remove_punctuation.py

@@ -1,4 +1,5 @@
 import string
+
 from ..core import BaseConverter
 
 

aisploit/converters/sequence.py

@@ -1,4 +1,5 @@
 from typing import Sequence
+
 from langchain_core.prompt_values import StringPromptValue
 
 from ..core import BaseConverter
@@ -11,8 +12,6 @@ def __init__(self, *, converters: Sequence[BaseConverter] = []) -> None:
     def _convert(self, prompt: str) -> str:
         converted_prompt = prompt
         for converter in self._converters:
-            converted_prompt = converter.convert(
-                StringPromptValue(text=converted_prompt)
-            ).to_string()
+            converted_prompt = converter.convert(StringPromptValue(text=converted_prompt)).to_string()
 
         return converted_prompt

aisploit/converters/stemming.py

@@ -1,6 +1,7 @@
 import nltk
-from nltk.tokenize import word_tokenize
 from nltk.stem.porter import PorterStemmer
+from nltk.tokenize import word_tokenize
+
 from ..core import BaseConverter
 
 

aisploit/converters/unicode_confusable.py

@@ -1,5 +1,7 @@
 import random
+
 from confusables import confusable_characters
+
 from ..core import BaseConverter
 
 

aisploit/core/__init__.py

@@ -1,10 +1,10 @@
-from .callbacks import BaseCallbackHandler, Callbacks, CallbackManager
+from .callbacks import BaseCallbackHandler, CallbackManager, Callbacks
 from .classifier import BaseClassifier, BaseTextClassifier, Score
-from .converter import BaseConverter, BaseChatModelConverter
+from .converter import BaseChatModelConverter, BaseConverter
 from .dataset import BaseDataset, YamlDeserializable
 from .generator import BaseGenerator
 from .job import BaseJob
-from .model import BaseLLM, BaseChatModel, BaseModel, BaseEmbeddings
+from .model import BaseChatModel, BaseEmbeddings, BaseLLM, BaseModel
 from .prompt import BasePromptValue
 from .report import BaseReport
 from .target import BaseTarget, Response

aisploit/core/callbacks.py

@@ -1,16 +1,14 @@
 from typing import Sequence
 
-from .prompt import BasePromptValue
 from .classifier import Score
+from .prompt import BasePromptValue
 from .target import Response
 
 
 class BaseCallbackHandler:
     """Base class for callback handlers."""
 
-    def on_redteam_attempt_start(
-        self, attempt: int, prompt: BasePromptValue, *, run_id: str
-    ):
+    def on_redteam_attempt_start(self, attempt: int, prompt: BasePromptValue, *, run_id: str):
         """Called when a red team attempt starts.
 
         Args:
@@ -20,9 +18,7 @@ def on_redteam_attempt_start(
         """
         pass
 
-    def on_redteam_attempt_end(
-        self, attempt: int, response: Response, score: Score, *, run_id: str
-    ):
+    def on_redteam_attempt_end(self, attempt: int, response: Response, score: Score, *, run_id: str):
         """Called when a red team attempt ends.
 
         Args:
@@ -68,7 +64,7 @@ def __init__(
         self,
         *,
         run_id: str,
-        callbacks: Sequence[BaseCallbackHandler] = [],
+        callbacks: Sequence[BaseCallbackHandler] | None,
     ) -> None:
         """Initialize the CallbackManager.
 
@@ -77,7 +73,7 @@ def __init__(
             callbacks (Sequence[BaseCallbackHandler], optional): The list of callback handlers. Defaults to [].
         """
         self.run_id = run_id
-        self._callbacks = callbacks
+        self._callbacks = callbacks or []
 
     def on_redteam_attempt_start(self, attempt: int, prompt: BasePromptValue):
         """Notify callback handlers when a red team attempt starts.
@@ -87,9 +83,7 @@ def on_redteam_attempt_start(self, attempt: int, prompt: BasePromptValue):
             prompt (BasePromptValue): The prompt value.
         """
         for cb in self._callbacks:
-            cb.on_redteam_attempt_start(
-                attempt=attempt, prompt=prompt, run_id=self.run_id
-            )
+            cb.on_redteam_attempt_start(attempt=attempt, prompt=prompt, run_id=self.run_id)
 
     def on_redteam_attempt_end(self, attempt: int, response: Response, score: Score):
         """Notify callback handlers when a red team attempt ends.
@@ -100,9 +94,7 @@ def on_redteam_attempt_end(self, attempt: int, response: Response, score: Score)
             score (Score): The score of the attempt.
         """
         for cb in self._callbacks:
-            cb.on_redteam_attempt_end(
-                attempt=attempt, response=response, score=score, run_id=self.run_id
-            )
+            cb.on_redteam_attempt_end(attempt=attempt, response=response, score=score, run_id=self.run_id)
 
     def on_scanner_plugin_start(self, name: str):
         """Notify callback handlers when a scanner plugin starts.

aisploit/core/classifier.py

@@ -1,6 +1,6 @@
-from typing import TypeVar, Generic
 from abc import ABC, abstractmethod
 from dataclasses import dataclass
+from typing import Generic, TypeVar
 
 T = TypeVar("T")
 Input = TypeVar("Input")

aisploit/core/converter.py

@@ -1,6 +1,8 @@
-from typing import Union
 from abc import ABC, abstractmethod
+from typing import Union
+
 from langchain_core.prompt_values import StringPromptValue
+
 from .model import BaseChatModel
 from .prompt import BasePromptValue
 

aisploit/core/dataset.py

@@ -1,7 +1,7 @@
-from abc import ABC
-import yaml
 from pathlib import Path
-from typing import Generic, Type, TypeVar, Sequence
+from typing import Generic, Sequence, Type, TypeVar
+
+import yaml
 
 T = TypeVar("T")
 
@@ -18,8 +18,8 @@ def __len__(self):
         return len(self._entries)
 
 
-class YamlDeserializable(ABC):
-    """Abstract base class for objects that can be deserialized from YAML."""
+class YamlDeserializable:
+    """Base class for objects that can be deserialized from YAML."""
 
     @classmethod
     def from_yaml_file(cls: Type[T], file: Path) -> T:
@@ -44,6 +44,6 @@ def from_yaml_file(cls: Type[T], file: Path) -> T:
             try:
                 yaml_data = yaml.safe_load(f)
             except yaml.YAMLError as exc:
-                raise ValueError(f"Invalid YAML file '{file}': {exc}")
+                raise ValueError(f"Invalid YAML file '{file}'") from exc
 
         return cls(**yaml_data)

aisploit/core/generator.py

@@ -1,5 +1,6 @@
-from typing import Generic, TypeVar
 from abc import ABC, abstractmethod
+from typing import Generic, TypeVar
+
 from .dataset import BaseDataset
 
 T = TypeVar("T")

aisploit/core/job.py

@@ -1,8 +1,7 @@
-from abc import ABC
 from uuid import uuid4
 
 
-class BaseJob(ABC):
+class BaseJob:
     def __init__(self, *, verbose=False) -> None:
         self.verbose = verbose
 

aisploit/core/model.py

@@ -1,9 +1,10 @@
 from abc import abstractmethod
 from typing import Union
+
+from langchain_core.embeddings import Embeddings
 from langchain_core.language_models import LanguageModelInput
 from langchain_core.messages import BaseMessage
 from langchain_core.runnables import Runnable
-from langchain_core.embeddings import Embeddings
 
 
 class BaseLLM(Runnable[LanguageModelInput, str]):

aisploit/core/report.py

@@ -1,9 +1,9 @@
-from typing import Any, Generic, TypeVar, List
-from datetime import datetime
 from abc import ABC, abstractmethod
+from datetime import datetime
 from pathlib import Path
-from jinja2 import Template
+from typing import Any, Generic, List, TypeVar
 
+from jinja2 import Template
 
 T = TypeVar("T")
 

aisploit/core/target.py

@@ -1,6 +1,6 @@
-from typing import Dict, Any
 from abc import ABC, abstractmethod
 from dataclasses import dataclass, field
+from typing import Any, Dict
 
 from .prompt import BasePromptValue
 

aisploit/datasets/__init__.py

@@ -1,5 +1,5 @@
-from .prompt import PromptDataset, Prompt, JailbreakPromptDataset
-from .sample import SampleDataset, Sample
+from .prompt import JailbreakPromptDataset, Prompt, PromptDataset
+from .sample import Sample, SampleDataset
 
 __all__ = [
     "PromptDataset",

aisploit/datasets/prompt.py

@@ -1,7 +1,7 @@
 import os
-from pathlib import Path
-from typing import Sequence, Optional
 from dataclasses import dataclass
+from pathlib import Path
+from typing import Optional, Sequence
 
 from ..core.dataset import BaseDataset, YamlDeserializable
 
@@ -31,9 +31,7 @@ def __init__(self, prompts: Sequence[Prompt]) -> None:
         self._entries = prompts
 
     @classmethod
-    def load_from_directory(
-        cls, path: Path, tags_filter: Optional[Sequence[str]] = None
-    ) -> "PromptDataset":
+    def load_from_directory(cls, path: Path, tags_filter: Optional[Sequence[str]] = None) -> "PromptDataset":
         """Create a JailbreakDataset instance by loading prompts from a directory.
 
         Args:
@@ -46,9 +44,7 @@ def load_from_directory(
         prompts = []
         for file_name in os.listdir(path):
             prompt = Prompt.from_yaml_file(path / file_name)
-            if not prompt.skip and (
-                not tags_filter or set(prompt.tags).intersection(tags_filter)
-            ):
+            if not prompt.skip and (not tags_filter or set(prompt.tags).intersection(tags_filter)):
                 prompts.append(prompt)
         return cls(prompts)
 

aisploit/datasets/sample.py

@@ -1,7 +1,7 @@
 import os
-from typing import Sequence, Optional
-from pathlib import Path
 from dataclasses import dataclass
+from pathlib import Path
+from typing import Optional, Sequence
 
 from ..core.dataset import BaseDataset, YamlDeserializable
 
@@ -29,9 +29,7 @@ def __init__(self, samples: Sequence[Sample]) -> None:
         self._entries = samples
 
     @classmethod
-    def load_from_directory(
-        cls, path: Path, tags_filter: Optional[Sequence[str]] = None
-    ) -> "SampleDataset":
+    def load_from_directory(cls, path: Path, tags_filter: Optional[Sequence[str]] = None) -> "SampleDataset":
         """Create a SampleDataset instance by loading samples from a directory.
 
         Args:
@@ -44,8 +42,6 @@ def load_from_directory(
         samples = []
         for file_name in os.listdir(path):
             sample = Sample.from_yaml_file(path / file_name)
-            if not sample.skip and (
-                not tags_filter or set(sample.tags).intersection(tags_filter)
-            ):
+            if not sample.skip and (not tags_filter or set(sample.tags).intersection(tags_filter)):
                 samples.append(sample)
         return cls(samples)

aisploit/demo/__init__.py

@@ -1,4 +1,4 @@
-from .gandalf import GandalfBot, GandalfScorer, GandalfLevel
+from .gandalf import GandalfBot, GandalfLevel, GandalfScorer
 from .rag import VectorStoreRAG
 
 __all__ = [