Conversation
Contributor
Author
This stack of pull requests is managed by Graphite. Learn more about stacking. |
Radheshg04
changed the title
|
|
Contributor
Author
|
@coderabbitai full review |
1 similar comment
Collaborator
|
@coderabbitai full review |
Contributor
✅ Actions performedFull review triggered. |
Contributor
📝 WalkthroughSummary by CodeRabbit
WalkthroughBedrock invoke routing gained model-aware detection; Bedrock now supports Titan and Cohere-style embeddings (including multiple encoding types) and expanded image generation/edit/variation payloads. Converters and response parsers were added/changed to map between Bedrock and Bifrost shapes; tests and transport routing were updated accordingly. Changes
Sequence Diagram(s)sequenceDiagram
actor Client
participant Transport as "Bedrock Transport"
participant Detector as "DetectInvokeRequestType"
participant Converter as "Bedrock->Bifrost Converter"
participant Bifrost as "Bifrost Service"
participant Provider as "Provider API"
participant RespConv as "Response Converter"
Client->>Transport: POST /model/{modelId}/invoke (body + modelId)
Transport->>Detector: DetectInvokeRequestType(body, modelId)
Detector-->>Transport: RequestType (Embedding / Image / Responses / Text)
Transport->>Converter: Convert to Bifrost request (by type)
Converter-->>Bifrost: BifrostRequest (Embedding / Image / ...)
Bifrost->>Provider: Call provider API
Provider-->>Bifrost: BifrostResponse
Bifrost->>RespConv: ToBedrock*Response (preserve raw when applicable)
RespConv-->>Transport: Bedrock invoke response
Transport-->>Client: HTTP 200 + body
sequenceDiagram
actor Client
participant Transport as "Bedrock Transport"
participant Detector as "DetectInvokeRequestType"
participant Converter as "Image/Embed Converter"
participant Provider as "Bedrock Provider (Titan/Stability/Cohere)"
participant RespConv as "Image/Embed Resp Converter"
Client->>Transport: POST /invoke (image gen/edit or embed)
Transport->>Detector: DetectInvokeRequestType(body, modelId)
Detector-->>Transport: RequestType = Image / Embedding / ...
Transport->>Converter: Build BifrostImageRequest or BifrostEmbeddingRequest
Converter-->>Provider: Call provider API
alt Embedding (Cohere: embeddings_by_type)
Provider-->>RespConv: Raw typed embeddings (floats/int8/uint8/base64)
RespConv-->>Transport: BifrostEmbeddingResponse with ExtraFields.RawResponse
else Image (Titan/Stability)
Provider-->>RespConv: BifrostImageResponse (base64 images, seeds, finish_reasons)
RespConv-->>Transport: BedrockImageGenerationResponse
end
Transport-->>Client: HTTP 200 + mapped response
Estimated code review effort🎯 4 (Complex) | ⏱️ ~60 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 docstrings
🧪 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}
Contributor
There was a problem hiding this comment.
Actionable comments posted: 5
🧹 Nitpick comments (1)
core/providers/bedrock/invoke.go (1)
851-857: Consider adding nil check for B64JSON field.When iterating over
resp.Data, the code appendsd.B64JSONdirectly without checking if it's empty. While this may be acceptable if the upstream guarantees non-empty values, a defensive check would prevent appending empty strings.🛡️ Optional defensive check
// Default: Titan Image Generator v1/v2, Nova Canvas — reconstruct from Bifrost data result := &BedrockImageGenerationResponse{} for _, d := range resp.Data { + if d.B64JSON != "" { result.Images = append(result.Images, d.B64JSON) + } } return result, nil🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@core/providers/bedrock/invoke.go` around lines 851 - 857, The loop that builds a BedrockImageGenerationResponse appends d.B64JSON from resp.Data without checking for empty or nil values; update the loop that iterates resp.Data (in the function returning *BedrockImageGenerationResponse) to only append non-empty B64JSON strings (e.g., check that d.B64JSON != "" or is not nil) before calling result.Images = append(...), so result.Images never receives empty entries.
🤖 Prompt for all review comments with AI agents
Verify each finding against the current code and only fix it if needed.
Inline comments:
In `@core/providers/bedrock/embedding.go`:
- Around line 104-118: The code mutates the caller's ExtraParams map by aliasing
it with extra and calling delete, making the converter stateful; fix by cloning
bifrostReq.Params.ExtraParams into a new local map before stripping recognized
Cohere fields so deletions don't affect the original request, then read from the
clone to set req.InputType and req.Truncate (and the similar block at lines
~124-153), ensuring the converter function remains a pure transformation (no
side effects).
- Around line 119-123: Detect and reject requests that specify multiple
embedding encodings: when reading extra["embedding_types"] into ss and assigning
req.EmbeddingTypes, if len(ss) > 1 return an error (or otherwise fail fast)
instead of accepting multiple entries; do the same guard in the other parsing
block that sets req.EmbeddingTypes (the second occurrence referenced alongside
ToBifrostEmbeddingResponse()), so callers cannot request multi-encoding
responses until ToBifrostEmbeddingResponse() is updated to preserve multiple
encodings.
In `@core/providers/bedrock/invoke.go`:
- Around line 886-891: The code returns a bedrockEmbeddingResp using
resp.Data[0].Embedding.EmbeddingArray without checking for nil; add a nil check
for resp.Data and resp.Data[0].Embedding before accessing EmbeddingArray in the
function that constructs the bedrockEmbeddingResp, and if nil return a clear
error (or handle gracefully) instead of dereferencing; ensure the check
references resp.Data[0].Embedding and then proceed to use EmbeddingArray only
when non-nil.
In `@tests/integrations/python/tests/test_bedrock.py`:
- Around line 48-60: The docstring and test function numbering are out of sync:
the docstring lists item 37 as "Stability AI image inpaint via invoke
(image+mask)" but the test function is named
test_38_invoke_stability_ai_inpaint, creating a gap; fix by either inserting a
new test with the symbol/name test_37_invoke_stability_ai_inpaint matching the
docstring or renumber the subsequent test functions (e.g., rename
test_38_invoke_stability_ai_inpaint → test_37_invoke_stability_ai_inpaint and
adjust any later test names/docstring entries) so the docstring list and
function names are consistent.
In `@transports/bifrost-http/integrations/bedrock.go`:
- Around line 199-245: RequestConverter handles ImageEditRequest and
ImageVariationRequest but the route configuration lacks matching response
converters; add ImageEditResponseConverter and ImageVariationResponseConverter
to the same route config that contains RequestConverter so responses are
converted back to Bedrock format. Implement each converter to handle the
corresponding Bifrost response -> Bedrock response transformation (mirror how
EmbeddingResponse/ImageGenerationResponse/ResponsesRequest are handled), wiring
them alongside the existing RequestConverter, ImageGenerationResponseConverter,
and EmbeddingResponseConverter so ImageEditRequest and ImageVariationRequest use
their proper converters.
---
Nitpick comments:
In `@core/providers/bedrock/invoke.go`:
- Around line 851-857: The loop that builds a BedrockImageGenerationResponse
appends d.B64JSON from resp.Data without checking for empty or nil values;
update the loop that iterates resp.Data (in the function returning
*BedrockImageGenerationResponse) to only append non-empty B64JSON strings (e.g.,
check that d.B64JSON != "" or is not nil) before calling result.Images =
append(...), so result.Images never receives empty entries.
🪄 Autofix (Beta)
Fix all unresolved CodeRabbit comments on this PR:
- Push a commit to this branch (recommended)
- Create a new PR with the fixes
ℹ️ Review info
⚙️ Run configuration
Configuration used: Organization UI
Review profile: CHILL
Plan: Pro
Run ID: a3b662e9-1cb8-455a-8945-587167f5bfbe
📒 Files selected for processing (10)
core/providers/bedrock/bedrock.gocore/providers/bedrock/embedding.gocore/providers/bedrock/images.gocore/providers/bedrock/invoke.gocore/providers/bedrock/types.gocore/schemas/images.gotests/integrations/python/config.ymltests/integrations/python/tests/test_bedrock.pytests/integrations/python/tests/utils/common.pytransports/bifrost-http/integrations/bedrock.go
👮 Files not reviewed due to content moderation or server errors (4)
- core/providers/bedrock/bedrock.go
- core/schemas/images.go
- core/providers/bedrock/images.go
- tests/integrations/python/config.yml
![greptile-apps[bot]](https://avatars.githubusercontent.com/in/867647?s=60&v=4){kind=small}
Contributor
There was a problem hiding this comment.
Your free trial has ended. If you'd like to continue receiving code reviews, you can add a payment method here.
Contributor
There was a problem hiding this comment.
Your free trial has ended. If you'd like to continue receiving code reviews, you can add a payment method here.
Contributor
Confidence Score: 5/5Safe to merge — all remaining findings are P2 or lower and do not affect the primary invoke endpoint code paths. Both open findings are P2: one is dead code in UnmarshalJSON (bypassed entirely on the invoke path via RawResponse), and one is a minor style inconsistency in model-ID handling that works correctly in practice. No P0 or P1 issues remain. core/schemas/embedding.go — unreachable int8/int32 unmarshal branches should be fixed for correct JSON round-trip semantics. Important Files Changed
Reviews (8): Last reviewed commit: "feat: add embeddings, image gen, edit an..." | Re-trigger Greptile |
Radheshg04
force-pushed
the
03-24-feat_add_embeddings_image_gen_edit_and_variation_to_bedrock_integration
branch
from
cc38138 to
b360580
Compare
Contributor
There was a problem hiding this comment.
Actionable comments posted: 5
Caution
Some comments are outside the diff and can’t be posted inline due to platform limitations.
⚠️ Outside diff range comments (1)
core/providers/bedrock/bedrock.go (1)
77-87:⚠️ Potential issue | 🟠 MajorBedrock must enforce HTTP/2 unconditionally.
Line 86 makes
ForceAttemptHTTP2configurable viaconfig.NetworkConfig.EnforceHTTP2, but the guideline explicitly requires Bedrock's HTTP transport to hardcodeForceAttemptHTTP2: true. This allows HTTP/2 to be disabled at runtime, which violates the architectural requirement for Bedrock provider support.🔧 Suggested fix
transport := &http.Transport{ Proxy: http.ProxyFromEnvironment, MaxConnsPerHost: config.NetworkConfig.MaxConnsPerHost, MaxIdleConns: schemas.DefaultMaxIdleConnsPerHost, MaxIdleConnsPerHost: schemas.DefaultMaxIdleConnsPerHost, IdleConnTimeout: 30 * time.Second, TLSHandshakeTimeout: 10 * time.Second, ResponseHeaderTimeout: requestTimeout, ExpectContinueTimeout: 1 * time.Second, - ForceAttemptHTTP2: config.NetworkConfig.EnforceHTTP2, + ForceAttemptHTTP2: true, }Note: Related tests in
transport_test.goverify optional HTTP/2 behavior and will require updates when this is corrected.🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@core/providers/bedrock/bedrock.go` around lines 77 - 87, The HTTP transport in the Bedrock provider currently sets ForceAttemptHTTP2 from config.NetworkConfig.EnforceHTTP2, allowing HTTP/2 to be disabled at runtime; change the transport construction (the transport variable where ForceAttemptHTTP2 is set) to hardcode ForceAttemptHTTP2: true and remove or ignore config.NetworkConfig.EnforceHTTP2 for this field (leave other network config fields as-is); update any tests or references that assert optional HTTP/2 behavior to reflect that ForceAttemptHTTP2 is always true.
♻️ Duplicate comments (1)
core/providers/bedrock/invoke.go (1)
878-901:⚠️ Potential issue | 🔴 CriticalDon't infer Cohere-vs-Titan shape from
len(resp.Data), and guard nil embeddings.A single-text Cohere request still expects the Cohere
embeddingsschema, but this code falls back to Titan'sembeddingshape whenever only one vector comes back. The same block also dereferencesd.Embedding/resp.Data[0].Embeddingwithout a nil check, so a partial provider response can panic the converter.Proposed fix
if len(resp.Data) == 0 { return &BedrockInvokeEmbeddingResp{InputTextTokenCount: tokenCount}, nil } - // Multi-embedding: reconstruct Cohere-format response - if len(resp.Data) > 1 { + model := resp.Model + if resp.ExtraFields.ModelRequested != "" { + model = resp.ExtraFields.ModelRequested + } + isCohere := strings.Contains(strings.ToLower(model), "cohere") + + if isCohere { floats := make([][]float32, 0, len(resp.Data)) for _, d := range resp.Data { + if d.Embedding == nil { + return nil, fmt.Errorf("embedding data is missing") + } floats = append(floats, d.Embedding.EmbeddingArray) } return &BedrockInvokeCohereEmbeddingResp{ Embeddings: floats, ResponseType: "embeddings_floats", }, nil } // Single-embedding: Titan format - if resp.Data[0].Embedding.EmbeddingArray == nil { + if resp.Data[0].Embedding == nil || resp.Data[0].Embedding.EmbeddingArray == nil { return &BedrockInvokeEmbeddingResp{InputTextTokenCount: tokenCount}, nil }🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@core/providers/bedrock/invoke.go` around lines 878 - 901, Do not branch on len(resp.Data) to choose Cohere vs Titan shapes and add nil guards: inspect the actual fields present in resp.Data instead of count (e.g., check whether each data item has a non-nil d.Embedding and non-nil d.Embedding.EmbeddingArray or whether a Cohere-style embeddings collection exists) and base the response type on those field presences; ensure every access to d.Embedding and resp.Data[0].Embedding is nil-checked before dereference to avoid panics, and return an appropriate BedrockInvokeCohereEmbeddingResp when Cohere-shaped fields are present, otherwise return BedrockInvokeEmbeddingResp (or empty token response) when embeddings are absent or nil.
🧹 Nitpick comments (1)
tests/integrations/python/tests/test_bedrock.py (1)
2088-2097: Add a native Titan text-generation invoke case to this suite.The new detector now branches on
inputText, but these/invoketests only coverinputTextfor embeddings. A smoke test for{inputText, textGenerationConfig}would lock in the completion path and catch this class of routing regression.🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@tests/integrations/python/tests/test_bedrock.py` around lines 2088 - 2097, Add a smoke test inside the TestBedrockInvokeEndpoint test suite that exercises the `/invoke` text-generation path by sending a native Titan-style payload containing both inputText and textGenerationConfig; specifically, add a test method (e.g., test_invoke_titan_text_generation) that constructs an invoke request with inputText and a minimal valid textGenerationConfig and asserts the call routes to the completion path and returns a successful text-generation-like response, ensuring the detector branch on inputText is covered.
🤖 Prompt for all review comments with AI agents
Verify each finding against the current code and only fix it if needed.
Inline comments:
In `@core/providers/bedrock/invoke.go`:
- Around line 440-445: The ImageGenerationParameters builder currently only sets
params.N from Titan's nested imageGenerationConfig, causing top-level n (e.g.
for prompt-only models like openai/gpt-image-1) to be ignored; update the
construction of params (and the similar block in the 456-484 region) to also
copy any top-level r.N into params.N when present (e.g. if r.N > 0) and only
override with imageGenerationConfig.N when that nested value exists, ensuring
prompt-only routes that set r.N are respected despite n being a known field that
doesn't fall into ExtraParams.
- Around line 498-501: The edit conversion currently creates
schemas.ImageEditParameters only with NegativePrompt and Seed, losing nested
imageGenerationConfig fields; update the code that builds params (where
ImageEditParameters is created) to extract and map
r.Params.imageGenerationConfig (or request.Params.imageGenerationConfig) into
the ImageEditParameters struct — include Seed, NumberOfImages (or Count), Size,
Quality, CfgScale (cfgScale) and any other generation settings used by
core/providers/bedrock/images.go when reconstructing native edit payloads; apply
this mapping consistently across all Titan/Nova edit branches so the round-trip
preserves seed, numberOfImages, size, quality and cfgScale.
- Around line 151-154: The current check in invoke.go that returns
schemas.EmbeddingRequest when sonic.Get(body, "inputText") exists incorrectly
routes Bedrock text-generation payloads (which also use top-level inputText) to
the embeddings branch; update the condition in the Titan embedding detection
(the sonic.Get(body, "inputText") / node.Exists() check) to only treat the
request as an embedding when inputText exists AND there is no text-generation
indicator (e.g. ensure sonic.Get(body, "textGenerationConfig") does NOT exist)
or by checking for an explicit embedding-specific field; modify the conditional
so embeddings are returned only when the request shape clearly indicates
embeddings, leaving requests with textGenerationConfig to continue down the
completions/path.
- Around line 621-650: The variation prompt is being placed into
extraParams["prompt"] which the Bedrock provider won't read; move the text
prompt into the ImageVariationInput's Input.Prompt field instead. In the
ImageVariation construction where you currently populate extraParams (using
r.ImageVariationParams.Text), set the prompt on the ImageVariationInput
(schemas.ImageVariationInput.Input.Prompt or the appropriate Prompt field on
that Input struct) when r.ImageVariationParams.Text != nil, and remove or avoid
setting extraParams["prompt"]; leave negativeText and similarityStrength in
extraParams if the provider expects them. This ensures the provider code that
reads request.Input.Prompt (in core/providers/bedrock/images.go) receives the
variation prompt.
In `@core/providers/bedrock/types.go`:
- Around line 1096-1098: The invoke path currently assumes a fixed [][]float32
response keyed by "embeddings_floats" which breaks when EmbeddingTypes is used
and Bedrock returns "embeddings_by_type"; update the invoke handling that
unmarshals the Cohere response (the logic that reads "embeddings_floats") to
instead accept a polymorphic payload (use json.RawMessage like the existing
Cohere embedding path in embedding.go lines 186-219), detect whether the
response contains "embeddings_by_type" vs "embeddings_floats", and when
"embeddings_by_type" is present unmarshal it into a map[string]json.RawMessage
and decode each typed array according to the requested EmbeddingTypes rather
than forcing [][]float32. Ensure any downstream consumers expect the parsed
typed arrays instead of only float32 arrays.
---
Outside diff comments:
In `@core/providers/bedrock/bedrock.go`:
- Around line 77-87: The HTTP transport in the Bedrock provider currently sets
ForceAttemptHTTP2 from config.NetworkConfig.EnforceHTTP2, allowing HTTP/2 to be
disabled at runtime; change the transport construction (the transport variable
where ForceAttemptHTTP2 is set) to hardcode ForceAttemptHTTP2: true and remove
or ignore config.NetworkConfig.EnforceHTTP2 for this field (leave other network
config fields as-is); update any tests or references that assert optional HTTP/2
behavior to reflect that ForceAttemptHTTP2 is always true.
---
Duplicate comments:
In `@core/providers/bedrock/invoke.go`:
- Around line 878-901: Do not branch on len(resp.Data) to choose Cohere vs Titan
shapes and add nil guards: inspect the actual fields present in resp.Data
instead of count (e.g., check whether each data item has a non-nil d.Embedding
and non-nil d.Embedding.EmbeddingArray or whether a Cohere-style embeddings
collection exists) and base the response type on those field presences; ensure
every access to d.Embedding and resp.Data[0].Embedding is nil-checked before
dereference to avoid panics, and return an appropriate
BedrockInvokeCohereEmbeddingResp when Cohere-shaped fields are present,
otherwise return BedrockInvokeEmbeddingResp (or empty token response) when
embeddings are absent or nil.
---
Nitpick comments:
In `@tests/integrations/python/tests/test_bedrock.py`:
- Around line 2088-2097: Add a smoke test inside the TestBedrockInvokeEndpoint
test suite that exercises the `/invoke` text-generation path by sending a native
Titan-style payload containing both inputText and textGenerationConfig;
specifically, add a test method (e.g., test_invoke_titan_text_generation) that
constructs an invoke request with inputText and a minimal valid
textGenerationConfig and asserts the call routes to the completion path and
returns a successful text-generation-like response, ensuring the detector branch
on inputText is covered.
🪄 Autofix (Beta)
Fix all unresolved CodeRabbit comments on this PR:
- Push a commit to this branch (recommended)
- Create a new PR with the fixes
ℹ️ Review info
⚙️ Run configuration
Configuration used: Organization UI
Review profile: CHILL
Plan: Pro
Run ID: 35847e4b-a4af-4804-a9a6-2a4ed4fdfb00
📒 Files selected for processing (10)
core/providers/bedrock/bedrock.gocore/providers/bedrock/embedding.gocore/providers/bedrock/images.gocore/providers/bedrock/invoke.gocore/providers/bedrock/types.gocore/schemas/images.gotests/integrations/python/config.ymltests/integrations/python/tests/test_bedrock.pytests/integrations/python/tests/utils/common.pytransports/bifrost-http/integrations/bedrock.go
✅ Files skipped from review due to trivial changes (1)
- core/providers/bedrock/embedding.go
🚧 Files skipped from review as they are similar to previous changes (3)
- core/providers/bedrock/images.go
- tests/integrations/python/tests/utils/common.py
- tests/integrations/python/config.yml
Radheshg04
force-pushed
the
03-24-feat_add_embeddings_image_gen_edit_and_variation_to_bedrock_integration
branch
from
b360580 to
be10d3f
Compare
Contributor
There was a problem hiding this comment.
Actionable comments posted: 6
Caution
Some comments are outside the diff and can’t be posted inline due to platform limitations.
⚠️ Outside diff range comments (1)
transports/bifrost-http/integrations/bedrock.go (1)
123-147:⚠️ Potential issue | 🟠 MajorReject non-text request types on
/invoke-with-response-stream.
GetHTTPRequestTypecan now return embeddings/images, but this route still only constructsResponsesRequestorTextCompletionRequestand only has stream converters for those two shapes. That leaves the request type, builtBifrostRequest, andStreamConfigout of sync for anything else.🛠️ Fail fast for unsupported streaming request types
RequestConverter: func(ctx *schemas.BifrostContext, req interface{}) (*schemas.BifrostRequest, error) { - if invokeReq, ok := req.(*bedrock.BedrockInvokeRequest); ok { - invokeReq.Stream = true - if invokeReq.IsMessagesRequest() { - // Messages-based → Responses path (streaming) - converseReq := invokeReq.ToBedrockConverseRequest() - responsesReq, err := converseReq.ToBifrostResponsesRequest(ctx) - if err != nil { - return nil, fmt.Errorf("failed to convert invoke messages stream request: %w", err) - } - return &schemas.BifrostRequest{ResponsesRequest: responsesReq}, nil - } - // Prompt-based → Text Completion path (streaming) - return &schemas.BifrostRequest{ - TextCompletionRequest: invokeReq.ToBifrostTextCompletionRequest(ctx), - }, nil - } - return nil, errors.New("invalid request type") + invokeReq, ok := req.(*bedrock.BedrockInvokeRequest) + if !ok { + return nil, errors.New("invalid request type") + } + invokeReq.Stream = true + + requestType, _ := ctx.Value(schemas.BifrostContextKeyHTTPRequestType).(schemas.RequestType) + switch requestType { + case schemas.ResponsesRequest: + converseReq := invokeReq.ToBedrockConverseRequest() + responsesReq, err := converseReq.ToBifrostResponsesRequest(ctx) + if err != nil { + return nil, fmt.Errorf("failed to convert invoke messages stream request: %w", err) + } + return &schemas.BifrostRequest{ResponsesRequest: responsesReq}, nil + case schemas.TextCompletionRequest: + return &schemas.BifrostRequest{ + TextCompletionRequest: invokeReq.ToBifrostTextCompletionRequest(ctx), + }, nil + default: + return nil, fmt.Errorf("%s is not supported on invoke-with-response-stream", requestType) + } }🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@transports/bifrost-http/integrations/bedrock.go` around lines 123 - 147, The /invoke-with-response-stream path must reject non-streaming/text request shapes: in RequestConverter, after casting to *bedrock.BedrockInvokeRequest (and before setting Stream = true or converting), validate the invokeReq request type (using existing helpers like IsMessagesRequest or a prompt/text check) and return a clear error for unsupported types (e.g., embeddings, images) instead of attempting to build a BifrostRequest; ensure the error path prevents constructing ResponsesRequest or TextCompletionRequest so GetHTTPRequestType, the built BifrostRequest and StreamConfig remain in sync (refer to GetHTTPRequestType, bedrock.BedrockInvokeRequest, IsMessagesRequest, ToBedrockConverseRequest, ToBifrostResponsesRequest, and ToBifrostTextCompletionRequest).
♻️ Duplicate comments (1)
core/providers/bedrock/invoke.go (1)
904-923:⚠️ Potential issue | 🔴 CriticalGuard both embedding reconstruction paths against nil
Embeddingpointers.The multi-item loop dereferences
d.Embeddingunconditionally, and the single-item branch still readsresp.Data[0].Embedding.EmbeddingArraybefore provingEmbeddingis non-nil. A partial provider response will panic here instead of returning a clear error.🛡️ Suggested fix
// Multi-embedding: reconstruct Cohere-format response if len(resp.Data) > 1 { floats := make([][]float32, 0, len(resp.Data)) - for _, d := range resp.Data { + for i, d := range resp.Data { + if d.Embedding == nil || d.Embedding.EmbeddingArray == nil { + return nil, fmt.Errorf("embedding data[%d] is missing", i) + } floats = append(floats, d.Embedding.EmbeddingArray) } return &BedrockInvokeCohereEmbeddingResp{ Embeddings: floats, ResponseType: "embeddings_floats", @@ // Single-embedding: Titan format - if resp.Data[0].Embedding.EmbeddingArray == nil { + if resp.Data[0].Embedding == nil || resp.Data[0].Embedding.EmbeddingArray == nil { return &BedrockInvokeEmbeddingResp{InputTextTokenCount: tokenCount}, nil }🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@core/providers/bedrock/invoke.go` around lines 904 - 923, The code dereferences resp.Data[].Embedding without nil checks causing panics on partial responses; update the multi-embedding loop (where you build floats for BedrockInvokeCohereEmbeddingResp) to check each d.Embedding for nil before accessing d.Embedding.EmbeddingArray (either skip nil entries or return a clear error), and in the single-embedding branch verify resp.Data[0].Embedding != nil before reading Embedding.EmbeddingArray and return an error (instead of panicking) so BedrockInvokeEmbeddingResp construction (and tokenCount handling) only happens when the Embedding pointer is non-nil.
🧹 Nitpick comments (1)
tests/integrations/python/tests/test_bedrock.py (1)
2359-2377: Use the configured OpenAI image model instead of hardcodingopenai/gpt-image-1.Hardcoding the model ID bypasses
providers.openai.image_generation, so this test can drift from the repo's fixtures/config and fail in environments that alias the OpenAI image model differently.♻️ Suggested refactor
`@skip_if_no_api_key`("openai") def test_39_invoke_openai_image_generation(self, bedrock_client): """Test Case 39: OpenAI gpt-image-1 via Bedrock invoke endpoint (cross-provider)""" print("\n=== Test 39: OpenAI gpt-image-1 via invoke ===") + model_id = format_provider_model("openai", get_model("openai", "image_generation")) body = { "prompt": "a gecko resting on a tropical leaf", "n": 1, "quality": "low", @@ response = bedrock_client.invoke_model( - modelId="openai/gpt-image-1", + modelId=model_id, contentType="application/json", accept="application/json", body=json.dumps(body), )Based on learnings: In tests under
tests/integrations/python, prefer using the OpenAI image generation modelgpt-image-1via the config keyproviders.openai.image_generationfor image-generation scenarios.🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@tests/integrations/python/tests/test_bedrock.py` around lines 2359 - 2377, The test test_39_invoke_openai_image_generation hardcodes modelId="openai/gpt-image-1", bypassing the repo config; update the test to read the configured OpenAI image model (providers.openai.image_generation) and pass that value into bedrock_client.invoke_model's modelId parameter so the test uses the configured alias instead of a literal "openai/gpt-image-1".
🤖 Prompt for all review comments with AI agents
Verify each finding against the current code and only fix it if needed.
Inline comments:
In `@core/providers/bedrock/images.go`:
- Around line 42-48: The current isPromptOnlyImageGenerationModel function uses
strings.Contains(m, "image") which is too broad; replace that logic with an
explicit allowlist of prompt-only families (e.g., "imagen", "vertex-imagen",
"google-imagen", or other confirmed prompt-only family substrings) and return
true only if the lowercased model string matches one of those allowlisted
substrings (use strings.Contains or prefix checks as appropriate). Update
isPromptOnlyImageGenerationModel to iterate over the allowlist and only classify
models as prompt-only when they match an allowlisted family, thereby excluding
"stability.stable-image-*" and "amazon.titan-image-*" which should not be
treated as prompt-only.
- Around line 162-172: The helper ToStabilityAIImageGenerationResponse reads
promoted fields response.FinishReasons and response.Seeds which can panic if the
embedded *ImageGenerationResponseParameters is nil; add a nil-check for
response.Parameters (or the embedded pointer) before reading those fields and
only assign result.FinishReasons/result.Seeds when non-nil, and update the
forward converter ToBifrostImageGenerationResponse to ensure it initializes and
populates the embedded ImageGenerationResponseParameters on Bedrock->Bifrost
conversion so metadata survives the round trip; reference the functions
ToStabilityAIImageGenerationResponse and ToBifrostImageGenerationResponse and
the fields FinishReasons and Seeds when making these changes.
In `@core/providers/bedrock/invoke.go`:
- Around line 576-622: The code currently assigns r.ExtraParams to
params.ExtraParams at the end, overwriting any synthesized entries like the
earlier cfgScale; change this to merge r.ExtraParams into params.ExtraParams
instead: ensure params.ExtraParams is non-nil (initialize if needed), then
iterate over r.ExtraParams and copy each key/value into params.ExtraParams so
existing synthesized keys (e.g., "cfgScale") are preserved; refer to
params.ExtraParams, r.ExtraParams, and the earlier cfgScale assignment when
making the change.
- Around line 160-165: The detection branch that currently only checks
sonic.Get(body, "texts") should be expanded to also check for non-empty Cohere
embedding fields "inputs" and "images" so these payloads return
schemas.EmbeddingRequest; update the logic in the same block (the code using
sonic.Get(body, "...") and node.Raw() checks) to test node existence and that
raw != "null" && raw != "[]" for "texts", "inputs", and "images" (treat any
non-empty value for those keys as an embedding request) and return
schemas.EmbeddingRequest when any of them is present and non-empty.
In `@tests/integrations/python/tests/test_bedrock.py`:
- Around line 2067-2076: The helper _assert_invoke_images currently assumes
response_body.get("images") is iterable and may pass for strings or dicts;
update it to first assert that images is a list (e.g., isinstance(images, list))
and raise a clear assertion if not, including the actual type and response keys,
before checking length and iterating; keep the subsequent length check
(len(images) >= min_images) and per-item string non-empty assertions intact so
only proper list-shaped image arrays pass.
- Around line 2200-2206: The test currently only checks for the presence of
"embedding" or "embeddings" in out; instead, inspect which key is present (use
the existing has_embedding/out checks) and then validate the payload shape: if
"embedding" is present assert it's a non-empty list of numeric types and record
its length as dim; if "embeddings" is present assert it's a non-empty list of
lists, that all inner lists are non-empty numeric vectors and that every inner
vector has the same length; fail with clear messages if types/lengths mismatch
so the Cohere invoke is validated for vector count and numeric shape.
---
Outside diff comments:
In `@transports/bifrost-http/integrations/bedrock.go`:
- Around line 123-147: The /invoke-with-response-stream path must reject
non-streaming/text request shapes: in RequestConverter, after casting to
*bedrock.BedrockInvokeRequest (and before setting Stream = true or converting),
validate the invokeReq request type (using existing helpers like
IsMessagesRequest or a prompt/text check) and return a clear error for
unsupported types (e.g., embeddings, images) instead of attempting to build a
BifrostRequest; ensure the error path prevents constructing ResponsesRequest or
TextCompletionRequest so GetHTTPRequestType, the built BifrostRequest and
StreamConfig remain in sync (refer to GetHTTPRequestType,
bedrock.BedrockInvokeRequest, IsMessagesRequest, ToBedrockConverseRequest,
ToBifrostResponsesRequest, and ToBifrostTextCompletionRequest).
---
Duplicate comments:
In `@core/providers/bedrock/invoke.go`:
- Around line 904-923: The code dereferences resp.Data[].Embedding without nil
checks causing panics on partial responses; update the multi-embedding loop
(where you build floats for BedrockInvokeCohereEmbeddingResp) to check each
d.Embedding for nil before accessing d.Embedding.EmbeddingArray (either skip nil
entries or return a clear error), and in the single-embedding branch verify
resp.Data[0].Embedding != nil before reading Embedding.EmbeddingArray and return
an error (instead of panicking) so BedrockInvokeEmbeddingResp construction (and
tokenCount handling) only happens when the Embedding pointer is non-nil.
---
Nitpick comments:
In `@tests/integrations/python/tests/test_bedrock.py`:
- Around line 2359-2377: The test test_39_invoke_openai_image_generation
hardcodes modelId="openai/gpt-image-1", bypassing the repo config; update the
test to read the configured OpenAI image model
(providers.openai.image_generation) and pass that value into
bedrock_client.invoke_model's modelId parameter so the test uses the configured
alias instead of a literal "openai/gpt-image-1".
🪄 Autofix (Beta)
Fix all unresolved CodeRabbit comments on this PR:
- Push a commit to this branch (recommended)
- Create a new PR with the fixes
ℹ️ Review info
⚙️ Run configuration
Configuration used: Organization UI
Review profile: CHILL
Plan: Pro
Run ID: 5c616370-f65c-4447-b051-acb0e344d225
📒 Files selected for processing (10)
core/providers/bedrock/bedrock.gocore/providers/bedrock/embedding.gocore/providers/bedrock/images.gocore/providers/bedrock/invoke.gocore/providers/bedrock/types.gocore/schemas/images.gotests/integrations/python/config.ymltests/integrations/python/tests/test_bedrock.pytests/integrations/python/tests/utils/common.pytransports/bifrost-http/integrations/bedrock.go
✅ Files skipped from review due to trivial changes (1)
- core/providers/bedrock/embedding.go
🚧 Files skipped from review as they are similar to previous changes (2)
- tests/integrations/python/config.yml
- tests/integrations/python/tests/utils/common.py
Radheshg04
force-pushed
the
03-24-feat_add_embeddings_image_gen_edit_and_variation_to_bedrock_integration
branch
from
be10d3f to
a70a006
Compare
Contributor
There was a problem hiding this comment.
Actionable comments posted: 6
♻️ Duplicate comments (2)
core/providers/bedrock/images.go (2)
159-175:⚠️ Potential issue | 🟠 MajorPreserve
finish_reasons/seedson the Bedrock→Bifrost hop too.This helper now reads the embedded params safely, but
ToBifrostImageGenerationResponsebelow still never initializesImageGenerationResponseParameters. Native Stability responses therefore losefinish_reasonsandseedsbefore this converter runs, so the invoke round-trip still drops the metadata.🛠️ Suggested change
func ToStabilityAIImageGenerationResponse(response *schemas.BifrostImageGenerationResponse) (*BedrockImageGenerationResponse, error) { if response == nil { return nil, fmt.Errorf("response is nil") } result := &BedrockImageGenerationResponse{} for _, d := range response.Data { result.Images = append(result.Images, d.B64JSON) } if response.ImageGenerationResponseParameters != nil { - result.FinishReasons = response.ImageGenerationResponseParameters.FinishReasons - result.Seeds = response.ImageGenerationResponseParameters.Seeds + result.FinishReasons = append([]*string(nil), response.ImageGenerationResponseParameters.FinishReasons...) + result.Seeds = append([]int(nil), response.ImageGenerationResponseParameters.Seeds...) } return result, nil }func ToBifrostImageGenerationResponse(response *BedrockImageGenerationResponse) *schemas.BifrostImageGenerationResponse { if response == nil { return nil } bifrostResponse := &schemas.BifrostImageGenerationResponse{} + if len(response.FinishReasons) > 0 || len(response.Seeds) > 0 { + bifrostResponse.ImageGenerationResponseParameters = &schemas.ImageGenerationResponseParameters{ + FinishReasons: append([]*string(nil), response.FinishReasons...), + Seeds: append([]int(nil), response.Seeds...), + } + } for index, image := range response.Images { bifrostResponse.Data = append(bifrostResponse.Data, schemas.ImageData{ B64JSON: image, Index: index,🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@core/providers/bedrock/images.go` around lines 159 - 175, The ToBifrostImageGenerationResponse path never initializes ImageGenerationResponseParameters, so finish_reasons and seeds are dropped; update the ToBifrostImageGenerationResponse implementation to allocate and populate response.ImageGenerationResponseParameters (setting FinishReasons and Seeds from the native Bedrock/Bifrost input) before returning, ensuring the metadata flows through the Bedrock→Bifrost→ToStabilityAIImageGenerationResponse round-trip; touch the ToBifrostImageGenerationResponse function and the ImageGenerationResponseParameters struct population logic so it mirrors how ToStabilityAIImageGenerationResponse reads those fields.
42-48:⚠️ Potential issue | 🟠 MajorNarrow the prompt-only model allowlist.
The comment says Stability AI is excluded here, but
strings.Contains(m, "image")still returns true forstability.stable-image-*(and also sweeps in Titan image families). That keeps this helper broader than its contract and can still misroute task-backed models as prompt-only generation.🛠️ Suggested change
func isPromptOnlyImageGenerationModel(model string) bool { m := strings.ToLower(model) - return strings.Contains(m, "image") + for _, family := range []string{"gpt-image", "imagen"} { + if strings.Contains(m, family) { + return true + } + } + return false }🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@core/providers/bedrock/images.go` around lines 42 - 48, The isPromptOnlyImageGenerationModel helper is too broad because it returns true for any model name containing "image" (catching stability.stable-image-* and Titan image families); update isPromptOnlyImageGenerationModel to narrow the allowlist by first excluding known non‑prompt-only vendors (e.g., if the lowercased model string contains "stability." or "titan" return false) and then only return true for explicit prompt-only identifiers (e.g., match known prompt-only tokens like "imagen", "vertex-imagen", or an exact allowlist of model names/patterns) so task-backed image models are not misclassified. Ensure changes are made inside isPromptOnlyImageGenerationModel to keep the contract and adjust any tests/places that rely on its result.
🧹 Nitpick comments (1)
core/providers/bedrock/invoke.go (1)
520-580: Normalize the new error strings to lowercase Go style.Several new messages embed camelCase identifiers or title-cased provider names (
inPaintingParams,taskType,Stability AI,IMAGE_VARIATION). Please normalize them before merge so error text stays consistent across providers. As per coding guidelines, "Error strings should be lowercase with no trailing punctuation (Go convention)."Also applies to: 602-616, 643-648
🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@core/providers/bedrock/invoke.go` around lines 520 - 580, Normalize the error strings in the image-edit switch cases by making them lowercase, removing trailing punctuation, and replacing camelCase/title-cased identifiers with plain lowercase words; update the fmt.Errorf messages that reference r.InPaintingParams, r.OutPaintingParams, r.BackgroundRemovalParams, and the default unsupported taskType error to use e.g. "inpainting params required for inpainting task", "failed to decode inpainting image: %w", "outpainting params required for outpainting task", "failed to decode outpainting image: %w", "background removal params required for background removal task", "failed to decode background removal image: %w", and "unsupported task type for image edit: %s" (also apply the same normalization to the other related messages at the other mentioned ranges 602-616 and 643-648).
🤖 Prompt for all review comments with AI agents
Verify each finding against the current code and only fix it if needed.
Inline comments:
In `@core/providers/bedrock/embedding.go`:
- Around line 122-126: The code currently assigns extra["embedding_types"] into
req.EmbeddingTypes without validation (see req.EmbeddingTypes usage) which can
lead to dropped embeddings if the response parser only handles certain
encodings; update the logic that parses and assigns embedding_types to either
(1) validate each value against the set of encodings the response handler can
round-trip (e.g., "float", "int8", "uint8", "binary", "ubinary") and return an
error for unsupported values, or (2) extend the Bedrock response parsing code to
preserve and decode all supported encoding variants before exposing them — make
the change in the embedding request parsing (where req.EmbeddingTypes is set)
and in the Bedrock response handler so both sides agree on the allowed/handled
encodings.
In `@core/providers/bedrock/invoke.go`:
- Around line 518-540: The inpainting branch (TaskTypeInpainting) currently
decodes image/mask and sets params.Type, params.NegativePrompt and params.Mask
but fails to forward Titan edit fields like maskPrompt, returnMask, and
outPaintingMode; modify the TaskTypeInpainting handling to copy
r.InPaintingParams.MaskPrompt, r.InPaintingParams.ReturnMask, and
r.InPaintingParams.OutPaintingMode (or their exact names) into the outgoing
params (or into params.ExtraParams if the schema requires) so they are preserved
during the bedrock→Bifrost conversion; apply the same forwarding logic to the
analogous branch around lines 542-564 so both edit paths preserve all known
Titan edit fields.
- Around line 687-694: The ImageGenerationConfig handling in the invoke path
only maps NumberOfImages and size, dropping seed, quality, and cfgScale; update
the block that processes r.ImageGenerationConfig so it forwards cfg.Seed,
cfg.Quality, and cfg.CfgScale into the outgoing params (preferably into
params.ExtraParams as well as any explicit params fields if available),
preserving nil checks and correct types/keys so variation requests retain these
values when the native Bedrock payload is built; reference
ImageGenerationConfig, params, and params.ExtraParams to locate and extend the
mapping.
- Around line 916-935: The code currently chooses the Cohere vs Titan embedding
response by batch size (len(resp.Data) > 1); change it to check the model name
instead so single-input Cohere responses still return the Cohere envelope. In
the embedding response logic around resp.Data, use the incoming request/model
identifier (e.g., request.Model or model variable used earlier in this file) to
detect Cohere models (match the Cohere model name pattern) and when the model is
Cohere return BedrockInvokeCohereEmbeddingResp with Embeddings and ResponseType;
otherwise return the Titan BedrockInvokeEmbeddingResp as before. Ensure you
still handle nil embedding arrays (return InputTextTokenCount) for non-Cohere
models as currently done.
- Around line 191-196: The image converters are dropping several fields: in the
image edit converter add extraction and mapping for InPaintingParams.MaskPrompt
and InPaintingParams.ReturnMask, and for OutPaintingParams add MaskPrompt,
ReturnMask, and OutPaintingMode; in ToBifrostImageVariationRequest ensure
ImageGenerationConfig.Seed, ImageGenerationConfig.Quality, and
ImageGenerationConfig.CfgScale are copied into the outgoing params (these are
already handled in the edit converter around the edit converter block at lines
~580–590, use the same mapping pattern). Update the relevant converters
(ToBifrostImageEditRequest / edit converter and ToBifrostImageVariationRequest)
to parse these fields from the incoming config and include them in the params
sent to Bifrost using the existing params key convention.
In `@transports/bifrost-http/integrations/bedrock.go`:
- Around line 123-125: GetHTTPRequestType currently calls
bedrock.DetectInvokeRequestType but doesn't block non-streamable types for the
/invoke-with-response-stream route; update GetHTTPRequestType so after calling
DetectInvokeRequestType(ctx.Request.Body(), modelID) it checks the returned type
and rejects non-streamable types (e.g., embeddings, image ops) by setting an
explicit error response (HTTP 400/unsupported media type) or returning a
dedicated invalid/reject RequestType value so the handler fails fast; reference
DetectInvokeRequestType and GetHTTPRequestType to add this guard and ensure only
streamable request types (text/messages) are allowed on
/invoke-with-response-stream.
---
Duplicate comments:
In `@core/providers/bedrock/images.go`:
- Around line 159-175: The ToBifrostImageGenerationResponse path never
initializes ImageGenerationResponseParameters, so finish_reasons and seeds are
dropped; update the ToBifrostImageGenerationResponse implementation to allocate
and populate response.ImageGenerationResponseParameters (setting FinishReasons
and Seeds from the native Bedrock/Bifrost input) before returning, ensuring the
metadata flows through the Bedrock→Bifrost→ToStabilityAIImageGenerationResponse
round-trip; touch the ToBifrostImageGenerationResponse function and the
ImageGenerationResponseParameters struct population logic so it mirrors how
ToStabilityAIImageGenerationResponse reads those fields.
- Around line 42-48: The isPromptOnlyImageGenerationModel helper is too broad
because it returns true for any model name containing "image" (catching
stability.stable-image-* and Titan image families); update
isPromptOnlyImageGenerationModel to narrow the allowlist by first excluding
known non‑prompt-only vendors (e.g., if the lowercased model string contains
"stability." or "titan" return false) and then only return true for explicit
prompt-only identifiers (e.g., match known prompt-only tokens like "imagen",
"vertex-imagen", or an exact allowlist of model names/patterns) so task-backed
image models are not misclassified. Ensure changes are made inside
isPromptOnlyImageGenerationModel to keep the contract and adjust any
tests/places that rely on its result.
---
Nitpick comments:
In `@core/providers/bedrock/invoke.go`:
- Around line 520-580: Normalize the error strings in the image-edit switch
cases by making them lowercase, removing trailing punctuation, and replacing
camelCase/title-cased identifiers with plain lowercase words; update the
fmt.Errorf messages that reference r.InPaintingParams, r.OutPaintingParams,
r.BackgroundRemovalParams, and the default unsupported taskType error to use
e.g. "inpainting params required for inpainting task", "failed to decode
inpainting image: %w", "outpainting params required for outpainting task",
"failed to decode outpainting image: %w", "background removal params required
for background removal task", "failed to decode background removal image: %w",
and "unsupported task type for image edit: %s" (also apply the same
normalization to the other related messages at the other mentioned ranges
602-616 and 643-648).
🪄 Autofix (Beta)
Fix all unresolved CodeRabbit comments on this PR:
- Push a commit to this branch (recommended)
- Create a new PR with the fixes
ℹ️ Review info
⚙️ Run configuration
Configuration used: Organization UI
Review profile: CHILL
Plan: Pro
Run ID: 4b8ac27d-6d52-4215-bb26-933a8a28c489
📒 Files selected for processing (10)
core/providers/bedrock/bedrock.gocore/providers/bedrock/embedding.gocore/providers/bedrock/images.gocore/providers/bedrock/invoke.gocore/providers/bedrock/types.gocore/schemas/images.gotests/integrations/python/config.ymltests/integrations/python/tests/test_bedrock.pytests/integrations/python/tests/utils/common.pytransports/bifrost-http/integrations/bedrock.go
✅ Files skipped from review due to trivial changes (2)
- tests/integrations/python/tests/utils/common.py
- tests/integrations/python/tests/test_bedrock.py
🚧 Files skipped from review as they are similar to previous changes (2)
- tests/integrations/python/config.yml
- core/providers/bedrock/bedrock.go
Radheshg04
force-pushed
the
03-24-feat_add_embeddings_image_gen_edit_and_variation_to_bedrock_integration
branch
from
a70a006 to
1a2fabd
Compare
Contributor
There was a problem hiding this comment.
Actionable comments posted: 2
♻️ Duplicate comments (5)
core/providers/bedrock/embedding.go (1)
122-126:⚠️ Potential issue | 🟠 Major
embedding_typescan request payloads this parser still drops.
ToBedrockCohereEmbeddingRequestnow forwards arbitraryembedding_types, butToBifrostEmbeddingResponseonly unmarshalsfloatandbase64. Bedrock Cohere returns keys likeint8,uint8,binary, andubinaryinembeddings_by_type, so those responses will silently lose data unless you either reject them up front or preserve them here.🛡️ Minimum safe guard
if v, ok := extra["embedding_types"]; ok { if ss, ok := v.([]string); ok { + for _, t := range ss { + if t != "float" { + return nil, fmt.Errorf("embedding_type %q is not supported yet", t) + } + } req.EmbeddingTypes = ss delete(extra, "embedding_types") } }For AWS Bedrock Cohere Embed models, what values are valid in `embedding_types`, and what keys are returned in `embeddings_by_type` for each value?Also applies to: 185-213
🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@core/providers/bedrock/embedding.go` around lines 122 - 126, The request parser ToBedrockCohereEmbeddingRequest currently allows arbitrary embedding_types but ToBifrostEmbeddingResponse only decodes "float" and "base64", causing silent data loss for types like "int8", "uint8", "binary", and "ubinary"; either validate/normalize embedding_types in ToBedrockCohereEmbeddingRequest to reject unsupported values up-front, or (preferable) extend ToBifrostEmbeddingResponse to parse and preserve additional keys returned in embeddings_by_type (handle int8/uint8 as integer byte arrays, binary/ubinary as raw bytes or base64-encoded strings) and fall back to preserving the raw payload for unknown types so nothing is dropped. Ensure the change updates both ToBedrockCohereEmbeddingRequest and ToBifrostEmbeddingResponse and includes tests for int8/uint8/binary/ubinary cases.transports/bifrost-http/integrations/bedrock.go (1)
123-125:⚠️ Potential issue | 🟠 MajorBlock embeddings and image ops on
/invoke-with-response-stream.This route now accepts
EmbeddingRequestandImage*Request, but only text/messages are emitted here as event-stream chunks. Returning a normal embedding/image payload from the Bedrock streaming endpoint breaks the streaming client contract instead of failing fast.🛠️ Minimal guard
requestType, _ := ctx.Value(schemas.BifrostContextKeyHTTPRequestType).(schemas.RequestType) +if requestType != schemas.ResponsesRequest && requestType != schemas.TextCompletionRequest { + return nil, fmt.Errorf("request type %s is not supported on invoke-with-response-stream", requestType) +} switch requestType { case schemas.EmbeddingRequest:Does AWS Bedrock InvokeModelWithResponseStream support embedding or image generation/edit/variation requests, or only request types that return an event stream?Also applies to: 132-173
🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@transports/bifrost-http/integrations/bedrock.go` around lines 123 - 125, The streaming route's GetHTTPRequestType currently uses bedrock.DetectInvokeRequestType but doesn't guard against EmbeddingRequest or any Image*Request, causing non-stream payloads to be emitted into an event-stream; update the GetHTTPRequestType implementation used for the /invoke-with-response-stream route to call bedrock.DetectInvokeRequestType(ctx.Request.Body(), modelID), check the returned request type for embedding or image variants (e.g., any EmbeddingRequest or Image*Request), and if detected return a non-stream/invalid type or trigger an immediate failure response so the handler fails fast; ensure only streaming-capable types are allowed to proceed to the response-stream path.core/providers/bedrock/invoke.go (3)
688-695:⚠️ Potential issue | 🟠 MajorVariation requests still lose part of
imageGenerationConfig.This block only forwards
numberOfImagesand size.seed,quality, andcfgScaleare still dropped here even though Lines 584-599 preserve them from the sameimageGenerationConfigin the edit path, so invoke variation requests do not round-trip Bedrock generation settings.🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@core/providers/bedrock/invoke.go` around lines 688 - 695, The variation-invoke path only forwards NumberOfImages and Size from r.ImageGenerationConfig to params, dropping seed, quality, and cfgScale which breaks round-tripping; update the same block that handles r.ImageGenerationConfig to also copy cfg.Seed, cfg.Quality, and cfg.CfgScale into the corresponding fields on params (guarding nils / zero values the same way NumberOfImages and Size are handled) so invoke variation requests preserve all imageGenerationConfig settings.
519-565:⚠️ Potential issue | 🟠 MajorCarry the remaining Titan edit knobs through the round-trip.
maskPrompt,returnMask, andoutPaintingModeare known nested fields oninPaintingParams/outPaintingParams, so they never fall intor.ExtraParams. These branches currently drop them entirely, which changes Bedrock edit behavior on invoke → Bifrost → Bedrock round-trips.💡 Suggested fix
if r.InPaintingParams.MaskImage != nil { maskBytes, err := base64.StdEncoding.DecodeString(*r.InPaintingParams.MaskImage) if err != nil { return nil, fmt.Errorf("failed to decode inpainting mask: %w", err) } params.Mask = maskBytes } + if r.InPaintingParams.MaskPrompt != nil || r.InPaintingParams.ReturnMask != nil { + if params.ExtraParams == nil { + params.ExtraParams = make(map[string]interface{}) + } + if r.InPaintingParams.MaskPrompt != nil { + params.ExtraParams["maskPrompt"] = *r.InPaintingParams.MaskPrompt + } + if r.InPaintingParams.ReturnMask != nil { + params.ExtraParams["returnMask"] = *r.InPaintingParams.ReturnMask + } + } case TaskTypeOutpainting: if r.OutPaintingParams == nil { return nil, fmt.Errorf("outPaintingParams required for OUTPAINTING task") @@ if r.OutPaintingParams.MaskImage != nil { maskBytes, err := base64.StdEncoding.DecodeString(*r.OutPaintingParams.MaskImage) if err != nil { return nil, fmt.Errorf("failed to decode outpainting mask: %w", err) } params.Mask = maskBytes } + if r.OutPaintingParams.MaskPrompt != nil || r.OutPaintingParams.ReturnMask != nil || r.OutPaintingParams.OutPaintingMode != nil { + if params.ExtraParams == nil { + params.ExtraParams = make(map[string]interface{}) + } + if r.OutPaintingParams.MaskPrompt != nil { + params.ExtraParams["maskPrompt"] = *r.OutPaintingParams.MaskPrompt + } + if r.OutPaintingParams.ReturnMask != nil { + params.ExtraParams["returnMask"] = *r.OutPaintingParams.ReturnMask + } + if r.OutPaintingParams.OutPaintingMode != nil { + params.ExtraParams["outPaintingMode"] = *r.OutPaintingParams.OutPaintingMode + } + }🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@core/providers/bedrock/invoke.go` around lines 519 - 565, The inpainting/outpainting branches drop Titan-specific edit knobs (maskPrompt, returnMask, outPaintingMode) by not copying them from r.InPaintingParams / r.OutPaintingParams into the outgoing params; update the TaskTypeInpainting and TaskTypeOutpainting handling to explicitly carry those fields through the round-trip by copying maskPrompt and returnMask from r.InPaintingParams to params for inpainting, and maskPrompt, returnMask and outPaintingMode from r.OutPaintingParams to params for outpainting (e.g. check for non-nil r.InPaintingParams.MaskPrompt / ReturnMask and r.OutPaintingParams.MaskPrompt / ReturnMask / OutPaintingMode and assign them to params.MaskPrompt / params.ReturnMask / params.OutPaintingMode) so these nested fields are preserved rather than lost to r.ExtraParams.
913-936:⚠️ Potential issue | 🟠 MajorChoose the Cohere envelope from the model, not the batch size.
A single-input Cohere invoke request still needs the Cohere response shape. Using
len(resp.Data) > 1here makes the response format depend on batch size, so one-item Cohere requests fall through to the Titanembeddingpayload.💡 Suggested fix
+ model := resp.Model + if resp.ExtraFields.ModelRequested != "" { + model = resp.ExtraFields.ModelRequested + } + + if strings.Contains(strings.ToLower(model), "cohere") { + floats := make([][]float32, 0, len(resp.Data)) + for _, d := range resp.Data { + floats = append(floats, d.Embedding.EmbeddingArray) + } + return &BedrockInvokeCohereEmbeddingResp{ + Embeddings: floats, + ResponseType: "embeddings_floats", + }, nil + } + if len(resp.Data) == 0 { return &BedrockInvokeEmbeddingResp{InputTextTokenCount: tokenCount}, nil } - - // Multi-embedding: reconstruct Cohere-format response - if len(resp.Data) > 1 { - floats := make([][]float32, 0, len(resp.Data)) - for _, d := range resp.Data { - floats = append(floats, d.Embedding.EmbeddingArray) - } - return &BedrockInvokeCohereEmbeddingResp{ - Embeddings: floats, - ResponseType: "embeddings_floats", - }, nil - }
🧹 Nitpick comments (2)
transports/bifrost-http/integrations/bedrock.go (1)
230-269: Fail closed on unknown invoke request types.The
defaultbranch currently treatsschemas.UnknownRequestand any future enum the same as text completion. That turns malformed or unsupported invoke bodies into the wrong provider path instead of surfacing a clear error.♻️ Suggested change
- default: - // TextCompletionRequest and any unrecognised type forwarded to text completion path + case schemas.TextCompletionRequest: return &schemas.BifrostRequest{ TextCompletionRequest: invokeReq.ToBifrostTextCompletionRequest(ctx), }, nil + default: + return nil, fmt.Errorf("unsupported bedrock invoke request type: %s", requestType) }🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@transports/bifrost-http/integrations/bedrock.go` around lines 230 - 269, The switch on requestType currently falls through to TextCompletionRequest in the default branch, causing unknown or unsupported schemas.RequestType values to be handled as text completions; update the switch in the conversion logic (the switch on requestType that returns *schemas.BifrostRequest from invokeReq) to explicitly handle schemas.UnknownRequest (and any other unsupported enum values) by returning an error (e.g., fmt.Errorf("unsupported invoke request type: %v", requestType)) instead of calling invokeReq.ToBifrostTextCompletionRequest(ctx), and keep the explicit cases (EmbeddingRequest, ImageGenerationRequest, ImageEditRequest, ImageVariationRequest, ResponsesRequest) unchanged so only known text completion types map to TextCompletionRequest.core/providers/bedrock/invoke.go (1)
521-607: Normalize the new error strings to lowercase.Several new
fmt.Errorfmessages start withinPaintingParams,OUTPAINTING, orStability AI. That breaks the repo's Go error-string convention.As per coding guidelines, "Error strings should be lowercase with no trailing punctuation (Go convention)."
Also applies to: 614-616, 643-649
🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@core/providers/bedrock/invoke.go` around lines 521 - 607, Normalize the new error strings to follow Go conventions (lowercase, no trailing punctuation): update the fmt.Errorf messages that reference in-painting/out-painting/background removal and Stability AI so they are all lowercase and phrased consistently. Specifically change the returns that use r.InPaintingParams, r.OutPaintingParams, r.BackgroundRemovalParams and the Stability AI path (the fmt.Errorf calls like "inPaintingParams required for INPAINTING task", "outPaintingParams required for OUTPAINTING task", "backgroundRemovalParams required for BACKGROUND_REMOVAL task", "unsupported taskType for image edit: %s", and "image field is required for Stability AI image edit"/decode messages) to lowercase forms such as "inpainting params required for inpainting task", "outpainting params required for outpainting task", "background removal params required for background removal task", "unsupported task type for image edit: %s", and "image field is required for stability ai image edit" (also normalize "Stability AI" -> "stability ai" in decode error messages).
🤖 Prompt for all review comments with AI agents
Verify each finding against the current code and only fix it if needed.
Inline comments:
In `@core/providers/bedrock/invoke.go`:
- Around line 386-390: The model id string (r.ModelID) must be percent-unescaped
before parsing so providers split correctly; call path.PathUnescape (or
url.PathUnescape) on r.ModelID and use the unescaped value when calling
schemas.ParseModelString and providerUtils.CheckAndSetDefaultProvider (where
r.ModelID is used, e.g., in the current block constructing
schemas.BifrostEmbeddingRequest and the similar blocks at the other
occurrences). Ensure any earlier DetectInvokeRequestType handling remains but
normalize r.ModelID right before ParseModelString to fix cases like
"openai%2Fgpt-image-1" not splitting into provider/model.
- Around line 886-891: The returned BedrockImageGenerationResponse (variable
result) is a superset and leaves optional fields (maskImage, error, seeds,
finish_reasons) as zero-values, so the Titan/Nova fallback emits empty/null
fields; fix by making those optional fields truly optional: add `omitempty` to
the JSON struct tags for maskImage, error, seeds, and finish_reasons in the
BedrockImageGenerationResponse definition (core/providers/bedrock/types.go) so
when result only sets Images (populated from resp.Data[].B64JSON) the extra
fields are omitted from the serialized response.
---
Duplicate comments:
In `@core/providers/bedrock/embedding.go`:
- Around line 122-126: The request parser ToBedrockCohereEmbeddingRequest
currently allows arbitrary embedding_types but ToBifrostEmbeddingResponse only
decodes "float" and "base64", causing silent data loss for types like "int8",
"uint8", "binary", and "ubinary"; either validate/normalize embedding_types in
ToBedrockCohereEmbeddingRequest to reject unsupported values up-front, or
(preferable) extend ToBifrostEmbeddingResponse to parse and preserve additional
keys returned in embeddings_by_type (handle int8/uint8 as integer byte arrays,
binary/ubinary as raw bytes or base64-encoded strings) and fall back to
preserving the raw payload for unknown types so nothing is dropped. Ensure the
change updates both ToBedrockCohereEmbeddingRequest and
ToBifrostEmbeddingResponse and includes tests for int8/uint8/binary/ubinary
cases.
In `@core/providers/bedrock/invoke.go`:
- Around line 688-695: The variation-invoke path only forwards NumberOfImages
and Size from r.ImageGenerationConfig to params, dropping seed, quality, and
cfgScale which breaks round-tripping; update the same block that handles
r.ImageGenerationConfig to also copy cfg.Seed, cfg.Quality, and cfg.CfgScale
into the corresponding fields on params (guarding nils / zero values the same
way NumberOfImages and Size are handled) so invoke variation requests preserve
all imageGenerationConfig settings.
- Around line 519-565: The inpainting/outpainting branches drop Titan-specific
edit knobs (maskPrompt, returnMask, outPaintingMode) by not copying them from
r.InPaintingParams / r.OutPaintingParams into the outgoing params; update the
TaskTypeInpainting and TaskTypeOutpainting handling to explicitly carry those
fields through the round-trip by copying maskPrompt and returnMask from
r.InPaintingParams to params for inpainting, and maskPrompt, returnMask and
outPaintingMode from r.OutPaintingParams to params for outpainting (e.g. check
for non-nil r.InPaintingParams.MaskPrompt / ReturnMask and
r.OutPaintingParams.MaskPrompt / ReturnMask / OutPaintingMode and assign them to
params.MaskPrompt / params.ReturnMask / params.OutPaintingMode) so these nested
fields are preserved rather than lost to r.ExtraParams.
In `@transports/bifrost-http/integrations/bedrock.go`:
- Around line 123-125: The streaming route's GetHTTPRequestType currently uses
bedrock.DetectInvokeRequestType but doesn't guard against EmbeddingRequest or
any Image*Request, causing non-stream payloads to be emitted into an
event-stream; update the GetHTTPRequestType implementation used for the
/invoke-with-response-stream route to call
bedrock.DetectInvokeRequestType(ctx.Request.Body(), modelID), check the returned
request type for embedding or image variants (e.g., any EmbeddingRequest or
Image*Request), and if detected return a non-stream/invalid type or trigger an
immediate failure response so the handler fails fast; ensure only
streaming-capable types are allowed to proceed to the response-stream path.
---
Nitpick comments:
In `@core/providers/bedrock/invoke.go`:
- Around line 521-607: Normalize the new error strings to follow Go conventions
(lowercase, no trailing punctuation): update the fmt.Errorf messages that
reference in-painting/out-painting/background removal and Stability AI so they
are all lowercase and phrased consistently. Specifically change the returns that
use r.InPaintingParams, r.OutPaintingParams, r.BackgroundRemovalParams and the
Stability AI path (the fmt.Errorf calls like "inPaintingParams required for
INPAINTING task", "outPaintingParams required for OUTPAINTING task",
"backgroundRemovalParams required for BACKGROUND_REMOVAL task", "unsupported
taskType for image edit: %s", and "image field is required for Stability AI
image edit"/decode messages) to lowercase forms such as "inpainting params
required for inpainting task", "outpainting params required for outpainting
task", "background removal params required for background removal task",
"unsupported task type for image edit: %s", and "image field is required for
stability ai image edit" (also normalize "Stability AI" -> "stability ai" in
decode error messages).
In `@transports/bifrost-http/integrations/bedrock.go`:
- Around line 230-269: The switch on requestType currently falls through to
TextCompletionRequest in the default branch, causing unknown or unsupported
schemas.RequestType values to be handled as text completions; update the switch
in the conversion logic (the switch on requestType that returns
*schemas.BifrostRequest from invokeReq) to explicitly handle
schemas.UnknownRequest (and any other unsupported enum values) by returning an
error (e.g., fmt.Errorf("unsupported invoke request type: %v", requestType))
instead of calling invokeReq.ToBifrostTextCompletionRequest(ctx), and keep the
explicit cases (EmbeddingRequest, ImageGenerationRequest, ImageEditRequest,
ImageVariationRequest, ResponsesRequest) unchanged so only known text completion
types map to TextCompletionRequest.
🪄 Autofix (Beta)
Fix all unresolved CodeRabbit comments on this PR:
- Push a commit to this branch (recommended)
- Create a new PR with the fixes
ℹ️ Review info
⚙️ Run configuration
Configuration used: Organization UI
Review profile: CHILL
Plan: Pro
Run ID: 3ad532ea-8c55-4f6f-b080-8c12dc6db0c7
📒 Files selected for processing (10)
core/providers/bedrock/bedrock.gocore/providers/bedrock/embedding.gocore/providers/bedrock/images.gocore/providers/bedrock/invoke.gocore/providers/bedrock/types.gocore/schemas/images.gotests/integrations/python/config.ymltests/integrations/python/tests/test_bedrock.pytests/integrations/python/tests/utils/common.pytransports/bifrost-http/integrations/bedrock.go
✅ Files skipped from review due to trivial changes (1)
- core/providers/bedrock/bedrock.go
🚧 Files skipped from review as they are similar to previous changes (3)
- tests/integrations/python/config.yml
- tests/integrations/python/tests/utils/common.py
- core/providers/bedrock/images.go
Radheshg04
force-pushed
the
03-24-feat_add_embeddings_image_gen_edit_and_variation_to_bedrock_integration
branch
from
1a2fabd to
512594b
Compare
Contributor
There was a problem hiding this comment.
Actionable comments posted: 4
♻️ Duplicate comments (3)
tests/integrations/python/tests/test_bedrock.py (1)
2464-2469:⚠️ Potential issue | 🟠 MajorValidate the
inputsembedding payload, not just the key name.This repeats the earlier key-only assertion pattern on the
inputsbranch:{"embedding": "..."}or{"embeddings": null}still passes, and a strayresultsfield is never rejected. Mirror the shape checks fromtest_32here and assert thatresultsis absent.Suggested fix
- has_embedding = "embedding" in out or "embeddings" in out - assert has_embedding, ( - f"Request was misrouted — expected 'embedding' or 'embeddings' key but got: {list(out.keys())}. " - f"If 'results' is present, the request was routed to text completion instead of embeddings." - ) + assert "results" not in out, ( + f"Request was misrouted to text completion. Response keys: {list(out.keys())}" + ) + if "embedding" in out: + _assert_invoke_embedding(out) + else: + embeddings = out.get("embeddings") + assert isinstance(embeddings, list) and len(embeddings) == len(body["inputs"]), ( + f"Expected {len(body['inputs'])} embeddings, got: {out}" + ) + for i, vector in enumerate(embeddings): + assert isinstance(vector, list) and len(vector) > 0, f"Embedding {i} is empty" + assert all(isinstance(v, (int, float)) for v in vector), ( + f"Embedding {i} must be numeric" + )🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@tests/integrations/python/tests/test_bedrock.py` around lines 2464 - 2469, The current check only asserts presence of "embedding" or "embeddings" in out (via has_embedding) which lets invalid shapes pass; update the "Cohere inputs" branch that sets has_embedding/out to also assert that "results" is not present and that the embedding payload has the expected shape: if "embedding" is present it must be a non-empty list/sequence of numbers, and if "embeddings" is present it must be a non-empty list of non-empty numeric lists (match the same shape validation used in test_32). Locate the assertion using has_embedding and the print "✓ Cohere inputs embedding response keys" and replace it with checks that reject None or scalar values, verify element types are numeric, and explicitly assert "results" not in out.core/providers/bedrock/invoke.go (2)
947-971:⚠️ Potential issue | 🔴 CriticalGuard
Embeddingbefore touchingEmbeddingArray.Line 961 dereferences
d.Embeddingin the Cohere loop, and Line 970 still dereferencesresp.Data[0].Embeddingdirectly. A missing embedding payload will panic here before the fallback runs, and the earlylen(resp.Data) == 0return also emits the Titan envelope even for Cohere models.🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@core/providers/bedrock/invoke.go` around lines 947 - 971, The loop and subsequent checks access d.Embedding and resp.Data[0].Embedding without nil guards which can panic and also cause the len(resp.Data)==0 early return to emit the Titan envelope for Cohere models; update the Cohere branch (where you build floats and return BedrockInvokeCohereEmbeddingResp) to first verify each resp.Data[i] and resp.Data[i].Embedding are non-nil and that Embedding.EmbeddingArray is non-nil (skip or bail to the fallback if any are nil), and likewise guard resp.Data[0].Embedding before using EmbeddingArray in the Titan branch (return BedrockInvokeEmbeddingResp as the safe fallback when missing); ensure the model detection logic still uses resp.ExtraFields.ModelRequested to determine Cohere vs Titan but only proceeds with Cohere output when those nil checks pass.
386-386:⚠️ Potential issue | 🟠 MajorUnescape
modelIdbeforeParseModelString.Lines 386, 446, 506, and 677 still feed the raw path value into
schemas.ParseModelString, even thoughDetectInvokeRequestTypealready has toPathUnescapethe same field. Requests likeopenai%2Fgpt-image-1orvertex%2F...will stay on the default Bedrock provider here, which breaks the new cross-provider/invokeflow.🛠️ Suggested fix
+func (r *BedrockInvokeRequest) normalizedModelID() string { + if unescaped, err := url.PathUnescape(r.ModelID); err == nil { + return unescaped + } + return r.ModelID +} + func (r *BedrockInvokeRequest) ToBifrostEmbeddingRequest(ctx *schemas.BifrostContext) *schemas.BifrostEmbeddingRequest { - provider, model := schemas.ParseModelString(r.ModelID, providerUtils.CheckAndSetDefaultProvider(ctx, schemas.Bedrock)) + modelID := r.normalizedModelID() + provider, model := schemas.ParseModelString(modelID, providerUtils.CheckAndSetDefaultProvider(ctx, schemas.Bedrock))Use the same helper in the image generation, image edit, and image variation converters too.
Also applies to: 446-446, 506-506, 677-677
🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@core/providers/bedrock/invoke.go` at line 386, Unescape the incoming model id before parsing: call the same PathUnescape helper used in DetectInvokeRequestType on r.ModelID and pass the unescaped value into schemas.ParseModelString instead of the raw path; update the four call sites that construct provider, model (the ParseModelString call at the top-level invoke handler and the image generation, image edit, and image variation converter functions) to use the unescaped modelId so cross-provider ids like "openai%2Fgpt-image-1" or "vertex%2F..." resolve correctly.
🧹 Nitpick comments (2)
tests/integrations/python/tests/test_bedrock.py (1)
358-360: Avoid dropping the config-driven invoke coverage here.This decorator disables the only non-streaming test that still exercises
get_model("bedrock", "text_completion").test_40covers Titan bareinputText, but it doesn't cover the configured Mistral invoke envelope or response conversion thattest_01was checking. If this path is temporarily broken, anxfailwith a tracking reason is safer than a permanent skip.🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@tests/integrations/python/tests/test_bedrock.py` around lines 358 - 360, Replace the unconditional skip on test_01_text_completion_invoke with a pytest xfail so the configured-bedrock invoke path (get_model("bedrock", "text_completion")) remains exercised; specifically, change the `@pytest.mark.skip` decorator on test_01_text_completion_invoke to `@pytest.mark.xfail`(reason="tracking: <ticket-or-issue-id> - preserve coverage of Mistral invoke/response conversion", strict=False) while keeping the `@skip_if_no_api_key`("bedrock") decorator so the test still runs when credentials are present but is marked xfail until the tracked issue is resolved.core/providers/bedrock/invoke.go (1)
520-606: Normalize the new error strings to lowercase Go style.These added errors still mix camelCase field names and enum-style uppercase fragments (
INPAINTING,BACKGROUND_REMOVAL,Stability AI,IMAGE_VARIATION), which makes logs inconsistent with the rest of the package.As per coding guidelines, "Error strings should be lowercase with no trailing punctuation (Go convention)."
Also applies to: 628-641, 668-695
🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@core/providers/bedrock/invoke.go` around lines 520 - 606, The error messages in the image edit switch (errors like "inPaintingParams required for INPAINTING task", "outPaintingParams required for OUTPAINTING task", "backgroundRemovalParams required for BACKGROUND_REMOVAL task", and the other occurrences around 628-641 and 668-695) must follow Go error-style: make them lowercase, avoid camelCase/ALL_CAPS fragments, and no trailing punctuation; update fmt.Errorf calls in the switch handling TaskTypeInpainting, TaskTypeOutpainting, TaskTypeBackgroundRemoval (and the other similar messages referenced) to use lowercase, space-separated words (e.g., "inpainting params required for inpainting task" or preferably "inpainting params required for image edit") and similarly normalize any provider/model names or enum fragments to lowercase in those error strings. Ensure you only change the string literals in the fmt.Errorf calls (leave variable names and logic intact).
🤖 Prompt for all review comments with AI agents
Verify each finding against the current code and only fix it if needed.
Inline comments:
In `@core/providers/bedrock/invoke.go`:
- Around line 199-201: The top-level image detection using sonic.Get(body,
"image") that returns schemas.ImageEditRequest is too broad and conflicts with
the converter which expects a Stability-only model inferred via params.Type;
update the detection so it only classifies a request as ImageEditRequest when
either params.Type explicitly equals the image-edit type or when the
provider/model namespace indicates a Stability image-edit model, otherwise fall
back to the generic inference path. Concretely, modify the branch that uses
sonic.Get(body, "image") to check params.Type (and/or the model/provider prefix)
before returning schemas.ImageEditRequest, and ensure the same logic is applied
to the other duplicate branch used around the converter (the block that inspects
params.Type) so openai/... and vertex/... flat edit bodies are not routed
incorrectly into the edit conversion path.
In `@tests/integrations/python/config.yml`:
- Around line 522-524: scenario_capabilities is missing the image_variation
entry even though provider_scenarios.bedrock.image_variation (and
providers.bedrock.image_variation) are enabled; add an image_variation key under
scenario_capabilities mapping to point to the same provider scenario flag so
config-driven lookups will find it. Specifically, update the
scenario_capabilities structure to include an image_variation entry that
references provider_scenarios.bedrock.image_variation (or the
providers.bedrock.image_variation model) analogous to the existing
image_generation and image_edit entries so the new scenario is wired end-to-end.
In `@tests/integrations/python/tests/test_bedrock.py`:
- Around line 2338-2343: The test currently sets the mask field to
BASE64_IMAGE_LARGE (sending the same image twice), so mask handling isn't
exercised; update the body dict used in this Stability inpaint test to use a
real mask fixture (e.g., BASE64_MASK_STABILITY or BASE64_IMAGE_MASK) instead of
BASE64_IMAGE_LARGE so the image+mask code path is actually exercised when
invoking the converter and verifying mask behavior.
- Around line 2072-2085: The helper _assert_invoke_images currently only checks
that each image is a non-empty string, which allows URLs or placeholders to slip
through; update _assert_invoke_images to validate that each entry in images is a
real base64-encoded payload by attempting to decode each string (e.g., using
base64.b64decode with validate=True), asserting the decode succeeds and yields a
non-trivial byte length (optionally strip a data URI prefix like "data:image/"
before decoding), and keep the existing min_images and error messages; reference
the function name _assert_invoke_images and the images variable when making
these changes.
---
Duplicate comments:
In `@core/providers/bedrock/invoke.go`:
- Around line 947-971: The loop and subsequent checks access d.Embedding and
resp.Data[0].Embedding without nil guards which can panic and also cause the
len(resp.Data)==0 early return to emit the Titan envelope for Cohere models;
update the Cohere branch (where you build floats and return
BedrockInvokeCohereEmbeddingResp) to first verify each resp.Data[i] and
resp.Data[i].Embedding are non-nil and that Embedding.EmbeddingArray is non-nil
(skip or bail to the fallback if any are nil), and likewise guard
resp.Data[0].Embedding before using EmbeddingArray in the Titan branch (return
BedrockInvokeEmbeddingResp as the safe fallback when missing); ensure the model
detection logic still uses resp.ExtraFields.ModelRequested to determine Cohere
vs Titan but only proceeds with Cohere output when those nil checks pass.
- Line 386: Unescape the incoming model id before parsing: call the same
PathUnescape helper used in DetectInvokeRequestType on r.ModelID and pass the
unescaped value into schemas.ParseModelString instead of the raw path; update
the four call sites that construct provider, model (the ParseModelString call at
the top-level invoke handler and the image generation, image edit, and image
variation converter functions) to use the unescaped modelId so cross-provider
ids like "openai%2Fgpt-image-1" or "vertex%2F..." resolve correctly.
In `@tests/integrations/python/tests/test_bedrock.py`:
- Around line 2464-2469: The current check only asserts presence of "embedding"
or "embeddings" in out (via has_embedding) which lets invalid shapes pass;
update the "Cohere inputs" branch that sets has_embedding/out to also assert
that "results" is not present and that the embedding payload has the expected
shape: if "embedding" is present it must be a non-empty list/sequence of
numbers, and if "embeddings" is present it must be a non-empty list of non-empty
numeric lists (match the same shape validation used in test_32). Locate the
assertion using has_embedding and the print "✓ Cohere inputs embedding response
keys" and replace it with checks that reject None or scalar values, verify
element types are numeric, and explicitly assert "results" not in out.
---
Nitpick comments:
In `@core/providers/bedrock/invoke.go`:
- Around line 520-606: The error messages in the image edit switch (errors like
"inPaintingParams required for INPAINTING task", "outPaintingParams required for
OUTPAINTING task", "backgroundRemovalParams required for BACKGROUND_REMOVAL
task", and the other occurrences around 628-641 and 668-695) must follow Go
error-style: make them lowercase, avoid camelCase/ALL_CAPS fragments, and no
trailing punctuation; update fmt.Errorf calls in the switch handling
TaskTypeInpainting, TaskTypeOutpainting, TaskTypeBackgroundRemoval (and the
other similar messages referenced) to use lowercase, space-separated words
(e.g., "inpainting params required for inpainting task" or preferably
"inpainting params required for image edit") and similarly normalize any
provider/model names or enum fragments to lowercase in those error strings.
Ensure you only change the string literals in the fmt.Errorf calls (leave
variable names and logic intact).
In `@tests/integrations/python/tests/test_bedrock.py`:
- Around line 358-360: Replace the unconditional skip on
test_01_text_completion_invoke with a pytest xfail so the configured-bedrock
invoke path (get_model("bedrock", "text_completion")) remains exercised;
specifically, change the `@pytest.mark.skip` decorator on
test_01_text_completion_invoke to `@pytest.mark.xfail`(reason="tracking:
<ticket-or-issue-id> - preserve coverage of Mistral invoke/response conversion",
strict=False) while keeping the `@skip_if_no_api_key`("bedrock") decorator so the
test still runs when credentials are present but is marked xfail until the
tracked issue is resolved.
🪄 Autofix (Beta)
Fix all unresolved CodeRabbit comments on this PR:
- Push a commit to this branch (recommended)
- Create a new PR with the fixes
ℹ️ Review info
⚙️ Run configuration
Configuration used: Organization UI
Review profile: CHILL
Plan: Pro
Run ID: ebe82cca-a5a2-4bc6-91aa-a7d44eeae62a
📒 Files selected for processing (10)
core/providers/bedrock/bedrock.gocore/providers/bedrock/embedding.gocore/providers/bedrock/images.gocore/providers/bedrock/invoke.gocore/providers/bedrock/types.gocore/schemas/images.gotests/integrations/python/config.ymltests/integrations/python/tests/test_bedrock.pytests/integrations/python/tests/utils/common.pytransports/bifrost-http/integrations/bedrock.go
✅ Files skipped from review due to trivial changes (1)
- core/providers/bedrock/types.go
🚧 Files skipped from review as they are similar to previous changes (1)
- core/providers/bedrock/embedding.go
Radheshg04
force-pushed
the
03-24-feat_add_embeddings_image_gen_edit_and_variation_to_bedrock_integration
branch
from
512594b to
227a943
Compare
Contributor
There was a problem hiding this comment.
Actionable comments posted: 2
♻️ Duplicate comments (4)
tests/integrations/python/tests/test_bedrock.py (3)
2464-2468:⚠️ Potential issue | 🟡 MinorAssert the
inputsembedding shape, not just the route.This regression test still passes on malformed outputs like
{"embeddings": {}}as long as the handler lands on the embedding branch. Reuse the same payload-shape checks astest_32so theinputsconverter is actually covered.🔧 Suggested fix
- has_embedding = "embedding" in out or "embeddings" in out - assert has_embedding, ( - f"Request was misrouted — expected 'embedding' or 'embeddings' key but got: {list(out.keys())}. " - f"If 'results' is present, the request was routed to text completion instead of embeddings." - ) + if "embedding" in out: + _assert_invoke_embedding(out) + else: + embeddings = out.get("embeddings") + assert isinstance(embeddings, list) and len(embeddings) == len(body["inputs"]), ( + f"Expected {len(body['inputs'])} embeddings, got: {out}" + ) + for i, vector in enumerate(embeddings): + assert isinstance(vector, list) and len(vector) > 0, f"Embedding {i} is empty" + assert all(isinstance(v, (int, float)) for v in vector), ( + f"Embedding {i} must be numeric" + ) print(f" ✓ Cohere inputs embedding response keys: {list(out.keys())}")🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@tests/integrations/python/tests/test_bedrock.py` around lines 2464 - 2468, The test only checks for presence of "embedding"/"embeddings" but not the actual shape; update the assertion to validate the embeddings payload shape (reuse the same checks as test_32) by verifying that out contains "inputs" (or "embeddings") and that each entry in out["inputs"] is a list/sequence of numeric vectors with expected length and that the number of input vectors matches the number of input items; specifically, in the test where variables like has_embedding and out are used, add validations that out["inputs"] is a list, len(out["inputs"]) equals the number of original inputs, and each element is a list/tuple of numbers (non-empty) to ensure the inputs converter produced real embedding arrays rather than empty dicts or malformed structures.
2072-2085:⚠️ Potential issue | 🟠 MajorValidate that
imagescontains real base64 payloads.This helper still accepts URLs or placeholder strings, so the invoke-image tests can pass even if the converter stops returning native Bedrock bytes.
🔧 Suggested fix
def _assert_invoke_images(response_body: dict, min_images: int = 1) -> None: """Assert that an invoke response contains at least min_images base64 images.""" images = response_body.get("images") or [] assert isinstance(images, list), ( f"Expected 'images' to be a list, got {type(images).__name__}. " f"Response keys: {list(response_body.keys())}" ) assert len(images) >= min_images, ( f"Expected at least {min_images} image(s) in response, got {len(images)}. " f"Response keys: {list(response_body.keys())}" ) for i, img in enumerate(images): assert isinstance(img, str) and len(img) > 0, f"Image {i} is not a non-empty string" + try: + decoded = base64.b64decode(img, validate=True) + except Exception as exc: + raise AssertionError(f"Image {i} is not valid base64") from exc + assert len(decoded) > 0, f"Image {i} decoded to empty bytes" print(f" ✓ {len(images)} image(s) returned")🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@tests/integrations/python/tests/test_bedrock.py` around lines 2072 - 2085, The _assert_invoke_images helper currently accepts any non-empty string; update it to verify each image is a real base64 payload by stripping an optional data URL prefix (e.g. "data:...;base64,") from each img string and attempting to base64-decode it (use a strict/validate mode if available), asserting the decode succeeds and yields non-zero bytes; keep the existing checks for list type and minimum count and reference _assert_invoke_images and the images variable when implementing this validation.
2338-2343:⚠️ Potential issue | 🟠 MajorUse a real mask fixture in the Stability inpaint case.
maskis currently the same payload asimage, so this never verifies the image+mask path. If the invoke converter ignores or mangles the mask, this test still passes as long as any image comes back. Swap in a dedicated Stability mask fixture here.🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@tests/integrations/python/tests/test_bedrock.py` around lines 2338 - 2343, The test builds the request body with both "image" and "mask" set to BASE64_IMAGE_LARGE, so it never exercises the image+mask inpaint path; replace the "mask" value in the body dict with a real Stability inpaint mask fixture (e.g., a new or existing constant like BASE64_MASK_STABILITY or STABILITY_INPAINT_MASK) and ensure that fixture contains a proper base64 mask image rather than the full source image; update any imports or fixture definitions so the test uses that mask when constructing body.core/providers/bedrock/embedding.go (1)
218-245:⚠️ Potential issue | 🟠 Major
embeddings_by_typestill loses the encoding identity.
core/schemas/embeddings.gohas no per-entry encoding field, so Lines 218-245 flattenint8vsbinaryanduint8vsubinaryinto indistinguishableEmbeddingStrvalues with duplicate indexes. Themap[...]loops also make their relative order non-deterministic. A consumer readingDatastill cannot reconstruct which encoding each row came from. Either carry the encoding name alongside each entry or fail fast on this generic multi-type path.🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@core/providers/bedrock/embedding.go` around lines 218 - 245, The code in the loop that builds bifrostResponse.Data (using typed.Int8, typed.Binary, typed.Uint8, typed.Ubinary and creating schemas.EmbeddingData / schemas.EmbeddingStruct entries) loses per-entry encoding identity and iterates maps (non-deterministic order), so consumers cannot reconstruct types; either add and populate an explicit encoding field on each EmbeddingData entry (e.g., set Encoding: "int8"/"binary"/"uint8"/"ubinary") and ensure indexes are unique/consistent, or if schemas.EmbeddingData cannot be changed, fail fast when multiple different typed slices are present; also replace the map[...] iteration with a deterministic ordered slice of (name,rows) pairs to preserve order (e.g., iterate a fixed list ["int8","binary",...] instead of ranging a map). Ensure error messages reflect the chosen behavior.
🧹 Nitpick comments (2)
tests/integrations/python/tests/test_bedrock.py (1)
2376-2394: Prefer the configured OpenAI image model here.Hardcoding
openai/gpt-image-1bypassesproviders.openai.image_generation, so this test can drift from the repo's configured deterministic image model.Based on learnings, prefer using the OpenAI image generation model 'gpt-image-1' via the config key providers.openai.image_generation for image-generation scenarios.♻️ Suggested refactor
+ model_id = format_provider_model("openai", get_model("openai", "image_generation")) response = bedrock_client.invoke_model( - modelId="openai/gpt-image-1", + modelId=model_id, contentType="application/json", accept="application/json", body=json.dumps(body), )🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@tests/integrations/python/tests/test_bedrock.py` around lines 2376 - 2394, The test test_39_invoke_openai_image_generation hardcodes "openai/gpt-image-1"; change it to read the configured OpenAI image model from providers.openai.image_generation and pass that value into bedrock_client.invoke_model's modelId instead. Locate the test function (test_39_invoke_openai_image_generation) and replace the literal modelId with the configured value (e.g., providers.openai.image_generation or the project config accessor that returns that key) so the test uses the repo's deterministic image model when calling bedrock_client.invoke_model.core/providers/bedrock/embedding.go (1)
4-4: Usesonicfor the new JSON work in this provider hot path.These
json.Unmarshal/json.Marshalcalls now sit directly on the embedding conversion path here. This repo reservesencoding/jsonfor cases that need stdlib-specific behavior.As per coding guidelines, "Use
github.com/bytedance/sonicfor JSON marshaling in hot paths; use standardencoding/jsononly when custom marshaling is needed."Also applies to: 196-250
🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@core/providers/bedrock/embedding.go` at line 4, Replace use of the stdlib encoding/json in this file with github.com/bytedance/sonic on the hot embedding conversion path: remove "encoding/json" from imports and add "github.com/bytedance/sonic", then change all json.Unmarshal/json.Marshal calls in this file (including the calls around the embedding conversion and the range indicated ~lines 196-250) to use sonic.Unmarshal and sonic.Marshal (preserve the same error handling and variable usages). Ensure the new import is used and no encoding/json references remain.
🤖 Prompt for all review comments with AI agents
Verify each finding against the current code and only fix it if needed.
Inline comments:
In `@core/providers/bedrock/embedding.go`:
- Around line 35-53: The code is silently mutating invalid provider params: when
extracting bifrostReq.Params.ExtraParams["normalize"] you remove it even if it's
not a bool, and when handling max_tokens you truncate float64 to int (e.g., 10.7
-> 10). Change the extraction logic in embedding.go so you only remove and set
titanReq.Normalize when the ExtraParams["normalize"] type-asserts to bool; if it
doesn't, leave the "normalize" entry in bifrostReq.Params.ExtraParams (or
forward it into titanReq.ExtraParams) so downstream validation can handle it.
Similarly, when processing max_tokens (the code around the max tokens handling),
only convert float64 to int if the float is an exact integral value (e.g., float
== float64(int)); if it has a fractional part, do not truncate—leave the
original value (or forward it) instead of silently converting. Ensure you update
the logic that builds titanReq.ExtraParams so invalid-type entries are preserved
rather than dropped or coerced.
In `@core/providers/bedrock/invoke.go`:
- Around line 974-983: The Cohere branch that builds floats in the
strings.Contains(strings.ToLower(model), "cohere") block can append nil entries
when d.Embedding.EmbeddingArray is nil; update the loop in that block (where
floats is built and d.Embedding.EmbeddingArray is accessed) to check for nil
before appending and skip any nil EmbeddingArray entries so the returned
BedrockInvokeCohereEmbeddingResp.Embeddings contains only non-nil []float32
slices; ensure ResponseType remains "embeddings_floats".
---
Duplicate comments:
In `@core/providers/bedrock/embedding.go`:
- Around line 218-245: The code in the loop that builds bifrostResponse.Data
(using typed.Int8, typed.Binary, typed.Uint8, typed.Ubinary and creating
schemas.EmbeddingData / schemas.EmbeddingStruct entries) loses per-entry
encoding identity and iterates maps (non-deterministic order), so consumers
cannot reconstruct types; either add and populate an explicit encoding field on
each EmbeddingData entry (e.g., set Encoding: "int8"/"binary"/"uint8"/"ubinary")
and ensure indexes are unique/consistent, or if schemas.EmbeddingData cannot be
changed, fail fast when multiple different typed slices are present; also
replace the map[...] iteration with a deterministic ordered slice of (name,rows)
pairs to preserve order (e.g., iterate a fixed list ["int8","binary",...]
instead of ranging a map). Ensure error messages reflect the chosen behavior.
In `@tests/integrations/python/tests/test_bedrock.py`:
- Around line 2464-2468: The test only checks for presence of
"embedding"/"embeddings" but not the actual shape; update the assertion to
validate the embeddings payload shape (reuse the same checks as test_32) by
verifying that out contains "inputs" (or "embeddings") and that each entry in
out["inputs"] is a list/sequence of numeric vectors with expected length and
that the number of input vectors matches the number of input items;
specifically, in the test where variables like has_embedding and out are used,
add validations that out["inputs"] is a list, len(out["inputs"]) equals the
number of original inputs, and each element is a list/tuple of numbers
(non-empty) to ensure the inputs converter produced real embedding arrays rather
than empty dicts or malformed structures.
- Around line 2072-2085: The _assert_invoke_images helper currently accepts any
non-empty string; update it to verify each image is a real base64 payload by
stripping an optional data URL prefix (e.g. "data:...;base64,") from each img
string and attempting to base64-decode it (use a strict/validate mode if
available), asserting the decode succeeds and yields non-zero bytes; keep the
existing checks for list type and minimum count and reference
_assert_invoke_images and the images variable when implementing this validation.
- Around line 2338-2343: The test builds the request body with both "image" and
"mask" set to BASE64_IMAGE_LARGE, so it never exercises the image+mask inpaint
path; replace the "mask" value in the body dict with a real Stability inpaint
mask fixture (e.g., a new or existing constant like BASE64_MASK_STABILITY or
STABILITY_INPAINT_MASK) and ensure that fixture contains a proper base64 mask
image rather than the full source image; update any imports or fixture
definitions so the test uses that mask when constructing body.
---
Nitpick comments:
In `@core/providers/bedrock/embedding.go`:
- Line 4: Replace use of the stdlib encoding/json in this file with
github.com/bytedance/sonic on the hot embedding conversion path: remove
"encoding/json" from imports and add "github.com/bytedance/sonic", then change
all json.Unmarshal/json.Marshal calls in this file (including the calls around
the embedding conversion and the range indicated ~lines 196-250) to use
sonic.Unmarshal and sonic.Marshal (preserve the same error handling and variable
usages). Ensure the new import is used and no encoding/json references remain.
In `@tests/integrations/python/tests/test_bedrock.py`:
- Around line 2376-2394: The test test_39_invoke_openai_image_generation
hardcodes "openai/gpt-image-1"; change it to read the configured OpenAI image
model from providers.openai.image_generation and pass that value into
bedrock_client.invoke_model's modelId instead. Locate the test function
(test_39_invoke_openai_image_generation) and replace the literal modelId with
the configured value (e.g., providers.openai.image_generation or the project
config accessor that returns that key) so the test uses the repo's deterministic
image model when calling bedrock_client.invoke_model.
🪄 Autofix (Beta)
Fix all unresolved CodeRabbit comments on this PR:
- Push a commit to this branch (recommended)
- Create a new PR with the fixes
ℹ️ Review info
⚙️ Run configuration
Configuration used: Organization UI
Review profile: CHILL
Plan: Pro
Run ID: d202d766-70e1-436d-93b3-1caff8ef7483
📒 Files selected for processing (10)
core/providers/bedrock/bedrock.gocore/providers/bedrock/embedding.gocore/providers/bedrock/images.gocore/providers/bedrock/invoke.gocore/providers/bedrock/types.gocore/schemas/images.gotests/integrations/python/config.ymltests/integrations/python/tests/test_bedrock.pytests/integrations/python/tests/utils/common.pytransports/bifrost-http/integrations/bedrock.go
🚧 Files skipped from review as they are similar to previous changes (1)
- tests/integrations/python/config.yml
Radheshg04
force-pushed
the
03-24-feat_add_embeddings_image_gen_edit_and_variation_to_bedrock_integration
branch
from
227a943 to
a5e7130
Compare
Contributor
There was a problem hiding this comment.
Actionable comments posted: 2
♻️ Duplicate comments (3)
tests/integrations/python/tests/test_bedrock.py (2)
2345-2348:⚠️ Potential issue | 🟠 MajorUse a real mask fixture in the Stability inpaint case.
maskcurrently reusesBASE64_IMAGE_LARGE, so this path never verifies that the converter forwards a distinct mask.🔧 Suggested fix
body = { "image": BASE64_IMAGE_LARGE, - "mask": BASE64_IMAGE_LARGE, + "mask": BASE64_TITAN_MASK_IMAGE, "prompt": "replace masked area with flowers", "output_format": "png", }🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@tests/integrations/python/tests/test_bedrock.py` around lines 2345 - 2348, The test currently sets the inpaint request body mask to the same constant BASE64_IMAGE_LARGE, so the converter never gets a distinct mask value; update the test in test_bedrock.py to use a dedicated mask fixture (e.g., BASE64_IMAGE_MASK or BASE64_MASK_LARGE) when building the body dict (key "mask") instead of BASE64_IMAGE_LARGE, and adjust or add assertions that the converter received/forwarded the distinct mask value (referencing the body variable and the mask field) to validate the inpaint path correctly.
2077-2090:⚠️ Potential issue | 🟠 MajorValidate that invoke images are actual base64 payloads.
This helper still lets URLs or placeholder strings pass, so a converter that stops returning native Bedrock image bytes would not fail the new image tests.
🔧 Suggested fix
def _assert_invoke_images(response_body: dict, min_images: int = 1) -> None: """Assert that an invoke response contains at least min_images base64 images.""" images = response_body.get("images") or [] assert isinstance(images, list), ( f"Expected 'images' to be a list, got {type(images).__name__}. " f"Response keys: {list(response_body.keys())}" ) assert len(images) >= min_images, ( f"Expected at least {min_images} image(s) in response, got {len(images)}. " f"Response keys: {list(response_body.keys())}" ) for i, img in enumerate(images): assert isinstance(img, str) and len(img) > 0, f"Image {i} is not a non-empty string" + decoded = base64.b64decode(img, validate=True) + assert len(decoded) > 0, f"Image {i} decoded to empty bytes" print(f" ✓ {len(images)} image(s) returned")🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@tests/integrations/python/tests/test_bedrock.py` around lines 2077 - 2090, The _assert_invoke_images helper currently only checks for non-empty strings so URLs/placeholders pass; update _assert_invoke_images to validate that each image string is valid base64 image data by: for each img in images ensure it's a str, then attempt to base64-decode it (use base64.b64decode(..., validate=True)) and assert decoding succeeds and yields non-empty bytes; if you expect data-URI prefixes, strip a leading "data:...;base64," before decoding. Reference the _assert_invoke_images function and add the base64 decode/validate logic and appropriate assertion messages.core/providers/bedrock/embedding.go (1)
35-53:⚠️ Potential issue | 🟡 MinorDon’t silently rewrite invalid provider params.
A non-
boolnormalizevalue is still dropped here, and a fractionalmax_tokenslike10.7is still coerced to10. Both cases change malformed caller input instead of preserving it for downstream validation.As per coding guidelines, "Implement converter functions as pure transformation functions with no HTTP calls, logging, or side effects."🔧 Suggested fix
if bifrostReq.Params != nil { titanReq.Dimensions = bifrostReq.Params.Dimensions + recognizedNormalize := false if normalize, ok := bifrostReq.Params.ExtraParams["normalize"]; ok { if b, ok := normalize.(bool); ok { titanReq.Normalize = &b + recognizedNormalize = true } } // Forward remaining extra params (excluding normalize which is now a first-class field) if len(bifrostReq.Params.ExtraParams) > 0 { extra := make(map[string]interface{}) for k, v := range bifrostReq.Params.ExtraParams { - if k != "normalize" { + if k != "normalize" || !recognizedNormalize { extra[k] = v } } if len(extra) > 0 { titanReq.ExtraParams = extra @@ if v, ok := extra["max_tokens"]; ok { switch n := v.(type) { case int: req.MaxTokens = &n delete(extra, "max_tokens") case float64: - i := int(n) - req.MaxTokens = &i - delete(extra, "max_tokens") + if n == float64(int(n)) { + i := int(n) + req.MaxTokens = &i + delete(extra, "max_tokens") + } } }Also applies to: 140-149
🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@core/providers/bedrock/embedding.go` around lines 35 - 53, The current converter mutates/filters out malformed provider params (e.g., non-bool "normalize" or fractional "max_tokens") instead of preserving them for downstream validation; update the conversion logic in the embedding converter that builds titanReq from bifrostReq.Params so it only moves "normalize" into titanReq.Normalize when the value is actually a bool (leave any non-bool "normalize" entry in ExtraParams), do not coerce numeric types (e.g., preserve floats instead of truncating to int for max_tokens), and otherwise copy all remaining ExtraParams verbatim into titanReq.ExtraParams; ensure this transformation is pure (no HTTP/logging/side-effects) and references the existing symbols bifrostReq.Params, titanReq.Normalize, titanReq.ExtraParams, and bifrostReq.Params.ExtraParams.
🧹 Nitpick comments (2)
core/schemas/embedding.go (1)
172-186: Unmarshal order limitation:[]float64will match before[]int8/[]int32branches.JSON arrays like
[1,2,3]will always unmarshal successfully as[]float64, preventing the integer type branches from being reached. Since JSON doesn't distinguish integer vs floating-point at the wire level, the current order means integer embeddings cannot be recovered from JSON with type fidelity.Based on code analysis, the integer fields (
EmbeddingInt8Array,EmbeddingInt32Array) are primarily populated programmatically by the Bedrock provider (not via JSON unmarshaling) and then serialized back to JSON. If these fields are not expected to round-trip through JSON deserialization, the current implementation is acceptable. However, if JSON round-trip support for integer embeddings is required, the unmarshal order should be reconsidered—for example, attempting stricter integer parsing before falling back to[]float64.🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@core/schemas/embedding.go` around lines 172 - 186, The current Unmarshal path tries []float64 before []int8/[]int32 so numeric JSON arrays like [1,2,3] always become []float64 and integer branches never run; to fix, either attempt integer parsing first (move the int8/int32 Unmarshal attempts above the []float64 branch) or parse into a neutral form (e.g., []interface{} or []json.Number) and inspect each element to decide whether to populate EmbeddingInt8Array or EmbeddingInt32Array (checking integer-ness and range) before falling back to EmbeddingFloatArray; update the Unmarshal calls and branching around Unmarshal(data, &int8Content)/Unmarshal(data, &int32Content) and the []float64 branch accordingly so integer embeddings can round-trip with type fidelity.core/providers/bedrock/invoke.go (1)
532-534: Error strings should be lowercase per Go convention.Several error messages start with uppercase letters. Per the coding guidelines, error strings should be lowercase with no trailing punctuation.
🔧 Proposed fix
case TaskTypeInpainting: if r.InPaintingParams == nil { - return nil, fmt.Errorf("inPaintingParams required for INPAINTING task") + return nil, fmt.Errorf("inPaintingParams required for inpainting task") } ... case TaskTypeOutpainting: if r.OutPaintingParams == nil { - return nil, fmt.Errorf("outPaintingParams required for OUTPAINTING task") + return nil, fmt.Errorf("outPaintingParams required for outpainting task") } ... case TaskTypeBackgroundRemoval: if r.BackgroundRemovalParams == nil { - return nil, fmt.Errorf("backgroundRemovalParams required for BACKGROUND_REMOVAL task") + return nil, fmt.Errorf("backgroundRemovalParams required for background_removal task") }As per coding guidelines: "Error strings should be lowercase with no trailing punctuation (Go convention)."
Also applies to: 567-569, 605-607
🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@core/providers/bedrock/invoke.go` around lines 532 - 534, Several returned error strings in invoke.go begin with uppercase letters or include trailing punctuation; update these to follow Go conventions by making them lowercase with no trailing punctuation. Specifically, edit the error returns in the validation checks for in/out/image paint params (the branches that check r.InPaintingParams, r.OutPaintingParams, r.Image2ImageParams or similar) that currently return messages referencing task names like "INPAINTING", "OUTPAINTING", or "IMAGE2IMAGE" and replace those messages with lowercase, punctuation-free strings (e.g., "inpainting params required for inpainting task" etc.) so all error messages conform to the lowercase/no-trailing-punctuation rule.
🤖 Prompt for all review comments with AI agents
Verify each finding against the current code and only fix it if needed.
Inline comments:
In `@core/bifrost.go`:
- Around line 903-905: The current hasExtraInputs check lets any provider bypass
typed input/embedding validation when ExtraParams contains "inputs" or "images";
restrict this to Bedrock-only by updating the hasExtraInputs assignment to also
verify the request is for Bedrock (e.g., check req.Provider == "bedrock" or
req.Model != nil && req.Model.Provider == "bedrock" as appropriate), so
hasExtraInputs becomes true only when ExtraParams["inputs"]/"images" exist AND
the request is a Bedrock request; keep the rest of the conditional using
isLargePayloadPassthrough(ctx) unchanged.
In `@core/providers/bedrock/types.go`:
- Around line 1111-1120: The fallback wrapper (BedrockInvokeCohereEmbeddingResp)
is dropping native Cohere fields (id, texts) used by
BedrockCohereEmbeddingResponse so downstream BedrockProvider.Embedding and
ToBedrockEmbeddingInvokeResponse lose those fields; update the reconstruction
path to preserve them by either adding id and texts to
BedrockInvokeCohereEmbeddingResp or by having ToBedrockEmbeddingInvokeResponse
copy id/texts from ExtraFields.RawResponse (the original Cohere payload) into
the rebuilt embeddings_floats/embeddings_by_type response; ensure code in
ToBedrockEmbeddingInvokeResponse and the embedding fallback branch that handles
embeddings_floats/embeddings_by_type explicitly transfers the id and texts
fields into the returned envelope so native fields are never dropped.
---
Duplicate comments:
In `@core/providers/bedrock/embedding.go`:
- Around line 35-53: The current converter mutates/filters out malformed
provider params (e.g., non-bool "normalize" or fractional "max_tokens") instead
of preserving them for downstream validation; update the conversion logic in the
embedding converter that builds titanReq from bifrostReq.Params so it only moves
"normalize" into titanReq.Normalize when the value is actually a bool (leave any
non-bool "normalize" entry in ExtraParams), do not coerce numeric types (e.g.,
preserve floats instead of truncating to int for max_tokens), and otherwise copy
all remaining ExtraParams verbatim into titanReq.ExtraParams; ensure this
transformation is pure (no HTTP/logging/side-effects) and references the
existing symbols bifrostReq.Params, titanReq.Normalize, titanReq.ExtraParams,
and bifrostReq.Params.ExtraParams.
In `@tests/integrations/python/tests/test_bedrock.py`:
- Around line 2345-2348: The test currently sets the inpaint request body mask
to the same constant BASE64_IMAGE_LARGE, so the converter never gets a distinct
mask value; update the test in test_bedrock.py to use a dedicated mask fixture
(e.g., BASE64_IMAGE_MASK or BASE64_MASK_LARGE) when building the body dict (key
"mask") instead of BASE64_IMAGE_LARGE, and adjust or add assertions that the
converter received/forwarded the distinct mask value (referencing the body
variable and the mask field) to validate the inpaint path correctly.
- Around line 2077-2090: The _assert_invoke_images helper currently only checks
for non-empty strings so URLs/placeholders pass; update _assert_invoke_images to
validate that each image string is valid base64 image data by: for each img in
images ensure it's a str, then attempt to base64-decode it (use
base64.b64decode(..., validate=True)) and assert decoding succeeds and yields
non-empty bytes; if you expect data-URI prefixes, strip a leading
"data:...;base64," before decoding. Reference the _assert_invoke_images function
and add the base64 decode/validate logic and appropriate assertion messages.
---
Nitpick comments:
In `@core/providers/bedrock/invoke.go`:
- Around line 532-534: Several returned error strings in invoke.go begin with
uppercase letters or include trailing punctuation; update these to follow Go
conventions by making them lowercase with no trailing punctuation. Specifically,
edit the error returns in the validation checks for in/out/image paint params
(the branches that check r.InPaintingParams, r.OutPaintingParams,
r.Image2ImageParams or similar) that currently return messages referencing task
names like "INPAINTING", "OUTPAINTING", or "IMAGE2IMAGE" and replace those
messages with lowercase, punctuation-free strings (e.g., "inpainting params
required for inpainting task" etc.) so all error messages conform to the
lowercase/no-trailing-punctuation rule.
In `@core/schemas/embedding.go`:
- Around line 172-186: The current Unmarshal path tries []float64 before
[]int8/[]int32 so numeric JSON arrays like [1,2,3] always become []float64 and
integer branches never run; to fix, either attempt integer parsing first (move
the int8/int32 Unmarshal attempts above the []float64 branch) or parse into a
neutral form (e.g., []interface{} or []json.Number) and inspect each element to
decide whether to populate EmbeddingInt8Array or EmbeddingInt32Array (checking
integer-ness and range) before falling back to EmbeddingFloatArray; update the
Unmarshal calls and branching around Unmarshal(data,
&int8Content)/Unmarshal(data, &int32Content) and the []float64 branch
accordingly so integer embeddings can round-trip with type fidelity.
🪄 Autofix (Beta)
Fix all unresolved CodeRabbit comments on this PR:
- Push a commit to this branch (recommended)
- Create a new PR with the fixes
ℹ️ Review info
⚙️ Run configuration
Configuration used: Organization UI
Review profile: CHILL
Plan: Pro
Run ID: b2b982c7-6a7b-430a-a361-7f2729ee126c
📒 Files selected for processing (12)
core/bifrost.gocore/providers/bedrock/bedrock.gocore/providers/bedrock/embedding.gocore/providers/bedrock/images.gocore/providers/bedrock/invoke.gocore/providers/bedrock/types.gocore/schemas/embedding.gocore/schemas/images.gotests/integrations/python/config.ymltests/integrations/python/tests/test_bedrock.pytests/integrations/python/tests/utils/common.pytransports/bifrost-http/integrations/bedrock.go
✅ Files skipped from review due to trivial changes (1)
- tests/integrations/python/config.yml
🚧 Files skipped from review as they are similar to previous changes (2)
- core/providers/bedrock/images.go
- transports/bifrost-http/integrations/bedrock.go
akshaydeo
deleted the
03-24-feat_add_embeddings_image_gen_edit_and_variation_to_bedrock_integration
branch
This was referenced Apr 8, 2026
Merged
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
4 participants
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Summary
Adds comprehensive support for Bedrock's native invoke API endpoint, enabling image generation, image editing, image variation, and embedding operations through the
/bedrock/model/{modelId}/invokeroute. This enhancement allows users to interact with Bedrock models using their native payload formats while maintaining cross-provider compatibility.Changes
DetectInvokeRequestTypeto analyze both request body and model ID to properly route embedding, image generation, image editing, and image variation requestsinputText) and Cohere (texts) formats, removing dependency on generic Cohere providerBedrockInvokeRequestwith conversion methods for all supported operation typesType of change
Affected areas
How to test
Run the new Bedrock invoke endpoint integration tests:
The tests cover:
Breaking changes
Related issues
Enhances Bedrock provider capabilities by supporting the full range of operations available through the native invoke API.
Security considerations
No new security implications. Uses existing authentication and authorization mechanisms for Bedrock access.
Checklist
docs/contributing/README.mdand followed the guidelines