From 2c2305d0192e0c8b78037c441e5fe4f8d7cccbf1 Mon Sep 17 00:00:00 2001
From: Dmitry Vyukov <dvyukov@google.com>
Date: Mon, 26 Jan 2026 15:07:40 +0100
Subject: pkg/aflow: keep LLM reply on tool calls

---
 pkg/aflow/flow_test.go                          | 3 +++
 pkg/aflow/llm_agent.go                          | 5 +++++
 pkg/aflow/testdata/TestWorkflow.llm.json        | 6 ++++++
 pkg/aflow/testdata/TestWorkflow.trajectory.json | 2 +-
 4 files changed, 15 insertions(+), 1 deletion(-)

(limited to 'pkg')

diff --git a/pkg/aflow/flow_test.go b/pkg/aflow/flow_test.go
index 203ff798c..2a03d82d0 100644
--- a/pkg/aflow/flow_test.go
+++ b/pkg/aflow/flow_test.go
@@ -187,6 +187,9 @@ func TestWorkflow(t *testing.T) {
 						},
 					},
 				},
+				{
+					Text: "Some non-thoughts reply along with tool calls",
+				},
 				{
 					Text:    "I am thinking I need to call some tools",
 					Thought: true,
diff --git a/pkg/aflow/llm_agent.go b/pkg/aflow/llm_agent.go
index 2c83be2eb..692443ac3 100644
--- a/pkg/aflow/llm_agent.go
+++ b/pkg/aflow/llm_agent.go
@@ -347,6 +347,11 @@ func (a *LLMAgent) parseResponse(resp *genai.GenerateContentResponse, span *traj
 	if strings.TrimSpace(reply) == "" {
 		reply = ""
 	}
+	// If there is any reply along with tool calls, append it to thoughts.
+	// Otherwise it won't show up in the trajectory anywhere.
+	if len(calls) != 0 && reply != "" {
+		span.Thoughts += "\n" + reply
+	}
 	return
 }
 
diff --git a/pkg/aflow/testdata/TestWorkflow.llm.json b/pkg/aflow/testdata/TestWorkflow.llm.json
index 681552571..d7e490daa 100644
--- a/pkg/aflow/testdata/TestWorkflow.llm.json
+++ b/pkg/aflow/testdata/TestWorkflow.llm.json
@@ -320,6 +320,9 @@
 							"name": "tool2"
 						}
 					},
+					{
+						"text": "Some non-thoughts reply along with tool calls"
+					},
 					{
 						"text": "I am thinking I need to call some tools",
 						"thought": true
@@ -523,6 +526,9 @@
 							"name": "tool2"
 						}
 					},
+					{
+						"text": "Some non-thoughts reply along with tool calls"
+					},
 					{
 						"text": "I am thinking I need to call some tools",
 						"thought": true
diff --git a/pkg/aflow/testdata/TestWorkflow.trajectory.json b/pkg/aflow/testdata/TestWorkflow.trajectory.json
index 3a4859bbf..ab3e49899 100644
--- a/pkg/aflow/testdata/TestWorkflow.trajectory.json
+++ b/pkg/aflow/testdata/TestWorkflow.trajectory.json
@@ -51,7 +51,7 @@
 		"Model": "model1",
 		"Started": "0001-01-01T00:00:05Z",
 		"Finished": "0001-01-01T00:00:06Z",
-		"Thoughts": "I am thinking I need to call some tools"
+		"Thoughts": "I am thinking I need to call some tools\nSome non-thoughts reply along with tool calls"
 	},
 	{
 		"Seq": 4,
-- 
cgit mrf-deployment