diff --git a/internal/util/messages_test.go b/internal/util/messages_test.go index 776853b..ab11b59 100644 --- a/internal/util/messages_test.go +++ b/internal/util/messages_test.go @@ -12,7 +12,7 @@ func TestMessagesPrepareBasic(t *testing.T) { if got == "" { t.Fatal("expected non-empty prompt") } - if got != "Hello" { + if got != "<|User|>Hello" { t.Fatalf("unexpected prompt: %q", got) } } @@ -55,7 +55,7 @@ func TestMessagesPrepareArrayTextVariants(t *testing.T) { }, } got := MessagesPrepare(messages) - if got != "line1\nline2" { + if got != "<|User|>line1\nline2" { t.Fatalf("unexpected content from text variants: %q", got) } } diff --git a/internal/util/util_edge_test.go b/internal/util/util_edge_test.go index 5d024a9..9f91602 100644 --- a/internal/util/util_edge_test.go +++ b/internal/util/util_edge_test.go @@ -162,13 +162,16 @@ func TestMessagesPrepareMergesConsecutiveSameRole(t *testing.T) { {"role": "user", "content": "World"}, } got := MessagesPrepare(messages) + if !strings.HasPrefix(got, "<|User|>") { + t.Fatalf("expected user marker at the start, got %q", got) + } if !strings.Contains(got, "Hello") || !strings.Contains(got, "World") { t.Fatalf("expected both messages, got %q", got) } - // Should be merged without <|User|> between them + // Should be merged into a single user turn with one marker at the start. count := strings.Count(got, "<|User|>") - if count != 0 { - t.Fatalf("expected no User marker for first message pair, got %d occurrences", count) + if count != 1 { + t.Fatalf("expected one User marker for the merged pair, got %d occurrences", count) } }