@@ -192,15 +192,15 @@ public void testParsingJsonBlockFromResponse() {
192
192
public void testParsingJsonBlockFromResponse2 () {
193
193
// Prepare the response with JSON block
194
194
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\" }" ;
196
196
String responseWithJsonBlock = "Some text```json" + jsonBlock + "```More text" ;
197
197
198
198
// Mock LLM response to not contain "thought" but contain "response" with JSON block
199
199
Map <String , String > llmResponse = new HashMap <>();
200
200
llmResponse .put ("response" , responseWithJsonBlock );
201
201
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 ));
204
204
205
205
// Create an MLAgent and run the MLChatAgentRunner
206
206
MLAgent mlAgent = createMLAgentWithTools ();
@@ -252,17 +252,23 @@ public void testRunWithIncludeOutputNotSet() {
252
252
@ Test
253
253
public void testRunWithIncludeOutputMLModel () {
254
254
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 ());
257
263
MLToolSpec firstToolSpec = MLToolSpec .builder ().name (FIRST_TOOL ).type (FIRST_TOOL ).build ();
258
264
MLToolSpec secondToolSpec = MLToolSpec .builder ().name (SECOND_TOOL ).type (SECOND_TOOL ).build ();
259
265
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 ();
266
272
mlChatAgentRunner .run (mlAgent , new HashMap <>(), agentActionListener );
267
273
Mockito .verify (agentActionListener ).onResponse (objectCaptor .capture ());
268
274
ModelTensorOutput modelTensorOutput = (ModelTensorOutput ) objectCaptor .getValue ();
@@ -579,8 +585,7 @@ private Answer generateToolResponseAsMLModelResult(String response, int type) {
579
585
ModelTensor modelTensor ;
580
586
if (type == 1 ) {
581
587
modelTensor = ModelTensor .builder ().dataAsMap (ImmutableMap .of ("return" , response )).build ();
582
- }
583
- else {
588
+ } else {
584
589
modelTensor = ModelTensor .builder ().result (response ).build ();
585
590
}
586
591
ModelTensors modelTensors = ModelTensors .builder ().mlModelTensors (Arrays .asList (modelTensor )).build ();
0 commit comments