Skip to content

Commit ccf9d6b

Browse files
Fix: Ensure consistent response format for PER agent (#4110) (#4117)
(cherry picked from commit 1f28c21) Co-authored-by: Pavan Yekbote <[email protected]>
1 parent 20f37bd commit ccf9d6b

File tree

3 files changed

+72
-102
lines changed

3 files changed

+72
-102
lines changed

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

Lines changed: 16 additions & 21 deletions
Original file line numberDiff line numberDiff line change
@@ -610,15 +610,15 @@ String extractJsonFromMarkdown(String response) {
610610
@VisibleForTesting
611611
void addToolsToPrompt(Map<String, Tool> tools, Map<String, String> allParams) {
612612
StringBuilder toolsPrompt = new StringBuilder(
613-
"In this environment, the executor agent only has access to the below tools. You must choose from only the following tools — no other tools are available. Do not use tools not listed here. \n"
613+
"In this environment, you have access to the tools listed below. Use these tools to create your plan, and do not reference or use any tools not listed here.\n"
614614
);
615615
int toolNumber = 0;
616616
for (Map.Entry<String, Tool> entry : tools.entrySet()) {
617617
String toolName = entry.getKey();
618618
String toolDescription = entry.getValue().getDescription();
619619
toolsPrompt.append(String.format("Tool %d - %s: %s\n\n", ++toolNumber, toolName, toolDescription));
620620
}
621-
toolsPrompt.append("No other tools are available. Do not invent tools.\n\n");
621+
toolsPrompt.append("No other tools are available. Do not invent tools. Only use tools to create the plan.\n\n");
622622
allParams.put(DEFAULT_PROMPT_TOOLS_FIELD, toolsPrompt.toString());
623623
populatePrompt(allParams);
624624
cleanUpResource(tools);
@@ -677,25 +677,20 @@ static List<ModelTensors> createModelTensors(
677677
String reactParentInteractionId
678678
) {
679679
List<ModelTensors> modelTensors = new ArrayList<>();
680-
modelTensors
681-
.add(
682-
ModelTensors
683-
.builder()
684-
.mlModelTensors(
685-
List
686-
.of(
687-
ModelTensor.builder().name(MLAgentExecutor.MEMORY_ID).result(sessionId).build(),
688-
ModelTensor.builder().name(MLAgentExecutor.PARENT_INTERACTION_ID).result(parentInteractionId).build(),
689-
ModelTensor.builder().name(EXECUTOR_AGENT_MEMORY_ID_FIELD).result(reactAgentMemoryId).build(),
690-
ModelTensor
691-
.builder()
692-
.name(EXECUTOR_AGENT_PARENT_INTERACTION_ID_FIELD)
693-
.result(reactParentInteractionId)
694-
.build()
695-
)
696-
)
697-
.build()
698-
);
680+
List<ModelTensor> tensors = new ArrayList<>();
681+
682+
tensors.add(ModelTensor.builder().name(MLAgentExecutor.MEMORY_ID).result(sessionId).build());
683+
tensors.add(ModelTensor.builder().name(MLAgentExecutor.PARENT_INTERACTION_ID).result(parentInteractionId).build());
684+
685+
if (reactAgentMemoryId != null && !reactAgentMemoryId.isEmpty()) {
686+
tensors.add(ModelTensor.builder().name(EXECUTOR_AGENT_MEMORY_ID_FIELD).result(reactAgentMemoryId).build());
687+
}
688+
689+
if (reactParentInteractionId != null && !reactParentInteractionId.isEmpty()) {
690+
tensors.add(ModelTensor.builder().name(EXECUTOR_AGENT_PARENT_INTERACTION_ID_FIELD).result(reactParentInteractionId).build());
691+
}
692+
693+
modelTensors.add(ModelTensors.builder().mlModelTensors(tensors).build());
699694
return modelTensors;
700695
}
701696

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

Lines changed: 42 additions & 77 deletions
Original file line numberDiff line numberDiff line change
@@ -30,7 +30,7 @@ public class PromptTemplate {
3030
+ "} \n"
3131
+ "Objective: ${parameters."
3232
+ USER_PROMPT_FIELD
33-
+ "} \n\n";
33+
+ "} \n\nRemember: Respond only in JSON format following the required schema.";
3434

3535
public static final String DEFAULT_REFLECT_PROMPT_TEMPLATE = "${parameters."
3636
+ DEFAULT_PROMPT_TOOLS_FIELD
@@ -49,7 +49,7 @@ public class PromptTemplate {
4949
+ "}] \n\n"
5050
+ "${parameters."
5151
+ REFLECT_PROMPT_FIELD
52-
+ "} \n\n.";
52+
+ "} \n\n.Remember: Respond only in JSON format following the required schema.";
5353

5454
public static final String DEFAULT_PLANNER_WITH_HISTORY_PROMPT_TEMPLATE = "${parameters."
5555
+ DEFAULT_PROMPT_TOOLS_FIELD
@@ -62,27 +62,23 @@ public class PromptTemplate {
6262
+ "}``` \n\n"
6363
+ "You have currently executed the following steps: \n[${parameters."
6464
+ COMPLETED_STEPS_FIELD
65-
+ "}] \n\n";
65+
+ "}] \n\nRemember: Respond only in JSON format following the required schema.";
6666

67-
// modify these -- ensure that it breaks down steps simpler
6867
public static final String DEFAULT_PLANNER_PROMPT =
69-
"For the given objective, generate a step-by-step plan composed of simple, self-contained tasks. The final step should directly yield the final answer. Avoid unnecessary steps.";
68+
"For the given objective, generate a step-by-step plan composed of simple, self-contained steps. The final step should directly yield the final answer. Avoid unnecessary steps.";
7069

71-
// modify these -- update your plan based on completed steps
7270
public static final String DEFAULT_REFLECT_PROMPT =
73-
"Update your plan based on the latest results. If the task is complete, return the final answer. Otherwise, include only the remaining steps — do not repeat previously completed ones.";
71+
"Update your plan based on the latest step results. If the task is complete, return the final answer. Otherwise, include only the remaining steps. Do not repeat previously completed steps.";
7472

7573
public static final String FINAL_RESULT_RESPONSE_INSTRUCTIONS =
7674
"""
77-
When you deliver your final result, include a comprehensive report. This report must:
78-
79-
1. List every analysis or step you performed.
80-
2. Summarize the inputs, methods, tools, and data used at each step.
81-
3. Include key findings from all intermediate steps — do NOT omit them.
82-
4. Clearly explain how the steps led to your final conclusion. Only mention the completed steps.
83-
5. Return the full analysis and conclusion in the 'result' field, even if some of this was mentioned earlier.
84-
85-
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.
75+
When you deliver your final result, include a comprehensive report. This report must:
76+
1. List every analysis or step you performed.
77+
2. Summarize the inputs, methods, tools, and data used at each step.
78+
3. Include key findings from all intermediate steps — do NOT omit them.
79+
4. Clearly explain how the steps led to your final conclusion. Only mention the completed steps.
80+
5. Return the full analysis and conclusion in the 'result' field, even if some of this was mentioned earlier. Ensure that special characters are escaped in the 'result' field.
81+
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.
8682
""";
8783

8884
public static final String PLAN_EXECUTE_REFLECT_RESPONSE_FORMAT = "Response Instructions: \n"
@@ -93,7 +89,7 @@ public class PromptTemplate {
9389
+ "\t\"result\": string \n"
9490
+ "}\n"
9591
+ "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"
96-
+ "Use \"result\" return the final response when you have enough information, leave it empty if you want to execute more steps \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"
9793
+ "Here are examples of valid responses following the required JSON schema:\n\n"
9894
+ "Example 1 - When you need to execute steps:\n"
9995
+ "{\n"
@@ -103,7 +99,7 @@ public class PromptTemplate {
10399
+ "Example 2 - When you have the final result:\n"
104100
+ "{\n"
105101
+ "\t\"steps\": [],\n"
106-
+ "\t\"result\": \"This is an example result\"\n"
102+
+ "\t\"result\": \"This is an example result\\n with escaped special characters\"\n"
107103
+ "}\n"
108104
+ "Important rules for the response:\n"
109105
+ "1. Do not use commas within individual steps \n"
@@ -112,67 +108,36 @@ public class PromptTemplate {
112108

113109
public static final String PLANNER_RESPONSIBILITY =
114110
"""
115-
You are a thoughtful and analytical agent working as the `Planner & Reflector Agent` in a Plan–Execute–Reflect framework. You collaborate with a separate `Executor Agent`, whose sole responsibility is to carry out specific Steps that you generate.
116-
117-
## Core Responsibilities
118-
- Receive a high-level objective or user goal and generate a clear, ordered sequence of simple executable Steps to complete the objective
119-
- Ensure each Step is self-contained, meaning it can be executed without any prior context
120-
- Each Step must specify exactly what to do, where to do it, and with which tools or parameters — avoid abstract instructions like “for each index” or “try something”
121-
- If a partially completed plan and its execution results are provided, update the plan accordingly:
122-
- Only include new Steps that still need to be executed
123-
- Do not repeat previously completed Steps unless their output is outdated, missing, or clearly insufficient
124-
- Use results from completed steps to avoid redundant or unnecessary follow-up actions
125-
- If the task is already complete, return the final answer instead of a new plan
126-
- If the available information is sufficient to provide a useful or partial answer, do so — do not over-plan or run unnecessary steps
127-
- Use only the tools provided to construct your plan. You will be provided a list of available tools for each objective. Use only these tools in your plan — do not invent new tool names, do not guess what tools might exist, and do not reference tools not explicitly listed. If no suitable tool is available, plan using reasoning or observations instead.
128-
- Always respond in JSON format
129-
130-
## Step Guidelines
131-
- Each Step must be simple, atomic, and concrete — suitable for execution by a separate agent
132-
- Avoid ambiguity: Steps should clearly define the **specific data sources, indexes, services, or parameters** to use
133-
- Do not include generic instructions that require iteration or interpretation (e.g., “for all indexes” or “check relevant logs”)
134-
- Do not add any superfluous steps — the result of the final step should directly answer the objective
135-
136-
### Bad Step Example: "Use the SearchIndexTool to sample documents from each index"
137-
138-
### Good Step Example: "Use the SearchIndexTool to sample documents for the index: index-name"
139-
140-
## Structural Expectations
141-
- Track what Steps you generate and why
142-
- Specify what tool or method each Step will likely require
143-
- Use execution results to guide re-planning or task completion decisions
144-
- Reuse prior results — do not re-fetch documents or metadata if they have already been retrieved
145-
- If further progress is unlikely based on tool limitations or available data, stop and return the best possible result to the user
146-
- Never rely on implicit knowledge, do not make make assumptions
147-
148-
Your goal is to produce a clean, efficient, and logically sound plan — or to adapt an existing one — to help the Executor Agent make steady progress toward the final answer. If no further progress can reasonably be made, summarize what has been learned and end the investigation.
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\"
149128
""";
150129

151-
// ask it to break down a large step into a smaller step
152130
public static final String EXECUTOR_RESPONSIBILITY =
153131
"""
154-
You are a dedicated helper agent working as the `Executor Agent` in a Plan–Execute–Reflect framework. In this setup, a separate `Planner & Reflector Agent` both creates an ordered list of discrete Steps and, after seeing your execution outputs, re-plans or refines those Steps as needed.
155-
156-
Your sole responsibility is to execute whatever Step you receive.
157-
158-
## Core Responsibilities
159-
- Receive a discrete Step and execute it completely
160-
- Run all necessary internal reasoning or tool calls
161-
- Return a single, consolidated response that fully addresses that Step
162-
- If previous context can help you answer the Step, reuse that information instead of calling tools again
163-
164-
## Critical Requirements
165-
- You must never return an empty response
166-
- Never end your reply with questions or requests for more information
167-
- If you search any index, always include the full raw documents in your output. Do not summarize—so that every piece of retrieved evidence remains visible. This is critical for the `Planner & Reflector Agent` to decide the next step.
168-
- If you cannot complete the Step, provide a clear explanation of what went wrong or what information was missing
169-
- Never rely on implicit knowledge, do not make make assumptions
170-
171-
## Efficiency Guidelines
172-
- Reuse previous context when applicable, stating what you're reusing and why
173-
- Use the most direct approach first
174-
- If a tool call fails, try alternative approaches before declaring failure
175-
- If a search request is complex, break it down into multiple simple search queries
176-
177-
Your response must be complete and actionable as-is.""";
132+
You are a precise and reliable executor agent in a plan-execute-reflect framework. Your job is to execute the given instruction provided by the planner and return a complete, actionable result.
133+
134+
Instructions:
135+
- Fully execute the given Step using the most relevant tools or reasoning.
136+
- Include all relevant raw tool outputs (e.g., full documents from searches) so the planner has complete information; do not summarize unless explicitly instructed.
137+
- Base your execution and conclusions only on the data and tool outputs available; do not rely on unstated knowledge or external facts.
138+
- If the available data is insufficient to complete the Step, summarize what was obtained so far and clearly state the additional information or access required to proceed (do not guess).
139+
- If unable to complete the Step, clearly explain what went wrong and what is needed to proceed.
140+
- Avoid making assumptions and relying on implicit knowledge.
141+
- Your response must be self-contained and ready for the planner to use without modification. Never end with a question.
142+
- Break complex searches into simpler queries when appropriate.""";
178143
}

ml-algorithms/src/test/java/org/opensearch/ml/engine/algorithms/agent/MLPlanExecuteAndReflectAgentRunnerTest.java

Lines changed: 14 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -237,7 +237,7 @@ public void testBasicExecution() {
237237

238238
ModelTensors firstModelTensors = mlModelOutputs.get(0);
239239
List<ModelTensor> firstModelTensorList = firstModelTensors.getMlModelTensors();
240-
assertEquals(4, firstModelTensorList.size());
240+
assertEquals(2, firstModelTensorList.size());
241241

242242
ModelTensor memoryIdTensor = firstModelTensorList.get(0);
243243
assertEquals("memory_id", memoryIdTensor.getName());
@@ -310,7 +310,7 @@ public void testExecutionWithHistory() {
310310

311311
ModelTensors firstModelTensors = mlModelOutputs.get(0);
312312
List<ModelTensor> firstModelTensorList = firstModelTensors.getMlModelTensors();
313-
assertEquals(4, firstModelTensorList.size());
313+
assertEquals(2, firstModelTensorList.size());
314314

315315
ModelTensor memoryIdTensor = firstModelTensorList.get(0);
316316
assertEquals("memory_id", memoryIdTensor.getName());
@@ -575,10 +575,11 @@ public void testAddToolsToPrompt() {
575575
mlPlanExecuteAndReflectAgentRunner.addToolsToPrompt(tools, testParams);
576576

577577
assertEquals(
578-
"In this environment, the executor agent only has access to the below tools. You must choose from only the following tools — no other tools are available. Do not use tools not listed here. \n"
578+
"In this environment, you have access to the tools listed below. Use these tools to create your plan, and do not reference or use any tools not listed here.\n"
579579
+ "Tool 1 - tool1: description1\n"
580580
+ "\n"
581-
+ "No other tools are available. Do not invent tools.\n\n",
581+
+ "No other tools are available. Do not invent tools. Only use tools to create the plan.\n"
582+
+ "\n",
582583
testParams.get(MLPlanExecuteAndReflectAgentRunner.DEFAULT_PROMPT_TOOLS_FIELD)
583584
);
584585
}
@@ -612,6 +613,15 @@ public void testCreateModelTensors() {
612613
assertEquals(parentInteractionId, tensors.getMlModelTensors().get(1).getResult());
613614
assertEquals(executorMemoryId, tensors.getMlModelTensors().get(2).getResult());
614615
assertEquals(executorParentId, tensors.getMlModelTensors().get(3).getResult());
616+
617+
result = MLPlanExecuteAndReflectAgentRunner.createModelTensors(sessionId, parentInteractionId, "", "");
618+
619+
assertNotNull(result);
620+
assertEquals(1, result.size());
621+
tensors = result.get(0);
622+
assertEquals(2, tensors.getMlModelTensors().size());
623+
assertEquals(sessionId, tensors.getMlModelTensors().get(0).getResult());
624+
assertEquals(parentInteractionId, tensors.getMlModelTensors().get(1).getResult());
615625
}
616626

617627
@Test

0 commit comments

Comments
 (0)