mirror of
https://github.com/mulgadc/spinifex.git
synced 2026-10-15 09:01:31 +00:00
feat(bedrock): Ochre runtime MVP — Converse over self-host + Anthropic-direct (#520)
* 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.
This commit is contained in:
@@ -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]
|
||||
|
||||
@@ -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."},
|
||||
}
|
||||
|
||||
@@ -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)
|
||||
|
||||
@@ -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)
|
||||
}
|
||||
@@ -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<N>:<M>" 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()),
|
||||
},
|
||||
}
|
||||
}
|
||||
@@ -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))
|
||||
})
|
||||
}
|
||||
}
|
||||
@@ -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:<vendor>" 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:<vendor>"
|
||||
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
|
||||
}
|
||||
@@ -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())
|
||||
}
|
||||
@@ -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{}
|
||||
@@ -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)
|
||||
}
|
||||
@@ -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)
|
||||
}
|
||||
}
|
||||
@@ -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)
|
||||
})
|
||||
}
|
||||
}
|
||||
@@ -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
|
||||
}
|
||||
}
|
||||
@@ -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))
|
||||
})
|
||||
}
|
||||
}
|
||||
@@ -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)
|
||||
}
|
||||
@@ -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())
|
||||
}
|
||||
@@ -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
|
||||
}
|
||||
@@ -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())
|
||||
}
|
||||
@@ -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())
|
||||
}
|
||||
@@ -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)
|
||||
}
|
||||
@@ -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
|
||||
}
|
||||
@@ -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)
|
||||
|
||||
@@ -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
|
||||
}
|
||||
|
||||
@@ -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))
|
||||
})
|
||||
}
|
||||
}
|
||||
+4
-3
@@ -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/<suite>.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/...
|
||||
|
||||
@@ -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
|
||||
}
|
||||
@@ -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
|
||||
}
|
||||
+24
-18
@@ -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),
|
||||
}
|
||||
}
|
||||
|
||||
|
||||
Reference in New Issue
Block a user