feat: add agentic callbacks template (#681)

components/agentic/callback_extra.go

@@ -1,85 +0,0 @@
-/*
- * Copyright 2025 CloudWeGo Authors
- *
- * Licensed under the Apache License, Version 2.0 (the "License");
- * you may not use this file except in compliance with the License.
- * You may obtain a copy of the License at
- *
- *     http://www.apache.org/licenses/LICENSE-2.0
- *
- * Unless required by applicable law or agreed to in writing, software
- * distributed under the License is distributed on an "AS IS" BASIS,
- * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
- * See the License for the specific language governing permissions and
- * limitations under the License.
- */
-
-// Package agentic defines callback payloads and configuration types for agentic models.
-package agentic
-
-import (
-	"github.com/cloudwego/eino/callbacks"
-	"github.com/cloudwego/eino/schema"
-)
-
-// Config is the config for the model.
-type Config struct {
-	// Model is the model name.
-	Model string
-	// Temperature is the temperature, which controls the randomness of the model.
-	Temperature float64
-	// TopP is the top p, which controls the diversity of the model.
-	TopP float64
-}
-
-// CallbackInput is the input for the model callback.
-type CallbackInput struct {
-	// Messages is the messages to be sent to the model.
-	Messages []*schema.AgenticMessage
-	// Tools is the tools to be used in the model.
-	Tools []*schema.ToolInfo
-	// ToolChoice controls which tool is called by the model.
-	ToolChoice *schema.ToolChoice
-	// Config is the config for the model.
-	Config *Config
-	// Extra is the extra information for the callback.
-	Extra map[string]any
-}
-
-// CallbackOutput is the output for the model callback.
-type CallbackOutput struct {
-	// Message is the message generated by the model.
-	Message *schema.AgenticMessage
-	// Config is the config for the model.
-	Config *Config
-	// Extra is the extra information for the callback.
-	Extra map[string]any
-}
-
-// ConvCallbackInput converts the callback input to the model callback input.
-func ConvCallbackInput(src callbacks.CallbackInput) *CallbackInput {
-	switch t := src.(type) {
-	case *CallbackInput: // when callback is triggered within component implementation, the input is usually already a typed *model.CallbackInput
-		return t
-	case []*schema.AgenticMessage: // when callback is injected by graph node, not the component implementation itself, the input is the input of Chat Model interface, which is []*schema.AgenticMessage
-		return &CallbackInput{
-			Messages: t,
-		}
-	default:
-		return nil
-	}
-}
-
-// ConvCallbackOutput converts the callback output to the model callback output.
-func ConvCallbackOutput(src callbacks.CallbackOutput) *CallbackOutput {
-	switch t := src.(type) {
-	case *CallbackOutput: // when callback is triggered within component implementation, the output is usually already a typed *model.CallbackOutput
-		return t
-	case *schema.AgenticMessage: // when callback is injected by graph node, not the component implementation itself, the output is the output of Chat Model interface, which is *schema.AgenticMessage
-		return &CallbackOutput{
-			Message: t,
-		}
-	default:
-		return nil
-	}
-}

components/agentic/callback_extra_test.go

@@ -1,35 +0,0 @@
-/*
- * Copyright 2025 CloudWeGo Authors
- *
- * Licensed under the Apache License, Version 2.0 (the "License");
- * you may not use this file except in compliance with the License.
- * You may obtain a copy of the License at
- *
- *     http://www.apache.org/licenses/LICENSE-2.0
- *
- * Unless required by applicable law or agreed to in writing, software
- * distributed under the License is distributed on an "AS IS" BASIS,
- * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
- * See the License for the specific language governing permissions and
- * limitations under the License.
- */
-
-package agentic
-
-import (
-	"testing"
-
-	"github.com/stretchr/testify/assert"
-
-	"github.com/cloudwego/eino/schema"
-)
-
-func TestConvModel(t *testing.T) {
-	assert.NotNil(t, ConvCallbackInput(&CallbackInput{}))
-	assert.NotNil(t, ConvCallbackInput([]*schema.AgenticMessage{}))
-	assert.Nil(t, ConvCallbackInput("asd"))
-
-	assert.NotNil(t, ConvCallbackOutput(&CallbackOutput{}))
-	assert.NotNil(t, ConvCallbackOutput(&schema.AgenticMessage{}))
-	assert.Nil(t, ConvCallbackOutput("asd"))
-}

components/agentic/interface.go

@@ -1,32 +0,0 @@
-/*
- * Copyright 2025 CloudWeGo Authors
- *
- * Licensed under the Apache License, Version 2.0 (the "License");
- * you may not use this file except in compliance with the License.
- * You may obtain a copy of the License at
- *
- *     http://www.apache.org/licenses/LICENSE-2.0
- *
- * Unless required by applicable law or agreed to in writing, software
- * distributed under the License is distributed on an "AS IS" BASIS,
- * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
- * See the License for the specific language governing permissions and
- * limitations under the License.
- */
-
-package agentic
-
-import (
-	"context"
-
-	"github.com/cloudwego/eino/schema"
-)
-
-type Model interface {
-	Generate(ctx context.Context, input []*schema.AgenticMessage, opts ...Option) (*schema.AgenticMessage, error)
-	Stream(ctx context.Context, input []*schema.AgenticMessage, opts ...Option) (*schema.StreamReader[*schema.AgenticMessage], error)
-
-	// WithTools returns a new Model instance with the specified tools bound.
-	// This method does not modify the current instance, making it safer for concurrent use.
-	WithTools(tools []*schema.ToolInfo) (Model, error)
-}

components/agentic/option.go

@@ -1,142 +0,0 @@
-/*
- * Copyright 2025 CloudWeGo Authors
- *
- * Licensed under the Apache License, Version 2.0 (the "License");
- * you may not use this file except in compliance with the License.
- * You may obtain a copy of the License at
- *
- *     http://www.apache.org/licenses/LICENSE-2.0
- *
- * Unless required by applicable law or agreed to in writing, software
- * distributed under the License is distributed on an "AS IS" BASIS,
- * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
- * See the License for the specific language governing permissions and
- * limitations under the License.
- */
-
-package agentic
-
-import (
-	"github.com/cloudwego/eino/schema"
-)
-
-// Options is the common options for the model.
-type Options struct {
-	// Temperature is the temperature for the model, which controls the randomness of the model.
-	Temperature *float64
-	// Model is the model name.
-	Model *string
-	// TopP is the top p for the model, which controls the diversity of the model.
-	TopP *float64
-	// Tools is a list of tools the model may call.
-	Tools []*schema.ToolInfo
-	// ToolChoice controls how the model call the tools.
-	ToolChoice *schema.ToolChoice
-	// AllowedTools is a list of allowed tools the model may call.
-	AllowedTools []*schema.AllowedTool
-}
-
-// Option is the call option for ChatModel component.
-type Option struct {
-	apply func(opts *Options)
-
-	implSpecificOptFn any
-}
-
-// WithTemperature is the option to set the temperature for the model.
-func WithTemperature(temperature float64) Option {
-	return Option{
-		apply: func(opts *Options) {
-			opts.Temperature = &temperature
-		},
-	}
-}
-
-// WithModel is the option to set the model name.
-func WithModel(name string) Option {
-	return Option{
-		apply: func(opts *Options) {
-			opts.Model = &name
-		},
-	}
-}
-
-// WithTopP is the option to set the top p for the model.
-func WithTopP(topP float64) Option {
-	return Option{
-		apply: func(opts *Options) {
-			opts.TopP = &topP
-		},
-	}
-}
-
-// WithTools is the option to set tools for the model.
-func WithTools(tools []*schema.ToolInfo) Option {
-	if tools == nil {
-		tools = []*schema.ToolInfo{}
-	}
-	return Option{
-		apply: func(opts *Options) {
-			opts.Tools = tools
-		},
-	}
-}
-
-// WithToolChoice is the option to set tool choice for the model.
-func WithToolChoice(toolChoice schema.ToolChoice, allowedTools ...*schema.AllowedTool) Option {
-	return Option{
-		apply: func(opts *Options) {
-			opts.ToolChoice = &toolChoice
-			opts.AllowedTools = allowedTools
-		},
-	}
-}
-
-// WrapImplSpecificOptFn is the option to wrap the implementation specific option function.
-func WrapImplSpecificOptFn[T any](optFn func(*T)) Option {
-	return Option{
-		implSpecificOptFn: optFn,
-	}
-}
-
-// GetCommonOptions extract model Options from Option list, optionally providing a base Options with default values.
-func GetCommonOptions(base *Options, opts ...Option) *Options {
-	if base == nil {
-		base = &Options{}
-	}
-
-	for i := range opts {
-		opt := opts[i]
-		if opt.apply != nil {
-			opt.apply(base)
-		}
-	}
-
-	return base
-}
-
-// GetImplSpecificOptions extract the implementation specific options from Option list, optionally providing a base options with default values.
-// e.g.
-//
-//	myOption := &MyOption{
-//		Field1: "default_value",
-//	}
-//
-//	myOption := model.GetImplSpecificOptions(myOption, opts...)
-func GetImplSpecificOptions[T any](base *T, opts ...Option) *T {
-	if base == nil {
-		base = new(T)
-	}
-
-	for i := range opts {
-		opt := opts[i]
-		if opt.implSpecificOptFn != nil {
-			optFn, ok := opt.implSpecificOptFn.(func(*T))
-			if ok {
-				optFn(base)
-			}
-		}
-	}
-
-	return base
-}

components/agentic/option_test.go

@@ -1,79 +0,0 @@
-/*
- * Copyright 2025 CloudWeGo Authors
- *
- * Licensed under the Apache License, Version 2.0 (the "License");
- * you may not use this file except in compliance with the License.
- * You may obtain a copy of the License at
- *
- *     http://www.apache.org/licenses/LICENSE-2.0
- *
- * Unless required by applicable law or agreed to in writing, software
- * distributed under the License is distributed on an "AS IS" BASIS,
- * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
- * See the License for the specific language governing permissions and
- * limitations under the License.
- */
-
-package agentic
-
-import (
-	"testing"
-
-	"github.com/stretchr/testify/assert"
-
-	"github.com/cloudwego/eino/schema"
-)
-
-func TestCommon(t *testing.T) {
-	o := GetCommonOptions(nil,
-		WithTools([]*schema.ToolInfo{{Name: "test"}}),
-		WithModel("test"),
-		WithTemperature(0.1),
-		WithToolChoice(schema.ToolChoiceAllowed, []*schema.AllowedTool{{FunctionToolName: "test"}}...),
-		WithTopP(0.1),
-	)
-	assert.Len(t, o.Tools, 1)
-	assert.Equal(t, "test", o.Tools[0].Name)
-	assert.Equal(t, "test", *o.Model)
-	assert.Equal(t, float64(0.1), *o.Temperature)
-	assert.Equal(t, schema.ToolChoiceAllowed, *o.ToolChoice)
-	assert.Equal(t, float64(0.1), *o.TopP)
-}
-
-func TestImplSpecificOpts(t *testing.T) {
-	type implSpecificOptions struct {
-		conf  string
-		index int
-	}
-
-	withConf := func(conf string) func(o *implSpecificOptions) {
-		return func(o *implSpecificOptions) {
-			o.conf = conf
-		}
-	}
-
-	withIndex := func(index int) func(o *implSpecificOptions) {
-		return func(o *implSpecificOptions) {
-			o.index = index
-		}
-	}
-
-	documentOption1 := WrapImplSpecificOptFn(withConf("test_conf"))
-	documentOption2 := WrapImplSpecificOptFn(withIndex(1))
-
-	implSpecificOpts := GetImplSpecificOptions(&implSpecificOptions{}, documentOption1, documentOption2)
-
-	assert.Equal(t, &implSpecificOptions{
-		conf:  "test_conf",
-		index: 1,
-	}, implSpecificOpts)
-	documentOption1 = WrapImplSpecificOptFn(withConf("test_conf"))
-	documentOption2 = WrapImplSpecificOptFn(withIndex(1))
-
-	implSpecificOpts = GetImplSpecificOptions(&implSpecificOptions{}, documentOption1, documentOption2)
-
-	assert.Equal(t, &implSpecificOptions{
-		conf:  "test_conf",
-		index: 1,
-	}, implSpecificOpts)
-}

components/model/callback_extra.go

@@ -31,15 +31,15 @@ type TokenUsage struct {
 	CompletionTokens int
 	// TotalTokens is the total number of tokens.
 	TotalTokens int
-	// CompletionTokensDetails is a breakdown of the completion tokens.
-	CompletionTokensDetails CompletionTokensDetails
+	// CompletionTokensDetails is breakdown of completion tokens.
+	CompletionTokensDetails CompletionTokensDetails `json:"completion_token_details"`
 }
 
 type CompletionTokensDetails struct {
 	// ReasoningTokens tokens generated by the model for reasoning.
 	// This is currently supported by OpenAI, Gemini, ARK and Qwen  chat models.
 	// For other models, this field will be 0.
-	ReasoningTokens int
+	ReasoningTokens int `json:"reasoning_tokens,omitempty"`
 }
 
 // PromptTokenDetails provides a breakdown of prompt token usage.
@@ -66,6 +66,8 @@ type Config struct {
 type CallbackInput struct {
 	// Messages is the messages to be sent to the model.
 	Messages []*schema.Message
+	// AgenticMessages is the agentic messages to be sent to the agentic model.
+	AgenticMessages []*schema.AgenticMessage
 	// Tools is the tools to be used in the model.
 	Tools []*schema.ToolInfo
 	// ToolChoice is the tool choice, which controls the tool to be used in the model.
@@ -80,6 +82,8 @@ type CallbackInput struct {
 type CallbackOutput struct {
 	// Message is the message generated by the model.
 	Message *schema.Message
+	// AgenticMessage is the agentic message generated by the agentic model.
+	AgenticMessage *schema.AgenticMessage
 	// Config is the config for the model.
 	Config *Config
 	// TokenUsage is the token usage of this request.
@@ -97,6 +101,10 @@ func ConvCallbackInput(src callbacks.CallbackInput) *CallbackInput {
 		return &CallbackInput{
 			Messages: t,
 		}
+	case []*schema.AgenticMessage: // when callback is injected by graph node, not the component implementation itself, the input is the input of Agentic Model interface, which is []*schema.AgenticMessage
+		return &CallbackInput{
+			AgenticMessages: t,
+		}
 	default:
 		return nil
 	}
@@ -111,6 +119,10 @@ func ConvCallbackOutput(src callbacks.CallbackOutput) *CallbackOutput {
 		return &CallbackOutput{
 			Message: t,
 		}
+	case *schema.AgenticMessage: // when callback is injected by graph node, not the component implementation itself, the output is the output of Agentic Model interface, which is *schema.AgenticMessage
+		return &CallbackOutput{
+			AgenticMessage: t,
+		}
 	default:
 		return nil
 	}

components/model/interface.go

@@ -89,3 +89,15 @@ type ToolCallingChatModel interface {
 	// This method does not modify the current instance, making it safer for concurrent use.
 	WithTools(tools []*schema.ToolInfo) (ToolCallingChatModel, error)
 }
+
+// AgenticModel defines the interface for agentic models that support AgenticMessage.
+// It provides methods for generating complete and streaming outputs, and supports
+// tool calling via the WithTools method.
+type AgenticModel interface {
+	Generate(ctx context.Context, input []*schema.AgenticMessage, opts ...Option) (*schema.AgenticMessage, error)
+	Stream(ctx context.Context, input []*schema.AgenticMessage, opts ...Option) (*schema.StreamReader[*schema.AgenticMessage], error)
+
+	// WithTools returns a new Model instance with the specified tools bound.
+	// This method does not modify the current instance, making it safer for concurrent use.
+	WithTools(tools []*schema.ToolInfo) (AgenticModel, error)
+}

components/model/option.go

@@ -22,21 +22,29 @@ import "github.com/cloudwego/eino/schema"
 type Options struct {
 	// Temperature is the temperature for the model, which controls the randomness of the model.
 	Temperature *float32
-	// MaxTokens is the max number of tokens, if reached the max tokens, the model will stop generating, and mostly return an finish reason of "length".
-	MaxTokens *int
 	// Model is the model name.
 	Model *string
 	// TopP is the top p for the model, which controls the diversity of the model.
 	TopP *float32
-	// Stop is the stop words for the model, which controls the stopping condition of the model.
-	Stop []string
 	// Tools is a list of tools the model may call.
 	Tools []*schema.ToolInfo
 	// ToolChoice controls which tool is called by the model.
 	ToolChoice *schema.ToolChoice
+
+	// Options only for chat model.
+
+	// MaxTokens is the max number of tokens, if reached the max tokens, the model will stop generating, and mostly return an finish reason of "length".
+	MaxTokens *int
 	// AllowedToolNames specifies a list of tool names that the model is allowed to call.
 	// This allows for constraining the model to a specific subset of the available tools.
 	AllowedToolNames []string
+	// Stop is the stop words for the model, which controls the stopping condition of the model.
+	Stop []string
+
+	// Options only for agentic model.
+
+	// AllowedTools is a list of allowed tools the model may call.
+	AllowedTools []*schema.AllowedTool
 }
 
 // Option is a call-time option for a ChatModel. Options are immutable and
@@ -59,6 +67,7 @@ func WithTemperature(temperature float32) Option {
 }
 
 // WithMaxTokens is the option to set the max tokens for the model.
+// Only available for ChatModel.
 func WithMaxTokens(maxTokens int) Option {
 	return Option{
 		apply: func(opts *Options) {
@@ -86,6 +95,7 @@ func WithTopP(topP float32) Option {
 }
 
 // WithStop is the option to set the stop words for the model.
+// Only available for ChatModel.
 func WithStop(stop []string) Option {
 	return Option{
 		apply: func(opts *Options) {
@@ -108,6 +118,7 @@ func WithTools(tools []*schema.ToolInfo) Option {
 
 // WithToolChoice sets the tool choice for the model. It also allows for providing a list of
 // tool names to constrain the model to a specific subset of the available tools.
+// Only available for ChatModel.
 func WithToolChoice(toolChoice schema.ToolChoice, allowedToolNames ...string) Option {
 	return Option{
 		apply: func(opts *Options) {
@@ -117,6 +128,18 @@ func WithToolChoice(toolChoice schema.ToolChoice, allowedToolNames ...string) Op
 	}
 }
 
+// WithAgenticToolChoice is the option to set tool choice for the agentic model.
+// Only available for AgenticModel.
+func WithAgenticToolChoice(toolChoice schema.ToolChoice, allowedTools ...*schema.AllowedTool) Option {
+	return Option{
+		apply: func(opts *Options) {
+			opts.ToolChoice = &toolChoice
+			opts.AllowedTools = allowedTools
+		},
+	}
+}
+
+// WrapImplSpecificOptFn is the option to wrap the implementation specific option function.
 // WrapImplSpecificOptFn wraps an implementation-specific option function into
 // an [Option] so it can be passed alongside standard options.
 //

components/model/option_test.go

@@ -82,6 +82,22 @@ func TestOptions(t *testing.T) {
 		convey.So(opts.Tools, convey.ShouldNotBeNil)
 		convey.So(len(opts.Tools), convey.ShouldEqual, 0)
 	})
+
+	convey.Convey("test agentic tool choice option", t, func() {
+		var (
+			toolChoice   = schema.ToolChoiceForced
+			allowedTools = []*schema.AllowedTool{
+				{FunctionToolName: "agentic_tool"},
+			}
+		)
+		opts := GetCommonOptions(
+			nil,
+			WithAgenticToolChoice(toolChoice, allowedTools...),
+		)
+
+		convey.So(opts.ToolChoice, convey.ShouldResemble, &toolChoice)
+		convey.So(opts.AllowedTools, convey.ShouldResemble, allowedTools)
+	})
 }
 
 type implOption struct {

components/prompt/agentic_chat_template.go

@@ -1,5 +1,5 @@
 /*
- * Copyright 2025 CloudWeGo Authors
+ * Copyright 2026 CloudWeGo Authors
  *
  * Licensed under the Apache License, Version 2.0 (the "License");
  * you may not use this file except in compliance with the License.
@@ -45,9 +45,9 @@ type DefaultAgenticChatTemplate struct {
 
 func (t *DefaultAgenticChatTemplate) Format(ctx context.Context, vs map[string]any, opts ...Option) (result []*schema.AgenticMessage, err error) {
 	ctx = callbacks.EnsureRunInfo(ctx, t.GetType(), components.ComponentOfAgenticPrompt)
-	ctx = callbacks.OnStart(ctx, &AgenticCallbackInput{
-		Variables: vs,
-		Templates: t.templates,
+	ctx = callbacks.OnStart(ctx, &CallbackInput{
+		Variables:        vs,
+		AgenticTemplates: t.templates,
 	})
 	defer func() {
 		if err != nil {
@@ -65,15 +65,15 @@ func (t *DefaultAgenticChatTemplate) Format(ctx context.Context, vs map[string]a
 		result = append(result, msgs...)
 	}
 
-	_ = callbacks.OnEnd(ctx, &AgenticCallbackOutput{
-		Result:    result,
-		Templates: t.templates,
+	_ = callbacks.OnEnd(ctx, &CallbackOutput{
+		AgenticResult:    result,
+		AgenticTemplates: t.templates,
 	})
 
 	return result, nil
 }
 
-// GetType returns the type of the chat template (Default).
+// GetType returns the type of the agentic template (DefaultAgentic).
 func (t *DefaultAgenticChatTemplate) GetType() string {
 	return "Default"
 }

components/prompt/agentic_chat_template_test.go

@@ -0,0 +1,124 @@
+/*
+ * Copyright 2026 CloudWeGo Authors
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package prompt
+
+import (
+	"context"
+	"errors"
+	"testing"
+
+	"github.com/cloudwego/eino/callbacks"
+	"github.com/cloudwego/eino/schema"
+	"github.com/stretchr/testify/assert"
+)
+
+type mockAgenticTemplate struct {
+	err error
+}
+
+func (m *mockAgenticTemplate) Format(ctx context.Context, vs map[string]any, formatType schema.FormatType) ([]*schema.AgenticMessage, error) {
+	if m.err != nil {
+		return nil, m.err
+	}
+	return []*schema.AgenticMessage{schema.UserAgenticMessage("mocked")}, nil
+}
+
+func TestFromAgenticMessages(t *testing.T) {
+	t.Run("create template", func(t *testing.T) {
+		tpl := schema.UserAgenticMessage("hello")
+		ft := schema.FString
+		at := FromAgenticMessages(ft, tpl)
+
+		assert.NotNil(t, at)
+		assert.Equal(t, ft, at.formatType)
+		assert.Len(t, at.templates, 1)
+		assert.Same(t, tpl, at.templates[0])
+	})
+}
+
+func TestDefaultAgenticTemplate_GetType(t *testing.T) {
+	t.Run("get type", func(t *testing.T) {
+		at := &DefaultAgenticChatTemplate{}
+		assert.Equal(t, "Default", at.GetType())
+	})
+}
+
+func TestDefaultAgenticTemplate_IsCallbacksEnabled(t *testing.T) {
+	t.Run("callbacks enabled", func(t *testing.T) {
+		at := &DefaultAgenticChatTemplate{}
+		assert.True(t, at.IsCallbacksEnabled())
+	})
+}
+
+func TestDefaultAgenticTemplate_Format(t *testing.T) {
+	t.Run("success", func(t *testing.T) {
+		// Mock callback handler
+		cb := callbacks.NewHandlerBuilder().
+			OnStartFn(func(ctx context.Context, info *callbacks.RunInfo, input callbacks.CallbackInput) context.Context {
+				assert.Equal(t, "Default", info.Type)
+				return ctx
+			}).
+			OnEndFn(func(ctx context.Context, info *callbacks.RunInfo, output callbacks.CallbackOutput) context.Context {
+				assert.Equal(t, "Default", info.Type)
+				return ctx
+			}).
+			OnErrorFn(func(ctx context.Context, info *callbacks.RunInfo, err error) context.Context {
+				assert.Fail(t, "unexpected error callback")
+				return ctx
+			}).
+			Build()
+
+		tpl := schema.UserAgenticMessage("hello {val}")
+		at := FromAgenticMessages(schema.FString, tpl)
+
+		ctx := context.Background()
+		ctx = callbacks.InitCallbacks(ctx, &callbacks.RunInfo{
+			Type:      "Default",
+			Component: "agentic_prompt",
+		}, cb)
+
+		res, err := at.Format(ctx, map[string]any{"val": "world"})
+		assert.NoError(t, err)
+		assert.Len(t, res, 1)
+		assert.Equal(t, "hello world", res[0].ContentBlocks[0].UserInputText.Text)
+	})
+
+	t.Run("template format error", func(t *testing.T) {
+		mockErr := errors.New("mock error")
+		mockTpl := &mockAgenticTemplate{err: mockErr}
+		at := FromAgenticMessages(schema.FString, mockTpl)
+
+		// Mock callback handler to verify OnError
+		cb := callbacks.NewHandlerBuilder().
+			OnErrorFn(func(ctx context.Context, info *callbacks.RunInfo, err error) context.Context {
+				assert.Equal(t, mockErr, err)
+				return ctx
+			}).
+			Build()
+
+		ctx := context.Background()
+		ctx = callbacks.InitCallbacks(ctx, &callbacks.RunInfo{
+			Type:      "Default",
+			Component: "agentic_prompt",
+		}, cb)
+
+		res, err := at.Format(ctx, map[string]any{})
+		assert.Error(t, err)
+		assert.Nil(t, res)
+		assert.Equal(t, mockErr, err)
+	})
+}

components/prompt/callback_extra.go

@@ -21,52 +21,14 @@ import (
 	"github.com/cloudwego/eino/schema"
 )
 
-type AgenticCallbackInput struct {
-	Variables map[string]any
-	Templates []schema.AgenticMessagesTemplate
-	Extra     map[string]any
-}
-
-type AgenticCallbackOutput struct {
-	Result    []*schema.AgenticMessage
-	Templates []schema.AgenticMessagesTemplate
-	Extra     map[string]any
-}
-
-// ConvAgenticCallbackInput converts the callback input to the agentic callback input.
-func ConvAgenticCallbackInput(src callbacks.CallbackInput) *AgenticCallbackInput {
-	switch t := src.(type) {
-	case *AgenticCallbackInput:
-		return t
-	case map[string]any:
-		return &AgenticCallbackInput{
-			Variables: t,
-		}
-	default:
-		return nil
-	}
-}
-
-// ConvAgenticCallbackOutput converts the callback output to the agentic callback output.
-func ConvAgenticCallbackOutput(src callbacks.CallbackOutput) *AgenticCallbackOutput {
-	switch t := src.(type) {
-	case *AgenticCallbackOutput:
-		return t
-	case []*schema.AgenticMessage:
-		return &AgenticCallbackOutput{
-			Result: t,
-		}
-	default:
-		return nil
-	}
-}
-
 // CallbackInput is the input for the callback.
 type CallbackInput struct {
 	// Variables is the variables for the callback.
 	Variables map[string]any
 	// Templates is the templates for the callback.
 	Templates []schema.MessagesTemplate
+	// AgenticTemplates is the agentic templates for the callback.
+	AgenticTemplates []schema.AgenticMessagesTemplate
 	// Extra is the extra information for the callback.
 	Extra map[string]any
 }
@@ -75,8 +37,12 @@ type CallbackInput struct {
 type CallbackOutput struct {
 	// Result is the result for the callback.
 	Result []*schema.Message
+	// AgenticResult is the agentic result for the callback.
+	AgenticResult []*schema.AgenticMessage
 	// Templates is the templates for the callback.
 	Templates []schema.MessagesTemplate
+	// AgenticTemplates is the agentic templates for the callback.
+	AgenticTemplates []schema.AgenticMessagesTemplate
 	// Extra is the extra information for the callback.
 	Extra map[string]any
 }
@@ -104,6 +70,10 @@ func ConvCallbackOutput(src callbacks.CallbackOutput) *CallbackOutput {
 		return &CallbackOutput{
 			Result: t,
 		}
+	case []*schema.AgenticMessage:
+		return &CallbackOutput{
+			AgenticResult: t,
+		}
 	default:
 		return nil
 	}

components/prompt/callback_extra_test.go

@@ -25,11 +25,28 @@ import (
 )
 
 func TestConvPrompt(t *testing.T) {
-	assert.NotNil(t, ConvCallbackInput(&CallbackInput{}))
+	assert.NotNil(t, ConvCallbackInput(&CallbackInput{
+		AgenticTemplates: []schema.AgenticMessagesTemplate{
+			&schema.AgenticMessage{},
+		},
+	}))
 	assert.NotNil(t, ConvCallbackInput(map[string]any{}))
 	assert.Nil(t, ConvCallbackInput("asd"))
 
-	assert.NotNil(t, ConvCallbackOutput(&CallbackOutput{}))
+	assert.NotNil(t, ConvCallbackOutput(&CallbackOutput{
+		AgenticResult: []*schema.AgenticMessage{
+			{},
+		},
+		AgenticTemplates: []schema.AgenticMessagesTemplate{
+			&schema.AgenticMessage{},
+		},
+	}))
 	assert.NotNil(t, ConvCallbackOutput([]*schema.Message{}))
+
+	agenticResult := []*schema.AgenticMessage{{}}
+	out := ConvCallbackOutput(agenticResult)
+	assert.NotNil(t, out)
+	assert.Equal(t, agenticResult, out.AgenticResult)
+
 	assert.Nil(t, ConvCallbackOutput("asd"))
 }

components/prompt/chat_template_agentic_test.go

@@ -1,111 +0,0 @@
-/*
- * Copyright 2025 CloudWeGo Authors
- *
- * Licensed under the Apache License, Version 2.0 (the "License");
- * you may not use this file except in compliance with the License.
- * You may obtain a copy of the License at
- *
- *     http://www.apache.org/licenses/LICENSE-2.0
- *
- * Unless required by applicable law or agreed to in writing, software
- * distributed under the License is distributed on an "AS IS" BASIS,
- * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
- * See the License for the specific language governing permissions and
- * limitations under the License.
- */
-
-package prompt
-
-import (
-	"context"
-	"testing"
-
-	"github.com/stretchr/testify/assert"
-
-	"github.com/cloudwego/eino/schema"
-)
-
-func TestAgenticFormat(t *testing.T) {
-	pyFmtTestTemplate := []schema.AgenticMessagesTemplate{
-		&schema.AgenticMessage{
-			Role: schema.AgenticRoleTypeUser,
-			ContentBlocks: []*schema.ContentBlock{
-				{Type: schema.ContentBlockTypeUserInputText, UserInputText: &schema.UserInputText{Text: "{context}"}},
-			},
-		},
-		schema.AgenticMessagesPlaceholder("chat_history", true),
-	}
-	jinja2TestTemplate := []schema.AgenticMessagesTemplate{
-		&schema.AgenticMessage{
-			Role: schema.AgenticRoleTypeUser,
-			ContentBlocks: []*schema.ContentBlock{
-				{Type: schema.ContentBlockTypeUserInputText, UserInputText: &schema.UserInputText{Text: "{{context}}"}},
-			},
-		},
-		schema.AgenticMessagesPlaceholder("chat_history", true),
-	}
-	goFmtTestTemplate := []schema.AgenticMessagesTemplate{
-		&schema.AgenticMessage{
-			Role: schema.AgenticRoleTypeUser,
-			ContentBlocks: []*schema.ContentBlock{
-				{Type: schema.ContentBlockTypeUserInputText, UserInputText: &schema.UserInputText{Text: "{{.context}}"}},
-			},
-		},
-		schema.AgenticMessagesPlaceholder("chat_history", true),
-	}
-	testValues := map[string]any{
-		"context": "it's beautiful day",
-		"chat_history": []*schema.AgenticMessage{
-			{
-				Role: schema.AgenticRoleTypeUser,
-				ContentBlocks: []*schema.ContentBlock{
-					{Type: schema.ContentBlockTypeUserInputText, UserInputText: &schema.UserInputText{Text: "1"}},
-				},
-			},
-			{
-				Role: schema.AgenticRoleTypeUser,
-				ContentBlocks: []*schema.ContentBlock{
-					{Type: schema.ContentBlockTypeUserInputText, UserInputText: &schema.UserInputText{Text: "2"}},
-				},
-			},
-		},
-	}
-	expected := []*schema.AgenticMessage{
-		{
-			Role: schema.AgenticRoleTypeUser,
-			ContentBlocks: []*schema.ContentBlock{
-				{Type: schema.ContentBlockTypeUserInputText, UserInputText: &schema.UserInputText{Text: "it's beautiful day"}},
-			},
-		},
-		{
-			Role: schema.AgenticRoleTypeUser,
-			ContentBlocks: []*schema.ContentBlock{
-				{Type: schema.ContentBlockTypeUserInputText, UserInputText: &schema.UserInputText{Text: "1"}},
-			},
-		},
-		{
-			Role: schema.AgenticRoleTypeUser,
-			ContentBlocks: []*schema.ContentBlock{
-				{Type: schema.ContentBlockTypeUserInputText, UserInputText: &schema.UserInputText{Text: "2"}},
-			},
-		},
-	}
-
-	// FString
-	chatTemplate := FromAgenticMessages(schema.FString, pyFmtTestTemplate...)
-	msgs, err := chatTemplate.Format(context.Background(), testValues)
-	assert.Nil(t, err)
-	assert.Equal(t, expected, msgs)
-
-	// Jinja2
-	chatTemplate = FromAgenticMessages(schema.Jinja2, jinja2TestTemplate...)
-	msgs, err = chatTemplate.Format(context.Background(), testValues)
-	assert.Nil(t, err)
-	assert.Equal(t, expected, msgs)
-
-	// GoTemplate
-	chatTemplate = FromAgenticMessages(schema.GoTemplate, goFmtTestTemplate...)
-	msgs, err = chatTemplate.Format(context.Background(), testValues)
-	assert.Nil(t, err)
-	assert.Equal(t, expected, msgs)
-}

components/prompt/interface.go

@@ -44,6 +44,7 @@ type ChatTemplate interface {
 	Format(ctx context.Context, vs map[string]any, opts ...Option) ([]*schema.Message, error)
 }
 
+// AgenticChatTemplate formats variables into a list of agentic messages according to a prompt schema.
 type AgenticChatTemplate interface {
 	Format(ctx context.Context, vs map[string]any, opts ...Option) ([]*schema.AgenticMessage, error)
 }

compose/chain.go

@@ -22,7 +22,6 @@ import (
 	"fmt"
 	"reflect"
 
-	"github.com/cloudwego/eino/components/agentic"
 	"github.com/cloudwego/eino/components/document"
 	"github.com/cloudwego/eino/components/embedding"
 	"github.com/cloudwego/eino/components/indexer"
@@ -181,7 +180,7 @@ func (c *Chain[I, O]) AppendChatModel(node model.BaseChatModel, opts ...GraphAdd
 //	model, err := openai.NewAgenticModel(ctx, config)
 //	if err != nil {...}
 //	chain.AppendAgenticModel(model)
-func (c *Chain[I, O]) AppendAgenticModel(node agentic.Model, opts ...GraphAddNodeOpt) *Chain[I, O] {
+func (c *Chain[I, O]) AppendAgenticModel(node model.AgenticModel, opts ...GraphAddNodeOpt) *Chain[I, O] {
 	gNode, options := toAgenticModelNode(node, opts...)
 	c.addNode(gNode, options)
 	return c

compose/chain_branch.go

@@ -20,7 +20,6 @@ import (
 	"context"
 	"fmt"
 
-	"github.com/cloudwego/eino/components/agentic"
 	"github.com/cloudwego/eino/components/document"
 	"github.com/cloudwego/eino/components/embedding"
 	"github.com/cloudwego/eino/components/indexer"
@@ -158,7 +157,7 @@ func (cb *ChainBranch) AddChatModel(key string, node model.BaseChatModel, opts .
 //	})
 //	cb.AddAgenticModel("agentic_model_key_1", model1)
 //	cb.AddAgenticModel("agentic_model_key_2", model2)
-func (cb *ChainBranch) AddAgenticModel(key string, node agentic.Model, opts ...GraphAddNodeOpt) *ChainBranch {
+func (cb *ChainBranch) AddAgenticModel(key string, node model.AgenticModel, opts ...GraphAddNodeOpt) *ChainBranch {
 	gNode, options := toAgenticModelNode(node, opts...)
 	return cb.addNode(key, gNode, options)
 }

compose/chain_parallel.go

@@ -19,7 +19,6 @@ package compose
 import (
 	"fmt"
 
-	"github.com/cloudwego/eino/components/agentic"
 	"github.com/cloudwego/eino/components/document"
 	"github.com/cloudwego/eino/components/embedding"
 	"github.com/cloudwego/eino/components/indexer"
@@ -84,7 +83,7 @@ func (p *Parallel) AddChatModel(outputKey string, node model.BaseChatModel, opts
 //
 //	p.AddAgenticModel("output_key1", model1)
 //	p.AddAgenticModel("output_key2", model2)
-func (p *Parallel) AddAgenticModel(outputKey string, node agentic.Model, opts ...GraphAddNodeOpt) *Parallel {
+func (p *Parallel) AddAgenticModel(outputKey string, node model.AgenticModel, opts ...GraphAddNodeOpt) *Parallel {
 	gNode, options := toAgenticModelNode(node, append(opts, WithOutputKey(outputKey))...)
 	return p.addNode(outputKey, gNode, options)
 }

compose/component_to_graph_node.go

@@ -18,7 +18,6 @@ package compose
 
 import (
 	"github.com/cloudwego/eino/components"
-	"github.com/cloudwego/eino/components/agentic"
 	"github.com/cloudwego/eino/components/document"
 	"github.com/cloudwego/eino/components/embedding"
 	"github.com/cloudwego/eino/components/indexer"
@@ -102,7 +101,7 @@ func toChatModelNode(node model.BaseChatModel, opts ...GraphAddNodeOpt) (*graphN
 		opts...)
 }
 
-func toAgenticModelNode(node agentic.Model, opts ...GraphAddNodeOpt) (*graphNode, *graphAddNodeOpts) {
+func toAgenticModelNode(node model.AgenticModel, opts ...GraphAddNodeOpt) (*graphNode, *graphAddNodeOpts) {
 	return toComponentNode(
 		node,
 		components.ComponentOfAgenticModel,

compose/graph.go

@@ -23,7 +23,6 @@ import (
 	"reflect"
 	"strings"
 
-	"github.com/cloudwego/eino/components/agentic"
 	"github.com/cloudwego/eino/components/document"
 	"github.com/cloudwego/eino/components/embedding"
 	"github.com/cloudwego/eino/components/indexer"
@@ -361,7 +360,7 @@ func (g *graph) AddChatModelNode(key string, node model.BaseChatModel, opts ...G
 //	})
 //
 //	graph.AddAgenticModelNode("agentic_model_node_key", model)
-func (g *graph) AddAgenticModelNode(key string, node agentic.Model, opts ...GraphAddNodeOpt) error {
+func (g *graph) AddAgenticModelNode(key string, node model.AgenticModel, opts ...GraphAddNodeOpt) error {
 	gNode, options := toAgenticModelNode(node, opts...)
 	return g.addNode(key, gNode, options)
 }

schema/agentic_message.go

@@ -66,15 +66,15 @@ const (
 )
 
 type AgenticMessage struct {
-	// ResponseMeta is the response metadata.
-	ResponseMeta *AgenticResponseMeta
-
 	// Role is the message role.
 	Role AgenticRoleType
 
 	// ContentBlocks is the list of content blocks.
 	ContentBlocks []*ContentBlock
 
+	// ResponseMeta is the response metadata.
+	ResponseMeta *AgenticResponseMeta
+
 	// Extra is the additional information.
 	Extra map[string]any
 }
@@ -541,7 +541,7 @@ func NewContentBlockChunk[T contentBlockVariant](content *T, meta *StreamingMeta
 	return block
 }
 
-// AgenticMessagesTemplate is the interface for messages template.
+// AgenticMessagesTemplate is the interface for agentic messages template.
 // It's used to render a template to a list of agentic messages.
 // e.g.
 //

utils/callbacks/template.go

@@ -55,17 +55,20 @@ func NewHandlerHelper() *HandlerHelper {
 //
 // then use the handler with runnable.Invoke(ctx, input, compose.WithCallbacks(handler))
 type HandlerHelper struct {
-	promptHandler      *PromptCallbackHandler
-	chatModelHandler   *ModelCallbackHandler
-	embeddingHandler   *EmbeddingCallbackHandler
-	indexerHandler     *IndexerCallbackHandler
-	retrieverHandler   *RetrieverCallbackHandler
-	loaderHandler      *LoaderCallbackHandler
-	transformerHandler *TransformerCallbackHandler
-	toolHandler        *ToolCallbackHandler
-	toolsNodeHandler   *ToolsNodeCallbackHandlers
-	agentHandler       *AgentCallbackHandler
-	composeTemplates   map[components.Component]callbacks.Handler
+	promptHandler           *PromptCallbackHandler
+	chatModelHandler        *ModelCallbackHandler
+	embeddingHandler        *EmbeddingCallbackHandler
+	indexerHandler          *IndexerCallbackHandler
+	retrieverHandler        *RetrieverCallbackHandler
+	loaderHandler           *LoaderCallbackHandler
+	transformerHandler      *TransformerCallbackHandler
+	toolHandler             *ToolCallbackHandler
+	toolsNodeHandler        *ToolsNodeCallbackHandlers
+	agenticPromptHandler    *AgenticPromptCallbackHandler
+	agenticModelHandler     *AgenticModelCallbackHandler
+	agenticToolsNodeHandler *AgenticToolsNodeCallbackHandlers
+	agentHandler      		*AgentCallbackHandler
+	composeTemplates        map[components.Component]callbacks.Handler
 }
 
 // Handler returns the callbacks.Handler created by HandlerHelper.
@@ -127,6 +130,24 @@ func (c *HandlerHelper) ToolsNode(handler *ToolsNodeCallbackHandlers) *HandlerHe
 	return c
 }
 
+// AgenticPrompt sets the agentic prompt handler for the handler helper, which will be called when the agentic prompt component is executed.
+func (c *HandlerHelper) AgenticPrompt(handler *AgenticPromptCallbackHandler) *HandlerHelper {
+	c.agenticPromptHandler = handler
+	return c
+}
+
+// AgenticModel sets the agentic chat model handler for the handler helper, which will be called when the agentic chat model component is executed.
+func (c *HandlerHelper) AgenticModel(handler *AgenticModelCallbackHandler) *HandlerHelper {
+	c.agenticModelHandler = handler
+	return c
+}
+
+// AgenticToolsNode sets the agentic tools node handler for the handler helper, which will be called when the agentic tools node is executed.
+func (c *HandlerHelper) AgenticToolsNode(handler *AgenticToolsNodeCallbackHandlers) *HandlerHelper {
+	c.agenticToolsNodeHandler = handler
+	return c
+}
+
 // Agent sets the agent handler for the handler helper, which will be called when the agent is executed.
 func (c *HandlerHelper) Agent(handler *AgentCallbackHandler) *HandlerHelper {
 	c.agentHandler = handler
@@ -161,8 +182,12 @@ func (c *handlerTemplate) OnStart(ctx context.Context, info *callbacks.RunInfo,
 	switch info.Component {
 	case components.ComponentOfPrompt:
 		return c.promptHandler.OnStart(ctx, info, prompt.ConvCallbackInput(input))
+	case components.ComponentOfAgenticPrompt:
+		return c.agenticPromptHandler.OnStart(ctx, info, prompt.ConvCallbackInput(input))
 	case components.ComponentOfChatModel:
 		return c.chatModelHandler.OnStart(ctx, info, model.ConvCallbackInput(input))
+	case components.ComponentOfAgenticModel:
+		return c.agenticModelHandler.OnStart(ctx, info, model.ConvCallbackInput(input))
 	case components.ComponentOfEmbedding:
 		return c.embeddingHandler.OnStart(ctx, info, embedding.ConvCallbackInput(input))
 	case components.ComponentOfIndexer:
@@ -177,6 +202,8 @@ func (c *handlerTemplate) OnStart(ctx context.Context, info *callbacks.RunInfo,
 		return c.toolHandler.OnStart(ctx, info, tool.ConvCallbackInput(input))
 	case compose.ComponentOfToolsNode:
 		return c.toolsNodeHandler.OnStart(ctx, info, convToolsNodeCallbackInput(input))
+	case compose.ComponentOfAgenticToolsNode:
+		return c.agenticToolsNodeHandler.OnStart(ctx, info, convAgenticToolsNodeCallbackInput(input))
 	case adk.ComponentOfAgent:
 		return c.agentHandler.OnStart(ctx, info, adk.ConvAgentCallbackInput(input))
 	case compose.ComponentOfGraph,
@@ -194,8 +221,12 @@ func (c *handlerTemplate) OnEnd(ctx context.Context, info *callbacks.RunInfo, ou
 	switch info.Component {
 	case components.ComponentOfPrompt:
 		return c.promptHandler.OnEnd(ctx, info, prompt.ConvCallbackOutput(output))
+	case components.ComponentOfAgenticPrompt:
+		return c.agenticPromptHandler.OnEnd(ctx, info, prompt.ConvCallbackOutput(output))
 	case components.ComponentOfChatModel:
 		return c.chatModelHandler.OnEnd(ctx, info, model.ConvCallbackOutput(output))
+	case components.ComponentOfAgenticModel:
+		return c.agenticModelHandler.OnEnd(ctx, info, model.ConvCallbackOutput(output))
 	case components.ComponentOfEmbedding:
 		return c.embeddingHandler.OnEnd(ctx, info, embedding.ConvCallbackOutput(output))
 	case components.ComponentOfIndexer:
@@ -210,6 +241,8 @@ func (c *handlerTemplate) OnEnd(ctx context.Context, info *callbacks.RunInfo, ou
 		return c.toolHandler.OnEnd(ctx, info, tool.ConvCallbackOutput(output))
 	case compose.ComponentOfToolsNode:
 		return c.toolsNodeHandler.OnEnd(ctx, info, convToolsNodeCallbackOutput(output))
+	case compose.ComponentOfAgenticToolsNode:
+		return c.agenticToolsNodeHandler.OnEnd(ctx, info, convAgenticToolsNodeCallbackOutput(output))
 	case adk.ComponentOfAgent:
 		return c.agentHandler.OnEnd(ctx, info, adk.ConvAgentCallbackOutput(output))
 	case compose.ComponentOfGraph,
@@ -227,8 +260,12 @@ func (c *handlerTemplate) OnError(ctx context.Context, info *callbacks.RunInfo,
 	switch info.Component {
 	case components.ComponentOfPrompt:
 		return c.promptHandler.OnError(ctx, info, err)
+	case components.ComponentOfAgenticPrompt:
+		return c.agenticPromptHandler.OnError(ctx, info, err)
 	case components.ComponentOfChatModel:
 		return c.chatModelHandler.OnError(ctx, info, err)
+	case components.ComponentOfAgenticModel:
+		return c.agenticModelHandler.OnError(ctx, info, err)
 	case components.ComponentOfEmbedding:
 		return c.embeddingHandler.OnError(ctx, info, err)
 	case components.ComponentOfIndexer:
@@ -243,6 +280,8 @@ func (c *handlerTemplate) OnError(ctx context.Context, info *callbacks.RunInfo,
 		return c.toolHandler.OnError(ctx, info, err)
 	case compose.ComponentOfToolsNode:
 		return c.toolsNodeHandler.OnError(ctx, info, err)
+	case compose.ComponentOfAgenticToolsNode:
+		return c.agenticToolsNodeHandler.OnError(ctx, info, err)
 	case compose.ComponentOfGraph,
 		compose.ComponentOfChain,
 		compose.ComponentOfLambda:
@@ -275,6 +314,11 @@ func (c *handlerTemplate) OnEndWithStreamOutput(ctx context.Context, info *callb
 			schema.StreamReaderWithConvert(output, func(item callbacks.CallbackOutput) (*model.CallbackOutput, error) {
 				return model.ConvCallbackOutput(item), nil
 			}))
+	case components.ComponentOfAgenticModel:
+		return c.agenticModelHandler.OnEndWithStreamOutput(ctx, info,
+			schema.StreamReaderWithConvert(output, func(item callbacks.CallbackOutput) (*model.CallbackOutput, error) {
+				return model.ConvCallbackOutput(item), nil
+			}))
 	case components.ComponentOfTool:
 		return c.toolHandler.OnEndWithStreamOutput(ctx, info,
 			schema.StreamReaderWithConvert(output, func(item callbacks.CallbackOutput) (*tool.CallbackOutput, error) {
@@ -285,6 +329,11 @@ func (c *handlerTemplate) OnEndWithStreamOutput(ctx context.Context, info *callb
 			schema.StreamReaderWithConvert(output, func(item callbacks.CallbackOutput) ([]*schema.Message, error) {
 				return convToolsNodeCallbackOutput(item), nil
 			}))
+	case compose.ComponentOfAgenticToolsNode:
+		return c.agenticToolsNodeHandler.OnEndWithStreamOutput(ctx, info,
+			schema.StreamReaderWithConvert(output, func(item callbacks.CallbackOutput) ([]*schema.AgenticMessage, error) {
+				return convAgenticToolsNodeCallbackOutput(item), nil
+			}))
 	case compose.ComponentOfGraph,
 		compose.ComponentOfChain,
 		compose.ComponentOfLambda:
@@ -295,6 +344,8 @@ func (c *handlerTemplate) OnEndWithStreamOutput(ctx context.Context, info *callb
 }
 
 // Needed checks if the callback handler is needed for the given timing.
+//
+//nolint:cyclop
 func (c *handlerTemplate) Needed(ctx context.Context, info *callbacks.RunInfo, timing callbacks.CallbackTiming) bool {
 	if info == nil {
 		return false
@@ -305,6 +356,10 @@ func (c *handlerTemplate) Needed(ctx context.Context, info *callbacks.RunInfo, t
 		if c.chatModelHandler != nil && c.chatModelHandler.Needed(ctx, info, timing) {
 			return true
 		}
+	case components.ComponentOfAgenticModel:
+		if c.agenticModelHandler != nil && c.agenticModelHandler.Needed(ctx, info, timing) {
+			return true
+		}
 	case components.ComponentOfEmbedding:
 		if c.embeddingHandler != nil && c.embeddingHandler.Needed(ctx, info, timing) {
 			return true
@@ -321,6 +376,10 @@ func (c *handlerTemplate) Needed(ctx context.Context, info *callbacks.RunInfo, t
 		if c.promptHandler != nil && c.promptHandler.Needed(ctx, info, timing) {
 			return true
 		}
+	case components.ComponentOfAgenticPrompt:
+		if c.agenticPromptHandler != nil && c.agenticPromptHandler.Needed(ctx, info, timing) {
+			return true
+		}
 	case components.ComponentOfRetriever:
 		if c.retrieverHandler != nil && c.retrieverHandler.Needed(ctx, info, timing) {
 			return true
@@ -337,6 +396,10 @@ func (c *handlerTemplate) Needed(ctx context.Context, info *callbacks.RunInfo, t
 		if c.toolsNodeHandler != nil && c.toolsNodeHandler.Needed(ctx, info, timing) {
 			return true
 		}
+	case compose.ComponentOfAgenticToolsNode:
+		if c.agenticToolsNodeHandler != nil && c.agenticToolsNodeHandler.Needed(ctx, info, timing) {
+			return true
+		}
 	case adk.ComponentOfAgent:
 		if c.agentHandler != nil && c.agentHandler.Needed(ctx, info, timing) {
 			return true
@@ -596,3 +659,94 @@ func (ch *AgentCallbackHandler) Needed(ctx context.Context, info *callbacks.RunI
 		return false
 	}
 }
+
+// AgenticPromptCallbackHandler is the handler for the agentic prompt callback.
+type AgenticPromptCallbackHandler struct {
+	// OnStart is the callback function for the start of the agentic prompt.
+	OnStart func(ctx context.Context, runInfo *callbacks.RunInfo, input *prompt.CallbackInput) context.Context
+	// OnEnd is the callback function for the end of the agentic prompt.
+	OnEnd func(ctx context.Context, runInfo *callbacks.RunInfo, output *prompt.CallbackOutput) context.Context
+	// OnError is the callback function for the error of the agentic prompt.
+	OnError func(ctx context.Context, runInfo *callbacks.RunInfo, err error) context.Context
+}
+
+// Needed checks if the callback handler is needed for the given timing.
+func (ch *AgenticPromptCallbackHandler) Needed(ctx context.Context, runInfo *callbacks.RunInfo, timing callbacks.CallbackTiming) bool {
+	switch timing {
+	case callbacks.TimingOnStart:
+		return ch.OnStart != nil
+	case callbacks.TimingOnEnd:
+		return ch.OnEnd != nil
+	case callbacks.TimingOnError:
+		return ch.OnError != nil
+	default:
+		return false
+	}
+}
+
+// AgenticModelCallbackHandler is the handler for the agentic chat model callback.
+type AgenticModelCallbackHandler struct {
+	OnStart               func(ctx context.Context, runInfo *callbacks.RunInfo, input *model.CallbackInput) context.Context
+	OnEnd                 func(ctx context.Context, runInfo *callbacks.RunInfo, output *model.CallbackOutput) context.Context
+	OnEndWithStreamOutput func(ctx context.Context, runInfo *callbacks.RunInfo, output *schema.StreamReader[*model.CallbackOutput]) context.Context
+	OnError               func(ctx context.Context, runInfo *callbacks.RunInfo, err error) context.Context
+}
+
+// Needed checks if the callback handler is needed for the given timing.
+func (ch *AgenticModelCallbackHandler) Needed(ctx context.Context, runInfo *callbacks.RunInfo, timing callbacks.CallbackTiming) bool {
+	switch timing {
+	case callbacks.TimingOnStart:
+		return ch.OnStart != nil
+	case callbacks.TimingOnEnd:
+		return ch.OnEnd != nil
+	case callbacks.TimingOnError:
+		return ch.OnError != nil
+	case callbacks.TimingOnEndWithStreamOutput:
+		return ch.OnEndWithStreamOutput != nil
+	default:
+		return false
+	}
+}
+
+// AgenticToolsNodeCallbackHandlers defines optional callbacks for the Agentic Tools node
+// lifecycle events.
+type AgenticToolsNodeCallbackHandlers struct {
+	OnStart               func(ctx context.Context, info *callbacks.RunInfo, input *schema.AgenticMessage) context.Context
+	OnEnd                 func(ctx context.Context, info *callbacks.RunInfo, input []*schema.AgenticMessage) context.Context
+	OnEndWithStreamOutput func(ctx context.Context, info *callbacks.RunInfo, output *schema.StreamReader[[]*schema.AgenticMessage]) context.Context
+	OnError               func(ctx context.Context, info *callbacks.RunInfo, err error) context.Context
+}
+
+// Needed reports whether a handler is registered for the given timing.
+func (ch *AgenticToolsNodeCallbackHandlers) Needed(ctx context.Context, runInfo *callbacks.RunInfo, timing callbacks.CallbackTiming) bool {
+	switch timing {
+	case callbacks.TimingOnStart:
+		return ch.OnStart != nil
+	case callbacks.TimingOnEnd:
+		return ch.OnEnd != nil
+	case callbacks.TimingOnEndWithStreamOutput:
+		return ch.OnEndWithStreamOutput != nil
+	case callbacks.TimingOnError:
+		return ch.OnError != nil
+	default:
+		return false
+	}
+}
+
+func convAgenticToolsNodeCallbackInput(src callbacks.CallbackInput) *schema.AgenticMessage {
+	switch t := src.(type) {
+	case *schema.AgenticMessage:
+		return t
+	default:
+		return nil
+	}
+}
+
+func convAgenticToolsNodeCallbackOutput(src callbacks.CallbackInput) []*schema.AgenticMessage {
+	switch t := src.(type) {
+	case []*schema.AgenticMessage:
+		return t
+	default:
+		return nil
+	}
+}

utils/callbacks/template_test.go

@@ -142,6 +142,58 @@ func TestNewComponentTemplate(t *testing.T) {
 					cnt++
 					return ctx
 				}).Build()).
+			AgenticModel(&AgenticModelCallbackHandler{
+				OnStart: func(ctx context.Context, runInfo *callbacks.RunInfo, input *model.CallbackInput) context.Context {
+					cnt++
+					return ctx
+				},
+				OnEnd: func(ctx context.Context, runInfo *callbacks.RunInfo, output *model.CallbackOutput) context.Context {
+					cnt++
+					return ctx
+				},
+				OnEndWithStreamOutput: func(ctx context.Context, runInfo *callbacks.RunInfo, output *schema.StreamReader[*model.CallbackOutput]) context.Context {
+					output.Close()
+					cnt++
+					return ctx
+				},
+				OnError: func(ctx context.Context, runInfo *callbacks.RunInfo, err error) context.Context {
+					cnt++
+					return ctx
+				},
+			}).
+			AgenticPrompt(&AgenticPromptCallbackHandler{
+				OnStart: func(ctx context.Context, runInfo *callbacks.RunInfo, input *prompt.CallbackInput) context.Context {
+					cnt++
+					return ctx
+				},
+				OnEnd: func(ctx context.Context, runInfo *callbacks.RunInfo, output *prompt.CallbackOutput) context.Context {
+					cnt++
+					return ctx
+				},
+				OnError: func(ctx context.Context, runInfo *callbacks.RunInfo, err error) context.Context {
+					cnt++
+					return ctx
+				},
+			}).
+			AgenticToolsNode(&AgenticToolsNodeCallbackHandlers{
+				OnStart: func(ctx context.Context, info *callbacks.RunInfo, input *schema.AgenticMessage) context.Context {
+					cnt++
+					return ctx
+				},
+				OnEnd: func(ctx context.Context, info *callbacks.RunInfo, input []*schema.AgenticMessage) context.Context {
+					cnt++
+					return ctx
+				},
+				OnEndWithStreamOutput: func(ctx context.Context, info *callbacks.RunInfo, output *schema.StreamReader[[]*schema.AgenticMessage]) context.Context {
+					output.Close()
+					cnt++
+					return ctx
+				},
+				OnError: func(ctx context.Context, info *callbacks.RunInfo, err error) context.Context {
+					cnt++
+					return ctx
+				},
+			}).
 			Handler()
 
 		types := []components.Component{
@@ -151,6 +203,9 @@ func TestNewComponentTemplate(t *testing.T) {
 			components.ComponentOfRetriever,
 			components.ComponentOfTool,
 			compose.ComponentOfLambda,
+			components.ComponentOfAgenticModel,
+			components.ComponentOfAgenticPrompt,
+			compose.ComponentOfAgenticToolsNode,
 		}
 
 		handler := tpl.Handler()
@@ -169,28 +224,28 @@ func TestNewComponentTemplate(t *testing.T) {
 			handler.OnEndWithStreamOutput(ctx, &callbacks.RunInfo{Component: typ}, sor)
 		}
 
-		assert.Equal(t, 22, cnt)
+		assert.Equal(t, 33, cnt)
 
 		ctx = context.Background()
 		ctx = callbacks.InitCallbacks(ctx, &callbacks.RunInfo{Component: components.ComponentOfTransformer}, handler)
 		callbacks.OnStart[any](ctx, nil)
-		assert.Equal(t, 22, cnt)
+		assert.Equal(t, 33, cnt)
 
 		ctx = callbacks.ReuseHandlers(ctx, &callbacks.RunInfo{Component: components.ComponentOfPrompt})
 		ctx = callbacks.OnStart[any](ctx, nil)
-		assert.Equal(t, 23, cnt)
+		assert.Equal(t, 34, cnt)
 
 		ctx = callbacks.ReuseHandlers(ctx, &callbacks.RunInfo{Component: components.ComponentOfIndexer})
 		callbacks.OnEnd[any](ctx, nil)
-		assert.Equal(t, 23, cnt)
+		assert.Equal(t, 34, cnt)
 
 		ctx = callbacks.ReuseHandlers(ctx, &callbacks.RunInfo{Component: components.ComponentOfEmbedding})
 		callbacks.OnError(ctx, nil)
-		assert.Equal(t, 24, cnt)
+		assert.Equal(t, 35, cnt)
 
 		ctx = callbacks.ReuseHandlers(ctx, &callbacks.RunInfo{Component: components.ComponentOfLoader})
 		callbacks.OnStart[any](ctx, nil)
-		assert.Equal(t, 24, cnt)
+		assert.Equal(t, 35, cnt)
 
 		tpl.Transformer(&TransformerCallbackHandler{
 			OnStart: func(ctx context.Context, runInfo *callbacks.RunInfo, input *document.TransformerCallbackInput) context.Context {
@@ -250,6 +305,37 @@ func TestNewComponentTemplate(t *testing.T) {
 					}
 				}
 			},
+		}).AgenticPrompt(&AgenticPromptCallbackHandler{
+			OnStart: func(ctx context.Context, runInfo *callbacks.RunInfo, input *prompt.CallbackInput) context.Context {
+				cnt++
+				return ctx
+			},
+			OnEnd: func(ctx context.Context, runInfo *callbacks.RunInfo, output *prompt.CallbackOutput) context.Context {
+				cnt++
+				return ctx
+			},
+			OnError: func(ctx context.Context, runInfo *callbacks.RunInfo, err error) context.Context {
+				cnt++
+				return ctx
+			},
+		}).AgenticToolsNode(&AgenticToolsNodeCallbackHandlers{
+			OnStart: func(ctx context.Context, info *callbacks.RunInfo, input *schema.AgenticMessage) context.Context {
+				cnt++
+				return ctx
+			},
+			OnEnd: func(ctx context.Context, info *callbacks.RunInfo, input []*schema.AgenticMessage) context.Context {
+				cnt++
+				return ctx
+			},
+			OnEndWithStreamOutput: func(ctx context.Context, info *callbacks.RunInfo, output *schema.StreamReader[[]*schema.AgenticMessage]) context.Context {
+				output.Close()
+				cnt++
+				return ctx
+			},
+			OnError: func(ctx context.Context, info *callbacks.RunInfo, err error) context.Context {
+				cnt++
+				return ctx
+			},
 		})
 
 		handler = tpl.Handler()
@@ -257,36 +343,222 @@ func TestNewComponentTemplate(t *testing.T) {
 		ctx = callbacks.InitCallbacks(ctx, &callbacks.RunInfo{Component: components.ComponentOfTransformer}, handler)
 
 		ctx = callbacks.OnStart[any](ctx, nil)
-		assert.Equal(t, 25, cnt)
+		assert.Equal(t, 36, cnt)
 
 		callbacks.OnEnd[any](ctx, nil)
-		assert.Equal(t, 26, cnt)
+		assert.Equal(t, 37, cnt)
 
 		ctx = callbacks.ReuseHandlers(ctx, &callbacks.RunInfo{Component: components.ComponentOfLoader})
 		callbacks.OnEnd[any](ctx, nil)
-		assert.Equal(t, 27, cnt)
+		assert.Equal(t, 38, cnt)
 
 		ctx = callbacks.ReuseHandlers(ctx, &callbacks.RunInfo{Component: compose.ComponentOfToolsNode})
 		callbacks.OnStart[any](ctx, nil)
-		assert.Equal(t, 28, cnt)
+		assert.Equal(t, 39, cnt)
 
 		sr, sw := schema.Pipe[any](0)
 		sw.Close()
 		callbacks.OnEndWithStreamOutput[any](ctx, sr)
-		assert.Equal(t, 29, cnt)
+		assert.Equal(t, 40, cnt)
 
 		sr1, sw1 := schema.Pipe[[]*schema.Message](1)
 		sw1.Send([]*schema.Message{{}}, nil)
 		sw1.Close()
 		callbacks.OnEndWithStreamOutput[[]*schema.Message](ctx, sr1)
-		assert.Equal(t, 30, cnt)
+		// Check AgenticModel stream
+		sir2, siw2 := schema.Pipe[callbacks.CallbackOutput](1)
+		siw2.Close()
+		handler.OnEndWithStreamOutput(ctx, &callbacks.RunInfo{Component: components.ComponentOfAgenticModel}, sir2)
+		assert.Equal(t, 42, cnt)
 
-		callbacks.OnError(ctx, nil)
-		assert.Equal(t, 30, cnt)
+		// Check AgenticToolsNode stream
+		sir3, siw3 := schema.Pipe[callbacks.CallbackOutput](1)
+		siw3.Close()
+		handler.OnEndWithStreamOutput(ctx, &callbacks.RunInfo{Component: compose.ComponentOfAgenticToolsNode}, sir3)
+		assert.Equal(t, 43, cnt)
 
 		ctx = callbacks.ReuseHandlers(ctx, nil)
 		callbacks.OnStart[any](ctx, nil)
-		assert.Equal(t, 30, cnt)
+		assert.Equal(t, 43, cnt)
+	})
+
+	t.Run("EdgeCases", func(t *testing.T) {
+		ctx := context.Background()
+		cnt := 0
+
+		// 1. Test Graph and Chain Setters and Execution
+		tpl := NewHandlerHelper().
+			Graph(callbacks.NewHandlerBuilder().
+				OnStartFn(func(ctx context.Context, info *callbacks.RunInfo, input callbacks.CallbackInput) context.Context {
+					cnt++
+					return ctx
+				}).Build()).
+			Chain(callbacks.NewHandlerBuilder().
+				OnEndFn(func(ctx context.Context, info *callbacks.RunInfo, output callbacks.CallbackOutput) context.Context {
+					cnt++
+					return ctx
+				}).Build())
+
+		h := tpl.Handler()
+
+		// Trigger Graph OnStart
+		h.OnStart(ctx, &callbacks.RunInfo{Component: compose.ComponentOfGraph}, nil)
+		assert.Equal(t, 1, cnt)
+
+		// Trigger Chain OnEnd
+		h.OnEnd(ctx, &callbacks.RunInfo{Component: compose.ComponentOfChain}, nil)
+		assert.Equal(t, 2, cnt)
+
+		// 2. Test Needed logic for Graph/Chain when handler is present/absent
+		// Graph is present (OnStart)
+		needed := h.(callbacks.TimingChecker).Needed(ctx, &callbacks.RunInfo{Component: compose.ComponentOfGraph}, callbacks.TimingOnStart)
+		assert.True(t, needed)
+
+		// Chain is present (OnEnd) - but we check OnStart which is not defined in the builder above?
+		// NewHandlerBuilder returns a handler that usually returns true for Needed if the specific func is not nil.
+		// Let's verify Chain OnStart is NOT needed because we only set OnEndFn.
+		needed = h.(callbacks.TimingChecker).Needed(ctx, &callbacks.RunInfo{Component: compose.ComponentOfChain}, callbacks.TimingOnStart)
+		assert.False(t, needed) // Should be false because OnStartFn wasn't set for Chain
+
+		// Lambda is NOT present
+		needed = h.(callbacks.TimingChecker).Needed(ctx, &callbacks.RunInfo{Component: compose.ComponentOfLambda}, callbacks.TimingOnStart)
+		assert.False(t, needed)
+
+		// 3. Test Conversion Fallbacks (Default cases)
+		// We need a handler with ToolsNode and AgenticToolsNode to test their conversion fallbacks
+		tpl2 := NewHandlerHelper().
+			ToolsNode(&ToolsNodeCallbackHandlers{
+				OnStart: func(ctx context.Context, info *callbacks.RunInfo, input *schema.Message) context.Context {
+					if input == nil {
+						cnt++
+					}
+					return ctx
+				},
+				OnEnd: func(ctx context.Context, info *callbacks.RunInfo, input []*schema.Message) context.Context {
+					if input == nil {
+						cnt++
+					}
+					return ctx
+				},
+			}).
+			AgenticToolsNode(&AgenticToolsNodeCallbackHandlers{
+				OnStart: func(ctx context.Context, info *callbacks.RunInfo, input *schema.AgenticMessage) context.Context {
+					if input == nil {
+						cnt++
+					}
+					return ctx
+				},
+				OnEnd: func(ctx context.Context, info *callbacks.RunInfo, input []*schema.AgenticMessage) context.Context {
+					if input == nil {
+						cnt++
+					}
+					return ctx
+				},
+			})
+
+		h2 := tpl2.Handler()
+
+		// Pass wrong type (string) to trigger default case in convToolsNodeCallbackInput -> returns nil
+		h2.OnStart(ctx, &callbacks.RunInfo{Component: compose.ComponentOfToolsNode}, "wrong-input-type")
+		assert.Equal(t, 3, cnt) // +1
+
+		// Pass wrong type to trigger default case in convToolsNodeCallbackOutput -> returns nil
+		h2.OnEnd(ctx, &callbacks.RunInfo{Component: compose.ComponentOfToolsNode}, "wrong-output-type")
+		assert.Equal(t, 4, cnt) // +1
+
+		// Pass wrong type to trigger default case in convAgenticToolsNodeCallbackInput -> returns nil
+		h2.OnStart(ctx, &callbacks.RunInfo{Component: compose.ComponentOfAgenticToolsNode}, "wrong-input-type")
+		assert.Equal(t, 5, cnt) // +1
+
+		// Pass wrong type to trigger default case in convAgenticToolsNodeCallbackOutput -> returns nil
+		h2.OnEnd(ctx, &callbacks.RunInfo{Component: compose.ComponentOfAgenticToolsNode}, "wrong-output-type")
+		assert.Equal(t, 6, cnt) // +1
+
+		// 4. Test Needed for Agentic components when handlers are Set vs Unset
+		// tpl2 has AgenticToolsNode set
+		needed = h2.(callbacks.TimingChecker).Needed(ctx, &callbacks.RunInfo{Component: compose.ComponentOfAgenticToolsNode}, callbacks.TimingOnStart)
+		assert.True(t, needed)
+
+		// tpl2 does NOT have AgenticModel set
+		needed = h2.(callbacks.TimingChecker).Needed(ctx, &callbacks.RunInfo{Component: components.ComponentOfAgenticModel}, callbacks.TimingOnStart)
+		assert.False(t, needed)
+
+		// Set it now
+		tpl2.AgenticModel(&AgenticModelCallbackHandler{
+			OnStart: func(ctx context.Context, runInfo *callbacks.RunInfo, input *model.CallbackInput) context.Context {
+				return ctx
+			},
+		})
+
+		needed = h2.(callbacks.TimingChecker).Needed(ctx, &callbacks.RunInfo{Component: components.ComponentOfAgenticModel}, callbacks.TimingOnStart)
+		assert.True(t, needed)
+
+		// Check invalid component
+		needed = h2.(callbacks.TimingChecker).Needed(ctx, &callbacks.RunInfo{Component: "UnknownComponent"}, callbacks.TimingOnStart)
+		assert.False(t, needed)
+
+		// Check RunInfo nil
+		needed = h2.(callbacks.TimingChecker).Needed(ctx, nil, callbacks.TimingOnStart)
+		assert.False(t, needed)
+
+		// 5. Test Needed for Transformer, Loader, Indexer, etc to ensure switch coverage
+		tpl3 := NewHandlerHelper().
+			Transformer(&TransformerCallbackHandler{OnStart: func(ctx context.Context, info *callbacks.RunInfo, input *document.TransformerCallbackInput) context.Context {
+				return ctx
+			}}).
+			Loader(&LoaderCallbackHandler{OnStart: func(ctx context.Context, info *callbacks.RunInfo, input *document.LoaderCallbackInput) context.Context {
+				return ctx
+			}}).
+			Indexer(&IndexerCallbackHandler{OnStart: func(ctx context.Context, info *callbacks.RunInfo, input *indexer.CallbackInput) context.Context {
+				return ctx
+			}}).
+			Retriever(&RetrieverCallbackHandler{OnStart: func(ctx context.Context, info *callbacks.RunInfo, input *retriever.CallbackInput) context.Context {
+				return ctx
+			}}).
+			Embedding(&EmbeddingCallbackHandler{OnStart: func(ctx context.Context, info *callbacks.RunInfo, input *embedding.CallbackInput) context.Context {
+				return ctx
+			}}).
+			Tool(&ToolCallbackHandler{OnStart: func(ctx context.Context, info *callbacks.RunInfo, input *tool.CallbackInput) context.Context {
+				return ctx
+			}})
+
+		h3 := tpl3.Handler()
+		checker := h3.(callbacks.TimingChecker)
+
+		assert.True(t, checker.Needed(ctx, &callbacks.RunInfo{Component: components.ComponentOfTransformer}, callbacks.TimingOnStart))
+		assert.True(t, checker.Needed(ctx, &callbacks.RunInfo{Component: components.ComponentOfLoader}, callbacks.TimingOnStart))
+		assert.True(t, checker.Needed(ctx, &callbacks.RunInfo{Component: components.ComponentOfIndexer}, callbacks.TimingOnStart))
+		assert.True(t, checker.Needed(ctx, &callbacks.RunInfo{Component: components.ComponentOfRetriever}, callbacks.TimingOnStart))
+		assert.True(t, checker.Needed(ctx, &callbacks.RunInfo{Component: components.ComponentOfEmbedding}, callbacks.TimingOnStart))
+		assert.True(t, checker.Needed(ctx, &callbacks.RunInfo{Component: components.ComponentOfTool}, callbacks.TimingOnStart))
+
+		// Verify False paths (by using a helper without them)
+		emptyH := NewHandlerHelper().Handler().(callbacks.TimingChecker)
+		assert.False(t, emptyH.Needed(ctx, &callbacks.RunInfo{Component: components.ComponentOfTransformer}, callbacks.TimingOnStart))
+		assert.False(t, emptyH.Needed(ctx, &callbacks.RunInfo{Component: components.ComponentOfLoader}, callbacks.TimingOnStart))
+		assert.False(t, emptyH.Needed(ctx, &callbacks.RunInfo{Component: components.ComponentOfIndexer}, callbacks.TimingOnStart))
+		assert.False(t, emptyH.Needed(ctx, &callbacks.RunInfo{Component: components.ComponentOfRetriever}, callbacks.TimingOnStart))
+		assert.False(t, emptyH.Needed(ctx, &callbacks.RunInfo{Component: components.ComponentOfEmbedding}, callbacks.TimingOnStart))
+		assert.False(t, emptyH.Needed(ctx, &callbacks.RunInfo{Component: components.ComponentOfTool}, callbacks.TimingOnStart))
+
+		// 6. Test Needed for remaining components (ChatModel, Prompt, AgenticPrompt)
+		tpl4 := NewHandlerHelper().
+			ChatModel(&ModelCallbackHandler{OnStart: func(ctx context.Context, runInfo *callbacks.RunInfo, input *model.CallbackInput) context.Context {
+				return ctx
+			}}).
+			Prompt(&PromptCallbackHandler{OnStart: func(ctx context.Context, runInfo *callbacks.RunInfo, input *prompt.CallbackInput) context.Context {
+				return ctx
+			}}).
+			AgenticPrompt(&AgenticPromptCallbackHandler{OnStart: func(ctx context.Context, runInfo *callbacks.RunInfo, input *prompt.CallbackInput) context.Context {
+				return ctx
+			}})
+
+		h4 := tpl4.Handler()
+		checker4 := h4.(callbacks.TimingChecker)
+
+		assert.True(t, checker4.Needed(ctx, &callbacks.RunInfo{Component: components.ComponentOfChatModel}, callbacks.TimingOnStart))
+		assert.True(t, checker4.Needed(ctx, &callbacks.RunInfo{Component: components.ComponentOfPrompt}, callbacks.TimingOnStart))
+		assert.True(t, checker4.Needed(ctx, &callbacks.RunInfo{Component: components.ComponentOfAgenticPrompt}, callbacks.TimingOnStart))
 	})
 }