pkg/aflow: handle empty LLM replies

4 files changed, 354 insertions, 10 deletions

diff --git a/pkg/aflow/flow_test.go b/pkg/aflow/flow_test.go
index 2c2f52374..30bf23bf0 100644
--- a/pkg/aflow/flow_test.go
+++ b/pkg/aflow/flow_test.go
@@ -422,6 +422,8 @@ func TestToolMisbehavior(t *testing.T) {
 					},
 				},
 			},
+			// LLM tries to get away w/o answering anything.
+			genai.NewPartFromText(""),
 			genai.NewPartFromText("Finally done"),
 		},
 	)
diff --git a/pkg/aflow/llm_agent.go b/pkg/aflow/llm_agent.go
index 406947e25..aceadd415 100644
--- a/pkg/aflow/llm_agent.go
+++ b/pkg/aflow/llm_agent.go
@@ -98,8 +98,14 @@ const llmMultipleToolsInstruction = `
 Prefer calling several tools at the same time to save round-trips.
 `
 
+const llmMissingReply = `You did not provide any final reply to the question. Please return something.
+Or did you want to call some other tools, but did not actually do that?
+`
+
 const llmMissingOutputs = `You did not call set-results tool.
 Please call set-results tool to provide results of the analysis.
+Note: if you already provided you final reply, you will need to provide it again after calling set-results tool.
+Or did you want to call some other tools, but did not actually do that?
 `
 
 type llmOutputs struct {
@@ -191,13 +197,18 @@ func (a *LLMAgent) chat(ctx *Context, cfg *genai.GenerateContentConfig, tools ma
 		}
 		req = append(req, resp.Candidates[0].Content)
 		if len(calls) == 0 {
-			// This is the final reply.
-			if a.Outputs == nil || outputs != nil {
-				return reply, outputs, nil
+			if a.Outputs != nil && outputs == nil {
+				// LLM did not call set-results.
+				req = append(req, genai.NewContentFromText(llmMissingOutputs, genai.RoleUser))
+				continue
 			}
-			// LLM did not call set-results.
-			req = append(req, genai.NewContentFromText(llmMissingOutputs, genai.RoleUser))
-			continue
+			if reply == "" {
+				// LLM did not provide any final reply.
+				req = append(req, genai.NewContentFromText(llmMissingReply, genai.RoleUser))
+				continue
+			}
+			// This is the final reply.
+			return reply, outputs, nil
 		}
 		// This is not the final reply, LLM asked to execute some tools.
 		// Append the current reply, and tool responses to the next request.
@@ -325,6 +336,9 @@ func (a *LLMAgent) parseResponse(resp *genai.GenerateContentResponse) (
 			reply += part.Text
 		}
 	}
+	if strings.TrimSpace(reply) == "" {
+		reply = ""
+	}
 	return
 }
 
diff --git a/pkg/aflow/testdata/TestToolMisbehavior.llm.json b/pkg/aflow/testdata/TestToolMisbehavior.llm.json
index e51e6ab8d..370954d98 100644
--- a/pkg/aflow/testdata/TestToolMisbehavior.llm.json
+++ b/pkg/aflow/testdata/TestToolMisbehavior.llm.json
@@ -600,7 +600,7 @@
 			{
 				"parts": [
 					{
-						"text": "You did not call set-results tool.\nPlease call set-results tool to provide results of the analysis.\n"
+						"text": "You did not call set-results tool.\nPlease call set-results tool to provide results of the analysis.\nNote: if you already provided you final reply, you will need to provide it again after calling set-results tool.\nOr did you want to call some other tools, but did not actually do that?\n"
 					}
 				],
 				"role": "user"
@@ -851,7 +851,7 @@
 			{
 				"parts": [
 					{
-						"text": "You did not call set-results tool.\nPlease call set-results tool to provide results of the analysis.\n"
+						"text": "You did not call set-results tool.\nPlease call set-results tool to provide results of the analysis.\nNote: if you already provided you final reply, you will need to provide it again after calling set-results tool.\nOr did you want to call some other tools, but did not actually do that?\n"
 					}
 				],
 				"role": "user"
@@ -903,5 +903,316 @@
 				"role": "user"
 			}
 		]
+	},
+	{
+		"Model": "model",
+		"Config": {
+			"systemInstruction": {
+				"parts": [
+					{
+						"text": "Do something!\nPrefer calling several tools at the same time to save round-trips.\n\n\nUse set-results tool to provide results of the analysis.\nIt must be called exactly once before the final reply.\nIgnore results of this tool.\n"
+					}
+				],
+				"role": "user"
+			},
+			"temperature": 1,
+			"tools": [
+				{
+					"functionDeclarations": [
+						{
+							"description": "tool description",
+							"name": "tool1",
+							"parametersJsonSchema": {
+								"additionalProperties": false,
+								"properties": {
+									"Tool1Arg": {
+										"description": "arg",
+										"type": "string"
+									}
+								},
+								"required": [
+									"Tool1Arg"
+								],
+								"type": "object"
+							},
+							"responseJsonSchema": {
+								"additionalProperties": false,
+								"type": "object"
+							}
+						}
+					]
+				},
+				{
+					"functionDeclarations": [
+						{
+							"description": "tool description",
+							"name": "tool2",
+							"parametersJsonSchema": {
+								"additionalProperties": false,
+								"properties": {
+									"Tool2Arg": {
+										"description": "arg",
+										"type": "integer"
+									}
+								},
+								"required": [
+									"Tool2Arg"
+								],
+								"type": "object"
+							},
+							"responseJsonSchema": {
+								"additionalProperties": false,
+								"properties": {
+									"Result": {
+										"description": "arg",
+										"type": "integer"
+									}
+								},
+								"required": [
+									"Result"
+								],
+								"type": "object"
+							}
+						}
+					]
+				},
+				{
+					"functionDeclarations": [
+						{
+							"description": "Use this tool to provide results of the analysis.",
+							"name": "set-results",
+							"parametersJsonSchema": {
+								"additionalProperties": false,
+								"properties": {
+									"AdditionalOutput": {
+										"description": "arg",
+										"type": "integer"
+									}
+								},
+								"required": [
+									"AdditionalOutput"
+								],
+								"type": "object"
+							},
+							"responseJsonSchema": {
+								"additionalProperties": false,
+								"properties": {
+									"AdditionalOutput": {
+										"description": "arg",
+										"type": "integer"
+									}
+								},
+								"required": [
+									"AdditionalOutput"
+								],
+								"type": "object"
+							}
+						}
+					]
+				}
+			],
+			"responseModalities": [
+				"TEXT"
+			]
+		},
+		"Request": [
+			{
+				"parts": [
+					{
+						"text": "Prompt"
+					}
+				],
+				"role": "user"
+			},
+			{
+				"parts": [
+					{
+						"functionCall": {
+							"id": "id1",
+							"args": {
+								"Tool1Arg": "string"
+							},
+							"name": "tool1"
+						}
+					},
+					{
+						"functionCall": {
+							"id": "id2",
+							"args": {
+								"Tool2Arg": "string-instead-of-int"
+							},
+							"name": "tool2"
+						}
+					},
+					{
+						"functionCall": {
+							"id": "id3",
+							"name": "tool2"
+						}
+					},
+					{
+						"functionCall": {
+							"id": "id4",
+							"args": {
+								"Tool2Arg": 0,
+								"Tool2Arg2": 100
+							},
+							"name": "tool2"
+						}
+					},
+					{
+						"functionCall": {
+							"id": "id5",
+							"args": {
+								"Arg": 0
+							},
+							"name": "tool3"
+						}
+					},
+					{
+						"functionCall": {
+							"id": "id6",
+							"args": {
+								"WrongArg": 0
+							},
+							"name": "set-results"
+						}
+					}
+				],
+				"role": "user"
+			},
+			{
+				"parts": [
+					{
+						"functionResponse": {
+							"id": "id1",
+							"name": "tool1"
+						}
+					},
+					{
+						"functionResponse": {
+							"id": "id2",
+							"name": "tool2",
+							"response": {
+								"error": "argument \"Tool2Arg\" has wrong type: got string, want int"
+							}
+						}
+					},
+					{
+						"functionResponse": {
+							"id": "id3",
+							"name": "tool2",
+							"response": {
+								"error": "missing argument \"Tool2Arg\""
+							}
+						}
+					},
+					{
+						"functionResponse": {
+							"id": "id4",
+							"name": "tool2",
+							"response": {
+								"Result": 42
+							}
+						}
+					},
+					{
+						"functionResponse": {
+							"id": "id5",
+							"name": "tool3",
+							"response": {
+								"error": "tool \"tool3\" does not exist, please correct the name"
+							}
+						}
+					},
+					{
+						"functionResponse": {
+							"id": "id6",
+							"name": "set-results",
+							"response": {
+								"error": "missing argument \"AdditionalOutput\""
+							}
+						}
+					}
+				],
+				"role": "user"
+			},
+			{
+				"parts": [
+					{
+						"text": "I am done"
+					}
+				],
+				"role": "user"
+			},
+			{
+				"parts": [
+					{
+						"text": "You did not call set-results tool.\nPlease call set-results tool to provide results of the analysis.\nNote: if you already provided you final reply, you will need to provide it again after calling set-results tool.\nOr did you want to call some other tools, but did not actually do that?\n"
+					}
+				],
+				"role": "user"
+			},
+			{
+				"parts": [
+					{
+						"functionCall": {
+							"id": "id1",
+							"args": {
+								"AdditionalOutput": 1
+							},
+							"name": "set-results"
+						}
+					},
+					{
+						"functionCall": {
+							"id": "id2",
+							"args": {
+								"AdditionalOutput": 2
+							},
+							"name": "set-results"
+						}
+					}
+				],
+				"role": "user"
+			},
+			{
+				"parts": [
+					{
+						"functionResponse": {
+							"id": "id1",
+							"name": "set-results",
+							"response": {
+								"AdditionalOutput": 1
+							}
+						}
+					},
+					{
+						"functionResponse": {
+							"id": "id2",
+							"name": "set-results",
+							"response": {
+								"AdditionalOutput": 2
+							}
+						}
+					}
+				],
+				"role": "user"
+			},
+			{
+				"parts": [
+					{}
+				],
+				"role": "user"
+			},
+			{
+				"parts": [
+					{
+						"text": "You did not provide any final reply to the question. Please return something.\nOr did you want to call some other tools, but did not actually do that?\n"
+					}
+				],
+				"role": "user"
+			}
+		]
 	}
 ]
\ No newline at end of file
diff --git a/pkg/aflow/testdata/TestToolMisbehavior.trajectory.json b/pkg/aflow/testdata/TestToolMisbehavior.trajectory.json
index 4bfe632d2..6a4888811 100644
--- a/pkg/aflow/testdata/TestToolMisbehavior.trajectory.json
+++ b/pkg/aflow/testdata/TestToolMisbehavior.trajectory.json
@@ -263,13 +263,30 @@
 		"Finished": "0001-01-01T00:00:26Z"
 	},
 	{
+		"Seq": 14,
+		"Nesting": 2,
+		"Type": "llm",
+		"Name": "smarty",
+		"Model": "model",
+		"Started": "0001-01-01T00:00:27Z"
+	},
+	{
+		"Seq": 14,
+		"Nesting": 2,
+		"Type": "llm",
+		"Name": "smarty",
+		"Model": "model",
+		"Started": "0001-01-01T00:00:27Z",
+		"Finished": "0001-01-01T00:00:28Z"
+	},
+	{
 		"Seq": 1,
 		"Nesting": 1,
 		"Type": "agent",
 		"Name": "smarty",
 		"Model": "model",
 		"Started": "0001-01-01T00:00:02Z",
-		"Finished": "0001-01-01T00:00:27Z",
+		"Finished": "0001-01-01T00:00:29Z",
 		"Results": {
 			"AdditionalOutput": 2
 		},
@@ -283,7 +300,7 @@
 		"Type": "flow",
 		"Name": "test",
 		"Started": "0001-01-01T00:00:01Z",
-		"Finished": "0001-01-01T00:00:28Z",
+		"Finished": "0001-01-01T00:00:30Z",
 		"Results": {
 			"AdditionalOutput": 2,
 			"Reply": "Finally done"