183 lines
3.4 KiB
Go
183 lines
3.4 KiB
Go
package client
|
|
|
|
import (
|
|
"errors"
|
|
"testing"
|
|
|
|
"github.com/google/uuid"
|
|
"github.com/stretchr/testify/suite"
|
|
)
|
|
|
|
type ToolTestSuite struct {
|
|
suite.Suite
|
|
|
|
handler ToolsHandlers
|
|
}
|
|
|
|
func (suite *ToolTestSuite) SetupTest() {
|
|
suite.handler = ToolsHandlers{
|
|
handlers: &map[string]ToolHandler{},
|
|
}
|
|
|
|
suite.handler.AddTool("a", func(info ToolHandlerInfo, args string, call ToolCall) (any, error) {
|
|
return args, nil
|
|
})
|
|
|
|
suite.handler.AddTool("error", func(info ToolHandlerInfo, args string, call ToolCall) (any, error) {
|
|
return false, errors.New("I will always error")
|
|
})
|
|
}
|
|
|
|
func (suite *ToolTestSuite) TestSingleToolCall() {
|
|
assert := suite.Assert()
|
|
require := suite.Require()
|
|
|
|
response, err := suite.handler.Handle(
|
|
ToolHandlerInfo{
|
|
UserId: uuid.Nil,
|
|
ImageId: uuid.Nil,
|
|
},
|
|
AgentAssistantToolCall{
|
|
Role: "assistant",
|
|
Content: "",
|
|
ToolCalls: []ToolCall{{
|
|
Index: 0,
|
|
Id: "1",
|
|
Function: FunctionCall{
|
|
Name: "a",
|
|
Arguments: "return",
|
|
},
|
|
}},
|
|
})
|
|
|
|
require.NoError(err, "Tool call shouldnt return an error")
|
|
|
|
assert.EqualValues(response, []AgentTextMessage{{
|
|
Role: "tool",
|
|
Content: "\"return\"",
|
|
ToolCallId: "1",
|
|
Name: "a",
|
|
}})
|
|
}
|
|
|
|
func (suite *ToolTestSuite) TestEmptyCall() {
|
|
require := suite.Require()
|
|
|
|
_, err := suite.handler.Handle(
|
|
ToolHandlerInfo{
|
|
UserId: uuid.Nil,
|
|
ImageId: uuid.Nil,
|
|
},
|
|
AgentAssistantToolCall{
|
|
Role: "assistant",
|
|
Content: "",
|
|
ToolCalls: []ToolCall{},
|
|
})
|
|
|
|
require.ErrorIs(err, NoToolCallError)
|
|
}
|
|
|
|
func (suite *ToolTestSuite) TestMultipleToolCalls() {
|
|
assert := suite.Assert()
|
|
require := suite.Require()
|
|
|
|
response, err := suite.handler.Handle(
|
|
ToolHandlerInfo{
|
|
UserId: uuid.Nil,
|
|
ImageId: uuid.Nil,
|
|
},
|
|
AgentAssistantToolCall{
|
|
Role: "assistant",
|
|
Content: "",
|
|
ToolCalls: []ToolCall{
|
|
{
|
|
Index: 0,
|
|
Id: "1",
|
|
Function: FunctionCall{
|
|
Name: "a",
|
|
Arguments: "first-call",
|
|
},
|
|
},
|
|
{
|
|
Index: 1,
|
|
Id: "2",
|
|
Function: FunctionCall{
|
|
Name: "a",
|
|
Arguments: "second-call",
|
|
},
|
|
},
|
|
},
|
|
})
|
|
|
|
require.NoError(err, "Tool call shouldnt return an error")
|
|
|
|
assert.EqualValues(response, []AgentTextMessage{
|
|
{
|
|
Role: "tool",
|
|
Content: "\"first-call\"",
|
|
ToolCallId: "1",
|
|
Name: "a",
|
|
},
|
|
{
|
|
Role: "tool",
|
|
Content: "\"second-call\"",
|
|
ToolCallId: "2",
|
|
Name: "a",
|
|
},
|
|
})
|
|
}
|
|
|
|
func (suite *ToolTestSuite) TestMultipleToolCallsWithErrors() {
|
|
assert := suite.Assert()
|
|
require := suite.Require()
|
|
|
|
response, err := suite.handler.Handle(
|
|
ToolHandlerInfo{
|
|
UserId: uuid.Nil,
|
|
ImageId: uuid.Nil,
|
|
},
|
|
AgentAssistantToolCall{
|
|
Role: "assistant",
|
|
Content: "",
|
|
ToolCalls: []ToolCall{
|
|
{
|
|
Index: 0,
|
|
Id: "1",
|
|
Function: FunctionCall{
|
|
Name: "error",
|
|
Arguments: "",
|
|
},
|
|
},
|
|
{
|
|
Index: 1,
|
|
Id: "2",
|
|
Function: FunctionCall{
|
|
Name: "a",
|
|
Arguments: "no-error",
|
|
},
|
|
},
|
|
},
|
|
})
|
|
|
|
require.NoError(err, "Tool call shouldnt return an error")
|
|
|
|
assert.EqualValues(response, []AgentTextMessage{
|
|
{
|
|
Role: "tool",
|
|
Content: "I will always error",
|
|
ToolCallId: "1",
|
|
Name: "error",
|
|
},
|
|
{
|
|
Role: "tool",
|
|
Content: "\"no-error\"",
|
|
ToolCallId: "2",
|
|
Name: "a",
|
|
},
|
|
})
|
|
}
|
|
|
|
func TestToolSuite(t *testing.T) {
|
|
suite.Run(t, &ToolTestSuite{})
|
|
}
|