Skip to content

Commit e0fe596

Browse files
opensearch-trigger-bot[bot]pyek-bot
opensearch-trigger-bot[bot]
and
pyek-bot
authored
[BUG] Allow user to control react agent max_interations value to prevent empty response (#3756) (#3784)
* fix: expose max_iteration for react Signed-off-by: Pavan Yekbote <[email protected]> * fix: defaults for agent execution and differentiate between step and step result Signed-off-by: Pavan Yekbote <[email protected]> * fix: return react agent id in agent response to expose more details Signed-off-by: Pavan Yekbote <[email protected]> * spotless Signed-off-by: Pavan Yekbote <[email protected]> * fix: remove test prompt from react system prompt Signed-off-by: Pavan Yekbote <[email protected]> * refactor: rename parameters exposed to user to executor Signed-off-by: Pavan Yekbote <[email protected]> * fix: give user complete control over planner system prompt Signed-off-by: Pavan Yekbote <[email protected]> --------- Signed-off-by: Pavan Yekbote <[email protected]> (cherry picked from commit 8339bcf) Co-authored-by: Pavan Yekbote <[email protected]>
1 parent 6c5c580 commit e0fe596

File tree

2 files changed

+73
-29
lines changed

2 files changed

+73
-29
lines changed

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

Lines changed: 70 additions & 26 deletions
Original file line numberDiff line numberDiff line change
@@ -17,6 +17,7 @@
1717
import static org.opensearch.ml.engine.algorithms.agent.AgentUtils.getMcpToolSpecs;
1818
import static org.opensearch.ml.engine.algorithms.agent.AgentUtils.getMlToolSpecs;
1919
import static org.opensearch.ml.engine.algorithms.agent.MLChatAgentRunner.LLM_INTERFACE;
20+
import static org.opensearch.ml.engine.algorithms.agent.MLChatAgentRunner.MAX_ITERATION;
2021
import static org.opensearch.ml.engine.algorithms.agent.MLChatAgentRunner.saveTraceData;
2122
import static org.opensearch.ml.engine.algorithms.agent.PromptTemplate.DEFAULT_PLANNER_PROMPT;
2223
import static org.opensearch.ml.engine.algorithms.agent.PromptTemplate.DEFAULT_PLANNER_PROMPT_TEMPLATE;
@@ -90,16 +91,19 @@ public class MLPlanExecuteAndReflectAgentRunner implements MLAgentRunner {
9091
private String plannerWithHistoryPromptTemplate;
9192

9293
// defaults
93-
private static final String DEFAULT_SYSTEM_PROMPT = "Always respond in JSON format.";
94-
private static final String DEFAULT_REACT_SYSTEM_PROMPT = "You are a helpful assistant.";
94+
private static final String DEFAULT_PLANNER_SYSTEM_PROMPT =
95+
"You are part of an OpenSearch cluster. When you deliver your final result, include a comprehensive report. This report MUST:\\n1. List every analysis or step you performed.\\n2. Summarize the inputs, methods, tools, and data used at each step.\\n3. Include key findings from all intermediate steps — do NOT omit them.\\n4. Clearly explain how the steps led to your final conclusion.\\n5. Return the full analysis and conclusion in the 'result' field, even if some of this was mentioned earlier.\\n\\nThe final response should be fully self-contained and detailed, allowing a user to understand the full investigation without needing to reference prior messages. Always respond in JSON format.";
96+
private static final String DEFAULT_EXECUTOR_SYSTEM_PROMPT =
97+
"You are a dedicated helper agent working as part of a plan‑execute‑reflect framework. Your role is to receive a discrete task, execute all necessary internal reasoning or tool calls, and return a single, final response that fully addresses the task. You must never return an empty response. If you are unable to complete the task or retrieve meaningful information, you must respond with a clear explanation of the issue or what was missing. Under no circumstances should you end your reply with a question or ask for more information. If you search any index, always include the raw documents in the final result instead of summarizing the content. This is critical to give visibility into what the query retrieved.";
9598
private static final String DEFAULT_NO_ESCAPE_PARAMS = "tool_configs,_tools";
9699
private static final String DEFAULT_MAX_STEPS_EXECUTED = "20";
97100
private static final int DEFAULT_MESSAGE_HISTORY_LIMIT = 10;
101+
private static final String DEFAULT_REACT_MAX_ITERATIONS = "20";
98102

99103
// fields
100104
public static final String PROMPT_FIELD = "prompt";
101105
public static final String USER_PROMPT_FIELD = "user_prompt";
102-
public static final String REACT_SYSTEM_PROMPT_FIELD = "react_system_prompt";
106+
public static final String EXECUTOR_SYSTEM_PROMPT_FIELD = "executor_system_prompt";
103107
public static final String STEPS_FIELD = "steps";
104108
public static final String COMPLETED_STEPS_FIELD = "completed_steps";
105109
public static final String PLANNER_PROMPT_FIELD = "planner_prompt";
@@ -109,18 +113,21 @@ public class MLPlanExecuteAndReflectAgentRunner implements MLAgentRunner {
109113
public static final String SYSTEM_PROMPT_FIELD = "system_prompt";
110114
public static final String QUESTION_FIELD = "question";
111115
public static final String MEMORY_ID_FIELD = "memory_id";
116+
public static final String PARENT_INTERACTION_ID_FIELD = "parent_interaction_id";
112117
public static final String TENANT_ID_FIELD = "tenant_id";
113118
public static final String RESULT_FIELD = "result";
114119
public static final String RESPONSE_FIELD = "response";
115120
public static final String STEP_RESULT_FIELD = "step_result";
116-
public static final String REACT_AGENT_ID_FIELD = "reAct_agent_id";
117-
public static final String REACT_AGENT_MEMORY_ID_FIELD = "reAct_agent_memory_id";
121+
public static final String EXECUTOR_AGENT_ID_FIELD = "executor_agent_id";
122+
public static final String EXECUTOR_AGENT_MEMORY_ID_FIELD = "executor_agent_memory_id";
123+
public static final String EXECUTOR_AGENT_PARENT_INTERACTION_ID_FIELD = "executor_agent_parent_interaction_id";
118124
public static final String NO_ESCAPE_PARAMS_FIELD = "no_escape_params";
119125
public static final String DEFAULT_PROMPT_TOOLS_FIELD = "tools_prompt";
120126
public static final String MAX_STEPS_EXECUTED_FIELD = "max_steps";
121127
public static final String PLANNER_PROMPT_TEMPLATE_FIELD = "planner_prompt_template";
122128
public static final String REFLECT_PROMPT_TEMPLATE_FIELD = "reflect_prompt_template";
123129
public static final String PLANNER_WITH_HISTORY_TEMPLATE_FIELD = "planner_with_history_template";
130+
public static final String EXECUTOR_MAX_ITERATIONS_FIELD = "executor_max_iterations";
124131

125132
public MLPlanExecuteAndReflectAgentRunner(
126133
Client client,
@@ -154,9 +161,7 @@ private void setupPromptParameters(Map<String, String> params) {
154161

155162
String userPrompt = params.get(QUESTION_FIELD);
156163
params.put(USER_PROMPT_FIELD, userPrompt);
157-
158-
String userSystemPrompt = params.getOrDefault(SYSTEM_PROMPT_FIELD, "");
159-
params.put(SYSTEM_PROMPT_FIELD, userSystemPrompt + DEFAULT_SYSTEM_PROMPT);
164+
params.put(SYSTEM_PROMPT_FIELD, params.getOrDefault(SYSTEM_PROMPT_FIELD, DEFAULT_PLANNER_SYSTEM_PROMPT));
160165

161166
if (params.get(PLANNER_PROMPT_FIELD) != null) {
162167
this.plannerPrompt = params.get(PLANNER_PROMPT_FIELD);
@@ -329,6 +334,8 @@ private void executePlanningLoop(
329334
parentInteractionId,
330335
finalResult,
331336
completedSteps.get(completedSteps.size() - 2),
337+
allParams.get(EXECUTOR_AGENT_MEMORY_ID_FIELD),
338+
allParams.get(EXECUTOR_AGENT_PARENT_INTERACTION_ID_FIELD),
332339
finalListener
333340
);
334341
return;
@@ -353,22 +360,31 @@ private void executePlanningLoop(
353360

354361
if (parseLLMOutput.get(RESULT_FIELD) != null) {
355362
String finalResult = parseLLMOutput.get(RESULT_FIELD);
356-
saveAndReturnFinalResult((ConversationIndexMemory) memory, parentInteractionId, finalResult, null, finalListener);
363+
saveAndReturnFinalResult(
364+
(ConversationIndexMemory) memory,
365+
parentInteractionId,
366+
allParams.get(EXECUTOR_AGENT_MEMORY_ID_FIELD),
367+
allParams.get(EXECUTOR_AGENT_PARENT_INTERACTION_ID_FIELD),
368+
finalResult,
369+
null,
370+
finalListener
371+
);
357372
} else {
358373
// todo: optimize double conversion of steps (string to list to string)
359374
List<String> steps = Arrays.stream(parseLLMOutput.get(STEPS_FIELD).split(", ")).toList();
360375
addSteps(steps, allParams, STEPS_FIELD);
361376

362377
String stepToExecute = steps.getFirst();
363-
String reActAgentId = allParams.get(REACT_AGENT_ID_FIELD);
378+
String reActAgentId = allParams.get(EXECUTOR_AGENT_ID_FIELD);
364379
Map<String, String> reactParams = new HashMap<>();
365380
reactParams.put(QUESTION_FIELD, stepToExecute);
366-
if (allParams.containsKey(REACT_AGENT_MEMORY_ID_FIELD)) {
367-
reactParams.put(MEMORY_ID_FIELD, allParams.get(REACT_AGENT_MEMORY_ID_FIELD));
381+
if (allParams.containsKey(EXECUTOR_AGENT_MEMORY_ID_FIELD)) {
382+
reactParams.put(MEMORY_ID_FIELD, allParams.get(EXECUTOR_AGENT_MEMORY_ID_FIELD));
368383
}
369384

370-
reactParams.put(SYSTEM_PROMPT_FIELD, allParams.getOrDefault(REACT_SYSTEM_PROMPT_FIELD, DEFAULT_REACT_SYSTEM_PROMPT));
385+
reactParams.put(SYSTEM_PROMPT_FIELD, allParams.getOrDefault(EXECUTOR_SYSTEM_PROMPT_FIELD, DEFAULT_EXECUTOR_SYSTEM_PROMPT));
371386
reactParams.put(LLM_RESPONSE_FILTER, allParams.get(LLM_RESPONSE_FILTER));
387+
reactParams.put(MAX_ITERATION, allParams.getOrDefault(EXECUTOR_MAX_ITERATIONS_FIELD, DEFAULT_REACT_MAX_ITERATIONS));
372388

373389
AgentMLInput agentInput = AgentMLInput
374390
.AgentMLInputBuilder()
@@ -387,13 +403,18 @@ private void executePlanningLoop(
387403

388404
// Process tensors in a single stream
389405
reactResult.getMlModelOutputs().stream().flatMap(output -> output.getMlModelTensors().stream()).forEach(tensor -> {
390-
if (MEMORY_ID_FIELD.equals(tensor.getName())) {
391-
results.put(MEMORY_ID_FIELD, tensor.getResult());
392-
} else {
393-
Map<String, ?> dataMap = tensor.getDataAsMap();
394-
if (dataMap != null && dataMap.containsKey(RESPONSE_FIELD)) {
395-
results.put(STEP_RESULT_FIELD, (String) dataMap.get(RESPONSE_FIELD));
396-
}
406+
switch (tensor.getName()) {
407+
case MEMORY_ID_FIELD:
408+
results.put(MEMORY_ID_FIELD, tensor.getResult());
409+
break;
410+
case PARENT_INTERACTION_ID_FIELD:
411+
results.put(PARENT_INTERACTION_ID_FIELD, tensor.getResult());
412+
break;
413+
default:
414+
Map<String, ?> dataMap = tensor.getDataAsMap();
415+
if (dataMap != null && dataMap.containsKey(RESPONSE_FIELD)) {
416+
results.put(STEP_RESULT_FIELD, (String) dataMap.get(RESPONSE_FIELD));
417+
}
397418
}
398419
});
399420

@@ -404,11 +425,16 @@ private void executePlanningLoop(
404425
// Only add memory_id to params if it exists and is not empty
405426
String reActMemoryId = results.get(MEMORY_ID_FIELD);
406427
if (reActMemoryId != null && !reActMemoryId.isEmpty()) {
407-
allParams.put(REACT_AGENT_MEMORY_ID_FIELD, reActMemoryId);
428+
allParams.put(EXECUTOR_AGENT_MEMORY_ID_FIELD, reActMemoryId);
429+
}
430+
431+
String reActParentInteractionId = results.get(PARENT_INTERACTION_ID_FIELD);
432+
if (reActParentInteractionId != null && !reActParentInteractionId.isEmpty()) {
433+
allParams.put(EXECUTOR_AGENT_PARENT_INTERACTION_ID_FIELD, reActParentInteractionId);
408434
}
409435

410-
completedSteps.add(stepToExecute);
411-
completedSteps.add(results.get(STEP_RESULT_FIELD));
436+
completedSteps.add(String.format("\nStep: %s\n", stepToExecute));
437+
completedSteps.add(String.format("\nStep Result: %s\n", results.get(STEP_RESULT_FIELD)));
412438

413439
saveTraceData(
414440
(ConversationIndexMemory) memory,
@@ -524,6 +550,8 @@ private void addSteps(List<String> steps, Map<String, String> allParams, String
524550
private void saveAndReturnFinalResult(
525551
ConversationIndexMemory memory,
526552
String parentInteractionId,
553+
String reactAgentMemoryId,
554+
String reactParentInteractionId,
527555
String finalResult,
528556
String input,
529557
ActionListener<Object> finalListener
@@ -536,7 +564,12 @@ private void saveAndReturnFinalResult(
536564
}
537565

538566
memory.getMemoryManager().updateInteraction(parentInteractionId, updateContent, ActionListener.wrap(res -> {
539-
List<ModelTensors> finalModelTensors = createModelTensors(memory.getConversationId(), parentInteractionId);
567+
List<ModelTensors> finalModelTensors = createModelTensors(
568+
memory.getConversationId(),
569+
parentInteractionId,
570+
reactAgentMemoryId,
571+
reactParentInteractionId
572+
);
540573
finalModelTensors
541574
.add(
542575
ModelTensors
@@ -553,7 +586,12 @@ private void saveAndReturnFinalResult(
553586
}));
554587
}
555588

556-
private static List<ModelTensors> createModelTensors(String sessionId, String parentInteractionId) {
589+
private static List<ModelTensors> createModelTensors(
590+
String sessionId,
591+
String parentInteractionId,
592+
String reactAgentMemoryId,
593+
String reactParentInteractionId
594+
) {
557595
List<ModelTensors> modelTensors = new ArrayList<>();
558596
modelTensors
559597
.add(
@@ -563,7 +601,13 @@ private static List<ModelTensors> createModelTensors(String sessionId, String pa
563601
List
564602
.of(
565603
ModelTensor.builder().name(MLAgentExecutor.MEMORY_ID).result(sessionId).build(),
566-
ModelTensor.builder().name(MLAgentExecutor.PARENT_INTERACTION_ID).result(parentInteractionId).build()
604+
ModelTensor.builder().name(MLAgentExecutor.PARENT_INTERACTION_ID).result(parentInteractionId).build(),
605+
ModelTensor.builder().name(EXECUTOR_AGENT_MEMORY_ID_FIELD).result(reactAgentMemoryId).build(),
606+
ModelTensor
607+
.builder()
608+
.name(EXECUTOR_AGENT_PARENT_INTERACTION_ID_FIELD)
609+
.result(reactParentInteractionId)
610+
.build()
567611
)
568612
)
569613
.build()

‎plugin/src/main/java/org/opensearch/ml/action/agents/TransportRegisterAgentAction.java

Lines changed: 3 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -98,12 +98,12 @@ private void registerAgent(MLAgent agent, ActionListener<MLRegisterAgentResponse
9898
return;
9999
}
100100

101-
// If the agent is a PLAN_EXECUTE_AND_REFLECT agent and does not have a reAct agent id, create a reAct agent
101+
// If the agent is a PLAN_EXECUTE_AND_REFLECT agent and does not have an executor agent id, create an executor (reAct) agent
102102
if (MLAgentType.from(mlAgent.getType()) == MLAgentType.PLAN_EXECUTE_AND_REFLECT
103-
&& !mlAgent.getParameters().containsKey(MLPlanExecuteAndReflectAgentRunner.REACT_AGENT_ID_FIELD)) {
103+
&& !mlAgent.getParameters().containsKey(MLPlanExecuteAndReflectAgentRunner.EXECUTOR_AGENT_ID_FIELD)) {
104104
createConversationAgent(mlAgent, tenantId, ActionListener.wrap(conversationAgentId -> {
105105
Map<String, String> parameters = new HashMap<>(mlAgent.getParameters());
106-
parameters.put(MLPlanExecuteAndReflectAgentRunner.REACT_AGENT_ID_FIELD, conversationAgentId);
106+
parameters.put(MLPlanExecuteAndReflectAgentRunner.EXECUTOR_AGENT_ID_FIELD, conversationAgentId);
107107
MLAgent updatedAgent = mlAgent.toBuilder().parameters(parameters).build();
108108
registerAgentToIndex(updatedAgent, tenantId, listener);
109109
}, listener::onFailure));

0 commit comments

Comments
 (0)