diff --git a/aibridge/config/config.go b/aibridge/config/config.go index acc1192319fb4..ee0c5fec8a672 100644 --- a/aibridge/config/config.go +++ b/aibridge/config/config.go @@ -70,10 +70,20 @@ type AWSBedrock struct { Protocol BedrockProtocol } +// ResolvedProtocol returns the configured protocol, mapping the empty value to +// the legacy InvokeModel protocol so existing providers keep the legacy +// behavior. +func (c AWSBedrock) ResolvedProtocol() BedrockProtocol { + if c.Protocol == "" { + return BedrockProtocolInvokeModel + } + return c.Protocol +} + // Validate verifies protocol-specific Bedrock configuration. func (c AWSBedrock) Validate() error { - switch c.Protocol { - case "", BedrockProtocolInvokeModel: + switch c.ResolvedProtocol() { + case BedrockProtocolInvokeModel: if c.Region == "" && c.BaseURL == "" { return xerrors.New("region or base url required") } diff --git a/aibridge/intercept/messages/base.go b/aibridge/intercept/messages/base.go index a0025d6881b99..1a9ea0dfb858a 100644 --- a/aibridge/intercept/messages/base.go +++ b/aibridge/intercept/messages/base.go @@ -133,14 +133,13 @@ func (i *interceptionBase) CorrelatingToolCallID() *string { // isBedrockMantle reports whether the interception targets the Bedrock mantle // protocol. func (i *interceptionBase) isBedrockMantle() bool { - return i.bedrock != nil && i.bedrock.Cfg.Protocol == aibconfig.BedrockProtocolMantle + return i.bedrock != nil && i.bedrock.Cfg.ResolvedProtocol() == aibconfig.BedrockProtocolMantle } // isBedrockInvokeModel reports whether the interception targets the Bedrock // InvokeModel protocol. func (i *interceptionBase) isBedrockInvokeModel() bool { - return i.bedrock != nil && - (i.bedrock.Cfg.Protocol == "" || i.bedrock.Cfg.Protocol == aibconfig.BedrockProtocolInvokeModel) + return i.bedrock != nil && i.bedrock.Cfg.ResolvedProtocol() == aibconfig.BedrockProtocolInvokeModel } func (i *interceptionBase) Model() string { @@ -164,7 +163,7 @@ func (i *interceptionBase) Model() string { } func (i *interceptionBase) baseTraceAttributes(r *http.Request, streaming bool) []attribute.KeyValue { - return []attribute.KeyValue{ + attrs := []attribute.KeyValue{ attribute.String(tracing.RequestPath, r.URL.Path), attribute.String(tracing.InterceptionID, i.id.String()), attribute.String(tracing.InitiatorID, aibcontext.ActorIDFromContext(r.Context())), @@ -173,6 +172,10 @@ func (i *interceptionBase) baseTraceAttributes(r *http.Request, streaming bool) attribute.Bool(tracing.Streaming, streaming), attribute.Bool(tracing.IsBedrock, i.bedrock != nil), } + if i.bedrock != nil { + attrs = append(attrs, attribute.String(tracing.BedrockProtocol, string(i.bedrock.Cfg.ResolvedProtocol()))) + } + return attrs } func (i *interceptionBase) injectTools() { diff --git a/aibridge/internal/integrationtest/trace_internal_test.go b/aibridge/internal/integrationtest/trace_internal_test.go index 23fcbdeae8b75..43719c3f177ff 100644 --- a/aibridge/internal/integrationtest/trace_internal_test.go +++ b/aibridge/internal/integrationtest/trace_internal_test.go @@ -187,6 +187,9 @@ func TestTraceAnthropic(t *testing.T) { attribute.Bool(tracing.Streaming, tc.streaming), attribute.Bool(tracing.IsBedrock, tc.bedrock), } + if tc.bedrock { + attrs = append(attrs, attribute.String(tracing.BedrockProtocol, string(config.BedrockProtocolInvokeModel))) + } require.Len(t, sr.Ended(), totalCount) verifyTraces(t, sr, tc.expect, attrs) @@ -311,6 +314,9 @@ func TestTraceAnthropicErr(t *testing.T) { attribute.Bool(tracing.Streaming, tc.streaming), attribute.Bool(tracing.IsBedrock, tc.bedrock), } + if tc.bedrock { + attrs = append(attrs, attribute.String(tracing.BedrockProtocol, string(config.BedrockProtocolInvokeModel))) + } verifyTraces(t, sr, tc.expect, attrs) }) @@ -422,6 +428,9 @@ func TestInjectedToolsTrace(t *testing.T) { } if tc.expectProvider == config.ProviderAnthropic { attrs = append(attrs, attribute.Bool(tracing.IsBedrock, tc.bedrock)) + if tc.bedrock { + attrs = append(attrs, attribute.String(tracing.BedrockProtocol, string(config.BedrockProtocolInvokeModel))) + } } verifyTraces(t, sr, []expectTrace{{"Intercept.ProcessRequest.ToolCall", 1, codes.Unset}}, attrs) diff --git a/aibridge/provider/anthropic.go b/aibridge/provider/anthropic.go index 9424d4c9601a5..460b38102cd8c 100644 --- a/aibridge/provider/anthropic.go +++ b/aibridge/provider/anthropic.go @@ -79,6 +79,9 @@ func NewAnthropic(ctx context.Context, cfg config.Anthropic, bedrockCfg *config. if runtimeCfg.Region == "" { runtimeCfg.Region = resolvedRegion } + if err := runtimeCfg.Validate(); err != nil { + return nil, xerrors.Errorf("bedrock config: %w", err) + } bedrock = &messages.BedrockRuntime{Cfg: runtimeCfg, Creds: creds} } diff --git a/aibridge/provider/anthropic_internal_test.go b/aibridge/provider/anthropic_internal_test.go index 9df6f0371843f..cdc8afe9148aa 100644 --- a/aibridge/provider/anthropic_internal_test.go +++ b/aibridge/provider/anthropic_internal_test.go @@ -114,6 +114,42 @@ func TestNewAnthropic_KeyResolution(t *testing.T) { } } +// NOTE: no t.Parallel() because the subtests use t.Setenv. +func TestNewAnthropic_BedrockRegionResolution(t *testing.T) { + t.Run("mantle_region_from_env", func(t *testing.T) { + t.Setenv("AWS_REGION", "us-west-2") + + p, err := NewAnthropic(context.Background(), config.Anthropic{}, &config.AWSBedrock{ + BaseURL: "https://bedrock-mantle.us-west-2.api.aws/anthropic", + Protocol: config.BedrockProtocolMantle, + AccessKey: "test-key", + AccessKeySecret: "test-secret", + }) + require.NoError(t, err) + require.NotNil(t, p.bedrock) + require.Equal(t, "us-west-2", p.bedrock.Cfg.Region) + }) + + t.Run("mantle_no_region_anywhere", func(t *testing.T) { + // Clear every source the AWS SDK consults for a region so none + // resolves, then confirm construction rejects the mantle provider. + t.Setenv("AWS_REGION", "") + t.Setenv("AWS_DEFAULT_REGION", "") + t.Setenv("AWS_PROFILE", "") + t.Setenv("AWS_CONFIG_FILE", "/dev/null") + t.Setenv("AWS_SHARED_CREDENTIALS_FILE", "/dev/null") + t.Setenv("AWS_EC2_METADATA_DISABLED", "true") + + _, err := NewAnthropic(context.Background(), config.Anthropic{}, &config.AWSBedrock{ + BaseURL: "https://proxy.internal", + Protocol: config.BedrockProtocolMantle, + AccessKey: "test-key", + AccessKeySecret: "test-secret", + }) + require.ErrorContains(t, err, "region required") + }) +} + func TestAnthropic_CreateInterceptor(t *testing.T) { t.Parallel() diff --git a/aibridge/tracing/tracing.go b/aibridge/tracing/tracing.go index 7adaf3f65e355..f01e3f1f33026 100644 --- a/aibridge/tracing/tracing.go +++ b/aibridge/tracing/tracing.go @@ -17,12 +17,13 @@ const ( // trace attribute key constants RequestPath = "request_path" - InterceptionID = "interception_id" - InitiatorID = "user_id" - Provider = "provider" - Model = "model" - Streaming = "streaming" - IsBedrock = "aws_bedrock" + InterceptionID = "interception_id" + InitiatorID = "user_id" + Provider = "provider" + Model = "model" + Streaming = "streaming" + IsBedrock = "aws_bedrock" + BedrockProtocol = "aws_bedrock_protocol" PassthroughURL = "passthrough_url" PassthroughUpstreamURL = "passthrough_upstream_url" diff --git a/cli/aibridged_internal_test.go b/cli/aibridged_internal_test.go index d8d02c131461b..59b2153c6f591 100644 --- a/cli/aibridged_internal_test.go +++ b/cli/aibridged_internal_test.go @@ -175,6 +175,8 @@ func TestBuildProviders(t *testing.T) { cfg.LegacyBedrock.Region = serpent.String("us-west-2") cfg.LegacyBedrock.AccessKey = serpent.String("AKID") cfg.LegacyBedrock.AccessKeySecret = serpent.String("secret") + cfg.LegacyBedrock.Model = serpent.String("anthropic.claude-3-5-sonnet-20241022-v2:0") + cfg.LegacyBedrock.SmallFastModel = serpent.String("anthropic.claude-3-5-haiku-20241022-v1:0") providers, err := buildFromEnv(t, cfg) require.NoError(t, err) @@ -191,6 +193,8 @@ func TestBuildProviders(t *testing.T) { cfg.LegacyBedrock.Region = serpent.String("us-west-2") cfg.LegacyBedrock.AccessKey = serpent.String("AKID") cfg.LegacyBedrock.AccessKeySecret = serpent.String("secret") + cfg.LegacyBedrock.Model = serpent.String("anthropic.claude-3-5-sonnet-20241022-v2:0") + cfg.LegacyBedrock.SmallFastModel = serpent.String("anthropic.claude-3-5-haiku-20241022-v1:0") providers, err := buildFromEnv(t, cfg) require.NoError(t, err) diff --git a/cli/server_aibridge_internal_test.go b/cli/server_aibridge_internal_test.go index 7a3732cef7784..781b642f938cd 100644 --- a/cli/server_aibridge_internal_test.go +++ b/cli/server_aibridge_internal_test.go @@ -727,6 +727,8 @@ func TestBuildProviderFromProtoSetsAPIDumpDir(t *testing.T) { Region: "us-east-1", AccessKey: "AKID", AccessKeySecret: "secret", + Model: "anthropic.claude-3-5-sonnet-20241022-v2:0", + SmallFastModel: "anthropic.claude-3-5-haiku-20241022-v1:0", }, }, expectedType: aibridge.ProviderAnthropic,