feat: add more AgentOption for react agent and a converter function to compose.Option

Change-Id: I3a069cd40b5fff9d701b35c4f1d383c303e2387f

Author: shentongmartin

_typos.toml

@@ -6,6 +6,7 @@ Invokable = "Invokable"
 invokable = "invokable"
 InvokableLambda = "InvokableLambda"
 InvokableRun = "InvokableRun"
+typ = "typ"
 
 [files]
 extend-exclude = ["go.mod", "go.sum", "check_branch_name.sh"]

components/tool/option.go

@@ -17,8 +17,8 @@
 package tool
 
 // Option defines call option for InvokableTool or StreamableTool component, which is part of component interface signature.
-// Each tool implementation could define its own options struct and option funcs within its own package,
-// then wrap the impl specific option funcs into this type, before passing to InvokableRun or StreamableRun.
+// Each tool implementation could define its own options struct and option functions within its own package,
+// then wrap the impl specific option functions into this type, before passing to InvokableRun or StreamableRun.
 type Option struct {
 	implSpecificOptFn any
 }

compose/graph_call_options.go

@@ -80,6 +80,24 @@ func (o Option) DesignateNodeWithPath(path ...*NodePath) Option {
 	return o
 }
 
+// DesignateNodePrependPath prepends the prefix to the path of the node(s) to which the option will be applied to.
+// Useful when you already have an Option designated to a graph's node, and now you want to add this graph as a subgraph.
+// e.g.
+// Your subgraph has a Node with key "A", and your subgraph's NodeKey is "sub_graph", you can specify option to A using:
+//
+// option := WithCallbacks(...).DesignateNode("A").DesignateNodePrependPath("sub_graph")
+// Note: as an End User, you probably don't need to use this method, as DesignateNodeWithPath will be sufficient in most use cases.
+// Note: as a Flow author, if you define your own Option type, and at the same time your flow can be exported to graph and added as GraphNode,
+// you can use this method to prepend your Option's designated path with the GraphNode's path.
+func (o Option) DesignateNodePrependPath(prefix *NodePath) Option {
+	for i := range o.paths {
+		p := o.paths[i]
+		p.path = append(prefix.path, p.path...)
+	}
+
+	return o
+}
+
 // WithEmbeddingOption is a functional option type for embedding component.
 // e.g.
 //

flow/agent/agent_option.go

@@ -19,14 +19,14 @@ package agent
 import "github.com/cloudwego/eino/compose"
 
 // AgentOption is the common option type for various agent and multi-agent implementations.
-// For options intended to use with underlying graph or components, use WithComposeOptions to specify.
 // For options intended to use with particular agent/multi-agent implementations, use WrapImplSpecificOptFn to specify.
 type AgentOption struct {
 	implSpecificOptFn any
 	composeOptions    []compose.Option
 }
 
 // GetComposeOptions returns all compose options from the given agent options.
+// Deprecated
 func GetComposeOptions(opts ...AgentOption) []compose.Option {
 	var result []compose.Option
 	for _, opt := range opts {
@@ -37,6 +37,7 @@ func GetComposeOptions(opts ...AgentOption) []compose.Option {
 }
 
 // WithComposeOptions returns an agent option that specifies compose options.
+// Deprecated: use option functions defined by each agent flow implementation instead.
 func WithComposeOptions(opts ...compose.Option) AgentOption {
 	return AgentOption{
 		composeOptions: opts,

flow/agent/multiagent/host/callback.go

@@ -39,6 +39,7 @@ type HandOffInfo struct {
 }
 
 // ConvertCallbackHandlers converts []host.MultiAgentCallback to callbacks.Handler.
+// Deprecated: use ConvertOptions to convert agent.AgentOption to compose.Option when adding MultiAgent's Graph to another Graph.
 func ConvertCallbackHandlers(handlers ...MultiAgentCallback) callbacks.Handler {
 	onChatModelEnd := func(ctx context.Context, info *callbacks.RunInfo, output *model.CallbackOutput) context.Context {
 		if output == nil || info == nil {
@@ -121,5 +122,76 @@ func convertCallbacks(opts ...agent.AgentOption) callbacks.Handler {
 	}
 
 	handlers := agentOptions.agentCallbacks
-	return ConvertCallbackHandlers(handlers...)
+
+	onChatModelEnd := func(ctx context.Context, info *callbacks.RunInfo, output *model.CallbackOutput) context.Context {
+		if output == nil || info == nil {
+			return ctx
+		}
+
+		msg := output.Message
+		if msg == nil || msg.Role != schema.Assistant || len(msg.ToolCalls) == 0 {
+			return ctx
+		}
+
+		agentName := msg.ToolCalls[0].Function.Name
+		argument := msg.ToolCalls[0].Function.Arguments
+
+		for _, cb := range handlers {
+			ctx = cb.OnHandOff(ctx, &HandOffInfo{
+				ToAgentName: agentName,
+				Argument:    argument,
+			})
+		}
+
+		return ctx
+	}
+
+	onChatModelEndWithStreamOutput := func(ctx context.Context, info *callbacks.RunInfo, output *schema.StreamReader[*model.CallbackOutput]) context.Context {
+		if output == nil || info == nil {
+			return ctx
+		}
+
+		defer output.Close()
+
+		var msgs []*schema.Message
+		for {
+			oneOutput, err := output.Recv()
+			if err == io.EOF {
+				break
+			}
+			if err != nil {
+				return ctx
+			}
+
+			msg := oneOutput.Message
+			if msg == nil {
+				continue
+			}
+
+			msgs = append(msgs, msg)
+		}
+
+		msg, err := schema.ConcatMessages(msgs)
+		if err != nil {
+			return ctx
+		}
+
+		if msg.Role != schema.Assistant || len(msg.ToolCalls) == 0 {
+			return ctx
+		}
+
+		for _, cb := range handlers {
+			ctx = cb.OnHandOff(ctx, &HandOffInfo{
+				ToAgentName: msg.ToolCalls[0].Function.Name,
+				Argument:    msg.ToolCalls[0].Function.Arguments,
+			})
+		}
+
+		return ctx
+	}
+
+	return template.NewHandlerHelper().ChatModel(&template.ModelCallbackHandler{
+		OnEnd:                 onChatModelEnd,
+		OnEndWithStreamOutput: onChatModelEndWithStreamOutput,
+	}).Handler()
 }

flow/agent/multiagent/host/compose_test.go

@@ -24,12 +24,15 @@ import (
 	"github.com/stretchr/testify/assert"
 	"go.uber.org/mock/gomock"
 
+	"github.com/cloudwego/eino/callbacks"
+	chatmodel "github.com/cloudwego/eino/components/model"
 	"github.com/cloudwego/eino/components/prompt"
 	"github.com/cloudwego/eino/compose"
 	"github.com/cloudwego/eino/flow/agent"
 	"github.com/cloudwego/eino/internal/generic"
 	"github.com/cloudwego/eino/internal/mock/components/model"
 	"github.com/cloudwego/eino/schema"
+	template "github.com/cloudwego/eino/utils/callbacks"
 )
 
 func TestHostMultiAgent(t *testing.T) {
@@ -48,6 +51,14 @@ func TestHostMultiAgent(t *testing.T) {
 
 	specialist2 := &Specialist{
 		Invokable: func(ctx context.Context, input []*schema.Message, opts ...agent.AgentOption) (*schema.Message, error) {
+			agentOpts := agent.GetImplSpecificOptions(&specialist2Options{}, opts...)
+			if agentOpts.mockOutput != nil {
+				return &schema.Message{
+					Role:    schema.Assistant,
+					Content: *agentOpts.mockOutput,
+				}, nil
+			}
+
 			return &schema.Message{
 				Role:    schema.Assistant,
 				Content: "specialist2 invoke answer",
@@ -92,11 +103,18 @@ func TestHostMultiAgent(t *testing.T) {
 			Content: "direct answer",
 		}
 
-		mockHostLLM.EXPECT().Generate(gomock.Any(), gomock.Any()).Return(directAnswerMsg, nil).Times(1)
+		mockHostLLM.EXPECT().Generate(gomock.Any(), gomock.Any(), gomock.Any()).
+			DoAndReturn(func(_ context.Context, input []*schema.Message, opts ...chatmodel.Option) (*schema.Message, error) {
+				modelOpts := chatmodel.GetCommonOptions(&chatmodel.Options{}, opts...)
+				assert.Equal(t, *modelOpts.Temperature, float32(0.7))
+				return directAnswerMsg, nil
+			}).
+			Times(1)
 
 		mockCallback := &mockAgentCallback{}
 
-		out, err := hostMA.Generate(ctx, nil, WithAgentCallbacks(mockCallback))
+		out, err := hostMA.Generate(ctx, nil, WithAgentCallbacks(mockCallback),
+			WithAgentModelOptions(hostMA.HostNodeKey(), chatmodel.WithTemperature(0.7)))
 		assert.NoError(t, err)
 		assert.Equal(t, "direct answer", out.Content)
 		assert.Empty(t, mockCallback.infos)
@@ -164,11 +182,18 @@ func TestHostMultiAgent(t *testing.T) {
 		}
 
 		mockHostLLM.EXPECT().Generate(gomock.Any(), gomock.Any()).Return(handOffMsg, nil).Times(1)
-		mockSpecialistLLM1.EXPECT().Generate(gomock.Any(), gomock.Any()).Return(specialistMsg, nil).Times(1)
+		mockSpecialistLLM1.EXPECT().Generate(gomock.Any(), gomock.Any(), gomock.Any()).
+			DoAndReturn(func(_ context.Context, input []*schema.Message, opts ...chatmodel.Option) (*schema.Message, error) {
+				modelOpts := chatmodel.GetCommonOptions(&chatmodel.Options{}, opts...)
+				assert.Equal(t, *modelOpts.Temperature, float32(0.7))
+				return specialistMsg, nil
+			}).
+			Times(1)
 
 		mockCallback := &mockAgentCallback{}
 
-		out, err := hostMA.Generate(ctx, nil, WithAgentCallbacks(mockCallback))
+		out, err := hostMA.Generate(ctx, nil, WithAgentCallbacks(mockCallback),
+			WithAgentModelOptions(specialist1.Name, chatmodel.WithTemperature(0.7)))
 		assert.NoError(t, err)
 		assert.Equal(t, "specialist 1 answer", out.Content)
 		assert.Equal(t, []*HandOffInfo{
@@ -379,16 +404,41 @@ func TestHostMultiAgent(t *testing.T) {
 			},
 		}
 
-		specialistMsg := &schema.Message{
+		specialist1Msg := &schema.Message{
 			Role:    schema.Assistant,
 			Content: "Beijing",
 		}
 
-		mockHostLLM.EXPECT().Generate(gomock.Any(), gomock.Any()).Return(handOffMsg, nil).Times(1)
-		mockSpecialistLLM1.EXPECT().Generate(gomock.Any(), gomock.Any()).Return(specialistMsg, nil).Times(1)
+		mockHostLLM.EXPECT().Generate(gomock.Any(), gomock.Any()).Return(handOffMsg, nil).Times(2)
+		mockSpecialistLLM1.EXPECT().Generate(gomock.Any(), gomock.Any(), gomock.Any()).
+			DoAndReturn(func(_ context.Context, input []*schema.Message, opts ...chatmodel.Option) (*schema.Message, error) {
+				modelOpts := chatmodel.GetCommonOptions(&chatmodel.Options{}, opts...)
+				assert.Equal(t, *modelOpts.Temperature, float32(0.7))
+				return specialist1Msg, nil
+			}).
+			Times(1)
 
 		mockCallback := &mockAgentCallback{}
 
+		var hostOutput, specialist1Output, specialist2Output string
+		hostModelCallback := template.NewHandlerHelper().ChatModel(&template.ModelCallbackHandler{
+			OnEnd: func(ctx context.Context, runInfo *callbacks.RunInfo, output *chatmodel.CallbackOutput) context.Context {
+				hostOutput = output.Message.ToolCalls[0].Function.Name
+				return ctx
+			},
+		}).Handler()
+		specialist1ModelCallback := template.NewHandlerHelper().ChatModel(&template.ModelCallbackHandler{
+			OnEnd: func(ctx context.Context, runInfo *callbacks.RunInfo, output *chatmodel.CallbackOutput) context.Context {
+				specialist1Output = output.Message.Content
+				return ctx
+			},
+		}).Handler()
+		specialist2LambdaCallback := template.NewHandlerHelper().Lambda(callbacks.NewHandlerBuilder().OnEndFn(
+			func(ctx context.Context, info *callbacks.RunInfo, output callbacks.CallbackOutput) context.Context {
+				specialist2Output = output.(*schema.Message).Content
+				return ctx
+			}).Build()).Handler()
+
 		hostMA, err := NewMultiAgent(ctx, &MultiAgentConfig{
 			Host: Host{
 				ChatModel: mockHostLLM,
@@ -409,7 +459,14 @@ func TestHostMultiAgent(t *testing.T) {
 			Compile(ctx)
 		assert.NoError(t, err)
 
-		out, err := fullGraph.Invoke(ctx, map[string]any{"country_name": "China"}, compose.WithCallbacks(ConvertCallbackHandlers(mockCallback)).DesignateNodeWithPath(compose.NewNodePath("host_ma_node", hostMA.HostNodeKey())))
+		convertedOptions := ConvertOptions(compose.NewNodePath("host_ma_node"), WithAgentCallbacks(mockCallback),
+			WithAgentModelOptions(specialist1.Name, chatmodel.WithTemperature(0.7)),
+			WithAgentModelCallbacks(hostMA.HostNodeKey(), hostModelCallback),
+			WithAgentModelCallbacks(specialist1.Name, specialist1ModelCallback),
+			WithSpecialistLambdaCallbacks(specialist2.Name, specialist2LambdaCallback),
+			WithSpecialistLambdaOptions(specialist2.Name, withSpecialist2MockOutput("mock_city_name")))
+
+		out, err := fullGraph.Invoke(ctx, map[string]any{"country_name": "China"}, convertedOptions...)
 		assert.NoError(t, err)
 		assert.Equal(t, "Beijing", out.Content)
 		assert.Equal(t, []*HandOffInfo{
@@ -418,6 +475,28 @@ func TestHostMultiAgent(t *testing.T) {
 				Argument:    `{"reason": "specialist 1 is the best"}`,
 			},
 		}, mockCallback.infos)
+		assert.Equal(t, hostOutput, specialist1.Name)
+		assert.Equal(t, specialist1Output, out.Content)
+		assert.Equal(t, specialist2Output, "")
+
+		handOffMsg.ToolCalls[0].Function.Name = specialist2.Name
+		handOffMsg.ToolCalls[0].Function.Arguments = `{"reason": "specialist 2 is even better"}`
+
+		out, err = fullGraph.Invoke(ctx, map[string]any{"country_name": "China"}, convertedOptions...)
+		assert.NoError(t, err)
+		assert.Equal(t, "mock_city_name", out.Content)
+		assert.Equal(t, []*HandOffInfo{
+			{
+				ToAgentName: specialist1.Name,
+				Argument:    `{"reason": "specialist 1 is the best"}`,
+			},
+			{
+				ToAgentName: specialist2.Name,
+				Argument:    `{"reason": "specialist 2 is even better"}`,
+			},
+		}, mockCallback.infos)
+		assert.Equal(t, hostOutput, specialist2.Name)
+		assert.Equal(t, specialist2Output, "mock_city_name")
 	})
 }
 
@@ -429,3 +508,13 @@ func (m *mockAgentCallback) OnHandOff(ctx context.Context, info *HandOffInfo) co
 	m.infos = append(m.infos, info)
 	return ctx
 }
+
+type specialist2Options struct {
+	mockOutput *string
+}
+
+func withSpecialist2MockOutput(mockOutput string) agent.AgentOption {
+	return agent.WrapImplSpecificOptFn(func(o *specialist2Options) {
+		o.mockOutput = &mockOutput
+	})
+}