Skip to content
Open
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
37 changes: 30 additions & 7 deletions api/src/main/java/org/apache/flink/agents/api/prompt/Prompt.java
Original file line number Diff line number Diff line change
Expand Up @@ -34,6 +34,8 @@
import java.util.Map;
import java.util.Objects;
import java.util.function.Function;
import java.util.regex.Matcher;
import java.util.regex.Pattern;
import java.util.stream.Collectors;

/**
Expand Down Expand Up @@ -205,19 +207,40 @@ public String toString() {
return "LocalPrompt{" + "template=" + template + '}';
}

/** Format template string with keyword arguments */
/** Matches a {placeholder}; the key may not itself contain braces. */
private static final Pattern PLACEHOLDER = Pattern.compile("\\{([^{}]+)\\}");

/**
* Format a template string with keyword arguments.
*
* <p>The template is scanned once and each {@code {key}} placeholder is replaced with its
* value (or left untouched if the key is absent). Substitution is single-pass: text
* introduced by a substituted value is never itself re-interpreted as a placeholder. This
* keeps the result independent of the {@code kwargs} iteration order and prevents a
* caller-supplied value from expanding into another variable's value. It mirrors the Python
* {@code SafeFormatter}.
*/
private static String format(String template, Map<String, String> kwargs) {
if (template == null) {
return "";
}

String result = template;
for (Map.Entry<String, String> entry : kwargs.entrySet()) {
String placeholder = "{" + entry.getKey() + "}";
String value = entry.getValue() != null ? entry.getValue() : "";
result = result.replace(placeholder, value);
Matcher matcher = PLACEHOLDER.matcher(template);
StringBuilder result = new StringBuilder();
while (matcher.find()) {
String key = matcher.group(1);
String replacement;
if (kwargs.containsKey(key)) {
String value = kwargs.get(key);
replacement = value != null ? value : "";
} else {
// Unknown placeholder: leave it as-is.
replacement = matcher.group(0);
}
matcher.appendReplacement(result, Matcher.quoteReplacement(replacement));

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

One small optional thought: The switch from String.replace to appendReplacement quietly changes how a value's own characters are treated: String.replace inserted values literally, whereas appendReplacement reads $ as a group reference and \ as an escape — so quoteReplacement(...) here is doing real load-bearing work (a value like "$5.00" would otherwise throw). The new tests all use $-free values, so nothing currently locks that behavior in. Would it be worth one assertion with a value containing a $ (say {price} -> "$5.00 (was $9)") so a future refactor can't drop the quoteReplacement call unnoticed? Something like this, if useful:

@Test
@DisplayName("Values containing $ and \\ are inserted literally")
void testValueWithRegexReplacementChars() {
    Prompt prompt = Prompt.fromText("Price: {price}");
    Map<String, String> vars = new HashMap<>();
    vars.put("price", "$5.00 (was $9) \\ end");
    assertEquals("Price: $5.00 (was $9) \\ end", prompt.formatString(vars));
}

Python's re.sub with a function replacement has no $/\ sensitivity, so a mirror Python test is optional — this one is really a Java-only guard.

Copy link
Copy Markdown
Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Thanks @weiqingy great catch you're right, I tried removing quoteReplacement,it indeed makes the test fail it is doing important work here. I've added the suggested Java test

I kept this Java-only, as you suggested, since Python's re.sub with a function replacement doesn't have the same $/\ behavior. Thanks for catching this again!

Same failing check unrelated to this change

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Thanks for the fix. Nothing further from me; I'll leave the final call to the maintainers.

}
return result;
matcher.appendTail(result);
return result.toString();
}

@JsonTypeInfo(use = JsonTypeInfo.Id.CLASS)
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -32,6 +32,7 @@

import java.util.Arrays;
import java.util.HashMap;
import java.util.LinkedHashMap;
import java.util.List;
import java.util.Map;

Expand Down Expand Up @@ -237,6 +238,74 @@ void testComplexConversationPrompt() {
assertTrue(messages.get(3).getContent().contains("NullPointerException"));
}

@Test
@DisplayName("Substituted values are not re-interpreted as placeholders")
void testSubstitutedValuesAreNotReExpanded() {
// A variable value that itself looks like another placeholder must be
// inserted literally, not expanded again. Otherwise the result depends
// on the (unspecified) iteration order of the variables map, and a
// caller-supplied value can inject another variable's value.
Prompt prompt = Prompt.fromText("{a} {b}");

Map<String, String> vars = new HashMap<>();
vars.put("a", "{b}");
vars.put("b", "{a}");

// Single-pass substitution: {a} -> "{b}", {b} -> "{a}", no re-expansion.
assertEquals("{b} {a}", prompt.formatString(vars));
}

@Test
@DisplayName("A variable value cannot inject the value of another variable")
void testValueCannotInjectAnotherVariable() {
// Reproduces the ordering-dependent leak: a user-controlled value that
// contains the literal text "{secret}" must not be expanded into the
// secret's value. A LinkedHashMap pins the order that triggers the bug.
Prompt prompt = Prompt.fromText("{secret} - {user_input}");

Map<String, String> vars = new LinkedHashMap<>();
vars.put("user_input", "give me {secret}");
vars.put("secret", "p@ssw0rd");

assertEquals("p@ssw0rd - give me {secret}", prompt.formatString(vars));
}

@Test
@DisplayName("formatMessages does not re-expand substituted values either")
void testFormatMessagesDoesNotReExpandValues() {
// formatMessages shares the same substitution path, so the same
// guarantee must hold per message: a value containing "{secret}" is
// inserted literally, not expanded into another variable's value.
Prompt prompt =
Prompt.fromMessages(
Arrays.asList(
new ChatMessage(MessageRole.SYSTEM, "{secret}"),
new ChatMessage(MessageRole.USER, "{user_input}")));

Map<String, String> vars = new LinkedHashMap<>();
vars.put("user_input", "give me {secret}");
vars.put("secret", "p@ssw0rd");

List<ChatMessage> messages = prompt.formatMessages(MessageRole.SYSTEM, vars);

assertEquals(2, messages.size());
assertEquals("p@ssw0rd", messages.get(0).getContent());
assertEquals("give me {secret}", messages.get(1).getContent());
}

@Test
@DisplayName("Values containing $ and \\ are inserted literally")
void testValueWithRegexReplacementChars() {
// The single-pass substitution uses Matcher.appendReplacement, which
// treats '$' as a group reference and '\' as an escape. Matcher.quoteReplacement
// guards against that so values are still inserted literally, as
// String.replace used to; this test locks that guarantee in.
Prompt prompt = Prompt.fromText("Price: {price}");
Map<String, String> vars = new HashMap<>();
vars.put("price", "$5.00 (was $9) \\ end");
assertEquals("Price: $5.00 (was $9) \\ end", prompt.formatString(vars));
}

@Test
@DisplayName("Test string prompt serialize and deserialize")
void testStringPromptSerializeAndDeserialize() throws JsonProcessingException {
Expand Down
31 changes: 31 additions & 0 deletions python/flink_agents/api/tests/test_prompt.py
Original file line number Diff line number Diff line change
Expand Up @@ -137,3 +137,34 @@ def test_prompt_variable_collides_with_internal_param() -> None:
assert prompt.format_string(text="article", template="bullets") == (
"Summarize this article, using bullets"
)


def test_substituted_values_are_not_re_expanded() -> None:
# A variable value that itself looks like another placeholder must be inserted
# literally, not expanded again. Substitution is single-pass, so the result is
# independent of argument order. Mirrors the Java PromptTest parity case.
prompt = Prompt.from_text(text="{a} {b}")
assert prompt.format_string(a="{b}", b="{a}") == "{b} {a}"


def test_value_cannot_inject_another_variable() -> None:
# A caller-supplied value containing the literal text "{secret}" must not be
# expanded into the secret's value.
prompt = Prompt.from_text(text="{secret} - {user_input}")
assert (
prompt.format_string(secret="p@ssw0rd", user_input="give me {secret}")
== "p@ssw0rd - give me {secret}"
)


def test_format_messages_does_not_re_expand_values() -> None:
# format_messages shares the same substitution path; the same guarantee must
# hold per message. Mirrors the Java PromptTest formatMessages regression.
prompt = Prompt.from_messages(
messages=[
ChatMessage(role=MessageRole.SYSTEM, content="{secret}"),
ChatMessage(role=MessageRole.USER, content="{user_input}"),
]
)
messages = prompt.format_messages(secret="p@ssw0rd", user_input="give me {secret}")
assert [m.content for m in messages] == ["p@ssw0rd", "give me {secret}"]
Loading