template.go 8.63 KB
Newer Older
Michael Yang's avatar
Michael Yang committed
1
2
3
4
5
6
7
package template

import (
	"bytes"
	"embed"
	"encoding/json"
	"errors"
Michael Yang's avatar
Michael Yang committed
8
	"fmt"
Michael Yang's avatar
Michael Yang committed
9
10
11
12
13
14
15
16
17
	"io"
	"math"
	"slices"
	"strings"
	"sync"
	"text/template"
	"text/template/parse"

	"github.com/agnivade/levenshtein"
Michael Yang's avatar
Michael Yang committed
18
	"github.com/ollama/ollama/api"
Michael Yang's avatar
Michael Yang committed
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
	"golang.org/x/exp/maps"
)

//go:embed index.json
var indexBytes []byte

//go:embed *.gotmpl
var templatesFS embed.FS

var templatesOnce = sync.OnceValues(func() ([]*named, error) {
	var templates []*named
	if err := json.Unmarshal(indexBytes, &templates); err != nil {
		return nil, err
	}

	for _, t := range templates {
		bts, err := templatesFS.ReadFile(t.Name + ".gotmpl")
		if err != nil {
			return nil, err
		}

		// normalize line endings
		t.Bytes = bytes.ReplaceAll(bts, []byte("\r\n"), []byte("\n"))
	}

	return templates, nil
})

type named struct {
	Name     string `json:"name"`
	Template string `json:"template"`
	Bytes    []byte
}

func (t named) Reader() io.Reader {
	return bytes.NewReader(t.Bytes)
}

func Named(s string) (*named, error) {
	templates, err := templatesOnce()
	if err != nil {
		return nil, err
	}

	var template *named
	score := math.MaxInt
	for _, t := range templates {
		if s := levenshtein.ComputeDistance(s, t.Template); s < score {
			score = s
			template = t
		}
	}

	if score < 100 {
		return template, nil
	}

	return nil, errors.New("no matching template found")
}

Michael Yang's avatar
Michael Yang committed
79
80
var DefaultTemplate, _ = Parse("{{ .Prompt }}")

Michael Yang's avatar
Michael Yang committed
81
82
83
84
85
type Template struct {
	*template.Template
	raw string
}

Michael Yang's avatar
Michael Yang committed
86
// response is a template node that can be added to templates that don't already have one
Michael Yang's avatar
Michael Yang committed
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
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"},
					},
				},
			},
		},
	},
Michael Yang's avatar
Michael Yang committed
103
104
}

Michael Yang's avatar
tools  
Michael Yang committed
105
106
107
108
109
110
111
var funcs = template.FuncMap{
	"json": func(v any) string {
		b, _ := json.Marshal(v)
		return string(b)
	},
}

Michael Yang's avatar
Michael Yang committed
112
func Parse(s string) (*Template, error) {
Michael Yang's avatar
tools  
Michael Yang committed
113
	tmpl := template.New("").Option("missingkey=zero").Funcs(funcs)
Michael Yang's avatar
Michael Yang committed
114
115

	tmpl, err := tmpl.Parse(s)
Michael Yang's avatar
Michael Yang committed
116
117
118
119
	if err != nil {
		return nil, err
	}

Michael Yang's avatar
Michael Yang committed
120
121
122
123
124
125
126
127
128
129
130
	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
Michael Yang's avatar
Michael Yang committed
131
132
133
134
}

func (t *Template) Vars() []string {
	var vars []string
Michael Yang's avatar
Michael Yang committed
135
136
	for _, tt := range t.Templates() {
		for _, n := range tt.Root.Nodes {
Michael Yang's avatar
tools  
Michael Yang committed
137
			vars = append(vars, Identifiers(n)...)
Michael Yang's avatar
Michael Yang committed
138
		}
Michael Yang's avatar
Michael Yang committed
139
140
141
142
143
144
145
146
147
148
149
150
	}

	set := make(map[string]struct{})
	for _, n := range vars {
		set[strings.ToLower(n)] = struct{}{}
	}

	vars = maps.Keys(set)
	slices.Sort(vars)
	return vars
}

Michael Yang's avatar
Michael Yang committed
151
152
type Values struct {
	Messages []api.Message
Michael Yang's avatar
tools  
Michael Yang committed
153
	Tools    []api.Tool
154
155
156

	// forceLegacy is a flag used to test compatibility with legacy templates
	forceLegacy bool
Michael Yang's avatar
Michael Yang committed
157
158
}

Michael Yang's avatar
tools  
Michael Yang committed
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
func (t *Template) Subtree(fn func(parse.Node) bool) *template.Template {
	var walk func(parse.Node) parse.Node
	walk = func(n parse.Node) parse.Node {
		if fn(n) {
			return n
		}

		switch t := n.(type) {
		case *parse.ListNode:
			for _, c := range t.Nodes {
				if n := walk(c); n != nil {
					return n
				}
			}
		case *parse.BranchNode:
			for _, n := range []*parse.ListNode{t.List, t.ElseList} {
				if n != nil {
					if n := walk(n); n != nil {
						return n
					}
				}
			}
		case *parse.IfNode:
			return walk(&t.BranchNode)
		case *parse.WithNode:
			return walk(&t.BranchNode)
		case *parse.RangeNode:
			return walk(&t.BranchNode)
		}

		return nil
	}

	if n := walk(t.Tree.Root); n != nil {
		return (&template.Template{
			Tree: &parse.Tree{
				Root: &parse.ListNode{
					Nodes: []parse.Node{n},
				},
			},
		}).Funcs(funcs)
	}

	return nil
}

Michael Yang's avatar
Michael Yang committed
205
func (t *Template) Execute(w io.Writer, v Values) error {
Michael Yang's avatar
Michael Yang committed
206
	system, messages := collate(v.Messages)
207
	if !v.forceLegacy && slices.Contains(t.Vars(), "messages") {
Michael Yang's avatar
Michael Yang committed
208
		return t.Template.Execute(w, map[string]any{
209
			"System":   system,
Michael Yang's avatar
Michael Yang committed
210
			"Messages": messages,
Michael Yang's avatar
tools  
Michael Yang committed
211
			"Tools":    v.Tools,
Michael Yang's avatar
Michael Yang committed
212
213
214
		})
	}

Michael Yang's avatar
Michael Yang committed
215
	system = ""
Michael Yang's avatar
Michael Yang committed
216
	var b bytes.Buffer
217
	var prompt, response string
Michael Yang's avatar
Michael Yang committed
218
	for _, m := range messages {
Michael Yang's avatar
tools  
Michael Yang committed
219
		execute := func() error {
Michael Yang's avatar
Michael Yang committed
220
			if err := t.Template.Execute(&b, map[string]any{
221
				"System":   system,
Michael Yang's avatar
Michael Yang committed
222
223
224
225
226
227
				"Prompt":   prompt,
				"Response": response,
			}); err != nil {
				return err
			}

228
			system = ""
Michael Yang's avatar
Michael Yang committed
229
230
			prompt = ""
			response = ""
Michael Yang's avatar
Michael Yang committed
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
			return nil
		}

		switch m.Role {
		case "system":
			if prompt != "" || response != "" {
				if err := execute(); err != nil {
					return err
				}
			}
			system = m.Content
		case "user":
			if response != "" {
				if err := execute(); err != nil {
					return err
				}
			}
			prompt = m.Content
		case "assistant":
			response = m.Content
Michael Yang's avatar
Michael Yang committed
251
252
253
254
		}
	}

	var cut bool
255
	nodes := deleteNode(t.Template.Root.Copy(), func(n parse.Node) bool {
Michael Yang's avatar
tools  
Michael Yang committed
256
257
		if field, ok := n.(*parse.FieldNode); ok && slices.Contains(field.Ident, "Response") {
			cut = true
Michael Yang's avatar
Michael Yang committed
258
259
260
261
262
		}

		return cut
	})

263
264
	tree := parse.Tree{Root: nodes.(*parse.ListNode)}
	if err := template.Must(template.New("").AddParseTree("", &tree)).Execute(&b, map[string]any{
Michael Yang's avatar
Michael Yang committed
265
		"System": system,
Michael Yang's avatar
Michael Yang committed
266
267
268
269
270
271
272
273
274
		"Prompt": prompt,
	}); err != nil {
		return err
	}

	_, err := io.Copy(w, &b)
	return err
}

Michael Yang's avatar
Michael Yang committed
275
// collate messages based on role. consecutive messages of the same role are merged
276
277
278
// into a single message. collate also collects and returns all system messages.
// collate mutates message content adding image tags ([img-%d]) as needed
func collate(msgs []api.Message) (string, []*api.Message) {
Michael Yang's avatar
Michael Yang committed
279
	var n int
280
281
282

	var system []string
	var collated []*api.Message
Michael Yang's avatar
Michael Yang committed
283
284
285
286
287
288
289
290
291
292
293
294
	for i := range msgs {
		msg := msgs[i]
		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++
		}

295
296
297
298
		if msg.Role == "system" {
			system = append(system, msg.Content)
		}

Michael Yang's avatar
Michael Yang committed
299
300
301
302
303
304
305
		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)
		}
	}

306
	return strings.Join(system, "\n\n"), collated
Michael Yang's avatar
Michael Yang committed
307
308
}

Michael Yang's avatar
tools  
Michael Yang committed
309
310
// Identifiers walks the node tree returning any identifiers it finds along the way
func Identifiers(n parse.Node) []string {
Michael Yang's avatar
Michael Yang committed
311
	switch n := n.(type) {
Michael Yang's avatar
tools  
Michael Yang committed
312
313
314
315
	case *parse.ListNode:
		var names []string
		for _, n := range n.Nodes {
			names = append(names, Identifiers(n)...)
Michael Yang's avatar
Michael Yang committed
316
		}
Michael Yang's avatar
tools  
Michael Yang committed
317

Michael Yang's avatar
Michael Yang committed
318
		return names
Michael Yang's avatar
tools  
Michael Yang committed
319
320
321
322
323
324
325
326
327
328
	case *parse.TemplateNode:
		return Identifiers(n.Pipe)
	case *parse.ActionNode:
		return Identifiers(n.Pipe)
	case *parse.BranchNode:
		names := Identifiers(n.Pipe)
		for _, n := range []*parse.ListNode{n.List, n.ElseList} {
			if n != nil {
				names = append(names, Identifiers(n)...)
			}
Michael Yang's avatar
Michael Yang committed
329
330
		}
		return names
Michael Yang's avatar
tools  
Michael Yang committed
331
332
333
334
	case *parse.IfNode:
		return Identifiers(&n.BranchNode)
	case *parse.RangeNode:
		return Identifiers(&n.BranchNode)
Michael Yang's avatar
Michael Yang committed
335
	case *parse.WithNode:
Michael Yang's avatar
tools  
Michael Yang committed
336
		return Identifiers(&n.BranchNode)
Michael Yang's avatar
Michael Yang committed
337
338
339
340
	case *parse.PipeNode:
		var names []string
		for _, c := range n.Cmds {
			for _, a := range c.Args {
Michael Yang's avatar
tools  
Michael Yang committed
341
				names = append(names, Identifiers(a)...)
Michael Yang's avatar
Michael Yang committed
342
343
344
345
346
			}
		}
		return names
	case *parse.FieldNode:
		return n.Ident
Michael Yang's avatar
tools  
Michael Yang committed
347
348
	case *parse.VariableNode:
		return n.Ident
Michael Yang's avatar
Michael Yang committed
349
350
351
352
	}

	return nil
}
353
354
355
356
357
358
359
360
361
362
363
364
365
366
367
368
369
370
371
372
373
374
375
376
377
378
379
380
381
382
383
384
385
386
387
388
389
390
391
392
393
394
395
396
397
398
399
400
401
402
403
404
405
406
407
408
409
410
411
412
413
414
415
416
417
418
419
420
421

// deleteNode walks the node list and deletes nodes that match the predicate
// this is currently to remove the {{ .Response }} node from templates
func deleteNode(n parse.Node, fn func(parse.Node) bool) parse.Node {
	var walk func(n parse.Node) parse.Node
	walk = func(n parse.Node) parse.Node {
		if fn(n) {
			return nil
		}

		switch t := n.(type) {
		case *parse.ListNode:
			var nodes []parse.Node
			for _, c := range t.Nodes {
				if n := walk(c); n != nil {
					nodes = append(nodes, n)
				}
			}

			t.Nodes = nodes
			return t
		case *parse.IfNode:
			t.BranchNode = *(walk(&t.BranchNode).(*parse.BranchNode))
		case *parse.WithNode:
			t.BranchNode = *(walk(&t.BranchNode).(*parse.BranchNode))
		case *parse.RangeNode:
			t.BranchNode = *(walk(&t.BranchNode).(*parse.BranchNode))
		case *parse.BranchNode:
			t.List = walk(t.List).(*parse.ListNode)
			if t.ElseList != nil {
				t.ElseList = walk(t.ElseList).(*parse.ListNode)
			}
		case *parse.ActionNode:
			n := walk(t.Pipe)
			if n == nil {
				return nil
			}

			t.Pipe = n.(*parse.PipeNode)
		case *parse.PipeNode:
			var commands []*parse.CommandNode
			for _, c := range t.Cmds {
				var args []parse.Node
				for _, a := range c.Args {
					if n := walk(a); n != nil {
						args = append(args, n)
					}
				}

				if len(args) == 0 {
					return nil
				}

				c.Args = args
				commands = append(commands, c)
			}

			if len(commands) == 0 {
				return nil
			}

			t.Cmds = commands
		}

		return n
	}

	return walk(n)
}