Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Fix argument pass #1941

Merged
merged 14 commits into from
Feb 2, 2024
Merged

Fix argument pass #1941

Show file tree
Hide file tree
Changes from 12 commits
Commits
Show all changes
14 commits
Select commit Hold shift + click to select a range
97f1199
add logs
xinyual Jan 27, 2024
1d9d3b7
option2FixBug
xinyual Jan 29, 2024
20aa5bb
remove useless log
xinyual Jan 29, 2024
70ad223
remove useless log
xinyual Jan 29, 2024
e86a863
fix spot less
xinyual Jan 29, 2024
57e62e3
change argument parsing
xinyual Jan 29, 2024
fd458c7
move common function to utils
xinyual Jan 30, 2024
a30111c
checkout for typo
xinyual Jan 30, 2024
d289673
remove useless code
xinyual Feb 1, 2024
3b29152
add UTs
xinyual Feb 1, 2024
a9a20ed
apply spotless
xinyual Feb 2, 2024
c7439f4
add more uts
xinyual Feb 2, 2024
f14a509
modify import
xinyual Feb 2, 2024
dea4cf6
protect original parameters
xinyual Feb 2, 2024
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
33 changes: 33 additions & 0 deletions ml-algorithms/src/main/java/org/opensearch/ml/engine/algorithms/agent/AgentUtils.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -15,6 +15,9 @@
import static org.opensearch.ml.engine.algorithms.agent.MLChatAgentRunner.TOOL_DESCRIPTIONS;
import static org.opensearch.ml.engine.algorithms.agent.MLChatAgentRunner.TOOL_NAMES;

import java.security.AccessController;
import java.security.PrivilegedActionException;
import java.security.PrivilegedExceptionAction;
import java.util.HashMap;
import java.util.List;
import java.util.Map;
Expand All @@ -23,6 +26,8 @@
import java.util.regex.Pattern;

import org.apache.commons.text.StringSubstitutor;
import org.opensearch.ml.common.output.model.ModelTensor;
import org.opensearch.ml.common.output.model.ModelTensorOutput;
import org.opensearch.ml.common.spi.tools.Tool;

public class AgentUtils {
Expand Down Expand Up @@ -152,4 +157,32 @@
throw new IllegalArgumentException("Model output is invalid");
}
}

public static String outputToOutputString(Object output) throws PrivilegedActionException {
String outputString;
if (output instanceof ModelTensorOutput) {
ModelTensor outputModel = ((ModelTensorOutput) output).getMlModelOutputs().get(0).getMlModelTensors().get(0);
if (outputModel.getDataAsMap() != null) {
outputString = AccessController
.doPrivileged((PrivilegedExceptionAction<String>) () -> gson.toJson(outputModel.getDataAsMap()));
} else {
outputString = outputModel.getResult();
}
} else if (output instanceof String) {
outputString = (String) output;
} else {
outputString = AccessController.doPrivileged((PrivilegedExceptionAction<String>) () -> gson.toJson(output));

Check warning on line 174 in ml-algorithms/src/main/java/org/opensearch/ml/engine/algorithms/agent/AgentUtils.java

View check run for this annotation

Codecov / codecov/patch

ml-algorithms/src/main/java/org/opensearch/ml/engine/algorithms/agent/AgentUtils.java#L174 

Added line #L174 was not covered by tests
}
return outputString;
}

public static String parseInputFromLLMReturn(Map<String, ?> retMap) {
Object actionInput = retMap.get("action_input");
if (actionInput instanceof Map) {
return gson.toJson(actionInput);
} else {
return String.valueOf(actionInput);
}

}
}
26 changes: 14 additions & 12 deletions ml-algorithms/src/main/java/org/opensearch/ml/engine/algorithms/agent/MLChatAgentRunner.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -8,10 +8,8 @@
import static org.opensearch.ml.common.conversation.ActionConstants.ADDITIONAL_INFO_FIELD;
import static org.opensearch.ml.common.conversation.ActionConstants.AI_RESPONSE_FIELD;
import static org.opensearch.ml.common.utils.StringUtils.gson;
import static org.opensearch.ml.engine.algorithms.agent.AgentUtils.extractModelResponseJson;
import static org.opensearch.ml.engine.algorithms.agent.AgentUtils.*;
Copy link
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Avoid import *

Copy link
Collaborator Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Already fix this.


import java.security.AccessController;
import java.security.PrivilegedExceptionAction;
import java.util.ArrayList;
import java.util.Collections;
import java.util.HashMap;
Expand Down Expand Up @@ -325,7 +323,7 @@
}
String thought = String.valueOf(dataAsMap.get("thought"));
String action = String.valueOf(dataAsMap.get("action"));
String actionInput = String.valueOf(dataAsMap.get("action_input"));
String actionInput = parseInputFromLLMReturn(dataAsMap);
String finalAnswer = (String) dataAsMap.get("final_answer");
if (!dataAsMap.containsKey("thought")) {
String response = (String) dataAsMap.get("response");
Expand All @@ -336,7 +334,7 @@
Map map = gson.fromJson(jsonBlock, Map.class);
thought = String.valueOf(map.get("thought"));
action = String.valueOf(map.get("action"));
actionInput = String.valueOf(map.get("action_input"));
actionInput = parseInputFromLLMReturn(map);

Check warning on line 337 in ml-algorithms/src/main/java/org/opensearch/ml/engine/algorithms/agent/MLChatAgentRunner.java

View check run for this annotation

Codecov / codecov/patch

ml-algorithms/src/main/java/org/opensearch/ml/engine/algorithms/agent/MLChatAgentRunner.java#L337 

Added line #L337 was not covered by tests
finalAnswer = (String) map.get("final_answer");
} else {
finalAnswer = response;
Expand Down Expand Up @@ -524,9 +522,7 @@
} else {
MLToolSpec toolSpec = toolSpecMap.get(lastAction.get());
if (toolSpec != null && toolSpec.isIncludeOutputInAgentResponse()) {
String outputString = output instanceof String
? (String) output
: AccessController.doPrivileged((PrivilegedExceptionAction<String>) () -> gson.toJson(output));
String outputString = outputToOutputString(output);

String toolOutputKey = String.format("%s.output", toolSpec.getType());
if (additionalInfo.get(toolOutputKey) != null) {
Expand All @@ -546,7 +542,13 @@
.singletonList(
ModelTensor
.builder()
.dataAsMap(ImmutableMap.of("response", lastThought.get() + "\nObservation: " + output))
.dataAsMap(
ImmutableMap
.of(
"response",
lastThought.get() + "\nObservation: " + outputToOutputString(output)
)
)
.build()
)
)
Expand All @@ -555,7 +557,7 @@

String toolResponse = tmpParameters.get("prompt.tool_response");
StringSubstitutor toolResponseSubstitutor = new StringSubstitutor(
ImmutableMap.of("observation", output),
ImmutableMap.of("observation", outputToOutputString(output)),
"${parameters.",
"}"
);
Expand All @@ -567,7 +569,7 @@
.conversationIndexMessageBuilder()
.type("ReAct")
.question(lastActionInput.get())
.response((String) output)
.response(outputToOutputString(output))
.finalAnswer(false)
.sessionId(sessionId)
.build();
Expand All @@ -582,7 +584,7 @@
newPrompt.set(substitutor.replace(finalPrompt));
tmpParameters.put(PROMPT, newPrompt.get());

sessionMsgAnswerBuilder.append("\nObservation: ").append(output);
sessionMsgAnswerBuilder.append("\nObservation: ").append(outputToOutputString(output));
cotModelTensors
.add(
ModelTensors
Expand Down
14 changes: 14 additions & 0 deletions ml-algorithms/src/main/java/org/opensearch/ml/engine/tools/AgentTool.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -5,6 +5,8 @@

package org.opensearch.ml.engine.tools;

import static org.opensearch.ml.common.utils.StringUtils.gson;

import java.util.Map;

import org.opensearch.action.ActionRequest;
Expand Down Expand Up @@ -51,6 +53,7 @@ public AgentTool(Client client, String agentId) {

@Override
public <T> void run(Map<String, String> parameters, ActionListener<T> listener) {
extractFromChatParameters(parameters);
AgentMLInput agentMLInput = AgentMLInput
.AgentMLInputBuilder()
.agentId(agentId)
Expand Down Expand Up @@ -135,4 +138,15 @@ public String getDefaultVersion() {
return null;
}
}

private void extractFromChatParameters(Map<String, String> parameters) {
Copy link
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Is this for chat only?

if (parameters.containsKey("input")) {
try {
Map<String, String> chatParameters = gson.fromJson(parameters.get("input"), Map.class);
parameters.putAll(chatParameters);
Copy link
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Why do this ? Looks it's possible to override the original params in parameters if they have same key

} catch (Exception exception) {
log.info("fail extract parameters from key 'input' due to" + exception.getMessage());
}
}
}
}
126 changes: 126 additions & 0 deletions ...rithms/src/test/java/org/opensearch/ml/engine/algorithms/agent/MLChatAgentRunnerTest.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -188,6 +188,86 @@ public void testParsingJsonBlockFromResponse() {
assertEquals("parsed final answer", modelTensor2.getResult());
}

@Test
public void testParsingJsonBlockFromResponse2() {
// Prepare the response with JSON block
String jsonBlock = "{\"thought\":\"parsed thought\", \"action\":\"parsed action\", "
+ "\"action_input\":\"parsed action input\", \"final_answer\":\"parsed final answer\"}";
String responseWithJsonBlock = "Some text```json" + jsonBlock + "```More text";

// Mock LLM response to not contain "thought" but contain "response" with JSON block
Map<String, String> llmResponse = new HashMap<>();
llmResponse.put("response", responseWithJsonBlock);
doAnswer(getLLMAnswer(llmResponse))
.when(client)
.execute(any(ActionType.class), any(ActionRequest.class), isA(ActionListener.class));

// Create an MLAgent and run the MLChatAgentRunner
MLAgent mlAgent = createMLAgentWithTools();
Map<String, String> params = new HashMap<>();
params.put(MLAgentExecutor.PARENT_INTERACTION_ID, "parent_interaction_id");
params.put("verbose", "true");
mlChatAgentRunner.run(mlAgent, params, agentActionListener);

// Capture the response passed to the listener
ArgumentCaptor<Object> responseCaptor = ArgumentCaptor.forClass(Object.class);
verify(agentActionListener).onResponse(responseCaptor.capture());

// Extract the captured response
Object capturedResponse = responseCaptor.getValue();
assertTrue(capturedResponse instanceof ModelTensorOutput);
ModelTensorOutput modelTensorOutput = (ModelTensorOutput) capturedResponse;

ModelTensor parentInteractionModelTensor = modelTensorOutput.getMlModelOutputs().get(0).getMlModelTensors().get(1);
ModelTensor modelTensor1 = modelTensorOutput.getMlModelOutputs().get(1).getMlModelTensors().get(0);
ModelTensor modelTensor2 = modelTensorOutput.getMlModelOutputs().get(2).getMlModelTensors().get(0);

// Verify that the parsed values from JSON block are correctly set
assertEquals("parent_interaction_id", parentInteractionModelTensor.getResult());
assertEquals("Thought: parsed thought", modelTensor1.getResult());
assertEquals("parsed final answer", modelTensor2.getResult());
}

@Test
public void testParsingJsonBlockFromResponse3() {
// Prepare the response with JSON block
String jsonBlock = "{\"thought\":\"parsed thought\", \"action\":\"parsed action\", "
+ "\"action_input\":{\"a\":\"n\"}, \"final_answer\":\"parsed final answer\"}";
String responseWithJsonBlock = "Some text```json" + jsonBlock + "```More text";

// Mock LLM response to not contain "thought" but contain "response" with JSON block
Map<String, String> llmResponse = new HashMap<>();
llmResponse.put("response", responseWithJsonBlock);
doAnswer(getLLMAnswer(llmResponse))
.when(client)
.execute(any(ActionType.class), any(ActionRequest.class), isA(ActionListener.class));

// Create an MLAgent and run the MLChatAgentRunner
MLAgent mlAgent = createMLAgentWithTools();
Map<String, String> params = new HashMap<>();
params.put(MLAgentExecutor.PARENT_INTERACTION_ID, "parent_interaction_id");
params.put("verbose", "true");
mlChatAgentRunner.run(mlAgent, params, agentActionListener);

// Capture the response passed to the listener
ArgumentCaptor<Object> responseCaptor = ArgumentCaptor.forClass(Object.class);
verify(agentActionListener).onResponse(responseCaptor.capture());

// Extract the captured response
Object capturedResponse = responseCaptor.getValue();
assertTrue(capturedResponse instanceof ModelTensorOutput);
ModelTensorOutput modelTensorOutput = (ModelTensorOutput) capturedResponse;

ModelTensor parentInteractionModelTensor = modelTensorOutput.getMlModelOutputs().get(0).getMlModelTensors().get(1);
ModelTensor modelTensor1 = modelTensorOutput.getMlModelOutputs().get(1).getMlModelTensors().get(0);
ModelTensor modelTensor2 = modelTensorOutput.getMlModelOutputs().get(2).getMlModelTensors().get(0);

// Verify that the parsed values from JSON block are correctly set
assertEquals("parent_interaction_id", parentInteractionModelTensor.getResult());
assertEquals("Thought: parsed thought", modelTensor1.getResult());
assertEquals("parsed final answer", modelTensor2.getResult());
}

@Test
public void testRunWithIncludeOutputNotSet() {
LLMSpec llmSpec = LLMSpec.builder().modelId("MODEL_ID").build();
Expand All @@ -209,6 +289,35 @@ public void testRunWithIncludeOutputNotSet() {
assertEquals("This is the final answer", agentOutput.get(0).getDataAsMap().get("response"));
}

@Test
public void testRunWithIncludeOutputMLModel() {
LLMSpec llmSpec = LLMSpec.builder().modelId("MODEL_ID").build();
Mockito
.doAnswer(generateToolResponseAsMLModelResult("First tool response", 1))
.when(firstTool)
.run(Mockito.anyMap(), toolListenerCaptor.capture());
Mockito
.doAnswer(generateToolResponseAsMLModelResult("Second tool response", 2))
.when(secondTool)
.run(Mockito.anyMap(), toolListenerCaptor.capture());
MLToolSpec firstToolSpec = MLToolSpec.builder().name(FIRST_TOOL).type(FIRST_TOOL).build();
MLToolSpec secondToolSpec = MLToolSpec.builder().name(SECOND_TOOL).type(SECOND_TOOL).build();
final MLAgent mlAgent = MLAgent
.builder()
.name("TestAgent")
.llm(llmSpec)
.memory(mlMemorySpec)
.tools(Arrays.asList(firstToolSpec, secondToolSpec))
.build();
mlChatAgentRunner.run(mlAgent, new HashMap<>(), agentActionListener);
Mockito.verify(agentActionListener).onResponse(objectCaptor.capture());
ModelTensorOutput modelTensorOutput = (ModelTensorOutput) objectCaptor.getValue();
List<ModelTensor> agentOutput = modelTensorOutput.getMlModelOutputs().get(0).getMlModelTensors();
assertEquals(1, agentOutput.size());
// Respond with last tool output
assertEquals("This is the final answer", agentOutput.get(0).getDataAsMap().get("response"));
}

@Test
public void testRunWithIncludeOutputSet() {
LLMSpec llmSpec = LLMSpec.builder().modelId("MODEL_ID").build();
Expand Down Expand Up @@ -512,6 +621,23 @@ private Answer generateToolResponse(String response) {
};
}

private Answer generateToolResponseAsMLModelResult(String response, int type) {
ModelTensor modelTensor;
if (type == 1) {
modelTensor = ModelTensor.builder().dataAsMap(ImmutableMap.of("return", response)).build();
} else {
modelTensor = ModelTensor.builder().result(response).build();
}
ModelTensors modelTensors = ModelTensors.builder().mlModelTensors(Arrays.asList(modelTensor)).build();
ModelTensorOutput mlModelTensorOutput = ModelTensorOutput.builder().mlModelOutputs(Arrays.asList(modelTensors)).build();

return invocation -> {
ActionListener<Object> listener = invocation.getArgument(1);
listener.onResponse(mlModelTensorOutput);
return null;
};
}

private Answer generateToolFailure(Exception e) {
return invocation -> {
ActionListener<Object> listener = invocation.getArgument(1);
Expand Down
28 changes: 20 additions & 8 deletions ml-algorithms/src/test/java/org/opensearch/ml/engine/tools/AgentToolTests.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -10,6 +10,7 @@
import static org.mockito.ArgumentMatchers.eq;
import static org.mockito.Mockito.doAnswer;
import static org.mockito.Mockito.verify;
import static org.opensearch.ml.common.utils.StringUtils.gson;
import static org.opensearch.ml.engine.tools.AgentTool.DEFAULT_DESCRIPTION;

import java.util.Arrays;
Expand All @@ -24,8 +25,6 @@
import org.opensearch.client.Client;
import org.opensearch.core.action.ActionListener;
import org.opensearch.ml.common.FunctionName;
import org.opensearch.ml.common.dataset.remote.RemoteInferenceInputDataSet;
import org.opensearch.ml.common.input.execute.agent.AgentMLInput;
import org.opensearch.ml.common.output.model.ModelTensor;
import org.opensearch.ml.common.output.model.ModelTensorOutput;
import org.opensearch.ml.common.output.model.ModelTensors;
Expand Down Expand Up @@ -67,13 +66,26 @@ public void setup() {
public void testAgenttestRunMethod() {
Map<String, String> parameters = new HashMap<>();
parameters.put("testKey", "testValue");
AgentMLInput agentMLInput = AgentMLInput
.AgentMLInputBuilder()
.agentId("agentId")
.functionName(FunctionName.AGENT)
.inputDataset(RemoteInferenceInputDataSet.builder().parameters(parameters).build())
.build();
doTestRunMethod(parameters);
}

@Test
public void testAgentWithChatAgentInput() {
Map<String, String> parameters = new HashMap<>();
parameters.put("testKey", "testValue");
Map<String, String> chatAgentInput = new HashMap<>();
chatAgentInput.put("input", gson.toJson(parameters));
doTestRunMethod(chatAgentInput);
}

@Test
public void testAgentWithChatAgentInputWrongFormat() {
Map<String, String> chatAgentInput = new HashMap<>();
chatAgentInput.put("input", "wrong format");
doTestRunMethod(chatAgentInput);
}

private void doTestRunMethod(Map<String, String> parameters) {
ModelTensor modelTensor = ModelTensor.builder().dataAsMap(ImmutableMap.of("thought", "thought 1", "action", "action1")).build();
ModelTensors modelTensors = ModelTensors.builder().mlModelTensors(Arrays.asList(modelTensor)).build();
ModelTensorOutput mlModelTensorOutput = ModelTensorOutput.builder().mlModelOutputs(Arrays.asList(modelTensors)).build();
Expand Down
Loading