From cee73671ffa61ac11fb4e54dfd9a38e136b23451 Mon Sep 17 00:00:00 2001 From: Josh Siviour Date: Mon, 13 Jul 2026 21:40:34 +1000 Subject: [PATCH] =?UTF-8?q?feat(bedrock):=20Ochre=20runtime=20MVP=20?= =?UTF-8?q?=E2=80=94=20Converse=20over=20self-host=20+=20Anthropic-direct?= =?UTF-8?q?=20(#520)?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit * feat(bedrock): Ochre runtime MVP — Converse over self-host + Anthropic-direct Add the AWS-Bedrock-compatible gateway services (bedrock, bedrock-runtime) following the EKS REST-JSON pattern: ListFoundationModels / GetFoundationModel over a static connectivity-tiered catalog, and non-streaming Converse routed per modelId to one of two provider backends — self-hosted vLLM (OpenAI wire) and Anthropic-direct (Anthropic Messages wire). Per-account provider API keys are a master-key-encrypted JetStream KV secret with an optional platform default; the advertised catalog gates external models on a resolvable credential. Self-host endpoints are pinned via config (OCHRE_VLLM_ENDPOINTS) behind an EndpointResolver seam. Adds gateway/bedrock/ (catalog, provider router, Anthropic + vLLM adapters, credential store), awserrors REST-JSON codes, and the bedrock E2E suite (list/get light; Converse gated on live backends). Daemon/NATS control plane and dynamic endpoint lifecycle are deferred to Phase 3. * test(bedrock): cover dispatch, router, credentials KV + endpoint config Add unit coverage for the Ochre Phase 1 diff: Bedrock_Request / BedrockRuntime_Request end-to-end dispatch, provider router, JSON response/error writers, upstream-status mapping, the credential-store KV round-trip, the Anthropic bound-provider path, and parseBedrockEndpoints. Lifts PR diff coverage above the 60% gate. --- docs/service-interfaces.yaml | 8 + spinifex/awserrors/awserrors.go | 20 ++ spinifex/awserrors/awserrors_test.go | 2 +- spinifex/gateway/bedrock.go | 128 +++++++++ spinifex/gateway/bedrock/anthropic.go | 251 ++++++++++++++++++ spinifex/gateway/bedrock/anthropic_test.go | 142 ++++++++++ spinifex/gateway/bedrock/catalog.go | 154 +++++++++++ spinifex/gateway/bedrock/catalog_test.go | 62 +++++ spinifex/gateway/bedrock/credentials.go | 130 +++++++++ spinifex/gateway/bedrock/credentials_test.go | 116 ++++++++ spinifex/gateway/bedrock/handler.go | 71 +++++ spinifex/gateway/bedrock/handler_test.go | 75 ++++++ spinifex/gateway/bedrock/httperrors.go | 24 ++ spinifex/gateway/bedrock/httperrors_test.go | 31 +++ spinifex/gateway/bedrock/provider.go | 86 ++++++ spinifex/gateway/bedrock/router_test.go | 104 ++++++++ spinifex/gateway/bedrock/vllm.go | 224 ++++++++++++++++ spinifex/gateway/bedrock/vllm_test.go | 73 +++++ spinifex/gateway/bedrock_request_test.go | 166 ++++++++++++ spinifex/gateway/bedrock_test.go | 48 ++++ spinifex/gateway/bedrockruntime.go | 113 ++++++++ spinifex/gateway/gateway.go | 20 +- spinifex/services/awsgw/awsgw.go | 76 ++++-- spinifex/services/awsgw/awsgw_bedrock_test.go | 60 +++++ tests/e2e/Makefile | 7 +- tests/e2e/bedrock/bedrock_test.go | 130 +++++++++ tests/e2e/bedrock/main_test.go | 61 +++++ tests/e2e/harness/aws.go | 42 +-- 28 files changed, 2383 insertions(+), 41 deletions(-) create mode 100644 spinifex/gateway/bedrock.go create mode 100644 spinifex/gateway/bedrock/anthropic.go create mode 100644 spinifex/gateway/bedrock/anthropic_test.go create mode 100644 spinifex/gateway/bedrock/catalog.go create mode 100644 spinifex/gateway/bedrock/catalog_test.go create mode 100644 spinifex/gateway/bedrock/credentials.go create mode 100644 spinifex/gateway/bedrock/credentials_test.go create mode 100644 spinifex/gateway/bedrock/handler.go create mode 100644 spinifex/gateway/bedrock/handler_test.go create mode 100644 spinifex/gateway/bedrock/httperrors.go create mode 100644 spinifex/gateway/bedrock/httperrors_test.go create mode 100644 spinifex/gateway/bedrock/provider.go create mode 100644 spinifex/gateway/bedrock/router_test.go create mode 100644 spinifex/gateway/bedrock/vllm.go create mode 100644 spinifex/gateway/bedrock/vllm_test.go create mode 100644 spinifex/gateway/bedrock_request_test.go create mode 100644 spinifex/gateway/bedrock_test.go create mode 100644 spinifex/gateway/bedrockruntime.go create mode 100644 spinifex/services/awsgw/awsgw_bedrock_test.go create mode 100644 tests/e2e/bedrock/bedrock_test.go create mode 100644 tests/e2e/bedrock/main_test.go diff --git a/docs/service-interfaces.yaml b/docs/service-interfaces.yaml index c7131da67..e12744e7f 100644 --- a/docs/service-interfaces.yaml +++ b/docs/service-interfaces.yaml @@ -300,6 +300,14 @@ suites: # spinifex_daemon. docker/crane/skopeo data-plane subtests skip when the # client binary or registry-host resolution is absent. + e2e-bedrock: + path: tests/e2e/bedrock + covers: [awsgw] + # Ochre (AWS-Bedrock) control plane: static catalog list/get served + # straight from awsgw (gateway/bedrock.go, gateway/bedrock/) — no NATS + # hop to spinifex_daemon. Converse (Anthropic-direct / self-host vLLM) + # subtests are gated on backend availability and skip by default. + e2e-ecs: path: tests/e2e/ecs covers: [awsgw, spinifex_daemon, vpcd, viperblockd, predastore] diff --git a/spinifex/awserrors/awserrors.go b/spinifex/awserrors/awserrors.go index 2da1430c3..08b8faab4 100644 --- a/spinifex/awserrors/awserrors.go +++ b/spinifex/awserrors/awserrors.go @@ -499,6 +499,17 @@ var ( ErrorELBv2IncompatibleProtocols = "IncompatibleProtocols" ErrorELBv2CertificateNotFound = "CertificateNotFound" ErrorELBv2SSLPolicyNotFound = "SSLPolicyNotFound" + + // Bedrock/bedrock-runtime error codes + ErrorValidationException = "ValidationException" + // ErrorResourceNotFoundException aliases ErrorEKSResourceNotFound: both + // services use the same wire code, so it is not a second ErrorLookup entry. + ErrorResourceNotFoundException = ErrorEKSResourceNotFound + ErrorAccessDeniedException = "AccessDeniedException" + ErrorThrottlingException = "ThrottlingException" + ErrorModelNotReadyException = "ModelNotReadyException" + ErrorServiceUnavailableException = "ServiceUnavailableException" + ErrorModelErrorException = "ModelErrorException" ) // ValidErrorCode returns the error code if it exists in ErrorLookup, @@ -1021,4 +1032,13 @@ var ErrorLookup = map[string]ErrorMessage{ ErrorELBv2IncompatibleProtocols: {HTTPCode: 400, Message: "The listener protocol is incompatible with the target group protocol."}, ErrorELBv2CertificateNotFound: {HTTPCode: 400, Message: "One or more specified certificates do not exist."}, ErrorELBv2SSLPolicyNotFound: {HTTPCode: 400, Message: "The specified SSL policy does not exist."}, + + // Bedrock/bedrock-runtime error codes. ResourceNotFoundException reuses the + // EKS entry (ErrorEKSResourceNotFound) — same wire code, same HTTP status. + ErrorValidationException: {HTTPCode: 400, Message: "The input fails to satisfy the constraints specified by the model or service."}, + ErrorAccessDeniedException: {HTTPCode: 403, Message: "You do not have sufficient access to perform this action."}, + ErrorThrottlingException: {HTTPCode: 429, Message: "The request was denied due to request throttling."}, + ErrorModelNotReadyException: {HTTPCode: 429, Message: "The model specified in the request is not ready to serve inference requests."}, + ErrorServiceUnavailableException: {HTTPCode: 503, Message: "The service isn't currently available. Try again later."}, + ErrorModelErrorException: {HTTPCode: 424, Message: "The request failed because of an error while running the model."}, } diff --git a/spinifex/awserrors/awserrors_test.go b/spinifex/awserrors/awserrors_test.go index 0cd6275d2..e7d9cef96 100644 --- a/spinifex/awserrors/awserrors_test.go +++ b/spinifex/awserrors/awserrors_test.go @@ -13,7 +13,7 @@ func TestErrorLookup_Structure(t *testing.T) { t.Fatalf("ErrorLookup unexpectedly small: %d entries", len(ErrorLookup)) } - validHTTP := map[int]bool{400: true, 403: true, 404: true, 409: true, 412: true, 413: true, 500: true, 501: true, 503: true} + validHTTP := map[int]bool{400: true, 403: true, 404: true, 409: true, 412: true, 413: true, 424: true, 429: true, 500: true, 501: true, 503: true} for code, msg := range ErrorLookup { if !validHTTP[msg.HTTPCode] { t.Errorf("%s has invalid HTTPCode %d", code, msg.HTTPCode) diff --git a/spinifex/gateway/bedrock.go b/spinifex/gateway/bedrock.go new file mode 100644 index 000000000..46ef4acbf --- /dev/null +++ b/spinifex/gateway/bedrock.go @@ -0,0 +1,128 @@ +package gateway + +import ( + "context" + "errors" + "io" + "log/slog" + "net/http" + "net/url" + "regexp" + + "github.com/aws/aws-sdk-go/service/bedrock" + "github.com/mulgadc/spinifex/spinifex/awserrors" + gateway_bedrock "github.com/mulgadc/spinifex/spinifex/gateway/bedrock" +) + +// bedrockRoute maps one HTTP method + path regex to an AWS action and handler. +type bedrockRoute struct { + method string + pattern *regexp.Regexp + action string + handler bedrockRouteHandler +} + +// bedrockRouteHandler invokes a per-action bedrock (control-plane) gateway +// function. params holds the regex capture groups, PathUnescape'd. resolver +// is gw.bedrockResolver(): the configured credential store, or a no-op +// fallback. +type bedrockRouteHandler func(ctx context.Context, accountID string, params []string, body []byte, resolver gateway_bedrock.CredentialResolver) (any, error) + +// bedrockRoutes is the dispatch table. More-specific paths must precede +// less-specific ones with the same prefix so the regex matcher picks the +// deeper route first. +var bedrockRoutes = []bedrockRoute{ + {"GET", regexp.MustCompile(`^/foundation-models$`), "ListFoundationModels", + func(ctx context.Context, acct string, p []string, b []byte, resolver gateway_bedrock.CredentialResolver) (any, error) { + return gateway_bedrock.ListFoundationModels(ctx, acct, resolver, new(bedrock.ListFoundationModelsInput)) + }}, + {"GET", regexp.MustCompile(`^/foundation-models/([^/]+)$`), "GetFoundationModel", + func(ctx context.Context, acct string, p []string, b []byte, _ gateway_bedrock.CredentialResolver) (any, error) { + return gateway_bedrock.GetFoundationModel(ctx, acct, p[0]) + }}, +} + +// lookupBedrockAction matches method+path against bedrockRoutes, returning the +// action, path params, and handler, or ("", nil, nil, false) on no match. +// path must be r.URL.EscapedPath(): captured params are PathUnescape'd before +// returning, mirroring lookupEKSAction. +func lookupBedrockAction(method, path string) (string, []string, bedrockRouteHandler, bool) { + for _, route := range bedrockRoutes { + if route.method != method { + continue + } + m := route.pattern.FindStringSubmatch(path) + if m == nil { + continue + } + var params []string + if len(m) > 1 { + params = make([]string, 0, len(m)-1) + for _, raw := range m[1:] { + decoded, err := url.PathUnescape(raw) + if err != nil { + slog.Debug("bedrock: bad percent-encoding in path param", "param", raw, "err", err) + decoded = raw + } + params = append(params, decoded) + } + } + return route.action, params, route.handler, true + } + return "", nil, nil, false +} + +// Bedrock_Request dispatches bedrock (control-plane) REST-JSON requests: +// resolves method+path to an action, reads the body, calls the handler, and +// serialises the output as JSON. +func (gw *GatewayConfig) Bedrock_Request(w http.ResponseWriter, r *http.Request) error { + action, params, handler, ok := lookupBedrockAction(r.Method, r.URL.EscapedPath()) + if !ok { + slog.DebugContext(r.Context(), "bedrock: no route for request", "method", r.Method, "path", r.URL.Path) + return errors.New(awserrors.ErrorInvalidAction) + } + + if err := gw.checkPolicy(r, "bedrock", action); err != nil { + return err + } + + if gw.NATSConn == nil { + return errors.New(awserrors.ErrorServerInternal) + } + + accountID, _ := r.Context().Value(ctxAccountID).(string) + if accountID == "" { + slog.ErrorContext(r.Context(), "Bedrock_Request: no account ID in auth context") + return errors.New(awserrors.ErrorServerInternal) + } + + body, err := io.ReadAll(r.Body) + if err != nil { + slog.ErrorContext(r.Context(), "Bedrock_Request: failed to read body", "err", err) + return errors.New(awserrors.ErrorInvalidParameterValue) + } + + output, err := handler(r.Context(), accountID, params, body, gw.bedrockResolver()) + if err != nil { + return err + } + + gateway_bedrock.WriteJSONResponse(w, output) + return nil +} + +// bedrockResolver returns gw.BedrockCredentials as a CredentialResolver, or +// the no-op fallback when no credential store is configured. +func (gw *GatewayConfig) bedrockResolver() gateway_bedrock.CredentialResolver { + if gw.BedrockCredentials != nil { + return gw.BedrockCredentials + } + return gateway_bedrock.NoopCredentialResolver +} + +// bedrockEndpointResolver returns an EndpointResolver over the configured +// pinned self-host endpoints (gw.BedrockEndpoints). A nil/empty map resolves +// nothing, so self-host models return ModelNotReady until configured. +func (gw *GatewayConfig) bedrockEndpointResolver() gateway_bedrock.EndpointResolver { + return gateway_bedrock.NewStaticEndpointResolver(gw.BedrockEndpoints) +} diff --git a/spinifex/gateway/bedrock/anthropic.go b/spinifex/gateway/bedrock/anthropic.go new file mode 100644 index 000000000..9b9ad292f --- /dev/null +++ b/spinifex/gateway/bedrock/anthropic.go @@ -0,0 +1,251 @@ +package gateway_bedrock + +import ( + "bytes" + "context" + "encoding/json" + "errors" + "io" + "log/slog" + "net/http" + "regexp" + "strings" + "time" + + "github.com/aws/aws-sdk-go/aws" + "github.com/aws/aws-sdk-go/service/bedrockruntime" + "github.com/mulgadc/spinifex/spinifex/awserrors" +) + +const ( + anthropicDefaultBaseURL = "https://api.anthropic.com" + anthropicMessagesPath = "/v1/messages" + anthropicAPIVersion = "2023-06-01" + anthropicDefaultMaxTokens = 4096 +) + +// anthropicModelIDPattern strips a "anthropic." vendor prefix and a trailing +// "-v:" Bedrock version suffix, e.g. +// "anthropic.claude-3-5-sonnet-20240620-v1:0" -> "claude-3-5-sonnet-20240620". +var anthropicModelIDPattern = regexp.MustCompile(`^anthropic\.(.+)-v\d+:\d+$`) + +// anthropicModelID maps a Bedrock modelId to its Anthropic model name. IDs +// that don't match the expected shape pass through unchanged. +func anthropicModelID(bedrockModelID string) string { + if m := anthropicModelIDPattern.FindStringSubmatch(bedrockModelID); m != nil { + return m[1] + } + return bedrockModelID +} + +// anthropicContentBlock is a single Messages API content block. Phase 1 is +// text-only, so only "text" blocks are produced or read. +type anthropicContentBlock struct { + Type string `json:"type"` + Text string `json:"text"` +} + +type anthropicMessage struct { + Role string `json:"role"` + Content []anthropicContentBlock `json:"content"` +} + +type anthropicRequest struct { + Model string `json:"model"` + MaxTokens int64 `json:"max_tokens"` + System string `json:"system,omitempty"` + Messages []anthropicMessage `json:"messages"` + Temperature *float64 `json:"temperature,omitempty"` + TopP *float64 `json:"top_p,omitempty"` + StopSequences []string `json:"stop_sequences,omitempty"` +} + +type anthropicUsage struct { + InputTokens int64 `json:"input_tokens"` + OutputTokens int64 `json:"output_tokens"` +} + +type anthropicResponse struct { + Content []anthropicContentBlock `json:"content"` + StopReason string `json:"stop_reason"` + Usage anthropicUsage `json:"usage"` +} + +// anthropicStopReasons maps Anthropic stop_reason values to Bedrock's +// StopReason enum. Unrecognised values (and tool/guardrail reasons Phase 1 +// doesn't model) fall back to end_turn. +var anthropicStopReasons = map[string]string{ + "end_turn": bedrockruntime.StopReasonEndTurn, + "max_tokens": bedrockruntime.StopReasonMaxTokens, + "stop_sequence": bedrockruntime.StopReasonStopSequence, + "tool_use": bedrockruntime.StopReasonToolUse, +} + +func mapAnthropicStopReason(reason string) string { + if mapped, ok := anthropicStopReasons[reason]; ok { + return mapped + } + return bedrockruntime.StopReasonEndTurn +} + +// anthropicProvider calls the Anthropic Messages API directly. httpClient +// and baseURL are injectable for tests. +type anthropicProvider struct { + httpClient *http.Client + baseURL string +} + +func newAnthropicProviderClient() *anthropicProvider { + return &anthropicProvider{ + httpClient: &http.Client{Timeout: providerHTTPTimeout}, + baseURL: anthropicDefaultBaseURL, + } +} + +// anthropicBoundProvider adapts anthropicProvider to Provider by baking in a +// resolved per-account (or platform-default) API key. +type anthropicBoundProvider struct { + inner *anthropicProvider + apiKey string +} + +var _ Provider = (*anthropicBoundProvider)(nil) + +func (b *anthropicBoundProvider) Converse(ctx context.Context, modelID string, input *bedrockruntime.ConverseInput) (*bedrockruntime.ConverseOutput, error) { + return b.inner.Converse(ctx, modelID, input, b.apiKey) +} + +// newAnthropicProvider returns a Provider that calls the Anthropic Messages +// API with apiKey. +func newAnthropicProvider(apiKey string) Provider { + return &anthropicBoundProvider{inner: newAnthropicProviderClient(), apiKey: apiKey} +} + +// Converse translates input to an Anthropic Messages request, calls the API +// with apiKey, and translates the response back to a Bedrock ConverseOutput. +func (p *anthropicProvider) Converse(ctx context.Context, modelID string, input *bedrockruntime.ConverseInput, apiKey string) (*bedrockruntime.ConverseOutput, error) { + reqBody, err := json.Marshal(buildAnthropicRequest(modelID, input)) + if err != nil { + slog.Error("anthropic: failed to marshal request", "model", modelID, "err", err) + return nil, errors.New(awserrors.ErrorInternalError) + } + + httpReq, err := http.NewRequestWithContext(ctx, http.MethodPost, p.baseURL+anthropicMessagesPath, bytes.NewReader(reqBody)) + if err != nil { + slog.Error("anthropic: failed to build request", "model", modelID, "err", err) + return nil, errors.New(awserrors.ErrorInternalError) + } + httpReq.Header.Set("x-api-key", apiKey) + httpReq.Header.Set("anthropic-version", anthropicAPIVersion) + httpReq.Header.Set("content-type", "application/json") + + start := time.Now() + resp, err := p.httpClient.Do(httpReq) + latency := time.Since(start) + if err != nil { + slog.Error("anthropic: request failed", "model", modelID, "err", err) + return nil, errors.New(awserrors.ErrorServiceUnavailableException) + } + defer resp.Body.Close() + + respBody, err := io.ReadAll(resp.Body) + if err != nil { + slog.Error("anthropic: failed to read response body", "model", modelID, "err", err) + return nil, errors.New(awserrors.ErrorServiceUnavailableException) + } + + if resp.StatusCode < 200 || resp.StatusCode >= 300 { + slog.Error("anthropic: upstream error", "model", modelID, "status", resp.StatusCode, "body", string(respBody)) + return nil, errors.New(mapUpstreamStatus(resp.StatusCode)) + } + + var ar anthropicResponse + if err := json.Unmarshal(respBody, &ar); err != nil { + slog.Error("anthropic: failed to parse response", "model", modelID, "err", err) + return nil, errors.New(awserrors.ErrorModelErrorException) + } + + return anthropicToConverseOutput(ar, latency), nil +} + +// buildAnthropicRequest maps a Bedrock ConverseInput to an Anthropic Messages +// request. Non-text content blocks are dropped (Phase 1 is text-only). +func buildAnthropicRequest(modelID string, input *bedrockruntime.ConverseInput) anthropicRequest { + maxTokens := int64(anthropicDefaultMaxTokens) + var temperature, topP *float64 + var stopSequences []string + if input.InferenceConfig != nil { + if input.InferenceConfig.MaxTokens != nil { + maxTokens = *input.InferenceConfig.MaxTokens + } + temperature = input.InferenceConfig.Temperature + topP = input.InferenceConfig.TopP + stopSequences = aws.StringValueSlice(input.InferenceConfig.StopSequences) + } + + var systemParts []string + for _, s := range input.System { + if s != nil && s.Text != nil { + systemParts = append(systemParts, *s.Text) + } + } + + messages := make([]anthropicMessage, 0, len(input.Messages)) + for _, m := range input.Messages { + if m == nil { + continue + } + var blocks []anthropicContentBlock + for _, c := range m.Content { + // TODO: Phase 1 is text-only; non-text content blocks are dropped. + if c == nil || c.Text == nil { + continue + } + blocks = append(blocks, anthropicContentBlock{Type: "text", Text: *c.Text}) + } + messages = append(messages, anthropicMessage{Role: aws.StringValue(m.Role), Content: blocks}) + } + + return anthropicRequest{ + Model: anthropicModelID(modelID), + MaxTokens: maxTokens, + System: strings.Join(systemParts, "\n"), + Messages: messages, + Temperature: temperature, + TopP: topP, + StopSequences: stopSequences, + } +} + +// anthropicToConverseOutput maps an Anthropic Messages response to a Bedrock +// ConverseOutput, measuring latency from the caller-supplied wall-clock delta. +func anthropicToConverseOutput(ar anthropicResponse, latency time.Duration) *bedrockruntime.ConverseOutput { + var content []*bedrockruntime.ContentBlock + for _, block := range ar.Content { + if block.Type != "text" { + continue + } + content = append(content, &bedrockruntime.ContentBlock{Text: aws.String(block.Text)}) + } + + inputTokens := ar.Usage.InputTokens + outputTokens := ar.Usage.OutputTokens + + return &bedrockruntime.ConverseOutput{ + Output: &bedrockruntime.ConverseOutput_{ + Message: &bedrockruntime.Message{ + Role: aws.String(bedrockruntime.ConversationRoleAssistant), + Content: content, + }, + }, + StopReason: aws.String(mapAnthropicStopReason(ar.StopReason)), + Usage: &bedrockruntime.TokenUsage{ + InputTokens: aws.Int64(inputTokens), + OutputTokens: aws.Int64(outputTokens), + TotalTokens: aws.Int64(inputTokens + outputTokens), + }, + Metrics: &bedrockruntime.ConverseMetrics{ + LatencyMs: aws.Int64(latency.Milliseconds()), + }, + } +} diff --git a/spinifex/gateway/bedrock/anthropic_test.go b/spinifex/gateway/bedrock/anthropic_test.go new file mode 100644 index 000000000..06901e199 --- /dev/null +++ b/spinifex/gateway/bedrock/anthropic_test.go @@ -0,0 +1,142 @@ +package gateway_bedrock + +import ( + "context" + "encoding/json" + "net/http" + "net/http/httptest" + "testing" + + "github.com/aws/aws-sdk-go/aws" + "github.com/aws/aws-sdk-go/service/bedrockruntime" + "github.com/mulgadc/spinifex/spinifex/awserrors" + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" +) + +func TestAnthropicProvider_Converse_MapsResponse(t *testing.T) { + var capturedBody anthropicRequest + var capturedAPIKey string + + ts := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + capturedAPIKey = r.Header.Get("x-api-key") + require.NoError(t, json.NewDecoder(r.Body).Decode(&capturedBody)) + w.Header().Set("Content-Type", "application/json") + w.WriteHeader(http.StatusOK) + _, _ = w.Write([]byte(`{ + "content": [{"type": "text", "text": "Hello there"}], + "stop_reason": "end_turn", + "usage": {"input_tokens": 12, "output_tokens": 5} + }`)) + })) + defer ts.Close() + + p := &anthropicProvider{httpClient: ts.Client(), baseURL: ts.URL} + + input := &bedrockruntime.ConverseInput{ + System: []*bedrockruntime.SystemContentBlock{{Text: aws.String("be nice")}}, + Messages: []*bedrockruntime.Message{ + { + Role: aws.String(bedrockruntime.ConversationRoleUser), + Content: []*bedrockruntime.ContentBlock{{Text: aws.String("hi")}}, + }, + }, + InferenceConfig: &bedrockruntime.InferenceConfiguration{ + MaxTokens: aws.Int64(256), + }, + } + + out, err := p.Converse(context.Background(), "anthropic.claude-3-5-sonnet-20240620-v1:0", input, "sk-test-key") + require.NoError(t, err) + + assert.Equal(t, "sk-test-key", capturedAPIKey) + assert.Equal(t, "claude-3-5-sonnet-20240620", capturedBody.Model) + assert.Equal(t, int64(256), capturedBody.MaxTokens) + assert.Equal(t, "be nice", capturedBody.System) + require.Len(t, capturedBody.Messages, 1) + assert.Equal(t, "user", capturedBody.Messages[0].Role) + + require.NotNil(t, out.Output.Message) + assert.Equal(t, "assistant", *out.Output.Message.Role) + require.Len(t, out.Output.Message.Content, 1) + assert.Equal(t, "Hello there", *out.Output.Message.Content[0].Text) + assert.Equal(t, int64(12), *out.Usage.InputTokens) + assert.Equal(t, int64(5), *out.Usage.OutputTokens) + assert.Equal(t, int64(17), *out.Usage.TotalTokens) + assert.Equal(t, bedrockruntime.StopReasonEndTurn, *out.StopReason) + require.NotNil(t, out.Metrics) + assert.GreaterOrEqual(t, *out.Metrics.LatencyMs, int64(0)) +} + +func TestAnthropicProvider_Converse_401ReturnsAccessDenied(t *testing.T) { + ts := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + w.WriteHeader(http.StatusUnauthorized) + _, _ = w.Write([]byte(`{"error": {"message": "invalid x-api-key"}}`)) + })) + defer ts.Close() + + p := &anthropicProvider{httpClient: ts.Client(), baseURL: ts.URL} + + input := &bedrockruntime.ConverseInput{ + Messages: []*bedrockruntime.Message{ + {Role: aws.String(bedrockruntime.ConversationRoleUser), Content: []*bedrockruntime.ContentBlock{{Text: aws.String("hi")}}}, + }, + } + + _, err := p.Converse(context.Background(), "anthropic.claude-3-5-sonnet-20240620-v1:0", input, "bad-key") + require.Error(t, err) + assert.Equal(t, awserrors.ErrorAccessDeniedException, err.Error()) +} + +// TestAnthropicBoundProvider_Converse exercises newAnthropicProvider, +// newAnthropicProviderClient, and anthropicBoundProvider.Converse (the +// Provider used by Router for the "provider:anthropic" tier). The real +// Anthropic base URL is hardcoded by newAnthropicProvider, so — same package +// as production code — the test reaches into the unexported inner client to +// redirect it at an httptest stub instead of the network. +func TestAnthropicBoundProvider_Converse(t *testing.T) { + ts := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + assert.Equal(t, "sk-bound-test", r.Header.Get("x-api-key")) + w.Header().Set("Content-Type", "application/json") + w.WriteHeader(http.StatusOK) + _, _ = w.Write([]byte(`{ + "content": [{"type": "text", "text": "bound response"}], + "stop_reason": "end_turn", + "usage": {"input_tokens": 3, "output_tokens": 2} + }`)) + })) + defer ts.Close() + + p := newAnthropicProvider("sk-bound-test") + bound, ok := p.(*anthropicBoundProvider) + require.True(t, ok) + require.NotNil(t, bound.inner) + bound.inner.httpClient = ts.Client() + bound.inner.baseURL = ts.URL + + input := &bedrockruntime.ConverseInput{ + Messages: []*bedrockruntime.Message{ + {Role: aws.String(bedrockruntime.ConversationRoleUser), Content: []*bedrockruntime.ContentBlock{{Text: aws.String("hi")}}}, + }, + } + + out, err := p.Converse(context.Background(), "anthropic.claude-3-5-sonnet-20240620-v1:0", input) + require.NoError(t, err) + require.NotNil(t, out.Output.Message) + assert.Equal(t, "bound response", *out.Output.Message.Content[0].Text) +} + +func TestAnthropicModelID(t *testing.T) { + cases := []struct { + name, in, want string + }{ + {"sonnet with version suffix", "anthropic.claude-3-5-sonnet-20240620-v1:0", "claude-3-5-sonnet-20240620"}, + {"haiku with version suffix", "anthropic.claude-3-haiku-20240307-v1:0", "claude-3-haiku-20240307"}, + {"non-matching id passes through", "meta.llama3-70b-instruct-v1:0", "meta.llama3-70b-instruct-v1:0"}, + } + for _, tc := range cases { + t.Run(tc.name, func(t *testing.T) { + assert.Equal(t, tc.want, anthropicModelID(tc.in)) + }) + } +} diff --git a/spinifex/gateway/bedrock/catalog.go b/spinifex/gateway/bedrock/catalog.go new file mode 100644 index 000000000..f416c7b5d --- /dev/null +++ b/spinifex/gateway/bedrock/catalog.go @@ -0,0 +1,154 @@ +package gateway_bedrock + +import ( + "context" + "errors" + "fmt" + "strings" + + "github.com/aws/aws-sdk-go/aws" + "github.com/aws/aws-sdk-go/service/bedrock" + "github.com/mulgadc/spinifex/spinifex/awserrors" +) + +// Tier flags on catalogEntry.Provider. "self-host" models run on Spinifex GPU +// compute; "provider:" models are served by calling the vendor's own +// API and require a resolvable per-account (or platform-default) credential. +const ( + tierSelfHost = "self-host" + providerPrefix = "provider:" + vendorAnthropic = "anthropic" + modelARNFormat = "arn:aws:bedrock:*::foundation-model/%s" +) + +// catalogEntry is one static catalog record. Model IDs mirror AWS exactly so +// existing SDK code and configs are drop-in. +type catalogEntry struct { + ModelID string + ModelName string + ProviderName string + Provider string // "self-host" or "provider:" + InputModalities []string + OutputModalities []string + ResponseStreamingSupported bool + InferenceTypesSupported []string + CustomizationsSupported []string +} + +// catalog is the static Phase-1 model set: one self-hosted open model and one +// Anthropic-direct model. Later phases extend this list. +var catalog = []catalogEntry{ + { + ModelID: "meta.llama3-70b-instruct-v1:0", + ModelName: "Llama 3 70B Instruct", + ProviderName: "Meta", + Provider: tierSelfHost, + InputModalities: []string{"TEXT"}, + OutputModalities: []string{"TEXT"}, + ResponseStreamingSupported: false, + InferenceTypesSupported: []string{"ON_DEMAND"}, + }, + { + ModelID: "anthropic.claude-3-5-sonnet-20240620-v1:0", + ModelName: "Claude 3.5 Sonnet", + ProviderName: "Anthropic", + Provider: providerPrefix + vendorAnthropic, + InputModalities: []string{"TEXT", "IMAGE"}, + OutputModalities: []string{"TEXT"}, + ResponseStreamingSupported: false, + InferenceTypesSupported: []string{"ON_DEMAND"}, + }, +} + +// CredentialResolver resolves accountID's usable provider credential for +// vendor: a per-account key, else an optional platform default. key is only +// meaningful when ok is true. +type CredentialResolver interface { + Resolve(ctx context.Context, accountID, vendor string) (key string, ok bool, err error) +} + +// tieredCatalog returns the catalog entries advertised to accountID: self-host +// entries always, provider entries only when resolver finds a usable credential. +func tieredCatalog(ctx context.Context, accountID string, resolver CredentialResolver) []catalogEntry { + var out []catalogEntry + for _, entry := range catalog { + if entry.Provider == tierSelfHost { + out = append(out, entry) + continue + } + vendor, ok := strings.CutPrefix(entry.Provider, providerPrefix) + if !ok { + continue + } + _, resolvable, err := resolver.Resolve(ctx, accountID, vendor) + if err != nil || !resolvable { + continue + } + out = append(out, entry) + } + return out +} + +func (e catalogEntry) toSummary() *bedrock.FoundationModelSummary { + return &bedrock.FoundationModelSummary{ + ModelArn: aws.String(modelARN(e.ModelID)), + ModelId: aws.String(e.ModelID), + ModelName: aws.String(e.ModelName), + ProviderName: aws.String(e.ProviderName), + InputModalities: aws.StringSlice(e.InputModalities), + OutputModalities: aws.StringSlice(e.OutputModalities), + ResponseStreamingSupported: aws.Bool(e.ResponseStreamingSupported), + InferenceTypesSupported: aws.StringSlice(e.InferenceTypesSupported), + CustomizationsSupported: aws.StringSlice(e.CustomizationsSupported), + } +} + +func (e catalogEntry) toDetails() *bedrock.FoundationModelDetails { + return &bedrock.FoundationModelDetails{ + ModelArn: aws.String(modelARN(e.ModelID)), + ModelId: aws.String(e.ModelID), + ModelName: aws.String(e.ModelName), + ProviderName: aws.String(e.ProviderName), + InputModalities: aws.StringSlice(e.InputModalities), + OutputModalities: aws.StringSlice(e.OutputModalities), + ResponseStreamingSupported: aws.Bool(e.ResponseStreamingSupported), + InferenceTypesSupported: aws.StringSlice(e.InferenceTypesSupported), + CustomizationsSupported: aws.StringSlice(e.CustomizationsSupported), + } +} + +func modelARN(modelID string) string { + return fmt.Sprintf(modelARNFormat, modelID) +} + +// lookupCatalogEntry finds a catalog entry by exact modelId, ignoring tier +// gating (used by the runtime router, which returns its own error class). +func lookupCatalogEntry(modelID string) (catalogEntry, bool) { + for _, entry := range catalog { + if entry.ModelID == modelID { + return entry, true + } + } + return catalogEntry{}, false +} + +// ListFoundationModels returns the deployment-tiered catalog: self-host +// entries always, provider entries only where a credential resolves. +func ListFoundationModels(ctx context.Context, accountID string, resolver CredentialResolver, _ *bedrock.ListFoundationModelsInput) (*bedrock.ListFoundationModelsOutput, error) { + entries := tieredCatalog(ctx, accountID, resolver) + summaries := make([]*bedrock.FoundationModelSummary, 0, len(entries)) + for _, entry := range entries { + summaries = append(summaries, entry.toSummary()) + } + return &bedrock.ListFoundationModelsOutput{ModelSummaries: summaries}, nil +} + +// GetFoundationModel looks up a single model by exact modelId, independent of +// tiering. Unknown models return ResourceNotFoundException. +func GetFoundationModel(_ context.Context, _ string, modelID string) (*bedrock.GetFoundationModelOutput, error) { + entry, ok := lookupCatalogEntry(modelID) + if !ok { + return nil, errors.New(awserrors.ErrorResourceNotFoundException) + } + return &bedrock.GetFoundationModelOutput{ModelDetails: entry.toDetails()}, nil +} diff --git a/spinifex/gateway/bedrock/catalog_test.go b/spinifex/gateway/bedrock/catalog_test.go new file mode 100644 index 000000000..f3a9d11df --- /dev/null +++ b/spinifex/gateway/bedrock/catalog_test.go @@ -0,0 +1,62 @@ +package gateway_bedrock + +import ( + "context" + "testing" + + "github.com/aws/aws-sdk-go/service/bedrock" + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" +) + +// stubResolver reports a resolvable credential (with a fixed key) only for +// vendors in ok. +type stubResolver struct { + ok map[string]bool +} + +func (s stubResolver) Resolve(_ context.Context, _, vendor string) (string, bool, error) { + if !s.ok[vendor] { + return "", false, nil + } + return "stub-key", true, nil +} + +func modelIDs(out *bedrock.ListFoundationModelsOutput) []string { + ids := make([]string, 0, len(out.ModelSummaries)) + for _, m := range out.ModelSummaries { + ids = append(ids, *m.ModelId) + } + return ids +} + +func TestListFoundationModels_SelfHostAlwaysIncluded(t *testing.T) { + out, err := ListFoundationModels(context.Background(), "000000000001", stubResolver{ok: map[string]bool{}}, &bedrock.ListFoundationModelsInput{}) + require.NoError(t, err) + assert.Contains(t, modelIDs(out), "meta.llama3-70b-instruct-v1:0") +} + +func TestListFoundationModels_ProviderIncludedWhenResolvable(t *testing.T) { + out, err := ListFoundationModels(context.Background(), "000000000001", stubResolver{ok: map[string]bool{"anthropic": true}}, &bedrock.ListFoundationModelsInput{}) + require.NoError(t, err) + assert.Contains(t, modelIDs(out), "anthropic.claude-3-5-sonnet-20240620-v1:0") +} + +func TestListFoundationModels_ProviderExcludedWhenUnresolvable(t *testing.T) { + out, err := ListFoundationModels(context.Background(), "000000000001", stubResolver{ok: map[string]bool{}}, &bedrock.ListFoundationModelsInput{}) + require.NoError(t, err) + assert.NotContains(t, modelIDs(out), "anthropic.claude-3-5-sonnet-20240620-v1:0") +} + +func TestGetFoundationModel_KnownModel(t *testing.T) { + out, err := GetFoundationModel(context.Background(), "000000000001", "meta.llama3-70b-instruct-v1:0") + require.NoError(t, err) + require.NotNil(t, out.ModelDetails) + assert.Equal(t, "meta.llama3-70b-instruct-v1:0", *out.ModelDetails.ModelId) +} + +func TestGetFoundationModel_UnknownModelReturnsNotFound(t *testing.T) { + _, err := GetFoundationModel(context.Background(), "000000000001", "does-not-exist") + require.Error(t, err) + assert.Equal(t, "ResourceNotFoundException", err.Error()) +} diff --git a/spinifex/gateway/bedrock/credentials.go b/spinifex/gateway/bedrock/credentials.go new file mode 100644 index 000000000..9c82319cb --- /dev/null +++ b/spinifex/gateway/bedrock/credentials.go @@ -0,0 +1,130 @@ +package gateway_bedrock + +import ( + "context" + "fmt" + "sync" + + handlers_iam "github.com/mulgadc/spinifex/spinifex/handlers/iam" + "github.com/nats-io/nats.go" +) + +// bedrockCredentialsBucket is the cluster-replicated KV bucket holding +// per-account provider API keys, encrypted at rest with the IAM master key. +const bedrockCredentialsBucket = "bedrock-credentials" //nolint:gosec // G101: KV bucket name, not a credential + +// bedrockCredentialsHistory keeps one revision; rotation overwrites in place. +const bedrockCredentialsHistory = 1 + +// credentialKey returns the KV key for accountID's vendor credential. +func credentialKey(accountID, vendor string) string { + return fmt.Sprintf("providers/%s/%s", accountID, vendor) +} + +// CredentialStore resolves per-account provider API keys from the +// bedrock-credentials JetStream KV bucket, falling back to a platform-wide +// default when an account has none. Keys are encrypted at rest. +type CredentialStore struct { + js nats.JetStreamContext + masterKey []byte + replicas int + platformDefaults map[string]string + + mu sync.Mutex + kv nats.KeyValue +} + +var _ CredentialResolver = (*CredentialStore)(nil) + +// NewCredentialStore constructs a CredentialStore. platformDefaults maps +// vendor to a platform-wide API key used when an account has no key of its +// own; pass an empty map to disable platform defaults entirely. +func NewCredentialStore(js nats.JetStreamContext, masterKey []byte, replicas int, platformDefaults map[string]string) *CredentialStore { + return &CredentialStore{ + js: js, + masterKey: masterKey, + replicas: replicas, + platformDefaults: platformDefaults, + } +} + +// bucket lazily opens (or creates) the bedrock-credentials KV bucket, +// caching the handle, mirroring gateway_ecrauth.openSigningBucket. +func (s *CredentialStore) bucket() (nats.KeyValue, error) { + s.mu.Lock() + defer s.mu.Unlock() + if s.kv != nil { + return s.kv, nil + } + kv, err := s.js.CreateKeyValue(&nats.KeyValueConfig{ + Bucket: bedrockCredentialsBucket, + History: bedrockCredentialsHistory, + Replicas: s.replicas, + }) + if err != nil { + if kv, err = s.js.KeyValue(bedrockCredentialsBucket); err != nil { + return nil, fmt.Errorf("open %s bucket: %w", bedrockCredentialsBucket, err) + } + } + s.kv = kv + return kv, nil +} + +// Resolve returns accountID's vendor API key: a per-account key if one is +// stored, else the platform default, else ("", false, nil). A nil js is +// tolerated so a store built with only platformDefaults (e.g. in tests) can +// still serve the default-only path without touching JetStream. +func (s *CredentialStore) Resolve(_ context.Context, accountID, vendor string) (string, bool, error) { + if s.js != nil { + kv, err := s.bucket() + if err != nil { + return "", false, err + } + entry, err := kv.Get(credentialKey(accountID, vendor)) + switch { + case err == nil: + plaintext, err := handlers_iam.DecryptSecret(string(entry.Value()), s.masterKey) + if err != nil { + return "", false, fmt.Errorf("decrypt credential for %s/%s: %w", accountID, vendor, err) + } + return plaintext, true, nil + case err != nats.ErrKeyNotFound: + return "", false, fmt.Errorf("kv get credential for %s/%s: %w", accountID, vendor, err) + } + } + if key, ok := s.platformDefaults[vendor]; ok && key != "" { + return key, true, nil + } + return "", false, nil +} + +// PutCredential encrypts and stores key as accountID's vendor credential. +func (s *CredentialStore) PutCredential(_ context.Context, accountID, vendor, key string) error { + kv, err := s.bucket() + if err != nil { + return err + } + ciphertext, err := handlers_iam.EncryptSecret(key, s.masterKey) + if err != nil { + return fmt.Errorf("encrypt credential for %s/%s: %w", accountID, vendor, err) + } + if _, err := kv.Put(credentialKey(accountID, vendor), []byte(ciphertext)); err != nil { + return fmt.Errorf("kv put credential for %s/%s: %w", accountID, vendor, err) + } + return nil +} + +// noopCredentialResolver is the zero-value fallback used when no +// CredentialStore is configured (e.g. unit tests of unrelated routes): it +// resolves no vendors. +type noopCredentialResolver struct{} + +var _ CredentialResolver = (*noopCredentialResolver)(nil) + +func (noopCredentialResolver) Resolve(_ context.Context, _, _ string) (string, bool, error) { + return "", false, nil +} + +// NoopCredentialResolver resolves no vendors. Used as the fallback wherever +// no CredentialStore is configured (nil GatewayConfig.BedrockCredentials). +var NoopCredentialResolver CredentialResolver = noopCredentialResolver{} diff --git a/spinifex/gateway/bedrock/credentials_test.go b/spinifex/gateway/bedrock/credentials_test.go new file mode 100644 index 000000000..4803d9ef9 --- /dev/null +++ b/spinifex/gateway/bedrock/credentials_test.go @@ -0,0 +1,116 @@ +package gateway_bedrock + +import ( + "bytes" + "context" + "testing" + + handlers_iam "github.com/mulgadc/spinifex/spinifex/handlers/iam" + "github.com/mulgadc/spinifex/spinifex/testutil" + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" +) + +// fixedMasterKey returns a deterministic 32-byte AES-256 key for tests. +func fixedMasterKey() []byte { + return bytes.Repeat([]byte{0x42}, 32) +} + +// TestCredentialStore_Resolve_NilJS covers the platform-default and +// vendor-miss paths of a store built with a nil JetStream context (as a +// live NATS server is out of scope here). Resolve must consult +// platformDefaults without dereferencing js. +func TestCredentialStore_Resolve_NilJS(t *testing.T) { + store := NewCredentialStore(nil, fixedMasterKey(), 1, map[string]string{"anthropic": "sk-platform-default"}) + + cases := []struct { + name string + vendor string + wantKey string + wantOK bool + }{ + {"platform default hit", "anthropic", "sk-platform-default", true}, + {"vendor miss", "openai", "", false}, + } + for _, tc := range cases { + t.Run(tc.name, func(t *testing.T) { + key, ok, err := store.Resolve(context.Background(), "000000000001", tc.vendor) + require.NoError(t, err) + assert.Equal(t, tc.wantOK, ok) + assert.Equal(t, tc.wantKey, key) + }) + } +} + +func TestCredentialStore_Resolve_NoPlatformDefaults(t *testing.T) { + store := NewCredentialStore(nil, fixedMasterKey(), 1, nil) + + key, ok, err := store.Resolve(context.Background(), "000000000001", "anthropic") + require.NoError(t, err) + assert.False(t, ok) + assert.Empty(t, key) +} + +func TestNoopCredentialResolver_ResolvesNothing(t *testing.T) { + key, ok, err := NoopCredentialResolver.Resolve(context.Background(), "000000000001", "anthropic") + require.NoError(t, err) + assert.False(t, ok) + assert.Empty(t, key) +} + +// TestCredentialStore_PutAndResolve_KV exercises the real JetStream KV path: +// bucket (lazy create), credentialKey, PutCredential (encrypt+Put), and the +// KV-hit branch of Resolve (Get+decrypt). +func TestCredentialStore_PutAndResolve_KV(t *testing.T) { + _, _, js := testutil.StartTestJetStream(t) + store := NewCredentialStore(js, fixedMasterKey(), 1, nil) + + ctx := context.Background() + require.NoError(t, store.PutCredential(ctx, "000000000001", "anthropic", "sk-test")) + + key, ok, err := store.Resolve(ctx, "000000000001", "anthropic") + require.NoError(t, err) + assert.True(t, ok) + assert.Equal(t, "sk-test", key) +} + +// TestCredentialStore_Resolve_KVMissFallsBackToPlatformDefault covers a KV +// miss (nats.ErrKeyNotFound) that falls through to the platform default. +func TestCredentialStore_Resolve_KVMissFallsBackToPlatformDefault(t *testing.T) { + _, _, js := testutil.StartTestJetStream(t) + store := NewCredentialStore(js, fixedMasterKey(), 1, map[string]string{"anthropic": "sk-platform-default"}) + + key, ok, err := store.Resolve(context.Background(), "999999999999", "anthropic") + require.NoError(t, err) + assert.True(t, ok) + assert.Equal(t, "sk-platform-default", key) +} + +// TestCredentialStore_Resolve_KVMissNoPlatformDefault covers a KV miss on an +// unknown account with no platform defaults configured: ("", false, nil). +func TestCredentialStore_Resolve_KVMissNoPlatformDefault(t *testing.T) { + _, _, js := testutil.StartTestJetStream(t) + store := NewCredentialStore(js, fixedMasterKey(), 1, nil) + + key, ok, err := store.Resolve(context.Background(), "999999999999", "anthropic") + require.NoError(t, err) + assert.False(t, ok) + assert.Empty(t, key) +} + +func TestCredentialKey(t *testing.T) { + assert.Equal(t, "providers/000000000001/anthropic", credentialKey("000000000001", "anthropic")) +} + +func TestEncryptDecryptSecret_RoundTrip(t *testing.T) { + key := fixedMasterKey() + plaintext := "sk-ant-test-secret" + + ciphertext, err := handlers_iam.EncryptSecret(plaintext, key) + require.NoError(t, err) + assert.NotEqual(t, plaintext, ciphertext) + + decrypted, err := handlers_iam.DecryptSecret(ciphertext, key) + require.NoError(t, err) + assert.Equal(t, plaintext, decrypted) +} diff --git a/spinifex/gateway/bedrock/handler.go b/spinifex/gateway/bedrock/handler.go new file mode 100644 index 000000000..fe827b4e3 --- /dev/null +++ b/spinifex/gateway/bedrock/handler.go @@ -0,0 +1,71 @@ +// Package gateway_bedrock is the HTTP-side glue between awsgw and the Ochre +// (bedrock / bedrock-runtime) handlers. Both speak AWS REST-JSON 1.1; error +// envelopes are JSON and Content-Type is application/x-amz-json-1.1 throughout. +package gateway_bedrock + +import ( + "encoding/json" + "fmt" + "log/slog" + "net/http" + "strings" + + "github.com/aws/aws-sdk-go/private/protocol/json/jsonutil" + "github.com/mulgadc/spinifex/spinifex/awserrors" +) + +// JSONContentType is the AWS REST-JSON 1.1 content type bedrock clients expect. +const JSONContentType = "application/x-amz-json-1.1" + +// jsonError is the AWS REST-JSON error envelope. SDKs key off __type for +// awserr.Code() and message for awserr.Message(). +type jsonError struct { + Type string `json:"__type"` + Message string `json:"message"` +} + +// generateJSONError marshals the AWS REST-JSON error envelope. The suffix +// "Exception" is appended idempotently — codes that already end in it are +// left unchanged to avoid the double "ExceptionException" the SDK rejects. +func generateJSONError(code, message string) []byte { + if !strings.HasSuffix(code, "Exception") { + code += "Exception" + } + body, err := json.Marshal(jsonError{Type: code, Message: message}) + if err != nil { + slog.Error("Failed to marshal bedrock error JSON", "code", code, "err", err) + return fmt.Appendf(nil, `{"__type":"InternalErrorException","message":%q}`, message) + } + return body +} + +// WriteJSONResponse serializes obj as AWS REST-JSON and writes a 200 response. +// Uses jsonutil.BuildJSON (the SDK's own restjson marshaler) to honour +// locationName tags that encoding/json does not understand. +func WriteJSONResponse(w http.ResponseWriter, obj any) { + body, err := jsonutil.BuildJSON(obj) + if err != nil { + slog.Error("Failed to marshal bedrock response JSON", "err", err) + WriteJSONError(w, awserrors.ErrorInternalError, "failed to marshal response", http.StatusInternalServerError) + return + } + w.Header().Set("Content-Type", JSONContentType) + w.WriteHeader(http.StatusOK) + if _, err := w.Write(body); err != nil { + slog.Error("Failed to write bedrock response", "err", err) + } +} + +// WriteJSONError emits the AWS REST-JSON error envelope with the given code, +// message, and HTTP status. +func WriteJSONError(w http.ResponseWriter, code, message string, httpStatus int) { + body := generateJSONError(code, message) + w.Header().Set("Content-Type", JSONContentType) + if httpStatus == 0 { + httpStatus = http.StatusInternalServerError + } + w.WriteHeader(httpStatus) + if _, err := w.Write(body); err != nil { + slog.Error("Failed to write bedrock error response", "err", err) + } +} diff --git a/spinifex/gateway/bedrock/handler_test.go b/spinifex/gateway/bedrock/handler_test.go new file mode 100644 index 000000000..2d58c706c --- /dev/null +++ b/spinifex/gateway/bedrock/handler_test.go @@ -0,0 +1,75 @@ +package gateway_bedrock + +import ( + "encoding/json" + "net/http" + "net/http/httptest" + "testing" + + "github.com/mulgadc/spinifex/spinifex/awserrors" + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" +) + +func TestWriteJSONResponse_WritesBodyAndContentType(t *testing.T) { + type payload struct { + Foo string `json:"foo"` + } + + w := httptest.NewRecorder() + WriteJSONResponse(w, payload{Foo: "bar"}) + + assert.Equal(t, http.StatusOK, w.Code) + assert.Equal(t, JSONContentType, w.Header().Get("Content-Type")) + + var out payload + require.NoError(t, json.Unmarshal(w.Body.Bytes(), &out)) + assert.Equal(t, "bar", out.Foo) +} + +func TestWriteJSONError_WritesEnvelopeAndStatus(t *testing.T) { + cases := []struct { + name string + code string + httpStatus int + wantStatus int + wantType string + }{ + {"explicit status", awserrors.ErrorValidationException, http.StatusBadRequest, http.StatusBadRequest, "ValidationException"}, + {"zero status defaults to 500", awserrors.ErrorInternalError, 0, http.StatusInternalServerError, "InternalErrorException"}, + } + for _, tc := range cases { + t.Run(tc.name, func(t *testing.T) { + w := httptest.NewRecorder() + WriteJSONError(w, tc.code, "boom", tc.httpStatus) + + assert.Equal(t, tc.wantStatus, w.Code) + assert.Equal(t, JSONContentType, w.Header().Get("Content-Type")) + + var env jsonError + require.NoError(t, json.Unmarshal(w.Body.Bytes(), &env)) + assert.Equal(t, tc.wantType, env.Type) + assert.Equal(t, "boom", env.Message) + }) + } +} + +func TestGenerateJSONError_AppendsExceptionSuffixIdempotently(t *testing.T) { + cases := []struct { + name string + code string + wantType string + }{ + {"bare code gets suffix", "Validation", "ValidationException"}, + {"already-suffixed code unchanged", "ValidationException", "ValidationException"}, + } + for _, tc := range cases { + t.Run(tc.name, func(t *testing.T) { + body := generateJSONError(tc.code, "msg") + var env jsonError + require.NoError(t, json.Unmarshal(body, &env)) + assert.Equal(t, tc.wantType, env.Type) + assert.Equal(t, "msg", env.Message) + }) + } +} diff --git a/spinifex/gateway/bedrock/httperrors.go b/spinifex/gateway/bedrock/httperrors.go new file mode 100644 index 000000000..a185fbceb --- /dev/null +++ b/spinifex/gateway/bedrock/httperrors.go @@ -0,0 +1,24 @@ +package gateway_bedrock + +import ( + "net/http" + + "github.com/mulgadc/spinifex/spinifex/awserrors" +) + +// mapUpstreamStatus maps an upstream provider HTTP status to a bedrock +// awserrors code, shared by the Anthropic and vLLM adapters. +func mapUpstreamStatus(status int) string { + switch { + case status == http.StatusUnauthorized || status == http.StatusForbidden: + return awserrors.ErrorAccessDeniedException + case status == http.StatusTooManyRequests: + return awserrors.ErrorThrottlingException + case status == http.StatusBadRequest: + return awserrors.ErrorValidationException + case status >= 500: + return awserrors.ErrorServiceUnavailableException + default: + return awserrors.ErrorModelErrorException + } +} diff --git a/spinifex/gateway/bedrock/httperrors_test.go b/spinifex/gateway/bedrock/httperrors_test.go new file mode 100644 index 000000000..d95193a86 --- /dev/null +++ b/spinifex/gateway/bedrock/httperrors_test.go @@ -0,0 +1,31 @@ +package gateway_bedrock + +import ( + "net/http" + "testing" + + "github.com/mulgadc/spinifex/spinifex/awserrors" + "github.com/stretchr/testify/assert" +) + +func TestMapUpstreamStatus(t *testing.T) { + cases := []struct { + name string + status int + want string + }{ + {"401 unauthorized -> access denied", http.StatusUnauthorized, awserrors.ErrorAccessDeniedException}, + {"403 forbidden -> access denied", http.StatusForbidden, awserrors.ErrorAccessDeniedException}, + {"429 too many requests -> throttling", http.StatusTooManyRequests, awserrors.ErrorThrottlingException}, + {"400 bad request -> validation", http.StatusBadRequest, awserrors.ErrorValidationException}, + {"500 internal server error -> service unavailable", http.StatusInternalServerError, awserrors.ErrorServiceUnavailableException}, + {"503 service unavailable -> service unavailable", http.StatusServiceUnavailable, awserrors.ErrorServiceUnavailableException}, + {"418 teapot -> model error", http.StatusTeapot, awserrors.ErrorModelErrorException}, + {"other unmapped status -> model error", http.StatusConflict, awserrors.ErrorModelErrorException}, + } + for _, tc := range cases { + t.Run(tc.name, func(t *testing.T) { + assert.Equal(t, tc.want, mapUpstreamStatus(tc.status)) + }) + } +} diff --git a/spinifex/gateway/bedrock/provider.go b/spinifex/gateway/bedrock/provider.go new file mode 100644 index 000000000..e684566bd --- /dev/null +++ b/spinifex/gateway/bedrock/provider.go @@ -0,0 +1,86 @@ +package gateway_bedrock + +import ( + "context" + "errors" + "strings" + "time" + + "github.com/aws/aws-sdk-go/service/bedrockruntime" + "github.com/mulgadc/spinifex/spinifex/awserrors" +) + +// providerHTTPTimeout bounds outbound calls to both provider backends. It is +// long because Phase 1 is synchronous (no streaming) and matches the +// platform's tolerance for a single Converse call. +const providerHTTPTimeout = 15 * time.Minute + +// Provider translates a Converse request into a backend's native wire format +// and back. vllmProvider serves self-hosted models over OpenAI chat +// completions; anthropicProvider (via newAnthropicProvider) serves Claude +// over the Anthropic Messages API with a per-call API key baked in. +type Provider interface { + Converse(ctx context.Context, modelID string, input *bedrockruntime.ConverseInput) (*bedrockruntime.ConverseOutput, error) +} + +// Router resolves a modelId to its catalog entry and dispatches to the +// matching Provider, resolving self-host endpoints and provider credentials +// as needed. +type Router struct { + resolver CredentialResolver + endpointResolver EndpointResolver +} + +// NewRouter constructs a Router. A nil resolver or endpointResolver falls +// back to a resolver/resolver that finds nothing, so a Router is always safe +// to use even before the real stores are wired in. +func NewRouter(resolver CredentialResolver, endpointResolver EndpointResolver) *Router { + if resolver == nil { + resolver = NoopCredentialResolver + } + if endpointResolver == nil { + endpointResolver = NewStaticEndpointResolver(nil) + } + return &Router{resolver: resolver, endpointResolver: endpointResolver} +} + +// Converse routes modelID to its provider via the catalog. Unknown modelIds +// and unresolvable vendors return ResourceNotFoundException; a vendor with no +// resolvable credential returns AccessDeniedException. +func (rt *Router) Converse(ctx context.Context, accountID, modelID string, input *bedrockruntime.ConverseInput) (*bedrockruntime.ConverseOutput, error) { + entry, ok := lookupCatalogEntry(modelID) + if !ok { + return nil, errors.New(awserrors.ErrorResourceNotFoundException) + } + + var p Provider + switch { + case entry.Provider == tierSelfHost: + p = newVLLMProvider(rt.endpointResolver) + case strings.HasPrefix(entry.Provider, providerPrefix): + switch strings.TrimPrefix(entry.Provider, providerPrefix) { + case vendorAnthropic: + key, ok, err := rt.resolver.Resolve(ctx, accountID, vendorAnthropic) + if err != nil { + return nil, err + } + if !ok { + return nil, errors.New(awserrors.ErrorAccessDeniedException) + } + p = newAnthropicProvider(key) + default: + return nil, errors.New(awserrors.ErrorResourceNotFoundException) + } + default: + return nil, errors.New(awserrors.ErrorResourceNotFoundException) + } + + return p.Converse(ctx, modelID, input) +} + +// Converse is the bedrock-runtime Converse entry point used by the gateway +// route table. resolver and endpointResolver may be nil; NewRouter supplies +// no-op fallbacks. +func Converse(ctx context.Context, accountID, modelID string, input *bedrockruntime.ConverseInput, resolver CredentialResolver, endpointResolver EndpointResolver) (*bedrockruntime.ConverseOutput, error) { + return NewRouter(resolver, endpointResolver).Converse(ctx, accountID, modelID, input) +} diff --git a/spinifex/gateway/bedrock/router_test.go b/spinifex/gateway/bedrock/router_test.go new file mode 100644 index 000000000..91e3c4730 --- /dev/null +++ b/spinifex/gateway/bedrock/router_test.go @@ -0,0 +1,104 @@ +package gateway_bedrock + +import ( + "context" + "net/http" + "net/http/httptest" + "testing" + + "github.com/aws/aws-sdk-go/aws" + "github.com/aws/aws-sdk-go/service/bedrockruntime" + "github.com/mulgadc/spinifex/spinifex/awserrors" + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" +) + +// stubCredentialResolver returns a fixed (key, ok, err) for every Resolve +// call, letting tests drive the Router's anthropic branch without a real +// CredentialStore. +type stubCredentialResolver struct { + key string + ok bool + err error +} + +func (s stubCredentialResolver) Resolve(_ context.Context, _, _ string) (string, bool, error) { + return s.key, s.ok, s.err +} + +func converseInput() *bedrockruntime.ConverseInput { + return &bedrockruntime.ConverseInput{ + Messages: []*bedrockruntime.Message{ + {Role: aws.String(bedrockruntime.ConversationRoleUser), Content: []*bedrockruntime.ContentBlock{{Text: aws.String("hello")}}}, + }, + } +} + +func TestRouter_Converse_SelfHostSuccess(t *testing.T) { + ts := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + w.Header().Set("Content-Type", "application/json") + w.WriteHeader(http.StatusOK) + _, _ = w.Write([]byte(`{ + "choices": [{"message": {"role": "assistant", "content": "hi"}, "finish_reason": "stop"}], + "usage": {"prompt_tokens": 1, "completion_tokens": 1} + }`)) + })) + defer ts.Close() + + modelID := "meta.llama3-70b-instruct-v1:0" + rt := NewRouter(nil, NewStaticEndpointResolver(map[string]string{modelID: ts.URL})) + + out, err := rt.Converse(context.Background(), "000000000001", modelID, converseInput()) + require.NoError(t, err) + require.NotNil(t, out.Output.Message) + assert.Equal(t, "hi", *out.Output.Message.Content[0].Text) +} + +func TestRouter_Converse_UnknownModelReturnsResourceNotFound(t *testing.T) { + rt := NewRouter(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) + _, 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()) +} + +func TestConverse_PackageEntryPoint_SelfHostSuccess(t *testing.T) { + ts := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + w.Header().Set("Content-Type", "application/json") + w.WriteHeader(http.StatusOK) + _, _ = w.Write([]byte(`{ + "choices": [{"message": {"role": "assistant", "content": "via package Converse"}, "finish_reason": "stop"}], + "usage": {"prompt_tokens": 1, "completion_tokens": 1} + }`)) + })) + defer ts.Close() + + modelID := "meta.llama3-70b-instruct-v1:0" + out, err := Converse(context.Background(), "000000000001", modelID, converseInput(), nil, NewStaticEndpointResolver(map[string]string{modelID: ts.URL})) + 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) + require.Error(t, err) + assert.Equal(t, awserrors.ErrorResourceNotFoundException, err.Error()) +} + +func TestNewRouter_NilArgumentsFallBackToNoops(t *testing.T) { + rt := NewRouter(nil, nil) + require.NotNil(t, rt) + + // Self-host model with no endpoint resolver configured resolves nothing, + // so it must report ModelNotReady rather than panic on a nil resolver. + _, err := rt.Converse(context.Background(), "000000000001", "meta.llama3-70b-instruct-v1:0", converseInput()) + require.Error(t, err) + assert.Equal(t, awserrors.ErrorModelNotReadyException, err.Error()) +} diff --git a/spinifex/gateway/bedrock/vllm.go b/spinifex/gateway/bedrock/vllm.go new file mode 100644 index 000000000..4b76bd8cf --- /dev/null +++ b/spinifex/gateway/bedrock/vllm.go @@ -0,0 +1,224 @@ +package gateway_bedrock + +import ( + "bytes" + "context" + "encoding/json" + "errors" + "io" + "log/slog" + "net/http" + "strings" + "time" + + "github.com/aws/aws-sdk-go/aws" + "github.com/aws/aws-sdk-go/service/bedrockruntime" + "github.com/mulgadc/spinifex/spinifex/awserrors" +) + +const vllmChatCompletionsPath = "/v1/chat/completions" + +// EndpointResolver resolves a self-hosted model's OpenAI-compatible base URL. +// A later task backs this with the daemon's endpoint registry. +type EndpointResolver interface { + Endpoint(ctx context.Context, modelID string) (baseURL string, ok bool, err error) +} + +// staticEndpointResolver is a fixed modelID->baseURL map, for tests and +// config-driven single-endpoint setups. +type staticEndpointResolver map[string]string + +var _ EndpointResolver = staticEndpointResolver(nil) + +// NewStaticEndpointResolver returns an EndpointResolver backed by a fixed +// modelID->baseURL map. A nil or empty map resolves nothing. +func NewStaticEndpointResolver(endpoints map[string]string) EndpointResolver { + return staticEndpointResolver(endpoints) +} + +func (s staticEndpointResolver) Endpoint(_ context.Context, modelID string) (string, bool, error) { + baseURL, ok := s[modelID] + return baseURL, ok, nil +} + +type vllmMessage struct { + Role string `json:"role"` + Content string `json:"content"` +} + +type vllmChatRequest struct { + Model string `json:"model"` + Messages []vllmMessage `json:"messages"` + MaxTokens *int64 `json:"max_tokens,omitempty"` + Temperature *float64 `json:"temperature,omitempty"` + TopP *float64 `json:"top_p,omitempty"` + Stop []string `json:"stop,omitempty"` +} + +type vllmChoice struct { + Message vllmMessage `json:"message"` + FinishReason string `json:"finish_reason"` +} + +type vllmUsage struct { + PromptTokens int64 `json:"prompt_tokens"` + CompletionTokens int64 `json:"completion_tokens"` +} + +type vllmChatResponse struct { + Choices []vllmChoice `json:"choices"` + Usage vllmUsage `json:"usage"` +} + +// vllmFinishReasons maps OpenAI finish_reason values to Bedrock's StopReason +// enum. Unrecognised values fall back to end_turn. +var vllmFinishReasons = map[string]string{ + "stop": bedrockruntime.StopReasonEndTurn, + "length": bedrockruntime.StopReasonMaxTokens, +} + +func mapVLLMFinishReason(reason string) string { + if mapped, ok := vllmFinishReasons[reason]; ok { + return mapped + } + return bedrockruntime.StopReasonEndTurn +} + +// vllmProvider serves self-hosted open-weight models over the OpenAI +// chat-completions wire. httpClient and endpointResolver are injectable for +// tests. +type vllmProvider struct { + endpointResolver EndpointResolver + httpClient *http.Client +} + +var _ Provider = (*vllmProvider)(nil) + +func newVLLMProvider(endpointResolver EndpointResolver) *vllmProvider { + return &vllmProvider{ + endpointResolver: endpointResolver, + httpClient: &http.Client{Timeout: providerHTTPTimeout}, + } +} + +// Converse resolves modelID's endpoint, translates input to an OpenAI +// chat-completions request, calls the endpoint, and translates the response +// back to a Bedrock ConverseOutput. +func (p *vllmProvider) Converse(ctx context.Context, modelID string, input *bedrockruntime.ConverseInput) (*bedrockruntime.ConverseOutput, error) { + baseURL, ok, err := p.endpointResolver.Endpoint(ctx, modelID) + if err != nil { + slog.Error("vllm: endpoint resolution failed", "model", modelID, "err", err) + return nil, errors.New(awserrors.ErrorServiceUnavailableException) + } + if !ok { + return nil, errors.New(awserrors.ErrorModelNotReadyException) + } + + reqBody, err := json.Marshal(buildVLLMRequest(modelID, input)) + if err != nil { + slog.Error("vllm: failed to marshal request", "model", modelID, "err", err) + return nil, errors.New(awserrors.ErrorInternalError) + } + + httpReq, err := http.NewRequestWithContext(ctx, http.MethodPost, baseURL+vllmChatCompletionsPath, bytes.NewReader(reqBody)) + if err != nil { + slog.Error("vllm: failed to build request", "model", modelID, "err", err) + return nil, errors.New(awserrors.ErrorInternalError) + } + httpReq.Header.Set("content-type", "application/json") + + start := time.Now() + resp, err := p.httpClient.Do(httpReq) + latency := time.Since(start) + if err != nil { + slog.Error("vllm: request failed", "model", modelID, "endpoint", baseURL, "err", err) + return nil, errors.New(awserrors.ErrorServiceUnavailableException) + } + defer resp.Body.Close() + + respBody, err := io.ReadAll(resp.Body) + if err != nil { + slog.Error("vllm: failed to read response body", "model", modelID, "err", err) + return nil, errors.New(awserrors.ErrorServiceUnavailableException) + } + + if resp.StatusCode < 200 || resp.StatusCode >= 300 { + slog.Error("vllm: upstream error", "model", modelID, "status", resp.StatusCode, "body", string(respBody)) + return nil, errors.New(mapUpstreamStatus(resp.StatusCode)) + } + + var cr vllmChatResponse + if err := json.Unmarshal(respBody, &cr); err != nil { + slog.Error("vllm: failed to parse response", "model", modelID, "err", err) + return nil, errors.New(awserrors.ErrorModelErrorException) + } + + return vllmToConverseOutput(cr, latency) +} + +// buildVLLMRequest maps a Bedrock ConverseInput to an OpenAI chat-completions +// request. Non-text content blocks are dropped (Phase 1 is text-only). +func buildVLLMRequest(modelID string, input *bedrockruntime.ConverseInput) vllmChatRequest { + var messages []vllmMessage + + var systemParts []string + for _, s := range input.System { + if s != nil && s.Text != nil { + systemParts = append(systemParts, *s.Text) + } + } + if len(systemParts) > 0 { + messages = append(messages, vllmMessage{Role: "system", Content: strings.Join(systemParts, "\n")}) + } + + for _, m := range input.Messages { + if m == nil { + continue + } + var textParts []string + for _, c := range m.Content { + // TODO: Phase 1 is text-only; non-text content blocks are dropped. + if c == nil || c.Text == nil { + continue + } + textParts = append(textParts, *c.Text) + } + messages = append(messages, vllmMessage{Role: aws.StringValue(m.Role), Content: strings.Join(textParts, "\n")}) + } + + req := vllmChatRequest{Model: modelID, Messages: messages} + if input.InferenceConfig != nil { + req.MaxTokens = input.InferenceConfig.MaxTokens + req.Temperature = input.InferenceConfig.Temperature + req.TopP = input.InferenceConfig.TopP + req.Stop = aws.StringValueSlice(input.InferenceConfig.StopSequences) + } + return req +} + +// vllmToConverseOutput maps an OpenAI chat-completions response to a Bedrock +// ConverseOutput, measuring latency from the caller-supplied wall-clock delta. +func vllmToConverseOutput(cr vllmChatResponse, latency time.Duration) (*bedrockruntime.ConverseOutput, error) { + if len(cr.Choices) == 0 { + return nil, errors.New(awserrors.ErrorModelErrorException) + } + choice := cr.Choices[0] + inputTokens := cr.Usage.PromptTokens + outputTokens := cr.Usage.CompletionTokens + + return &bedrockruntime.ConverseOutput{ + Output: &bedrockruntime.ConverseOutput_{ + Message: &bedrockruntime.Message{ + Role: aws.String(bedrockruntime.ConversationRoleAssistant), + Content: []*bedrockruntime.ContentBlock{{Text: aws.String(choice.Message.Content)}}, + }, + }, + StopReason: aws.String(mapVLLMFinishReason(choice.FinishReason)), + Usage: &bedrockruntime.TokenUsage{ + InputTokens: aws.Int64(inputTokens), + OutputTokens: aws.Int64(outputTokens), + TotalTokens: aws.Int64(inputTokens + outputTokens), + }, + Metrics: &bedrockruntime.ConverseMetrics{LatencyMs: aws.Int64(latency.Milliseconds())}, + }, nil +} diff --git a/spinifex/gateway/bedrock/vllm_test.go b/spinifex/gateway/bedrock/vllm_test.go new file mode 100644 index 000000000..e6ed4441e --- /dev/null +++ b/spinifex/gateway/bedrock/vllm_test.go @@ -0,0 +1,73 @@ +package gateway_bedrock + +import ( + "context" + "encoding/json" + "net/http" + "net/http/httptest" + "testing" + + "github.com/aws/aws-sdk-go/aws" + "github.com/aws/aws-sdk-go/service/bedrockruntime" + "github.com/mulgadc/spinifex/spinifex/awserrors" + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" +) + +func TestVLLMProvider_Converse_MapsResponse(t *testing.T) { + var captured vllmChatRequest + + ts := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + require.NoError(t, json.NewDecoder(r.Body).Decode(&captured)) + w.Header().Set("Content-Type", "application/json") + w.WriteHeader(http.StatusOK) + _, _ = w.Write([]byte(`{ + "choices": [{"message": {"role": "assistant", "content": "hi there"}, "finish_reason": "stop"}], + "usage": {"prompt_tokens": 8, "completion_tokens": 3} + }`)) + })) + defer ts.Close() + + modelID := "meta.llama3-70b-instruct-v1:0" + p := newVLLMProvider(NewStaticEndpointResolver(map[string]string{modelID: ts.URL})) + p.httpClient = ts.Client() + + input := &bedrockruntime.ConverseInput{ + Messages: []*bedrockruntime.Message{ + {Role: aws.String(bedrockruntime.ConversationRoleUser), Content: []*bedrockruntime.ContentBlock{{Text: aws.String("hello")}}}, + }, + InferenceConfig: &bedrockruntime.InferenceConfiguration{MaxTokens: aws.Int64(128)}, + } + + out, err := p.Converse(context.Background(), modelID, input) + require.NoError(t, err) + + assert.Equal(t, modelID, captured.Model) + require.NotNil(t, captured.MaxTokens) + assert.Equal(t, int64(128), *captured.MaxTokens) + require.Len(t, captured.Messages, 1) + assert.Equal(t, "user", captured.Messages[0].Role) + assert.Equal(t, "hello", captured.Messages[0].Content) + + require.NotNil(t, out.Output.Message) + require.Len(t, out.Output.Message.Content, 1) + assert.Equal(t, "hi there", *out.Output.Message.Content[0].Text) + assert.Equal(t, int64(8), *out.Usage.InputTokens) + assert.Equal(t, int64(3), *out.Usage.OutputTokens) + assert.Equal(t, int64(11), *out.Usage.TotalTokens) + assert.Equal(t, bedrockruntime.StopReasonEndTurn, *out.StopReason) +} + +func TestVLLMProvider_Converse_UnresolvedEndpointReturnsModelNotReady(t *testing.T) { + p := newVLLMProvider(NewStaticEndpointResolver(nil)) + + input := &bedrockruntime.ConverseInput{ + Messages: []*bedrockruntime.Message{ + {Role: aws.String(bedrockruntime.ConversationRoleUser), Content: []*bedrockruntime.ContentBlock{{Text: aws.String("hello")}}}, + }, + } + + _, err := p.Converse(context.Background(), "meta.llama3-70b-instruct-v1:0", input) + require.Error(t, err) + assert.Equal(t, awserrors.ErrorModelNotReadyException, err.Error()) +} diff --git a/spinifex/gateway/bedrock_request_test.go b/spinifex/gateway/bedrock_request_test.go new file mode 100644 index 000000000..ddcf49417 --- /dev/null +++ b/spinifex/gateway/bedrock_request_test.go @@ -0,0 +1,166 @@ +package gateway + +import ( + "context" + "encoding/json" + "net/http" + "net/http/httptest" + "strings" + "testing" + + "github.com/aws/aws-sdk-go/service/bedrock" + "github.com/aws/aws-sdk-go/service/bedrockruntime" + "github.com/mulgadc/spinifex/spinifex/awserrors" + "github.com/mulgadc/spinifex/spinifex/testutil" + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" +) + +const bedrockTestAccount = "123456789012" +const bedrockTestLlamaModelID = "meta.llama3-70b-instruct-v1:0" + +// newVLLMStub stands up an httptest server that answers the OpenAI +// chat-completions wire, mirroring gateway/bedrock/vllm_test.go's stub. +func newVLLMStub(t *testing.T) *httptest.Server { + t.Helper() + ts := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + w.Header().Set("Content-Type", "application/json") + w.WriteHeader(http.StatusOK) + _, _ = w.Write([]byte(`{ + "choices": [{"message": {"role": "assistant", "content": "hi from vllm"}, "finish_reason": "stop"}], + "usage": {"prompt_tokens": 4, "completion_tokens": 2} + }`)) + })) + t.Cleanup(ts.Close) + return ts +} + +// newBedrockRequestGateway builds a GatewayConfig with a real NATS connection +// (satisfying Bedrock_Request/BedrockRuntime_Request's nil-check only — no +// NATS subject handling is required for these routes), no IAMService (so +// checkPolicy is a no-op), and BedrockEndpoints pinned at the given vLLM stub. +func newBedrockRequestGateway(t *testing.T, vllmURL string) *GatewayConfig { + t.Helper() + _, nc, _ := testutil.StartTestJetStream(t) + return &GatewayConfig{ + NATSConn: nc, + DisableLogging: true, + BedrockEndpoints: map[string]string{ + bedrockTestLlamaModelID: vllmURL, + }, + } +} + +func bedrockRequestWithAccount(method, path, body string) *http.Request { + req := httptest.NewRequest(method, path, strings.NewReader(body)) + ctx := context.WithValue(req.Context(), ctxAccountID, bedrockTestAccount) + return req.WithContext(ctx) +} + +func TestBedrockRequest_ListFoundationModels(t *testing.T) { + ts := newVLLMStub(t) + gw := newBedrockRequestGateway(t, ts.URL) + + req := bedrockRequestWithAccount(http.MethodGet, "/foundation-models", "") + w := httptest.NewRecorder() + require.NoError(t, gw.Bedrock_Request(w, req)) + require.Equal(t, http.StatusOK, w.Code) + + var out bedrock.ListFoundationModelsOutput + require.NoError(t, json.Unmarshal(w.Body.Bytes(), &out)) + + var ids []string + for _, m := range out.ModelSummaries { + ids = append(ids, *m.ModelId) + } + assert.Contains(t, ids, bedrockTestLlamaModelID) +} + +func TestBedrockRequest_GetFoundationModel(t *testing.T) { + ts := newVLLMStub(t) + gw := newBedrockRequestGateway(t, ts.URL) + + req := bedrockRequestWithAccount(http.MethodGet, "/foundation-models/"+bedrockTestLlamaModelID, "") + w := httptest.NewRecorder() + require.NoError(t, gw.Bedrock_Request(w, req)) + require.Equal(t, http.StatusOK, w.Code) + + var out bedrock.GetFoundationModelOutput + require.NoError(t, json.Unmarshal(w.Body.Bytes(), &out)) + require.NotNil(t, out.ModelDetails) + assert.Equal(t, bedrockTestLlamaModelID, *out.ModelDetails.ModelId) +} + +func TestBedrockRequest_GetFoundationModel_NotFound(t *testing.T) { + ts := newVLLMStub(t) + gw := newBedrockRequestGateway(t, ts.URL) + + req := bedrockRequestWithAccount(http.MethodGet, "/foundation-models/does.not-exist-v1:0", "") + w := httptest.NewRecorder() + err := gw.Bedrock_Request(w, req) + require.Error(t, err) + assert.Equal(t, awserrors.ErrorResourceNotFoundException, err.Error()) +} + +func TestBedrockRequest_UnknownRouteReturnsInvalidAction(t *testing.T) { + ts := newVLLMStub(t) + gw := newBedrockRequestGateway(t, ts.URL) + + req := bedrockRequestWithAccount(http.MethodDelete, "/foundation-models", "") + w := httptest.NewRecorder() + err := gw.Bedrock_Request(w, req) + require.Error(t, err) + assert.Equal(t, awserrors.ErrorInvalidAction, err.Error()) +} + +func TestBedrockRequest_MissingAccountIDReturnsServerInternal(t *testing.T) { + ts := newVLLMStub(t) + gw := newBedrockRequestGateway(t, ts.URL) + + req := httptest.NewRequest(http.MethodGet, "/foundation-models", nil) + w := httptest.NewRecorder() + err := gw.Bedrock_Request(w, req) + require.Error(t, err) + assert.Equal(t, awserrors.ErrorServerInternal, err.Error()) +} + +func TestBedrockRuntimeRequest_Converse(t *testing.T) { + ts := newVLLMStub(t) + gw := newBedrockRequestGateway(t, ts.URL) + + body := `{"messages":[{"role":"user","content":[{"text":"hello"}]}]}` + req := bedrockRequestWithAccount(http.MethodPost, "/model/"+bedrockTestLlamaModelID+"/converse", body) + w := httptest.NewRecorder() + require.NoError(t, gw.BedrockRuntime_Request(w, req)) + require.Equal(t, http.StatusOK, w.Code) + + var out bedrockruntime.ConverseOutput + require.NoError(t, json.Unmarshal(w.Body.Bytes(), &out)) + require.NotNil(t, out.Output) + require.NotNil(t, out.Output.Message) + assert.Equal(t, bedrockruntime.ConversationRoleAssistant, *out.Output.Message.Role) + require.Len(t, out.Output.Message.Content, 1) + assert.Equal(t, "hi from vllm", *out.Output.Message.Content[0].Text) +} + +func TestBedrockRuntimeRequest_MalformedBodyReturnsValidationException(t *testing.T) { + ts := newVLLMStub(t) + gw := newBedrockRequestGateway(t, ts.URL) + + req := bedrockRequestWithAccount(http.MethodPost, "/model/"+bedrockTestLlamaModelID+"/converse", "{not-json") + w := httptest.NewRecorder() + err := gw.BedrockRuntime_Request(w, req) + require.Error(t, err) + assert.Equal(t, awserrors.ErrorValidationException, err.Error()) +} + +func TestBedrockRuntimeRequest_MissingAccountIDReturnsServerInternal(t *testing.T) { + ts := newVLLMStub(t) + gw := newBedrockRequestGateway(t, ts.URL) + + req := httptest.NewRequest(http.MethodPost, "/model/"+bedrockTestLlamaModelID+"/converse", strings.NewReader("{}")) + w := httptest.NewRecorder() + err := gw.BedrockRuntime_Request(w, req) + require.Error(t, err) + assert.Equal(t, awserrors.ErrorServerInternal, err.Error()) +} diff --git a/spinifex/gateway/bedrock_test.go b/spinifex/gateway/bedrock_test.go new file mode 100644 index 000000000..9a52cc718 --- /dev/null +++ b/spinifex/gateway/bedrock_test.go @@ -0,0 +1,48 @@ +package gateway + +import ( + "testing" + + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" +) + +func TestLookupBedrockAction_ResolvesKnownRoutes(t *testing.T) { + cases := []struct { + method, path string + wantAction string + wantParams []string + }{ + {"GET", "/foundation-models", "ListFoundationModels", nil}, + {"GET", "/foundation-models/anthropic.claude-3-5-sonnet-20240620-v1:0", "GetFoundationModel", []string{"anthropic.claude-3-5-sonnet-20240620-v1:0"}}, + } + for _, tc := range cases { + t.Run(tc.method+"_"+tc.path, func(t *testing.T) { + action, params, handler, ok := lookupBedrockAction(tc.method, tc.path) + require.True(t, ok, "expected route to match for %s %s", tc.method, tc.path) + require.NotNil(t, handler) + assert.Equal(t, tc.wantAction, action) + assert.Equal(t, tc.wantParams, params) + }) + } +} + +func TestLookupBedrockAction_UnknownReturnsFalse(t *testing.T) { + _, _, handler, ok := lookupBedrockAction("DELETE", "/foundation-models") + assert.False(t, ok) + assert.Nil(t, handler) +} + +func TestLookupBedrockRuntimeAction_ResolvesConverse(t *testing.T) { + action, params, handler, ok := lookupBedrockRuntimeAction("POST", "/model/meta.llama3-70b-instruct-v1:0/converse") + require.True(t, ok) + require.NotNil(t, handler) + assert.Equal(t, "Converse", action) + assert.Equal(t, []string{"meta.llama3-70b-instruct-v1:0"}, params) +} + +func TestLookupBedrockRuntimeAction_UnknownReturnsFalse(t *testing.T) { + _, _, handler, ok := lookupBedrockRuntimeAction("GET", "/model/foo/converse") + assert.False(t, ok) + assert.Nil(t, handler) +} diff --git a/spinifex/gateway/bedrockruntime.go b/spinifex/gateway/bedrockruntime.go new file mode 100644 index 000000000..e9d0c1bf0 --- /dev/null +++ b/spinifex/gateway/bedrockruntime.go @@ -0,0 +1,113 @@ +package gateway + +import ( + "context" + "encoding/json" + "errors" + "io" + "log/slog" + "net/http" + "net/url" + "regexp" + + "github.com/aws/aws-sdk-go/service/bedrockruntime" + "github.com/mulgadc/spinifex/spinifex/awserrors" + gateway_bedrock "github.com/mulgadc/spinifex/spinifex/gateway/bedrock" +) + +// bedrockRuntimeRoute maps one HTTP method + path regex to an AWS action and handler. +type bedrockRuntimeRoute struct { + method string + pattern *regexp.Regexp + action string + handler bedrockRuntimeRouteHandler +} + +// bedrockRuntimeRouteHandler invokes a per-action bedrock-runtime (data-plane) +// gateway function. params holds the regex capture groups, PathUnescape'd. +// resolver is gw.bedrockResolver() (credential store or no-op); endpoints is +// gw.bedrockEndpointResolver() over the configured pinned self-host endpoints. +type bedrockRuntimeRouteHandler func(ctx context.Context, accountID string, params []string, body []byte, resolver gateway_bedrock.CredentialResolver, endpoints gateway_bedrock.EndpointResolver) (any, error) + +// bedrockRuntimeRoutes is the dispatch table. +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) (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) + }}, +} + +// lookupBedrockRuntimeAction matches method+path against bedrockRuntimeRoutes, +// returning the action, path params, and handler, or ("", nil, nil, false) on +// no match. path must be r.URL.EscapedPath(): captured params are +// PathUnescape'd before returning, mirroring lookupEKSAction. +func lookupBedrockRuntimeAction(method, path string) (string, []string, bedrockRuntimeRouteHandler, bool) { + for _, route := range bedrockRuntimeRoutes { + if route.method != method { + continue + } + m := route.pattern.FindStringSubmatch(path) + if m == nil { + continue + } + var params []string + if len(m) > 1 { + params = make([]string, 0, len(m)-1) + for _, raw := range m[1:] { + decoded, err := url.PathUnescape(raw) + if err != nil { + slog.Debug("bedrock-runtime: bad percent-encoding in path param", "param", raw, "err", err) + decoded = raw + } + params = append(params, decoded) + } + } + return route.action, params, route.handler, true + } + return "", nil, nil, false +} + +// BedrockRuntime_Request dispatches bedrock-runtime (data-plane) REST-JSON +// requests: resolves method+path to an action, reads the body, calls the +// handler, and serialises the output as JSON. +func (gw *GatewayConfig) BedrockRuntime_Request(w http.ResponseWriter, r *http.Request) error { + action, params, handler, ok := lookupBedrockRuntimeAction(r.Method, r.URL.EscapedPath()) + if !ok { + slog.DebugContext(r.Context(), "bedrock-runtime: no route for request", "method", r.Method, "path", r.URL.Path) + return errors.New(awserrors.ErrorInvalidAction) + } + + if err := gw.checkPolicy(r, "bedrock-runtime", action); err != nil { + return err + } + + if gw.NATSConn == nil { + return errors.New(awserrors.ErrorServerInternal) + } + + accountID, _ := r.Context().Value(ctxAccountID).(string) + if accountID == "" { + slog.ErrorContext(r.Context(), "BedrockRuntime_Request: no account ID in auth context") + return errors.New(awserrors.ErrorServerInternal) + } + + body, err := io.ReadAll(r.Body) + if err != nil { + slog.ErrorContext(r.Context(), "BedrockRuntime_Request: failed to read body", "err", err) + return errors.New(awserrors.ErrorInvalidParameterValue) + } + + output, err := handler(r.Context(), accountID, params, body, gw.bedrockResolver(), gw.bedrockEndpointResolver()) + if err != nil { + return err + } + + gateway_bedrock.WriteJSONResponse(w, output) + return nil +} diff --git a/spinifex/gateway/gateway.go b/spinifex/gateway/gateway.go index 02fe41fe0..561decc78 100644 --- a/spinifex/gateway/gateway.go +++ b/spinifex/gateway/gateway.go @@ -20,6 +20,7 @@ import ( "github.com/mulgadc/predastore/pkg/iampolicy" "github.com/mulgadc/predastore/ratelimit" "github.com/mulgadc/spinifex/spinifex/awserrors" + gateway_bedrock "github.com/mulgadc/spinifex/spinifex/gateway/bedrock" gateway_ecr "github.com/mulgadc/spinifex/spinifex/gateway/ecr" gateway_ecrauth "github.com/mulgadc/spinifex/spinifex/gateway/ecrauth" "github.com/mulgadc/spinifex/spinifex/gateway/policy" @@ -110,6 +111,14 @@ type GatewayConfig struct { // in unit tests of unrelated routes). ECRTokenIssuer *gateway_ecrauth.Issuer ECRTokenVerifier *gateway_ecrauth.Verifier + // BedrockCredentials resolves per-account provider API keys for bedrock + // routes. Nil falls back to no external providers (self-host models only). + BedrockCredentials *gateway_bedrock.CredentialStore + // BedrockEndpoints maps a self-hosted modelId to its OpenAI-compatible base + // URL. Phase 1 endpoints are pinned/always-resident, so this is static + // config; a later phase backs it with the daemon's dynamic endpoint + // registry. Nil/empty means no self-hosted models are reachable. + BedrockEndpoints map[string]string } var supportedServices = map[string]bool{ @@ -123,6 +132,8 @@ var supportedServices = map[string]bool{ "acm": true, "tagging": true, "spinifex": true, + "bedrock": true, + "bedrock-runtime": true, } // EC2ErrorResponse is the EC2 query-API error envelope. @@ -340,6 +351,10 @@ func (gw *GatewayConfig) Request(w http.ResponseWriter, r *http.Request) { err = gw.ELBv2_Request(w, r) case "eks": err = gw.EKS_Request(w, r) + case "bedrock": + err = gw.Bedrock_Request(w, r) + case "bedrock-runtime": + err = gw.BedrockRuntime_Request(w, r) case "ecs": err = gw.ECS_Request(w, r) case "ecr": @@ -501,8 +516,9 @@ func (gw *GatewayConfig) ErrorHandler(w http.ResponseWriter, r *http.Request, er errorMsg.HTTPCode = 500 } - // EKS, ECR, ACM, ECS, and tagging use AWS JSON 1.1; query/XML services fall through. - if svc == "eks" || svc == "ecr" || svc == "acm" || svc == "ecs" || svc == "tagging" { + // EKS, ECR, ACM, ECS, tagging, and bedrock/bedrock-runtime use AWS JSON 1.1; + // query/XML services fall through. + if svc == "eks" || svc == "ecr" || svc == "acm" || svc == "ecs" || svc == "tagging" || svc == "bedrock" || svc == "bedrock-runtime" { body := GenerateEKSErrorResponse(err.Error(), errorMsg.Message, requestId) slog.Debug("Generated JSON error response", "service", svc, "error", err.Error(), "json", string(body), "requestId", requestId) w.Header().Set("Content-Type", eksJSONContentType) diff --git a/spinifex/services/awsgw/awsgw.go b/spinifex/services/awsgw/awsgw.go index 696c2e505..d25e70e53 100644 --- a/spinifex/services/awsgw/awsgw.go +++ b/spinifex/services/awsgw/awsgw.go @@ -18,6 +18,7 @@ import ( "github.com/mulgadc/spinifex/spinifex/admin" "github.com/mulgadc/spinifex/spinifex/config" "github.com/mulgadc/spinifex/spinifex/gateway" + gateway_bedrock "github.com/mulgadc/spinifex/spinifex/gateway/bedrock" gateway_ecr "github.com/mulgadc/spinifex/spinifex/gateway/ecr" gateway_ecrauth "github.com/mulgadc/spinifex/spinifex/gateway/ecrauth" "github.com/mulgadc/spinifex/spinifex/handlers/ecr" @@ -283,24 +284,40 @@ func launchService(config *config.ClusterConfig) error { registryHost = nodeConfig.AdvertiseIP } + // Bedrock provider credentials: per-account keys live in the + // bedrock-credentials KV bucket; OCHRE_ANTHROPIC_API_KEY seeds an optional + // platform-wide default so accounts work before onboarding their own key. + bedrockPlatformDefaults := map[string]string{} + if key := os.Getenv("OCHRE_ANTHROPIC_API_KEY"); key != "" { + bedrockPlatformDefaults["anthropic"] = key + } + bedrockCredentials := gateway_bedrock.NewCredentialStore(js, masterKey, len(config.Nodes), bedrockPlatformDefaults) + + // Bedrock self-host endpoints: Phase 1 models are pinned, so their + // OpenAI-compatible base URLs come from static config. OCHRE_VLLM_ENDPOINTS + // is a comma-separated list of modelId=baseURL pairs. + bedrockEndpoints := parseBedrockEndpoints(os.Getenv("OCHRE_VLLM_ENDPOINTS")) + gw := gateway.GatewayConfig{ - Debug: nodeConfig.AWSGW.Debug, - DisableLogging: false, - NATSConn: natsConn, - Config: nodeConfig.AWSGW.Config, - ExpectedNodes: len(config.Nodes), - Region: nodeConfig.Region, - InternalSuffix: config.AWS.InternalSuffix, - RegistryPort: registryPort, - RegistryHost: registryHost, - AZ: nodeConfig.AZ, - IAMService: iamService, - STSService: stsService, - Version: version, - Commit: commit, - ECRRegistry: ecrRegistry, - ECRTokenIssuer: gateway_ecrauth.NewIssuer(signingKey, ecrAudience), - ECRTokenVerifier: gateway_ecrauth.NewVerifier(verifyKeys, ecrAudience), + Debug: nodeConfig.AWSGW.Debug, + DisableLogging: false, + NATSConn: natsConn, + Config: nodeConfig.AWSGW.Config, + ExpectedNodes: len(config.Nodes), + Region: nodeConfig.Region, + InternalSuffix: config.AWS.InternalSuffix, + RegistryPort: registryPort, + RegistryHost: registryHost, + AZ: nodeConfig.AZ, + IAMService: iamService, + STSService: stsService, + Version: version, + Commit: commit, + ECRRegistry: ecrRegistry, + ECRTokenIssuer: gateway_ecrauth.NewIssuer(signingKey, ecrAudience), + ECRTokenVerifier: gateway_ecrauth.NewVerifier(verifyKeys, ecrAudience), + BedrockCredentials: bedrockCredentials, + BedrockEndpoints: bedrockEndpoints, } // Rotate the ECR signing key on a 30-day cadence, retaining the previous keys @@ -448,3 +465,28 @@ func findBootstrapFile(baseDir string) string { func isConcreteRegistryHost(host string) bool { return host != "" && host != "0.0.0.0" && host != "::" } + +// parseBedrockEndpoints parses a comma-separated list of modelId=baseURL pairs +// into a map for the bedrock self-host endpoint resolver. Malformed or empty +// entries are skipped. Returns nil for empty input. +func parseBedrockEndpoints(raw string) map[string]string { + raw = strings.TrimSpace(raw) + if raw == "" { + return nil + } + endpoints := map[string]string{} + for pair := range strings.SplitSeq(raw, ",") { + modelID, baseURL, ok := strings.Cut(pair, "=") + modelID = strings.TrimSpace(modelID) + baseURL = strings.TrimSpace(baseURL) + if !ok || modelID == "" || baseURL == "" { + slog.Warn("awsgw: skipping malformed OCHRE_VLLM_ENDPOINTS entry", "entry", pair) + continue + } + endpoints[modelID] = baseURL + } + if len(endpoints) == 0 { + return nil + } + return endpoints +} diff --git a/spinifex/services/awsgw/awsgw_bedrock_test.go b/spinifex/services/awsgw/awsgw_bedrock_test.go new file mode 100644 index 000000000..508172991 --- /dev/null +++ b/spinifex/services/awsgw/awsgw_bedrock_test.go @@ -0,0 +1,60 @@ +package awsgw + +import ( + "testing" + + "github.com/stretchr/testify/assert" +) + +func TestParseBedrockEndpoints(t *testing.T) { + cases := []struct { + name string + raw string + want map[string]string + }{ + {"empty input returns nil", "", nil}, + {"whitespace-only input returns nil", " ", nil}, + { + "single pair", + "meta.llama3-70b-instruct-v1:0=http://vllm.internal:8000", + map[string]string{"meta.llama3-70b-instruct-v1:0": "http://vllm.internal:8000"}, + }, + { + "multiple pairs", + "meta.llama3-70b-instruct-v1:0=http://vllm1:8000,mistral.mixtral-8x7b-v1:0=http://vllm2:8000", + map[string]string{ + "meta.llama3-70b-instruct-v1:0": "http://vllm1:8000", + "mistral.mixtral-8x7b-v1:0": "http://vllm2:8000", + }, + }, + { + "whitespace around pairs and entries trimmed", + " meta.llama3-70b-instruct-v1:0 = http://vllm.internal:8000 , mistral.mixtral-8x7b-v1:0=http://vllm2:8000 ", + map[string]string{ + "meta.llama3-70b-instruct-v1:0": "http://vllm.internal:8000", + "mistral.mixtral-8x7b-v1:0": "http://vllm2:8000", + }, + }, + { + "malformed entry with no equals sign is skipped", + "no-equals,meta.llama3-70b-instruct-v1:0=http://vllm.internal:8000", + map[string]string{"meta.llama3-70b-instruct-v1:0": "http://vllm.internal:8000"}, + }, + { + "malformed entry with empty modelId is skipped", + "=http://vllm.internal:8000,meta.llama3-70b-instruct-v1:0=http://vllm.internal:8000", + map[string]string{"meta.llama3-70b-instruct-v1:0": "http://vllm.internal:8000"}, + }, + { + "malformed entry with empty baseURL is skipped", + "meta.llama3-70b-instruct-v1:0=,mistral.mixtral-8x7b-v1:0=http://vllm2:8000", + map[string]string{"mistral.mixtral-8x7b-v1:0": "http://vllm2:8000"}, + }, + {"all entries malformed returns nil", "no-equals,=u,m=", nil}, + } + for _, tc := range cases { + t.Run(tc.name, func(t *testing.T) { + assert.Equal(t, tc.want, parseBedrockEndpoints(tc.raw)) + }) + } +} diff --git a/tests/e2e/Makefile b/tests/e2e/Makefile index 58faa0a2e..ca10092f8 100644 --- a/tests/e2e/Makefile +++ b/tests/e2e/Makefile @@ -1,4 +1,4 @@ -PKG ?= ./cert/... ./lb/... ./eks/... ./ecs/... ./iam/... ./ecr/... ./multinode/... ./baremetal/... ./reboot/... ./tofu/... +PKG ?= ./cert/... ./lb/... ./eks/... ./ecs/... ./iam/... ./ecr/... ./bedrock/... ./multinode/... ./baremetal/... ./reboot/... ./tofu/... ARTIFACT_DIR ?= /tmp/spinifex-e2e-artifacts TIMEOUT ?= 30m MODE ?= single @@ -11,9 +11,9 @@ SUITE_TAG ?= $(notdir $(patsubst %/...,%,$(firstword $(PKG)))) # named suite into BIN_DIR/.test so CI can ship binaries (not source + # go toolchain) to the remote VM. BIN_DIR ?= $(CURDIR)/_bin -SUITES_BUILD ?= cert lb eks ecs single iam ecr gpu natuplink +SUITES_BUILD ?= cert lb eks ecs single iam ecr bedrock gpu natuplink -.PHONY: e2e e2e-cert e2e-lb e2e-eks e2e-ecs e2e-single e2e-iam e2e-ecr e2e-multinode e2e-bm e2e-reboot e2e-tofu e2e-gpu e2e-natuplink e2e-all clean e2e-build +.PHONY: e2e e2e-cert e2e-lb e2e-eks e2e-ecs e2e-single e2e-iam e2e-ecr e2e-bedrock e2e-multinode e2e-bm e2e-reboot e2e-tofu e2e-gpu e2e-natuplink e2e-all clean e2e-build ifeq ($(GOTESTSUM),1) E2E_CMD = gotestsum \ @@ -39,6 +39,7 @@ e2e-ecs: ; $(MAKE) e2e PKG=./ecs/... e2e-single: ; $(MAKE) e2e PKG=./single/... e2e-iam: ; $(MAKE) e2e PKG=./iam/... e2e-ecr: ; $(MAKE) e2e PKG=./ecr/... +e2e-bedrock: ; $(MAKE) e2e PKG=./bedrock/... e2e-multinode: ; $(MAKE) e2e PKG=./multinode/... MODE=multinode e2e-bm: ; $(MAKE) e2e PKG=./baremetal/... MODE=baremetal e2e-reboot: ; $(MAKE) e2e PKG=./reboot/... diff --git a/tests/e2e/bedrock/bedrock_test.go b/tests/e2e/bedrock/bedrock_test.go new file mode 100644 index 000000000..15d57b7e1 --- /dev/null +++ b/tests/e2e/bedrock/bedrock_test.go @@ -0,0 +1,130 @@ +//go:build e2e + +package bedrock + +import ( + "os" + "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/tests/e2e/harness" + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" +) + +// Catalog model IDs mirror the static Phase-1 catalog in gateway_bedrock. +const ( + modelSelfHostLlama = "meta.llama3-70b-instruct-v1:0" + modelAnthropic = "anthropic.claude-3-5-sonnet-20240620-v1:0" + modelBogus = "does.not-exist-v1:0" +) + +// TestListFoundationModels confirms the self-host model is always advertised; +// the Anthropic model is only asserted on when present, since its advertisement +// depends on cluster credential configuration. +func TestListFoundationModels(t *testing.T) { + f := requireBedrockFixture(t) + + out, err := f.AWS.Bedrock.ListFoundationModels(&bedrock.ListFoundationModelsInput{}) + require.NoError(t, err, "list-foundation-models") + + ids := make(map[string]*bedrock.FoundationModelSummary, len(out.ModelSummaries)) + for _, m := range out.ModelSummaries { + ids[aws.StringValue(m.ModelId)] = m + } + + assert.Contains(t, ids, modelSelfHostLlama, "self-host model must always be advertised") + + if anthropic, ok := ids[modelAnthropic]; ok { + assert.Equal(t, "Anthropic", aws.StringValue(anthropic.ProviderName), + "anthropic model advertised with unexpected provider name") + } +} + +// TestGetFoundationModel round-trips a known model and asserts +// ResourceNotFoundException for an unknown one. +func TestGetFoundationModel(t *testing.T) { + f := requireBedrockFixture(t) + + t.Run("known model", func(t *testing.T) { + out, err := f.AWS.Bedrock.GetFoundationModel(&bedrock.GetFoundationModelInput{ + ModelIdentifier: aws.String(modelSelfHostLlama), + }) + require.NoError(t, err, "get-foundation-model %s", modelSelfHostLlama) + require.NotNil(t, out.ModelDetails, "empty ModelDetails") + assert.Equal(t, modelSelfHostLlama, aws.StringValue(out.ModelDetails.ModelId)) + }) + + t.Run("unknown model", func(t *testing.T) { + harness.ExpectError(t, "ResourceNotFoundException", func() error { + _, e := f.AWS.Bedrock.GetFoundationModel(&bedrock.GetFoundationModelInput{ + ModelIdentifier: aws.String(modelBogus), + }) + return e + }) + }) +} + +// TestConverseSelfHost exercises the real vLLM Converse path. Gated behind +// OCHRE_E2E_SELFHOST=1 — it needs a GPU-backed vLLM endpoint configured on the +// gateway via OCHRE_VLLM_ENDPOINTS at boot. +func TestConverseSelfHost(t *testing.T) { + if os.Getenv("OCHRE_E2E_SELFHOST") != "1" { + t.Skip("OCHRE_E2E_SELFHOST!=1; skipping real self-host Converse") + } + f := requireBedrockFixture(t) + + out, err := f.AWS.BedrockRuntime.Converse(&bedrockruntime.ConverseInput{ + ModelId: aws.String(modelSelfHostLlama), + Messages: []*bedrockruntime.Message{{ + Role: aws.String(bedrockruntime.ConversationRoleUser), + Content: []*bedrockruntime.ContentBlock{{Text: aws.String("Say hello in one word.")}}, + }}, + InferenceConfig: &bedrockruntime.InferenceConfiguration{MaxTokens: aws.Int64(64)}, + }) + require.NoError(t, err, "converse self-host") + require.NotNil(t, out.Output, "nil output") + require.NotNil(t, out.Output.Message, "nil output message") + assert.NotEmpty(t, converseText(out.Output.Message), "empty assistant text") + require.NotNil(t, out.Usage, "nil usage") + assert.Greater(t, aws.Int64Value(out.Usage.OutputTokens), int64(0), "OutputTokens must be > 0") +} + +// TestConverseAnthropic exercises the real Anthropic-direct Converse path. +// Gated behind OCHRE_E2E_ANTHROPIC=1 — it needs a live Anthropic API key +// configured on the cluster (OCHRE_ANTHROPIC_API_KEY or a per-account KV entry). +func TestConverseAnthropic(t *testing.T) { + if os.Getenv("OCHRE_E2E_ANTHROPIC") != "1" { + t.Skip("OCHRE_E2E_ANTHROPIC!=1; skipping real Anthropic Converse") + } + f := requireBedrockFixture(t) + + out, err := f.AWS.BedrockRuntime.Converse(&bedrockruntime.ConverseInput{ + ModelId: aws.String(modelAnthropic), + Messages: []*bedrockruntime.Message{{ + Role: aws.String(bedrockruntime.ConversationRoleUser), + Content: []*bedrockruntime.ContentBlock{{Text: aws.String("Say hello in one word.")}}, + }}, + InferenceConfig: &bedrockruntime.InferenceConfiguration{MaxTokens: aws.Int64(64)}, + }) + require.NoError(t, err, "converse anthropic") + require.NotNil(t, out.Output, "nil output") + require.NotNil(t, out.Output.Message, "nil output message") + assert.NotEmpty(t, converseText(out.Output.Message), "empty assistant text") + require.NotNil(t, out.Usage, "nil usage") + assert.Greater(t, aws.Int64Value(out.Usage.InputTokens), int64(0), "InputTokens must be > 0") + assert.Greater(t, aws.Int64Value(out.Usage.OutputTokens), int64(0), "OutputTokens must be > 0") +} + +// converseText concatenates the text content blocks of a Converse output message. +func converseText(msg *bedrockruntime.Message) string { + var text string + for _, c := range msg.Content { + if c != nil && c.Text != nil { + text += *c.Text + } + } + return text +} diff --git a/tests/e2e/bedrock/main_test.go b/tests/e2e/bedrock/main_test.go new file mode 100644 index 000000000..d92bc2a8d --- /dev/null +++ b/tests/e2e/bedrock/main_test.go @@ -0,0 +1,61 @@ +//go:build e2e + +// Package bedrock is the Ochre (AWS-Bedrock) E2E suite. Hermetic API-shape and +// provider-translation parity (Anthropic Messages / vLLM OpenAI wire formats) +// is covered by the gateway_bedrock unit tests; this suite exercises +// ListFoundationModels/GetFoundationModel against any running cluster and +// gates real Converse calls on backend availability (self-host GPU or a live +// Anthropic key), so CI parity holds without requiring either to be present. +package bedrock + +import ( + "os" + "sync" + "testing" + + "github.com/mulgadc/spinifex/tests/e2e/harness" +) + +// Package-scoped singleton fixture, built lazily by requireBedrockFixture and +// torn down by TestMain after the run. +var ( + pkgFixOnce sync.Once + pkgFix *Fixture + pkgFixErr error +) + +// Fixture carries per-process state shared across every Test* in this package. +type Fixture struct { + Env *harness.Env + AWS *harness.AWSClient + Account string +} + +func TestMain(m *testing.M) { + os.Exit(m.Run()) +} + +// requireBedrockFixture returns the package-scoped Fixture, building it on +// first call. Skips the calling test if SPINIFEX_E2E is unset. +func requireBedrockFixture(t *testing.T) *Fixture { + t.Helper() + pkgFixOnce.Do(func() { + if os.Getenv("SPINIFEX_E2E") == "" { + return + } + env := harness.LoadEnv(t) + awsCli := harness.NewAWSClient(t, env) + pkgFix = &Fixture{ + Env: env, + AWS: awsCli, + Account: harness.IAMAccountID(t, awsCli), + } + }) + if pkgFixErr != nil { + t.Fatalf("bedrock fixture init failed: %v", pkgFixErr) + } + if pkgFix == nil { + t.Skip("bedrock fixture unavailable (SPINIFEX_E2E unset)") + } + return pkgFix +} diff --git a/tests/e2e/harness/aws.go b/tests/e2e/harness/aws.go index 2b2d0da24..5aa807dfc 100644 --- a/tests/e2e/harness/aws.go +++ b/tests/e2e/harness/aws.go @@ -14,6 +14,8 @@ import ( "github.com/aws/aws-sdk-go/aws/credentials" "github.com/aws/aws-sdk-go/aws/session" "github.com/aws/aws-sdk-go/service/acm" + "github.com/aws/aws-sdk-go/service/bedrock" + "github.com/aws/aws-sdk-go/service/bedrockruntime" "github.com/aws/aws-sdk-go/service/ec2" "github.com/aws/aws-sdk-go/service/ec2/ec2iface" "github.com/aws/aws-sdk-go/service/ecr" @@ -30,15 +32,17 @@ type AWSClient struct { // EC2 is wrapped so RunInstances auto-serialises against cluster capacity // (see capacityRetryEC2). EC2Conf is the underlying SDK client shared by // both; kept for call sites that read or mutate its Config. - EC2 ec2iface.EC2API - EC2Conf *ec2.EC2 - ELBv2 *elbv2.ELBV2 - IAM *iam.IAM - STS *sts.STS - EKS *eks.EKS - ACM *acm.ACM - ECR *ecr.ECR - ECS *ecs.ECS + EC2 ec2iface.EC2API + EC2Conf *ec2.EC2 + ELBv2 *elbv2.ELBV2 + IAM *iam.IAM + STS *sts.STS + EKS *eks.EKS + ACM *acm.ACM + ECR *ecr.ECR + ECS *ecs.ECS + Bedrock *bedrock.Bedrock + BedrockRuntime *bedrockruntime.BedrockRuntime } // NewAWSClient builds clients pointed at the spinifex gateway using the @@ -128,15 +132,17 @@ func newAWSClient(t *testing.T, env *Env, accessKey, secretKey, sessionToken str rawEC2 := ec2.New(sess) return &AWSClient{ - EC2: &capacityRetryEC2{EC2API: rawEC2}, - EC2Conf: rawEC2, - ELBv2: elbv2.New(sess), - IAM: iam.New(sess), - STS: sts.New(sess), - EKS: eks.New(sess), - ACM: acm.New(sess), - ECR: ecr.New(sess), - ECS: ecs.New(sess), + EC2: &capacityRetryEC2{EC2API: rawEC2}, + EC2Conf: rawEC2, + ELBv2: elbv2.New(sess), + IAM: iam.New(sess), + STS: sts.New(sess), + EKS: eks.New(sess), + ACM: acm.New(sess), + ECR: ecr.New(sess), + ECS: ecs.New(sess), + Bedrock: bedrock.New(sess), + BedrockRuntime: bedrockruntime.New(sess), } }