diff --git a/.golangci.yml b/.golangci.yml index 860ef2bd9..48cc499d3 100644 --- a/.golangci.yml +++ b/.golangci.yml @@ -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: diff --git a/spinifex/gateway/bedrock.go b/spinifex/gateway/bedrock.go index 096d89a86..087394046 100644 --- a/spinifex/gateway/bedrock.go +++ b/spinifex/gateway/bedrock.go @@ -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 +} diff --git a/spinifex/gateway/bedrock/anthropic.go b/spinifex/gateway/bedrock/anthropic.go index ee107f1d8..5e03cde01 100644 --- a/spinifex/gateway/bedrock/anthropic.go +++ b/spinifex/gateway/bedrock/anthropic.go @@ -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) diff --git a/spinifex/gateway/bedrock/contract_test.go b/spinifex/gateway/bedrock/contract_test.go index ca0e5423c..e068180ac 100644 --- a/spinifex/gateway/bedrock/contract_test.go +++ b/spinifex/gateway/bedrock/contract_test.go @@ -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) }) diff --git a/spinifex/gateway/bedrock/converse_stream.go b/spinifex/gateway/bedrock/converse_stream.go index eb9273975..76e8c2e16 100644 --- a/spinifex/gateway/bedrock/converse_stream.go +++ b/spinifex/gateway/bedrock/converse_stream.go @@ -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) diff --git a/spinifex/gateway/bedrock/converse_stream_test.go b/spinifex/gateway/bedrock/converse_stream_test.go index 2d6f2b4d4..fa533859d 100644 --- a/spinifex/gateway/bedrock/converse_stream_test.go +++ b/spinifex/gateway/bedrock/converse_stream_test.go @@ -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) diff --git a/spinifex/gateway/bedrock/guardrail.go b/spinifex/gateway/bedrock/guardrail.go index fe00b257b..070a4b888 100644 --- a/spinifex/gateway/bedrock/guardrail.go +++ b/spinifex/gateway/bedrock/guardrail.go @@ -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 { diff --git a/spinifex/gateway/bedrock/guardrail_filter.go b/spinifex/gateway/bedrock/guardrail_filter.go index 8dad20a8c..c6a04688f 100644 --- a/spinifex/gateway/bedrock/guardrail_filter.go +++ b/spinifex/gateway/bedrock/guardrail_filter.go @@ -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 } diff --git a/spinifex/gateway/bedrock/guardrail_filter_test.go b/spinifex/gateway/bedrock/guardrail_filter_test.go index 2131bda64..5147c01d3 100644 --- a/spinifex/gateway/bedrock/guardrail_filter_test.go +++ b/spinifex/gateway/bedrock/guardrail_filter_test.go @@ -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)) } diff --git a/spinifex/gateway/bedrock/guardrail_inference.go b/spinifex/gateway/bedrock/guardrail_inference.go index 53058d14a..092e15d5e 100644 --- a/spinifex/gateway/bedrock/guardrail_inference.go +++ b/spinifex/gateway/bedrock/guardrail_inference.go @@ -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 diff --git a/spinifex/gateway/bedrock/guardrail_inference_test.go b/spinifex/gateway/bedrock/guardrail_inference_test.go index ce40476d2..2eb884a1d 100644 --- a/spinifex/gateway/bedrock/guardrail_inference_test.go +++ b/spinifex/gateway/bedrock/guardrail_inference_test.go @@ -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)) diff --git a/spinifex/gateway/bedrock/guardrail_invoke.go b/spinifex/gateway/bedrock/guardrail_invoke.go index 3d3b6278f..0d43503da 100644 --- a/spinifex/gateway/bedrock/guardrail_invoke.go +++ b/spinifex/gateway/bedrock/guardrail_invoke.go @@ -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 diff --git a/spinifex/gateway/bedrock/guardrail_topic_embed.go b/spinifex/gateway/bedrock/guardrail_topic_embed.go new file mode 100644 index 000000000..7a993e442 --- /dev/null +++ b/spinifex/gateway/bedrock/guardrail_topic_embed.go @@ -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 +} diff --git a/spinifex/gateway/bedrock/guardrail_topic_embed_test.go b/spinifex/gateway/bedrock/guardrail_topic_embed_test.go new file mode 100644 index 000000000..46555f53e --- /dev/null +++ b/spinifex/gateway/bedrock/guardrail_topic_embed_test.go @@ -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) +} diff --git a/spinifex/gateway/bedrock/invoke.go b/spinifex/gateway/bedrock/invoke.go index 91490479f..0e010a617 100644 --- a/spinifex/gateway/bedrock/invoke.go +++ b/spinifex/gateway/bedrock/invoke.go @@ -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 } diff --git a/spinifex/gateway/bedrock/invoke_anthropic.go b/spinifex/gateway/bedrock/invoke_anthropic.go index 12a74d6b0..1043d64b2 100644 --- a/spinifex/gateway/bedrock/invoke_anthropic.go +++ b/spinifex/gateway/bedrock/invoke_anthropic.go @@ -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) diff --git a/spinifex/gateway/bedrock/invoke_llama.go b/spinifex/gateway/bedrock/invoke_llama.go index 662d07d1f..c85aaa2e7 100644 --- a/spinifex/gateway/bedrock/invoke_llama.go +++ b/spinifex/gateway/bedrock/invoke_llama.go @@ -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) diff --git a/spinifex/gateway/bedrock/invoke_stream.go b/spinifex/gateway/bedrock/invoke_stream.go index 49b07a2c9..119c51f1f 100644 --- a/spinifex/gateway/bedrock/invoke_stream.go +++ b/spinifex/gateway/bedrock/invoke_stream.go @@ -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 } diff --git a/spinifex/gateway/bedrock/invoke_stream_test.go b/spinifex/gateway/bedrock/invoke_stream_test.go index 0f5a441c5..2146c98b4 100644 --- a/spinifex/gateway/bedrock/invoke_stream_test.go +++ b/spinifex/gateway/bedrock/invoke_stream_test.go @@ -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) diff --git a/spinifex/gateway/bedrock/invoke_test.go b/spinifex/gateway/bedrock/invoke_test.go index 163b85fbf..eb06d0025 100644 --- a/spinifex/gateway/bedrock/invoke_test.go +++ b/spinifex/gateway/bedrock/invoke_test.go @@ -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) diff --git a/spinifex/gateway/bedrock/provider.go b/spinifex/gateway/bedrock/provider.go index e453e9989..355c1892e 100644 --- a/spinifex/gateway/bedrock/provider.go +++ b/spinifex/gateway/bedrock/provider.go @@ -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 diff --git a/spinifex/gateway/bedrock/provisioned_inference_test.go b/spinifex/gateway/bedrock/provisioned_inference_test.go index b81bf347d..f03ec2414 100644 --- a/spinifex/gateway/bedrock/provisioned_inference_test.go +++ b/spinifex/gateway/bedrock/provisioned_inference_test.go @@ -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) diff --git a/spinifex/gateway/bedrock/provisioned_stream_inference_test.go b/spinifex/gateway/bedrock/provisioned_stream_inference_test.go index 5deccfb5e..9a4a10f9d 100644 --- a/spinifex/gateway/bedrock/provisioned_stream_inference_test.go +++ b/spinifex/gateway/bedrock/provisioned_stream_inference_test.go @@ -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) diff --git a/spinifex/gateway/bedrock/router_test.go b/spinifex/gateway/bedrock/router_test.go index 96d1c9544..285e2f9cd 100644 --- a/spinifex/gateway/bedrock/router_test.go +++ b/spinifex/gateway/bedrock/router_test.go @@ -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 diff --git a/spinifex/gateway/bedrock/vllm.go b/spinifex/gateway/bedrock/vllm.go index 0a5708c5c..8bbe61802 100644 --- a/spinifex/gateway/bedrock/vllm.go +++ b/spinifex/gateway/bedrock/vllm.go @@ -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) diff --git a/spinifex/gateway/bedrockagentruntime.go b/spinifex/gateway/bedrockagentruntime.go index d2e7867db..f64ba9446 100644 --- a/spinifex/gateway/bedrockagentruntime.go +++ b/spinifex/gateway/bedrockagentruntime.go @@ -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) diff --git a/spinifex/gateway/bedrockruntime.go b/spinifex/gateway/bedrockruntime.go index 6e775c440..4a073c37e 100644 --- a/spinifex/gateway/bedrockruntime.go +++ b/spinifex/gateway/bedrockruntime.go @@ -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 } diff --git a/spinifex/gateway/gateway.go b/spinifex/gateway/gateway.go index fa56ffb6b..60a449a4c 100644 --- a/spinifex/gateway/gateway.go +++ b/spinifex/gateway/gateway.go @@ -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 diff --git a/spinifex/services/awsgw/awsgw.go b/spinifex/services/awsgw/awsgw.go index f39b0c5f4..e8639b531 100644 --- a/spinifex/services/awsgw/awsgw.go +++ b/spinifex/services/awsgw/awsgw.go @@ -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,