Skip to content

Commit a9a20ed

Browse files
committed
apply spotless
Signed-off-by: xinyual <xinyual@amazon.com>
1 parent 3b29152 commit a9a20ed

File tree

2 files changed

+19
-17
lines changed

2 files changed

+19
-17
lines changed

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

+18-13
Original file line numberDiff line numberDiff line change
@@ -192,15 +192,15 @@ public void testParsingJsonBlockFromResponse() {
192192
public void testParsingJsonBlockFromResponse2() {
193193
// Prepare the response with JSON block
194194
String jsonBlock = "{\"thought\":\"parsed thought\", \"action\":\"parsed action\", "
195-
+ "\"action_input\":\"parsed action input\", \"final_answer\":\"parsed final answer\"}";
195+
+ "\"action_input\":\"parsed action input\", \"final_answer\":\"parsed final answer\"}";
196196
String responseWithJsonBlock = "Some text```json" + jsonBlock + "```More text";
197197

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

205205
// Create an MLAgent and run the MLChatAgentRunner
206206
MLAgent mlAgent = createMLAgentWithTools();
@@ -252,17 +252,23 @@ public void testRunWithIncludeOutputNotSet() {
252252
@Test
253253
public void testRunWithIncludeOutputMLModel() {
254254
LLMSpec llmSpec = LLMSpec.builder().modelId("MODEL_ID").build();
255-
Mockito.doAnswer(generateToolResponseAsMLModelResult("First tool response", 1)).when(firstTool).run(Mockito.anyMap(), toolListenerCaptor.capture());
256-
Mockito.doAnswer(generateToolResponseAsMLModelResult("Second tool response", 2)).when(secondTool).run(Mockito.anyMap(), toolListenerCaptor.capture());
255+
Mockito
256+
.doAnswer(generateToolResponseAsMLModelResult("First tool response", 1))
257+
.when(firstTool)
258+
.run(Mockito.anyMap(), toolListenerCaptor.capture());
259+
Mockito
260+
.doAnswer(generateToolResponseAsMLModelResult("Second tool response", 2))
261+
.when(secondTool)
262+
.run(Mockito.anyMap(), toolListenerCaptor.capture());
257263
MLToolSpec firstToolSpec = MLToolSpec.builder().name(FIRST_TOOL).type(FIRST_TOOL).build();
258264
MLToolSpec secondToolSpec = MLToolSpec.builder().name(SECOND_TOOL).type(SECOND_TOOL).build();
259265
final MLAgent mlAgent = MLAgent
260-
.builder()
261-
.name("TestAgent")
262-
.llm(llmSpec)
263-
.memory(mlMemorySpec)
264-
.tools(Arrays.asList(firstToolSpec, secondToolSpec))
265-
.build();
266+
.builder()
267+
.name("TestAgent")
268+
.llm(llmSpec)
269+
.memory(mlMemorySpec)
270+
.tools(Arrays.asList(firstToolSpec, secondToolSpec))
271+
.build();
266272
mlChatAgentRunner.run(mlAgent, new HashMap<>(), agentActionListener);
267273
Mockito.verify(agentActionListener).onResponse(objectCaptor.capture());
268274
ModelTensorOutput modelTensorOutput = (ModelTensorOutput) objectCaptor.getValue();
@@ -579,8 +585,7 @@ private Answer generateToolResponseAsMLModelResult(String response, int type) {
579585
ModelTensor modelTensor;
580586
if (type == 1) {
581587
modelTensor = ModelTensor.builder().dataAsMap(ImmutableMap.of("return", response)).build();
582-
}
583-
else {
588+
} else {
584589
modelTensor = ModelTensor.builder().result(response).build();
585590
}
586591
ModelTensors modelTensors = ModelTensors.builder().mlModelTensors(Arrays.asList(modelTensor)).build();

ml-algorithms/src/test/java/org/opensearch/ml/engine/tools/AgentToolTests.java

+1-4
Original file line numberDiff line numberDiff line change
@@ -25,8 +25,6 @@
2525
import org.opensearch.client.Client;
2626
import org.opensearch.core.action.ActionListener;
2727
import org.opensearch.ml.common.FunctionName;
28-
import org.opensearch.ml.common.dataset.remote.RemoteInferenceInputDataSet;
29-
import org.opensearch.ml.common.input.execute.agent.AgentMLInput;
3028
import org.opensearch.ml.common.output.model.ModelTensor;
3129
import org.opensearch.ml.common.output.model.ModelTensorOutput;
3230
import org.opensearch.ml.common.output.model.ModelTensors;
@@ -79,8 +77,7 @@ public void testAgentWithChatAgentInput() {
7977
doTestRunMethod(chatAgentInput);
8078
}
8179

82-
private void doTestRunMethod(Map<String, String> parameters)
83-
{
80+
private void doTestRunMethod(Map<String, String> parameters) {
8481
ModelTensor modelTensor = ModelTensor.builder().dataAsMap(ImmutableMap.of("thought", "thought 1", "action", "action1")).build();
8582
ModelTensors modelTensors = ModelTensors.builder().mlModelTensors(Arrays.asList(modelTensor)).build();
8683
ModelTensorOutput mlModelTensorOutput = ModelTensorOutput.builder().mlModelOutputs(Arrays.asList(modelTensors)).build();

0 commit comments

Comments
 (0)