Add Bedrock auth support for OpenAI providers

Author: akichannnnn

docs/providers.md

@@ -96,6 +96,8 @@ model := provider.ChatModel("gpt-4o-mini")
 | Option | Default | Description |
 |--------|---------|-------------|
 | `WithAPIKey(key)` | `""` | API key sent as `Authorization: Bearer <key>` |
+| `WithBedrockRegion(region)` | disabled | Use AWS SigV4 with the default AWS credential chain for Bedrock's OpenAI-compatible endpoint |
+| `WithBedrockCredentials(region, accessKeyID, secretAccessKey, sessionToken)` | disabled | Use AWS SigV4 with static AWS credentials for Bedrock |
 | `WithBaseURL(url)` | `https://api.openai.com/v1` | Base URL for API requests |
 | `WithHTTPClient(client)` | `&http.Client{}` | Custom HTTP client (for proxies, timeouts, etc.) |
 
@@ -130,6 +132,12 @@ provider := completions.New(
     completions.WithBaseURL("https://your-resource.openai.azure.com/openai/deployments/gpt-4o"),
 )
 
+// Amazon Bedrock OpenAI-compatible endpoint (AWS credentials)
+provider := completions.New(
+    completions.WithBedrockRegion("us-east-1"),
+    completions.WithBaseURL("https://bedrock-mantle.us-east-1.api.aws/v1"),
+)
+
 // Local (Ollama, vLLM, etc.)
 provider := completions.New(
     completions.WithBaseURL("http://localhost:11434/v1"),
@@ -198,6 +206,8 @@ model := provider.ChatModel("gpt-4o-mini")
 | Option | Default | Description |
 |--------|---------|-------------|
 | `WithAPIKey(key)` | `""` | API key sent as `Authorization: Bearer <key>` |
+| `WithBedrockRegion(region)` | disabled | Use AWS SigV4 with the default AWS credential chain for Bedrock's OpenAI-compatible endpoint |
+| `WithBedrockCredentials(region, accessKeyID, secretAccessKey, sessionToken)` | disabled | Use AWS SigV4 with static AWS credentials for Bedrock |
 | `WithBaseURL(url)` | `https://api.openai.com/v1` | Base URL for API requests |
 | `WithHTTPClient(client)` | `&http.Client{}` | Custom HTTP client |
 
@@ -218,6 +228,8 @@ provider := responses.New(
     responses.WithAPIKey("sk-or-v1-..."),
     responses.WithBaseURL("https://openrouter.ai/api/v1"),
 )
+
+Amazon Bedrock's OpenAI-compatible Responses endpoint also works with the same provider when you configure a Bedrock Mantle base URL and either `WithBedrockRegion(...)` or `WithBedrockCredentials(...)`.
 model := provider.ChatModel("openai/o4-mini")
 ```
 
@@ -777,6 +789,8 @@ vec, err := sdk.Embed(ctx, "Hello world",
 | Option | Default | Description |
 |--------|---------|-------------|
 | `WithAPIKey(key)` | `""` | API key sent as `Authorization: Bearer <key>` |
+| `WithBedrockRegion(region)` | disabled | Use AWS SigV4 with the default AWS credential chain for Bedrock's OpenAI-compatible endpoint |
+| `WithBedrockCredentials(region, accessKeyID, secretAccessKey, sessionToken)` | disabled | Use AWS SigV4 with static AWS credentials for Bedrock |
 | `WithBaseURL(url)` | `https://api.openai.com/v1` | Base URL for API requests |
 | `WithHTTPClient(client)` | `&http.Client{}` | Custom HTTP client |
 

go.mod

@@ -3,13 +3,27 @@ module github.com/memohai/twilight-ai
 go 1.25.7
 
 require (
+	github.com/aws/aws-sdk-go-v2 v1.41.5
+	github.com/aws/aws-sdk-go-v2/config v1.32.14
+	github.com/aws/aws-sdk-go-v2/credentials v1.19.14
 	github.com/google/jsonschema-go v0.4.2
 	github.com/google/uuid v1.6.0
 	github.com/gorilla/websocket v1.5.3
 	github.com/modelcontextprotocol/go-sdk v1.5.0
 )
 
 require (
+	github.com/aws/aws-sdk-go-v2/feature/ec2/imds v1.18.21 // indirect
+	github.com/aws/aws-sdk-go-v2/internal/configsources v1.4.21 // indirect
+	github.com/aws/aws-sdk-go-v2/internal/endpoints/v2 v2.7.21 // indirect
+	github.com/aws/aws-sdk-go-v2/internal/ini v1.8.6 // indirect
+	github.com/aws/aws-sdk-go-v2/service/internal/accept-encoding v1.13.7 // indirect
+	github.com/aws/aws-sdk-go-v2/service/internal/presigned-url v1.13.21 // indirect
+	github.com/aws/aws-sdk-go-v2/service/signin v1.0.9 // indirect
+	github.com/aws/aws-sdk-go-v2/service/sso v1.30.15 // indirect
+	github.com/aws/aws-sdk-go-v2/service/ssooidc v1.35.19 // indirect
+	github.com/aws/aws-sdk-go-v2/service/sts v1.41.10 // indirect
+	github.com/aws/smithy-go v1.24.2 // indirect
 	github.com/segmentio/asm v1.2.1 // indirect
 	github.com/segmentio/encoding v0.5.4 // indirect
 	github.com/yosida95/uritemplate/v3 v3.0.2 // indirect

go.sum

@@ -1,3 +1,31 @@
+github.com/aws/aws-sdk-go-v2 v1.41.5 h1:dj5kopbwUsVUVFgO4Fi5BIT3t4WyqIDjGKCangnV/yY=
+github.com/aws/aws-sdk-go-v2 v1.41.5/go.mod h1:mwsPRE8ceUUpiTgF7QmQIJ7lgsKUPQOUl3o72QBrE1o=
+github.com/aws/aws-sdk-go-v2/config v1.32.14 h1:opVIRo/ZbbI8OIqSOKmpFaY7IwfFUOCCXBsUpJOwDdI=
+github.com/aws/aws-sdk-go-v2/config v1.32.14/go.mod h1:U4/V0uKxh0Tl5sxmCBZ3AecYny4UNlVmObYjKuuaiOo=
+github.com/aws/aws-sdk-go-v2/credentials v1.19.14 h1:n+UcGWAIZHkXzYt87uMFBv/l8THYELoX6gVcUvgl6fI=
+github.com/aws/aws-sdk-go-v2/credentials v1.19.14/go.mod h1:cJKuyWB59Mqi0jM3nFYQRmnHVQIcgoxjEMAbLkpr62w=
+github.com/aws/aws-sdk-go-v2/feature/ec2/imds v1.18.21 h1:NUS3K4BTDArQqNu2ih7yeDLaS3bmHD0YndtA6UP884g=
+github.com/aws/aws-sdk-go-v2/feature/ec2/imds v1.18.21/go.mod h1:YWNWJQNjKigKY1RHVJCuupeWDrrHjRqHm0N9rdrWzYI=
+github.com/aws/aws-sdk-go-v2/internal/configsources v1.4.21 h1:Rgg6wvjjtX8bNHcvi9OnXWwcE0a2vGpbwmtICOsvcf4=
+github.com/aws/aws-sdk-go-v2/internal/configsources v1.4.21/go.mod h1:A/kJFst/nm//cyqonihbdpQZwiUhhzpqTsdbhDdRF9c=
+github.com/aws/aws-sdk-go-v2/internal/endpoints/v2 v2.7.21 h1:PEgGVtPoB6NTpPrBgqSE5hE/o47Ij9qk/SEZFbUOe9A=
+github.com/aws/aws-sdk-go-v2/internal/endpoints/v2 v2.7.21/go.mod h1:p+hz+PRAYlY3zcpJhPwXlLC4C+kqn70WIHwnzAfs6ps=
+github.com/aws/aws-sdk-go-v2/internal/ini v1.8.6 h1:qYQ4pzQ2Oz6WpQ8T3HvGHnZydA72MnLuFK9tJwmrbHw=
+github.com/aws/aws-sdk-go-v2/internal/ini v1.8.6/go.mod h1:O3h0IK87yXci+kg6flUKzJnWeziQUKciKrLjcatSNcY=
+github.com/aws/aws-sdk-go-v2/service/internal/accept-encoding v1.13.7 h1:5EniKhLZe4xzL7a+fU3C2tfUN4nWIqlLesfrjkuPFTY=
+github.com/aws/aws-sdk-go-v2/service/internal/accept-encoding v1.13.7/go.mod h1:x0nZssQ3qZSnIcePWLvcoFisRXJzcTVvYpAAdYX8+GI=
+github.com/aws/aws-sdk-go-v2/service/internal/presigned-url v1.13.21 h1:c31//R3xgIJMSC8S6hEVq+38DcvUlgFY0FM6mSI5oto=
+github.com/aws/aws-sdk-go-v2/service/internal/presigned-url v1.13.21/go.mod h1:r6+pf23ouCB718FUxaqzZdbpYFyDtehyZcmP5KL9FkA=
+github.com/aws/aws-sdk-go-v2/service/signin v1.0.9 h1:QKZH0S178gCmFEgst8hN0mCX1KxLgHBKKY/CLqwP8lg=
+github.com/aws/aws-sdk-go-v2/service/signin v1.0.9/go.mod h1:7yuQJoT+OoH8aqIxw9vwF+8KpvLZ8AWmvmUWHsGQZvI=
+github.com/aws/aws-sdk-go-v2/service/sso v1.30.15 h1:lFd1+ZSEYJZYvv9d6kXzhkZu07si3f+GQ1AaYwa2LUM=
+github.com/aws/aws-sdk-go-v2/service/sso v1.30.15/go.mod h1:WSvS1NLr7JaPunCXqpJnWk1Bjo7IxzZXrZi1QQCkuqM=
+github.com/aws/aws-sdk-go-v2/service/ssooidc v1.35.19 h1:dzztQ1YmfPrxdrOiuZRMF6fuOwWlWpD2StNLTceKpys=
+github.com/aws/aws-sdk-go-v2/service/ssooidc v1.35.19/go.mod h1:YO8TrYtFdl5w/4vmjL8zaBSsiNp3w0L1FfKVKenZT7w=
+github.com/aws/aws-sdk-go-v2/service/sts v1.41.10 h1:p8ogvvLugcR/zLBXTXrTkj0RYBUdErbMnAFFp12Lm/U=
+github.com/aws/aws-sdk-go-v2/service/sts v1.41.10/go.mod h1:60dv0eZJfeVXfbT1tFJinbHrDfSJ2GZl4Q//OSSNAVw=
+github.com/aws/smithy-go v1.24.2 h1:FzA3bu/nt/vDvmnkg+R8Xl46gmzEDam6mZ1hzmwXFng=
+github.com/aws/smithy-go v1.24.2/go.mod h1:YE2RhdIuDbA5E5bTdciG9KrW3+TiEONeUWCqxX9i1Fc=
 github.com/golang-jwt/jwt/v5 v5.3.1 h1:kYf81DTWFe7t+1VvL7eS+jKFVWaUnK9cB1qbwn63YCY=
 github.com/golang-jwt/jwt/v5 v5.3.1/go.mod h1:fxCRLWMO43lRc8nhHWY6LGqRcf+1gQWArsqaEUEa5bE=
 github.com/google/go-cmp v0.7.0 h1:wk8382ETsv4JYUZwIsn6YpYiWiBsYLSJiTsyBybVuN8=

internal/utils/aws.go

@@ -0,0 +1,118 @@
+package utils
+
+import (
+	"bytes"
+	"context"
+	"crypto/sha256"
+	"encoding/hex"
+	"fmt"
+	"io"
+	"net/http"
+	"sync"
+	"time"
+
+	"github.com/aws/aws-sdk-go-v2/aws"
+	v4 "github.com/aws/aws-sdk-go-v2/aws/signer/v4"
+	"github.com/aws/aws-sdk-go-v2/config"
+	"github.com/aws/aws-sdk-go-v2/credentials"
+)
+
+const (
+	emptySHA256 = "e3b0c44298fc1c149afbf4c8996fb92427ae41e4649b934ca495991b7852b855"
+
+	bedrockService = "bedrock"
+)
+
+type lazyCredentialsProvider struct {
+	region string
+
+	once sync.Once
+	cp   aws.CredentialsProvider
+	err  error
+}
+
+func NewBedrockDefaultCredentialsPreparer(region string) func(*http.Request) error {
+	return NewSigV4Preparer(bedrockService, region, &lazyCredentialsProvider{region: region})
+}
+
+func NewBedrockStaticCredentialsPreparer(region, accessKeyID, secretAccessKey, sessionToken string) func(*http.Request) error {
+	cp := aws.NewCredentialsCache(credentials.NewStaticCredentialsProvider(accessKeyID, secretAccessKey, sessionToken))
+	return NewSigV4Preparer(bedrockService, region, cp)
+}
+
+func NewSigV4Preparer(service, region string, cp aws.CredentialsProvider) func(*http.Request) error {
+	signer := v4.NewSigner()
+
+	return func(req *http.Request) error {
+		payloadHash, err := requestPayloadHash(req)
+		if err != nil {
+			return err
+		}
+
+		creds, err := cp.Retrieve(req.Context())
+		if err != nil {
+			return fmt.Errorf("retrieve AWS credentials: %w", err)
+		}
+
+		if err := signer.SignHTTP(req.Context(), creds, req, payloadHash, service, region, time.Now().UTC()); err != nil {
+			return fmt.Errorf("sign AWS request: %w", err)
+		}
+
+		return nil
+	}
+}
+
+func (p *lazyCredentialsProvider) Retrieve(ctx context.Context) (aws.Credentials, error) {
+	p.once.Do(func() {
+		cfg, err := config.LoadDefaultConfig(ctx, config.WithRegion(p.region))
+		if err != nil {
+			p.err = err
+			return
+		}
+		p.cp = cfg.Credentials
+	})
+	if p.err != nil {
+		return aws.Credentials{}, p.err
+	}
+	return p.cp.Retrieve(ctx)
+}
+
+func requestPayloadHash(req *http.Request) (string, error) {
+	if req.Body == nil {
+		return emptySHA256, nil
+	}
+
+	if req.GetBody != nil {
+		body, err := req.GetBody()
+		if err != nil {
+			return "", fmt.Errorf("clone request body: %w", err)
+		}
+		defer body.Close()
+		return hashReader(body)
+	}
+
+	data, err := io.ReadAll(req.Body)
+	if err != nil {
+		return "", fmt.Errorf("read request body: %w", err)
+	}
+
+	req.Body = io.NopCloser(bytes.NewReader(data))
+	req.GetBody = func() (io.ReadCloser, error) {
+		return io.NopCloser(bytes.NewReader(data)), nil
+	}
+
+	return hashBytes(data), nil
+}
+
+func hashReader(r io.Reader) (string, error) {
+	h := sha256.New()
+	if _, err := io.Copy(h, r); err != nil {
+		return "", fmt.Errorf("hash request body: %w", err)
+	}
+	return hex.EncodeToString(h.Sum(nil)), nil
+}
+
+func hashBytes(data []byte) string {
+	sum := sha256.Sum256(data)
+	return hex.EncodeToString(sum[:])
+}

internal/utils/fetch.go

@@ -17,6 +17,7 @@ type RequestOptions struct {
 	Headers map[string]string
 	Query   map[string]string
 	Body    any
+	Prepare func(*http.Request) error
 }
 
 type APIError struct {
@@ -86,6 +87,12 @@ func BuildRequest(ctx context.Context, opts *RequestOptions) (*http.Request, err
 		req.Header.Set(k, v)
 	}
 
+	if opts.Prepare != nil {
+		if err := opts.Prepare(req); err != nil {
+			return nil, fmt.Errorf("prepare request: %w", err)
+		}
+	}
+
 	return req, nil
 }
 

provider/openai/completions/completions.go

@@ -15,9 +15,10 @@ import (
 const defaultBaseURL = "https://api.openai.com/v1"
 
 type Provider struct {
-	apiKey     string
-	baseURL    string
-	httpClient *http.Client
+	apiKey         string
+	baseURL        string
+	httpClient     *http.Client
+	prepareRequest func(*http.Request) error
 }
 
 type Option func(*Provider)
@@ -28,6 +29,22 @@ func WithAPIKey(apiKey string) Option {
 	}
 }
 
+// WithBedrockRegion enables AWS SigV4 authentication for Amazon Bedrock's
+// OpenAI-compatible endpoint using the default AWS credential chain.
+func WithBedrockRegion(region string) Option {
+	return func(p *Provider) {
+		p.prepareRequest = utils.NewBedrockDefaultCredentialsPreparer(region)
+	}
+}
+
+// WithBedrockCredentials enables AWS SigV4 authentication for Amazon Bedrock's
+// OpenAI-compatible endpoint using static credentials.
+func WithBedrockCredentials(region, accessKeyID, secretAccessKey, sessionToken string) Option {
+	return func(p *Provider) {
+		p.prepareRequest = utils.NewBedrockStaticCredentialsPreparer(region, accessKeyID, secretAccessKey, sessionToken)
+	}
+}
+
 func WithBaseURL(baseURL string) Option {
 	return func(p *Provider) {
 		p.baseURL = baseURL
@@ -60,7 +77,8 @@ func (p *Provider) ListModels(ctx context.Context) ([]sdk.Model, error) {
 		Method:  http.MethodGet,
 		BaseURL: p.baseURL,
 		Path:    "/models",
-		Headers: utils.AuthHeader(p.apiKey),
+		Headers: p.authHeaders(),
+		Prepare: p.prepareRequest,
 	})
 	if err != nil {
 		return nil, fmt.Errorf("openai: list models request failed: %w", err)
@@ -83,7 +101,8 @@ func (p *Provider) Test(ctx context.Context) *sdk.ProviderTestResult {
 		BaseURL: p.baseURL,
 		Path:    "/models",
 		Query:   map[string]string{"limit": "1"},
-		Headers: utils.AuthHeader(p.apiKey),
+		Headers: p.authHeaders(),
+		Prepare: p.prepareRequest,
 	})
 	if err != nil {
 		return classifyError(err)
@@ -96,7 +115,8 @@ func (p *Provider) TestModel(ctx context.Context, modelID string) (*sdk.ModelTes
 		Method:  http.MethodGet,
 		BaseURL: p.baseURL,
 		Path:    "/models/" + modelID,
-		Headers: utils.AuthHeader(p.apiKey),
+		Headers: p.authHeaders(),
+		Prepare: p.prepareRequest,
 	})
 	if err == nil {
 		return &sdk.ModelTestResult{Supported: true, Message: "supported"}, nil
@@ -112,7 +132,8 @@ func (p *Provider) TestModel(ctx context.Context, modelID string) (*sdk.ModelTes
 		Method:  http.MethodPost,
 		BaseURL: p.baseURL,
 		Path:    "/chat/completions",
-		Headers: utils.AuthHeader(p.apiKey),
+		Headers: p.authHeaders(),
+		Prepare: p.prepareRequest,
 		Body: map[string]any{
 			"model":      modelID,
 			"messages":   []map[string]string{{"role": "user", "content": "hi"}},
@@ -147,7 +168,8 @@ func (p *Provider) DoGenerate(ctx context.Context, params sdk.GenerateParams) (*
 		Method:  http.MethodPost,
 		BaseURL: p.baseURL,
 		Path:    "/chat/completions",
-		Headers: utils.AuthHeader(p.apiKey),
+		Headers: p.authHeaders(),
+		Prepare: p.prepareRequest,
 		Body:    req,
 	})
 	if err != nil {
@@ -407,7 +429,8 @@ func (p *Provider) DoStream(ctx context.Context, params sdk.GenerateParams) (*sd
 			Method:  http.MethodPost,
 			BaseURL: p.baseURL,
 			Path:    "/chat/completions",
-			Headers: utils.AuthHeader(p.apiKey),
+			Headers: p.authHeaders(),
+			Prepare: p.prepareRequest,
 			Body:    req,
 		}, func(ev *utils.SSEEvent) error {
 			if ev.Data == "[DONE]" {
@@ -444,6 +467,16 @@ func (p *Provider) DoStream(ctx context.Context, params sdk.GenerateParams) (*sd
 	return &sdk.StreamResult{Stream: ch}, nil
 }
 
+func (p *Provider) authHeaders() map[string]string {
+	if p.prepareRequest != nil {
+		return nil
+	}
+	if p.apiKey == "" {
+		return nil
+	}
+	return utils.AuthHeader(p.apiKey)
+}
+
 type streamingToolCall struct {
 	id       string
 	name     string