wehub-resource-sync
823 lines
21 KiB
Go
823 lines
21 KiB
Go
/*
|
|
* Copyright 2024 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 react
|
|
|
|
import (
|
|
"context"
|
|
"errors"
|
|
"fmt"
|
|
"io"
|
|
"math/rand"
|
|
"testing"
|
|
|
|
"github.com/bytedance/sonic"
|
|
"github.com/stretchr/testify/assert"
|
|
"go.uber.org/mock/gomock"
|
|
|
|
"github.com/cloudwego/eino/components/model"
|
|
"github.com/cloudwego/eino/components/tool"
|
|
"github.com/cloudwego/eino/compose"
|
|
"github.com/cloudwego/eino/flow/agent"
|
|
mockModel "github.com/cloudwego/eino/internal/mock/components/model"
|
|
"github.com/cloudwego/eino/schema"
|
|
template "github.com/cloudwego/eino/utils/callbacks"
|
|
)
|
|
|
|
func TestReact(t *testing.T) {
|
|
ctx := context.Background()
|
|
|
|
fakeTool := &fakeToolGreetForTest{
|
|
tarCount: 3,
|
|
}
|
|
|
|
info, err := fakeTool.Info(ctx)
|
|
assert.NoError(t, err)
|
|
|
|
ctrl := gomock.NewController(t)
|
|
cm := mockModel.NewMockChatModel(ctrl)
|
|
|
|
times := 0
|
|
cm.EXPECT().Generate(gomock.Any(), gomock.Any(), gomock.Any()).
|
|
DoAndReturn(func(ctx context.Context, input []*schema.Message, opts ...model.Option) (*schema.Message, error) {
|
|
times++
|
|
if times <= 2 {
|
|
info, _ := fakeTool.Info(ctx)
|
|
|
|
return schema.AssistantMessage("hello max",
|
|
[]schema.ToolCall{
|
|
{
|
|
ID: randStr(),
|
|
Function: schema.FunctionCall{
|
|
Name: info.Name,
|
|
Arguments: fmt.Sprintf(`{"name": "%s", "hh": "123"}`, randStr()),
|
|
},
|
|
},
|
|
}),
|
|
nil
|
|
}
|
|
|
|
return schema.AssistantMessage("bye", nil), nil
|
|
}).AnyTimes()
|
|
cm.EXPECT().BindTools(gomock.Any()).Return(nil).AnyTimes()
|
|
|
|
a, err := NewAgent(ctx, &AgentConfig{
|
|
Model: cm,
|
|
ToolsConfig: compose.ToolsNodeConfig{
|
|
Tools: []tool.BaseTool{fakeTool},
|
|
},
|
|
MessageModifier: func(ctx context.Context, input []*schema.Message) []*schema.Message {
|
|
assert.Equal(t, len(input), times*2+1)
|
|
return input
|
|
},
|
|
MaxStep: 40,
|
|
})
|
|
assert.Nil(t, err)
|
|
|
|
out, err := a.Generate(ctx, []*schema.Message{
|
|
{
|
|
Role: schema.User,
|
|
Content: "Use greet tool to continuously say hello until you get a bye response, greet names in the following order: max, bob, alice, john, marry, joe, ken, lily, please start directly! please start directly! please start directly!",
|
|
},
|
|
}, agent.WithComposeOptions(compose.WithCallbacks(callbackForTest)))
|
|
assert.Nil(t, err)
|
|
|
|
if out != nil {
|
|
t.Log(out.Content)
|
|
}
|
|
|
|
// test return directly
|
|
times = 0
|
|
a, err = NewAgent(ctx, &AgentConfig{
|
|
Model: cm,
|
|
ToolsConfig: compose.ToolsNodeConfig{
|
|
Tools: []tool.BaseTool{fakeTool},
|
|
},
|
|
MessageModifier: func(ctx context.Context, input []*schema.Message) []*schema.Message {
|
|
assert.Equal(t, len(input), times*2+1)
|
|
return input
|
|
},
|
|
MaxStep: 40,
|
|
ToolReturnDirectly: map[string]struct{}{info.Name: {}},
|
|
})
|
|
assert.Nil(t, err)
|
|
|
|
out, err = a.Generate(ctx, []*schema.Message{
|
|
{
|
|
Role: schema.User,
|
|
Content: "Use greet tool to continuously say hello until you get a bye response, greet names in the following order: max, bob, alice, john, marry, joe, ken, lily, please start directly! please start directly! please start directly!",
|
|
},
|
|
}, agent.WithComposeOptions(compose.WithCallbacks(callbackForTest)))
|
|
assert.Nil(t, err)
|
|
|
|
if out != nil {
|
|
t.Log(out.Content)
|
|
}
|
|
}
|
|
|
|
func TestReactWithMessageRewriterAndModifier(t *testing.T) {
|
|
ctx := context.Background()
|
|
|
|
ctrl := gomock.NewController(t)
|
|
cm := mockModel.NewMockToolCallingChatModel(ctrl)
|
|
|
|
// This test simulates a single Generate call with a long history.
|
|
// The MessageRewriter should shorten the history.
|
|
// The MessageModifier should add a system prompt.
|
|
|
|
cm.EXPECT().Generate(gomock.Any(), gomock.Any(), gomock.Any()).
|
|
DoAndReturn(func(ctx context.Context, input []*schema.Message, opts ...model.Option) (*schema.Message, error) {
|
|
// Check messages passed to the model.
|
|
// Expected: [system prompt, user: "message 2", assistant: "response 2"]
|
|
assert.Len(t, input, 3)
|
|
assert.Equal(t, schema.System, input[0].Role)
|
|
assert.Equal(t, "system prompt", input[0].Content)
|
|
assert.Equal(t, schema.User, input[1].Role)
|
|
assert.Equal(t, "message 2", input[1].Content)
|
|
assert.Equal(t, schema.Assistant, input[2].Role)
|
|
assert.Equal(t, "response 2", input[2].Content)
|
|
return schema.AssistantMessage("final response", nil), nil
|
|
}).Times(1)
|
|
cm.EXPECT().WithTools(gomock.Any()).Return(cm, nil).AnyTimes()
|
|
|
|
ra, err := NewAgent(ctx, &AgentConfig{
|
|
ToolCallingModel: cm,
|
|
MessageRewriter: func(ctx context.Context, messages []*schema.Message) []*schema.Message {
|
|
// Keep only the last 2 messages if history is longer.
|
|
assert.Len(t, messages, 4) // user1, assistant1, user2, assistant2
|
|
if len(messages) > 2 {
|
|
return messages[len(messages)-2:]
|
|
}
|
|
return messages
|
|
},
|
|
MessageModifier: func(ctx context.Context, messages []*schema.Message) []*schema.Message {
|
|
// messages should be the result from rewriter
|
|
assert.Len(t, messages, 2) // user2, assistant2
|
|
|
|
// Add a system prompt
|
|
res := make([]*schema.Message, 0, len(messages)+1)
|
|
res = append(res, schema.SystemMessage("system prompt"))
|
|
res = append(res, messages...)
|
|
return res
|
|
},
|
|
})
|
|
assert.NoError(t, err)
|
|
|
|
// Simulate a conversation history
|
|
history := []*schema.Message{
|
|
schema.UserMessage("message 1"),
|
|
schema.AssistantMessage("response 1", nil),
|
|
schema.UserMessage("message 2"),
|
|
schema.AssistantMessage("response 2", nil),
|
|
}
|
|
|
|
// Run the react agent
|
|
finalMsg, err := ra.Generate(ctx, history)
|
|
assert.NoError(t, err)
|
|
assert.Equal(t, "final response", finalMsg.Content)
|
|
}
|
|
|
|
func TestReactStream(t *testing.T) {
|
|
ctx := context.Background()
|
|
|
|
fakeTool := &fakeToolGreetForTest{
|
|
tarCount: 20,
|
|
}
|
|
|
|
fakeStreamTool := &fakeStreamToolGreetForTest{
|
|
tarCount: 20,
|
|
}
|
|
|
|
ctrl := gomock.NewController(t)
|
|
cm := mockModel.NewMockChatModel(ctrl)
|
|
|
|
times := 0
|
|
cm.EXPECT().BindTools(gomock.Any()).Return(nil).AnyTimes()
|
|
cm.EXPECT().Stream(gomock.Any(), gomock.Any(), gomock.Any()).
|
|
DoAndReturn(func(ctx context.Context, input []*schema.Message, opts ...model.Option) (
|
|
*schema.StreamReader[*schema.Message], error) {
|
|
sr, sw := schema.Pipe[*schema.Message](1)
|
|
defer sw.Close()
|
|
|
|
info, _ := fakeTool.Info(ctx)
|
|
streamInfo, _ := fakeStreamTool.Info(ctx)
|
|
|
|
times++
|
|
if times <= 2 {
|
|
sw.Send(schema.AssistantMessage("hello max",
|
|
[]schema.ToolCall{
|
|
{
|
|
ID: randStr(),
|
|
Function: schema.FunctionCall{
|
|
Name: info.Name,
|
|
Arguments: fmt.Sprintf(`{"name": "%s", "hh": "tool"}`, randStr()),
|
|
},
|
|
},
|
|
}),
|
|
nil)
|
|
return sr, nil
|
|
} else if times == 3 {
|
|
sw.Send(schema.AssistantMessage("hello max",
|
|
[]schema.ToolCall{
|
|
{
|
|
ID: randStr(),
|
|
Function: schema.FunctionCall{
|
|
Name: streamInfo.Name,
|
|
Arguments: fmt.Sprintf(`{"name": "%s", "hh": "stream tool"}`, randStr()),
|
|
},
|
|
},
|
|
}),
|
|
nil)
|
|
return sr, nil
|
|
} else if times == 4 { // parallel tool call
|
|
sw.Send(schema.AssistantMessage("hello max",
|
|
[]schema.ToolCall{
|
|
{
|
|
ID: randStr(),
|
|
Function: schema.FunctionCall{
|
|
Name: info.Name,
|
|
Arguments: fmt.Sprintf(`{"name": "%s", "hh": "tool"}`, randStr()),
|
|
},
|
|
},
|
|
{
|
|
ID: randStr(),
|
|
Function: schema.FunctionCall{
|
|
Name: streamInfo.Name,
|
|
Arguments: fmt.Sprintf(`{"name": "%s", "hh": "stream tool"}`, randStr()),
|
|
},
|
|
},
|
|
}),
|
|
nil)
|
|
return sr, nil
|
|
}
|
|
|
|
sw.Send(schema.AssistantMessage("bye", nil), nil)
|
|
return sr, nil
|
|
}).AnyTimes()
|
|
|
|
a, err := NewAgent(ctx, &AgentConfig{
|
|
Model: cm,
|
|
ToolsConfig: compose.ToolsNodeConfig{
|
|
Tools: []tool.BaseTool{fakeTool, fakeStreamTool},
|
|
},
|
|
|
|
MaxStep: 40,
|
|
})
|
|
assert.Nil(t, err)
|
|
|
|
out, err := a.Stream(ctx, []*schema.Message{
|
|
{
|
|
Role: schema.User,
|
|
Content: "Use greet tool to continuously say hello until you get a bye response, greet names in the following order: max, bob, alice, john, marry, joe, ken, lily, please start directly! please start directly! please start directly!",
|
|
},
|
|
}, agent.WithComposeOptions(compose.WithCallbacks(callbackForTest)))
|
|
if err != nil {
|
|
t.Fatal(err)
|
|
}
|
|
|
|
defer out.Close()
|
|
|
|
msgs := make([]*schema.Message, 0)
|
|
for {
|
|
msg, err := out.Recv()
|
|
if err != nil {
|
|
if errors.Is(err, io.EOF) {
|
|
break
|
|
}
|
|
|
|
t.Fatal(err)
|
|
}
|
|
|
|
msgs = append(msgs, msg)
|
|
}
|
|
|
|
assert.Equal(t, 1, len(msgs))
|
|
|
|
msg, err := schema.ConcatMessages(msgs)
|
|
if err != nil {
|
|
t.Fatal(err)
|
|
}
|
|
|
|
t.Log(msg.Content)
|
|
|
|
info, err := fakeStreamTool.Info(ctx)
|
|
assert.NoError(t, err)
|
|
|
|
// test return directly
|
|
a, err = NewAgent(ctx, &AgentConfig{
|
|
Model: cm,
|
|
ToolsConfig: compose.ToolsNodeConfig{
|
|
Tools: []tool.BaseTool{fakeTool, fakeStreamTool},
|
|
},
|
|
|
|
MaxStep: 40,
|
|
ToolReturnDirectly: map[string]struct{}{info.Name: {}}, // one of the two tools is return directly
|
|
})
|
|
assert.Nil(t, err)
|
|
|
|
times = 0
|
|
out, err = a.Stream(ctx, []*schema.Message{
|
|
{
|
|
Role: schema.User,
|
|
Content: "Use greet tool to continuously say hello until you get a bye response, greet names in the following order: max, bob, alice, john, marry, joe, ken, lily, please start directly! please start directly! please start directly!",
|
|
},
|
|
}, agent.WithComposeOptions(compose.WithCallbacks(callbackForTest)))
|
|
if err != nil {
|
|
t.Fatal(err)
|
|
}
|
|
|
|
defer out.Close()
|
|
|
|
msgs = make([]*schema.Message, 0)
|
|
for {
|
|
msg, err := out.Recv()
|
|
if err != nil {
|
|
if errors.Is(err, io.EOF) {
|
|
break
|
|
}
|
|
|
|
t.Fatal(err)
|
|
}
|
|
|
|
msgs = append(msgs, msg)
|
|
}
|
|
|
|
assert.Equal(t, 1, len(msgs))
|
|
|
|
msg, err = schema.ConcatMessages(msgs)
|
|
if err != nil {
|
|
t.Fatal(err)
|
|
}
|
|
|
|
t.Log(msg.Content)
|
|
|
|
// return directly tool call within parallel tool calls
|
|
out, err = a.Stream(ctx, []*schema.Message{
|
|
{
|
|
Role: schema.User,
|
|
Content: "Use greet tool to continuously say hello until you get a bye response, greet names in the following order: max, bob, alice, john, marry, joe, ken, lily, please start directly! please start directly! please start directly!",
|
|
},
|
|
}, agent.WithComposeOptions(compose.WithCallbacks(callbackForTest)))
|
|
assert.NoError(t, err)
|
|
|
|
defer out.Close()
|
|
|
|
msgs = make([]*schema.Message, 0)
|
|
for {
|
|
msg, err := out.Recv()
|
|
if err != nil {
|
|
if errors.Is(err, io.EOF) {
|
|
break
|
|
}
|
|
assert.NoError(t, err)
|
|
}
|
|
|
|
msgs = append(msgs, msg)
|
|
}
|
|
|
|
assert.Equal(t, 1, len(msgs))
|
|
|
|
msg, err = schema.ConcatMessages(msgs)
|
|
assert.NoError(t, err)
|
|
|
|
t.Log("parallel tool call with return directly: ", msg.Content)
|
|
}
|
|
|
|
func TestReactWithModifier(t *testing.T) {
|
|
ctx := context.Background()
|
|
|
|
fakeTool := &fakeToolGreetForTest{}
|
|
ctrl := gomock.NewController(t)
|
|
cm := mockModel.NewMockChatModel(ctrl)
|
|
|
|
times := 0
|
|
cm.EXPECT().Generate(gomock.Any(), gomock.Any(), gomock.Any()).
|
|
DoAndReturn(func(ctx context.Context, input []*schema.Message, opts ...model.Option) (*schema.Message, error) {
|
|
times++
|
|
if times <= 2 {
|
|
info, _ := fakeTool.Info(ctx)
|
|
|
|
return schema.AssistantMessage("hello max",
|
|
[]schema.ToolCall{
|
|
{
|
|
ID: randStr(),
|
|
Function: schema.FunctionCall{
|
|
Name: info.Name,
|
|
Arguments: fmt.Sprintf(`{"name": "%s", "hh": "123"}`, randStr()),
|
|
},
|
|
},
|
|
}),
|
|
nil
|
|
}
|
|
|
|
return schema.AssistantMessage("bye", nil), nil
|
|
}).AnyTimes()
|
|
cm.EXPECT().BindTools(gomock.Any()).Return(nil).AnyTimes()
|
|
|
|
a, err := NewAgent(ctx, &AgentConfig{
|
|
Model: cm,
|
|
ToolsConfig: compose.ToolsNodeConfig{
|
|
Tools: []tool.BaseTool{fakeTool},
|
|
},
|
|
MessageModifier: func(ctx context.Context, input []*schema.Message) []*schema.Message {
|
|
res := make([]*schema.Message, 0, len(input)+1)
|
|
|
|
res = append(res, schema.SystemMessage("you are a helpful assistant"))
|
|
res = append(res, input...)
|
|
return res
|
|
},
|
|
|
|
MaxStep: 40,
|
|
})
|
|
|
|
assert.Nil(t, err)
|
|
|
|
out, err := a.Generate(ctx, []*schema.Message{
|
|
{
|
|
Role: schema.User,
|
|
Content: "hello",
|
|
},
|
|
}, agent.WithComposeOptions(compose.WithCallbacks(callbackForTest)))
|
|
if err != nil {
|
|
t.Fatal(err)
|
|
}
|
|
|
|
if out != nil {
|
|
t.Log(out.Content)
|
|
}
|
|
}
|
|
|
|
func TestAgentInGraph(t *testing.T) {
|
|
t.Run("agent generate in chain", func(t *testing.T) {
|
|
ctx := context.Background()
|
|
|
|
fakeTool := &fakeToolGreetForTest{}
|
|
ctrl := gomock.NewController(t)
|
|
cm := mockModel.NewMockChatModel(ctrl)
|
|
|
|
times := 0
|
|
cm.EXPECT().Generate(gomock.Any(), gomock.Any(), gomock.Any()).
|
|
DoAndReturn(func(ctx context.Context, input []*schema.Message, opts ...model.Option) (*schema.Message, error) {
|
|
|
|
times += 1
|
|
if times <= 2 {
|
|
info, _ := fakeTool.Info(ctx)
|
|
|
|
return schema.AssistantMessage("hello max",
|
|
[]schema.ToolCall{
|
|
{
|
|
ID: randStr(),
|
|
Function: schema.FunctionCall{
|
|
Name: info.Name,
|
|
Arguments: fmt.Sprintf(`{"name": "%s", "hh": "123"}`, randStr()),
|
|
},
|
|
},
|
|
}),
|
|
nil
|
|
}
|
|
|
|
return schema.AssistantMessage("bye", nil), nil
|
|
|
|
}).Times(3)
|
|
cm.EXPECT().BindTools(gomock.Any()).Return(nil).AnyTimes()
|
|
|
|
a, err := NewAgent(ctx, &AgentConfig{
|
|
Model: cm,
|
|
ToolsConfig: compose.ToolsNodeConfig{
|
|
Tools: []tool.BaseTool{fakeTool, &fakeStreamToolGreetForTest{}},
|
|
},
|
|
|
|
MaxStep: 40,
|
|
})
|
|
assert.Nil(t, err)
|
|
|
|
chain := compose.NewChain[[]*schema.Message, string]()
|
|
agentLambda, err := compose.AnyLambda(a.Generate, a.Stream, nil, nil)
|
|
assert.Nil(t, err)
|
|
|
|
chain.
|
|
AppendLambda(agentLambda).
|
|
AppendLambda(compose.InvokableLambda(func(ctx context.Context, input *schema.Message) (string, error) {
|
|
t.Log("got agent response: ", input.Content)
|
|
return input.Content, nil
|
|
}))
|
|
r, err := chain.Compile(ctx)
|
|
assert.Nil(t, err)
|
|
|
|
res, err := r.Invoke(ctx, []*schema.Message{{Role: schema.User, Content: "hello"}},
|
|
compose.WithCallbacks(callbackForTest))
|
|
assert.Nil(t, err)
|
|
|
|
t.Log(res)
|
|
})
|
|
|
|
t.Run("agent stream in chain", func(t *testing.T) {
|
|
|
|
fakeStreamTool := &fakeStreamToolGreetForTest{}
|
|
ctx := context.Background()
|
|
ctrl := gomock.NewController(t)
|
|
cm := mockModel.NewMockChatModel(ctrl)
|
|
|
|
times := 0
|
|
cm.EXPECT().Stream(gomock.Any(), gomock.Any(), gomock.Any()).
|
|
DoAndReturn(func(ctx context.Context, input []*schema.Message, opts ...model.Option) (
|
|
*schema.StreamReader[*schema.Message], error) {
|
|
sr, sw := schema.Pipe[*schema.Message](1)
|
|
defer sw.Close()
|
|
|
|
times += 1
|
|
if times <= 2 {
|
|
info, _ := fakeStreamTool.Info(ctx)
|
|
sw.Send(schema.AssistantMessage("hello max",
|
|
[]schema.ToolCall{
|
|
{
|
|
ID: randStr(),
|
|
Function: schema.FunctionCall{
|
|
Name: info.Name,
|
|
Arguments: fmt.Sprintf(`{"name": "%s", "hh": "123"}`, randStr()),
|
|
},
|
|
},
|
|
}),
|
|
nil)
|
|
return sr, nil
|
|
}
|
|
|
|
sw.Send(schema.AssistantMessage("bye", nil), nil)
|
|
return sr, nil
|
|
}).Times(3)
|
|
cm.EXPECT().BindTools(gomock.Any()).Return(nil).AnyTimes()
|
|
|
|
a, err := NewAgent(ctx, &AgentConfig{
|
|
Model: cm,
|
|
ToolsConfig: compose.ToolsNodeConfig{
|
|
Tools: []tool.BaseTool{&fakeToolGreetForTest{}, fakeStreamTool},
|
|
},
|
|
|
|
MaxStep: 40,
|
|
})
|
|
assert.Nil(t, err)
|
|
|
|
chain := compose.NewChain[[]*schema.Message, string]()
|
|
agentGraph, opts := a.ExportGraph()
|
|
assert.Nil(t, err)
|
|
|
|
chain.
|
|
AppendGraph(agentGraph, opts...).
|
|
AppendLambda(compose.InvokableLambda(func(ctx context.Context, input *schema.Message) (string, error) {
|
|
t.Log("got agent response: ", input.Content)
|
|
return input.Content, nil
|
|
}))
|
|
r, err := chain.Compile(ctx)
|
|
assert.Nil(t, err)
|
|
|
|
outStream, err := r.Stream(ctx, []*schema.Message{{Role: schema.User, Content: "hello"}},
|
|
compose.WithCallbacks(callbackForTest))
|
|
if err != nil {
|
|
t.Fatal(err)
|
|
}
|
|
|
|
defer outStream.Close()
|
|
|
|
msg := ""
|
|
for {
|
|
msgItem, err := outStream.Recv()
|
|
if err != nil {
|
|
if errors.Is(err, io.EOF) {
|
|
break
|
|
}
|
|
|
|
t.Fatal(err)
|
|
}
|
|
|
|
msg += msgItem
|
|
}
|
|
|
|
t.Log(msg)
|
|
})
|
|
|
|
}
|
|
|
|
func TestWithTools(t *testing.T) {
|
|
ctx := context.Background()
|
|
|
|
fakeTool := &fakeToolGreetForTest{
|
|
tarCount: 2,
|
|
}
|
|
fakeStreamTool := &fakeStreamToolGreetForTest{
|
|
tarCount: 2,
|
|
}
|
|
|
|
ctrl := gomock.NewController(t)
|
|
cm := mockModel.NewMockToolCallingChatModel(ctrl)
|
|
|
|
times := 0
|
|
cm.EXPECT().Generate(gomock.Any(), gomock.Any(), gomock.Any()).
|
|
DoAndReturn(func(ctx context.Context, input []*schema.Message, opts ...model.Option) (*schema.Message, error) {
|
|
times++
|
|
if times <= 1 {
|
|
info, _ := fakeTool.Info(ctx)
|
|
return schema.AssistantMessage("calling tool",
|
|
[]schema.ToolCall{
|
|
{
|
|
ID: randStr(),
|
|
Function: schema.FunctionCall{
|
|
Name: info.Name,
|
|
Arguments: `{"name": "test"}`,
|
|
},
|
|
},
|
|
}),
|
|
nil
|
|
}
|
|
return schema.AssistantMessage("done", nil), nil
|
|
}).AnyTimes()
|
|
|
|
cm.EXPECT().Stream(gomock.Any(), gomock.Any(), gomock.Any()).
|
|
DoAndReturn(func(ctx context.Context, input []*schema.Message, opts ...model.Option) (
|
|
*schema.StreamReader[*schema.Message], error) {
|
|
sr, sw := schema.Pipe[*schema.Message](1)
|
|
defer sw.Close()
|
|
|
|
times++
|
|
if times <= 2 {
|
|
info, _ := fakeStreamTool.Info(ctx)
|
|
sw.Send(schema.AssistantMessage("calling stream tool",
|
|
[]schema.ToolCall{
|
|
{
|
|
ID: randStr(),
|
|
Function: schema.FunctionCall{
|
|
Name: info.Name,
|
|
Arguments: `{"name": "test"}`,
|
|
},
|
|
},
|
|
}),
|
|
nil)
|
|
return sr, nil
|
|
}
|
|
|
|
sw.Send(schema.AssistantMessage("stream done", nil), nil)
|
|
return sr, nil
|
|
}).AnyTimes()
|
|
|
|
// Test WithTools function
|
|
toolOptions, err := WithTools(ctx, fakeTool, fakeStreamTool)
|
|
assert.NoError(t, err)
|
|
assert.Len(t, toolOptions, 2, "WithTools should return exactly 2 options")
|
|
|
|
// Create agent without tools in config
|
|
a, err := NewAgent(ctx, &AgentConfig{
|
|
ToolCallingModel: cm,
|
|
MaxStep: 10,
|
|
})
|
|
assert.NoError(t, err)
|
|
|
|
// Test Generate with WithTools options
|
|
times = 0
|
|
msg, err := a.Generate(ctx, []*schema.Message{
|
|
schema.UserMessage("test generate with tools"),
|
|
}, toolOptions...)
|
|
assert.NoError(t, err)
|
|
assert.Equal(t, "done", msg.Content)
|
|
|
|
// Test Stream with WithTools options
|
|
times = 0
|
|
stream, err := a.Stream(ctx, []*schema.Message{
|
|
schema.UserMessage("test stream with tools"),
|
|
}, toolOptions...)
|
|
assert.NoError(t, err)
|
|
|
|
defer stream.Close()
|
|
msgs := make([]*schema.Message, 0)
|
|
for {
|
|
msg, err := stream.Recv()
|
|
if err != nil {
|
|
if errors.Is(err, io.EOF) {
|
|
break
|
|
}
|
|
assert.NoError(t, err)
|
|
}
|
|
msgs = append(msgs, msg)
|
|
}
|
|
|
|
assert.Len(t, msgs, 1)
|
|
concatMsg, err := schema.ConcatMessages(msgs)
|
|
assert.NoError(t, err)
|
|
assert.Equal(t, "stream done", concatMsg.Content)
|
|
|
|
// Test error case - tool Info() returns error
|
|
errorTool := &errorToolForTest{}
|
|
_, err = WithTools(ctx, errorTool)
|
|
assert.Error(t, err)
|
|
assert.Contains(t, err.Error(), "info error")
|
|
}
|
|
|
|
// Helper tool for testing error cases
|
|
type errorToolForTest struct{}
|
|
|
|
func (t *errorToolForTest) Info(_ context.Context) (*schema.ToolInfo, error) {
|
|
return nil, errors.New("info error")
|
|
}
|
|
|
|
func (t *errorToolForTest) InvokableRun(_ context.Context, _ string, _ ...tool.Option) (string, error) {
|
|
return "", nil
|
|
}
|
|
|
|
type fakeStreamToolGreetForTest struct {
|
|
tarCount int
|
|
curCount int
|
|
}
|
|
|
|
func (t *fakeStreamToolGreetForTest) StreamableRun(_ context.Context, argumentsInJSON string, _ ...tool.Option) (
|
|
*schema.StreamReader[string], error) {
|
|
p := &fakeToolInput{}
|
|
err := sonic.UnmarshalString(argumentsInJSON, p)
|
|
if err != nil {
|
|
return nil, err
|
|
}
|
|
|
|
if t.curCount >= t.tarCount {
|
|
s := schema.StreamReaderFromArray([]string{`{"say": "bye"}`})
|
|
return s, nil
|
|
}
|
|
t.curCount++
|
|
s := schema.StreamReaderFromArray([]string{fmt.Sprintf(`{"say": "hello %v"}`, p.Name)})
|
|
return s, nil
|
|
}
|
|
|
|
type fakeToolGreetForTest struct {
|
|
tarCount int
|
|
curCount int
|
|
}
|
|
|
|
func (t *fakeToolGreetForTest) Info(_ context.Context) (*schema.ToolInfo, error) {
|
|
return &schema.ToolInfo{
|
|
Name: "greet",
|
|
Desc: "greet with name",
|
|
ParamsOneOf: schema.NewParamsOneOfByParams(
|
|
map[string]*schema.ParameterInfo{
|
|
"name": {
|
|
Desc: "user name who to greet",
|
|
Required: true,
|
|
Type: schema.String,
|
|
},
|
|
}),
|
|
}, nil
|
|
}
|
|
|
|
func (t *fakeStreamToolGreetForTest) Info(_ context.Context) (*schema.ToolInfo, error) {
|
|
return &schema.ToolInfo{
|
|
Name: "greet in stream",
|
|
Desc: "greet with name in stream",
|
|
ParamsOneOf: schema.NewParamsOneOfByParams(
|
|
map[string]*schema.ParameterInfo{
|
|
"name": {
|
|
Desc: "user name who to greet",
|
|
Required: true,
|
|
Type: schema.String,
|
|
},
|
|
}),
|
|
}, nil
|
|
}
|
|
|
|
func (t *fakeToolGreetForTest) InvokableRun(_ context.Context, argumentsInJSON string, _ ...tool.Option) (string, error) {
|
|
p := &fakeToolInput{}
|
|
err := sonic.UnmarshalString(argumentsInJSON, p)
|
|
if err != nil {
|
|
return "", err
|
|
}
|
|
|
|
if t.curCount >= t.tarCount {
|
|
return `{"say": "bye"}`, nil
|
|
}
|
|
|
|
t.curCount++
|
|
return fmt.Sprintf(`{"say": "hello %v"}`, p.Name), nil
|
|
}
|
|
|
|
type fakeToolInput struct {
|
|
Name string `json:"name"`
|
|
}
|
|
|
|
func randStr() string {
|
|
seeds := []rune("this is a seed")
|
|
b := make([]rune, 8)
|
|
for i := range b {
|
|
b[i] = seeds[rand.Intn(len(seeds))]
|
|
}
|
|
return string(b)
|
|
}
|
|
|
|
var callbackForTest = BuildAgentCallback(&template.ModelCallbackHandler{}, &template.ToolCallbackHandler{})
|