Unverified Commit 626af2d8 authored by Devon Rifkin's avatar Devon Rifkin Committed by GitHub
Browse files

template: fix args-as-json rendering (#13636)

In #13525, I accidentally broke templates' ability to automatically
render tool call function arguments as JSON.

We do need these to be proper maps because we need templates to be able
to call range, which can't be done on custom types.
parent 76912c06
...@@ -381,6 +381,28 @@ func (t templateTools) String() string { ...@@ -381,6 +381,28 @@ func (t templateTools) String() string {
return string(bts) return string(bts)
} }
// templateArgs is a map type with JSON string output for templates.
type templateArgs map[string]any
func (t templateArgs) String() string {
if t == nil {
return "{}"
}
bts, _ := json.Marshal(t)
return string(bts)
}
// templateProperties is a map type with JSON string output for templates.
type templateProperties map[string]api.ToolProperty
func (t templateProperties) String() string {
if t == nil {
return "{}"
}
bts, _ := json.Marshal(t)
return string(bts)
}
// templateTool is a template-compatible representation of api.Tool // templateTool is a template-compatible representation of api.Tool
// with Properties as a regular map for template ranging. // with Properties as a regular map for template ranging.
type templateTool struct { type templateTool struct {
...@@ -396,11 +418,11 @@ type templateToolFunction struct { ...@@ -396,11 +418,11 @@ type templateToolFunction struct {
} }
type templateToolFunctionParameters struct { type templateToolFunctionParameters struct {
Type string `json:"type"` Type string `json:"type"`
Defs any `json:"$defs,omitempty"` Defs any `json:"$defs,omitempty"`
Items any `json:"items,omitempty"` Items any `json:"items,omitempty"`
Required []string `json:"required,omitempty"` Required []string `json:"required,omitempty"`
Properties map[string]api.ToolProperty `json:"properties"` Properties templateProperties `json:"properties"`
} }
// templateToolCall is a template-compatible representation of api.ToolCall // templateToolCall is a template-compatible representation of api.ToolCall
...@@ -413,7 +435,7 @@ type templateToolCall struct { ...@@ -413,7 +435,7 @@ type templateToolCall struct {
type templateToolCallFunction struct { type templateToolCallFunction struct {
Index int Index int
Name string Name string
Arguments map[string]any Arguments templateArgs
} }
// templateMessage is a template-compatible representation of api.Message // templateMessage is a template-compatible representation of api.Message
...@@ -446,7 +468,7 @@ func convertToolsForTemplate(tools api.Tools) templateTools { ...@@ -446,7 +468,7 @@ func convertToolsForTemplate(tools api.Tools) templateTools {
Defs: tool.Function.Parameters.Defs, Defs: tool.Function.Parameters.Defs,
Items: tool.Function.Parameters.Items, Items: tool.Function.Parameters.Items,
Required: tool.Function.Parameters.Required, Required: tool.Function.Parameters.Required,
Properties: tool.Function.Parameters.Properties.ToMap(), Properties: templateProperties(tool.Function.Parameters.Properties.ToMap()),
}, },
}, },
} }
...@@ -468,7 +490,7 @@ func convertMessagesForTemplate(messages []*api.Message) []*templateMessage { ...@@ -468,7 +490,7 @@ func convertMessagesForTemplate(messages []*api.Message) []*templateMessage {
Function: templateToolCallFunction{ Function: templateToolCallFunction{
Index: tc.Function.Index, Index: tc.Function.Index,
Name: tc.Function.Name, Name: tc.Function.Name,
Arguments: tc.Function.Arguments.ToMap(), Arguments: templateArgs(tc.Function.Arguments.ToMap()),
}, },
}) })
} }
......
...@@ -613,3 +613,159 @@ func TestCollate(t *testing.T) { ...@@ -613,3 +613,159 @@ func TestCollate(t *testing.T) {
}) })
} }
} }
func TestTemplateArgumentsJSON(t *testing.T) {
// Test that {{ .Function.Arguments }} outputs valid JSON, not map[key:value]
tmpl := `{{- range .Messages }}{{- range .ToolCalls }}{{ .Function.Arguments }}{{- end }}{{- end }}`
template, err := Parse(tmpl)
if err != nil {
t.Fatal(err)
}
args := api.NewToolCallFunctionArguments()
args.Set("location", "Tokyo")
args.Set("unit", "celsius")
var buf bytes.Buffer
err = template.Execute(&buf, Values{
Messages: []api.Message{{
Role: "assistant",
ToolCalls: []api.ToolCall{{
Function: api.ToolCallFunction{
Name: "get_weather",
Arguments: args,
},
}},
}},
})
if err != nil {
t.Fatal(err)
}
got := buf.String()
// Should be valid JSON, not "map[location:Tokyo unit:celsius]"
if strings.HasPrefix(got, "map[") {
t.Errorf("Arguments output as Go map format: %s", got)
}
var parsed map[string]any
if err := json.Unmarshal([]byte(got), &parsed); err != nil {
t.Errorf("Arguments not valid JSON: %s, error: %v", got, err)
}
}
func TestTemplatePropertiesJSON(t *testing.T) {
// Test that {{ .Function.Parameters.Properties }} outputs valid JSON
// Note: template must reference .Messages to trigger the modern code path that converts Tools
tmpl := `{{- range .Messages }}{{- end }}{{- range .Tools }}{{ .Function.Parameters.Properties }}{{- end }}`
template, err := Parse(tmpl)
if err != nil {
t.Fatal(err)
}
props := api.NewToolPropertiesMap()
props.Set("location", api.ToolProperty{Type: api.PropertyType{"string"}, Description: "City name"})
var buf bytes.Buffer
err = template.Execute(&buf, Values{
Messages: []api.Message{{Role: "user", Content: "test"}},
Tools: api.Tools{{
Type: "function",
Function: api.ToolFunction{
Name: "get_weather",
Description: "Get weather",
Parameters: api.ToolFunctionParameters{
Type: "object",
Properties: props,
},
},
}},
})
if err != nil {
t.Fatal(err)
}
got := buf.String()
// Should be valid JSON, not "map[location:{...}]"
if strings.HasPrefix(got, "map[") {
t.Errorf("Properties output as Go map format: %s", got)
}
var parsed map[string]any
if err := json.Unmarshal([]byte(got), &parsed); err != nil {
t.Errorf("Properties not valid JSON: %s, error: %v", got, err)
}
}
func TestTemplateArgumentsRange(t *testing.T) {
// Test that we can range over Arguments in templates
tmpl := `{{- range .Messages }}{{- range .ToolCalls }}{{- range $k, $v := .Function.Arguments }}{{ $k }}={{ $v }};{{- end }}{{- end }}{{- end }}`
template, err := Parse(tmpl)
if err != nil {
t.Fatal(err)
}
args := api.NewToolCallFunctionArguments()
args.Set("city", "Tokyo")
var buf bytes.Buffer
err = template.Execute(&buf, Values{
Messages: []api.Message{{
Role: "assistant",
ToolCalls: []api.ToolCall{{
Function: api.ToolCallFunction{
Name: "get_weather",
Arguments: args,
},
}},
}},
})
if err != nil {
t.Fatal(err)
}
got := buf.String()
if got != "city=Tokyo;" {
t.Errorf("Range over Arguments failed, got: %s, want: city=Tokyo;", got)
}
}
func TestTemplatePropertiesRange(t *testing.T) {
// Test that we can range over Properties in templates
// Note: template must reference .Messages to trigger the modern code path that converts Tools
tmpl := `{{- range .Messages }}{{- end }}{{- range .Tools }}{{- range $name, $prop := .Function.Parameters.Properties }}{{ $name }}:{{ $prop.Type }};{{- end }}{{- end }}`
template, err := Parse(tmpl)
if err != nil {
t.Fatal(err)
}
props := api.NewToolPropertiesMap()
props.Set("location", api.ToolProperty{Type: api.PropertyType{"string"}})
var buf bytes.Buffer
err = template.Execute(&buf, Values{
Messages: []api.Message{{Role: "user", Content: "test"}},
Tools: api.Tools{{
Type: "function",
Function: api.ToolFunction{
Name: "get_weather",
Parameters: api.ToolFunctionParameters{
Type: "object",
Properties: props,
},
},
}},
})
if err != nil {
t.Fatal(err)
}
got := buf.String()
if got != "location:string;" {
t.Errorf("Range over Properties failed, got: %s, want: location:string;", got)
}
}
Markdown is supported
0% or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment