llm/chatml_test.go


1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28

package llm

import (
	"strings"
	"testing"
)

func TestChatML(t *testing.T) {
	session := Session{
		Messages: []ChatMLer{
			Message{
				Role:    "user",
				Content: "hello",
			},
			Message{
				Role: "assistant",
			},
		},
	}

	expected := `<|im_start|>user
hello<|im_end|>
<|im_start|>assistant`

	if strings.TrimSpace(session.ChatML()) != strings.TrimSpace(expected) {
		t.Errorf("Expected\n\n%s\n\ngot\n\n%s", expected, session.ChatML())
	}
}