feat: improve AgenticToolChoice (#684)

Author: mrh997

components/model/option.go

@@ -28,11 +28,11 @@ type Options struct {
 	TopP *float32
 	// 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.
+	// Options only available for chat model.
 
+	// ToolChoice controls which tool is called by the model.
+	ToolChoice *schema.ToolChoice
 	// 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.
@@ -41,10 +41,10 @@ type Options struct {
 	// Stop is the stop words for the model, which controls the stopping condition of the model.
 	Stop []string
 
-	// Options only for agentic model.
+	// Options only available for agentic model.
 
-	// AllowedTools is a list of allowed tools the model may call.
-	AllowedTools []*schema.AllowedTool
+	// AgenticToolChoice controls how the agentic model calls tools.
+	AgenticToolChoice *schema.AgenticToolChoice
 }
 
 // Option is the call option for ChatModel component.
@@ -127,11 +127,10 @@ 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 {
+func WithAgenticToolChoice(toolChoice *schema.AgenticToolChoice) Option {
 	return Option{
 		apply: func(opts *Options) {
-			opts.ToolChoice = &toolChoice
-			opts.AllowedTools = allowedTools
+			opts.AgenticToolChoice = toolChoice
 		},
 	}
 }

components/model/option_test.go

@@ -92,11 +92,18 @@ func TestOptions(t *testing.T) {
 		)
 		opts := GetCommonOptions(
 			nil,
-			WithAgenticToolChoice(toolChoice, allowedTools...),
+			WithAgenticToolChoice(&schema.AgenticToolChoice{
+				Type: toolChoice,
+				Forced: &schema.AgenticForcedToolChoice{
+					Tools: allowedTools,
+				},
+			}),
 		)
 
-		convey.So(opts.ToolChoice, convey.ShouldResemble, &toolChoice)
-		convey.So(opts.AllowedTools, convey.ShouldResemble, allowedTools)
+		convey.So(opts.AgenticToolChoice, convey.ShouldNotBeNil)
+		convey.So(opts.AgenticToolChoice.Type, convey.ShouldEqual, toolChoice)
+		convey.So(opts.AgenticToolChoice.Forced, convey.ShouldNotBeNil)
+		convey.So(opts.AgenticToolChoice.Forced.Tools, convey.ShouldResemble, allowedTools)
 	})
 }
 

compose/workflow.go

@@ -89,18 +89,36 @@ func (wf *Workflow[I, O]) AddChatModelNode(key string, chatModel model.BaseChatM
 	return wf.initNode(key)
 }
 
+// AddAgenticModelNode adds an agentic model node and returns it.
+func (wf *Workflow[I, O]) AddAgenticModelNode(key string, agenticModel model.AgenticModel, opts ...GraphAddNodeOpt) *WorkflowNode {
+	_ = wf.g.AddAgenticModelNode(key, agenticModel, opts...)
+	return wf.initNode(key)
+}
+
 // AddChatTemplateNode adds a chat template node and returns it.
 func (wf *Workflow[I, O]) AddChatTemplateNode(key string, chatTemplate prompt.ChatTemplate, opts ...GraphAddNodeOpt) *WorkflowNode {
 	_ = wf.g.AddChatTemplateNode(key, chatTemplate, opts...)
 	return wf.initNode(key)
 }
 
+// AddAgenticChatTemplateNode adds an agentic chat template node and returns it.
+func (wf *Workflow[I, O]) AddAgenticChatTemplateNode(key string, chatTemplate prompt.AgenticChatTemplate, opts ...GraphAddNodeOpt) *WorkflowNode {
+	_ = wf.g.AddAgenticChatTemplateNode(key, chatTemplate, opts...)
+	return wf.initNode(key)
+}
+
 // AddToolsNode adds a tools node and returns it.
 func (wf *Workflow[I, O]) AddToolsNode(key string, tools *ToolsNode, opts ...GraphAddNodeOpt) *WorkflowNode {
 	_ = wf.g.AddToolsNode(key, tools, opts...)
 	return wf.initNode(key)
 }
 
+// AddAgenticToolsNode adds an agentic tools node and returns it.
+func (wf *Workflow[I, O]) AddAgenticToolsNode(key string, tools *AgenticToolsNode, opts ...GraphAddNodeOpt) *WorkflowNode {
+	_ = wf.g.AddAgenticToolsNode(key, tools, opts...)
+	return wf.initNode(key)
+}
+
 // AddRetrieverNode adds a retriever node and returns it.
 func (wf *Workflow[I, O]) AddRetrieverNode(key string, retriever retriever.Retriever, opts ...GraphAddNodeOpt) *WorkflowNode {
 	_ = wf.g.AddRetrieverNode(key, retriever, opts...)

schema/tool.go

@@ -55,6 +55,31 @@ const (
 	ToolChoiceForced ToolChoice = "forced"
 )
 
+type AgenticToolChoice struct {
+	// Type is the tool choice mode.
+	Type ToolChoice
+	// Allowed optionally specifies the list of tools that the model is permitted to call.
+	Allowed *AgenticAllowedToolChoice
+	// Forced optionally specifies the list of tools that the model is required to call.
+	Forced *AgenticForcedToolChoice
+}
+
+// AgenticAllowedToolChoice specifies a list of allowed tools for the model.
+type AgenticAllowedToolChoice struct {
+	// Tools is the list of allowed tools for the model to call.
+	// Optional.
+	Tools []*AllowedTool
+}
+
+// AgenticForcedToolChoice specifies a list of tools that the model must call.
+type AgenticForcedToolChoice struct {
+	// Tools is the list of tools that the model must call.
+	// Optional.
+	Tools []*AllowedTool
+}
+
+// AllowedTool represents a tool that the model is allowed or forced to call.
+// Exactly one of FunctionToolName, MCPTool, or ServerTool must be specified.
 type AllowedTool struct {
 	// FunctionToolName is the name of the function tool.
 	FunctionToolName string
@@ -64,15 +89,17 @@ type AllowedTool struct {
 	ServerTool *AllowedServerTool
 }
 
+// AllowedMCPTool contains the information for identifying an MCP tool.
 type AllowedMCPTool struct {
 	// ServerLabel is the label of the MCP server.
 	ServerLabel string
-	// The name of the MCP tool.
+	// Name is the name of the MCP tool.
 	Name string
 }
 
+// AllowedServerTool contains the information for identifying a server tool.
 type AllowedServerTool struct {
-	// The name of the server tool.
+	// Name is the name of the server tool.
 	Name string
 }