diff --git a/server/images.go b/server/images.go index 58fb87dccb..d6cceff4c8 100644 --- a/server/images.go +++ b/server/images.go @@ -26,6 +26,7 @@ import ( "github.com/ollama/ollama/fs/ggml" "github.com/ollama/ollama/parser" "github.com/ollama/ollama/template" + "github.com/ollama/ollama/thinking" "github.com/ollama/ollama/types/model" "github.com/ollama/ollama/version" ) @@ -113,7 +114,7 @@ func (m *Model) Capabilities() []model.Capability { } // Check for thinking capability - openingTag, closingTag := inferThinkingTags(m.Template.Template) + openingTag, closingTag := thinking.InferTags(m.Template.Template) if openingTag != "" && closingTag != "" { capabilities = append(capabilities, model.CapabilityThinking) } diff --git a/server/routes.go b/server/routes.go index d03ac2ece9..70cb6cef9d 100644 --- a/server/routes.go +++ b/server/routes.go @@ -37,6 +37,7 @@ import ( "github.com/ollama/ollama/server/internal/client/ollama" "github.com/ollama/ollama/server/internal/registry" "github.com/ollama/ollama/template" + "github.com/ollama/ollama/thinking" "github.com/ollama/ollama/tools" "github.com/ollama/ollama/types/errtypes" "github.com/ollama/ollama/types/model" @@ -282,10 +283,10 @@ func (s *Server) GenerateHandler(c *gin.Context) { prompt = b.String() } - var thinkingState *ThinkingParser - openingTag, closingTag := inferThinkingTags(m.Template.Template) + var thinkingState *thinking.Parser + openingTag, closingTag := thinking.InferTags(m.Template.Template) if req.Think != nil && *req.Think && openingTag != "" && closingTag != "" { - thinkingState = &ThinkingParser{ + thinkingState = &thinking.Parser{ OpeningTag: openingTag, ClosingTag: closingTag, } @@ -1522,10 +1523,10 @@ func (s *Server) ChatHandler(c *gin.Context) { return } - var thinkingState *ThinkingParser - openingTag, closingTag := inferThinkingTags(m.Template.Template) + var thinkingState *thinking.Parser + openingTag, closingTag := thinking.InferTags(m.Template.Template) if req.Think != nil && *req.Think && openingTag != "" && closingTag != "" { - thinkingState = &ThinkingParser{ + thinkingState = &thinking.Parser{ OpeningTag: openingTag, ClosingTag: closingTag, } @@ -1676,7 +1677,7 @@ func filterThinkTags(msgs []api.Message, m *Model) []api.Message { // change the user output), we should probably perform this filtering // for all thinking models (not just qwen3 & deepseek-r1) since it tends // to save tokens and improve quality. - thinkingState := &ThinkingParser{ + thinkingState := &thinking.Parser{ OpeningTag: "", ClosingTag: "", } diff --git a/server/thinking.go b/thinking/parser.go similarity index 59% rename from server/thinking.go rename to thinking/parser.go index 4ef3c1848a..a4d05e35a3 100644 --- a/server/thinking.go +++ b/thinking/parser.go @@ -1,9 +1,7 @@ -package server +package thinking import ( "strings" - "text/template" - "text/template/parse" "unicode" ) @@ -46,7 +44,7 @@ func (s thinkingState) String() string { } } -type ThinkingParser struct { +type Parser struct { state thinkingState OpeningTag string ClosingTag string @@ -56,7 +54,7 @@ type ThinkingParser struct { // AddContent returns the thinking content and the non-thinking content that // should be immediately sent to the user. It will internally buffer if it needs // to see more raw content to disambiguate -func (s *ThinkingParser) AddContent(content string) (string, string) { +func (s *Parser) AddContent(content string) (string, string) { s.acc.WriteString(content) var thinkingSb, remainingSb strings.Builder @@ -76,7 +74,7 @@ func (s *ThinkingParser) AddContent(content string) (string, string) { } // the additional bool return is true iff we should continue eating -func eat(s *ThinkingParser) (string, string, bool) { +func eat(s *Parser) (string, string, bool) { switch s.state { case thinkingState_LookingForOpening: trimmed := strings.TrimLeftFunc(s.acc.String(), unicode.IsSpace) @@ -171,130 +169,3 @@ func overlap(s, delim string) int { } return 0 } - -func templateVisit(n parse.Node, enterFn func(parse.Node) bool, exitFn func(parse.Node)) { - if n == nil { - return - } - shouldContinue := enterFn(n) - if !shouldContinue { - return - } - switch x := n.(type) { - case *parse.ListNode: - for _, c := range x.Nodes { - templateVisit(c, enterFn, exitFn) - } - case *parse.BranchNode: - if x.Pipe != nil { - templateVisit(x.Pipe, enterFn, exitFn) - } - if x.List != nil { - templateVisit(x.List, enterFn, exitFn) - } - if x.ElseList != nil { - templateVisit(x.ElseList, enterFn, exitFn) - } - case *parse.ActionNode: - templateVisit(x.Pipe, enterFn, exitFn) - case *parse.WithNode: - templateVisit(&x.BranchNode, enterFn, exitFn) - case *parse.RangeNode: - templateVisit(&x.BranchNode, enterFn, exitFn) - case *parse.IfNode: - templateVisit(&x.BranchNode, enterFn, exitFn) - case *parse.TemplateNode: - templateVisit(x.Pipe, enterFn, exitFn) - case *parse.PipeNode: - for _, c := range x.Cmds { - templateVisit(c, enterFn, exitFn) - } - case *parse.CommandNode: - for _, a := range x.Args { - templateVisit(a, enterFn, exitFn) - } - // text, field, number, etc. are leaves – nothing to recurse into - } - if exitFn != nil { - exitFn(n) - } -} - -// We use a heuristic to infer the tags that surround thinking traces: -// We look for a range node that iterates over "Messages" and then look for a -// reference to "Thinking" like `{{.Thinking}}`. We then go up to the nearest -// ListNode and take the first and last TextNodes as the opening and closing -// tags. -func inferThinkingTags(t *template.Template) (string, string) { - ancestors := []parse.Node{} - - openingTag := "" - closingTag := "" - - enterFn := func(n parse.Node) bool { - ancestors = append(ancestors, n) - - switch x := n.(type) { - case *parse.FieldNode: - if len(x.Ident) > 0 && x.Ident[0] == "Thinking" { - var mostRecentRange *parse.RangeNode - for i := len(ancestors) - 1; i >= 0; i-- { - if r, ok := ancestors[i].(*parse.RangeNode); ok { - mostRecentRange = r - break - } - } - if mostRecentRange == nil || !rangeUsesField(mostRecentRange, "Messages") { - return true - } - - // TODO(drifkin): to be more robust, check that it's in the action - // part, not the `if`'s pipeline part. We do match on the nearest list - // that starts and ends with text nodes, which makes this not strictly - // necessary for our heuristic - - // go up to the nearest ancestor that is a *parse.ListNode - for i := len(ancestors) - 1; i >= 0; i-- { - if l, ok := ancestors[i].(*parse.ListNode); ok { - firstNode := l.Nodes[0] - if t, ok := firstNode.(*parse.TextNode); ok { - openingTag = strings.TrimSpace(t.String()) - } - lastNode := l.Nodes[len(l.Nodes)-1] - if t, ok := lastNode.(*parse.TextNode); ok { - closingTag = strings.TrimSpace(t.String()) - } - - break - } - } - } - } - - return true - } - - exitFn := func(n parse.Node) { - ancestors = ancestors[:len(ancestors)-1] - } - - templateVisit(t.Root, enterFn, exitFn) - - return openingTag, closingTag -} - -// checks to see if the given field name is present in the pipeline of the given range node -func rangeUsesField(rangeNode *parse.RangeNode, field string) bool { - found := false - enterFn := func(n parse.Node) bool { - switch x := n.(type) { - case *parse.FieldNode: - if x.Ident[0] == field { - found = true - } - } - return true - } - templateVisit(rangeNode.BranchNode.Pipe, enterFn, nil) - return found -} diff --git a/server/thinking_test.go b/thinking/parser_test.go similarity index 66% rename from server/thinking_test.go rename to thinking/parser_test.go index 90d3f961db..78c297cd95 100644 --- a/server/thinking_test.go +++ b/thinking/parser_test.go @@ -1,8 +1,7 @@ -package server +package thinking import ( "testing" - "text/template" ) func TestExtractThinking(t *testing.T) { @@ -26,7 +25,7 @@ func TestExtractThinking(t *testing.T) { }, } for i, tt := range tests { - parser := ThinkingParser{ + parser := Parser{ OpeningTag: "", ClosingTag: "", } @@ -259,7 +258,7 @@ func TestThinkingStreaming(t *testing.T) { } for _, c := range cases { - parser := ThinkingParser{ + parser := Parser{ OpeningTag: "", ClosingTag: "", } @@ -277,127 +276,3 @@ func TestThinkingStreaming(t *testing.T) { } } } - -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 }} - {{ .Thinking }} - {{ end }} - {{ end }} - `, - wantOpeningTag: "", - wantClosingTag: "", - }, - { - 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 }} - {{ .Thinking }} - {{ 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 XML tags: - -{{- range .Tools }} -{"type": "function", "function": {{ .Function }}} -{{- end }} - - -For each function call, return a json object with function name and arguments within XML tags: - -{"name": , "arguments": } - -{{- 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 }} -{{ .Thinking }} -{{ end }} -{{ if .Content }}{{ .Content }} -{{- else if .ToolCalls }} -{{ range .ToolCalls }}{"name": "{{ .Function.Name }}", "arguments": {{ .Function.Arguments }}} -{{ end }} -{{- end }}{{ if not $last }}<|im_end|> -{{ end }} -{{- else if eq .Role "tool" }}<|im_start|>user - -{{ .Content }} -<|im_end|> -{{ end }} -{{- if and (ne .Role "assistant") $last }}<|im_start|>assistant -{{ end }} -{{- end }} - `, - wantOpeningTag: "", - wantClosingTag: "", - }, - } - for _, c := range cases { - tmpl := template.Must(template.New("test").Parse(c.tmplString)) - openingTag, closingTag := inferThinkingTags(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) - } - } -} diff --git a/thinking/template.go b/thinking/template.go new file mode 100644 index 0000000000..20bd65ec12 --- /dev/null +++ b/thinking/template.go @@ -0,0 +1,134 @@ +package thinking + +import ( + "strings" + "text/template" + "text/template/parse" +) + +func templateVisit(n parse.Node, enterFn func(parse.Node) bool, exitFn func(parse.Node)) { + if n == nil { + return + } + shouldContinue := enterFn(n) + if !shouldContinue { + return + } + switch x := n.(type) { + case *parse.ListNode: + for _, c := range x.Nodes { + templateVisit(c, enterFn, exitFn) + } + case *parse.BranchNode: + if x.Pipe != nil { + templateVisit(x.Pipe, enterFn, exitFn) + } + if x.List != nil { + templateVisit(x.List, enterFn, exitFn) + } + if x.ElseList != nil { + templateVisit(x.ElseList, enterFn, exitFn) + } + case *parse.ActionNode: + templateVisit(x.Pipe, enterFn, exitFn) + case *parse.WithNode: + templateVisit(&x.BranchNode, enterFn, exitFn) + case *parse.RangeNode: + templateVisit(&x.BranchNode, enterFn, exitFn) + case *parse.IfNode: + templateVisit(&x.BranchNode, enterFn, exitFn) + case *parse.TemplateNode: + templateVisit(x.Pipe, enterFn, exitFn) + case *parse.PipeNode: + for _, c := range x.Cmds { + templateVisit(c, enterFn, exitFn) + } + case *parse.CommandNode: + for _, a := range x.Args { + templateVisit(a, enterFn, exitFn) + } + // text, field, number, etc. are leaves – nothing to recurse into + } + if exitFn != nil { + exitFn(n) + } +} + +// InferTags uses a heuristic to infer the tags that surround thinking traces: +// We look for a range node that iterates over "Messages" and then look for a +// reference to "Thinking" like `{{.Thinking}}`. We then go up to the nearest +// ListNode and take the first and last TextNodes as the opening and closing +// tags. +func InferTags(t *template.Template) (string, string) { + ancestors := []parse.Node{} + + openingTag := "" + closingTag := "" + + enterFn := func(n parse.Node) bool { + ancestors = append(ancestors, n) + + switch x := n.(type) { + case *parse.FieldNode: + if len(x.Ident) > 0 && x.Ident[0] == "Thinking" { + var mostRecentRange *parse.RangeNode + for i := len(ancestors) - 1; i >= 0; i-- { + if r, ok := ancestors[i].(*parse.RangeNode); ok { + mostRecentRange = r + break + } + } + if mostRecentRange == nil || !rangeUsesField(mostRecentRange, "Messages") { + return true + } + + // TODO(drifkin): to be more robust, check that it's in the action + // part, not the `if`'s pipeline part. We do match on the nearest list + // that starts and ends with text nodes, which makes this not strictly + // necessary for our heuristic + + // go up to the nearest ancestor that is a *parse.ListNode + for i := len(ancestors) - 1; i >= 0; i-- { + if l, ok := ancestors[i].(*parse.ListNode); ok { + firstNode := l.Nodes[0] + if t, ok := firstNode.(*parse.TextNode); ok { + openingTag = strings.TrimSpace(t.String()) + } + lastNode := l.Nodes[len(l.Nodes)-1] + if t, ok := lastNode.(*parse.TextNode); ok { + closingTag = strings.TrimSpace(t.String()) + } + + break + } + } + } + } + + return true + } + + exitFn := func(n parse.Node) { + ancestors = ancestors[:len(ancestors)-1] + } + + templateVisit(t.Root, enterFn, exitFn) + + return openingTag, closingTag +} + +// checks to see if the given field name is present in the pipeline of the given range node +func rangeUsesField(rangeNode *parse.RangeNode, field string) bool { + found := false + enterFn := func(n parse.Node) bool { + switch x := n.(type) { + case *parse.FieldNode: + if x.Ident[0] == field { + found = true + } + } + return true + } + templateVisit(rangeNode.BranchNode.Pipe, enterFn, nil) + return found +} diff --git a/thinking/template_test.go b/thinking/template_test.go new file mode 100644 index 0000000000..e63558e281 --- /dev/null +++ b/thinking/template_test.go @@ -0,0 +1,130 @@ +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 }} + {{ .Thinking }} + {{ end }} + {{ end }} + `, + wantOpeningTag: "", + wantClosingTag: "", + }, + { + 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 }} + {{ .Thinking }} + {{ 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 XML tags: + +{{- range .Tools }} +{"type": "function", "function": {{ .Function }}} +{{- end }} + + +For each function call, return a json object with function name and arguments within XML tags: + +{"name": , "arguments": } + +{{- 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 }} +{{ .Thinking }} +{{ end }} +{{ if .Content }}{{ .Content }} +{{- else if .ToolCalls }} +{{ range .ToolCalls }}{"name": "{{ .Function.Name }}", "arguments": {{ .Function.Arguments }}} +{{ end }} +{{- end }}{{ if not $last }}<|im_end|> +{{ end }} +{{- else if eq .Role "tool" }}<|im_start|>user + +{{ .Content }} +<|im_end|> +{{ end }} +{{- if and (ne .Role "assistant") $last }}<|im_start|>assistant +{{ end }} +{{- end }} + `, + wantOpeningTag: "", + wantClosingTag: "", + }, + } + 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) + } + } +}