Skip to content

feat: add alias-level region/ARN overrides and context-aware model family resolution for Bedrock #4182

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
Pratham-Mishra04 merged 1 commit into from
Jun 9, 2026
Merged

feat: add alias-level region/ARN overrides and context-aware model family resolution for Bedrock #4182

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
65 changes: 31 additions & 34 deletions core/providers/bedrock/bedrock.go
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -205,7 +205,7 @@ var retryableBedrockExceptions = map[string]int{
// Returns the response body, request latency, or an error if the request fails.
func (provider *BedrockProvider) completeRequest(ctx *schemas.BifrostContext, jsonData []byte, path string, key schemas.Key, model string) ([]byte, time.Duration, map[string]string, *schemas.BifrostError) {
config := key.BedrockKeyConfig
region := resolveBedrockRegion(key, model)
region := resolveBedrockRegion(ctx, key, model)

// Create the request with the JSON body
requestURL := fmt.Sprintf("https://bedrock-runtime.%s.amazonaws.com/model/%s", region, path)
Expand Down Expand Up @@ -434,7 +434,7 @@ func (provider *BedrockProvider) completeAgentRuntimeRequest(ctx *schemas.Bifros
// Returns the response body and an error if the request fails.
func (provider *BedrockProvider) makeStreamingRequest(ctx *schemas.BifrostContext, jsonData []byte, key schemas.Key, model string, action string) (*http.Response, *schemas.BifrostError) {
// Parse region and path in one pass to avoid running the regex twice.
path, region := provider.getModelPathAndRegion(action, model, key)
path, region := provider.getModelPathAndRegion(ctx, action, model, key)

// Create HTTP request for streaming
requestURL := fmt.Sprintf("https://bedrock-runtime.%s.amazonaws.com/model/%s", region, path)
Expand Down Expand Up @@ -860,7 +860,7 @@ func (provider *BedrockProvider) TextCompletion(ctx *schemas.BifrostContext, key
return nil, bifrostErr
}

path, _ := provider.getModelPathAndRegion("invoke", request.Model, key)
path, _ := provider.getModelPathAndRegion(ctx, "invoke", request.Model, key)
body, latency, providerResponseHeaders, err := provider.completeRequest(ctx, jsonData, path, key, request.Model)
if providerResponseHeaders != nil {
ctx.SetValue(schemas.BifrostContextKeyProviderResponseHeaders, providerResponseHeaders)
Expand All @@ -872,14 +872,14 @@ func (provider *BedrockProvider) TextCompletion(ctx *schemas.BifrostContext, key
// Handle model-specific response conversion
var bifrostResponse *schemas.BifrostTextCompletionResponse
switch {
case schemas.IsAnthropicModel(request.Model):
case schemas.IsAnthropicModelFamily(ctx, request.Model):
var response BedrockAnthropicTextResponse
if err := sonic.Unmarshal(body, &response); err != nil {
return nil, providerUtils.NewBifrostOperationError("error parsing anthropic response", err)
}
bifrostResponse = response.ToBifrostTextCompletionResponse()

case schemas.IsMistralModel(request.Model):
case schemas.IsMistralModelFamily(ctx, request.Model):
var response BedrockMistralTextResponse
if err := sonic.Unmarshal(body, &response); err != nil {
return nil, providerUtils.NewBifrostOperationError("error parsing mistral response", err)
Expand Down Expand Up @@ -1089,7 +1089,7 @@ func (provider *BedrockProvider) ChatCompletion(ctx *schemas.BifrostContext, key
}

// Format the path with proper model identifier
path, _ := provider.getModelPathAndRegion("converse", request.Model, key)
path, _ := provider.getModelPathAndRegion(ctx, "converse", request.Model, key)

// Create the signed request
responseBody, latency, providerResponseHeaders, bifrostErr := provider.completeRequest(ctx, jsonData, path, key, request.Model)
Expand Down Expand Up @@ -1477,7 +1477,7 @@ func (provider *BedrockProvider) Responses(ctx *schemas.BifrostContext, key sche
}

// Format the path with proper model identifier
path, _ := provider.getModelPathAndRegion("converse", request.Model, key)
path, _ := provider.getModelPathAndRegion(ctx, "converse", request.Model, key)

// Create the signed request
responseBody, latency, providerResponseHeaders, bifrostErr := provider.completeRequest(ctx, jsonData, path, key, request.Model)
Expand Down Expand Up @@ -1837,7 +1837,7 @@ func (provider *BedrockProvider) Embedding(ctx *schemas.BifrostContext, key sche
}

// Determine model type
modelType, err := DetermineEmbeddingModelType(request.Model)
modelType, err := DetermineEmbeddingModelType(ctx, request.Model)
if err != nil {
return nil, providerUtils.NewConfigurationError(err.Error())
}
Expand All @@ -1861,7 +1861,7 @@ func (provider *BedrockProvider) Embedding(ctx *schemas.BifrostContext, key sche
if bifrostError != nil {
return nil, bifrostError
}
path, _ = provider.getModelPathAndRegion("invoke", request.Model, key)
path, _ = provider.getModelPathAndRegion(ctx, "invoke", request.Model, key)
rawResponse, latency, providerResponseHeaders, bifrostError = provider.completeRequest(ctx, jsonData, path, key, request.Model)

case "cohere":
Expand All @@ -1874,7 +1874,7 @@ func (provider *BedrockProvider) Embedding(ctx *schemas.BifrostContext, key sche
if bifrostError != nil {
return nil, bifrostError
}
path, _ = provider.getModelPathAndRegion("invoke", request.Model, key)
path, _ = provider.getModelPathAndRegion(ctx, "invoke", request.Model, key)
rawResponse, latency, providerResponseHeaders, bifrostError = provider.completeRequest(ctx, jsonData, path, key, request.Model)

default:
Expand Down Expand Up @@ -2027,7 +2027,7 @@ func (provider *BedrockProvider) ImageGeneration(ctx *schemas.BifrostContext, ke
var providerResponseHeaders map[string]string
var path string

path, _ = provider.getModelPathAndRegion("invoke", request.Model, key)
path, _ = provider.getModelPathAndRegion(ctx, "invoke", request.Model, key)

jsonData, bifrostError = providerUtils.CheckContextAndGetRequestBody(
ctx,
Expand Down Expand Up @@ -2100,7 +2100,7 @@ func (provider *BedrockProvider) ImageEdit(ctx *schemas.BifrostContext, key sche
var bifrostError *schemas.BifrostError

// Stability AI routing and task-type inference use the actual model ID.
path, _ := provider.getModelPathAndRegion("invoke", request.Model, key)
path, _ := provider.getModelPathAndRegion(ctx, "invoke", request.Model, key)

jsonData, bifrostError = providerUtils.CheckContextAndGetRequestBody(
ctx,
Expand Down Expand Up @@ -2182,7 +2182,7 @@ func (provider *BedrockProvider) ImageVariation(ctx *schemas.BifrostContext, key
}

// Make API request (same URL as image generation)
path, _ := provider.getModelPathAndRegion("invoke", request.Model, key)
path, _ := provider.getModelPathAndRegion(ctx, "invoke", request.Model, key)
rawResponse, latency, providerResponseHeaders, bifrostError := provider.completeRequest(ctx, jsonData, path, key, request.Model)
if providerResponseHeaders != nil {
ctx.SetValue(schemas.BifrostContextKeyProviderResponseHeaders, providerResponseHeaders)
Expand Down Expand Up @@ -3575,32 +3575,29 @@ func (provider *BedrockProvider) BatchResults(ctx *schemas.BifrostContext, keys
return batchResultsResp, nil
}

// resolveBedrockRegion returns the AWS region to use for a request.
// the priority is: model string region > key configured region > default region
func resolveBedrockRegion(key schemas.Key, model string) string {
if region, _ := parseBedrockRegionAndModel(model); region != "" {
return region
}
if key.BedrockKeyConfig != nil && key.BedrockKeyConfig.Region != nil && key.BedrockKeyConfig.Region.GetValue() != "" {
return key.BedrockKeyConfig.Region.GetValue()
}
return DefaultBedrockRegion
}

// getModelPathAndRegion is a helper that calls parseBedrockRegionAndModel
// once and returns both the request path and the AWS signing region
func (provider *BedrockProvider) getModelPathAndRegion(basePath, model string, key schemas.Key) (path, region string) {
// once and returns both the request path and the AWS signing region.
// Honors per-alias Region and BedrockAliasCfg.InferenceProfileARN overrides
// via the resolved alias in ctx.
func (provider *BedrockProvider) getModelPathAndRegion(ctx *schemas.BifrostContext, basePath, model string, key schemas.Key) (path, region string) {
r, bareModel := parseBedrockRegionAndModel(model)
if r == "" {
if key.BedrockKeyConfig != nil && key.BedrockKeyConfig.Region != nil && key.BedrockKeyConfig.Region.GetValue() != "" {
r = key.BedrockKeyConfig.Region.GetValue()
} else {
r = DefaultBedrockRegion
if ra := schemas.GetResolvedAlias(ctx); ra != nil && ra.Config != nil && ra.Config.Region != nil {
if v := ra.Config.Region.GetValue(); v != "" {
r = v
}
}
if r == "" {
if key.BedrockKeyConfig != nil && key.BedrockKeyConfig.Region != nil && key.BedrockKeyConfig.Region.GetValue() != "" {
r = key.BedrockKeyConfig.Region.GetValue()
} else {
r = DefaultBedrockRegion
}
}
}
p := fmt.Sprintf("%s/%s", bareModel, basePath)
if key.BedrockKeyConfig != nil && key.BedrockKeyConfig.ARN != nil && key.BedrockKeyConfig.ARN.GetValue() != "" {
encodedModelIdentifier := url.PathEscape(fmt.Sprintf("%s/%s", key.BedrockKeyConfig.ARN.GetValue(), bareModel))
if arn := resolveBedrockARN(ctx, key); arn != "" {
encodedModelIdentifier := url.PathEscape(fmt.Sprintf("%s/%s", arn, bareModel))
p = fmt.Sprintf("%s/%s", encodedModelIdentifier, basePath)
}
return p, r
Expand All @@ -3627,7 +3624,7 @@ func (provider *BedrockProvider) CountTokens(ctx *schemas.BifrostContext, key sc
}

// Format the path with proper model identifier
path, _ := provider.getModelPathAndRegion("count-tokens", request.Model, key)
path, _ := provider.getModelPathAndRegion(ctx, "count-tokens", request.Model, key)

// Send the request
responseBody, latency, providerResponseHeaders, bifrostErr := provider.completeRequest(ctx, jsonData, path, key, request.Model)
Expand Down
4 changes: 2 additions & 2 deletions core/providers/bedrock/chat.go
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -25,7 +25,7 @@ func ToBedrockChatCompletionRequest(ctx *schemas.BifrostContext, bifrostReq *sch
}

input := bifrostReq.Input
if schemas.IsAnthropicModel(bifrostReq.Model) && ctx.Value(schemas.BifrostContextKeySupportsAssistantPrefill) == false {
if schemas.IsAnthropicModelFamily(ctx, bifrostReq.Model) && ctx.Value(schemas.BifrostContextKeySupportsAssistantPrefill) == false {
trimmed := len(input)
for trimmed > 0 && input[trimmed-1].Role == schemas.ChatMessageRoleAssistant {
trimmed--
Expand All @@ -46,7 +46,7 @@ func ToBedrockChatCompletionRequest(ctx *schemas.BifrostContext, bifrostReq *sch
// Trim trailing whitespace from the last assistant message text blocks
// (only for Anthropic models which use text-based prefill)
lastMsgIndex := len(bedrockReq.Messages) - 1
if schemas.IsAnthropicModel(bifrostReq.Model) && lastMsgIndex >= 0 && bedrockReq.Messages[lastMsgIndex].Role == BedrockMessageRoleAssistant {
if schemas.IsAnthropicModelFamily(ctx, bifrostReq.Model) && lastMsgIndex >= 0 && bedrockReq.Messages[lastMsgIndex].Role == BedrockMessageRoleAssistant {
blocks := bedrockReq.Messages[lastMsgIndex].Content
for j := len(blocks) - 1; j >= 0; j-- {
if blocks[j].Text != nil {
Expand Down
13 changes: 8 additions & 5 deletions core/providers/bedrock/embedding.go
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -3,7 +3,6 @@ package bedrock
import (
"encoding/json"
"fmt"
"strings"

"github.com/maximhq/bifrost/core/schemas"
)
Expand Down Expand Up @@ -159,12 +158,16 @@ func ToBedrockCohereEmbeddingRequest(bifrostReq *schemas.BifrostEmbeddingRequest
return req, nil
}

// DetermineEmbeddingModelType determines the embedding model type from the model name
func DetermineEmbeddingModelType(model string) (string, error) {
// DetermineEmbeddingModelType determines the embedding model type for the
// current attempt. It consults the resolved alias family first
// (model_family / model_name / model_id / alias key) and falls back to the
// substring detectors against the wire model — so an alias to an opaque
// Bedrock deployment that's tagged with the right family routes correctly.
func DetermineEmbeddingModelType(ctx *schemas.BifrostContext, model string) (string, error) {
switch {
case strings.Contains(model, "amazon.titan-embed-text"):
case schemas.IsTitanModelFamily(ctx, model):
return "titan", nil
case strings.Contains(model, "cohere.embed"):
case schemas.IsCohereModelFamily(ctx, model):
return "cohere", nil
default:
return "", fmt.Errorf("unsupported embedding model: %s", model)
Expand Down
6 changes: 3 additions & 3 deletions core/providers/bedrock/invoke.go
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -956,7 +956,7 @@ func ToBedrockInvokeImagesResponse(ctx *schemas.BifrostContext, resp *schemas.Bi
// Bedrock invoke API response format.
// Single-embedding (Titan) responses use: {"embedding": [...], "inputTextTokenCount": N}
// Multi-embedding (Cohere) responses use: {"embeddings": [[...],[...]], "response_type": "embeddings_floats"}
func ToBedrockEmbeddingInvokeResponse(resp *schemas.BifrostEmbeddingResponse) (interface{}, error) {
func ToBedrockEmbeddingInvokeResponse(ctx *schemas.BifrostContext, resp *schemas.BifrostEmbeddingResponse) (interface{}, error) {
if resp == nil {
return nil, fmt.Errorf("bifrost embedding response is nil")
}
Expand All @@ -975,7 +975,7 @@ func ToBedrockEmbeddingInvokeResponse(resp *schemas.BifrostEmbeddingResponse) (i
return &BedrockInvokeEmbeddingResp{InputTextTokenCount: tokenCount}, nil
}

// Use model name to distinguish Cohere from Titan — not batch size.
// Use the resolved family to distinguish Cohere from Titan — not batch size.
// A single-input Cohere request must still return the Cohere envelope format.
model := resp.Model
if model == "" {
Expand All @@ -986,7 +986,7 @@ func ToBedrockEmbeddingInvokeResponse(resp *schemas.BifrostEmbeddingResponse) (i
}
}

if strings.Contains(strings.ToLower(model), "cohere") {
if schemas.IsCohereModelFamily(ctx, model) {
floats := make([][]float32, 0, len(resp.Data))
for _, d := range resp.Data {
float32Emb := make([]float32, len(d.Embedding.EmbeddingArray))
Expand Down
8 changes: 4 additions & 4 deletions core/providers/bedrock/mantle.go
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -60,7 +60,7 @@ func (provider *BedrockProvider) chatCompletionViaMantle(
key schemas.Key,
request *schemas.BifrostChatRequest,
) (*schemas.BifrostChatResponse, *schemas.BifrostError) {
region := resolveBedrockRegion(key, request.Model)
region := resolveBedrockRegion(ctx, key, request.Model)
url := mantleURL(region, "chat/completions")

// Build extraHeaders: always start with network-config headers, then overlay SigV4 if needed.
Expand Down Expand Up @@ -106,7 +106,7 @@ func (provider *BedrockProvider) chatCompletionStreamViaMantle(
key schemas.Key,
request *schemas.BifrostChatRequest,
) (chan *schemas.BifrostStreamChunk, *schemas.BifrostError) {
region := resolveBedrockRegion(key, request.Model)
region := resolveBedrockRegion(ctx, key, request.Model)
url := mantleURL(region, "chat/completions")

// Bearer: identical to Groq / any OpenAI-compatible provider.
Expand Down Expand Up @@ -162,7 +162,7 @@ func (provider *BedrockProvider) responsesViaMantle(
key schemas.Key,
request *schemas.BifrostResponsesRequest,
) (*schemas.BifrostResponsesResponse, *schemas.BifrostError) {
region := resolveBedrockRegion(key, request.Model)
region := resolveBedrockRegion(ctx, key, request.Model)
url := mantleURL(region, "responses")

extraHeaders := make(map[string]string, len(provider.networkConfig.ExtraHeaders))
Expand Down Expand Up @@ -204,7 +204,7 @@ func (provider *BedrockProvider) responsesStreamViaMantle(
key schemas.Key,
request *schemas.BifrostResponsesRequest,
) (chan *schemas.BifrostStreamChunk, *schemas.BifrostError) {
region := resolveBedrockRegion(key, request.Model)
region := resolveBedrockRegion(ctx, key, request.Model)
url := mantleURL(region, "responses")

// Bearer: identical to Groq / any OpenAI-compatible provider.
Expand Down
Loading
Loading