update message processing

template/template.go

@@ -5,6 +5,7 @@ import (
 	"embed"
 	"encoding/json"
 	"errors"
+	"fmt"
 	"io"
 	"math"
 	"slices"
@@ -14,6 +15,7 @@ import (
 	"text/template/parse"
 
 	"github.com/agnivade/levenshtein"
+	"github.com/ollama/ollama/api"
 	"golang.org/x/exp/maps"
 )
 
@@ -74,30 +76,78 @@ func Named(s string) (*named, error) {
 	return nil, errors.New("no matching template found")
 }
 
+var DefaultTemplate, _ = Parse("{{ .Prompt }}")
+
 type Template struct {
 	*template.Template
 	raw string
 }
 
+var response = parse.ActionNode{
+	NodeType: parse.NodeAction,
+	Pipe: &parse.PipeNode{
+		NodeType: parse.NodePipe,
+		Cmds: []*parse.CommandNode{
+			{
+				NodeType: parse.NodeCommand,
+				Args: []parse.Node{
+					&parse.FieldNode{
+						NodeType: parse.NodeField,
+						Ident:    []string{"Response"},
+					},
+				},
+			},
+		},
+	},
+}
+
+func Parse(s string) (*Template, error) {
+	tmpl := template.New("").Option("missingkey=zero").Funcs(template.FuncMap{
+		"toJson": func(v any) string {
+			b, err := json.Marshal(v)
+			if err != nil {
+				return ""
+			}
+
+			return string(b)
+		},
+		"isLastMessage": func(s []*api.Message, m *api.Message) bool {
+			for i := len(s) - 1; i >= 0; i-- {
+				if m.Role != s[i].Role {
+					continue
+				}
+
+				return m == s[i]
+			}
+
+			return false
+		},
+	})
+
+	tmpl, err := tmpl.Parse(s)
+	if err != nil {
+		return nil, err
+	}
+
+	t := Template{Template: tmpl, raw: s}
+	if vars := t.Vars(); !slices.Contains(vars, "messages") && !slices.Contains(vars, "response") {
+		// touch up the template and append {{ .Response }}
+		tmpl.Tree.Root.Nodes = append(tmpl.Tree.Root.Nodes, &response)
+	}
+
+	return &t, nil
+}
+
 func (t *Template) String() string {
 	return t.raw
 }
 
-var DefaultTemplate, _ = Parse("{{ .Prompt }}")
-
-func Parse(s string) (*Template, error) {
-	t, err := template.New("").Option("missingkey=zero").Parse(s)
-	if err != nil {
-		return nil, err
-	}
-
-	return &Template{Template: t, raw: s}, nil
-}
-
 func (t *Template) Vars() []string {
 	var vars []string
-	for _, n := range t.Tree.Root.Nodes {
-		vars = append(vars, parseNode(n)...)
+	for _, tt := range t.Templates() {
+		for _, n := range tt.Root.Nodes {
+			vars = append(vars, parseNode(n)...)
+		}
 	}
 
 	set := make(map[string]struct{})
@@ -110,6 +160,97 @@ func (t *Template) Vars() []string {
 	return vars
 }
 
+type Values struct {
+	Messages []api.Message
+}
+
+func (t *Template) Execute(w io.Writer, v Values) error {
+	system, collated := collate(v.Messages)
+	if slices.Contains(t.Vars(), "messages") {
+		return t.Template.Execute(w, map[string]any{
+			"System":   system,
+			"Messages": collated,
+		})
+	}
+
+	var b bytes.Buffer
+	var prompt, response string
+	for i, m := range collated {
+		if m.Role == "user" {
+			prompt = m.Content
+		} else {
+			response = m.Content
+		}
+
+		if i != len(collated)-1 && prompt != "" && response != "" {
+			if err := t.Template.Execute(&b, map[string]any{
+				"System":   "",
+				"Prompt":   prompt,
+				"Response": response,
+			}); err != nil {
+				return err
+			}
+
+			prompt = ""
+			response = ""
+		}
+	}
+
+	var cut bool
+	tree := t.Template.Copy()
+	// for the last message, cut everything after "{{ .Response }}"
+	tree.Root.Nodes = slices.DeleteFunc(tree.Root.Nodes, func(n parse.Node) bool {
+		if slices.Contains(parseNode(n), "Response") {
+			cut = true
+		}
+
+		return cut
+	})
+
+	if err := template.Must(template.New("").AddParseTree("", tree)).Execute(&b, map[string]any{
+		"System": system,
+		"Prompt": prompt,
+	}); err != nil {
+		return err
+	}
+
+	_, err := io.Copy(w, &b)
+	return err
+}
+
+func collate(msgs []api.Message) (system string, collated []*api.Message) {
+	var n int
+	for i := range msgs {
+		msg := msgs[i]
+		if msg.Role == "system" {
+			if system != "" {
+				system += "\n\n"
+			}
+
+			system += msg.Content
+			continue
+		}
+
+		for range msg.Images {
+			imageTag := fmt.Sprintf("[img-%d]", n)
+			if !strings.Contains(msg.Content, "[img]") {
+				msg.Content = strings.TrimSpace("[img] " + msg.Content)
+			}
+
+			msg.Content = strings.Replace(msg.Content, "[img]", imageTag, 1)
+			n++
+		}
+
+		if len(collated) > 0 && collated[len(collated)-1].Role == msg.Role {
+			collated[len(collated)-1].Content += "\n\n" + msg.Content
+		} else {
+			collated = append(collated, &msg)
+		}
+	}
+
+	return
+}
+
 func parseNode(n parse.Node) []string {
 	switch n := n.(type) {
 	case *parse.ActionNode:
@@ -152,6 +293,8 @@ func parseNode(n parse.Node) []string {
 		return names
 	case *parse.FieldNode:
 		return n.Ident
+	case *parse.TemplateNode:
+		return parseNode(n.Pipe)
 	}
 
 	return nil

template/template_test.go

@@ -11,6 +11,7 @@ import (
 	"testing"
 	"text/template"
 
+	"github.com/ollama/ollama/api"
 	"github.com/ollama/ollama/llm"
 )
 
@@ -64,13 +65,12 @@ func TestParse(t *testing.T) {
 		template string
 		vars     []string
 	}{
-		{"{{ .Prompt }}", []string{"prompt"}},
-		{"{{ .System }} {{ .Prompt }}", []string{"prompt", "system"}},
+		{"{{ .Prompt }}", []string{"prompt", "response"}},
+		{"{{ .System }} {{ .Prompt }}", []string{"prompt", "response", "system"}},
 		{"{{ .System }} {{ .Prompt }} {{ .Response }}", []string{"prompt", "response", "system"}},
-		{"{{ with .Tools }}{{ . }}{{ end }} {{ .System }} {{ .Prompt }}", []string{"prompt", "system", "tools"}},
+		{"{{ with .Tools }}{{ . }}{{ end }} {{ .System }} {{ .Prompt }}", []string{"prompt", "response", "system", "tools"}},
 		{"{{ range .Messages }}{{ .Role }} {{ .Content }}{{ end }}", []string{"content", "messages", "role"}},
 		{"{{ range .Messages }}{{ if eq .Role \"system\" }}SYSTEM: {{ .Content }}{{ else if eq .Role \"user\" }}USER: {{ .Content }}{{ else if eq .Role \"assistant\" }}ASSISTANT: {{ .Content }}{{ end }}{{ end }}", []string{"content", "messages", "role"}},
-		{"{{ .Prompt }} {{ .Suffix }}", []string{"prompt", "suffix"}},
 	}
 
 	for _, tt := range cases {
@@ -87,3 +87,148 @@ func TestParse(t *testing.T) {
 		})
 	}
 }
+
+func TestExecuteWithMessages(t *testing.T) {
+	cases := []struct {
+		templates []string
+		values    Values
+		expected  string
+	}{
+		{
+			[]string{
+				`[INST] {{ if .System }}{{ .System }}{{ print "\n\n" }}{{ end }}{{ .Prompt }}[/INST] `,
+				`[INST] {{ if .System }}{{ .System }}{{ print "\n\n" }}{{ end }}{{ .Prompt }}[/INST] {{ .Response }}`,
+				`{{- range .Messages }}
+{{- if eq .Role "user" }}[INST] {{ if and (isLastMessage $.Messages .) $.System }}{{ $.System }}{{ print "\n\n" }}
+{{- end }}{{ .Content }}[/INST] {{ else if eq .Role "assistant" }}{{ .Content }}
+{{- end }}
+{{- end }}`,
+			},
+			Values{
+				Messages: []api.Message{
+					{Role: "user", Content: "Hello friend!"},
+					{Role: "assistant", Content: "Hello human!"},
+					{Role: "user", Content: "Yay!"},
+				},
+			},
+			`[INST] Hello friend![/INST] Hello human![INST] Yay![/INST] `,
+		},
+		{
+			[]string{
+				`[INST] {{ if .System }}{{ .System }}{{ print "\n\n" }}{{ end }}{{ .Prompt }}[/INST] `,
+				`[INST] {{ if .System }}{{ .System }}{{ print "\n\n" }}{{ end }}{{ .Prompt }}[/INST] {{ .Response }}`,
+				`
+{{- range .Messages }}
+{{- if eq .Role "user" }}[INST] {{ if and (isLastMessage $.Messages .) $.System }}{{ $.System }}{{ print "\n\n" }}
+{{- end }}{{ .Content }}[/INST] {{ else if eq .Role "assistant" }}{{ .Content }}
+{{- end }}
+{{- end }}`,
+			},
+			Values{
+				Messages: []api.Message{
+					{Role: "system", Content: "You are a helpful assistant!"},
+					{Role: "user", Content: "Hello friend!"},
+					{Role: "assistant", Content: "Hello human!"},
+					{Role: "user", Content: "Yay!"},
+				},
+			},
+			`[INST] Hello friend![/INST] Hello human![INST] You are a helpful assistant!
+
+Yay![/INST] `,
+		},
+		{
+			[]string{
+				`{{ if .System }}<|im_start|>system
+{{ .System }}<|im_end|>
+{{ end }}{{ if .Prompt }}<|im_start|>user
+{{ .Prompt }}<|im_end|>
+{{ end }}<|im_start|>assistant
+{{ .Response }}<|im_end|>
+`,
+				`
+{{- range .Messages }}
+{{- if and (eq .Role "user") (isLastMessage $.Messages .) $.System }}<|im_start|>system
+{{ $.System }}<|im_end|>{{ print "\n" }}
+{{- end }}<|im_start|>{{ .Role }}
+{{ .Content }}<|im_end|>{{ print "\n" }}
+{{- end }}<|im_start|>assistant
+`,
+			},
+			Values{
+				Messages: []api.Message{
+					{Role: "system", Content: "You are a helpful assistant!"},
+					{Role: "user", Content: "Hello friend!"},
+					{Role: "assistant", Content: "Hello human!"},
+					{Role: "user", Content: "Yay!"},
+				},
+			},
+			`<|im_start|>user
+Hello friend!<|im_end|>
+<|im_start|>assistant
+Hello human!<|im_end|>
+<|im_start|>system
+You are a helpful assistant!<|im_end|>
+<|im_start|>user
+Yay!<|im_end|>
+<|im_start|>assistant
+`,
+		},
+		{
+			[]string{
+				`{{ if .Prompt }}Question: {{ .Prompt }}
+
+{{ end }}Answer: {{ .Response }}
+
+`,
+				`
+{{- range .Messages }}
+{{- if eq .Role "user" }}Question: {{ .Content }}{{ print "\n\n" }}
+{{- else if eq .Role "assistant" }}Answer: {{ .Content }}{{ print "\n\n" }}
+{{- end }}
+{{- end }}Answer: `,
+			},
+			Values{
+				Messages: []api.Message{
+					{Role: "user", Content: "What's in this image?", Images: []api.ImageData{[]byte("")}},
+					{Role: "assistant", Content: "It's a hot dog."},
+					{Role: "user", Content: "What's in _this_ image?"},
+					{Role: "user", Images: []api.ImageData{[]byte("")}},
+					{Role: "user", Content: "Is it a hot dog?"},
+				},
+			},
+			`Question: [img-0] What's in this image?
+
+Answer: It's a hot dog.
+
+Question: What's in _this_ image?
+
+[img-1]
+
+Is it a hot dog?
+
+Answer: `,
+		},
+	}
+
+	for _, tt := range cases {
+		t.Run("", func(t *testing.T) {
+			for _, tmpl := range tt.templates {
+				t.Run("", func(t *testing.T) {
+					tmpl, err := Parse(tmpl)
+					if err != nil {
+						t.Fatal(err)
+					}
+
+					var b bytes.Buffer
+					if err := tmpl.Execute(&b, tt.values); err != nil {
+						t.Fatal(err)
+					}
+
+					if b.String() != tt.expected {
+						t.Errorf("expected\n%s,\ngot\n%s", tt.expected, b.String())
+					}
+				})
+			}
+		})
+	}
+}