-
Notifications
You must be signed in to change notification settings - Fork 1.2k
feat: add ProjectID to Bedrock and Bedrock Mantle key configs for Mantle project scoping
#5131
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.
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
Changes from all commits
File filter
Filter by extension
Conversations
Jump to
Diff view
Diff view
There are no files selected for viewing
| Original file line number | Diff line number | Diff line change |
|---|---|---|
|
|
@@ -5,6 +5,7 @@ import ( | |
| "context" | ||
| "fmt" | ||
| "io" | ||
| "maps" | ||
| "net/http" | ||
| "strings" | ||
|
|
||
|
|
@@ -13,6 +14,40 @@ import ( | |
| schemas "github.com/maximhq/bifrost/core/schemas" | ||
| ) | ||
|
|
||
| const ( | ||
| // MantleOpenAIProjectHeader selects a Bedrock Mantle project on the OpenAI-compatible surface | ||
| // (chat/completions, responses, /models). AWS routes to the account's default project when absent. | ||
| MantleOpenAIProjectHeader = "OpenAI-Project" | ||
| // MantleAnthropicProjectHeader selects a Bedrock Mantle project on the native-Anthropic surface | ||
| // (/anthropic/v1/messages). | ||
| MantleAnthropicProjectHeader = "anthropic-workspace-id" | ||
| ) | ||
|
|
||
| // WithMantleProject returns headers with the given Mantle project header set when projectID is | ||
| // non-empty, letting AWS fall back to the account's default project when it is empty. It never | ||
| // mutates base (which may be the shared networkConfig.ExtraHeaders map). The project header is a | ||
| // plain (non x-amz-*) header, so it does not need to be part of the SigV4 SignedHeaders. | ||
| func WithMantleProject(base map[string]string, headerName, projectID string) map[string]string { | ||
| if projectID == "" { | ||
| return base | ||
| } | ||
| out := maps.Clone(base) | ||
| if out == nil { | ||
| out = make(map[string]string, 1) | ||
| } | ||
| out[headerName] = projectID | ||
| return out | ||
| } | ||
|
|
||
| // resolveMantleProjectID returns the Bedrock project configured for the mantle sub-surface of the | ||
| // Bedrock provider, or "" when none is set (AWS then routes to the account's default project). | ||
| func resolveMantleProjectID(key schemas.Key) string { | ||
|
impoiler marked this conversation as resolved.
|
||
| if key.BedrockKeyConfig != nil && key.BedrockKeyConfig.ProjectID != nil { | ||
| return key.BedrockKeyConfig.ProjectID.GetValue() | ||
| } | ||
| return "" | ||
| } | ||
|
Comment on lines
+44
to
+49
Contributor
There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more.
|
||
|
|
||
| // isMantleModel reports whether a model should be routed via the Bedrock Mantle | ||
| // OpenAI-compatible endpoint. OpenAI-family (gpt-*) and Gemma 4 models are mantle-only | ||
| // (they have no Converse equivalent). Gemma 3 is intentionally excluded: it only supports | ||
|
|
@@ -131,7 +166,7 @@ func (provider *BedrockProvider) mantleChatCompletions( | |
| url, | ||
| request, | ||
| openai.BearerAuthHeader(key), | ||
| provider.networkConfig.ExtraHeaders, | ||
| WithMantleProject(provider.networkConfig.ExtraHeaders, MantleOpenAIProjectHeader, resolveMantleProjectID(key)), | ||
| providerUtils.ShouldSendBackRawRequest(ctx, provider.sendBackRawRequest), | ||
| providerUtils.ShouldSendBackRawResponse(ctx, provider.sendBackRawResponse), | ||
| provider.GetProviderKey(), | ||
|
|
@@ -165,7 +200,7 @@ func (provider *BedrockProvider) mantleChatCompletionsStream( | |
|
|
||
| return openai.HandleOpenAIChatCompletionStreaming( | ||
| ctx, provider.mantleStreamingClient, url, request, | ||
| openai.BearerAuthHeader(key), provider.networkConfig.ExtraHeaders, | ||
| openai.BearerAuthHeader(key), WithMantleProject(provider.networkConfig.ExtraHeaders, MantleOpenAIProjectHeader, resolveMantleProjectID(key)), | ||
| provider.networkConfig.StreamIdleTimeoutInSeconds, | ||
| providerUtils.ShouldSendBackRawRequest(ctx, provider.sendBackRawRequest), | ||
| providerUtils.ShouldSendBackRawResponse(ctx, provider.sendBackRawResponse), | ||
|
|
@@ -206,7 +241,7 @@ func (provider *BedrockProvider) mantleResponses( | |
| url, | ||
| request, | ||
| openai.BearerAuthHeader(key), | ||
| provider.networkConfig.ExtraHeaders, | ||
| WithMantleProject(provider.networkConfig.ExtraHeaders, MantleOpenAIProjectHeader, resolveMantleProjectID(key)), | ||
| providerUtils.ShouldSendBackRawRequest(ctx, provider.sendBackRawRequest), | ||
| providerUtils.ShouldSendBackRawResponse(ctx, provider.sendBackRawResponse), | ||
| provider.GetProviderKey(), | ||
|
|
@@ -240,7 +275,7 @@ func (provider *BedrockProvider) mantleResponsesStream( | |
|
|
||
| return openai.HandleOpenAIResponsesStreaming( | ||
| ctx, provider.mantleStreamingClient, url, request, | ||
| openai.BearerAuthHeader(key), provider.networkConfig.ExtraHeaders, | ||
| openai.BearerAuthHeader(key), WithMantleProject(provider.networkConfig.ExtraHeaders, MantleOpenAIProjectHeader, resolveMantleProjectID(key)), | ||
| provider.networkConfig.StreamIdleTimeoutInSeconds, | ||
| providerUtils.ShouldSendBackRawRequest(ctx, provider.sendBackRawRequest), | ||
| providerUtils.ShouldSendBackRawResponse(ctx, provider.sendBackRawResponse), | ||
|
|
||
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -0,0 +1,75 @@ | ||
| package bedrock | ||
|
|
||
| import ( | ||
| "testing" | ||
|
|
||
| schemas "github.com/maximhq/bifrost/core/schemas" | ||
| ) | ||
|
|
||
| // TestWithMantleProject verifies the project header is added only when a project ID is present, | ||
| // the target header name is honoured, and the base map is never mutated. | ||
| func TestWithMantleProject(t *testing.T) { | ||
| t.Run("empty project returns base unchanged", func(t *testing.T) { | ||
| base := map[string]string{"X-Custom": "v"} | ||
| got := WithMantleProject(base, MantleOpenAIProjectHeader, "") | ||
| if _, ok := got[MantleOpenAIProjectHeader]; ok { | ||
| t.Fatalf("expected no project header when project ID is empty, got %v", got) | ||
| } | ||
| // Empty project must return the base map as-is (default-project behaviour). | ||
| if len(got) != 1 || got["X-Custom"] != "v" { | ||
| t.Fatalf("expected base returned unchanged, got %v", got) | ||
| } | ||
| }) | ||
|
|
||
| t.Run("OpenAI project header set", func(t *testing.T) { | ||
| base := map[string]string{"X-Custom": "v"} | ||
| got := WithMantleProject(base, MantleOpenAIProjectHeader, "proj_abc") | ||
| if got[MantleOpenAIProjectHeader] != "proj_abc" { | ||
| t.Fatalf("expected %s=proj_abc, got %v", MantleOpenAIProjectHeader, got) | ||
| } | ||
| if got["X-Custom"] != "v" { | ||
| t.Fatalf("existing headers must be preserved, got %v", got) | ||
| } | ||
| // base must not be mutated. | ||
| if _, ok := base[MantleOpenAIProjectHeader]; ok { | ||
| t.Fatalf("base map was mutated: %v", base) | ||
| } | ||
| }) | ||
|
|
||
| t.Run("Anthropic workspace header set", func(t *testing.T) { | ||
| got := WithMantleProject(nil, MantleAnthropicProjectHeader, "proj_xyz") | ||
| if got[MantleAnthropicProjectHeader] != "proj_xyz" { | ||
| t.Fatalf("expected %s=proj_xyz, got %v", MantleAnthropicProjectHeader, got) | ||
| } | ||
| }) | ||
|
|
||
| t.Run("nil base with empty project stays nil", func(t *testing.T) { | ||
| if got := WithMantleProject(nil, MantleOpenAIProjectHeader, ""); got != nil { | ||
| t.Fatalf("expected nil when base is nil and project is empty, got %v", got) | ||
| } | ||
| }) | ||
| } | ||
|
|
||
| // TestResolveMantleProjectID verifies precedence of the BedrockKeyConfig.ProjectID field. | ||
| func TestResolveMantleProjectID(t *testing.T) { | ||
| tests := []struct { | ||
| name string | ||
| key schemas.Key | ||
| want string | ||
| }{ | ||
| {name: "no bedrock config", key: schemas.Key{}, want: ""}, | ||
| {name: "config without project", key: schemas.Key{BedrockKeyConfig: &schemas.BedrockKeyConfig{}}, want: ""}, | ||
| { | ||
| name: "config with project", | ||
| key: schemas.Key{BedrockKeyConfig: &schemas.BedrockKeyConfig{ProjectID: schemas.NewSecretVar("proj_abc")}}, | ||
| want: "proj_abc", | ||
| }, | ||
| } | ||
| for _, tt := range tests { | ||
| t.Run(tt.name, func(t *testing.T) { | ||
| if got := resolveMantleProjectID(tt.key); got != tt.want { | ||
| t.Fatalf("resolveMantleProjectID = %q, want %q", got, tt.want) | ||
| } | ||
| }) | ||
| } | ||
| } |
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -0,0 +1,32 @@ | ||
| package bedrockmantle | ||
|
|
||
| import ( | ||
| "testing" | ||
|
|
||
| schemas "github.com/maximhq/bifrost/core/schemas" | ||
| ) | ||
|
|
||
| // TestResolveProjectID verifies the BedrockMantleKeyConfig.ProjectID field is honoured and that an | ||
| // absent project resolves to "" (AWS default project). | ||
| func TestResolveProjectID(t *testing.T) { | ||
| tests := []struct { | ||
| name string | ||
| key schemas.Key | ||
| want string | ||
| }{ | ||
| {name: "no mantle config", key: schemas.Key{}, want: ""}, | ||
| {name: "config without project", key: schemas.Key{BedrockMantleKeyConfig: &schemas.BedrockMantleKeyConfig{}}, want: ""}, | ||
| { | ||
| name: "config with project", | ||
| key: schemas.Key{BedrockMantleKeyConfig: &schemas.BedrockMantleKeyConfig{ProjectID: schemas.NewSecretVar("proj_xyz")}}, | ||
| want: "proj_xyz", | ||
| }, | ||
| } | ||
| for _, tt := range tests { | ||
| t.Run(tt.name, func(t *testing.T) { | ||
| if got := resolveProjectID(tt.key); got != tt.want { | ||
| t.Fatalf("resolveProjectID = %q, want %q", got, tt.want) | ||
| } | ||
| }) | ||
| } | ||
| } |
Uh oh!
There was an error while loading. Please reload this page.