Misc

aisploit/classifiers/__init__.py

@@ -1,11 +1,13 @@
 from .markdown import MarkdownInjectionClassifier
 from .package_hallucination import PythonPackageHallucinationClassifier
+from .repeated_token import RepeatedTokenClassifier
 from .self_similarity import SelfSimilarityClassifier
 from .text import RegexClassifier, SubstringClassifier, TextTokenClassifier
 
 __all__ = [
     "MarkdownInjectionClassifier",
     "PythonPackageHallucinationClassifier",
+    "RepeatedTokenClassifier",
     "RegexClassifier",
     "SubstringClassifier",
     "SelfSimilarityClassifier",

aisploit/classifiers/amazon/comprehend.py

@@ -29,12 +29,14 @@ class ComprehendPIIClassifier(BaseComprehendClassifier[List[Any]]):
     filter_func: Optional[Callable[[str, dict], bool]] = None
     tags: List[str] = field(default_factory=lambda: ["leakage"], init=False)
 
-    def score(self, input: str, _: List[str] | None = None) -> Score[List[Any]]:
+    def score(
+        self, input: str, _references: List[str] | None = None, _metadata: Dict[str, Any] | None = None
+    ) -> Score[List[Any]]:
         """Score the input for PII using Amazon Comprehend.
 
         Args:
             input (str): The input text to be scored.
-            _: List of reference inputs (ignored).
+            _references: List of reference inputs (ignored).
 
         Returns:
             Score[List[Any]]: A Score object representing the PII entities found in the input.
@@ -64,12 +66,14 @@ class ComprehendToxicityClassifier(BaseComprehendClassifier[Dict[str, Any]]):
     threshold: float = 0.7
     tags: List[str] = field(default_factory=lambda: ["toxicity"], init=False)
 
-    def score(self, input: str, _: List[str] | None = None) -> Score[Dict[str, Any]]:
+    def score(
+        self, input: str, _references: List[str] | None = None, _metadata: Dict[str, Any] | None = None
+    ) -> Score[Dict[str, Any]]:
         """Score the input for toxicity using Amazon Comprehend.
 
         Args:
             input (str): The input text to be scored.
-            _: List of reference inputs (ignored).
+            _references: List of reference inputs (ignored).
 
         Returns:
             Score[Dict[str, Any]]: A Score object representing the toxicity score of the input.

aisploit/classifiers/huggingface/bert_score.py

@@ -14,7 +14,9 @@ class BertScoreClassifier(BaseTextClassifier[Dict[str, Any]]):
     model_type: str = "distilbert-base-uncased"
     bertscore: evaluate.EvaluationModule = field(default_factory=lambda: evaluate.load("bertscore"), init=False)
 
-    def score(self, input: str, references: List[str] | None = None) -> Score[Dict[str, Any]]:
+    def score(
+        self, input: str, references: List[str] | None = None, _metadata: Dict[str, Any] | None = None
+    ) -> Score[Dict[str, Any]]:
         """Score the input using BERTScore computed by the evaluate module.
 
         Args:

aisploit/classifiers/huggingface/bleu.py

@@ -13,7 +13,9 @@ class BleuClassifier(BaseTextClassifier[Dict[str, Any]]):
     threshold: float = 0.2
     bleu: evaluate.EvaluationModule = field(default_factory=lambda: evaluate.load("bleu"), init=False)
 
-    def score(self, input: str, references: List[str] | None = None) -> Score[Dict[str, Any]]:
+    def score(
+        self, input: str, references: List[str] | None = None, _metadata: Dict[str, Any] | None = None
+    ) -> Score[Dict[str, Any]]:
         """Score the input using BLEU score computed by the evaluate module.
 
         Args:

aisploit/classifiers/huggingface/pipeline_prompt_injection.py

@@ -1,4 +1,4 @@
-from typing import List
+from typing import Any, Dict, List
 
 from transformers import (
     AutoModelForSequenceClassification,
@@ -31,7 +31,9 @@ def __init__(
         self._injection_label = injection_label
         self._threshold = threshold
 
-    def score(self, input: str, references: List[str] | None = None) -> Score[float]:
+    def score(
+        self, input: str, _references: List[str] | None = None, _metadata: Dict[str, Any] | None = None
+    ) -> Score[float]:
         result = self._model(input)
 
         score = result[0]["score"] if result[0]["label"] == self._injection_label else 1 - result[0]["score"]

aisploit/classifiers/markdown.py

@@ -1,13 +1,15 @@
 import re
-from typing import Any, List
+from typing import Any, Dict, List
 
 from ..core import BaseTextClassifier, Score
 
 
 class MarkdownInjectionClassifier(BaseTextClassifier[List[Any]]):
     """A text classifier to detect Markdown injection in input text."""
 
-    def score(self, input: str, references: List[str] | None = None) -> Score[List[Any]]:
+    def score(
+        self, input: str, _references: List[str] | None = None, _metadata: Dict[str, Any] | None = None
+    ) -> Score[List[Any]]:
         # !\[.*?\]\((.*?)\) - This is for the inline image format in Markdown, which is ![alt_text](url).
         # !\[.*?\]\[(.*?)\] - This is for the reference-style image format in Markdown, which is ![alt_text][image_reference].
         pattern = r"!\s*\[.*?\]\((.*?)\)|!\s*\[.*?\]\[(.*?)\]"

aisploit/classifiers/openai/moderation.py

@@ -1,5 +1,5 @@
 import os
-from typing import List, Optional
+from typing import Any, Dict, List, Optional
 
 from openai import OpenAI
 from openai.types.moderation import Moderation
@@ -20,7 +20,9 @@ def __init__(
 
         self._client = OpenAI(api_key=api_key)
 
-    def score(self, input: str, _: List[str] | None = None) -> Score[Moderation]:
+    def score(
+        self, input: str, _references: List[str] | None = None, _metadata: Dict[str, Any] | None = None
+    ) -> Score[Moderation]:
         """Score the input using the OpenAI Moderations API.
 
         Args:

aisploit/classifiers/package_hallucination.py

@@ -1,6 +1,6 @@
 import re
 from dataclasses import dataclass, field
-from typing import List
+from typing import Any, Dict, List
 
 import requests
 from stdlib_list import stdlib_list
@@ -20,7 +20,9 @@ class PythonPackageHallucinationClassifier(BaseTextClassifier[List[str]]):
     def __post_init__(self) -> None:
         self.libraries = stdlib_list(self.python_version)
 
-    def score(self, input: str, references: List[str] | None = None) -> Score[List[str]]:
+    def score(
+        self, input: str, _references: List[str] | None = None, _metadata: Dict[str, Any] | None = None
+    ) -> Score[List[str]]:
         """
         Scores the input based on the presence of hallucinated Python package names.
 

aisploit/classifiers/presidio/presidio_analyser.py

@@ -1,5 +1,5 @@
 from dataclasses import dataclass, field
-from typing import Callable, List, Optional
+from typing import Any, Callable, Dict, List, Optional
 
 from presidio_analyzer import AnalyzerEngine, EntityRecognizer, RecognizerResult
 
@@ -25,12 +25,14 @@ def __post_init__(self) -> None:
         for recognizer in self.additional_recognizers:
             self._analyzer.registry.add_recognizer(recognizer=recognizer)
 
-    def score(self, input: str, _: List[str] | None = None) -> Score[List[RecognizerResult]]:
+    def score(
+        self, input: str, _references: List[str] | None = None, _metadata: Dict[str, Any] | None = None
+    ) -> Score[List[RecognizerResult]]:
         """Score the input text for Personally Identifiable Information (PII) entities.
 
         Args:
             input (str): The input text to be scored.
-            _: List[str], optional): Ignored parameter. Defaults to None.
+            _references: List[str], optional): Ignored parameter. Defaults to None.
 
         Returns:
             Score[List[RecognizerResult]]: A Score object representing the results of PII detection.

aisploit/classifiers/repeated_token.py

@@ -0,0 +1,27 @@
+from dataclasses import dataclass
+from typing import Any, Dict, List
+
+from ..core import BaseTextClassifier, Score
+
+
+@dataclass
+class RepeatedTokenClassifier(BaseTextClassifier[bool]):
+    def score(
+        self, input: str, _references: List[str] | None = None, metadata: Dict[str, Any] | None = None
+    ) -> Score[bool]:
+        if not metadata:
+            raise ValueError("metadata  is missing")
+
+        repeated_token = metadata.get("repeated_token")
+
+        if not repeated_token:
+            raise ValueError("metadata with repeated token is missing")
+
+        # TODO
+
+        return Score[bool](
+            flagged=False,
+            value=False,
+            description="TODO",
+            explanation="TODO",
+        )

aisploit/classifiers/self_similarity.py

@@ -17,7 +17,9 @@ class SelfSimilarityClassifier(BaseTextClassifier[Dict[str, Any]]):
     aggregation: Literal["mean", "min"] = "mean"
     tags: List[str] = field(default_factory=lambda: ["hallucination"], init=False)
 
-    def score(self, input: str, references: List[str] | None = None) -> Score[Dict[str, Any]]:
+    def score(
+        self, input: str, references: List[str] | None = None, _metadata: Dict[str, Any] | None = None
+    ) -> Score[Dict[str, Any]]:
         """Score the input text based on its self-similarity to reference texts.
 
         Args:

aisploit/classifiers/text.py

@@ -1,6 +1,6 @@
 import re
 from dataclasses import dataclass
-from typing import List
+from typing import Any, Dict, List
 
 from ..core import BaseTextClassifier, Score
 
@@ -18,7 +18,9 @@ def __init__(self, *, pattern: re.Pattern, flag_matches=True) -> None:
         self._pattern = pattern
         self._flag_matches = flag_matches
 
-    def score(self, input: str, references: List[str] | None = None) -> Score[bool]:
+    def score(
+        self, input: str, _references: List[str] | None = None, _metadata: Dict[str, Any] | None = None
+    ) -> Score[bool]:
         """Score the input based on the regular expression pattern.
 
         Args:
@@ -62,7 +64,9 @@ def __init__(self, *, substring: str, ignore_case=True, flag_matches=True) -> No
 class TextTokenClassifier(BaseTextClassifier[bool]):
     token: str
 
-    def score(self, input: str, references: List[str] | None = None) -> Score[bool]:
+    def score(
+        self, input: str, _references: List[str] | None = None, _metadata: Dict[str, Any] | None = None
+    ) -> Score[bool]:
         return Score[bool](
             flagged=self.token in input,
             value=self.token in input,

aisploit/core/classifier.py

@@ -1,6 +1,6 @@
 from abc import ABC, abstractmethod
 from dataclasses import dataclass, field
-from typing import Generic, List, TypeVar
+from typing import Any, Dict, Generic, List, TypeVar
 
 T = TypeVar("T")
 Input = TypeVar("Input")
@@ -30,12 +30,15 @@ class BaseClassifier(ABC, Generic[T, Input]):
     tags: List[str] = field(default_factory=list, init=False)
 
     @abstractmethod
-    def score(self, input: Input, references: List[Input] | None = None) -> Score[T]:
+    def score(
+        self, input: Input, references: List[Input] | None = None, metadata: Dict[str, Any] | None = None
+    ) -> Score[T]:
         """Score the input and return a Score object.
 
         Args:
             input (Input): The input to be scored.
             references (List[Input], optional): List of reference inputs. Defaults to None.
+            metadata (Dict[str, Any], optional): Additional metadata for scoring. Defaults to {}.
 
         Returns:
             Score[T]: A Score object representing the score of the input.

aisploit/demo/gandalf.py

@@ -1,7 +1,7 @@
 import textwrap
 from dataclasses import dataclass
 from enum import Enum
-from typing import List
+from typing import Any, Dict, List
 
 import requests
 from langchain_core.messages import HumanMessage, SystemMessage
@@ -90,7 +90,9 @@ def __init__(self, level: GandalfLevel, chat_model: BaseChatModel) -> None:
         self._endpoint = f"{GANDALF_API_ENDPOINT}/guess-password"
         self._chain = chat_model | StrOutputParser()
 
-    def score(self, input: str, _: List[str] | None = None) -> Score[str]:
+    def score(
+        self, input: str, _references: List[str] | None = None, _metadata: Dict[str, Any] | None = None
+    ) -> Score[str]:
         score_description = (
             "When password is found in input, the score is the password itself. " "If nothing found, the score is empty"
         )

aisploit/scanner/__init__.py

@@ -1,10 +1,10 @@
 from .job import ScannerJob
-from .plugin import Plugin, SendPromptsPlugin
+from .plugin import BasePlugin, SendPromptsPlugin
 from .report import Issue, IssueCategory
 
 __all__ = [
     "ScannerJob",
-    "Plugin",
+    "BasePlugin",
     "SendPromptsPlugin",
     "Issue",
     "IssueCategory",

aisploit/scanner/job.py

@@ -3,15 +3,15 @@
 
 from tqdm.auto import tqdm
 
-from .plugin import Plugin
+from .plugin import BasePlugin
 from .report import Issue, ScanReport
 from ..core import BaseJob, BaseTarget, CallbackManager, Callbacks
 
 
 @dataclass
 class ScannerJob(BaseJob):
     target: BaseTarget
-    plugins: Sequence[Plugin]
+    plugins: Sequence[BasePlugin]
     callbacks: Callbacks = field(default_factory=list)
 
     def execute(self, *, run_id: Optional[str] = None, tags: Optional[Sequence[str]] = None) -> ScanReport: