mirror of
https://github.com/ollama/ollama.git
synced 2025-09-12 20:03:49 +02:00
131 lines
3.1 KiB
Go
131 lines
3.1 KiB
Go
package thinking
|
|
|
|
import (
|
|
"testing"
|
|
"text/template"
|
|
)
|
|
|
|
func TestInferThinkingTags(t *testing.T) {
|
|
cases := []struct {
|
|
desc string
|
|
tmplString string
|
|
wantOpeningTag string
|
|
wantClosingTag string
|
|
}{
|
|
{
|
|
desc: "basic",
|
|
tmplString: `
|
|
{{ if .Thinking}}
|
|
/think
|
|
{{ end }}
|
|
{{- range $i, $_ := .Messages }}
|
|
{{- $last := eq (len (slice $.Messages $i)) 1 -}}
|
|
{{ if and $last .Thinking }}
|
|
<think>{{ .Thinking }}</think>
|
|
{{ end }}
|
|
{{ end }}
|
|
`,
|
|
wantOpeningTag: "<think>",
|
|
wantClosingTag: "</think>",
|
|
},
|
|
{
|
|
desc: "doubly nested range",
|
|
tmplString: `
|
|
{{ if .Thinking}}
|
|
/think
|
|
{{ end }}
|
|
{{- range $i, $_ := .Messages }}
|
|
{{- range $j, $_ := .NotMessages }}
|
|
{{- $last := eq (len (slice $.Messages $i)) 1 -}}
|
|
{{ if and $last .Thinking }}
|
|
<think>{{ .Thinking }}</think>
|
|
{{ end }}
|
|
{{ end }}
|
|
{{ end }}
|
|
`,
|
|
wantOpeningTag: "",
|
|
wantClosingTag: "",
|
|
},
|
|
{
|
|
desc: "whitespace is trimmed",
|
|
tmplString: `
|
|
{{ if .Thinking}}
|
|
/think
|
|
{{ end }}
|
|
{{- range $i, $_ := .Messages }}
|
|
{{- $last := eq (len (slice $.Messages $i)) 1 -}}
|
|
{{ if and $last .Thinking }}
|
|
Some text before {{ .Thinking }} Some text after
|
|
{{ end }}
|
|
{{ end }}
|
|
`,
|
|
wantOpeningTag: "Some text before",
|
|
wantClosingTag: "Some text after",
|
|
},
|
|
{
|
|
desc: "qwen3",
|
|
tmplString: `
|
|
{{- if or .System .Tools .Thinking }}<|im_start|>system
|
|
{{- if .System }}
|
|
{{ .System }}
|
|
{{- end }}
|
|
{{- if .Tools }}
|
|
|
|
# Tools
|
|
|
|
You may call one or more functions to assist with the user query.
|
|
|
|
You are provided with function signatures within <tools></tools> XML tags:
|
|
<tools>
|
|
{{- range .Tools }}
|
|
{"type": "function", "function": {{ .Function }}}
|
|
{{- end }}
|
|
</tools>
|
|
|
|
For each function call, return a json object with function name and arguments within <tool_call></tool_call> XML tags:
|
|
<tool_call>
|
|
{"name": <function-name>, "arguments": <args-json-object>}
|
|
</tool_call>
|
|
{{- end }}
|
|
{{- if .Thinking }}
|
|
/think
|
|
{{- else }}
|
|
/no_think
|
|
{{- end }}<|im_end|>
|
|
{{ end }}
|
|
{{- range $i, $_ := .Messages }}
|
|
{{- $last := eq (len (slice $.Messages $i)) 1 -}}
|
|
{{- if eq .Role "user" }}<|im_start|>user
|
|
{{ .Content }}<|im_end|>
|
|
{{ else if eq .Role "assistant" }}<|im_start|>assistant
|
|
{{ if and $last .Thinking }}
|
|
<think>{{ .Thinking }}</think>
|
|
{{ end }}
|
|
{{ if .Content }}{{ .Content }}
|
|
{{- else if .ToolCalls }}<tool_call>
|
|
{{ range .ToolCalls }}{"name": "{{ .Function.Name }}", "arguments": {{ .Function.Arguments }}}
|
|
{{ end }}</tool_call>
|
|
{{- end }}{{ if not $last }}<|im_end|>
|
|
{{ end }}
|
|
{{- else if eq .Role "tool" }}<|im_start|>user
|
|
<tool_response>
|
|
{{ .Content }}
|
|
</tool_response><|im_end|>
|
|
{{ end }}
|
|
{{- if and (ne .Role "assistant") $last }}<|im_start|>assistant
|
|
{{ end }}
|
|
{{- end }}
|
|
`,
|
|
wantOpeningTag: "<think>",
|
|
wantClosingTag: "</think>",
|
|
},
|
|
}
|
|
for _, c := range cases {
|
|
tmpl := template.Must(template.New("test").Parse(c.tmplString))
|
|
openingTag, closingTag := InferTags(tmpl)
|
|
if openingTag != c.wantOpeningTag || closingTag != c.wantClosingTag {
|
|
t.Errorf("case %q: got (%q,%q), want (%q,%q)", c.desc, openingTag, closingTag, c.wantOpeningTag, c.wantClosingTag)
|
|
}
|
|
}
|
|
}
|