feat: add Bedrock invoke endpoint support for images and embeddings

core/bifrost.go

@@ -900,7 +900,9 @@ func (bifrost *Bifrost) EmbeddingRequest(ctx *schemas.BifrostContext, req *schem
 			},
 		}
 	}
-	if (req.Input == nil || (req.Input.Text == nil && req.Input.Texts == nil && req.Input.Embedding == nil && req.Input.Embeddings == nil)) && !isLargePayloadPassthrough(ctx) {
+	hasExtraInputs := req.Params != nil && req.Params.ExtraParams != nil &&
+		(req.Params.ExtraParams["inputs"] != nil || req.Params.ExtraParams["images"] != nil)
+	if (req.Input == nil || (req.Input.Text == nil && req.Input.Texts == nil && req.Input.Embedding == nil && req.Input.Embeddings == nil)) && !hasExtraInputs && !isLargePayloadPassthrough(ctx) {
 		return nil, &schemas.BifrostError{
 			IsBifrostError: false,
 			Error: &schemas.ErrorField{

core/providers/bedrock/bedrock.go

@@ -26,7 +26,6 @@ import (
 	"github.com/bytedance/sonic"
 	"github.com/google/uuid"
 	"github.com/maximhq/bifrost/core/providers/anthropic"
-	"github.com/maximhq/bifrost/core/providers/cohere"
 	providerUtils "github.com/maximhq/bifrost/core/providers/utils"
 	schemas "github.com/maximhq/bifrost/core/schemas"
 )
@@ -1763,12 +1762,25 @@ func (provider *BedrockProvider) Embedding(ctx *schemas.BifrostContext, key sche
 		bifrostResponse.Model = request.Model
 
 	case "cohere":
-		var cohereResp cohere.CohereEmbeddingResponse
+		var cohereResp BedrockCohereEmbeddingResponse
 		if err := sonic.Unmarshal(rawResponse, &cohereResp); err != nil {
 			return nil, providerUtils.EnrichError(ctx, providerUtils.NewBifrostOperationError("error parsing Cohere embedding response", err, providerName), jsonData, rawResponse, provider.sendBackRawRequest, provider.sendBackRawResponse)
 		}
-		bifrostResponse = cohereResp.ToBifrostEmbeddingResponse()
+		converted, convErr := cohereResp.ToBifrostEmbeddingResponse()
+		if convErr != nil {
+			return nil, providerUtils.EnrichError(ctx, providerUtils.NewBifrostOperationError("error parsing Cohere embedding response", convErr, providerName), jsonData, rawResponse, provider.sendBackRawRequest, provider.sendBackRawResponse)
+		}
+		bifrostResponse = converted
 		bifrostResponse.Model = request.Model
+		// For embeddings_by_type responses preserve the raw Bedrock payload so the
+		// invoke-endpoint converter can return all encoding variants verbatim, since
+		// the internal BifrostEmbeddingResponse only has float32 and string fields.
+		if cohereResp.ResponseType == "embeddings_by_type" {
+			var rawResponseData interface{}
+			if err := sonic.Unmarshal(rawResponse, &rawResponseData); err == nil {
+				bifrostResponse.ExtraFields.RawResponse = rawResponseData
+			}
+		}
 	}
 
 	// Set ExtraFields

core/providers/bedrock/embedding.go

@@ -1,10 +1,10 @@
 package bedrock
 
 import (
+	"encoding/json"
 	"fmt"
 	"strings"
 
-	"github.com/maximhq/bifrost/core/providers/cohere"
 	"github.com/maximhq/bifrost/core/schemas"
 )
 
@@ -19,11 +19,6 @@ func ToBedrockTitanEmbeddingRequest(bifrostReq *schemas.BifrostEmbeddingRequest)
 		return nil, fmt.Errorf("no input text provided for embedding")
 	}
 
-	// Validate dimensions parameter - Titan models do not support it
-	if bifrostReq.Params != nil && bifrostReq.Params.Dimensions != nil {
-		return nil, fmt.Errorf("amazon Titan embedding models do not support custom dimensions parameter")
-	}
-
 	titanReq := &BedrockTitanEmbeddingRequest{}
 
 	// Set input text
@@ -36,8 +31,26 @@ func ToBedrockTitanEmbeddingRequest(bifrostReq *schemas.BifrostEmbeddingRequest)
 		}
 		titanReq.InputText = embeddingText
 	}
+
 	if bifrostReq.Params != nil {
-		titanReq.ExtraParams = bifrostReq.Params.ExtraParams
+		titanReq.Dimensions = bifrostReq.Params.Dimensions
+		if normalize, ok := bifrostReq.Params.ExtraParams["normalize"]; ok {
+			if b, ok := normalize.(bool); ok {
+				titanReq.Normalize = &b
+			}
+		}
+		// 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" {
+					extra[k] = v
+				}
+			}
+			if len(extra) > 0 {
+				titanReq.ExtraParams = extra
+			}
+		}
 	}
 
 	return titanReq, nil
@@ -69,20 +82,81 @@ func (response *BedrockTitanEmbeddingResponse) ToBifrostEmbeddingResponse() *sch
 	return bifrostResponse
 }
 
-// ToBedrockCohereEmbeddingRequest converts a Bifrost embedding request to Bedrock Cohere format
-// Reuses the Cohere converter since the format is identical
-func ToBedrockCohereEmbeddingRequest(bifrostReq *schemas.BifrostEmbeddingRequest) (*cohere.CohereEmbeddingRequest, error) {
+// ToBedrockCohereEmbeddingRequest converts a Bifrost embedding request to Bedrock Cohere format.
+// Unlike the direct Cohere API, Bedrock does not accept a "model" field in the request body.
+func ToBedrockCohereEmbeddingRequest(bifrostReq *schemas.BifrostEmbeddingRequest) (*BedrockCohereEmbeddingRequest, error) {
 	if bifrostReq == nil {
 		return nil, fmt.Errorf("bifrost embedding request is nil")
 	}
+	if bifrostReq.Input == nil {
+		return nil, fmt.Errorf("no input provided for embedding")
+	}
 
-	// Reuse Cohere's converter - the format is identical for Bedrock
-	cohereReq := cohere.ToCohereEmbeddingRequest(bifrostReq)
-	if cohereReq == nil {
-		return nil, fmt.Errorf("failed to convert to Cohere embedding request")
+	req := &BedrockCohereEmbeddingRequest{}
+
+	// Map texts
+	if bifrostReq.Input.Text != nil {
+		req.Texts = []string{*bifrostReq.Input.Text}
+	} else if len(bifrostReq.Input.Texts) > 0 {
+		req.Texts = bifrostReq.Input.Texts
 	}
 
-	return cohereReq, nil
+	if bifrostReq.Params != nil {
+		extra := make(map[string]interface{}, len(bifrostReq.Params.ExtraParams))
+		for k, v := range bifrostReq.Params.ExtraParams {
+			extra[k] = v
+		}
+
+		if v, ok := extra["input_type"]; ok {
+			if s, ok := v.(string); ok {
+				req.InputType = s
+				delete(extra, "input_type")
+			}
+		}
+		if v, ok := extra["truncate"]; ok {
+			if s, ok := v.(string); ok {
+				req.Truncate = &s
+				delete(extra, "truncate")
+			}
+		}
+		if v, ok := extra["embedding_types"]; ok {
+			if ss, ok := v.([]string); ok {
+				req.EmbeddingTypes = ss
+				delete(extra, "embedding_types")
+			}
+		}
+		if v, ok := extra["images"]; ok {
+			if ss, ok := v.([]string); ok {
+				req.Images = ss
+				delete(extra, "images")
+			}
+		}
+		if v, ok := extra["inputs"]; ok {
+			if inputs, ok := v.([]BedrockCohereEmbeddingInput); ok {
+				req.Inputs = inputs
+				delete(extra, "inputs")
+			}
+		}
+		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 bifrostReq.Params.Dimensions != nil {
+			req.OutputDimension = bifrostReq.Params.Dimensions
+		}
+		if len(extra) > 0 {
+			req.ExtraParams = extra
+		}
+	}
+
+	return req, nil
 }
 
 // DetermineEmbeddingModelType determines the embedding model type from the model name
@@ -96,3 +170,102 @@ func DetermineEmbeddingModelType(model string) (string, error) {
 		return "", fmt.Errorf("unsupported embedding model: %s", model)
 	}
 }
+
+// ToBifrostEmbeddingResponse converts a BedrockCohereEmbeddingResponse to Bifrost format.
+// Bedrock returns embeddings as a raw [][]float32 when response_type is "embeddings_floats"
+// (the default, when no embedding_types are requested), and as a typed object when
+// response_type is "embeddings_by_type".
+func (r *BedrockCohereEmbeddingResponse) ToBifrostEmbeddingResponse() (*schemas.BifrostEmbeddingResponse, error) {
+	if r == nil {
+		return nil, fmt.Errorf("nil Bedrock Cohere embedding response")
+	}
+
+	bifrostResponse := &schemas.BifrostEmbeddingResponse{Object: "list"}
+
+	switch r.ResponseType {
+	case "embeddings_by_type":
+		// Object form: {"float": [[...]], "int8": [[...]], "uint8": [[...]], "binary": [[...]], "ubinary": [[...]], "base64": [...]}
+		var typed struct {
+			Float   [][]float32 `json:"float"`
+			Base64  []string    `json:"base64"`
+			Int8    [][]int8    `json:"int8"`
+			Uint8   [][]int32   `json:"uint8"`  // int32 avoids []byte→base64 JSON issue
+			Binary  [][]int8    `json:"binary"`
+			Ubinary [][]int32   `json:"ubinary"` // int32 avoids []byte→base64 JSON issue
+		}
+		if err := json.Unmarshal(r.Embeddings, &typed); err != nil {
+			return nil, fmt.Errorf("error parsing embeddings_by_type: %w", err)
+		}
+		if typed.Float != nil {
+			for i, emb := range typed.Float {
+				float64Emb := make([]float64, len(emb))
+				for j, v := range emb {
+					float64Emb[j] = float64(v)
+				}
+				bifrostResponse.Data = append(bifrostResponse.Data, schemas.EmbeddingData{
+					Object:    "embedding",
+					Index:     i,
+					Embedding: schemas.EmbeddingStruct{EmbeddingArray: float64Emb},
+				})
+			}
+		}
+		if typed.Base64 != nil {
+			for i, emb := range typed.Base64 {
+				e := emb
+				bifrostResponse.Data = append(bifrostResponse.Data, schemas.EmbeddingData{
+					Object:    "embedding",
+					Index:     i,
+					Embedding: schemas.EmbeddingStruct{EmbeddingStr: &e},
+				})
+			}
+		}
+		for i, emb := range typed.Int8 {
+			bifrostResponse.Data = append(bifrostResponse.Data, schemas.EmbeddingData{
+				Object:    "embedding",
+				Index:     i,
+				Embedding: schemas.EmbeddingStruct{EmbeddingInt8Array: emb},
+			})
+		}
+		for i, emb := range typed.Binary {
+			bifrostResponse.Data = append(bifrostResponse.Data, schemas.EmbeddingData{
+				Object:    "embedding",
+				Index:     i,
+				Embedding: schemas.EmbeddingStruct{EmbeddingInt8Array: emb},
+			})
+		}
+		for i, emb := range typed.Uint8 {
+			bifrostResponse.Data = append(bifrostResponse.Data, schemas.EmbeddingData{
+				Object:    "embedding",
+				Index:     i,
+				Embedding: schemas.EmbeddingStruct{EmbeddingInt32Array: emb},
+			})
+		}
+		for i, emb := range typed.Ubinary {
+			bifrostResponse.Data = append(bifrostResponse.Data, schemas.EmbeddingData{
+				Object:    "embedding",
+				Index:     i,
+				Embedding: schemas.EmbeddingStruct{EmbeddingInt32Array: emb},
+			})
+		}
+
+	default:
+		// Default / "embeddings_floats": raw array form [[...], [...]]
+		var floats [][]float32
+		if err := json.Unmarshal(r.Embeddings, &floats); err != nil {
+			return nil, fmt.Errorf("error parsing embeddings_floats: %w", err)
+		}
+		for i, emb := range floats {
+			float64Emb := make([]float64, len(emb))
+			for j, v := range emb {
+				float64Emb[j] = float64(v)
+			}
+			bifrostResponse.Data = append(bifrostResponse.Data, schemas.EmbeddingData{
+				Object:    "embedding",
+				Index:     i,
+				Embedding: schemas.EmbeddingStruct{EmbeddingArray: float64Emb},
+			})
+		}
+	}
+
+	return bifrostResponse, nil
+}

core/providers/bedrock/images.go

@@ -39,6 +39,14 @@ func isStabilityAIModel(model string) bool {
 	return strings.Contains(strings.ToLower(model), "stability.")
 }
 
+// isPromptOnlyImageGenerationModel returns true for image generation models that use a flat
+// {"prompt": "..."} payload (no taskType field). Covers Vertex Imagen and similar models.
+// Stability AI is excluded here — it's handled separately because it also supports image edit.
+func isPromptOnlyImageGenerationModel(model string) bool {
+	m := strings.ToLower(model)
+	return strings.Contains(m, "image")
+}
+
 // ToStabilityAIImageGenerationRequest converts a Bifrost image generation request to the Stability AI
 // flat request format used by Bedrock (stability.stable-image-* models).
 func ToStabilityAIImageGenerationRequest(request *schemas.BifrostImageGenerationRequest) (*StabilityAIImageGenerationRequest, error) {
@@ -148,6 +156,24 @@ func ToBedrockImageGenerationRequest(request *schemas.BifrostImageGenerationRequ
 
 }
 
+// ToStabilityAIImageGenerationResponse converts a BifrostImageGenerationResponse back to
+// the native Bedrock invoke API response format used by Stability AI models.
+// Stability AI models use the same BedrockImageGenerationResponse format as Titan/Nova Canvas.
+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
+	}
+	return result, nil
+}
+
 // ToBedrockImageVariationRequest converts a Bifrost image variation request to a Bedrock image variation request
 func ToBedrockImageVariationRequest(request *schemas.BifrostImageVariationRequest) (*BedrockImageVariationRequest, error) {
 	if request == nil {
@@ -699,6 +725,13 @@ func ToBifrostImageGenerationResponse(response *BedrockImageGenerationResponse)
 
 	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,