Merge branch 'main' into drifkin/array-head-count-simple

thinking/parser_test.go

+package thinking
+
+import (
+	"testing"
+)
+
+func TestExtractThinking(t *testing.T) {
+	tests := []struct {
+		in, wantContent, wantThink string
+	}{
+		{
+			in:          "<think> internal </think> world",
+			wantThink:   "internal ",
+			wantContent: "world",
+		},
+		{
+			in:          "<think>a</think><think>b</think>c",
+			wantThink:   "a",
+			wantContent: "<think>b</think>c",
+		},
+		{
+			in:          "no think",
+			wantThink:   "",
+			wantContent: "no think",
+		},
+	}
+	for i, tt := range tests {
+		parser := Parser{
+			OpeningTag: "<think>",
+			ClosingTag: "</think>",
+		}
+		gotThinking, gotContent := parser.AddContent(tt.in)
+		if gotContent != tt.wantContent || gotThinking != tt.wantThink {
+			t.Errorf("case %d: got (%q,%q), want (%q,%q)", i, gotThinking, gotContent, tt.wantThink, tt.wantContent)
+		}
+	}
+}
+
+func TestThinkingStreaming(t *testing.T) {
+	type step struct {
+		input          string
+		wantThinking   string
+		wantContent    string
+		wantStateAfter thinkingState
+	}
+
+	cases := []struct {
+		desc  string
+		skip  bool
+		steps []step
+	}{
+		{
+			desc: "content without a thinking tag",
+			steps: []step{
+				{
+					input:          "  abc",
+					wantThinking:   "",
+					wantContent:    "  abc",
+					wantStateAfter: thinkingState_ThinkingDone,
+				},
+			},
+		},
+		{
+			desc: "content before a thinking tag nerfs the thinking tag",
+			steps: []step{
+				{
+					input:          "  abc <think>def</think> ghi",
+					wantThinking:   "",
+					wantContent:    "  abc <think>def</think> ghi",
+					wantStateAfter: thinkingState_ThinkingDone,
+				},
+			},
+		},
+		{
+			desc: "building up a thinking tag partially",
+			steps: []step{
+				{
+					input:          "  <th",
+					wantThinking:   "",
+					wantContent:    "",
+					wantStateAfter: thinkingState_LookingForOpening,
+				},
+				{
+					input:          "in",
+					wantThinking:   "",
+					wantContent:    "",
+					wantStateAfter: thinkingState_LookingForOpening,
+				},
+				{
+					input:          "k>a",
+					wantThinking:   "a",
+					wantContent:    "",
+					wantStateAfter: thinkingState_Thinking,
+				},
+			},
+		},
+		{
+			desc: "partial closing tag",
+			steps: []step{
+				{
+					input:          "<think>abc</th",
+					wantThinking:   "abc",
+					wantContent:    "",
+					wantStateAfter: thinkingState_Thinking,
+				},
+				{
+					input:          "ink>def",
+					wantThinking:   "",
+					wantContent:    "def",
+					wantStateAfter: thinkingState_ThinkingDone,
+				},
+			},
+		},
+		{
+			desc: "partial closing tag fakeout",
+			steps: []step{
+				{
+					input:          "<think>abc</th",
+					wantThinking:   "abc",
+					wantContent:    "",
+					wantStateAfter: thinkingState_Thinking,
+				},
+				{
+					input:          "ing>def",
+					wantThinking:   "</thing>def",
+					wantContent:    "",
+					wantStateAfter: thinkingState_Thinking,
+				},
+				{
+					input:          "ghi</thi",
+					wantThinking:   "ghi",
+					wantContent:    "",
+					wantStateAfter: thinkingState_Thinking,
+				},
+				{
+					input:          "nk>jkl",
+					wantThinking:   "",
+					wantContent:    "jkl",
+					wantStateAfter: thinkingState_ThinkingDone,
+				},
+			},
+		},
+		{
+			desc: "whitespace after thinking tag",
+			steps: []step{
+				{
+					input:          "  <think>abc</think>\n\ndef",
+					wantThinking:   "abc",
+					wantContent:    "def",
+					wantStateAfter: thinkingState_ThinkingDone,
+				},
+			},
+		},
+		{
+			desc: "whitespace after thinking tag (incremental)",
+			steps: []step{
+				{
+					input:          "  <think>abc</think>",
+					wantThinking:   "abc",
+					wantContent:    "",
+					wantStateAfter: thinkingState_ThinkingDoneEatingWhitespace,
+				},
+				{
+					input:          "\n\ndef",
+					wantThinking:   "",
+					wantContent:    "def",
+					wantStateAfter: thinkingState_ThinkingDone,
+				},
+			},
+		},
+		{
+			desc: "whitespace after thinking tag with content and more whitespace",
+			steps: []step{
+				{
+					input:          "  <think>abc</think>\n\ndef ",
+					wantThinking:   "abc",
+					wantContent:    "def ",
+					wantStateAfter: thinkingState_ThinkingDone,
+				},
+				{
+					input:          " ghi",
+					wantThinking:   "",
+					wantContent:    " ghi",
+					wantStateAfter: thinkingState_ThinkingDone,
+				},
+			},
+		},
+		{
+			desc: "token by token",
+			steps: []step{
+				{
+					input:          "<think>",
+					wantThinking:   "",
+					wantContent:    "",
+					wantStateAfter: thinkingState_ThinkingStartedEatingWhitespace,
+				},
+				{
+					input:          "\n",
+					wantThinking:   "",
+					wantContent:    "",
+					wantStateAfter: thinkingState_ThinkingStartedEatingWhitespace,
+				},
+				{
+					input:          "</think>",
+					wantThinking:   "",
+					wantContent:    "",
+					wantStateAfter: thinkingState_ThinkingDoneEatingWhitespace,
+				},
+				{
+					input:          "\n\n",
+					wantThinking:   "",
+					wantContent:    "",
+					wantStateAfter: thinkingState_ThinkingDoneEatingWhitespace,
+				},
+				{
+					input:          "Hi",
+					wantThinking:   "",
+					wantContent:    "Hi",
+					wantStateAfter: thinkingState_ThinkingDone,
+				},
+				{
+					input:          " there",
+					wantThinking:   "",
+					wantContent:    " there",
+					wantStateAfter: thinkingState_ThinkingDone,
+				},
+			},
+		},
+		{
+			desc: "leading thinking whitespace",
+			steps: []step{
+				{
+					input:          "  <think>   \t ",
+					wantThinking:   "",
+					wantContent:    "",
+					wantStateAfter: thinkingState_ThinkingStartedEatingWhitespace,
+				},
+				{
+					input:          "  these are some ",
+					wantThinking:   "these are some ",
+					wantContent:    "",
+					wantStateAfter: thinkingState_Thinking,
+				},
+				{
+					input:          "thoughts </think>  ",
+					wantThinking:   "thoughts ",
+					wantContent:    "",
+					wantStateAfter: thinkingState_ThinkingDoneEatingWhitespace,
+				},
+				{
+					input:          "  more content",
+					wantThinking:   "",
+					wantContent:    "more content",
+					wantStateAfter: thinkingState_ThinkingDone,
+				},
+			},
+		},
+	}
+
+	for _, c := range cases {
+		parser := Parser{
+			OpeningTag: "<think>",
+			ClosingTag: "</think>",
+		}
+		if c.skip {
+			continue
+		}
+		for i, step := range c.steps {
+			thinking, content := parser.AddContent(step.input)
+			if content != step.wantContent || thinking != step.wantThinking {
+				t.Errorf("case %q (step %d): got (%q,%q), want (%q,%q)", c.desc, i, content, thinking, step.wantContent, step.wantThinking)
+			}
+			if parser.state != step.wantStateAfter {
+				t.Errorf("case %q (step %d): got state %s, want %s", c.desc, i, parser.state, step.wantStateAfter)
+			}
+		}
+	}
+}

thinking/template.go

+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
+}

thinking/template_test.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)
+		}
+	}
+}

tools/template.go

+package tools
+
+import (
+	"bytes"
+	"log/slog"
+	"slices"
+	"strings"
+	"text/template"
+	"text/template/parse"
+)
+
+// parseTag finds the tool calling tag from a Go template
+// often <tool_call> [TOOL_CALL] or similar by finding the
+// first text node after .ToolCalls and returning the content
+// if no tag is found, return "{" to indicate that json objects
+// should be attempted to be parsed as tool calls
+func parseTag(tmpl *template.Template) string {
+	if tmpl == nil || tmpl.Tree == nil {
+		slog.Debug("template or tree is nil")
+		return "{"
+	}
+
+	tc := findToolCallNode(tmpl.Tree.Root.Nodes)
+	if tc == nil {
+		return "{"
+	}
+
+	tn := findTextNode(tc.List.Nodes)
+	if tn == nil {
+		return "{"
+	}
+
+	tag := string(tn.Text)
+	tag = strings.ReplaceAll(tag, "\r\n", "\n")
+
+	// avoid parsing { onwards as this may be a tool call
+	// however keep '{' as a prefix if there is no tag
+	// so that all json objects will be attempted to
+	// be parsed as tool calls
+	tag, _, _ = strings.Cut(tag, "{")
+	tag = strings.TrimSpace(tag)
+	if tag == "" {
+		tag = "{"
+	}
+
+	return tag
+}
+
+// findToolCallNode searches for and returns an IfNode with .ToolCalls
+func findToolCallNode(nodes []parse.Node) *parse.IfNode {
+	isToolCallsNode := func(n *parse.IfNode) bool {
+		for _, cmd := range n.Pipe.Cmds {
+			for _, arg := range cmd.Args {
+				if field, ok := arg.(*parse.FieldNode); ok {
+					if slices.Contains(field.Ident, "ToolCalls") {
+						return true
+					}
+				}
+			}
+		}
+		return false
+	}
+
+	for _, node := range nodes {
+		switch n := node.(type) {
+		case *parse.IfNode:
+			if isToolCallsNode(n) {
+				return n
+			}
+			// Recursively search in nested IfNodes
+			if result := findToolCallNode(n.List.Nodes); result != nil {
+				return result
+			}
+			if n.ElseList != nil {
+				if result := findToolCallNode(n.ElseList.Nodes); result != nil {
+					return result
+				}
+			}
+		case *parse.ListNode:
+			if result := findToolCallNode(n.Nodes); result != nil {
+				return result
+			}
+		case *parse.RangeNode:
+			if result := findToolCallNode(n.List.Nodes); result != nil {
+				return result
+			}
+			if n.ElseList != nil {
+				if result := findToolCallNode(n.ElseList.Nodes); result != nil {
+					return result
+				}
+			}
+		case *parse.WithNode:
+			if result := findToolCallNode(n.List.Nodes); result != nil {
+				return result
+			}
+			if n.ElseList != nil {
+				if result := findToolCallNode(n.ElseList.Nodes); result != nil {
+					return result
+				}
+			}
+		}
+	}
+	return nil
+}
+
+// findTextNode does a depth-first search for the first text content in nodes,
+// stopping at template constructs to avoid parsing text after the tool calls
+func findTextNode(nodes []parse.Node) *parse.TextNode {
+	for _, node := range nodes {
+		switch n := node.(type) {
+		case *parse.TextNode:
+			// skip whitespace-only text nodes
+			if len(bytes.TrimSpace(n.Text)) == 0 {
+				continue
+			}
+			return n
+		case *parse.IfNode:
+			if text := findTextNode(n.List.Nodes); text != nil {
+				return text
+			}
+			if n.ElseList != nil {
+				if text := findTextNode(n.ElseList.Nodes); text != nil {
+					return text
+				}
+			}
+			return nil
+		case *parse.ListNode:
+			if text := findTextNode(n.Nodes); text != nil {
+				return text
+			}
+		case *parse.RangeNode:
+			if text := findTextNode(n.List.Nodes); text != nil {
+				return text
+			}
+			if n.ElseList != nil {
+				if text := findTextNode(n.ElseList.Nodes); text != nil {
+					return text
+				}
+			}
+			return nil
+		case *parse.WithNode:
+			if text := findTextNode(n.List.Nodes); text != nil {
+				return text
+			}
+			if n.ElseList != nil {
+				if text := findTextNode(n.ElseList.Nodes); text != nil {
+					return text
+				}
+			}
+			return nil
+		case *parse.ActionNode:
+			return nil
+		}
+	}
+	return nil
+}

tools/template_test.go

+package tools
+
+import (
+	"testing"
+	"text/template"
+)
+
+func TestParseTag(t *testing.T) {
+	cases := []struct {
+		name     string
+		template string
+		want     string
+	}{
+		{
+			name:     "empty",
+			template: "",
+			want:     "{",
+		},
+		{
+			name:     "no tag",
+			template: "{{if .ToolCalls}}{{end}}",
+			want:     "{",
+		},
+		{
+			name:     "no tag with range",
+			template: "{{if .ToolCalls}}{{range .ToolCalls}}{{ . }}{{end}}{{end}}",
+			want:     "{",
+		},
+		{
+			name:     "tool call with json format",
+			template: "{{if .ToolCalls}}```json\n{{end}}",
+			want:     "```json",
+		},
+		{
+			name:     "square brackets",
+			template: "{{if .ToolCalls}}[{{range .ToolCalls}}{{ . }}{{end}}]{{end}}",
+			want:     "[",
+		},
+		{
+			name:     "square brackets with whitespace",
+			template: "{{if .ToolCalls}}\n [ {{range .ToolCalls}}{{ . }}{{end}}]{{end}}",
+			want:     "[",
+		},
+		{
+			name:     "tailing ]",
+			template: "{{if .ToolCalls}}{{range .ToolCalls}}{{ . }}{{end}}]{{end}}",
+			want:     "{",
+		},
+		{
+			name:     "whitespace only",
+			template: "{{if .ToolCalls}} {{range .ToolCalls}}{{ . }}{{end}}{{end}}",
+			want:     "{",
+		},
+		{
+			name:     "whitespace only in range",
+			template: "{{if .ToolCalls}}{{range .ToolCalls}}\n{{ . }}\n{{end}}{{end}}",
+			want:     "{",
+		},
+		{
+			name:     "json objects",
+			template: `{{if .ToolCalls}}{{range .ToolCalls}}{"name": "{{ .Function.Name }}", "arguments": {{ .Function.Arguments }}}{{end}}{{end}}`,
+			want:     "{",
+		},
+		{
+			name:     "json objects with whitespace",
+			template: "{{if .ToolCalls}}{{range .ToolCalls}}\n{\"name\": \"{{ .Function.Name }}\", \"arguments\": {{ .Function.Arguments }}}{{end}}{{end}}",
+			want:     "{",
+		},
+		{
+			name:     "json objects with CRLF",
+			template: "{{if .ToolCalls}}{{range .ToolCalls}}\r\n{\"name\": \"{{ .Function.Name }}\", \"arguments\": {{ .Function.Arguments }}}{{end}}{{end}}",
+			want:     "{",
+		},
+		{
+			name:     "json objects with whitespace before and after range",
+			template: "{{if .ToolCalls}}\n{{range .ToolCalls}}\n{\"name\": \"{{ .Function.Name }}\", \"arguments\": {{ .Function.Arguments }}}\r\n{{end}}\r\n{{end}}",
+			want:     "{",
+		},
+		{
+			name:     "before and after range",
+			template: "{{if .ToolCalls}}<|tool▁calls▁begin|>{{range .ToolCalls}}<|tool▁call▁begin|>functionget_current_weather\n```json\n{\"location\": \"Tokyo\"}\n```<|tool▁call▁end|>\n{{end}}<|tool▁calls▁end|>{{end}}",
+			want:     "<|tool▁calls▁begin|>",
+		},
+		{
+			name:     "after range",
+			template: "{{if .ToolCalls}}{{range .ToolCalls}}<tool_call>{\"name\": \"{{ .Function.Name }}\", \"arguments\": {{ .Function.Arguments }}}</tool_call>{{end}}{{end}}",
+			want:     "<tool_call>",
+		},
+		{
+			name:     "after range with leading whitespace before range",
+			template: "{{if .ToolCalls}}\n{{range .ToolCalls}}<tool_call>{\"name\": \"{{ .Function.Name }}\", \"arguments\": {{ .Function.Arguments }}}</tool_call>{{end}}{{end}}",
+			want:     "<tool_call>",
+		},
+		{
+			name:     "tool call in range with {",
+			template: `{{if .ToolCalls}}{{range .ToolCalls}}<tool_call>{"name": "{{ .Function.Name }}", "arguments": {{ .Function.Arguments }}}<tool_call>{{end}}{{end}}`,
+			want:     "<tool_call>",
+		},
+		{
+			name:     "tool call with multiple text nodes",
+			template: "{{if .ToolCalls}}First text{{if .Something}}inner{{end}}Second text{{end}}",
+			want:     "First text",
+		},
+		{
+			name:     "action tag",
+			template: "{{if .ToolCalls}}Action: ```json{{end}}",
+			want:     "Action: ```json",
+		},
+		{
+			name:     "incomplete functools bracket",
+			template: "{{if .ToolCalls}}functools[{{end}}",
+			want:     "functools[",
+		},
+		{
+			name:     "uppercase tool call with incomplete bracket",
+			template: "{{if .ToolCalls}}[TOOL_CALL] [{{end}}",
+			want:     "[TOOL_CALL] [",
+		},
+		{
+			name:     "uppercase tool call with adjacent bracket",
+			template: "{{if .ToolCalls}}[TOOL_CALL][{{end}}",
+			want:     "[TOOL_CALL][",
+		},
+	}
+
+	for _, tc := range cases {
+		t.Run(tc.name, func(t *testing.T) {
+			tmpl, err := template.New("test").Parse(tc.template)
+			if err != nil && tc.template != "" {
+				t.Fatalf("failed to parse template: %v", err)
+			}
+
+			got := parseTag(tmpl)
+			if got != tc.want {
+				t.Errorf("got text %q, want %q", got, tc.want)
+			}
+		})
+	}
+}

tools/tools.go

+package tools
+
+import (
+	"bytes"
+	"encoding/json"
+	"strings"
+	"text/template"
+
+	"github.com/ollama/ollama/api"
+)
+
+type toolsState int
+
+const (
+	toolsState_LookingForTag toolsState = iota
+	toolsState_ToolCalling
+	toolsState_Done
+)
+
+type Parser struct {
+	tag   string
+	tools []api.Tool
+
+	state  toolsState
+	buffer []byte
+	n      int
+}
+
+// NewParser creates a new tool call parser from a model's chat
+// template and a list of provided tools.
+func NewParser(tmpl *template.Template, tools []api.Tool) *Parser {
+	return NewParserWithTag(tools, parseTag(tmpl))
+}
+
+func NewParserWithTag(tools []api.Tool, tag string) *Parser {
+	return &Parser{
+		tag:   tag,
+		tools: tools,
+	}
+}
+
+// Add processes a string input to parse tool calls and content that
+// should be sent back to the user.
+func (p *Parser) Add(s string) (calls []api.ToolCall, content string) {
+	if p.state == toolsState_Done {
+		return nil, s
+	}
+
+	p.buffer = append(p.buffer, s...)
+
+	if p.state == toolsState_LookingForTag {
+		i, found := p.findTag()
+		if i == -1 {
+			content = string(p.buffer)
+			p.buffer = []byte{}
+		} else {
+			content = string(p.buffer[:i])
+			p.buffer = p.buffer[i:]
+		}
+
+		// for models where { or [ are used as tool calling
+		// tags, we only support parsing tools if the first non-
+		// whitespace character is { or [
+		if p.tag == "{" || p.tag == "[" {
+			if strings.TrimSpace(content) != "" {
+				p.state = toolsState_Done
+				return nil, content + string(p.buffer)
+			}
+		}
+
+		if !found {
+			return nil, content
+		}
+
+		p.state = toolsState_ToolCalling
+	}
+
+	for {
+		call := p.parseToolCall()
+		if call == nil {
+			break
+		}
+
+		calls = append(calls, *call)
+	}
+
+	if p.done() {
+		p.state = toolsState_Done
+		content = string(p.buffer)
+		p.buffer = []byte{}
+	}
+
+	return calls, content
+}
+
+// findTag searches the buffer to find and handle a tool calling tag
+// returning true if the tag was found and false otherwise, and
+// a string content signaling any content that should be sent back to the user
+func (p *Parser) findTag() (int, bool) {
+	// First check for complete substring anywhere in s
+	if i := bytes.Index(p.buffer, []byte(p.tag)); i > -1 {
+		return i, true
+	}
+
+	// Then check for partial suffix overlap
+	max := min(len(p.buffer), len(p.tag))
+	for i := max; i > 0; i-- {
+		if bytes.HasSuffix(p.buffer, []byte(p.tag[:i])) {
+			return len(p.buffer) - i, false
+		}
+	}
+	return -1, false
+}
+
+// parseToolCall finds the next complete tool call in the buffer
+// incrementing n and advancing the buffer.
+func (p *Parser) parseToolCall() *api.ToolCall {
+	var tool *api.Tool
+	var end int = len(p.buffer)
+	var i int
+
+	// find tool name
+	for _, t := range p.tools {
+		n := t.Function.Name
+		if i = bytes.Index(p.buffer, []byte(n)); i != -1 {
+			if i+len(n) < end {
+				tool = &t
+				end = i + len(n)
+			}
+		}
+	}
+
+	if tool == nil {
+		return nil
+	}
+
+	// only look for arguments if the tool has parameters
+	args := map[string]any{}
+	if len(tool.Function.Parameters.Properties) > 0 {
+		if args, i = p.findArguments(*tool); args == nil {
+			return nil
+		}
+
+		if i > end {
+			end = i
+		}
+	}
+
+	tc := &api.ToolCall{
+		Function: api.ToolCallFunction{
+			Name:      tool.Function.Name,
+			Arguments: args,
+			Index:     p.n,
+		},
+	}
+
+	p.n++
+	p.buffer = p.buffer[end:]
+	return tc
+}
+
+// findArguments returns the first object that appears to be
+// arguments for the provided tool, returning nil
+func (p *Parser) findArguments(tool api.Tool) (map[string]any, int) {
+	if len(p.buffer) == 0 {
+		return nil, 0
+	}
+
+	// no arguments to parse
+	if len(tool.Function.Parameters.Properties) == 0 {
+		return nil, 0
+	}
+
+	var braces int
+	var start int = -1
+	var end int
+	var object []byte
+
+	// find any outer json object
+	for i, c := range p.buffer {
+		if c == '{' {
+			braces++
+			if start == -1 {
+				start = i
+			}
+		}
+
+		if c == '}' {
+			if start != -1 {
+				braces--
+				if braces == 0 {
+					end = i + 1
+					object = p.buffer[start:end]
+					break
+				}
+			}
+		}
+	}
+
+	if braces > 0 {
+		return nil, 0
+	}
+
+	var data map[string]any
+
+	// not valid json
+	if err := json.Unmarshal(object, &data); err != nil {
+		return nil, 0
+	}
+
+	var find func(obj any) map[string]any
+	find = func(obj any) map[string]any {
+		switch obj := obj.(type) {
+		case map[string]any:
+			found := true
+			for key := range obj {
+				if _, exists := tool.Function.Parameters.Properties[key]; !exists {
+					found = false
+					break
+				}
+			}
+
+			if found {
+				return obj
+			}
+
+			for _, value := range obj {
+				if result := find(value); result != nil {
+					return result
+				}
+			}
+		case []any:
+			for _, item := range obj {
+				if result := find(item); result != nil {
+					return result
+				}
+			}
+		}
+
+		return nil
+	}
+
+	result := find(data)
+	if result != nil {
+		return result, end
+	}
+
+	return nil, 0
+}
+
+// done checks if the parser is done parsing by looking
+// for closing tag. currently only } and ] are supported
+// for closing tags as {} or [] pairs may not always
+// represent tool calls and we need to send the content back
+func (p *Parser) done() bool {
+	var open, close rune
+	switch p.tag {
+	case "{":
+		open, close = '{', '}'
+	case "[":
+		open, close = '[', ']'
+	default:
+		return false
+	}
+
+	var count int
+	for _, c := range p.buffer {
+		if c == byte(open) {
+			count++
+		} else if c == byte(close) {
+			count--
+			if count == 0 {
+				return true
+			}
+		}
+	}
+
+	return false
+}
+
+// Content returns any remaining content that
+// should be sent to the user. This should be the empty string
+// string unless the tag is { or [ and a tool call was not found
+func (p *Parser) Content() string {
+	if p.n > 0 {
+		return ""
+	}
+
+	if p.tag == "{" || p.tag == "[" {
+		return string(p.buffer)
+	}
+
+	return ""
+}

tools/tools_test.go

+package tools
+
+import (
+	"testing"
+	"text/template"
+
+	"github.com/google/go-cmp/cmp"
+	"github.com/ollama/ollama/api"
+)
+
+func TestParser(t *testing.T) {
+	qwen, err := template.New("qwen").Parse(`{{if .ToolCalls}}<tool_call>{{range .ToolCalls}}{"name": "{{.Function.Name}}", "arguments": {{.Function.Arguments}}}{{end}}</tool_call>{{end}}`)
+	if err != nil {
+		t.Fatalf("Failed to parse template: %v", err)
+	}
+
+	deepseek, err := template.New("deepseek").Parse("{{if .ToolCalls}}<|tool▁calls▁begin|>{{range .ToolCalls}}<|tool▁call▁begin|>function<|tool▁sep|>get_current_weather\n```json\n{\"location\": \"Tokyo\"}\n```<|tool▁call▁end|>{{end}}<|tool▁calls▁end|><|end▁of▁sentence|>{{end}}")
+	if err != nil {
+		t.Fatalf("Failed to parse template: %v", err)
+	}
+
+	json, err := template.New("json").Parse(`{{if .ToolCalls}}{{range .ToolCalls}}{"name": "{{.Function.Name}}", "arguments": {{.Function.Arguments}}}{{end}}{{end}}`)
+	if err != nil {
+		t.Fatalf("Failed to parse template: %v", err)
+	}
+
+	mistral, err := template.New("mistral").Parse(`{{if .ToolCalls}}[TOOL_CALLS] [{{range .ToolCalls}}{"name": "{{.Function.Name}}", "arguments": {{.Function.Arguments}}}{{end}}][/TOOL_CALLS]{{end}}`)
+	if err != nil {
+		t.Fatalf("Failed to parse template: %v", err)
+	}
+
+	list, err := template.New("list").Parse(`{{if .ToolCalls}}[{{range .ToolCalls}}{"name": "{{.Function.Name}}", "arguments": {{.Function.Arguments}}}{{end}}]{{end}}`)
+	if err != nil {
+		t.Fatalf("Failed to parse template: %v", err)
+	}
+
+	tools := []api.Tool{
+		{
+			Type: "function",
+			Function: api.ToolFunction{
+				Name:        "get_temperature",
+				Description: "Retrieve the temperature for a given location",
+				Parameters: struct {
+					Type       string   `json:"type"`
+					Defs       any      `json:"$defs,omitempty"`
+					Items      any      `json:"items,omitempty"`
+					Required   []string `json:"required"`
+					Properties map[string]struct {
+						Type        api.PropertyType `json:"type"`
+						Items       any              `json:"items,omitempty"`
+						Description string           `json:"description"`
+						Enum        []any            `json:"enum,omitempty"`
+					} `json:"properties"`
+				}{
+					Type: "object",
+					Properties: map[string]struct {
+						Type        api.PropertyType `json:"type"`
+						Items       any              `json:"items,omitempty"`
+						Description string           `json:"description"`
+						Enum        []any            `json:"enum,omitempty"`
+					}{
+						"format": {
+							Type:        api.PropertyType{"string"},
+							Description: "The format to return the temperature in",
+							Enum:        []any{"fahrenheit", "celsius"},
+						},
+						"city": {
+							Type:        api.PropertyType{"string"},
+							Description: "The city to get the temperature for",
+						},
+					},
+				},
+			},
+		},
+		{
+			Type: "function",
+			Function: api.ToolFunction{
+				Name:        "get_conditions",
+				Description: "Retrieve the current weather conditions for a given location",
+				Parameters: struct {
+					Type       string   `json:"type"`
+					Defs       any      `json:"$defs,omitempty"`
+					Items      any      `json:"items,omitempty"`
+					Required   []string `json:"required"`
+					Properties map[string]struct {
+						Type        api.PropertyType `json:"type"`
+						Items       any              `json:"items,omitempty"`
+						Description string           `json:"description"`
+						Enum        []any            `json:"enum,omitempty"`
+					} `json:"properties"`
+				}{
+					Type: "object",
+					Properties: map[string]struct {
+						Type        api.PropertyType `json:"type"`
+						Items       any              `json:"items,omitempty"`
+						Description string           `json:"description"`
+						Enum        []any            `json:"enum,omitempty"`
+					}{
+						"location": {
+							Type:        api.PropertyType{"string"},
+							Description: "The location to get the weather conditions for",
+						},
+					},
+				},
+			},
+		},
+		{
+			Type: "function",
+			Function: api.ToolFunction{
+				Name:        "say_hello",
+				Description: "Say hello",
+			},
+		},
+	}
+
+	tests := []struct {
+		name    string
+		inputs  []string
+		tmpl    *template.Template
+		content string
+		calls   []api.ToolCall
+	}{
+		{
+			name:    "no tool calls - just text",
+			inputs:  []string{"Hello, how can I help you today?"},
+			content: "Hello, how can I help you today?",
+			tmpl:    qwen,
+			calls:   nil,
+		},
+		{
+			name:    "empty input",
+			inputs:  []string{""},
+			content: "",
+			tmpl:    qwen,
+			calls:   nil,
+		},
+		{
+			name:    "tool call",
+			inputs:  []string{`<tool_call>{"name": "get_conditions", "arguments": {"location": "San Francisco"}}</tool_call>`},
+			content: "",
+			tmpl:    qwen,
+			calls: []api.ToolCall{
+				{
+					Function: api.ToolCallFunction{
+						Index: 0,
+						Name:  "get_conditions",
+						Arguments: api.ToolCallFunctionArguments{
+							"location": "San Francisco",
+						},
+					},
+				},
+			},
+		},
+		{
+			name:    "invalid arguments",
+			inputs:  []string{`<tool_call>{"name": "get_conditions", "arguments": {"city": "San Francisco"}}</tool_call>`},
+			content: "",
+			tmpl:    qwen,
+			calls:   nil,
+		},
+		{
+			name:    "missing args",
+			inputs:  []string{`<tool_call>{"name": "get_conditions"}</tool_call>`},
+			content: "",
+			tmpl:    qwen,
+			calls:   nil,
+		},
+		{
+			name:    "text before tool call",
+			inputs:  []string{`Let me check the weather. <tool_call>{"name": "get_temperature", "arguments": {"city": "New York"}}</tool_call>`},
+			content: "Let me check the weather. ",
+			tmpl:    qwen,
+			calls: []api.ToolCall{
+				{
+					Function: api.ToolCallFunction{
+						Index: 0,
+						Name:  "get_temperature",
+						Arguments: api.ToolCallFunctionArguments{
+							"city": "New York",
+						},
+					},
+				},
+			},
+		},
+		{
+			name:    "qwen no args tool call",
+			inputs:  []string{`Let me say hello to the user. I'll use the say_hello tool <tool_call>{"name": "say_hello"}</tool_call>`},
+			content: "Let me say hello to the user. I'll use the say_hello tool ",
+			tmpl:    qwen,
+			calls: []api.ToolCall{
+				{
+					Function: api.ToolCallFunction{
+						Index:     0,
+						Name:      "say_hello",
+						Arguments: api.ToolCallFunctionArguments{},
+					},
+				},
+			},
+		},
+		{
+			name:    "qwen no args with text",
+			inputs:  []string{"Let me say hello to the user. I'll use the say_hello tool. "},
+			content: "Let me say hello to the user. I'll use the say_hello tool. ",
+			tmpl:    qwen,
+			calls:   nil,
+		},
+		{
+			name:    "two tool calls in a list",
+			inputs:  []string{`[TOOL_CALLS] [{"name": "get_temperature", "arguments": {"city": "London", "format": "fahrenheit"}}, {"name": "get_conditions", "arguments": {"location": "Tokyo"}}][/TOOL_CALLS]`},
+			content: "",
+			tmpl:    mistral,
+			calls: []api.ToolCall{
+				{
+					Function: api.ToolCallFunction{
+						Index: 0,
+						Name:  "get_temperature",
+						Arguments: api.ToolCallFunctionArguments{
+							"city":   "London",
+							"format": "fahrenheit",
+						},
+					},
+				},
+				{
+					Function: api.ToolCallFunction{
+						Index: 1,
+						Name:  "get_conditions",
+						Arguments: api.ToolCallFunctionArguments{
+							"location": "Tokyo",
+						},
+					},
+				},
+			},
+		},
+		{
+			name:    "qwen two tool calls",
+			inputs:  []string{`Okay, let's call both tools! <tool_call>{"name": "get_temperature", "arguments": {"city": "London", "format": "fahrenheit"}}</tool_call><tool_call>{"name": "get_conditions", "arguments": {"location": "Tokyo"}}</tool_call>`},
+			content: "Okay, let's call both tools! ",
+			tmpl:    qwen,
+			calls: []api.ToolCall{
+				{
+					Function: api.ToolCallFunction{
+						Index: 0,
+						Name:  "get_temperature",
+						Arguments: api.ToolCallFunctionArguments{
+							"city":   "London",
+							"format": "fahrenheit",
+						},
+					},
+				},
+				{
+					Function: api.ToolCallFunction{
+						Index: 1,
+						Name:  "get_conditions",
+						Arguments: api.ToolCallFunctionArguments{
+							"location": "Tokyo",
+						},
+					},
+				},
+			},
+		},
+		{
+			name:    "qwen two tool calls one with no args",
+			inputs:  []string{`Let me check the weather. <tool_call>{"name": "say_hello"}</tool_call><tool_call>{"name": "get_conditions", "arguments": {"location": "Tokyo"}}`},
+			content: "Let me check the weather. ",
+			tmpl:    qwen,
+			calls: []api.ToolCall{
+				{
+					Function: api.ToolCallFunction{
+						Index:     0,
+						Name:      "say_hello",
+						Arguments: api.ToolCallFunctionArguments{},
+					},
+				},
+				{
+					Function: api.ToolCallFunction{
+						Index: 1,
+						Name:  "get_conditions",
+						Arguments: api.ToolCallFunctionArguments{
+							"location": "Tokyo",
+						},
+					},
+				},
+			},
+		},
+		{
+			name:    "deepseek",
+			inputs:  []string{"<think>Wait, I need to call a tool</think><|tool▁calls▁begin|><|tool▁call▁begin|>function<|tool▁sep|>get_temperature\n```json\n{\"city\": \"Tokyo\"}\n```<|tool▁call▁end|><|tool▁calls▁end|><|end▁of▁sentence|>"},
+			content: "<think>Wait, I need to call a tool</think>",
+			tmpl:    deepseek,
+			calls: []api.ToolCall{
+				{
+					Function: api.ToolCallFunction{
+						Index: 0,
+						Name:  "get_temperature",
+						Arguments: api.ToolCallFunctionArguments{
+							"city": "Tokyo",
+						},
+					},
+				},
+			},
+		},
+		{
+			name: "deepseek incremental",
+			inputs: []string{
+				"<think>Wait",
+				", I need",
+				" to call",
+				" a tool</think><|too",
+				"l▁calls▁begin",
+				"|>",
+				"<|tool▁call▁begin|>function<|tool▁sep|>get_temperature\n",
+				"```json\n",
+				"{\"city\": \"Tokyo\"}\n",
+				"```",
+				"<|tool▁c", "all▁end|>",
+				"<|tool▁calls▁end|>",
+				"<|end▁of▁sentence|>",
+			},
+			content: "<think>Wait, I need to call a tool</think>",
+			tmpl:    deepseek,
+			calls: []api.ToolCall{
+				{
+					Function: api.ToolCallFunction{
+						Index: 0,
+						Name:  "get_temperature",
+						Arguments: api.ToolCallFunctionArguments{
+							"city": "Tokyo",
+						},
+					},
+				},
+			},
+		},
+		{
+			name: "json",
+			inputs: []string{
+				"{",
+				"\"name\": \"get_temperature\",",
+				"\"arguments\": {",
+				"\"city\": \"Tokyo\"",
+				"}",
+				"}",
+			},
+			content: "",
+			tmpl:    json,
+			calls: []api.ToolCall{
+				{
+					Function: api.ToolCallFunction{
+						Index: 0,
+						Name:  "get_temperature",
+						Arguments: api.ToolCallFunctionArguments{
+							"city": "Tokyo",
+						},
+					},
+				},
+			},
+		},
+		{
+			name: "json maybe a tool call",
+			inputs: []string{
+				"{",
+				"\"name\": \"get_temperature\",",
+				"\"arguments\": {",
+			},
+			content: "",
+			tmpl:    json,
+			calls:   nil,
+		},
+		{
+			name: "json not a tool call",
+			inputs: []string{
+				"{",
+				"\"name\": \"search\", ",
+				"\"arguments\": {",
+				"\"query\": \"What is the capital of Canada?\"",
+				"}",
+				"}",
+			},
+			content: "{\"name\": \"search\", \"arguments\": {\"query\": \"What is the capital of Canada?\"}}",
+			tmpl:    json,
+			calls:   nil,
+		},
+		{
+			name: "json object followed by tool call",
+			inputs: []string{
+				"{\"name\": \"jeff\"}",
+				"{\"name\": \"get_conditions\", \"arguments\": {\"location\": \"San Francisco\"}}",
+			},
+			content: "{\"name\": \"jeff\"}{\"name\": \"get_conditions\", \"arguments\": {\"location\": \"San Francisco\"}}",
+			tmpl:    json,
+		},
+		{
+			name: "json object followed by tool call split",
+			inputs: []string{
+				"{\"name\": \"jeff\"} {",
+				"\"name\": \"get_conditions\", \"arguments\": {\"location\": \"San Francisco\"}}",
+			},
+			content: "{\"name\": \"jeff\"} {\"name\": \"get_conditions\", \"arguments\": {\"location\": \"San Francisco\"}}",
+			tmpl:    json,
+		},
+		{
+			name: "json code",
+			inputs: []string{
+				"for { fmt.Println(\"hello\") }",
+			},
+			content: "for { fmt.Println(\"hello\") }",
+			tmpl:    json,
+		},
+		{
+			name: "json no args tool call",
+			inputs: []string{
+				"{\"name\": \"say_hello\"}",
+			},
+			content: "",
+			tmpl:    json,
+			calls: []api.ToolCall{
+				{
+					Function: api.ToolCallFunction{
+						Index:     0,
+						Name:      "say_hello",
+						Arguments: api.ToolCallFunctionArguments{},
+					},
+				},
+			},
+		},
+		{
+			name: "json no args no tool call",
+			inputs: []string{
+				"I'll use the say_hello tool to say hello to the user.",
+			},
+			content: "I'll use the say_hello tool to say hello to the user.",
+			tmpl:    json,
+			calls:   nil,
+		},
+
+		// TODO (jmorganca): this is a false positive, we should
+		// not be parsing this as a tool call
+		{
+			name: "json no args false positive",
+			inputs: []string{
+				`{say_hello!!!}`,
+			},
+			content: "",
+			tmpl:    json,
+			calls: []api.ToolCall{
+				{
+					Function: api.ToolCallFunction{
+						Index:     0,
+						Name:      "say_hello",
+						Arguments: api.ToolCallFunctionArguments{},
+					},
+				},
+			},
+		},
+		{
+			name: "list multiple",
+			inputs: []string{
+				"[",
+				"{",
+				"\"name\": \"get_temperature\", ",
+				"\"arguments\": {",
+				"\"city\": \"London\"",
+				"}",
+				"},",
+				"{",
+				"\"name\": \"get_conditions\", ",
+				"\"arguments\": {",
+				"\"location\": \"Tokyo\"",
+				"}",
+				"}]",
+			},
+			content: "",
+			tmpl:    list,
+			calls: []api.ToolCall{
+				{
+					Function: api.ToolCallFunction{
+						Index: 0,
+						Name:  "get_temperature",
+						Arguments: api.ToolCallFunctionArguments{
+							"city": "London",
+						},
+					},
+				},
+				{
+					Function: api.ToolCallFunction{
+						Index: 1,
+						Name:  "get_conditions",
+						Arguments: api.ToolCallFunctionArguments{
+							"location": "Tokyo",
+						},
+					},
+				},
+			},
+		},
+		{
+			name: "list partial",
+			inputs: []string{
+				"[{",
+				"\"name\": \"get_conditions\", ",
+				"\"arguments\": {",
+				"\"location\": \"Tokyo\"",
+				"}",
+				"}",
+			},
+			content: "",
+			tmpl:    list,
+			calls: []api.ToolCall{
+				{
+					Function: api.ToolCallFunction{
+						Index: 0,
+						Name:  "get_conditions",
+						Arguments: api.ToolCallFunctionArguments{
+							"location": "Tokyo",
+						},
+					},
+				},
+			},
+		},
+		{
+			name: "list invalid",
+			inputs: []string{
+				"[",
+				"{",
+				"\"name\": \"search\", ",
+				"\"arguments\": {",
+				"\"query\": \"What is the capital of Canada?\"",
+				"}",
+				"}",
+			},
+			content: "",
+			tmpl:    list,
+			calls:   nil,
+		},
+		{
+			name: "list trailing ]",
+			inputs: []string{
+				"[",
+				"{",
+				"\"name\": \"get_conditions\", ",
+				"\"arguments\": {",
+				"\"location\": \"Tokyo\"",
+				"}",
+				"}",
+				"]",
+				"]",
+			},
+			content: "",
+			tmpl:    list,
+			calls: []api.ToolCall{
+				{
+					Function: api.ToolCallFunction{
+						Index: 0,
+						Name:  "get_conditions",
+						Arguments: api.ToolCallFunctionArguments{
+							"location": "Tokyo",
+						},
+					},
+				},
+			},
+		},
+		{
+			name: "list not a tool call",
+			inputs: []string{
+				"[special",
+				" del",
+				"ivery]",
+			},
+			content: "[special delivery]",
+			tmpl:    list,
+			calls:   nil,
+		},
+		{
+			name: "list with no arguments",
+			inputs: []string{
+				"[",
+				"{",
+				"\"name\": \"say_hello\"",
+				"}",
+			},
+			content: "",
+			tmpl:    list,
+			calls: []api.ToolCall{
+				{
+					Function: api.ToolCallFunction{
+						Index:     0,
+						Name:      "say_hello",
+						Arguments: api.ToolCallFunctionArguments{},
+					},
+				},
+			},
+		},
+	}
+
+	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
+			parser := NewParser(tt.tmpl, tools)
+
+			var calls []api.ToolCall
+			var content string
+			for _, input := range tt.inputs {
+				tcs, c := parser.Add(input)
+				calls = append(calls, tcs...)
+				content += c
+			}
+
+			if content != tt.content {
+				t.Errorf("Expected content %q, got %q", tt.content, content)
+			}
+
+			if len(calls) != len(tt.calls) {
+				t.Fatalf("Expected %d tool calls, got %d", len(tt.calls), len(calls))
+			}
+
+			for i, want := range tt.calls {
+				if diff := cmp.Diff(calls[i], want); diff != "" {
+					t.Errorf("Tool call %d mismatch (-got +want):\n%s", i, diff)
+				}
+			}
+		})
+	}
+}
+
+func TestDone(t *testing.T) {
+	tests := []struct {
+		name   string
+		tag    string
+		buffer []byte
+		want   bool
+	}{
+		{
+			name:   "empty",
+			tag:    "<tool_call>",
+			buffer: []byte{},
+			want:   false,
+		},
+		{
+			name:   "empty",
+			tag:    "<tool_call>",
+			buffer: []byte{},
+			want:   false,
+		},
+		{
+			name:   "json open",
+			tag:    "{",
+			buffer: []byte("{\"name\": \"get_weather\""),
+			want:   false,
+		},
+		{
+			name:   "json closed",
+			tag:    "{",
+			buffer: []byte("{\"name\": \"get_weather\"}"),
+			want:   true,
+		},
+		{
+			name:   "json empty",
+			tag:    "{",
+			buffer: []byte("{}"),
+			want:   true,
+		},
+		{
+			name:   "list open",
+			tag:    "[",
+			buffer: []byte("[{\"name\": \"get_weather\""),
+			want:   false,
+		},
+		{
+			name:   "list closed",
+			tag:    "[",
+			buffer: []byte("[{\"name\": \"get_weather\"}]"),
+			want:   true,
+		},
+		{
+			name:   "list empty",
+			tag:    "[",
+			buffer: []byte("[]"),
+			want:   true,
+		},
+	}
+
+	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
+			parser := &Parser{
+				tag:    tt.tag,
+				buffer: tt.buffer,
+			}
+			got := parser.done()
+			if got != tt.want {
+				t.Errorf("done() = %t, want %t", got, tt.want)
+			}
+		})
+	}
+}
+
+func TestContent(t *testing.T) {
+	tests := []struct {
+		name    string
+		tag     string
+		content []byte
+		want    string
+		n       int
+	}{
+		{
+			name:    "empty",
+			content: []byte{},
+			tag:     "{",
+			want:    "",
+			n:       0,
+		},
+		{
+			name:    "tag",
+			tag:     "<tool_call>",
+			content: []byte("<tool_call>{\"name\": \"get_temperature\""),
+			want:    "",
+			n:       0,
+		},
+		{
+			name:    "json object",
+			tag:     "{",
+			content: []byte("{\"name\": \"get_temperature\"}"),
+			want:    "{\"name\": \"get_temperature\"}",
+			n:       0,
+		},
+		{
+			name:    "json object after called",
+			tag:     "{",
+			content: []byte("{\"hello\": \"world\"}"),
+			want:    "{\"hello\": \"world\"}",
+			n:       0,
+		},
+		{
+			name:    "json object after called",
+			tag:     "{",
+			content: []byte("{\"hello\": \"world\"}"),
+			want:    "",
+			n:       1,
+		},
+		{
+			name:    "list",
+			tag:     "[",
+			content: []byte("[{\"name\": \"get_temperature\"}]"),
+			want:    "[{\"name\": \"get_temperature\"}]",
+			n:       0,
+		},
+		{
+			name:    "code",
+			tag:     "{",
+			content: []byte("{ fmt.Println(\"hello\")"),
+			want:    "{ fmt.Println(\"hello\")",
+			n:       0,
+		},
+	}
+
+	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
+			parser := &Parser{
+				tag:    tt.tag,
+				buffer: tt.content,
+				n:      tt.n,
+			}
+			got := parser.Content()
+			if got != tt.want {
+				t.Errorf("Content() = %q, want %q", got, tt.want)
+			}
+		})
+	}
+}
+
+func TestFindTag(t *testing.T) {
+	cases := []struct {
+		name   string
+		buffer []byte
+		tag    string
+		i      int
+		found  bool
+	}{
+		{
+			name:   "no overlap",
+			buffer: []byte("hello world"),
+			tag:    "<tool_call>",
+			i:      -1,
+			found:  false,
+		},
+		{
+			name:   "full overlap",
+			buffer: []byte("<tool_call>"),
+			tag:    "<tool_call>",
+			i:      0,
+			found:  true,
+		},
+		{
+			name:   "whitespace",
+			buffer: []byte("    <tool_call>\n {\"name\": \"bob\"}"),
+			tag:    "<tool_call>",
+			i:      4,
+			found:  true,
+		},
+		{
+			name:   "over",
+			buffer: []byte("<tool_call>{\"name\""),
+			tag:    "<tool_call>",
+			i:      0,
+			found:  true,
+		},
+		{
+			name:   "partial overlap",
+			buffer: []byte("text <tool_call>"),
+			tag:    "<tool_call>",
+			i:      5,
+			found:  true,
+		},
+		{
+			name:   "overlap with extra",
+			buffer: []byte("<tool_calls><tool_call>"),
+			tag:    "<tool_calls>",
+			i:      0,
+			found:  true,
+		},
+		{
+			name:   "delimiter longer than string",
+			buffer: []byte("<tool>"),
+			tag:    "<tool_call>",
+			i:      -1,
+			found:  false,
+		},
+		{
+			name:   "empty string",
+			buffer: []byte{},
+			tag:    "<tool_call>",
+			i:      -1,
+			found:  false,
+		},
+		{
+			name:   "single char overlap",
+			buffer: []byte("test<"),
+			tag:    "<tool_call>",
+			i:      4,
+			found:  false,
+		},
+		{
+			name:   "partial tool call",
+			buffer: []byte("hello <tool_"),
+			tag:    "<tool_call>",
+			i:      6,
+			found:  false,
+		},
+		{
+			name:   "square bracket",
+			buffer: []byte("calling tools: ["),
+			tag:    "[",
+			i:      15,
+			found:  true,
+		},
+		{
+			name:   "bracket",
+			buffer: []byte("{\"name\": \"bob\""),
+			tag:    "{",
+			i:      0,
+			found:  true,
+		},
+		{
+			name:   "bracket with whitespace",
+			buffer: []byte("\n\n{\n\"name\": \"bob\""),
+			tag:    "{",
+			i:      2,
+			found:  true,
+		},
+	}
+
+	for _, tt := range cases {
+		t.Run(tt.name, func(t *testing.T) {
+			parser := &Parser{
+				tag:    tt.tag,
+				buffer: tt.buffer,
+				n:      0,
+			}
+			i, found := parser.findTag()
+			if i != tt.i {
+				t.Errorf("findTag(%q, %q) = %d; want %d", tt.buffer, tt.tag, i, tt.i)
+			}
+			if found != tt.found {
+				t.Errorf("findTag(%q, %q) = %t; want %t", tt.buffer, tt.tag, found, tt.found)
+			}
+		})
+	}
+}
+
+func TestFindArguments(t *testing.T) {
+	tool := api.Tool{
+		Type: "function",
+		Function: api.ToolFunction{
+			Name:        "get_temperature",
+			Description: "Retrieve the temperature for a given location",
+			Parameters: struct {
+				Type       string   `json:"type"`
+				Defs       any      `json:"$defs,omitempty"`
+				Items      any      `json:"items,omitempty"`
+				Required   []string `json:"required"`
+				Properties map[string]struct {
+					Type        api.PropertyType `json:"type"`
+					Items       any              `json:"items,omitempty"`
+					Description string           `json:"description"`
+					Enum        []any            `json:"enum,omitempty"`
+				} `json:"properties"`
+			}{
+				Type: "object",
+				Properties: map[string]struct {
+					Type        api.PropertyType `json:"type"`
+					Items       any              `json:"items,omitempty"`
+					Description string           `json:"description"`
+					Enum        []any            `json:"enum,omitempty"`
+				}{
+					"format": {
+						Type:        api.PropertyType{"string"},
+						Description: "The format to return the temperature in",
+						Enum:        []any{"fahrenheit", "celsius"},
+					},
+					"location": {
+						Type:        api.PropertyType{"string"},
+						Description: "The location to get the temperature for",
+					},
+				},
+			},
+		},
+	}
+
+	tool2 := api.Tool{
+		Type: "function",
+		Function: api.ToolFunction{
+			Name:        "say_hello",
+			Description: "Say hello to the user",
+		},
+	}
+
+	tests := []struct {
+		name   string
+		buffer []byte
+		want   map[string]any
+		tool   api.Tool
+	}{
+		{
+			name:   "empty string",
+			buffer: []byte{},
+			want:   nil,
+			tool:   tool,
+		},
+		{
+			name:   "whitespace only",
+			buffer: []byte("   \n\t  "),
+			want:   nil,
+			tool:   tool,
+		},
+		{
+			name:   "unbalanced braces - missing closing",
+			buffer: []byte(`{"format": "fahrenheit", "location": "San Francisco"`),
+			want:   nil,
+			tool:   tool,
+		},
+		{
+			name:   "unbalanced braces - extra closing",
+			buffer: []byte(`{"format": "fahrenheit"}}`),
+			want: map[string]any{
+				"format": "fahrenheit",
+			},
+			tool: tool,
+		},
+		{
+			name:   "invalid JSON",
+			buffer: []byte(`{format: fahrenheit, location: "San Francisco"}`),
+			want:   nil,
+			tool:   tool,
+		},
+		{
+			name:   "valid json",
+			buffer: []byte(`{"name": "get_temperature", "arguments": {"format": "fahrenheit", "location": "San Francisco, CA"}}`),
+			want: map[string]any{
+				"format":   "fahrenheit",
+				"location": "San Francisco, CA",
+			},
+			tool: tool,
+		},
+		{
+			name:   "valid arguments with special tokens",
+			buffer: []byte(`[tool]get_temperature[args]{"format": "fahrenheit", "location": "San Francisco, CA"}[end]`),
+			want: map[string]any{
+				"format":   "fahrenheit",
+				"location": "San Francisco, CA",
+			},
+			tool: tool,
+		},
+		{
+			name:   "valid arguments in array",
+			buffer: []byte(`[{"arguments": {"format": "fahrenheit", "location": "San Francisco, CA"}}`),
+			want: map[string]any{
+				"format":   "fahrenheit",
+				"location": "San Francisco, CA",
+			},
+			tool: tool,
+		},
+		{
+			name:   "nested deep",
+			buffer: []byte(`{"function": {"name": "get_temperature", "arguments": {"format": "fahrenheit", "location": "San Francisco, CA"}}}`),
+			want: map[string]any{
+				"format":   "fahrenheit",
+				"location": "San Francisco, CA",
+			},
+			tool: tool,
+		},
+		{
+			name:   "one arg",
+			buffer: []byte(`get_temperature({"location": "San Francisco, CA"})`),
+			want: map[string]any{
+				"location": "San Francisco, CA",
+			},
+			tool: tool,
+		},
+		{
+			name:   "two args",
+			buffer: []byte(`[{"name": "get_temperature", "arguments": {"location": "San Francisco, CA", "format": "fahrenheit"}}, {"name": "get_weather", "arguments": {"location": "San Francisco, CA", "format": "fahrenheit"}}]`),
+			want: map[string]any{
+				"location": "San Francisco, CA",
+				"format":   "fahrenheit",
+			},
+			tool: tool,
+		},
+		{
+			name:   "no args",
+			buffer: []byte(`{"name": "say_hello"}`),
+			want:   nil,
+			tool:   tool2,
+		},
+		{
+			name:   "deepseek",
+			buffer: []byte("<|tool▁calls▁begin|><|tool▁call▁begin|>function<|tool▁sep|>get_temperature\n```json\n{\"location\": \"Tokyo\"}\n```<|tool▁call▁end|><|tool▁calls▁end|><|end▁of▁sentence|>"),
+			want: map[string]any{
+				"location": "Tokyo",
+			},
+			tool: tool,
+		},
+	}
+
+	for _, tt := range tests {
+		parser := &Parser{
+			buffer: tt.buffer,
+			tools:  []api.Tool{tool, tool2},
+		}
+
+		t.Run(tt.name, func(t *testing.T) {
+			got, _ := parser.findArguments(tool)
+
+			if diff := cmp.Diff(got, tt.want); diff != "" {
+				t.Errorf("scanArguments() args mismatch (-got +want):\n%s", diff)
+			}
+		})
+	}
+}

types/model/capability.go

@@ -8,6 +8,7 @@ const (
 	CapabilityInsert     = Capability("insert")
 	CapabilityVision     = Capability("vision")
 	CapabilityEmbedding  = Capability("embedding")
+	CapabilityThinking   = Capability("thinking")
 )
 
 func (c Capability) String() string {