Add unicode tags converter

hupe1980 · Apr 29, 2024 · 5c9e7d9 · 5c9e7d9
1 parent 9cce92d
commit 5c9e7d9
Show file tree

Hide file tree

Showing 4 changed files with 71 additions and 2 deletions.
diff --git a/aisploit/converters/__init__.py b/aisploit/converters/__init__.py
@@ -13,6 +13,7 @@
 from .stemming import StemmingConverter
 from .translation import TranslationConverter
 from .unicode_confusable import UnicodeConfusableConverter
+from .unicode_tags import UnicodeTagsConverter
 
 __all__ = [
     "Base64Converter",
@@ -30,4 +31,5 @@
     "StemmingConverter",
     "TranslationConverter",
     "UnicodeConfusableConverter",
+    "UnicodeTagsConverter",
 ]
diff --git a/aisploit/converters/unicode_tags.py b/aisploit/converters/unicode_tags.py
@@ -0,0 +1,23 @@
+from dataclasses import dataclass
+
+from ..core import BaseConverter
+
+
+@dataclass
+class UnicodeTagsConverter(BaseConverter):
+    prefix: str = ""
+    suffix: str = ""
+    add_sequence_markers: bool = False
+
+    def _convert(self, prompt: str) -> str:
+        encoded = ""
+
+        if self.add_sequence_markers:
+            encoded += chr(0xE0001)
+
+        encoded = ''.join(chr(0xE0000 + ord(ch)) for ch in prompt)
+
+        if self.add_sequence_markers:
+            encoded += chr(0xE007F)
+
+        return self.prefix + encoded + self.suffix
diff --git a/aisploit/scanner/plugins/prompt_injection.py b/aisploit/scanner/plugins/prompt_injection.py
@@ -31,7 +31,6 @@ class PromptInjectionPlugin(SendPromptsPlugin):
 
     def create_prompts(self) -> Sequence[str | Prompt]:
         """Create prompts using templates from the specified dataset."""
-
         prompts = []
         for data in self.dataset:
             prompt_template = PromptTemplate.from_template(data.template)

diff --git a/examples/converter.ipynb b/examples/converter.ipynb
@@ -24,7 +24,7 @@
     }
    ],
    "source": [
-    "from IPython.display import display, Markdown\n",
+    "from IPython.display import display, Markdown, HTML\n",
     "import os\n",
     "import textwrap\n",
     "from dotenv import load_dotenv\n",
@@ -41,6 +41,7 @@
     "    StemmingConverter,\n",
     "    UnicodeConfusableConverter,\n",
     "    TranslationConverter,\n",
+    "    UnicodeTagsConverter,\n",
     ")\n",
     "from aisploit.models import ChatOpenAI\n",
     "\n",
@@ -58,6 +59,50 @@
     ")"
    ]
   },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "## UnicodeTagsConverter"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 4,
+   "metadata": {},
+   "outputs": [
+    {
+     "data": {
+      "text/html": [
+       "<textarea rows=\"1\" cols=\"30\">START󠁗󠁨󠁡󠁴󠀠󠁩󠁳󠀠󠁴󠁨󠁥󠀠󠁣󠁡󠁰󠁩󠁴󠁡󠁬󠀠󠁯󠁦󠀠󠁆󠁲󠁡󠁮󠁣󠁥󠀿END</textarea>"
+      ],
+      "text/plain": [
+       "<IPython.core.display.HTML object>"
+      ]
+     },
+     "metadata": {},
+     "output_type": "display_data"
+    },
+    {
+     "data": {
+      "text/plain": [
+       "AIMessage(content=\"It seems like you've made a mistake in your input. Could you please provide more context or details?\", response_metadata={'token_usage': {'completion_tokens': 21, 'prompt_tokens': 9, 'total_tokens': 30}, 'model_name': 'gpt-4', 'system_fingerprint': None, 'finish_reason': 'stop', 'logprobs': None}, id='run-17e2154f-9504-45c4-ac36-f8b1b4f7ce01-0')"
+      ]
+     },
+     "execution_count": 4,
+     "metadata": {},
+     "output_type": "execute_result"
+    }
+   ],
+   "source": [
+    "converter = UnicodeTagsConverter(prefix=\"START\", suffix=\"END\")\n",
+    "prompt = converter.convert(\"What is the capital of France?\")\n",
+    "\n",
+    "display(HTML(f'<textarea rows=\"1\" cols=\"30\">{prompt.to_string()}</textarea>'))\n",
+    "\n",
+    "chat_model.invoke(prompt)"
+   ]
+  },
   {
    "cell_type": "markdown",
    "metadata": {},