Add prefix

Author: hupe1980

aisploit/converters/base64.py

@@ -1,8 +1,17 @@
 import base64
+from dataclasses import dataclass
 
 from ..core import BaseConverter
 
 
+@dataclass
 class Base64Converter(BaseConverter):
+    prefix: str = "The following instruction is base64 encoded:"
+    add_prefix: bool = True
+
     def _convert(self, prompt: str) -> str:
-        return base64.b64encode(prompt.encode("utf-8")).decode("utf-8")
+        b64_string = base64.b64encode(prompt.encode("utf-8")).decode("utf-8")
+        if self.add_prefix and not b64_string == "":
+            return f"{self.prefix} {b64_string}"
+        else:
+            return b64_string

examples/converter.ipynb

@@ -72,7 +72,7 @@
     {
      "data": {
       "text/markdown": [
-       "VGhpcyBpcyBhIHNhbXBsZSBwcm9tcHQu"
+       "The following instruction is base64 encoded: IVt0ZXN0XShodHRwOi8vdGVzdC5kZSk="
       ],
       "text/plain": [
        "<IPython.core.display.Markdown object>"

tests/converter/test_base64.py

@@ -1,4 +1,3 @@
-import base64
 
 from langchain_core.prompt_values import StringPromptValue
 
@@ -8,34 +7,33 @@
 def test_convert_single_prompt():
     converter = Base64Converter()
     prompt = "Test"
-    expected_output = base64.b64encode(prompt.encode("utf-8")).decode("utf-8")
-    assert converter.convert(StringPromptValue(text=prompt)) == StringPromptValue(
-        text=expected_output
-    )
+    expected = "The following instruction is base64 encoded: VGVzdA=="
+    assert converter.convert(StringPromptValue(text=prompt)).to_string() == expected
+
+
+def test_convert_single_prompt_no_prefix():
+    converter = Base64Converter(add_prefix=False)
+    prompt = "Test"
+    expected = "VGVzdA=="
+    assert converter.convert(StringPromptValue(text=prompt)).to_string() == expected
 
 
 def test_convert_multiple_prompts():
     converter = Base64Converter()
     prompt = "Test1"
-    expected_output = base64.b64encode(prompt.encode("utf-8")).decode("utf-8")
-    assert converter.convert(StringPromptValue(text=prompt)) == StringPromptValue(
-        text=expected_output
-    )
+    expected="The following instruction is base64 encoded: VGVzdDE="
+    assert converter.convert(StringPromptValue(text=prompt)).to_string() == expected
 
 
 def test_convert_empty_prompt():
     converter = Base64Converter()
     prompt = ""
-    expected_output = base64.b64encode(prompt.encode("utf-8")).decode("utf-8")
-    assert converter.convert(StringPromptValue(text=prompt)) == StringPromptValue(
-        text=expected_output
-    )
+    expected = ""
+    assert converter.convert(StringPromptValue(text=prompt)).to_string() == expected
 
 
 def test_convert_with_unicode_characters():
     converter = Base64Converter()
     prompt = "äöüß"
-    expected_output = base64.b64encode(prompt.encode("utf-8")).decode("utf-8")
-    assert converter.convert(StringPromptValue(text=prompt)) == StringPromptValue(
-        text=expected_output
-    )
+    expected = "The following instruction is base64 encoded: w6TDtsO8w58="
+    assert converter.convert(StringPromptValue(text=prompt)).to_string() == expected