Skip to content

Commit e48ad21

Browse files
committed
Add AWS Bedrock multi-provider support
1 parent 18a10d4 commit e48ad21

8 files changed

Lines changed: 346 additions & 9 deletions

File tree

Lines changed: 54 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,54 @@
1+
# --------------------------------------------------------------------
2+
# Copyright (c) 2026, WSO2 LLC. (https://www.wso2.com).
3+
#
4+
# Licensed under the Apache License, Version 2.0 (the "License");
5+
# you may not use this file except in compliance with the License.
6+
# You may obtain a copy of the License at
7+
#
8+
# http://www.apache.org/licenses/LICENSE-2.0
9+
# --------------------------------------------------------------------
10+
11+
# --------------------------------------------------------------------
12+
# OpenAI -> AWS Bedrock LLM Proxy (single-provider)
13+
#
14+
# Exposes an OpenAI-shaped /chat/completions endpoint that the
15+
# `openai-to-bedrock-transformer` policy translates into Bedrock's Converse API format
16+
# (path /model/{modelId}/converse[-stream], body rewrite) and forwards to the
17+
# existing `bedrock-provider` upstream. For streaming requests the policy
18+
# decodes Bedrock's binary Amazon event-stream and re-emits OpenAI SSE, so
19+
# OpenAI clients get `data: {...}` chunks terminated by `data: [DONE]`.
20+
#
21+
# Single-provider mode: there is no router in front of the policy, so nothing
22+
# sets metadata["selected_provider"] and the translator runs on every request.
23+
#
24+
# The provider must already be deployed (see bedrock-provider.yaml) with its
25+
# Bearer key on upstream.auth (the policy does NOT inject the Bedrock key).
26+
# The provider example's api-key-auth policy expects the proxy to send the
27+
# issued provider loopback key in X-API-Key.
28+
# --------------------------------------------------------------------
29+
30+
apiVersion: gateway.api-platform.wso2.com/v1
31+
kind: LlmProxy
32+
metadata:
33+
name: openai-to-bedrock
34+
spec:
35+
displayName: OpenAI to Bedrock Proxy
36+
version: v1.0
37+
context: /openai-bedrock
38+
provider:
39+
id: bedrock-provider
40+
auth:
41+
type: api-key
42+
header: X-API-Key
43+
value: REPLACE_WITH_BEDROCK_PROVIDER_LOOPBACK_KEY
44+
policies:
45+
- name: openai-to-bedrock-transformer
46+
version: v1
47+
paths:
48+
- path: /chat/completions
49+
methods: [POST]
50+
params:
51+
# Bedrock puts the model id in the URL path; the translator builds
52+
# /model/<model>/converse[-stream] from this param. Use an inference
53+
# profile id your key can invoke (e.g. us.* / eu.*).
54+
model: us.amazon.nova-lite-v1:0
Lines changed: 73 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,73 @@
1+
# --------------------------------------------------------------------
2+
# Copyright (c) 2026, WSO2 LLC. (https://www.wso2.com).
3+
#
4+
# Licensed under the Apache License, Version 2.0 (the "License");
5+
# you may not use this file except in compliance with the License.
6+
# You may obtain a copy of the License at
7+
#
8+
# http://www.apache.org/licenses/LICENSE-2.0
9+
# --------------------------------------------------------------------
10+
11+
# --------------------------------------------------------------------
12+
# AWS Bedrock LLM Provider Configuration (native Converse API)
13+
#
14+
# Uses the built-in `awsbedrock` template and points to the region-specific
15+
# Bedrock runtime endpoint. AWS Bedrock now supports Bearer-token auth
16+
# (Amazon Bedrock API keys) in addition to the legacy SigV4 flow -- this
17+
# provider uses the Bearer key. Replace the placeholder with a real key and
18+
# set the URL to the region your key and model access live in.
19+
#
20+
# The model id is carried in the URL path (/model/{modelId}/converse), so the
21+
# awsbedrock template extracts request/response model via location: pathParam
22+
# (a RE2 capture group), for both /converse and /converse-stream.
23+
#
24+
# Auth: upstream.auth is the gateway -> Bedrock vendor auth. The header is
25+
# Authorization and the value is sent verbatim, so the "Bearer " prefix MUST
26+
# be included here (the portal auto-prepends it; raw YAML does not). Proxies in
27+
# front of this provider inherit this vendor auth automatically.
28+
# --------------------------------------------------------------------
29+
30+
apiVersion: gateway.api-platform.wso2.com/v1
31+
kind: LlmProvider
32+
metadata:
33+
name: bedrock-provider
34+
spec:
35+
displayName: AWS Bedrock Provider
36+
version: v1.0
37+
template: awsbedrock
38+
context: /bedrock
39+
upstream:
40+
# Match the region your key/model access is in. For the no-AWS mock
41+
# (bedrock-mock.py), use: http://host.docker.internal:9010
42+
url: https://bedrock-runtime.us-east-1.amazonaws.com
43+
auth:
44+
type: api-key
45+
header: Authorization
46+
# Bedrock expects: Authorization: Bearer <key> -- sent verbatim, so the
47+
# "Bearer " prefix MUST be included here. The mock ignores the value.
48+
value: Bearer REPLACE_WITH_BEDROCK_API_KEY
49+
policies:
50+
- name: api-key-auth
51+
version: v1
52+
paths:
53+
- path: /model/{modelId}/converse
54+
methods: [POST]
55+
params:
56+
key: X-API-Key
57+
in: header
58+
- path: /model/{modelId}/converse-stream
59+
methods: [POST]
60+
params:
61+
key: X-API-Key
62+
in: header
63+
accessControl:
64+
mode: deny_all
65+
exceptions:
66+
- path: /model/{modelId}/converse
67+
methods: [POST]
68+
- path: /model/{modelId}/converse-stream
69+
methods: [POST]
70+
71+
# The proxy in front of this provider authenticates over the internal loopback
72+
# route with an API key issued for this LlmProvider (see bedrock-testing.md,
73+
# "Issue a consumer key"). Set the issued value as the proxy's provider.auth.value.

gateway/gateway-controller/default-llm-provider-templates/awsbedrock-template.yaml

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -36,7 +36,7 @@ spec:
3636
identifier: $.usage.totalTokens
3737
requestModel:
3838
location: pathParam
39-
identifier: (?<=model/)[a-zA-Z0-9.:-]+(?=/)
39+
identifier: model/([A-Za-z0-9.:-]+)/
4040
responseModel:
4141
location: pathParam
42-
identifier: (?<=model/)[a-zA-Z0-9.:-]+(?=/)
42+
identifier: model/([A-Za-z0-9.:-]+)/

gateway/gateway-controller/pkg/utils/llm_transformer.go

Lines changed: 11 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -179,15 +179,22 @@ func (t *LLMProviderTransformer) transformProxy(proxy *api.LLMProxyConfiguration
179179
if err != nil {
180180
return nil, fmt.Errorf("failed to get context for additional provider '%s': %w", ap.Id, err)
181181
}
182-
addURL := fmt.Sprintf("%s://%s:%d%s",
183-
constants.SchemeHTTP, constants.LocalhostIP, t.routerConfig.ListenerPort, addCtx)
184-
defs = append(defs, api.UpstreamDefinition{
182+
// Named upstream definition URLs are host-only. Keep the provider context
183+
// in basePath so dynamic provider routing does not drop it and send the
184+
// loopback request to the proxy listener root.
185+
addURL := fmt.Sprintf("%s://%s:%d",
186+
constants.SchemeHTTP, constants.LocalhostIP, t.routerConfig.ListenerPort)
187+
def := api.UpstreamDefinition{
185188
Name: name,
186189
Upstreams: []struct {
187190
Url string `json:"url" yaml:"url"`
188191
Weight *int `json:"weight,omitempty" yaml:"weight,omitempty"`
189192
}{{Url: addURL}},
190-
})
193+
}
194+
if addCtx != "" && addCtx != constants.BASE_PATH {
195+
def.BasePath = &addCtx
196+
}
197+
defs = append(defs, def)
191198
}
192199
spec.UpstreamDefinitions = &defs
193200
}

gateway/gateway-controller/pkg/utils/llm_transformer_multiprovider_test.go

Lines changed: 4 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -119,6 +119,10 @@ func TestLLMProviderTransformer_TransformProxy_AdditionalProviderAuthIsCondition
119119
require.NotNil(t, result.Spec.UpstreamDefinitions)
120120
require.Len(t, *result.Spec.UpstreamDefinitions, 1)
121121
assert.Equal(t, "anthropic-provider", (*result.Spec.UpstreamDefinitions)[0].Name)
122+
require.NotNil(t, (*result.Spec.UpstreamDefinitions)[0].BasePath)
123+
assert.Equal(t, "/anthropic-provider", *(*result.Spec.UpstreamDefinitions)[0].BasePath)
124+
require.Len(t, (*result.Spec.UpstreamDefinitions)[0].Upstreams, 1)
125+
assert.Equal(t, "http://127.0.0.1:8080", (*result.Spec.UpstreamDefinitions)[0].Upstreams[0].Url)
122126

123127
var chatOp *api.Operation
124128
for i := range result.Spec.Operations {

gateway/system-policies/analytics/analytics.go

Lines changed: 57 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -7,9 +7,11 @@ import (
77
"fmt"
88
"log/slog"
99
"reflect"
10+
"regexp"
1011
"sort"
1112
"strconv"
1213
"strings"
14+
"sync"
1315

1416
policy "github.com/wso2/api-platform/sdk/core/policy/v1alpha2"
1517
"github.com/wso2/api-platform/sdk/core/utils"
@@ -568,7 +570,7 @@ func (a *AnalyticsPolicy) OnResponseBodyChunk(_ context.Context, ctx *policy.Res
568570
// Streaming responses are SSE; the last data event carries usage fields.
569571
tokenInfo, err := extractLLMProviderAnalyticsInfoFromBytes(
570572
template, ctx.RequestHeaders, ctx.ResponseHeaders,
571-
requestBodyBytes, accumulated,
573+
requestBodyBytes, accumulated, ctx.RequestPath,
572574
)
573575
if err != nil {
574576
slog.Warn("Failed to extract LLM token info from streaming response", "error", err)
@@ -662,7 +664,7 @@ func extractLLMProviderAnalyticsInfo(template map[string]interface{}, ctx *polic
662664

663665
return extractLLMProviderAnalyticsInfoFromBytes(
664666
template, ctx.RequestHeaders, ctx.ResponseHeaders,
665-
requestBodyBytes, responseBodyBytes,
667+
requestBodyBytes, responseBodyBytes, ctx.RequestPath,
666668
)
667669
}
668670

@@ -676,6 +678,7 @@ func extractLLMProviderAnalyticsInfoFromBytes(
676678
template map[string]interface{},
677679
requestHeaders, responseHeaders *policy.Headers,
678680
requestBodyBytes, responseBodyBytes []byte,
681+
requestPath string,
679682
) (*LLMProviderAnalyticsInfo, error) {
680683
var responseJSON map[string]interface{}
681684
if len(responseBodyBytes) > 0 {
@@ -693,14 +696,15 @@ func extractLLMProviderAnalyticsInfoFromBytes(
693696
_ = json.Unmarshal(requestBodyBytes, &requestJSON)
694697
}
695698

696-
return extractLLMAnalyticsFromJSON(template, requestHeaders, responseHeaders, requestJSON, responseJSON)
699+
return extractLLMAnalyticsFromJSON(template, requestHeaders, responseHeaders, requestJSON, responseJSON, requestPath)
697700
}
698701

699702
// extractLLMAnalyticsFromJSON is the core extraction logic operating on pre-parsed JSON maps.
700703
func extractLLMAnalyticsFromJSON(
701704
template map[string]interface{},
702705
requestHeaders, responseHeaders *policy.Headers,
703706
requestJSON, responseJSON map[string]interface{},
707+
requestPath string,
704708
) (*LLMProviderAnalyticsInfo, error) {
705709
if template == nil {
706710
return nil, fmt.Errorf("template is nil")
@@ -754,6 +758,14 @@ func extractLLMAnalyticsFromJSON(
754758
}
755759
}
756760
return nil, fmt.Errorf("header %s not found", identifier)
761+
case "pathparam":
762+
// Model ids for providers like AWS Bedrock and Gemini live in the
763+
// request URL path (e.g. /model/{modelId}/converse), not the body or
764+
// a header. The identifier is a regex whose first capture group (when
765+
// present) is the value; otherwise the whole match is used. The path
766+
// is available for both buffered and streaming responses, so this
767+
// meters the model in either case.
768+
return extractPathParam(requestPath, identifier)
757769
default:
758770
return nil, fmt.Errorf("unsupported location %s", location)
759771
}
@@ -806,6 +818,48 @@ func extractLLMAnalyticsFromJSON(
806818
return info, nil
807819
}
808820

821+
// pathParamRegexCache memoises compiled pathParam identifiers so the response
822+
// data path does not recompile a provider template's regex on every request.
823+
var pathParamRegexCache sync.Map // map[string]*regexp.Regexp
824+
825+
// extractPathParam applies a pathParam identifier (a regex) to the request URL
826+
// path and returns the first capture group when the pattern defines one,
827+
// otherwise the whole match. AWS Bedrock and Gemini carry the model id in the
828+
// path (e.g. /model/{modelId}/converse), which is why request/response model is
829+
// declared with `location: pathParam`. The path is present for both buffered
830+
// and streaming responses, so the model meters in either case. Go's regexp is
831+
// RE2 — identifiers must use a capture group, not lookaround.
832+
func extractPathParam(requestPath, pattern string) (string, error) {
833+
if requestPath == "" {
834+
return "", fmt.Errorf("request path not available")
835+
}
836+
re, err := compilePathParamRegex(pattern)
837+
if err != nil {
838+
return "", err
839+
}
840+
match := re.FindStringSubmatch(requestPath)
841+
if match == nil {
842+
return "", fmt.Errorf("pathParam identifier did not match request path")
843+
}
844+
if len(match) > 1 {
845+
return match[1], nil
846+
}
847+
return match[0], nil
848+
}
849+
850+
// compilePathParamRegex compiles (and caches) a pathParam identifier regex.
851+
func compilePathParamRegex(pattern string) (*regexp.Regexp, error) {
852+
if cached, ok := pathParamRegexCache.Load(pattern); ok {
853+
return cached.(*regexp.Regexp), nil
854+
}
855+
re, err := regexp.Compile(pattern)
856+
if err != nil {
857+
return nil, fmt.Errorf("invalid pathParam identifier %q: %w", pattern, err)
858+
}
859+
pathParamRegexCache.Store(pattern, re)
860+
return re, nil
861+
}
862+
809863
// populateTokenAnalyticsMetadata copies LLM token fields into an analytics metadata map.
810864
func populateTokenAnalyticsMetadata(analyticsMetadata map[string]any, tokenInfo *LLMProviderAnalyticsInfo) {
811865
if tokenInfo.PromptTokens != nil {

0 commit comments

Comments
 (0)