mirror of
https://github.com/mulgadc/spinifex.git
synced 2026-10-15 17:11:30 +00:00
feat: semantic topic detection for guardrails via embedding similarity (#908)
* feat: semantic topic detection for guardrails via embedding similarity Guardrail topic policy previously only blocked on an exact word-boundary match against a denied topic's name or listed examples, so a paraphrase or synonym for a blocked topic sailed straight through. This adds an embedding-similarity check: each denied topic's name, definition, and examples are embedded and cached, incoming text is embedded and compared by max cosine similarity against those phrase vectors, and a topic blocks when that similarity clears a conservative threshold. The literal exact-match check still runs first and always blocks on its own. If no embedder is configured, or an embed call fails, the check falls back to the existing literal match alone rather than ever failing open. * fix: fail closed when guardrail topic embedding is unavailable The topic policy's semantic check silently fell back to its literal-only matcher whenever the embedder was configured but temporarily unreachable (for example during model cold-start), so a paraphrase of a denied topic could pass through unblocked while the embedding endpoint was down. Now that case is distinguished from an unconfigured embedder: if literal matching already caught the request, that block still wins outright. Otherwise, when the embedder is wired but the embed call itself fails, the guardrail check returns an error and the caller aborts the request rather than letting unverified content through.
This commit is contained in:
@@ -107,6 +107,7 @@ linters:
|
||||
- G116 # bidi chars — bidichk is more thorough
|
||||
- G117 # use of net/http serve functions (we handle TLS separately)
|
||||
- G703 # path traversal taint analysis — false positives on config file loading
|
||||
- G704 # SSRF taint analysis — false positives on resolver-derived endpoint URLs
|
||||
- G705 # XSS taint analysis — false positives on internal error responses
|
||||
|
||||
govet:
|
||||
|
||||
@@ -327,3 +327,10 @@ func (gw *GatewayConfig) bedrockEndpointResolver() gateway_bedrock.EndpointResol
|
||||
}
|
||||
return gateway_bedrock.NewStaticEndpointResolver(gw.BedrockEndpoints)
|
||||
}
|
||||
|
||||
// bedrockEmbedder returns gw.BedrockEmbedder, or nil when unconfigured. A nil
|
||||
// Embedder is a valid value throughout gateway_bedrock's guardrail engine:
|
||||
// it simply falls back to topicPolicy's literal matcher.
|
||||
func (gw *GatewayConfig) bedrockEmbedder() gateway_bedrock.Embedder {
|
||||
return gw.BedrockEmbedder
|
||||
}
|
||||
|
||||
@@ -276,7 +276,7 @@ func (p *anthropicProvider) ConverseStream(ctx context.Context, modelID string,
|
||||
// than merely stopping our own reads — a disconnected client must free
|
||||
// the Anthropic-side generation, not just this goroutine.
|
||||
reqCtx, cancel := context.WithCancel(ctx)
|
||||
httpReq, err := http.NewRequestWithContext(reqCtx, http.MethodPost, p.baseURL+anthropicMessagesPath, bytes.NewReader(reqBody)) //nolint:gosec // G704: p.baseURL is the hardcoded Anthropic API endpoint (or an httptest stub in tests), never user input
|
||||
httpReq, err := http.NewRequestWithContext(reqCtx, http.MethodPost, p.baseURL+anthropicMessagesPath, bytes.NewReader(reqBody))
|
||||
if err != nil {
|
||||
cancel()
|
||||
slog.Error("anthropic stream: failed to build request", "model", modelID, "err", err)
|
||||
@@ -287,7 +287,7 @@ func (p *anthropicProvider) ConverseStream(ctx context.Context, modelID string,
|
||||
httpReq.Header.Set("Content-Type", "application/json")
|
||||
httpReq.Header.Set("Accept", "text/event-stream")
|
||||
|
||||
resp, err := p.httpClient.Do(httpReq) //nolint:gosec // G704: httpReq targets p.baseURL, not user input
|
||||
resp, err := p.httpClient.Do(httpReq)
|
||||
if err != nil {
|
||||
cancel()
|
||||
slog.Error("anthropic stream: request failed", "model", modelID, "err", err)
|
||||
|
||||
@@ -56,7 +56,7 @@ func TestContract_ConverseStream_SelfHost_RealSDKConsumerDecodesFrames(t *testin
|
||||
if !assert.NoError(t, err) {
|
||||
return
|
||||
}
|
||||
err = ConverseStream(r.Context(), w, "000000000001", m[1], body, nil, endpoints, nil, grantAll{}, nil, nil)
|
||||
err = ConverseStream(r.Context(), w, "000000000001", m[1], body, nil, endpoints, nil, grantAll{}, nil, nil, nil)
|
||||
assert.NoError(t, err)
|
||||
})
|
||||
|
||||
@@ -200,7 +200,7 @@ func TestContract_InvokeModelWithResponseStream_SelfHost_RealSDKConsumerDecodesF
|
||||
if !assert.NoError(t, err) {
|
||||
return
|
||||
}
|
||||
err = InvokeModelWithResponseStream(r.Context(), w, "000000000001", m[1], body, nil, endpoints, r.Header.Get("Content-Type"), nil, grantAll{}, nil, "", "", nil)
|
||||
err = InvokeModelWithResponseStream(r.Context(), w, "000000000001", m[1], body, nil, endpoints, r.Header.Get("Content-Type"), nil, grantAll{}, nil, "", "", nil, nil)
|
||||
assert.NoError(t, err)
|
||||
})
|
||||
|
||||
|
||||
@@ -100,8 +100,10 @@ type converseStreamSource interface {
|
||||
// an awserrors code for the normal ErrorHandler envelope; once the first
|
||||
// frame is written it always returns nil, surfacing later failures as an
|
||||
// in-band exception event instead. provisioned and guardrails may be nil,
|
||||
// which fails closed on a PT ARN or GuardrailConfig respectively.
|
||||
func ConverseStream(ctx context.Context, w http.ResponseWriter, accountID, modelID string, body []byte, resolver CredentialResolver, endpointResolver EndpointResolver, recorder Recorder, access AccessResolver, provisioned *ProvisionedStore, guardrails *GuardrailStore) error {
|
||||
// which fails closed on a PT ARN or GuardrailConfig respectively. embedder
|
||||
// drives topicPolicy's semantic match; nil falls back to the literal
|
||||
// matcher.
|
||||
func ConverseStream(ctx context.Context, w http.ResponseWriter, accountID, modelID string, body []byte, resolver CredentialResolver, endpointResolver EndpointResolver, recorder Recorder, access AccessResolver, provisioned *ProvisionedStore, guardrails *GuardrailStore, embedder Embedder) error {
|
||||
if recorder == nil {
|
||||
recorder = NoopRecorder
|
||||
}
|
||||
@@ -143,7 +145,7 @@ func ConverseStream(ctx context.Context, w http.ResponseWriter, accountID, model
|
||||
traceEnabled := gc != nil && aws.StringValue(gc.Trace) == bedrockruntime.GuardrailTraceEnabled
|
||||
if gc != nil {
|
||||
ident, version := aws.StringValue(gc.GuardrailIdentifier), aws.StringValue(gc.GuardrailVersion)
|
||||
blocked, message, _, assessments, err := enforceGuardrail(ctx, guardrails, accountID, ident, version,
|
||||
blocked, message, _, assessments, err := enforceGuardrail(ctx, guardrails, embedder, accountID, ident, version,
|
||||
bedrockruntime.GuardrailContentSourceInput, streamGuardrailTexts(input))
|
||||
if err != nil {
|
||||
return err
|
||||
@@ -154,7 +156,7 @@ func ConverseStream(ctx context.Context, w http.ResponseWriter, accountID, model
|
||||
}
|
||||
}
|
||||
|
||||
src, err := NewRouter(resolver, endpointResolver, recorder, access, provisioned, guardrails).ConverseStream(ctx, accountID, modelID, input)
|
||||
src, err := NewRouter(resolver, endpointResolver, recorder, access, provisioned, guardrails, embedder).ConverseStream(ctx, accountID, modelID, input)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
@@ -169,7 +171,7 @@ func ConverseStream(ctx context.Context, w http.ResponseWriter, accountID, model
|
||||
// forward unassessed model text.
|
||||
if gc != nil {
|
||||
ident, version := aws.StringValue(gc.GuardrailIdentifier), aws.StringValue(gc.GuardrailVersion)
|
||||
src = newGuardrailStreamSource(src, guardrails, accountID, ident, version, traceEnabled, inputAssessments)
|
||||
src = newGuardrailStreamSource(src, guardrails, embedder, accountID, ident, version, traceEnabled, inputAssessments)
|
||||
}
|
||||
|
||||
fw, err := newFrameWriter(w)
|
||||
|
||||
@@ -154,7 +154,7 @@ func TestConverseStream_UnknownModelReturnsResourceNotFoundPreHeader(t *testing.
|
||||
rec := httptest.NewRecorder()
|
||||
body := []byte(`{"messages":[{"role":"user","content":[{"text":"hi"}]}]}`)
|
||||
|
||||
err := ConverseStream(context.Background(), rec, "000000000001", "does.not-exist-v1:0", body, nil, nil, nil, grantAll{}, nil, nil)
|
||||
err := ConverseStream(context.Background(), rec, "000000000001", "does.not-exist-v1:0", body, nil, nil, nil, grantAll{}, nil, nil, nil)
|
||||
require.Error(t, err)
|
||||
assert.Equal(t, awserrors.ErrorResourceNotFoundException, err.Error())
|
||||
// A pre-stream failure must not have written anything: the gateway's
|
||||
@@ -164,7 +164,7 @@ func TestConverseStream_UnknownModelReturnsResourceNotFoundPreHeader(t *testing.
|
||||
|
||||
func TestConverseStream_MalformedBodyReturnsValidationException(t *testing.T) {
|
||||
rec := httptest.NewRecorder()
|
||||
err := ConverseStream(context.Background(), rec, "000000000001", "meta.llama3-2-1b-instruct-v1:0", []byte("{not-json"), nil, nil, nil, grantAll{}, nil, nil)
|
||||
err := ConverseStream(context.Background(), rec, "000000000001", "meta.llama3-2-1b-instruct-v1:0", []byte("{not-json"), nil, nil, nil, grantAll{}, nil, nil, nil)
|
||||
require.Error(t, err)
|
||||
assert.Equal(t, awserrors.ErrorValidationException, err.Error())
|
||||
}
|
||||
@@ -186,7 +186,7 @@ func TestConverseStream_SelfHostHappyPath_WritesFramedTaxonomy(t *testing.T) {
|
||||
})
|
||||
require.NoError(t, err)
|
||||
|
||||
err = ConverseStream(context.Background(), rec, "000000000001", modelID, body, nil, NewStaticEndpointResolver(map[string]string{modelID: ts.URL}), nil, grantAll{}, nil, nil)
|
||||
err = ConverseStream(context.Background(), rec, "000000000001", modelID, body, nil, NewStaticEndpointResolver(map[string]string{modelID: ts.URL}), nil, grantAll{}, nil, nil, nil)
|
||||
require.NoError(t, err)
|
||||
|
||||
assert.Equal(t, http.StatusOK, rec.Code)
|
||||
@@ -300,7 +300,7 @@ func TestConverseStream_ClientDisconnectAbortsUpstreamRequest(t *testing.T) {
|
||||
done := make(chan struct{})
|
||||
go func() {
|
||||
defer close(done)
|
||||
_ = ConverseStream(ctx, rec, "000000000001", modelID, body, nil, NewStaticEndpointResolver(map[string]string{modelID: ts.URL}), nil, grantAll{}, nil, nil)
|
||||
_ = ConverseStream(ctx, rec, "000000000001", modelID, body, nil, NewStaticEndpointResolver(map[string]string{modelID: ts.URL}), nil, grantAll{}, nil, nil, nil)
|
||||
}()
|
||||
|
||||
select {
|
||||
@@ -336,7 +336,7 @@ func TestConverseStream_NonFlusherWriter_ReturnsPreHeaderError(t *testing.T) {
|
||||
})
|
||||
require.NoError(t, err)
|
||||
|
||||
err = ConverseStream(context.Background(), w, "000000000001", modelID, body, nil, NewStaticEndpointResolver(map[string]string{modelID: ts.URL}), nil, grantAll{}, nil, nil)
|
||||
err = ConverseStream(context.Background(), w, "000000000001", modelID, body, nil, NewStaticEndpointResolver(map[string]string{modelID: ts.URL}), nil, grantAll{}, nil, nil, nil)
|
||||
require.Error(t, err)
|
||||
assert.Equal(t, awserrors.ErrorInternalError, err.Error())
|
||||
assert.False(t, w.wroteHeader)
|
||||
|
||||
@@ -533,8 +533,10 @@ func CreateGuardrailVersion(ctx context.Context, accountID string, store *Guardr
|
||||
// input.Source, honouring input.GuardrailVersion (DRAFT or a numbered
|
||||
// snapshot), and returns the aws-sdk-go bedrockruntime shape the runtime op
|
||||
// hands back to the caller. Unknown/foreign guardrail or version both report
|
||||
// ResourceNotFoundException, the same as the control-plane ops.
|
||||
func ApplyGuardrail(ctx context.Context, accountID string, store *GuardrailStore, input *bedrockruntime.ApplyGuardrailInput) (*bedrockruntime.ApplyGuardrailOutput, error) {
|
||||
// ResourceNotFoundException, the same as the control-plane ops. embedder
|
||||
// drives topicPolicy's semantic match; a nil embedder falls back to the
|
||||
// literal matcher (see assessTopicPolicy).
|
||||
func ApplyGuardrail(ctx context.Context, accountID string, store *GuardrailStore, embedder Embedder, input *bedrockruntime.ApplyGuardrailInput) (*bedrockruntime.ApplyGuardrailOutput, error) {
|
||||
if input == nil || aws.StringValue(input.GuardrailIdentifier) == "" ||
|
||||
aws.StringValue(input.GuardrailVersion) == "" || aws.StringValue(input.Source) == "" {
|
||||
return nil, errors.New(awserrors.ErrorValidationException)
|
||||
@@ -570,7 +572,10 @@ func ApplyGuardrail(ctx context.Context, accountID string, store *GuardrailStore
|
||||
texts = append(texts, aws.StringValue(c.Text.Text))
|
||||
}
|
||||
|
||||
action, assessments, outputTexts, usage := applyGuardrailPolicies(view, texts, aws.StringValue(input.Source))
|
||||
action, assessments, outputTexts, usage, err := applyGuardrailPolicies(ctx, embedder, view, texts, aws.StringValue(input.Source))
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
|
||||
outputs := make([]*bedrockruntime.GuardrailOutputContent, 0, len(outputTexts))
|
||||
if action == bedrockruntime.GuardrailActionGuardrailIntervened {
|
||||
|
||||
@@ -1,6 +1,8 @@
|
||||
package gateway_bedrock
|
||||
|
||||
import (
|
||||
"context"
|
||||
"errors"
|
||||
"fmt"
|
||||
"regexp"
|
||||
"slices"
|
||||
@@ -8,6 +10,7 @@ import (
|
||||
"github.com/aws/aws-sdk-go/aws"
|
||||
"github.com/aws/aws-sdk-go/service/bedrock"
|
||||
"github.com/aws/aws-sdk-go/service/bedrockruntime"
|
||||
"github.com/mulgadc/spinifex/spinifex/awserrors"
|
||||
)
|
||||
|
||||
// piiPattern binds one AWS PII entity type to the regex that detects it.
|
||||
@@ -232,15 +235,52 @@ func scaffoldContentPolicyAssessment(cfg *bedrock.GuardrailContentPolicyConfig)
|
||||
return &bedrockruntime.GuardrailContentPolicyAssessment{Filters: []*bedrockruntime.GuardrailContentFilter{}}
|
||||
}
|
||||
|
||||
// assessTopicPolicy checks texts against cfg's denied topics. Matching is
|
||||
// deterministic (no classifier): a topic's Name and each of its Examples are
|
||||
// matched via matchesWord's word-boundary semantics. The long-form
|
||||
// Definition is never matched directly — it reads as prose, not a phrase,
|
||||
// and substring-matching it produces erratic hits. Each topic is appended at
|
||||
// most once even if several of its phrases hit.
|
||||
func assessTopicPolicy(cfg *bedrock.GuardrailTopicPolicyConfig, texts []string) (*bedrockruntime.GuardrailTopicPolicyAssessment, bool) {
|
||||
// literalTopicHit is the deterministic exact-match check: topic's Name and
|
||||
// each of its Examples are matched via matchesWord's word-boundary
|
||||
// semantics. The long-form Definition is never matched literally — it reads
|
||||
// as prose, not a phrase, and substring-matching it produces erratic hits.
|
||||
// This always runs regardless of semantic availability, so an exact phrase
|
||||
// hit is never lost to an embedder outage.
|
||||
func literalTopicHit(topic *bedrock.GuardrailTopicConfig, texts []string) bool {
|
||||
phrases := []string{}
|
||||
if name := aws.StringValue(topic.Name); name != "" {
|
||||
phrases = append(phrases, name)
|
||||
}
|
||||
for _, ex := range topic.Examples {
|
||||
if example := aws.StringValue(ex); example != "" {
|
||||
phrases = append(phrases, example)
|
||||
}
|
||||
}
|
||||
return slices.ContainsFunc(phrases, func(phrase string) bool { return matchesWord(texts, phrase) })
|
||||
}
|
||||
|
||||
// assessTopicPolicy checks texts against cfg's denied topics. A topic
|
||||
// blocks on literalTopicHit's exact-match, or when any input text's
|
||||
// embedding reaches topicSimilarityThreshold cosine similarity against any
|
||||
// of the topic's Name/Definition/Examples phrase vectors (Definition is
|
||||
// finally used here, unlike the literal path). embedder == nil is a
|
||||
// deliberately literal-only deployment (tests, or no embedder wired) and is
|
||||
// not an error. embedder != nil whose Embed call fails means the policy
|
||||
// could not be semantically evaluated: if literal matching already blocked
|
||||
// the request that block wins outright, otherwise this returns an error so
|
||||
// the caller fails closed instead of passing unverified content through.
|
||||
func assessTopicPolicy(ctx context.Context, embedder Embedder, cfg *bedrock.GuardrailTopicPolicyConfig, texts []string) (*bedrockruntime.GuardrailTopicPolicyAssessment, bool, error) {
|
||||
if cfg == nil {
|
||||
return nil, false
|
||||
return nil, false, nil
|
||||
}
|
||||
|
||||
var textVectors [][]float32
|
||||
semanticOK := false
|
||||
unverified := false
|
||||
if len(cfg.TopicsConfig) > 0 {
|
||||
vectors, err := embedGuardrailTexts(ctx, embedder, DefaultEmbeddingModel, texts)
|
||||
switch {
|
||||
case err != nil:
|
||||
unverified = true
|
||||
case vectors != nil:
|
||||
textVectors = vectors
|
||||
semanticOK = true
|
||||
}
|
||||
}
|
||||
|
||||
blocked := false
|
||||
@@ -250,17 +290,17 @@ func assessTopicPolicy(cfg *bedrock.GuardrailTopicPolicyConfig, texts []string)
|
||||
continue
|
||||
}
|
||||
|
||||
phrases := []string{}
|
||||
if name := aws.StringValue(topic.Name); name != "" {
|
||||
phrases = append(phrases, name)
|
||||
}
|
||||
for _, ex := range topic.Examples {
|
||||
if example := aws.StringValue(ex); example != "" {
|
||||
phrases = append(phrases, example)
|
||||
hit := literalTopicHit(topic, texts)
|
||||
if !hit && semanticOK {
|
||||
vectors, err := topicVectors(ctx, embedder, DefaultEmbeddingModel, topic)
|
||||
switch {
|
||||
case err != nil:
|
||||
unverified = true
|
||||
case vectors != nil:
|
||||
hit = topicSemanticHit(textVectors, vectors, topicSimilarityThreshold(topic))
|
||||
}
|
||||
}
|
||||
|
||||
if !slices.ContainsFunc(phrases, func(phrase string) bool { return matchesWord(texts, phrase) }) {
|
||||
if !hit {
|
||||
continue
|
||||
}
|
||||
|
||||
@@ -272,7 +312,15 @@ func assessTopicPolicy(cfg *bedrock.GuardrailTopicPolicyConfig, texts []string)
|
||||
blocked = true
|
||||
}
|
||||
|
||||
return &bedrockruntime.GuardrailTopicPolicyAssessment{Topics: topics}, blocked
|
||||
assessment := &bedrockruntime.GuardrailTopicPolicyAssessment{Topics: topics}
|
||||
switch {
|
||||
case blocked:
|
||||
return assessment, true, nil
|
||||
case unverified:
|
||||
return assessment, false, errors.New(awserrors.ErrorServiceUnavailableException)
|
||||
default:
|
||||
return assessment, false, nil
|
||||
}
|
||||
}
|
||||
|
||||
// scaffoldContextualGroundingPolicyAssessment is
|
||||
@@ -318,24 +366,31 @@ func guardrailUsage(view guardrailView, texts []string) *bedrockruntime.Guardrai
|
||||
return usage
|
||||
}
|
||||
|
||||
// applyGuardrailPolicies is the pure filter engine: it evaluates view's
|
||||
// policies over texts for source (INPUT or OUTPUT) and returns the
|
||||
// aws-sdk-go bedrockruntime shape's pieces. wordPolicy,
|
||||
// sensitiveInformationPolicy and topicPolicy are enforced deterministically;
|
||||
// applyGuardrailPolicies is the filter engine: it evaluates view's policies
|
||||
// over texts for source (INPUT or OUTPUT) and returns the aws-sdk-go
|
||||
// bedrockruntime shape's pieces. wordPolicy and sensitiveInformationPolicy
|
||||
// are enforced deterministically; topicPolicy adds embedding-similarity
|
||||
// scoring over the literal match (assessTopicPolicy, using embedder);
|
||||
// content/contextualGrounding evaluate to NONE (see the scaffold* helpers,
|
||||
// which need a classifier). A BLOCK anywhere makes the overall action
|
||||
// GUARDRAIL_INTERVENED and short-circuits redaction, since the caller
|
||||
// substitutes the guardrail's blocked messaging instead of the text. source
|
||||
// is accepted (not yet branched on) for parity with AWS's per-source
|
||||
// substitutes the guardrail's blocked messaging instead of the text. A
|
||||
// non-nil error means assessTopicPolicy could not semantically verify texts
|
||||
// and nothing else blocked outright: the caller must fail the request
|
||||
// rather than use the zero-value action/assessments returned alongside it.
|
||||
// source is accepted (not yet branched on) for parity with AWS's per-source
|
||||
// assessment shape; the deterministic policies apply identically to INPUT
|
||||
// and OUTPUT text until a content-policy classifier needs to tell them apart
|
||||
// via InputStrength/OutputStrength.
|
||||
func applyGuardrailPolicies(view guardrailView, texts []string, source string) (string, []*bedrockruntime.GuardrailAssessment, []string, *bedrockruntime.GuardrailUsage) {
|
||||
func applyGuardrailPolicies(ctx context.Context, embedder Embedder, view guardrailView, texts []string, source string) (string, []*bedrockruntime.GuardrailAssessment, []string, *bedrockruntime.GuardrailUsage, error) {
|
||||
_ = source
|
||||
|
||||
wordAssessment, wordBlocked := assessWordPolicy(view.WordPolicy, texts)
|
||||
piiAssessment, piiBlocked := assessSensitiveInformationPolicy(view.SensitiveInformationPolicy, texts)
|
||||
topicAssessment, topicBlocked := assessTopicPolicy(view.TopicPolicy, texts)
|
||||
topicAssessment, topicBlocked, topicErr := assessTopicPolicy(ctx, embedder, view.TopicPolicy, texts)
|
||||
if topicErr != nil {
|
||||
return "", nil, nil, nil, topicErr
|
||||
}
|
||||
|
||||
action := bedrockruntime.GuardrailActionNone
|
||||
outputs := texts
|
||||
@@ -354,5 +409,5 @@ func applyGuardrailPolicies(view guardrailView, texts []string, source string) (
|
||||
ContextualGroundingPolicy: scaffoldContextualGroundingPolicyAssessment(view.ContextualGroundingPolicy),
|
||||
}
|
||||
|
||||
return action, []*bedrockruntime.GuardrailAssessment{assessment}, outputs, guardrailUsage(view, texts)
|
||||
return action, []*bedrockruntime.GuardrailAssessment{assessment}, outputs, guardrailUsage(view, texts), nil
|
||||
}
|
||||
|
||||
@@ -2,6 +2,7 @@ package gateway_bedrock
|
||||
|
||||
import (
|
||||
"context"
|
||||
"errors"
|
||||
"testing"
|
||||
|
||||
"github.com/aws/aws-sdk-go/aws"
|
||||
@@ -144,7 +145,8 @@ func TestApplyGuardrailPolicies_WordAndPII(t *testing.T) {
|
||||
|
||||
for _, tc := range cases {
|
||||
t.Run(tc.name, func(t *testing.T) {
|
||||
action, assessments, outputs, usage := applyGuardrailPolicies(tc.view, tc.texts, tc.source)
|
||||
action, assessments, outputs, usage, err := applyGuardrailPolicies(context.Background(), nil, tc.view, tc.texts, tc.source)
|
||||
require.NoError(t, err)
|
||||
assert.Equal(t, tc.wantAction, action)
|
||||
require.Len(t, assessments, 1)
|
||||
require.NotNil(t, usage)
|
||||
@@ -193,7 +195,8 @@ func TestApplyGuardrailPolicies_CustomRegex(t *testing.T) {
|
||||
},
|
||||
}
|
||||
|
||||
action, assessments, outputs, _ := applyGuardrailPolicies(view, []string{"account 123456789012 is active"}, bedrockruntime.GuardrailContentSourceInput)
|
||||
action, assessments, outputs, _, err := applyGuardrailPolicies(context.Background(), nil, view, []string{"account 123456789012 is active"}, bedrockruntime.GuardrailContentSourceInput)
|
||||
require.NoError(t, err)
|
||||
assert.Equal(t, bedrockruntime.GuardrailActionNone, action)
|
||||
require.Len(t, outputs, 1)
|
||||
assert.Equal(t, "account {account-id} is active", outputs[0])
|
||||
@@ -208,7 +211,8 @@ func TestApplyGuardrailPolicies_CustomRegex(t *testing.T) {
|
||||
},
|
||||
},
|
||||
}
|
||||
action, _, _, _ = applyGuardrailPolicies(viewBlock, []string{"account 123456789012 is active"}, bedrockruntime.GuardrailContentSourceInput)
|
||||
action, _, _, _, err = applyGuardrailPolicies(context.Background(), nil, viewBlock, []string{"account 123456789012 is active"}, bedrockruntime.GuardrailContentSourceInput)
|
||||
require.NoError(t, err)
|
||||
assert.Equal(t, bedrockruntime.GuardrailActionGuardrailIntervened, action)
|
||||
}
|
||||
|
||||
@@ -235,7 +239,8 @@ func TestApplyGuardrailPolicies_ScaffoldPolicies(t *testing.T) {
|
||||
},
|
||||
}
|
||||
|
||||
action, assessments, outputs, _ := applyGuardrailPolicies(view, []string{"our competitor's product is inferior"}, bedrockruntime.GuardrailContentSourceInput)
|
||||
action, assessments, outputs, _, err := applyGuardrailPolicies(context.Background(), nil, view, []string{"our competitor's product is inferior"}, bedrockruntime.GuardrailContentSourceInput)
|
||||
require.NoError(t, err)
|
||||
assert.Equal(t, bedrockruntime.GuardrailActionNone, action)
|
||||
assert.Equal(t, []string{"our competitor's product is inferior"}, outputs)
|
||||
|
||||
@@ -251,7 +256,8 @@ func TestApplyGuardrailPolicies_ScaffoldPolicies(t *testing.T) {
|
||||
// policy the guardrail never configured has no assessment entry at all,
|
||||
// rather than an empty scaffold one.
|
||||
func TestApplyGuardrailPolicies_UnconfiguredPoliciesAreOmitted(t *testing.T) {
|
||||
_, assessments, _, _ := applyGuardrailPolicies(guardrailView{}, []string{"hello"}, bedrockruntime.GuardrailContentSourceInput)
|
||||
_, assessments, _, _, err := applyGuardrailPolicies(context.Background(), nil, guardrailView{}, []string{"hello"}, bedrockruntime.GuardrailContentSourceInput)
|
||||
require.NoError(t, err)
|
||||
require.Len(t, assessments, 1)
|
||||
assert.Nil(t, assessments[0].WordPolicy)
|
||||
assert.Nil(t, assessments[0].SensitiveInformationPolicy)
|
||||
@@ -338,7 +344,8 @@ func TestApplyGuardrailPolicies_TopicPolicy(t *testing.T) {
|
||||
|
||||
for _, tc := range cases {
|
||||
t.Run(tc.name, func(t *testing.T) {
|
||||
action, assessments, _, _ := applyGuardrailPolicies(tc.view, tc.texts, tc.source)
|
||||
action, assessments, _, _, err := applyGuardrailPolicies(context.Background(), nil, tc.view, tc.texts, tc.source)
|
||||
require.NoError(t, err)
|
||||
assert.Equal(t, tc.wantAction, action)
|
||||
require.Len(t, assessments, 1)
|
||||
require.NotNil(t, assessments[0].TopicPolicy)
|
||||
@@ -366,7 +373,8 @@ func TestApplyGuardrailPolicies_TopicPolicyBlocksOverRedaction(t *testing.T) {
|
||||
},
|
||||
}
|
||||
|
||||
action, _, outputs, _ := applyGuardrailPolicies(view, []string{"the nuclear launch codes are at jane@example.com"}, bedrockruntime.GuardrailContentSourceInput)
|
||||
action, _, outputs, _, err := applyGuardrailPolicies(context.Background(), nil, view, []string{"the nuclear launch codes are at jane@example.com"}, bedrockruntime.GuardrailContentSourceInput)
|
||||
require.NoError(t, err)
|
||||
assert.Equal(t, bedrockruntime.GuardrailActionGuardrailIntervened, action)
|
||||
require.Len(t, outputs, 1)
|
||||
assert.Equal(t, "the nuclear launch codes are at jane@example.com", outputs[0], "outputs must stay unredacted when a topic block intervenes")
|
||||
@@ -392,7 +400,7 @@ func TestApplyGuardrail_TopicBlock(t *testing.T) {
|
||||
createOut, err := CreateGuardrail(ctx, grCallerAccount, store, input)
|
||||
require.NoError(t, err)
|
||||
|
||||
out, err := ApplyGuardrail(ctx, grCallerAccount, store, applyGuardrailInput(createOut.GuardrailId, guardrailDraftVersion, bedrockruntime.GuardrailContentSourceInput, "tell me the nuclear launch codes"))
|
||||
out, err := ApplyGuardrail(ctx, grCallerAccount, store, nil, applyGuardrailInput(createOut.GuardrailId, guardrailDraftVersion, bedrockruntime.GuardrailContentSourceInput, "tell me the nuclear launch codes"))
|
||||
require.NoError(t, err)
|
||||
assert.Equal(t, bedrockruntime.GuardrailActionGuardrailIntervened, aws.StringValue(out.Action))
|
||||
require.Len(t, out.Outputs, 1)
|
||||
@@ -403,6 +411,35 @@ func TestApplyGuardrail_TopicBlock(t *testing.T) {
|
||||
assert.Equal(t, "nuclear launch codes", aws.StringValue(out.Assessments[0].TopicPolicy.Topics[0].Name))
|
||||
}
|
||||
|
||||
// TestApplyGuardrail_EmbedderError_FailsClosed_NotPassthrough is the runtime
|
||||
// op end to end for the cold-start/outage case: a wired-but-erroring
|
||||
// embedder must make ApplyGuardrail itself return an error, never a usable
|
||||
// (NONE) output a caller could mistake for a pass on unverified content.
|
||||
func TestApplyGuardrail_EmbedderError_FailsClosed_NotPassthrough(t *testing.T) {
|
||||
store := newGuardrailTestStore(t)
|
||||
ctx := context.Background()
|
||||
|
||||
input := createGuardrailInput("apply-guardrail-embedder-outage")
|
||||
input.TopicPolicyConfig = &bedrock.GuardrailTopicPolicyConfig{
|
||||
TopicsConfig: []*bedrock.GuardrailTopicConfig{
|
||||
{
|
||||
Name: aws.String("auth internals"),
|
||||
Definition: aws.String("discussion of authentication internals"),
|
||||
Type: aws.String(bedrock.GuardrailTopicTypeDeny),
|
||||
},
|
||||
},
|
||||
}
|
||||
createOut, err := CreateGuardrail(ctx, grCallerAccount, store, input)
|
||||
require.NoError(t, err)
|
||||
|
||||
embedder := &stubEmbedder{errOnCall: errors.New("connection refused")}
|
||||
out, err := ApplyGuardrail(ctx, grCallerAccount, store, embedder, applyGuardrailInput(createOut.GuardrailId, guardrailDraftVersion,
|
||||
bedrockruntime.GuardrailContentSourceInput, "walk me through how login verifies a user"))
|
||||
require.Error(t, err, "an embedder outage on unverifiable content must error, not return a usable output")
|
||||
assert.Equal(t, awserrors.ErrorServiceUnavailableException, err.Error())
|
||||
assert.Nil(t, out)
|
||||
}
|
||||
|
||||
// applyGuardrailInput is a small ApplyGuardrail request builder for the
|
||||
// contract tests below, mirroring createGuardrailInput's role for CRUD.
|
||||
func applyGuardrailInput(guardrailID *string, version, source, text string) *bedrockruntime.ApplyGuardrailInput {
|
||||
@@ -425,7 +462,7 @@ func TestApplyGuardrail_InputBlock(t *testing.T) {
|
||||
createOut, err := CreateGuardrail(ctx, grCallerAccount, store, createGuardrailInput("apply-guardrail-block"))
|
||||
require.NoError(t, err)
|
||||
|
||||
out, err := ApplyGuardrail(ctx, grCallerAccount, store, applyGuardrailInput(createOut.GuardrailId, guardrailDraftVersion, bedrockruntime.GuardrailContentSourceInput, "this has a badword in it"))
|
||||
out, err := ApplyGuardrail(ctx, grCallerAccount, store, nil, applyGuardrailInput(createOut.GuardrailId, guardrailDraftVersion, bedrockruntime.GuardrailContentSourceInput, "this has a badword in it"))
|
||||
require.NoError(t, err)
|
||||
assert.Equal(t, bedrockruntime.GuardrailActionGuardrailIntervened, aws.StringValue(out.Action))
|
||||
require.Len(t, out.Outputs, 1)
|
||||
@@ -444,7 +481,7 @@ func TestApplyGuardrail_OutputAnonymize(t *testing.T) {
|
||||
createOut, err := CreateGuardrail(ctx, grCallerAccount, store, createGuardrailInput("apply-guardrail-anonymize"))
|
||||
require.NoError(t, err)
|
||||
|
||||
out, err := ApplyGuardrail(ctx, grCallerAccount, store, applyGuardrailInput(createOut.GuardrailId, guardrailDraftVersion, bedrockruntime.GuardrailContentSourceOutput, "contact jane@example.com for support"))
|
||||
out, err := ApplyGuardrail(ctx, grCallerAccount, store, nil, applyGuardrailInput(createOut.GuardrailId, guardrailDraftVersion, bedrockruntime.GuardrailContentSourceOutput, "contact jane@example.com for support"))
|
||||
require.NoError(t, err)
|
||||
assert.Equal(t, bedrockruntime.GuardrailActionNone, aws.StringValue(out.Action))
|
||||
require.Len(t, out.Outputs, 1)
|
||||
@@ -457,14 +494,14 @@ func TestApplyGuardrail_UnknownOrForeignGuardrail(t *testing.T) {
|
||||
store := newGuardrailTestStore(t)
|
||||
ctx := context.Background()
|
||||
|
||||
_, err := ApplyGuardrail(ctx, grCallerAccount, store, applyGuardrailInput(aws.String("does-not-exist"), guardrailDraftVersion, bedrockruntime.GuardrailContentSourceInput, "hello"))
|
||||
_, err := ApplyGuardrail(ctx, grCallerAccount, store, nil, applyGuardrailInput(aws.String("does-not-exist"), guardrailDraftVersion, bedrockruntime.GuardrailContentSourceInput, "hello"))
|
||||
require.Error(t, err)
|
||||
assert.True(t, awserrors.IsErrorCode(err, awserrors.ErrorResourceNotFoundException))
|
||||
|
||||
createOut, err := CreateGuardrail(ctx, grCallerAccount, store, createGuardrailInput("apply-guardrail-foreign"))
|
||||
require.NoError(t, err)
|
||||
|
||||
_, err = ApplyGuardrail(ctx, grOtherCaller, store, applyGuardrailInput(createOut.GuardrailId, guardrailDraftVersion, bedrockruntime.GuardrailContentSourceInput, "hello"))
|
||||
_, err = ApplyGuardrail(ctx, grOtherCaller, store, nil, applyGuardrailInput(createOut.GuardrailId, guardrailDraftVersion, bedrockruntime.GuardrailContentSourceInput, "hello"))
|
||||
require.Error(t, err)
|
||||
assert.True(t, awserrors.IsErrorCode(err, awserrors.ErrorResourceNotFoundException))
|
||||
}
|
||||
@@ -495,21 +532,21 @@ func TestApplyGuardrail_VersionResolution(t *testing.T) {
|
||||
require.NoError(t, err)
|
||||
|
||||
// The numbered version still enforces the original "badword" blocklist.
|
||||
outSnap, err := ApplyGuardrail(ctx, grCallerAccount, store, applyGuardrailInput(createOut.GuardrailId, aws.StringValue(verOut.Version), bedrockruntime.GuardrailContentSourceInput, "this has a badword in it"))
|
||||
outSnap, err := ApplyGuardrail(ctx, grCallerAccount, store, nil, applyGuardrailInput(createOut.GuardrailId, aws.StringValue(verOut.Version), bedrockruntime.GuardrailContentSourceInput, "this has a badword in it"))
|
||||
require.NoError(t, err)
|
||||
assert.Equal(t, bedrockruntime.GuardrailActionGuardrailIntervened, aws.StringValue(outSnap.Action))
|
||||
|
||||
// The DRAFT no longer blocks "badword" but does block "newword".
|
||||
outDraftOld, err := ApplyGuardrail(ctx, grCallerAccount, store, applyGuardrailInput(createOut.GuardrailId, guardrailDraftVersion, bedrockruntime.GuardrailContentSourceInput, "this has a badword in it"))
|
||||
outDraftOld, err := ApplyGuardrail(ctx, grCallerAccount, store, nil, applyGuardrailInput(createOut.GuardrailId, guardrailDraftVersion, bedrockruntime.GuardrailContentSourceInput, "this has a badword in it"))
|
||||
require.NoError(t, err)
|
||||
assert.Equal(t, bedrockruntime.GuardrailActionNone, aws.StringValue(outDraftOld.Action))
|
||||
|
||||
outDraftNew, err := ApplyGuardrail(ctx, grCallerAccount, store, applyGuardrailInput(createOut.GuardrailId, guardrailDraftVersion, bedrockruntime.GuardrailContentSourceInput, "this has a newword in it"))
|
||||
outDraftNew, err := ApplyGuardrail(ctx, grCallerAccount, store, nil, applyGuardrailInput(createOut.GuardrailId, guardrailDraftVersion, bedrockruntime.GuardrailContentSourceInput, "this has a newword in it"))
|
||||
require.NoError(t, err)
|
||||
assert.Equal(t, bedrockruntime.GuardrailActionGuardrailIntervened, aws.StringValue(outDraftNew.Action))
|
||||
|
||||
// An unknown numbered version is not-found.
|
||||
_, err = ApplyGuardrail(ctx, grCallerAccount, store, applyGuardrailInput(createOut.GuardrailId, "99", bedrockruntime.GuardrailContentSourceInput, "hello"))
|
||||
_, err = ApplyGuardrail(ctx, grCallerAccount, store, nil, applyGuardrailInput(createOut.GuardrailId, "99", bedrockruntime.GuardrailContentSourceInput, "hello"))
|
||||
require.Error(t, err)
|
||||
assert.True(t, awserrors.IsErrorCode(err, awserrors.ErrorResourceNotFoundException))
|
||||
}
|
||||
|
||||
@@ -45,13 +45,18 @@ func loadGuardrailView(ctx context.Context, store *GuardrailStore, accountID, id
|
||||
// enforceGuardrail is the single resolve->load->filter path Converse and
|
||||
// ConverseStream both call, mirroring ApplyGuardrail's own sequence. An
|
||||
// unresolvable or foreign guardrail fails closed: ResourceNotFoundException.
|
||||
func enforceGuardrail(ctx context.Context, store *GuardrailStore, accountID, ident, version, source string, texts []string) (blocked bool, blockedMessage string, redactedTexts []string, assessments []*bedrockruntime.GuardrailAssessment, err error) {
|
||||
// embedder drives topicPolicy's semantic match; a nil embedder falls back to
|
||||
// the literal matcher (see assessTopicPolicy).
|
||||
func enforceGuardrail(ctx context.Context, store *GuardrailStore, embedder Embedder, accountID, ident, version, source string, texts []string) (blocked bool, blockedMessage string, redactedTexts []string, assessments []*bedrockruntime.GuardrailAssessment, err error) {
|
||||
view, err := loadGuardrailView(ctx, store, accountID, ident, version)
|
||||
if err != nil {
|
||||
return false, "", texts, nil, err
|
||||
}
|
||||
|
||||
action, gassessments, outputs, _ := applyGuardrailPolicies(view, texts, source)
|
||||
action, gassessments, outputs, _, err := applyGuardrailPolicies(ctx, embedder, view, texts, source)
|
||||
if err != nil {
|
||||
return false, "", texts, nil, err
|
||||
}
|
||||
if action == bedrockruntime.GuardrailActionGuardrailIntervened {
|
||||
message := view.BlockedInputMessaging
|
||||
if source == bedrockruntime.GuardrailContentSourceOutput {
|
||||
@@ -242,6 +247,7 @@ type guardrailStreamSource struct {
|
||||
inner converseStreamSource
|
||||
|
||||
store *GuardrailStore
|
||||
embedder Embedder
|
||||
accountID string
|
||||
ident string
|
||||
version string
|
||||
@@ -261,8 +267,9 @@ type guardrailStreamSource struct {
|
||||
// newGuardrailStreamSource constructs a guardrailStreamSource. inputAssessments
|
||||
// is the already-computed result of the pre-stream INPUT check, carried
|
||||
// through so the trailing metadata event's trace can report both halves.
|
||||
func newGuardrailStreamSource(inner converseStreamSource, store *GuardrailStore, accountID, ident, version string, trace bool, inputAssessments []*bedrockruntime.GuardrailAssessment) *guardrailStreamSource {
|
||||
return &guardrailStreamSource{inner: inner, store: store, accountID: accountID, ident: ident, version: version, trace: trace, inputAssessments: inputAssessments}
|
||||
// embedder drives topicPolicy's semantic match on the OUTPUT check below.
|
||||
func newGuardrailStreamSource(inner converseStreamSource, store *GuardrailStore, embedder Embedder, accountID, ident, version string, trace bool, inputAssessments []*bedrockruntime.GuardrailAssessment) *guardrailStreamSource {
|
||||
return &guardrailStreamSource{inner: inner, store: store, embedder: embedder, accountID: accountID, ident: ident, version: version, trace: trace, inputAssessments: inputAssessments}
|
||||
}
|
||||
|
||||
var _ converseStreamSource = (*guardrailStreamSource)(nil)
|
||||
@@ -290,7 +297,7 @@ func (g *guardrailStreamSource) assess(ctx context.Context) {
|
||||
return
|
||||
}
|
||||
g.assessed = true
|
||||
blocked, message, redacted, assessments, err := enforceGuardrail(ctx, g.store, g.accountID, g.ident, g.version,
|
||||
blocked, message, redacted, assessments, err := enforceGuardrail(ctx, g.store, g.embedder, g.accountID, g.ident, g.version,
|
||||
bedrockruntime.GuardrailContentSourceOutput, []string{g.text.String()})
|
||||
if err != nil {
|
||||
g.assessErr = err
|
||||
|
||||
@@ -62,7 +62,7 @@ func TestRouter_Converse_GuardrailInputBlock_SkipsBackend(t *testing.T) {
|
||||
|
||||
modelID := "meta.llama3-2-1b-instruct-v1:0"
|
||||
ts, hits := countingVLLMServer(t, "hi")
|
||||
rt := NewRouter(nil, NewStaticEndpointResolver(map[string]string{modelID: ts.URL}), nil, grantAll{}, nil, store)
|
||||
rt := NewRouter(nil, NewStaticEndpointResolver(map[string]string{modelID: ts.URL}), nil, grantAll{}, nil, store, nil)
|
||||
|
||||
out, err := rt.Converse(ctx, grCallerAccount, modelID, guardedConverseInput("this has a badword in it", createOut.GuardrailId, false))
|
||||
require.NoError(t, err)
|
||||
@@ -84,7 +84,7 @@ func TestRouter_Converse_GuardrailOutputBlock(t *testing.T) {
|
||||
|
||||
modelID := "meta.llama3-2-1b-instruct-v1:0"
|
||||
ts, _ := countingVLLMServer(t, "this has a badword in it")
|
||||
rt := NewRouter(nil, NewStaticEndpointResolver(map[string]string{modelID: ts.URL}), nil, grantAll{}, nil, store)
|
||||
rt := NewRouter(nil, NewStaticEndpointResolver(map[string]string{modelID: ts.URL}), nil, grantAll{}, nil, store, nil)
|
||||
|
||||
out, err := rt.Converse(ctx, grCallerAccount, modelID, guardedConverseInput("hello", createOut.GuardrailId, false))
|
||||
require.NoError(t, err)
|
||||
@@ -103,7 +103,7 @@ func TestRouter_Converse_GuardrailOutputAnonymize(t *testing.T) {
|
||||
|
||||
modelID := "meta.llama3-2-1b-instruct-v1:0"
|
||||
ts, _ := countingVLLMServer(t, "contact jane@example.com for support")
|
||||
rt := NewRouter(nil, NewStaticEndpointResolver(map[string]string{modelID: ts.URL}), nil, grantAll{}, nil, store)
|
||||
rt := NewRouter(nil, NewStaticEndpointResolver(map[string]string{modelID: ts.URL}), nil, grantAll{}, nil, store, nil)
|
||||
|
||||
out, err := rt.Converse(ctx, grCallerAccount, modelID, guardedConverseInput("hello", createOut.GuardrailId, false))
|
||||
require.NoError(t, err)
|
||||
@@ -118,7 +118,7 @@ func TestRouter_Converse_GuardrailOutputAnonymize(t *testing.T) {
|
||||
func TestRouter_Converse_NoGuardrailConfig_Regression(t *testing.T) {
|
||||
modelID := "meta.llama3-2-1b-instruct-v1:0"
|
||||
ts, hits := countingVLLMServer(t, "hi there")
|
||||
rt := NewRouter(nil, NewStaticEndpointResolver(map[string]string{modelID: ts.URL}), nil, grantAll{}, nil, nil)
|
||||
rt := NewRouter(nil, NewStaticEndpointResolver(map[string]string{modelID: ts.URL}), nil, grantAll{}, nil, nil, nil)
|
||||
|
||||
out, err := rt.Converse(context.Background(), "000000000001", modelID, converseInput())
|
||||
require.NoError(t, err)
|
||||
@@ -134,7 +134,7 @@ func TestRouter_Converse_UnknownOrForeignGuardrailReturnsResourceNotFound(t *tes
|
||||
ctx := context.Background()
|
||||
modelID := "meta.llama3-2-1b-instruct-v1:0"
|
||||
ts, hits := countingVLLMServer(t, "hi")
|
||||
rt := NewRouter(nil, NewStaticEndpointResolver(map[string]string{modelID: ts.URL}), nil, grantAll{}, nil, store)
|
||||
rt := NewRouter(nil, NewStaticEndpointResolver(map[string]string{modelID: ts.URL}), nil, grantAll{}, nil, store, nil)
|
||||
|
||||
_, err := rt.Converse(ctx, grCallerAccount, modelID, guardedConverseInput("hello", aws.String("does-not-exist"), false))
|
||||
require.Error(t, err)
|
||||
@@ -158,7 +158,7 @@ func TestRouter_Converse_TraceEnabledSurfacesAssessment(t *testing.T) {
|
||||
|
||||
modelID := "meta.llama3-2-1b-instruct-v1:0"
|
||||
ts, _ := countingVLLMServer(t, "contact jane@example.com for support")
|
||||
rt := NewRouter(nil, NewStaticEndpointResolver(map[string]string{modelID: ts.URL}), nil, grantAll{}, nil, store)
|
||||
rt := NewRouter(nil, NewStaticEndpointResolver(map[string]string{modelID: ts.URL}), nil, grantAll{}, nil, store, nil)
|
||||
|
||||
out, err := rt.Converse(ctx, grCallerAccount, modelID, guardedConverseInput("hello", createOut.GuardrailId, true))
|
||||
require.NoError(t, err)
|
||||
@@ -177,7 +177,7 @@ func TestRouter_Converse_TraceDisabledOmitsAssessment(t *testing.T) {
|
||||
|
||||
modelID := "meta.llama3-2-1b-instruct-v1:0"
|
||||
ts, _ := countingVLLMServer(t, "contact jane@example.com for support")
|
||||
rt := NewRouter(nil, NewStaticEndpointResolver(map[string]string{modelID: ts.URL}), nil, grantAll{}, nil, store)
|
||||
rt := NewRouter(nil, NewStaticEndpointResolver(map[string]string{modelID: ts.URL}), nil, grantAll{}, nil, store, nil)
|
||||
|
||||
out, err := rt.Converse(ctx, grCallerAccount, modelID, guardedConverseInput("hello", createOut.GuardrailId, false))
|
||||
require.NoError(t, err)
|
||||
@@ -233,7 +233,7 @@ func TestConverseStream_GuardrailInputBlock_SkipsBackend(t *testing.T) {
|
||||
body, err := json.Marshal(guardrailStreamConverseInput("this has a badword in it", createOut.GuardrailId, false))
|
||||
require.NoError(t, err)
|
||||
|
||||
err = ConverseStream(ctx, rec, grCallerAccount, modelID, body, nil, NewStaticEndpointResolver(map[string]string{modelID: ts.URL}), nil, grantAll{}, nil, store)
|
||||
err = ConverseStream(ctx, rec, grCallerAccount, modelID, body, nil, NewStaticEndpointResolver(map[string]string{modelID: ts.URL}), nil, grantAll{}, nil, store, nil)
|
||||
require.NoError(t, err)
|
||||
assert.Equal(t, int32(0), hits.Load(), "the backend stream must never be started on an INPUT block")
|
||||
|
||||
@@ -260,7 +260,7 @@ func TestConverseStream_UnknownOrForeignGuardrailReturnsResourceNotFoundPreHeade
|
||||
body, err := json.Marshal(guardrailStreamConverseInput("hello", aws.String("does-not-exist"), false))
|
||||
require.NoError(t, err)
|
||||
|
||||
err = ConverseStream(ctx, rec, grCallerAccount, modelID, body, nil, nil, nil, grantAll{}, nil, store)
|
||||
err = ConverseStream(ctx, rec, grCallerAccount, modelID, body, nil, nil, nil, grantAll{}, nil, store, nil)
|
||||
require.Error(t, err)
|
||||
assert.True(t, awserrors.IsErrorCode(err, awserrors.ErrorResourceNotFoundException))
|
||||
// A pre-stream failure must not have written anything.
|
||||
@@ -280,7 +280,7 @@ func TestConverseStream_NoGuardrailConfig_Regression(t *testing.T) {
|
||||
body, err := json.Marshal(converseStreamInput())
|
||||
require.NoError(t, err)
|
||||
|
||||
err = ConverseStream(context.Background(), rec, "000000000001", modelID, body, nil, NewStaticEndpointResolver(map[string]string{modelID: ts.URL}), nil, grantAll{}, nil, nil)
|
||||
err = ConverseStream(context.Background(), rec, "000000000001", modelID, body, nil, NewStaticEndpointResolver(map[string]string{modelID: ts.URL}), nil, grantAll{}, nil, nil, nil)
|
||||
require.NoError(t, err)
|
||||
|
||||
frames := decodeAllFrames(t, rec.Body.Bytes())
|
||||
@@ -325,7 +325,7 @@ func TestGuardrailStreamSource_OutputBlock_BuffersAndReplacesText(t *testing.T)
|
||||
require.NoError(t, err)
|
||||
|
||||
inner := &fakeConverseStreamSource{events: guardrailStreamFixtureEvents("this has a ", "badword in it")}
|
||||
src := newGuardrailStreamSource(inner, store, grCallerAccount, aws.StringValue(createOut.GuardrailId), guardrailDraftVersion, false, nil)
|
||||
src := newGuardrailStreamSource(inner, store, nil, grCallerAccount, aws.StringValue(createOut.GuardrailId), guardrailDraftVersion, false, nil)
|
||||
|
||||
events := drainConverseStream(t, src)
|
||||
// The two raw deltas collapse into exactly one guarded delta.
|
||||
@@ -353,7 +353,7 @@ func TestGuardrailStreamSource_OutputAnonymize_RedactsAccumulatedText(t *testing
|
||||
require.NoError(t, err)
|
||||
|
||||
inner := &fakeConverseStreamSource{events: guardrailStreamFixtureEvents("contact jane@", "example.com for support")}
|
||||
src := newGuardrailStreamSource(inner, store, grCallerAccount, aws.StringValue(createOut.GuardrailId), guardrailDraftVersion, false, nil)
|
||||
src := newGuardrailStreamSource(inner, store, nil, grCallerAccount, aws.StringValue(createOut.GuardrailId), guardrailDraftVersion, false, nil)
|
||||
|
||||
events := drainConverseStream(t, src)
|
||||
require.Len(t, events, 5)
|
||||
@@ -370,7 +370,7 @@ func TestGuardrailStreamSource_TraceEnabledSurfacesAssessment(t *testing.T) {
|
||||
|
||||
inputAssessments := []*bedrockruntime.GuardrailAssessment{{}}
|
||||
inner := &fakeConverseStreamSource{events: guardrailStreamFixtureEvents("contact jane@example.com for support")}
|
||||
src := newGuardrailStreamSource(inner, store, grCallerAccount, aws.StringValue(createOut.GuardrailId), guardrailDraftVersion, true, inputAssessments)
|
||||
src := newGuardrailStreamSource(inner, store, nil, grCallerAccount, aws.StringValue(createOut.GuardrailId), guardrailDraftVersion, true, inputAssessments)
|
||||
|
||||
events := drainConverseStream(t, src)
|
||||
metadata := events[len(events)-1]
|
||||
@@ -408,7 +408,7 @@ func TestInvokeRouter_GuardrailInputBlock_Llama_SkipsBackend(t *testing.T) {
|
||||
|
||||
modelID := "meta.llama3-2-1b-instruct-v1:0"
|
||||
ts, hits := countingLlamaCompletionsServer(t, "hi")
|
||||
rt := NewInvokeRouter(nil, NewStaticEndpointResolver(map[string]string{modelID: ts.URL}), nil, grantAll{}, nil, store)
|
||||
rt := NewInvokeRouter(nil, NewStaticEndpointResolver(map[string]string{modelID: ts.URL}), nil, grantAll{}, nil, store, nil)
|
||||
|
||||
respBody, contentType, err := rt.InvokeModel(ctx, grCallerAccount, modelID, []byte(`{"prompt":"this has a badword in it"}`), aws.StringValue(createOut.GuardrailId), guardrailDraftVersion)
|
||||
require.NoError(t, err)
|
||||
@@ -432,7 +432,7 @@ func TestInvokeRouter_GuardrailInputBlock_Anthropic_SkipsBackend(t *testing.T) {
|
||||
require.NoError(t, err)
|
||||
|
||||
modelID := "anthropic.claude-3-5-sonnet-20240620-v1:0"
|
||||
rt := NewInvokeRouter(stubCredentialResolver{key: "sk-test", ok: true}, nil, nil, grantAll{}, nil, store)
|
||||
rt := NewInvokeRouter(stubCredentialResolver{key: "sk-test", ok: true}, nil, nil, grantAll{}, nil, store, nil)
|
||||
|
||||
body := []byte(`{"anthropic_version":"bedrock-2023-05-31","max_tokens":100,"messages":[{"role":"user","content":[{"type":"text","text":"this has a badword in it"}]}]}`)
|
||||
respBody, contentType, err := rt.InvokeModel(ctx, grCallerAccount, modelID, body, aws.StringValue(createOut.GuardrailId), guardrailDraftVersion)
|
||||
@@ -452,7 +452,7 @@ func TestInvokeRouter_GuardrailOutputBlock_Llama(t *testing.T) {
|
||||
|
||||
modelID := "meta.llama3-2-1b-instruct-v1:0"
|
||||
ts, _ := countingLlamaCompletionsServer(t, "this has a badword in it")
|
||||
rt := NewInvokeRouter(nil, NewStaticEndpointResolver(map[string]string{modelID: ts.URL}), nil, grantAll{}, nil, store)
|
||||
rt := NewInvokeRouter(nil, NewStaticEndpointResolver(map[string]string{modelID: ts.URL}), nil, grantAll{}, nil, store, nil)
|
||||
|
||||
respBody, _, err := rt.InvokeModel(ctx, grCallerAccount, modelID, []byte(`{"prompt":"hello"}`), aws.StringValue(createOut.GuardrailId), guardrailDraftVersion)
|
||||
require.NoError(t, err)
|
||||
@@ -471,7 +471,7 @@ func TestInvokeRouter_GuardrailOutputAnonymize_Llama(t *testing.T) {
|
||||
|
||||
modelID := "meta.llama3-2-1b-instruct-v1:0"
|
||||
ts, _ := countingLlamaCompletionsServer(t, "contact jane@example.com for support")
|
||||
rt := NewInvokeRouter(nil, NewStaticEndpointResolver(map[string]string{modelID: ts.URL}), nil, grantAll{}, nil, store)
|
||||
rt := NewInvokeRouter(nil, NewStaticEndpointResolver(map[string]string{modelID: ts.URL}), nil, grantAll{}, nil, store, nil)
|
||||
|
||||
respBody, _, err := rt.InvokeModel(ctx, grCallerAccount, modelID, []byte(`{"prompt":"hello"}`), aws.StringValue(createOut.GuardrailId), guardrailDraftVersion)
|
||||
require.NoError(t, err)
|
||||
@@ -486,7 +486,7 @@ func TestInvokeRouter_GuardrailOutputAnonymize_Llama(t *testing.T) {
|
||||
func TestInvokeRouter_NoGuardrailHeader_Regression(t *testing.T) {
|
||||
modelID := "meta.llama3-2-1b-instruct-v1:0"
|
||||
ts, hits := countingLlamaCompletionsServer(t, "hi there")
|
||||
rt := NewInvokeRouter(nil, NewStaticEndpointResolver(map[string]string{modelID: ts.URL}), nil, grantAll{}, nil, nil)
|
||||
rt := NewInvokeRouter(nil, NewStaticEndpointResolver(map[string]string{modelID: ts.URL}), nil, grantAll{}, nil, nil, nil)
|
||||
|
||||
respBody, _, err := rt.InvokeModel(context.Background(), "000000000001", modelID, []byte(`{"prompt":"hello"}`), "", "")
|
||||
require.NoError(t, err)
|
||||
@@ -502,7 +502,7 @@ func TestInvokeRouter_UnknownOrForeignGuardrailReturnsResourceNotFound(t *testin
|
||||
ctx := context.Background()
|
||||
modelID := "meta.llama3-2-1b-instruct-v1:0"
|
||||
ts, hits := countingLlamaCompletionsServer(t, "hi")
|
||||
rt := NewInvokeRouter(nil, NewStaticEndpointResolver(map[string]string{modelID: ts.URL}), nil, grantAll{}, nil, store)
|
||||
rt := NewInvokeRouter(nil, NewStaticEndpointResolver(map[string]string{modelID: ts.URL}), nil, grantAll{}, nil, store, nil)
|
||||
|
||||
_, _, err := rt.InvokeModel(ctx, grCallerAccount, modelID, []byte(`{"prompt":"hello"}`), "does-not-exist", guardrailDraftVersion)
|
||||
require.Error(t, err)
|
||||
@@ -562,7 +562,7 @@ func TestInvokeStreamRouter_GuardrailInputBlock_SkipsBackend(t *testing.T) {
|
||||
}))
|
||||
defer ts.Close()
|
||||
|
||||
rt := NewInvokeStreamRouter(nil, NewStaticEndpointResolver(map[string]string{modelID: ts.URL}), grantAll{}, nil, store)
|
||||
rt := NewInvokeStreamRouter(nil, NewStaticEndpointResolver(map[string]string{modelID: ts.URL}), grantAll{}, nil, store, nil)
|
||||
|
||||
src, err := rt.InvokeModelWithResponseStream(ctx, grCallerAccount, modelID, []byte(`{"prompt":"this has a badword in it"}`), aws.StringValue(createOut.GuardrailId), guardrailDraftVersion)
|
||||
require.NoError(t, err)
|
||||
@@ -593,7 +593,7 @@ func TestInvokeStreamRouter_GuardrailOutputBlock_Buffered(t *testing.T) {
|
||||
}))
|
||||
defer ts.Close()
|
||||
|
||||
rt := NewInvokeStreamRouter(nil, NewStaticEndpointResolver(map[string]string{modelID: ts.URL}), grantAll{}, nil, store)
|
||||
rt := NewInvokeStreamRouter(nil, NewStaticEndpointResolver(map[string]string{modelID: ts.URL}), grantAll{}, nil, store, nil)
|
||||
|
||||
src, err := rt.InvokeModelWithResponseStream(ctx, grCallerAccount, modelID, []byte(`{"prompt":"hello"}`), aws.StringValue(createOut.GuardrailId), guardrailDraftVersion)
|
||||
require.NoError(t, err)
|
||||
@@ -623,7 +623,7 @@ func TestInvokeStreamRouter_GuardrailOutputAnonymize_Buffered(t *testing.T) {
|
||||
}))
|
||||
defer ts.Close()
|
||||
|
||||
rt := NewInvokeStreamRouter(nil, NewStaticEndpointResolver(map[string]string{modelID: ts.URL}), grantAll{}, nil, store)
|
||||
rt := NewInvokeStreamRouter(nil, NewStaticEndpointResolver(map[string]string{modelID: ts.URL}), grantAll{}, nil, store, nil)
|
||||
|
||||
src, err := rt.InvokeModelWithResponseStream(ctx, grCallerAccount, modelID, []byte(`{"prompt":"hello"}`), aws.StringValue(createOut.GuardrailId), guardrailDraftVersion)
|
||||
require.NoError(t, err)
|
||||
@@ -649,7 +649,7 @@ func TestInvokeStreamRouter_NoGuardrailHeader_Regression(t *testing.T) {
|
||||
defer ts.Close()
|
||||
|
||||
modelID := "meta.llama3-2-1b-instruct-v1:0"
|
||||
rt := NewInvokeStreamRouter(nil, NewStaticEndpointResolver(map[string]string{modelID: ts.URL}), grantAll{}, nil, nil)
|
||||
rt := NewInvokeStreamRouter(nil, NewStaticEndpointResolver(map[string]string{modelID: ts.URL}), grantAll{}, nil, nil, nil)
|
||||
|
||||
src, err := rt.InvokeModelWithResponseStream(context.Background(), "000000000001", modelID, []byte(`{"prompt":"hello"}`), "", "")
|
||||
require.NoError(t, err)
|
||||
@@ -713,7 +713,7 @@ func TestGuardrailStreamSource_TraceDisabledOmitsAssessment(t *testing.T) {
|
||||
require.NoError(t, err)
|
||||
|
||||
inner := &fakeConverseStreamSource{events: guardrailStreamFixtureEvents("contact jane@example.com for support")}
|
||||
src := newGuardrailStreamSource(inner, store, grCallerAccount, aws.StringValue(createOut.GuardrailId), guardrailDraftVersion, false, nil)
|
||||
src := newGuardrailStreamSource(inner, store, nil, grCallerAccount, aws.StringValue(createOut.GuardrailId), guardrailDraftVersion, false, nil)
|
||||
|
||||
events := drainConverseStream(t, src)
|
||||
metadata := events[len(events)-1]
|
||||
@@ -805,7 +805,7 @@ func TestInvokeRouter_GuardrailOutputPath_DeletedMidRequest_FailsClosed(t *testi
|
||||
}))
|
||||
defer ts.Close()
|
||||
|
||||
rt := NewInvokeRouter(nil, NewStaticEndpointResolver(map[string]string{modelID: ts.URL}), nil, grantAll{}, nil, store)
|
||||
rt := NewInvokeRouter(nil, NewStaticEndpointResolver(map[string]string{modelID: ts.URL}), nil, grantAll{}, nil, store, nil)
|
||||
|
||||
respBody, _, err := rt.InvokeModel(ctx, grCallerAccount, modelID, []byte(`{"prompt":"hello"}`), aws.StringValue(createOut.GuardrailId), guardrailDraftVersion)
|
||||
require.Error(t, err)
|
||||
@@ -828,7 +828,7 @@ func TestGuardrailInvokeStreamSource_OutputBlock_UndecodableChunk_Llama(t *testi
|
||||
[]byte(`{"generation":"partial safe text "}`),
|
||||
[]byte(`{"generation": 123}`),
|
||||
}}
|
||||
g := newGuardrailInvokeStreamSource(inner, store, grCallerAccount, aws.StringValue(createOut.GuardrailId), guardrailDraftVersion, tierSelfHost)
|
||||
g := newGuardrailInvokeStreamSource(inner, store, nil, grCallerAccount, aws.StringValue(createOut.GuardrailId), guardrailDraftVersion, tierSelfHost)
|
||||
|
||||
chunks := drainInvokeStream(t, g)
|
||||
require.Len(t, chunks, 2, "the guarded blocked sequence replaces the raw buffer wholesale")
|
||||
@@ -863,7 +863,7 @@ func TestGuardrailInvokeStreamSource_OutputBlock_UndecodableChunk_Anthropic(t *t
|
||||
[]byte(`{"type":"content_block_delta","index":0,"delta":{"type":"text_delta","text":"partial safe text "}}`),
|
||||
[]byte(`{"type":"content_block_delta","index":0,"delta":{"type":"text_delta","text":123}}`),
|
||||
}}
|
||||
g := newGuardrailInvokeStreamSource(inner, store, grCallerAccount, aws.StringValue(createOut.GuardrailId), guardrailDraftVersion, backend)
|
||||
g := newGuardrailInvokeStreamSource(inner, store, nil, grCallerAccount, aws.StringValue(createOut.GuardrailId), guardrailDraftVersion, backend)
|
||||
|
||||
chunks := drainInvokeStream(t, g)
|
||||
require.Len(t, chunks, 5, "anthropicGuardedStreamChunks' minimal event sequence")
|
||||
@@ -904,7 +904,7 @@ func TestGuardrailInvokeStreamSource_EmptyCompletion_PassesThroughUnblocked_Regr
|
||||
[]byte(`{"generation":"","prompt_token_count":1,"generation_token_count":0,"stop_reason":"stop"}`),
|
||||
}
|
||||
inner := &fakeInvokeStreamSource{chunks: append([][]byte{}, rawChunks...)}
|
||||
g := newGuardrailInvokeStreamSource(inner, store, grCallerAccount, aws.StringValue(createOut.GuardrailId), guardrailDraftVersion, tierSelfHost)
|
||||
g := newGuardrailInvokeStreamSource(inner, store, nil, grCallerAccount, aws.StringValue(createOut.GuardrailId), guardrailDraftVersion, tierSelfHost)
|
||||
|
||||
chunks := drainInvokeStream(t, g)
|
||||
require.Len(t, chunks, len(rawChunks))
|
||||
|
||||
@@ -402,6 +402,7 @@ type guardrailInvokeStreamSource struct {
|
||||
backend string
|
||||
|
||||
store *GuardrailStore
|
||||
embedder Embedder
|
||||
accountID string
|
||||
ident string
|
||||
version string
|
||||
@@ -413,8 +414,9 @@ type guardrailInvokeStreamSource struct {
|
||||
}
|
||||
|
||||
// newGuardrailInvokeStreamSource constructs a guardrailInvokeStreamSource.
|
||||
func newGuardrailInvokeStreamSource(inner invokeStreamSource, store *GuardrailStore, accountID, ident, version, backend string) *guardrailInvokeStreamSource {
|
||||
return &guardrailInvokeStreamSource{inner: inner, store: store, accountID: accountID, ident: ident, version: version, backend: backend}
|
||||
// embedder drives topicPolicy's semantic match on the OUTPUT check below.
|
||||
func newGuardrailInvokeStreamSource(inner invokeStreamSource, store *GuardrailStore, embedder Embedder, accountID, ident, version, backend string) *guardrailInvokeStreamSource {
|
||||
return &guardrailInvokeStreamSource{inner: inner, store: store, embedder: embedder, accountID: accountID, ident: ident, version: version, backend: backend}
|
||||
}
|
||||
|
||||
var _ invokeStreamSource = (*guardrailInvokeStreamSource)(nil)
|
||||
@@ -482,7 +484,7 @@ func (g *guardrailInvokeStreamSource) Next(ctx context.Context) ([]byte, bool, e
|
||||
}
|
||||
|
||||
original := g.text.String()
|
||||
blocked, message, redacted, _, err := enforceGuardrail(ctx, g.store, g.accountID, g.ident, g.version,
|
||||
blocked, message, redacted, _, err := enforceGuardrail(ctx, g.store, g.embedder, g.accountID, g.ident, g.version,
|
||||
bedrockruntime.GuardrailContentSourceOutput, []string{original})
|
||||
if err != nil {
|
||||
return nil, false, err
|
||||
|
||||
@@ -0,0 +1,163 @@
|
||||
package gateway_bedrock
|
||||
|
||||
import (
|
||||
"context"
|
||||
"crypto/sha256"
|
||||
"encoding/hex"
|
||||
"log/slog"
|
||||
"math"
|
||||
"strings"
|
||||
"sync"
|
||||
|
||||
"github.com/aws/aws-sdk-go/aws"
|
||||
"github.com/aws/aws-sdk-go/service/bedrock"
|
||||
)
|
||||
|
||||
// defaultTopicSimilarityThreshold is the cosine similarity a topic's
|
||||
// embedded phrases must reach against an input text's embedding to block.
|
||||
// Conservative starting value calibrated for nomic-embed-text-v1.5 -- retune
|
||||
// here if live traffic shows the false positive/negative rate drifting.
|
||||
const defaultTopicSimilarityThreshold = 0.42
|
||||
|
||||
// topicSimilarityThreshold returns topic's similarity threshold. AWS's
|
||||
// GuardrailTopicConfig wire shape has no field to override this per topic,
|
||||
// so every topic currently gets the same default; this indirection is the
|
||||
// seam a future per-topic override plugs into.
|
||||
func topicSimilarityThreshold(_ *bedrock.GuardrailTopicConfig) float64 {
|
||||
return defaultTopicSimilarityThreshold
|
||||
}
|
||||
|
||||
// topicVectorCache holds each denied topic's per-phrase embedding vectors,
|
||||
// keyed by a hash of the embedding model id plus the topic's own phrase set.
|
||||
// An unchanged topic -- across applies, or shared by the DRAFT and a
|
||||
// snapshot taken from it -- is embedded at most once per process lifetime; a
|
||||
// mutated DRAFT invalidates itself automatically via a changed hash.
|
||||
var topicVectorCache sync.Map // string -> [][]float32
|
||||
|
||||
// topicPhrases returns the phrase set embedded for topic: its Name,
|
||||
// Definition, and every Example, trimmed and deduplicated so an identical
|
||||
// phrase is never embedded twice.
|
||||
func topicPhrases(topic *bedrock.GuardrailTopicConfig) []string {
|
||||
seen := make(map[string]bool)
|
||||
var phrases []string
|
||||
add := func(s string) {
|
||||
s = strings.TrimSpace(s)
|
||||
if s == "" || seen[s] {
|
||||
return
|
||||
}
|
||||
seen[s] = true
|
||||
phrases = append(phrases, s)
|
||||
}
|
||||
add(aws.StringValue(topic.Name))
|
||||
add(aws.StringValue(topic.Definition))
|
||||
for _, ex := range topic.Examples {
|
||||
add(aws.StringValue(ex))
|
||||
}
|
||||
return phrases
|
||||
}
|
||||
|
||||
// topicCacheKey hashes modelID together with phrases, so vectors for a
|
||||
// different embedding model, or a topic whose phrase set changed, never
|
||||
// collide in topicVectorCache.
|
||||
func topicCacheKey(modelID string, phrases []string) string {
|
||||
h := sha256.New()
|
||||
h.Write([]byte(modelID))
|
||||
for _, p := range phrases {
|
||||
h.Write([]byte{0})
|
||||
h.Write([]byte(p))
|
||||
}
|
||||
return hex.EncodeToString(h.Sum(nil))
|
||||
}
|
||||
|
||||
// topicVectors returns topic's cached per-phrase embedding vectors,
|
||||
// embedding and caching them on first use. A nil, nil return means embedder
|
||||
// is nil or topic has no phrases -- the caller relies on the literal matcher
|
||||
// alone, not an error. A non-nil error means embedder was available but the
|
||||
// embed call itself failed: the caller must fail closed, not fall back.
|
||||
func topicVectors(ctx context.Context, embedder Embedder, modelID string, topic *bedrock.GuardrailTopicConfig) ([][]float32, error) {
|
||||
if embedder == nil || topic == nil {
|
||||
return nil, nil
|
||||
}
|
||||
phrases := topicPhrases(topic)
|
||||
if len(phrases) == 0 {
|
||||
return nil, nil
|
||||
}
|
||||
|
||||
key := topicCacheKey(modelID, phrases)
|
||||
if cached, ok := topicVectorCache.Load(key); ok {
|
||||
vectors, _ := cached.([][]float32)
|
||||
return vectors, nil
|
||||
}
|
||||
|
||||
vectors, err := embedder.Embed(ctx, modelID, phrases)
|
||||
if err != nil {
|
||||
slog.Warn("guardrail: topic embedding unavailable, failing closed on unverified content",
|
||||
"topic", aws.StringValue(topic.Name), "model", modelID, "err", err)
|
||||
return nil, err
|
||||
}
|
||||
topicVectorCache.Store(key, vectors)
|
||||
return vectors, nil
|
||||
}
|
||||
|
||||
// cosineSimilarity returns the cosine similarity of a and b. A length
|
||||
// mismatch or a zero-magnitude vector returns 0 rather than NaN, so a
|
||||
// malformed embedding never compares true against a positive threshold.
|
||||
func cosineSimilarity(a, b []float32) float64 {
|
||||
if len(a) != len(b) || len(a) == 0 {
|
||||
return 0
|
||||
}
|
||||
var dot, magA, magB float64
|
||||
for i := range a {
|
||||
dot += float64(a[i]) * float64(b[i])
|
||||
magA += float64(a[i]) * float64(a[i])
|
||||
magB += float64(b[i]) * float64(b[i])
|
||||
}
|
||||
if magA == 0 || magB == 0 {
|
||||
return 0
|
||||
}
|
||||
return dot / (math.Sqrt(magA) * math.Sqrt(magB))
|
||||
}
|
||||
|
||||
// maxCosineSimilarity returns the highest cosine similarity between vector
|
||||
// and any vector in candidates, or 0 for an empty candidate set.
|
||||
func maxCosineSimilarity(vector []float32, candidates [][]float32) float64 {
|
||||
best := 0.0
|
||||
for _, c := range candidates {
|
||||
if s := cosineSimilarity(vector, c); s > best {
|
||||
best = s
|
||||
}
|
||||
}
|
||||
return best
|
||||
}
|
||||
|
||||
// embedGuardrailTexts embeds texts once against modelID so every topic in a
|
||||
// policy reuses the same input vectors instead of re-embedding per topic. A
|
||||
// nil, nil return means embedder is nil -- the caller relies on the literal
|
||||
// matcher alone, which is not an error, just an unconfigured deployment. A
|
||||
// non-nil error means embedder was available but the call failed: the
|
||||
// caller must fail closed rather than silently pass unverified content.
|
||||
func embedGuardrailTexts(ctx context.Context, embedder Embedder, modelID string, texts []string) ([][]float32, error) {
|
||||
if embedder == nil {
|
||||
slog.Warn("guardrail: no embedder configured, falling back to literal topic matching")
|
||||
return nil, nil
|
||||
}
|
||||
vectors, err := embedder.Embed(ctx, modelID, texts)
|
||||
if err != nil {
|
||||
slog.Warn("guardrail: input embedding unavailable, failing closed on unverified content", "model", modelID, "err", err)
|
||||
return nil, err
|
||||
}
|
||||
return vectors, nil
|
||||
}
|
||||
|
||||
// topicSemanticHit reports whether any textVector reaches threshold cosine
|
||||
// similarity against any of topicPhraseVectors -- max-cosine per phrase,
|
||||
// not a single centroid, so one strong phrase match is enough even when the
|
||||
// topic's examples are otherwise diverse.
|
||||
func topicSemanticHit(textVectors, topicPhraseVectors [][]float32, threshold float64) bool {
|
||||
for _, tv := range textVectors {
|
||||
if maxCosineSimilarity(tv, topicPhraseVectors) >= threshold {
|
||||
return true
|
||||
}
|
||||
}
|
||||
return false
|
||||
}
|
||||
@@ -0,0 +1,201 @@
|
||||
package gateway_bedrock
|
||||
|
||||
//test:in-package: exercises assessTopicPolicy/applyGuardrailPolicies (the
|
||||
// unexported filter engine) directly, matching guardrail_filter_test.go's
|
||||
// existing in-package convention for this same engine in this package.
|
||||
|
||||
import (
|
||||
"context"
|
||||
"errors"
|
||||
"testing"
|
||||
|
||||
"github.com/aws/aws-sdk-go/aws"
|
||||
"github.com/aws/aws-sdk-go/service/bedrock"
|
||||
"github.com/aws/aws-sdk-go/service/bedrockruntime"
|
||||
"github.com/mulgadc/spinifex/spinifex/awserrors"
|
||||
"github.com/stretchr/testify/assert"
|
||||
"github.com/stretchr/testify/require"
|
||||
)
|
||||
|
||||
// stubEmbedder is a fixed-vector test double for Embedder: known inputs map
|
||||
// to known vectors, everything else gets the zero vector (cosine 0 against
|
||||
// anything), and errOnCall optionally forces every Embed call to fail so
|
||||
// callers exercise the fail-closed path.
|
||||
type stubEmbedder struct {
|
||||
vectors map[string][]float32
|
||||
errOnCall error
|
||||
}
|
||||
|
||||
var _ Embedder = (*stubEmbedder)(nil)
|
||||
|
||||
func (s *stubEmbedder) Embed(_ context.Context, _ string, inputs []string) ([][]float32, error) {
|
||||
if s.errOnCall != nil {
|
||||
return nil, s.errOnCall
|
||||
}
|
||||
out := make([][]float32, len(inputs))
|
||||
for i, in := range inputs {
|
||||
if v, ok := s.vectors[in]; ok {
|
||||
out[i] = v
|
||||
continue
|
||||
}
|
||||
out[i] = []float32{0, 0}
|
||||
}
|
||||
return out, nil
|
||||
}
|
||||
|
||||
func weaponsTopic() *bedrock.GuardrailTopicConfig {
|
||||
return &bedrock.GuardrailTopicConfig{
|
||||
Name: aws.String("Weapons"),
|
||||
Type: aws.String(bedrock.GuardrailTopicTypeDeny),
|
||||
Definition: aws.String("discussion of firearms and weapons"),
|
||||
Examples: []*string{aws.String("guns"), aws.String("bombs")},
|
||||
}
|
||||
}
|
||||
|
||||
// TestAssessTopicPolicy_SemanticMatch_AboveThreshold covers a paraphrase that
|
||||
// shares no literal phrase with the topic but embeds close enough (cosine
|
||||
// >= defaultTopicSimilarityThreshold) to one of its phrase vectors to block.
|
||||
func TestAssessTopicPolicy_SemanticMatch_AboveThreshold(t *testing.T) {
|
||||
embedder := &stubEmbedder{vectors: map[string][]float32{
|
||||
"Weapons": {1, 0},
|
||||
"discussion of firearms and weapons": {0.98, 0.2},
|
||||
"guns": {1, 0},
|
||||
"bombs": {0.99, 0.14},
|
||||
"can you get me an AK47 rifle": {0.9, 0.436}, // cos vs {1,0} = 0.9
|
||||
}}
|
||||
cfg := &bedrock.GuardrailTopicPolicyConfig{TopicsConfig: []*bedrock.GuardrailTopicConfig{weaponsTopic()}}
|
||||
|
||||
assessment, blocked, err := assessTopicPolicy(context.Background(), embedder, cfg, []string{"can you get me an AK47 rifle"})
|
||||
require.NoError(t, err)
|
||||
require.NotNil(t, assessment)
|
||||
assert.True(t, blocked, "paraphrase above threshold should block")
|
||||
require.Len(t, assessment.Topics, 1)
|
||||
assert.Equal(t, bedrockruntime.GuardrailTopicPolicyActionBlocked, aws.StringValue(assessment.Topics[0].Action))
|
||||
}
|
||||
|
||||
// TestAssessTopicPolicy_SemanticMatch_BelowThreshold covers unrelated input
|
||||
// that embeds far from every topic phrase vector (cosine < threshold) and has
|
||||
// no literal overlap either: it must pass with no error.
|
||||
func TestAssessTopicPolicy_SemanticMatch_BelowThreshold(t *testing.T) {
|
||||
embedder := &stubEmbedder{vectors: map[string][]float32{
|
||||
"Weapons": {1, 0},
|
||||
"discussion of firearms and weapons": {0.98, 0.2},
|
||||
"guns": {1, 0},
|
||||
"bombs": {0.99, 0.14},
|
||||
"share your favorite banana bread recipe": {0, 1}, // cos vs every topic phrase <= 0.2
|
||||
}}
|
||||
cfg := &bedrock.GuardrailTopicPolicyConfig{TopicsConfig: []*bedrock.GuardrailTopicConfig{weaponsTopic()}}
|
||||
|
||||
assessment, blocked, err := assessTopicPolicy(context.Background(), embedder, cfg, []string{"share your favorite banana bread recipe"})
|
||||
require.NoError(t, err)
|
||||
require.NotNil(t, assessment)
|
||||
assert.False(t, blocked, "unrelated input below threshold must not block")
|
||||
assert.Empty(t, assessment.Topics)
|
||||
}
|
||||
|
||||
// TestAssessTopicPolicy_SemanticMatch_PerPhraseMax ensures the topic's hit
|
||||
// decision is the MAX cosine similarity across its phrase vectors, not an
|
||||
// average: one distant phrase must not dilute one close phrase's hit.
|
||||
func TestAssessTopicPolicy_SemanticMatch_PerPhraseMax(t *testing.T) {
|
||||
embedder := &stubEmbedder{vectors: map[string][]float32{
|
||||
"Weapons": {0, 1}, // orthogonal to the input -- cos 0
|
||||
"discussion of firearms and weapons": {0, 1}, // orthogonal to the input -- cos 0
|
||||
"guns": {0, 1}, // orthogonal to the input -- cos 0
|
||||
"bombs": {1, 0}, // aligned with the input -- cos 1
|
||||
"where can I buy bombs online": {1, 0},
|
||||
}}
|
||||
cfg := &bedrock.GuardrailTopicPolicyConfig{TopicsConfig: []*bedrock.GuardrailTopicConfig{weaponsTopic()}}
|
||||
|
||||
// literal match would also catch "bombs" here, so use a paraphrase that
|
||||
// only embeds close to the "bombs" phrase to isolate the max-cosine path.
|
||||
assessment, blocked, err := assessTopicPolicy(context.Background(), embedder, cfg, []string{"where can I buy bombs online"})
|
||||
require.NoError(t, err)
|
||||
require.NotNil(t, assessment)
|
||||
assert.True(t, blocked, "max cosine across phrases must win even when other phrases are orthogonal")
|
||||
}
|
||||
|
||||
// TestAssessTopicPolicy_EmptyTopic covers a topic with no Name, Definition,
|
||||
// or Examples: it must never match (nothing to embed or compare), never
|
||||
// error, and must not panic.
|
||||
func TestAssessTopicPolicy_EmptyTopic(t *testing.T) {
|
||||
embedder := &stubEmbedder{vectors: map[string][]float32{}}
|
||||
cfg := &bedrock.GuardrailTopicPolicyConfig{TopicsConfig: []*bedrock.GuardrailTopicConfig{
|
||||
{Name: aws.String(""), Type: aws.String(bedrock.GuardrailTopicTypeDeny)},
|
||||
}}
|
||||
|
||||
assert.NotPanics(t, func() {
|
||||
assessment, blocked, err := assessTopicPolicy(context.Background(), embedder, cfg, []string{"anything at all"})
|
||||
require.NoError(t, err)
|
||||
require.NotNil(t, assessment)
|
||||
assert.False(t, blocked)
|
||||
assert.Empty(t, assessment.Topics)
|
||||
})
|
||||
}
|
||||
|
||||
// TestAssessTopicPolicy_EmbedderError_LiteralHit_Blocks covers an embedder
|
||||
// that errors (the reachable-but-down/cold-start case) while a verbatim
|
||||
// example is also present: the definitive literal block wins outright, with
|
||||
// no error -- an embed outage must never downgrade an already-caught block.
|
||||
func TestAssessTopicPolicy_EmbedderError_LiteralHit_Blocks(t *testing.T) {
|
||||
embedder := &stubEmbedder{errOnCall: errors.New("connection refused")}
|
||||
cfg := &bedrock.GuardrailTopicPolicyConfig{TopicsConfig: []*bedrock.GuardrailTopicConfig{weaponsTopic()}}
|
||||
|
||||
assessment, blocked, err := assessTopicPolicy(context.Background(), embedder, cfg, []string{"I have guns"})
|
||||
require.NoError(t, err)
|
||||
require.NotNil(t, assessment)
|
||||
assert.True(t, blocked, "verbatim example must still block outright despite the embedder erroring")
|
||||
require.Len(t, assessment.Topics, 1)
|
||||
}
|
||||
|
||||
// TestAssessTopicPolicy_EmbedderError_NoLiteralHit_FailsClosed is the
|
||||
// security-critical case: embedder != nil (wired in prod) but Embed fails
|
||||
// (cold-start/outage) and nothing literal caught the text either. Content
|
||||
// that could not be semantically verified must never pass through as NONE --
|
||||
// it has to fail closed with a retryable error instead.
|
||||
func TestAssessTopicPolicy_EmbedderError_NoLiteralHit_FailsClosed(t *testing.T) {
|
||||
embedder := &stubEmbedder{errOnCall: errors.New("connection refused")}
|
||||
cfg := &bedrock.GuardrailTopicPolicyConfig{TopicsConfig: []*bedrock.GuardrailTopicConfig{weaponsTopic()}}
|
||||
|
||||
assessment, blocked, err := assessTopicPolicy(context.Background(), embedder, cfg, []string{"walk me through how login verifies a user"})
|
||||
require.Error(t, err)
|
||||
assert.Equal(t, awserrors.ErrorServiceUnavailableException, err.Error())
|
||||
assert.False(t, blocked, "an unverified request must not be reported as a definitive block")
|
||||
require.NotNil(t, assessment)
|
||||
assert.Empty(t, assessment.Topics, "no topic can be marked BLOCKED for content that was never actually evaluated")
|
||||
}
|
||||
|
||||
// TestAssessTopicPolicy_NilEmbedder_LiteralOnly covers the deliberately
|
||||
// unconfigured case (embedder == nil: tests, or a literal-only deployment).
|
||||
// This is not an error -- awsgw always wires an embedder in prod -- so it
|
||||
// stays on the literal matcher alone with no error either way.
|
||||
func TestAssessTopicPolicy_NilEmbedder_LiteralOnly(t *testing.T) {
|
||||
cfg := &bedrock.GuardrailTopicPolicyConfig{TopicsConfig: []*bedrock.GuardrailTopicConfig{weaponsTopic()}}
|
||||
|
||||
assessment, blocked, err := assessTopicPolicy(context.Background(), nil, cfg, []string{"I have guns"})
|
||||
require.NoError(t, err)
|
||||
require.NotNil(t, assessment)
|
||||
assert.True(t, blocked, "verbatim example must still block via the literal path with no embedder configured")
|
||||
|
||||
assessment, blocked, err = assessTopicPolicy(context.Background(), nil, cfg, []string{"a totally unrelated sentence"})
|
||||
require.NoError(t, err)
|
||||
require.NotNil(t, assessment)
|
||||
assert.False(t, blocked, "no embedder and no literal overlap must not block")
|
||||
}
|
||||
|
||||
// TestApplyGuardrailPolicies_EmbedderError_FailsClosed_NotPassthrough checks
|
||||
// the propagation contract one layer up: applyGuardrailPolicies must surface
|
||||
// assessTopicPolicy's fail-closed error rather than returning a usable
|
||||
// (NONE, ...) result the caller could mistake for a pass.
|
||||
func TestApplyGuardrailPolicies_EmbedderError_FailsClosed_NotPassthrough(t *testing.T) {
|
||||
embedder := &stubEmbedder{errOnCall: errors.New("connection refused")}
|
||||
view := guardrailView{TopicPolicy: &bedrock.GuardrailTopicPolicyConfig{TopicsConfig: []*bedrock.GuardrailTopicConfig{weaponsTopic()}}}
|
||||
|
||||
action, assessments, outputs, usage, err := applyGuardrailPolicies(context.Background(), embedder,
|
||||
view, []string{"walk me through how login verifies a user"}, bedrockruntime.GuardrailContentSourceInput)
|
||||
require.Error(t, err)
|
||||
assert.Equal(t, awserrors.ErrorServiceUnavailableException, err.Error())
|
||||
assert.Empty(t, action, "an error result must not carry a usable action a caller could treat as NONE")
|
||||
assert.Nil(t, assessments)
|
||||
assert.Nil(t, outputs)
|
||||
assert.Nil(t, usage)
|
||||
}
|
||||
@@ -59,13 +59,17 @@ type InvokeRouter struct {
|
||||
// fails closed on any guardrail identifier (ResourceNotFoundException)
|
||||
// rather than proceeding unguarded; a request with none is unaffected.
|
||||
guardrails *GuardrailStore
|
||||
// embedder drives topicPolicy's semantic match. Nil falls back to the
|
||||
// guardrail engine's literal matcher (see assessTopicPolicy).
|
||||
embedder Embedder
|
||||
}
|
||||
|
||||
// NewInvokeRouter constructs an InvokeRouter. A nil resolver, endpointResolver,
|
||||
// or recorder falls back to a no-op implementation, and a nil access falls back
|
||||
// to denying every model, so an InvokeRouter is always safe to use even before
|
||||
// the real stores are wired in. A nil provisioned disables PT ARN acceptance.
|
||||
func NewInvokeRouter(resolver CredentialResolver, endpointResolver EndpointResolver, recorder Recorder, access AccessResolver, provisioned *ProvisionedStore, guardrails *GuardrailStore) *InvokeRouter {
|
||||
// A nil embedder leaves topicPolicy on the literal matcher alone.
|
||||
func NewInvokeRouter(resolver CredentialResolver, endpointResolver EndpointResolver, recorder Recorder, access AccessResolver, provisioned *ProvisionedStore, guardrails *GuardrailStore, embedder Embedder) *InvokeRouter {
|
||||
if resolver == nil {
|
||||
resolver = NoopCredentialResolver
|
||||
}
|
||||
@@ -78,7 +82,7 @@ func NewInvokeRouter(resolver CredentialResolver, endpointResolver EndpointResol
|
||||
if access == nil {
|
||||
access = DenyAllAccessResolver
|
||||
}
|
||||
return &InvokeRouter{resolver: resolver, endpointResolver: endpointResolver, recorder: recorder, access: access, provisioned: provisioned, guardrails: guardrails}
|
||||
return &InvokeRouter{resolver: resolver, endpointResolver: endpointResolver, recorder: recorder, access: access, provisioned: provisioned, guardrails: guardrails, embedder: embedder}
|
||||
}
|
||||
|
||||
// InvokeModel routes modelID to its family adapter via the catalog. Unknown
|
||||
@@ -182,7 +186,7 @@ func (rt *InvokeRouter) InvokeModel(ctx context.Context, accountID, modelID stri
|
||||
}
|
||||
var blocked bool
|
||||
var message string
|
||||
blocked, message, _, _, err = enforceGuardrail(ctx, rt.guardrails, accountID, guardrailIdent, guardrailVersion,
|
||||
blocked, message, _, _, err = enforceGuardrail(ctx, rt.guardrails, rt.embedder, accountID, guardrailIdent, guardrailVersion,
|
||||
bedrockruntime.GuardrailContentSourceInput, texts)
|
||||
if err != nil {
|
||||
return nil, "", err
|
||||
@@ -218,7 +222,7 @@ func (rt *InvokeRouter) InvokeModel(ctx context.Context, accountID, modelID stri
|
||||
var messageOut string
|
||||
var redacted []string
|
||||
var gerr error
|
||||
blockedOut, messageOut, redacted, _, gerr = enforceGuardrail(ctx, rt.guardrails, accountID, guardrailIdent, guardrailVersion,
|
||||
blockedOut, messageOut, redacted, _, gerr = enforceGuardrail(ctx, rt.guardrails, rt.embedder, accountID, guardrailIdent, guardrailVersion,
|
||||
bedrockruntime.GuardrailContentSourceOutput, texts)
|
||||
if gerr != nil {
|
||||
err = gerr
|
||||
@@ -240,8 +244,10 @@ func (rt *InvokeRouter) InvokeModel(ctx context.Context, accountID, modelID stri
|
||||
// provisioned may be nil; NewInvokeRouter supplies no-op (and, for access,
|
||||
// deny-all; for provisioned, PT-ARN-rejecting) fallbacks. guardrailIdent/
|
||||
// guardrailVersion come from the request's X-Amzn-Bedrock-Guardrail* headers.
|
||||
func InvokeModel(ctx context.Context, accountID, modelID string, body []byte, resolver CredentialResolver, endpointResolver EndpointResolver, recorder Recorder, access AccessResolver, provisioned *ProvisionedStore, guardrailIdent, guardrailVersion string, guardrails *GuardrailStore) ([]byte, string, error) {
|
||||
return NewInvokeRouter(resolver, endpointResolver, recorder, access, provisioned, guardrails).InvokeModel(ctx, accountID, modelID, body, guardrailIdent, guardrailVersion)
|
||||
// embedder drives topicPolicy's semantic match; nil falls back to the
|
||||
// literal matcher.
|
||||
func InvokeModel(ctx context.Context, accountID, modelID string, body []byte, resolver CredentialResolver, endpointResolver EndpointResolver, recorder Recorder, access AccessResolver, provisioned *ProvisionedStore, guardrailIdent, guardrailVersion string, guardrails *GuardrailStore, embedder Embedder) ([]byte, string, error) {
|
||||
return NewInvokeRouter(resolver, endpointResolver, recorder, access, provisioned, guardrails, embedder).InvokeModel(ctx, accountID, modelID, body, guardrailIdent, guardrailVersion)
|
||||
}
|
||||
|
||||
// InvokeStreamAdapter is the optional streaming capability an InvokeAdapter
|
||||
@@ -266,14 +272,18 @@ type InvokeStreamRouter struct {
|
||||
// guardrails resolves an InvokeModelWithResponseStream request's
|
||||
// guardrail header, mirroring InvokeRouter.guardrails.
|
||||
guardrails *GuardrailStore
|
||||
// embedder drives topicPolicy's semantic match. Nil falls back to the
|
||||
// guardrail engine's literal matcher (see assessTopicPolicy).
|
||||
embedder Embedder
|
||||
}
|
||||
|
||||
// NewInvokeStreamRouter constructs an InvokeStreamRouter. A nil resolver or
|
||||
// endpointResolver falls back to a resolver/resolver that finds nothing, and a
|
||||
// nil access falls back to denying every model, so an InvokeStreamRouter is
|
||||
// always safe to use even before the real stores are wired in. A nil
|
||||
// provisioned disables PT ARN acceptance.
|
||||
func NewInvokeStreamRouter(resolver CredentialResolver, endpointResolver EndpointResolver, access AccessResolver, provisioned *ProvisionedStore, guardrails *GuardrailStore) *InvokeStreamRouter {
|
||||
// provisioned disables PT ARN acceptance. A nil embedder leaves topicPolicy
|
||||
// on the literal matcher alone.
|
||||
func NewInvokeStreamRouter(resolver CredentialResolver, endpointResolver EndpointResolver, access AccessResolver, provisioned *ProvisionedStore, guardrails *GuardrailStore, embedder Embedder) *InvokeStreamRouter {
|
||||
if resolver == nil {
|
||||
resolver = NoopCredentialResolver
|
||||
}
|
||||
@@ -283,7 +293,7 @@ func NewInvokeStreamRouter(resolver CredentialResolver, endpointResolver Endpoin
|
||||
if access == nil {
|
||||
access = DenyAllAccessResolver
|
||||
}
|
||||
return &InvokeStreamRouter{resolver: resolver, endpointResolver: endpointResolver, access: access, provisioned: provisioned, guardrails: guardrails}
|
||||
return &InvokeStreamRouter{resolver: resolver, endpointResolver: endpointResolver, access: access, provisioned: provisioned, guardrails: guardrails, embedder: embedder}
|
||||
}
|
||||
|
||||
// InvokeModelWithResponseStream routes modelID to its family adapter via the
|
||||
@@ -360,7 +370,7 @@ func (rt *InvokeStreamRouter) InvokeModelWithResponseStream(ctx context.Context,
|
||||
}
|
||||
return nil, errors.New(awserrors.ErrorValidationException)
|
||||
}
|
||||
blocked, message, _, _, gerr := enforceGuardrail(ctx, rt.guardrails, accountID, guardrailIdent, guardrailVersion,
|
||||
blocked, message, _, _, gerr := enforceGuardrail(ctx, rt.guardrails, rt.embedder, accountID, guardrailIdent, guardrailVersion,
|
||||
bedrockruntime.GuardrailContentSourceInput, texts)
|
||||
if gerr != nil {
|
||||
if selfHostRelease != nil {
|
||||
@@ -396,7 +406,7 @@ func (rt *InvokeStreamRouter) InvokeModelWithResponseStream(ctx context.Context,
|
||||
src = newSlotReleasingInvokeSource(src, selfHostRelease)
|
||||
}
|
||||
if guardrailIdent != "" {
|
||||
src = newGuardrailInvokeStreamSource(src, rt.guardrails, accountID, guardrailIdent, guardrailVersion, entry.Provider)
|
||||
src = newGuardrailInvokeStreamSource(src, rt.guardrails, rt.embedder, accountID, guardrailIdent, guardrailVersion, entry.Provider)
|
||||
}
|
||||
return src, nil
|
||||
}
|
||||
|
||||
@@ -138,7 +138,7 @@ func (a *anthropicInvokeAdapter) InvokeModelWithResponseStream(ctx context.Conte
|
||||
// than merely stopping our own reads — a disconnected client must free
|
||||
// the Anthropic-side generation, not just this goroutine.
|
||||
reqCtx, cancel := context.WithCancel(ctx)
|
||||
httpReq, err := http.NewRequestWithContext(reqCtx, http.MethodPost, a.baseURL+anthropicMessagesPath, bytes.NewReader(reqBody)) //nolint:gosec // G704: a.baseURL is the hardcoded Anthropic API endpoint (or an httptest stub in tests), never user input
|
||||
httpReq, err := http.NewRequestWithContext(reqCtx, http.MethodPost, a.baseURL+anthropicMessagesPath, bytes.NewReader(reqBody))
|
||||
if err != nil {
|
||||
cancel()
|
||||
slog.Error("anthropic invoke-stream: failed to build request", "model", modelID, "err", err)
|
||||
@@ -149,7 +149,7 @@ func (a *anthropicInvokeAdapter) InvokeModelWithResponseStream(ctx context.Conte
|
||||
httpReq.Header.Set("Content-Type", "application/json")
|
||||
httpReq.Header.Set("Accept", "text/event-stream")
|
||||
|
||||
resp, err := a.httpClient.Do(httpReq) //nolint:gosec // G704: httpReq targets a.baseURL, not user input
|
||||
resp, err := a.httpClient.Do(httpReq)
|
||||
if err != nil {
|
||||
cancel()
|
||||
slog.Error("anthropic invoke-stream: request failed", "model", modelID, "err", err)
|
||||
|
||||
@@ -273,7 +273,7 @@ func (a *llamaInvokeAdapter) InvokeModelWithResponseStream(ctx context.Context,
|
||||
// than merely stopping our own reads — a disconnected client must free
|
||||
// the vLLM-side generation, not just this goroutine.
|
||||
reqCtx, cancel := context.WithCancel(ctx)
|
||||
httpReq, err := http.NewRequestWithContext(reqCtx, http.MethodPost, baseURL+llamaCompletionsPath, bytes.NewReader(reqBody)) //nolint:gosec // G704: baseURL is a resolved pinned self-host endpoint, not user input
|
||||
httpReq, err := http.NewRequestWithContext(reqCtx, http.MethodPost, baseURL+llamaCompletionsPath, bytes.NewReader(reqBody))
|
||||
if err != nil {
|
||||
cancel()
|
||||
slog.Error("llama invoke-stream: failed to build request", "model", modelID, "err", err)
|
||||
@@ -282,7 +282,7 @@ func (a *llamaInvokeAdapter) InvokeModelWithResponseStream(ctx context.Context,
|
||||
httpReq.Header.Set("Content-Type", "application/json")
|
||||
httpReq.Header.Set("Accept", "text/event-stream")
|
||||
|
||||
resp, err := a.httpClient.Do(httpReq) //nolint:gosec // G704: httpReq targets the resolved pinned self-host endpoint, not user input
|
||||
resp, err := a.httpClient.Do(httpReq)
|
||||
if err != nil {
|
||||
cancel()
|
||||
slog.Error("llama invoke-stream: request failed", "model", modelID, "endpoint", baseURL, "err", err)
|
||||
|
||||
@@ -32,8 +32,9 @@ type invokeStreamSource interface {
|
||||
// NewInvokeStreamRouter and the internal NoopRecorder fallback keep this call
|
||||
// safe either way. provisioned may be nil, disabling PT ARN acceptance (any
|
||||
// PT ARN then reads as an unknown modelId). guardrailIdent/guardrailVersion
|
||||
// come from the request's X-Amzn-Bedrock-Guardrail* headers.
|
||||
func InvokeModelWithResponseStream(ctx context.Context, w http.ResponseWriter, accountID, modelID string, body []byte, resolver CredentialResolver, endpointResolver EndpointResolver, requestContentType string, recorder Recorder, access AccessResolver, provisioned *ProvisionedStore, guardrailIdent, guardrailVersion string, guardrails *GuardrailStore) error {
|
||||
// come from the request's X-Amzn-Bedrock-Guardrail* headers. embedder drives
|
||||
// topicPolicy's semantic match; nil falls back to the literal matcher.
|
||||
func InvokeModelWithResponseStream(ctx context.Context, w http.ResponseWriter, accountID, modelID string, body []byte, resolver CredentialResolver, endpointResolver EndpointResolver, requestContentType string, recorder Recorder, access AccessResolver, provisioned *ProvisionedStore, guardrailIdent, guardrailVersion string, guardrails *GuardrailStore, embedder Embedder) error {
|
||||
if recorder == nil {
|
||||
recorder = NoopRecorder
|
||||
}
|
||||
@@ -46,7 +47,7 @@ func InvokeModelWithResponseStream(ctx context.Context, w http.ResponseWriter, a
|
||||
_, recordModelID, _ := resolveInferenceTarget(ctx, accountID, modelID, provisioned)
|
||||
entry, _ := lookupCatalogEntry(recordModelID) // InvokeStreamRouter below re-validates; only its Provider tag is needed here.
|
||||
|
||||
src, err := NewInvokeStreamRouter(resolver, endpointResolver, access, provisioned, guardrails).InvokeModelWithResponseStream(ctx, accountID, modelID, body, guardrailIdent, guardrailVersion)
|
||||
src, err := NewInvokeStreamRouter(resolver, endpointResolver, access, provisioned, guardrails, embedder).InvokeModelWithResponseStream(ctx, accountID, modelID, body, guardrailIdent, guardrailVersion)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
|
||||
@@ -119,7 +119,7 @@ func TestPumpInvokeStream_RecordsInvocationOnUpstreamFault(t *testing.T) {
|
||||
|
||||
func TestInvokeModelWithResponseStream_UnknownModelReturnsResourceNotFoundPreHeader(t *testing.T) {
|
||||
rec := httptest.NewRecorder()
|
||||
err := InvokeModelWithResponseStream(context.Background(), rec, "000000000001", "does.not-exist-v1:0", []byte(`{}`), nil, nil, "application/json", nil, grantAll{}, nil, "", "", nil)
|
||||
err := InvokeModelWithResponseStream(context.Background(), rec, "000000000001", "does.not-exist-v1:0", []byte(`{}`), nil, nil, "application/json", nil, grantAll{}, nil, "", "", nil, nil)
|
||||
require.Error(t, err)
|
||||
assert.Equal(t, awserrors.ErrorResourceNotFoundException, err.Error())
|
||||
assert.Equal(t, 0, rec.Body.Len())
|
||||
@@ -137,7 +137,7 @@ func TestInvokeModelWithResponseStream_SelfHostHappyPath_WritesChunkFrames(t *te
|
||||
rec := httptest.NewRecorder()
|
||||
body := []byte(`{"prompt":"hello","max_gen_len":128}`)
|
||||
|
||||
err := InvokeModelWithResponseStream(context.Background(), rec, "000000000001", modelID, body, nil, NewStaticEndpointResolver(map[string]string{modelID: ts.URL}), "application/json", nil, grantAll{}, nil, "", "", nil)
|
||||
err := InvokeModelWithResponseStream(context.Background(), rec, "000000000001", modelID, body, nil, NewStaticEndpointResolver(map[string]string{modelID: ts.URL}), "application/json", nil, grantAll{}, nil, "", "", nil, nil)
|
||||
require.NoError(t, err)
|
||||
|
||||
assert.Equal(t, http.StatusOK, rec.Code)
|
||||
@@ -162,7 +162,7 @@ func TestInvokeModelWithResponseStream_NonFlusherWriter_ReturnsPreHeaderError(t
|
||||
w := &nonFlusherWriter{}
|
||||
body := []byte(`{"prompt":"hello"}`)
|
||||
|
||||
err := InvokeModelWithResponseStream(context.Background(), w, "000000000001", modelID, body, nil, NewStaticEndpointResolver(map[string]string{modelID: ts.URL}), "application/json", nil, grantAll{}, nil, "", "", nil)
|
||||
err := InvokeModelWithResponseStream(context.Background(), w, "000000000001", modelID, body, nil, NewStaticEndpointResolver(map[string]string{modelID: ts.URL}), "application/json", nil, grantAll{}, nil, "", "", nil, nil)
|
||||
require.Error(t, err)
|
||||
assert.Equal(t, awserrors.ErrorInternalError, err.Error())
|
||||
assert.False(t, w.wroteHeader)
|
||||
|
||||
@@ -51,7 +51,7 @@ func TestInvokeRouter_SelfHostSuccess(t *testing.T) {
|
||||
defer ts.Close()
|
||||
|
||||
modelID := "meta.llama3-2-1b-instruct-v1:0"
|
||||
rt := NewInvokeRouter(nil, NewStaticEndpointResolver(map[string]string{modelID: ts.URL}), nil, grantAll{}, nil, nil)
|
||||
rt := NewInvokeRouter(nil, NewStaticEndpointResolver(map[string]string{modelID: ts.URL}), nil, grantAll{}, nil, nil, nil)
|
||||
|
||||
respBody, contentType, err := rt.InvokeModel(context.Background(), "000000000001", modelID, []byte(`{"prompt":"hello"}`), "", "")
|
||||
require.NoError(t, err)
|
||||
@@ -63,21 +63,21 @@ func TestInvokeRouter_SelfHostSuccess(t *testing.T) {
|
||||
}
|
||||
|
||||
func TestInvokeRouter_UnknownModelReturnsResourceNotFound(t *testing.T) {
|
||||
rt := NewInvokeRouter(nil, nil, nil, grantAll{}, nil, nil)
|
||||
rt := NewInvokeRouter(nil, nil, nil, grantAll{}, nil, nil, nil)
|
||||
_, _, err := rt.InvokeModel(context.Background(), "000000000001", "does.not-exist-v1:0", []byte(`{}`), "", "")
|
||||
require.Error(t, err)
|
||||
assert.Equal(t, awserrors.ErrorResourceNotFoundException, err.Error())
|
||||
}
|
||||
|
||||
func TestInvokeRouter_AnthropicNoCredentialReturnsAccessDenied(t *testing.T) {
|
||||
rt := NewInvokeRouter(stubCredentialResolver{ok: false}, nil, nil, grantAll{}, nil, nil)
|
||||
rt := NewInvokeRouter(stubCredentialResolver{ok: false}, nil, nil, grantAll{}, nil, nil, nil)
|
||||
_, _, err := rt.InvokeModel(context.Background(), "000000000001", "anthropic.claude-3-5-sonnet-20240620-v1:0", []byte(`{}`), "", "")
|
||||
require.Error(t, err)
|
||||
assert.Equal(t, awserrors.ErrorAccessDeniedException, err.Error())
|
||||
}
|
||||
|
||||
func TestInvokeRouter_SelfHostNoEndpointReturnsModelNotReady(t *testing.T) {
|
||||
rt := NewInvokeRouter(nil, nil, nil, grantAll{}, nil, nil)
|
||||
rt := NewInvokeRouter(nil, nil, nil, grantAll{}, nil, nil, nil)
|
||||
_, _, err := rt.InvokeModel(context.Background(), "000000000001", "meta.llama3-2-1b-instruct-v1:0", []byte(`{"prompt":"hello"}`), "", "")
|
||||
require.Error(t, err)
|
||||
assert.Equal(t, awserrors.ErrorModelNotReadyException, err.Error())
|
||||
@@ -95,7 +95,7 @@ func TestInvokeModel_PackageEntryPoint_SelfHostSuccess(t *testing.T) {
|
||||
defer ts.Close()
|
||||
|
||||
modelID := "meta.llama3-2-1b-instruct-v1:0"
|
||||
respBody, contentType, err := InvokeModel(context.Background(), "000000000001", modelID, []byte(`{"prompt":"hello"}`), nil, NewStaticEndpointResolver(map[string]string{modelID: ts.URL}), nil, grantAll{}, nil, "", "", nil)
|
||||
respBody, contentType, err := InvokeModel(context.Background(), "000000000001", modelID, []byte(`{"prompt":"hello"}`), nil, NewStaticEndpointResolver(map[string]string{modelID: ts.URL}), nil, grantAll{}, nil, "", "", nil, nil)
|
||||
require.NoError(t, err)
|
||||
assert.Equal(t, "application/json", contentType)
|
||||
|
||||
@@ -105,7 +105,7 @@ func TestInvokeModel_PackageEntryPoint_SelfHostSuccess(t *testing.T) {
|
||||
}
|
||||
|
||||
func TestInvokeModel_PackageEntryPoint_UnknownModel(t *testing.T) {
|
||||
_, _, err := InvokeModel(context.Background(), "000000000001", "does.not-exist-v1:0", []byte(`{}`), nil, nil, nil, grantAll{}, nil, "", "", nil)
|
||||
_, _, err := InvokeModel(context.Background(), "000000000001", "does.not-exist-v1:0", []byte(`{}`), nil, nil, nil, grantAll{}, nil, "", "", nil, nil)
|
||||
require.Error(t, err)
|
||||
assert.Equal(t, awserrors.ErrorResourceNotFoundException, err.Error())
|
||||
}
|
||||
@@ -119,7 +119,7 @@ func TestInvokeStreamRouter_SelfHostSuccess(t *testing.T) {
|
||||
defer ts.Close()
|
||||
|
||||
modelID := "meta.llama3-2-1b-instruct-v1:0"
|
||||
rt := NewInvokeStreamRouter(nil, NewStaticEndpointResolver(map[string]string{modelID: ts.URL}), grantAll{}, nil, nil)
|
||||
rt := NewInvokeStreamRouter(nil, NewStaticEndpointResolver(map[string]string{modelID: ts.URL}), grantAll{}, nil, nil, nil)
|
||||
|
||||
src, err := rt.InvokeModelWithResponseStream(context.Background(), "000000000001", modelID, []byte(`{"prompt":"hello"}`), "", "")
|
||||
require.NoError(t, err)
|
||||
@@ -130,21 +130,21 @@ func TestInvokeStreamRouter_SelfHostSuccess(t *testing.T) {
|
||||
}
|
||||
|
||||
func TestInvokeStreamRouter_UnknownModelReturnsResourceNotFound(t *testing.T) {
|
||||
rt := NewInvokeStreamRouter(nil, nil, grantAll{}, nil, nil)
|
||||
rt := NewInvokeStreamRouter(nil, nil, grantAll{}, nil, nil, nil)
|
||||
_, err := rt.InvokeModelWithResponseStream(context.Background(), "000000000001", "does.not-exist-v1:0", []byte(`{}`), "", "")
|
||||
require.Error(t, err)
|
||||
assert.Equal(t, awserrors.ErrorResourceNotFoundException, err.Error())
|
||||
}
|
||||
|
||||
func TestInvokeStreamRouter_AnthropicNoCredentialReturnsAccessDenied(t *testing.T) {
|
||||
rt := NewInvokeStreamRouter(stubCredentialResolver{ok: false}, nil, grantAll{}, nil, nil)
|
||||
rt := NewInvokeStreamRouter(stubCredentialResolver{ok: false}, nil, grantAll{}, nil, nil, nil)
|
||||
_, err := rt.InvokeModelWithResponseStream(context.Background(), "000000000001", "anthropic.claude-3-5-sonnet-20240620-v1:0", []byte(`{}`), "", "")
|
||||
require.Error(t, err)
|
||||
assert.Equal(t, awserrors.ErrorAccessDeniedException, err.Error())
|
||||
}
|
||||
|
||||
func TestInvokeStreamRouter_SelfHostNoEndpointReturnsModelNotReady(t *testing.T) {
|
||||
rt := NewInvokeStreamRouter(nil, nil, grantAll{}, nil, nil)
|
||||
rt := NewInvokeStreamRouter(nil, nil, grantAll{}, nil, nil, nil)
|
||||
_, err := rt.InvokeModelWithResponseStream(context.Background(), "000000000001", "meta.llama3-2-1b-instruct-v1:0", []byte(`{"prompt":"hello"}`), "", "")
|
||||
require.Error(t, err)
|
||||
assert.Equal(t, awserrors.ErrorModelNotReadyException, err.Error())
|
||||
@@ -165,7 +165,7 @@ func TestInvokeRouter_SelfHostUnhandledFamilyReturnsValidationException(t *testi
|
||||
}))
|
||||
defer ts.Close()
|
||||
|
||||
rt := NewInvokeRouter(nil, NewStaticEndpointResolver(map[string]string{modelID: ts.URL}), nil, grantAll{}, nil, nil)
|
||||
rt := NewInvokeRouter(nil, NewStaticEndpointResolver(map[string]string{modelID: ts.URL}), nil, grantAll{}, nil, nil, nil)
|
||||
_, _, err := rt.InvokeModel(context.Background(), "000000000001", modelID, []byte(`{"prompt":"hello"}`), "", "")
|
||||
require.Error(t, err)
|
||||
assert.Equal(t, awserrors.ErrorValidationException, err.Error())
|
||||
@@ -186,7 +186,7 @@ func TestInvokeStreamRouter_SelfHostUnhandledFamilyReturnsValidationException(t
|
||||
}))
|
||||
defer ts.Close()
|
||||
|
||||
rt := NewInvokeStreamRouter(nil, NewStaticEndpointResolver(map[string]string{modelID: ts.URL}), grantAll{}, nil, nil)
|
||||
rt := NewInvokeStreamRouter(nil, NewStaticEndpointResolver(map[string]string{modelID: ts.URL}), grantAll{}, nil, nil, nil)
|
||||
_, err := rt.InvokeModelWithResponseStream(context.Background(), "000000000001", modelID, []byte(`{"prompt":"hello"}`), "", "")
|
||||
require.Error(t, err)
|
||||
assert.Equal(t, awserrors.ErrorValidationException, err.Error())
|
||||
@@ -211,7 +211,7 @@ func TestInvokeRouter_SelfHostGatedAtCapacity(t *testing.T) {
|
||||
}))
|
||||
defer ts.Close()
|
||||
|
||||
rt := NewInvokeRouter(nil, NewStaticEndpointResolver(map[string]string{modelID: ts.URL}), nil, grantAll{}, nil, nil)
|
||||
rt := NewInvokeRouter(nil, NewStaticEndpointResolver(map[string]string{modelID: ts.URL}), nil, grantAll{}, nil, nil, nil)
|
||||
|
||||
release, ok := selfHostLimiter.Acquire(admissionKey("", modelID), 1)
|
||||
require.True(t, ok)
|
||||
@@ -244,7 +244,7 @@ func TestInvokeStreamRouter_SelfHostGatedAtCapacity(t *testing.T) {
|
||||
}))
|
||||
defer ts.Close()
|
||||
|
||||
rt := NewInvokeStreamRouter(nil, NewStaticEndpointResolver(map[string]string{modelID: ts.URL}), grantAll{}, nil, nil)
|
||||
rt := NewInvokeStreamRouter(nil, NewStaticEndpointResolver(map[string]string{modelID: ts.URL}), grantAll{}, nil, nil, nil)
|
||||
|
||||
release, ok := selfHostLimiter.Acquire(admissionKey("", modelID), 1)
|
||||
require.True(t, ok)
|
||||
|
||||
@@ -43,12 +43,16 @@ type Router struct {
|
||||
// closed on any GuardrailConfig (ResourceNotFoundException) rather than
|
||||
// proceeding unguarded; a request with none is unaffected either way.
|
||||
guardrails *GuardrailStore
|
||||
// embedder drives topicPolicy's semantic match. Nil falls back to the
|
||||
// guardrail engine's literal matcher (see assessTopicPolicy).
|
||||
embedder Embedder
|
||||
}
|
||||
|
||||
// NewRouter constructs a Router. A nil resolver, endpointResolver, or
|
||||
// recorder falls back to a no-op, and a nil access denies every model, so a
|
||||
// Router is always safe to use even before the real stores are wired in.
|
||||
func NewRouter(resolver CredentialResolver, endpointResolver EndpointResolver, recorder Recorder, access AccessResolver, provisioned *ProvisionedStore, guardrails *GuardrailStore) *Router {
|
||||
// Router is always safe to use even before the real stores are wired in. A
|
||||
// nil embedder leaves topicPolicy on the literal matcher alone.
|
||||
func NewRouter(resolver CredentialResolver, endpointResolver EndpointResolver, recorder Recorder, access AccessResolver, provisioned *ProvisionedStore, guardrails *GuardrailStore, embedder Embedder) *Router {
|
||||
if resolver == nil {
|
||||
resolver = NoopCredentialResolver
|
||||
}
|
||||
@@ -61,7 +65,7 @@ func NewRouter(resolver CredentialResolver, endpointResolver EndpointResolver, r
|
||||
if access == nil {
|
||||
access = DenyAllAccessResolver
|
||||
}
|
||||
return &Router{resolver: resolver, endpointResolver: endpointResolver, recorder: recorder, access: access, provisioned: provisioned, guardrails: guardrails}
|
||||
return &Router{resolver: resolver, endpointResolver: endpointResolver, recorder: recorder, access: access, provisioned: provisioned, guardrails: guardrails, embedder: embedder}
|
||||
}
|
||||
|
||||
// Converse routes modelID to its provider via the catalog. Unknown modelIds
|
||||
@@ -174,7 +178,7 @@ func (rt *Router) Converse(ctx context.Context, accountID, modelID string, input
|
||||
ident, version := aws.StringValue(gc.GuardrailIdentifier), aws.StringValue(gc.GuardrailVersion)
|
||||
var blockedIn bool
|
||||
var messageIn string
|
||||
blockedIn, messageIn, _, inputAssessments, err = enforceGuardrail(ctx, rt.guardrails, accountID, ident, version,
|
||||
blockedIn, messageIn, _, inputAssessments, err = enforceGuardrail(ctx, rt.guardrails, rt.embedder, accountID, ident, version,
|
||||
bedrockruntime.GuardrailContentSourceInput, converseGuardrailTexts(input))
|
||||
if err != nil {
|
||||
return nil, err
|
||||
@@ -194,7 +198,7 @@ func (rt *Router) Converse(ctx context.Context, accountID, modelID string, input
|
||||
}
|
||||
|
||||
ident, version := aws.StringValue(gc.GuardrailIdentifier), aws.StringValue(gc.GuardrailVersion)
|
||||
blockedOut, messageOut, redacted, outputAssessments, gerr := enforceGuardrail(ctx, rt.guardrails, accountID, ident, version,
|
||||
blockedOut, messageOut, redacted, outputAssessments, gerr := enforceGuardrail(ctx, rt.guardrails, rt.embedder, accountID, ident, version,
|
||||
bedrockruntime.GuardrailContentSourceOutput, converseOutputTexts(out))
|
||||
if gerr != nil {
|
||||
err = gerr
|
||||
@@ -215,8 +219,8 @@ func (rt *Router) Converse(ctx context.Context, accountID, modelID string, input
|
||||
// Converse is the bedrock-runtime Converse entry point used by the gateway
|
||||
// route table. All dependency params may be nil; NewRouter supplies safe
|
||||
// fallbacks (deny-all access, PT-ARN- and GuardrailConfig-rejecting).
|
||||
func Converse(ctx context.Context, accountID, modelID string, input *bedrockruntime.ConverseInput, resolver CredentialResolver, endpointResolver EndpointResolver, recorder Recorder, access AccessResolver, provisioned *ProvisionedStore, guardrails *GuardrailStore) (*bedrockruntime.ConverseOutput, error) {
|
||||
return NewRouter(resolver, endpointResolver, recorder, access, provisioned, guardrails).Converse(ctx, accountID, modelID, input)
|
||||
func Converse(ctx context.Context, accountID, modelID string, input *bedrockruntime.ConverseInput, resolver CredentialResolver, endpointResolver EndpointResolver, recorder Recorder, access AccessResolver, provisioned *ProvisionedStore, guardrails *GuardrailStore, embedder Embedder) (*bedrockruntime.ConverseOutput, error) {
|
||||
return NewRouter(resolver, endpointResolver, recorder, access, provisioned, guardrails, embedder).Converse(ctx, accountID, modelID, input)
|
||||
}
|
||||
|
||||
// ConverseStreamProvider is the optional streaming capability a Provider may
|
||||
|
||||
@@ -88,7 +88,7 @@ func TestRouter_Converse_ProvisionedThroughputARN_ResolvesPinnedEndpointForRecor
|
||||
arn := createPTCommitment(t, store, ptCallerAccount)
|
||||
|
||||
spy := &spyEndpointResolver{baseURL: ts.URL}
|
||||
rt := NewRouter(nil, spy, nil, grantAll{}, store, nil)
|
||||
rt := NewRouter(nil, spy, nil, grantAll{}, store, nil, nil)
|
||||
|
||||
out, err := rt.Converse(context.Background(), ptCallerAccount, arn, converseInput())
|
||||
require.NoError(t, err)
|
||||
@@ -110,7 +110,7 @@ func TestRouter_Converse_BareModelID_StillUsesGlobalShorthand(t *testing.T) {
|
||||
store := newProvisionedTestStore(t, newStubEndpointProvisioner())
|
||||
|
||||
spy := &spyEndpointResolver{baseURL: ts.URL}
|
||||
rt := NewRouter(nil, spy, nil, grantAll{}, store, nil)
|
||||
rt := NewRouter(nil, spy, nil, grantAll{}, store, nil, nil)
|
||||
|
||||
_, err := rt.Converse(context.Background(), ptCallerAccount, selfHostTestModel, converseInput())
|
||||
require.NoError(t, err)
|
||||
@@ -128,7 +128,7 @@ func TestRouter_Converse_ProvisionedThroughputARN_UnknownCommitment(t *testing.T
|
||||
store := newProvisionedTestStore(t, newStubEndpointProvisioner())
|
||||
arn := FormatProvisionedModelARN(ptTestRegion, ptCallerAccount, "does-not-exist")
|
||||
|
||||
rt := NewRouter(nil, &spyEndpointResolver{}, nil, grantAll{}, store, nil)
|
||||
rt := NewRouter(nil, &spyEndpointResolver{}, nil, grantAll{}, store, nil, nil)
|
||||
_, err := rt.Converse(context.Background(), ptCallerAccount, arn, converseInput())
|
||||
require.Error(t, err)
|
||||
assert.Equal(t, awserrors.ErrorResourceNotFoundException, err.Error())
|
||||
@@ -142,7 +142,7 @@ func TestRouter_Converse_ProvisionedThroughputARN_ForeignAccount(t *testing.T) {
|
||||
arn := createPTCommitment(t, store, ptCallerAccount)
|
||||
|
||||
spy := &spyEndpointResolver{baseURL: "http://unused:8000"}
|
||||
rt := NewRouter(nil, spy, nil, grantAll{}, store, nil)
|
||||
rt := NewRouter(nil, spy, nil, grantAll{}, store, nil, nil)
|
||||
|
||||
_, err := rt.Converse(context.Background(), ptOtherCaller, arn, converseInput())
|
||||
require.Error(t, err)
|
||||
@@ -169,7 +169,7 @@ func TestInvokeRouter_InvokeModel_ProvisionedThroughputARN_ResolvesPinnedEndpoin
|
||||
arn := createPTCommitment(t, store, ptCallerAccount)
|
||||
|
||||
spy := &spyEndpointResolver{baseURL: ts.URL}
|
||||
rt := NewInvokeRouter(nil, spy, nil, grantAll{}, store, nil)
|
||||
rt := NewInvokeRouter(nil, spy, nil, grantAll{}, store, nil, nil)
|
||||
|
||||
respBody, contentType, err := rt.InvokeModel(context.Background(), ptCallerAccount, arn, []byte(`{"prompt":"hello"}`), "", "")
|
||||
require.NoError(t, err)
|
||||
@@ -197,7 +197,7 @@ func TestInvokeRouter_InvokeModel_BareModelID_StillUsesGlobalShorthand(t *testin
|
||||
|
||||
store := newProvisionedTestStore(t, newStubEndpointProvisioner())
|
||||
spy := &spyEndpointResolver{baseURL: ts.URL}
|
||||
rt := NewInvokeRouter(nil, spy, nil, grantAll{}, store, nil)
|
||||
rt := NewInvokeRouter(nil, spy, nil, grantAll{}, store, nil, nil)
|
||||
|
||||
_, _, err := rt.InvokeModel(context.Background(), ptCallerAccount, selfHostTestModel, []byte(`{"prompt":"hello"}`), "", "")
|
||||
require.NoError(t, err)
|
||||
|
||||
@@ -48,7 +48,7 @@ func TestRouter_ConverseStream_ProvisionedThroughputARN_ResolvesPinnedEndpointFo
|
||||
arn := createPTCommitment(t, store, ptCallerAccount)
|
||||
|
||||
spy := &spyEndpointResolver{baseURL: ts.URL}
|
||||
rt := NewRouter(nil, spy, nil, grantAll{}, store, nil)
|
||||
rt := NewRouter(nil, spy, nil, grantAll{}, store, nil, nil)
|
||||
|
||||
src, err := rt.ConverseStream(context.Background(), ptCallerAccount, arn, converseStreamInput())
|
||||
require.NoError(t, err)
|
||||
@@ -71,7 +71,7 @@ func TestRouter_ConverseStream_BareModelID_StillUsesGlobalShorthand(t *testing.T
|
||||
store := newProvisionedTestStore(t, newStubEndpointProvisioner())
|
||||
|
||||
spy := &spyEndpointResolver{baseURL: ts.URL}
|
||||
rt := NewRouter(nil, spy, nil, grantAll{}, store, nil)
|
||||
rt := NewRouter(nil, spy, nil, grantAll{}, store, nil, nil)
|
||||
|
||||
src, err := rt.ConverseStream(context.Background(), ptCallerAccount, selfHostTestModel, converseStreamInput())
|
||||
require.NoError(t, err)
|
||||
@@ -89,7 +89,7 @@ func TestRouter_ConverseStream_ProvisionedThroughputARN_UnknownCommitment(t *tes
|
||||
store := newProvisionedTestStore(t, newStubEndpointProvisioner())
|
||||
arn := FormatProvisionedModelARN(ptTestRegion, ptCallerAccount, "does-not-exist")
|
||||
|
||||
rt := NewRouter(nil, &spyEndpointResolver{}, nil, grantAll{}, store, nil)
|
||||
rt := NewRouter(nil, &spyEndpointResolver{}, nil, grantAll{}, store, nil, nil)
|
||||
_, err := rt.ConverseStream(context.Background(), ptCallerAccount, arn, converseStreamInput())
|
||||
require.Error(t, err)
|
||||
assert.Equal(t, awserrors.ErrorResourceNotFoundException, err.Error())
|
||||
@@ -104,7 +104,7 @@ func TestRouter_ConverseStream_ProvisionedThroughputARN_ForeignAccount(t *testin
|
||||
arn := createPTCommitment(t, store, ptCallerAccount)
|
||||
|
||||
spy := &spyEndpointResolver{baseURL: "http://unused:8000"}
|
||||
rt := NewRouter(nil, spy, nil, grantAll{}, store, nil)
|
||||
rt := NewRouter(nil, spy, nil, grantAll{}, store, nil, nil)
|
||||
|
||||
_, err := rt.ConverseStream(context.Background(), ptOtherCaller, arn, converseStreamInput())
|
||||
require.Error(t, err)
|
||||
@@ -121,7 +121,7 @@ func TestInvokeStreamRouter_InvokeModelWithResponseStream_ProvisionedThroughputA
|
||||
arn := createPTCommitment(t, store, ptCallerAccount)
|
||||
|
||||
spy := &spyEndpointResolver{baseURL: ts.URL}
|
||||
rt := NewInvokeStreamRouter(nil, spy, grantAll{}, store, nil)
|
||||
rt := NewInvokeStreamRouter(nil, spy, grantAll{}, store, nil, nil)
|
||||
|
||||
src, err := rt.InvokeModelWithResponseStream(context.Background(), ptCallerAccount, arn, []byte(`{"prompt":"hello"}`), "", "")
|
||||
require.NoError(t, err)
|
||||
@@ -142,7 +142,7 @@ func TestInvokeStreamRouter_InvokeModelWithResponseStream_BareModelID_StillUsesG
|
||||
store := newProvisionedTestStore(t, newStubEndpointProvisioner())
|
||||
|
||||
spy := &spyEndpointResolver{baseURL: ts.URL}
|
||||
rt := NewInvokeStreamRouter(nil, spy, grantAll{}, store, nil)
|
||||
rt := NewInvokeStreamRouter(nil, spy, grantAll{}, store, nil, nil)
|
||||
|
||||
src, err := rt.InvokeModelWithResponseStream(context.Background(), ptCallerAccount, selfHostTestModel, []byte(`{"prompt":"hello"}`), "", "")
|
||||
require.NoError(t, err)
|
||||
@@ -160,7 +160,7 @@ func TestInvokeStreamRouter_InvokeModelWithResponseStream_ProvisionedThroughputA
|
||||
store := newProvisionedTestStore(t, newStubEndpointProvisioner())
|
||||
arn := FormatProvisionedModelARN(ptTestRegion, ptCallerAccount, "does-not-exist")
|
||||
|
||||
rt := NewInvokeStreamRouter(nil, &spyEndpointResolver{}, grantAll{}, store, nil)
|
||||
rt := NewInvokeStreamRouter(nil, &spyEndpointResolver{}, grantAll{}, store, nil, nil)
|
||||
_, err := rt.InvokeModelWithResponseStream(context.Background(), ptCallerAccount, arn, []byte(`{"prompt":"hello"}`), "", "")
|
||||
require.Error(t, err)
|
||||
assert.Equal(t, awserrors.ErrorResourceNotFoundException, err.Error())
|
||||
@@ -174,7 +174,7 @@ func TestInvokeStreamRouter_InvokeModelWithResponseStream_ProvisionedThroughputA
|
||||
arn := createPTCommitment(t, store, ptCallerAccount)
|
||||
|
||||
spy := &spyEndpointResolver{baseURL: "http://unused:8000"}
|
||||
rt := NewInvokeStreamRouter(nil, spy, grantAll{}, store, nil)
|
||||
rt := NewInvokeStreamRouter(nil, spy, grantAll{}, store, nil, nil)
|
||||
|
||||
_, err := rt.InvokeModelWithResponseStream(context.Background(), ptOtherCaller, arn, []byte(`{"prompt":"hello"}`), "", "")
|
||||
require.Error(t, err)
|
||||
|
||||
@@ -47,7 +47,7 @@ func TestRouter_Converse_SelfHostSuccess(t *testing.T) {
|
||||
defer ts.Close()
|
||||
|
||||
modelID := "meta.llama3-2-1b-instruct-v1:0"
|
||||
rt := NewRouter(nil, NewStaticEndpointResolver(map[string]string{modelID: ts.URL}), nil, grantAll{}, nil, nil)
|
||||
rt := NewRouter(nil, NewStaticEndpointResolver(map[string]string{modelID: ts.URL}), nil, grantAll{}, nil, nil, nil)
|
||||
|
||||
out, err := rt.Converse(context.Background(), "000000000001", modelID, converseInput())
|
||||
require.NoError(t, err)
|
||||
@@ -56,14 +56,14 @@ func TestRouter_Converse_SelfHostSuccess(t *testing.T) {
|
||||
}
|
||||
|
||||
func TestRouter_Converse_UnknownModelReturnsResourceNotFound(t *testing.T) {
|
||||
rt := NewRouter(nil, nil, nil, grantAll{}, nil, nil)
|
||||
rt := NewRouter(nil, nil, nil, grantAll{}, nil, nil, nil)
|
||||
_, err := rt.Converse(context.Background(), "000000000001", "does.not-exist-v1:0", converseInput())
|
||||
require.Error(t, err)
|
||||
assert.Equal(t, awserrors.ErrorResourceNotFoundException, err.Error())
|
||||
}
|
||||
|
||||
func TestRouter_Converse_AnthropicNoCredentialReturnsAccessDenied(t *testing.T) {
|
||||
rt := NewRouter(stubCredentialResolver{ok: false}, nil, nil, grantAll{}, nil, nil)
|
||||
rt := NewRouter(stubCredentialResolver{ok: false}, nil, nil, grantAll{}, nil, nil, nil)
|
||||
_, err := rt.Converse(context.Background(), "000000000001", "anthropic.claude-3-5-sonnet-20240620-v1:0", converseInput())
|
||||
require.Error(t, err)
|
||||
assert.Equal(t, awserrors.ErrorAccessDeniedException, err.Error())
|
||||
@@ -81,14 +81,14 @@ func TestConverse_PackageEntryPoint_SelfHostSuccess(t *testing.T) {
|
||||
defer ts.Close()
|
||||
|
||||
modelID := "meta.llama3-2-1b-instruct-v1:0"
|
||||
out, err := Converse(context.Background(), "000000000001", modelID, converseInput(), nil, NewStaticEndpointResolver(map[string]string{modelID: ts.URL}), nil, grantAll{}, nil, nil)
|
||||
out, err := Converse(context.Background(), "000000000001", modelID, converseInput(), nil, NewStaticEndpointResolver(map[string]string{modelID: ts.URL}), nil, grantAll{}, nil, nil, nil)
|
||||
require.NoError(t, err)
|
||||
require.NotNil(t, out.Output.Message)
|
||||
assert.Equal(t, "via package Converse", *out.Output.Message.Content[0].Text)
|
||||
}
|
||||
|
||||
func TestConverse_PackageEntryPoint_UnknownModel(t *testing.T) {
|
||||
_, err := Converse(context.Background(), "000000000001", "does.not-exist-v1:0", converseInput(), nil, nil, nil, grantAll{}, nil, nil)
|
||||
_, err := Converse(context.Background(), "000000000001", "does.not-exist-v1:0", converseInput(), nil, nil, nil, grantAll{}, nil, nil, nil)
|
||||
require.Error(t, err)
|
||||
assert.Equal(t, awserrors.ErrorResourceNotFoundException, err.Error())
|
||||
}
|
||||
@@ -110,7 +110,7 @@ func TestRouter_ConverseStream_SelfHostSuccess(t *testing.T) {
|
||||
defer ts.Close()
|
||||
|
||||
modelID := "meta.llama3-2-1b-instruct-v1:0"
|
||||
rt := NewRouter(nil, NewStaticEndpointResolver(map[string]string{modelID: ts.URL}), nil, grantAll{}, nil, nil)
|
||||
rt := NewRouter(nil, NewStaticEndpointResolver(map[string]string{modelID: ts.URL}), nil, grantAll{}, nil, nil, nil)
|
||||
|
||||
src, err := rt.ConverseStream(context.Background(), "000000000001", modelID, converseStreamInput())
|
||||
require.NoError(t, err)
|
||||
@@ -121,28 +121,28 @@ func TestRouter_ConverseStream_SelfHostSuccess(t *testing.T) {
|
||||
}
|
||||
|
||||
func TestRouter_ConverseStream_UnknownModelReturnsResourceNotFound(t *testing.T) {
|
||||
rt := NewRouter(nil, nil, nil, grantAll{}, nil, nil)
|
||||
rt := NewRouter(nil, nil, nil, grantAll{}, nil, nil, nil)
|
||||
_, err := rt.ConverseStream(context.Background(), "000000000001", "does.not-exist-v1:0", converseStreamInput())
|
||||
require.Error(t, err)
|
||||
assert.Equal(t, awserrors.ErrorResourceNotFoundException, err.Error())
|
||||
}
|
||||
|
||||
func TestRouter_ConverseStream_AnthropicNoCredentialReturnsAccessDenied(t *testing.T) {
|
||||
rt := NewRouter(stubCredentialResolver{ok: false}, nil, nil, grantAll{}, nil, nil)
|
||||
rt := NewRouter(stubCredentialResolver{ok: false}, nil, nil, grantAll{}, nil, nil, nil)
|
||||
_, err := rt.ConverseStream(context.Background(), "000000000001", "anthropic.claude-3-5-sonnet-20240620-v1:0", converseStreamInput())
|
||||
require.Error(t, err)
|
||||
assert.Equal(t, awserrors.ErrorAccessDeniedException, err.Error())
|
||||
}
|
||||
|
||||
func TestRouter_ConverseStream_SelfHostNoEndpointReturnsModelNotReady(t *testing.T) {
|
||||
rt := NewRouter(nil, nil, nil, grantAll{}, nil, nil)
|
||||
rt := NewRouter(nil, nil, nil, grantAll{}, nil, nil, nil)
|
||||
_, err := rt.ConverseStream(context.Background(), "000000000001", "meta.llama3-2-1b-instruct-v1:0", converseStreamInput())
|
||||
require.Error(t, err)
|
||||
assert.Equal(t, awserrors.ErrorModelNotReadyException, err.Error())
|
||||
}
|
||||
|
||||
func TestNewRouter_NilArgumentsFallBackToNoops(t *testing.T) {
|
||||
rt := NewRouter(nil, nil, nil, grantAll{}, nil, nil)
|
||||
rt := NewRouter(nil, nil, nil, grantAll{}, nil, nil, nil)
|
||||
require.NotNil(t, rt)
|
||||
|
||||
// Self-host model with no endpoint resolver configured resolves nothing,
|
||||
@@ -169,7 +169,7 @@ func TestRouter_Converse_SelfHostGatedAtCapacity(t *testing.T) {
|
||||
}))
|
||||
defer ts.Close()
|
||||
|
||||
rt := NewRouter(nil, NewStaticEndpointResolver(map[string]string{modelID: ts.URL}), nil, grantAll{}, nil, nil)
|
||||
rt := NewRouter(nil, NewStaticEndpointResolver(map[string]string{modelID: ts.URL}), nil, grantAll{}, nil, nil, nil)
|
||||
|
||||
// Occupy the endpoint's only admission slot directly, standing in for a
|
||||
// concurrent in-flight request on the same (account, model) key.
|
||||
@@ -201,7 +201,7 @@ func TestRouter_ConverseStream_SelfHostGatedAtCapacity(t *testing.T) {
|
||||
}))
|
||||
defer ts.Close()
|
||||
|
||||
rt := NewRouter(nil, NewStaticEndpointResolver(map[string]string{modelID: ts.URL}), nil, grantAll{}, nil, nil)
|
||||
rt := NewRouter(nil, NewStaticEndpointResolver(map[string]string{modelID: ts.URL}), nil, grantAll{}, nil, nil, nil)
|
||||
|
||||
release, ok := selfHostLimiter.Acquire(admissionKey("", modelID), 1)
|
||||
require.True(t, ok)
|
||||
@@ -237,7 +237,7 @@ func TestRouter_ConverseStream_SelfHostGatedAtCapacity(t *testing.T) {
|
||||
// never ThrottlingException — the managed-provider branch must never reach
|
||||
// admitSelfHost.
|
||||
func TestRouter_Converse_AnthropicPathIsNeverGated(t *testing.T) {
|
||||
rt := NewRouter(stubCredentialResolver{ok: false}, nil, nil, grantAll{}, nil, nil)
|
||||
rt := NewRouter(stubCredentialResolver{ok: false}, nil, nil, grantAll{}, nil, nil, nil)
|
||||
|
||||
const n = 20
|
||||
errs := make([]error, n)
|
||||
@@ -280,7 +280,7 @@ func TestRouter_Converse_ProvisionedThroughputARN_CapacityMultipliesByModelUnits
|
||||
}))
|
||||
defer ts.Close()
|
||||
|
||||
rt := NewRouter(nil, NewStaticEndpointResolver(map[string]string{modelID: ts.URL}), nil, grantAll{}, store, nil)
|
||||
rt := NewRouter(nil, NewStaticEndpointResolver(map[string]string{modelID: ts.URL}), nil, grantAll{}, store, nil, nil)
|
||||
key := admissionKey(ptCallerAccount, modelID)
|
||||
|
||||
// Occupy 1 of the 3 units directly (matching MaxConcurrency alone). If
|
||||
|
||||
@@ -308,7 +308,7 @@ func (p *vllmProvider) ConverseStream(ctx context.Context, modelID string, input
|
||||
// than merely stopping our own reads — a disconnected client must free
|
||||
// the vLLM-side generation, not just this goroutine.
|
||||
reqCtx, cancel := context.WithCancel(ctx)
|
||||
httpReq, err := http.NewRequestWithContext(reqCtx, http.MethodPost, baseURL+vllmChatCompletionsPath, bytes.NewReader(reqBody)) //nolint:gosec // G704: baseURL is a resolved pinned self-host endpoint, not user input
|
||||
httpReq, err := http.NewRequestWithContext(reqCtx, http.MethodPost, baseURL+vllmChatCompletionsPath, bytes.NewReader(reqBody))
|
||||
if err != nil {
|
||||
cancel()
|
||||
slog.Error("vllm stream: failed to build request", "model", modelID, "err", err)
|
||||
@@ -317,7 +317,7 @@ func (p *vllmProvider) ConverseStream(ctx context.Context, modelID string, input
|
||||
httpReq.Header.Set("Content-Type", "application/json")
|
||||
httpReq.Header.Set("Accept", "text/event-stream")
|
||||
|
||||
resp, err := p.httpClient.Do(httpReq) //nolint:gosec // G704: httpReq targets the resolved pinned self-host endpoint, not user input
|
||||
resp, err := p.httpClient.Do(httpReq)
|
||||
if err != nil {
|
||||
cancel()
|
||||
slog.Error("vllm stream: request failed", "model", modelID, "endpoint", baseURL, "err", err)
|
||||
|
||||
@@ -153,7 +153,7 @@ func (gw *GatewayConfig) BedrockAgentRuntime_Request(w http.ResponseWriter, r *h
|
||||
}
|
||||
|
||||
converse := func(ctx context.Context, acct, modelID string, input *bedrockruntime.ConverseInput) (*bedrockruntime.ConverseOutput, error) {
|
||||
return gateway_bedrock.Converse(ctx, acct, modelID, input, gw.bedrockResolver(), gw.bedrockEndpointResolver(), gw.bedrockRecorder(), gw.bedrockAccessResolver(), gw.bedrockProvisionedStore(), gw.bedrockGuardrailStore())
|
||||
return gateway_bedrock.Converse(ctx, acct, modelID, input, gw.bedrockResolver(), gw.bedrockEndpointResolver(), gw.bedrockRecorder(), gw.bedrockAccessResolver(), gw.bedrockProvisionedStore(), gw.bedrockGuardrailStore(), gw.bedrockEmbedder())
|
||||
}
|
||||
|
||||
output, err := handler(r.Context(), accountID, params, body, gw.BedrockAgentKB, gw.BedrockAgentVector, converse)
|
||||
|
||||
@@ -37,28 +37,29 @@ type bedrockRuntimeRoute struct {
|
||||
// endpoints; recorder is gw.bedrockRecorder() (invocation recorder or no-op);
|
||||
// access is gw.bedrockAccessResolver() (grant store or deny-all); provisioned
|
||||
// is gw.bedrockProvisionedStore(), consulted when modelId is a PT ARN;
|
||||
// guardrails is gw.bedrockGuardrailStore().
|
||||
type bedrockRuntimeRouteHandler func(ctx context.Context, accountID string, params []string, body []byte, resolver gateway_bedrock.CredentialResolver, endpoints gateway_bedrock.EndpointResolver, recorder gateway_bedrock.Recorder, access gateway_bedrock.AccessResolver, provisioned *gateway_bedrock.ProvisionedStore, guardrails *gateway_bedrock.GuardrailStore) (any, error)
|
||||
// guardrails is gw.bedrockGuardrailStore(); embedder is gw.bedrockEmbedder(),
|
||||
// driving guardrail topicPolicy's semantic match.
|
||||
type bedrockRuntimeRouteHandler func(ctx context.Context, accountID string, params []string, body []byte, resolver gateway_bedrock.CredentialResolver, endpoints gateway_bedrock.EndpointResolver, recorder gateway_bedrock.Recorder, access gateway_bedrock.AccessResolver, provisioned *gateway_bedrock.ProvisionedStore, guardrails *gateway_bedrock.GuardrailStore, embedder gateway_bedrock.Embedder) (any, error)
|
||||
|
||||
// bedrockRuntimeRoutes is the dispatch table. InvokeModel has no handler
|
||||
// function here: BedrockRuntime_Request special-cases its action to bypass
|
||||
// the JSON-marshaling dispatch below, since its response is raw bytes.
|
||||
var bedrockRuntimeRoutes = []bedrockRuntimeRoute{
|
||||
{"POST", regexp.MustCompile(`^/model/([^/]+)/converse$`), "Converse",
|
||||
func(ctx context.Context, acct string, p []string, b []byte, resolver gateway_bedrock.CredentialResolver, endpoints gateway_bedrock.EndpointResolver, recorder gateway_bedrock.Recorder, access gateway_bedrock.AccessResolver, provisioned *gateway_bedrock.ProvisionedStore, guardrails *gateway_bedrock.GuardrailStore) (any, error) {
|
||||
func(ctx context.Context, acct string, p []string, b []byte, resolver gateway_bedrock.CredentialResolver, endpoints gateway_bedrock.EndpointResolver, recorder gateway_bedrock.Recorder, access gateway_bedrock.AccessResolver, provisioned *gateway_bedrock.ProvisionedStore, guardrails *gateway_bedrock.GuardrailStore, embedder gateway_bedrock.Embedder) (any, error) {
|
||||
input := new(bedrockruntime.ConverseInput)
|
||||
if len(b) > 0 {
|
||||
if err := json.Unmarshal(b, input); err != nil {
|
||||
return nil, errors.New(awserrors.ErrorValidationException)
|
||||
}
|
||||
}
|
||||
return gateway_bedrock.Converse(ctx, acct, p[0], input, resolver, endpoints, recorder, access, provisioned, guardrails)
|
||||
return gateway_bedrock.Converse(ctx, acct, p[0], input, resolver, endpoints, recorder, access, provisioned, guardrails, embedder)
|
||||
}},
|
||||
{"POST", regexp.MustCompile(`^/model/([^/]+)/invoke$`), "InvokeModel", nil},
|
||||
{"POST", regexp.MustCompile(`^/model/([^/]+)/converse-stream$`), "ConverseStream", nil},
|
||||
{"POST", regexp.MustCompile(`^/model/([^/]+)/invoke-with-response-stream$`), "InvokeModelWithResponseStream", nil},
|
||||
{"POST", regexp.MustCompile(`^/guardrail/([^/]+)/version/([^/]+)/apply$`), "ApplyGuardrail",
|
||||
func(ctx context.Context, acct string, p []string, b []byte, _ gateway_bedrock.CredentialResolver, _ gateway_bedrock.EndpointResolver, _ gateway_bedrock.Recorder, _ gateway_bedrock.AccessResolver, _ *gateway_bedrock.ProvisionedStore, guardrails *gateway_bedrock.GuardrailStore) (any, error) {
|
||||
func(ctx context.Context, acct string, p []string, b []byte, _ gateway_bedrock.CredentialResolver, _ gateway_bedrock.EndpointResolver, _ gateway_bedrock.Recorder, _ gateway_bedrock.AccessResolver, _ *gateway_bedrock.ProvisionedStore, guardrails *gateway_bedrock.GuardrailStore, embedder gateway_bedrock.Embedder) (any, error) {
|
||||
input := new(bedrockruntime.ApplyGuardrailInput)
|
||||
if len(b) > 0 {
|
||||
if err := json.Unmarshal(b, input); err != nil {
|
||||
@@ -67,7 +68,7 @@ var bedrockRuntimeRoutes = []bedrockRuntimeRoute{
|
||||
}
|
||||
input.GuardrailIdentifier = aws.String(p[0])
|
||||
input.GuardrailVersion = aws.String(p[1])
|
||||
return gateway_bedrock.ApplyGuardrail(ctx, acct, guardrails, input)
|
||||
return gateway_bedrock.ApplyGuardrail(ctx, acct, guardrails, embedder, input)
|
||||
}},
|
||||
}
|
||||
|
||||
@@ -156,7 +157,7 @@ func (gw *GatewayConfig) BedrockRuntime_Request(w http.ResponseWriter, r *http.R
|
||||
if action == "InvokeModel" {
|
||||
guardrailIdent := r.Header.Get(bedrockGuardrailIdentifierHeader)
|
||||
guardrailVersion := r.Header.Get(bedrockGuardrailVersionHeader)
|
||||
respBody, contentType, err := gateway_bedrock.InvokeModel(r.Context(), accountID, params[0], body, gw.bedrockResolver(), gw.bedrockEndpointResolver(), gw.bedrockRecorder(), gw.bedrockAccessResolver(), gw.bedrockProvisionedStore(), guardrailIdent, guardrailVersion, gw.bedrockGuardrailStore())
|
||||
respBody, contentType, err := gateway_bedrock.InvokeModel(r.Context(), accountID, params[0], body, gw.bedrockResolver(), gw.bedrockEndpointResolver(), gw.bedrockRecorder(), gw.bedrockAccessResolver(), gw.bedrockProvisionedStore(), guardrailIdent, guardrailVersion, gw.bedrockGuardrailStore(), gw.bedrockEmbedder())
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
@@ -171,15 +172,15 @@ func (gw *GatewayConfig) BedrockRuntime_Request(w http.ResponseWriter, r *http.R
|
||||
// (-> ErrorHandler); once streaming starts they always return nil,
|
||||
// surfacing any further failure as an in-band exception event.
|
||||
if action == "ConverseStream" {
|
||||
return gateway_bedrock.ConverseStream(r.Context(), w, accountID, params[0], body, gw.bedrockResolver(), gw.bedrockEndpointResolver(), gw.bedrockRecorder(), gw.bedrockAccessResolver(), gw.bedrockProvisionedStore(), gw.bedrockGuardrailStore())
|
||||
return gateway_bedrock.ConverseStream(r.Context(), w, accountID, params[0], body, gw.bedrockResolver(), gw.bedrockEndpointResolver(), gw.bedrockRecorder(), gw.bedrockAccessResolver(), gw.bedrockProvisionedStore(), gw.bedrockGuardrailStore(), gw.bedrockEmbedder())
|
||||
}
|
||||
if action == "InvokeModelWithResponseStream" {
|
||||
guardrailIdent := r.Header.Get(bedrockGuardrailIdentifierHeader)
|
||||
guardrailVersion := r.Header.Get(bedrockGuardrailVersionHeader)
|
||||
return gateway_bedrock.InvokeModelWithResponseStream(r.Context(), w, accountID, params[0], body, gw.bedrockResolver(), gw.bedrockEndpointResolver(), r.Header.Get("Content-Type"), gw.bedrockRecorder(), gw.bedrockAccessResolver(), gw.bedrockProvisionedStore(), guardrailIdent, guardrailVersion, gw.bedrockGuardrailStore())
|
||||
return gateway_bedrock.InvokeModelWithResponseStream(r.Context(), w, accountID, params[0], body, gw.bedrockResolver(), gw.bedrockEndpointResolver(), r.Header.Get("Content-Type"), gw.bedrockRecorder(), gw.bedrockAccessResolver(), gw.bedrockProvisionedStore(), guardrailIdent, guardrailVersion, gw.bedrockGuardrailStore(), gw.bedrockEmbedder())
|
||||
}
|
||||
|
||||
output, err := handler(r.Context(), accountID, params, body, gw.bedrockResolver(), gw.bedrockEndpointResolver(), gw.bedrockRecorder(), gw.bedrockAccessResolver(), gw.bedrockProvisionedStore(), gw.bedrockGuardrailStore())
|
||||
output, err := handler(r.Context(), accountID, params, body, gw.bedrockResolver(), gw.bedrockEndpointResolver(), gw.bedrockRecorder(), gw.bedrockAccessResolver(), gw.bedrockProvisionedStore(), gw.bedrockGuardrailStore(), gw.bedrockEmbedder())
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
|
||||
@@ -180,6 +180,10 @@ type GatewayConfig struct {
|
||||
// and friends). Nil falls back to an unconfigured store, under which
|
||||
// reads/writes error rather than panic.
|
||||
BedrockGuardrails *gateway_bedrock.GuardrailStore
|
||||
// BedrockEmbedder drives guardrail topicPolicy's semantic match, reusing
|
||||
// the same embedding endpoint the Ochre KB path resolves models through.
|
||||
// Nil leaves topicPolicy on its literal (word-boundary) matcher alone.
|
||||
BedrockEmbedder gateway_bedrock.Embedder
|
||||
|
||||
// SignupMaxAccounts caps how many accounts /admin/CreateAccount will allow
|
||||
// to exist. Zero means uncapped, which is the behaviour of every cluster
|
||||
|
||||
@@ -402,6 +402,11 @@ func launchService(config *config.ClusterConfig) error {
|
||||
bedrockEndpointSvc, bedrockEndpoints, 0,
|
||||
handlers_bedrock.WithColdStartWait(parseColdStartWait(os.Getenv("OCHRE_COLD_START_WAIT"))))
|
||||
|
||||
// Guardrail topicPolicy's semantic match reuses this same endpoint
|
||||
// resolver (the one every self-hosted model, including the embedding
|
||||
// model, already resolves through) rather than standing up a second one.
|
||||
bedrockEmbedder := gateway_bedrock.NewEmbedder(bedrockEndpointResolver)
|
||||
|
||||
// Bedrock provisioned throughput: commitment metadata lives in the
|
||||
// bedrock-provisioned KV bucket (gateway control plane), while the pinned
|
||||
// endpoint it commits to is requested through the same NATS endpoint
|
||||
@@ -484,6 +489,7 @@ func launchService(config *config.ClusterConfig) error {
|
||||
BedrockAccessAdmin: bedrockAccess,
|
||||
BedrockProvisioned: bedrockProvisioned,
|
||||
BedrockGuardrails: bedrockGuardrails,
|
||||
BedrockEmbedder: bedrockEmbedder,
|
||||
SignupMaxAccounts: signupMaxAccounts,
|
||||
BedrockAgentKB: bedrockAgentKB,
|
||||
BedrockAgentDataSources: bedrockAgentDataSources,
|
||||
|
||||
Reference in New Issue
Block a user