Skip to content

Commit d4d3ac9

Browse files
committed
Prompt refactor for planner agent (ml-commons)
Signed-off-by: yyfamazon <yyf@amazon.com>
1 parent be86ad0 commit d4d3ac9

3 files changed

Lines changed: 124 additions & 77 deletions

File tree

ml-algorithms/src/main/java/org/opensearch/ml/engine/algorithms/agent/MLPlanExecuteAndReflectAgentRunner.java

Lines changed: 44 additions & 17 deletions
Original file line numberDiff line numberDiff line change
@@ -34,14 +34,15 @@
3434
import static org.opensearch.ml.engine.algorithms.agent.MLChatAgentRunner.saveTraceData;
3535
import static org.opensearch.ml.engine.algorithms.agent.PromptTemplate.DEFAULT_PLANNER_PROMPT;
3636
import static org.opensearch.ml.engine.algorithms.agent.PromptTemplate.DEFAULT_PLANNER_PROMPT_TEMPLATE;
37+
import static org.opensearch.ml.engine.algorithms.agent.PromptTemplate.DEFAULT_PLANNER_SYSTEM_PROMPT_PREFIX;
3738
import static org.opensearch.ml.engine.algorithms.agent.PromptTemplate.DEFAULT_PLANNER_WITH_HISTORY_PROMPT_TEMPLATE;
3839
import static org.opensearch.ml.engine.algorithms.agent.PromptTemplate.DEFAULT_REFLECT_PROMPT;
3940
import static org.opensearch.ml.engine.algorithms.agent.PromptTemplate.DEFAULT_REFLECT_PROMPT_TEMPLATE;
4041
import static org.opensearch.ml.engine.algorithms.agent.PromptTemplate.EXECUTOR_RESPONSIBILITY;
4142
import static org.opensearch.ml.engine.algorithms.agent.PromptTemplate.FINAL_RESULT_RESPONSE_INSTRUCTIONS;
4243
import static org.opensearch.ml.engine.algorithms.agent.PromptTemplate.MAX_STEP_SUMMARY_PER_SYSTEM_PROMPT;
43-
import static org.opensearch.ml.engine.algorithms.agent.PromptTemplate.PLANNER_RESPONSIBILITY;
44-
import static org.opensearch.ml.engine.algorithms.agent.PromptTemplate.PLAN_EXECUTE_REFLECT_RESPONSE_FORMAT;
44+
import static org.opensearch.ml.engine.algorithms.agent.PromptTemplate.getCorePlanningInstructions;
45+
import static org.opensearch.ml.engine.algorithms.agent.PromptTemplate.getPlanExecuteReflectResponseFormat;
4546

4647
import java.util.ArrayList;
4748
import java.util.HashMap;
@@ -120,8 +121,8 @@ public class MLPlanExecuteAndReflectAgentRunner implements MLAgentRunner {
120121
private String plannerWithHistoryPromptTemplate;
121122

122123
@VisibleForTesting
123-
static final String DEFAULT_PLANNER_SYSTEM_PROMPT = PLANNER_RESPONSIBILITY + PLAN_EXECUTE_REFLECT_RESPONSE_FORMAT
124-
+ FINAL_RESULT_RESPONSE_INSTRUCTIONS;
124+
static final String DEFAULT_PLANNER_SYSTEM_PROMPT = DEFAULT_PLANNER_SYSTEM_PROMPT_PREFIX + getCorePlanningInstructions()
125+
+ getPlanExecuteReflectResponseFormat() + FINAL_RESULT_RESPONSE_INSTRUCTIONS;
125126

126127
@VisibleForTesting
127128
static final String DEFAULT_EXECUTOR_SYSTEM_PROMPT = EXECUTOR_RESPONSIBILITY;
@@ -141,6 +142,9 @@ public class MLPlanExecuteAndReflectAgentRunner implements MLAgentRunner {
141142
public static final String PLAN_EXECUTE_REFLECT_RESPONSE_FORMAT_FIELD = "plan_execute_reflect_response_format";
142143
public static final String PROMPT_TEMPLATE_FIELD = "prompt_template";
143144
public static final String SYSTEM_PROMPT_FIELD = "system_prompt";
145+
public static final String RESULT_EXPAND_OVERRIDE = "result_expand_and_override";
146+
public static final String IMPORTANT_RULES_EXPAND = "important_rules_expand";
147+
public static final String PLANNER_SYSTEM_PROMPT_PREFIX = "planner_system_prompt_prefix";
144148
public static final String QUESTION_FIELD = "question";
145149
public static final String MEMORY_ID_FIELD = "memory_id";
146150
public static final String PARENT_INTERACTION_ID_FIELD = "parent_interaction_id";
@@ -199,23 +203,38 @@ public MLPlanExecuteAndReflectAgentRunner(
199203
}
200204

201205
@VisibleForTesting
202-
void setupPromptParameters(Map<String, String> params) {
206+
void setupPromptParameters(Map<String, String> params, boolean injectDate, String currentDateTime) {
203207
// populated depending on whether LLM is asked to plan or re-evaluate
204208
// removed here, so that error is thrown in case this field is not populated
205209
params.remove(PROMPT_FIELD);
206210

207211
String userPrompt = params.get(QUESTION_FIELD);
208212
params.put(USER_PROMPT_FIELD, userPrompt);
209213

210-
boolean injectDate = Boolean.parseBoolean(params.getOrDefault(INJECT_DATETIME_FIELD, "false"));
211-
String dateFormat = params.get(DATETIME_FORMAT_FIELD);
212-
String currentDateTime = injectDate ? getCurrentDateTime(dateFormat) : "";
214+
String clientBusinessPrompt = params.get(SYSTEM_PROMPT_FIELD);
213215

214-
String plannerSystemPrompt = params.getOrDefault(SYSTEM_PROMPT_FIELD, DEFAULT_PLANNER_SYSTEM_PROMPT);
215-
if (injectDate) {
216-
plannerSystemPrompt = String.format("%s\n\n%s", plannerSystemPrompt, currentDateTime);
216+
if (clientBusinessPrompt != null && !clientBusinessPrompt.isEmpty()) { // Replace the whole prompt if client specifies system prompt template
217+
params.put(SYSTEM_PROMPT_FIELD, clientBusinessPrompt);
218+
} else { // Using the default system prompt template if client does not specify.
219+
params.put(SYSTEM_PROMPT_FIELD, DEFAULT_PLANNER_SYSTEM_PROMPT);
220+
}
221+
222+
String resultExpandOverride = params.get(RESULT_EXPAND_OVERRIDE);
223+
String importantRulesExpand = params.get(IMPORTANT_RULES_EXPAND);
224+
String plannerSystemPromptPrefix = params.getOrDefault(PLANNER_SYSTEM_PROMPT_PREFIX, DEFAULT_PLANNER_SYSTEM_PROMPT_PREFIX);
225+
226+
// Append / mutate the system prompt through these 3 parameters: plannerSystemPromptPrefix, resultExpandOverride and importantRulesExpand
227+
if (plannerSystemPromptPrefix != null && !plannerSystemPromptPrefix.isEmpty()
228+
&& resultExpandOverride != null && !resultExpandOverride.isEmpty()
229+
&& importantRulesExpand != null && !importantRulesExpand.isEmpty()) {
230+
StringBuilder stringBuilder = new StringBuilder();
231+
stringBuilder.append(plannerSystemPromptPrefix).append("\n");
232+
stringBuilder.append(getCorePlanningInstructions()).append("\n");
233+
stringBuilder.append(getPlanExecuteReflectResponseFormat(resultExpandOverride, importantRulesExpand));
234+
String finalPlannerPrompt = stringBuilder.toString();
235+
236+
params.put(SYSTEM_PROMPT_FIELD, finalPlannerPrompt);
217237
}
218-
params.put(SYSTEM_PROMPT_FIELD, plannerSystemPrompt);
219238

220239
String executorSystemPrompt = params.getOrDefault(EXECUTOR_SYSTEM_PROMPT_FIELD, DEFAULT_EXECUTOR_SYSTEM_PROMPT);
221240
if (injectDate) {
@@ -245,7 +264,7 @@ void setupPromptParameters(Map<String, String> params) {
245264
this.plannerWithHistoryPromptTemplate = params.get(PLANNER_WITH_HISTORY_TEMPLATE_FIELD);
246265
}
247266

248-
params.put(PLAN_EXECUTE_REFLECT_RESPONSE_FORMAT_FIELD, PLAN_EXECUTE_REFLECT_RESPONSE_FORMAT);
267+
params.put(PLAN_EXECUTE_REFLECT_RESPONSE_FORMAT_FIELD, getPlanExecuteReflectResponseFormat());
249268

250269
params.put(NO_ESCAPE_PARAMS_FIELD, DEFAULT_NO_ESCAPE_PARAMS);
251270

@@ -264,8 +283,12 @@ void setupPromptParameters(Map<String, String> params) {
264283
}
265284

266285
@VisibleForTesting
267-
void usePlannerPromptTemplate(Map<String, String> params) {
268-
params.put(PROMPT_TEMPLATE_FIELD, this.plannerPromptTemplate);
286+
void usePlannerPromptTemplate(Map<String, String> params, boolean injectDate, String currentDateTime) {
287+
String template = this.plannerPromptTemplate;
288+
if (injectDate) {
289+
template = template + "\n\n" + currentDateTime;
290+
}
291+
params.put(PROMPT_TEMPLATE_FIELD, template);
269292
populatePrompt(params);
270293
}
271294

@@ -297,10 +320,14 @@ public void run(MLAgent mlAgent, Map<String, String> apiParams, ActionListener<O
297320
allParams.put(TENANT_ID_FIELD, mlAgent.getTenantId());
298321
log.debug("MLPlanExecuteAndReflectAgentRunner called with allParams: {}", allParams);
299322

300-
setupPromptParameters(allParams);
323+
boolean injectDate = Boolean.parseBoolean(allParams.getOrDefault(INJECT_DATETIME_FIELD, "false"));
324+
String dateFormat = allParams.get(DATETIME_FORMAT_FIELD);
325+
String currentDateTime = injectDate ? getCurrentDateTime(dateFormat) : "";
326+
327+
setupPromptParameters(allParams, injectDate, currentDateTime);
301328

302329
// planner prompt for the first call
303-
usePlannerPromptTemplate(allParams);
330+
usePlannerPromptTemplate(allParams, injectDate, currentDateTime);
304331

305332
// Token tracking: resolve model metadata, then proceed with memory setup and execution.
306333
String llmModelId = mlAgent.getLlm().getModelId();

ml-algorithms/src/main/java/org/opensearch/ml/engine/algorithms/agent/PromptTemplate.java

Lines changed: 69 additions & 45 deletions
Original file line numberDiff line numberDiff line change
@@ -81,51 +81,75 @@ public class PromptTemplate {
8181
6. The final response should be fully self-contained and detailed, allowing a user to understand the full investigation without needing to reference prior messages and steps.
8282
""";
8383

84-
public static final String PLAN_EXECUTE_REFLECT_RESPONSE_FORMAT = "Response Instructions: \n"
85-
+ "Only respond in JSON format. Always follow the given response instructions. Do not return any content that does not follow the response instructions. Do not add anything before or after the expected JSON. \n"
86-
+ "Always respond with a valid JSON object that strictly follows the below schema:\n"
87-
+ "{\n"
88-
+ "\t\"steps\": array[string], \n"
89-
+ "\t\"result\": string \n"
90-
+ "}\n"
91-
+ "Use \"steps\" to return an array of strings where each string is a step to complete the objective, leave it empty if you know the final result. Please wrap each step in quotes and escape any special characters within the string. \n"
92-
+ "Use \"result\" return the final response when you have enough information, leave it empty if you want to execute more steps. Please escape any special characters within the result. \n"
93-
+ "Here are examples of valid responses following the required JSON schema:\n\n"
94-
+ "Example 1 - When you need to execute steps:\n"
95-
+ "{\n"
96-
+ "\t\"steps\": [\"This is an example step\", \"this is another example step\"],\n"
97-
+ "\t\"result\": \"\"\n"
98-
+ "}\n\n"
99-
+ "Example 2 - When you have the final result:\n"
100-
+ "{\n"
101-
+ "\t\"steps\": [],\n"
102-
+ "\t\"result\": \"This is an example result\\n with escaped special characters\"\n"
103-
+ "}\n"
104-
+ "Important rules for the response:\n"
105-
+ "1. Do not use commas within individual steps \n"
106-
+ "2. Do not add any content before or after the JSON \n"
107-
+ "3. Only respond with a pure JSON object \n\n";
108-
109-
public static final String PLANNER_RESPONSIBILITY =
110-
"""
111-
You are a thoughtful and analytical planner agent in a plan-execute-reflect framework. Your job is to design a clear, step-by-step plan for a given objective.
112-
113-
Instructions:
114-
- Break the objective into an ordered list of atomic, self-contained Steps that, if executed, will lead to the final result or complete the objective.
115-
- Each Step must state what to do, where, and which tool/parameters would be used. You do not execute tools, only reference them for planning.
116-
- Use only the provided tools; do not invent or assume tools. If no suitable tool applies, use reasoning or observations instead.
117-
- Base your plan only on the data and information explicitly provided; do not rely on unstated knowledge or external facts.
118-
- If there is insufficient information to create a complete plan, summarize what is known so far and clearly state what additional information is required to proceed.
119-
- Stop and summarize if the task is complete or further progress is unlikely.
120-
- Avoid vague instructions; be specific about data sources, indexes, or parameters.
121-
- Never make assumptions or rely on implicit knowledge.
122-
- Respond only in JSON format.
123-
124-
Step examples:
125-
Good example: \"Use Tool to sample documents from index: 'my-index'\"
126-
Bad example: \"Use Tool to sample documents from each index\"
127-
Bad example: \"Use Tool to sample documents from all indices\"
128-
""";
84+
public static String getPlanExecuteReflectResponseFormat(String resultExpandAndOverride, String importantRulesExpand) {
85+
String resultExample = "Here are examples of valid responses following the required JSON schema:\n\n"
86+
+ "Example 1 - When you need to execute steps:\n"
87+
+ "{\n"
88+
+ "\t\"steps\": [\"This is an example step\", \"this is another example step\"],\n"
89+
+ "\t\"result\": \"\"\n"
90+
+ "}\n\n"
91+
+ "Example 2 - When you have the final result:\n"
92+
+ "{\n"
93+
+ "\t\"steps\": [],\n"
94+
+ "\t\"result\": \"This is an example result\\n with escaped special characters\"\n"
95+
+ "}\n";
96+
97+
return "# Response Format\n\n"
98+
+ "## JSON Response Requirements\n"
99+
+ "Only respond in JSON format. Always follow the given response instructions. Do not return any content that does not follow the response instructions. Do not add anything before or after the expected JSON\n\n"
100+
+ "Always respond with a valid JSON object that strictly follows the below schema:\n"
101+
+ "```json\n"
102+
+ "{\n"
103+
+ " \"steps\": array[string],\n"
104+
+ " \"result\": string\n"
105+
+ "}\n"
106+
+ "```\n"
107+
+ "\n"
108+
+ "- Use \"steps\" to return an array of strings where each string is a step to complete the objective, leave it empty if you know the final result. Please wrap each step in quotes and escape any special characters within the string\n"
109+
+ "- Use \"result\" to return the final response when you have enough information, leave it empty if you want to execute more steps. "
110+
+ (resultExpandAndOverride == null ? "" : resultExpandAndOverride)
111+
+ "\n"
112+
+ (resultExpandAndOverride == null ? resultExample : "")
113+
+ "## Critical Rules\n"
114+
+ "1. Do not use commas within individual steps\n"
115+
+ "2. Do not add any content before or after the JSON\n"
116+
+ "3. Only respond with a pure JSON object\n"
117+
+ (importantRulesExpand == null ? "" : importantRulesExpand)
118+
+ "\n";
119+
}
120+
121+
public static String getPlanExecuteReflectResponseFormat() {
122+
return getPlanExecuteReflectResponseFormat(null, null);
123+
}
124+
125+
public static final String DEFAULT_PLANNER_SYSTEM_PROMPT_PREFIX =
126+
"# Investigation Planner Agent\n\nYou are a thoughtful and analytical planner agent in a plan-execute-reflect framework. Your job is to design a clear, step-by-step plan for a given objective.\n\n";
127+
128+
public static String getCorePlanningInstructions() {
129+
return """
130+
# Instructions
131+
132+
## Core Planning Rules
133+
- Break the objective into an ordered list of atomic, self-contained Steps that, if executed, will lead to the final result or complete the objective
134+
- Each Step must state what to do, where, and which tool/parameters would be used. You do not execute tools, only reference them for planning
135+
- Use only the provided tools; do not invent or assume tools. If no suitable tool applies, use reasoning or observations instead
136+
- Base your plan only on the data and information explicitly provided; do not rely on unstated knowledge or external facts
137+
- If there is insufficient information to create a complete plan, summarize what is known so far and clearly state what additional information is required to proceed
138+
- Stop and summarize if the task is complete or further progress is unlikely
139+
- Avoid vague instructions; be specific about data sources, indexes, or parameters
140+
- Never make assumptions or rely on implicit knowledge
141+
- Respond only in JSON format
142+
"""
143+
+ """
144+
## Step Examples
145+
**Good example:** "Use Tool to sample documents from index: 'my-index'"
146+
147+
**Bad example:** "Use Tool to sample documents from each index"
148+
149+
**Bad example:** "Use Tool to sample documents from all indices"
150+
151+
""";
152+
}
129153

130154
public static final String EXECUTOR_RESPONSIBILITY =
131155
"""

0 commit comments

Comments
 (0)