wehub-resource-sync
2418 lines
72 KiB
Go
2418 lines
72 KiB
Go
/*
|
|
* 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 adk
|
|
|
|
import (
|
|
"context"
|
|
"encoding/json"
|
|
"errors"
|
|
"testing"
|
|
"time"
|
|
|
|
"github.com/stretchr/testify/assert"
|
|
"github.com/stretchr/testify/require"
|
|
"go.uber.org/mock/gomock"
|
|
|
|
"github.com/cloudwego/eino/components/model"
|
|
"github.com/cloudwego/eino/components/tool"
|
|
"github.com/cloudwego/eino/compose"
|
|
mockModel "github.com/cloudwego/eino/internal/mock/components/model"
|
|
"github.com/cloudwego/eino/schema"
|
|
)
|
|
|
|
// TestChatModelAgentRun tests the Run method of ChatModelAgent
|
|
func TestChatModelAgentRun(t *testing.T) {
|
|
// Basic test for Run method
|
|
t.Run("BasicFunctionality", func(t *testing.T) {
|
|
ctx := context.Background()
|
|
|
|
// Create a mock chat model
|
|
ctrl := gomock.NewController(t)
|
|
cm := mockModel.NewMockToolCallingChatModel(ctrl)
|
|
|
|
// Set up expectations for the mock model
|
|
cm.EXPECT().Generate(gomock.Any(), gomock.Any(), gomock.Any()).
|
|
Return(schema.AssistantMessage("Hello, I am an AI assistant.", nil), nil).
|
|
Times(1)
|
|
|
|
// Create a ChatModelAgent
|
|
agent, err := NewChatModelAgent(ctx, &ChatModelAgentConfig{
|
|
Name: "TestAgent",
|
|
Description: "Test agent for unit testing",
|
|
Instruction: "You are a helpful assistant.",
|
|
Model: cm,
|
|
})
|
|
assert.NoError(t, err)
|
|
assert.NotNil(t, agent)
|
|
|
|
// Run the agent
|
|
input := &AgentInput{
|
|
Messages: []Message{
|
|
schema.UserMessage("Hello, who are you?"),
|
|
},
|
|
}
|
|
iterator := agent.Run(ctx, input)
|
|
assert.NotNil(t, iterator)
|
|
|
|
// Get the event from the iterator
|
|
event, ok := iterator.Next()
|
|
assert.True(t, ok)
|
|
assert.NotNil(t, event)
|
|
assert.Nil(t, event.Err)
|
|
assert.NotNil(t, event.Output)
|
|
assert.NotNil(t, event.Output.MessageOutput)
|
|
|
|
// Verify the message content
|
|
msg := event.Output.MessageOutput.Message
|
|
assert.NotNil(t, msg)
|
|
assert.Equal(t, "Hello, I am an AI assistant.", msg.Content)
|
|
|
|
// No more events
|
|
_, ok = iterator.Next()
|
|
assert.False(t, ok)
|
|
})
|
|
|
|
t.Run("BasicChatModelWithAgentMiddleware", func(t *testing.T) {
|
|
ctx := context.Background()
|
|
|
|
// Create a mock chat model
|
|
ctrl := gomock.NewController(t)
|
|
cm := mockModel.NewMockToolCallingChatModel(ctrl)
|
|
|
|
// Set up expectations for the mock model
|
|
cm.EXPECT().Generate(gomock.Any(), gomock.Any(), gomock.Any()).
|
|
Return(schema.AssistantMessage("Hello, I am an AI assistant.", nil), nil).
|
|
Times(1)
|
|
|
|
afterChatModelExecuted := false
|
|
|
|
// Create a ChatModelAgent
|
|
agent, err := NewChatModelAgent(ctx, &ChatModelAgentConfig{
|
|
Name: "TestAgent",
|
|
Description: "Test agent for unit testing",
|
|
Instruction: "You are a helpful assistant.",
|
|
Model: cm,
|
|
Middlewares: []AgentMiddleware{
|
|
{
|
|
BeforeChatModel: func(ctx context.Context, state *ChatModelAgentState) error {
|
|
state.Messages = append(state.Messages, schema.UserMessage("m"))
|
|
return nil
|
|
},
|
|
AfterChatModel: func(ctx context.Context, state *ChatModelAgentState) error {
|
|
assert.Len(t, state.Messages, 4)
|
|
afterChatModelExecuted = true
|
|
return nil
|
|
},
|
|
},
|
|
},
|
|
})
|
|
assert.NoError(t, err)
|
|
assert.NotNil(t, agent)
|
|
|
|
// Run the agent
|
|
input := &AgentInput{
|
|
Messages: []Message{
|
|
schema.UserMessage("Hello, who are you?"),
|
|
},
|
|
}
|
|
iterator := agent.Run(ctx, input)
|
|
assert.NotNil(t, iterator)
|
|
|
|
// Get the event from the iterator
|
|
event, ok := iterator.Next()
|
|
assert.True(t, ok)
|
|
assert.NotNil(t, event)
|
|
assert.Nil(t, event.Err)
|
|
assert.NotNil(t, event.Output)
|
|
assert.NotNil(t, event.Output.MessageOutput)
|
|
|
|
// Verify the message content
|
|
msg := event.Output.MessageOutput.Message
|
|
assert.NotNil(t, msg)
|
|
assert.Equal(t, "Hello, I am an AI assistant.", msg.Content)
|
|
|
|
// No more events
|
|
_, ok = iterator.Next()
|
|
assert.False(t, ok)
|
|
|
|
assert.True(t, afterChatModelExecuted)
|
|
})
|
|
|
|
t.Run("AfterChatModel_NoTools_ModifyDoesNotAffectEvent", func(t *testing.T) {
|
|
ctx := context.Background()
|
|
|
|
ctrl := gomock.NewController(t)
|
|
cm := mockModel.NewMockToolCallingChatModel(ctrl)
|
|
|
|
cm.EXPECT().Generate(gomock.Any(), gomock.Any(), gomock.Any()).
|
|
Return(schema.AssistantMessage("original content", nil), nil).
|
|
Times(1)
|
|
|
|
var capturedMessages []*schema.Message
|
|
|
|
agent, err := NewChatModelAgent(ctx, &ChatModelAgentConfig{
|
|
Name: "TestAgent",
|
|
Description: "Test agent for AfterChatModel NoTools scenario",
|
|
Instruction: "You are a helpful assistant.",
|
|
Model: cm,
|
|
Middlewares: []AgentMiddleware{
|
|
{
|
|
AfterChatModel: func(ctx context.Context, state *ChatModelAgentState) error {
|
|
capturedMessages = make([]*schema.Message, len(state.Messages))
|
|
copy(capturedMessages, state.Messages)
|
|
state.Messages = append(state.Messages, schema.AssistantMessage("appended content", nil))
|
|
return nil
|
|
},
|
|
},
|
|
},
|
|
})
|
|
assert.NoError(t, err)
|
|
|
|
input := &AgentInput{
|
|
Messages: []Message{schema.UserMessage("Hello")},
|
|
}
|
|
iterator := agent.Run(ctx, input)
|
|
|
|
event, ok := iterator.Next()
|
|
assert.True(t, ok)
|
|
assert.NotNil(t, event)
|
|
assert.Nil(t, event.Err)
|
|
assert.NotNil(t, event.Output)
|
|
assert.NotNil(t, event.Output.MessageOutput)
|
|
|
|
msg := event.Output.MessageOutput.Message
|
|
assert.NotNil(t, msg)
|
|
assert.Equal(t, "original content", msg.Content)
|
|
|
|
_, ok = iterator.Next()
|
|
assert.False(t, ok)
|
|
|
|
assert.Len(t, capturedMessages, 3)
|
|
})
|
|
|
|
t.Run("AfterChatModel_ReAct_ModifyAffectsFlow", func(t *testing.T) {
|
|
ctx := context.Background()
|
|
|
|
ctrl := gomock.NewController(t)
|
|
cm := mockModel.NewMockToolCallingChatModel(ctrl)
|
|
|
|
generateCount := 0
|
|
cm.EXPECT().Generate(gomock.Any(), gomock.Any(), gomock.Any()).
|
|
DoAndReturn(func(ctx context.Context, msgs []*schema.Message, opts ...model.Option) (*schema.Message, error) {
|
|
generateCount++
|
|
if generateCount == 1 {
|
|
return schema.AssistantMessage("first response with tool call", []schema.ToolCall{
|
|
{ID: "tc1", Function: schema.FunctionCall{Name: "test_tool", Arguments: "{}"}},
|
|
}), nil
|
|
}
|
|
return schema.AssistantMessage("final response", nil), nil
|
|
}).AnyTimes()
|
|
cm.EXPECT().WithTools(gomock.Any()).Return(cm, nil).AnyTimes()
|
|
|
|
toolCalled := false
|
|
testTool := &fakeToolForTest{tarCount: 0}
|
|
|
|
agent, err := NewChatModelAgent(ctx, &ChatModelAgentConfig{
|
|
Name: "TestAgent",
|
|
Description: "Test agent for AfterChatModel ReAct scenario",
|
|
Instruction: "You are a helpful assistant.",
|
|
Model: cm,
|
|
ToolsConfig: ToolsConfig{
|
|
ToolsNodeConfig: compose.ToolsNodeConfig{
|
|
Tools: []tool.BaseTool{testTool},
|
|
},
|
|
},
|
|
Middlewares: []AgentMiddleware{
|
|
{
|
|
AfterChatModel: func(ctx context.Context, state *ChatModelAgentState) error {
|
|
lastMsg := state.Messages[len(state.Messages)-1]
|
|
if len(lastMsg.ToolCalls) > 0 {
|
|
toolCalled = true
|
|
state.Messages[len(state.Messages)-1] = schema.AssistantMessage("modified to remove tool call", nil)
|
|
}
|
|
return nil
|
|
},
|
|
},
|
|
},
|
|
})
|
|
assert.NoError(t, err)
|
|
|
|
input := &AgentInput{
|
|
Messages: []Message{schema.UserMessage("Hello")},
|
|
}
|
|
iterator := agent.Run(ctx, input)
|
|
|
|
var events []*AgentEvent
|
|
for {
|
|
event, ok := iterator.Next()
|
|
if !ok {
|
|
break
|
|
}
|
|
events = append(events, event)
|
|
}
|
|
|
|
assert.True(t, toolCalled)
|
|
assert.Equal(t, 1, generateCount)
|
|
|
|
assert.Equal(t, 1, len(events))
|
|
event := events[0]
|
|
assert.NotNil(t, event.Output)
|
|
assert.NotNil(t, event.Output.MessageOutput)
|
|
assert.Equal(t, "first response with tool call", event.Output.MessageOutput.Message.Content)
|
|
assert.Len(t, event.Output.MessageOutput.Message.ToolCalls, 1)
|
|
})
|
|
|
|
t.Run("AfterChatModel_ReAct_AppendToolCall_AffectsFlow", func(t *testing.T) {
|
|
ctx := context.Background()
|
|
|
|
ctrl := gomock.NewController(t)
|
|
cm := mockModel.NewMockToolCallingChatModel(ctrl)
|
|
|
|
generateCount := 0
|
|
cm.EXPECT().Generate(gomock.Any(), gomock.Any(), gomock.Any()).
|
|
DoAndReturn(func(ctx context.Context, msgs []*schema.Message, opts ...model.Option) (*schema.Message, error) {
|
|
generateCount++
|
|
if generateCount == 1 {
|
|
return schema.AssistantMessage("first response no tool", nil), nil
|
|
}
|
|
return schema.AssistantMessage("final response", nil), nil
|
|
}).AnyTimes()
|
|
cm.EXPECT().WithTools(gomock.Any()).Return(cm, nil).AnyTimes()
|
|
|
|
testTool := &fakeToolForTest{tarCount: 0}
|
|
|
|
agent, err := NewChatModelAgent(ctx, &ChatModelAgentConfig{
|
|
Name: "TestAgent",
|
|
Description: "Test agent for AfterChatModel ReAct append tool call",
|
|
Instruction: "You are a helpful assistant.",
|
|
Model: cm,
|
|
ToolsConfig: ToolsConfig{
|
|
ToolsNodeConfig: compose.ToolsNodeConfig{
|
|
Tools: []tool.BaseTool{testTool},
|
|
},
|
|
},
|
|
Middlewares: []AgentMiddleware{
|
|
{
|
|
AfterChatModel: func(ctx context.Context, state *ChatModelAgentState) error {
|
|
if generateCount == 1 {
|
|
state.Messages[len(state.Messages)-1] = schema.AssistantMessage("modified with tool call", []schema.ToolCall{
|
|
{ID: "tc1", Function: schema.FunctionCall{Name: "test_tool", Arguments: "{}"}},
|
|
})
|
|
}
|
|
return nil
|
|
},
|
|
},
|
|
},
|
|
})
|
|
assert.NoError(t, err)
|
|
|
|
input := &AgentInput{
|
|
Messages: []Message{schema.UserMessage("Hello")},
|
|
}
|
|
iterator := agent.Run(ctx, input)
|
|
|
|
var events []*AgentEvent
|
|
for {
|
|
event, ok := iterator.Next()
|
|
if !ok {
|
|
break
|
|
}
|
|
events = append(events, event)
|
|
}
|
|
|
|
assert.Equal(t, 2, generateCount)
|
|
|
|
assert.Equal(t, 3, len(events))
|
|
|
|
event0 := events[0]
|
|
assert.NotNil(t, event0.Output)
|
|
assert.NotNil(t, event0.Output.MessageOutput)
|
|
assert.Equal(t, "first response no tool", event0.Output.MessageOutput.Message.Content)
|
|
assert.Empty(t, event0.Output.MessageOutput.Message.ToolCalls)
|
|
|
|
event2 := events[2]
|
|
assert.NotNil(t, event2.Output)
|
|
assert.NotNil(t, event2.Output.MessageOutput)
|
|
assert.Equal(t, "final response", event2.Output.MessageOutput.Message.Content)
|
|
})
|
|
|
|
// Test with streaming output
|
|
t.Run("StreamOutput", func(t *testing.T) {
|
|
ctx := context.Background()
|
|
|
|
// Create a mock chat model
|
|
ctrl := gomock.NewController(t)
|
|
cm := mockModel.NewMockToolCallingChatModel(ctrl)
|
|
|
|
// Create a stream reader for the mock response
|
|
sr := schema.StreamReaderFromArray([]*schema.Message{
|
|
schema.AssistantMessage("Hello", nil),
|
|
schema.AssistantMessage(", I am", nil),
|
|
schema.AssistantMessage(" an AI assistant.", nil),
|
|
})
|
|
|
|
// Set up expectations for the mock model
|
|
cm.EXPECT().Stream(gomock.Any(), gomock.Any(), gomock.Any()).
|
|
Return(sr, nil).
|
|
Times(1)
|
|
|
|
// Create a ChatModelAgent
|
|
agent, err := NewChatModelAgent(ctx, &ChatModelAgentConfig{
|
|
Name: "TestAgent",
|
|
Description: "Test agent for unit testing",
|
|
Instruction: "You are a helpful assistant.",
|
|
Model: cm,
|
|
})
|
|
assert.NoError(t, err)
|
|
assert.NotNil(t, agent)
|
|
|
|
// Run the agent with streaming enabled
|
|
input := &AgentInput{
|
|
Messages: []Message{schema.UserMessage("Hello, who are you?")},
|
|
EnableStreaming: true,
|
|
}
|
|
iterator := agent.Run(ctx, input)
|
|
assert.NotNil(t, iterator)
|
|
|
|
// Get the event from the iterator
|
|
event, ok := iterator.Next()
|
|
assert.True(t, ok)
|
|
assert.NotNil(t, event)
|
|
assert.Nil(t, event.Err)
|
|
assert.NotNil(t, event.Output)
|
|
assert.NotNil(t, event.Output.MessageOutput)
|
|
assert.True(t, event.Output.MessageOutput.IsStreaming)
|
|
|
|
// No more events
|
|
_, ok = iterator.Next()
|
|
assert.False(t, ok)
|
|
})
|
|
|
|
// Test error handling
|
|
t.Run("ErrorHandling", func(t *testing.T) {
|
|
ctx := context.Background()
|
|
|
|
// Create a mock chat model
|
|
ctrl := gomock.NewController(t)
|
|
cm := mockModel.NewMockToolCallingChatModel(ctrl)
|
|
|
|
// Set up expectations for the mock model to return an error
|
|
cm.EXPECT().Generate(gomock.Any(), gomock.Any(), gomock.Any()).
|
|
Return(nil, errors.New("model error")).
|
|
Times(1)
|
|
|
|
// Create a ChatModelAgent
|
|
agent, err := NewChatModelAgent(ctx, &ChatModelAgentConfig{
|
|
Name: "TestAgent",
|
|
Description: "Test agent for unit testing",
|
|
Instruction: "You are a helpful assistant.",
|
|
Model: cm,
|
|
})
|
|
assert.NoError(t, err)
|
|
assert.NotNil(t, agent)
|
|
|
|
// Run the agent
|
|
input := &AgentInput{
|
|
Messages: []Message{schema.UserMessage("Hello, who are you?")},
|
|
}
|
|
iterator := agent.Run(ctx, input)
|
|
assert.NotNil(t, iterator)
|
|
|
|
// Get the event from the iterator, should contain an error
|
|
event, ok := iterator.Next()
|
|
assert.True(t, ok)
|
|
assert.NotNil(t, event)
|
|
assert.NotNil(t, event.Err)
|
|
assert.Contains(t, event.Err.Error(), "model error")
|
|
|
|
// No more events
|
|
_, ok = iterator.Next()
|
|
assert.False(t, ok)
|
|
})
|
|
|
|
// Test with tools
|
|
t.Run("WithTools", func(t *testing.T) {
|
|
ctx := context.Background()
|
|
|
|
// Create a fake tool for testing
|
|
fakeTool := &fakeToolForTest{
|
|
tarCount: 1,
|
|
}
|
|
|
|
info, err := fakeTool.Info(ctx)
|
|
assert.NoError(t, err)
|
|
|
|
// Create a mock chat model
|
|
ctrl := gomock.NewController(t)
|
|
cm := mockModel.NewMockToolCallingChatModel(ctrl)
|
|
|
|
// Set up expectations for the mock model
|
|
cm.EXPECT().Generate(gomock.Any(), gomock.Any(), gomock.Any()).
|
|
Return(schema.AssistantMessage("Using tool",
|
|
[]schema.ToolCall{
|
|
{
|
|
ID: "tool-call-1",
|
|
Function: schema.FunctionCall{
|
|
Name: info.Name,
|
|
Arguments: `{"name": "test user"}`,
|
|
},
|
|
}}), nil).
|
|
Times(1)
|
|
cm.EXPECT().Generate(gomock.Any(), gomock.Any(), gomock.Any()).
|
|
Return(schema.AssistantMessage("Task completed", nil), nil).
|
|
Times(1)
|
|
cm.EXPECT().WithTools(gomock.Any()).Return(cm, nil).AnyTimes()
|
|
|
|
// Create a ChatModelAgent with tools
|
|
agent, err := NewChatModelAgent(ctx, &ChatModelAgentConfig{
|
|
Name: "TestAgent",
|
|
Description: "Test agent for unit testing",
|
|
Instruction: "You are a helpful assistant.",
|
|
Model: cm,
|
|
ToolsConfig: ToolsConfig{
|
|
ToolsNodeConfig: compose.ToolsNodeConfig{
|
|
Tools: []tool.BaseTool{fakeTool},
|
|
},
|
|
},
|
|
})
|
|
assert.NoError(t, err)
|
|
assert.NotNil(t, agent)
|
|
|
|
// Run the agent
|
|
input := &AgentInput{
|
|
Messages: []Message{schema.UserMessage("Use the test tool")},
|
|
}
|
|
iterator := agent.Run(ctx, input)
|
|
assert.NotNil(t, iterator)
|
|
|
|
// Get events from the iterator
|
|
// First event should be the model output with tool call
|
|
event1, ok := iterator.Next()
|
|
assert.True(t, ok)
|
|
assert.NotNil(t, event1)
|
|
assert.Nil(t, event1.Err)
|
|
assert.NotNil(t, event1.Output)
|
|
assert.NotNil(t, event1.Output.MessageOutput)
|
|
assert.Equal(t, schema.Assistant, event1.Output.MessageOutput.Role)
|
|
|
|
// Second event should be the tool output
|
|
event2, ok := iterator.Next()
|
|
assert.True(t, ok)
|
|
assert.NotNil(t, event2)
|
|
assert.Nil(t, event2.Err)
|
|
assert.NotNil(t, event2.Output)
|
|
assert.NotNil(t, event2.Output.MessageOutput)
|
|
assert.Equal(t, schema.Tool, event2.Output.MessageOutput.Role)
|
|
|
|
// Third event should be the final model output
|
|
event3, ok := iterator.Next()
|
|
assert.True(t, ok)
|
|
assert.NotNil(t, event3)
|
|
assert.Nil(t, event3.Err)
|
|
assert.NotNil(t, event3.Output)
|
|
assert.NotNil(t, event3.Output.MessageOutput)
|
|
assert.Equal(t, schema.Assistant, event3.Output.MessageOutput.Role)
|
|
|
|
// No more events
|
|
_, ok = iterator.Next()
|
|
assert.False(t, ok)
|
|
})
|
|
|
|
t.Run("WrapToolCall_ToolMiddleware", func(t *testing.T) {
|
|
ctx := context.Background()
|
|
|
|
ctrl := gomock.NewController(t)
|
|
cm := mockModel.NewMockToolCallingChatModel(ctrl)
|
|
|
|
fakeTool := &fakeToolForTest{tarCount: 1}
|
|
info, err := fakeTool.Info(ctx)
|
|
assert.NoError(t, err)
|
|
|
|
generateCount := 0
|
|
cm.EXPECT().Generate(gomock.Any(), gomock.Any(), gomock.Any()).
|
|
DoAndReturn(func(ctx context.Context, msgs []*schema.Message, opts ...model.Option) (*schema.Message, error) {
|
|
generateCount++
|
|
if generateCount == 1 {
|
|
return schema.AssistantMessage("calling tool", []schema.ToolCall{
|
|
{ID: "tc1", Function: schema.FunctionCall{Name: info.Name, Arguments: `{"name":"test"}`}},
|
|
}), nil
|
|
}
|
|
return schema.AssistantMessage("done", nil), nil
|
|
}).AnyTimes()
|
|
cm.EXPECT().WithTools(gomock.Any()).Return(cm, nil).AnyTimes()
|
|
|
|
var capturedToolName string
|
|
var capturedArgs string
|
|
middlewareInvoked := false
|
|
|
|
agent, err := NewChatModelAgent(ctx, &ChatModelAgentConfig{
|
|
Name: "TestAgent",
|
|
Description: "Test WrapToolCall middleware",
|
|
Instruction: "You are a helpful assistant.",
|
|
Model: cm,
|
|
ToolsConfig: ToolsConfig{
|
|
ToolsNodeConfig: compose.ToolsNodeConfig{
|
|
Tools: []tool.BaseTool{fakeTool},
|
|
},
|
|
},
|
|
Middlewares: []AgentMiddleware{
|
|
{
|
|
WrapToolCall: compose.ToolMiddleware{
|
|
Invokable: func(next compose.InvokableToolEndpoint) compose.InvokableToolEndpoint {
|
|
return func(ctx context.Context, input *compose.ToolInput) (*compose.ToolOutput, error) {
|
|
middlewareInvoked = true
|
|
capturedToolName = input.Name
|
|
capturedArgs = input.Arguments
|
|
return next(ctx, input)
|
|
}
|
|
},
|
|
},
|
|
},
|
|
},
|
|
})
|
|
assert.NoError(t, err)
|
|
|
|
input := &AgentInput{
|
|
Messages: []Message{schema.UserMessage("use tool")},
|
|
}
|
|
iterator := agent.Run(ctx, input)
|
|
|
|
var events []*AgentEvent
|
|
for {
|
|
event, ok := iterator.Next()
|
|
if !ok {
|
|
break
|
|
}
|
|
events = append(events, event)
|
|
}
|
|
|
|
assert.True(t, middlewareInvoked, "WrapToolCall middleware should have been invoked")
|
|
assert.Equal(t, info.Name, capturedToolName)
|
|
assert.Equal(t, `{"name":"test"}`, capturedArgs)
|
|
assert.Equal(t, 3, len(events))
|
|
assert.Equal(t, schema.Assistant, events[0].Output.MessageOutput.Role)
|
|
assert.Equal(t, schema.Tool, events[1].Output.MessageOutput.Role)
|
|
assert.Equal(t, schema.Assistant, events[2].Output.MessageOutput.Role)
|
|
assert.Equal(t, "done", events[2].Output.MessageOutput.Message.Content)
|
|
})
|
|
|
|
t.Run("WrapToolCall_MultipleMiddlewares_Order", func(t *testing.T) {
|
|
ctx := context.Background()
|
|
|
|
ctrl := gomock.NewController(t)
|
|
cm := mockModel.NewMockToolCallingChatModel(ctrl)
|
|
|
|
fakeTool := &fakeToolForTest{tarCount: 1}
|
|
info, err := fakeTool.Info(ctx)
|
|
assert.NoError(t, err)
|
|
|
|
generateCount := 0
|
|
cm.EXPECT().Generate(gomock.Any(), gomock.Any(), gomock.Any()).
|
|
DoAndReturn(func(ctx context.Context, msgs []*schema.Message, opts ...model.Option) (*schema.Message, error) {
|
|
generateCount++
|
|
if generateCount == 1 {
|
|
return schema.AssistantMessage("calling tool", []schema.ToolCall{
|
|
{ID: "tc1", Function: schema.FunctionCall{Name: info.Name, Arguments: `{}`}},
|
|
}), nil
|
|
}
|
|
return schema.AssistantMessage("done", nil), nil
|
|
}).AnyTimes()
|
|
cm.EXPECT().WithTools(gomock.Any()).Return(cm, nil).AnyTimes()
|
|
|
|
var order []string
|
|
|
|
agent, err := NewChatModelAgent(ctx, &ChatModelAgentConfig{
|
|
Name: "TestAgent",
|
|
Description: "Test WrapToolCall middleware order",
|
|
Instruction: "You are a helpful assistant.",
|
|
Model: cm,
|
|
ToolsConfig: ToolsConfig{
|
|
ToolsNodeConfig: compose.ToolsNodeConfig{
|
|
Tools: []tool.BaseTool{fakeTool},
|
|
},
|
|
},
|
|
Middlewares: []AgentMiddleware{
|
|
{
|
|
WrapToolCall: compose.ToolMiddleware{
|
|
Invokable: func(next compose.InvokableToolEndpoint) compose.InvokableToolEndpoint {
|
|
return func(ctx context.Context, input *compose.ToolInput) (*compose.ToolOutput, error) {
|
|
order = append(order, "mw1_before")
|
|
out, err := next(ctx, input)
|
|
order = append(order, "mw1_after")
|
|
return out, err
|
|
}
|
|
},
|
|
},
|
|
},
|
|
{
|
|
WrapToolCall: compose.ToolMiddleware{
|
|
Invokable: func(next compose.InvokableToolEndpoint) compose.InvokableToolEndpoint {
|
|
return func(ctx context.Context, input *compose.ToolInput) (*compose.ToolOutput, error) {
|
|
order = append(order, "mw2_before")
|
|
out, err := next(ctx, input)
|
|
order = append(order, "mw2_after")
|
|
return out, err
|
|
}
|
|
},
|
|
},
|
|
},
|
|
},
|
|
})
|
|
assert.NoError(t, err)
|
|
|
|
input := &AgentInput{
|
|
Messages: []Message{schema.UserMessage("use tool")},
|
|
}
|
|
iterator := agent.Run(ctx, input)
|
|
|
|
for {
|
|
_, ok := iterator.Next()
|
|
if !ok {
|
|
break
|
|
}
|
|
}
|
|
|
|
// First registered middleware is outermost
|
|
assert.Equal(t, []string{"mw1_before", "mw2_before", "mw2_after", "mw1_after"}, order)
|
|
})
|
|
|
|
t.Run("WrapToolCall_ModifyResult", func(t *testing.T) {
|
|
ctx := context.Background()
|
|
|
|
ctrl := gomock.NewController(t)
|
|
cm := mockModel.NewMockToolCallingChatModel(ctrl)
|
|
|
|
fakeTool := &fakeToolForTest{tarCount: 1}
|
|
info, err := fakeTool.Info(ctx)
|
|
assert.NoError(t, err)
|
|
|
|
generateCount := 0
|
|
cm.EXPECT().Generate(gomock.Any(), gomock.Any(), gomock.Any()).
|
|
DoAndReturn(func(ctx context.Context, msgs []*schema.Message, opts ...model.Option) (*schema.Message, error) {
|
|
generateCount++
|
|
if generateCount == 1 {
|
|
return schema.AssistantMessage("calling tool", []schema.ToolCall{
|
|
{ID: "tc1", Function: schema.FunctionCall{Name: info.Name, Arguments: `{"name":"test"}`}},
|
|
}), nil
|
|
}
|
|
return schema.AssistantMessage("done", nil), nil
|
|
}).AnyTimes()
|
|
cm.EXPECT().WithTools(gomock.Any()).Return(cm, nil).AnyTimes()
|
|
|
|
agent, err := NewChatModelAgent(ctx, &ChatModelAgentConfig{
|
|
Name: "TestAgent",
|
|
Description: "Test WrapToolCall modify result",
|
|
Instruction: "You are a helpful assistant.",
|
|
Model: cm,
|
|
ToolsConfig: ToolsConfig{
|
|
ToolsNodeConfig: compose.ToolsNodeConfig{
|
|
Tools: []tool.BaseTool{fakeTool},
|
|
},
|
|
},
|
|
Middlewares: []AgentMiddleware{
|
|
{
|
|
WrapToolCall: compose.ToolMiddleware{
|
|
Invokable: func(next compose.InvokableToolEndpoint) compose.InvokableToolEndpoint {
|
|
return func(ctx context.Context, input *compose.ToolInput) (*compose.ToolOutput, error) {
|
|
out, err := next(ctx, input)
|
|
if err != nil {
|
|
return out, err
|
|
}
|
|
out.Result = "modified: " + out.Result
|
|
return out, nil
|
|
}
|
|
},
|
|
},
|
|
},
|
|
},
|
|
})
|
|
assert.NoError(t, err)
|
|
|
|
input := &AgentInput{
|
|
Messages: []Message{schema.UserMessage("use tool")},
|
|
}
|
|
iterator := agent.Run(ctx, input)
|
|
|
|
var events []*AgentEvent
|
|
for {
|
|
event, ok := iterator.Next()
|
|
if !ok {
|
|
break
|
|
}
|
|
events = append(events, event)
|
|
}
|
|
|
|
assert.Equal(t, 3, len(events))
|
|
// Tool result event should have the modified content
|
|
toolEvent := events[1]
|
|
assert.Equal(t, schema.Tool, toolEvent.Output.MessageOutput.Role)
|
|
assert.Contains(t, toolEvent.Output.MessageOutput.Message.Content, "modified:")
|
|
})
|
|
|
|
t.Run("WrapToolCall_StreamableMiddleware", func(t *testing.T) {
|
|
ctx := context.Background()
|
|
|
|
ctrl := gomock.NewController(t)
|
|
cm := mockModel.NewMockToolCallingChatModel(ctrl)
|
|
|
|
streamTool := &fakeStreamToolForTest{tarCount: 1}
|
|
info, err := streamTool.Info(ctx)
|
|
assert.NoError(t, err)
|
|
|
|
generateCount := 0
|
|
cm.EXPECT().Generate(gomock.Any(), gomock.Any(), gomock.Any()).
|
|
DoAndReturn(func(ctx context.Context, msgs []*schema.Message, opts ...model.Option) (*schema.Message, error) {
|
|
generateCount++
|
|
if generateCount == 1 {
|
|
return schema.AssistantMessage("calling tool", []schema.ToolCall{
|
|
{ID: "tc1", Function: schema.FunctionCall{Name: info.Name, Arguments: `{"name":"test"}`}},
|
|
}), nil
|
|
}
|
|
return schema.AssistantMessage("done", nil), nil
|
|
}).AnyTimes()
|
|
cm.EXPECT().WithTools(gomock.Any()).Return(cm, nil).AnyTimes()
|
|
|
|
middlewareInvoked := false
|
|
|
|
agent, err := NewChatModelAgent(ctx, &ChatModelAgentConfig{
|
|
Name: "TestAgent",
|
|
Description: "Test Streamable WrapToolCall middleware",
|
|
Instruction: "You are a helpful assistant.",
|
|
Model: cm,
|
|
ToolsConfig: ToolsConfig{
|
|
ToolsNodeConfig: compose.ToolsNodeConfig{
|
|
Tools: []tool.BaseTool{streamTool},
|
|
},
|
|
},
|
|
Middlewares: []AgentMiddleware{
|
|
{
|
|
WrapToolCall: compose.ToolMiddleware{
|
|
Streamable: func(next compose.StreamableToolEndpoint) compose.StreamableToolEndpoint {
|
|
return func(ctx context.Context, input *compose.ToolInput) (*compose.StreamToolOutput, error) {
|
|
middlewareInvoked = true
|
|
return next(ctx, input)
|
|
}
|
|
},
|
|
},
|
|
},
|
|
},
|
|
})
|
|
assert.NoError(t, err)
|
|
|
|
input := &AgentInput{
|
|
Messages: []Message{schema.UserMessage("use tool")},
|
|
}
|
|
iterator := agent.Run(ctx, input)
|
|
|
|
var events []*AgentEvent
|
|
for {
|
|
event, ok := iterator.Next()
|
|
if !ok {
|
|
break
|
|
}
|
|
events = append(events, event)
|
|
}
|
|
|
|
assert.True(t, middlewareInvoked, "Streamable WrapToolCall middleware should have been invoked")
|
|
assert.Equal(t, 3, len(events))
|
|
assert.Equal(t, schema.Tool, events[1].Output.MessageOutput.Role)
|
|
})
|
|
|
|
t.Run("WrapToolCall_EnhancedInvokableMiddleware", func(t *testing.T) {
|
|
ctx := context.Background()
|
|
|
|
ctrl := gomock.NewController(t)
|
|
cm := mockModel.NewMockToolCallingChatModel(ctrl)
|
|
|
|
enhTool := &fakeEnhancedInvokableToolForTest{}
|
|
info, err := enhTool.Info(ctx)
|
|
assert.NoError(t, err)
|
|
|
|
generateCount := 0
|
|
cm.EXPECT().Generate(gomock.Any(), gomock.Any(), gomock.Any()).
|
|
DoAndReturn(func(ctx context.Context, msgs []*schema.Message, opts ...model.Option) (*schema.Message, error) {
|
|
generateCount++
|
|
if generateCount == 1 {
|
|
return schema.AssistantMessage("calling tool", []schema.ToolCall{
|
|
{ID: "tc1", Function: schema.FunctionCall{Name: info.Name, Arguments: `{}`}},
|
|
}), nil
|
|
}
|
|
return schema.AssistantMessage("done", nil), nil
|
|
}).AnyTimes()
|
|
cm.EXPECT().WithTools(gomock.Any()).Return(cm, nil).AnyTimes()
|
|
|
|
middlewareInvoked := false
|
|
var capturedToolName string
|
|
|
|
agent, err := NewChatModelAgent(ctx, &ChatModelAgentConfig{
|
|
Name: "TestAgent",
|
|
Description: "Test EnhancedInvokable WrapToolCall middleware",
|
|
Instruction: "You are a helpful assistant.",
|
|
Model: cm,
|
|
ToolsConfig: ToolsConfig{
|
|
ToolsNodeConfig: compose.ToolsNodeConfig{
|
|
Tools: []tool.BaseTool{enhTool},
|
|
},
|
|
},
|
|
Middlewares: []AgentMiddleware{
|
|
{
|
|
WrapToolCall: compose.ToolMiddleware{
|
|
EnhancedInvokable: func(next compose.EnhancedInvokableToolEndpoint) compose.EnhancedInvokableToolEndpoint {
|
|
return func(ctx context.Context, input *compose.ToolInput) (*compose.EnhancedInvokableToolOutput, error) {
|
|
middlewareInvoked = true
|
|
capturedToolName = input.Name
|
|
return next(ctx, input)
|
|
}
|
|
},
|
|
},
|
|
},
|
|
},
|
|
})
|
|
assert.NoError(t, err)
|
|
|
|
input := &AgentInput{
|
|
Messages: []Message{schema.UserMessage("use tool")},
|
|
}
|
|
iterator := agent.Run(ctx, input)
|
|
|
|
var events []*AgentEvent
|
|
for {
|
|
event, ok := iterator.Next()
|
|
if !ok {
|
|
break
|
|
}
|
|
events = append(events, event)
|
|
}
|
|
|
|
assert.True(t, middlewareInvoked, "EnhancedInvokable WrapToolCall middleware should have been invoked")
|
|
assert.Equal(t, info.Name, capturedToolName)
|
|
assert.Equal(t, 3, len(events))
|
|
assert.Equal(t, schema.Tool, events[1].Output.MessageOutput.Role)
|
|
})
|
|
|
|
t.Run("WrapToolCall_EnhancedStreamableMiddleware", func(t *testing.T) {
|
|
ctx := context.Background()
|
|
|
|
ctrl := gomock.NewController(t)
|
|
cm := mockModel.NewMockToolCallingChatModel(ctrl)
|
|
|
|
enhStreamTool := &fakeEnhancedStreamableToolForTest{}
|
|
info, err := enhStreamTool.Info(ctx)
|
|
assert.NoError(t, err)
|
|
|
|
generateCount := 0
|
|
cm.EXPECT().Generate(gomock.Any(), gomock.Any(), gomock.Any()).
|
|
DoAndReturn(func(ctx context.Context, msgs []*schema.Message, opts ...model.Option) (*schema.Message, error) {
|
|
generateCount++
|
|
if generateCount == 1 {
|
|
return schema.AssistantMessage("calling tool", []schema.ToolCall{
|
|
{ID: "tc1", Function: schema.FunctionCall{Name: info.Name, Arguments: `{}`}},
|
|
}), nil
|
|
}
|
|
return schema.AssistantMessage("done", nil), nil
|
|
}).AnyTimes()
|
|
cm.EXPECT().WithTools(gomock.Any()).Return(cm, nil).AnyTimes()
|
|
|
|
middlewareInvoked := false
|
|
|
|
agent, err := NewChatModelAgent(ctx, &ChatModelAgentConfig{
|
|
Name: "TestAgent",
|
|
Description: "Test EnhancedStreamable WrapToolCall middleware",
|
|
Instruction: "You are a helpful assistant.",
|
|
Model: cm,
|
|
ToolsConfig: ToolsConfig{
|
|
ToolsNodeConfig: compose.ToolsNodeConfig{
|
|
Tools: []tool.BaseTool{enhStreamTool},
|
|
},
|
|
},
|
|
Middlewares: []AgentMiddleware{
|
|
{
|
|
WrapToolCall: compose.ToolMiddleware{
|
|
EnhancedStreamable: func(next compose.EnhancedStreamableToolEndpoint) compose.EnhancedStreamableToolEndpoint {
|
|
return func(ctx context.Context, input *compose.ToolInput) (*compose.EnhancedStreamableToolOutput, error) {
|
|
middlewareInvoked = true
|
|
return next(ctx, input)
|
|
}
|
|
},
|
|
},
|
|
},
|
|
},
|
|
})
|
|
assert.NoError(t, err)
|
|
|
|
input := &AgentInput{
|
|
Messages: []Message{schema.UserMessage("use tool")},
|
|
}
|
|
iterator := agent.Run(ctx, input)
|
|
|
|
var events []*AgentEvent
|
|
for {
|
|
event, ok := iterator.Next()
|
|
if !ok {
|
|
break
|
|
}
|
|
events = append(events, event)
|
|
}
|
|
|
|
assert.True(t, middlewareInvoked, "EnhancedStreamable WrapToolCall middleware should have been invoked")
|
|
assert.Equal(t, 3, len(events))
|
|
assert.Equal(t, schema.Tool, events[1].Output.MessageOutput.Role)
|
|
})
|
|
}
|
|
|
|
// TestExitTool tests the Exit tool functionality
|
|
func TestExitTool(t *testing.T) {
|
|
ctx := context.Background()
|
|
|
|
// Create a mock controller
|
|
ctrl := gomock.NewController(t)
|
|
defer ctrl.Finish()
|
|
|
|
// Create a mock chat model
|
|
cm := mockModel.NewMockToolCallingChatModel(ctrl)
|
|
|
|
// Set up expectations for the mock model
|
|
// First call: model generates a message with Exit tool call
|
|
cm.EXPECT().Generate(gomock.Any(), gomock.Any(), gomock.Any()).
|
|
Return(schema.AssistantMessage("I'll exit with a final result",
|
|
[]schema.ToolCall{
|
|
{
|
|
ID: "tool-call-1",
|
|
Function: schema.FunctionCall{
|
|
Name: "exit",
|
|
Arguments: `{"final_result": "This is the final result"}`},
|
|
},
|
|
}), nil).
|
|
Times(1)
|
|
|
|
// Model should implement WithTools
|
|
cm.EXPECT().WithTools(gomock.Any()).Return(cm, nil).AnyTimes()
|
|
|
|
// Create an agent with the Exit tool
|
|
agent, err := NewChatModelAgent(ctx, &ChatModelAgentConfig{
|
|
Name: "TestAgent",
|
|
Description: "Test agent with Exit tool",
|
|
Instruction: "You are a helpful assistant.",
|
|
Model: cm,
|
|
Exit: &ExitTool{},
|
|
})
|
|
assert.NoError(t, err)
|
|
assert.NotNil(t, agent)
|
|
|
|
// Run the agent
|
|
input := &AgentInput{
|
|
Messages: []Message{
|
|
schema.UserMessage("Please exit with a final result"),
|
|
},
|
|
}
|
|
iterator := agent.Run(ctx, input)
|
|
assert.NotNil(t, iterator)
|
|
|
|
// First event: model output with tool call
|
|
event1, ok := iterator.Next()
|
|
assert.True(t, ok)
|
|
assert.NotNil(t, event1)
|
|
assert.Nil(t, event1.Err)
|
|
assert.NotNil(t, event1.Output)
|
|
assert.NotNil(t, event1.Output.MessageOutput)
|
|
assert.Equal(t, schema.Assistant, event1.Output.MessageOutput.Role)
|
|
|
|
// Second event: tool output (Exit)
|
|
event2, ok := iterator.Next()
|
|
assert.True(t, ok)
|
|
assert.NotNil(t, event2)
|
|
assert.Nil(t, event2.Err)
|
|
assert.NotNil(t, event2.Output)
|
|
assert.NotNil(t, event2.Output.MessageOutput)
|
|
assert.Equal(t, schema.Tool, event2.Output.MessageOutput.Role)
|
|
|
|
// Verify the action is Exit
|
|
assert.NotNil(t, event2.Action)
|
|
assert.True(t, event2.Action.Exit)
|
|
|
|
// Verify the final result
|
|
assert.Equal(t, "This is the final result", event2.Output.MessageOutput.Message.Content)
|
|
|
|
// No more events
|
|
_, ok = iterator.Next()
|
|
assert.False(t, ok)
|
|
}
|
|
|
|
func TestParallelReturnDirectlyToolCall(t *testing.T) {
|
|
ctx := context.Background()
|
|
// Create a mock controller
|
|
ctrl := gomock.NewController(t)
|
|
defer ctrl.Finish()
|
|
|
|
// Create a mock chat model
|
|
cm := mockModel.NewMockToolCallingChatModel(ctrl)
|
|
|
|
// Set up expectations for the mock model
|
|
// First call: model generates a message with Exit tool call
|
|
cm.EXPECT().Generate(gomock.Any(), gomock.Any(), gomock.Any()).
|
|
Return(schema.AssistantMessage("I'll exit with a final result",
|
|
[]schema.ToolCall{
|
|
{
|
|
ID: "tool-call-1",
|
|
Function: schema.FunctionCall{Name: "tool1"},
|
|
},
|
|
{
|
|
ID: "tool-call-2",
|
|
Function: schema.FunctionCall{Name: "tool2"},
|
|
},
|
|
{
|
|
ID: "tool-call-3",
|
|
Function: schema.FunctionCall{Name: "tool3"},
|
|
},
|
|
}), nil).
|
|
Times(1)
|
|
|
|
// Model should implement WithTools
|
|
cm.EXPECT().WithTools(gomock.Any()).Return(cm, nil).AnyTimes()
|
|
|
|
// Create an agent with the Exit tool
|
|
agent, err := NewChatModelAgent(ctx, &ChatModelAgentConfig{
|
|
Name: "TestAgent",
|
|
Description: "Test agent with Exit tool",
|
|
Instruction: "You are a helpful assistant.",
|
|
Model: cm,
|
|
ToolsConfig: ToolsConfig{
|
|
ToolsNodeConfig: compose.ToolsNodeConfig{
|
|
Tools: []tool.BaseTool{
|
|
&myTool{name: "tool1", desc: "tool1", waitTime: time.Millisecond},
|
|
&myTool{name: "tool2", desc: "tool2", waitTime: 10 * time.Millisecond},
|
|
&myTool{name: "tool3", desc: "tool3", waitTime: 100 * time.Millisecond},
|
|
},
|
|
},
|
|
ReturnDirectly: map[string]bool{
|
|
"tool1": true,
|
|
},
|
|
},
|
|
})
|
|
assert.NoError(t, err)
|
|
assert.NotNil(t, agent)
|
|
|
|
r := NewRunner(ctx, RunnerConfig{
|
|
Agent: agent,
|
|
})
|
|
iter := r.Query(ctx, "")
|
|
times := 0
|
|
for {
|
|
e, ok := iter.Next()
|
|
if !ok {
|
|
assert.Equal(t, 4, times)
|
|
break
|
|
}
|
|
if times == 3 {
|
|
assert.Equal(t, "tool1", e.Output.MessageOutput.Message.ToolName)
|
|
}
|
|
times++
|
|
}
|
|
}
|
|
|
|
func TestConcurrentSameToolSendToolGenActionUsesToolCallID(t *testing.T) {
|
|
ctx := context.Background()
|
|
ctrl := gomock.NewController(t)
|
|
defer ctrl.Finish()
|
|
|
|
cm := mockModel.NewMockToolCallingChatModel(ctrl)
|
|
|
|
cm.EXPECT().WithTools(gomock.Any()).Return(cm, nil).AnyTimes()
|
|
|
|
cm.EXPECT().Generate(gomock.Any(), gomock.Any(), gomock.Any()).
|
|
Return(schema.AssistantMessage("tools", []schema.ToolCall{
|
|
{ID: "id1", Function: schema.FunctionCall{Name: "action_tool", Arguments: "A"}},
|
|
{ID: "id2", Function: schema.FunctionCall{Name: "action_tool", Arguments: "B"}},
|
|
}), nil).
|
|
Times(1)
|
|
|
|
cm.EXPECT().Generate(gomock.Any(), gomock.Any(), gomock.Any()).
|
|
Return(schema.AssistantMessage("done", nil), nil).
|
|
Times(1)
|
|
|
|
agent, err := NewChatModelAgent(ctx, &ChatModelAgentConfig{
|
|
Name: "TestAgent",
|
|
Description: "Agent with action tool",
|
|
Instruction: "",
|
|
Model: cm,
|
|
ToolsConfig: ToolsConfig{ToolsNodeConfig: compose.ToolsNodeConfig{Tools: []tool.BaseTool{actionTool{}}}},
|
|
})
|
|
assert.NoError(t, err)
|
|
|
|
iter := agent.Run(ctx, &AgentInput{Messages: []Message{schema.UserMessage("go")}})
|
|
seen := map[string]bool{}
|
|
for {
|
|
e, ok := iter.Next()
|
|
if !ok {
|
|
break
|
|
}
|
|
if e.Output != nil && e.Output.MessageOutput != nil && e.Output.MessageOutput.Message != nil && e.Output.MessageOutput.Message.Role == schema.Tool {
|
|
if e.Action != nil && e.Action.CustomizedAction != nil {
|
|
if s, ok := e.Action.CustomizedAction.(string); ok {
|
|
seen[s] = true
|
|
}
|
|
}
|
|
}
|
|
}
|
|
assert.True(t, seen["A"])
|
|
assert.True(t, seen["B"])
|
|
}
|
|
|
|
type myTool struct {
|
|
name string
|
|
desc string
|
|
waitTime time.Duration
|
|
}
|
|
|
|
func (m *myTool) Info(ctx context.Context) (*schema.ToolInfo, error) {
|
|
return &schema.ToolInfo{
|
|
Name: m.name,
|
|
Desc: m.desc,
|
|
}, nil
|
|
}
|
|
|
|
func (m *myTool) InvokableRun(ctx context.Context, argumentsInJSON string, opts ...tool.Option) (string, error) {
|
|
time.Sleep(m.waitTime)
|
|
return "success", nil
|
|
}
|
|
|
|
type actionTool struct{}
|
|
|
|
func (a actionTool) Info(ctx context.Context) (*schema.ToolInfo, error) {
|
|
return &schema.ToolInfo{Name: "action_tool", Desc: "action tool"}, nil
|
|
}
|
|
|
|
func (a actionTool) InvokableRun(ctx context.Context, argumentsInJSON string, _ ...tool.Option) (string, error) {
|
|
_ = SendToolGenAction(ctx, "action_tool", &AgentAction{CustomizedAction: argumentsInJSON})
|
|
return "ok", nil
|
|
}
|
|
|
|
type streamActionTool struct{}
|
|
|
|
func (s streamActionTool) Info(ctx context.Context) (*schema.ToolInfo, error) {
|
|
return &schema.ToolInfo{Name: "action_tool_stream", Desc: "action stream tool"}, nil
|
|
}
|
|
|
|
func (s streamActionTool) StreamableRun(ctx context.Context, argumentsInJSON string, _ ...tool.Option) (*schema.StreamReader[string], error) {
|
|
_ = SendToolGenAction(ctx, "action_tool_stream", &AgentAction{CustomizedAction: argumentsInJSON})
|
|
sr, sw := schema.Pipe[string](1)
|
|
go func() {
|
|
defer sw.Close()
|
|
_ = sw.Send("o", nil)
|
|
_ = sw.Send("k", nil)
|
|
}()
|
|
return sr, nil
|
|
}
|
|
|
|
type legacyStreamActionTool struct{}
|
|
|
|
func (s legacyStreamActionTool) Info(ctx context.Context) (*schema.ToolInfo, error) {
|
|
return &schema.ToolInfo{Name: "legacy_action_tool_stream", Desc: "legacy action stream tool"}, nil
|
|
}
|
|
|
|
func (s legacyStreamActionTool) StreamableRun(ctx context.Context, argumentsInJSON string, _ ...tool.Option) (*schema.StreamReader[string], error) {
|
|
_ = compose.ProcessState(ctx, func(ctx context.Context, st *State) error {
|
|
st.setToolGenAction("legacy_action_tool_stream", &AgentAction{CustomizedAction: argumentsInJSON})
|
|
return nil
|
|
})
|
|
sr, sw := schema.Pipe[string](1)
|
|
go func() {
|
|
defer sw.Close()
|
|
_ = sw.Send("o", nil)
|
|
_ = sw.Send("k", nil)
|
|
}()
|
|
return sr, nil
|
|
}
|
|
|
|
// TestChatModelAgentOutputKey tests the outputKey configuration and setOutputToSession function
|
|
func TestChatModelAgentOutputKey(t *testing.T) {
|
|
// Test outputKey configuration - stores output in session
|
|
t.Run("OutputKeyStoresInSession", func(t *testing.T) {
|
|
for i := 0; i < 1000; i++ {
|
|
|
|
}
|
|
ctx := context.Background()
|
|
|
|
// Create a mock chat model
|
|
ctrl := gomock.NewController(t)
|
|
cm := mockModel.NewMockToolCallingChatModel(ctrl)
|
|
|
|
// Set up expectations for the mock model
|
|
cm.EXPECT().Generate(gomock.Any(), gomock.Any(), gomock.Any()).
|
|
Return(schema.AssistantMessage("Hello, I am an AI assistant.", nil), nil).
|
|
Times(1)
|
|
|
|
// Create a ChatModelAgent with outputKey configured
|
|
agent, err := NewChatModelAgent(ctx, &ChatModelAgentConfig{
|
|
Name: "TestAgent",
|
|
Description: "Test agent for unit testing",
|
|
Instruction: "You are a helpful assistant.",
|
|
Model: cm,
|
|
OutputKey: "agent_output", // This should store output in session
|
|
})
|
|
assert.NoError(t, err)
|
|
assert.NotNil(t, agent)
|
|
|
|
// Initialize a run context to enable session storage
|
|
input := &AgentInput{
|
|
Messages: []Message{
|
|
schema.UserMessage("Hello, who are you?"),
|
|
},
|
|
}
|
|
ctx, runCtx := initRunCtx(ctx, "TestAgent", input)
|
|
assert.NotNil(t, runCtx)
|
|
assert.NotNil(t, runCtx.Session)
|
|
|
|
// Run the agent
|
|
iterator := agent.Run(ctx, input)
|
|
assert.NotNil(t, iterator)
|
|
|
|
// Get the event from the iterator
|
|
event, ok := iterator.Next()
|
|
assert.True(t, ok)
|
|
assert.NotNil(t, event)
|
|
assert.Nil(t, event.Err)
|
|
|
|
// Verify the message content
|
|
msg := event.Output.MessageOutput.Message
|
|
assert.Equal(t, "Hello, I am an AI assistant.", msg.Content)
|
|
|
|
// Verify that the output was stored in the session
|
|
time.AfterFunc(100*time.Millisecond, func() {
|
|
sessionValues := GetSessionValues(ctx)
|
|
assert.Contains(t, sessionValues, "agent_output")
|
|
assert.Equal(t, "Hello, I am an AI assistant.", sessionValues["agent_output"])
|
|
})
|
|
|
|
// No more events
|
|
_, ok = iterator.Next()
|
|
assert.False(t, ok)
|
|
})
|
|
|
|
// Test outputKey configuration with streaming - stores concatenated output in session
|
|
t.Run("OutputKeyWithStreamingStoresInSession", func(t *testing.T) {
|
|
ctx := context.Background()
|
|
|
|
// Create a mock chat model
|
|
ctrl := gomock.NewController(t)
|
|
cm := mockModel.NewMockToolCallingChatModel(ctrl)
|
|
|
|
// Create a stream reader for the mock response
|
|
sr := schema.StreamReaderFromArray([]*schema.Message{
|
|
schema.AssistantMessage("Hello", nil),
|
|
schema.AssistantMessage(", I am", nil),
|
|
schema.AssistantMessage(" an AI assistant.", nil),
|
|
})
|
|
|
|
// Set up expectations for the mock model
|
|
cm.EXPECT().Stream(gomock.Any(), gomock.Any(), gomock.Any()).
|
|
Return(sr, nil).
|
|
Times(1)
|
|
|
|
// Create a ChatModelAgent with outputKey configured
|
|
agent, err := NewChatModelAgent(ctx, &ChatModelAgentConfig{
|
|
Name: "TestAgent",
|
|
Description: "Test agent for unit testing",
|
|
Instruction: "You are a helpful assistant.",
|
|
Model: cm,
|
|
OutputKey: "agent_output", // This should store concatenated stream in session
|
|
})
|
|
assert.NoError(t, err)
|
|
assert.NotNil(t, agent)
|
|
|
|
// Initialize a run context to enable session storage
|
|
input := &AgentInput{
|
|
Messages: []Message{schema.UserMessage("Hello, who are you?")},
|
|
EnableStreaming: true,
|
|
}
|
|
ctx, runCtx := initRunCtx(ctx, "TestAgent", input)
|
|
assert.NotNil(t, runCtx)
|
|
assert.NotNil(t, runCtx.Session)
|
|
|
|
// Run the agent
|
|
iterator := agent.Run(ctx, input)
|
|
assert.NotNil(t, iterator)
|
|
|
|
// Get the event from the iterator
|
|
event, ok := iterator.Next()
|
|
assert.True(t, ok)
|
|
assert.NotNil(t, event)
|
|
assert.Nil(t, event.Err)
|
|
assert.True(t, event.Output.MessageOutput.IsStreaming)
|
|
|
|
time.AfterFunc(100*time.Millisecond, func() {
|
|
// Verify that the concatenated output was stored in the session
|
|
sessionValues := GetSessionValues(ctx)
|
|
assert.Contains(t, sessionValues, "agent_output")
|
|
assert.Equal(t, "Hello, I am an AI assistant.", sessionValues["agent_output"])
|
|
})
|
|
|
|
// No more events
|
|
_, ok = iterator.Next()
|
|
assert.False(t, ok)
|
|
})
|
|
|
|
// Test setOutputToSession function directly - regular message
|
|
t.Run("SetOutputToSessionRegularMessage", func(t *testing.T) {
|
|
ctx := context.Background()
|
|
|
|
// Initialize a run context to enable session storage
|
|
input := &AgentInput{
|
|
Messages: []Message{schema.UserMessage("test")},
|
|
}
|
|
ctx, runCtx := initRunCtx(ctx, "TestAgent", input)
|
|
assert.NotNil(t, runCtx)
|
|
assert.NotNil(t, runCtx.Session)
|
|
|
|
// Test with a regular message
|
|
msg := schema.AssistantMessage("Test response", nil)
|
|
err := setOutputToSession(ctx, msg, nil, "test_output")
|
|
assert.NoError(t, err)
|
|
|
|
// Verify the message content was stored
|
|
sessionValues := GetSessionValues(ctx)
|
|
assert.Contains(t, sessionValues, "test_output")
|
|
assert.Equal(t, "Test response", sessionValues["test_output"])
|
|
})
|
|
|
|
// Test setOutputToSession function directly - streaming message
|
|
t.Run("SetOutputToSessionStreamingMessage", func(t *testing.T) {
|
|
ctx := context.Background()
|
|
|
|
// Initialize a run context to enable session storage
|
|
input := &AgentInput{
|
|
Messages: []Message{schema.UserMessage("test")},
|
|
}
|
|
ctx, runCtx := initRunCtx(ctx, "TestAgent", input)
|
|
assert.NotNil(t, runCtx)
|
|
assert.NotNil(t, runCtx.Session)
|
|
|
|
// Test with a streaming message
|
|
sr := schema.StreamReaderFromArray([]*schema.Message{
|
|
schema.AssistantMessage("Stream", nil),
|
|
schema.AssistantMessage(" response", nil),
|
|
schema.AssistantMessage(" content", nil),
|
|
})
|
|
err := setOutputToSession(ctx, nil, sr, "test_output")
|
|
assert.NoError(t, err)
|
|
|
|
// Verify the concatenated stream content was stored
|
|
sessionValues := GetSessionValues(ctx)
|
|
assert.Contains(t, sessionValues, "test_output")
|
|
assert.Equal(t, "Stream response content", sessionValues["test_output"])
|
|
})
|
|
|
|
// Test setOutputToSession function directly - error case
|
|
t.Run("SetOutputToSessionErrorCase", func(t *testing.T) {
|
|
ctx := context.Background()
|
|
|
|
// Initialize a run context to enable session storage
|
|
input := &AgentInput{
|
|
Messages: []Message{schema.UserMessage("test")},
|
|
}
|
|
ctx, runCtx := initRunCtx(ctx, "TestAgent", input)
|
|
assert.NotNil(t, runCtx)
|
|
assert.NotNil(t, runCtx.Session)
|
|
|
|
// Test with an invalid stream (simulate error)
|
|
// Create a stream that will fail when concatenated
|
|
sr := schema.StreamReaderFromArray([]*schema.Message{
|
|
schema.AssistantMessage("test", nil),
|
|
})
|
|
// Close the stream to simulate an error condition
|
|
sr.Close()
|
|
|
|
// This should return an error because the stream is closed
|
|
err := setOutputToSession(ctx, nil, sr, "test_output")
|
|
// Note: The actual behavior may vary depending on the stream implementation
|
|
// Some streams may not error when closed, so we'll accept either outcome
|
|
if err != nil {
|
|
// If there's an error, verify nothing was stored
|
|
sessionValues := GetSessionValues(ctx)
|
|
assert.NotContains(t, sessionValues, "test_output")
|
|
} else {
|
|
// If no error, verify the content was stored
|
|
sessionValues := GetSessionValues(ctx)
|
|
assert.Contains(t, sessionValues, "test_output")
|
|
assert.Equal(t, "test", sessionValues["test_output"])
|
|
}
|
|
})
|
|
|
|
// Test outputKey with React workflow (tools enabled)
|
|
t.Run("OutputKeyWithReactWorkflow", func(t *testing.T) {
|
|
ctx := context.Background()
|
|
|
|
// Create a mock chat model
|
|
ctrl := gomock.NewController(t)
|
|
cm := mockModel.NewMockToolCallingChatModel(ctrl)
|
|
|
|
// Create a simple tool for testing
|
|
fakeTool := &fakeToolForTest{
|
|
tarCount: 1,
|
|
}
|
|
|
|
// Set up expectations for the mock model - it will generate a final response
|
|
cm.EXPECT().Generate(gomock.Any(), gomock.Any(), gomock.Any()).
|
|
Return(schema.AssistantMessage("Final response from React workflow", nil), nil).
|
|
Times(1)
|
|
// Model should implement WithTools
|
|
cm.EXPECT().WithTools(gomock.Any()).Return(cm, nil).AnyTimes()
|
|
|
|
// Create a ChatModelAgent with outputKey and tools configured
|
|
agent, err := NewChatModelAgent(ctx, &ChatModelAgentConfig{
|
|
Name: "TestAgent",
|
|
Description: "Test agent with tools",
|
|
Instruction: "You are a helpful assistant.",
|
|
Model: cm,
|
|
OutputKey: "agent_output",
|
|
ToolsConfig: ToolsConfig{
|
|
ToolsNodeConfig: compose.ToolsNodeConfig{
|
|
Tools: []tool.BaseTool{fakeTool},
|
|
},
|
|
},
|
|
})
|
|
assert.NoError(t, err)
|
|
assert.NotNil(t, agent)
|
|
|
|
// Initialize a run context to enable session storage
|
|
input := &AgentInput{
|
|
Messages: []Message{schema.UserMessage("Use the tool")},
|
|
}
|
|
ctx, runCtx := initRunCtx(ctx, "TestAgent", input)
|
|
assert.NotNil(t, runCtx)
|
|
assert.NotNil(t, runCtx.Session)
|
|
|
|
// Run the agent
|
|
iterator := agent.Run(ctx, input)
|
|
assert.NotNil(t, iterator)
|
|
|
|
// Get the event from the iterator
|
|
event, ok := iterator.Next()
|
|
assert.True(t, ok)
|
|
assert.NotNil(t, event)
|
|
assert.Nil(t, event.Err)
|
|
|
|
// Verify the message content
|
|
msg := event.Output.MessageOutput.Message
|
|
assert.Equal(t, "Final response from React workflow", msg.Content)
|
|
|
|
// Verify that the output was stored in the session
|
|
time.AfterFunc(time.Millisecond*10, func() {
|
|
sessionValues := GetSessionValues(ctx)
|
|
assert.Contains(t, sessionValues, "agent_output")
|
|
assert.Equal(t, "Final response from React workflow", sessionValues["agent_output"])
|
|
})
|
|
|
|
// No more events
|
|
_, ok = iterator.Next()
|
|
assert.False(t, ok)
|
|
})
|
|
}
|
|
|
|
func TestConcurrentSameStreamToolSendToolGenActionUsesToolCallID(t *testing.T) {
|
|
ctx := context.Background()
|
|
ctrl := gomock.NewController(t)
|
|
defer ctrl.Finish()
|
|
|
|
cm := mockModel.NewMockToolCallingChatModel(ctrl)
|
|
|
|
cm.EXPECT().WithTools(gomock.Any()).Return(cm, nil).AnyTimes()
|
|
|
|
cm.EXPECT().Generate(gomock.Any(), gomock.Any(), gomock.Any()).
|
|
Return(schema.AssistantMessage("tools", []schema.ToolCall{
|
|
{ID: "sid1", Function: schema.FunctionCall{Name: "action_tool_stream", Arguments: "SA"}},
|
|
{ID: "sid2", Function: schema.FunctionCall{Name: "action_tool_stream", Arguments: "SB"}},
|
|
}), nil).
|
|
Times(1)
|
|
|
|
cm.EXPECT().Generate(gomock.Any(), gomock.Any(), gomock.Any()).
|
|
Return(schema.AssistantMessage("done", nil), nil).
|
|
Times(1)
|
|
|
|
agent, err := NewChatModelAgent(ctx, &ChatModelAgentConfig{
|
|
Name: "TestAgent",
|
|
Description: "Agent with stream action tool",
|
|
Instruction: "",
|
|
Model: cm,
|
|
ToolsConfig: ToolsConfig{ToolsNodeConfig: compose.ToolsNodeConfig{Tools: []tool.BaseTool{streamActionTool{}}}},
|
|
})
|
|
assert.NoError(t, err)
|
|
|
|
iter := agent.Run(ctx, &AgentInput{Messages: []Message{schema.UserMessage("go")}})
|
|
seen := map[string]bool{}
|
|
for {
|
|
e, ok := iter.Next()
|
|
if !ok {
|
|
break
|
|
}
|
|
if e.Output != nil && e.Output.MessageOutput != nil {
|
|
if e.Output.MessageOutput.IsStreaming {
|
|
if e.Action != nil && e.Action.CustomizedAction != nil {
|
|
if s, ok := e.Action.CustomizedAction.(string); ok {
|
|
seen[s] = true
|
|
}
|
|
}
|
|
}
|
|
}
|
|
}
|
|
assert.True(t, seen["SA"])
|
|
assert.True(t, seen["SB"])
|
|
}
|
|
|
|
func TestStreamToolLegacyNameKeyFallback(t *testing.T) {
|
|
ctx := context.Background()
|
|
ctrl := gomock.NewController(t)
|
|
defer ctrl.Finish()
|
|
|
|
cm := mockModel.NewMockToolCallingChatModel(ctrl)
|
|
cm.EXPECT().WithTools(gomock.Any()).Return(cm, nil).AnyTimes()
|
|
|
|
cm.EXPECT().Generate(gomock.Any(), gomock.Any(), gomock.Any()).
|
|
Return(schema.AssistantMessage("tools", []schema.ToolCall{
|
|
{ID: "lsid1", Function: schema.FunctionCall{Name: "legacy_action_tool_stream", Arguments: "LA"}},
|
|
}), nil).
|
|
Times(1)
|
|
|
|
cm.EXPECT().Generate(gomock.Any(), gomock.Any(), gomock.Any()).
|
|
Return(schema.AssistantMessage("done", nil), nil).
|
|
Times(1)
|
|
|
|
agent, err := NewChatModelAgent(ctx, &ChatModelAgentConfig{
|
|
Name: "TestAgent",
|
|
Description: "Agent with legacy stream action tool",
|
|
Instruction: "",
|
|
Model: cm,
|
|
ToolsConfig: ToolsConfig{ToolsNodeConfig: compose.ToolsNodeConfig{Tools: []tool.BaseTool{legacyStreamActionTool{}}}},
|
|
})
|
|
assert.NoError(t, err)
|
|
|
|
iter := agent.Run(ctx, &AgentInput{Messages: []Message{schema.UserMessage("go")}})
|
|
found := false
|
|
for {
|
|
e, ok := iter.Next()
|
|
if !ok {
|
|
break
|
|
}
|
|
if e.Output != nil && e.Output.MessageOutput != nil && e.Output.MessageOutput.IsStreaming {
|
|
if e.Action != nil && e.Action.CustomizedAction != nil {
|
|
if s, ok := e.Action.CustomizedAction.(string); ok {
|
|
found = (s == "LA")
|
|
}
|
|
}
|
|
}
|
|
}
|
|
assert.True(t, found)
|
|
}
|
|
|
|
func TestChatModelAgent_ToolResultMiddleware_EmitsFinalResult(t *testing.T) {
|
|
originalResult := "original_result"
|
|
modifiedResult := "modified_by_middleware"
|
|
|
|
resultModifyingMiddleware := compose.ToolMiddleware{
|
|
Invokable: func(next compose.InvokableToolEndpoint) compose.InvokableToolEndpoint {
|
|
return func(ctx context.Context, input *compose.ToolInput) (*compose.ToolOutput, error) {
|
|
output, err := next(ctx, input)
|
|
if err != nil {
|
|
return nil, err
|
|
}
|
|
output.Result = modifiedResult
|
|
return output, nil
|
|
}
|
|
},
|
|
Streamable: func(next compose.StreamableToolEndpoint) compose.StreamableToolEndpoint {
|
|
return func(ctx context.Context, input *compose.ToolInput) (*compose.StreamToolOutput, error) {
|
|
output, err := next(ctx, input)
|
|
if err != nil {
|
|
return nil, err
|
|
}
|
|
output.Result = schema.StreamReaderFromArray([]string{modifiedResult})
|
|
return output, nil
|
|
}
|
|
},
|
|
}
|
|
|
|
t.Run("Invoke", func(t *testing.T) {
|
|
ctx := context.Background()
|
|
testTool := &simpleToolForMiddlewareTest{name: "test_tool", result: originalResult}
|
|
|
|
ctrl := gomock.NewController(t)
|
|
cm := mockModel.NewMockToolCallingChatModel(ctrl)
|
|
|
|
info, err := testTool.Info(ctx)
|
|
assert.NoError(t, err)
|
|
|
|
cm.EXPECT().Generate(gomock.Any(), gomock.Any(), gomock.Any()).
|
|
Return(schema.AssistantMessage("",
|
|
[]schema.ToolCall{
|
|
{
|
|
ID: "tool-call-1",
|
|
Function: schema.FunctionCall{
|
|
Name: info.Name,
|
|
Arguments: `{"input": "test"}`,
|
|
},
|
|
},
|
|
}), nil).
|
|
Times(1)
|
|
cm.EXPECT().Generate(gomock.Any(), gomock.Any(), gomock.Any()).
|
|
Return(schema.AssistantMessage("final response", nil), nil).
|
|
Times(1)
|
|
cm.EXPECT().WithTools(gomock.Any()).Return(cm, nil).AnyTimes()
|
|
|
|
agent, err := NewChatModelAgent(ctx, &ChatModelAgentConfig{
|
|
Name: "test_agent",
|
|
Description: "test agent with middleware",
|
|
Model: cm,
|
|
ToolsConfig: ToolsConfig{
|
|
ToolsNodeConfig: compose.ToolsNodeConfig{
|
|
Tools: []tool.BaseTool{testTool},
|
|
ToolCallMiddlewares: []compose.ToolMiddleware{resultModifyingMiddleware},
|
|
},
|
|
},
|
|
})
|
|
assert.NoError(t, err)
|
|
|
|
r := NewRunner(ctx, RunnerConfig{Agent: agent, EnableStreaming: false, CheckPointStore: newBridgeStore()})
|
|
it := r.Run(ctx, []Message{schema.UserMessage("call the tool")})
|
|
|
|
var toolResultEvents []*AgentEvent
|
|
for {
|
|
ev, ok := it.Next()
|
|
if !ok {
|
|
break
|
|
}
|
|
if ev.Output != nil && ev.Output.MessageOutput != nil &&
|
|
ev.Output.MessageOutput.Message != nil &&
|
|
ev.Output.MessageOutput.Message.Role == schema.Tool {
|
|
toolResultEvents = append(toolResultEvents, ev)
|
|
}
|
|
}
|
|
|
|
assert.NotEmpty(t, toolResultEvents, "should have at least one tool result event")
|
|
for _, ev := range toolResultEvents {
|
|
assert.Equal(t, modifiedResult, ev.Output.MessageOutput.Message.Content,
|
|
"tool result event should contain the middleware-modified result, not the original")
|
|
assert.NotEqual(t, originalResult, ev.Output.MessageOutput.Message.Content,
|
|
"tool result event should NOT contain the original result")
|
|
}
|
|
})
|
|
|
|
t.Run("Stream", func(t *testing.T) {
|
|
ctx := context.Background()
|
|
testTool := &simpleToolForMiddlewareTest{name: "test_tool_stream", result: originalResult}
|
|
|
|
ctrl := gomock.NewController(t)
|
|
cm := mockModel.NewMockToolCallingChatModel(ctrl)
|
|
|
|
info, err := testTool.Info(ctx)
|
|
assert.NoError(t, err)
|
|
|
|
cm.EXPECT().Stream(gomock.Any(), gomock.Any(), gomock.Any()).
|
|
Return(schema.StreamReaderFromArray([]*schema.Message{
|
|
schema.AssistantMessage("", []schema.ToolCall{
|
|
{
|
|
ID: "tool-call-1",
|
|
Function: schema.FunctionCall{
|
|
Name: info.Name,
|
|
Arguments: `{"input": "test"}`,
|
|
},
|
|
},
|
|
}),
|
|
}), nil).
|
|
Times(1)
|
|
cm.EXPECT().Stream(gomock.Any(), gomock.Any(), gomock.Any()).
|
|
Return(schema.StreamReaderFromArray([]*schema.Message{
|
|
schema.AssistantMessage("final response", nil),
|
|
}), nil).
|
|
Times(1)
|
|
cm.EXPECT().WithTools(gomock.Any()).Return(cm, nil).AnyTimes()
|
|
|
|
agent, err := NewChatModelAgent(ctx, &ChatModelAgentConfig{
|
|
Name: "test_agent",
|
|
Description: "test agent with middleware",
|
|
Model: cm,
|
|
ToolsConfig: ToolsConfig{
|
|
ToolsNodeConfig: compose.ToolsNodeConfig{
|
|
Tools: []tool.BaseTool{testTool},
|
|
ToolCallMiddlewares: []compose.ToolMiddleware{resultModifyingMiddleware},
|
|
},
|
|
},
|
|
})
|
|
assert.NoError(t, err)
|
|
|
|
r := NewRunner(ctx, RunnerConfig{Agent: agent, EnableStreaming: true, CheckPointStore: newBridgeStore()})
|
|
it := r.Run(ctx, []Message{schema.UserMessage("call the tool")})
|
|
|
|
var toolResultContents []string
|
|
for {
|
|
ev, ok := it.Next()
|
|
if !ok {
|
|
break
|
|
}
|
|
if ev.Output != nil && ev.Output.MessageOutput != nil {
|
|
if ev.Output.MessageOutput.Message != nil &&
|
|
ev.Output.MessageOutput.Message.Role == schema.Tool {
|
|
toolResultContents = append(toolResultContents, ev.Output.MessageOutput.Message.Content)
|
|
}
|
|
if ev.Output.MessageOutput.IsStreaming &&
|
|
ev.Output.MessageOutput.MessageStream != nil &&
|
|
ev.Output.MessageOutput.Role == schema.Tool {
|
|
var msgs []*schema.Message
|
|
for {
|
|
msg, err := ev.Output.MessageOutput.MessageStream.Recv()
|
|
if err != nil {
|
|
break
|
|
}
|
|
msgs = append(msgs, msg)
|
|
}
|
|
if len(msgs) > 0 {
|
|
concated, err := schema.ConcatMessages(msgs)
|
|
if err == nil {
|
|
toolResultContents = append(toolResultContents, concated.Content)
|
|
}
|
|
}
|
|
}
|
|
}
|
|
}
|
|
|
|
assert.NotEmpty(t, toolResultContents, "should have at least one tool result event")
|
|
for _, content := range toolResultContents {
|
|
assert.Equal(t, modifiedResult, content,
|
|
"tool result event should contain the middleware-modified result, not the original")
|
|
assert.NotEqual(t, originalResult, content,
|
|
"tool result event should NOT contain the original result")
|
|
}
|
|
})
|
|
}
|
|
|
|
type simpleToolForMiddlewareTest struct {
|
|
name string
|
|
result string
|
|
}
|
|
|
|
func (s *simpleToolForMiddlewareTest) Info(_ context.Context) (*schema.ToolInfo, error) {
|
|
return &schema.ToolInfo{
|
|
Name: s.name,
|
|
Desc: "simple tool",
|
|
ParamsOneOf: schema.NewParamsOneOfByParams(
|
|
map[string]*schema.ParameterInfo{
|
|
"input": {
|
|
Desc: "input",
|
|
Required: true,
|
|
Type: schema.String,
|
|
},
|
|
}),
|
|
}, nil
|
|
}
|
|
|
|
func (s *simpleToolForMiddlewareTest) InvokableRun(_ context.Context, _ string, _ ...tool.Option) (string, error) {
|
|
return s.result, nil
|
|
}
|
|
|
|
func (s *simpleToolForMiddlewareTest) StreamableRun(_ context.Context, _ string, _ ...tool.Option) (*schema.StreamReader[string], error) {
|
|
return schema.StreamReaderFromArray([]string{s.result}), nil
|
|
}
|
|
|
|
func TestGetComposeOptions(t *testing.T) {
|
|
t.Run("WithChatModelOptions", func(t *testing.T) {
|
|
ctx := context.Background()
|
|
ctrl := gomock.NewController(t)
|
|
cm := mockModel.NewMockToolCallingChatModel(ctrl)
|
|
|
|
var capturedTemperature float32
|
|
cm.EXPECT().Generate(gomock.Any(), gomock.Any(), gomock.Any()).
|
|
DoAndReturn(func(ctx context.Context, msgs []*schema.Message, opts ...model.Option) (*schema.Message, error) {
|
|
options := model.GetCommonOptions(&model.Options{}, opts...)
|
|
if options.Temperature != nil {
|
|
capturedTemperature = *options.Temperature
|
|
}
|
|
return schema.AssistantMessage("response", nil), nil
|
|
}).Times(1)
|
|
|
|
agent, err := NewChatModelAgent(ctx, &ChatModelAgentConfig{
|
|
Name: "TestAgent",
|
|
Description: "Test agent",
|
|
Model: cm,
|
|
})
|
|
assert.NoError(t, err)
|
|
|
|
temp := float32(0.7)
|
|
iter := agent.Run(ctx, &AgentInput{Messages: []Message{schema.UserMessage("test")}},
|
|
WithChatModelOptions([]model.Option{model.WithTemperature(temp)}))
|
|
for {
|
|
_, ok := iter.Next()
|
|
if !ok {
|
|
break
|
|
}
|
|
}
|
|
|
|
assert.Equal(t, temp, capturedTemperature, "Temperature should be passed through WithChatModelOptions")
|
|
})
|
|
|
|
t.Run("WithToolOptions", func(t *testing.T) {
|
|
ctx := context.Background()
|
|
ctrl := gomock.NewController(t)
|
|
cm := mockModel.NewMockToolCallingChatModel(ctrl)
|
|
|
|
var toolOptionsCaptured bool
|
|
testTool := &toolOptionCapturingTool{
|
|
name: "test_tool",
|
|
onRun: func(opts []tool.Option) {
|
|
if len(opts) > 0 {
|
|
toolOptionsCaptured = true
|
|
}
|
|
},
|
|
}
|
|
info, _ := testTool.Info(ctx)
|
|
|
|
cm.EXPECT().WithTools(gomock.Any()).Return(cm, nil).AnyTimes()
|
|
cm.EXPECT().Generate(gomock.Any(), gomock.Any(), gomock.Any()).
|
|
Return(schema.AssistantMessage("Using tool", []schema.ToolCall{
|
|
{ID: "call1", Function: schema.FunctionCall{Name: info.Name, Arguments: "{}"}},
|
|
}), nil).Times(1)
|
|
cm.EXPECT().Generate(gomock.Any(), gomock.Any(), gomock.Any()).
|
|
Return(schema.AssistantMessage("done", nil), nil).Times(1)
|
|
|
|
agent, err := NewChatModelAgent(ctx, &ChatModelAgentConfig{
|
|
Name: "TestAgent",
|
|
Description: "Test agent",
|
|
Model: cm,
|
|
ToolsConfig: ToolsConfig{
|
|
ToolsNodeConfig: compose.ToolsNodeConfig{
|
|
Tools: []tool.BaseTool{testTool},
|
|
},
|
|
},
|
|
})
|
|
assert.NoError(t, err)
|
|
|
|
iter := agent.Run(ctx, &AgentInput{Messages: []Message{schema.UserMessage("test")}},
|
|
WithToolOptions([]tool.Option{testToolOption("test_value")}))
|
|
for {
|
|
_, ok := iter.Next()
|
|
if !ok {
|
|
break
|
|
}
|
|
}
|
|
|
|
assert.True(t, toolOptionsCaptured, "Tool options should be passed through WithToolOptions")
|
|
})
|
|
|
|
}
|
|
|
|
type toolOptionCapturingTool struct {
|
|
name string
|
|
onRun func(opts []tool.Option)
|
|
}
|
|
|
|
func (t *toolOptionCapturingTool) Info(_ context.Context) (*schema.ToolInfo, error) {
|
|
return &schema.ToolInfo{Name: t.name, Desc: t.name + " description"}, nil
|
|
}
|
|
|
|
func (t *toolOptionCapturingTool) InvokableRun(_ context.Context, _ string, opts ...tool.Option) (string, error) {
|
|
if t.onRun != nil {
|
|
t.onRun(opts)
|
|
}
|
|
return t.name + " result", nil
|
|
}
|
|
|
|
type testToolOptions struct {
|
|
value string
|
|
}
|
|
|
|
func testToolOption(value string) tool.Option {
|
|
return tool.WrapImplSpecificOptFn(func(o *testToolOptions) {
|
|
o.value = value
|
|
})
|
|
}
|
|
|
|
type errorTool struct {
|
|
infoErr error
|
|
}
|
|
|
|
func (e *errorTool) Info(_ context.Context) (*schema.ToolInfo, error) {
|
|
return nil, e.infoErr
|
|
}
|
|
|
|
func (e *errorTool) InvokableRun(_ context.Context, _ string, _ ...tool.Option) (string, error) {
|
|
return "", nil
|
|
}
|
|
|
|
func TestChatModelAgent_PrepareExecContextError(t *testing.T) {
|
|
t.Run("Run_WithToolInfoError_ReturnsError", func(t *testing.T) {
|
|
ctx := context.Background()
|
|
|
|
ctrl := gomock.NewController(t)
|
|
cm := mockModel.NewMockToolCallingChatModel(ctrl)
|
|
|
|
expectedErr := errors.New("tool info error")
|
|
errTool := &errorTool{infoErr: expectedErr}
|
|
|
|
agent, err := NewChatModelAgent(ctx, &ChatModelAgentConfig{
|
|
Name: "TestAgent",
|
|
Description: "Test agent",
|
|
Model: cm,
|
|
ToolsConfig: ToolsConfig{
|
|
ToolsNodeConfig: compose.ToolsNodeConfig{
|
|
Tools: []tool.BaseTool{errTool},
|
|
},
|
|
},
|
|
})
|
|
assert.NoError(t, err)
|
|
|
|
iter := agent.Run(ctx, &AgentInput{Messages: []Message{schema.UserMessage("test")}})
|
|
|
|
event, ok := iter.Next()
|
|
assert.True(t, ok)
|
|
assert.NotNil(t, event.Err)
|
|
assert.Contains(t, event.Err.Error(), "tool info error")
|
|
|
|
_, ok = iter.Next()
|
|
assert.False(t, ok)
|
|
})
|
|
|
|
t.Run("Resume_WithToolInfoError_ReturnsError", func(t *testing.T) {
|
|
ctx := context.Background()
|
|
|
|
ctrl := gomock.NewController(t)
|
|
cm := mockModel.NewMockToolCallingChatModel(ctrl)
|
|
|
|
expectedErr := errors.New("tool info error for resume")
|
|
errTool := &errorTool{infoErr: expectedErr}
|
|
|
|
agent, err := NewChatModelAgent(ctx, &ChatModelAgentConfig{
|
|
Name: "TestAgent",
|
|
Description: "Test agent",
|
|
Model: cm,
|
|
ToolsConfig: ToolsConfig{
|
|
ToolsNodeConfig: compose.ToolsNodeConfig{
|
|
Tools: []tool.BaseTool{errTool},
|
|
},
|
|
},
|
|
})
|
|
assert.NoError(t, err)
|
|
|
|
iter := agent.Resume(ctx, &ResumeInfo{
|
|
InterruptState: []byte("dummy"),
|
|
EnableStreaming: false,
|
|
})
|
|
|
|
event, ok := iter.Next()
|
|
assert.True(t, ok)
|
|
assert.NotNil(t, event.Err)
|
|
assert.Contains(t, event.Err.Error(), "tool info error for resume")
|
|
|
|
_, ok = iter.Next()
|
|
assert.False(t, ok)
|
|
})
|
|
}
|
|
|
|
// fakeEnhancedInvokableToolForTest implements tool.EnhancedInvokableTool for testing.
|
|
type fakeEnhancedInvokableToolForTest struct{}
|
|
|
|
func (t *fakeEnhancedInvokableToolForTest) Info(_ context.Context) (*schema.ToolInfo, error) {
|
|
return &schema.ToolInfo{
|
|
Name: "test_enhanced_invokable_tool",
|
|
Desc: "test enhanced invokable tool",
|
|
}, nil
|
|
}
|
|
|
|
func (t *fakeEnhancedInvokableToolForTest) InvokableRun(_ context.Context, _ *schema.ToolArgument, _ ...tool.Option) (*schema.ToolResult, error) {
|
|
return &schema.ToolResult{
|
|
Parts: []schema.ToolOutputPart{
|
|
{Type: schema.ToolPartTypeText, Text: "enhanced invokable result"},
|
|
},
|
|
}, nil
|
|
}
|
|
|
|
// fakeEnhancedStreamableToolForTest implements tool.EnhancedStreamableTool for testing.
|
|
type fakeEnhancedStreamableToolForTest struct{}
|
|
|
|
func (t *fakeEnhancedStreamableToolForTest) Info(_ context.Context) (*schema.ToolInfo, error) {
|
|
return &schema.ToolInfo{
|
|
Name: "test_enhanced_streamable_tool",
|
|
Desc: "test enhanced streamable tool",
|
|
}, nil
|
|
}
|
|
|
|
func (t *fakeEnhancedStreamableToolForTest) StreamableRun(_ context.Context, _ *schema.ToolArgument, _ ...tool.Option) (*schema.StreamReader[*schema.ToolResult], error) {
|
|
sr := schema.StreamReaderFromArray([]*schema.ToolResult{
|
|
{Parts: []schema.ToolOutputPart{{Type: schema.ToolPartTypeText, Text: "enhanced streamable result"}}},
|
|
})
|
|
return sr, nil
|
|
}
|
|
|
|
func TestPreprocessComposeCheckpoint_MigrateErrorIsReturned(t *testing.T) {
|
|
in := []byte("prefix\u0015" + stateGobNameV080 + "suffix")
|
|
_, err := preprocessComposeCheckpoint(in)
|
|
assert.Error(t, err)
|
|
}
|
|
|
|
func TestNewChatModelAgent_FailoverConfigValidation(t *testing.T) {
|
|
ctx := context.Background()
|
|
cm := &fakeChatModel{
|
|
callbacksEnabled: true,
|
|
generate: func(_ context.Context, _ []*schema.Message, _ ...model.Option) (*schema.Message, error) {
|
|
return schema.AssistantMessage("ok", nil), nil
|
|
},
|
|
stream: func(_ context.Context, _ []*schema.Message, _ ...model.Option) (*schema.StreamReader[*schema.Message], error) {
|
|
return schema.StreamReaderFromArray([]*schema.Message{schema.AssistantMessage("ok", nil)}), nil
|
|
},
|
|
}
|
|
|
|
t.Run("missing GetFailoverModel", func(t *testing.T) {
|
|
_, err := NewChatModelAgent(ctx, &ChatModelAgentConfig{
|
|
Name: "TestAgent",
|
|
Description: "test",
|
|
Model: cm,
|
|
ModelFailoverConfig: &ModelFailoverConfig[*schema.Message]{
|
|
ShouldFailover: func(context.Context, *schema.Message, error) bool { return true },
|
|
},
|
|
})
|
|
assert.Error(t, err)
|
|
assert.Contains(t, err.Error(), "ModelFailoverConfig.GetFailoverModel")
|
|
})
|
|
|
|
t.Run("missing ShouldFailover", func(t *testing.T) {
|
|
_, err := NewChatModelAgent(ctx, &ChatModelAgentConfig{
|
|
Name: "TestAgent",
|
|
Description: "test",
|
|
Model: cm,
|
|
ModelFailoverConfig: &ModelFailoverConfig[*schema.Message]{
|
|
GetFailoverModel: func(_ context.Context, _ *FailoverContext[*schema.Message]) (model.BaseChatModel, []*schema.Message, error) {
|
|
return cm, nil, nil
|
|
},
|
|
},
|
|
})
|
|
assert.Error(t, err)
|
|
assert.Contains(t, err.Error(), "ModelFailoverConfig.ShouldFailover")
|
|
})
|
|
}
|
|
|
|
// aliasCaptureTool captures the raw arguments JSON received by the tool.
|
|
type aliasCaptureTool struct {
|
|
name string
|
|
params map[string]*schema.ParameterInfo
|
|
receivedArgs string
|
|
}
|
|
|
|
func (t *aliasCaptureTool) Info(_ context.Context) (*schema.ToolInfo, error) {
|
|
return &schema.ToolInfo{
|
|
Name: t.name,
|
|
Desc: t.name + " tool",
|
|
ParamsOneOf: schema.NewParamsOneOfByParams(t.params),
|
|
}, nil
|
|
}
|
|
|
|
func (t *aliasCaptureTool) InvokableRun(_ context.Context, argumentsInJSON string, _ ...tool.Option) (string, error) {
|
|
t.receivedArgs = argumentsInJSON
|
|
return "ok", nil
|
|
}
|
|
|
|
func TestToolAliasesPropagation(t *testing.T) {
|
|
t.Run("prepareExecContext_propagates_ToolAliases", func(t *testing.T) {
|
|
ctx := context.Background()
|
|
ctrl := gomock.NewController(t)
|
|
|
|
captureTool := &aliasCaptureTool{
|
|
name: "grep",
|
|
params: map[string]*schema.ParameterInfo{
|
|
"pattern": {Type: schema.String, Desc: "regex pattern"},
|
|
"path": {Type: schema.String, Desc: "search path"},
|
|
},
|
|
}
|
|
|
|
cm := mockModel.NewMockToolCallingChatModel(ctrl)
|
|
generateCount := 0
|
|
cm.EXPECT().Generate(gomock.Any(), gomock.Any(), gomock.Any()).
|
|
DoAndReturn(func(ctx context.Context, msgs []*schema.Message, opts ...model.Option) (*schema.Message, error) {
|
|
generateCount++
|
|
if generateCount == 1 {
|
|
return &schema.Message{
|
|
Role: schema.Assistant,
|
|
ToolCalls: []schema.ToolCall{
|
|
{
|
|
ID: "call_1",
|
|
Function: schema.FunctionCall{
|
|
Name: "grep",
|
|
Arguments: `{"grep_content": "TODO", "path": "/src"}`,
|
|
},
|
|
},
|
|
},
|
|
}, nil
|
|
}
|
|
return schema.AssistantMessage("done", nil), nil
|
|
}).AnyTimes()
|
|
cm.EXPECT().WithTools(gomock.Any()).Return(cm, nil).AnyTimes()
|
|
|
|
agent, err := NewChatModelAgent(ctx, &ChatModelAgentConfig{
|
|
Name: "test",
|
|
Instruction: "test",
|
|
Model: cm,
|
|
ToolsConfig: ToolsConfig{
|
|
ToolsNodeConfig: compose.ToolsNodeConfig{
|
|
Tools: []tool.BaseTool{captureTool},
|
|
ToolAliases: map[string]compose.ToolAliasConfig{
|
|
"grep": {
|
|
ArgumentsAliases: map[string][]string{
|
|
"pattern": {"grep_content"},
|
|
},
|
|
},
|
|
},
|
|
},
|
|
},
|
|
})
|
|
require.NoError(t, err)
|
|
|
|
iter := agent.Run(ctx, &AgentInput{Messages: []Message{schema.UserMessage("search for TODOs")}})
|
|
for {
|
|
_, ok := iter.Next()
|
|
if !ok {
|
|
break
|
|
}
|
|
}
|
|
|
|
require.NotEmpty(t, captureTool.receivedArgs, "tool should have been called")
|
|
var args map[string]any
|
|
err = json.Unmarshal([]byte(captureTool.receivedArgs), &args)
|
|
require.NoError(t, err)
|
|
assert.Equal(t, "TODO", args["pattern"], "alias 'grep_content' should be remapped to 'pattern'")
|
|
assert.NotContains(t, args, "grep_content", "alias key should not be present after remapping")
|
|
assert.Equal(t, "/src", args["path"])
|
|
})
|
|
|
|
t.Run("applyBeforeAgent_preserves_ToolAliases_when_handler_modifies_tools", func(t *testing.T) {
|
|
ctx := context.Background()
|
|
ctrl := gomock.NewController(t)
|
|
|
|
captureTool := &aliasCaptureTool{
|
|
name: "grep",
|
|
params: map[string]*schema.ParameterInfo{
|
|
"pattern": {Type: schema.String, Desc: "regex pattern"},
|
|
},
|
|
}
|
|
|
|
extraTool := &aliasCaptureTool{
|
|
name: "extra_tool",
|
|
params: map[string]*schema.ParameterInfo{
|
|
"input": {Type: schema.String},
|
|
},
|
|
}
|
|
|
|
cm := mockModel.NewMockToolCallingChatModel(ctrl)
|
|
generateCount := 0
|
|
cm.EXPECT().Generate(gomock.Any(), gomock.Any(), gomock.Any()).
|
|
DoAndReturn(func(ctx context.Context, msgs []*schema.Message, opts ...model.Option) (*schema.Message, error) {
|
|
generateCount++
|
|
if generateCount == 1 {
|
|
return &schema.Message{
|
|
Role: schema.Assistant,
|
|
ToolCalls: []schema.ToolCall{
|
|
{
|
|
ID: "call_1",
|
|
Function: schema.FunctionCall{
|
|
Name: "grep",
|
|
Arguments: `{"grep_content": "FIXME"}`,
|
|
},
|
|
},
|
|
},
|
|
}, nil
|
|
}
|
|
return schema.AssistantMessage("done", nil), nil
|
|
}).AnyTimes()
|
|
cm.EXPECT().WithTools(gomock.Any()).Return(cm, nil).AnyTimes()
|
|
|
|
handler := &testToolsHandler{
|
|
BaseChatModelAgentMiddleware: &BaseChatModelAgentMiddleware{},
|
|
tools: []tool.BaseTool{extraTool},
|
|
}
|
|
|
|
agent, err := NewChatModelAgent(ctx, &ChatModelAgentConfig{
|
|
Name: "test",
|
|
Instruction: "test",
|
|
Model: cm,
|
|
ToolsConfig: ToolsConfig{
|
|
ToolsNodeConfig: compose.ToolsNodeConfig{
|
|
Tools: []tool.BaseTool{captureTool},
|
|
ToolAliases: map[string]compose.ToolAliasConfig{
|
|
"grep": {
|
|
ArgumentsAliases: map[string][]string{
|
|
"pattern": {"grep_content"},
|
|
},
|
|
},
|
|
},
|
|
},
|
|
},
|
|
Handlers: []ChatModelAgentMiddleware{handler},
|
|
})
|
|
require.NoError(t, err)
|
|
|
|
iter := agent.Run(ctx, &AgentInput{Messages: []Message{schema.UserMessage("search for FIXMEs")}})
|
|
for {
|
|
_, ok := iter.Next()
|
|
if !ok {
|
|
break
|
|
}
|
|
}
|
|
|
|
require.NotEmpty(t, captureTool.receivedArgs, "tool should have been called")
|
|
var args map[string]any
|
|
err = json.Unmarshal([]byte(captureTool.receivedArgs), &args)
|
|
require.NoError(t, err)
|
|
assert.Equal(t, "FIXME", args["pattern"], "alias 'grep_content' should be remapped to 'pattern' even after handler rebuild")
|
|
assert.NotContains(t, args, "grep_content", "alias key should not be present after remapping")
|
|
})
|
|
|
|
t.Run("name_alias_propagated_through_prepareExecContext", func(t *testing.T) {
|
|
ctx := context.Background()
|
|
ctrl := gomock.NewController(t)
|
|
|
|
captureTool := &aliasCaptureTool{
|
|
name: "grep",
|
|
params: map[string]*schema.ParameterInfo{
|
|
"pattern": {Type: schema.String},
|
|
},
|
|
}
|
|
|
|
cm := mockModel.NewMockToolCallingChatModel(ctrl)
|
|
generateCount := 0
|
|
cm.EXPECT().Generate(gomock.Any(), gomock.Any(), gomock.Any()).
|
|
DoAndReturn(func(ctx context.Context, msgs []*schema.Message, opts ...model.Option) (*schema.Message, error) {
|
|
generateCount++
|
|
if generateCount == 1 {
|
|
return &schema.Message{
|
|
Role: schema.Assistant,
|
|
ToolCalls: []schema.ToolCall{
|
|
{
|
|
ID: "call_1",
|
|
Function: schema.FunctionCall{
|
|
Name: "search_content",
|
|
Arguments: `{"pattern": "TODO"}`,
|
|
},
|
|
},
|
|
},
|
|
}, nil
|
|
}
|
|
return schema.AssistantMessage("done", nil), nil
|
|
}).AnyTimes()
|
|
cm.EXPECT().WithTools(gomock.Any()).Return(cm, nil).AnyTimes()
|
|
|
|
agent, err := NewChatModelAgent(ctx, &ChatModelAgentConfig{
|
|
Name: "test",
|
|
Instruction: "test",
|
|
Model: cm,
|
|
ToolsConfig: ToolsConfig{
|
|
ToolsNodeConfig: compose.ToolsNodeConfig{
|
|
Tools: []tool.BaseTool{captureTool},
|
|
ToolAliases: map[string]compose.ToolAliasConfig{
|
|
"grep": {
|
|
NameAliases: []string{"search_content"},
|
|
},
|
|
},
|
|
},
|
|
},
|
|
})
|
|
require.NoError(t, err)
|
|
|
|
iter := agent.Run(ctx, &AgentInput{Messages: []Message{schema.UserMessage("search")}})
|
|
for {
|
|
_, ok := iter.Next()
|
|
if !ok {
|
|
break
|
|
}
|
|
}
|
|
|
|
require.NotEmpty(t, captureTool.receivedArgs, "tool should have been called via name alias 'search_content'")
|
|
var args map[string]any
|
|
err = json.Unmarshal([]byte(captureTool.receivedArgs), &args)
|
|
require.NoError(t, err)
|
|
assert.Equal(t, "TODO", args["pattern"])
|
|
})
|
|
|
|
t.Run("handler_adds_tool_matching_preexisting_ToolAliases_with_no_initial_tools", func(t *testing.T) {
|
|
ctx := context.Background()
|
|
ctrl := gomock.NewController(t)
|
|
|
|
captureTool := &aliasCaptureTool{
|
|
name: "grep",
|
|
params: map[string]*schema.ParameterInfo{
|
|
"pattern": {Type: schema.String, Desc: "regex pattern"},
|
|
},
|
|
}
|
|
|
|
cm := mockModel.NewMockToolCallingChatModel(ctrl)
|
|
generateCount := 0
|
|
cm.EXPECT().Generate(gomock.Any(), gomock.Any(), gomock.Any()).
|
|
DoAndReturn(func(ctx context.Context, msgs []*schema.Message, opts ...model.Option) (*schema.Message, error) {
|
|
generateCount++
|
|
if generateCount == 1 {
|
|
return &schema.Message{
|
|
Role: schema.Assistant,
|
|
ToolCalls: []schema.ToolCall{
|
|
{
|
|
ID: "call_1",
|
|
Function: schema.FunctionCall{
|
|
Name: "grep",
|
|
Arguments: `{"grep_content": "BUG"}`,
|
|
},
|
|
},
|
|
},
|
|
}, nil
|
|
}
|
|
return schema.AssistantMessage("done", nil), nil
|
|
}).AnyTimes()
|
|
cm.EXPECT().WithTools(gomock.Any()).Return(cm, nil).AnyTimes()
|
|
|
|
handler := &testToolsHandler{
|
|
BaseChatModelAgentMiddleware: &BaseChatModelAgentMiddleware{},
|
|
tools: []tool.BaseTool{captureTool},
|
|
}
|
|
|
|
agent, err := NewChatModelAgent(ctx, &ChatModelAgentConfig{
|
|
Name: "test",
|
|
Instruction: "test",
|
|
Model: cm,
|
|
ToolsConfig: ToolsConfig{
|
|
ToolsNodeConfig: compose.ToolsNodeConfig{
|
|
ToolAliases: map[string]compose.ToolAliasConfig{
|
|
"grep": {
|
|
ArgumentsAliases: map[string][]string{
|
|
"pattern": {"grep_content"},
|
|
},
|
|
},
|
|
},
|
|
},
|
|
},
|
|
Handlers: []ChatModelAgentMiddleware{handler},
|
|
})
|
|
require.NoError(t, err)
|
|
|
|
iter := agent.Run(ctx, &AgentInput{Messages: []Message{schema.UserMessage("find bugs")}})
|
|
for {
|
|
_, ok := iter.Next()
|
|
if !ok {
|
|
break
|
|
}
|
|
}
|
|
|
|
require.NotEmpty(t, captureTool.receivedArgs, "tool added by handler should have been called")
|
|
var args map[string]any
|
|
err = json.Unmarshal([]byte(captureTool.receivedArgs), &args)
|
|
require.NoError(t, err)
|
|
assert.Equal(t, "BUG", args["pattern"], "alias 'grep_content' should be remapped to 'pattern' for handler-added tool")
|
|
assert.NotContains(t, args, "grep_content")
|
|
})
|
|
}
|