restore AGUI context for legacy interface agent (#4720)

Signed-off-by: Jiaping Zeng <jpz@amazon.com>

Author: jiapingzeng

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

@@ -1347,6 +1347,13 @@ void processAgentInput(AgentMLInput agentMLInput, MLAgent mlAgent) {
                 remoteDataSet.getParameters().putAll(parameters);
             } else {
                 // For old-style AG_UI agents without model field
+                // Prepend context to question if available
+                if (agentType == MLAgentType.AG_UI) {
+                    String context = remoteDataSet.getParameters().get(AGUI_PARAM_CONTEXT);
+                    if (context != null && !context.isEmpty()) {
+                        question = "Context: " + context + "\nQuestion: " + question;
+                    }
+                }
                 remoteDataSet.getParameters().putIfAbsent(QUESTION, question);
             }
         } catch (Exception e) {

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

@@ -7,6 +7,7 @@
 
 import static org.junit.Assert.*;
 import static org.mockito.Mockito.*;
+import static org.opensearch.ml.common.agui.AGUIConstants.AGUI_PARAM_CONTEXT;
 import static org.opensearch.ml.engine.algorithms.agent.MLAgentExecutor.QUESTION;
 
 import java.io.IOException;
@@ -595,33 +596,97 @@ public void test_ProcessAgentInput_AGUIAgent_WithoutContext() {
 
     @Test
     public void test_ProcessAgentInput_AGUIAgent_WithContext_LegacyInterface() {
-        // AGUI agent with legacy LLM interface
-        // Context has already been appended by AGUIInputConverter before reaching MLAgentExecutor
+        // AGUI agent with legacy LLM interface (no model field)
+        // Context is passed via AGUI_PARAM_CONTEXT and should be prepended to question
         MLAgent agent = MLAgent
             .builder()
             .name("agui_agent_legacy_context")
             .type(MLAgentType.AG_UI.name())
             .llm(LLMSpec.builder().modelId("gpt-4").build())
             .build();
 
-        // Simulate message with context already appended (as done by AGUIInputConverter)
         ContentBlock textBlock = new ContentBlock();
         textBlock.setType(ContentType.TEXT);
-        textBlock.setText("Context:\n- Location: San Francisco\n\nWhat is the weather?");
+        textBlock.setText("What is the weather?");
         Message message = new Message("user", Collections.singletonList(textBlock));
         AgentInput agentInput = new AgentInput();
         agentInput.setInput(Collections.singletonList(message));
         AgentMLInput agentMLInput = new AgentMLInput("test", null, FunctionName.AGENT, agentInput, null, false);
 
+        // Set context via params (as AGUIInputConverter stores it)
+        Map<String, String> params = new HashMap<>();
+        params.put(AGUI_PARAM_CONTEXT, "[{\"description\":\"Location\",\"value\":\"San Francisco\"}]");
+        agentMLInput.setInputDataset(new RemoteInferenceInputDataSet(params));
+
         mlAgentExecutor.processAgentInput(agentMLInput, agent);
 
-        // Verify question contains context that was already appended
+        // Verify question contains context prepended by the new code path
         RemoteInferenceInputDataSet dataset = (RemoteInferenceInputDataSet) agentMLInput.getInputDataset();
         String question = dataset.getParameters().get(QUESTION);
         Assert.assertNotNull(question);
-        Assert.assertTrue(question.contains("Context:"));
+        Assert.assertTrue(question.startsWith("Context: "));
         Assert.assertTrue(question.contains("San Francisco"));
+        Assert.assertTrue(question.contains("Question: "));
+        Assert.assertTrue(question.contains("What is the weather?"));
+    }
+
+    @Test
+    public void test_ProcessAgentInput_AGUIAgent_NoContext_LegacyInterface() {
+        // AGUI agent with legacy LLM interface, no context param
+        MLAgent agent = MLAgent
+            .builder()
+            .name("agui_agent_legacy_no_context")
+            .type(MLAgentType.AG_UI.name())
+            .llm(LLMSpec.builder().modelId("gpt-4").build())
+            .build();
+
+        ContentBlock textBlock = new ContentBlock();
+        textBlock.setType(ContentType.TEXT);
+        textBlock.setText("What is the weather?");
+        Message message = new Message("user", Collections.singletonList(textBlock));
+        AgentInput agentInput = new AgentInput();
+        agentInput.setInput(Collections.singletonList(message));
+        AgentMLInput agentMLInput = new AgentMLInput("test", null, FunctionName.AGENT, agentInput, null, false);
+
+        mlAgentExecutor.processAgentInput(agentMLInput, agent);
+
+        RemoteInferenceInputDataSet dataset = (RemoteInferenceInputDataSet) agentMLInput.getInputDataset();
+        String question = dataset.getParameters().get(QUESTION);
+        Assert.assertNotNull(question);
         Assert.assertTrue(question.contains("What is the weather?"));
+        Assert.assertFalse(question.contains("Context:"));
+    }
+
+    @Test
+    public void test_ProcessAgentInput_AGUIAgent_EmptyContext_LegacyInterface() {
+        // AGUI agent with legacy LLM interface, empty context string
+        MLAgent agent = MLAgent
+            .builder()
+            .name("agui_agent_legacy_bad_context")
+            .type(MLAgentType.AG_UI.name())
+            .llm(LLMSpec.builder().modelId("gpt-4").build())
+            .build();
+
+        ContentBlock textBlock = new ContentBlock();
+        textBlock.setType(ContentType.TEXT);
+        textBlock.setText("What is the weather?");
+        Message message = new Message("user", Collections.singletonList(textBlock));
+        AgentInput agentInput = new AgentInput();
+        agentInput.setInput(Collections.singletonList(message));
+        AgentMLInput agentMLInput = new AgentMLInput("test", null, FunctionName.AGENT, agentInput, null, false);
+
+        // Empty context should not be prepended
+        Map<String, String> params = new HashMap<>();
+        params.put(AGUI_PARAM_CONTEXT, "");
+        agentMLInput.setInputDataset(new RemoteInferenceInputDataSet(params));
+
+        mlAgentExecutor.processAgentInput(agentMLInput, agent);
+
+        RemoteInferenceInputDataSet dataset = (RemoteInferenceInputDataSet) agentMLInput.getInputDataset();
+        String question = dataset.getParameters().get(QUESTION);
+        Assert.assertNotNull(question);
+        Assert.assertTrue(question.contains("What is the weather?"));
+        Assert.assertFalse(question.contains("Context:"));
     }
 
     @Test