feat: support for image generation #980
Conversation
📝 WalkthroughSummary by CodeRabbit
✏️ Tip: You can customize this high-level summary in your review settings. WalkthroughAdds end-to-end image generation: new image request/stream schemas and Bifrost APIs, provider image methods (OpenAI/Azure + stubs), streaming accumulator and chunk reassembly, semantic-cache image keys/metadata, HTTP image endpoints (streaming and non‑streaming), UI image streaming hook/component, and many image-focused tests/utilities. Changes
Sequence Diagram(s)sequenceDiagram
actor Client
participant HTTP as HTTP Handler
participant Bifrost as Bifrost Core
participant Cache as Semantic Cache
participant Provider as Provider (OpenAI/Azure)
participant Stream as Streaming Framework
Client->>HTTP: POST /v1/images/generations (prompt, provider, model, params)
HTTP->>Bifrost: ImageGenerationRequest(ctx, req)
Bifrost->>Cache: lookup(getImageCacheKey)
Cache-->>Bifrost: hit/miss
alt Cache hit
Bifrost-->>HTTP: return cached BifrostImageGenerationResponse
HTTP-->>Client: 200 OK (images)
else Cache miss
Bifrost->>Provider: ImageGeneration(ctx, key, req)
Provider->>Provider: vendor API call
Provider-->>Bifrost: BifrostImageGenerationResponse
Bifrost->>Cache: store response
Bifrost-->>HTTP: return response
HTTP-->>Client: 200 OK (images)
end
sequenceDiagram
actor Client
participant HTTP as HTTP Handler
participant Bifrost as Bifrost Core
participant Provider as Provider (OpenAI/Azure)
participant Stream as Streaming Framework
participant Accum as Accumulator
Client->>HTTP: POST /v1/images/generations?stream=true (prompt,...)
HTTP->>Bifrost: ImageGenerationStreamRequest(ctx, req)
Bifrost->>Provider: ImageGenerationStream(ctx, postHook, key, req)
Provider->>Stream: SSE stream of image chunks
loop per chunk
Stream->>Accum: addImageStreamChunk(requestID, chunk)
Accum-->>HTTP: partial BifrostImageGenerationStreamResponse
HTTP-->>Client: SSE data: {partial chunk}
end
Provider->>Stream: completion event (+usage)
Stream->>Accum: addImageStreamChunk(final=true)
Accum->>Accum: assemble images (buildCompleteImageFromImageStreamChunks)
Accum-->>HTTP: final BifrostImageGenerationStreamResponse (complete + usage)
HTTP-->>Client: SSE data: {complete response}
Estimated code review effort🎯 4 (Complex) | ⏱️ ~45 minutes Areas needing focused review:
Poem
Pre-merge checks and finishing touches❌ Failed checks (1 warning)
✅ Passed checks (4 passed)
✨ Finishing touches
🧪 Generate unit tests (beta)
Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out. Comment |
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
There was a problem hiding this comment.
Actionable comments posted: 15
Caution
Some comments are outside the diff and can’t be posted inline due to platform limitations.
⚠️ Outside diff range comments (2)
core/schemas/responses.go (1)
525-573: Fix UnmarshalJSON ordering to avoid misclassifying computer tool outputs as image-generation outputsThe new
ResponsesImageGenerationCallOutputintegration introduces a bug in theUnmarshalJSONmethod:
- For JSON objects like a computer tool screenshot output
{"type": "computer_screenshot", "file_id": "...", "image_url": "..."},
sonic.Unmarshalsucceeds (missing/unknown fields simply leave struct fields at their zero values) because the unmarshaling does not fail.- As a result, the image-generation branch always "wins" for any object-shaped payload, and the code never reaches the
ResponsesComputerToolCallOutputDatabranch. Existing computer tool outputs will be parsed as image-generation outputs with an emptyResult.Fix this by only treating the payload as an image-generation output when it actually contains a
resultfield, and otherwise fall through to the computer-tool path:func (output *ResponsesToolMessageOutputStruct) UnmarshalJSON(data []byte) error { var str string if err := sonic.Unmarshal(data, &str); err == nil { output.ResponsesToolCallOutputStr = &str return nil } var array []ResponsesMessageContentBlock if err := sonic.Unmarshal(data, &array); err == nil { output.ResponsesFunctionToolCallOutputBlocks = array return nil } - var imageGenerationCallOutput ResponsesImageGenerationCallOutput - if err := sonic.Unmarshal(data, &imageGenerationCallOutput); err == nil { - output.ResponsesImageGenerationCallOutput = &imageGenerationCallOutput - return nil - } - var computerToolCallOutput ResponsesComputerToolCallOutputData - if err := sonic.Unmarshal(data, &computerToolCallOutput); err == nil { - output.ResponsesComputerToolCallOutput = &computerToolCallOutput - return nil - } + // Peek at the object to distinguish image-generation vs computer tool outputs. + var raw map[string]any + if err := sonic.Unmarshal(data, &raw); err == nil { + if _, hasResult := raw["result"]; hasResult { + var imageGenerationCallOutput ResponsesImageGenerationCallOutput + if err := sonic.Unmarshal(data, &imageGenerationCallOutput); err == nil { + output.ResponsesImageGenerationCallOutput = &imageGenerationCallOutput + return nil + } + } + } + + var computerToolCallOutput ResponsesComputerToolCallOutputData + if err := sonic.Unmarshal(data, &computerToolCallOutput); err == nil { + output.ResponsesComputerToolCallOutput = &computerToolCallOutput + return nil + } return fmt.Errorf("responses tool message output struct is neither a string nor an array of responses message content blocks nor a computer tool call output data nor an image generation call output") }This preserves the intended priority for image-generation outputs while ensuring computer tool outputs continue to be parsed correctly.
core/bifrost.go (1)
1806-1863: Missing ImageGenerationRequest handling in prepareFallbackRequest.The
prepareFallbackRequestmethod handles fallbacks for all request types exceptImageGenerationRequest. This will cause fallbacks to silently fail for image generation requests.if req.TranscriptionRequest != nil { tmp := *req.TranscriptionRequest tmp.Provider = fallback.Provider 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 }
🧹 Nitpick comments (27)
core/schemas/provider.go (1)
234-239: ImageGeneration methods correctly extend the Provider interface; consider aligning gating config laterThe new
ImageGeneration/ImageGenerationStreamsignatures are consistent with existing operations (e.g., Chat/Speech) and look good. As a follow-up, you may want to extendAllowedRequests/CustomProviderConfig.IsOperationAllowedand any request-type-based routing so image-generation and its stream can be selectively enabled/disabled like the other operations.core/providers/cohere/cohere.go (1)
802-810: Use a stream-specific request type for ImageGenerationStream unsupported errorsThe unsupported stubs are fine, but
ImageGenerationStreamcurrently reportsschemas.ImageGenerationRequestas the request type inNewUnsupportedOperationError. For clearer observability and any future per-operation gating, consider using a dedicatedImageGenerationStreamRequestconstant here, mirroring howTextCompletionStream/ChatCompletionStreamuse their own request types.plugins/semanticcache/main.go (1)
14-14: getImageCacheKey implementation is fine; consider expanding inputs if you need stricter separation
getImageCacheKey’s use ofxxhashover prompt + main image params is sound, and returning""whenreq/Inputis nil gives a clear “no cache key” signal to call-sites.If you eventually need stricter cache separation (e.g., by provider/model or provider-specific
ExtraParams), you may want to fold those into the hash input or compose this helper with the existingCacheByModel/CacheByProviderbehavior at call-sites.Also applies to: 748-781
core/internal/testutil/image_generation_stream.go (1)
15-136: Streaming test harness looks solid; double-check event type strings and consider a buffereddoneThe concurrency pattern (wg + mutex +
donesignaling) is reasonable, and the test does a good job asserting partial vs completion behavior. Two minor points:
The checks against
imgResp.Type("image_generation.completed"andstrings.Contains(..., "partial_image")) are tightly coupled to whatever valuesBifrostImageGenerationStreamResponse.Typeactually uses. It’s worth confirming those are in sync (or, if you later introduce typed constants for image stream events, using them here to avoid drift).If you want to avoid any chance of the collector goroutine blocking when the test times out or aborts early, you could make
donea buffered channel of size 1.core/providers/openrouter/openrouter.go (1)
271-279: Unsupported image-generation stubs align with existing pattern; consider using a stream request typeThe new
ImageGeneration/ImageGenerationStreammethods correctly advertise lack of support viaNewUnsupportedOperationErrorand match the style of other unsupported operations. ForImageGenerationStream, you may want to use a dedicated streamRequestTypeto keep metrics/logging distinct from non-stream image generation, similar to other streamed operations.core/image_generation_tools_test.go (1)
1-512: Comprehensive image-generation tool tests; consider guarding against nil Error in assertionsThis suite does a nice job covering happy paths, argument validation (missing/empty prompt, invalid JSON, missing required fields), and the new
ResponsesImageGeneration*output serialization behavior, including priority semantics.One small robustness tweak: in the tests that expect failures you directly access
bifrostErr.Error.Messageafterassert.NotNil(t, bifrostErr). If future changes ever construct aBifrostErrorwith a nilErrorfield, those tests will panic instead of failing cleanly. Adding arequire.NotNil(t, bifrostErr.Error)before inspecting.Messagein those cases would make the tests more defensive without changing intent.core/providers/vertex/vertex.go (1)
894-902: Vertex image-generation stubs are consistent; stream method could use a dedicated request typeThe new
ImageGeneration/ImageGenerationStreammethods correctly signal that Vertex doesn’t support image generation yet and match existing unsupported-operation patterns. ForImageGenerationStream, consider swappingschemas.ImageGenerationRequestfor a stream-specificRequestTypeso observability and any future per-operation controls distinguish stream vs non-stream image generation.core/schemas/responses.go (1)
1367-1407: Image-generation streaming response types fit cleanly into the existing enumThe added
ResponsesStreamResponseTypeImageGenerationCall*constants follow the existing naming scheme for streamed events and give you a clear, typed surface for image-generation-related responses. They look good as-is; when convenient, it would be worth using these constants (rather than raw strings) wherever you check image-generation stream event types to avoid drift.core/internal/testutil/image_generation_fallback.go (1)
50-60: Consider stricter error handling for fallback validation.The current error handling logs a warning but doesn't fail the test when fallbacks are configured but the request still fails. This could mask real fallback mechanism bugs. Consider adding an assertion that when fallbacks are configured and valid, the request should eventually succeed.
if bifrostErr != nil { // Check if error indicates fallback was attempted errorMsg := GetErrorMessage(bifrostErr) if strings.Contains(strings.ToLower(errorMsg), "fallback") { t.Logf("✅ Fallback mechanism triggered (expected behavior)") } else { - // If we have fallbacks configured, the request should succeed via fallback - // If it still fails, log it but don't fail the test (provider-specific) - t.Logf("⚠️ Request failed even with fallbacks: %v", errorMsg) + // If fallbacks are configured and valid, the request should succeed + // This might indicate a fallback mechanism issue + t.Errorf("⚠️ Request failed even with fallbacks configured: %v", errorMsg) } return }ui/components/chat/ImageMessage.tsx (1)
23-24: Consider using the image'sindexproperty as the key.Using array index (
idx) as the React key is an anti-pattern when list items can be reordered. Theimagesarray items already have anindexproperty that could provide a more stable key.- {images.map((img, idx) => ( - <Card key={idx} className="overflow-hidden"> + {images.map((img) => ( + <Card key={img.index} className="overflow-hidden">transports/bifrost-http/handlers/inference.go (2)
1191-1197: Movedefer cancel()after the streaming check for consistency.Other handlers (e.g.,
textCompletion,chatCompletion,speech) placedefer cancel()after the streaming check, not before. While calling cancel twice is safe, this placement is inconsistent with the established pattern and could cause confusion.// Convert context bifrostCtx, cancel := lib.ConvertToBifrostContext(ctx, h.handlerStore.ShouldAllowDirectKeys()) if bifrostCtx == nil { SendError(ctx, fasthttp.StatusInternalServerError, "Failed to convert context") return } - defer cancel() // Handle streaming image generation if req.BifrostParams.Stream != nil && *req.BifrostParams.Stream { h.handleStreamingImageGeneration(ctx, bifrostReq, bifrostCtx, cancel) return } + + defer cancel() // Ensure cleanup on function exit // Execute request resp, bifrostErr := h.client.ImageGenerationRequest(*bifrostCtx, bifrostReq)
1165-1165: Fix typo: "witout" → "without".- // Continue witout extra params + // Continue without extra paramscore/internal/testutil/image_generation_cache.go (2)
28-38: Consider including fallbacks for comprehensive cache testing.The request omits
testConfig.ImageGenerationFallbacks, which is available in the config (peraccount.golines 65-66). While this may be intentional to test cache behavior without fallback complexity, including fallbacks would ensure cache key generation accounts for them properly.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"), }, + Fallbacks: testConfig.ImageGenerationFallbacks, }
94-104: Single-image comparison may miss edge cases.The validation only compares the first image's
B64JSON. Ifn > 1is ever used, consider iterating through all images to ensure complete cache integrity.framework/streaming/types.go (1)
258-269: Consider usingImageGenerationRequestfor consistency.Other stream type handlers in
ToBifrostResponseuse the non-stream request type for the final accumulated response (e.g.,StreamTypeAudiousesschemas.SpeechRequest,StreamTypeTranscriptionusesschemas.TranscriptionRequest). For consistency, this should likely useschemas.ImageGenerationRequestinstead ofschemas.ImageGenerationStreamRequest.case StreamTypeImage: imageResp := p.Data.ImageGenerationOutput if imageResp == nil { imageResp = &schemas.BifrostImageGenerationResponse{} } resp.ImageGenerationResponse = imageResp resp.ImageGenerationResponse.ExtraFields = schemas.BifrostResponseExtraFields{ - RequestType: schemas.ImageGenerationStreamRequest, + RequestType: schemas.ImageGenerationRequest, Provider: p.Provider, ModelRequested: p.Model, Latency: p.Data.Latency, }core/providers/openai/openai_image_test.go (2)
13-14: Consider renaming test function for clarity.The function name
TestImageGenerationStreamingRequestConversionis misleading since it tests general request conversion (viaToOpenAIImageGenerationRequest), not streaming-specific behavior.-// TestImageGenerationStreamingRequestConversion -func TestImageGenerationStreamingRequestConversion(t *testing.T) { +// TestImageGenerationRequestConversion tests Bifrost to OpenAI request conversion +func TestImageGenerationRequestConversion(t *testing.T) {
245-252: Handle potential unmarshal errors in test verification.The
json.Unmarshalerror is silently ignored. If unmarshaling fails, the subsequenttt.verifycall would operate on an empty map, potentially hiding issues.t.Run(tt.name, func(t *testing.T) { t.Parallel() jsonBytes, _ := sonic.Marshal(tt.chunk) var data map[string]interface{} - json.Unmarshal(jsonBytes, &data) + if err := json.Unmarshal(jsonBytes, &data); err != nil { + t.Fatalf("Failed to unmarshal: %v", err) + } tt.verify(t, data) })plugins/semanticcache/utils.go (2)
88-88: Fix typo in comment.The comment has a typo: "styele" should be "style".
- // Special handling for image generation (hash = prompt + size + quality + styele + n) + // Special handling for image generation (hash = prompt + size + quality + style + n)
516-517: Consider adding nil check for defensive coding.The code accesses
req.ImageGenerationRequest.Inputwithout verifying thatreq.ImageGenerationRequestis not nil. While the request type check suggests it should be populated, a defensive nil check would prevent potential panics if the request is malformed.case schemas.ImageGenerationRequest, schemas.ImageGenerationStreamRequest: + if req.ImageGenerationRequest == nil { + return nil + } return req.ImageGenerationRequest.Inputframework/streaming/images.go (2)
137-163: Consolidate repeated last chunk access.The code accesses
acc.ImageStreamChunks[len(acc.ImageStreamChunks)-1]three times (lines 138, 150, 158). Consider extracting this to a single variable for clarity and efficiency.// Update token usage from final chunk if available if len(acc.ImageStreamChunks) > 0 { lastChunk := acc.ImageStreamChunks[len(acc.ImageStreamChunks)-1] if lastChunk.Delta != nil && lastChunk.Delta.Usage != nil { data.TokenUsage = &schemas.BifrostLLMUsage{ PromptTokens: lastChunk.Delta.Usage.PromptTokens, CompletionTokens: 0, // Image generation doesn't have completion tokens TotalTokens: lastChunk.Delta.Usage.TotalTokens, } } - } - - // Update cost from final chunk if available - if len(acc.ImageStreamChunks) > 0 { - lastChunk := acc.ImageStreamChunks[len(acc.ImageStreamChunks)-1] if lastChunk.Cost != nil { data.Cost = lastChunk.Cost } - } - - // Update semantic cache debug from final chunk if available - if len(acc.ImageStreamChunks) > 0 { - lastChunk := acc.ImageStreamChunks[len(acc.ImageStreamChunks)-1] if lastChunk.SemanticCacheDebug != nil { data.CacheDebug = lastChunk.SemanticCacheDebug } data.FinishReason = lastChunk.FinishReason }
202-206: Remove or implement empty conditional block.This block contains only a comment and no actual code. Either remove it entirely or implement the intended usage extraction logic.
- // Extract usage if available - if result.ImageGenerationStreamResponse.Usage != nil { - // Note: ImageUsage doesn't directly map to BifrostLLMUsage, but we can store it - // The actual usage will be extracted in processAccumulatedImageStreamingChunks - }ui/hooks/useImageStream.ts (3)
107-179: Stale closure risk withoptionsinprocessChunkdependency.The
optionsobject in the dependency array can cause issues if the consumer passes a new object reference on each render (common with inline callbacks). Consider usinguseRefto hold callbacks or usinguseCallbackfor the options callbacks at the call site.+ const optionsRef = useRef(options); + optionsRef.current = options; + // Process incoming chunk const processChunk = useCallback( (chunk: ImageStreamChunk) => { // ... existing logic ... if (type === "error" || error) { const errorMsg = error?.message || "Unknown streaming error"; setState((prev) => ({ ...prev, error: errorMsg, isStreaming: false })); - options.onError?.(errorMsg); + optionsRef.current.onError?.(errorMsg); return; } // ... rest of function, similar changes for onComplete ... }, - [buildImageFromChunks, options] + [buildImageFromChunks] );
138-141: Progress estimation is coarse; consider documenting the assumption.The hardcoded assumption of ~10 chunks per image may not reflect actual streaming behavior. Consider making this configurable or adding a comment explaining the estimate.
245-250: Consider logging parse errors in development only.The
console.errorfor parse failures could be noisy in production. Consider gating this log or using a more structured error handling approach.try { const chunk: ImageStreamChunk = JSON.parse(data); processChunk(chunk); } catch (parseError) { - console.error("Failed to parse SSE chunk:", parseError); + if (process.env.NODE_ENV === "development") { + console.error("Failed to parse SSE chunk:", parseError); + } }core/bifrost.go (1)
1192-1238: Consider extracting shared parameter mapping logic.The parameter extraction code for
size,quality,style,response_format, andnis nearly identical betweenExecuteImageGenerationToolandExecuteResponsesImageGenerationTool. Consider extracting to a helper function.// Helper to extract image generation parameters from tool call arguments func extractImageGenerationParams(args map[string]interface{}) *schemas.ImageGenerationParameters { var params *schemas.ImageGenerationParameters if size, ok := args["size"].(string); ok && size != "" { if params == nil { params = &schemas.ImageGenerationParameters{} } params.Size = &size } if quality, ok := args["quality"].(string); ok && quality != "" { if params == nil { params = &schemas.ImageGenerationParameters{} } params.Quality = &quality } // ... similar for style, response_format, n return params }Also applies to: 1342-1388
core/providers/openai/images.go (2)
365-366: ChunkIndex initialized to -1 is unusual but intentional.The
chunkIndexstarts at -1 and is incremented to 0 before first use at line 456. This works correctly but is non-obvious.Consider initializing to 0 and restructuring the increment logic for clarity:
- chunkIndex := -1 + chunkIndex := 0 // ... later in the loop ... - chunkIndex++ - chunk := b64Data[offset:end] - bifrostChunk.PartialB64 = chunk - bifrostChunk.ChunkIndex = chunkIndex + chunk := b64Data[offset:end] + bifrostChunk.PartialB64 = chunk + bifrostChunk.ChunkIndex = chunkIndex + chunkIndex++
422-427: TODO comment indicates incomplete usage tracking.The comment suggests usage tracking is not fully implemented for streaming responses.
Would you like help implementing proper usage tracking for streaming image generation?
📜 Review details
Configuration used: CodeRabbit UI
Review profile: CHILL
Plan: Pro
📒 Files selected for processing (48)
core/bifrost.go(6 hunks)core/image_generation_tools_test.go(1 hunks)core/images_test.go(1 hunks)core/internal/testutil/account.go(2 hunks)core/internal/testutil/image_generation.go(1 hunks)core/internal/testutil/image_generation_cache.go(1 hunks)core/internal/testutil/image_generation_errors.go(1 hunks)core/internal/testutil/image_generation_fallback.go(1 hunks)core/internal/testutil/image_generation_load.go(1 hunks)core/internal/testutil/image_generation_stream.go(1 hunks)core/internal/testutil/tests.go(1 hunks)core/providers/anthropic/anthropic.go(4 hunks)core/providers/azure/azure.go(1 hunks)core/providers/azure/images.go(1 hunks)core/providers/bedrock/bedrock.go(6 hunks)core/providers/cerebras/cerebras.go(1 hunks)core/providers/cohere/cohere.go(3 hunks)core/providers/elevenlabs/elevenlabs.go(3 hunks)core/providers/gemini/gemini.go(4 hunks)core/providers/groq/groq.go(1 hunks)core/providers/mistral/mistral.go(1 hunks)core/providers/ollama/ollama.go(1 hunks)core/providers/openai/images.go(1 hunks)core/providers/openai/openai.go(8 hunks)core/providers/openai/openai_image_test.go(1 hunks)core/providers/openrouter/openrouter.go(1 hunks)core/providers/parasail/parasail.go(1 hunks)core/providers/perplexity/perplexity.go(1 hunks)core/providers/sgl/sgl.go(1 hunks)core/providers/utils/utils.go(2 hunks)core/providers/vertex/vertex.go(1 hunks)core/schemas/bifrost.go(10 hunks)core/schemas/chatcompletions.go(1 hunks)core/schemas/images.go(1 hunks)core/schemas/mux.go(1 hunks)core/schemas/provider.go(1 hunks)core/schemas/responses.go(3 hunks)core/utils.go(1 hunks)framework/streaming/accumulator.go(10 hunks)framework/streaming/images.go(1 hunks)framework/streaming/types.go(5 hunks)plugins/semanticcache/main.go(4 hunks)plugins/semanticcache/search.go(2 hunks)plugins/semanticcache/stream.go(1 hunks)plugins/semanticcache/utils.go(7 hunks)transports/bifrost-http/handlers/inference.go(6 hunks)ui/components/chat/ImageMessage.tsx(1 hunks)ui/hooks/useImageStream.ts(1 hunks)
🧰 Additional context used
📓 Path-based instructions (1)
**
⚙️ CodeRabbit configuration file
always check the stack if there is one for the current PR. do not give localized reviews for the PR, always see all changes in the light of the whole stack of PRs (if there is a stack, if there is no stack you can continue to make localized suggestions/reviews)
Files:
plugins/semanticcache/stream.gocore/schemas/provider.gocore/internal/testutil/account.gocore/providers/groq/groq.gocore/internal/testutil/image_generation_errors.gocore/providers/ollama/ollama.gocore/providers/azure/azure.gocore/providers/cerebras/cerebras.gocore/schemas/images.goui/components/chat/ImageMessage.tsxcore/providers/mistral/mistral.gocore/schemas/chatcompletions.gocore/providers/sgl/sgl.gocore/schemas/responses.gocore/providers/cohere/cohere.gocore/providers/anthropic/anthropic.gocore/providers/openai/openai_image_test.gocore/providers/utils/utils.gocore/providers/perplexity/perplexity.gocore/internal/testutil/image_generation_cache.gocore/internal/testutil/image_generation_fallback.goui/hooks/useImageStream.tsframework/streaming/images.gocore/image_generation_tools_test.goframework/streaming/accumulator.gocore/providers/openai/openai.gocore/images_test.gocore/internal/testutil/image_generation.gotransports/bifrost-http/handlers/inference.gocore/internal/testutil/image_generation_load.goplugins/semanticcache/search.gocore/providers/vertex/vertex.gocore/providers/openrouter/openrouter.gocore/providers/bedrock/bedrock.gocore/providers/gemini/gemini.gocore/schemas/mux.goplugins/semanticcache/utils.gocore/internal/testutil/tests.gocore/schemas/bifrost.gocore/internal/testutil/image_generation_stream.gocore/utils.goplugins/semanticcache/main.gocore/bifrost.gocore/providers/azure/images.gocore/providers/elevenlabs/elevenlabs.gocore/providers/parasail/parasail.gocore/providers/openai/images.goframework/streaming/types.go
🧬 Code graph analysis (34)
core/schemas/provider.go (3)
core/schemas/account.go (1)
Key(8-17)core/schemas/images.go (2)
BifrostImageGenerationRequest(4-11)BifrostImageGenerationResponse(33-40)core/schemas/bifrost.go (2)
BifrostError(378-387)BifrostStream(340-348)
core/internal/testutil/account.go (1)
core/schemas/bifrost.go (1)
Fallback(132-135)
core/providers/groq/groq.go (4)
core/schemas/images.go (2)
BifrostImageGenerationRequest(4-11)BifrostImageGenerationResponse(33-40)core/providers/utils/utils.go (1)
NewUnsupportedOperationError(391-402)core/schemas/bifrost.go (2)
ImageGenerationRequest(100-100)BifrostStream(340-348)core/schemas/provider.go (1)
PostHookRunner(204-204)
core/internal/testutil/image_generation_errors.go (2)
core/internal/testutil/account.go (1)
ComprehensiveTestConfig(49-68)core/schemas/images.go (3)
BifrostImageGenerationRequest(4-11)ImageGenerationInput(18-20)ImageGenerationParameters(22-30)
core/schemas/images.go (2)
core/schemas/provider.go (1)
Provider(207-240)core/schemas/bifrost.go (4)
ModelProvider(32-32)Fallback(132-135)BifrostResponseExtraFields(306-315)BifrostError(378-387)
ui/components/chat/ImageMessage.tsx (2)
ui/components/ui/card.tsx (1)
Card(50-50)ui/components/ui/skeleton.tsx (1)
Skeleton(7-7)
core/providers/mistral/mistral.go (3)
core/schemas/images.go (2)
BifrostImageGenerationRequest(4-11)BifrostImageGenerationResponse(33-40)core/schemas/bifrost.go (2)
BifrostError(378-387)ImageGenerationRequest(100-100)core/providers/utils/utils.go (1)
NewUnsupportedOperationError(391-402)
core/schemas/chatcompletions.go (2)
core/providers/gemini/types.go (1)
Type(778-778)ui/lib/types/logs.ts (1)
Function(141-146)
core/providers/sgl/sgl.go (3)
core/schemas/images.go (2)
BifrostImageGenerationRequest(4-11)BifrostImageGenerationResponse(33-40)core/providers/utils/utils.go (1)
NewUnsupportedOperationError(391-402)core/schemas/bifrost.go (2)
ImageGenerationRequest(100-100)BifrostStream(340-348)
core/providers/cohere/cohere.go (4)
core/providers/utils/utils.go (3)
ProcessAndSendResponse(535-565)GetBifrostResponseForStreamResponse(787-819)NewUnsupportedOperationError(391-402)core/schemas/images.go (2)
BifrostImageGenerationRequest(4-11)BifrostImageGenerationResponse(33-40)core/schemas/bifrost.go (3)
BifrostError(378-387)ImageGenerationRequest(100-100)BifrostStream(340-348)core/schemas/provider.go (1)
PostHookRunner(204-204)
core/providers/anthropic/anthropic.go (3)
core/providers/utils/utils.go (3)
ProcessAndSendResponse(535-565)GetBifrostResponseForStreamResponse(787-819)NewUnsupportedOperationError(391-402)core/schemas/bifrost.go (2)
BifrostError(378-387)BifrostStream(340-348)core/schemas/provider.go (1)
PostHookRunner(204-204)
core/providers/openai/openai_image_test.go (2)
core/schemas/images.go (4)
BifrostImageGenerationRequest(4-11)ImageGenerationInput(18-20)ImageGenerationParameters(22-30)BifrostImageGenerationResponse(33-40)core/providers/openai/images.go (4)
OpenAIImageGenerationRequest(27-37)ToOpenAIImageGenerationRequest(102-114)OpenAIImageGenerationResponse(40-52)OpenAIImageGenerationUsage(54-63)
core/providers/utils/utils.go (1)
core/schemas/images.go (1)
BifrostImageGenerationStreamResponse(55-65)
core/internal/testutil/image_generation_cache.go (3)
core/internal/testutil/account.go (1)
ComprehensiveTestConfig(49-68)core/schemas/images.go (3)
BifrostImageGenerationRequest(4-11)ImageGenerationInput(18-20)ImageGenerationParameters(22-30)ui/lib/types/logs.ts (1)
CacheDebug(202-211)
core/internal/testutil/image_generation_fallback.go (3)
core/internal/testutil/account.go (1)
ComprehensiveTestConfig(49-68)core/schemas/images.go (3)
BifrostImageGenerationRequest(4-11)ImageGenerationInput(18-20)ImageGenerationParameters(22-30)core/schemas/bifrost.go (1)
ImageGenerationRequest(100-100)
ui/hooks/useImageStream.ts (2)
ui/lib/store/apis/baseApi.ts (1)
getTokenFromStorage(10-25)ui/lib/utils/port.ts (1)
getEndpointUrl(117-127)
framework/streaming/accumulator.go (2)
framework/streaming/types.go (4)
ImageStreamChunk(99-109)StreamAccumulator(112-124)ChatStreamChunk(75-84)ResponsesStreamChunk(87-96)core/schemas/bifrost.go (1)
ImageGenerationStreamRequest(101-101)
core/providers/openai/openai.go (1)
core/providers/utils/utils.go (2)
ProcessAndSendResponse(535-565)GetBifrostResponseForStreamResponse(787-819)
core/images_test.go (2)
core/schemas/images.go (7)
BifrostImageGenerationRequest(4-11)ImageGenerationInput(18-20)ImageGenerationParameters(22-30)BifrostImageGenerationResponse(33-40)ImageData(42-47)ImageUsage(49-52)BifrostImageGenerationStreamResponse(55-65)core/schemas/bifrost.go (1)
OpenAI(35-35)
core/internal/testutil/image_generation.go (3)
core/internal/testutil/account.go (1)
ComprehensiveTestConfig(49-68)core/schemas/images.go (3)
BifrostImageGenerationRequest(4-11)ImageGenerationInput(18-20)ImageGenerationParameters(22-30)core/schemas/bifrost.go (1)
ImageGenerationRequest(100-100)
transports/bifrost-http/handlers/inference.go (4)
core/schemas/images.go (3)
ImageGenerationInput(18-20)ImageGenerationParameters(22-30)BifrostImageGenerationRequest(4-11)core/schemas/utils.go (1)
ParseModelString(23-36)core/schemas/bifrost.go (4)
ModelProvider(32-32)ImageGenerationRequest(100-100)BifrostStream(340-348)BifrostError(378-387)transports/bifrost-http/lib/ctx.go (1)
ConvertToBifrostContext(72-283)
core/providers/vertex/vertex.go (4)
core/schemas/images.go (2)
BifrostImageGenerationRequest(4-11)BifrostImageGenerationResponse(33-40)core/schemas/bifrost.go (2)
BifrostError(378-387)ImageGenerationRequest(100-100)core/providers/utils/utils.go (1)
NewUnsupportedOperationError(391-402)core/schemas/provider.go (1)
PostHookRunner(204-204)
core/providers/openrouter/openrouter.go (4)
core/schemas/images.go (2)
BifrostImageGenerationRequest(4-11)BifrostImageGenerationResponse(33-40)core/providers/utils/utils.go (1)
NewUnsupportedOperationError(391-402)core/schemas/bifrost.go (2)
ImageGenerationRequest(100-100)BifrostStream(340-348)core/schemas/provider.go (1)
PostHookRunner(204-204)
core/providers/bedrock/bedrock.go (3)
core/providers/utils/utils.go (3)
ProcessAndSendResponse(535-565)GetBifrostResponseForStreamResponse(787-819)NewUnsupportedOperationError(391-402)core/schemas/bifrost.go (2)
BifrostError(378-387)BifrostStream(340-348)core/schemas/provider.go (1)
PostHookRunner(204-204)
core/providers/gemini/gemini.go (4)
core/providers/utils/utils.go (3)
ProcessAndSendResponse(535-565)GetBifrostResponseForStreamResponse(787-819)NewUnsupportedOperationError(391-402)core/schemas/images.go (2)
BifrostImageGenerationRequest(4-11)BifrostImageGenerationResponse(33-40)core/schemas/bifrost.go (3)
BifrostError(378-387)ImageGenerationRequest(100-100)BifrostStream(340-348)core/schemas/provider.go (1)
PostHookRunner(204-204)
core/schemas/mux.go (2)
core/schemas/chatcompletions.go (1)
ChatToolTypeImageGeneration(199-199)core/schemas/responses.go (1)
ResponsesToolImageGeneration(1318-1329)
plugins/semanticcache/utils.go (2)
core/schemas/bifrost.go (3)
RequestType(83-83)ImageGenerationRequest(100-100)ImageGenerationStreamRequest(101-101)core/schemas/images.go (2)
ImageGenerationInput(18-20)ImageGenerationParameters(22-30)
core/internal/testutil/tests.go (6)
core/internal/testutil/image_generation.go (1)
RunImageGenerationTest(13-86)core/internal/testutil/image_generation_stream.go (1)
RunImageGenerationStreamTest(16-137)core/internal/testutil/image_generation_fallback.go (1)
RunImageGenerationFallbackTest(14-80)core/internal/testutil/image_generation_cache.go (1)
RunImageGenerationCacheTest(14-140)core/internal/testutil/image_generation_errors.go (1)
RunImageGenerationErrorTest(14-154)core/internal/testutil/image_generation_load.go (3)
RunImageGenerationLoadTest(16-101)RunImageGenerationStreamLoadTest(104-211)RunImageGenerationCacheLoadTest(214-307)
core/internal/testutil/image_generation_stream.go (2)
core/schemas/images.go (4)
BifrostImageGenerationRequest(4-11)ImageGenerationInput(18-20)ImageGenerationParameters(22-30)BifrostImageGenerationStreamResponse(55-65)core/schemas/bifrost.go (2)
ImageGenerationStreamRequest(101-101)BifrostStream(340-348)
core/utils.go (1)
core/schemas/bifrost.go (6)
TextCompletionStreamRequest(90-90)ChatCompletionStreamRequest(92-92)ResponsesStreamRequest(94-94)SpeechStreamRequest(97-97)TranscriptionStreamRequest(99-99)ImageGenerationStreamRequest(101-101)
core/bifrost.go (6)
core/schemas/bifrost.go (6)
BifrostError(378-387)OpenAI(35-35)ModelProvider(32-32)ImageGenerationRequest(100-100)RequestType(83-83)BifrostStream(340-348)core/schemas/utils.go (2)
ParseModelString(23-36)Ptr(16-18)core/schemas/images.go (4)
BifrostImageGenerationRequest(4-11)ImageGenerationInput(18-20)ImageGenerationParameters(22-30)BifrostImageGenerationResponse(33-40)core/schemas/provider.go (1)
Provider(207-240)core/utils.go (1)
Ptr(51-53)core/schemas/responses.go (3)
ResponsesMessage(310-322)ResponsesImageGenerationCall(762-764)ResponsesMessageTypeImageGenerationCall(299-299)
core/providers/parasail/parasail.go (5)
core/schemas/account.go (1)
Key(8-17)core/schemas/images.go (2)
BifrostImageGenerationRequest(4-11)BifrostImageGenerationResponse(33-40)core/providers/utils/utils.go (1)
NewUnsupportedOperationError(391-402)core/schemas/bifrost.go (2)
ImageGenerationRequest(100-100)BifrostStream(340-348)core/schemas/provider.go (1)
PostHookRunner(204-204)
core/providers/openai/images.go (4)
core/providers/openai/openai.go (1)
ParseOpenAIError(2080-2108)core/schemas/images.go (4)
ImageGenerationParameters(22-30)ImageData(42-47)ImageUsage(49-52)BifrostImageGenerationStreamResponse(55-65)core/schemas/bifrost.go (8)
BifrostError(378-387)OpenAI(35-35)RequestType(83-83)BifrostResponseExtraFields(306-315)BifrostStream(340-348)ModelProvider(32-32)ErrorField(396-403)BifrostErrorExtraFields(445-449)core/providers/utils/utils.go (7)
CheckOperationAllowed(408-420)NewBifrostOperationError(449-460)SetExtraHeaders(178-208)ShouldSendBackRawResponse(482-487)ProcessAndSendResponse(535-565)GetBifrostResponseForStreamResponse(787-819)ProcessAndSendError(607-653)
framework/streaming/types.go (2)
core/schemas/images.go (2)
BifrostImageGenerationResponse(33-40)BifrostImageGenerationStreamResponse(55-65)core/schemas/provider.go (1)
Provider(207-240)
There was a problem hiding this comment.
Actionable comments posted: 2
🧹 Nitpick comments (8)
core/schemas/responses.go (1)
526-582: New image-generation tool output variant is wired correctlyThe added
ResponsesImageGenerationCallOutputfield and the updated marshal/unmarshal logic cleanly integrate an image-generation variant without breaking existing string/array/computer-output handling.framework/streaming/types.go (1)
16-19: Align image stream RequestType with other streaming responses
ToBifrostResponseuses non-stream request types (TextCompletionRequest,ChatCompletionRequest,ResponsesRequest,SpeechRequest,TranscriptionRequest) for aggregated streaming results, but the new image branch setsRequestTypetoImageGenerationStreamRequest. For consistency and simpler downstream handling (metrics, routing, caching), it should likely useImageGenerationRequestinstead.- case StreamTypeImage: - imageResp := p.Data.ImageGenerationOutput - if imageResp == nil { - imageResp = &schemas.BifrostImageGenerationResponse{} - } - resp.ImageGenerationResponse = imageResp - resp.ImageGenerationResponse.ExtraFields = schemas.BifrostResponseExtraFields{ - RequestType: schemas.ImageGenerationStreamRequest, - Provider: p.Provider, - ModelRequested: p.Model, - Latency: p.Data.Latency, - } + case StreamTypeImage: + imageResp := p.Data.ImageGenerationOutput + if imageResp == nil { + imageResp = &schemas.BifrostImageGenerationResponse{} + } + resp.ImageGenerationResponse = imageResp + resp.ImageGenerationResponse.ExtraFields = schemas.BifrostResponseExtraFields{ + RequestType: schemas.ImageGenerationRequest, + Provider: p.Provider, + ModelRequested: p.Model, + Latency: p.Data.Latency, + }Also applies to: 28-48, 98-109, 119-119, 258-269
core/images_test.go (2)
14-92: Parallel subtests in for-range loops: OK on Go 1.22+, consider defensive rebind
TestImageGenerationRequestSerializationandTestImageGenerationRequestDeserializationusefor _, tt := range tests { t.Run(tt.name, func(t *testing.T) { t.Parallel(); ... tt ... }) }. On Go 1.22+ the loop variable is captured by value; on older Go versions this pattern is racy. If you still support pre‑1.22, addtc := ttinside the loop and usetcin the closure to avoid accidental test flakiness.Also applies to: 95-161
198-255: Assert onjson.Unmarshalresult inTestImageStreamResponseSerializationThe test currently ignores the error from
json.Unmarshal(jsonBytes, &data). If the JSON shape drifts, the test could panic later on type assertions rather than failing with a clear message. Consider checking the error and failing explicitly, similar to the other tests in this file.- jsonBytes, _ := sonic.Marshal(tt.chunk) - var data map[string]interface{} - json.Unmarshal(jsonBytes, &data) + jsonBytes, err := sonic.Marshal(tt.chunk) + if err != nil { + t.Fatalf("failed to marshal chunk: %v", err) + } + var data map[string]interface{} + if err := json.Unmarshal(jsonBytes, &data); err != nil { + t.Fatalf("failed to unmarshal JSON: %v", err) + }framework/streaming/accumulator.go (1)
19-24: Image stream pooling and accumulator wiring look correct; consider using typed reset in CleanupThe new
imageStreamChunkPool,getImageStreamChunk/putImageStreamChunk,ImageStreamChunksslice,addImageStreamChunk, and prewarming inNewAccumulatormirror the existing patterns for other chunk types and should behave correctly for image streams.One small consistency nit: in
Cleanupyou currently return image chunks witha.imageStreamChunkPool.Put(chunk)rather thana.putImageStreamChunk(chunk), so any fields on those structs aren’t cleared before reuse. This matches the existing pattern for other chunk types, but if you later rely on partially populated chunks or add more fields, using the typedput*helpers here (for image and possibly the other chunk kinds) would make the pooling behavior more robust.Also applies to: 101-116, 118-130, 221-237, 259-261, 488-508
framework/streaming/images.go (1)
13-92: Image reconstruction logic is solid; optionally propagate usage into the final response
buildCompleteImageFromImageStreamChunkscorrectly:
- sorts by
ImageIndexthenChunkIndex,- accumulates
PartialB64per image,- preserves per-image
RevisedPrompt,- and derives a stable
IDandModel.Right now the returned
BifrostImageGenerationResponseleavesUsagenil, even though usage is later derived from the last chunk inprocessAccumulatedImageStreamingChunks. If you want the reconstructed response to be fully self-contained, consider also wiring thatImageUsageintofinalResponse.Usagewhen it’s available from the chunks.core/bifrost.go (2)
1130-1281: Tool → image request mapping is sound but duplicated; consider extracting a helper
ExecuteImageGenerationToolnicely:
- validates non-empty
argumentsandprompt,- defaults to
openai/dall-e-3and optionally parsesmodelviaParseModelString,- maps
size,quality,style,response_format, andnintoImageGenerationParameters,- calls
ImageGenerationRequest, and- returns a
toolroleChatMessagecontaining marshaledimageResp.DataandToolCallID.The same arg‑to‑
BifrostImageGenerationRequestmapping logic appears again inExecuteResponsesImageGenerationTool. To keep both call paths in sync as you add parameters (e.g.,user,extra_params), consider extracting a small helper likebuildImageGenerationRequestFromArgs(args map[string]any) (*schemas.BifrostImageGenerationRequest, *schemas.BifrostError)and reusing it in both executors.
2955-3020: Public image-generation APIs follow existing patterns; just ensure validateRequest covers the new types
ImageGenerationRequest/ImageGenerationStreamRequest:
- perform nil and prompt validation similar to other public methods,
- populate a pooled
BifrostRequestwithRequestTypeandImageGenerationRequest,- route through
handleRequest/handleStreamRequest,- and in the non-streaming case, defensively guard against
nilor missingImageGenerationResponse.This is consistent with the other entrypoints. Just make sure
validateRequest(and helpers likeIsStreamRequestType) have been updated elsewhere to treatImageGenerationRequestandImageGenerationStreamRequestthe same way as the existing request types (provider/model required, etc.); otherwise, you may want to add provider/model checks here for symmetry.Minor nit: the comment
// ImageGenerationStreamRequest sends a image generation stream requestshould read “an image generation…”.
📜 Review details
Configuration used: CodeRabbit UI
Review profile: CHILL
Plan: Pro
📒 Files selected for processing (28)
core/bifrost.go(7 hunks)core/images_test.go(1 hunks)core/internal/testutil/account.go(2 hunks)core/providers/anthropic/anthropic.go(4 hunks)core/providers/bedrock/bedrock.go(6 hunks)core/providers/cerebras/cerebras.go(1 hunks)core/providers/cohere/cohere.go(3 hunks)core/providers/elevenlabs/elevenlabs.go(3 hunks)core/providers/gemini/gemini.go(4 hunks)core/providers/groq/groq.go(1 hunks)core/providers/mistral/mistral.go(1 hunks)core/providers/ollama/ollama.go(1 hunks)core/providers/openai/openai.go(8 hunks)core/providers/openrouter/openrouter.go(1 hunks)core/providers/parasail/parasail.go(1 hunks)core/providers/perplexity/perplexity.go(1 hunks)core/providers/sgl/sgl.go(1 hunks)core/providers/utils/utils.go(2 hunks)core/providers/vertex/vertex.go(1 hunks)core/schemas/mux.go(1 hunks)core/schemas/provider.go(1 hunks)core/schemas/responses.go(3 hunks)core/utils.go(1 hunks)framework/streaming/accumulator.go(10 hunks)framework/streaming/images.go(1 hunks)framework/streaming/types.go(5 hunks)transports/bifrost-http/handlers/inference.go(6 hunks)ui/components/chat/ImageMessage.tsx(1 hunks)
🚧 Files skipped from review as they are similar to previous changes (12)
- core/schemas/mux.go
- core/utils.go
- core/schemas/provider.go
- core/providers/perplexity/perplexity.go
- core/providers/utils/utils.go
- core/providers/mistral/mistral.go
- core/internal/testutil/account.go
- ui/components/chat/ImageMessage.tsx
- core/providers/openai/openai.go
- core/providers/parasail/parasail.go
- core/providers/openrouter/openrouter.go
- core/providers/vertex/vertex.go
🧰 Additional context used
📓 Path-based instructions (1)
**
⚙️ CodeRabbit configuration file
always check the stack if there is one for the current PR. do not give localized reviews for the PR, always see all changes in the light of the whole stack of PRs (if there is a stack, if there is no stack you can continue to make localized suggestions/reviews)
Files:
core/providers/groq/groq.gocore/providers/sgl/sgl.gotransports/bifrost-http/handlers/inference.gocore/providers/elevenlabs/elevenlabs.gocore/images_test.gocore/providers/ollama/ollama.gocore/schemas/responses.goframework/streaming/types.gocore/providers/gemini/gemini.gocore/providers/bedrock/bedrock.gocore/providers/cohere/cohere.gocore/providers/anthropic/anthropic.gocore/providers/cerebras/cerebras.goframework/streaming/accumulator.goframework/streaming/images.gocore/bifrost.go
🧬 Code graph analysis (10)
core/providers/groq/groq.go (4)
core/schemas/images.go (2)
BifrostImageGenerationRequest(4-11)BifrostImageGenerationResponse(33-40)core/providers/utils/utils.go (1)
NewUnsupportedOperationError(391-402)core/schemas/bifrost.go (2)
ImageGenerationRequest(100-100)ImageGenerationStreamRequest(101-101)core/schemas/provider.go (1)
PostHookRunner(277-277)
core/providers/elevenlabs/elevenlabs.go (1)
core/providers/utils/utils.go (3)
ProcessAndSendResponse(535-565)GetBifrostResponseForStreamResponse(777-809)NewUnsupportedOperationError(391-402)
core/images_test.go (2)
core/schemas/images.go (7)
BifrostImageGenerationRequest(4-11)ImageGenerationInput(18-20)ImageGenerationParameters(22-30)BifrostImageGenerationResponse(33-40)ImageData(42-47)ImageUsage(49-52)BifrostImageGenerationStreamResponse(55-65)core/schemas/bifrost.go (1)
OpenAI(35-35)
core/providers/ollama/ollama.go (4)
core/schemas/images.go (2)
BifrostImageGenerationRequest(4-11)BifrostImageGenerationResponse(33-40)core/providers/utils/utils.go (1)
NewUnsupportedOperationError(391-402)core/schemas/bifrost.go (3)
ImageGenerationRequest(100-100)BifrostStream(340-348)ImageGenerationStreamRequest(101-101)core/schemas/provider.go (1)
PostHookRunner(277-277)
framework/streaming/types.go (3)
core/schemas/chatcompletions.go (1)
BifrostLLMUsage(648-655)core/schemas/bifrost.go (4)
BifrostError(378-387)BifrostCacheDebug(318-332)BifrostResponseExtraFields(306-315)RequestType(83-83)core/schemas/images.go (2)
BifrostImageGenerationResponse(33-40)BifrostImageGenerationStreamResponse(55-65)
core/providers/gemini/gemini.go (4)
core/providers/utils/utils.go (3)
ProcessAndSendResponse(535-565)GetBifrostResponseForStreamResponse(777-809)NewUnsupportedOperationError(391-402)core/schemas/images.go (2)
BifrostImageGenerationRequest(4-11)BifrostImageGenerationResponse(33-40)core/schemas/bifrost.go (3)
BifrostError(378-387)ImageGenerationRequest(100-100)BifrostStream(340-348)core/schemas/provider.go (1)
PostHookRunner(277-277)
core/providers/bedrock/bedrock.go (5)
core/providers/utils/utils.go (3)
ProcessAndSendResponse(535-565)GetBifrostResponseForStreamResponse(777-809)NewUnsupportedOperationError(391-402)core/schemas/account.go (1)
Key(8-17)core/schemas/images.go (2)
BifrostImageGenerationRequest(4-11)BifrostImageGenerationResponse(33-40)core/schemas/bifrost.go (3)
BifrostError(378-387)ImageGenerationRequest(100-100)BifrostStream(340-348)core/schemas/provider.go (1)
PostHookRunner(277-277)
core/providers/cohere/cohere.go (4)
core/providers/utils/utils.go (2)
GetBifrostResponseForStreamResponse(777-809)NewUnsupportedOperationError(391-402)core/schemas/images.go (2)
BifrostImageGenerationRequest(4-11)BifrostImageGenerationResponse(33-40)core/schemas/bifrost.go (4)
BifrostError(378-387)ImageGenerationRequest(100-100)BifrostStream(340-348)ImageGenerationStreamRequest(101-101)core/schemas/provider.go (1)
PostHookRunner(277-277)
core/providers/anthropic/anthropic.go (4)
core/schemas/account.go (1)
Key(8-17)core/schemas/images.go (2)
BifrostImageGenerationRequest(4-11)BifrostImageGenerationResponse(33-40)core/schemas/bifrost.go (4)
BifrostError(378-387)ImageGenerationRequest(100-100)BifrostStream(340-348)ImageGenerationStreamRequest(101-101)core/schemas/provider.go (1)
PostHookRunner(277-277)
framework/streaming/accumulator.go (2)
framework/streaming/types.go (4)
ImageStreamChunk(99-109)StreamAccumulator(112-124)ChatStreamChunk(75-84)ResponsesStreamChunk(87-96)core/schemas/bifrost.go (1)
ImageGenerationStreamRequest(101-101)
🔇 Additional comments (23)
core/providers/cerebras/cerebras.go (1)
206-215: ImageGeneration unsupported stubs are consistent and correctThese two methods cleanly implement the provider interface for image generation while marking Cerebras as unsupported. The request-type enums passed to
NewUnsupportedOperationError(ImageGenerationRequestvs.ImageGenerationStreamRequest) match the non‑stream/stream split and are consistent with the existing Embedding/Speech/Transcription patterns in this provider. No changes needed.core/providers/groq/groq.go (1)
247-255: Image generation stubs correctly signal unsupported operationsThe new
ImageGenerationandImageGenerationStreammethods are consistent with other providers’ unsupported stubs and use the correct request types inNewUnsupportedOperationError.core/providers/elevenlabs/elevenlabs.go (2)
407-425: Updated streaming helper usage is consistent with new signatureThe additional
nilargument inGetBifrostResponseForStreamResponsekeepsSpeechStreamaligned with the extended helper signature while preserving existing behavior.
655-663: Explicit image-generation stubs align Elevenlabs with provider interfaceThe new
ImageGenerationandImageGenerationStreammethods correctly advertise that ElevenLabs does not support image generation via standardized unsupported-operation errors.core/providers/sgl/sgl.go (1)
215-223: SGL image-generation stubs are consistent and non-disruptiveThe added
ImageGeneration/ImageGenerationStreammethods correctly surface unsupported image generation for SGL via standard errors without affecting existing functionality.transports/bifrost-http/handlers/inference.go (1)
141-152: HTTP request schema and route for image generation match existing patternsThe
imageParamsKnownFieldsmap,ImageGenerationHTTPRequeststruct, and/v1/images/generationsroute wiring mirror the existing completion/embeddings/speech patterns and look structurally sound.Also applies to: 190-195, 306-307
core/providers/ollama/ollama.go (1)
218-226: Ollama image-generation stubs correctly expose unsupported featureThe added
ImageGeneration/ImageGenerationStreammethods cleanly report that Ollama does not support image generation via standardized errors.core/providers/bedrock/bedrock.go (2)
653-654: Streaming helper updates for Bedrock are consistent with extended signatureAll Bedrock streaming paths now pass six arguments to
GetBifrostResponseForStreamResponse, with the image-generation stream slot set tonil, maintaining prior behavior for text/chat/responses streams.Also applies to: 873-884, 1032-1033, 1125-1126
1251-1259: Bedrock image-generation stubs cleanly signal unsupported operationsThe new
ImageGenerationandImageGenerationStreammethods correctly useImageGenerationRequest/ImageGenerationStreamRequestwithprovider.GetProviderKey(), aligning with the common unsupported-operation pattern.core/providers/cohere/cohere.go (2)
486-490: Streaming response wrapper arity update looks correctBoth ChatCompletionStream and ResponsesStream now call
GetBifrostResponseForStreamResponse(nil, response, nil, nil, nil, nil)/(..., nil, nil, response, nil, nil), which matches the extended signature and keeps the non-image paths unchanged.Also applies to: 716-720
804-812: Cohere image-generation stubs correctly signal unsupported operations
ImageGenerationandImageGenerationStreamcleanly returnNewUnsupportedOperationErrorwithImageGenerationRequestandImageGenerationStreamRequestrespectively, aligning with other unsupported methods in this provider and the new image-generation API surface.core/images_test.go (2)
257-392: Streaming chunk accumulation tests provide solid coverageThe three subtests around single-image multi-chunk, multi-image interleaved chunks, and out-of-order chunk sorting exercise the key reconstruction behaviors for
PartialB64and ordering byChunkIndex/image index. This is a good level of coverage for the new image streaming accumulator.
421-477: Cache-key tests and helper align with intended components
generateTestCacheKeynow concatenates prompt, size, quality, style, andnviafmt.Sprintf("%d", *req.Params.N), and the tests verify both determinism for identical requests and key changes for each param toggle. This matches the semantic cache requirements for image generation.Also applies to: 479-514
core/providers/gemini/gemini.go (2)
604-606: Streaming PostHook wrapper calls updated consistently for new response variantAll four
ProcessAndSendResponsecall sites inSpeechStreamandTranscriptionStreamnow invokeGetBifrostResponseForStreamResponse(nil, nil, nil, response, nil, nil)/(..., nil, response, nil)as appropriate. This matches the extended factory signature and keeps speech/transcription behavior unchanged.Also applies to: 627-628, 862-863, 890-891
897-905: Gemini image-generation methods correctly report unsupported operations
ImageGenerationandImageGenerationStreamreturnNewUnsupportedOperationErrorwithImageGenerationRequestandImageGenerationStreamRequestrespectively, keeping Gemini’s API surface compatible with the core image-generation interfaces while clearly signaling lack of support.core/providers/anthropic/anthropic.go (2)
606-607: Anthropic streaming helpers correctly adapted to extended stream response factoryBoth the chat and responses streaming helpers now use
GetBifrostResponseForStreamResponse(nil, response, nil, nil, nil, nil)and(..., nil, response, nil, nil, nil)respectively, which matches the new image-aware signature while leaving existing behaviors unchanged.Also applies to: 632-633, 937-941
983-991: Anthropic image-generation stubs mirror other unsupported operations
ImageGenerationandImageGenerationStreamcleanly returnNewUnsupportedOperationErrorforImageGenerationRequestandImageGenerationStreamRequest, keeping the provider interface complete for the new request types without implying support Anthropic doesn’t have.framework/streaming/accumulator.go (1)
352-372: Image streaming dispatch integrates cleanly into ProcessStreamingResponseAdding
isImageStreaming := requestType == schemas.ImageGenerationStreamRequestand delegating toprocessImageStreamingResponseextends the accumulator’s routing without affecting existing chat/audio/responses paths. This keeps the centralized dispatch logic easy to follow as new stream types are introduced.framework/streaming/images.go (2)
94-174: Accumulated image streaming data covers latency, usage, cost, and cache debug
processAccumulatedImageStreamingChunkscorrectly:
- protects the accumulator with
mu,- computes latency from
StartTimestamp/FinalTimestamp,- attaches the reconstructed image response,
- derives
TokenUsagefrom the last chunk’sUsage,- and surfaces cost, cache debug info, and
FinishReasonfrom the final chunk.The cleanup-on-final-chunk path (
cleanupStreamAccumulatorin the defer) is a nice touch to avoid leaking pooled chunks. This looks good.
176-278: processImageStreamingResponse correctly wires per-chunk metadata and final aggregationThe image streaming handler:
- extracts the request ID and provider/model via
GetResponseFields,- builds a pooled
ImageStreamChunk, copying the stream delta to avoid stack aliasing,- marks error vs completed finish reasons appropriately,
- records pricing and semantic cache debug on the final chunk, and
- uses the
IsCompleteflag under lock so accumulated chunks are processed exactly once per request.The separation between per-chunk delta responses (returned as
StreamResponseTypeDelta) and the final aggregatedStreamResponseTypeFinalwithAccumulatedDatais clear and matches the rest of the streaming framework.core/bifrost.go (3)
1878-1883: Fallback support extended correctly to image-generation requestsThe new
ImageGenerationRequestbranch inprepareFallbackRequestcopies the original request and overridesProvider/Modelwith the fallback’s, consistent with how text/chat/responses/speech/transcription fallbacks are handled. This should make image generation participate in the same fallback flow without surprises.
2555-2560: Provider dispatch now includes image-generation and image-generation-stream requestsAdding
ImageGenerationRequestandImageGenerationStreamRequestcases tohandleProviderRequestandhandleProviderStreamRequestensures image flows go through the same centralized dispatch as other request types and benefit from retries, plugins, and fallbacks. This aligns well with the rest of the core routing.Also applies to: 2585-2587
2748-2758: BifrostRequest pool reset includes image-generation fieldExtending
resetBifrostRequestto clearImageGenerationRequestkeeps the request pooling safe for reuse now that image-generation is part of the schema, and avoids stale pointers leaking between calls.
Pratham-Mishra04
left a comment
Pratham-Mishra04
left a comment
There was a problem hiding this comment.
added comments
Radheshg04
force-pushed
the
feature/image-generation-950
branch
from
b5899e6 to
f50dccc
Compare
There was a problem hiding this comment.
Actionable comments posted: 0
Caution
Some comments are outside the diff and can’t be posted inline due to platform limitations.
⚠️ Outside diff range comments (1)
core/internal/testutil/account.go (1)
541-563: Enable image generation scenarios in AllProviderConfigs for OpenAI.The OpenAI entry sets
ImageGenerationModel: "dall-e-2"but doesn't enableImageGeneration: trueorImageGenerationStream: truein theScenariosstruct. This means image generation tests won't run when usingAllProviderConfigs.ImageGenerationModel: "dall-e-2", Scenarios: TestScenarios{ TextCompletion: false, // Not supported TextCompletionStream: false, // Not supported SimpleChat: true, CompletionStream: true, MultiTurnConversation: true, ToolCalls: true, MultipleToolCalls: true, End2EndToolCalling: true, AutomaticFunctionCall: true, ImageURL: true, ImageBase64: true, MultipleImages: true, CompleteEnd2End: true, SpeechSynthesis: true, // OpenAI supports TTS SpeechSynthesisStream: true, // OpenAI supports streaming TTS Transcription: true, // OpenAI supports STT with Whisper TranscriptionStream: true, // OpenAI supports streaming STT Embedding: true, Reasoning: true, // OpenAI supports reasoning via o1 models ListModels: true, + ImageGeneration: true, // OpenAI supports image generation + ImageGenerationStream: true, // OpenAI supports streaming image generation },
♻️ Duplicate comments (1)
transports/bifrost-http/handlers/inference.go (1)
1184-1203: Fix premature context cancellation in imageGeneration streaming path
defer cancel()is executed before the stream check, so for streaming requests the context is canceled as soon asimageGenerationreturns, which will terminate the upstream image stream almost immediately instead of only on client disconnect (unlike the other stream handlers).Align this with
textCompletion,chatCompletion,responses, andspeechby deferringcancelonly for the non‑stream path:// Convert context bifrostCtx, cancel := lib.ConvertToBifrostContext(ctx, h.handlerStore.ShouldAllowDirectKeys()) if bifrostCtx == nil { SendError(ctx, fasthttp.StatusInternalServerError, "Failed to convert context") return } - defer cancel() - - // Handle streaming image generation - if req.BifrostParams.Stream != nil && *req.BifrostParams.Stream { + + // Handle streaming image generation + if req.Stream != nil && *req.Stream { if req.ResponseFormat != nil && *req.ResponseFormat == "url" { SendError(ctx, fasthttp.StatusBadRequest, "streaming images must be requested in base64") return } h.handleStreamingImageGeneration(ctx, bifrostReq, bifrostCtx, cancel) return } + + // Non-streaming: cancel on handler exit + defer cancel() // Execute request resp, bifrostErr := h.client.ImageGenerationRequest(*bifrostCtx, bifrostReq)This preserves the cancellable context for the duration of the stream and only cancels when
handleStreamingResponsedetects client disconnects.
🧹 Nitpick comments (9)
transports/bifrost-http/handlers/inference.go (1)
1003-1006: Avoid logging full streaming chunk payload on marshal errorsIncluding
%vofchunkcan dump large base64 blobs (and potentially sensitive content) into logs, especially for image streams.Consider logging only high‑level metadata or a truncated representation instead:
- chunkJSON, err := sonic.Marshal(chunk) - if err != nil { - logger.Warn(fmt.Sprintf("Failed to marshal streaming response: %v, chunk: %v", err, chunk)) - continue - } + chunkJSON, err := sonic.Marshal(chunk) + if err != nil { + logger.Warn(fmt.Sprintf("Failed to marshal streaming response: %v (omitting chunk payload)", err)) + continue + }core/providers/utils/utils.go (1)
586-596: Consider adding image generation stream response handling to error functions for consistency.
ProcessAndSendBifrostErrorandProcessAndSendErrordon't propagateImageGenerationStreamResponsewhenprocessedResponseis non-nil, unlikeProcessAndSendResponse. While error scenarios typically don't carry image data, the inconsistency could cause issues if a post-hook transforms an error response into a partial image response.core/internal/testutil/image_generation.go (1)
100-105: Consider increasing the minimum base64 length threshold.A 100-byte threshold is extremely low for validating image data. A 1024x1024 image in base64 would typically be several hundred kilobytes to megabytes. This could allow corrupt or placeholder responses to pass validation.
// Validate base64 if present if imageData.B64JSON != "" { - if len(imageData.B64JSON) < 100 { + if len(imageData.B64JSON) < 10000 { // ~7KB minimum for a reasonable image t.Errorf("❌ Base64 image data too short: %d bytes", len(imageData.B64JSON)) } }core/internal/testutil/test_retry_conditions.go (1)
845-881: EmptyImageGenerationCondition logic is sound; only a minor comment nitThe retry condition correctly covers nil responses, missing
ImageGenerationResponse, emptyData, and entries lacking bothB64JSONandURL, consistent with other modality “empty” checks. The implementation looks good; consider fixing the comment typo (“image image”) for clarity.core/internal/testutil/image_generation_load.go (1)
213-307: Avoid sharing a single request instance across cache-load goroutinesIn the cache load test, the same
requestpointer is reused in all goroutines:
- If
client.ImageGenerationRequestor downstream code mutates the request (e.g., addsRawRequestBodyor other fields), this can introduce data races when running tests with-race.- It’s safer to build a fresh request per goroutine (while reusing
cachePromptand common params) to ensure immutability of arguments.You can keep the initial “populate cache” request as-is, and construct a per-goroutine request for the parallel section, e.g.:
- for i := 0; i < numConcurrentRequests; i++ { + for i := 0; i < numConcurrentRequests; i++ { wg.Add(1) go func() { defer wg.Done() - reqStart := time.Now() - response, bifrostErr := client.ImageGenerationRequest(ctx, request) + // Build a fresh request per goroutine to avoid shared mutable state. + localReq := &schemas.BifrostImageGenerationRequest{ + Provider: testConfig.Provider, + Model: testConfig.ImageGenerationModel, + Input: &schemas.ImageGenerationInput{ + Prompt: cachePrompt, + }, + Params: &schemas.ImageGenerationParameters{ + Size: bifrost.Ptr("1024x1024"), + ResponseFormat: bifrost.Ptr("b64_json"), + }, + } + + reqStart := time.Now() + response, bifrostErr := client.ImageGenerationRequest(ctx, localReq)core/internal/testutil/response_validation.go (1)
1053-1122: ImageGeneration field checks are solid; consider extending technical expectations laterThe helper functions:
- Enforce that
Datais non-empty and every image has eitherB64JSONorURL.- Optionally enforce a minimum number of images via
ProviderSpecific["min_images"].- Record
expected_sizeand validate latency whenShouldHaveLatencyis set.- Collect useful metrics (image count, URL vs B64 counts, usage tokens).
This is a good baseline. If you later need stricter technical validation for images, you could also honor
ResponseExpectations.ShouldHaveUsageStats/ShouldHaveModelsimilarly to chat/text, but that’s an optional enhancement, not a blocker.core/internal/testutil/test_retry_framework.go (3)
161-249: ImageGeneration retry interfaces and configs integrate cleanly with existing patternsAdding
ImageGenerationRetryConditionandImageGenerationRetryConfig, plusDefaultImageGenerationRetryConfigand the"ImageGeneration"/"ImageGenerationStream"cases inGetTestRetryConfigForScenario, follows the same structure used for other modalities (chat, speech, embeddings, etc.). UsingEmptyImageGenerationCondition+GenericResponseConditionas defaults makes sense and aligns with the new image-specific emptiness check.The only minor nit is that
DefaultImageGenerationRetryConfigreturns the genericTestRetryConfigwhile there is also anImageGenerationRetryConfigtype; this matches existing naming for other helpers (e.g., speech/transcription) but can be slightly confusing at first glance. Not blocking, just something to keep in mind for future naming consistency.Also applies to: 926-940, 1176-1177
1875-2049: WithImageGenerationRetry behavior matches other typed retry helpers
WithImageGenerationRetry:
- Runs the operation, validates via
ValidateImageGenerationResponsefirst, and returns immediately on success.- Always retries on validation failures (primary purpose), with exponential backoff, and prioritizes timeout errors via
isTimeoutError.- Incorporates image-specific retry conditions through
checkImageGenerationRetryConditions.- For error-only cases, it always retries non-structural errors at least once, mirroring the transcription/embedding retry behavior.
- On exhausting attempts, it emits
OnFinalFailand returns a syntheticBifrostErrorto force test failure.This is consistent with the rest of the framework and looks correct. There is some inevitable duplication with other
With*TestRetryfunctions; if this pattern grows further, it might be worth considering a generic helper, but that’s an optional future refactor.
2697-2845: Image-generation stream retry wrapper matches other stream validators; small comment/name nit
ImageGenerationStreamValidationResultandWithImageGenerationStreamRetryclosely mirror existing speech/chat/responses streaming validation wrappers: they retry on stream setup errors (with special handling for timeouts), validate the full stream via a callback, and always retry on validation failures with exponential backoff and appropriate logging/OnFinalFailhooks. Behavior is consistent and looks good.Minor nit: the function comment refers to
WithImageGenerationStreamValidationRetry, but the function is namedWithImageGenerationStreamRetry. Updating the comment to match the actual name would avoid confusion.
📜 Review details
Configuration used: CodeRabbit UI
Review profile: CHILL
Plan: Pro
📒 Files selected for processing (14)
core/internal/testutil/account.go(3 hunks)core/internal/testutil/image_generation.go(1 hunks)core/internal/testutil/image_generation_load.go(1 hunks)core/internal/testutil/image_generation_stream.go(1 hunks)core/internal/testutil/response_validation.go(2 hunks)core/internal/testutil/test_retry_conditions.go(1 hunks)core/internal/testutil/test_retry_framework.go(6 hunks)core/internal/testutil/validation_presets.go(2 hunks)core/providers/openai/images.go(1 hunks)core/providers/openai/openai.go(9 hunks)core/providers/openai/openai_test.go(2 hunks)core/providers/openai/types.go(1 hunks)core/providers/utils/utils.go(3 hunks)transports/bifrost-http/handlers/inference.go(7 hunks)
🚧 Files skipped from review as they are similar to previous changes (2)
- core/internal/testutil/image_generation_stream.go
- core/providers/openai/images.go
🧰 Additional context used
📓 Path-based instructions (1)
**
⚙️ CodeRabbit configuration file
always check the stack if there is one for the current PR. do not give localized reviews for the PR, always see all changes in the light of the whole stack of PRs (if there is a stack, if there is no stack you can continue to make localized suggestions/reviews)
Files:
core/internal/testutil/image_generation.gocore/internal/testutil/test_retry_conditions.gocore/internal/testutil/account.gocore/providers/openai/openai_test.gocore/providers/openai/types.gocore/providers/openai/openai.gocore/internal/testutil/response_validation.gocore/internal/testutil/validation_presets.gotransports/bifrost-http/handlers/inference.gocore/internal/testutil/image_generation_load.gocore/internal/testutil/test_retry_framework.gocore/providers/utils/utils.go
🧠 Learnings (1)
📚 Learning: 2025-12-09T17:07:42.007Z
Learnt from: qwerty-dvorak
Repo: maximhq/bifrost PR: 1006
File: core/schemas/account.go:9-18
Timestamp: 2025-12-09T17:07:42.007Z
Learning: In core/schemas/account.go, the HuggingFaceKeyConfig field within the Key struct is currently unused and reserved for future Hugging Face inference endpoint deployments. Do not flag this field as missing from OpenAPI documentation or require its presence in the API spec until the feature is actively implemented and used. When the feature is added, update the OpenAPI docs accordingly; otherwise, treat this field as non-breaking and not part of the current API surface.
Applied to files:
core/internal/testutil/image_generation.gocore/internal/testutil/test_retry_conditions.gocore/internal/testutil/account.gocore/providers/openai/openai_test.gocore/providers/openai/types.gocore/providers/openai/openai.gocore/internal/testutil/response_validation.gocore/internal/testutil/validation_presets.gotransports/bifrost-http/handlers/inference.gocore/internal/testutil/image_generation_load.gocore/internal/testutil/test_retry_framework.gocore/providers/utils/utils.go
🧬 Code graph analysis (9)
core/internal/testutil/test_retry_conditions.go (2)
core/schemas/bifrost.go (1)
BifrostResponse(262-274)core/internal/testutil/test_retry_framework.go (1)
TestRetryContext(174-179)
core/internal/testutil/account.go (1)
core/schemas/bifrost.go (1)
Fallback(132-135)
core/providers/openai/types.go (1)
core/providers/openai/images.go (1)
ImageGenerationEventType(20-20)
core/providers/openai/openai.go (1)
core/providers/utils/utils.go (2)
ProcessAndSendResponse(535-566)GetBifrostResponseForStreamResponse(788-820)
core/internal/testutil/response_validation.go (2)
core/schemas/images.go (1)
BifrostImageGenerationResponse(33-40)core/internal/testutil/error_parser.go (3)
ParseBifrostError(47-104)FormatErrorConcise(353-356)LogError(359-366)
core/internal/testutil/validation_presets.go (1)
core/internal/testutil/response_validation.go (1)
ResponseExpectations(18-43)
transports/bifrost-http/handlers/inference.go (8)
core/schemas/images.go (3)
ImageGenerationInput(18-20)ImageGenerationParameters(22-30)BifrostImageGenerationRequest(4-11)transports/bifrost-http/lib/middleware.go (1)
ChainMiddlewares(12-24)transports/bifrost-http/handlers/utils.go (2)
SendError(35-44)SendBifrostError(47-62)core/schemas/utils.go (1)
ParseModelString(23-36)core/schemas/models.go (1)
Model(109-129)core/schemas/provider.go (1)
Provider(207-240)core/schemas/bifrost.go (5)
ModelProvider(32-32)ImageGenerationRequest(100-100)BifrostStream(340-348)BifrostError(378-387)ImageGenerationStreamRequest(101-101)transports/bifrost-http/lib/ctx.go (1)
ConvertToBifrostContext(72-283)
core/internal/testutil/image_generation_load.go (3)
core/internal/testutil/account.go (1)
ComprehensiveTestConfig(49-68)core/schemas/images.go (4)
BifrostImageGenerationRequest(4-11)ImageGenerationInput(18-20)ImageGenerationParameters(22-30)BifrostImageGenerationStreamResponse(55-65)core/schemas/bifrost.go (2)
ImageGenerationRequest(100-100)ImageGenerationStreamRequest(101-101)
core/internal/testutil/test_retry_framework.go (5)
core/schemas/images.go (1)
BifrostImageGenerationResponse(33-40)core/internal/testutil/test_retry_conditions.go (2)
EmptyImageGenerationCondition(851-851)GenericResponseCondition(370-370)core/internal/testutil/response_validation.go (2)
ResponseExpectations(18-43)ValidateImageGenerationResponse(241-275)core/internal/testutil/utils.go (1)
GetErrorMessage(606-639)core/schemas/bifrost.go (1)
BifrostStream(340-348)
🔇 Additional comments (26)
transports/bifrost-http/handlers/inference.go (4)
5-25: Imports cleanup looks goodImport set is minimal and all packages are used; the commented
schemasimport has been properly cleaned up.
141-152: Image params whitelist is consistent
imageParamsKnownFieldscorrectly captures the expected top-level fields for/v1/images/generations, so unknown keys will flow intoExtraParamsas intended.
190-194: HTTP request wrapper for images matches existing patternsEmbedding
ImageGenerationInput,ImageGenerationParameters, andBifrostParamsis consistent with other request types and gives the expected JSON shape.
1212-1222: Streaming image handler wiring is consistent
handleStreamingImageGenerationmirrors the existing streaming helpers (text/chat/responses/speech), delegating toImageGenerationStreamRequestand the shared SSE handler, so no further changes are needed once the context lifecycle is fixed in the caller.core/providers/openai/openai_test.go (1)
40-40: LGTM!The test configuration correctly enables image generation testing with
dall-e-3and enables both streaming and non-streaming test scenarios. This aligns well with the newImageGenerationandImageGenerationStreamcapabilities added to the OpenAI provider.Also applies to: 63-64
core/providers/utils/utils.go (2)
555-556: LGTM!The
ProcessAndSendResponsefunction correctly propagatesImageGenerationStreamResponseto the unified stream response, maintaining consistency with other response types (text completion, chat, responses, speech, transcription).
794-817: LGTM!The
GetBifrostResponseForStreamResponsefunction is correctly extended with the newimageGenerationStreamResponseparameter. The switch-case handling follows the established pattern for other response types.core/providers/openai/openai.go (8)
550-550: LGTM!The text completion streaming handlers correctly pass
nilfor the newimageGenerationStreamResponseparameter toGetBifrostResponseForStreamResponse.Also applies to: 570-570
938-945: LGTM!The responses-to-chat-completions fallback streaming path correctly passes
nilfor the image generation parameter.
1013-1013: LGTM!The chat completion streaming handler correctly passes
nilfor the image generation parameter at both chunk emission points.Also applies to: 1034-1034
1366-1366: LGTM!The responses streaming handler correctly passes
nilfor the image generation parameter.Also applies to: 1373-1373
1741-1741: LGTM!The speech streaming handler correctly passes
nilfor the image generation parameter.Also applies to: 1745-1745
2013-2013: LGTM!The transcription streaming handler correctly passes
nilfor the image generation parameter.Also applies to: 2017-2017
2138-2164: Cannot complete verification due to repository access failure.The original review comment requests verification of the
doRequestmethod's existence and correctness in handling image generation API calls. This verification cannot be performed without access to the repository code. The review comment's approval is conditional on this verification being completed.
2169-2204: Verify that streaming dependencies are properly defined.The
ImageGenerationStreammethod callsHandleOpenAIImageGenerationStreaming,simulateImageStreaming, and referencesStreamingEnabledImageModels. These symbols do not appear in the public package documentation and should be verified to exist in the codebase and be correctly implemented.core/internal/testutil/account.go (2)
44-45: LGTM!The
TestScenariosstruct is correctly extended withImageGenerationandImageGenerationStreamflags, following the established pattern for other test scenarios.
66-67: LGTM!The
ComprehensiveTestConfigstruct correctly addsImageGenerationModelandImageGenerationFallbacksfields, consistent with other capability-specific configurations (e.g.,TranscriptionModel/TranscriptionFallbacks).core/internal/testutil/image_generation.go (3)
1-10: LGTM!Clean imports and package declaration for the new image generation test utility.
24-46: LGTM!Retry configuration is properly set up using the test framework's retry infrastructure, addressing flakiness concerns for CI. The
WithImageGenerationRetrywrapper is correctly used.
48-77: LGTM!The test operation function correctly builds the
BifrostImageGenerationRequestwith appropriate defaults and handles the response/error flow properly. The nil response is correctly converted to aBifrostError.core/internal/testutil/validation_presets.go (2)
188-203: LGTM!The
ImageGenerationExpectationsfunction correctly configures validation expectations for image generation scenarios. The settings appropriately disable text content and choices expectations while enabling metadata validation.
305-311: LGTM!The
GetExpectationsForScenarioswitch case for"ImageGeneration"correctly handles custom parameters with proper type assertions and provides sensible defaults (1 image, 1024x1024 size).core/internal/testutil/image_generation_load.go (2)
15-101: ImageGeneration load test is well-structured and concurrency-safeThe load test uses
t.Run+t.Parallel,WaitGroup, atomics, and a mutex for shared metrics correctly. Success/error accounting and basic latency/RPS stats look reasonable for exercising the image-generation path under modest concurrency.
103-211: Verify completion Type string and confirm stream-drain safetyTwo issues require verification before merge:
The completion check uses
Type == "image_generation.completed". The schema documentation mentions values like"image.chunk"/"image.complete"/"error". Confirm this string matches the actual completion marker fromBifrostImageGenerationStreamResponseincore/schemas/images.go, or replace with a shared constant to avoid silent failures where all streams count as incomplete.The loop breaks immediately upon seeing the completion event without draining remaining channel items. Verify the producer guarantees the channel closes immediately after sending the completion event. If bookkeeping events can follow, drain the channel fully to prevent blocking the producer goroutine.
core/internal/testutil/response_validation.go (1)
240-275: Image-generation validator follows existing validation pattern correctly
ValidateImageGenerationResponsemirrors the other modality validators: it treats non-nilerrand nil responses as failures, routes errors throughParseBifrostError/LogError, delegates to image-specific field validation, collects metrics, and logs vialogValidationResults. The control flow and error handling are consistent and look good.core/providers/openai/types.go (1)
162-208: OpenAI image generation structs are correctly aligned with current API specificationThe request, response, usage, and streaming event structs precisely match the OpenAI images/generations API schema (2025): created timestamp, data array with URL/b64_json/revised_prompt, background/output_format/size/quality fields, and usage tracking with input_tokens_details.
If other provider request types in this codebase expose a Bifrost-specific
Fallbacksfield for provider→Bifrost round-trips, consider adding it toOpenAIImageGenerationRequestas well for consistency. This is optional and depends on whether your image generation fallback flow requires it.
|
Hello @Radheshg04, can you pls resolve merge conflicts on this |
There was a problem hiding this comment.
Actionable comments posted: 2
🧹 Nitpick comments (1)
core/providers/anthropic/anthropic.go (1)
632-633: Updated stream-response wrapper call is correct, but the “many nils” API is brittle.
Passing the new 6th arg asnilmatches the updatedGetBifrostResponseForStreamResponsesignature; consider a small helper (per stream type) to avoid future arg-order mistakes.Also applies to: 662-663
📜 Review details
Configuration used: CodeRabbit UI
Review profile: CHILL
Plan: Pro
📒 Files selected for processing (2)
core/providers/anthropic/anthropic.go(4 hunks)core/providers/openai/openai.go(9 hunks)
🧰 Additional context used
📓 Path-based instructions (1)
**
⚙️ CodeRabbit configuration file
always check the stack if there is one for the current PR. do not give localized reviews for the PR, always see all changes in the light of the whole stack of PRs (if there is a stack, if there is no stack you can continue to make localized suggestions/reviews)
Files:
core/providers/openai/openai.gocore/providers/anthropic/anthropic.go
🧠 Learnings (2)
📚 Learning: 2025-12-09T17:07:42.007Z
Learnt from: qwerty-dvorak
Repo: maximhq/bifrost PR: 1006
File: core/schemas/account.go:9-18
Timestamp: 2025-12-09T17:07:42.007Z
Learning: In core/schemas/account.go, the HuggingFaceKeyConfig field within the Key struct is currently unused and reserved for future Hugging Face inference endpoint deployments. Do not flag this field as missing from OpenAPI documentation or require its presence in the API spec until the feature is actively implemented and used. When the feature is added, update the OpenAPI docs accordingly; otherwise, treat this field as non-breaking and not part of the current API surface.
Applied to files:
core/providers/openai/openai.gocore/providers/anthropic/anthropic.go
📚 Learning: 2025-12-11T11:58:25.307Z
Learnt from: Pratham-Mishra04
Repo: maximhq/bifrost PR: 1000
File: core/providers/openai/responses.go:42-84
Timestamp: 2025-12-11T11:58:25.307Z
Learning: In core/providers/openai/responses.go (and related OpenAI response handling), document and enforce the API format constraint: if ResponsesReasoning != nil and the response contains content blocks, all content blocks should be treated as reasoning blocks by default. Implement type guards or parsing logic accordingly, and add unit tests to verify that when ResponsesReasoning is non-nil, content blocks are labeled as reasoning blocks. Include clear comments in the code explaining the rationale and ensure downstream consumers rely on this behavior.
Applied to files:
core/providers/openai/openai.go
🧬 Code graph analysis (2)
core/providers/openai/openai.go (1)
core/providers/utils/utils.go (1)
GetBifrostResponseForStreamResponse(788-820)
core/providers/anthropic/anthropic.go (3)
core/providers/utils/utils.go (3)
ProcessAndSendResponse(535-566)GetBifrostResponseForStreamResponse(788-820)NewUnsupportedOperationError(391-402)core/schemas/bifrost.go (2)
BifrostError(378-387)BifrostStream(340-348)core/schemas/provider.go (1)
PostHookRunner(204-204)
🔇 Additional comments (3)
core/providers/anthropic/anthropic.go (2)
981-982: Responses stream wrapper update looks consistent with new image-stream slot.
These call sites correctly populate only theresponsesStreamResponseposition and passnilfor the new image-stream parameter.Also applies to: 984-985
1026-1035: TheImageGenerationStreamparameter type is correct. The interface (core/schemas/provider.go:238-239) defines the parameter as*BifrostImageGenerationRequest, and the Anthropic provider's implementation matches this exactly. All other providers consistently use the same parameter type. The error constantschemas.ImageGenerationStreamRequestis for error messaging only, not a request type.core/providers/openai/openai.go (1)
578-578: LGTM! Mechanical updates correctly handle signature change.All calls to
GetBifrostResponseForStreamResponsehave been correctly updated to passnilfor the newimageGenerationStreamResponseparameter. This maintains consistency across all existing streaming handlers while accommodating the new image generation streaming support.Also applies to: 602-602, 983-983, 990-990, 1064-1064, 1089-1089, 1434-1434, 1441-1441, 1817-1817, 1821-1821, 2089-2089, 2093-2093
Radheshg04
force-pushed
the
feature/image-generation-950
branch
from
bb970dd to
f50dccc
Compare
There was a problem hiding this comment.
Actionable comments posted: 20
Caution
Some comments are outside the diff and can’t be posted inline due to platform limitations.
⚠️ Outside diff range comments (1)
core/internal/testutil/tests.go (1)
24-62: UpdateprintTestSummaryto include new image-generation scenarios.Right now the runner executes image-generation tests, but the summary table doesn’t mention them, so it will under-report enabled coverage.
Also applies to: 73-101
♻️ Duplicate comments (3)
transports/bifrost-http/handlers/inference.go (1)
1260-1270: Fix premature context cancellation in streaming image generation handler.Line 1260 calls
defer cancel()before checking if the request is streaming (line 1263). For streaming requests, this means the context will be canceled as soon as the handler returns at line 1269, terminating the upstream image generation stream prematurely instead of only on client disconnect.This issue was previously flagged in past review comments. Align with the pattern used in other streaming handlers (
textCompletion,chatCompletion,responses,speech):Apply this diff:
// Convert context bifrostCtx, cancel := lib.ConvertToBifrostContext(ctx, h.handlerStore.ShouldAllowDirectKeys()) if bifrostCtx == nil { SendError(ctx, fasthttp.StatusInternalServerError, "Failed to convert context") return } - defer cancel() // Handle streaming image generation if req.BifrostParams.Stream != nil && *req.BifrostParams.Stream { if req.ResponseFormat != nil && *req.ResponseFormat == "url" { SendError(ctx, fasthttp.StatusBadRequest, "streaming images must be requested in base64") return } h.handleStreamingImageGeneration(ctx, bifrostReq, bifrostCtx, cancel) return } + // Non-streaming: cancel on handler exit + defer cancel() + // Execute request resp, bifrostErr := h.client.ImageGenerationRequest(*bifrostCtx, bifrostReq)core/providers/openai/openai.go (1)
2154-2179: Raw request handling gaps in ImageGeneration (+ missingsendBackRawRequestin streaming) — still unresolved.
This matches prior review feedback: non-streaming image generation doesn’t appear to honor raw request/response flags, and the streaming call toHandleOpenAIImageGenerationStreamingpasses onlysendBackRawResponse(nosendBackRawRequest).Also applies to: 2204-2219
core/bifrost.go (1)
1411-1570: Responses image-gen tool outputTypelooks mislabeled (call vs output).
This still setsoutputType := schemas.ResponsesMessageTypeImageGenerationCallwhile building an output payload (ResponsesImageGenerationCallOutput). If downstream switches onResponsesMessage.Type, this can break routing. (This matches the earlier review note.)
🧹 Nitpick comments (19)
ui/components/chat/ImageMessage.tsx (1)
5-57: Use stable keys + avoid “Loading… undefined%”; small accessibility/UX tweaks.
Right nowkey={idx}can cause remounts if the array is reordered; andstreamProgressbeing optional can render asundefined%.interface ImageMessageProps { @@ isStreaming?: boolean; streamProgress?: number; // 0-100 } export const ImageMessage: React.FC<ImageMessageProps> = ({ images, isStreaming, - streamProgress + streamProgress = 0, }) => { return ( - <div className="grid grid-cols-2 gap-4 my-4"> + <div className="grid grid-cols-1 sm:grid-cols-2 gap-4 my-4"> {images.map((img, idx) => ( - <Card key={idx} className="overflow-hidden"> + <Card key={img.index} className="overflow-hidden"> @@ - <div className="absolute bottom-2 left-2 text-sm text-muted-foreground"> - Loading... {streamProgress}% + <div className="absolute bottom-2 left-2 text-sm text-muted-foreground"> + Loading... {Math.max(0, Math.min(100, streamProgress))}% </div> @@ <img src={img.url || `data:image/png;base64,${img.b64_json}`} - alt={img.revised_prompt || `Generated image ${idx + 1}`} + alt={img.revised_prompt || `Generated image ${img.index + 1}`} className="w-full h-auto" loading="lazy" />ui/hooks/useImageStream.ts (1)
88-176: Free per-image chunk buffers after completion to avoid unbounded memory growth.
Once an image is completed and assembled, keeping all chunk strings around is unnecessary.if (type === "image_generation.completed") { const completedImage = buildImageFromChunks(index); + // Free chunk buffers for this image after assembly + imageChunksRef.current.delete(index);plugins/semanticcache/utils.go (2)
383-416: Avoid duplicating large image payloads in metadata (responsevsimage_b64).
metadata["response"]already contains the marshaledBifrostResponse(likely including base64). Addingimage_b64duplicates large data and can inflate vectorstore metadata size.
- Consider storing only
image_urls/revised_prompts(and maybeimage_b64_sha256/ byte length) and rely onresponsefor the canonical payload.
655-661: Normalize input + params extraction is sensible; consider including all cache-relevant params in key.This normalization/extraction looks consistent with other request types. Just ensure
getImageCacheKey(...)accounts for any params that materially change output (e.g., response_format, extra params) so the short-circuit hash doesn’t under-key.Also applies to: 960-986
core/internal/testutil/test_retry_framework.go (1)
161-165: Retry config surface is OK; consider reducing type explosion over time.This mirrors other modalities, but it’s trending toward a lot of near-duplicate retry wrappers/configs.
Also applies to: 241-249
core/internal/testutil/account.go (1)
21-68: Align image-generation enablement: model is set, but scenario flags aren’t.OpenAI sets
ImageGenerationModel: "dall-e-2"(Line 524) butTestScenariosdoesn’t setImageGeneration/ImageGenerationStreamto true. Depending on howRunImageGenerationTestgates execution, this can cause confusing behavior (tests run while summary reports “skipped”, or vice versa).
- Either set
Scenarios.ImageGeneration = true(+ stream if supported) inAllProviderConfigs, or make the scenario functions/summary consistently key offImageGenerationModelpresence (and drop the flags).Also applies to: 514-546
core/internal/testutil/image_generation_stream.go (1)
62-119: Consider enhancing stream validation checks.The current validation (lines 62-119) checks for:
- Received data (line 64)
- Completion event (lines 93-95)
- Stream errors (line 66)
However, the validation could be more thorough to match the documented streaming behavior:
- Track and verify that partial base64 chunks (
PartialB64) are received during streaming- Validate that usage data is present on the final completion chunk
- Verify chunk ordering via
ChunkIndex- Ensure
RevisedPromptappears on first chunk if presentConsider adding these checks inside the stream validation loop:
if response.BifrostImageGenerationStreamResponse != nil { receivedData = true imgResp := response.BifrostImageGenerationStreamResponse + // Track partial chunks + if imgResp.PartialB64 != "" { + // Validate partial base64 chunk + } + if imgResp.Type == string(openai.ImageGenerationCompleted) { hasCompleted = true + // Validate final chunk has usage + if imgResp.Usage == nil { + validationErrors = append(validationErrors, "Final chunk missing usage data") + } } }core/images_test.go (1)
452-477: Cache-key helper is collision-prone (no separators) — ensure it mirrors prod logic.
generateTestCacheKeyconcatenates prompt+params without delimiters, so different tuples can collide ("ab"+"c"vs"a"+"bc"). Since this is meant to emulate the semantic cache key, either import/use the real generator (preferred) or add separators and update expectations accordingly.@@ sb.WriteString(req.Input.Prompt) if req.Params != nil { if req.Params.Size != nil { - sb.WriteString(*req.Params.Size) + sb.WriteString("|size=") + sb.WriteString(*req.Params.Size) } if req.Params.Quality != nil { - sb.WriteString(*req.Params.Quality) + sb.WriteString("|quality=") + sb.WriteString(*req.Params.Quality) } if req.Params.Style != nil { - sb.WriteString(*req.Params.Style) + sb.WriteString("|style=") + sb.WriteString(*req.Params.Style) } if req.Params.N != nil { - sb.WriteString(fmt.Sprintf("%d", *req.Params.N)) + sb.WriteString("|n=") + sb.WriteString(fmt.Sprintf("%d", *req.Params.N)) } }core/providers/bedrock/bedrock.go (1)
1296-1311:getModelPath: consider settingdeploymenteven when ARN is absent (if that’s intended).
Right nowdeploymentis only set whenDeployments[model]exists andARN != nil. If callers rely onModelDeploymentfor metrics/routing, you may wantdeployment = inferenceProfileIDeven without ARN (while leavingpathunchanged).core/internal/testutil/validation_presets.go (2)
188-203: Remove stray “add this function” comment; keep presets clean.
305-312:customParamstype assertions may be too strict (oftenfloat64for numbers).
IfcustomParamscan come from decoded JSON,min_imageswill typically befloat64notint, so the override won’t apply. Consider accepting both.core/internal/testutil/image_generation_errors.go (1)
25-153: Avoid emojis + reduce repeatedstrings.ToLowercalls (test output hygiene).
Not a blocker, but the current output can get noisy in CI, and you callstrings.ToLower(errorMsg)repeatedly per assertion.core/internal/testutil/image_generation_load.go (1)
15-101: Make prompts truly unique to avoid collisions under concurrency. Second-level timestamps can collide across goroutines.- Prompt: "A test image for load testing - " + time.Now().Format("20060102150405"), + Prompt: fmt.Sprintf("A test image for load testing - %s - g=%d r=%d", time.Now().Format(time.RFC3339Nano), goroutineID, j),(Will require adding
fmtimport.)framework/streaming/accumulator.go (1)
380-411: Prefer usingputImageStreamChunkinCleanup()for consistent zeroing. Right nowCleanup()puts pooled chunks back without clearing fields.- for _, chunk := range accumulator.ImageStreamChunks { - a.imageStreamChunkPool.Put(chunk) - } + for _, chunk := range accumulator.ImageStreamChunks { + a.putImageStreamChunk(chunk) + }core/providers/openai/images.go (1)
61-96: Consider settingExtraFields.RequestTypeon image responses for consistency. TodayToBifrostImageResponsesets Provider/Latency but leavesRequestTypeempty (unless set elsewhere).ExtraFields: schemas.BifrostResponseExtraFields{ + RequestType: schemas.ImageGenerationRequest, Provider: schemas.OpenAI, Latency: latency.Milliseconds(), },core/bifrost.go (1)
3160-3225: AddExtraFieldscontext to validation errors for consistency/debuggability.
Unlike other public APIs, the early validation errors here (nil req / missing prompt / nil provider response) don’t populateExtraFields(request type/provider/model). This makes logs/metrics harder to correlate.Suggested diff:
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, + }, } }core/providers/azure/images.go (3)
15-47: Structs look reasonable, but consider consistency with OpenAI/Azure response shapes.
One note:AzureImageResponse.Usageis not taggedomitempty, so it’ll marshal asnullif absent; confirm Azure always returns usage for images in your targeted API versions.
81-107: Parameter mapping ignoresExtraParams(intentional?).
schemas.ImageGenerationParametersincludesExtraParams map[string]interface{}but it’s not forwarded. If Azure supports provider-specific image params via extra fields, consider plumbing them through (or explicitly documenting they’re ignored for Azure images).
171-229: Streaming auth/header handling: avoid sending empty Bearer tokens; validate request nil.
- If
ctx.Value(AzureAuthorizationTokenKey)exists but is empty, you’ll sendAuthorization: Bearerand skip api-key; considerauthToken != "".- Also consider a defensive
if request == nil/request.Input == nilearly return for symmetry with non-streaming.
📜 Review details
Configuration used: CodeRabbit UI
Review profile: CHILL
Plan: Pro
📒 Files selected for processing (53)
core/bifrost.go(7 hunks)core/image_generation_tools_test.go(1 hunks)core/images_test.go(1 hunks)core/internal/testutil/account.go(3 hunks)core/internal/testutil/image_generation.go(1 hunks)core/internal/testutil/image_generation_cache.go(1 hunks)core/internal/testutil/image_generation_errors.go(1 hunks)core/internal/testutil/image_generation_fallback.go(1 hunks)core/internal/testutil/image_generation_load.go(1 hunks)core/internal/testutil/image_generation_stream.go(1 hunks)core/internal/testutil/response_validation.go(2 hunks)core/internal/testutil/test_retry_conditions.go(1 hunks)core/internal/testutil/test_retry_framework.go(6 hunks)core/internal/testutil/tests.go(1 hunks)core/internal/testutil/validation_presets.go(2 hunks)core/providers/anthropic/anthropic.go(4 hunks)core/providers/azure/images.go(1 hunks)core/providers/bedrock/bedrock.go(6 hunks)core/providers/cerebras/cerebras.go(1 hunks)core/providers/cohere/cohere.go(3 hunks)core/providers/elevenlabs/elevenlabs.go(3 hunks)core/providers/gemini/gemini.go(4 hunks)core/providers/groq/groq.go(1 hunks)core/providers/mistral/mistral.go(1 hunks)core/providers/ollama/ollama.go(1 hunks)core/providers/openai/images.go(1 hunks)core/providers/openai/openai.go(9 hunks)core/providers/openai/openai_image_test.go(1 hunks)core/providers/openai/openai_test.go(2 hunks)core/providers/openai/types.go(1 hunks)core/providers/openrouter/openrouter.go(1 hunks)core/providers/parasail/parasail.go(1 hunks)core/providers/perplexity/perplexity.go(1 hunks)core/providers/sgl/sgl.go(1 hunks)core/providers/utils/utils.go(3 hunks)core/providers/vertex/vertex.go(1 hunks)core/schemas/bifrost.go(10 hunks)core/schemas/chatcompletions.go(1 hunks)core/schemas/images.go(1 hunks)core/schemas/mux.go(1 hunks)core/schemas/provider.go(1 hunks)core/schemas/responses.go(3 hunks)core/utils.go(1 hunks)framework/streaming/accumulator.go(10 hunks)framework/streaming/images.go(1 hunks)framework/streaming/types.go(5 hunks)plugins/semanticcache/main.go(4 hunks)plugins/semanticcache/search.go(2 hunks)plugins/semanticcache/stream.go(1 hunks)plugins/semanticcache/utils.go(7 hunks)transports/bifrost-http/handlers/inference.go(7 hunks)ui/components/chat/ImageMessage.tsx(1 hunks)ui/hooks/useImageStream.ts(1 hunks)
🚧 Files skipped from review as they are similar to previous changes (19)
- plugins/semanticcache/stream.go
- core/providers/groq/groq.go
- plugins/semanticcache/search.go
- core/providers/openai/openai_test.go
- core/utils.go
- core/schemas/provider.go
- core/schemas/responses.go
- core/providers/openai/types.go
- plugins/semanticcache/main.go
- core/providers/ollama/ollama.go
- core/providers/cerebras/cerebras.go
- core/providers/mistral/mistral.go
- core/internal/testutil/test_retry_conditions.go
- framework/streaming/types.go
- core/internal/testutil/image_generation_cache.go
- core/providers/parasail/parasail.go
- core/schemas/chatcompletions.go
- core/providers/elevenlabs/elevenlabs.go
- core/image_generation_tools_test.go
🧰 Additional context used
📓 Path-based instructions (1)
**
⚙️ CodeRabbit configuration file
always check the stack if there is one for the current PR. do not give localized reviews for the PR, always see all changes in the light of the whole stack of PRs (if there is a stack, if there is no stack you can continue to make localized suggestions/reviews)
Files:
transports/bifrost-http/handlers/inference.gocore/providers/openrouter/openrouter.gocore/internal/testutil/image_generation_fallback.gocore/internal/testutil/validation_presets.gocore/images_test.gocore/internal/testutil/image_generation_stream.goui/hooks/useImageStream.tscore/providers/cohere/cohere.gocore/providers/perplexity/perplexity.gocore/providers/bedrock/bedrock.gocore/schemas/images.gocore/internal/testutil/image_generation.gocore/providers/sgl/sgl.goframework/streaming/images.gocore/providers/vertex/vertex.gocore/providers/gemini/gemini.gocore/internal/testutil/tests.gocore/providers/openai/openai.gocore/internal/testutil/image_generation_errors.gocore/internal/testutil/response_validation.gocore/schemas/mux.goui/components/chat/ImageMessage.tsxcore/providers/utils/utils.gocore/internal/testutil/account.gocore/providers/anthropic/anthropic.gocore/providers/openai/openai_image_test.gocore/internal/testutil/image_generation_load.gocore/providers/azure/images.goplugins/semanticcache/utils.gocore/schemas/bifrost.gocore/internal/testutil/test_retry_framework.gocore/providers/openai/images.gocore/bifrost.goframework/streaming/accumulator.go
🧠 Learnings (3)
📚 Learning: 2025-12-09T17:07:42.007Z
Learnt from: qwerty-dvorak
Repo: maximhq/bifrost PR: 1006
File: core/schemas/account.go:9-18
Timestamp: 2025-12-09T17:07:42.007Z
Learning: In core/schemas/account.go, the HuggingFaceKeyConfig field within the Key struct is currently unused and reserved for future Hugging Face inference endpoint deployments. Do not flag this field as missing from OpenAPI documentation or require its presence in the API spec until the feature is actively implemented and used. When the feature is added, update the OpenAPI docs accordingly; otherwise, treat this field as non-breaking and not part of the current API surface.
Applied to files:
transports/bifrost-http/handlers/inference.gocore/providers/openrouter/openrouter.gocore/internal/testutil/image_generation_fallback.gocore/internal/testutil/validation_presets.gocore/images_test.gocore/internal/testutil/image_generation_stream.gocore/providers/cohere/cohere.gocore/providers/perplexity/perplexity.gocore/providers/bedrock/bedrock.gocore/schemas/images.gocore/internal/testutil/image_generation.gocore/providers/sgl/sgl.goframework/streaming/images.gocore/providers/vertex/vertex.gocore/providers/gemini/gemini.gocore/internal/testutil/tests.gocore/providers/openai/openai.gocore/internal/testutil/image_generation_errors.gocore/internal/testutil/response_validation.gocore/schemas/mux.gocore/providers/utils/utils.gocore/internal/testutil/account.gocore/providers/anthropic/anthropic.gocore/providers/openai/openai_image_test.gocore/internal/testutil/image_generation_load.gocore/providers/azure/images.goplugins/semanticcache/utils.gocore/schemas/bifrost.gocore/internal/testutil/test_retry_framework.gocore/providers/openai/images.gocore/bifrost.goframework/streaming/accumulator.go
📚 Learning: 2025-12-12T08:25:02.629Z
Learnt from: Pratham-Mishra04
Repo: maximhq/bifrost PR: 1000
File: transports/bifrost-http/integrations/router.go:709-712
Timestamp: 2025-12-12T08:25:02.629Z
Learning: In transports/bifrost-http/**/*.go, update streaming response handling to align with OpenAI Responses API: use typed SSE events such as response.created, response.output_text.delta, response.done, etc., and do not rely on the legacy data: [DONE] termination marker. Note that data: [DONE] is only used by the older Chat Completions and Text Completions streaming APIs. Ensure parsers, writers, and tests distinguish SSE events from the [DONE] sentinel and handle each event type accordingly for correct stream termination and progress updates.
Applied to files:
transports/bifrost-http/handlers/inference.go
📚 Learning: 2025-12-11T11:58:25.307Z
Learnt from: Pratham-Mishra04
Repo: maximhq/bifrost PR: 1000
File: core/providers/openai/responses.go:42-84
Timestamp: 2025-12-11T11:58:25.307Z
Learning: In core/providers/openai/responses.go (and related OpenAI response handling), document and enforce the API format constraint: if ResponsesReasoning != nil and the response contains content blocks, all content blocks should be treated as reasoning blocks by default. Implement type guards or parsing logic accordingly, and add unit tests to verify that when ResponsesReasoning is non-nil, content blocks are labeled as reasoning blocks. Include clear comments in the code explaining the rationale and ensure downstream consumers rely on this behavior.
Applied to files:
core/providers/openai/openai.gocore/providers/openai/openai_image_test.gocore/providers/openai/images.go
🧬 Code graph analysis (26)
core/providers/openrouter/openrouter.go (4)
core/schemas/images.go (2)
BifrostImageGenerationRequest(4-11)BifrostImageGenerationResponse(33-40)core/schemas/bifrost.go (4)
BifrostError(383-392)ImageGenerationRequest(100-100)BifrostStream(345-353)ImageGenerationStreamRequest(101-101)core/providers/utils/utils.go (1)
NewUnsupportedOperationError(458-469)core/schemas/provider.go (1)
PostHookRunner(279-279)
core/internal/testutil/image_generation_fallback.go (3)
core/internal/testutil/account.go (1)
ComprehensiveTestConfig(49-68)core/schemas/images.go (3)
BifrostImageGenerationRequest(4-11)ImageGenerationInput(18-20)ImageGenerationParameters(22-30)core/schemas/bifrost.go (1)
ImageGenerationRequest(100-100)
core/internal/testutil/validation_presets.go (1)
core/internal/testutil/response_validation.go (1)
ResponseExpectations(18-43)
core/images_test.go (2)
core/schemas/images.go (7)
BifrostImageGenerationRequest(4-11)ImageGenerationInput(18-20)ImageGenerationParameters(22-30)BifrostImageGenerationResponse(33-40)ImageData(42-47)ImageUsage(49-52)BifrostImageGenerationStreamResponse(55-65)core/schemas/bifrost.go (1)
OpenAI(35-35)
core/internal/testutil/image_generation_stream.go (5)
core/internal/testutil/account.go (1)
ComprehensiveTestConfig(49-68)core/internal/testutil/test_retry_framework.go (3)
TestRetryContext(174-179)WithImageGenerationStreamRetry(2707-2845)ImageGenerationStreamValidationResult(2697-2703)core/schemas/images.go (4)
BifrostImageGenerationRequest(4-11)ImageGenerationInput(18-20)ImageGenerationParameters(22-30)BifrostImageGenerationStreamResponse(55-65)core/schemas/bifrost.go (2)
BifrostStream(345-353)ImageGenerationStreamRequest(101-101)core/providers/openai/images.go (1)
ImageGenerationCompleted(24-24)
ui/hooks/useImageStream.ts (2)
ui/lib/store/apis/baseApi.ts (1)
getTokenFromStorage(10-25)ui/lib/utils/port.ts (1)
getEndpointUrl(117-127)
core/providers/cohere/cohere.go (2)
core/providers/utils/utils.go (3)
ProcessAndSendResponse(612-643)GetBifrostResponseForStreamResponse(855-887)NewUnsupportedOperationError(458-469)core/schemas/bifrost.go (2)
BifrostError(383-392)BifrostStream(345-353)
core/providers/perplexity/perplexity.go (4)
core/schemas/images.go (2)
BifrostImageGenerationRequest(4-11)BifrostImageGenerationResponse(33-40)core/schemas/bifrost.go (3)
BifrostError(383-392)ImageGenerationRequest(100-100)BifrostStream(345-353)core/providers/utils/utils.go (1)
NewUnsupportedOperationError(458-469)core/schemas/provider.go (1)
PostHookRunner(279-279)
core/providers/bedrock/bedrock.go (4)
core/providers/utils/utils.go (1)
ProcessAndSendResponse(612-643)core/schemas/images.go (2)
BifrostImageGenerationRequest(4-11)BifrostImageGenerationResponse(33-40)core/schemas/bifrost.go (4)
BifrostError(383-392)ImageGenerationRequest(100-100)BifrostStream(345-353)ImageGenerationStreamRequest(101-101)core/schemas/provider.go (1)
PostHookRunner(279-279)
core/schemas/images.go (2)
core/schemas/provider.go (1)
Provider(282-315)core/schemas/bifrost.go (4)
ModelProvider(32-32)Fallback(136-139)BifrostResponseExtraFields(310-320)BifrostError(383-392)
core/providers/sgl/sgl.go (5)
core/schemas/account.go (1)
Key(8-17)core/schemas/images.go (2)
BifrostImageGenerationRequest(4-11)BifrostImageGenerationResponse(33-40)core/schemas/bifrost.go (3)
BifrostError(383-392)ImageGenerationRequest(100-100)ImageGenerationStreamRequest(101-101)core/providers/utils/utils.go (1)
NewUnsupportedOperationError(458-469)core/schemas/provider.go (1)
PostHookRunner(279-279)
framework/streaming/images.go (3)
framework/streaming/types.go (1)
ImageStreamChunk(104-114)core/schemas/images.go (3)
BifrostImageGenerationResponse(33-40)ImageData(42-47)BifrostImageGenerationStreamResponse(55-65)core/utils.go (2)
GetResponseFields(216-223)IsFinalChunk(198-213)
core/providers/vertex/vertex.go (4)
core/schemas/images.go (2)
BifrostImageGenerationRequest(4-11)BifrostImageGenerationResponse(33-40)core/schemas/bifrost.go (2)
BifrostError(383-392)ImageGenerationRequest(100-100)core/providers/utils/utils.go (1)
NewUnsupportedOperationError(458-469)core/schemas/provider.go (1)
PostHookRunner(279-279)
core/providers/gemini/gemini.go (2)
core/providers/utils/utils.go (3)
ProcessAndSendResponse(612-643)GetBifrostResponseForStreamResponse(855-887)NewUnsupportedOperationError(458-469)core/schemas/bifrost.go (2)
BifrostError(383-392)BifrostStream(345-353)
core/internal/testutil/tests.go (6)
core/internal/testutil/image_generation.go (1)
RunImageGenerationTest(13-126)core/internal/testutil/image_generation_stream.go (1)
RunImageGenerationStreamTest(17-133)core/internal/testutil/image_generation_fallback.go (1)
RunImageGenerationFallbackTest(14-80)core/internal/testutil/image_generation_cache.go (1)
RunImageGenerationCacheTest(14-140)core/internal/testutil/image_generation_errors.go (1)
RunImageGenerationErrorTest(14-154)core/internal/testutil/image_generation_load.go (3)
RunImageGenerationLoadTest(16-101)RunImageGenerationStreamLoadTest(104-211)RunImageGenerationCacheLoadTest(214-307)
core/providers/openai/openai.go (3)
core/providers/utils/utils.go (2)
ProcessAndSendResponse(612-643)GetBifrostResponseForStreamResponse(855-887)core/schemas/bifrost.go (3)
BifrostError(383-392)ImageGenerationRequest(100-100)BifrostStream(345-353)core/schemas/provider.go (2)
Provider(282-315)PostHookRunner(279-279)
core/internal/testutil/image_generation_errors.go (2)
core/internal/testutil/account.go (1)
ComprehensiveTestConfig(49-68)core/schemas/images.go (3)
BifrostImageGenerationRequest(4-11)ImageGenerationInput(18-20)ImageGenerationParameters(22-30)
core/internal/testutil/response_validation.go (3)
core/schemas/images.go (1)
BifrostImageGenerationResponse(33-40)core/internal/testutil/error_parser.go (3)
ParseBifrostError(47-104)FormatErrorConcise(353-356)LogError(359-366)plugins/mocker/main.go (1)
Latency(125-129)
core/schemas/mux.go (3)
core/providers/gemini/types.go (1)
Type(780-780)core/schemas/chatcompletions.go (1)
ChatToolTypeImageGeneration(241-241)core/schemas/responses.go (1)
ResponsesToolImageGeneration(1336-1347)
core/providers/utils/utils.go (1)
core/schemas/images.go (1)
BifrostImageGenerationStreamResponse(55-65)
core/internal/testutil/account.go (1)
core/schemas/bifrost.go (1)
Fallback(136-139)
core/providers/azure/images.go (7)
core/providers/azure/azure.go (2)
AzureProvider(22-28)AzureAuthorizationTokenKey(19-19)core/schemas/images.go (6)
BifrostImageGenerationRequest(4-11)BifrostImageGenerationResponse(33-40)ImageGenerationParameters(22-30)ImageData(42-47)ImageUsage(49-52)BifrostImageGenerationStreamResponse(55-65)core/schemas/bifrost.go (5)
BifrostError(383-392)ImageGenerationRequest(100-100)BifrostResponseExtraFields(310-320)Azure(36-36)BifrostStream(345-353)core/providers/openai/images.go (1)
ToBifrostImageResponse(62-96)core/schemas/provider.go (2)
Provider(282-315)PostHookRunner(279-279)core/utils.go (1)
Ptr(51-53)core/providers/azure/types.go (1)
AzureAPIVersionDefault(4-4)
core/schemas/bifrost.go (9)
core/schemas/models.go (2)
BifrostListModelsRequest(23-34)Model(109-129)core/schemas/textcompletions.go (2)
BifrostTextCompletionRequest(10-17)BifrostTextCompletionResponse(69-77)core/schemas/chatcompletions.go (2)
BifrostChatRequest(12-19)BifrostChatResponse(26-41)core/schemas/responses.go (3)
BifrostResponsesRequest(32-39)BifrostResponsesResponse(45-85)BifrostResponsesStreamResponse(1452-1491)core/schemas/embedding.go (2)
BifrostEmbeddingRequest(9-16)BifrostEmbeddingResponse(22-28)core/schemas/speech.go (2)
BifrostSpeechRequest(9-16)BifrostSpeechResponse(22-29)core/schemas/transcriptions.go (2)
BifrostTranscriptionRequest(3-10)BifrostTranscriptionResponse(16-26)core/schemas/images.go (2)
BifrostImageGenerationRequest(4-11)BifrostImageGenerationResponse(33-40)core/schemas/provider.go (1)
Provider(282-315)
core/internal/testutil/test_retry_framework.go (5)
core/schemas/images.go (1)
BifrostImageGenerationResponse(33-40)core/internal/testutil/test_retry_conditions.go (2)
EmptyImageGenerationCondition(851-851)GenericResponseCondition(370-370)core/internal/testutil/response_validation.go (2)
ResponseExpectations(18-43)ValidateImageGenerationResponse(241-275)core/internal/testutil/utils.go (1)
GetErrorMessage(606-639)core/schemas/bifrost.go (1)
BifrostStream(345-353)
core/providers/openai/images.go (4)
core/providers/openai/types.go (3)
OpenAIImageGenerationRequest(292-302)OpenAIImageGenerationResponse(305-317)OpenAIImageStreamResponse(331-336)core/schemas/bifrost.go (8)
BifrostResponseExtraFields(310-320)OpenAI(35-35)BifrostError(383-392)ImageGenerationRequest(100-100)ModelProvider(32-32)ErrorField(401-408)BifrostErrorExtraFields(450-454)RequestType(83-83)core/providers/utils/utils.go (1)
SetExtraHeaders(179-209)core/providers/openai/errors.go (1)
ParseOpenAIError(10-42)
framework/streaming/accumulator.go (2)
framework/streaming/types.go (4)
ImageStreamChunk(104-114)StreamAccumulator(117-129)ChatStreamChunk(78-88)ResponsesStreamChunk(91-101)core/schemas/bifrost.go (1)
ImageGenerationStreamRequest(101-101)
🔇 Additional comments (33)
core/providers/vertex/vertex.go (1)
1409-1417: ImageGeneration/ImageGenerationStream stubs are consistent with “unsupported provider” pattern.core/providers/perplexity/perplexity.go (1)
255-263: Perplexity ImageGeneration/ImageGenerationStream stubs look correct (and request types match).core/providers/gemini/gemini.go (2)
1424-1432: Gemini ImageGeneration/ImageGenerationStream stubs look correct (and stream uses ImageGenerationStreamRequest).
1115-1142: Unable to verify: The specified file and function references could not be located in the repository.The review comment references
core/providers/gemini/gemini.gowith specific line ranges (1115-1142, 1385-1418) and claims thatGetBifrostResponseForStreamResponsecalls have inconsistent arity (5 vs 6 arguments). However, these files and function references could not be found during verification. Please confirm:
- The correct file path for the Gemini provider implementation
- The actual function name and its current signature
- Whether these changes have been committed to the branch
Once confirmed, the arity mismatch can be properly verified.
ui/hooks/useImageStream.ts (1)
217-253: Backend sends single-line JSON only—no multi-line data blocks to handle.The backend uses
sonic.Marshal()for compact JSON serialization (line 1073, inference.go) and writes responses asdata: {json}\n\nin a single line (lines 1096, 1102). Image generation streams do not receiveevent:lines at all (onlyBifrostResponsesStreamResponsesetsincludeEventType = true). The frontend code correctly handles the actual format the backend sends; multi-line SSE payloads are not a concern here.plugins/semanticcache/utils.go (1)
175-179: Embedding text extraction for image prompts looks good; keep metadata+prompt consistent.Line 337-344’s prompt validation + normalization is a solid safeguard for cache/embedding consistency.
Also applies to: 336-345
core/providers/anthropic/anthropic.go (2)
1027-1035: ImageGeneration and ImageGenerationStream method signatures correctly match the Provider interface.Both method signatures in AnthropicProvider implementation are properly aligned with the Provider interface definitions. The ImageGeneration and ImageGenerationStream methods have the correct parameter order and return types, and returning
NewUnsupportedOperationErrorfor unimplemented operations is the appropriate pattern.
632-633: AllGetBifrostResponseForStreamResponse(...)parameter positions are correct and semantically appropriate.The function signature has 6 parameters:
textCompletionResponse,chatResponse,responsesStreamResponse,speechStreamResponse,transcriptionStreamResponse,imageGenerationStreamResponse. All callsites correctly map their response objects to the appropriate parameter position—chat responses to position 2 (lines 632, 662) and responses stream responses to position 3 (lines 981, 984). The newimageGenerationStreamResponseparameter is properly set tonilacross all existing callsites.core/internal/testutil/test_retry_framework.go (1)
926-940: No nil-safety bug exists here. BothEmptyImageGenerationConditionandGenericResponseConditionexplicitly handle nil responses: the former returnstrue, "response is nil"(line 856), and the latter returnsfalse, ""(line 372-374). ThecheckStreamRetryConditionsfunction intentionally passesnilfor streaming operations, as documented in its comment (line 718): "Pass nil response since streaming doesn't have a single response." This design is safe and deliberate.Likely an incorrect or invalid review comment.
core/providers/utils/utils.go (1)
632-632: LGTM! Image generation streaming support correctly integrated.The changes consistently extend the streaming infrastructure to support image generation:
- Line 632 properly propagates
ImageGenerationStreamResponseinProcessAndSendResponse- Line 861 extends
GetBifrostResponseForStreamResponsesignature with the new parameter- Lines 882-884 add the corresponding case handling
This follows the established pattern used for other streaming response types (speech, transcription, etc.).
Also applies to: 861-861, 882-884
core/providers/openrouter/openrouter.go (1)
285-293: LGTM! Unsupported operations correctly declared.Both
ImageGenerationandImageGenerationStreammethods appropriately return unsupported operation errors usingNewUnsupportedOperationError, consistent with other unsupported operations in the OpenRouter provider (lines 262, 267, 272, 277, 282).core/providers/cohere/cohere.go (2)
486-486: LGTM! Streaming calls correctly updated with new parameter.Lines 486, 489, 728, and 731 correctly add the sixth parameter (nil for
imageGenerationStreamResponse) toGetBifrostResponseForStreamResponsecalls, aligning with the extended signature incore/providers/utils/utils.go.Also applies to: 489-489, 728-728, 731-731
821-829: LGTM! Unsupported image generation operations correctly declared.Both
ImageGenerationandImageGenerationStreammethods appropriately return unsupported operation errors, consistent with other unsupported operations in the Cohere provider (lines 802-819).core/internal/testutil/image_generation.go (1)
24-79: LGTM! Retry framework correctly integrated.The test properly uses the retry framework as requested in past review comments:
- Line 24:
GetTestRetryConfigForScenarioto get retry configuration- Lines 39-46: Sets up
ImageGenerationRetryConfigwith retry handlers- Line 79:
WithImageGenerationRetrywrapper wraps the operationThis addresses the previous concern about flaky tests in CI.
core/schemas/images.go (1)
3-65: LGTM! Image generation schemas well-structured.The schema definitions are well-designed and consistent:
BifrostImageGenerationRequestfollows the pattern of other Bifrost request types (includes Provider, Model, Input, Params, Fallbacks, RawRequestBody)GetRawRequestBody()implements theRequestBodyGetterinterface correctly (lines 14-16)- Parameter struct (lines 22-30) includes all standard OpenAI image generation parameters (n, size, quality, style, response_format, user)
- Response structs (lines 33-52) align with OpenAI's response format
- Streaming response (lines 55-65) properly supports chunked delivery with type, index, chunk_index, partial base64, and usage fields
- All JSON tags follow Go conventions
core/internal/testutil/image_generation_fallback.go (1)
14-80: Add retry framework wrapper to fallback test.Based on past review comments, all new tests should use the retry framework to avoid flakiness in CI. This test currently executes the request directly (line 46) without retry logic.
Wrap the request execution with the retry framework:
t.Run("ImageGenerationFallback", func(t *testing.T) { if os.Getenv("SKIP_PARALLEL_TESTS") != "true" { t.Parallel() } + retryConfig := GetTestRetryConfigForScenario("ImageGenerationFallback", testConfig) + retryContext := TestRetryContext{ + ScenarioName: "ImageGenerationFallback", + ExpectedBehavior: map[string]interface{}{ + "fallback_enabled": true, + }, + TestMetadata: map[string]interface{}{ + "provider": testConfig.Provider, + "model": testConfig.ImageGenerationModel, + }, + } + + imageGenerationRetryConfig := ImageGenerationRetryConfig{ + MaxAttempts: retryConfig.MaxAttempts, + BaseDelay: retryConfig.BaseDelay, + MaxDelay: retryConfig.MaxDelay, + Conditions: []ImageGenerationRetryCondition{}, + OnRetry: retryConfig.OnRetry, + OnFinalFail: retryConfig.OnFinalFail, + } + // Create request with primary provider that should fail, triggering fallback - // Note: This test assumes the primary provider will fail (e.g., invalid key) - // In practice, you might need to configure a failing provider for this test request := &schemas.BifrostImageGenerationRequest{ Provider: testConfig.Provider, Model: testConfig.ImageGenerationModel, Input: &schemas.ImageGenerationInput{ Prompt: "A beautiful mountain landscape at dawn", }, Params: &schemas.ImageGenerationParameters{ Size: bifrost.Ptr("1024x1024"), ResponseFormat: bifrost.Ptr("b64_json"), }, Fallbacks: testConfig.ImageGenerationFallbacks, } - response, bifrostErr := client.ImageGenerationRequest(ctx, request) + response, bifrostErr := WithImageGenerationRetry(t, imageGenerationRetryConfig, retryContext, GetExpectationsForScenario("ImageGenerationFallback", testConfig, map[string]interface{}{}), "ImageGenerationFallback", func() (*schemas.BifrostImageGenerationResponse, *schemas.BifrostError) { + return client.ImageGenerationRequest(ctx, request) + })Likely an incorrect or invalid review comment.
core/providers/bedrock/bedrock.go (2)
667-668: Updated streaming wrapper arity looks consistent with newGetBifrostResponseForStreamResponsesignature.Also applies to: 894-909, 1066-1067, 1160-1161
1286-1294: Bedrock image generation “unsupported” stubs are fine for interface parity.core/providers/sgl/sgl.go (1)
223-231: SGL “unsupported image generation” stubs look good for provider-surface consistency.framework/streaming/images.go (1)
164-172: Bug:FinishReasonassignment is incorrectly gated bySemanticCacheDebug != nil.
Finish reason should be set independently of cache debug presence.@@ - if len(acc.ImageStreamChunks) > 0 { - lastChunk := acc.ImageStreamChunks[len(acc.ImageStreamChunks)-1] - if lastChunk.SemanticCacheDebug != nil { - data.CacheDebug = lastChunk.SemanticCacheDebug - } - data.FinishReason = lastChunk.FinishReason - } + if len(acc.ImageStreamChunks) > 0 { + lastChunk := acc.ImageStreamChunks[len(acc.ImageStreamChunks)-1] + if lastChunk.SemanticCacheDebug != nil { + data.CacheDebug = lastChunk.SemanticCacheDebug + } + data.FinishReason = lastChunk.FinishReason + }Likely an incorrect or invalid review comment.
core/internal/testutil/response_validation.go (2)
240-275: Nice parity with other validators (error handling + metrics + logging).
1064-1109: Fix potential nil-map panic whenProviderSpecificis unset.expectations.ProviderSpecific["..."]will panic ifProviderSpecific == nil.func validateImageGenerationFields(t *testing.T, response *schemas.BifrostImageGenerationResponse, expectations ResponseExpectations, result *ValidationResult) { @@ // Check minimum number of images if specified - if minImages, ok := expectations.ProviderSpecific["min_images"].(int); ok { - actualCount := len(response.Data) - if actualCount < minImages { - result.Passed = false - result.Errors = append(result.Errors, - fmt.Sprintf("Too few images: got %d, expected at least %d", actualCount, minImages)) - } else { - result.MetricsCollected["image_count"] = actualCount - } - } + if expectations.ProviderSpecific != nil { + if minImages, ok := expectations.ProviderSpecific["min_images"].(int); ok { + actualCount := len(response.Data) + if actualCount < minImages { + result.Passed = false + result.Errors = append(result.Errors, + fmt.Sprintf("Too few images: got %d, expected at least %d", actualCount, minImages)) + } else { + result.MetricsCollected["image_count"] = actualCount + } + } + } @@ // Validate image size if specified - if expectedSize, ok := expectations.ProviderSpecific["expected_size"].(string); ok { - result.MetricsCollected["expected_size"] = expectedSize - // Note: Actual size validation would require downloading/decoding images - } + if expectations.ProviderSpecific != nil { + if expectedSize, ok := expectations.ProviderSpecific["expected_size"].(string); ok { + result.MetricsCollected["expected_size"] = expectedSize + // Note: Actual size validation would require downloading/decoding images + } + }Likely an incorrect or invalid review comment.
framework/streaming/accumulator.go (1)
468-515: Confirm Go version compatibility forfor range 1000. This requires Go 1.22+. If the repo isn’t on 1.22+, use a traditional loop.- for range 1000 { + for i := 0; i < 1000; i++ { a.chatStreamChunkPool.Put(&ChatStreamChunk{}) a.responsesStreamChunkPool.Put(&ResponsesStreamChunk{}) a.audioStreamChunkPool.Put(&AudioStreamChunk{}) a.transcriptionStreamChunkPool.Put(&TranscriptionStreamChunk{}) a.imageStreamChunkPool.Put(&ImageStreamChunk{}) }core/schemas/bifrost.go (1)
141-307: Schema wiring for image generation request/response/stream looks consistent with existing modalities.Verification ask: please ensure the image generation response path consistently populates
BifrostResponseExtraFields.RequestType(andModelRequestedif expected) the same way other modalities do.Also applies to: 343-375
core/providers/openai/images.go (1)
20-31: Go version check:maps.Copyrequires Go 1.21+. Please confirm repo toolchain supports this.core/bifrost.go (5)
17-17: Confirmgithub.com/bytedance/sonicbehavior matches expectations (number decoding, map semantics).
You rely onsonic.Unmarshalintomap[string]interface{}and then do type switches (float64/int/int64). Please confirm sonic’s decoding rules here match your assumptions (especially for integers and large values).
2006-2011: Fallback propagation forImageGenerationRequestlooks correct.
Cloning the request and overriding provider/model mirrors the other request types.
2748-2754: Provider routing forImageGenerationRequestlooks fine.
Straight-through call toprovider.ImageGenerationand wiring intoresponse.ImageGenerationResponsematches existing patterns.
2784-2786: Streaming routing forImageGenerationStreamRequestlooks fine.
Routes toprovider.ImageGenerationStreamand passes the post-hook runner consistently with other stream request types.
2962-2963: Good: pooled request reset clearsImageGenerationRequest.
Prevents cross-request leakage when reusingschemas.BifrostRequestfrom the pool.core/providers/azure/images.go (3)
52-79: Non-streamingImageGenerationflow is clean and sets useful ExtraFields.
Validation → transform → DoRequest → map response mirrors other providers well.
109-145: Response mapping: double-check token field semantics.
You mapPromptTokens = InputTokensandTotalTokens = TotalTokens. If Azure’s image usage differs from OpenAI’s (or doesn’t provide these fields consistently), ensure this doesn’t mislead callers.
146-169: DoRequest path looks fine; ensurecompleteRequestmatches Azure images requirements.
Given Azure images typically use deployment-based routing, confirmprovider.completeRequest(..., "images/generations", ..., azureRequest.Model, ...)selects the correct deployment/model for images (and doesn’t require a different path or payload shape).
|
Hey @Radheshg04, PR looks solid but few checkboxes before we merge this
|
There was a problem hiding this comment.
Actionable comments posted: 9
🧹 Nitpick comments (12)
core/providers/azure/types.go (1)
49-69: Consider extracting anonymous struct for better reusability.The
Datafield uses an anonymous struct. If this data structure is used elsewhere (e.g., in tests or response handling), extracting it as a named type improves maintainability.+// AzureImageData represents a single image result from Azure. +type AzureImageData struct { + URL string `json:"url,omitempty"` + B64JSON string `json:"b64_json,omitempty"` + RevisedPrompt string `json:"revised_prompt,omitempty"` +} + // AzureImageResponse is the struct for Image Generation responses by Azure. type AzureImageResponse struct { Created int64 `json:"created"` - Data []struct { - URL string `json:"url,omitempty"` - B64JSON string `json:"b64_json,omitempty"` - RevisedPrompt string `json:"revised_prompt,omitempty"` - } `json:"data"` + Data []AzureImageData `json:"data"` Usage *AzureImageGenerationUsage `json:"usage"` }ui/hooks/useImageStream.ts (2)
173-179:optionsin dependency array may cause unnecessary re-renders.If
optionsobject is recreated on each render (common pattern),processChunkwill be recreated, potentially causing stale closure issues or unnecessary work.Consider memoizing callbacks from options or using refs:
+ const onCompleteRef = useRef(options.onComplete); + const onErrorRef = useRef(options.onError); + + useEffect(() => { + onCompleteRef.current = options.onComplete; + onErrorRef.current = options.onError; + }, [options.onComplete, options.onError]); // In processChunk: - options.onError?.(errorMsg); + onErrorRef.current?.(errorMsg);
251-256: Silent failure on JSON parse errors may hide issues.Parse errors are logged but not surfaced to the user. Consider whether malformed chunks should trigger error state.
try { const chunk: ImageStreamChunk = JSON.parse(data); processChunk(chunk); } catch (parseError) { console.error("Failed to parse SSE chunk:", parseError); + // Optionally: count errors and fail if too many }framework/streaming/images.go (1)
14-92: Guard against nil*ImageStreamChunkentries before sorting.
sort.Slicedereferenceschunks[i]/chunks[j]; a nil entry would panic. If nils are impossible by construction, a short comment/assert would still help.func (a *Accumulator) buildCompleteImageFromImageStreamChunks(chunks []*ImageStreamChunk) *schemas.BifrostImageGenerationResponse { + // Defensive: drop nil entries to avoid panics in sort/comparators. + filtered := chunks[:0] + for _, c := range chunks { + if c != nil { + filtered = append(filtered, c) + } + } + chunks = filtered + // Sort chunks by ImageIndex, then ChunkIndex sort.Slice(chunks, func(i, j int) bool { if chunks[i].ImageIndex != chunks[j].ImageIndex { return chunks[i].ImageIndex < chunks[j].ImageIndex } return chunks[i].ChunkIndex < chunks[j].ChunkIndex })core/providers/azure/azure.go (1)
846-908: Azure ImageGenerationStream: consider aligning deployment lookup helper.
This usesdeployment := key.AzureKeyConfig.Deployments[request.Model]while non-streaming usesgetModelDeployment. Not wrong, but it’s easy to drift; consider usinggetModelDeploymentin both for consistent error messages and future behavior.core/providers/openai/openai.go (1)
2156-2231: Non-streaming ImageGeneration: considerCheckAndDecodeBodyfor parity with other ops.
Most other operations decode body (gzip, etc.) viaproviderUtils.CheckAndDecodeBody; here you unmarshalhttpResp.Body()directly. This can be fine, but it’s an inconsistency that can bite if compression is enabled.core/schemas/bifrost.go (2)
185-202: Set helpers: ok, but consider a defensive “else” for misconfigured requests.*
If multiple request pointers are ever set, theseswitch { case ... != nil }helpers will silently update only the first match. Optional: add a debug log or panic in dev builds to catch invariant violations.Also applies to: 204-221, 223-240, 242-259
263-305: GetExtraFields(): good extension, but returning a new empty struct can hide bugs.
Returning&BifrostResponseExtraFields{}makes it hard to detect “no response payload set” issues; consider returningnil(and handle nil at callsites) or logging when no known response is present.core/providers/openai/images.go (1)
24-50: Request mapping is fine; consider passing through provider-specific extra params if intended.
mapImageParamsignoresschemas.ImageGenerationParameters.ExtraParams(if OpenAI supports any additional image params in your gateway). If that’s intentional, ok; otherwise you may need an allowlist mapping to OpenAI request fields.framework/streaming/accumulator.go (1)
382-413: UseputImageStreamChunk(and other put helpers) in Cleanup() for consistency.*
Cleanup()currently returns chunks viapool.Putdirectly; ifCleanup()is ever called and the accumulator reused, stale fields could bleed. Low risk ifCleanup()is strictly shutdown-only, but easy to make consistent.- for _, chunk := range accumulator.ImageStreamChunks { - a.imageStreamChunkPool.Put(chunk) - } + for _, chunk := range accumulator.ImageStreamChunks { + a.putImageStreamChunk(chunk) + }core/bifrost.go (2)
870-935: Public ImageGeneration APIs: wire-up looks correct; add ExtraFields on validation errors for consistency.*
Right now nil/empty-prompt errors don’t includeBifrostErrorExtraFields{RequestType, Provider, ModelRequested}like other public APIs do. This makes client-side debugging and log correlation harder.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, + }, } }
870-911: Nit: grammar in doc comments (“a image” -> “an image”).
📜 Review details
Configuration used: CodeRabbit UI
Review profile: CHILL
Plan: Pro
📒 Files selected for processing (19)
core/bifrost.go(5 hunks)core/images_test.go(1 hunks)core/internal/testutil/image_generation.go(1 hunks)core/internal/testutil/image_generation_load.go(1 hunks)core/providers/azure/azure.go(2 hunks)core/providers/azure/images.go(1 hunks)core/providers/azure/types.go(1 hunks)core/providers/gemini/gemini.go(8 hunks)core/providers/openai/images.go(1 hunks)core/providers/openai/openai.go(10 hunks)core/providers/openai/openai_image_test.go(1 hunks)core/providers/openai/types.go(1 hunks)core/providers/vertex/vertex.go(1 hunks)core/schemas/bifrost.go(10 hunks)framework/streaming/accumulator.go(10 hunks)framework/streaming/images.go(1 hunks)plugins/semanticcache/utils.go(7 hunks)transports/bifrost-http/handlers/inference.go(7 hunks)ui/hooks/useImageStream.ts(1 hunks)
🚧 Files skipped from review as they are similar to previous changes (4)
- core/providers/openai/openai_image_test.go
- core/providers/openai/types.go
- transports/bifrost-http/handlers/inference.go
- core/providers/vertex/vertex.go
🧰 Additional context used
📓 Path-based instructions (1)
**
⚙️ CodeRabbit configuration file
always check the stack if there is one for the current PR. do not give localized reviews for the PR, always see all changes in the light of the whole stack of PRs (if there is a stack, if there is no stack you can continue to make localized suggestions/reviews)
Files:
core/providers/gemini/gemini.goframework/streaming/images.goframework/streaming/accumulator.gocore/providers/azure/images.gocore/internal/testutil/image_generation.gocore/schemas/bifrost.gocore/internal/testutil/image_generation_load.gocore/providers/azure/azure.goplugins/semanticcache/utils.gocore/bifrost.gocore/images_test.gocore/providers/openai/images.gocore/providers/azure/types.goui/hooks/useImageStream.tscore/providers/openai/openai.go
🧠 Learnings (2)
📚 Learning: 2025-12-09T17:07:42.007Z
Learnt from: qwerty-dvorak
Repo: maximhq/bifrost PR: 1006
File: core/schemas/account.go:9-18
Timestamp: 2025-12-09T17:07:42.007Z
Learning: In core/schemas/account.go, the HuggingFaceKeyConfig field within the Key struct is currently unused and reserved for future Hugging Face inference endpoint deployments. Do not flag this field as missing from OpenAPI documentation or require its presence in the API spec until the feature is actively implemented and used. When the feature is added, update the OpenAPI docs accordingly; otherwise, treat this field as non-breaking and not part of the current API surface.
Applied to files:
core/providers/gemini/gemini.goframework/streaming/images.goframework/streaming/accumulator.gocore/providers/azure/images.gocore/internal/testutil/image_generation.gocore/schemas/bifrost.gocore/internal/testutil/image_generation_load.gocore/providers/azure/azure.goplugins/semanticcache/utils.gocore/bifrost.gocore/images_test.gocore/providers/openai/images.gocore/providers/azure/types.gocore/providers/openai/openai.go
📚 Learning: 2025-12-11T11:58:25.307Z
Learnt from: Pratham-Mishra04
Repo: maximhq/bifrost PR: 1000
File: core/providers/openai/responses.go:42-84
Timestamp: 2025-12-11T11:58:25.307Z
Learning: In core/providers/openai/responses.go (and related OpenAI response handling), document and enforce the API format constraint: if ResponsesReasoning != nil and the response contains content blocks, all content blocks should be treated as reasoning blocks by default. Implement type guards or parsing logic accordingly, and add unit tests to verify that when ResponsesReasoning is non-nil, content blocks are labeled as reasoning blocks. Include clear comments in the code explaining the rationale and ensure downstream consumers rely on this behavior.
Applied to files:
core/providers/openai/images.gocore/providers/openai/openai.go
🧬 Code graph analysis (8)
framework/streaming/images.go (3)
framework/streaming/types.go (1)
ImageStreamChunk(104-114)core/schemas/images.go (3)
BifrostImageGenerationResponse(33-40)ImageData(42-47)BifrostImageGenerationStreamResponse(55-65)core/utils.go (2)
GetResponseFields(216-223)IsFinalChunk(198-213)
core/schemas/bifrost.go (4)
ui/lib/types/config.ts (1)
RequestType(122-134)core/schemas/models.go (2)
BifrostListModelsRequest(23-34)Model(109-129)core/schemas/images.go (3)
BifrostImageGenerationRequest(4-11)BifrostImageGenerationResponse(33-40)BifrostImageGenerationStreamResponse(55-65)core/schemas/provider.go (1)
Provider(282-315)
core/internal/testutil/image_generation_load.go (3)
core/internal/testutil/account.go (1)
ComprehensiveTestConfig(49-68)core/schemas/images.go (3)
BifrostImageGenerationRequest(4-11)ImageGenerationInput(18-20)ImageGenerationParameters(22-30)core/schemas/bifrost.go (2)
ImageGenerationRequest(98-98)ImageGenerationStreamRequest(99-99)
plugins/semanticcache/utils.go (2)
core/schemas/bifrost.go (3)
RequestType(83-83)ImageGenerationRequest(98-98)ImageGenerationStreamRequest(99-99)core/schemas/images.go (2)
ImageGenerationInput(18-20)ImageGenerationParameters(22-30)
core/bifrost.go (2)
core/schemas/bifrost.go (6)
ImageGenerationRequest(98-98)BifrostError(381-390)RequestType(83-83)ImageGenerationStreamRequest(99-99)BifrostStream(343-351)BifrostRequest(150-161)core/schemas/images.go (2)
BifrostImageGenerationRequest(4-11)BifrostImageGenerationResponse(33-40)
core/images_test.go (2)
core/schemas/images.go (7)
BifrostImageGenerationRequest(4-11)ImageGenerationInput(18-20)ImageGenerationParameters(22-30)BifrostImageGenerationResponse(33-40)ImageData(42-47)ImageUsage(49-52)BifrostImageGenerationStreamResponse(55-65)core/schemas/bifrost.go (1)
OpenAI(35-35)
core/providers/openai/images.go (2)
core/providers/openai/types.go (2)
OpenAIImageGenerationRequest(294-304)OpenAIImageGenerationResponse(307-319)core/schemas/images.go (5)
BifrostImageGenerationRequest(4-11)ImageGenerationParameters(22-30)BifrostImageGenerationResponse(33-40)ImageData(42-47)ImageUsage(49-52)
ui/hooks/useImageStream.ts (2)
ui/lib/store/apis/baseApi.ts (1)
getTokenFromStorage(10-25)ui/lib/utils/port.ts (1)
getEndpointUrl(117-127)
🔇 Additional comments (34)
core/images_test.go (5)
14-98: Well-structured serialization tests with proper error handling.The test cases cover both full request serialization and minimal request omission of nil fields. Error handling improvements from past reviews have been incorporated correctly with
t.Fatalfon unmarshal failures and safe type assertions withokchecks.
100-167: Deserialization tests cover key scenarios appropriately.Good coverage of valid JSON parsing, invalid JSON error handling, and missing optional fields behavior. The
wantErrpattern is properly implemented.
207-280: Streaming response serialization tests are thorough.Tests correctly verify both partial chunks and completed chunks with usage data. Type assertions include proper
okchecks andt.Fatalfcalls on failure.
282-417: Stream chunk accumulation tests effectively validate reconstruction logic.The tests cover single-image multi-chunk, multi-image parallel chunks, and out-of-order chunk sorting. This aligns well with the streaming accumulator requirements in the PR objectives.
477-502: Cache key generation helper correctly usesfmt.Sprintffor integer conversion.The previous issue with
string(rune(*req.Params.N))has been addressed by usingfmt.Sprintf("%d", *req.Params.N)on line 497.plugins/semanticcache/utils.go (4)
87-94: Image generation hash handling correctly short-circuits for dedicated cache key.The special handling for
ImageGenerationRequestandImageGenerationStreamRequestappropriately returns a dedicated image cache key. The nil check on line 90 prevents potential panics.Note: Past review comment about typo "styele" appears to be addressed - the comment now correctly reads "style".
175-179: Image generation parameter extraction follows established patterns.Consistent with other request type handlers in the switch statement.
336-345: Image generation text extraction correctly normalizes the prompt.The implementation properly handles the nil checks and returns the normalized prompt for embedding generation.
960-986: Image generation parameter extraction is complete and follows conventions.All parameters from
ImageGenerationParametersstruct are correctly extracted to metadata, including theExtraParamsmap copy.core/providers/azure/types.go (1)
37-47: Azure image request type is well-structured.Fields align with Azure OpenAI image generation API parameters. Optional fields correctly use pointer types with
omitemptytags.ui/hooks/useImageStream.ts (4)
54-77: State management and reset logic are correctly implemented.Initial state is well-defined and the reset callback properly clears all refs and state.
79-86: Cancel function correctly handles abort controller.Properly aborts in-flight requests and nullifies the ref.
107-179: Chunk processing logic handles all stream event types correctly.Error handling, chunk accumulation, and completion detection are well-implemented. The progress estimation using
expectedImagesRef.current * 10is reasonable for UX purposes.
181-194: Stream initiation correctly cancels existing requests first.The past review concern about concurrent SSE readers has been addressed -
cancel()is called beforereset()on line 184, and placeholder images are pre-seeded on line 192.core/providers/azure/images.go (4)
10-23: Request conversion is clean with proper nil guards.The function correctly handles nil input and delegates parameter mapping to a helper function.
25-36: Parameter mapping helper is appropriately scoped as unexported.Direct pointer assignments are correct here since both source and target use the same pointer types.
38-73: Response conversion correctly maps all fields including ExtraFields.The implementation properly:
- Handles nil response
- Maps image data with correct indexing
- Conditionally creates usage struct
- Populates ExtraFields with provider and latency metadata
54-60: TheInputTokenstoPromptTokensmapping is consistent with the OpenAI provider and follows the standard pattern across the codebase. Both Azure and OpenAI providers use identical field mappings for their respective API responses, so no changes are needed.core/providers/gemini/gemini.go (2)
494-500: Streaming response builder arity update looks consistent.
All updatedGetBifrostResponseForStreamResponse(...)call sites appear to pass the correct “slot” plus the new trailingnilfor images.Also applies to: 802-803, 812-813, 1115-1116, 1141-1142, 1385-1386, 1417-1418
1424-1432: Unsupported Gemini image-generation stubs look correct.
ReturningNewUnsupportedOperationErrorwithschemas.ImageGenerationRequest/schemas.ImageGenerationStreamRequestmatches the provider’s “not supported” pattern.core/internal/testutil/image_generation.go (1)
12-126: Test helper wiring + basic validations look good (and uses retry wrapper).
TheWithImageGenerationRetry(...)usage and the stronger base64 length threshold are good for CI stability and catching truncation.core/providers/openai/openai.go (1)
17-18: The dependencygithub.com/google/uuid v1.6.0is already properly pinned incore/go.modand aligns with existing project dependencies. No action needed.core/schemas/bifrost.go (2)
139-161: BifrostRequest image field wiring looks consistent.
ImageGenerationRequestis included in the “only one field” set and correctly plumbed throughGetRequestFields().Also applies to: 163-183
341-373: Stream marshaling for image chunks is correct (single-non-nil selection).
The addedBifrostImageGenerationStreamResponsebranch is consistent with existing marshal behavior.framework/streaming/accumulator.go (4)
124-136: Accumulator integration for image chunks looks consistent.
Initialization, append order, start/final timestamps, and pool-return on cleanup all follow the patterns used for other stream types.Also applies to: 227-243, 245-270
346-380: Routing by RequestType is fine—event “type” naming still needs one canonical contract.
This dispatch will work, but only if upstream consistently setsExtraFields.RequestTypeand stream chunks use the expectedBifrostImageGenerationStreamResponse.Typevalues.
470-517: Pool prewarm includes image chunks: nice.
Prewarming the new pool matches the existing perf strategy.
13-123: All ImageStreamChunk fields are properly reset in putImageStreamChunk.The struct contains: Timestamp, Delta, FinishReason, ChunkIndex, ImageIndex, ErrorDetails, Cost, SemanticCacheDebug, and TokenUsage. Your reset logic covers all of these—nil'ing all pointers and zeroing integer and time fields. Unlike ChatStreamChunk, ImageStreamChunk has no RawResponse field, so the current reset is complete and safe for object pool reuse.
core/internal/testutil/image_generation_load.go (2)
15-101: Load test concurrency approach looks fine.
Requests are constructed per call, counters are atomic, and duration aggregation is mutex-protected.
220-325: Good: avoids shared request pointer races for cache-load goroutines.
The per-goroutinelocalRequestavoids the data race risk that often shows up in load tests.core/bifrost.go (3)
1758-1763: Fallback propagation for image generation is consistent with other request types.
Provider/model override is correctly applied for fallbacks.
2500-2506: Provider dispatch for image generation (sync + stream) looks correct.
The new cases route through the provider interface as expected.Also applies to: 2536-2538
2705-2715: Request pool reset includes ImageGenerationRequest: good.
Clearing the new pointer avoids stale reuse from the pool.core/providers/openai/images.go (1)
52-87: The fieldsRequestTypeandModelRequestedare already being populated in theImageGenerationhandler (lines 2220–2222 ofcore/providers/openai/openai.go), immediately after callingToBifrostImageResponse:bifrostResp.ExtraFields.Provider = provider.GetProviderKey() bifrostResp.ExtraFields.ModelRequested = openaiRequest.Model bifrostResp.ExtraFields.RequestType = schemas.ImageGenerationRequestThis is the established pattern across the codebase (see cerebras, vertex, and other OpenAI handlers). The conversion function returns a bare response; the handler enriches
ExtraFields. No changes needed inToBifrostImageResponse.Likely an incorrect or invalid review comment.
… for image generation for non-go consumers
There was a problem hiding this comment.
Actionable comments posted: 1
Caution
Some comments are outside the diff and can’t be posted inline due to platform limitations.
⚠️ Outside diff range comments (2)
transports/bifrost-http/handlers/inference.go (1)
1038-1124: Don’t log full streaming chunks; also makeincludeEventTypesticky to avoid emitting[DONE]for Responses streams.func (h *CompletionHandler) handleStreamingResponse(ctx *fasthttp.RequestCtx, getStream func() (chan *schemas.BifrostStream, *schemas.BifrostError), cancel context.CancelFunc) { @@ - var includeEventType bool + // Sticky: once we detect Responses streaming, we must not emit legacy `[DONE]`. + var includeEventType bool @@ for chunk := range stream { if chunk == nil { continue } - includeEventType = false - if chunk.BifrostResponsesStreamResponse != nil || - (chunk.BifrostError != nil && chunk.BifrostError.ExtraFields.RequestType == schemas.ResponsesStreamRequest) { - includeEventType = true - } + // Sticky decision for this stream (don’t flip-flop per chunk). + if !includeEventType { + if chunk.BifrostResponsesStreamResponse != nil || + (chunk.BifrostError != nil && chunk.BifrostError.ExtraFields.RequestType == schemas.ResponsesStreamRequest) { + includeEventType = true + } + } @@ chunkJSON, err := sonic.Marshal(chunk) if err != nil { - logger.Warn(fmt.Sprintf("Failed to marshal streaming response: %v, chunk: %v", err, chunk)) + // Avoid logging full chunk (may include large base64 payloads / sensitive content) + logger.Warn(fmt.Sprintf("Failed to marshal streaming response: %v", err)) continue } @@ if !includeEventType { // Send the [DONE] marker to indicate the end of the stream (only for non-responses APIs) if _, err := fmt.Fprint(w, "data: [DONE]\n\n"); err != nil { logger.Warn(fmt.Sprintf("Failed to write SSE [DONE] marker: %v", err)) cancel() // Client disconnected (write error), cancel upstream stream return } }Based on learnings, don’t use
[DONE]for Responses API streaming.ui/lib/types/config.ts (1)
121-154: UpdateAllowedRequestsschema and form initializers to includeimage_generationandimage_generation_streamfields.The TypeScript type definition includes both new fields, but multiple schema and initialization locations are missing them:
- ui/lib/types/schemas.ts (line ~310):
allowedRequestsSchemais missingimage_generationandimage_generation_stream- ui/lib/schemas/providerForm.ts (line ~55):
AllowedRequestsSchemais missing both new fields pluslist_models- ui/app/workspace/providers/dialogs/addNewCustomProviderSheet.tsx (line ~43): Default values missing both new fields
- ui/app/workspace/providers/fragments/apiStructureFormFragment.tsx (line ~40): Default values missing both new fields
Forms will fail validation when the new fields are submitted, and existing providers won't have these fields initialized. Add
image_generation: z.boolean()andimage_generation_stream: z.boolean()to both Zod schemas and both form default initializations.
♻️ Duplicate comments (1)
transports/bifrost-http/handlers/inference.go (1)
222-227: (Optional) Keep request unmarshalling consistent with Chat/Responses if you run into embedded-pointer + sonic edge cases.
🧹 Nitpick comments (1)
core/providers/openai/openai.go (1)
2165-2170: Minor: Redundant provider name retrieval.Lines 2165 and 2170 both retrieve the provider key/name. Consider removing the redundant call.
Apply this diff:
if err := providerUtils.CheckOperationAllowed(schemas.OpenAI, provider.customProviderConfig, schemas.ImageGenerationRequest); err != nil { return nil, err // Handle error } providerKey := provider.GetProviderKey() openaiRequest := ToOpenAIImageGenerationRequest(req) if openaiRequest == nil { return nil, providerUtils.NewBifrostOperationError("invalid request: input is required", nil, providerKey) } - providerName := provider.GetProviderKey() + providerName := providerKey
📜 Review details
Configuration used: CodeRabbit UI
Review profile: CHILL
Plan: Pro
📒 Files selected for processing (13)
core/changelog.md(1 hunks)core/internal/testutil/image_generation_load.go(1 hunks)core/providers/azure/azure.go(2 hunks)core/providers/openai/openai.go(10 hunks)core/schemas/images.go(1 hunks)framework/changelog.md(1 hunks)plugins/semanticcache/changelog.md(1 hunks)transports/bifrost-http/handlers/inference.go(7 hunks)transports/bifrost-http/integrations/openai.go(1 hunks)transports/changelog.md(1 hunks)ui/app/workspace/providers/fragments/allowedRequestsFields.tsx(2 hunks)ui/lib/constants/logs.ts(3 hunks)ui/lib/types/config.ts(2 hunks)
✅ Files skipped from review due to trivial changes (1)
- framework/changelog.md
🧰 Additional context used
📓 Path-based instructions (1)
**
⚙️ CodeRabbit configuration file
always check the stack if there is one for the current PR. do not give localized reviews for the PR, always see all changes in the light of the whole stack of PRs (if there is a stack, if there is no stack you can continue to make localized suggestions/reviews)
Files:
transports/bifrost-http/integrations/openai.gocore/changelog.mdui/lib/constants/logs.tsui/app/workspace/providers/fragments/allowedRequestsFields.tsxplugins/semanticcache/changelog.mdcore/providers/openai/openai.goui/lib/types/config.tscore/internal/testutil/image_generation_load.gocore/schemas/images.gotransports/changelog.mdtransports/bifrost-http/handlers/inference.gocore/providers/azure/azure.go
🧠 Learnings (4)
📚 Learning: 2025-12-09T17:07:42.007Z
Learnt from: qwerty-dvorak
Repo: maximhq/bifrost PR: 1006
File: core/schemas/account.go:9-18
Timestamp: 2025-12-09T17:07:42.007Z
Learning: In core/schemas/account.go, the HuggingFaceKeyConfig field within the Key struct is currently unused and reserved for future Hugging Face inference endpoint deployments. Do not flag this field as missing from OpenAPI documentation or require its presence in the API spec until the feature is actively implemented and used. When the feature is added, update the OpenAPI docs accordingly; otherwise, treat this field as non-breaking and not part of the current API surface.
Applied to files:
transports/bifrost-http/integrations/openai.gocore/providers/openai/openai.gocore/internal/testutil/image_generation_load.gocore/schemas/images.gotransports/bifrost-http/handlers/inference.gocore/providers/azure/azure.go
📚 Learning: 2025-12-12T08:25:02.629Z
Learnt from: Pratham-Mishra04
Repo: maximhq/bifrost PR: 1000
File: transports/bifrost-http/integrations/router.go:709-712
Timestamp: 2025-12-12T08:25:02.629Z
Learning: In transports/bifrost-http/**/*.go, update streaming response handling to align with OpenAI Responses API: use typed SSE events such as response.created, response.output_text.delta, response.done, etc., and do not rely on the legacy data: [DONE] termination marker. Note that data: [DONE] is only used by the older Chat Completions and Text Completions streaming APIs. Ensure parsers, writers, and tests distinguish SSE events from the [DONE] sentinel and handle each event type accordingly for correct stream termination and progress updates.
Applied to files:
transports/bifrost-http/integrations/openai.gotransports/bifrost-http/handlers/inference.go
📚 Learning: 2025-12-11T11:58:25.307Z
Learnt from: Pratham-Mishra04
Repo: maximhq/bifrost PR: 1000
File: core/providers/openai/responses.go:42-84
Timestamp: 2025-12-11T11:58:25.307Z
Learning: In core/providers/openai/responses.go (and related OpenAI response handling), document and enforce the API format constraint: if ResponsesReasoning != nil and the response contains content blocks, all content blocks should be treated as reasoning blocks by default. Implement type guards or parsing logic accordingly, and add unit tests to verify that when ResponsesReasoning is non-nil, content blocks are labeled as reasoning blocks. Include clear comments in the code explaining the rationale and ensure downstream consumers rely on this behavior.
Applied to files:
core/providers/openai/openai.go
📚 Learning: 2025-12-14T14:43:30.902Z
Learnt from: Radheshg04
Repo: maximhq/bifrost PR: 980
File: core/providers/openai/images.go:10-22
Timestamp: 2025-12-14T14:43:30.902Z
Learning: Enforce the OpenAI image generation SSE event type values across the OpenAI image flow in the repository: use "image_generation.partial_image" for partial chunks, "image_generation.completed" for the final result, and "error" for errors. Apply this consistently in schemas, constants, tests, accumulator routing, and UI code within core/providers/openai (and related Go files) to ensure uniform event typing and avoid mismatches.
Applied to files:
core/providers/openai/openai.go
🧬 Code graph analysis (5)
transports/bifrost-http/integrations/openai.go (2)
core/providers/openai/types.go (1)
OpenAIImageGenerationRequest(294-304)core/schemas/models.go (1)
Model(109-129)
core/providers/openai/openai.go (5)
core/providers/utils/utils.go (2)
GetBifrostResponseForStreamResponse(855-887)ReleaseStreamingResponse(845-852)core/schemas/images.go (4)
BifrostImageGenerationRequest(4-11)BifrostImageGenerationResponse(33-40)BifrostImageGenerationStreamResponse(55-65)ImageUsage(49-52)core/providers/openai/images.go (6)
ToOpenAIImageGenerationRequest(25-37)ToBifrostImageResponse(53-87)StreamingEnabledImageModels(18-22)ImageGenerationCompleted(12-12)ImageGenerationPartial(11-11)ImageGenerationChunkSize(15-15)core/providers/openai/errors.go (1)
ParseOpenAIError(10-42)core/providers/openai/types.go (3)
OpenAIImageGenerationResponse(307-319)OpenAIImageGenerationRequest(294-304)OpenAIImageStreamResponse(333-338)
core/schemas/images.go (2)
core/schemas/provider.go (1)
Provider(282-315)core/schemas/bifrost.go (4)
ModelProvider(32-32)Fallback(134-137)BifrostResponseExtraFields(308-318)BifrostError(381-390)
transports/bifrost-http/handlers/inference.go (5)
core/schemas/images.go (3)
ImageGenerationInput(18-20)ImageGenerationParameters(22-30)BifrostImageGenerationRequest(4-11)core/schemas/utils.go (1)
ParseModelString(23-36)core/schemas/provider.go (1)
Provider(282-315)core/schemas/bifrost.go (5)
ModelProvider(32-32)ImageGenerationRequest(98-98)BifrostStream(343-351)BifrostError(381-390)ImageGenerationStreamRequest(99-99)transports/bifrost-http/lib/ctx.go (1)
ConvertToBifrostContext(72-283)
core/providers/azure/azure.go (7)
core/schemas/account.go (2)
Key(8-17)AzureKeyConfig(21-25)core/schemas/images.go (2)
BifrostImageGenerationRequest(4-11)BifrostImageGenerationResponse(33-40)core/providers/azure/images.go (1)
ToAzureImageRequest(11-23)core/providers/azure/types.go (1)
AzureImageResponse(50-58)core/schemas/utils.go (1)
Ptr(16-18)core/utils.go (1)
Ptr(51-53)core/providers/openai/openai.go (1)
HandleOpenAIImageGenerationStreaming(2281-2550)
🔇 Additional comments (16)
plugins/semanticcache/changelog.md (1)
1-2: Changelog entry looks fine; just ensure the leading-formatting matches the repo’s changelog conventions.transports/changelog.md (1)
7-8: Changelog update matches the transport-layer change.ui/app/workspace/providers/fragments/allowedRequestsFields.tsx (1)
21-37: UI request-type + endpoint placeholder additions are consistent with/v1/images/generations.Also applies to: 61-76
transports/bifrost-http/integrations/openai.go (1)
49-64: Good: Azure pre-hook now normalizes image-generation model names the same way as other OpenAI request types.core/changelog.md (1)
6-7: Changelog entry aligns with the new core image-generation surface.ui/lib/constants/logs.ts (1)
28-42: Nice: request-type labels/colors are wired for image generation and streaming.Also applies to: 81-137
transports/bifrost-http/handlers/inference.go (3)
141-152: Known-fields list for image extra-params extraction looks correct.
358-371: Route registration matches the intended OpenAI-compatible path.
1282-1293: Streaming plumbing reuses the shared SSE writer consistently.core/providers/azure/azure.go (2)
787-851: LGTM! Past issues resolved.The non-streaming ImageGeneration implementation correctly:
- Uses the deployment-scoped path
openai/deployments/{deployment}/images/generations(Line 818).- Returns
bifrostErrdirectly on deployment resolution failure (Line 812).- Maps
request.Modelinto the Bifrost response (Line 834).
853-915: LGTM! Streaming implementation follows Azure patterns.The ImageGenerationStream method correctly:
- Validates the key configuration.
- Resolves the deployment and constructs the deployment-scoped URL.
- Handles Bearer token vs. api-key authentication.
- Enforces streaming support per model via
StreamingEnabledImageModels.- Delegates to the OpenAI-compatible streaming handler with appropriate parameters.
core/providers/openai/openai.go (2)
2233-2280: LGTM! Past issues resolved.The ImageGenerationStream method now correctly:
- Guards against nil request (Line 2243-2245).
- Uses
ImageGenerationStreamRequestfor URL building (Line 2267).- Passes
sendBackRawRequestparameter to the streaming handler (Line 2271).
2462-2539: LGTM! Completion chunk logic fixed.The streaming chunk handling now correctly:
- Increments
chunkIndexfor each base64 chunk (Line 2472).- Increments
chunkIndexagain for the completion chunk (Line 2503), ensuring it has a unique, strictly increasing index.- Returns after sending the completion chunk (Line 2538), preventing further processing.
This resolves the past issue where the completion chunk reused the same index and the goroutine continued scanning.
core/internal/testutil/image_generation_load.go (3)
17-102: LGTM! Concurrent load test implementation is safe.The function correctly:
- Creates a fresh request within each goroutine iteration (Lines 47-58), avoiding data races.
- Uses mutex-protected updates for
totalDuration(Lines 64-66).- Uses atomic operations for success/error counters.
- Validates a minimum 80% success rate.
104-222: LGTM! Stream load test correctly manages context lifecycle.The function now correctly:
- Creates a fresh
cancelfunction per stream (Line 148).- Calls
cancel()explicitly on error paths (Lines 152, 159, 173, 181).- Calls
cancel()after stream processing completes (Line 187).- Uses the canonical
string(openai.ImageGenerationCompleted)constant for completion checks (Line 179).This resolves the past issues with
defer cancel()placement and hardcoded event type strings.
224-329: LGTM! Cache load test avoids data races.The function now correctly creates a fresh
localRequestwithin each goroutine (Lines 275-285), preventing data races on the shared request pointer. The cache hit/miss tracking and metrics aggregation are properly mutex-protected or use atomic operations.
There was a problem hiding this comment.
Actionable comments posted: 0
🧹 Nitpick comments (2)
framework/changelog.md (1)
1-1: Clarify the scope and add detail to the changelog entry.The entry "feat: add image generation streaming accumulation support" is quite vague and underspecifies what framework changes are included. Based on the PR objectives, the framework likely contributes more than just streaming accumulation (e.g., streaming handlers, chunk reassembly, streaming finalization logic).
Consider expanding the entry to be more descriptive of the actual framework changes, such as:
feat: add image generation streaming support with accumulation and chunk reassemblyor
feat: add streaming framework for image generation with base64 chunk accumulation and progressive deliveryThis ensures the changelog is self-documenting and helps users and maintainers understand the scope of framework changes at a glance.
Please verify that this entry comprehensively captures all framework-level changes related to image generation in this PR. If there are other changes (e.g., new handler registration, serialization/deserialization helpers, or observability hooks), they should be reflected in the changelog entry as well.
core/changelog.md (1)
1-1: Improve changelog entry specificity.The changelog entry is quite generic and doesn't reflect the comprehensive scope of this PR. Based on the PR objectives, this release includes streaming support, semantic caching integration, UI components, and provider implementations—not just request/response support.
Consider expanding the entry to better communicate the scope to users:
-feat: added image generation request and response support +feat: image generation support with streaming, semantic caching, and provider implementationsOr break it into multiple entries if the changelog format supports detailed feature descriptions.
Please verify:
- That this changelog entry follows your project's format conventions (see the full
core/changelog.mdfile for consistency).- That corresponding updates are being made to
frameworks/changelog.md,plugins/semanticcache/changelog.md, andtransports/changelog.mdas noted in the pre-merge checklist.
📜 Review details
Configuration used: CodeRabbit UI
Review profile: CHILL
Plan: Pro
📒 Files selected for processing (6)
core/changelog.md(1 hunks)framework/changelog.md(1 hunks)plugins/semanticcache/changelog.md(1 hunks)transports/bifrost-http/handlers/inference.go(7 hunks)transports/bifrost-http/integrations/router.go(6 hunks)transports/changelog.md(1 hunks)
🚧 Files skipped from review as they are similar to previous changes (2)
- transports/changelog.md
- plugins/semanticcache/changelog.md
🧰 Additional context used
📓 Path-based instructions (1)
**
⚙️ CodeRabbit configuration file
always check the stack if there is one for the current PR. do not give localized reviews for the PR, always see all changes in the light of the whole stack of PRs (if there is a stack, if there is no stack you can continue to make localized suggestions/reviews)
Files:
transports/bifrost-http/integrations/router.goframework/changelog.mdcore/changelog.mdtransports/bifrost-http/handlers/inference.go
🧠 Learnings (2)
📚 Learning: 2025-12-09T17:07:42.007Z
Learnt from: qwerty-dvorak
Repo: maximhq/bifrost PR: 1006
File: core/schemas/account.go:9-18
Timestamp: 2025-12-09T17:07:42.007Z
Learning: In core/schemas/account.go, the HuggingFaceKeyConfig field within the Key struct is currently unused and reserved for future Hugging Face inference endpoint deployments. Do not flag this field as missing from OpenAPI documentation or require its presence in the API spec until the feature is actively implemented and used. When the feature is added, update the OpenAPI docs accordingly; otherwise, treat this field as non-breaking and not part of the current API surface.
Applied to files:
transports/bifrost-http/integrations/router.gotransports/bifrost-http/handlers/inference.go
📚 Learning: 2025-12-12T08:25:02.629Z
Learnt from: Pratham-Mishra04
Repo: maximhq/bifrost PR: 1000
File: transports/bifrost-http/integrations/router.go:709-712
Timestamp: 2025-12-12T08:25:02.629Z
Learning: In transports/bifrost-http/**/*.go, update streaming response handling to align with OpenAI Responses API: use typed SSE events such as response.created, response.output_text.delta, response.done, etc., and do not rely on the legacy data: [DONE] termination marker. Note that data: [DONE] is only used by the older Chat Completions and Text Completions streaming APIs. Ensure parsers, writers, and tests distinguish SSE events from the [DONE] sentinel and handle each event type accordingly for correct stream termination and progress updates.
Applied to files:
transports/bifrost-http/integrations/router.gotransports/bifrost-http/handlers/inference.go
🧬 Code graph analysis (2)
transports/bifrost-http/integrations/router.go (2)
core/schemas/images.go (2)
BifrostImageGenerationResponse(33-40)BifrostImageGenerationStreamResponse(55-65)core/schemas/bifrost.go (2)
ImageGenerationRequest(98-98)ImageGenerationStreamRequest(99-99)
transports/bifrost-http/handlers/inference.go (4)
core/schemas/images.go (3)
ImageGenerationInput(18-20)ImageGenerationParameters(22-30)BifrostImageGenerationRequest(4-11)core/schemas/utils.go (1)
ParseModelString(23-36)core/schemas/bifrost.go (2)
ModelProvider(32-32)ImageGenerationRequest(98-98)transports/bifrost-http/lib/ctx.go (1)
ConvertToBifrostContext(72-283)
🔇 Additional comments (4)
transports/bifrost-http/integrations/router.go (1)
133-139: LGTM—image generation converters are well integrated.The new
ImageGenerationResponseConverterandImageGenerationStreamResponseConvertertypes follow the established pattern for other modalities. Integration intoStreamConfig,RouteConfig, non-streaming dispatch (lines 579-601), streaming dispatch (lines 653-654), and the central streaming flow (lines 839-840) is consistent with existing request types. PostCallback support and nil checks are properly included.Also applies to: 192-192, 221-221, 579-601, 653-654, 839-840
transports/bifrost-http/handlers/inference.go (3)
141-152: LGTM—imageParamsKnownFields is complete.The map correctly includes all known parameters for image generation: model, prompt, fallbacks, stream, n, size, quality, style, response_format, and user. This aligns with
ImageGenerationParametersfields in core/schemas/images.go.
1205-1282: LGTM—image generation handler is well structured.The
imageGenerationhandler follows the established pattern for other modalities:
- Parses model in
provider/modelformat and validates prompt (lines 1215-1224)- Extracts extra params using
imageParamsKnownFields(lines 1226-1237)- Parses fallbacks (lines 1238-1243)
- Properly manages context lifecycle: calls
cancel()on nil check (line 1257), callscancel()before validation error return (line 1265), and deferscancel()for non-streaming path (line 1272)- Validates streaming constraints: rejects
"url"response format when streaming is requested (lines 1264-1266)
1284-1295: LGTM—streaming helper follows the established pattern.
handleStreamingImageGenerationcreates a cancellable stream context, wraps the Bifrost call in a closure, and delegates to the sharedhandleStreamingResponsehelper. This matches the structure ofhandleStreamingChatCompletion,handleStreamingResponses, and other streaming helpers.
…rix and updated schema docs
There was a problem hiding this comment.
Actionable comments posted: 4
🧹 Nitpick comments (1)
docs/quickstart/go-sdk/multimodal.mdx (1)
65-75: Nil/empty handling: tighten the example to avoid silent no-op.Right now, if
len(response.Data) > 0butresponse.Data[0].URL == nil, nothing is printed; consider adding anelseto log/print that the response format wasn’turl(or that the provider returned b64).
📜 Review details
Configuration used: CodeRabbit UI
Review profile: CHILL
Plan: Pro
📒 Files selected for processing (4)
core/schemas/images.go(1 hunks)docs/features/unified-interface.mdx(1 hunks)docs/quickstart/gateway/multimodal.mdx(2 hunks)docs/quickstart/go-sdk/multimodal.mdx(2 hunks)
🧰 Additional context used
📓 Path-based instructions (1)
**
⚙️ CodeRabbit configuration file
always check the stack if there is one for the current PR. do not give localized reviews for the PR, always see all changes in the light of the whole stack of PRs (if there is a stack, if there is no stack you can continue to make localized suggestions/reviews)
Files:
docs/quickstart/gateway/multimodal.mdxcore/schemas/images.godocs/features/unified-interface.mdxdocs/quickstart/go-sdk/multimodal.mdx
🧠 Learnings (1)
📚 Learning: 2025-12-09T17:07:42.007Z
Learnt from: qwerty-dvorak
Repo: maximhq/bifrost PR: 1006
File: core/schemas/account.go:9-18
Timestamp: 2025-12-09T17:07:42.007Z
Learning: In core/schemas/account.go, the HuggingFaceKeyConfig field within the Key struct is currently unused and reserved for future Hugging Face inference endpoint deployments. Do not flag this field as missing from OpenAPI documentation or require its presence in the API spec until the feature is actively implemented and used. When the feature is added, update the OpenAPI docs accordingly; otherwise, treat this field as non-breaking and not part of the current API surface.
Applied to files:
core/schemas/images.go
🔇 Additional comments (9)
docs/quickstart/gateway/multimodal.mdx (2)
3-3: Doc description update looks good; keep capability list consistent with the rest of the docs IA
If other pages use a canonical set/order of “vision / image / speech / transcription” capabilities, consider matching that wording/order here to avoid drift.
49-85: Update created timestamp and clarify unified model naming in request/responseThe example response shows an old epoch timestamp (
created: 1713833628); update to a current or placeholder value like0or a realistic recent timestamp to prevent user confusion about staleness.Regarding the model field: the example correctly shows the unified format (
openai/dall-e-3) in the request and provider-native format (dall-e-3) in the response—this matches Bifrost's design. However, to reduce user confusion when copy-pasting, consider adding a brief comment noting that Bifrost normalizes the unified model name to the provider-native format in responses.Streaming image generation is implemented in Bifrost (via
"stream": truewith base64 response format), but adding a streaming example appears to be outside the current PR scope focused on basic image generation. Defer streaming documentation to a follow-up if planned.docs/features/unified-interface.mdx (1)
88-89: LGTM! Table header correctly extended.The new "Image Generation" and "Image Generation (stream)" columns are properly added to the provider support matrix header.
core/schemas/images.go (6)
1-16: LGTM! Request structure follows Bifrost patterns.The
BifrostImageGenerationRequestcorrectly follows the established Bifrost request pattern with Provider, Model, Input, Params, and Fallbacks. TheGetRawRequestBodymethod properly implements the interface for raw body access.
18-20: LGTM! Clean input structure.The
ImageGenerationInputstruct is appropriately simple with just the requiredPromptfield.
32-40: LGTM! Response structure is consistent.The
BifrostImageGenerationResponsefollows the established Bifrost response pattern with standard fields (ID, Created, Model, Usage, ExtraFields).
42-47: LGTM! ImageData structure aligns with standards.The
ImageDatastruct correctly represents individual image results withURLorB64JSONoptions,RevisedPrompt, andIndexfields, consistent with OpenAI's image API format.
54-65: LGTM! Well-designed streaming response structure.The
BifrostImageGenerationStreamResponseis well-structured for progressive image delivery with clear semantics:
Typedifferentiates event typesIndexidentifies which image (for multi-image requests)ChunkIndexmaintains chunk orderingPartialB64carries base64 chunksThe optional fields (
RevisedPrompt,Usage,Error) are appropriately placed for their respective events.
49-52: ImageUsage struct fields are correctly aligned with provider implementations.The
ImageUsagestruct withPromptTokensandTotalTokensmatches the actual usage information returned by both OpenAI and Azure image generation APIs. Both providers map their API responses (InputTokens→PromptTokens,TotalTokens→TotalTokens) directly to these fields. NoCompletionTokensfield is needed for image generation, as image APIs only return input and total token counts, not completion tokens.
There was a problem hiding this comment.
Actionable comments posted: 1
♻️ Duplicate comments (1)
docs/quickstart/go-sdk/multimodal.mdx (1)
3-3: Add an immediate provider-matrix link (the “across various providers” claim is easy to over-read).You already link to Unified Interface later (Line 382-384), but it’s far from the top-level capability claim. Suggest adding a parenthetical link right in the description so readers don’t assume all providers support image generation.
-description: "Process multiple types of content including images, audio, and text with AI models. Bifrost supports vision analysis, image generation, speech synthesis, and audio transcription across various providers." +description: "Process multiple types of content including images, audio, and text with AI models. Bifrost supports vision analysis, image generation, speech synthesis, and audio transcription (see Unified Interface provider support) across various providers."
📜 Review details
Configuration used: CodeRabbit UI
Review profile: CHILL
Plan: Pro
📒 Files selected for processing (3)
core/schemas/images.go(1 hunks)docs/features/unified-interface.mdx(1 hunks)docs/quickstart/go-sdk/multimodal.mdx(2 hunks)
🚧 Files skipped from review as they are similar to previous changes (1)
- core/schemas/images.go
🧰 Additional context used
📓 Path-based instructions (1)
**
⚙️ CodeRabbit configuration file
always check the stack if there is one for the current PR. do not give localized reviews for the PR, always see all changes in the light of the whole stack of PRs (if there is a stack, if there is no stack you can continue to make localized suggestions/reviews)
Files:
docs/features/unified-interface.mdxdocs/quickstart/go-sdk/multimodal.mdx
🔇 Additional comments (1)
docs/quickstart/go-sdk/multimodal.mdx (1)
48-89: The code example compiles correctly; no changes needed. TheImageDatafields (URL,B64JSON,RevisedPrompt) are defined asstringtypes inschemas.ImageData, not*stringpointers. Similarly,ImageUsage.TotalTokensis anint, not a pointer. The documentation example correctly treats these as non-pointer types and will compile without the suggested nil-safety changes.
| r.Model = setAzureModelName(r.Model, deploymentIDStr) | ||
| case *openai.OpenAIEmbeddingRequest: | ||
| r.Model = setAzureModelName(r.Model, deploymentIDStr) | ||
| case *openai.OpenAIImageGenerationRequest: |
There was a problem hiding this comment.
you will also need to add a route for image gen right?
There was a problem hiding this comment.
yes, I’ve added it now and wired it up properly. should be all set
There was a problem hiding this comment.
Actionable comments posted: 1
🧹 Nitpick comments (1)
core/providers/openai/openai_image_test.go (1)
199-207: Minor: Consider restructuring the nil response test case.The "nil response returns nil" test case passes a nil OpenAI response, and the validate function checks
resp != nil. While this works, the pattern is slightly awkward since you're validating a nil pointer. Consider using thewantNilpattern fromTestImageGenerationStreamingRequestConversionfor consistency:{ name: "nil response returns nil", openai: nil, wantNil: true, },Then handle it in the test runner similar to lines 99-105.
📜 Review details
Configuration used: Organization UI
Review profile: CHILL
Plan: Pro
📒 Files selected for processing (4)
core/providers/openai/images.go(1 hunks)core/providers/openai/openai_image_test.go(1 hunks)core/providers/openai/types.go(1 hunks)transports/bifrost-http/integrations/openai.go(2 hunks)
🧰 Additional context used
📓 Path-based instructions (1)
**
⚙️ CodeRabbit configuration file
always check the stack if there is one for the current PR. do not give localized reviews for the PR, always see all changes in the light of the whole stack of PRs (if there is a stack, if there is no stack you can continue to make localized suggestions/reviews)
Files:
transports/bifrost-http/integrations/openai.gocore/providers/openai/images.gocore/providers/openai/openai_image_test.gocore/providers/openai/types.go
🧠 Learnings (5)
📚 Learning: 2025-12-09T17:07:42.007Z
Learnt from: qwerty-dvorak
Repo: maximhq/bifrost PR: 1006
File: core/schemas/account.go:9-18
Timestamp: 2025-12-09T17:07:42.007Z
Learning: In core/schemas/account.go, the HuggingFaceKeyConfig field within the Key struct is currently unused and reserved for future Hugging Face inference endpoint deployments. Do not flag this field as missing from OpenAPI documentation or require its presence in the API spec until the feature is actively implemented and used. When the feature is added, update the OpenAPI docs accordingly; otherwise, treat this field as non-breaking and not part of the current API surface.
Applied to files:
transports/bifrost-http/integrations/openai.gocore/providers/openai/images.gocore/providers/openai/openai_image_test.gocore/providers/openai/types.go
📚 Learning: 2025-12-12T08:25:02.629Z
Learnt from: Pratham-Mishra04
Repo: maximhq/bifrost PR: 1000
File: transports/bifrost-http/integrations/router.go:709-712
Timestamp: 2025-12-12T08:25:02.629Z
Learning: In transports/bifrost-http/**/*.go, update streaming response handling to align with OpenAI Responses API: use typed SSE events such as response.created, response.output_text.delta, response.done, etc., and do not rely on the legacy data: [DONE] termination marker. Note that data: [DONE] is only used by the older Chat Completions and Text Completions streaming APIs. Ensure parsers, writers, and tests distinguish SSE events from the [DONE] sentinel and handle each event type accordingly for correct stream termination and progress updates.
Applied to files:
transports/bifrost-http/integrations/openai.go
📚 Learning: 2025-12-11T11:58:25.307Z
Learnt from: Pratham-Mishra04
Repo: maximhq/bifrost PR: 1000
File: core/providers/openai/responses.go:42-84
Timestamp: 2025-12-11T11:58:25.307Z
Learning: In core/providers/openai/responses.go (and related OpenAI response handling), document and enforce the API format constraint: if ResponsesReasoning != nil and the response contains content blocks, all content blocks should be treated as reasoning blocks by default. Implement type guards or parsing logic accordingly, and add unit tests to verify that when ResponsesReasoning is non-nil, content blocks are labeled as reasoning blocks. Include clear comments in the code explaining the rationale and ensure downstream consumers rely on this behavior.
Applied to files:
core/providers/openai/images.gocore/providers/openai/openai_image_test.gocore/providers/openai/types.go
📚 Learning: 2025-12-14T14:43:30.902Z
Learnt from: Radheshg04
Repo: maximhq/bifrost PR: 980
File: core/providers/openai/images.go:10-22
Timestamp: 2025-12-14T14:43:30.902Z
Learning: Enforce the OpenAI image generation SSE event type values across the OpenAI image flow in the repository: use "image_generation.partial_image" for partial chunks, "image_generation.completed" for the final result, and "error" for errors. Apply this consistently in schemas, constants, tests, accumulator routing, and UI code within core/providers/openai (and related Go files) to ensure uniform event typing and avoid mismatches.
Applied to files:
core/providers/openai/images.gocore/providers/openai/openai_image_test.gocore/providers/openai/types.go
📚 Learning: 2025-12-15T10:16:21.909Z
Learnt from: qwerty-dvorak
Repo: maximhq/bifrost PR: 1006
File: core/providers/huggingface/huggingface_test.go:12-63
Timestamp: 2025-12-15T10:16:21.909Z
Learning: In provider tests under core/providers/<provider>/*_test.go, do not require or flag the use of defer for Shutdown(); instead call client.Shutdown() at the end of each test function. This pattern appears consistent across all provider tests. Apply this rule only within this path; for other tests or resources, defer may still be appropriate.
Applied to files:
core/providers/openai/openai_image_test.go
🧬 Code graph analysis (4)
transports/bifrost-http/integrations/openai.go (4)
core/providers/openai/types.go (1)
OpenAIImageGenerationRequest(294-302)transports/bifrost-http/integrations/router.go (6)
RouteConfig(207-226)RouteConfigTypeOpenAI(199-199)RequestConverter(82-82)ImageGenerationResponseConverter(135-135)ErrorConverter(143-143)StreamConfig(186-194)core/schemas/bifrost.go (4)
BifrostRequest(150-161)ImageGenerationRequest(98-98)OpenAI(35-35)BifrostError(381-390)core/schemas/images.go (2)
BifrostImageGenerationResponse(33-40)BifrostImageGenerationStreamResponse(55-65)
core/providers/openai/images.go (3)
core/providers/openai/types.go (2)
OpenAIImageGenerationRequest(294-302)OpenAIImageGenerationResponse(305-317)core/schemas/images.go (6)
BifrostImageGenerationRequest(4-11)ImageGenerationParameters(22-30)BifrostImageGenerationResponse(33-40)ImageData(42-47)ImageUsage(49-52)ImageGenerationInput(18-20)core/schemas/utils.go (2)
ParseModelString(23-36)ParseFallbacks(53-68)
core/providers/openai/openai_image_test.go (5)
core/schemas/images.go (4)
BifrostImageGenerationRequest(4-11)ImageGenerationInput(18-20)ImageGenerationParameters(22-30)BifrostImageGenerationResponse(33-40)core/providers/openai/types.go (3)
OpenAIImageGenerationRequest(294-302)OpenAIImageGenerationResponse(305-317)OpenAIImageGenerationUsage(319-328)core/schemas/bifrost.go (2)
OpenAI(35-35)Azure(36-36)core/providers/openai/images.go (2)
ToOpenAIImageGenerationRequest(25-39)ToBifrostImageResponse(42-76)core/providers/azure/images.go (1)
ToBifrostImageResponse(39-73)
core/providers/openai/types.go (1)
core/schemas/images.go (1)
ImageGenerationParameters(22-30)
🔇 Additional comments (11)
transports/bifrost-http/integrations/openai.go (2)
60-61: LGTM!The Azure endpoint pre-hook case for
OpenAIImageGenerationRequestfollows the established pattern for other request types, correctly applying thesetAzureModelNamelogic with the deployment ID.
378-425: Image generation endpoint configuration looks well-structured.The route configuration correctly:
- Exposes the standard OpenAI image generation paths including Azure deployment support
- Uses
ToBifrostImageGenerationRequest()for request conversion- Returns
resp.Typeas the SSE event type in the streaming converter (lines 414, 417), which aligns with the OpenAI spec ("image_generation.partial_image","image_generation.completed"). Based on learnings, this is the correct approach for typed SSE events.core/providers/openai/openai_image_test.go (3)
14-108: Well-structured table-driven test for request conversion.The test comprehensively covers:
- Full parameter mapping
- Nil request and nil input edge cases
- Nil params scenario
Good use of parallel subtests and validation functions.
110-139: LGTM!JSON serialization test properly handles the marshal error and verifies key fields are present in the output. The previous review comment about ignoring the
sonic.Marshalerror has been addressed.
236-420: Comprehensive test coverage for Bifrost request conversion.Excellent coverage including:
- Provider prefix parsing (
openai/,azure/)- Fallback chain parsing with mixed providers
- Partial parameters preservation
- Empty vs nil fallbacks distinction
The test cases thoroughly validate the
ToBifrostImageGenerationRequestmethod.core/providers/openai/types.go (1)
304-336: LGTM!The response types correctly model the OpenAI image generation API:
OpenAIImageGenerationResponsewith inlineDatastruct matching the API shapeOpenAIImageGenerationUsagewith detailed token breakdown includingInputTokensDetailsOpenAIImageStreamResponseusing the typedImageGenerationEventTypefor event classificationcore/providers/openai/images.go (5)
10-16: LGTM! Constants align with OpenAI specification.The SSE event type constants
"image_generation.partial_image"and"image_generation.completed"correctly match the OpenAI specification, as confirmed in previous review discussions. The 128KB chunk size is appropriate for base64 streaming.Based on learnings, these values should be used consistently across schemas, tests, accumulator routing, and UI code.
18-22: LGTM!The
StreamingEnabledImageModelsmap correctly identifies that onlygpt-image-1supports streaming, whiledall-e-2anddall-e-3do not. Unknown models will returnfalseby default (safe behavior).
24-39: LGTM!Clean nil-safe conversion with proper parameter mapping. The struct copy at line 36 is safe since
ImageGenerationParameterscontains only primitives and pointers.
41-76: LGTM!The response conversion correctly:
- Maps OpenAI's
InputTokensto Bifrost'sPromptTokens(intentional naming difference)- Preserves image index ordering
- Converts latency to milliseconds for consistency
- Generates a unique ID for the Bifrost response
78-91: LGTM!The method correctly:
- Parses provider prefix from model string (e.g.,
"azure/dall-e-3"→Azure,"dall-e-3")- Takes a pointer to the embedded
ImageGenerationParametersfor theParamsfield- Uses
ParseFallbacksfor proper fallback chain construction
|
|
||
| type ImageGenerationEventType string | ||
|
|
||
| // OpenAIImageGenerationRequest is the struct for Image Generation requests by OpenAI. | ||
| type OpenAIImageGenerationRequest struct { | ||
| Model string `json:"model"` | ||
| Prompt string `json:"prompt"` | ||
|
|
||
| schemas.ImageGenerationParameters | ||
|
|
||
| Stream *bool `json:"stream,omitempty"` | ||
| Fallbacks []string `json:"fallbacks,omitempty"` | ||
| } |
There was a problem hiding this comment.
🧩 Analysis chain
🏁 Script executed:
#!/bin/bash
# Search for StreamingRequest interface definition
echo "=== Looking for StreamingRequest interface ==="
rg -n "type StreamingRequest" --type=go
echo ""
echo "=== Looking for IsStreamingRequested implementations ==="
rg -n "func.*IsStreamingRequested" --type=go
echo ""
echo "=== Looking for OpenAIImageGenerationRequest ==="
rg -n "type OpenAIImageGenerationRequest" --type=go -A 10
echo ""
echo "=== Checking if OpenAIImageGenerationRequest has IsStreamingRequested ==="
rg -n "func.*OpenAIImageGenerationRequest.*IsStreamingRequested" --type=goRepository: maximhq/bifrost
Length of output: 2167
🏁 Script executed:
#!/bin/bash
# Get StreamingRequest interface definition
echo "=== StreamingRequest Interface Definition ==="
rg -n "type StreamingRequest" --type=go -A 5
echo ""
echo "=== Example IsStreamingRequested implementation ==="
rg -n "func.*OpenAIChatRequest.*IsStreamingRequested" --type=go -A 3
echo ""
echo "=== Check if OpenAIImageGenerationRequest is used anywhere ==="
rg -n "OpenAIImageGenerationRequest" --type=go | head -20
echo ""
echo "=== Check git log for context on related types ==="
git log --oneline -20 --allRepository: maximhq/bifrost
Length of output: 3419
🏁 Script executed:
#!/bin/bash
# Check the router integration for OpenAIImageGenerationRequest handling
echo "=== Router integration with OpenAIImageGenerationRequest ==="
rg -n "OpenAIImageGenerationRequest" transports/bifrost-http/integrations/openai.go -B 3 -A 3 | head -30
echo ""
echo "=== Check if OpenAIImageGenerationRequest is used as StreamingRequest ==="
grep -n "case \*openai.OpenAIImageGenerationRequest" transports/bifrost-http/integrations/openai.go -A 10Repository: maximhq/bifrost
Length of output: 1413
🏁 Script executed:
#!/bin/bash
# Check if router calls IsStreamingRequested on requests
echo "=== Check router.go for IsStreamingRequested calls ==="
rg -n "IsStreamingRequested" transports/bifrost-http/integrations/router.go -B 2 -A 2
echo ""
echo "=== Check broader handler context for streaming request handling ==="
rg -n "StreamingRequest" transports/bifrost-http/integrations/ --type=go -B 2 -A 2
echo ""
echo "=== Look at recent commit for image generation ==="
git log -1 --format="%H %s" 502f360
echo ""
echo "=== Check if there's a PR stack or related context ==="
git log --oneline --graph --all | head -30Repository: maximhq/bifrost
Length of output: 3760
Add IsStreamingRequested method to OpenAIImageGenerationRequest.
Image generation requests are routed through streaming detection logic that requires this method. Without it, streaming requests will be mishandled as non-streaming. All other request types with Stream fields implement this method consistently.
// IsStreamingRequested implements the StreamingRequest interface
func (r *OpenAIImageGenerationRequest) IsStreamingRequested() bool {
return r.Stream != nil && *r.Stream
}🤖 Prompt for AI Agents
In core/providers/openai/types.go around lines 290 to 302, the
OpenAIImageGenerationRequest type lacks the IsStreamingRequested method required
by the StreamingRequest interface; add a pointer-receiver method named
IsStreamingRequested that returns true only when Stream is non-nil and true
(i.e., return r.Stream != nil && *r.Stream) so streaming detection treats image
generation requests correctly.
|
@Pratham-Mishra04 openai just released new image gen models, should we add support now or defer that work to a later version? |
2 participants
Summary
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.
Changes
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:
Rollout Plan:
Type of change
Affected areas
Testing Plan
Unit Tests
Integration Tests
Load Tests
How to test
End-to-end image generation
Expected Outcome:
{ "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:
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)
replace github.com/maximhq/bifrost/core => ../../coreScreenshots/Recordings
N/A
Breaking changes
Related issues
Closes #950
Security considerations
N/A
Checklist
docs/contributing/README.mdand followed the guidelines