feat: support for image generation

core/internal/testutil/account.go

@@ -41,6 +41,8 @@ type TestScenarios struct {
 	Embedding             bool // Embedding functionality
 	Reasoning             bool // Reasoning/thinking functionality via Responses API
 	ListModels            bool // List available models functionality
+	ImageGeneration       bool // Image generation functionality
+	ImageGenerationStream bool // Streaming image generation functionality
 }
 
 // ComprehensiveTestConfig extends TestConfig with additional scenarios
@@ -61,6 +63,8 @@ type ComprehensiveTestConfig struct {
 	SpeechSynthesisFallbacks []schemas.Fallback // for speech synthesis tests
 	EmbeddingFallbacks       []schemas.Fallback // for embedding tests
 	SkipReason               string             // Reason to skip certain tests
+	ImageGenerationModel     string             // Model for image generation
+	ImageGenerationFallbacks []schemas.Fallback // Fallbacks for image generation
 }
 
 // ComprehensiveTestAccount provides a test implementation of the Account interface for comprehensive testing.
@@ -517,6 +521,7 @@ var AllProviderConfigs = []ComprehensiveTestConfig{
 		PromptCachingModel:   "gpt-4.1",
 		TranscriptionModel:   "whisper-1",
 		SpeechSynthesisModel: "tts-1",
+		ImageGenerationModel: "dall-e-2",
 		Scenarios: TestScenarios{
 			TextCompletion:        false, // Not supported
 			TextCompletionStream:  false, // Not supported

core/internal/testutil/image_generation.go

@@ -0,0 +1,126 @@
+package testutil
+
+import (
+	"context"
+	"os"
+	"testing"
+
+	bifrost "github.com/maximhq/bifrost/core"
+	"github.com/maximhq/bifrost/core/schemas"
+)
+
+// RunImageGenerationTest executes the end-to-end image generation test (non-streaming)
+func RunImageGenerationTest(t *testing.T, client *bifrost.Bifrost, ctx context.Context, testConfig ComprehensiveTestConfig) {
+	if testConfig.ImageGenerationModel == "" {
+		t.Logf("Image generation not configured for provider %s", testConfig.Provider)
+		return
+	}
+
+	t.Run("ImageGeneration", func(t *testing.T) {
+		if os.Getenv("SKIP_PARALLEL_TESTS") != "true" {
+			t.Parallel()
+		}
+
+		retryConfig := GetTestRetryConfigForScenario("ImageGeneration", testConfig)
+		retryContext := TestRetryContext{
+			ScenarioName:     "ImageGeneration",
+			ExpectedBehavior: map[string]interface{}{},
+			TestMetadata: map[string]interface{}{
+				"provider": testConfig.Provider,
+				"model":    testConfig.ImageGenerationModel,
+			},
+		}
+
+		expectations := GetExpectationsForScenario("ImageGeneration", testConfig, map[string]interface{}{
+			"min_images":    1,
+			"expected_size": "1024x1024",
+		})
+
+		imageGenerationRetryConfig := ImageGenerationRetryConfig{
+			MaxAttempts: retryConfig.MaxAttempts,
+			BaseDelay:   retryConfig.BaseDelay,
+			MaxDelay:    retryConfig.MaxDelay,
+			Conditions:  []ImageGenerationRetryCondition{},
+			OnRetry:     retryConfig.OnRetry,
+			OnFinalFail: retryConfig.OnFinalFail,
+		}
+		// Test basic image generation
+		imageGenerationOperation := func() (*schemas.BifrostImageGenerationResponse, *schemas.BifrostError) {
+			request := &schemas.BifrostImageGenerationRequest{
+				Provider: testConfig.Provider,
+				Model:    testConfig.ImageGenerationModel,
+				Input: &schemas.ImageGenerationInput{
+					Prompt: "A serene Japanese garden with cherry blossoms in spring",
+				},
+				Params: &schemas.ImageGenerationParameters{
+					Size:           bifrost.Ptr("1024x1024"),
+					Quality:        bifrost.Ptr("standard"),
+					ResponseFormat: bifrost.Ptr("b64_json"),
+					N:              bifrost.Ptr(1),
+				},
+				Fallbacks: testConfig.ImageGenerationFallbacks,
+			}
+
+			response, err := client.ImageGenerationRequest(ctx, request)
+			if err != nil {
+				return nil, err
+			}
+			if response != nil {
+				return response, nil
+			}
+			return nil, &schemas.BifrostError{
+				IsBifrostError: true,
+				Error: &schemas.ErrorField{
+					Message: "No image generation response returned",
+				},
+			}
+		}
+
+		imageGenerationResponse, imageGenerationError := WithImageGenerationRetry(t, imageGenerationRetryConfig, retryContext, expectations, "ImageGeneration", imageGenerationOperation)
+
+		if imageGenerationError != nil {
+			t.Fatalf("❌ Image generation failed: %v", GetErrorMessage(imageGenerationError))
+		}
+
+		// Validate response
+		if imageGenerationResponse == nil {
+			t.Fatal("❌ Image generation returned nil response")
+		}
+
+		if len(imageGenerationResponse.Data) == 0 {
+			t.Fatal("❌ Image generation returned no image data")
+		}
+
+		// Validate first image
+		imageData := imageGenerationResponse.Data[0]
+		if imageData.B64JSON == "" && imageData.URL == "" {
+			t.Fatal("❌ Image data missing both b64_json and URL")
+		}
+
+		// Validate base64 if present
+		if imageData.B64JSON != "" {
+			if len(imageData.B64JSON) < 100 {
+				t.Errorf("❌ Base64 image data too short: %d bytes", len(imageData.B64JSON))
+			}
+		}
+
+		// Validate usage if present
+		if imageGenerationResponse.Usage != nil {
+			if imageGenerationResponse.Usage.TotalTokens == 0 {
+				t.Logf("⚠️  Usage total_tokens is 0 (may be provider-specific)")
+			}
+		}
+
+		// Validate extra fields
+		if imageGenerationResponse.ExtraFields.Provider == "" {
+			t.Error("❌ ExtraFields.Provider is empty")
+		}
+
+		if imageGenerationResponse.ExtraFields.ModelRequested == "" {
+			t.Error("❌ ExtraFields.ModelRequested is empty")
+		}
+
+		t.Logf("✅ Image generation successful: ID=%s, Provider=%s, Model=%s, Images=%d",
+			imageGenerationResponse.ID, imageGenerationResponse.ExtraFields.Provider, imageGenerationResponse.ExtraFields.ModelRequested, len(imageGenerationResponse.Data))
+	})
+}

core/internal/testutil/image_generation_cache.go

@@ -0,0 +1,140 @@
+package testutil
+
+import (
+	"context"
+	"os"
+	"testing"
+	"time"
+
+	bifrost "github.com/maximhq/bifrost/core"
+	"github.com/maximhq/bifrost/core/schemas"
+)
+
+// RunImageGenerationCacheTest tests cache hit/miss scenarios
+func RunImageGenerationCacheTest(t *testing.T, client *bifrost.Bifrost, ctx context.Context, testConfig ComprehensiveTestConfig) {
+	if testConfig.ImageGenerationModel == "" {
+		t.Logf("Image generation cache test skipped: not configured for provider %s", testConfig.Provider)
+		return
+	}
+
+	t.Run("ImageGenerationCache", func(t *testing.T) {
+		if os.Getenv("SKIP_PARALLEL_TESTS") != "true" {
+			t.Parallel()
+		}
+
+		// Use a unique prompt for cache testing
+		cacheTestPrompt := "A unique test image for cache validation - " + time.Now().Format("20060102150405")
+
+		request := &schemas.BifrostImageGenerationRequest{
+			Provider: testConfig.Provider,
+			Model:    testConfig.ImageGenerationModel,
+			Input: &schemas.ImageGenerationInput{
+				Prompt: cacheTestPrompt,
+			},
+			Params: &schemas.ImageGenerationParameters{
+				Size:           bifrost.Ptr("1024x1024"),
+				ResponseFormat: bifrost.Ptr("b64_json"),
+			},
+		}
+
+		// First request - should be a cache miss
+		start1 := time.Now()
+		response1, err1 := client.ImageGenerationRequest(ctx, request)
+		duration1 := time.Since(start1)
+
+		if err1 != nil {
+			t.Fatalf("❌ First image generation request failed: %v", GetErrorMessage(err1))
+		}
+
+		if response1 == nil || len(response1.Data) == 0 {
+			t.Fatal("❌ First request returned no image data")
+		}
+
+		// Check cache debug info if available
+		cacheHit1 := false
+		if response1.ExtraFields.CacheDebug != nil {
+			cacheHit1 = response1.ExtraFields.CacheDebug.CacheHit
+		}
+
+		if cacheHit1 {
+			t.Logf("⚠️  First request was a cache hit (unexpected, but may be valid)")
+		} else {
+			t.Logf("✅ First request was a cache miss (expected)")
+		}
+
+		// Second request with same prompt - should be a cache hit
+		start2 := time.Now()
+		response2, err2 := client.ImageGenerationRequest(ctx, request)
+		duration2 := time.Since(start2)
+
+		if err2 != nil {
+			t.Fatalf("❌ Second image generation request failed: %v", GetErrorMessage(err2))
+		}
+
+		if response2 == nil || len(response2.Data) == 0 {
+			t.Fatal("❌ Second request returned no image data")
+		}
+
+		// Check cache debug info
+		cacheHit2 := false
+		if response2.ExtraFields.CacheDebug != nil {
+			cacheHit2 = response2.ExtraFields.CacheDebug.CacheHit
+		}
+
+		if cacheHit2 {
+			t.Logf("✅ Second request was a cache hit (expected)")
+
+			// Cache hit should be faster
+			if duration2 < duration1 {
+				t.Logf("✅ Cache hit was faster: %v vs %v", duration2, duration1)
+			} else {
+				t.Logf("⚠️  Cache hit was not faster: %v vs %v (may be due to network variance)", duration2, duration1)
+			}
+
+			// Validate cached response matches original
+			if len(response1.Data) == len(response2.Data) {
+				// Compare image data (should be identical for cache hit)
+				if response1.Data[0].B64JSON != "" && response2.Data[0].B64JSON != "" {
+					if response1.Data[0].B64JSON == response2.Data[0].B64JSON {
+						t.Logf("✅ Cached image data matches original")
+					} else {
+						t.Errorf("❌ Cached image data does not match original")
+					}
+				}
+			}
+		} else {
+			t.Logf("⚠️  Second request was a cache miss (cache may not be enabled or TTL expired)")
+		}
+
+		// Test with different prompt - should be cache miss
+		request2 := &schemas.BifrostImageGenerationRequest{
+			Provider: testConfig.Provider,
+			Model:    testConfig.ImageGenerationModel,
+			Input: &schemas.ImageGenerationInput{
+				Prompt: "A different prompt for cache miss test",
+			},
+			Params: &schemas.ImageGenerationParameters{
+				Size:           bifrost.Ptr("1024x1024"),
+				ResponseFormat: bifrost.Ptr("b64_json"),
+			},
+		}
+
+		response3, err3 := client.ImageGenerationRequest(ctx, request2)
+		if err3 != nil {
+			t.Fatalf("❌ Third image generation request failed: %v", GetErrorMessage(err3))
+		}
+
+		cacheHit3 := false
+		if response3.ExtraFields.CacheDebug != nil {
+			cacheHit3 = response3.ExtraFields.CacheDebug.CacheHit
+		}
+
+		if cacheHit3 {
+			t.Logf("⚠️  Different prompt was a cache hit (unexpected)")
+		} else {
+			t.Logf("✅ Different prompt was a cache miss (expected)")
+		}
+
+		t.Logf("✅ Cache test completed: First=%v, Second=%v, Different=%v", cacheHit1, cacheHit2, cacheHit3)
+	})
+}