@@ -228,6 +228,46 @@ public void testParsingJsonBlockFromResponse2() {
228
228
assertEquals ("parsed final answer" , modelTensor2 .getResult ());
229
229
}
230
230
231
+ @ Test
232
+ public void testParsingJsonBlockFromResponse3 () {
233
+ // Prepare the response with JSON block
234
+ String jsonBlock = "{\" thought\" :\" parsed thought\" , \" action\" :\" parsed action\" , "
235
+ + "\" action_input\" :{\" a\" :\" n\" }, \" final_answer\" :\" parsed final answer\" }" ;
236
+ String responseWithJsonBlock = "Some text```json" + jsonBlock + "```More text" ;
237
+
238
+ // Mock LLM response to not contain "thought" but contain "response" with JSON block
239
+ Map <String , String > llmResponse = new HashMap <>();
240
+ llmResponse .put ("response" , responseWithJsonBlock );
241
+ doAnswer (getLLMAnswer (llmResponse ))
242
+ .when (client )
243
+ .execute (any (ActionType .class ), any (ActionRequest .class ), isA (ActionListener .class ));
244
+
245
+ // Create an MLAgent and run the MLChatAgentRunner
246
+ MLAgent mlAgent = createMLAgentWithTools ();
247
+ Map <String , String > params = new HashMap <>();
248
+ params .put (MLAgentExecutor .PARENT_INTERACTION_ID , "parent_interaction_id" );
249
+ params .put ("verbose" , "true" );
250
+ mlChatAgentRunner .run (mlAgent , params , agentActionListener );
251
+
252
+ // Capture the response passed to the listener
253
+ ArgumentCaptor <Object > responseCaptor = ArgumentCaptor .forClass (Object .class );
254
+ verify (agentActionListener ).onResponse (responseCaptor .capture ());
255
+
256
+ // Extract the captured response
257
+ Object capturedResponse = responseCaptor .getValue ();
258
+ assertTrue (capturedResponse instanceof ModelTensorOutput );
259
+ ModelTensorOutput modelTensorOutput = (ModelTensorOutput ) capturedResponse ;
260
+
261
+ ModelTensor parentInteractionModelTensor = modelTensorOutput .getMlModelOutputs ().get (0 ).getMlModelTensors ().get (1 );
262
+ ModelTensor modelTensor1 = modelTensorOutput .getMlModelOutputs ().get (1 ).getMlModelTensors ().get (0 );
263
+ ModelTensor modelTensor2 = modelTensorOutput .getMlModelOutputs ().get (2 ).getMlModelTensors ().get (0 );
264
+
265
+ // Verify that the parsed values from JSON block are correctly set
266
+ assertEquals ("parent_interaction_id" , parentInteractionModelTensor .getResult ());
267
+ assertEquals ("Thought: parsed thought" , modelTensor1 .getResult ());
268
+ assertEquals ("parsed final answer" , modelTensor2 .getResult ());
269
+ }
270
+
231
271
@ Test
232
272
public void testRunWithIncludeOutputNotSet () {
233
273
LLMSpec llmSpec = LLMSpec .builder ().modelId ("MODEL_ID" ).build ();
0 commit comments