feat: added image generation support for bedrock#1384
Conversation
|
Caution Review failedThe pull request is closed. 📝 WalkthroughSummary by CodeRabbit
✏️ Tip: You can customize this high-level summary in your review settings. WalkthroughImplements Bedrock image generation: new types, request/response converters and helpers, model-type detection, request construction/signing/execution, model-specific response parsing and enrichment, contextual error propagation (optionally with raw data), and docs marking non-stream image support. Changes
Sequence DiagramsequenceDiagram
participant Client
participant BedrockProvider
participant ModelDetector
participant Converter
participant BedrockAPI
participant ResponseParser
Client->>BedrockProvider: ImageGeneration(Bifrost request)
BedrockProvider->>BedrockProvider: Validate operation & config
BedrockProvider->>ModelDetector: DetermineImageGenModelType(model)
ModelDetector-->>BedrockProvider: resolved model type
BedrockProvider->>Converter: ToBedrockImageGenerationRequest(request)
Converter-->>BedrockProvider: BedrockImageGenerationRequest
BedrockProvider->>BedrockAPI: signAndExecute(model, request)
BedrockAPI-->>BedrockProvider: BedrockImageGenerationResponse (with latency)
BedrockProvider->>ResponseParser: ToBifrostImageGenerationResponse(response)
ResponseParser-->>BedrockProvider: BifrostImageGenerationResponse
BedrockProvider->>BedrockProvider: Enrich response (provider, model, deployment, latency, requestType)
BedrockProvider-->>Client: BifrostImageGenerationResponse or error
Estimated Code Review Effort🎯 3 (Moderate) | ⏱️ ~20 minutes Poem
🚥 Pre-merge checks | ✅ 3✅ Passed checks (3 passed)
✏️ Tip: You can configure your own custom pre-merge checks in the settings. ✨ 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 |
🧪 Test Suite AvailableThis PR can be tested by a repository admin. |
There was a problem hiding this comment.
Actionable comments posted: 5
🤖 Fix all issues with AI agents
In `@core/providers/bedrock/bedrock.go`:
- Line 1539: The comment "ImageGeneration is not supported by the Bedrock
provider" is stale—update the comment near the ImageGeneration symbol in the
Bedrock provider implementation (e.g., the ImageGeneration-related function/type
in bedrock.go) to reflect that image generation is now supported; replace the
old negation with a concise positive statement indicating that the Bedrock
provider supports ImageGeneration and, if relevant, note any limitations or
required flags/options.
- Around line 1579-1607: The final "return nil,
providerUtils.NewUnsupportedOperationError(schemas.ImageGenerationRequest,
schemas.Bedrock)" is unreachable because the earlier switch on modelType already
returns in its case and the outer switch has a default return; remove this
trailing unreachable return or restructure the switch so that unsupported-path
handling is the single exit (e.g., move the NewUnsupportedOperationError into
the outer switch default or convert the inner switch to set bifrostResponse and
break instead of returning). Update references: modelType switch,
ToBifrostImageGenerationResponse, providerUtils.NewUnsupportedOperationError,
schemas.ImageGenerationRequest, schemas.Bedrock.
In `@core/providers/bedrock/images.go`:
- Around line 79-85: Remove the debug print in the loop over response.Images:
delete the fmt.Println(image) call inside the for index, image := range
response.Images loop that populates bifrostResponse.Data with schemas.ImageData;
after removing the print, ensure the fmt import is removed if it becomes unused
so the file compiles cleanly.
- Around line 24-43: The code accesses request.Params.Size before verifying
request.Params is non-nil which can panic; update the logic in the image size
parsing block (where request.Params.Size is read and parsed, and
bedrockReq.ImageGenerationConfig.Width/Height are set) to first check that
request.Params != nil and then check Params.Size != nil; move or add the nil
guard around the existing size parsing (including the strings.Split,
strconv.Atoi and assignments to bedrockReq.ImageGenerationConfig) so all
accesses to request.Params are safe.
- Around line 17-22: The code dereferences request.Input.Prompt when building
BedrockNovaCanvasImageGenerationRequest which can panic if request.Input is nil;
update the image generation function to validate request.Input (and optionally
request.Input.Prompt) before use, returning an error (or handling the missing
input) if nil, then construct BedrockTextToImageParams using the validated
prompt; ensure references to BedrockNovaCanvasImageGenerationRequest,
BedrockTextToImageParams, and request.Input.Prompt are updated accordingly.
🧹 Nitpick comments (3)
core/providers/bedrock/types.go (2)
585-592: Consider addingBedrockprefix for naming consistency.The
ImageGenerationConfigstruct lacks theBedrockprefix that other types in this file consistently use (e.g.,BedrockNovaCanvasImageGenerationRequest,BedrockTextToImageParams). Consider renaming toBedrockImageGenerationConfigfor consistency.
594-599: Consider using pointer types for optional response fields.The
MaskImageandErrorfields are documented as optional but declared as non-pointer strings. This means they'll serialize as empty strings ("") rather than being omitted when not present. Consider using*stringwithomitemptytags if omission is preferred.Suggested change
// BedrockImageGenerationResponse represents a Bedrock image generation response type BedrockImageGenerationResponse struct { Images []string `json:"images"` // list of Base64 encoded images - MaskImage string `json:"maskImage"` // Base64 encoded mask image (optional) - Error string `json:"error"` // error message (if present) + MaskImage *string `json:"maskImage,omitempty"` // Base64 encoded mask image (optional) + Error *string `json:"error,omitempty"` // error message (if present) }core/providers/bedrock/bedrock.go (1)
1557-1572: Consider checking for Bedrock API error in response.The
BedrockImageGenerationResponsehas anErrorfield. After unmarshaling at line 1582, consider checking ifimageResp.Erroris non-empty and returning an appropriate error to the caller rather than silently returning a response that may contain an error message.Suggested addition after unmarshal
var imageResp BedrockImageGenerationResponse if err := sonic.Unmarshal(rawResponse, &imageResp); err != nil { return nil, providerUtils.EnrichError(ctx, providerUtils.NewBifrostOperationError("error parsing image generation response", err, providerName), jsonData, rawResponse, provider.sendBackRawRequest, provider.sendBackRawResponse) } + if imageResp.Error != "" { + return nil, providerUtils.EnrichError(ctx, providerUtils.NewBifrostOperationError(imageResp.Error, nil, providerName), jsonData, rawResponse, provider.sendBackRawRequest, provider.sendBackRawResponse) + } bifrostResponse = ToBifrostImageGenerationResponse(&imageResp)
90b9d93 to
7b490fa
Compare
There was a problem hiding this comment.
Actionable comments posted: 1
🤖 Fix all issues with AI agents
In `@core/providers/bedrock/bedrock.go`:
- Around line 1542-1556: Add an early validation guard in
BedrockProvider.ImageGeneration to fast-fail on nil or invalid image-generation
requests before calling DetermineImageGenModelType or
CheckContextAndGetRequestBody: check that request is non-nil, request.Input is
non-nil, and the prompt (request.Input.Prompt) is not empty (or return a Bifrost
configuration/request error via
providerUtils.NewConfigurationError/NewRequestError). Place this check
immediately after verifying key.BedrockKeyConfig and before the
DetermineImageGenModelType call so downstream conversions cannot dereference nil
values.
♻️ Duplicate comments (1)
core/providers/bedrock/images.go (1)
12-22: Validaterequest.Inputand prompt before dereference.Line 20 dereferences
request.Input.Prompt; ifInputis nil (or prompt empty) this can panic or produce invalid requests. Add a guard after therequest == nilcheck.🛠️ Suggested fix
if request == nil { return nil, fmt.Errorf("request is nil") } + if request.Input == nil || strings.TrimSpace(request.Input.Prompt) == "" { + return nil, fmt.Errorf("prompt cannot be empty") + }Based on learnings, ensure image-generation inputs are non-nil before conversion.
🧹 Nitpick comments (1)
core/providers/bedrock/images.go (1)
57-58: Use the repo pointer helper for width/height.Consider replacing
&width/&heightwith the repo’s pointer helper (e.g.,bifrost.Ptr) for consistency.♻️ Suggested refactor
- bedrockReq.ImageGenerationConfig.Width = &width - bedrockReq.ImageGenerationConfig.Height = &height + bedrockReq.ImageGenerationConfig.Width = bifrost.Ptr(width) + bedrockReq.ImageGenerationConfig.Height = bifrost.Ptr(height)Based on learnings, prefer
bifrost.Ptr(...)over the address operator for pointer creation.
There was a problem hiding this comment.
Actionable comments posted: 1
🤖 Fix all issues with AI agents
In `@core/providers/bedrock/bedrock.go`:
- Around line 1552-1574: The code calls
DetermineImageGenModelType(request.Model) before mapping deployment aliases,
causing valid custom deployment names (from BedrockKeyConfig.Deployments) to be
treated as unsupported; change the flow to resolve the deployment alias first
(use the existing provider.getModelPath or the deployment-mapping logic that
looks up BedrockKeyConfig.Deployments) and then run DetermineImageGenModelType
against the resolved deployment/model string (or check both request.Model and
the resolved deployment) before the switch; update the logic around
DetermineImageGenModelType, getModelPath, and completeRequest to use the
resolved deployment variable so supported image generators are recognized.
♻️ Duplicate comments (1)
core/providers/bedrock/images.go (1)
12-22: Add a provider-level nil guard forrequest.Input.
request.Input.Promptis dereferenced without ensuringrequest.Inputis non-nil. If any non-HTTP entrypoint bypasses the handler validation, this can panic. Add a guard in the converter to fail fast.🔧 Proposed fix
func ToBedrockNovaCanvasImageGenerationRequest(request *schemas.BifrostImageGenerationRequest) (*BedrockNovaCanvasImageGenerationRequest, error) { if request == nil { return nil, fmt.Errorf("request is nil") } + if request.Input == nil { + return nil, fmt.Errorf("request input is nil") + } bedrockReq := &BedrockNovaCanvasImageGenerationRequest{ TaskType: "TEXT_IMAGE", TextToImageParams: BedrockTextToImageParams{ Text: request.Input.Prompt, }, }Based on learnings, please keep a provider-level nil guard here.
7b490fa to
e9ba497
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)
docs/providers/supported-providers/bedrock.mdx (1)
586-615: Fix section numbering after inserting Image Generation.Line 586 introduces “# 6. Batch API”, but the later headings still start at “# 6. Files API” (Line 615) and continue with
#7/#8/#9, causing duplicate numbering. Please renumber the subsequent sections.📝 Suggested update
-# 6. Files API +# 7. Files API ... -# 7. List Models +# 8. List Models ... -# 8. AWS Authentication & Configuration +# 9. AWS Authentication & Configuration ... -# 9. Error Handling +# 10. Error Handling
♻️ Duplicate comments (2)
core/providers/bedrock/bedrock.go (2)
1542-1551: Add an early guard for nil request/Input or empty prompt.Line 1552 reads
request.Model, and downstream conversion accessesrequest.Input.Prompt. If non-HTTP entry points bypass validation, this can panic or send an empty prompt. Add a fast-fail guard before model type detection.🛠️ Proposed fix
providerName := provider.GetProviderKey() + if request == nil || request.Input == nil || strings.TrimSpace(request.Input.Prompt) == "" { + return nil, providerUtils.NewBifrostOperationError("prompt cannot be empty", nil, providerName) + } if key.BedrockKeyConfig == nil { return nil, providerUtils.NewConfigurationError("bedrock key config is not provided", providerName) }Based on learnings, add provider-layer validation even if upstream handlers validate input.
1552-1571: Resolve deployment alias before determining image model type.Line 1552 derives
modelTypefromrequest.Model. If a deployment alias is configured, it won’t match the expected substrings and valid models will be rejected. Resolve the deployment first (or check both) and detect the model type from the resolved deployment.🔧 Suggested flow
- modelType := DetermineImageGenModelType(request.Model) + path, deployment := provider.getModelPath("invoke", request.Model, key) + modelType := DetermineImageGenModelType(deployment) var rawResponse []byte var jsonData []byte var bifrostError *schemas.BifrostError var latency time.Duration - var path string - var deployment string switch modelType { case "titan-image-generator-v1:0", "nova-canvas-v1:0", "titan-image-generator-v2:0": jsonData, bifrostError = providerUtils.CheckContextAndGetRequestBody( ctx, request, func() (any, error) { return ToBedrockNovaCanvasImageGenerationRequest(request) }, provider.GetProviderKey()) if bifrostError != nil { return nil, bifrostError } - path, deployment = provider.getModelPath("invoke", request.Model, key) rawResponse, latency, bifrostError = provider.completeRequest(ctx, jsonData, path, key) default: return nil, providerUtils.NewConfigurationError("unsupported image generation model type", providerName) }
🧹 Nitpick comments (1)
core/providers/bedrock/images.go (1)
62-63: Prefer pointer helper for width/height.Lines 62–63 use
&width/&height. The codebase convention is to use pointer helpers for consistency (e.g.,schemas.Ptr).♻️ Suggested tweak
- bedrockReq.ImageGenerationConfig.Width = &width - bedrockReq.ImageGenerationConfig.Height = &height + bedrockReq.ImageGenerationConfig.Width = schemas.Ptr(width) + bedrockReq.ImageGenerationConfig.Height = schemas.Ptr(height)Based on learnings, prefer pointer helpers over address-of operators.
e9ba497 to
a9761e2
Compare
There was a problem hiding this comment.
Actionable comments posted: 1
🤖 Fix all issues with AI agents
In `@core/providers/bedrock/bedrock.go`:
- Around line 1583-1605: After unmarshaling into BedrockImageGenerationResponse,
validate the response.Error field and treat non-empty/nonnull error payloads as
failures: if imageResp.Error (or the actual error field on
BedrockImageGenerationResponse) is present, return an enriched Bifrost error
using providerUtils.EnrichError/providerUtils.NewBifrostOperationError (similar
to the existing parse error handling) instead of continuing to map with
ToBifrostImageGenerationResponse; do this immediately after sonic.Unmarshal and
before calling ToBifrostImageGenerationResponse or setting
bifrostResponse.ExtraFields so raw request/response handling remains unchanged.
🧹 Nitpick comments (1)
docs/providers/supported-providers/bedrock.mdx (1)
522-585: Clarify the expectedsizeformat to reduce request errors.The converter expects
WIDTHxHEIGHT; documenting this avoids avoidable failures.✍️ Suggested doc tweak
| `size` | `imageGenerationConfig.width` & `imageGenerationConfig.height` | +**Note:** `size` must be in `WIDTHxHEIGHT` format (e.g., `1024x1024`).
a9761e2 to
f3aec9f
Compare
There was a problem hiding this comment.
Actionable comments posted: 1
🤖 Fix all issues with AI agents
In `@docs/providers/supported-providers/bedrock.mdx`:
- Around line 21-28: Update the provider attribution for Nova in the Bedrock
support table: replace the table row label "Nova (Anthropic)" with "Nova
(Amazon)" to match the model ID prefixes like "amazon.nova-canvas-v1" and
"us.amazon.nova-pro-v1" used elsewhere; also scan for any other occurrences of
the "Nova (Anthropic)" string in the same document and update them to "Nova
(Amazon)" to keep documentation consistent.
♻️ Duplicate comments (1)
core/providers/bedrock/bedrock.go (1)
1547-1553: Add provider-layer guard for nil/empty image generation requests.
request.Modelis accessed before validation; if this provider is invoked outside the HTTP handler flow, a nilrequest/request.Input(or empty prompt) can panic or build an invalid payload. Add a fast-fail guard nearproviderName.🛠️ Proposed fix
providerName := provider.GetProviderKey() + if request == nil || request.Input == nil || strings.TrimSpace(request.Input.Prompt) == "" { + return nil, providerUtils.NewBifrostOperationError("prompt cannot be empty", nil, providerName) + } if key.BedrockKeyConfig == nil { return nil, providerUtils.NewConfigurationError("bedrock key config is not provided", providerName) }Based on learnings, provider-layer guards should prevent nil input panics.
🧹 Nitpick comments (2)
docs/providers/supported-providers/bedrock.mdx (1)
528-537: DocumentcfgScalemapping for image generation.The implementation supports
cfgScaleviaextra_params, but it’s not shown in the parameter mapping table. Adding it helps users discover the supported tuning knob.✏️ Suggested edit
| `style` | `textToImageParams.style` | | `size` | `imageGenerationConfig.width` & `imageGenerationConfig.height` | +| `cfgScale` (extra_params.cfgScale) | `imageGenerationConfig.cfgScale` |core/providers/bedrock/images.go (1)
12-70: Trim size dimensions before parsing to tolerate whitespace.Inputs like
1024 x 1024will currently failAtoidue to embedded spaces. Trimming each dimension makes the parsing more forgiving without changing valid cases.♻️ Optional tweak
if request.Params.Size != nil && strings.TrimSpace(strings.ToLower(*request.Params.Size)) != "auto" { size := strings.Split(strings.TrimSpace(strings.ToLower(*request.Params.Size)), "x") if len(size) != 2 { return nil, fmt.Errorf("invalid size format: expected 'WIDTHxHEIGHT', got %q", *request.Params.Size) } - width, err := strconv.Atoi(size[0]) + widthStr := strings.TrimSpace(size[0]) + heightStr := strings.TrimSpace(size[1]) + width, err := strconv.Atoi(widthStr) if err != nil { return nil, fmt.Errorf("invalid width in size %q: %w", *request.Params.Size, err) } - height, err := strconv.Atoi(size[1]) + height, err := strconv.Atoi(heightStr) if err != nil { return nil, fmt.Errorf("invalid height in size %q: %w", *request.Params.Size, err) }
a82ce14 to
374c452
Compare
There was a problem hiding this comment.
Actionable comments posted: 1
🤖 Fix all issues with AI agents
In `@core/providers/bedrock/images.go`:
- Around line 45-63: Validate parsed width and height are positive before
assigning to bedrockReq.ImageGenerationConfig: after parsing size into width and
height from request.Params.Size and before calling schemas.Ptr(...), check that
width > 0 and height > 0 and return a descriptive error (e.g., "invalid size:
width and height must be positive integers") if not; ensure the error references
the original request.Params.Size value for clarity and keep the existing error
wrapping for parse failures.
♻️ Duplicate comments (1)
docs/providers/supported-providers/bedrock.mdx (1)
21-28: Align Nova attribution with the model IDs shown in this doc.The table labels Nova as Anthropic, but examples use
amazon.nova-*, which is inconsistent within the same document.Suggested edit
-| **Nova (Anthropic)** | ✅ | ✅ | ❌ | ❌ | ✅ | +| **Nova (Amazon)** | ✅ | ✅ | ❌ | ❌ | ✅ |
374c452 to
68b896c
Compare
68b896c to
9325a35
Compare
Merge activity
|

Summary
Add support for image generation in the Bedrock provider, enabling the use of Amazon's image generation models like Nova Canvas and Titan Image Generator.
Changes
ImageGenerationfunction in the Bedrock providernova-canvas-v1:0titan-image-generator-v1:0titan-image-generator-v2:0Type of change
Affected areas
How to test
Breaking changes
Related issues
N/A
Security considerations
This implementation maintains the existing authentication and authorization mechanisms for Bedrock API calls.
Checklist
docs/contributing/README.mdand followed the guidelines