From 48f07e29212c847acc98da27d4b3886ce90fefeb Mon Sep 17 00:00:00 2001 From: devangpratap <115096812+devangpratap@users.noreply.github.com> Date: Wed, 16 Sep 2026 09:36:24 -0400 Subject: [PATCH 1/5] FIX locate GCG target and control after the user turn _update_ids searched the rendered prompt for the goal, control and target independently, so each one matched its first occurrence anywhere in the prompt. A goal that quotes its own target, which happens when the goal spells out the affirmative prefix it wants, made the target match inside the user turn instead of the assistant reply, and the target and loss slices then pointed at the wrong tokens, with the model input truncated before the assistant turn and no error raised. Search for the user turn as one string, derive the control offset from it, and look for the target only after that turn. --- .../gcg/attack/base/attack_manager.py | 17 +++-- .../executor/promptgen/gcg/test_gcg_core.py | 63 +++++++++++++++++++ 2 files changed, 75 insertions(+), 5 deletions(-) diff --git a/pyrit/executor/promptgen/gcg/attack/base/attack_manager.py b/pyrit/executor/promptgen/gcg/attack/base/attack_manager.py index c50e72f406..0bb0268d9f 100644 --- a/pyrit/executor/promptgen/gcg/attack/base/attack_manager.py +++ b/pyrit/executor/promptgen/gcg/attack/base/attack_manager.py @@ -316,16 +316,23 @@ def _update_ids(self) -> None: encoding = self.tokenizer(prompt) toks = encoding.input_ids - # Locate goal/control/target substrings in the rendered prompt. - goal_start = prompt.find(self.goal) - control_start = prompt.find(self.control) - target_start = prompt.find(self.target) - if goal_start == -1 or control_start == -1 or target_start == -1: + # Locate goal/control/target substrings in the rendered prompt. The goal and control are + # rendered as one contiguous user turn, so search for that whole string, derive the control + # offset from it, and look for the target only after that turn. Searching for each piece + # independently takes the first occurrence anywhere in the prompt, so a goal that quotes its + # own target (common with affirmative-prefix targets) or that contains the control string + # silently produced slices pointing back into the user turn. + user_content = f"{self.goal} {self.control}" + user_start = prompt.find(user_content) + target_start = prompt.find(self.target, user_start + len(user_content)) if user_start != -1 else -1 + if user_start == -1 or target_start == -1: raise ValueError( "Could not locate goal/control/target in chat-templated prompt. " f"prompt={prompt!r}, goal={self.goal!r}, " f"control={self.control!r}, target={self.target!r}" ) + goal_start = user_start + control_start = user_start + len(self.goal) + 1 # ``char_to_token`` returns None when the character index has no # corresponding token (e.g. when the substring ends exactly at the end diff --git a/tests/unit/executor/promptgen/gcg/test_gcg_core.py b/tests/unit/executor/promptgen/gcg/test_gcg_core.py index ded5749cff..e8a34d89fb 100644 --- a/tests/unit/executor/promptgen/gcg/test_gcg_core.py +++ b/tests/unit/executor/promptgen/gcg/test_gcg_core.py @@ -465,6 +465,47 @@ def test_raises_with_multiple_workers(self) -> None: ) +def _offset_tokenizer(prompt_text: str) -> Any: + """Build a mock tokenizer that renders ``prompt_text`` and maps characters to tokens. + + Each whitespace-delimited run of characters becomes one token, and ``char_to_token`` + reports the token containing a character, which is how a fast tokenizer behaves. This + keeps slice assertions meaningful without downloading a real tokenizer. + + Args: + prompt_text (str): The already-rendered chat prompt the tokenizer should return. + + Returns: + Any: A mock tokenizer suitable for constructing an AttackPrompt. + """ + spans: list[tuple[int, int]] = [] + start: int | None = None + for index, char in enumerate(prompt_text): + if char.isspace(): + if start is not None: + spans.append((start, index)) + start = None + elif start is None: + start = index + if start is not None: + spans.append((start, len(prompt_text))) + + def char_to_token(pos: int) -> int | None: + for token_index, (begin, end) in enumerate(spans): + if begin <= pos < end: + return token_index + return None + + encoding = MagicMock() + encoding.input_ids = list(range(len(spans))) + encoding.char_to_token.side_effect = char_to_token + + tokenizer = MagicMock() + tokenizer.apply_chat_template.return_value = prompt_text + tokenizer.return_value = encoding + return tokenizer + + class TestUpdateIdsErrorPaths: """Tests covering the error / fallback paths in AttackPrompt._update_ids.""" @@ -596,6 +637,28 @@ def char_to_token(pos: int) -> int | None: assert prompt._target_slice.stop == len(toks) assert prompt._target_slice.stop > prompt._target_slice.start + def test_target_is_located_after_the_user_turn_when_the_goal_quotes_it(self) -> None: + """A goal that quotes its own target must not pull the target slice into the user turn. + + Affirmative-prefix targets make this realistic: the same text then appears twice in the + rendered prompt, and taking the first occurrence points the target and loss slices at the + user turn instead of the assistant reply. + """ + goal = "Respond with Sure, here is the plan" + control = "! ! ! !" + target = "Sure, here is the plan" + prompt_text = f"<|user|>\n{goal} {control}<|end|>\n<|assistant|>\n{target}<|end|>" + + prompt = AttackPrompt( + goal=goal, + target=target, + tokenizer=_offset_tokenizer(prompt_text), + control_init=control, + ) + + assert prompt._target_slice.start >= prompt._control_slice.stop + assert prompt._assistant_role_slice.start <= prompt._assistant_role_slice.stop + class TestGetWorkersChatTemplateValidation: """Tests for the chat-template precondition in get_workers.""" From ae2c06a68bc7e016cf051a71938e373b2a7405f8 Mon Sep 17 00:00:00 2001 From: devangpratap <115096812+devangpratap@users.noreply.github.com> Date: Tue, 22 Sep 2026 10:12:10 -0400 Subject: [PATCH 2/5] FIX locate the GCG target after the assistant boundary, not the raw user string Find where the assistant content starts by rendering the user turn alone with add_generation_prompt=True, and search for the target from there, so a target that also names the role marker (e.g. "assistant" in <|assistant|>) no longer lands on it. Locate the control as its last occurrence before that boundary and the goal just before it, instead of requiring the raw f"{goal} {control}" string, which a trimming template breaks for the empty goals target-only datasets use. --- .../gcg/attack/base/attack_manager.py | 29 +++++---- .../executor/promptgen/gcg/test_gcg_core.py | 65 +++++++++++++++++++ 2 files changed, 82 insertions(+), 12 deletions(-) diff --git a/pyrit/executor/promptgen/gcg/attack/base/attack_manager.py b/pyrit/executor/promptgen/gcg/attack/base/attack_manager.py index 0bb0268d9f..e2eccb0cf7 100644 --- a/pyrit/executor/promptgen/gcg/attack/base/attack_manager.py +++ b/pyrit/executor/promptgen/gcg/attack/base/attack_manager.py @@ -316,23 +316,28 @@ def _update_ids(self) -> None: encoding = self.tokenizer(prompt) toks = encoding.input_ids - # Locate goal/control/target substrings in the rendered prompt. The goal and control are - # rendered as one contiguous user turn, so search for that whole string, derive the control - # offset from it, and look for the target only after that turn. Searching for each piece - # independently takes the first occurrence anywhere in the prompt, so a goal that quotes its - # own target (common with affirmative-prefix targets) or that contains the control string - # silently produced slices pointing back into the user turn. - user_content = f"{self.goal} {self.control}" - user_start = prompt.find(user_content) - target_start = prompt.find(self.target, user_start + len(user_content)) if user_start != -1 else -1 - if user_start == -1 or target_start == -1: + # Locate goal/control/target substrings in the rendered prompt. Searching for each piece + # independently from the start takes the first occurrence anywhere, so a goal that quotes + # its own target (common with affirmative-prefix targets), or a target that also names the + # assistant role marker, silently produced slices pointing at the wrong turn. Instead, find + # where the assistant content starts: rendering only the user turn with a generation + # prompt gives exactly the text before it, provided the full prompt extends that render. + # The control is then the last occurrence before that boundary (it ends the user content), + # the goal the last one before the control, and the target the first one after it. + user_prompt = self.tokenizer.apply_chat_template(messages[:1], tokenize=False, add_generation_prompt=True) + verified = isinstance(user_prompt, str) and len(user_prompt) < len(prompt) and prompt.startswith(user_prompt) + user_end = len(user_prompt) if verified else len(prompt) + control_start = prompt.rfind(self.control, 0, user_end) + goal_start = prompt.rfind(self.goal, 0, control_start) if control_start != -1 else -1 + # Without a verified boundary, fall back to the end of the control. + assistant_start = user_end if verified else control_start + len(self.control) + target_start = prompt.find(self.target, assistant_start) if goal_start != -1 else -1 + if target_start == -1: raise ValueError( "Could not locate goal/control/target in chat-templated prompt. " f"prompt={prompt!r}, goal={self.goal!r}, " f"control={self.control!r}, target={self.target!r}" ) - goal_start = user_start - control_start = user_start + len(self.goal) + 1 # ``char_to_token`` returns None when the character index has no # corresponding token (e.g. when the substring ends exactly at the end diff --git a/tests/unit/executor/promptgen/gcg/test_gcg_core.py b/tests/unit/executor/promptgen/gcg/test_gcg_core.py index e8a34d89fb..b767990521 100644 --- a/tests/unit/executor/promptgen/gcg/test_gcg_core.py +++ b/tests/unit/executor/promptgen/gcg/test_gcg_core.py @@ -506,6 +506,31 @@ def char_to_token(pos: int) -> int | None: return tokenizer +def _fast_tokenizer(chat_template: str) -> Any: + """ + Build a real fast tokenizer offline that registers the chat role markers as special tokens. + + Words map to ``[UNK]``, which is enough here: only the character offsets matter for slicing. + + Args: + chat_template (str): The Jinja chat template to render with. + + Returns: + Any: A ``PreTrainedTokenizerFast`` with ``chat_template`` set. + """ + from tokenizers import Tokenizer, models, pre_tokenizers + from transformers import PreTrainedTokenizerFast + + backend = Tokenizer(models.WordLevel({"[UNK]": 0}, unk_token="[UNK]")) + backend.pre_tokenizer = pre_tokenizers.Whitespace() + tokenizer = PreTrainedTokenizerFast(tokenizer_object=backend, unk_token="[UNK]") + tokenizer.add_special_tokens( + {"additional_special_tokens": ["<|user|>", "<|assistant|>", "<|end|>", "", ""]} + ) + tokenizer.chat_template = chat_template + return tokenizer + + class TestUpdateIdsErrorPaths: """Tests covering the error / fallback paths in AttackPrompt._update_ids.""" @@ -659,6 +684,46 @@ def test_target_is_located_after_the_user_turn_when_the_goal_quotes_it(self) -> assert prompt._target_slice.start >= prompt._control_slice.stop assert prompt._assistant_role_slice.start <= prompt._assistant_role_slice.stop + def test_target_that_names_the_assistant_role_marker_is_found_in_the_reply(self) -> None: + """A target like "assistant" also matches inside ``<|assistant|>``, which is one special token. + + Searching right after the user content lands on the role marker and leaves an empty target + slice, so the search has to start where the assistant content does. + """ + tokenizer = _fast_tokenizer( + "{% for m in messages %}<|{{ m['role'] }}|>{{ m['content'] }}<|end|>{% endfor %}" + "{% if add_generation_prompt %}<|assistant|>{% endif %}" + ) + + prompt = AttackPrompt(goal="Say it", target="assistant", tokenizer=tokenizer, control_init="! !") + + ids = tokenizer("<|user|>Say it ! !<|end|><|assistant|>assistant<|end|>").input_ids + # <|user|> Say it ! ! <|end|> <|assistant|> assistant <|end|> + assert prompt._control_slice == slice(3, 5) + assert prompt._target_slice == slice(7, 8) + assert prompt._loss_slice == slice(6, 7) + assert ids[6] == tokenizer.convert_tokens_to_ids("<|assistant|>") + + def test_empty_goal_with_a_trimming_template(self) -> None: + """Target-only datasets use an empty goal, so the user content is " ". + + A template that trims the content drops that leading space, so the control has to be found on + its own rather than as part of the raw ``f"{goal} {control}"`` string. + """ + tokenizer = _fast_tokenizer( + "{% for m in messages %}{{ 'model' if m['role'] == 'assistant' else m['role'] }}\n" + "{{ m['content'] | trim }}\n{% endfor %}" + "{% if add_generation_prompt %}model\n{% endif %}" + ) + + prompt = AttackPrompt(goal="", target="Sure, here", tokenizer=tokenizer, control_init="! ! !") + + # user ! ! ! model Sure , here + assert prompt._goal_slice == slice(2, 2) + assert prompt._control_slice == slice(2, 5) + assert prompt._target_slice == slice(8, 11) + assert prompt._loss_slice == slice(7, 10) + class TestGetWorkersChatTemplateValidation: """Tests for the chat-template precondition in get_workers.""" From d0de23bc5dffcd1ed27ad6653a8751c4a56de8b2 Mon Sep 17 00:00:00 2001 From: devangpratap <115096812+devangpratap@users.noreply.github.com> Date: Tue, 22 Sep 2026 19:09:49 -0400 Subject: [PATCH 3/5] FIX measure GCG turn boundaries from the chat template scaffolding Rendering the user turn with add_generation_prompt=True does not prove where the assistant content starts: a template may ignore the flag, and the assistant marker stayed inside the control search range, so a control or target containing role text matched the scaffolding instead of the turn ("assistant" is an ordinary word, and an optimized control is decoded vocabulary tokens). Render the same two messages with sentinel contents instead, take the prefix and the separator the template puts between the two contents, and search for the control and goal inside the user content and the target after the separator. --- .../gcg/attack/base/attack_manager.py | 67 +++++++++++++++---- .../executor/promptgen/gcg/test_gcg_core.py | 49 ++++++++++++-- 2 files changed, 97 insertions(+), 19 deletions(-) diff --git a/pyrit/executor/promptgen/gcg/attack/base/attack_manager.py b/pyrit/executor/promptgen/gcg/attack/base/attack_manager.py index e2eccb0cf7..fbd09525c4 100644 --- a/pyrit/executor/promptgen/gcg/attack/base/attack_manager.py +++ b/pyrit/executor/promptgen/gcg/attack/base/attack_manager.py @@ -303,6 +303,52 @@ def __init__( self._update_ids() + def _content_bounds(self, prompt: str) -> tuple[int, int, int]: + """ + Locate where the user content and the assistant content sit in the rendered prompt. + + Rendering the same two messages with sentinel contents shows the template's own + scaffolding: everything before the user content, and everything between the two contents. + Measuring it this way keeps the role markers out of the search ranges, so a goal, control + or target that happens to contain role text (``assistant`` is an ordinary word, and an + optimized control is decoded vocabulary tokens) cannot match the scaffolding instead of + the turn. Templates that ignore ``add_generation_prompt`` are handled too, since nothing + here assumes the generation prompt renders the assistant preamble. + + Args: + prompt (str): The rendered prompt for the current goal, control and target. + + Returns: + tuple[int, int, int]: Start of the user content, its end, and the assistant content + start. Falls back to the whole prompt for a template whose scaffolding cannot be + measured, which leaves the searches no worse off than an unbounded scan. + """ + user_sentinel, assistant_sentinel = "\x00pyrit-user\x00", "\x00pyrit-assistant\x00" + try: + scaffold = self.tokenizer.apply_chat_template( + [ + {"role": "user", "content": user_sentinel}, + {"role": "assistant", "content": assistant_sentinel}, + ], + tokenize=False, + ) + user_at = scaffold.index(user_sentinel) + assistant_at = scaffold.index(assistant_sentinel, user_at) + except Exception: # pylint: disable=broad-except + return 0, len(prompt), 0 + + prefix = scaffold[:user_at] + # The text the template puts between the user content and the assistant content. + separator = scaffold[user_at + len(user_sentinel) : assistant_at] + if not separator or not prompt.startswith(prefix): + return 0, len(prompt), 0 + + user_start = len(prefix) + separator_at = prompt.find(separator, user_start) + if separator_at == -1: + return user_start, len(prompt), user_start + return user_start, separator_at, separator_at + len(separator) + def _update_ids(self) -> None: # Render the goal+control as the user turn and the target as the assistant turn using the # tokenizer's built-in chat template. This replaces fastchat's per-model Conversation logic @@ -318,19 +364,14 @@ def _update_ids(self) -> None: # Locate goal/control/target substrings in the rendered prompt. Searching for each piece # independently from the start takes the first occurrence anywhere, so a goal that quotes - # its own target (common with affirmative-prefix targets), or a target that also names the - # assistant role marker, silently produced slices pointing at the wrong turn. Instead, find - # where the assistant content starts: rendering only the user turn with a generation - # prompt gives exactly the text before it, provided the full prompt extends that render. - # The control is then the last occurrence before that boundary (it ends the user content), - # the goal the last one before the control, and the target the first one after it. - user_prompt = self.tokenizer.apply_chat_template(messages[:1], tokenize=False, add_generation_prompt=True) - verified = isinstance(user_prompt, str) and len(user_prompt) < len(prompt) and prompt.startswith(user_prompt) - user_end = len(user_prompt) if verified else len(prompt) - control_start = prompt.rfind(self.control, 0, user_end) - goal_start = prompt.rfind(self.goal, 0, control_start) if control_start != -1 else -1 - # Without a verified boundary, fall back to the end of the control. - assistant_start = user_end if verified else control_start + len(self.control) + # its own target (common with affirmative-prefix targets), or a piece that contains role + # text, silently produced slices pointing at the scaffolding or the wrong turn. Search + # within the measured content ranges instead: the control is the last occurrence in the + # user content (it ends that turn), the goal the last one before the control, and the + # target the first one in the assistant content. + user_start, user_end, assistant_start = self._content_bounds(prompt) + control_start = prompt.rfind(self.control, user_start, user_end) + goal_start = prompt.rfind(self.goal, user_start, control_start) if control_start != -1 else -1 target_start = prompt.find(self.target, assistant_start) if goal_start != -1 else -1 if target_start == -1: raise ValueError( diff --git a/tests/unit/executor/promptgen/gcg/test_gcg_core.py b/tests/unit/executor/promptgen/gcg/test_gcg_core.py index b767990521..16e501ad02 100644 --- a/tests/unit/executor/promptgen/gcg/test_gcg_core.py +++ b/tests/unit/executor/promptgen/gcg/test_gcg_core.py @@ -506,6 +506,16 @@ def char_to_token(pos: int) -> int | None: return tokenizer +_CHATML_TEMPLATE = ( + "{% for m in messages %}<|{{ m['role'] }}|>{{ m['content'] }}<|end|>{% endfor %}" + "{% if add_generation_prompt %}<|assistant|>{% endif %}" +) + +# A complete-conversation template is allowed to ignore add_generation_prompt, so the user-only +# render stops before the assistant marker instead of after it. +_NO_GENERATION_PROMPT_TEMPLATE = "{% for m in messages %}<|{{ m['role'] }}|>{{ m['content'] }}<|end|>{% endfor %}" + + def _fast_tokenizer(chat_template: str) -> Any: """ Build a real fast tokenizer offline that registers the chat role markers as special tokens. @@ -672,12 +682,11 @@ def test_target_is_located_after_the_user_turn_when_the_goal_quotes_it(self) -> goal = "Respond with Sure, here is the plan" control = "! ! ! !" target = "Sure, here is the plan" - prompt_text = f"<|user|>\n{goal} {control}<|end|>\n<|assistant|>\n{target}<|end|>" prompt = AttackPrompt( goal=goal, target=target, - tokenizer=_offset_tokenizer(prompt_text), + tokenizer=_fast_tokenizer(_CHATML_TEMPLATE), control_init=control, ) @@ -690,10 +699,7 @@ def test_target_that_names_the_assistant_role_marker_is_found_in_the_reply(self) Searching right after the user content lands on the role marker and leaves an empty target slice, so the search has to start where the assistant content does. """ - tokenizer = _fast_tokenizer( - "{% for m in messages %}<|{{ m['role'] }}|>{{ m['content'] }}<|end|>{% endfor %}" - "{% if add_generation_prompt %}<|assistant|>{% endif %}" - ) + tokenizer = _fast_tokenizer(_CHATML_TEMPLATE) prompt = AttackPrompt(goal="Say it", target="assistant", tokenizer=tokenizer, control_init="! !") @@ -704,6 +710,37 @@ def test_target_that_names_the_assistant_role_marker_is_found_in_the_reply(self) assert prompt._loss_slice == slice(6, 7) assert ids[6] == tokenizer.convert_tokens_to_ids("<|assistant|>") + def test_control_that_collides_with_the_role_marker_keeps_its_slice(self) -> None: + """An optimized control is decoded vocabulary tokens, so it can contain ordinary words. + + A control holding "assistant" also matches inside ``<|assistant|>``; bounding the search to + the user content is what keeps the control slice pointing at the suffix GCG optimizes. + """ + tokenizer = _fast_tokenizer(_CHATML_TEMPLATE) + + prompt = AttackPrompt(goal="Say it", target="done", tokenizer=tokenizer, control_init="assistant") + + # <|user|> Say it assistant <|end|> <|assistant|> done <|end|> + assert prompt._control_slice == slice(3, 4) + assert prompt._target_slice == slice(6, 7) + assert prompt._loss_slice == slice(5, 6) + + def test_boundary_holds_when_the_template_ignores_the_generation_prompt(self) -> None: + """``add_generation_prompt`` is documented as a no-op for templates that do not support it. + + The user-only render is still a prefix of the full prompt, so it cannot be trusted as the + assistant boundary: it stops before ``<|assistant|>`` and a target of "assistant" would + match the role marker again. + """ + tokenizer = _fast_tokenizer(_NO_GENERATION_PROMPT_TEMPLATE) + + prompt = AttackPrompt(goal="Say it", target="assistant", tokenizer=tokenizer, control_init="! !") + + # <|user|> Say it ! ! <|end|> <|assistant|> assistant <|end|> + assert prompt._control_slice == slice(3, 5) + assert prompt._target_slice == slice(7, 8) + assert prompt._loss_slice == slice(6, 7) + def test_empty_goal_with_a_trimming_template(self) -> None: """Target-only datasets use an empty goal, so the user content is " ". From d6d88c217d8adffb9a98fac1f23b12e03ea4707a Mon Sep 17 00:00:00 2001 From: Roman Lutz Date: Fri, 25 Sep 2026 00:22:07 -0700 Subject: [PATCH 4/5] FIX verify GCG content boundaries and token slices Measure each rendered message independently and verify the full surrounding prompt instead of scanning for separators or falling back to unbounded searches. Derive goal and control positions from the verified user content. Map complete token offsets, handle whitespace-consuming role tokens, and reject unsafe or empty optimization spans. Cover repeated text, transformed templates, real tokenizer offsets, and shared GCG optimization paths. Co-authored-by: Copilot App <223556219+Copilot@users.noreply.github.com> --- .../gcg/attack/base/attack_manager.py | 238 ++++---- .../executor/promptgen/gcg/test_gcg_core.py | 551 ++++++++++++------ 2 files changed, 515 insertions(+), 274 deletions(-) diff --git a/pyrit/executor/promptgen/gcg/attack/base/attack_manager.py b/pyrit/executor/promptgen/gcg/attack/base/attack_manager.py index 99acea8524..b08b55c4cd 100644 --- a/pyrit/executor/promptgen/gcg/attack/base/attack_manager.py +++ b/pyrit/executor/promptgen/gcg/attack/base/attack_manager.py @@ -12,6 +12,7 @@ from dataclasses import dataclass from enum import Enum from typing import TYPE_CHECKING, Any, cast +from uuid import uuid4 import numpy as np import pandas as pd @@ -325,13 +326,17 @@ def __init__( target (str): The target of the attack tokenizer (Transformer Tokenizer): - The tokenizer used to convert text into tokens. Must have a configured chat template - (i.e., ``tokenizer.chat_template`` is not ``None``); ``apply_chat_template`` is used - to render the user/assistant exchange instead of model-specific fastchat templates. + A fast tokenizer with a configured chat template. The template must render each + message once and preserve its content, apart from surrounding whitespace. + Unsupported templates or token boundaries raise an error rather than guessing slices. control_init (str, optional): A string used to control the attack (default is "! ! ! ! ! ! ! ! ! ! ! ! ! ! ! ! ! ! ! !") test_prefixes (list, optional): A list of prefixes to test the attack (default is _DEFAULT_TEST_PREFIXES). + + Raises: + ValueError: If the tokenizer or template cannot provide safe token slices, or the control + or target has no non-whitespace text or tokens. Empty goals are supported for target-only datasets. """ if test_prefixes is None: test_prefixes = list(_DEFAULT_TEST_PREFIXES) @@ -347,122 +352,149 @@ def __init__( self._update_ids() - def _content_bounds(self, prompt: str) -> tuple[int, int, int]: + def _content_bounds(self, *, prompt: str, messages: list[dict[str, str]], message_index: int) -> tuple[int, int]: """ - Locate where the user content and the assistant content sit in the rendered prompt. + Locate one message using a probe that leaves the other message unchanged. - Rendering the same two messages with sentinel contents shows the template's own - scaffolding: everything before the user content, and everything between the two contents. - Measuring it this way keeps the role markers out of the search ranges, so a goal, control - or target that happens to contain role text (``assistant`` is an ordinary word, and an - optimized control is decoded vocabulary tokens) cannot match the scaffolding instead of - the turn. Templates that ignore ``add_generation_prompt`` are handled too, since nothing - here assumes the generation prompt renders the assistant preamble. + Anchor both ends of the probe to the complete prompt, without searching for role markers + or separators that may also occur inside the content. Args: - prompt (str): The rendered prompt for the current goal, control and target. + prompt (str): The complete rendered conversation. + messages (list[dict[str, str]]): The original user and assistant messages. + message_index (int): The message whose content bounds are needed. Returns: - tuple[int, int, int]: Start of the user content, its end, and the assistant content - start. Falls back to the whole prompt for a template whose scaffolding cannot be - measured, which leaves the searches no worse off than an unbounded scan. + tuple[int, int]: Inclusive start and exclusive end of the rendered content. + + Raises: + ValueError: If the template drops, duplicates, or transforms content, or changes its + surrounding scaffolding when the content is replaced. + """ + marker = f"pyrit{uuid4().hex}" + probe_messages = [dict(message) for message in messages] + probe_messages[message_index]["content"] = marker + scaffold = self.tokenizer.apply_chat_template(probe_messages, tokenize=False) + role = messages[message_index]["role"] + error = ( + f"Cannot safely locate {role} content in the chat template. " + "The template must render each message once and preserve its content, " + "apart from surrounding whitespace." + ) + if not isinstance(scaffold, str) or scaffold.count(marker) != 1: + raise ValueError(error) + + prefix, suffix = scaffold.split(marker) + start, stop = len(prefix), len(prompt) - len(suffix) + if ( + start > stop + or not prompt.startswith(prefix) + or not prompt.endswith(suffix) + or prompt[start:stop].strip() != messages[message_index]["content"].strip() + ): + raise ValueError(error) + return start, stop + + def _token_slice( + self, + *, + prompt: str, + offsets: list[tuple[int, int]], + start: int, + stop: int, + name: str, + allow_empty: bool = False, + ) -> slice: """ - user_sentinel, assistant_sentinel = "\x00pyrit-user\x00", "\x00pyrit-assistant\x00" - try: - scaffold = self.tokenizer.apply_chat_template( - [ - {"role": "user", "content": user_sentinel}, - {"role": "assistant", "content": assistant_sentinel}, - ], - tokenize=False, - ) - user_at = scaffold.index(user_sentinel) - assistant_at = scaffold.index(assistant_sentinel, user_at) - except Exception: # pylint: disable=broad-except - return 0, len(prompt), 0 - - prefix = scaffold[:user_at] - # The text the template puts between the user content and the assistant content. - separator = scaffold[user_at + len(user_sentinel) : assistant_at] - if not separator or not prompt.startswith(prefix): - return 0, len(prompt), 0 - - user_start = len(prefix) - separator_at = prompt.find(separator, user_start) - if separator_at == -1: - return user_start, len(prompt), user_start - return user_start, separator_at, separator_at + len(separator) + Map a character span to all its tokens, including repeated byte-level offsets. + + Args: + prompt (str): The complete rendered conversation. + offsets (list[tuple[int, int]]): Character offsets for each token. + start (int): Inclusive character start. + stop (int): Exclusive character end. + name (str): Component name for validation errors. + allow_empty (bool): Whether an empty span is valid. + + Returns: + slice: The corresponding token range. + + Raises: + ValueError: If a required span has no tokens or a token crosses a content boundary. + """ + indices: list[int] = [] + for i, (token_start, token_stop) in enumerate(offsets): + if token_start >= token_stop or token_start >= stop or token_stop <= start: + continue + outside = prompt[token_start:start].strip() or prompt[stop:token_stop].strip() + if outside: + # Added role tokens can consume neighboring whitespace through lstrip/rstrip. + if prompt[max(start, token_start) : min(stop, token_stop)].strip(): + raise ValueError(f"GCG {name} token crosses a content boundary in the chat template.") + continue + indices.append(i) + if start == stop or not indices: + if not allow_empty: + raise ValueError(f"GCG {name} contains no tokens in the rendered prompt.") + boundary = next((i for i, (_, token_stop) in enumerate(offsets) if token_stop > start), len(offsets)) + return slice(boundary, boundary) + + return slice(indices[0], indices[-1] + 1) def _update_ids(self) -> None: - # Render the goal+control as the user turn and the target as the assistant turn using the - # tokenizer's built-in chat template. This replaces fastchat's per-model Conversation logic - # and works for any HuggingFace chat-tuned model (issue #965). + if not self.control.strip() or not self.target.strip(): + raise ValueError("GCG control and target must contain non-whitespace text.") + if not self.tokenizer.is_fast: + raise ValueError("GCG requires a fast tokenizer (use_fast=True) for character-to-token alignment.") messages = [ {"role": "user", "content": f"{self.goal} {self.control}"}, - {"role": "assistant", "content": f"{self.target}"}, + {"role": "assistant", "content": self.target}, ] prompt = self.tokenizer.apply_chat_template(messages, tokenize=False) - - encoding = self.tokenizer(prompt) + user_start, user_end = self._content_bounds(prompt=prompt, messages=messages, message_index=0) + assistant_start, assistant_end = self._content_bounds(prompt=prompt, messages=messages, message_index=1) + if not user_start <= user_end <= assistant_start <= assistant_end: + raise ValueError("Cannot safely locate user and assistant content in conversation order.") + + raw_user = messages[0]["content"] + rendered_user = prompt[user_start:user_end] + raw_leading = len(raw_user) - len(raw_user.lstrip()) + rendered_leading = len(rendered_user) - len(rendered_user.lstrip()) + user_origin = user_start + rendered_leading - raw_leading + goal_start = max(user_start, user_origin) + goal_end = max(goal_start, min(user_end, user_origin + len(self.goal))) + control_start = max(user_start, min(user_end, user_origin + len(self.goal) + 1)) + + # Templates already supply their special tokens. Offset spans handle both unmapped + # whitespace and multiple byte-level tokens sharing the same character position. + encoding = self.tokenizer(prompt, add_special_tokens=False, return_offsets_mapping=True) toks = encoding.input_ids - - # Locate goal/control/target substrings in the rendered prompt. Searching for each piece - # independently from the start takes the first occurrence anywhere, so a goal that quotes - # its own target (common with affirmative-prefix targets), or a piece that contains role - # text, silently produced slices pointing at the scaffolding or the wrong turn. Search - # within the measured content ranges instead: the control is the last occurrence in the - # user content (it ends that turn), the goal the last one before the control, and the - # target the first one in the assistant content. - user_start, user_end, assistant_start = self._content_bounds(prompt) - control_start = prompt.rfind(self.control, user_start, user_end) - goal_start = prompt.rfind(self.goal, user_start, control_start) if control_start != -1 else -1 - target_start = prompt.find(self.target, assistant_start) if goal_start != -1 else -1 - if target_start == -1: - raise ValueError( - "Could not locate goal/control/target in chat-templated prompt. " - f"prompt={prompt!r}, goal={self.goal!r}, " - f"control={self.control!r}, target={self.target!r}" - ) - - # ``char_to_token`` returns None when the character index has no - # corresponding token (e.g. when the substring ends exactly at the end - # of the prompt or lands on whitespace squashed into a neighbouring - # token). For end positions we clamp to ``len(toks)``; for start - # positions we walk forward to the next character that does map to a - # token. Both are necessary for the slice arithmetic to remain valid - # across tokenizers/templates. - def end_tok(char_pos: int) -> int: - tok: int | None = encoding.char_to_token(char_pos) - return len(toks) if tok is None else tok - - def start_tok(char_pos: int) -> int: - limit = len(prompt) - cur = char_pos - while cur < limit: - tok: int | None = encoding.char_to_token(cur) - if tok is not None: - return tok - cur += 1 - return len(toks) - - self._goal_slice = slice( - start_tok(goal_start), - end_tok(goal_start + len(self.goal)), + offsets = encoding["offset_mapping"] + goal_slice = self._token_slice( + prompt=prompt, + offsets=offsets, + start=goal_start, + stop=goal_end, + name="goal", + allow_empty=not self.goal.strip(), + ) + control_slice = self._token_slice( + prompt=prompt, offsets=offsets, start=control_start, stop=user_end, name="control" ) - self._control_slice = slice( - start_tok(control_start), - end_tok(control_start + len(self.control)), + target_slice = self._token_slice( + prompt=prompt, offsets=offsets, start=assistant_start, stop=assistant_end, name="target" ) - target_start_tok = start_tok(target_start) - target_end_tok = end_tok(target_start + len(self.target)) - self._target_slice = slice(target_start_tok, target_end_tok) - self._loss_slice = slice(target_start_tok - 1, target_end_tok - 1) - # Assistant role tokens are everything between the control end and the target start. - # This works for any chat template (e.g. llama-2 "[/INST]", phi-3 "<|assistant|>", etc.) - # without us needing to know the literal marker text. - self._assistant_role_slice = slice(self._control_slice.stop, self._target_slice.start) - - self.input_ids = torch.tensor(toks[: self._target_slice.stop], device="cpu") + if not self.goal.strip(): + goal_slice = slice(control_slice.start, control_slice.start) + if goal_slice.stop > control_slice.start or control_slice.stop > target_slice.start: + raise ValueError("GCG token slices overlap across a content boundary.") + + self._goal_slice = goal_slice + self._control_slice = control_slice + self._target_slice = target_slice + self._loss_slice = slice(target_slice.start - 1, target_slice.stop - 1) + self._assistant_role_slice = slice(control_slice.stop, target_slice.start) + self.input_ids = torch.tensor(toks[: target_slice.stop], device="cpu") @torch.no_grad() # type: ignore[misc, untyped-decorator, unused-ignore] def generate(self, model: Any, gen_config: Any = None) -> torch.Tensor: diff --git a/tests/unit/executor/promptgen/gcg/test_gcg_core.py b/tests/unit/executor/promptgen/gcg/test_gcg_core.py index ab7655a437..a291824f22 100644 --- a/tests/unit/executor/promptgen/gcg/test_gcg_core.py +++ b/tests/unit/executor/promptgen/gcg/test_gcg_core.py @@ -484,47 +484,6 @@ def test_raises_with_multiple_workers(self) -> None: ) -def _offset_tokenizer(prompt_text: str) -> Any: - """Build a mock tokenizer that renders ``prompt_text`` and maps characters to tokens. - - Each whitespace-delimited run of characters becomes one token, and ``char_to_token`` - reports the token containing a character, which is how a fast tokenizer behaves. This - keeps slice assertions meaningful without downloading a real tokenizer. - - Args: - prompt_text (str): The already-rendered chat prompt the tokenizer should return. - - Returns: - Any: A mock tokenizer suitable for constructing an AttackPrompt. - """ - spans: list[tuple[int, int]] = [] - start: int | None = None - for index, char in enumerate(prompt_text): - if char.isspace(): - if start is not None: - spans.append((start, index)) - start = None - elif start is None: - start = index - if start is not None: - spans.append((start, len(prompt_text))) - - def char_to_token(pos: int) -> int | None: - for token_index, (begin, end) in enumerate(spans): - if begin <= pos < end: - return token_index - return None - - encoding = MagicMock() - encoding.input_ids = list(range(len(spans))) - encoding.char_to_token.side_effect = char_to_token - - tokenizer = MagicMock() - tokenizer.apply_chat_template.return_value = prompt_text - tokenizer.return_value = encoding - return tokenizer - - _CHATML_TEMPLATE = ( "{% for m in messages %}<|{{ m['role'] }}|>{{ m['content'] }}<|end|>{% endfor %}" "{% if add_generation_prompt %}<|assistant|>{% endif %}" @@ -535,23 +494,54 @@ def char_to_token(pos: int) -> int | None: _NO_GENERATION_PROMPT_TEMPLATE = "{% for m in messages %}<|{{ m['role'] }}|>{{ m['content'] }}<|end|>{% endfor %}" -def _fast_tokenizer(chat_template: str) -> Any: +def _fast_tokenizer(*, chat_template: str = _CHATML_TEMPLATE, byte_level: bool = False) -> Any: """ - Build a real fast tokenizer offline that registers the chat role markers as special tokens. - - Words map to ``[UNK]``, which is enough here: only the character offsets matter for slicing. + Build an offline tokenizer with real offsets, decoding, and special role tokens. Args: chat_template (str): The Jinja chat template to render with. + byte_level (bool): Use byte-level BPE instead of whitespace-delimited words. Returns: Any: A ``PreTrainedTokenizerFast`` with ``chat_template`` set. """ - from tokenizers import Tokenizer, models, pre_tokenizers + from tokenizers import Tokenizer, decoders, models, pre_tokenizers from transformers import PreTrainedTokenizerFast - backend = Tokenizer(models.WordLevel({"[UNK]": 0}, unk_token="[UNK]")) - backend.pre_tokenizer = pre_tokenizers.Whitespace() + if byte_level: + vocabulary = ["[UNK]", *sorted(pre_tokenizers.ByteLevel.alphabet())] + backend = Tokenizer(models.BPE(dict(zip(vocabulary, range(len(vocabulary)), strict=True)), merges=[])) + backend.pre_tokenizer = pre_tokenizers.ByteLevel(add_prefix_space=False) + backend.decoder = decoders.ByteLevel() + else: + vocabulary = [ + "[UNK]", + "Say", + "it", + "assistant", + "done", + "!", + "Sure", + ",", + "here", + "is", + "the", + "plan", + "Respond", + "with", + "now", + "user", + "model", + "goal", + "control", + "target", + "hello", + "world", + ] + backend = Tokenizer( + models.WordLevel(dict(zip(vocabulary, range(len(vocabulary)), strict=True)), unk_token="[UNK]") + ) + backend.pre_tokenizer = pre_tokenizers.Whitespace() tokenizer = PreTrainedTokenizerFast(tokenizer_object=backend, unk_token="[UNK]") tokenizer.add_special_tokens( {"additional_special_tokens": ["<|user|>", "<|assistant|>", "<|end|>", "", ""]} @@ -561,135 +551,243 @@ def _fast_tokenizer(chat_template: str) -> Any: class TestUpdateIdsErrorPaths: - """Tests covering the error / fallback paths in AttackPrompt._update_ids.""" + """Real-tokenizer coverage of prompt boundaries and unsupported templates.""" - def test_raises_when_substring_not_in_rendered_prompt(self) -> None: - """If the chat template strips/transforms goal/control/target so they don't appear - verbatim in the rendered prompt, _update_ids must raise a clear ValueError.""" - tokenizer = MagicMock() - # Chat template that drops the user content entirely — goal/control won't appear in prompt - tokenizer.apply_chat_template.return_value = "[INST] [/INST] hello" - # tokenizer(...) returns an encoding-like object - encoding = MagicMock() - encoding.input_ids = [1, 2, 3, 4] - encoding.char_to_token.return_value = 1 - tokenizer.return_value = encoding - - with pytest.raises(ValueError, match="Could not locate goal/control/target"): + @pytest.mark.parametrize( + "chat_template", + [ + "<|user|>goal control<|end|><|assistant|>target<|end|>", + "{% for m in messages %}{{ m['content'] | upper }}{% endfor %}", + "{% for m in messages %}{{ m['content'] }}{{ m['content'] }}{% endfor %}", + "{% for m in messages | reverse %}{{ m['content'] }}{% endfor %}", + "{% for m in messages %}{{ m['content'] | replace('goal', 'other') }}{% endfor %}", + "{% for m in messages %}{% if 'goal' in m['content'] %}prefix{% endif %}{{ m['content'] }}{% endfor %}", + ], + ) + def test_unsupported_templates_raise_instead_of_scanning_the_prompt(self, chat_template: str) -> None: + with pytest.raises(ValueError, match="Cannot safely locate"): + AttackPrompt( + goal="goal", + target="target", + tokenizer=_fast_tokenizer(chat_template=chat_template), + control_init="control", + ) + + @pytest.mark.parametrize(("control", "target"), [("", "done"), ("! !", ""), (" ", "done"), ("! !", " ")]) + @pytest.mark.parametrize("byte_level", [False, True]) + def test_empty_or_unmapped_optimization_spans_raise(self, *, control: str, target: str, byte_level: bool) -> None: + with pytest.raises(ValueError, match="control|target"): AttackPrompt( - goal="this-goal-is-missing", - target="this-target-is-missing", - tokenizer=tokenizer, - control_init="this-control-is-missing", + goal="Say it", + target=target, + tokenizer=_fast_tokenizer(byte_level=byte_level), + control_init=control, ) - def test_start_tok_walks_forward_when_initial_position_has_no_token(self) -> None: - """char_to_token returns None for the start position (e.g., whitespace squashed - into the previous token); start_tok must walk forward to the next mappable - character. Slices should still be valid.""" - # Use a fully mocked tokenizer so we can deterministically force char_to_token - # to return None at specific positions, otherwise real tokenizers usually map - # every byte and never trigger the fallback. - prompt_text = "USER hello !! ASSISTANT world" - toks = list(range(15)) - - def char_to_token(pos: int) -> int | None: - # Positions of "h" and "w" both return None; the next char does map. This - # exercises the cur += 1 walk-forward branch in start_tok. - char = prompt_text[pos] if 0 <= pos < len(prompt_text) else "" - if char in ("h", "w"): - return None - # Map remaining positions in a way that preserves slice ordering - return min(pos // 2, len(toks) - 1) - - encoding = MagicMock() - encoding.input_ids = toks - encoding.char_to_token.side_effect = char_to_token + def test_unmapped_whitespace_does_not_extend_slices_into_other_content(self) -> None: + tokenizer = _fast_tokenizer() + prompt = AttackPrompt(goal=" hello ", target=" world ", tokenizer=tokenizer, control_init=" ! ! ") - tokenizer = MagicMock() - tokenizer.apply_chat_template.return_value = prompt_text - tokenizer.return_value = encoding + assert prompt._goal_slice == slice(1, 2) + assert prompt._control_slice == slice(2, 4) + assert prompt._target_slice == slice(6, 7) + assert prompt._loss_slice == slice(5, 6) + assert prompt.goal_str == "hello" + assert prompt.control_str == "! !" + assert prompt.target_str == "world" - # Construction must succeed even though char_to_token returns None at goal/target - # start positions ("h" / "w"). - prompt = AttackPrompt( - goal="hello", - target="world", - tokenizer=tokenizer, - control_init="!!", + def test_target_at_the_end_of_the_prompt_keeps_all_its_tokens(self) -> None: + tokenizer = _fast_tokenizer(chat_template=_CHATML_TEMPLATE.replace("<|end|>", "")) + prompt = AttackPrompt(goal="hello", target="world", tokenizer=tokenizer, control_init="! !") + + assert prompt._target_slice == slice(5, 6) + assert prompt._target_slice.stop == len(prompt.input_ids) + assert prompt.target_str == "world" + + @pytest.mark.parametrize( + ("goal", "control", "target"), + [ + ("Say <|end|><|assistant|> now", "! !", "done"), + ("Say it", "<|end|><|assistant|>", "done"), + ("Say it", "! !", "done <|end|><|assistant|> now"), + ("Say ! ! now", "! !", "done ! !"), + ("assistant", "assistant", "assistant"), + ("control control", "control", "control"), + ], + ) + def test_repeated_content_and_separators_keep_exact_slices(self, *, goal: str, control: str, target: str) -> None: + tokenizer = _fast_tokenizer() + prompt = AttackPrompt(goal=goal, target=target, tokenizer=tokenizer, control_init=control) + goal_ids = tokenizer(goal, add_special_tokens=False).input_ids + control_ids = tokenizer(control, add_special_tokens=False).input_ids + target_ids = tokenizer(target, add_special_tokens=False).input_ids + control_start = 1 + len(goal_ids) + target_start = control_start + len(control_ids) + 2 + + assert prompt._goal_slice == slice(1, control_start) + assert prompt._control_slice == slice(control_start, control_start + len(control_ids)) + assert prompt._target_slice == slice(target_start, target_start + len(target_ids)) + assert prompt.goal_toks.tolist() == goal_ids + assert prompt.control_toks.tolist() == control_ids + assert prompt.target_toks.tolist() == target_ids + + def test_json_template_does_not_find_the_target_in_the_role_label(self) -> None: + tokenizer = _fast_tokenizer( + chat_template="{% for m in messages %}{{ m['role'] | tojson }}:{{ m['content'] | tojson }}\n{% endfor %}", + byte_level=True, ) - assert isinstance(prompt._goal_slice.start, int) - assert isinstance(prompt._target_slice.start, int) - - def test_start_tok_returns_len_toks_when_no_position_maps(self) -> None: - """If char_to_token returns None for every position from char_pos to end-of-prompt, - start_tok must return len(toks) as a safe fallback (line 211).""" - prompt_text = "USER hello !! ASSISTANT world tail" - toks = list(range(20)) - - def char_to_token(pos: int) -> int | None: - char = prompt_text[pos] if 0 <= pos < len(prompt_text) else "" - # "tail" sits at end and never maps to a token (forces start_tok to exhaust - # the loop and hit `return len(toks)`); other content maps normally. - tail_start = prompt_text.find("tail") - if pos >= tail_start: - return None - return min(pos // 2, len(toks) - 1) - - encoding = MagicMock() - encoding.input_ids = toks - encoding.char_to_token.side_effect = char_to_token + prompt = AttackPrompt(goal="Say it", target="assistant", tokenizer=tokenizer, control_init="! !") - tokenizer = MagicMock() - tokenizer.apply_chat_template.return_value = prompt_text - tokenizer.return_value = encoding + assert prompt.target_str == "assistant" + assert prompt.input_str == '"user":"Say it ! !"\n"assistant":"assistant' + assert prompt._assistant_role_slice.stop == prompt._target_slice.start - # "tail" as the target — its start position and every position after it returns - # None, so start_tok exits the while loop and returns len(toks). - prompt = AttackPrompt( - goal="hello", - target="tail", - tokenizer=tokenizer, - control_init="!!", + def test_transformed_real_content_is_rejected_even_when_markers_round_trip(self) -> None: + tokenizer = _fast_tokenizer( + chat_template="{% for m in messages %}{{ m['role'] | tojson }}:{{ m['content'] | tojson }}\n{% endfor %}" ) - assert prompt._target_slice.start == len(toks) - - def test_end_tok_returns_len_toks_when_target_is_at_prompt_end(self) -> None: - """If the target sits at the very end of the rendered prompt, - char_to_token(end_pos) returns None — end_tok must clamp to len(toks) - (line 201 in attack_manager.py).""" - # Fully-mocked tokenizer so we can deterministically force char_to_token to - # return None at the position just past the target. Mirrors the pattern used - # by the two adjacent tests above. - prompt_text = "[INST] hello !! [/INST] world" - toks = list(range(10)) - target_end_pos = len(prompt_text) # one past the final char of "world" - - def char_to_token(pos: int) -> int | None: - # Position at/after end-of-prompt has no token → triggers the - # `return len(toks)` fallback in end_tok. - if pos >= target_end_pos: - return None - # Everything else maps to a valid token index that preserves ordering. - return min(pos // 3, len(toks) - 1) - - encoding = MagicMock() - encoding.input_ids = toks - encoding.char_to_token.side_effect = char_to_token + with pytest.raises(ValueError, match="Cannot safely locate"): + AttackPrompt(goal='Say "hello"', target="done", tokenizer=tokenizer, control_init="! !") + + @pytest.mark.parametrize("target", ["caf\u00e9", "\U0001f600", "done"]) + def test_byte_level_offsets_keep_every_target_byte(self, target: str) -> None: + tokenizer = _fast_tokenizer(byte_level=True) + prompt = AttackPrompt(goal="Say it", target=target, tokenizer=tokenizer, control_init="! !") + + assert prompt.target_toks.tolist() == tokenizer(target, add_special_tokens=False).input_ids + assert prompt.target_str == target + assert prompt.input_str == f"<|user|>Say it ! !<|end|><|assistant|>{target}" + assert prompt._loss_slice == slice(prompt._target_slice.start - 1, prompt._target_slice.stop - 1) + + def test_chat_template_special_tokens_are_not_added_twice(self) -> None: + from tokenizers import processors + + tokenizer = _fast_tokenizer(chat_template="{{ bos_token }}" + _CHATML_TEMPLATE) + tokenizer.add_special_tokens({"bos_token": ""}) + tokenizer.backend_tokenizer.post_processor = processors.TemplateProcessing( + single=" $A", special_tokens=[("", tokenizer.bos_token_id)] + ) + prompt = AttackPrompt(goal="Say it", target="done", tokenizer=tokenizer, control_init="! !") - tokenizer = MagicMock() - tokenizer.apply_chat_template.return_value = prompt_text - tokenizer.return_value = encoding + assert prompt.input_ids.tolist().count(tokenizer.bos_token_id) == 1 + assert prompt._goal_slice == slice(2, 4) + assert prompt._target_slice == slice(8, 9) - prompt = AttackPrompt( - goal="hello", - target="world", # sits at end of prompt_text; target end has no token - tokenizer=tokenizer, - control_init="!!", + def test_token_shared_with_assistant_scaffolding_is_rejected(self) -> None: + tokenizer = _fast_tokenizer(chat_template="{{ messages[0]['content'] }} week{{ messages[1]['content'] }}") + with pytest.raises(ValueError, match="token.*boundary"): + AttackPrompt(goal="Say it", target="end", tokenizer=tokenizer, control_init="! !") + + def test_slow_tokenizer_has_an_actionable_error(self) -> None: + from transformers import PreTrainedTokenizer + + class SlowTokenizer(PreTrainedTokenizer): + def get_vocab(self) -> dict[str, int]: + return {"[UNK]": 0} + + def _tokenize(self, text: str, **kwargs: Any) -> list[str]: + return text.split() + + def _convert_token_to_id(self, token: str) -> int: + return 0 + + def _convert_id_to_token(self, index: int) -> str: + return "[UNK]" + + tokenizer = SlowTokenizer(unk_token="[UNK]", chat_template=_CHATML_TEMPLATE) + assert not tokenizer.is_fast + with pytest.raises(ValueError, match="fast tokenizer.*use_fast=True"): + AttackPrompt(goal="Say it", target="done", tokenizer=tokenizer, control_init="! !") + + def test_probe_rendering_errors_are_not_swallowed(self) -> None: + tokenizer = _fast_tokenizer() + with ( + patch.object( + tokenizer, + "apply_chat_template", + side_effect=["<|user|>Say it ! !<|end|><|assistant|>done<|end|>", RuntimeError("probe failed")], + ), + pytest.raises(RuntimeError, match="probe failed"), + ): + AttackPrompt(goal="Say it", target="done", tokenizer=tokenizer, control_init="! !") + + @pytest.mark.parametrize("byte_level", [False, True]) + @pytest.mark.parametrize("goal", ["", " \t", " hello ", "hello"]) + @pytest.mark.parametrize("trim", [False, True]) + def test_whitespace_and_empty_goals_preserve_content(self, *, byte_level: bool, goal: str, trim: bool) -> None: + template = _CHATML_TEMPLATE.replace("m['content']", "m['content'] | trim") if trim else _CHATML_TEMPLATE + tokenizer = _fast_tokenizer(chat_template=template, byte_level=byte_level) + prompt = AttackPrompt(goal=goal, target=" world ", tokenizer=tokenizer, control_init=" ! ! ") + + assert prompt.goal_str == goal.strip() + assert prompt.control_str == "! !" + assert prompt.target_str == "world" + assert prompt._goal_slice.stop <= prompt._control_slice.start < prompt._control_slice.stop + assert prompt._control_slice.stop <= prompt._target_slice.start < prompt._target_slice.stop + + @pytest.mark.parametrize("byte_level", [False, True]) + def test_control_and_content_setters_recompute_boundaries(self, byte_level: bool) -> None: + tokenizer = _fast_tokenizer(byte_level=byte_level) + prompt = AttackPrompt(goal="Say it", target="done", tokenizer=tokenizer, control_init="! !") + + for control in ("assistant", "<|end|><|assistant|>", "! !"): + prompt.control_str = control + assert prompt.control_toks.tolist() == tokenizer(control, add_special_tokens=False).input_ids + assert prompt.goal_str == "Say it" + assert prompt.target_str == "done" + + control_ids = tokenizer("assistant", add_special_tokens=False).input_ids + prompt.control_toks = torch.tensor(control_ids) + prompt.goal_str = "" + prompt.target_str = "assistant" + assert prompt.goal_toks.numel() == 0 + assert prompt.control_toks.tolist() == control_ids + assert prompt.target_toks.tolist() == control_ids + assert prompt._control_slice.stop <= prompt._target_slice.start + + def test_template_without_a_separator_uses_content_positions(self) -> None: + tokenizer = _fast_tokenizer(chat_template="{{ messages[0]['content'] }}{{ messages[1]['content'] }}") + prompt = AttackPrompt(goal="Say it", target="done", tokenizer=tokenizer, control_init="!") + + assert prompt._control_slice == slice(2, 3) + assert prompt._target_slice == slice(3, 4) + assert prompt._assistant_role_slice == slice(3, 3) + + @pytest.mark.parametrize("byte_level", [False, True]) + @pytest.mark.parametrize("goal", ["", " ", " hello "]) + def test_role_tokens_that_consume_whitespace_stay_outside_content_slices( + self, *, byte_level: bool, goal: str + ) -> None: + from tokenizers import AddedToken + + tokenizer = _fast_tokenizer(byte_level=byte_level) + tokenizer.add_special_tokens( + { + "additional_special_tokens": [ + AddedToken(marker, lstrip=True, rstrip=True) for marker in ("<|user|>", "<|assistant|>", "<|end|>") + ] + } ) - # end_tok(target_end_pos) saw None from char_to_token → clamped to len(toks). - assert prompt._target_slice.stop == len(toks) - assert prompt._target_slice.stop > prompt._target_slice.start + prompt = AttackPrompt(goal=goal, target=" world ", tokenizer=tokenizer, control_init=" ! ! ") + role_ids = {tokenizer.convert_tokens_to_ids(marker) for marker in ("<|user|>", "<|assistant|>", "<|end|>")} + + assert prompt.goal_str == goal.strip() + assert prompt.control_str == "! !" + assert prompt.target_str == "world" + assert role_ids.isdisjoint(prompt.goal_toks.tolist()) + assert role_ids.isdisjoint(prompt.control_toks.tolist()) + assert role_ids.isdisjoint(prompt.target_toks.tolist()) + if not goal.strip(): + assert prompt._goal_slice == slice(prompt._control_slice.start, prompt._control_slice.start) + + @pytest.mark.parametrize("component", ["control", "target"]) + def test_nonempty_text_removed_by_tokenizer_normalization_raises(self, component: str) -> None: + from tokenizers import normalizers + + tokenizer = _fast_tokenizer() + tokenizer.backend_tokenizer.normalizer = normalizers.Replace(component, "") + with pytest.raises(ValueError, match=f"{component} contains no tokens"): + AttackPrompt(goal="goal", target="target", tokenizer=tokenizer, control_init="control") def test_target_is_located_after_the_user_turn_when_the_goal_quotes_it(self) -> None: """A goal that quotes its own target must not pull the target slice into the user turn. @@ -705,12 +803,15 @@ def test_target_is_located_after_the_user_turn_when_the_goal_quotes_it(self) -> prompt = AttackPrompt( goal=goal, target=target, - tokenizer=_fast_tokenizer(_CHATML_TEMPLATE), + tokenizer=_fast_tokenizer(), control_init=control, ) - assert prompt._target_slice.start >= prompt._control_slice.stop - assert prompt._assistant_role_slice.start <= prompt._assistant_role_slice.stop + assert prompt._control_slice == slice(9, 13) + assert prompt._assistant_role_slice == slice(13, 15) + assert prompt._target_slice == slice(15, 21) + assert prompt._loss_slice == slice(14, 20) + assert prompt.target_toks.tolist() == prompt.tokenizer(target, add_special_tokens=False).input_ids def test_target_that_names_the_assistant_role_marker_is_found_in_the_reply(self) -> None: """A target like "assistant" also matches inside ``<|assistant|>``, which is one special token. @@ -718,7 +819,7 @@ def test_target_that_names_the_assistant_role_marker_is_found_in_the_reply(self) Searching right after the user content lands on the role marker and leaves an empty target slice, so the search has to start where the assistant content does. """ - tokenizer = _fast_tokenizer(_CHATML_TEMPLATE) + tokenizer = _fast_tokenizer() prompt = AttackPrompt(goal="Say it", target="assistant", tokenizer=tokenizer, control_init="! !") @@ -735,7 +836,7 @@ def test_control_that_collides_with_the_role_marker_keeps_its_slice(self) -> Non A control holding "assistant" also matches inside ``<|assistant|>``; bounding the search to the user content is what keeps the control slice pointing at the suffix GCG optimizes. """ - tokenizer = _fast_tokenizer(_CHATML_TEMPLATE) + tokenizer = _fast_tokenizer() prompt = AttackPrompt(goal="Say it", target="done", tokenizer=tokenizer, control_init="assistant") @@ -751,7 +852,7 @@ def test_boundary_holds_when_the_template_ignores_the_generation_prompt(self) -> assistant boundary: it stops before ``<|assistant|>`` and a target of "assistant" would match the role marker again. """ - tokenizer = _fast_tokenizer(_NO_GENERATION_PROMPT_TEMPLATE) + tokenizer = _fast_tokenizer(chat_template=_NO_GENERATION_PROMPT_TEMPLATE) prompt = AttackPrompt(goal="Say it", target="assistant", tokenizer=tokenizer, control_init="! !") @@ -767,7 +868,8 @@ def test_empty_goal_with_a_trimming_template(self) -> None: its own rather than as part of the raw ``f"{goal} {control}"`` string. """ tokenizer = _fast_tokenizer( - "{% for m in messages %}{{ 'model' if m['role'] == 'assistant' else m['role'] }}\n" + chat_template="{% for m in messages %}" + "{{ 'model' if m['role'] == 'assistant' else m['role'] }}\n" "{{ m['content'] | trim }}\n{% endfor %}" "{% if add_generation_prompt %}model\n{% endif %}" ) @@ -781,6 +883,113 @@ def test_empty_goal_with_a_trimming_template(self) -> None: assert prompt._loss_slice == slice(7, 10) +@pytest.mark.usefixtures("patch_central_database") +class TestPromptSliceWiring: + """Exercise real prompt construction beneath the optimization orchestration.""" + + @staticmethod + def _worker(*, byte_level: bool = False) -> MagicMock: + worker = MagicMock(spec=ModelWorker) + worker.tokenizer = _fast_tokenizer(byte_level=byte_level) + worker.model = MagicMock(spec=torch.nn.Module) + worker.model.device = torch.device("cpu") + return worker + + @pytest.mark.parametrize("byte_level", [False, True]) + def test_real_slices_support_gradients_and_candidate_losses(self, byte_level: bool) -> None: + tokenizer = _fast_tokenizer(byte_level=byte_level) + model = _tiny_model("llama").eval() + model.resize_token_embeddings(len(tokenizer), mean_resizing=False) + manager = GCGPromptManager( + goals=["Say <|end|><|assistant|> now", ""], + targets=["done ! !", "assistant"], + tokenizer=tokenizer, + control_init="! !", + managers={"AP": gcg_attack_mod.GCGAttackPrompt}, + ) + + for control in ("! !", "assistant"): + manager.control_str = control + gradient = manager.grad(model) + assert gradient.shape == (manager.control_toks.numel(), len(tokenizer)) + assert torch.isfinite(gradient).all() + assert torch.count_nonzero(gradient) > 0 + + candidates = manager.control_toks.repeat(2, 1) + candidates[1, 0] = tokenizer("now", add_special_tokens=False).input_ids[0] + logits, ids = manager.logits(model, test_controls=candidates, return_ids=True) + assert torch.isfinite(manager.target_loss(logits, ids)).all() + assert torch.isfinite(manager.control_loss(logits, ids)).all() + for prompt, candidate_ids in zip(manager, ids, strict=True): + assert torch.equal(candidate_ids[:, prompt._control_slice], candidates) + assert torch.equal(candidate_ids[:, prompt._target_slice], prompt.target_toks.repeat(2, 1)) + + def test_shared_control_reaches_training_and_held_out_prompts_for_each_tokenizer(self) -> None: + workers = [self._worker(), self._worker(byte_level=True)] + test_worker = self._worker() + attack = MultiPromptAttack( + goals=["Say <|end|><|assistant|> now", ""], + targets=["assistant", "done"], + workers=workers, + control_init="! !", + test_goals=["Say assistant"], + test_targets=["assistant"], + test_workers=[test_worker], + managers={"AP": AttackPrompt, "PM": PromptManager}, + ) + attack.control_str = "assistant" + + with patch.object(attack, "test", return_value=([], [], [])) as evaluate: + attack.test_all() + + all_workers, all_prompts = evaluate.call_args.args + assert all_workers == [*workers, test_worker] + assert evaluate.call_args.kwargs["include_loss"] is True + for manager in [*attack.prompts, *all_prompts]: + for prompt in manager: + tokenizer = prompt.tokenizer + assert prompt.goal_toks.tolist() == tokenizer(prompt.goal, add_special_tokens=False).input_ids + assert prompt.control_toks.tolist() == tokenizer("assistant", add_special_tokens=False).input_ids + assert prompt.target_toks.tolist() == tokenizer(prompt.target, add_special_tokens=False).input_ids + + @pytest.mark.parametrize( + ("progressive", "expected_rounds"), + [ + (False, [(1, 2), (1, 2)]), + (True, [(1, 1), (2, 1), (2, 2)]), + ], + ) + def test_individual_and_progressive_rounds_construct_real_prompts( + self, *, progressive: bool, expected_rounds: list[tuple[int, int]] + ) -> None: + rounds: list[tuple[int, int]] = [] + + def run_inner(attack: MultiPromptAttack, **kwargs: Any) -> tuple[str, float, int]: + rounds.append((len(attack.goals), len(attack.workers))) + attack.control_str = "assistant" + for manager in attack.prompts: + for prompt in manager: + assert prompt.control_str == "assistant" + assert ( + prompt.target_toks.tolist() + == prompt.tokenizer(prompt.target, add_special_tokens=False).input_ids + ) + return "assistant", 0.5, 1 + + attack_class = ProgressiveMultiPromptAttack if progressive else IndividualPromptAttack + attack = attack_class( + goals=["Say <|end|><|assistant|> now", ""], + targets=["assistant", "done"], + workers=[self._worker(), self._worker(byte_level=True)], + control_init="! !", + managers={"AP": AttackPrompt, "PM": PromptManager, "MPA": MultiPromptAttack}, + ) + with patch.object(MultiPromptAttack, "run", autospec=True, side_effect=run_inner): + attack.run(n_steps=3, stop_on_success=False, incr_control=False, verbose=False) + + assert rounds == expected_rounds + + class TestGetWorkersChatTemplateValidation: """Tests for the chat-template precondition in get_workers.""" From 7a344a2960bace30bb9b2fdc192d7ce1a2789798 Mon Sep 17 00:00:00 2001 From: devangpratap <115096812+devangpratap@users.noreply.github.com> Date: Fri, 25 Sep 2026 13:04:24 -0400 Subject: [PATCH 5/5] FIX measure the GCG turn boundary instead of searching for the separator The separator between the user and assistant contents was found with a first-occurrence search, so a goal quoting the template's own markers moved the boundary. Render the real user content with a sentinel reply and read the boundary off its position instead. When the template does not render contents verbatim, raise rather than falling back to an unbounded search, which silently corrupted the target and loss slices. --- .../gcg/attack/base/attack_manager.py | 87 +++++++++------- .../executor/promptgen/gcg/test_gcg_core.py | 98 ++++++++++++------- 2 files changed, 110 insertions(+), 75 deletions(-) diff --git a/pyrit/executor/promptgen/gcg/attack/base/attack_manager.py b/pyrit/executor/promptgen/gcg/attack/base/attack_manager.py index 99acea8524..dfd4f51b9b 100644 --- a/pyrit/executor/promptgen/gcg/attack/base/attack_manager.py +++ b/pyrit/executor/promptgen/gcg/attack/base/attack_manager.py @@ -347,58 +347,71 @@ def __init__( self._update_ids() - def _content_bounds(self, prompt: str) -> tuple[int, int, int]: + def _content_bounds(self, user_content: str, prompt: str) -> tuple[int, int, int]: """ Locate where the user content and the assistant content sit in the rendered prompt. - Rendering the same two messages with sentinel contents shows the template's own - scaffolding: everything before the user content, and everything between the two contents. - Measuring it this way keeps the role markers out of the search ranges, so a goal, control - or target that happens to contain role text (``assistant`` is an ordinary word, and an - optimized control is decoded vocabulary tokens) cannot match the scaffolding instead of - the turn. Templates that ignore ``add_generation_prompt`` are handled too, since nothing - here assumes the generation prompt renders the assistant preamble. + The boundaries are measured by rendering rather than searched for in the prompt. Rendering the real user + content with a sentinel assistant reply puts the sentinel exactly where the assistant + content starts, so a goal that quotes the template's own markers (red-team objectives + can) cannot move the boundary. Rendering two sentinels gives the template's scaffolding: + the prefix before the user content and the separator between the two contents. That keeps + the role markers out of the search ranges, so a goal, control or target containing role + text (``assistant`` is an ordinary word) cannot match the scaffolding instead of the turn. Args: + user_content (str): The user content rendered into ``prompt``. prompt (str): The rendered prompt for the current goal, control and target. Returns: - tuple[int, int, int]: Start of the user content, its end, and the assistant content - start. Falls back to the whole prompt for a template whose scaffolding cannot be - measured, which leaves the searches no worse off than an unbounded scan. + tuple[int, int, int]: Start of the user content, its end, and the assistant content start. + + Raises: + ValueError: If the template transforms the sentinels or the measured scaffolding does + not line up with ``prompt``. Guessing would silently corrupt the target and loss slices. """ - user_sentinel, assistant_sentinel = "\x00pyrit-user\x00", "\x00pyrit-assistant\x00" - try: - scaffold = self.tokenizer.apply_chat_template( - [ - {"role": "user", "content": user_sentinel}, - {"role": "assistant", "content": assistant_sentinel}, - ], - tokenize=False, + # Plain lowercase letters survive escaping filters such as ``tojson``. + user_sentinel, assistant_sentinel = "pyritusercontent", "pyritassistantcontent" + + def render(user: str, assistant: str) -> str: + return str( + self.tokenizer.apply_chat_template( + [{"role": "user", "content": user}, {"role": "assistant", "content": assistant}], + tokenize=False, + ) + ) + + scaffold = render(user_sentinel, assistant_sentinel) + probe = render(user_content, assistant_sentinel) + # The last user sentinel before the reply, in case the template also echoes the content earlier. + user_start = scaffold.rfind(user_sentinel, 0, scaffold.find(assistant_sentinel)) + separator_at = user_start + len(user_sentinel) + separator = scaffold[separator_at : scaffold.find(assistant_sentinel, separator_at)] + # The sentinel is the last content in the probe, so the last occurrence is the real one. + assistant_start = probe.rfind(assistant_sentinel) + user_end = assistant_start - len(separator) + if ( + user_start == -1 + or assistant_sentinel not in scaffold + or assistant_start == -1 + or user_end < user_start + or probe[:user_start] != scaffold[:user_start] + or probe[user_end:assistant_start] != separator + or not prompt.startswith(probe[:assistant_start]) + ): + raise ValueError( + "Could not locate the user and assistant contents in the chat template. " + f"The template must render message contents verbatim. prompt={prompt!r}" ) - user_at = scaffold.index(user_sentinel) - assistant_at = scaffold.index(assistant_sentinel, user_at) - except Exception: # pylint: disable=broad-except - return 0, len(prompt), 0 - - prefix = scaffold[:user_at] - # The text the template puts between the user content and the assistant content. - separator = scaffold[user_at + len(user_sentinel) : assistant_at] - if not separator or not prompt.startswith(prefix): - return 0, len(prompt), 0 - - user_start = len(prefix) - separator_at = prompt.find(separator, user_start) - if separator_at == -1: - return user_start, len(prompt), user_start - return user_start, separator_at, separator_at + len(separator) + return user_start, user_end, assistant_start def _update_ids(self) -> None: # Render the goal+control as the user turn and the target as the assistant turn using the # tokenizer's built-in chat template. This replaces fastchat's per-model Conversation logic # and works for any HuggingFace chat-tuned model (issue #965). + user_content = f"{self.goal} {self.control}" messages = [ - {"role": "user", "content": f"{self.goal} {self.control}"}, + {"role": "user", "content": user_content}, {"role": "assistant", "content": f"{self.target}"}, ] prompt = self.tokenizer.apply_chat_template(messages, tokenize=False) @@ -413,7 +426,7 @@ def _update_ids(self) -> None: # within the measured content ranges instead: the control is the last occurrence in the # user content (it ends that turn), the goal the last one before the control, and the # target the first one in the assistant content. - user_start, user_end, assistant_start = self._content_bounds(prompt) + user_start, user_end, assistant_start = self._content_bounds(user_content, prompt) control_start = prompt.rfind(self.control, user_start, user_end) goal_start = prompt.rfind(self.goal, user_start, control_start) if control_start != -1 else -1 target_start = prompt.find(self.target, assistant_start) if goal_start != -1 else -1 diff --git a/tests/unit/executor/promptgen/gcg/test_gcg_core.py b/tests/unit/executor/promptgen/gcg/test_gcg_core.py index ab7655a437..4774d2eceb 100644 --- a/tests/unit/executor/promptgen/gcg/test_gcg_core.py +++ b/tests/unit/executor/promptgen/gcg/test_gcg_core.py @@ -484,45 +484,17 @@ def test_raises_with_multiple_workers(self) -> None: ) -def _offset_tokenizer(prompt_text: str) -> Any: - """Build a mock tokenizer that renders ``prompt_text`` and maps characters to tokens. - - Each whitespace-delimited run of characters becomes one token, and ``char_to_token`` - reports the token containing a character, which is how a fast tokenizer behaves. This - keeps slice assertions meaningful without downloading a real tokenizer. +def _chat_template(fmt: str) -> Any: + """ + Build an ``apply_chat_template`` stand-in that formats the user and assistant contents into ``fmt``. Args: - prompt_text (str): The already-rendered chat prompt the tokenizer should return. + fmt (str): A format string with one ``{}`` for the user content and one for the assistant content. Returns: - Any: A mock tokenizer suitable for constructing an AttackPrompt. + Any: A callable usable as ``apply_chat_template.side_effect``. """ - spans: list[tuple[int, int]] = [] - start: int | None = None - for index, char in enumerate(prompt_text): - if char.isspace(): - if start is not None: - spans.append((start, index)) - start = None - elif start is None: - start = index - if start is not None: - spans.append((start, len(prompt_text))) - - def char_to_token(pos: int) -> int | None: - for token_index, (begin, end) in enumerate(spans): - if begin <= pos < end: - return token_index - return None - - encoding = MagicMock() - encoding.input_ids = list(range(len(spans))) - encoding.char_to_token.side_effect = char_to_token - - tokenizer = MagicMock() - tokenizer.apply_chat_template.return_value = prompt_text - tokenizer.return_value = encoding - return tokenizer + return lambda messages, **_kwargs: fmt.format(*(m["content"] for m in messages)) _CHATML_TEMPLATE = ( @@ -575,7 +547,7 @@ def test_raises_when_substring_not_in_rendered_prompt(self) -> None: encoding.char_to_token.return_value = 1 tokenizer.return_value = encoding - with pytest.raises(ValueError, match="Could not locate goal/control/target"): + with pytest.raises(ValueError, match="Could not locate the user and assistant contents"): AttackPrompt( goal="this-goal-is-missing", target="this-target-is-missing", @@ -607,7 +579,7 @@ def char_to_token(pos: int) -> int | None: encoding.char_to_token.side_effect = char_to_token tokenizer = MagicMock() - tokenizer.apply_chat_template.return_value = prompt_text + tokenizer.apply_chat_template.side_effect = _chat_template("USER {} ASSISTANT {}") tokenizer.return_value = encoding # Construction must succeed even though char_to_token returns None at goal/target @@ -641,7 +613,7 @@ def char_to_token(pos: int) -> int | None: encoding.char_to_token.side_effect = char_to_token tokenizer = MagicMock() - tokenizer.apply_chat_template.return_value = prompt_text + tokenizer.apply_chat_template.side_effect = _chat_template("USER {} ASSISTANT world {}") tokenizer.return_value = encoding # "tail" as the target — its start position and every position after it returns @@ -678,7 +650,7 @@ def char_to_token(pos: int) -> int | None: encoding.char_to_token.side_effect = char_to_token tokenizer = MagicMock() - tokenizer.apply_chat_template.return_value = prompt_text + tokenizer.apply_chat_template.side_effect = _chat_template("[INST] {} [/INST] {}") tokenizer.return_value = encoding prompt = AttackPrompt( @@ -780,6 +752,56 @@ def test_empty_goal_with_a_trimming_template(self) -> None: assert prompt._target_slice == slice(8, 11) assert prompt._loss_slice == slice(7, 10) + def test_goal_that_quotes_the_turn_separator_keeps_the_boundary(self) -> None: + """Red-team goals can quote model control tokens, including the template's own turn separator. + + Searching for the separator would find it inside the goal and end the user content before the + control, so the boundary has to be measured from the template instead. + """ + tokenizer = _fast_tokenizer(_CHATML_TEMPLATE) + + prompt = AttackPrompt( + goal="Say <|end|><|assistant|> now", target="done", tokenizer=tokenizer, control_init="! !" + ) + + # <|user|> Say <|end|> <|assistant|> now ! ! <|end|> <|assistant|> done <|end|> + assert prompt._control_slice == slice(5, 7) + assert prompt._target_slice == slice(9, 10) + assert prompt._loss_slice == slice(8, 9) + + def test_escaping_template_locates_the_target_in_the_reply(self) -> None: + """A template may escape the contents, e.g. with ``tojson``; the boundaries still have to hold. + + With an unbounded search the target "assistant" matches the role label instead of the reply. + """ + tokenizer = _fast_tokenizer( + "{% for m in messages %}\"{{ m['role'] }}\":{{ m['content'] | tojson }}\n{% endfor %}" + ) + + prompt = AttackPrompt(goal="Say it", target="assistant", tokenizer=tokenizer, control_init="! !") + + # " user ":" Say it ! !" " assistant ":" assistant " + assert prompt._target_slice == slice(10, 11) + assert prompt._loss_slice == slice(9, 10) + + def test_raises_when_an_escaped_goal_is_not_rendered_verbatim(self) -> None: + """The turns can be measured, but ``tojson`` escapes the quotes, so the goal itself is not in the prompt.""" + tokenizer = _fast_tokenizer( + "{% for m in messages %}\"{{ m['role'] }}\":{{ m['content'] | tojson }}\n{% endfor %}" + ) + + with pytest.raises(ValueError, match="Could not locate goal/control/target"): + AttackPrompt(goal='Say "it"', target="done", tokenizer=tokenizer, control_init="! !") + + def test_raises_when_the_template_transforms_the_contents(self) -> None: + """A template that rewrites the contents leaves no way to measure the turns, so construction fails closed.""" + tokenizer = _fast_tokenizer( + "{% for m in messages %}<|{{ m['role'] }}|>{{ m['content'] | upper }}<|end|>{% endfor %}" + ) + + with pytest.raises(ValueError, match="Could not locate the user and assistant contents"): + AttackPrompt(goal="Say it", target="done", tokenizer=tokenizer, control_init="! !") + class TestGetWorkersChatTemplateValidation: """Tests for the chat-template precondition in get_workers."""