diff options
| -rw-r--r-- | internal/image/doc.go | 2 | ||||
| -rw-r--r-- | internal/image/download.go | 223 | ||||
| -rw-r--r-- | internal/image/download_test.go | 144 | ||||
| -rw-r--r-- | internal/image/gemini.go | 599 | ||||
| -rw-r--r-- | internal/image/gemini_test.go | 333 | ||||
| -rw-r--r-- | internal/image/prompt.go | 167 | ||||
| -rw-r--r-- | internal/image/prompt_test.go | 60 | ||||
| -rw-r--r-- | internal/image/registry.go | 74 | ||||
| -rw-r--r-- | internal/image/styles.go | 100 | ||||
| -rw-r--r-- | internal/image/styles_test.go | 49 | ||||
| -rw-r--r-- | internal/image/test_helpers_test.go | 38 | ||||
| -rw-r--r-- | internal/image/types.go | 113 | ||||
| -rw-r--r-- | internal/image/types_test.go | 101 |
13 files changed, 2003 insertions, 0 deletions
diff --git a/internal/image/doc.go b/internal/image/doc.go new file mode 100644 index 0000000..ae5d586 --- /dev/null +++ b/internal/image/doc.go @@ -0,0 +1,2 @@ +// Package image provides provider-neutral image generation and download helpers. +package image diff --git a/internal/image/download.go b/internal/image/download.go new file mode 100644 index 0000000..7a196fa --- /dev/null +++ b/internal/image/download.go @@ -0,0 +1,223 @@ +package image + +import ( + "context" + "fmt" + "io" + "os" + "path/filepath" + "strings" +) + +// DownloadOptions configures image download behavior. +type DownloadOptions struct { + OutputDir string + OverwriteExisting bool + CreateDir bool + FileNamePattern string + MaxSizeBytes int64 +} + +// DefaultDownloadOptions returns sensible defaults for image downloads. +func DefaultDownloadOptions() *DownloadOptions { + return &DownloadOptions{ + OutputDir: "./images", + OverwriteExisting: false, + CreateDir: true, + FileNamePattern: "{word}_{source}", + MaxSizeBytes: 10 * 1024 * 1024, + } +} + +// Downloader handles image downloads from search results. +type Downloader struct { + provider ImageProvider + options *DownloadOptions +} + +// NewDownloader creates a new image downloader. +func NewDownloader(provider ImageProvider, options *DownloadOptions) *Downloader { + if options == nil { + options = DefaultDownloadOptions() + } + return &Downloader{ + provider: provider, + options: options, + } +} + +// DownloadImage downloads a single image to the specified path. +func (d *Downloader) DownloadImage(ctx context.Context, result *SearchResult, outputPath string) (err error) { + if d == nil || d.provider == nil { + return fmt.Errorf("image provider is required") + } + if result == nil { + return fmt.Errorf("search result is required") + } + + dir := filepath.Dir(outputPath) + if d.options != nil && d.options.CreateDir && dir != "" && dir != "." { + if err := os.MkdirAll(dir, 0o755); err != nil { + return fmt.Errorf("create output dir %q: %w", dir, err) + } + } + + if d.options != nil && !d.options.OverwriteExisting { + if _, err := os.Stat(outputPath); err == nil { + return fmt.Errorf("output file exists: %s", outputPath) + } + } + + reader, err := d.provider.Download(ctx, result.URL) + if err != nil { + return fmt.Errorf("download %q: %w", result.URL, err) + } + defer func() { + _ = reader.Close() + }() + + file, err := os.Create(outputPath) + if err != nil { + return fmt.Errorf("create output file %q: %w", outputPath, err) + } + defer func() { + if closeErr := file.Close(); err == nil && closeErr != nil { + err = fmt.Errorf("close output file %q: %w", outputPath, closeErr) + } + }() + + if d.options != nil && d.options.MaxSizeBytes > 0 { + written, copyErr := io.CopyN(file, reader, d.options.MaxSizeBytes) + if copyErr != nil && copyErr != io.EOF { + _ = os.Remove(outputPath) + return fmt.Errorf("write output file %q: %w", outputPath, copyErr) + } + + if written == d.options.MaxSizeBytes { + var probe [1]byte + if n, probeErr := reader.Read(probe[:]); n > 0 || probeErr != io.EOF { + _ = os.Remove(outputPath) + return fmt.Errorf("image exceeds max size %d bytes", d.options.MaxSizeBytes) + } + } + } else { + if _, err = io.Copy(file, reader); err != nil { + _ = os.Remove(outputPath) + return fmt.Errorf("write output file %q: %w", outputPath, err) + } + } + + if err := file.Sync(); err != nil { + return fmt.Errorf("sync output file %q: %w", outputPath, err) + } + + if attribution := d.provider.GetAttribution(result); attribution != "" { + attrPath := strings.TrimSuffix(outputPath, filepath.Ext(outputPath)) + "_attribution.txt" + if err := os.WriteFile(attrPath, []byte(attribution), 0o644); err != nil { + fmt.Fprintf(os.Stderr, "Warning: failed to save attribution: %v\n", err) + } + } + + return nil +} + +// DownloadBestMatch downloads the best matching image for a query. +func (d *Downloader) DownloadBestMatch(ctx context.Context, query string) (*SearchResult, string, error) { + opts := DefaultSearchOptions(query) + opts.PerPage = 5 + return d.DownloadBestMatchWithOptions(ctx, opts) +} + +// DownloadBestMatchWithOptions downloads the best matching image for given search options. +func (d *Downloader) DownloadBestMatchWithOptions(ctx context.Context, opts *SearchOptions) (*SearchResult, string, error) { + if d == nil || d.provider == nil { + return nil, "", fmt.Errorf("image provider is required") + } + if opts == nil { + return nil, "", fmt.Errorf("search options are required") + } + + searchOpts := *opts + searchOpts.PerPage = 5 + + results, err := d.provider.Search(ctx, &searchOpts) + if err != nil { + return nil, "", fmt.Errorf("search images: %w", err) + } + if len(results) == 0 { + return nil, "", fmt.Errorf("no images found for %q", opts.Query) + } + + for i, result := range results { + filename := d.generateFileName(opts.Query, &result, i) + outputDir := "./images" + if d.options != nil && d.options.OutputDir != "" { + outputDir = d.options.OutputDir + } + outputPath := filepath.Join(outputDir, filename) + + if err := d.DownloadImage(ctx, &result, outputPath); err == nil { + return &result, outputPath, nil + } else { + fmt.Fprintf(os.Stderr, "Warning: failed to download image %d: %v\n", i+1, err) + } + } + + return nil, "", fmt.Errorf("no downloadable images found for %q", opts.Query) +} + +func (d *Downloader) generateFileName(word string, result *SearchResult, index int) string { + filename := "" + if d != nil && d.options != nil { + filename = d.options.FileNamePattern + } + if filename == "" { + filename = "{word}_{source}" + } + + filename = strings.ReplaceAll(filename, "{word}", sanitizeFileName(word)) + if result != nil { + filename = strings.ReplaceAll(filename, "{source}", result.Source) + filename = strings.ReplaceAll(filename, "{id}", result.ID) + } + filename = strings.ReplaceAll(filename, "{index}", fmt.Sprintf("%d", index)) + + ext := "" + if result != nil { + ext = filepath.Ext(result.URL) + if strings.HasPrefix(result.URL, geminiDataPrefix) { + ext = ".png" + } else if ext == "" || len(ext) > 5 { + ext = ".jpg" + } + } + + if filepath.Ext(filename) == "" { + filename += ext + } + + return filename +} + +func sanitizeFileName(name string) string { + replacer := strings.NewReplacer( + "/", "_", + "\\", "_", + ":", "_", + "*", "_", + "?", "_", + "\"", "_", + "<", "_", + ">", "_", + "|", "_", + " ", "_", + ".", "_", + ) + + sanitized := replacer.Replace(name) + if len(sanitized) > 50 { + sanitized = sanitized[:50] + } + + return sanitized +} diff --git a/internal/image/download_test.go b/internal/image/download_test.go new file mode 100644 index 0000000..588d875 --- /dev/null +++ b/internal/image/download_test.go @@ -0,0 +1,144 @@ +package image + +import ( + "context" + "io" + "os" + "path/filepath" + "strings" + "testing" +) + +type mockDownloaderProvider struct { + results []SearchResult + searchErr error + payload string + attribution string + searchQueries []string + downloadURLs []string +} + +func (m *mockDownloaderProvider) Name() string { return "mock" } + +func (m *mockDownloaderProvider) Search(_ context.Context, opts *SearchOptions) ([]SearchResult, error) { + if opts != nil { + m.searchQueries = append(m.searchQueries, opts.Query) + } + if m.searchErr != nil { + return nil, m.searchErr + } + return append([]SearchResult(nil), m.results...), nil +} + +func (m *mockDownloaderProvider) Download(_ context.Context, url string) (io.ReadCloser, error) { + m.downloadURLs = append(m.downloadURLs, url) + return io.NopCloser(strings.NewReader(m.payload)), nil +} + +func (m *mockDownloaderProvider) GetAttribution(*SearchResult) string { + return m.attribution +} + +func TestDownloaderGenerateFileName_DataURIUsesPNG(t *testing.T) { + t.Parallel() + + d := NewDownloader(&mockDownloaderProvider{}, &DownloadOptions{ + FileNamePattern: "{word}_{source}", + }) + result := &SearchResult{ + URL: "data:image/png;base64,AAAA", + Source: Gemini, + } + + if got := d.generateFileName("ябълка", result, 0); got != "ябълка_gemini.png" { + t.Fatalf("generateFileName() = %q, want %q", got, "ябълка_gemini.png") + } +} + +func TestDownloadImageWritesAttribution(t *testing.T) { + t.Parallel() + + provider := &mockDownloaderProvider{ + payload: "image-bytes", + attribution: "attribution text", + } + d := NewDownloader(provider, &DownloadOptions{ + OutputDir: t.TempDir(), + CreateDir: true, + OverwriteExisting: false, + FileNamePattern: "{word}_{source}", + MaxSizeBytes: 10 * 1024 * 1024, + }) + + outputPath := filepath.Join(d.options.OutputDir, "ябълка_gemini.png") + if err := d.DownloadImage(context.Background(), &SearchResult{ + URL: "https://example.com/image.png", + Source: Gemini, + ID: "1", + }, outputPath); err != nil { + t.Fatalf("DownloadImage() error = %v", err) + } + + data, err := os.ReadFile(outputPath) + if err != nil { + t.Fatalf("ReadFile() error = %v", err) + } + if string(data) != "image-bytes" { + t.Fatalf("downloaded file = %q, want %q", string(data), "image-bytes") + } + + attrPath := strings.TrimSuffix(outputPath, filepath.Ext(outputPath)) + "_attribution.txt" + attr, err := os.ReadFile(attrPath) + if err != nil { + t.Fatalf("ReadFile(attribution) error = %v", err) + } + if string(attr) != "attribution text" { + t.Fatalf("attribution = %q, want %q", string(attr), "attribution text") + } +} + +func TestDownloadBestMatchWithOptions(t *testing.T) { + t.Parallel() + + provider := &mockDownloaderProvider{ + results: []SearchResult{ + { + ID: "1", + URL: "https://example.com/image1.jpg", + Source: Gemini, + }, + }, + payload: "image-bytes", + attribution: "attribution text", + } + d := NewDownloader(provider, &DownloadOptions{ + OutputDir: t.TempDir(), + CreateDir: true, + OverwriteExisting: true, + FileNamePattern: "{word}_{source}", + MaxSizeBytes: 10 * 1024 * 1024, + }) + + result, path, err := d.DownloadBestMatchWithOptions(context.Background(), &SearchOptions{Query: "ябълка"}) + if err != nil { + t.Fatalf("DownloadBestMatchWithOptions() error = %v", err) + } + if result == nil || result.ID != "1" { + t.Fatalf("DownloadBestMatchWithOptions() result = %+v, want ID 1", result) + } + if !strings.HasSuffix(path, ".jpg") { + t.Fatalf("DownloadBestMatchWithOptions() path = %q, want jpg suffix", path) + } + if _, err := os.Stat(path); err != nil { + t.Fatalf("downloaded file missing: %v", err) + } +} + +func TestDownloadImageRejectsNilResult(t *testing.T) { + t.Parallel() + + d := NewDownloader(&mockDownloaderProvider{}, nil) + if err := d.DownloadImage(context.Background(), nil, filepath.Join(t.TempDir(), "out.png")); err == nil { + t.Fatal("expected error for nil result") + } +} diff --git a/internal/image/gemini.go b/internal/image/gemini.go new file mode 100644 index 0000000..2ccd6fa --- /dev/null +++ b/internal/image/gemini.go @@ -0,0 +1,599 @@ +package image + +import ( + "bytes" + "context" + "crypto/md5" + "encoding/base64" + "encoding/hex" + "fmt" + "image" + _ "image/jpeg" + "image/png" + "io" + "net/http" + "strings" + "time" + + "google.golang.org/genai" + + "codeberg.org/snonux/comicforge/internal/apicircuit" + "codeberg.org/snonux/comicforge/internal/httpctx" +) + +const ( + // DefaultGeminiImageModel is the Gemini image model used for image generation. + DefaultGeminiImageModel = "gemini-3.1-flash-image-preview" + + // DefaultGeminiTextModel is the Gemini text model used for scene generation. + DefaultGeminiTextModel = "gemini-2.5-flash" + + geminiAspectRatio = "4:3" + geminiDataPrefix = "data:image/png;base64," + geminiSource = Gemini + maxCustomPrompt = 4000 +) + +// GeminiConfig holds the settings needed to build a Gemini-backed image provider. +type GeminiConfig struct { + APIKey string + Model string + TextModel string +} + +// GeminiProvider implements ImageProvider for Google Gemini image generation. +type GeminiProvider struct { + client *genai.Client + initErr error + config *GeminiConfig + lastPrompt string + + // PromptCallback runs after the prompt is generated and before image creation. + PromptCallback func(prompt string) +} + +var _ ImageProvider = (*GeminiProvider)(nil) + +var imageHTTPClient = httpctx.ImageDownloadHTTPClient() +var newGeminiClient = httpctx.NewGenAIClient +var geminiGenerateText = func(ctx context.Context, c *GeminiProvider, model, systemPrompt, userPrompt string, temperature float32, maxOutputTokens int32) (string, error) { + return c.generateText(ctx, model, systemPrompt, userPrompt, temperature, maxOutputTokens) +} +var geminiGenerateImage = func(ctx context.Context, c *GeminiProvider, prompt, aspectRatio string) ([]byte, string, error) { + return c.generateImage(ctx, prompt, aspectRatio) +} + +// NewGeminiProvider creates a new Gemini image provider. +func NewGeminiProvider(config *GeminiConfig) *GeminiProvider { + normalized := normalizeGeminiConfig(config) + client := &GeminiProvider{config: normalized} + + if normalized.APIKey == "" { + return client + } + + genaiClient, err := newGeminiClient(context.Background(), &genai.ClientConfig{ + APIKey: normalized.APIKey, + Backend: genai.BackendGeminiAPI, + }) + if err != nil { + client.initErr = err + return client + } + + client.client = genaiClient + return client +} + +// Search generates an educational image for the requested word or phrase. +func (c *GeminiProvider) Search(ctx context.Context, opts *SearchOptions) ([]SearchResult, error) { + ctx, cancel := httpctx.WithTimeoutUnlessSet(ctx, httpctx.OperationTimeoutDefault) + defer cancel() + + if err := c.ensureReady(); err != nil { + return nil, err + } + if opts == nil { + return nil, &SearchError{ + Provider: geminiSource, + Code: "INVALID_OPTIONS", + Message: "search options are required", + } + } + + prompt, translatedWord, err := c.buildPrompt(ctx, opts) + if err != nil { + return nil, err + } + + c.lastPrompt = prompt + if c.PromptCallback != nil { + c.PromptCallback(prompt) + } + + aspectRatio := geminiAspectRatio + if opts.AspectRatio != "" { + aspectRatio = opts.AspectRatio + } + + fmt.Printf("Gemini Image Generation Prompt (%d chars): %s\n", len(prompt), prompt) + fmt.Printf("Gemini Image Generation: Using model %q with aspect ratio %q\n", c.modelName(), aspectRatio) + + var imageBytes []byte + var mimeType string + if len(opts.ReferenceImages) > 0 { + imageBytes, mimeType, err = c.generateImageWithRefs(ctx, prompt, aspectRatio, opts.ReferenceImages) + } else { + imageBytes, mimeType, err = geminiGenerateImage(ctx, c, prompt, aspectRatio) + } + if err != nil { + if searchErr, ok := err.(*SearchError); ok { + return nil, searchErr + } + return nil, &SearchError{ + Provider: geminiSource, + Code: "API_ERROR", + Message: fmt.Sprintf("failed to generate image: %v", err), + } + } + + dataURL, err := encodeDataURL(imageBytes, mimeType) + if err != nil { + return nil, err + } + width, height, err := decodedImageDimensions(imageBytes) + if err != nil { + return nil, err + } + + description := fmt.Sprintf("Generated educational image for %s", opts.Query) + if translatedWord != "" { + description = fmt.Sprintf("%s (%s)", description, translatedWord) + } + + result := SearchResult{ + ID: c.generateImageID(opts.Query), + URL: dataURL, + ThumbnailURL: dataURL, + Width: width, + Height: height, + Description: description, + Attribution: "Generated by Google Gemini Nano Banana", + Source: geminiSource, + } + + return []SearchResult{result}, nil +} + +// Download returns the image bytes for a data URI or a remote URL. +func (c *GeminiProvider) Download(ctx context.Context, url string) (io.ReadCloser, error) { + if strings.HasPrefix(url, geminiDataPrefix) { + return decodeDataURL(url) + } + + req, err := http.NewRequestWithContext(ctx, http.MethodGet, url, nil) + if err != nil { + return nil, err + } + + resp, err := imageHTTPClient.Do(req) + if err != nil { + return nil, err + } + + if resp.StatusCode != http.StatusOK { + if closeErr := resp.Body.Close(); closeErr != nil { + return nil, fmt.Errorf("HTTP %d: %s (failed to close response body: %v)", resp.StatusCode, resp.Status, closeErr) + } + return nil, fmt.Errorf("HTTP %d: %s", resp.StatusCode, resp.Status) + } + + return resp.Body, nil +} + +// GetAttribution returns attribution text for the generated image. +func (c *GeminiProvider) GetAttribution(result *SearchResult) string { + width, height := 0, 0 + if result != nil { + width = result.Width + height = result.Height + } + + var attribution strings.Builder + attribution.WriteString("Image generated by Google Gemini Nano Banana\n\n") + fmt.Fprintf(&attribution, "Model: %s\n", c.modelName()) + fmt.Fprintf(&attribution, "Text model: %s\n", c.textModelName()) + fmt.Fprintf(&attribution, "Aspect ratio: %s\n", geminiAspectRatio) + fmt.Fprintf(&attribution, "Size: %dx%d\n", width, height) + if result != nil && result.Description != "" { + fmt.Fprintf(&attribution, "Result: %s\n", result.Description) + } + fmt.Fprintf(&attribution, "\nPrompt used:\n%s\n", c.lastPrompt) + fmt.Fprintf(&attribution, "\nGenerated at: %s\n", time.Now().Format("2006-01-02 15:04:05")) + return attribution.String() +} + +// Name returns the provider name. +func (c *GeminiProvider) Name() string { + return geminiSource +} + +// LastPrompt returns the most recent image prompt. +func (c *GeminiProvider) LastPrompt() string { + return c.lastPrompt +} + +// SetPromptCallback registers a callback that runs after prompt generation. +func (c *GeminiProvider) SetPromptCallback(callback func(prompt string)) { + c.PromptCallback = callback +} + +func (c *GeminiProvider) ensureReady() error { + if c == nil || c.config == nil { + return &SearchError{ + Provider: geminiSource, + Code: "NO_CONFIG", + Message: "Gemini client not initialized", + } + } + if c.config.APIKey == "" { + return &SearchError{ + Provider: geminiSource, + Code: "NO_API_KEY", + Message: "Google API key not configured", + } + } + if c.initErr != nil { + return &SearchError{ + Provider: geminiSource, + Code: "CLIENT_INIT_FAILED", + Message: fmt.Sprintf("failed to initialize client: %v", c.initErr), + } + } + if c.client == nil { + return &SearchError{ + Provider: geminiSource, + Code: "CLIENT_NOT_READY", + Message: "Gemini client not initialized", + } + } + + return nil +} + +func (c *GeminiProvider) resolveTranslation(_ context.Context, opts *SearchOptions, translation string) (string, error) { + if translation != "" { + fmt.Printf("Using provided translation: %s -> %s\n", opts.Query, translation) + return translation, nil + } + + return opts.Query, nil +} + +func (c *GeminiProvider) resolvePrompt(ctx context.Context, opts *SearchOptions, translatedWord string) (string, error) { + if customPrompt := strings.TrimSpace(opts.CustomPrompt); customPrompt != "" { + if len(customPrompt) > maxCustomPrompt { + customPrompt = customPrompt[:maxCustomPrompt-3] + "..." + } + fmt.Printf("Using custom prompt: %s\n", customPrompt) + return customPrompt, nil + } + + return c.createEducationalPrompt(ctx, opts.Query, translatedWord), nil +} + +func (c *GeminiProvider) buildPrompt(ctx context.Context, opts *SearchOptions) (string, string, error) { + if opts == nil { + return "", "", &SearchError{ + Provider: geminiSource, + Code: "INVALID_OPTIONS", + Message: "search options are required", + } + } + + translation := strings.TrimSpace(opts.Translation) + if customPrompt := strings.TrimSpace(opts.CustomPrompt); customPrompt != "" { + if len(customPrompt) > maxCustomPrompt { + customPrompt = customPrompt[:maxCustomPrompt-3] + "..." + } + fmt.Printf("Using custom prompt: %s\n", customPrompt) + return customPrompt, translation, nil + } + + translatedWord, err := c.resolveTranslation(ctx, opts, translation) + if err != nil { + return "", "", err + } + + prompt, err := c.resolvePrompt(ctx, opts, translatedWord) + if err != nil { + return "", "", err + } + + return prompt, translatedWord, nil +} + +// createEducationalPrompt generates a prompt optimized for image generation. +func (c *GeminiProvider) createEducationalPrompt(ctx context.Context, query, translation string) string { + subject := promptSubject(translation, query) + + scene, err := c.generateSceneDescription(ctx, query, translation) + if err != nil { + fmt.Printf(" Failed to generate scene: %v, using basic prompt\n", err) + scene = "" + } + if scene != "" { + scene = sanitizeSceneDescription(scene) + if !usableSceneDescription(scene) { + fmt.Printf(" Scene response was too short or generic, using basic prompt\n") + scene = "" + } + } + + selectedStyle := chooseArtisticStyle() + if selectedStyle == defaultArtisticStyle { + fmt.Printf(" No artistic styles available, using generic prompt\n") + } + fmt.Printf(" Using image style: %s\n", selectedStyle) + + return buildEducationalPrompt(selectedStyle, scene, subject) +} + +func (c *GeminiProvider) generateSceneDescription(ctx context.Context, query, translation string) (string, error) { + fmt.Printf("Gemini Scene Generation: Creating scene for %q (%s)\n", query, translation) + + scene, err := geminiGenerateText( + ctx, + c, + c.textModelName(), + "You are helping create educational flashcards for language learning. Generate a brief, vivid scene description that incorporates the given English word in a memorable, contextual way. The scene should be visually interesting and help with memory retention. Keep it to 1-2 sentences, focusing on visual elements that can be illustrated. The subject (the English word) should be the clear focal point of the image, prominent and centered.", + fmt.Sprintf("Create a scene description for the English word %q that would make a memorable flashcard image. Make sure %q is the main focus and most prominent element in the scene.", translation, translation), + 0.7, + 100, + ) + if err != nil { + return "", fmt.Errorf("scene generation failed: %w", err) + } + scene = sanitizeSceneDescription(scene) + if !usableSceneDescription(scene) { + return "", fmt.Errorf("scene generation returned unusable content") + } + + fmt.Printf("Generated scene: %s\n", scene) + return scene, nil +} + +func (c *GeminiProvider) generateText(ctx context.Context, model, systemPrompt, userPrompt string, temperature float32, maxOutputTokens int32) (string, error) { + temp := temperature + resp, err := apicircuit.Execute(nil, c.Name(), apicircuit.CapabilityText, func() (*genai.GenerateContentResponse, error) { + return c.client.Models.GenerateContent(ctx, model, []*genai.Content{ + genai.NewContentFromText(userPrompt, genai.RoleUser), + }, &genai.GenerateContentConfig{ + SystemInstruction: genai.NewContentFromText(systemPrompt, genai.RoleUser), + Temperature: &temp, + MaxOutputTokens: maxOutputTokens, + }) + }) + if err != nil { + return "", fmt.Errorf("gemini API error: %w", err) + } + + text := strings.TrimSpace(resp.Text()) + if text == "" { + return "", fmt.Errorf("no response received") + } + + return text, nil +} + +func (c *GeminiProvider) generateImage(ctx context.Context, prompt, aspectRatio string) ([]byte, string, error) { + if aspectRatio == "" { + aspectRatio = geminiAspectRatio + } + + cfg := &genai.GenerateContentConfig{ + ResponseModalities: []string{string(genai.ModalityImage)}, + ImageConfig: &genai.ImageConfig{ + AspectRatio: aspectRatio, + }, + } + + resp, err := apicircuit.Execute(nil, c.Name(), apicircuit.CapabilityImage, func() (*genai.GenerateContentResponse, error) { + return c.client.Models.GenerateContent(ctx, c.modelName(), []*genai.Content{ + genai.NewContentFromText(prompt, genai.RoleUser), + }, cfg) + }) + if err != nil { + return nil, "", &SearchError{ + Provider: geminiSource, + Code: "API_ERROR", + Message: fmt.Sprintf("failed to generate image: %v", err), + } + } + + imageBytes, mimeType, err := extractGeneratedImage(resp) + if err != nil { + return nil, "", &SearchError{ + Provider: geminiSource, + Code: "NO_RESULTS", + Message: err.Error(), + } + } + + return imageBytes, mimeType, nil +} + +func (c *GeminiProvider) generateImageWithRefs(ctx context.Context, prompt, aspectRatio string, refs [][]byte) ([]byte, string, error) { + if aspectRatio == "" { + aspectRatio = geminiAspectRatio + } + + cfg := &genai.GenerateContentConfig{ + ResponseModalities: []string{string(genai.ModalityImage)}, + ImageConfig: &genai.ImageConfig{AspectRatio: aspectRatio}, + } + + parts := make([]*genai.Part, 0, len(refs)+1) + for _, ref := range refs { + if len(ref) > 0 { + parts = append(parts, &genai.Part{ + InlineData: &genai.Blob{MIMEType: "image/png", Data: ref}, + }) + } + } + refNote := fmt.Sprintf( + "The %d reference image(s) above show the exact character appearance that must be preserved. "+ + "Every character, animal, or object must look identical in the new image. Now generate:\n\n", + len(refs), + ) + parts = append(parts, &genai.Part{Text: refNote + prompt}) + + resp, err := apicircuit.Execute(nil, c.Name(), apicircuit.CapabilityImage, func() (*genai.GenerateContentResponse, error) { + return c.client.Models.GenerateContent(ctx, c.modelName(), []*genai.Content{ + { + Role: string(genai.RoleUser), + Parts: parts, + }, + }, cfg) + }) + if err != nil { + return nil, "", &SearchError{ + Provider: geminiSource, + Code: "API_ERROR", + Message: fmt.Sprintf("failed to generate image with refs: %v", err), + } + } + + imageBytes, mimeType, err := extractGeneratedImage(resp) + if err != nil { + return nil, "", &SearchError{ + Provider: geminiSource, + Code: "NO_RESULTS", + Message: err.Error(), + } + } + + return imageBytes, mimeType, nil +} + +func extractGeneratedImage(response *genai.GenerateContentResponse) ([]byte, string, error) { + if response == nil { + return nil, "", fmt.Errorf("no response from Gemini") + } + + for _, candidate := range response.Candidates { + if candidate == nil || candidate.Content == nil { + continue + } + + for _, part := range candidate.Content.Parts { + if part == nil || part.InlineData == nil || len(part.InlineData.Data) == 0 { + continue + } + + mimeType := part.InlineData.MIMEType + if mimeType == "" { + mimeType = "image/png" + } + + return append([]byte(nil), part.InlineData.Data...), mimeType, nil + } + } + + return nil, "", fmt.Errorf("no image data returned from Gemini") +} + +func encodeDataURL(imageBytes []byte, mimeType string) (string, error) { + if len(imageBytes) == 0 { + return "", fmt.Errorf("no image bytes returned") + } + + normalizedBytes, err := normalizePNG(imageBytes, mimeType) + if err != nil { + return "", err + } + + return geminiDataPrefix + base64.StdEncoding.EncodeToString(normalizedBytes), nil +} + +func decodeDataURL(url string) (io.ReadCloser, error) { + header, payload, ok := strings.Cut(url, ",") + if !ok || !strings.HasPrefix(header, "data:") || !strings.Contains(header, ";base64") { + return nil, fmt.Errorf("unsupported data URI: %s", url) + } + + data, err := base64.StdEncoding.DecodeString(payload) + if err != nil { + return nil, fmt.Errorf("decode data URI: %w", err) + } + + return io.NopCloser(bytes.NewReader(data)), nil +} + +func normalizePNG(imageBytes []byte, mimeType string) ([]byte, error) { + if strings.EqualFold(strings.TrimSpace(mimeType), "image/png") { + return append([]byte(nil), imageBytes...), nil + } + + img, _, err := image.Decode(bytes.NewReader(imageBytes)) + if err != nil { + return nil, fmt.Errorf("decode generated image: %w", err) + } + + var buffer bytes.Buffer + if err := png.Encode(&buffer, img); err != nil { + return nil, fmt.Errorf("encode generated image as png: %w", err) + } + + return buffer.Bytes(), nil +} + +func decodedImageDimensions(imageBytes []byte) (int, int, error) { + cfg, _, err := image.DecodeConfig(bytes.NewReader(imageBytes)) + if err != nil { + return 0, 0, fmt.Errorf("decode generated image dimensions: %w", err) + } + + return cfg.Width, cfg.Height, nil +} + +func (c *GeminiProvider) generateImageID(word string) string { + hash := md5.Sum([]byte(word)) + return hex.EncodeToString(hash[:])[:8] +} + +func (c *GeminiProvider) modelName() string { + if c == nil || c.config == nil || strings.TrimSpace(c.config.Model) == "" { + return DefaultGeminiImageModel + } + + return c.config.Model +} + +func (c *GeminiProvider) textModelName() string { + if c == nil || c.config == nil || strings.TrimSpace(c.config.TextModel) == "" { + return DefaultGeminiTextModel + } + + return c.config.TextModel +} + +func normalizeGeminiConfig(config *GeminiConfig) *GeminiConfig { + normalized := &GeminiConfig{} + if config != nil { + *normalized = *config + } + + normalized.APIKey = strings.TrimSpace(normalized.APIKey) + normalized.Model = strings.TrimSpace(normalized.Model) + normalized.TextModel = strings.TrimSpace(normalized.TextModel) + + if normalized.Model == "" { + normalized.Model = DefaultGeminiImageModel + } + if normalized.TextModel == "" { + normalized.TextModel = DefaultGeminiTextModel + } + + return normalized +} diff --git a/internal/image/gemini_test.go b/internal/image/gemini_test.go new file mode 100644 index 0000000..8c34c6e --- /dev/null +++ b/internal/image/gemini_test.go @@ -0,0 +1,333 @@ +package image + +import ( + "bytes" + "context" + "encoding/base64" + "io" + "net/http" + "net/http/httptest" + "strings" + "testing" + + "google.golang.org/genai" +) + +func TestNewGeminiProvider(t *testing.T) { + t.Parallel() + + client := NewGeminiProvider(&GeminiConfig{APIKey: "test-key"}) + if client == nil { + t.Fatal("expected client") + } + if client.config == nil { + t.Fatal("expected normalized config") + } + if client.config.Model != DefaultGeminiImageModel { + t.Fatalf("expected default model %q, got %q", DefaultGeminiImageModel, client.config.Model) + } + if client.config.TextModel != DefaultGeminiTextModel { + t.Fatalf("expected default text model %q, got %q", DefaultGeminiTextModel, client.config.TextModel) + } + if client.Name() != Gemini { + t.Fatalf("Name() = %q, want %q", client.Name(), Gemini) + } +} + +func TestGeminiProvider_NoAPIKey(t *testing.T) { + client := NewGeminiProvider(&GeminiConfig{}) + + _, err := client.Search(context.Background(), DefaultSearchOptions("ябълка")) + if err == nil { + t.Fatal("expected error for missing API key") + } + + searchErr, ok := err.(*SearchError) + if !ok { + t.Fatalf("expected SearchError, got %T", err) + } + if searchErr.Code != "NO_API_KEY" { + t.Fatalf("expected NO_API_KEY error, got %s", searchErr.Code) + } +} + +func TestGeminiProvider_Search_CustomPromptSkipsTextGeneration(t *testing.T) { + originalText := geminiGenerateText + originalImage := geminiGenerateImage + t.Cleanup(func() { + geminiGenerateText = originalText + geminiGenerateImage = or |
