Misc

aisploit/converter/__init__.py

@@ -1,4 +1,5 @@
 from .base64 import Base64Converter
+from .gender import GenderConverter
 from .join import JoinConverter
 from .keyboard_typo import (
     KeyboardTypoConverter,
@@ -10,6 +11,7 @@
 
 __all__ = [
     "Base64Converter",
+    "GenderConverter",
     "JoinConverter",
     "KeyboardTypoConverter",
     "KEYBOARD_NEIGHBORS_QWERTY",

aisploit/converter/gender.py

@@ -0,0 +1,30 @@
+import textwrap
+from langchain_core.prompts import ChatPromptTemplate
+from langchain_core.output_parsers import StrOutputParser
+from ..core import BaseConverter, BaseChatModel
+
+_template = ChatPromptTemplate.from_template(
+    textwrap.dedent(
+        """
+        Swap the gender in the given text.
+        Input text:
+        {input}
+
+        Instructions:
+        1. Identify gender-specific terms such as pronouns (e.g., he, she), names, titles, and other gendered language in the input text.
+        2. Replace the identified gender-specific terms with their gender-swapped counterparts. For example:
+            - Replace "he" with "she" and vice versa.
+            - Replace male names with female names and vice versa.
+            - Adjust any associated verbs, adjectives, or other language to maintain grammatical correctness.
+        3. Ensure that the gender-swapped text retains the original context and meaning as much as possible.
+        """
+    )
+)
+
+
+class GenderConverter(BaseConverter):
+    def __init__(self, *, chat_model: BaseChatModel, prompt=_template) -> None:
+        self._chain = prompt | chat_model | StrOutputParser()
+
+    def _convert(self, prompt: str) -> str:
+        return self._chain.invoke({"input": prompt})

aisploit/converter/join.py

@@ -13,8 +13,3 @@ def _convert(self, prompt: str) -> str:
         words = prompt.split()
         joined_words = [self.separator.join(word) for word in words]
         return " ".join(joined_words)
-
-    def __repr__(self) -> str:
-        return (
-            f"<{self.__module__}.{self.__class__.__name__}(separator={self.separator})>"
-        )

aisploit/core/__init__.py

@@ -1,6 +1,7 @@
 from .callbacks import BaseCallbackHandler, Callbacks, CallbackManager
 from .classifier import BaseClassifier, Score
 from .converter import BaseConverter
+from .dataset import BaseDataset, YamlDeserializable
 from .job import BaseJob
 from .model import BaseLLM, BaseChatModel, BaseModel, BaseEmbeddings
 from .prompt import BasePromptValue
@@ -15,7 +16,9 @@
     "BaseClassifier",
     "Score",
     "BaseConverter",
-    "BaseJob",
+    "BaseDataset",
+    "YamlDeserializable",
+    "Dataset" "BaseJob",
     "BaseLLM",
     "BaseChatModel",
     "BaseModel",

aisploit/core/classifier.py

@@ -8,14 +8,25 @@
 
 @dataclass(frozen=True)
 class Score(Generic[T]):
+    """A class representing a score."""
+
     flagged: bool
     value: T
     description: str = ""
     explanation: str = ""
 
 
 class BaseClassifier(ABC, Generic[T]):
+    """An abstract base class for classifiers."""
+
     @abstractmethod
     def score_text(self, text: str) -> Score[T]:
-        """Score the text and return a Score object."""
+        """Score the text and return a Score object.
+
+        Args:
+            text (str): The text to be scored.
+
+        Returns:
+            Score[T]: A Score object representing the score of the text.
+        """
         pass

aisploit/core/converter.py

@@ -1,3 +1,4 @@
+from typing import Union
 from abc import ABC, abstractmethod
 from langchain_core.prompt_values import StringPromptValue
 from .prompt import BasePromptValue
@@ -18,7 +19,7 @@ def _convert(self, prompt: str) -> str:
         """
         pass
 
-    def convert(self, prompt: BasePromptValue) -> BasePromptValue:
+    def convert(self, prompt: Union[str, BasePromptValue]) -> BasePromptValue:
         """Converts the prompt value.
 
         Args:
@@ -27,8 +28,11 @@ def convert(self, prompt: BasePromptValue) -> BasePromptValue:
         Returns:
             BasePromptValue: The converted prompt value.
         """
+        if isinstance(prompt, str):
+            return StringPromptValue(text=self._convert(prompt))
+
         if isinstance(prompt, StringPromptValue):
-            prompt = StringPromptValue(text=self._convert(prompt.text))
+            return StringPromptValue(text=self._convert(prompt.text))
 
         return prompt
 

aisploit/dataset/dataset.py → aisploit/core/dataset.py

@@ -6,7 +6,7 @@
 T = TypeVar("T")
 
 
-class Dataset(Generic[T]):
+class BaseDataset(Generic[T]):
     """Generic dataset class."""
 
     _prompts: Sequence[T]

aisploit/core/generator.py

@@ -0,0 +1,8 @@
+from abc import ABC, abstractmethod
+from .dataset import BaseDataset
+
+
+class BaseGenerator(ABC):
+    @abstractmethod
+    def generate_dataset(self) -> BaseDataset:
+        pass

aisploit/core/report.py

@@ -1,4 +1,5 @@
 from typing import Any, Generic, TypeVar, List
+from datetime import datetime
 from abc import ABC, abstractmethod
 from pathlib import Path
 from jinja2 import Template
@@ -8,12 +9,24 @@
 
 
 class BaseReport(Generic[T], ABC):
+    """A generic base class for reports."""
+
     _entries: List[T]
 
     def __init__(self, *, run_id: str) -> None:
+        """Initialize the BaseReport.
+
+        Args:
+            run_id (str): The ID of the report run.
+        """
         self.run_id = run_id
+        self._created_at = datetime.now()
         self._entries = []
 
+    @property
+    def created_at(self) -> datetime:
+        return self._created_at
+
     def __iter__(self):
         return iter(self._entries)
 
@@ -22,9 +35,19 @@ def __len__(self):
 
     @abstractmethod
     def _ipython_display_(self):
+        """Display the report in IPython."""
         pass
 
     def _render_template(self, template_path: Path, **kwargs: Any) -> str:
+        """Render the report template.
+
+        Args:
+            template_path (Path): The path to the template file.
+            **kwargs (Any): Additional keyword arguments to pass to the template.
+
+        Returns:
+            str: The rendered report as a string.
+        """
         with open(template_path, "r", encoding="utf8") as tpl_file:
             template = Template(tpl_file.read())
             return template.render(**kwargs)

aisploit/dataset/__init__.py

@@ -1,9 +1,7 @@
-from .dataset import Dataset
 from .prompt import PromptDataset, Prompt, JailbreakPromptDataset
 from .sample import SampleDataset, Sample
 
 __all__ = [
-    "Dataset",
     "PromptDataset",
     "Prompt",
     "JailbreakPromptDataset",

aisploit/dataset/prompt.py

@@ -3,7 +3,7 @@
 from typing import Sequence, Optional
 from dataclasses import dataclass
 
-from .dataset import Dataset, YamlDeserializable
+from ..core.dataset import BaseDataset, YamlDeserializable
 
 
 @dataclass
@@ -19,7 +19,7 @@ class Prompt(YamlDeserializable):
     template: str
 
 
-class PromptDataset(Dataset[Prompt]):
+class PromptDataset(BaseDataset[Prompt]):
     """Dataset for prompts."""
 
     def __init__(self, prompts: Sequence[Prompt]) -> None:

aisploit/dataset/sample.py

@@ -3,7 +3,7 @@
 from pathlib import Path
 from dataclasses import dataclass
 
-from .dataset import Dataset, YamlDeserializable
+from ..core.dataset import BaseDataset, YamlDeserializable
 
 
 @dataclass
@@ -17,7 +17,7 @@ class Sample(YamlDeserializable):
     tags: Sequence[str]
 
 
-class SampleDataset(Dataset[Sample]):
+class SampleDataset(BaseDataset[Sample]):
     """Dataset for samples."""
 
     def __init__(self, samples: Sequence[Sample]) -> None:

aisploit/scanner/report.py

@@ -45,7 +45,7 @@ def to_markdown(self, *, template_path=TEMPLATES_PATH / "report.md") -> str:
         return self._render_template(
             template_path=template_path,
             run_id=self.run_id,
-            report=self,
+            created_at=self.created_at,
             issues_by_category=issues_by_category,
         )
 

aisploit/scanner/templates/report.md

@@ -1,5 +1,7 @@
-## ScanReport 
-> RunID: {{ run_id }}
+## ScanReport
+| RunID | CreatedAt |
+|-------|-----------| 
+|{{ run_id }}|{{ created_at }}|
 
 {% if issues_by_category|length == 0 %}
 No issues!
@@ -8,10 +10,10 @@ No issues!
 ### {{ category.name }} issues
 > {{ category.description }}
 
-| Prompt | Converter | Response |
-|--------|-----------|----------|
+| Prompt | Converter | Response | RTT (seconds) |
+|--------|-----------|----------|---------------|
 {% for issue in issues -%}
-|{{ issue.send_report_entry.prompt }}|{{ issue.send_report_entry.converter }}|{{ issue.send_report_entry.response }}|
+|{{ issue.send_report_entry.prompt }}|{{ issue.send_report_entry.converter }}|{{ issue.send_report_entry.response }}| {{ issue.send_report_entry.round_trip_time }} |
 {% endfor %}
 {% endfor %}
 {% endif %}

aisploit/sender/report.py

@@ -1,4 +1,4 @@
-from typing import List, Optional
+from typing import Optional
 from dataclasses import dataclass
 from datetime import datetime
 
@@ -13,6 +13,10 @@ class SendReportEntry:
     start_time: datetime
     end_time: datetime
 
+    @property
+    def round_trip_time(self) -> float:
+        return (self.end_time - self.start_time).total_seconds()
+
 
 class SendReport(BaseReport[SendReportEntry]):
     def __init__(self, *, run_id: str) -> None:

examples/classifier.ipynb

@@ -1,5 +1,12 @@
 {
  "cells": [
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "# AISploit Classifier"
+   ]
+  },
   {
    "cell_type": "code",
    "execution_count": 1,