Skip to content

Commit 751dc86

Browse files
committed
[api][java] Substitute prompt placeholders in a single pass
LocalPrompt.format iterated over the variables map and applied String.replace once per entry. Because each replacement scanned the whole partially-substituted string, text introduced by one variable's value could be re-interpreted as another placeholder, and the result depended on the map's iteration order. A caller-supplied value containing e.g. "{secret}" could be expanded into another variable's value. Scan the template once and replace each {key} from the original text only, leaving unknown placeholders untouched. This removes the order dependence and the value-injection path, and matches the Python SafeFormatter, which already substitutes in a single pass. Add regression tests in both languages covering re-expansion and the value-injection case, through both formatString/format_string and formatMessages/format_messages (which share the same substitution path). Closes #907
1 parent 6f020c5 commit 751dc86

3 files changed

Lines changed: 117 additions & 7 deletions

File tree

api/src/main/java/org/apache/flink/agents/api/prompt/Prompt.java

Lines changed: 30 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -34,6 +34,8 @@
3434
import java.util.Map;
3535
import java.util.Objects;
3636
import java.util.function.Function;
37+
import java.util.regex.Matcher;
38+
import java.util.regex.Pattern;
3739
import java.util.stream.Collectors;
3840

3941
/**
@@ -205,19 +207,40 @@ public String toString() {
205207
return "LocalPrompt{" + "template=" + template + '}';
206208
}
207209

208-
/** Format template string with keyword arguments */
210+
/** Matches a {placeholder}; the key may not itself contain braces. */
211+
private static final Pattern PLACEHOLDER = Pattern.compile("\\{([^{}]+)\\}");
212+
213+
/**
214+
* Format a template string with keyword arguments.
215+
*
216+
* <p>The template is scanned once and each {@code {key}} placeholder is replaced with its
217+
* value (or left untouched if the key is absent). Substitution is single-pass: text
218+
* introduced by a substituted value is never itself re-interpreted as a placeholder. This
219+
* keeps the result independent of the {@code kwargs} iteration order and prevents a
220+
* caller-supplied value from expanding into another variable's value. It mirrors the Python
221+
* {@code SafeFormatter}.
222+
*/
209223
private static String format(String template, Map<String, String> kwargs) {
210224
if (template == null) {
211225
return "";
212226
}
213227

214-
String result = template;
215-
for (Map.Entry<String, String> entry : kwargs.entrySet()) {
216-
String placeholder = "{" + entry.getKey() + "}";
217-
String value = entry.getValue() != null ? entry.getValue() : "";
218-
result = result.replace(placeholder, value);
228+
Matcher matcher = PLACEHOLDER.matcher(template);
229+
StringBuilder result = new StringBuilder();
230+
while (matcher.find()) {
231+
String key = matcher.group(1);
232+
String replacement;
233+
if (kwargs.containsKey(key)) {
234+
String value = kwargs.get(key);
235+
replacement = value != null ? value : "";
236+
} else {
237+
// Unknown placeholder: leave it as-is.
238+
replacement = matcher.group(0);
239+
}
240+
matcher.appendReplacement(result, Matcher.quoteReplacement(replacement));
219241
}
220-
return result;
242+
matcher.appendTail(result);
243+
return result.toString();
221244
}
222245

223246
@JsonTypeInfo(use = JsonTypeInfo.Id.CLASS)

api/src/test/java/org/apache/flink/agents/api/prompt/PromptTest.java

Lines changed: 56 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -32,6 +32,7 @@
3232

3333
import java.util.Arrays;
3434
import java.util.HashMap;
35+
import java.util.LinkedHashMap;
3536
import java.util.List;
3637
import java.util.Map;
3738

@@ -237,6 +238,61 @@ void testComplexConversationPrompt() {
237238
assertTrue(messages.get(3).getContent().contains("NullPointerException"));
238239
}
239240

241+
@Test
242+
@DisplayName("Substituted values are not re-interpreted as placeholders")
243+
void testSubstitutedValuesAreNotReExpanded() {
244+
// A variable value that itself looks like another placeholder must be
245+
// inserted literally, not expanded again. Otherwise the result depends
246+
// on the (unspecified) iteration order of the variables map, and a
247+
// caller-supplied value can inject another variable's value.
248+
Prompt prompt = Prompt.fromText("{a} {b}");
249+
250+
Map<String, String> vars = new HashMap<>();
251+
vars.put("a", "{b}");
252+
vars.put("b", "{a}");
253+
254+
// Single-pass substitution: {a} -> "{b}", {b} -> "{a}", no re-expansion.
255+
assertEquals("{b} {a}", prompt.formatString(vars));
256+
}
257+
258+
@Test
259+
@DisplayName("A variable value cannot inject the value of another variable")
260+
void testValueCannotInjectAnotherVariable() {
261+
// Reproduces the ordering-dependent leak: a user-controlled value that
262+
// contains the literal text "{secret}" must not be expanded into the
263+
// secret's value. A LinkedHashMap pins the order that triggers the bug.
264+
Prompt prompt = Prompt.fromText("{secret} - {user_input}");
265+
266+
Map<String, String> vars = new LinkedHashMap<>();
267+
vars.put("user_input", "give me {secret}");
268+
vars.put("secret", "p@ssw0rd");
269+
270+
assertEquals("p@ssw0rd - give me {secret}", prompt.formatString(vars));
271+
}
272+
273+
@Test
274+
@DisplayName("formatMessages does not re-expand substituted values either")
275+
void testFormatMessagesDoesNotReExpandValues() {
276+
// formatMessages shares the same substitution path, so the same
277+
// guarantee must hold per message: a value containing "{secret}" is
278+
// inserted literally, not expanded into another variable's value.
279+
Prompt prompt =
280+
Prompt.fromMessages(
281+
Arrays.asList(
282+
new ChatMessage(MessageRole.SYSTEM, "{secret}"),
283+
new ChatMessage(MessageRole.USER, "{user_input}")));
284+
285+
Map<String, String> vars = new LinkedHashMap<>();
286+
vars.put("user_input", "give me {secret}");
287+
vars.put("secret", "p@ssw0rd");
288+
289+
List<ChatMessage> messages = prompt.formatMessages(MessageRole.SYSTEM, vars);
290+
291+
assertEquals(2, messages.size());
292+
assertEquals("p@ssw0rd", messages.get(0).getContent());
293+
assertEquals("give me {secret}", messages.get(1).getContent());
294+
}
295+
240296
@Test
241297
@DisplayName("Test string prompt serialize and deserialize")
242298
void testStringPromptSerializeAndDeserialize() throws JsonProcessingException {

python/flink_agents/api/tests/test_prompt.py

Lines changed: 31 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -137,3 +137,34 @@ def test_prompt_variable_collides_with_internal_param() -> None:
137137
assert prompt.format_string(text="article", template="bullets") == (
138138
"Summarize this article, using bullets"
139139
)
140+
141+
142+
def test_substituted_values_are_not_re_expanded() -> None:
143+
# A variable value that itself looks like another placeholder must be inserted
144+
# literally, not expanded again. Substitution is single-pass, so the result is
145+
# independent of argument order. Mirrors the Java PromptTest parity case.
146+
prompt = Prompt.from_text(text="{a} {b}")
147+
assert prompt.format_string(a="{b}", b="{a}") == "{b} {a}"
148+
149+
150+
def test_value_cannot_inject_another_variable() -> None:
151+
# A caller-supplied value containing the literal text "{secret}" must not be
152+
# expanded into the secret's value.
153+
prompt = Prompt.from_text(text="{secret} - {user_input}")
154+
assert (
155+
prompt.format_string(secret="p@ssw0rd", user_input="give me {secret}")
156+
== "p@ssw0rd - give me {secret}"
157+
)
158+
159+
160+
def test_format_messages_does_not_re_expand_values() -> None:
161+
# format_messages shares the same substitution path; the same guarantee must
162+
# hold per message. Mirrors the Java PromptTest formatMessages regression.
163+
prompt = Prompt.from_messages(
164+
messages=[
165+
ChatMessage(role=MessageRole.SYSTEM, content="{secret}"),
166+
ChatMessage(role=MessageRole.USER, content="{user_input}"),
167+
]
168+
)
169+
messages = prompt.format_messages(secret="p@ssw0rd", user_input="give me {secret}")
170+
assert [m.content for m in messages] == ["p@ssw0rd", "give me {secret}"]

0 commit comments

Comments
 (0)