From a6b31b55c5f08e2bebce3c353a1541fbfc414383 Mon Sep 17 00:00:00 2001 From: Pratham-Mishra04 Date: Tue, 2 Jun 2026 18:00:46 +0530 Subject: [PATCH] feat: add bedrock alias config support --- core/providers/bedrock/bedrock.go | 65 +++++++------ core/providers/bedrock/chat.go | 4 +- core/providers/bedrock/embedding.go | 13 ++- core/providers/bedrock/invoke.go | 6 +- core/providers/bedrock/mantle.go | 8 +- core/providers/bedrock/region_test.go | 92 ++++++++++++++++++- core/providers/bedrock/responses.go | 28 +++--- core/providers/bedrock/utils.go | 66 ++++++++++--- core/schemas/account.go | 44 ++++++++- core/schemas/utils.go | 13 +++ .../bifrost-http/integrations/bedrock.go | 2 +- 11 files changed, 262 insertions(+), 79 deletions(-) diff --git a/core/providers/bedrock/bedrock.go b/core/providers/bedrock/bedrock.go index 8b7c169b04e..8457a5bd132 100644 --- a/core/providers/bedrock/bedrock.go +++ b/core/providers/bedrock/bedrock.go @@ -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) @@ -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) @@ -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) @@ -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) @@ -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) @@ -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) @@ -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()) } @@ -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": @@ -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: @@ -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, @@ -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, @@ -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) @@ -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 @@ -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) diff --git a/core/providers/bedrock/chat.go b/core/providers/bedrock/chat.go index ed089e70986..175b5973c20 100644 --- a/core/providers/bedrock/chat.go +++ b/core/providers/bedrock/chat.go @@ -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-- @@ -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 { diff --git a/core/providers/bedrock/embedding.go b/core/providers/bedrock/embedding.go index cb4ef19e88b..b049bcfe611 100644 --- a/core/providers/bedrock/embedding.go +++ b/core/providers/bedrock/embedding.go @@ -3,7 +3,6 @@ package bedrock import ( "encoding/json" "fmt" - "strings" "github.com/maximhq/bifrost/core/schemas" ) @@ -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) diff --git a/core/providers/bedrock/invoke.go b/core/providers/bedrock/invoke.go index 65a8665af93..3edcd4e12c3 100644 --- a/core/providers/bedrock/invoke.go +++ b/core/providers/bedrock/invoke.go @@ -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") } @@ -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 == "" { @@ -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)) diff --git a/core/providers/bedrock/mantle.go b/core/providers/bedrock/mantle.go index 5c85a1cbdcd..932a9df5442 100644 --- a/core/providers/bedrock/mantle.go +++ b/core/providers/bedrock/mantle.go @@ -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. @@ -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. @@ -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)) @@ -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. diff --git a/core/providers/bedrock/region_test.go b/core/providers/bedrock/region_test.go index cbd8b3d1ccb..55bdf538e9f 100644 --- a/core/providers/bedrock/region_test.go +++ b/core/providers/bedrock/region_test.go @@ -59,7 +59,7 @@ func TestGetModelPathStripsRegion(t *testing.T) { } for _, tc := range cases { t.Run(tc.model, func(t *testing.T) { - got, _ := provider.getModelPathAndRegion(tc.basePath, tc.model, key) + got, _ := provider.getModelPathAndRegion(nil, tc.basePath, tc.model, key) assert.Equal(t, tc.wantPath, got) }) } @@ -91,12 +91,98 @@ func TestGetModelPathStripsRegionWithARN(t *testing.T) { } for _, tc := range cases { t.Run(tc.model, func(t *testing.T) { - got, _ := provider.getModelPathAndRegion("converse", tc.model, key) + got, _ := provider.getModelPathAndRegion(nil, "converse", tc.model, key) assert.Equal(t, tc.wantPath, got) }) } } +// TestResolveBedrockRegion_AliasOverride verifies the per-alias Region +// override slots between the model-string prefix (highest priority) and the +// key-level Region (lower priority). +func TestResolveBedrockRegion_AliasOverride(t *testing.T) { + keyRegion := "us-east-1" + aliasRegion := "us-west-2" + key := schemas.Key{ + BedrockKeyConfig: &schemas.BedrockKeyConfig{ + Region: schemas.NewEnvVar(keyRegion), + }, + } + + // Build ctx carrying an alias with Region override. + ctx := schemas.NewBifrostContext(nil, schemas.NoDeadline) + ctx.SetValue(schemas.BifrostContextKeyResolvedAlias, &schemas.ResolvedAlias{ + Key: "best-claude", + Config: &schemas.AliasConfig{ + ModelID: "anthropic.claude-3-5-sonnet-20241022-v2:0", + Region: schemas.NewEnvVar(aliasRegion), + }, + }) + + // Bare model — alias.Region wins over key.Region. + if got := resolveBedrockRegion(ctx, key, "anthropic.claude-3-5-sonnet-20241022-v2:0"); got != aliasRegion { + t.Errorf("alias override should win over key region: got %q, want %q", got, aliasRegion) + } + + // Model string with explicit region prefix — wins over alias override. + if got := resolveBedrockRegion(ctx, key, "eu-west-1/anthropic.claude-v2"); got != "eu-west-1" { + t.Errorf("model-string region should win over alias override: got %q", got) + } + + // No alias in ctx — falls through to key.Region. + emptyCtx := schemas.NewBifrostContext(nil, schemas.NoDeadline) + if got := resolveBedrockRegion(emptyCtx, key, "anthropic.claude-v2"); got != keyRegion { + t.Errorf("no alias: should use key.Region: got %q, want %q", got, keyRegion) + } +} + +// TestResolveBedrockARN_AliasOverride verifies the BedrockAliasCfg +// InferenceProfileARN override takes precedence over key.BedrockKeyConfig.ARN. +func TestResolveBedrockARN_AliasOverride(t *testing.T) { + keyARN := "arn:aws:bedrock:us-east-1:1234567890:resource-config/default" + aliasARN := "arn:aws:bedrock:us-east-1:1234567890:inference-profile/us.anthropic.claude-3-7-sonnet" + key := schemas.Key{ + BedrockKeyConfig: &schemas.BedrockKeyConfig{ + ARN: schemas.NewEnvVar(keyARN), + }, + } + + // No alias — falls back to key.ARN. + if got := resolveBedrockARN(nil, key); got != keyARN { + t.Errorf("nil ctx: got %q, want key ARN %q", got, keyARN) + } + + // Alias with InferenceProfileARN override wins. + ctx := schemas.NewBifrostContext(nil, schemas.NoDeadline) + ctx.SetValue(schemas.BifrostContextKeyResolvedAlias, &schemas.ResolvedAlias{ + Key: "best-claude", + Config: &schemas.AliasConfig{ + ModelID: "anthropic.claude-3-7-sonnet-20250219-v1:0", + BedrockAliasCfg: &schemas.BedrockAliasCfg{ + InferenceProfileARN: schemas.NewEnvVar(aliasARN), + }, + }, + }) + if got := resolveBedrockARN(ctx, key); got != aliasARN { + t.Errorf("alias override should win: got %q, want %q", got, aliasARN) + } + + // Empty alias ARN — falls through to key.ARN. + ctx2 := schemas.NewBifrostContext(nil, schemas.NoDeadline) + ctx2.SetValue(schemas.BifrostContextKeyResolvedAlias, &schemas.ResolvedAlias{ + Key: "x", + Config: &schemas.AliasConfig{ + ModelID: "x", + BedrockAliasCfg: &schemas.BedrockAliasCfg{ + InferenceProfileARN: schemas.NewEnvVar(""), + }, + }, + }) + if got := resolveBedrockARN(ctx2, key); got != keyARN { + t.Errorf("empty alias ARN should fall through to key ARN: got %q, want %q", got, keyARN) + } +} + func TestResolveBedrockRegion(t *testing.T) { configuredRegion := "ap-southeast-1" key := schemas.Key{ @@ -119,7 +205,7 @@ func TestResolveBedrockRegion(t *testing.T) { } for _, tc := range cases { t.Run(tc.desc, func(t *testing.T) { - got := resolveBedrockRegion(tc.key, tc.model) + got := resolveBedrockRegion(nil, tc.key, tc.model) assert.Equal(t, tc.wantRegion, got) }) } diff --git a/core/providers/bedrock/responses.go b/core/providers/bedrock/responses.go index e081f6bacca..4132021b1cb 100644 --- a/core/providers/bedrock/responses.go +++ b/core/providers/bedrock/responses.go @@ -1921,7 +1921,7 @@ func (request *BedrockConverseRequest) ToBifrostResponsesRequest(ctx *schemas.Bi continue } bifrostReq.Params.Tools = append(bifrostReq.Params.Tools, schemas.ResponsesTool{Type: toolType}) - } else if tool.CachePoint != nil && !schemas.IsNovaModel(bifrostReq.Model) { + } else if tool.CachePoint != nil && !schemas.IsNovaModelFamily(ctx, bifrostReq.Model) { // add cache control to last tool in tools array if len(bifrostReq.Params.Tools) > 0 { bifrostReq.Params.Tools[len(bifrostReq.Params.Tools)-1].CacheControl = &schemas.CacheControl{ @@ -2018,7 +2018,7 @@ func (request *BedrockConverseRequest) ToBifrostResponsesRequest(ctx *schemas.Bi if request.InferenceConfig != nil && request.InferenceConfig.MaxTokens != nil { defaultMaxTokens = *request.InferenceConfig.MaxTokens } - if schemas.IsAnthropicModel(bifrostReq.Model) { + if schemas.IsAnthropicModelFamily(ctx, bifrostReq.Model) { minBudgetTokens = anthropic.MinimumReasoningMaxTokens } effort := providerUtils.GetReasoningEffortFromBudgetTokens(maxTokens, minBudgetTokens, defaultMaxTokens) @@ -2151,7 +2151,7 @@ func ToBedrockResponsesRequest(ctx *schemas.BifrostContext, bifrostReq *schemas. // map bifrost messages to bedrock messages using the new conversion method if bifrostReq.Input != nil { 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 != nil && *input[trimmed-1].Role == schemas.ResponsesInputMessageRoleAssistant { trimmed-- @@ -2180,7 +2180,7 @@ func ToBedrockResponsesRequest(ctx *schemas.BifrostContext, bifrostReq *schemas. // 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 { @@ -2217,15 +2217,15 @@ func ToBedrockResponsesRequest(ctx *schemas.BifrostContext, bifrostReq *schemas. // setting it to default max tokens tokenBudget = anthropic.MinimumReasoningMaxTokens } - if schemas.IsAnthropicModel(bifrostReq.Model) && tokenBudget < anthropic.MinimumReasoningMaxTokens { + if schemas.IsAnthropicModelFamily(ctx, bifrostReq.Model) && tokenBudget < anthropic.MinimumReasoningMaxTokens { return nil, fmt.Errorf("reasoning.max_tokens must be >= %d for anthropic", anthropic.MinimumReasoningMaxTokens) } - if schemas.IsAnthropicModel(bifrostReq.Model) { + if schemas.IsAnthropicModelFamily(ctx, bifrostReq.Model) { bedrockReq.AdditionalModelRequestFields.Set("thinking", map[string]any{ "type": "enabled", "budget_tokens": tokenBudget, }) - } else if schemas.IsNovaModel(bifrostReq.Model) { + } else if schemas.IsNovaModelFamily(ctx, bifrostReq.Model) { minBudgetTokens := MinimumReasoningMaxTokens modelDefaultMaxTokens := providerUtils.GetMaxOutputTokensOrDefault(bifrostReq.Model, DefaultCompletionMaxTokens) defaultMaxTokens := modelDefaultMaxTokens @@ -2258,7 +2258,7 @@ func ToBedrockResponsesRequest(ctx *schemas.BifrostContext, bifrostReq *schemas. } } else { if bifrostReq.Params.Reasoning.Effort != nil && *bifrostReq.Params.Reasoning.Effort != "none" { - if schemas.IsNovaModel(bifrostReq.Model) { + if schemas.IsNovaModelFamily(ctx, bifrostReq.Model) { effort := *bifrostReq.Params.Reasoning.Effort typeStr := "enabled" switch effort { @@ -2283,7 +2283,7 @@ func ToBedrockResponsesRequest(ctx *schemas.BifrostContext, bifrostReq *schemas. } bedrockReq.AdditionalModelRequestFields.Set("reasoningConfig", config) - } else if schemas.IsAnthropicModel(bifrostReq.Model) { + } else if schemas.IsAnthropicModelFamily(ctx, bifrostReq.Model) { if anthropic.SupportsAdaptiveThinking(bifrostReq.Model) { // Opus 4.6+: adaptive thinking + output_config.effort effort := anthropic.MapBifrostEffortToAnthropic(*bifrostReq.Params.Reasoning.Effort) @@ -2338,11 +2338,11 @@ func ToBedrockResponsesRequest(ctx *schemas.BifrostContext, bifrostReq *schemas. }) } } else { - if schemas.IsAnthropicModel(bifrostReq.Model) { + if schemas.IsAnthropicModelFamily(ctx, bifrostReq.Model) { bedrockReq.AdditionalModelRequestFields.Set("thinking", map[string]any{ "type": "disabled", }) - } else if schemas.IsNovaModel(bifrostReq.Model) { + } else if schemas.IsNovaModelFamily(ctx, bifrostReq.Model) { bedrockReq.AdditionalModelRequestFields.Set("reasoningConfig", map[string]any{ "type": "disabled", }) @@ -2449,7 +2449,7 @@ func ToBedrockResponsesRequest(ctx *schemas.BifrostContext, bifrostReq *schemas. } bedrockTools = append(bedrockTools, bedrockTool) - if tool.CacheControl != nil && !schemas.IsNovaModel(bifrostReq.Model) { + if tool.CacheControl != nil && !schemas.IsNovaModelFamily(ctx, bifrostReq.Model) { bedrockTools = append(bedrockTools, BedrockTool{ CachePoint: &BedrockCachePoint{ Type: BedrockCachePointTypeDefault, @@ -2480,7 +2480,7 @@ func ToBedrockResponsesRequest(ctx *schemas.BifrostContext, bifrostReq *schemas. // behavior. See per-model support matrix at // https://docs.aws.amazon.com/bedrock/latest/APIReference/API_runtime_ToolChoice.html // (mirrors the gate in convertToolConfigFromFiltered for ChatCompletions). - if bedrockToolChoice != nil && bedrockToolChoice.Tool != nil && schemas.IsLlamaModel(bifrostReq.Model) { + if bedrockToolChoice != nil && bedrockToolChoice.Tool != nil && schemas.IsLlamaModelFamily(ctx, bifrostReq.Model) { bedrockToolChoice = nil } if bedrockToolChoice != nil { @@ -2510,7 +2510,7 @@ func ToBedrockResponsesRequest(ctx *schemas.BifrostContext, bifrostReq *schemas. thinkingEnabled := bifrostReq.Params.Reasoning != nil && (bifrostReq.Params.Reasoning.MaxTokens != nil || (bifrostReq.Params.Reasoning.Effort != nil && *bifrostReq.Params.Reasoning.Effort != "none")) - if !schemas.IsLlamaModel(bifrostReq.Model) && !thinkingEnabled { + if !schemas.IsLlamaModelFamily(ctx, bifrostReq.Model) && !thinkingEnabled { bedrockReq.ToolConfig.ToolChoice = &BedrockToolChoice{ Tool: &BedrockToolChoiceTool{ Name: responsesStructuredOutputTool.ToolSpec.Name, diff --git a/core/providers/bedrock/utils.go b/core/providers/bedrock/utils.go index 8ec4ee51cfe..bf86916acfd 100644 --- a/core/providers/bedrock/utils.go +++ b/core/providers/bedrock/utils.go @@ -41,6 +41,43 @@ func parseBedrockRegionAndModel(model string) (region, bareModel string) { return "", model } +// resolveBedrockRegion returns the AWS region to use for a request. +// Priority: model-string region prefix > alias-level Region > key-level +// BedrockKeyConfig.Region > DefaultBedrockRegion. The model-string prefix +// stays highest since it's the most explicit signal — when an admin types a +// region into their model ID they expect that to win. +func resolveBedrockRegion(ctx *schemas.BifrostContext, key schemas.Key, model string) string { + if region, _ := parseBedrockRegionAndModel(model); region != "" { + return region + } + if ra := schemas.GetResolvedAlias(ctx); ra != nil && ra.Config != nil && ra.Config.Region != nil { + if v := ra.Config.Region.GetValue(); v != "" { + return v + } + } + if key.BedrockKeyConfig != nil && key.BedrockKeyConfig.Region != nil && key.BedrockKeyConfig.Region.GetValue() != "" { + return key.BedrockKeyConfig.Region.GetValue() + } + return DefaultBedrockRegion +} + +// resolveBedrockARN returns the inference-profile / resource ARN prepended +// to the Bedrock URL path. Priority: alias-level BedrockAliasCfg +// InferenceProfileARN > key-level BedrockKeyConfig.ARN. Returns empty when +// neither is set, in which case getModelPathAndRegion emits the bare model +// path. +func resolveBedrockARN(ctx *schemas.BifrostContext, key schemas.Key) string { + if ra := schemas.GetResolvedAlias(ctx); ra != nil && ra.Config != nil && ra.Config.BedrockAliasCfg != nil && ra.Config.BedrockAliasCfg.InferenceProfileARN != nil { + if v := ra.Config.BedrockAliasCfg.InferenceProfileARN.GetValue(); v != "" { + return v + } + } + if key.BedrockKeyConfig != nil && key.BedrockKeyConfig.ARN != nil { + return key.BedrockKeyConfig.ARN.GetValue() + } + return "" +} + var ( invalidCharRegex = regexp.MustCompile(`[^a-zA-Z0-9\s\-\(\)\[\]]`) multiSpaceRegex = regexp.MustCompile(`\s{2,}`) @@ -157,7 +194,7 @@ func bedrockAliasToolName(ctx context.Context, name string) string { } alias := hash + "_" + semanticName - if bifrostCtx, ok := ctx.(*schemas.BifrostContext); ok && alias != name { + if bifrostCtx, ok := ctx.(*schemas.BifrostContext); ok && bifrostCtx != nil && alias != name { aliases, _ := bifrostCtx.Value(bedrockToolNameAliasKey{}).(map[string]string) if aliases == nil { aliases = make(map[string]string) @@ -170,7 +207,7 @@ func bedrockAliasToolName(ctx context.Context, name string) string { // bedrockRestoreToolName maps a Bedrock wire-name alias back to the caller's tool name. func bedrockRestoreToolName(ctx context.Context, name string) string { - if bifrostCtx, ok := ctx.(*schemas.BifrostContext); ok { + if bifrostCtx, ok := ctx.(*schemas.BifrostContext); ok && bifrostCtx != nil { if aliases, _ := bifrostCtx.Value(bedrockToolNameAliasKey{}).(map[string]string); aliases != nil { if original, ok := aliases[name]; ok { return original @@ -254,7 +291,7 @@ func convertChatParameters(ctx *schemas.BifrostContext, bifrostReq *schemas.Bifr // setting it to default max tokens tokenBudget = anthropic.MinimumReasoningMaxTokens } - if schemas.IsAnthropicModel(bifrostReq.Model) { + if schemas.IsAnthropicModelFamily(ctx, bifrostReq.Model) { if tokenBudget < anthropic.MinimumReasoningMaxTokens { return fmt.Errorf("reasoning.max_tokens must be >= %d for anthropic", anthropic.MinimumReasoningMaxTokens) } @@ -262,7 +299,7 @@ func convertChatParameters(ctx *schemas.BifrostContext, bifrostReq *schemas.Bifr "type": "enabled", "budget_tokens": tokenBudget, }) - } else if schemas.IsNovaModel(bifrostReq.Model) { + } else if schemas.IsNovaModelFamily(ctx, bifrostReq.Model) { minBudgetTokens := MinimumReasoningMaxTokens modelDefaultMaxTokens := providerUtils.GetMaxOutputTokensOrDefault(bifrostReq.Model, DefaultCompletionMaxTokens) defaultMaxTokens := modelDefaultMaxTokens @@ -319,7 +356,7 @@ func convertChatParameters(ctx *schemas.BifrostContext, bifrostReq *schemas.Bifr } } } - if schemas.IsNovaModel(bifrostReq.Model) { + if schemas.IsNovaModelFamily(ctx, bifrostReq.Model) { effort := *bifrostReq.Params.Reasoning.Effort typeStr := "enabled" switch effort { @@ -343,7 +380,7 @@ func convertChatParameters(ctx *schemas.BifrostContext, bifrostReq *schemas.Bifr } bedrockReq.AdditionalModelRequestFields.Set("reasoningConfig", config) - } else if schemas.IsAnthropicModel(bifrostReq.Model) { + } else if schemas.IsAnthropicModelFamily(ctx, bifrostReq.Model) { if anthropic.SupportsAdaptiveThinking(bifrostReq.Model) { // Opus 4.6+: adaptive thinking + output_config.effort effort := anthropic.MapBifrostEffortToAnthropic(*bifrostReq.Params.Reasoning.Effort) @@ -371,11 +408,11 @@ func convertChatParameters(ctx *schemas.BifrostContext, bifrostReq *schemas.Bifr } } } else { - if schemas.IsAnthropicModel(bifrostReq.Model) { + if schemas.IsAnthropicModelFamily(ctx, bifrostReq.Model) { bedrockReq.AdditionalModelRequestFields.Set("thinking", map[string]any{ "type": "disabled", }) - } else if schemas.IsNovaModel(bifrostReq.Model) { + } else if schemas.IsNovaModelFamily(ctx, bifrostReq.Model) { bedrockReq.AdditionalModelRequestFields.Set("reasoningConfig", map[string]any{ "type": "disabled", }) @@ -408,7 +445,7 @@ func convertChatParameters(ctx *schemas.BifrostContext, bifrostReq *schemas.Bifr thinkingEnabled := bifrostReq.Params.Reasoning != nil && (bifrostReq.Params.Reasoning.MaxTokens != nil || (bifrostReq.Params.Reasoning.Effort != nil && *bifrostReq.Params.Reasoning.Effort != "none")) - if !schemas.IsLlamaModel(bifrostReq.Model) && !thinkingEnabled { + if !schemas.IsLlamaModelFamily(ctx, bifrostReq.Model) && !thinkingEnabled { bedrockReq.ToolConfig.ToolChoice = &BedrockToolChoice{ Tool: &BedrockToolChoiceTool{ Name: responseFormatTool.ToolSpec.Name, @@ -1677,7 +1714,12 @@ func convertToolConfig(model string, params *schemas.ChatParameters) *BedrockToo // pre-filtered tool set. convertChatParameters uses this to avoid filtering // twice (once here, once in collectBedrockServerTools). The public // convertToolConfig entry point is a thin wrapper preserved for tests. -func convertToolConfigFromFiltered(ctx context.Context, model string, params *schemas.ChatParameters, filtered []schemas.ChatTool) *BedrockToolConfig { +// +// ctx is the BifrostContext (not context.Context) so the family gates inside +// this function can consult the resolved alias and honor explicit +// AliasConfig.ModelFamily overrides. Test paths may pass nil — family +// detection then falls back to substring matching on model. +func convertToolConfigFromFiltered(ctx *schemas.BifrostContext, model string, params *schemas.ChatParameters, filtered []schemas.ChatTool) *BedrockToolConfig { if params == nil { return nil } @@ -1717,7 +1759,7 @@ func convertToolConfigFromFiltered(ctx context.Context, model string, params *sc } bedrockTools = append(bedrockTools, bedrockTool) - if tool.CacheControl != nil && !schemas.IsNovaModel(model) { + if tool.CacheControl != nil && !schemas.IsNovaModelFamily(ctx, model) { bedrockTools = append(bedrockTools, BedrockTool{ CachePoint: &BedrockCachePoint{ Type: BedrockCachePointTypeDefault, @@ -1774,7 +1816,7 @@ func convertToolConfigFromFiltered(ctx context.Context, model string, params *sc // behavior. See per-model support matrix at // https://docs.aws.amazon.com/bedrock/latest/APIReference/API_runtime_ToolChoice.html // (mirrors the synthetic-tool gate in convertChatParameters). - if toolChoice != nil && toolChoice.Tool != nil && schemas.IsLlamaModel(model) { + if toolChoice != nil && toolChoice.Tool != nil && schemas.IsLlamaModelFamily(ctx, model) { toolChoice = nil } if toolChoice != nil { diff --git a/core/schemas/account.go b/core/schemas/account.go index 6a6e1bb519e..a23f57ce74e 100644 --- a/core/schemas/account.go +++ b/core/schemas/account.go @@ -156,6 +156,7 @@ const ( ModelFamilyMistral ModelFamily = "mistral" ModelFamilyCohere ModelFamily = "cohere" ModelFamilyGemini ModelFamily = "gemini" + ModelFamilyLlama ModelFamily = "llama" ModelFamilyNova ModelFamily = "nova" ModelFamilyTitan ModelFamily = "titan" ) @@ -167,7 +168,8 @@ func (mf *ModelFamily) IsValid() bool { } switch *mf { case ModelFamilyAnthropic, ModelFamilyOpenAI, ModelFamilyMistral, - ModelFamilyCohere, ModelFamilyGemini, ModelFamilyNova, ModelFamilyTitan: + ModelFamilyCohere, ModelFamilyGemini, ModelFamilyLlama, + ModelFamilyNova, ModelFamilyTitan: return true } return false @@ -359,8 +361,14 @@ func ResolveFamily(ctx *BifrostContext, fallbackModel string) ModelFamily { return ModelFamilyMistral case IsGeminiModel(s): return ModelFamilyGemini + case IsLlamaModel(s): + return ModelFamilyLlama case IsNovaModel(s): return ModelFamilyNova + case IsTitanModel(s): + return ModelFamilyTitan + case IsCohereModel(s): + return ModelFamilyCohere } } return "" @@ -376,6 +384,40 @@ func IsAnthropicModelFamily(ctx *BifrostContext, model string) bool { return ResolveFamily(ctx, model) == ModelFamilyAnthropic } +// IsMistralModelFamily reports whether the current attempt resolves to the +// Mistral model family. See IsAnthropicModelFamily for usage notes. +func IsMistralModelFamily(ctx *BifrostContext, model string) bool { + return ResolveFamily(ctx, model) == ModelFamilyMistral +} + +// IsLlamaModelFamily reports whether the current attempt resolves to the +// Llama model family. Used by Bedrock to gate tool_choice handling — AWS +// Bedrock Converse rejects toolConfig.toolChoice.tool on Meta Llama variants. +func IsLlamaModelFamily(ctx *BifrostContext, model string) bool { + return ResolveFamily(ctx, model) == ModelFamilyLlama +} + +// IsNovaModelFamily reports whether the current attempt resolves to the +// Amazon Nova model family. Used by Bedrock to gate cache-point insertion +// and tool shaping that differs from Anthropic. +func IsNovaModelFamily(ctx *BifrostContext, model string) bool { + return ResolveFamily(ctx, model) == ModelFamilyNova +} + +// IsCohereModelFamily reports whether the current attempt resolves to the +// Cohere model family. Used by Bedrock to pick the Cohere request/response +// shape for embeddings (vs. the Titan envelope). +func IsCohereModelFamily(ctx *BifrostContext, model string) bool { + return ResolveFamily(ctx, model) == ModelFamilyCohere +} + +// IsTitanModelFamily reports whether the current attempt resolves to the +// Amazon Titan model family. Used by Bedrock to pick the Titan embedding +// request/response envelope. +func IsTitanModelFamily(ctx *BifrostContext, model string) bool { + return ResolveFamily(ctx, model) == ModelFamilyTitan +} + // ResolveConfig returns the AliasConfig for the given user-facing model name, // or nil if no alias matches. Case-insensitive fallback matches Resolve. func (ka KeyAliases) ResolveConfig(model string) *AliasConfig { diff --git a/core/schemas/utils.go b/core/schemas/utils.go index 6ab1b9a3a79..ac55a928fb7 100644 --- a/core/schemas/utils.go +++ b/core/schemas/utils.go @@ -1433,6 +1433,19 @@ func IsImagenModel(model string) bool { return strings.Contains(strings.ToLower(model), "imagen") } +// IsCohereModel checks if the model is a Cohere model. Matches the Bedrock +// identifier prefix ("cohere.embed-*", "cohere.command-*") which is the wire +// shape that flows through alias resolution. +func IsCohereModel(model string) bool { + return strings.Contains(model, "cohere") +} + +// IsTitanModel checks if the model is an Amazon Titan model. Matches the +// Bedrock identifier prefix ("amazon.titan-*"). +func IsTitanModel(model string) bool { + return strings.Contains(model, "titan") +} + // List of grok reasoning models var grokReasoningModels = []string{ "grok-3", diff --git a/transports/bifrost-http/integrations/bedrock.go b/transports/bifrost-http/integrations/bedrock.go index efe2045261f..89b784e4400 100644 --- a/transports/bifrost-http/integrations/bedrock.go +++ b/transports/bifrost-http/integrations/bedrock.go @@ -255,7 +255,7 @@ func createBedrockInvokeRouteConfig(pathPrefix string, handlerStore lib.HandlerS return bedrock.ToBedrockInvokeMessagesResponse(ctx, resp) }, EmbeddingResponseConverter: func(ctx *schemas.BifrostContext, resp *schemas.BifrostEmbeddingResponse) (interface{}, error) { - return bedrock.ToBedrockEmbeddingInvokeResponse(resp) + return bedrock.ToBedrockEmbeddingInvokeResponse(ctx, resp) }, ImageGenerationResponseConverter: func(ctx *schemas.BifrostContext, resp *schemas.BifrostImageGenerationResponse) (interface{}, error) { return bedrock.ToBedrockInvokeImagesResponse(ctx, resp)