From d59813eb7ec44e8f57f75ece502a2360ab939816 Mon Sep 17 00:00:00 2001 From: Paul Buetow Date: Wed, 1 Apr 2026 19:59:35 +0300 Subject: z8: wire Nano Banana into processor --- internal/processor/processor.go | 144 +++++++++++++++++++++++------------ internal/processor/processor_test.go | 109 ++++++++++++++++++++++++++ 2 files changed, 206 insertions(+), 47 deletions(-) (limited to 'internal/processor') diff --git a/internal/processor/processor.go b/internal/processor/processor.go index 935ce60..51f530e 100644 --- a/internal/processor/processor.go +++ b/internal/processor/processor.go @@ -30,6 +30,14 @@ type Processor struct { phoneticFetcher *phonetic.Fetcher } +var newOpenAIImageClient = func(config *image.OpenAIConfig) image.ImageSearcher { + return image.NewOpenAIClient(config) +} + +var newNanoBananaImageClient = func(config *image.NanoBananaConfig) image.ImageSearcher { + return image.NewNanoBananaClient(config) +} + // NewProcessor creates a new word processor func NewProcessor(flags *cli.Flags) *Processor { openAIKey := cli.GetOpenAIKey() @@ -371,41 +379,9 @@ func (p *Processor) generateAudioWithVoiceAndFilenameInDir(word, voice, filename // downloadImagesWithTranslation downloads images for a word func (p *Processor) downloadImagesWithTranslation(word, translationText string) error { - // Create image searcher based on provider - var searcher image.ImageSearcher - - switch p.flags.ImageAPI { - case "openai": - // Create OpenAI image configuration - openaiConfig := &image.OpenAIConfig{ - APIKey: cli.GetOpenAIKey(), - Model: p.flags.OpenAIImageModel, - Size: p.flags.OpenAIImageSize, - Quality: p.flags.OpenAIImageQuality, - Style: p.flags.OpenAIImageStyle, - } - - // Use config file values if not overridden by flags - if p.flags.OpenAIImageModel == "dall-e-2" && viper.IsSet("image.openai_model") { - openaiConfig.Model = viper.GetString("image.openai_model") - } - if p.flags.OpenAIImageSize == "512x512" && viper.IsSet("image.openai_size") { - openaiConfig.Size = viper.GetString("image.openai_size") - } - if p.flags.OpenAIImageQuality == "standard" && viper.IsSet("image.openai_quality") { - openaiConfig.Quality = viper.GetString("image.openai_quality") - } - if p.flags.OpenAIImageStyle == "natural" && viper.IsSet("image.openai_style") { - openaiConfig.Style = viper.GetString("image.openai_style") - } - - searcher = image.NewOpenAIClient(openaiConfig) - if openaiConfig.APIKey == "" { - return fmt.Errorf("OpenAI API key is required for image generation") - } - - default: - return fmt.Errorf("unknown image provider: %s", p.flags.ImageAPI) + searcher, err := p.newImageSearcher() + if err != nil { + return err } // Find existing card directory or create new one @@ -436,18 +412,7 @@ func (p *Processor) downloadImagesWithTranslation(word, translationText string) } fmt.Printf(" Downloaded: %s\n", path) - // If using OpenAI, save the prompt - if p.flags.ImageAPI == "openai" { - if openaiClient, ok := searcher.(*image.OpenAIClient); ok { - usedPrompt := openaiClient.GetLastPrompt() - if usedPrompt != "" { - promptFile := filepath.Join(wordDir, "image_prompt.txt") - if err := os.WriteFile(promptFile, []byte(usedPrompt), 0644); err != nil { - fmt.Printf(" Warning: Failed to save image prompt: %v\n", err) - } - } - } - } + p.saveImagePrompt(wordDir, searcher) return nil } @@ -551,6 +516,9 @@ func (p *Processor) RunGUIMode() error { // User explicitly set a different output directory guiConfig.OutputDir = p.flags.OutputDir } + if guiConfig.GoogleAPIKey == "" { + guiConfig.GoogleAPIKey = cli.GetGoogleAPIKey() + } // Otherwise, gui.New will use its own default (XDG state directory) // Create and run GUI application @@ -577,6 +545,88 @@ func (p *Processor) guiConfigForRunMode() *gui.Config { } } +func (p *Processor) newImageSearcher() (image.ImageSearcher, error) { + switch p.flags.ImageAPI { + case "openai": + return p.newOpenAIImageSearcher() + case "nanobanana": + return p.newNanoBananaImageSearcher() + default: + return nil, fmt.Errorf("unknown image provider: %s", p.flags.ImageAPI) + } +} + +func (p *Processor) newOpenAIImageSearcher() (image.ImageSearcher, error) { + openaiConfig := &image.OpenAIConfig{ + APIKey: cli.GetOpenAIKey(), + Model: p.flags.OpenAIImageModel, + Size: p.flags.OpenAIImageSize, + Quality: p.flags.OpenAIImageQuality, + Style: p.flags.OpenAIImageStyle, + } + + if p.flags.OpenAIImageModel == "dall-e-2" && viper.IsSet("image.openai_model") { + openaiConfig.Model = viper.GetString("image.openai_model") + } + if p.flags.OpenAIImageSize == "512x512" && viper.IsSet("image.openai_size") { + openaiConfig.Size = viper.GetString("image.openai_size") + } + if p.flags.OpenAIImageQuality == "standard" && viper.IsSet("image.openai_quality") { + openaiConfig.Quality = viper.GetString("image.openai_quality") + } + if p.flags.OpenAIImageStyle == "natural" && viper.IsSet("image.openai_style") { + openaiConfig.Style = viper.GetString("image.openai_style") + } + + if openaiConfig.APIKey == "" { + return nil, fmt.Errorf("OpenAI API key is required for image generation") + } + + return newOpenAIImageClient(openaiConfig), nil +} + +func (p *Processor) newNanoBananaImageSearcher() (image.ImageSearcher, error) { + nanoBananaConfig := &image.NanoBananaConfig{ + APIKey: cli.GetGoogleAPIKey(), + Model: p.flags.NanoBananaModel, + TextModel: p.flags.NanoBananaTextModel, + } + + if p.flags.NanoBananaModel == image.DefaultNanoBananaModel && viper.IsSet("image.nanobanana_model") { + nanoBananaConfig.Model = viper.GetString("image.nanobanana_model") + } + if p.flags.NanoBananaTextModel == image.DefaultNanoBananaTextModel && viper.IsSet("image.nanobanana_text_model") { + nanoBananaConfig.TextModel = viper.GetString("image.nanobanana_text_model") + } + + if nanoBananaConfig.APIKey == "" { + return nil, fmt.Errorf("Google API key is required for image generation") + } + + return newNanoBananaImageClient(nanoBananaConfig), nil +} + +func (p *Processor) saveImagePrompt(wordDir string, searcher image.ImageSearcher) { + type promptGetter interface { + GetLastPrompt() string + } + + promptSource, ok := searcher.(promptGetter) + if !ok { + return + } + + usedPrompt := promptSource.GetLastPrompt() + if usedPrompt == "" { + return + } + + promptFile := filepath.Join(wordDir, "image_prompt.txt") + if err := os.WriteFile(promptFile, []byte(usedPrompt), 0644); err != nil { + fmt.Printf(" Warning: Failed to save image prompt: %v\n", err) + } +} + // Helper methods func (p *Processor) findOrCreateWordDirectory(word string) string { diff --git a/internal/processor/processor_test.go b/internal/processor/processor_test.go index 7ae475a..982393d 100644 --- a/internal/processor/processor_test.go +++ b/internal/processor/processor_test.go @@ -1,16 +1,61 @@ package processor import ( + "context" + "io" "os" "path/filepath" + "strings" "testing" "codeberg.org/snonux/totalrecall/internal/cli" "codeberg.org/snonux/totalrecall/internal/gui" + "codeberg.org/snonux/totalrecall/internal/image" "codeberg.org/snonux/totalrecall/internal/phonetic" "github.com/spf13/viper" ) +type stubImageSearcher struct { + lastPrompt string + searchErr error +} + +func (s *stubImageSearcher) Search(ctx context.Context, opts *image.SearchOptions) ([]image.SearchResult, error) { + if s.searchErr != nil { + return nil, s.searchErr + } + + s.lastPrompt = "stub nanobanana prompt" + return []image.SearchResult{ + { + ID: "stub-image", + URL: "data:image/png;base64,AAAA", + ThumbnailURL: "data:image/png;base64,AAAA", + Width: 1, + Height: 1, + Description: "stub image", + Attribution: "stub attribution", + Source: "nanobanana", + }, + }, nil +} + +func (s *stubImageSearcher) Download(ctx context.Context, url string) (io.ReadCloser, error) { + return io.NopCloser(strings.NewReader("mock image data")), nil +} + +func (s *stubImageSearcher) GetAttribution(result *image.SearchResult) string { + return result.Attribution +} + +func (s *stubImageSearcher) Name() string { + return "nanobanana" +} + +func (s *stubImageSearcher) GetLastPrompt() string { + return s.lastPrompt +} + func TestNewProcessor(t *testing.T) { t.Setenv("OPENAI_API_KEY", "test-openai-key") t.Setenv("GOOGLE_API_KEY", "test-google-key") @@ -156,6 +201,9 @@ func TestGUIConfigForRunModeUsesNanoBananaDefaultWhenImageAPIIsNotSpecified(t *t if guiConfig.AudioFormat != "wav" { t.Fatalf("guiConfig.AudioFormat = %q, want %q", guiConfig.AudioFormat, "wav") } + if guiConfig.GoogleAPIKey != "test-google-key" { + t.Fatalf("guiConfig.GoogleAPIKey = %q, want %q", guiConfig.GoogleAPIKey, "test-google-key") + } } func TestGUIConfigForRunModeHonorsExplicitImageAPI(t *testing.T) { @@ -178,6 +226,67 @@ func TestGUIConfigForRunModeHonorsExplicitImageAPI(t *testing.T) { if guiConfig.ImageProvider != "openai" { t.Fatalf("guiConfig.ImageProvider = %q, want %q", guiConfig.ImageProvider, "openai") } + if guiConfig.GoogleAPIKey != "test-google-key" { + t.Fatalf("guiConfig.GoogleAPIKey = %q, want %q", guiConfig.GoogleAPIKey, "test-google-key") + } +} + +func TestDownloadImagesWithTranslationUsesNanoBananaConfigAndSavesPrompt(t *testing.T) { + t.Setenv("OPENAI_API_KEY", "test-openai-key") + t.Setenv("GOOGLE_API_KEY", "test-google-key") + + originalConfig := viper.New() + *originalConfig = *viper.GetViper() + defer func() { + *viper.GetViper() = *originalConfig + }() + viper.Reset() + viper.Set("image.nanobanana_model", "custom-image-model") + viper.Set("image.nanobanana_text_model", "custom-text-model") + + originalConstructor := newNanoBananaImageClient + stubSearcher := &stubImageSearcher{} + capturedConfig := new(image.NanoBananaConfig) + newNanoBananaImageClient = func(config *image.NanoBananaConfig) image.ImageSearcher { + *capturedConfig = *config + return stubSearcher + } + t.Cleanup(func() { + newNanoBananaImageClient = originalConstructor + }) + + flags := cli.NewFlags() + flags.OutputDir = t.TempDir() + flags.ImageAPI = "nanobanana" + flags.ImageAPISpecified = true + + p := NewProcessor(flags) + if err := p.downloadImagesWithTranslation("ябълка", "apple"); err != nil { + t.Fatalf("downloadImagesWithTranslation() unexpected error: %v", err) + } + + if capturedConfig.APIKey != "test-google-key" { + t.Fatalf("NanoBanana APIKey = %q, want %q", capturedConfig.APIKey, "test-google-key") + } + if capturedConfig.Model != "custom-image-model" { + t.Fatalf("NanoBanana Model = %q, want %q", capturedConfig.Model, "custom-image-model") + } + if capturedConfig.TextModel != "custom-text-model" { + t.Fatalf("NanoBanana TextModel = %q, want %q", capturedConfig.TextModel, "custom-text-model") + } + + wordDir := p.findCardDirectory("ябълка") + if wordDir == "" { + t.Fatal("expected word directory to be created") + } + + promptData, err := os.ReadFile(filepath.Join(wordDir, "image_prompt.txt")) + if err != nil { + t.Fatalf("expected prompt file: %v", err) + } + if got := strings.TrimSpace(string(promptData)); got != stubSearcher.GetLastPrompt() { + t.Fatalf("prompt file = %q, want %q", got, stubSearcher.GetLastPrompt()) + } } func TestProcessSingleWord_InvalidWord(t *testing.T) { -- cgit v1.2.3