46 lines
1.3 KiB
Go
46 lines
1.3 KiB
Go
|
package aws_test
|
||
|
|
|
||
|
|
import (
|
||
|
|
"testing"
|
||
|
|
|
||
|
|
aws "github.com/songquanpeng/one-api/relay/adaptor/aws/llama3"
|
||
|
|
relaymodel "github.com/songquanpeng/one-api/relay/model"
|
||
|
|
"github.com/stretchr/testify/assert"
|
||
|
|
)
|
||
|
|
|
||
|
|
func TestRenderPrompt(t *testing.T) {
|
||
|
|
messages := []relaymodel.Message{
|
||
|
|
{
|
||
|
|
Role: "user",
|
||
|
|
Content: "What's your name?",
|
||
|
|
},
|
||
|
|
}
|
||
|
|
prompt := aws.RenderPrompt(messages)
|
||
|
|
expected := `<|begin_of_text|><|start_header_id|>user<|end_header_id|>What's your name?<|eot_id|><|start_header_id|>assistant<|end_header_id|>
|
||
|
|
`
|
||
|
|
assert.Equal(t, expected, prompt)
|
||
|
|
|
||
|
|
messages = []relaymodel.Message{
|
||
|
|
{
|
||
|
|
Role: "system",
|
||
|
|
Content: "Your name is Kat. You are a detective.",
|
||
|
|
},
|
||
|
|
{
|
||
|
|
Role: "user",
|
||
|
|
Content: "What's your name?",
|
||
|
|
},
|
||
|
|
{
|
||
|
|
Role: "assistant",
|
||
|
|
Content: "Kat",
|
||
|
|
},
|
||
|
|
{
|
||
|
|
Role: "user",
|
||
|
|
Content: "What's your job?",
|
||
|
|
},
|
||
|
|
}
|
||
|
|
prompt = aws.RenderPrompt(messages)
|
||
|
|
expected = `<|begin_of_text|><|start_header_id|>system<|end_header_id|>Your name is Kat. You are a detective.<|eot_id|><|start_header_id|>user<|end_header_id|>What's your name?<|eot_id|><|start_header_id|>assistant<|end_header_id|>Kat<|eot_id|><|start_header_id|>user<|end_header_id|>What's your job?<|eot_id|><|start_header_id|>assistant<|end_header_id|>
|
||
|
|
`
|
||
|
|
assert.Equal(t, expected, prompt)
|
||
|
|
}
|