From ac9ffca2ab33e59c42b17988f90811c36b84d272 Mon Sep 17 00:00:00 2001 From: Paul Buetow Date: Tue, 21 Apr 2026 10:53:01 +0300 Subject: p7: make image prompt truncation UTF-8 safe --- internal/image/gemini_prompt.go | 2 +- internal/image/gemini_test.go | 50 +++++++++++++++++++++++++++++++++++++++++ internal/image/prompt.go | 37 ++++++++++++++++++++++++++++-- internal/image/prompt_test.go | 34 ++++++++++++++++++++++++++++ 4 files changed, 120 insertions(+), 3 deletions(-) (limited to 'internal') diff --git a/internal/image/gemini_prompt.go b/internal/image/gemini_prompt.go index 5932174..6d84ee9 100644 --- a/internal/image/gemini_prompt.go +++ b/internal/image/gemini_prompt.go @@ -9,7 +9,7 @@ import ( func normalizeCustomPrompt(prompt string) string { prompt = strings.TrimSpace(prompt) if len(prompt) > maxCustomPrompt { - prompt = prompt[:maxCustomPrompt-3] + "..." + prompt = truncateToByteLimit(prompt, maxCustomPrompt) } return prompt } diff --git a/internal/image/gemini_test.go b/internal/image/gemini_test.go index aae8014..91e0ce4 100644 --- a/internal/image/gemini_test.go +++ b/internal/image/gemini_test.go @@ -9,6 +9,7 @@ import ( "net/http/httptest" "strings" "testing" + "unicode/utf8" "google.golang.org/genai" ) @@ -189,6 +190,55 @@ func TestGeminiProvider_Search_CustomPromptIsTruncated(t *testing.T) { } } +func TestGeminiProvider_Search_CustomPromptTruncatesUTF8Safely(t *testing.T) { + originalText := geminiGenerateText + originalImage := geminiGenerateImage + t.Cleanup(func() { + geminiGenerateText = originalText + geminiGenerateImage = originalImage + }) + + geminiGenerateText = func(context.Context, *GeminiProvider, string, string, string, float32, int32) (string, error) { + t.Fatal("unexpected text generation for truncated custom prompt") + return "", nil + } + + var gotPrompt string + geminiGenerateImage = func(_ context.Context, _ *GeminiProvider, prompt, _ string) ([]byte, string, error) { + gotPrompt = prompt + return mustJPEGBytes(t), "image/jpeg", nil + } + + client := NewGeminiProvider(&GeminiConfig{APIKey: "test-key"}) + longPrompt := strings.Repeat("яa", 2500) + if len(longPrompt) <= maxCustomPrompt { + t.Fatal("test prompt did not exceed truncation limit") + } + + results, err := client.Search(context.Background(), &SearchOptions{ + Query: "ябълка", + CustomPrompt: longPrompt, + }) + if err != nil { + t.Fatalf("Search() unexpected error: %v", err) + } + if !utf8.ValidString(gotPrompt) { + t.Fatalf("Search() returned invalid UTF-8 prompt: %q", gotPrompt) + } + if len(gotPrompt) > maxCustomPrompt { + t.Fatalf("prompt length = %d, want <= %d", len(gotPrompt), maxCustomPrompt) + } + if !strings.HasSuffix(gotPrompt, "...") { + t.Fatalf("prompt = %q, want ellipsis suffix", gotPrompt) + } + if client.LastPrompt() != gotPrompt { + t.Fatalf("LastPrompt() = %q, want %q", client.LastPrompt(), gotPrompt) + } + if len(results) != 1 { + t.Fatalf("expected 1 result, got %d", len(results)) + } +} + func TestGeminiProvider_GenerateImage_UsesCustomPrompt(t *testing.T) { originalText := geminiGenerateText originalImage := geminiGenerateImage diff --git a/internal/image/prompt.go b/internal/image/prompt.go index addf936..c67c8b3 100644 --- a/internal/image/prompt.go +++ b/internal/image/prompt.go @@ -6,6 +6,7 @@ import ( ) const maxImagePromptChars = 1000 +const promptEllipsis = "..." func promptSubject(englishTranslation, fallback string) string { subject := normalizePromptText(englishTranslation) @@ -95,6 +96,38 @@ func withTerminalPunctuation(text string) string { } } +func truncateToByteLimit(text string, maxBytes int) string { + text = strings.TrimSpace(text) + if text == "" || maxBytes <= 0 { + return "" + } + if len(text) <= maxBytes { + return text + } + + if maxBytes <= len(promptEllipsis) { + return truncateToRuneBoundary(text, maxBytes) + } + + return truncateToRuneBoundary(text, maxBytes-len(promptEllipsis)) + promptEllipsis +} + +func truncateToRuneBoundary(text string, maxBytes int) string { + if maxBytes <= 0 { + return "" + } + + end := 0 + for i := range text { + if i > maxBytes { + break + } + end = i + } + + return text[:end] +} + // buildEducationalPrompt assembles the final image-generation prompt from a // pre-chosen artistic style, an optional scene description, and the word subject. func buildEducationalPrompt(style, scene, subject string) string { @@ -128,7 +161,7 @@ func buildEducationalPrompt(style, scene, subject string) string { ) maxSceneLen := maxImagePromptChars - len(template) if maxSceneLen > 3 && len(scene) > maxSceneLen { - scene = scene[:maxSceneLen] + "..." + scene = truncateToByteLimit(scene, maxSceneLen) } prompt = fmt.Sprintf( "Generate a %s flashcard image illustrating \"%s\". Scene: %s "+ @@ -149,7 +182,7 @@ func buildEducationalPrompt(style, scene, subject string) string { } if len(prompt) > maxImagePromptChars { - prompt = prompt[:997] + "..." + prompt = truncateToByteLimit(prompt, maxImagePromptChars) } return prompt diff --git a/internal/image/prompt_test.go b/internal/image/prompt_test.go index 73eb30e..957c84c 100644 --- a/internal/image/prompt_test.go +++ b/internal/image/prompt_test.go @@ -1,8 +1,10 @@ package image import ( + "fmt" "strings" "testing" + "unicode/utf8" ) func TestPromptSubjectUsesTranslationFirst(t *testing.T) { @@ -58,3 +60,35 @@ func TestFallbackVisualDirectionForSingleWord(t *testing.T) { t.Fatalf("fallbackVisualDirection() = %q", got) } } + +func TestBuildEducationalPromptTruncatesLongSceneWithoutBreakingUTF8(t *testing.T) { + t.Parallel() + + style := "Photorealism" + subject := "ябълка" + for { + template := fmt.Sprintf( + "Generate a %s flashcard image illustrating \"%s\". Scene: "+ + "The image should be educational and suitable for language learning flashcards. "+ + "Requirements: The main subject or concept must be clearly visible, centered, well lit, and easy to identify.", + style, subject, + ) + if (maxImagePromptChars-len(template))%2 == 1 { + break + } + subject += "x" + } + + scene := strings.Repeat("Ярка сцена с ябълки и хора, ", 80) + got := buildEducationalPrompt(style, scene, subject) + + if !utf8.ValidString(got) { + t.Fatalf("buildEducationalPrompt() returned invalid UTF-8: %q", got) + } + if len(got) > maxImagePromptChars { + t.Fatalf("buildEducationalPrompt() length = %d, want <= %d", len(got), maxImagePromptChars) + } + if !strings.Contains(got, "...") { + t.Fatalf("buildEducationalPrompt() = %q, want truncated scene marker", got) + } +} -- cgit v1.2.3