282 lines
8.7 KiB
Go
282 lines
8.7 KiB
Go
package server
|
|
|
|
import (
|
|
"bytes"
|
|
"encoding/json"
|
|
"fmt"
|
|
"os"
|
|
"path/filepath"
|
|
"strings"
|
|
"testing"
|
|
|
|
"github.com/google/go-cmp/cmp"
|
|
|
|
"github.com/ollama/ollama/api"
|
|
"github.com/ollama/ollama/template"
|
|
)
|
|
|
|
func readFile(t *testing.T, base, name string) *bytes.Buffer {
|
|
t.Helper()
|
|
|
|
bts, err := os.ReadFile(filepath.Join(base, name))
|
|
if err != nil {
|
|
t.Fatal(err)
|
|
}
|
|
|
|
return bytes.NewBuffer(bts)
|
|
}
|
|
|
|
func TestParseToolCalls(t *testing.T) {
|
|
p := filepath.Join("testdata", "tools")
|
|
t1 := api.ToolCall{
|
|
Function: api.ToolCallFunction{
|
|
Name: "get_current_weather",
|
|
Arguments: api.ToolCallFunctionArguments{
|
|
"format": "fahrenheit",
|
|
"location": "San Francisco, CA",
|
|
},
|
|
},
|
|
}
|
|
t2 := api.ToolCall{
|
|
Function: api.ToolCallFunction{
|
|
Name: "get_current_weather",
|
|
Arguments: api.ToolCallFunctionArguments{
|
|
"format": "celsius",
|
|
"location": "Toronto, Canada",
|
|
},
|
|
},
|
|
}
|
|
|
|
cases := []struct {
|
|
name string
|
|
model string
|
|
output string
|
|
token string
|
|
expected []api.ToolCall
|
|
wantErr bool
|
|
}{
|
|
{
|
|
name: "mistral invalid json",
|
|
model: "mistral",
|
|
output: `[TOOL_CALLS] [{"name": "get_current_weather", "arguments": {"format":"fahrenheit","location":"San Francisco, CA"}},{"name": "get_curren}]`,
|
|
token: "[TOOL_CALLS]",
|
|
expected: []api.ToolCall{},
|
|
wantErr: true,
|
|
},
|
|
{
|
|
name: "mistral valid json",
|
|
model: "mistral",
|
|
output: `[TOOL_CALLS] [{"name": "get_current_weather", "arguments": {"format":"fahrenheit","location":"San Francisco, CA"}},{"name": "get_current_weather", "arguments": {"format":"celsius","location":"Toronto, Canada"}}]`,
|
|
token: "[TOOL_CALLS]",
|
|
expected: []api.ToolCall{t1, t2},
|
|
wantErr: false,
|
|
},
|
|
{
|
|
name: "mistral incomplete json",
|
|
model: "mistral",
|
|
output: `[TOOL_CALLS] [{"name": "get_current_weather", "arguments": {"format":"fahrenheit","location":"San Francisco, `,
|
|
token: "[TOOL_CALLS]",
|
|
expected: []api.ToolCall{},
|
|
wantErr: true,
|
|
},
|
|
{
|
|
name: "mistral without tool token",
|
|
model: "mistral",
|
|
output: `I'm not aware of that information. However, I can suggest searching for the weather using the "get_current_weather" function:
|
|
|
|
[{"name": "get_current_weather", "arguments": {"format":"fahrenheit","location":"San Francisco, CA"}},{"name": "get_current_weather", "arguments": {"format":"celsius","location":"Toronto, Canada"}}]`,
|
|
token: "[TOOL_CALLS]",
|
|
expected: []api.ToolCall{},
|
|
wantErr: true,
|
|
},
|
|
{
|
|
name: "mistral without tool token - tool first",
|
|
model: "mistral",
|
|
output: `[{"name": "get_current_weather", "arguments": {"format":"fahrenheit","location":"San Francisco, CA"}},{"name": "get_current_weather", "arguments": {"format":"celsius","location":"Toronto, Canada"}}]`,
|
|
token: "[TOOL_CALLS]",
|
|
expected: []api.ToolCall{t1, t2},
|
|
wantErr: false,
|
|
},
|
|
{
|
|
name: "command-r-plus with json block",
|
|
model: "command-r-plus",
|
|
output: "Action: ```json" + `
|
|
[
|
|
{
|
|
"tool_name": "get_current_weather",
|
|
"parameters": {
|
|
"format": "fahrenheit",
|
|
"location": "San Francisco, CA"
|
|
}
|
|
},
|
|
{
|
|
"tool_name": "get_current_weather",
|
|
"parameters": {
|
|
"format": "celsius",
|
|
"location": "Toronto, Canada"
|
|
}
|
|
}
|
|
]
|
|
` + "```",
|
|
token: "Action:",
|
|
expected: []api.ToolCall{t1, t2},
|
|
wantErr: false,
|
|
},
|
|
{
|
|
name: "firefunction with functools",
|
|
model: "firefunction",
|
|
output: ` functools[{"name": "get_current_weather", "arguments": {"format":"fahrenheit","location":"San Francisco, CA"}},{"name": "get_current_weather", "arguments": {"format":"celsius","location":"Toronto, Canada"}}]`,
|
|
token: "functools",
|
|
expected: []api.ToolCall{t1, t2},
|
|
wantErr: false,
|
|
},
|
|
{
|
|
name: "llama3 with tool call tags",
|
|
model: "llama3-groq-tool-use",
|
|
output: `<tool_call>
|
|
{"name": "get_current_weather", "arguments": {"format":"fahrenheit","location":"San Francisco, CA"}}
|
|
</tool_call>`,
|
|
token: "<tool_call>",
|
|
expected: []api.ToolCall{t1},
|
|
wantErr: false,
|
|
},
|
|
{
|
|
name: "xlam with tool_calls wrapper",
|
|
model: "xlam",
|
|
output: `{"tool_calls": [{"name": "get_current_weather", "arguments": {"format":"fahrenheit","location":"San Francisco, CA"}},{"name": "get_current_weather", "arguments": {"format":"celsius","location":"Toronto, Canada"}}]}`,
|
|
token: "",
|
|
expected: []api.ToolCall{t1, t2},
|
|
wantErr: false,
|
|
},
|
|
{
|
|
name: "qwen with single tool call",
|
|
model: "qwen2.5-coder",
|
|
output: `<tool_call>{"name": "get_current_weather", "arguments": {"format":"fahrenheit","location":"San Francisco, CA"}}</tool_call>`,
|
|
token: "<tool_call>",
|
|
expected: []api.ToolCall{t1},
|
|
wantErr: false,
|
|
},
|
|
{
|
|
name: "qwen with invalid tool token",
|
|
model: "qwen2.5-coder",
|
|
output: `[{"name": "get_current_weather", "arguments": {"format":"fahrenheit","location":"San Francisco, CA"}}, {"name": "get_current_weather", "arguments": {"format":"celsius","location":"Toronto, Canada"}}]`,
|
|
token: "[TOOL_CALLS]",
|
|
expected: []api.ToolCall{t1, t2},
|
|
wantErr: false,
|
|
},
|
|
{
|
|
name: "qwen with no tool calls",
|
|
model: "qwen2.5-coder",
|
|
output: " The weather in San Francisco, CA is 70°F and in Toronto, Canada is 20°C.",
|
|
token: "",
|
|
expected: []api.ToolCall{},
|
|
wantErr: true,
|
|
},
|
|
}
|
|
|
|
var tools []api.Tool
|
|
if err := json.Unmarshal(readFile(t, p, "tools.json").Bytes(), &tools); err != nil {
|
|
t.Fatal(err)
|
|
}
|
|
|
|
var messages []api.Message
|
|
if err := json.Unmarshal(readFile(t, p, "messages.json").Bytes(), &messages); err != nil {
|
|
t.Fatal(err)
|
|
}
|
|
|
|
for _, tt := range cases {
|
|
t.Run(tt.model, func(t *testing.T) {
|
|
tmpl, err := template.Parse(readFile(t, p, fmt.Sprintf("%s.gotmpl", tt.model)).String())
|
|
if err != nil {
|
|
t.Fatal(err)
|
|
}
|
|
|
|
t.Run("template", func(t *testing.T) {
|
|
var actual bytes.Buffer
|
|
if err := tmpl.Execute(&actual, template.Values{Tools: tools, Messages: messages}); err != nil {
|
|
t.Fatal(err)
|
|
}
|
|
|
|
if diff := cmp.Diff(actual.String(), readFile(t, p, fmt.Sprintf("%s.out", tt.model)).String()); diff != "" {
|
|
t.Errorf("mismatch (-got +want):\n%s", diff)
|
|
}
|
|
})
|
|
|
|
t.Run("parse", func(t *testing.T) {
|
|
m := &Model{Template: tmpl}
|
|
tmpl, ok := ToolTemplate(m)
|
|
if !ok {
|
|
t.Fatal("no tool template found")
|
|
}
|
|
got := []api.ToolCall{}
|
|
tokens := strings.Fields(tt.output)
|
|
sb := strings.Builder{}
|
|
success := false
|
|
for _, tok := range tokens {
|
|
sb.WriteString(" " + tok)
|
|
toolCalls, partial, err := ParseToolCalls(sb.String(), tt.token, tmpl)
|
|
if err == nil {
|
|
success = true
|
|
}
|
|
if partial {
|
|
continue
|
|
}
|
|
got = append(got, toolCalls...)
|
|
sb.Reset()
|
|
}
|
|
|
|
if !tt.wantErr {
|
|
if diff := cmp.Diff(got, tt.expected); diff != "" {
|
|
t.Errorf("mismatch (-got +want):\n%s", diff)
|
|
}
|
|
}
|
|
if !success && !tt.wantErr {
|
|
t.Errorf("expected success but got errors")
|
|
}
|
|
})
|
|
})
|
|
}
|
|
}
|
|
|
|
func TestParseObjects(t *testing.T) {
|
|
tests := []struct {
|
|
input string
|
|
want []map[string]any
|
|
}{
|
|
{
|
|
input: `[{"name": "get_current_weather", "arguments": {"format":"fahrenheit","location":"San Francisco, CA"}},{"name": "get_current_weather", "arguments": {"format":"celsius","location":"Toronto, Canada"}}]`,
|
|
want: []map[string]any{
|
|
{"name": "get_current_weather", "arguments": map[string]any{"format": "fahrenheit", "location": "San Francisco, CA"}},
|
|
{"name": "get_current_weather", "arguments": map[string]any{"format": "celsius", "location": "Toronto, Canada"}},
|
|
},
|
|
},
|
|
{
|
|
input: `<some_token>{"name": "get_current_weather", "arguments": {"format":"fahrenheit","location":"San Francisco, CA"}} </toolcall>`,
|
|
want: []map[string]any{
|
|
{"name": "get_current_weather", "arguments": map[string]any{"format": "fahrenheit", "location": "San Francisco, CA"}},
|
|
},
|
|
},
|
|
{
|
|
input: `<some_token>{"name": "get_current_weather", "arguments": {"format":"fahrenheit","location":"San Francisco, CA"}} </toolcall> <toolcall>{"name": "get_current_weather", "arguments": {"format":"celsius","location":"Toronto, ON"}} </toolcall>`,
|
|
want: []map[string]any{
|
|
{"name": "get_current_weather", "arguments": map[string]any{"format": "fahrenheit", "location": "San Francisco, CA"}},
|
|
{"name": "get_current_weather", "arguments": map[string]any{"format": "celsius", "location": "Toronto, ON"}},
|
|
},
|
|
},
|
|
{
|
|
input: `{"name": "get_current_weather", "arguments": `,
|
|
want: nil,
|
|
},
|
|
}
|
|
|
|
for _, tc := range tests {
|
|
t.Run(tc.input, func(t *testing.T) {
|
|
got := parseObjects(tc.input)
|
|
|
|
if diff := cmp.Diff(got, tc.want); diff != "" {
|
|
t.Errorf("mismatch (-got +want):\n%s", diff)
|
|
}
|
|
})
|
|
}
|
|
}
|