89 lines
2.5 KiB
Go
89 lines
2.5 KiB
Go
package model
|
|
|
|
import "testing"
|
|
|
|
func TestRouter_Route_PerRole(t *testing.T) {
|
|
t.Parallel()
|
|
r := NewRouter("default",
|
|
WithRoleModel(RoleUser, "user-model"),
|
|
WithRoleModel(RoleAssistant, "assistant-model"),
|
|
WithVisionModel("vision-model"),
|
|
)
|
|
cases := []struct {
|
|
role Role
|
|
want string
|
|
}{
|
|
{RoleUser, "user-model"},
|
|
{RoleAssistant, "assistant-model"},
|
|
{RoleSystem, "default"},
|
|
{Role("unknown"), "default"},
|
|
}
|
|
for _, c := range cases {
|
|
c := c
|
|
t.Run(string(c.role), func(t *testing.T) {
|
|
t.Parallel()
|
|
if got := r.Route(c.role); got != c.want {
|
|
t.Fatalf("Route(%q) = %q, want %q", c.role, got, c.want)
|
|
}
|
|
})
|
|
}
|
|
}
|
|
|
|
func TestRouter_RouteRequest_Text(t *testing.T) {
|
|
t.Parallel()
|
|
r := NewRouter("default",
|
|
WithRoleModel(RoleUser, "user-model"),
|
|
WithRoleModel(RoleAssistant, "assistant-model"),
|
|
WithVisionModel("vision-model"),
|
|
)
|
|
req := Request{Messages: []Message{TextMessage(RoleAssistant, "hi")}}
|
|
model, needsVision := r.RouteRequest(req)
|
|
if needsVision {
|
|
t.Errorf("needsVision = true, want false for text-only")
|
|
}
|
|
if model != "assistant-model" {
|
|
t.Errorf("model = %q, want assistant-model (last-message role)", model)
|
|
}
|
|
}
|
|
|
|
func TestRouter_RouteRequest_Vision(t *testing.T) {
|
|
t.Parallel()
|
|
r := NewRouter("default",
|
|
WithRoleModel(RoleUser, "user-model"),
|
|
WithVisionModel("vision-model"),
|
|
)
|
|
req := Request{Messages: []Message{ImageMessage(RoleUser, "", "data:image/png;base64,abc")}}
|
|
model, needsVision := r.RouteRequest(req)
|
|
if !needsVision {
|
|
t.Errorf("needsVision = false, want true for image content")
|
|
}
|
|
if model != "vision-model" {
|
|
t.Errorf("model = %q, want vision-model", model)
|
|
}
|
|
}
|
|
|
|
func TestRouter_RouteRequest_VisionUnconfigured(t *testing.T) {
|
|
t.Parallel()
|
|
r := NewRouter("default", WithRoleModel(RoleUser, "user-model"))
|
|
req := Request{Messages: []Message{ImageMessage(RoleUser, "", "data:image/png;base64,abc")}}
|
|
model, needsVision := r.RouteRequest(req)
|
|
if !needsVision {
|
|
t.Errorf("needsVision = false, want true when image present even without a vision model")
|
|
}
|
|
if model != "user-model" {
|
|
t.Errorf("model = %q, want user-model (role fallback)", model)
|
|
}
|
|
}
|
|
|
|
func TestRouter_RouteRequest_EmptyMessages(t *testing.T) {
|
|
t.Parallel()
|
|
r := NewRouter("default", WithRoleModel(RoleUser, "user-model"))
|
|
model, needsVision := r.RouteRequest(Request{})
|
|
if needsVision {
|
|
t.Errorf("needsVision = true, want false")
|
|
}
|
|
if model != "user-model" {
|
|
t.Errorf("model = %q, want user-model (default role for empty messages)", model)
|
|
}
|
|
}
|