diff --git a/pyrit/converter/prompt_template_converter.py b/pyrit/converter/prompt_template_converter.py index c2fba426bd..3419cd61c0 100644 --- a/pyrit/converter/prompt_template_converter.py +++ b/pyrit/converter/prompt_template_converter.py @@ -1,6 +1,7 @@ # Copyright (c) Microsoft Corporation. # Licensed under the MIT license. +import hashlib import re from pyrit.converter.converter import Converter, ConverterResult @@ -16,6 +17,10 @@ class PromptTemplateConverter(Converter): framing (e.g. ``TASK is '{{ prompt }}'``) as well as wrapping the prompt in surrounding content, such as hiding it in an HTML comment for indirect prompt injection (e.g. ``

Visible text

``). + + The converter identifier stores a fixed-size hash of the configured template, + rather than the template text. The deprecated ``TaskFramingConverter`` retains + its legacy identifier. """ SUPPORTED_INPUT_TYPES = ("text",) @@ -56,9 +61,10 @@ def _build_identifier(self) -> ComponentIdentifier: Returns: ComponentIdentifier: The identifier for this converter. """ + template_hash = hashlib.sha256(self._template.encode("utf-8")).hexdigest()[:16] return self._create_identifier( params={ - "template": self._template, + "template_hash": template_hash, "strip_characters": self._strip_characters, }, ) diff --git a/tests/unit/converter/test_prompt_template_converter.py b/tests/unit/converter/test_prompt_template_converter.py index 7bce0f4416..8864c7a223 100644 --- a/tests/unit/converter/test_prompt_template_converter.py +++ b/tests/unit/converter/test_prompt_template_converter.py @@ -96,12 +96,24 @@ def test_input_output_types(): assert converter.output_supported("text") is True -def test_identifier_includes_template_and_strip_characters(): +def test_identifier_includes_template_hash_and_strip_characters(): converter = PromptTemplateConverter(template="T: {{ prompt }}", strip_characters="'") identifier = converter.get_identifier() assert identifier.class_name == "PromptTemplateConverter" - assert identifier.params["template"] == "T: {{ prompt }}" + assert identifier.params["template_hash"] == "1d6417505e0571d2" assert identifier.params["strip_characters"] == "'" + assert "template" not in identifier.params + + +def test_identifier_keeps_hash(): + converter = PromptTemplateConverter(template="T: {{ prompt }}", strip_characters="'") + assert converter.get_identifier().hash == "f7a37b712954fc0dd196ee9481ca7a1e0779841b768171bb6af146a046dde693" + + +def test_identifier_matches_for_equivalent_templates(): + first = PromptTemplateConverter(template="A {{ prompt }}").get_identifier() + second = PromptTemplateConverter(template="A {{ prompt }}").get_identifier() + assert first == second def test_identifier_differs_by_template(): @@ -110,6 +122,19 @@ def test_identifier_differs_by_template(): assert first.hash != second.hash +def test_identifier_differs_by_strip_characters(): + first = PromptTemplateConverter(template="A {{ prompt }}", strip_characters="'").get_identifier() + second = PromptTemplateConverter(template="A {{ prompt }}", strip_characters='"').get_identifier() + assert first.hash != second.hash + + +def test_long_template_identifier_uses_fixed_size_hash(): + template = f"{'A' * 10_000} {{{{ prompt }}}}" + identifier = PromptTemplateConverter(template=template).get_identifier() + assert identifier.params["template_hash"] == "a280c02c55a808af" + assert template not in identifier.model_dump_json() + + # Markdown hidden-text recipes from the converters doc: put the prompt on one line, # backslash-escape ``\`` and the delimiter, then insert it with PromptTemplateConverter. _MARKDOWN_RECIPES = { diff --git a/tests/unit/converter/test_task_framing_converter.py b/tests/unit/converter/test_task_framing_converter.py index c9df10c5e8..7296aa79e2 100644 --- a/tests/unit/converter/test_task_framing_converter.py +++ b/tests/unit/converter/test_task_framing_converter.py @@ -56,3 +56,10 @@ def test_identifier_keeps_task_template_param_name(): assert identifier.params["task_template"] == "TASK is '{{ prompt }}'" assert identifier.params["strip_characters"] == "'" assert "template" not in identifier.params + assert "template_hash" not in identifier.params + + +def test_identifier_keeps_legacy_hash(): + # Pinned to the v1.1.0 value so the deprecated converter keeps matching identifiers already stored. + identifier = _make(strip_characters="'").get_identifier() + assert identifier.hash == "b680cca8b216deec74b66758dd6bdc005fe21e593fbd0a9a975fe40e7e9fe56d"