feat: support for image generation (#980)

This PR introduces image generation support to Bifrost and wires it into the existing gateway infrastructure. It adds core schemas, provider extensions, HTTP transport, streaming plumbing, cache integration, and initial UI support to address #950.

Implementation is in progress. Completed functionality is marked, and remaining work will be completed in follow-up commits before marking this PR ready for review.

Scoped Changes:
- [x] New /v1/images/generations endpoint (OpenAI-compatible)
- [x] Image generation via Chat Completion API (tool use pattern)
- [x] Image generation via Responses API (native support)
- [x] Streaming image delivery (base64 chunks)
- [x] Semantic caching for image generation
- [x] UI components for image rendering
- [x] Provider implementations: OpenAI DALL-E, Azure DALL-E

Rollout Plan:

- [x] Phase 1: Core schema and provider implementation (OpenAI + Azure)
- [x] Phase 2: HTTP transport and non-streaming endpoint
- [x] Phase 3: Streaming support and accumulator
- [x] Phase 4: Semantic cache integration (Base64 storage, 5min TTL)
- [x] Phase 5: UI components and documentation

- [ ] Bug fix
- [x] Feature
- [ ] Refactor
- [x] Documentation
- [ ] Chore/CI

- [x] Core (Go)
- [x] Transports (HTTP)
- [x] Providers/Integrations
- [x] Plugins
- [x] UI (Next.js)
- [x] Docs

Unit Tests
- [x] Schema serialization/deserialization
- [x] Request transformation (Bifrost → OpenAI format)
- [x] Response transformation (OpenAI → Bifrost format)
- [x] Stream chunk accumulation
- [x] Cache key generation

Integration Tests

- [x] End-to-end image generation (non-streaming)
- [x] End-to-end streaming image generation
- [x] Fallback to secondary provider
- [x] Cache hit/miss scenarios
- [x] Error handling (rate limits, invalid prompts)

Load Tests

- [x] Concurrent image generation requests
- [x] Stream memory usage under load
- [x] Cache performance at scale

End-to-end image generation
```sh
curl -X POST http://localhost:8080/v1/images/generations \
  -H "Content-Type: application/json" \
  -H "Authorization: Bearer $API_KEY" \
  -d '{
    "model": "openai/dall-e-3",
    "prompt": "A serene Japanese garden with cherry blossoms",
    "size": "1024x1024",
    "quality": "high",
    "response_format": "b64_json"
  }'
  ```
  Expected Outcome:
  ```json
  {
  "id": "abc123",
  "created": 1699999999,
  "model": "dall-e-3",
  "data": [
    {
      "b64_json": "iVBORw0KGgo...",
      "revised_prompt": "A tranquil Japanese garden featuring blooming cherry blossom trees...",
      "index": 0
    }
  ],
  "usage": {
    "prompt_tokens": 15,
    "total_tokens": 15
  },
  "extra_fields": {
    "provider": "openai",
    "latency_ms": 8500,
    "cache_debug": null
  }
}
```

Unit Tests:
```sh
go test -v github.com/maximhq/bifrost/core/providers/openai -run TestImage
go test -v github.com/maximhq/bifrost/core -run TestImage
```
Expected outcome: PASS
Note: Might need to replace remote paths with local modules in go.mod (Once PR is merged, this directive is not needed)
```sh
replace github.com/maximhq/bifrost/core => ../../core
```

N/A

- [ ] Yes
- [x] No

Closes #950

N/A

- [x] I read `docs/contributing/README.md` and followed the guidelines
- [x] I added/updated tests where appropriate
- [x] I updated documentation where needed
- [x] I verified builds succeed (Go and UI)
- [ ] I verified the CI pipeline passes locally if applicable

Author: Radheshg04

core/bifrost.go

@@ -1005,6 +1005,94 @@ func (bifrost *Bifrost) TranscriptionStreamRequest(ctx *schemas.BifrostContext,
 	return bifrost.handleStreamRequest(ctx, bifrostReq)
 }
 
+// ImageGenerationRequest sends an image generation request to the specified provider.
+func (bifrost *Bifrost) ImageGenerationRequest(ctx *schemas.BifrostContext,
+	req *schemas.BifrostImageGenerationRequest) (*schemas.BifrostImageGenerationResponse, *schemas.BifrostError) {
+	if req == nil {
+		return nil, &schemas.BifrostError{
+			IsBifrostError: false,
+			Error: &schemas.ErrorField{
+				Message: "image generation request is nil",
+			},
+			ExtraFields: schemas.BifrostErrorExtraFields{
+				RequestType: schemas.ImageGenerationRequest,
+			},
+		}
+	}
+	if req.Input == nil || req.Input.Prompt == "" {
+		return nil, &schemas.BifrostError{
+			IsBifrostError: false,
+			Error: &schemas.ErrorField{
+				Message: "prompt not provided for image generation request",
+			},
+			ExtraFields: schemas.BifrostErrorExtraFields{
+				RequestType:    schemas.ImageGenerationRequest,
+				Provider:       req.Provider,
+				ModelRequested: req.Model,
+			},
+		}
+	}
+
+	bifrostReq := bifrost.getBifrostRequest()
+	bifrostReq.RequestType = schemas.ImageGenerationRequest
+	bifrostReq.ImageGenerationRequest = req
+
+	response, err := bifrost.handleRequest(ctx, bifrostReq)
+	if err != nil {
+		return nil, err
+	}
+	if response == nil || response.ImageGenerationResponse == nil {
+		return nil, &schemas.BifrostError{
+			IsBifrostError: false,
+			Error: &schemas.ErrorField{
+				Message: "received nil response from provider",
+			},
+			ExtraFields: schemas.BifrostErrorExtraFields{
+				RequestType:    schemas.ImageGenerationRequest,
+				Provider:       req.Provider,
+				ModelRequested: req.Model,
+			},
+		}
+	}
+
+	return response.ImageGenerationResponse, nil
+}
+
+// ImageGenerationStreamRequest sends an image generation stream request to the specified provider.
+func (bifrost *Bifrost) ImageGenerationStreamRequest(ctx *schemas.BifrostContext,
+	req *schemas.BifrostImageGenerationRequest) (chan *schemas.BifrostStream, *schemas.BifrostError) {
+	if req == nil {
+		return nil, &schemas.BifrostError{
+			IsBifrostError: false,
+			Error: &schemas.ErrorField{
+				Message: "image generation stream request is nil",
+			},
+			ExtraFields: schemas.BifrostErrorExtraFields{
+				RequestType: schemas.ImageGenerationStreamRequest,
+			},
+		}
+	}
+	if req.Input == nil || req.Input.Prompt == "" {
+		return nil, &schemas.BifrostError{
+			IsBifrostError: false,
+			Error: &schemas.ErrorField{
+				Message: "prompt not provided for image generation stream request",
+			},
+			ExtraFields: schemas.BifrostErrorExtraFields{
+				RequestType:    schemas.ImageGenerationStreamRequest,
+				Provider:       req.Provider,
+				ModelRequested: req.Model,
+			},
+		}
+	}
+
+	bifrostReq := bifrost.getBifrostRequest()
+	bifrostReq.RequestType = schemas.ImageGenerationStreamRequest
+	bifrostReq.ImageGenerationRequest = req
+
+	return bifrost.handleStreamRequest(ctx, bifrostReq)
+}
+
 // BatchCreateRequest creates a new batch job for asynchronous processing.
 func (bifrost *Bifrost) BatchCreateRequest(ctx *schemas.BifrostContext, req *schemas.BifrostBatchCreateRequest) (*schemas.BifrostBatchCreateResponse, *schemas.BifrostError) {
 	if req == nil {
@@ -2320,6 +2408,12 @@ func (bifrost *Bifrost) prepareFallbackRequest(req *schemas.BifrostRequest, fall
 		tmp.Model = fallback.Model
 		fallbackReq.TranscriptionRequest = &tmp
 	}
+	if req.ImageGenerationRequest != nil {
+		tmp := *req.ImageGenerationRequest
+		tmp.Provider = fallback.Provider
+		tmp.Model = fallback.Model
+		fallbackReq.ImageGenerationRequest = &tmp
+	}
 
 	return &fallbackReq
 }
@@ -2761,6 +2855,9 @@ func (bifrost *Bifrost) tryStreamRequest(ctx *schemas.BifrostContext, req *schem
 					if streamMsg.BifrostTranscriptionStreamResponse != nil {
 						bifrostResponse.TranscriptionStreamResponse = streamMsg.BifrostTranscriptionStreamResponse
 					}
+					if streamMsg.BifrostImageGenerationStreamResponse != nil {
+						bifrostResponse.ImageGenerationStreamResponse = streamMsg.BifrostImageGenerationStreamResponse
+					}
 
 					// Run post hooks on the stream message
 					processedResponse, processedError := pipelinePostHookRunner(ctx, bifrostResponse, streamMsg.BifrostError)
@@ -2772,6 +2869,7 @@ func (bifrost *Bifrost) tryStreamRequest(ctx *schemas.BifrostContext, req *schem
 						streamResponse.BifrostResponsesStreamResponse = processedResponse.ResponsesStreamResponse
 						streamResponse.BifrostSpeechStreamResponse = processedResponse.SpeechStreamResponse
 						streamResponse.BifrostTranscriptionStreamResponse = processedResponse.TranscriptionStreamResponse
+						streamResponse.BifrostImageGenerationStreamResponse = processedResponse.ImageGenerationStreamResponse
 					}
 					if processedError != nil {
 						streamResponse.BifrostError = processedError
@@ -3246,6 +3344,12 @@ func (bifrost *Bifrost) handleProviderRequest(provider schemas.Provider, req *Ch
 			return nil, bifrostError
 		}
 		response.TranscriptionResponse = transcriptionResponse
+	case schemas.ImageGenerationRequest:
+		imageResponse, bifrostError := provider.ImageGeneration(req.Context, key, req.BifrostRequest.ImageGenerationRequest)
+		if bifrostError != nil {
+			return nil, bifrostError
+		}
+		response.ImageGenerationResponse = imageResponse
 	case schemas.FileUploadRequest:
 		fileUploadResponse, bifrostError := provider.FileUpload(req.Context, key, req.BifrostRequest.FileUploadRequest)
 		if bifrostError != nil {
@@ -3336,6 +3440,8 @@ func (bifrost *Bifrost) handleProviderStreamRequest(provider schemas.Provider, r
 		return provider.SpeechStream(req.Context, postHookRunner, key, req.BifrostRequest.SpeechRequest)
 	case schemas.TranscriptionStreamRequest:
 		return provider.TranscriptionStream(req.Context, postHookRunner, key, req.BifrostRequest.TranscriptionRequest)
+	case schemas.ImageGenerationStreamRequest:
+		return provider.ImageGenerationStream(req.Context, postHookRunner, key, req.BifrostRequest.ImageGenerationRequest)
 	default:
 		_, model, _ := req.BifrostRequest.GetRequestFields()
 		return nil, &schemas.BifrostError{
@@ -3651,6 +3757,7 @@ func resetBifrostRequest(req *schemas.BifrostRequest) {
 	req.EmbeddingRequest = nil
 	req.SpeechRequest = nil
 	req.TranscriptionRequest = nil
+	req.ImageGenerationRequest = nil
 	req.FileUploadRequest = nil
 	req.FileListRequest = nil
 	req.FileRetrieveRequest = nil

core/changelog.md

@@ -1,3 +1 @@
-- chore: added case-insensitive helper methods for header and query parameter lookups in HTTPRequest
-- feat: added support for path parameter lookups in HTTPRequest
-- fix: missing request type in error response for anthropic SDK integration
+feat: added image generation request and response support

core/internal/testutil/account.go

@@ -44,6 +44,8 @@ type TestScenarios struct {
 	Reasoning             bool // Reasoning/thinking functionality via Responses API
 	PromptCaching         bool // Prompt caching functionality
 	ListModels            bool // List available models functionality
+	ImageGeneration       bool // Image generation functionality
+	ImageGenerationStream bool // Streaming image generation functionality
 	BatchCreate           bool // Batch API create functionality
 	BatchList             bool // Batch API list functionality
 	BatchRetrieve         bool // Batch API retrieve functionality
@@ -79,6 +81,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
 	ExternalTTSProvider      schemas.ModelProvider  // External TTS provider to use for testing
 	ExternalTTSModel         string                 // External TTS model to use for testing
 	BatchExtraParams         map[string]interface{} // Extra params for batch operations (e.g., role_arn, output_s3_uri for Bedrock)
@@ -223,6 +227,7 @@ func (account *ComprehensiveTestAccount) GetKeysForProvider(ctx context.Context,
 						"gpt-4o-backup":          "gpt-4o-3",
 						"claude-opus-4-5":        "claude-opus-4-5",
 						"o1":                     "o1",
+						"gpt-image-1":            "gpt-image-1",
 						"text-embedding-ada-002": "text-embedding-ada-002",
 					},
 					ClientID:     bifrost.Ptr(os.Getenv("AZURE_CLIENT_ID")),
@@ -250,7 +255,7 @@ func (account *ComprehensiveTestAccount) GetKeysForProvider(ctx context.Context,
 		return []schemas.Key{
 			{
 				Value:  os.Getenv("VERTEX_API_KEY"),
-				Models: []string{"text-multilingual-embedding-002", "google/gemini-2.0-flash-001"},
+				Models: []string{"text-multilingual-embedding-002", "google/gemini-2.0-flash-001", "gemini-2.5-flash-image", "imagen-4.0-generate-001"},
 				Weight: 1.0,
 				VertexKeyConfig: &schemas.VertexKeyConfig{
 					ProjectID:       os.Getenv("VERTEX_PROJECT_ID"),
@@ -669,6 +674,7 @@ var AllProviderConfigs = []ComprehensiveTestConfig{
 		PromptCachingModel:   "gpt-4.1",
 		TranscriptionModel:   "whisper-1",
 		SpeechSynthesisModel: "tts-1",
+		ImageGenerationModel: "gpt-image-1",
 		ChatAudioModel:       "gpt-4o-mini-audio-preview",
 		Scenarios: TestScenarios{
 			TextCompletion:        false, // Not supported
@@ -688,6 +694,8 @@ var AllProviderConfigs = []ComprehensiveTestConfig{
 			SpeechSynthesisStream: true, // OpenAI supports streaming TTS
 			Transcription:         true, // OpenAI supports STT with Whisper
 			TranscriptionStream:   true, // OpenAI supports streaming STT
+			ImageGeneration:       true, // OpenAI supports image generation with DALL-E
+			ImageGenerationStream: true, // OpenAI supports streaming image generation
 			Embedding:             true,
 			Reasoning:             true, // OpenAI supports reasoning via o1 models
 			ListModels:            true,
@@ -730,6 +738,8 @@ var AllProviderConfigs = []ComprehensiveTestConfig{
 			Transcription:         false, // Not supported
 			TranscriptionStream:   false, // Not supported
 			Embedding:             false,
+			ImageGeneration:       false,
+			ImageGenerationStream: false,
 			ListModels:            true,
 			BatchCreate:           true, // Anthropic supports batch API
 			BatchList:             true, // Anthropic supports batch API
@@ -764,6 +774,8 @@ var AllProviderConfigs = []ComprehensiveTestConfig{
 			Transcription:         false, // Not supported
 			TranscriptionStream:   false, // Not supported
 			Embedding:             true,
+			ImageGeneration:       false,
+			ImageGenerationStream: false,
 			ListModels:            true,
 			BatchCreate:           true, // Bedrock supports batch via Model Invocation Jobs (requires S3 config)
 			BatchList:             true, // Bedrock supports listing batch jobs
@@ -797,6 +809,8 @@ var AllProviderConfigs = []ComprehensiveTestConfig{
 			ImageBase64:           false, // Check if supported
 			MultipleImages:        false, // Check if supported
 			CompleteEnd2End:       true,
+			ImageGeneration:       false,
+			ImageGenerationStream: false,
 			SpeechSynthesis:       false, // Not supported
 			SpeechSynthesisStream: false, // Not supported
 			Transcription:         false, // Not supported
@@ -815,6 +829,7 @@ var AllProviderConfigs = []ComprehensiveTestConfig{
 		ChatAudioModel:       "gpt-4o-mini-audio-preview",
 		TranscriptionModel:   "whisper-1",
 		SpeechSynthesisModel: "gpt-4o-mini-tts",
+		ImageGenerationModel: "gpt-image-1",
 		Scenarios: TestScenarios{
 			TextCompletion:        false, // Not supported
 			SimpleChat:            true,
@@ -833,6 +848,8 @@ var AllProviderConfigs = []ComprehensiveTestConfig{
 			Transcription:         true,  // Supported via whisper-1
 			TranscriptionStream:   false, // Not properly supported yet by Azure
 			Embedding:             true,
+			ImageGeneration:       true,
+			ImageGenerationStream: true,
 			ListModels:            true,
 			BatchCreate:           true, // Azure supports batch API
 			BatchList:             true, // Azure supports batch API
@@ -851,9 +868,10 @@ var AllProviderConfigs = []ComprehensiveTestConfig{
 		},
 	},
 	{
-		Provider:  schemas.Vertex,
-		ChatModel: "gemini-pro",
-		TextModel: "", // Vertex focuses on chat
+		Provider:             schemas.Vertex,
+		ChatModel:            "gemini-pro",
+		TextModel:            "", // Vertex focuses on chat
+		ImageGenerationModel: "imagen-4.0-generate-001",
 		Scenarios: TestScenarios{
 			TextCompletion:        false, // Not typical
 			SimpleChat:            true,
@@ -867,6 +885,8 @@ var AllProviderConfigs = []ComprehensiveTestConfig{
 			ImageBase64:           true,
 			MultipleImages:        true,
 			CompleteEnd2End:       true,
+			ImageGeneration:       true,
+			ImageGenerationStream: false,
 			SpeechSynthesis:       false, // Not supported
 			SpeechSynthesisStream: false, // Not supported
 			Transcription:         false, // Not supported
@@ -900,6 +920,8 @@ var AllProviderConfigs = []ComprehensiveTestConfig{
 			Transcription:         true,  // Supported via voxtral-mini-latest
 			TranscriptionStream:   true,  // Supported via voxtral-mini-latest
 			Embedding:             true,
+			ImageGeneration:       false,
+			ImageGenerationStream: false,
 			ListModels:            true,
 		},
 		Fallbacks: []schemas.Fallback{
@@ -928,6 +950,8 @@ var AllProviderConfigs = []ComprehensiveTestConfig{
 			Transcription:         false, // Not supported
 			TranscriptionStream:   false, // Not supported
 			Embedding:             false,
+			ImageGeneration:       false,
+			ImageGenerationStream: false,
 			ListModels:            true,
 		},
 		Fallbacks: []schemas.Fallback{
@@ -956,6 +980,8 @@ var AllProviderConfigs = []ComprehensiveTestConfig{
 			Transcription:         false, // Not supported
 			TranscriptionStream:   false, // Not supported
 			Embedding:             false,
+			ImageGeneration:       false,
+			ImageGenerationStream: false,
 			ListModels:            true,
 		},
 		Fallbacks: []schemas.Fallback{
@@ -997,6 +1023,7 @@ var AllProviderConfigs = []ComprehensiveTestConfig{
 		TranscriptionModel:   "gemini-2.5-flash",
 		SpeechSynthesisModel: "gemini-2.5-flash-preview-tts",
 		EmbeddingModel:       "text-embedding-004",
+		ImageGenerationModel: "imagen-4.0-generate-001",
 		Scenarios: TestScenarios{
 			TextCompletion:        false, // Not supported
 			SimpleChat:            true,
@@ -1015,6 +1042,8 @@ var AllProviderConfigs = []ComprehensiveTestConfig{
 			Transcription:         true,
 			TranscriptionStream:   true,
 			Embedding:             true,
+			ImageGeneration:       true,
+			ImageGenerationStream: false,
 			ListModels:            true,
 			BatchCreate:           true,
 			BatchList:             true,
@@ -1048,6 +1077,8 @@ var AllProviderConfigs = []ComprehensiveTestConfig{
 			ImageBase64:           true,
 			MultipleImages:        true,
 			CompleteEnd2End:       true,
+			ImageGeneration:       false,
+			ImageGenerationStream: false,
 			SpeechSynthesis:       false,
 			SpeechSynthesisStream: false,
 			Transcription:         false,
@@ -1066,6 +1097,7 @@ var AllProviderConfigs = []ComprehensiveTestConfig{
 		EmbeddingModel:       "sambanova/intfloat/e5-mistral-7b-instruct",
 		TranscriptionModel:   "fal-ai/openai/whisper-large-v3",
 		SpeechSynthesisModel: "fal-ai/hexgrad/Kokoro-82M",
+		ImageGenerationModel: "fal-ai/fal-ai/flux-2",
 		Scenarios: TestScenarios{
 			TextCompletion:        false,
 			TextCompletionStream:  false,
@@ -1082,6 +1114,8 @@ var AllProviderConfigs = []ComprehensiveTestConfig{
 			MultipleImages:        true,
 			CompleteEnd2End:       true,
 			Embedding:             true,
+			ImageGeneration:       true,
+			ImageGenerationStream: true,
 			Transcription:         true,
 			TranscriptionStream:   false,
 			SpeechSynthesis:       true,
@@ -1094,9 +1128,10 @@ var AllProviderConfigs = []ComprehensiveTestConfig{
 		},
 	},
 	{
-		Provider:  schemas.XAI,
-		ChatModel: "grok-4-0709",
-		TextModel: "", // XAI focuses on chat
+		Provider:             schemas.XAI,
+		ChatModel:            "grok-4-0709",
+		TextModel:            "", // XAI focuses on chat
+		ImageGenerationModel: "grok-2-image",
 		Scenarios: TestScenarios{
 			TextCompletion:        false, // Not typical
 			SimpleChat:            true,
@@ -1116,6 +1151,8 @@ var AllProviderConfigs = []ComprehensiveTestConfig{
 			TranscriptionStream:   false, // Not supported
 			Embedding:             false, // Not supported
 			ListModels:            true,
+			ImageGeneration:       true,
+			ImageGenerationStream: false,
 		},
 	},
 }