81 lines
1.9 KiB
Go
81 lines
1.9 KiB
Go
|
package openai
|
||
|
|
|
||
|
|
import "testing"
|
||
|
|
|
||
|
|
func TestNormalizeModelID(t *testing.T) {
|
||
|
|
tests := []struct {
|
||
|
|
input string
|
||
|
|
want string
|
||
|
|
}{
|
||
|
|
{"gpt-4", "gpt-4"},
|
||
|
|
{"openai/gpt-4", "gpt-4"},
|
||
|
|
{"anthropic/claude-3", "claude-3"},
|
||
|
|
{"", ""},
|
||
|
|
{" ", ""},
|
||
|
|
{"a/b/c", "c"},
|
||
|
|
}
|
||
|
|
for _, tc := range tests {
|
||
|
|
got := NormalizeModelID(tc.input)
|
||
|
|
if got != tc.want {
|
||
|
|
t.Errorf("NormalizeModelID(%q) = %q, want %q", tc.input, got, tc.want)
|
||
|
|
}
|
||
|
|
}
|
||
|
|
}
|
||
|
|
|
||
|
|
func TestBuildPromptFromMessages(t *testing.T) {
|
||
|
|
messages := []interface{}{
|
||
|
|
map[string]interface{}{"role": "system", "content": "You are helpful."},
|
||
|
|
map[string]interface{}{"role": "user", "content": "Hello"},
|
||
|
|
map[string]interface{}{"role": "assistant", "content": "Hi there"},
|
||
|
|
}
|
||
|
|
got := BuildPromptFromMessages(messages)
|
||
|
|
if got == "" {
|
||
|
|
t.Fatal("expected non-empty prompt")
|
||
|
|
}
|
||
|
|
containsSystem := false
|
||
|
|
containsUser := false
|
||
|
|
containsAssistant := false
|
||
|
|
for i := 0; i < len(got)-10; i++ {
|
||
|
|
if got[i:i+6] == "System" {
|
||
|
|
containsSystem = true
|
||
|
|
}
|
||
|
|
if got[i:i+4] == "User" {
|
||
|
|
containsUser = true
|
||
|
|
}
|
||
|
|
if got[i:i+9] == "Assistant" {
|
||
|
|
containsAssistant = true
|
||
|
|
}
|
||
|
|
}
|
||
|
|
if !containsSystem || !containsUser || !containsAssistant {
|
||
|
|
t.Errorf("prompt missing sections: system=%v user=%v assistant=%v\n%s",
|
||
|
|
containsSystem, containsUser, containsAssistant, got)
|
||
|
|
}
|
||
|
|
}
|
||
|
|
|
||
|
|
func TestToolsToSystemText(t *testing.T) {
|
||
|
|
tools := []interface{}{
|
||
|
|
map[string]interface{}{
|
||
|
|
"type": "function",
|
||
|
|
"function": map[string]interface{}{
|
||
|
|
"name": "get_weather",
|
||
|
|
"description": "Get weather",
|
||
|
|
"parameters": map[string]interface{}{"type": "object"},
|
||
|
|
},
|
||
|
|
},
|
||
|
|
}
|
||
|
|
got := ToolsToSystemText(tools, nil)
|
||
|
|
if got == "" {
|
||
|
|
t.Fatal("expected non-empty tools text")
|
||
|
|
}
|
||
|
|
if len(got) < 10 {
|
||
|
|
t.Errorf("tools text too short: %q", got)
|
||
|
|
}
|
||
|
|
}
|
||
|
|
|
||
|
|
func TestToolsToSystemTextEmpty(t *testing.T) {
|
||
|
|
got := ToolsToSystemText(nil, nil)
|
||
|
|
if got != "" {
|
||
|
|
t.Errorf("expected empty string for no tools, got %q", got)
|
||
|
|
}
|
||
|
|
}
|