prompt.go 2 KB
Newer Older
1
2
3
package server

import (
Michael Yang's avatar
Michael Yang committed
4
5
	"bytes"
	"context"
6
7
	"log/slog"

8
	"github.com/ollama/ollama/api"
Michael Yang's avatar
Michael Yang committed
9
	"github.com/ollama/ollama/llm"
Michael Yang's avatar
Michael Yang committed
10
	"github.com/ollama/ollama/template"
11
12
)

Michael Yang's avatar
Michael Yang committed
13
14
15
16
17
type tokenizeFunc func(context.Context, string) ([]int, error)

// chatPrompt accepts a list of messages and returns the prompt and images that should be used for the next chat turn.
// chatPrompt truncates any messages that exceed the context window of the model, making sure to always include 1) the
// latest message and 2) system messages
Michael Yang's avatar
tools  
Michael Yang committed
18
func chatPrompt(ctx context.Context, m *Model, tokenize tokenizeFunc, opts *api.Options, msgs []api.Message, tools []api.Tool) (prompt string, images []llm.ImageData, _ error) {
Michael Yang's avatar
Michael Yang committed
19
	var system []api.Message
Michael Yang's avatar
Michael Yang committed
20
	// always include the last message
Michael Yang's avatar
Michael Yang committed
21
	n := len(msgs) - 1
Michael Yang's avatar
Michael Yang committed
22
	// in reverse, find all messages that fit into context window
Michael Yang's avatar
Michael Yang committed
23
	for i := n - 1; i >= 0; i-- {
Michael Yang's avatar
Michael Yang committed
24
25
26
27
28
29
30
		system = make([]api.Message, 0)
		for j := range i {
			if msgs[j].Role == "system" {
				system = append(system, msgs[j])
			}
		}

Michael Yang's avatar
Michael Yang committed
31
		var b bytes.Buffer
Michael Yang's avatar
tools  
Michael Yang committed
32
		if err := m.Template.Execute(&b, template.Values{Messages: append(system, msgs[i:]...), Tools: tools}); err != nil {
Michael Yang's avatar
Michael Yang committed
33
			return "", nil, err
34
35
		}

Michael Yang's avatar
Michael Yang committed
36
		s, err := tokenize(ctx, b.String())
37
		if err != nil {
Michael Yang's avatar
Michael Yang committed
38
			return "", nil, err
39
40
		}

Michael Yang's avatar
Michael Yang committed
41
		c := len(s)
Michael Yang's avatar
Michael Yang committed
42
		if m.ProjectorPaths != nil {
Michael Yang's avatar
Michael Yang committed
43
			for _, m := range msgs[i:] {
Michael Yang's avatar
Michael Yang committed
44
45
				// images are represented as 768 sized embeddings
				// TODO: get embedding length from project metadata
Michael Yang's avatar
Michael Yang committed
46
47
				c += 768 * len(m.Images)
			}
48
49
		}

Michael Yang's avatar
Michael Yang committed
50
		if c > opts.NumCtx {
Michael Yang's avatar
Michael Yang committed
51
			slog.Debug("truncating input messages which exceed context length", "truncated", len(msgs[i:]))
52
			break
Michael Yang's avatar
Michael Yang committed
53
54
		} else {
			n = i
55
		}
Michael Yang's avatar
Michael Yang committed
56
	}
57

Michael Yang's avatar
Michael Yang committed
58
	// truncate any messages that do not fit into the context window
Michael Yang's avatar
Michael Yang committed
59
	var b bytes.Buffer
Michael Yang's avatar
tools  
Michael Yang committed
60
	if err := m.Template.Execute(&b, template.Values{Messages: append(system, msgs[n:]...), Tools: tools}); err != nil {
Michael Yang's avatar
Michael Yang committed
61
		return "", nil, err
62
63
	}

Michael Yang's avatar
Michael Yang committed
64
65
66
67
68
69
	for _, m := range msgs[n:] {
		for _, i := range m.Images {
			images = append(images, llm.ImageData{
				ID:   len(images),
				Data: i,
			})
70
71
72
		}
	}

Michael Yang's avatar
Michael Yang committed
73
	return b.String(), images, nil
74
}