template: preprocess message and collect system

This commit is contained in:
Michael Yang
2024-07-12 11:48:06 -07:00
parent 179737feb7
commit 36c87c433b
2 changed files with 23 additions and 67 deletions

View File

@@ -102,22 +102,8 @@ var response = parse.ActionNode{
},
}
var funcs = template.FuncMap{
// contents returns the contents of messages with an optional role filter
"contents": func(v []*api.Message, role ...string) string {
var parts []string
for _, m := range v {
if len(role) == 0 || role[0] == "" || m.Role == role[0] {
parts = append(parts, m.Content)
}
}
return strings.Join(parts, "\n\n")
},
}
func Parse(s string) (*Template, error) {
tmpl := template.New("").Option("missingkey=zero").Funcs(funcs)
tmpl := template.New("").Option("missingkey=zero")
tmpl, err := tmpl.Parse(s)
if err != nil {
@@ -163,15 +149,16 @@ type Values struct {
}
func (t *Template) Execute(w io.Writer, v Values) error {
collated := collate(v.Messages)
system, collated := collate(v.Messages)
if !v.forceLegacy && slices.Contains(t.Vars(), "messages") {
return t.Template.Execute(w, map[string]any{
"System": system,
"Messages": collated,
})
}
var b bytes.Buffer
var system, prompt, response string
var prompt, response string
for i, m := range collated {
switch m.Role {
case "system":
@@ -223,11 +210,13 @@ func (t *Template) Execute(w io.Writer, v Values) error {
}
// collate messages based on role. consecutive messages of the same role are merged
// into a single message. collate also pulls out and merges messages with Role == "system"
// which are templated separately. As a side effect, it mangles message content adding image
// tags ([img-%d]) as needed
func collate(msgs []api.Message) (collated []*api.Message) {
// 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) {
var n int
var system []string
var collated []*api.Message
for i := range msgs {
msg := msgs[i]
for range msg.Images {
@@ -240,6 +229,10 @@ func collate(msgs []api.Message) (collated []*api.Message) {
n++
}
if msg.Role == "system" {
system = append(system, msg.Content)
}
if len(collated) > 0 && collated[len(collated)-1].Role == msg.Role {
collated[len(collated)-1].Content += "\n\n" + msg.Content
} else {
@@ -247,7 +240,7 @@ func collate(msgs []api.Message) (collated []*api.Message) {
}
}
return
return strings.Join(system, "\n\n"), collated
}
func parseNode(n parse.Node) []string {