tool uses original input (opensearch-project#2022)

* tool uses original input Signed-off-by: Jing Zhang <jngz@amazon.com> * fix tool input in trace Signed-off-by: Jing Zhang <jngz@amazon.com> --------- Signed-off-by: Jing Zhang <jngz@amazon.com>
HenryL27 · Feb 6, 2024 · ef90bf0 · ef90bf0
1 parent 429ac86
commit ef90bf0
Show file tree

Hide file tree

Showing 2 changed files with 38 additions and 1 deletion.
diff --git a/ml-algorithms/src/main/java/org/opensearch/ml/engine/algorithms/agent/MLChatAgentRunner.java b/ml-algorithms/src/main/java/org/opensearch/ml/engine/algorithms/agent/MLChatAgentRunner.java
@@ -486,7 +486,12 @@ private void runReAct(
                         if (toolSpecParams != null) {
                             toolParams.putAll(toolSpecParams);
                         }
-                        toolParams.put("input", actionInput);
+                        if (tools.get(action).useOriginalInput()) {
+                            toolParams.put("input", question);
+                            lastActionInput.set(question);
+                        } else {
+                            toolParams.put("input", actionInput);
+                        }
                         if (tools.get(action).validate(toolParams)) {
                             try {
                                 String finalAction = action;

diff --git a/...rithms/src/test/java/org/opensearch/ml/engine/algorithms/agent/MLChatAgentRunnerTest.java b/...rithms/src/test/java/org/opensearch/ml/engine/algorithms/agent/MLChatAgentRunnerTest.java
@@ -12,6 +12,7 @@
 import static org.mockito.ArgumentMatchers.anyString;
 import static org.mockito.ArgumentMatchers.isA;
 import static org.mockito.Mockito.doAnswer;
+import static org.mockito.Mockito.doReturn;
 import static org.mockito.Mockito.never;
 import static org.mockito.Mockito.verify;
 import static org.mockito.Mockito.when;
@@ -114,6 +115,8 @@ public class MLChatAgentRunnerTest {
     private ArgumentCaptor<ActionListener<CreateInteractionResponse>> conversationIndexMemoryCapture;
     @Captor
     private ArgumentCaptor<ActionListener<UpdateResponse>> mlMemoryManagerCapture;
+    @Captor
+    private ArgumentCaptor<Map<String, String>> ToolParamsCapture;
 
     @Before
     @SuppressWarnings("unchecked")
@@ -690,6 +693,35 @@ public void testToolParameters() {
         assertNotNull(modelTensorOutput);
     }
 
+    @Test
+    public void testToolUseOriginalInput() {
+        // Mock tool validation to return false.
+        when(firstTool.validate(any())).thenReturn(true);
+
+        // Create an MLAgent with a tool including two parameters.
+        MLAgent mlAgent = createMLAgentWithTools();
+
+        // Create parameters for the agent.
+        Map<String, String> params = createAgentParamsWithAction(FIRST_TOOL, "someInput");
+        params.put("question", "raw input");
+        doReturn(true).when(firstTool).useOriginalInput();
+
+        // Run the MLChatAgentRunner.
+        mlChatAgentRunner.run(mlAgent, params, agentActionListener);
+
+        // Verify that the tool's run method was called.
+        verify(firstTool).run(any(), any());
+        // Verify the size of parameters passed in the tool run method.
+        ArgumentCaptor argumentCaptor = ArgumentCaptor.forClass(Map.class);
+        verify(firstTool).run((Map<String, String>) argumentCaptor.capture(), any());
+        assertEquals(3, ((Map) argumentCaptor.getValue()).size());
+        assertEquals("raw input", ((Map<?, ?>) argumentCaptor.getValue()).get("input"));
+
+        Mockito.verify(agentActionListener).onResponse(objectCaptor.capture());
+        ModelTensorOutput modelTensorOutput = (ModelTensorOutput) objectCaptor.getValue();
+        assertNotNull(modelTensorOutput);
+    }
+
     @Test
     public void testSaveLastTraceFailure() {
         // Mock tool validation to return true.