feat(tool/utils): add EnhancedStreamableTool implementation

- Add EnhancedStreamFunc and OptionableEnhancedStreamFunc type definitions
- Implement InferEnhancedStreamTool and InferOptionableEnhancedStreamTool functions
- Implement NewEnhancedStreamTool and newOptionableEnhancedStreamTool functions
- Add enhancedStreamableTool struct implementing tool.EnhancedStreamableTool interface
- Add unit tests for all new functions~

Author: zhuangjie1125

components/tool/utils/streamable_func.go

@@ -155,3 +155,101 @@ func (s *streamableTool[T, D]) getToolName() string {
 
 	return s.info.Name
 }
+
+// EnhancedStreamFunc is the function type for the enhanced streamable tool.
+type EnhancedStreamFunc[T any] func(ctx context.Context, input T) (output *schema.StreamReader[*schema.ToolResult], err error)
+
+// OptionableEnhancedStreamFunc is the function type for the enhanced streamable tool with tool option.
+type OptionableEnhancedStreamFunc[T any] func(ctx context.Context, input T, opts ...tool.Option) (output *schema.StreamReader[*schema.ToolResult], err error)
+
+// InferEnhancedStreamTool creates an EnhancedStreamableTool from a given function by inferring the ToolInfo from the function's request parameters.
+// End-user can pass a SchemaCustomizerFn in opts to customize the go struct tag parsing process, overriding default behavior.
+func InferEnhancedStreamTool[T any](toolName, toolDesc string, s EnhancedStreamFunc[T], opts ...Option) (tool.EnhancedStreamableTool, error) {
+	ti, err := goStruct2ToolInfo[T](toolName, toolDesc, opts...)
+	if err != nil {
+		return nil, err
+	}
+
+	return NewEnhancedStreamTool(ti, s, opts...), nil
+}
+
+// InferOptionableEnhancedStreamTool creates an EnhancedStreamableTool from a given function by inferring the ToolInfo from the function's request parameters, with tool option.
+func InferOptionableEnhancedStreamTool[T any](toolName, toolDesc string, s OptionableEnhancedStreamFunc[T], opts ...Option) (tool.EnhancedStreamableTool, error) {
+	ti, err := goStruct2ToolInfo[T](toolName, toolDesc, opts...)
+	if err != nil {
+		return nil, err
+	}
+
+	return newOptionableEnhancedStreamTool(ti, s, opts...), nil
+}
+
+// NewEnhancedStreamTool Create an enhanced streaming tool, where the input is in JSON format and output is *schema.StreamReader[*schema.ToolResult].
+func NewEnhancedStreamTool[T any](desc *schema.ToolInfo, s EnhancedStreamFunc[T], opts ...Option) tool.EnhancedStreamableTool {
+	return newOptionableEnhancedStreamTool(desc,
+		func(ctx context.Context, input T, _ ...tool.Option) (output *schema.StreamReader[*schema.ToolResult], err error) {
+			return s(ctx, input)
+		},
+		opts...)
+}
+
+func newOptionableEnhancedStreamTool[T any](desc *schema.ToolInfo, s OptionableEnhancedStreamFunc[T], opts ...Option) tool.EnhancedStreamableTool {
+	to := getToolOptions(opts...)
+
+	return &enhancedStreamableTool[T]{
+		info: desc,
+		um:   to.um,
+		Fn:   s,
+	}
+}
+
+type enhancedStreamableTool[T any] struct {
+	info *schema.ToolInfo
+
+	um UnmarshalArguments
+
+	Fn OptionableEnhancedStreamFunc[T]
+}
+
+func (s *enhancedStreamableTool[T]) Info(ctx context.Context) (*schema.ToolInfo, error) {
+	return s.info, nil
+}
+
+func (s *enhancedStreamableTool[T]) StreamableRun(ctx context.Context, toolArgument *schema.ToolArgument, opts ...tool.Option) (
+	outStream *schema.StreamReader[*schema.ToolResult], err error) {
+
+	var inst T
+	if s.um != nil {
+		var val any
+		val, err = s.um(ctx, toolArgument.TextArgument)
+		if err != nil {
+			return nil, fmt.Errorf("[EnhancedLocalStreamFunc] failed to unmarshal arguments, toolName=%s, err=%w", s.getToolName(), err)
+		}
+
+		gt, ok := val.(T)
+		if !ok {
+			return nil, fmt.Errorf("[EnhancedLocalStreamFunc] type err, toolName=%s, expected=%T, given=%T", s.getToolName(), inst, val)
+		}
+		inst = gt
+	} else {
+		inst = generic.NewInstance[T]()
+
+		err = sonic.UnmarshalString(toolArgument.TextArgument, &inst)
+		if err != nil {
+			return nil, fmt.Errorf("[EnhancedLocalStreamFunc] failed to unmarshal arguments in json, toolName=%s, err=%w", s.getToolName(), err)
+		}
+	}
+
+	return s.Fn(ctx, inst, opts...)
+}
+
+func (s *enhancedStreamableTool[T]) GetType() string {
+	return snakeToCamel(s.getToolName())
+}
+
+func (s *enhancedStreamableTool[T]) getToolName() string {
+	if s.info == nil {
+		return ""
+	}
+
+	return s.info.Name
+}

components/tool/utils/streamable_func_test.go

@@ -167,3 +167,163 @@ func TestInferStreamTool(t *testing.T) {
 		}
 	}
 }
+
+type EnhancedStreamInput struct {
+	Query string `json:"query" jsonschema:"description=the search query"`
+}
+
+func TestNewEnhancedStreamTool(t *testing.T) {
+	ctx := context.Background()
+
+	t.Run("simple_case", func(t *testing.T) {
+		tl := NewEnhancedStreamTool[*EnhancedStreamInput](
+			&schema.ToolInfo{
+				Name: "enhanced_stream_search",
+				Desc: "search with enhanced stream output",
+				ParamsOneOf: schema.NewParamsOneOfByParams(
+					map[string]*schema.ParameterInfo{
+						"query": {
+							Type: "string",
+							Desc: "the search query",
+						},
+					}),
+			},
+			func(ctx context.Context, input *EnhancedStreamInput) (*schema.StreamReader[*schema.ToolResult], error) {
+				sr, sw := schema.Pipe[*schema.ToolResult](2)
+				sw.Send(&schema.ToolResult{
+					Parts: []schema.ToolOutputPart{
+						{Type: schema.ToolPartTypeText, Text: "result for: " + input.Query},
+					},
+				}, nil)
+				sw.Send(&schema.ToolResult{
+					Parts: []schema.ToolOutputPart{
+						{Type: schema.ToolPartTypeText, Text: "more results"},
+					},
+				}, nil)
+				sw.Close()
+				return sr, nil
+			},
+		)
+
+		info, err := tl.Info(ctx)
+		assert.NoError(t, err)
+		assert.Equal(t, "enhanced_stream_search", info.Name)
+
+		sr, err := tl.StreamableRun(ctx, &schema.ToolArgument{TextArgument: `{"query":"test"}`})
+		assert.NoError(t, err)
+		defer sr.Close()
+
+		idx := 0
+		for {
+			m, err := sr.Recv()
+			if errors.Is(err, io.EOF) {
+				break
+			}
+			assert.NoError(t, err)
+
+			if idx == 0 {
+				assert.Len(t, m.Parts, 1)
+				assert.Equal(t, schema.ToolPartTypeText, m.Parts[0].Type)
+				assert.Equal(t, "result for: test", m.Parts[0].Text)
+			} else {
+				assert.Len(t, m.Parts, 1)
+				assert.Equal(t, "more results", m.Parts[0].Text)
+			}
+			idx++
+		}
+		assert.Equal(t, 2, idx)
+	})
+}
+
+type FakeEnhancedStreamOption struct {
+	Prefix string
+}
+
+func FakeWithEnhancedStreamOption(prefix string) tool.Option {
+	return tool.WrapImplSpecificOptFn(func(t *FakeEnhancedStreamOption) {
+		t.Prefix = prefix
+	})
+}
+
+func fakeEnhancedStreamFunc(ctx context.Context, input EnhancedStreamInput) (*schema.StreamReader[*schema.ToolResult], error) {
+	return schema.StreamReaderFromArray([]*schema.ToolResult{
+		{
+			Parts: []schema.ToolOutputPart{
+				{Type: schema.ToolPartTypeText, Text: "result: " + input.Query},
+			},
+		},
+	}), nil
+}
+
+func fakeOptionableEnhancedStreamFunc(ctx context.Context, input EnhancedStreamInput, opts ...tool.Option) (*schema.StreamReader[*schema.ToolResult], error) {
+	baseOpt := &FakeEnhancedStreamOption{
+		Prefix: "default",
+	}
+	option := tool.GetImplSpecificOptions(baseOpt, opts...)
+
+	return schema.StreamReaderFromArray([]*schema.ToolResult{
+		{
+			Parts: []schema.ToolOutputPart{
+				{Type: schema.ToolPartTypeText, Text: option.Prefix + ": " + input.Query},
+			},
+		},
+	}), nil
+}
+
+func TestInferEnhancedStreamTool(t *testing.T) {
+	ctx := context.Background()
+
+	t.Run("infer_enhanced_stream_tool", func(t *testing.T) {
+		tl, err := InferEnhancedStreamTool("infer_enhanced_stream", "test infer enhanced stream tool", fakeEnhancedStreamFunc)
+		assert.NoError(t, err)
+
+		info, err := tl.Info(ctx)
+		assert.NoError(t, err)
+		assert.Equal(t, "infer_enhanced_stream", info.Name)
+
+		sr, err := tl.StreamableRun(ctx, &schema.ToolArgument{TextArgument: `{"query":"hello"}`})
+		assert.NoError(t, err)
+		defer sr.Close()
+
+		m, err := sr.Recv()
+		assert.NoError(t, err)
+		assert.Len(t, m.Parts, 1)
+		assert.Equal(t, "result: hello", m.Parts[0].Text)
+	})
+}
+
+func TestInferOptionableEnhancedStreamTool(t *testing.T) {
+	ctx := context.Background()
+
+	t.Run("infer_optionable_enhanced_stream_tool", func(t *testing.T) {
+		tl, err := InferOptionableEnhancedStreamTool("infer_optionable_enhanced_stream", "test infer optionable enhanced stream tool", fakeOptionableEnhancedStreamFunc)
+		assert.NoError(t, err)
+
+		info, err := tl.Info(ctx)
+		assert.NoError(t, err)
+		assert.Equal(t, "infer_optionable_enhanced_stream", info.Name)
+
+		sr, err := tl.StreamableRun(ctx, &schema.ToolArgument{TextArgument: `{"query":"world"}`}, FakeWithEnhancedStreamOption("custom"))
+		assert.NoError(t, err)
+		defer sr.Close()
+
+		m, err := sr.Recv()
+		assert.NoError(t, err)
+		assert.Len(t, m.Parts, 1)
+		assert.Equal(t, "custom: world", m.Parts[0].Text)
+	})
+
+	t.Run("infer_optionable_enhanced_stream_tool_default_option", func(t *testing.T) {
+		tl, err := InferOptionableEnhancedStreamTool("infer_optionable_enhanced_stream", "test infer optionable enhanced stream tool", fakeOptionableEnhancedStreamFunc)
+		assert.NoError(t, err)
+
+		sr, err := tl.StreamableRun(ctx, &schema.ToolArgument{TextArgument: `{"query":"test"}`})
+		assert.NoError(t, err)
+		defer sr.Close()
+
+		m, err := sr.Recv()
+		assert.NoError(t, err)
+		assert.Len(t, m.Parts, 1)
+		assert.Equal(t, "default: test", m.Parts[0].Text)
+	})
+}