Skip to content
Open
Show file tree
Hide file tree
Changes from all commits
Commits
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
58 changes: 55 additions & 3 deletions aibridge/config/config.go
Original file line number Diff line number Diff line change
Expand Up @@ -3,6 +3,8 @@ package config
import (
"time"

"golang.org/x/xerrors"

"github.com/coder/coder/v2/aibridge/keypool"
)

Expand All @@ -25,13 +27,34 @@ type Anthropic struct {
SendActorHeaders bool
}

// BedrockProtocol selects which AWS Bedrock wire protocol a provider targets.
type BedrockProtocol string

const (
// BedrockProtocolInvokeModel is the legacy InvokeModel protocol
// (bedrock-runtime.{region}.amazonaws.com), which translates the native
// Messages request into Bedrock's InvokeModel format. It is the default
// for the zero value.
BedrockProtocolInvokeModel BedrockProtocol = "invoke-model"
// BedrockProtocolMantle is the mantle protocol
// (bedrock-mantle.{region}.api.aws/anthropic/v1/messages). It is a
// passthrough: the gateway forwards the native Messages request body
// unchanged and only applies AWS SigV4 signing (service bedrock-mantle).
BedrockProtocolMantle BedrockProtocol = "mantle"
)

type AWSBedrock struct {
Region string
AccessKey, AccessKeySecret string
Model, SmallFastModel string
// If set, requests will be sent to this URL instead of the default AWS Bedrock endpoint
// (https://bedrock-runtime.{region}.amazonaws.com).
// This is useful for routing requests through a proxy or for testing.
// BaseURL configures the upstream Bedrock endpoint.
//
// For InvokeModel, it is optional. When empty, requests use the default
// https://bedrock-runtime.{region}.amazonaws.com endpoint. Set it to route
// InvokeModel requests through a proxy or test server.
//
// For mantle, it is required and must be the Messages API prefix without
// /v1/messages, e.g. https://bedrock-mantle.{region}.api.aws/anthropic.
BaseURL string
// RoleARN, when set, is assumed via STS before calling Bedrock. The base
// identity (static keys or the AWS SDK default credential chain, e.g.
Expand All @@ -42,6 +65,35 @@ type AWSBedrock struct {
// It is meaningful only alongside RoleARN and must match the
// sts:ExternalId condition on the target role's trust policy.
ExternalID string
// Protocol selects the Bedrock wire protocol. The zero value behaves as
// BedrockProtocolInvokeModel.
Protocol BedrockProtocol
}

// Validate verifies protocol-specific Bedrock configuration.
func (c AWSBedrock) Validate() error {
switch c.Protocol {
case "", BedrockProtocolInvokeModel:
if c.Region == "" && c.BaseURL == "" {
return xerrors.New("region or base url required")
}
if c.Model == "" {
return xerrors.New("model required")
}
if c.SmallFastModel == "" {
return xerrors.New("small fast model required")
}
case BedrockProtocolMantle:
if c.Region == "" {
return xerrors.New("region required")
}
if c.BaseURL == "" {
return xerrors.New("base_url required")
}
default:
return xerrors.Errorf("unknown bedrock protocol: %q", c.Protocol)
}
return nil
}

// OpenAI carries configuration for an OpenAI provider.
Expand Down
113 changes: 113 additions & 0 deletions aibridge/config/config_test.go
Original file line number Diff line number Diff line change
@@ -0,0 +1,113 @@
package config_test

import (
"testing"

"github.com/stretchr/testify/require"

"github.com/coder/coder/v2/aibridge/config"
)

func TestAWSBedrockValidate(t *testing.T) {
t.Parallel()

tests := []struct {
name string
cfg config.AWSBedrock
errorMsg string
}{
{
name: "invoke model valid",
cfg: config.AWSBedrock{
Region: "us-east-1",
Model: "anthropic.claude-sonnet",
SmallFastModel: "anthropic.claude-haiku",
},
},
{
name: "invoke model valid with base url instead of region",
cfg: config.AWSBedrock{
BaseURL: "https://bedrock-runtime.example.com",
Model: "anthropic.claude-sonnet",
SmallFastModel: "anthropic.claude-haiku",
},
},
{
name: "invoke model missing region and base url",
cfg: config.AWSBedrock{
Model: "anthropic.claude-sonnet",
SmallFastModel: "anthropic.claude-haiku",
},
errorMsg: "region or base url required",
},
{
name: "invoke model missing model",
cfg: config.AWSBedrock{
Region: "us-east-1",
SmallFastModel: "anthropic.claude-haiku",
},
errorMsg: "model required",
},
{
name: "invoke model missing small fast model",
cfg: config.AWSBedrock{
Region: "us-east-1",
Model: "anthropic.claude-sonnet",
},
errorMsg: "small fast model required",
},
{
name: "unknown protocol rejected",
cfg: config.AWSBedrock{
Protocol: config.BedrockProtocol("unknown"),
},
errorMsg: "unknown bedrock protocol",
},
{
name: "mantle valid official api prefix",
cfg: config.AWSBedrock{
Region: "us-east-1",
BaseURL: "https://bedrock-mantle.us-east-1.api.aws/anthropic",
Protocol: config.BedrockProtocolMantle,
},
},
{
name: "mantle valid proxy api prefix",
cfg: config.AWSBedrock{
Region: "us-east-1",
BaseURL: "https://proxy.internal/proxy",
Protocol: config.BedrockProtocolMantle,
},
},
{
name: "mantle missing region",
cfg: config.AWSBedrock{
BaseURL: "https://bedrock-mantle.us-east-1.api.aws",
Protocol: config.BedrockProtocolMantle,
},
errorMsg: "region required",
},
{
name: "mantle missing base url",
cfg: config.AWSBedrock{
Region: "us-east-1",
Protocol: config.BedrockProtocolMantle,
},
errorMsg: "base_url required",
},
}

for _, tt := range tests {
t.Run(tt.name, func(t *testing.T) {
t.Parallel()

err := tt.cfg.Validate()
if tt.errorMsg != "" {
require.Error(t, err)
require.Contains(t, err.Error(), tt.errorMsg)
return
}
require.NoError(t, err)
})
}
}
Loading
Loading