From 43e1df1d9b982f24e3ccba50cf8881eed86d8994 Mon Sep 17 00:00:00 2001
From: Dmitry Vyukov <dvyukov@google.com>
Date: Mon, 26 Jan 2026 17:59:37 +0100
Subject: pkg/aflow: handle input token overflow for LLM tools

Handle LLM tool input token overflow by removing the last tool reply,
and replacing it with an order to answer right now.
I've seen an LLM tool went into too deap research and in the end
just overflowed input tokens. It could provide at least some answer instead.
---
 pkg/aflow/testdata/TestLLMTool.llm.json | 214 +++++++++++++++++++++++++++++++-
 1 file changed, 212 insertions(+), 2 deletions(-)

(limited to 'pkg/aflow/testdata/TestLLMTool.llm.json')

diff --git a/pkg/aflow/testdata/TestLLMTool.llm.json b/pkg/aflow/testdata/TestLLMTool.llm.json
index 7579aa970..705c011e2 100644
--- a/pkg/aflow/testdata/TestLLMTool.llm.json
+++ b/pkg/aflow/testdata/TestLLMTool.llm.json
@@ -396,7 +396,7 @@
 				"parts": [
 					{
 						"functionCall": {
-							"id": "id2",
+							"id": "id3",
 							"args": {
 								"Something": "subtool input 2"
 							},
@@ -410,12 +410,222 @@
 				"parts": [
 					{
 						"functionResponse": {
-							"id": "id2",
+							"id": "id3",
+							"name": "researcher-tool"
+						}
+					}
+				],
+				"role": "user"
+			}
+		]
+	},
+	{
+		"Model": "sub-agent-model",
+		"Config": {
+			"systemInstruction": {
+				"parts": [
+					{
+						"text": "researcher instruction\nPrefer calling several tools at the same time to save round-trips.\n"
+					}
+				],
+				"role": "user"
+			},
+			"temperature": 2,
+			"tools": [
+				{
+					"functionDeclarations": [
+						{
+							"description": "researcher-tool description",
+							"name": "researcher-tool",
+							"parametersJsonSchema": {
+								"additionalProperties": false,
+								"properties": {
+									"Something": {
+										"description": "something",
+										"type": "string"
+									}
+								},
+								"required": [
+									"Something"
+								],
+								"type": "object"
+							},
+							"responseJsonSchema": {
+								"additionalProperties": false,
+								"type": "object"
+							}
+						}
+					]
+				}
+			],
+			"responseModalities": [
+				"TEXT"
+			]
+		},
+		"Request": [
+			{
+				"parts": [
+					{
+						"text": "But really?"
+					}
+				],
+				"role": "user"
+			},
+			{
+				"parts": [
+					{
+						"functionCall": {
+							"id": "id3",
+							"args": {
+								"Something": "subtool input 2"
+							},
+							"name": "researcher-tool"
+						}
+					}
+				],
+				"role": "user"
+			},
+			{
+				"parts": [
+					{
+						"functionResponse": {
+							"id": "id3",
 							"name": "researcher-tool"
 						}
 					}
 				],
 				"role": "user"
+			},
+			{
+				"parts": [
+					{
+						"functionCall": {
+							"id": "id4",
+							"args": {
+								"Something": "subtool input 3"
+							},
+							"name": "researcher-tool"
+						}
+					}
+				],
+				"role": "user"
+			},
+			{
+				"parts": [
+					{
+						"functionResponse": {
+							"id": "id4",
+							"name": "researcher-tool"
+						}
+					}
+				],
+				"role": "user"
+			}
+		]
+	},
+	{
+		"Model": "sub-agent-model",
+		"Config": {
+			"systemInstruction": {
+				"parts": [
+					{
+						"text": "researcher instruction\nPrefer calling several tools at the same time to save round-trips.\n"
+					}
+				],
+				"role": "user"
+			},
+			"temperature": 2,
+			"tools": [
+				{
+					"functionDeclarations": [
+						{
+							"description": "researcher-tool description",
+							"name": "researcher-tool",
+							"parametersJsonSchema": {
+								"additionalProperties": false,
+								"properties": {
+									"Something": {
+										"description": "something",
+										"type": "string"
+									}
+								},
+								"required": [
+									"Something"
+								],
+								"type": "object"
+							},
+							"responseJsonSchema": {
+								"additionalProperties": false,
+								"type": "object"
+							}
+						}
+					]
+				}
+			],
+			"toolConfig": {
+				"functionCallingConfig": {
+					"mode": "NONE"
+				}
+			},
+			"responseModalities": [
+				"TEXT"
+			]
+		},
+		"Request": [
+			{
+				"parts": [
+					{
+						"text": "But really?"
+					}
+				],
+				"role": "user"
+			},
+			{
+				"parts": [
+					{
+						"functionCall": {
+							"id": "id3",
+							"args": {
+								"Something": "subtool input 2"
+							},
+							"name": "researcher-tool"
+						}
+					}
+				],
+				"role": "user"
+			},
+			{
+				"parts": [
+					{
+						"functionResponse": {
+							"id": "id3",
+							"name": "researcher-tool"
+						}
+					}
+				],
+				"role": "user"
+			},
+			{
+				"parts": [
+					{
+						"functionCall": {
+							"id": "id4",
+							"args": {
+								"Something": "subtool input 3"
+							},
+							"name": "researcher-tool"
+						}
+					}
+				],
+				"role": "user"
+			},
+			{
+				"parts": [
+					{
+						"text": "\nProvide a best-effort answer to the original question with all of the information\nyou have so far without calling any more tools!\n"
+					}
+				],
+				"role": "user"
 			}
 		]
 	},
-- 
cgit mrf-deployment