Skip to content

Commit c0ec411

Browse files
committed
FIX: keep template conditions until their parameters are rendered
Trusted templates are rendered once at load with only the path variables. The placeholder for a missing parameter is truthy, so an `{% if conversation_context %}` block was decided at load and its tags were dropped; a later render with `conversation_context=None` could not remove it. Crescendo and TAP without a prepended conversation sent the adversarial chat a "prior conversation" section that read "None". `render_template_value_silent` already kept a template unrendered when a `for` collection was missing. It now does the same for any parameter read by an `if`/`elif` test, an inline `x if test else y`, or a loop filter, found from the parsed template instead of a regular expression.
1 parent f916d8e commit c0ec411

3 files changed

Lines changed: 130 additions & 9 deletions

File tree

‎pyrit/models/seeds/seed.py‎

Lines changed: 56 additions & 9 deletions
Original file line numberDiff line numberDiff line change
@@ -9,6 +9,7 @@
99

1010
from __future__ import annotations
1111

12+
import functools
1213
import logging
1314
import re
1415
import uuid
@@ -23,7 +24,7 @@
2324
from pyrit.models.seeds.seed_origin import SeedOrigin
2425

2526
if TYPE_CHECKING:
26-
from collections.abc import Iterator
27+
from collections.abc import Callable, Iterator
2728
from pathlib import Path
2829

2930
logger = logging.getLogger(__name__)
@@ -81,25 +82,66 @@ def __repr__(self) -> str:
8182
"""
8283
return f"{{{{ {self._undefined_name} }}}}" if self._undefined_name else ""
8384

85+
# A placeholder cannot decide a branch or a loop: answering now would drop the
86+
# {% if %} or {% for %} tags, and the later render could not decide again.
8487
def __iter__(self) -> Iterator[object]:
8588
"""
86-
Return an empty iterator to prevent iteration over undefined variables.
89+
Defer rendering instead of iterating over an unresolved variable.
8790
88-
Returns:
89-
Iterator[object]: Empty iterator.
91+
Raises:
92+
_DeferRenderError: Always.
9093
9194
"""
92-
return iter([])
95+
raise _DeferRenderError(self._undefined_name)
9396

9497
def __bool__(self) -> bool:
9598
"""
96-
Evaluate as truthy to avoid falsey-branch side effects.
99+
Defer rendering instead of testing an unresolved variable.
97100
98-
Returns:
99-
bool: Always True.
101+
Raises:
102+
_DeferRenderError: Always.
103+
104+
"""
105+
raise _DeferRenderError(self._undefined_name)
100106

107+
def __eq__(self, other: object) -> bool:
101108
"""
102-
return True # Ensures it doesn't evaluate to False
109+
Defer rendering instead of comparing an unresolved variable.
110+
111+
Raises:
112+
_DeferRenderError: Always.
113+
114+
"""
115+
raise _DeferRenderError(self._undefined_name)
116+
117+
def __ne__(self, other: object) -> bool:
118+
"""
119+
Defer rendering instead of comparing an unresolved variable.
120+
121+
Raises:
122+
_DeferRenderError: Always.
123+
124+
"""
125+
raise _DeferRenderError(self._undefined_name)
126+
127+
__hash__ = Undefined.__hash__
128+
129+
130+
class _DeferRenderError(Exception):
131+
"""Raised when an unresolved variable would decide a branch or a loop."""
132+
133+
134+
def _deferring(function: Callable[..., Any]) -> Callable[..., Any]:
135+
# Jinja tests such as `is defined` and the `default` filter check the value's type, not its truth.
136+
# functools.wraps keeps Jinja's pass_environment marker, so the value may not be the first argument.
137+
@functools.wraps(function)
138+
def deferring(*args: Any, **kwargs: Any) -> Any:
139+
for arg in args:
140+
if isinstance(arg, PartialUndefined):
141+
raise _DeferRenderError(arg._undefined_name)
142+
return function(*args, **kwargs)
143+
144+
return deferring
103145

104146

105147
class Seed(BaseModel):
@@ -223,11 +265,16 @@ def render_template_value_silent(self, **kwargs: Any) -> str:
223265

224266
# Create a Jinja template with PartialUndefined placeholders
225267
env = SandboxedEnvironment(undefined=PartialUndefined)
268+
env.tests = {name: _deferring(test) for name, test in env.tests.items()}
269+
env.filters["default"] = env.filters["d"] = _deferring(env.filters["default"])
226270
is_jinja_template = env.from_string(self.value)
227271

228272
try:
229273
# Render the template with the provided kwargs
230274
return is_jinja_template.render(**kwargs)
275+
except _DeferRenderError:
276+
# A missing parameter decides a branch or a loop - preserve the template as-is
277+
return self.value
231278
except Exception as e:
232279
logger.error("Error rendering template: %s", e)
233280
return self.value

‎tests/unit/executor/attack/multi_turn/test_crescendo.py‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -710,6 +710,7 @@ async def test_setup_sets_adversarial_chat_system_prompt(
710710
call_args = mock_adversarial_chat.set_system_prompt_async.call_args
711711
assert "Test objective" in call_args.kwargs["system_prompt"]
712712
assert "15" in call_args.kwargs["system_prompt"] # Check for the max_turns value
713+
assert "Prior Conversation Context" not in call_args.kwargs["system_prompt"]
713714
assert call_args.kwargs["conversation_id"] == basic_context.session.adversarial_chat_conversation_id
714715

715716
async def test_setup_handles_prepended_conversation_with_refusal(

‎tests/unit/models/test_seed.py‎

Lines changed: 73 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -9,6 +9,9 @@
99

1010
import numpy as np
1111
import pytest
12+
import yaml
13+
from jinja2 import StrictUndefined
14+
from jinja2.sandbox import SandboxedEnvironment
1215
from PIL import Image
1316
from scipy.io import wavfile
1417

@@ -204,6 +207,76 @@ def test_render_template_value_silent_blocks_ssti_via_endraw_injection():
204207
assert "__class__" not in result or result == raw_wrapped
205208

206209

210+
@pytest.mark.parametrize(
211+
("template_value", "parameters", "expected"),
212+
[
213+
("{% if flag %}A{% else %}B{% endif %}", {"flag": True}, "A"),
214+
("{% if flag %}A{% else %}B{% endif %}", {"flag": False}, "B"),
215+
("{% if flag %}A{% else %}B{% endif %}", {"flag": None}, "B"),
216+
("{% if other %}A{% elif flag %}B{% else %}C{% endif %}", {"other": False, "flag": False}, "C"),
217+
('{{ "A" if flag else "B" }}', {"flag": False}, "B"),
218+
("{% set local = flag %}{% if local %}A{% else %}B{% endif %}", {"flag": False}, "B"),
219+
("{% if flag is defined %}A{% else %}B{% endif %}", {}, "B"),
220+
("{{ flag | default('B') }}", {}, "B"),
221+
("{% if flag == 'a' %}A{% else %}B{% endif %}", {"flag": "b"}, "B"),
222+
("{% if flag is filter %}A{% else %}B{% endif %}", {"flag": "upper"}, "A"),
223+
("{% if flag is test %}A{% else %}B{% endif %}", {"flag": "nope"}, "B"),
224+
("{% for item in ['A', 'B'] if flag %}{{ item }}{% endfor %}", {"flag": False}, ""),
225+
(
226+
"{% macro show(rows) %}{% for row in rows %}[{{ row }}]{% endfor %}{% endmacro %}{{ show(items) }}",
227+
{"items": ["x", "y"]},
228+
"[x][y]",
229+
),
230+
],
231+
)
232+
def test_seed_prompt_keeps_template_whose_missing_parameter_decides_a_branch(template_value, parameters, expected):
233+
template = SeedPrompt(value=template_value, data_type="text", is_jinja_template=True)
234+
235+
assert template.value == template_value
236+
assert template.render_template_value(**parameters) == expected
237+
238+
239+
def test_render_template_value_silent_renders_condition_once_its_parameters_are_provided():
240+
template = SeedPrompt(
241+
value="{% if flag %}{{ datasets_path }} {{ prompt }}{% endif %}",
242+
data_type="text",
243+
is_jinja_template=True,
244+
)
245+
246+
assert template.render_template_value_silent(flag=True, datasets_path="data") == "data {{ prompt }}"
247+
248+
249+
def test_render_template_value_silent_renders_if_guard_on_loop_variable():
250+
seed = SeedPrompt(
251+
value="{% for item in items %}{% if item %}[{{ item }}]{% endif %}{% endfor %}",
252+
data_type="text",
253+
is_jinja_template=True,
254+
)
255+
256+
assert seed.render_template_value_silent(items=["a", "", "b"]) == "[a][b]"
257+
258+
259+
_CONVERSATION_CONTEXT_TEMPLATES = sorted(
260+
path
261+
for path in pathlib.Path(DATASETS_PATH, "executors").rglob("*.yaml")
262+
if "{% if conversation_context %}" in path.read_text(encoding="utf-8")
263+
)
264+
265+
266+
@pytest.mark.parametrize("template_path", _CONVERSATION_CONTEXT_TEMPLATES, ids=lambda path: path.stem)
267+
@pytest.mark.parametrize("conversation_context", [None, "<context>"])
268+
def test_loaded_template_renders_like_its_source(template_path, conversation_context):
269+
seed_prompt = SeedPrompt.from_yaml_file(template_path)
270+
parameters = {name: f"<{name}>" for name in seed_prompt.parameters or []}
271+
parameters["conversation_context"] = conversation_context
272+
source = yaml.safe_load(template_path.read_text(encoding="utf-8"))["value"]
273+
274+
expected = SandboxedEnvironment(undefined=StrictUndefined).from_string(source).render(**parameters)
275+
276+
assert seed_prompt.render_template_value(**parameters) == expected
277+
assert ("<context>" in expected) == (conversation_context is not None)
278+
279+
207280
def test_seed_group_untrusted_auto_escapes():
208281
group = SeedGroup(seeds=[{"value": '{{ "".__class__ }}', "data_type": "text"}])
209282
seed = group.prompts[0]

0 commit comments

Comments
 (0)