make agent only return the last message

pkg/agent/agent.go

@@ -38,10 +38,11 @@ var (
 )
 
 type Session struct {
-	toolManager    ToolManager
-	llm            LLmProvider
-	messages       []any
-	conversationID string
+	toolManager           ToolManager
+	llm                   LLmProvider
+	messages              []any
+	conversationID        string
+	returnOnlyLastMessage bool
 }
 
 type Option func(*Session) error
@@ -91,6 +92,13 @@ func WithConversationID(id string) Option {
 	}
 }
 
+func WithReturnOnlyLastMessage() Option {
+	return func(a *Session) error {
+		a.returnOnlyLastMessage = true
+		return nil
+	}
+}
+
 func NewSession(developerInstructions string, llm LLmProvider, options ...Option) (*Session, error) {
 	agent := &Session{
 		llm:      llm,
@@ -125,18 +133,27 @@ func (a *Session) Query(ctx context.Context, query string) (string, error) {
 		Content: query,
 	}, nil)
 
-	stringBuilder := strings.Builder{}
+	var (
+		strBuilder  strings.Builder
+		lastMessage string
+	)
 	respChan := a.startLoop(ctx)
-
 	for r := range respChan {
 		if r.err != nil {
 			return "", r.err
 		}
-		stringBuilder.WriteString(r.response)
-		stringBuilder.WriteString("\n")
+		if a.returnOnlyLastMessage {
+			lastMessage = r.response
+		} else {
+			strBuilder.WriteString(r.response)
+			strBuilder.WriteString("\n")
+		}
+	}
+	if a.returnOnlyLastMessage {
+		return lastMessage, nil
+	} else {
+		return strBuilder.String(), nil
 	}
-
-	return stringBuilder.String(), nil
 }
 
 func (a *Session) startLoop(ctx context.Context) chan agentOutput {

pkg/agent/agent_test.go

@@ -586,3 +586,97 @@ func TestSession_ConversationIDMessageRetrieval(t *testing.T) {
 		assert.Contains(t, err.Error(), "conversationID can not be empty")
 	})
 }
+
+func TestSession_WithReturnOnlyLastMessage(t *testing.T) {
+	ctrl := gomock.NewController(t)
+	defer ctrl.Finish()
+
+	systemPrompt := "You are a helpful assistant"
+	testQuery := "Tell me about the weather"
+
+	// First LLM response with tool call
+	firstResponse := LLMResponse{
+		Content: []ContentResponse{
+			{
+				Text: "I'll check the weather for you",
+			},
+		},
+		Tools: []ToolResponseObject{
+			{
+				Name: "get_weather",
+				Input: map[string]any{
+					"location": "default",
+				},
+				ToolID: "test",
+			},
+		},
+	}
+
+	// Final LLM response without tool call
+	finalResponse := LLMResponse{
+		Content: []ContentResponse{
+			{
+				Text: "The weather is sunny today",
+			},
+		},
+	}
+
+	mockToolManager := NewMockToolManager(ctrl)
+	mockLLmHandler := NewMockLLmProvider(ctrl)
+
+	// Setup expectations
+	var toolInfoList []ToolInfo
+	if err := json.Unmarshal([]byte(toolList), &toolInfoList); err != nil {
+		t.Fatal(err)
+	}
+	mockToolManager.EXPECT().ToolList().Return(toolInfoList)
+	mockToolManager.EXPECT().
+		CallTool(gomock.Any(), "get_weather", map[string]any{"location": "default"}).
+		Return("Weather: Sunny, 25°C", nil)
+
+	// We expect two LLM calls - one that returns a tool call, and one that gives the final response
+	gomock.InOrder(
+		mockLLmHandler.EXPECT().
+			ProvideResponse(gomock.Any(), gomock.Any()).
+			Return(firstResponse, nil),
+		mockLLmHandler.EXPECT().
+			ProvideResponse(gomock.Any(), gomock.Any()).
+			Return(finalResponse, nil),
+	)
+
+	// Test with returnOnlyLastMessage enabled
+	session, err := NewSession(systemPrompt, mockLLmHandler, WithToolManager(mockToolManager), WithReturnOnlyLastMessage())
+	require.NoError(t, err)
+
+	response, err := session.Query(context.Background(), testQuery)
+	require.NoError(t, err)
+
+	// Should only contain the final response, not the first one
+	assert.Equal(t, "The weather is sunny today", response)
+	assert.NotContains(t, response, "I'll check the weather for you")
+
+	// Test without returnOnlyLastMessage (default behavior)
+	session2, err := NewSession(systemPrompt, mockLLmHandler, WithToolManager(mockToolManager))
+	require.NoError(t, err)
+
+	mockToolManager.EXPECT().ToolList().Return(toolInfoList)
+	mockToolManager.EXPECT().
+		CallTool(gomock.Any(), "get_weather", map[string]any{"location": "default"}).
+		Return("Weather: Sunny, 25°C", nil)
+
+	gomock.InOrder(
+		mockLLmHandler.EXPECT().
+			ProvideResponse(gomock.Any(), gomock.Any()).
+			Return(firstResponse, nil),
+		mockLLmHandler.EXPECT().
+			ProvideResponse(gomock.Any(), gomock.Any()).
+			Return(finalResponse, nil),
+	)
+
+	response2, err := session2.Query(context.Background(), testQuery)
+	require.NoError(t, err)
+
+	// Should contain both responses concatenated with newlines
+	assert.Contains(t, response2, "I'll check the weather for you")
+	assert.Contains(t, response2, "The weather is sunny today")
+}

pkg/engine/actions/executables/agent/exec.go

@@ -13,11 +13,12 @@ import (
 )
 
 type Config struct {
-	ToolConfigs    []ToolConfig `json:"toolConfigs" yaml:"toolConfigs"`
-	SystemPrompt   string       `json:"systemPrompt" yaml:"systemPrompt"`
-	UserPrompt     string       `json:"userPrompt" yaml:"userPrompt"`
-	IntegrationID  string       `json:"integrationID" yaml:"integrationID"`
-	ConversationID string       `json:"conversationID" yaml:"conversationID"`
+	ToolConfigs       []ToolConfig `json:"toolConfigs" yaml:"toolConfigs"`
+	SystemPrompt      string       `json:"systemPrompt" yaml:"systemPrompt"`
+	UserPrompt        string       `json:"userPrompt" yaml:"userPrompt"`
+	IntegrationID     string       `json:"integrationID" yaml:"integrationID"`
+	ConversationID    string       `json:"conversationID" yaml:"conversationID"`
+	ReturnLastMessage bool         `json:"returnLastMessage" yaml:"returnLastMessage"`
 }
 type MCPServerConfig struct {
 	Endpoint string   `json:"endpoint" yaml:"endpoint"`
@@ -62,6 +63,9 @@ func (a *Agent) Execute(ctx context.Context, modifiedConfig string) (interface{}
 	if newConfig.ConversationID != "" {
 		options = append(options, agent.WithConversationID(newConfig.ConversationID))
 	}
+	if newConfig.ReturnLastMessage {
+		options = append(options, agent.WithReturnOnlyLastMessage())
+	}
 	session, err := agent.NewSession(
 		newConfig.SystemPrompt,
 		a.integration,

pkg/engine/requestctx/templates.go

@@ -13,20 +13,6 @@ import (
 	"github.com/Servflow/servflow/pkg/engine/secrets/secretmanager"
 )
 
-//var FunctionMap = template.FuncMap{
-//	"strip":        tmplStripText,
-//	"secretForKey": secretmanager.SecretForKey,
-//	"jsonout":      jsonOut,
-//	"pluck":        tmplPluck,
-//	"escape":       stringEscape, // more idiomatic name
-//	"stringescape": stringEscape, // keep for backward compatibility
-//	"jsonraw":      jsonRaw,
-//	"join":         tmplJoin,
-//	"hash":         tmplHash,
-//	"now":          now,
-//	"secret":       secret,
-//}
-
 func getFuncMap(funcMap template.FuncMap) template.FuncMap {
 	m := template.FuncMap{
 		"strip":        tmplStripText,