diff --git a/spinifex/gateway/bedrock/handler.go b/spinifex/gateway/bedrock/handler.go index fe827b4e3..51787261d 100644 --- a/spinifex/gateway/bedrock/handler.go +++ b/spinifex/gateway/bedrock/handler.go @@ -56,6 +56,17 @@ func WriteJSONResponse(w http.ResponseWriter, obj any) { } } +// WriteRawResponse writes body verbatim with contentType and a 200 status, +// bypassing JSON marshaling. Used by InvokeModel, whose response is the +// provider-native body rather than a struct. +func WriteRawResponse(w http.ResponseWriter, body []byte, contentType string) { + w.Header().Set("Content-Type", contentType) + w.WriteHeader(http.StatusOK) + if _, err := w.Write(body); err != nil { + slog.Error("Failed to write bedrock raw 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) { diff --git a/spinifex/gateway/bedrock/invoke.go b/spinifex/gateway/bedrock/invoke.go new file mode 100644 index 000000000..d1872cdbb --- /dev/null +++ b/spinifex/gateway/bedrock/invoke.go @@ -0,0 +1,80 @@ +package gateway_bedrock + +import ( + "context" + "errors" + "strings" + + "github.com/mulgadc/spinifex/spinifex/awserrors" +) + +// InvokeAdapter translates a Bedrock InvokeModel raw request body into a +// backend's native wire format and back, returning the response bytes +// verbatim with their content-type. Unlike Provider (Converse), the wire +// shape is per-family and is not unified by the gateway. +type InvokeAdapter interface { + InvokeModel(ctx context.Context, modelID string, body []byte) (respBody []byte, contentType string, err error) +} + +// InvokeRouter resolves a modelId to its catalog entry and dispatches to the +// matching InvokeAdapter, resolving self-host endpoints and provider +// credentials as needed. +type InvokeRouter struct { + resolver CredentialResolver + endpointResolver EndpointResolver +} + +// NewInvokeRouter constructs an InvokeRouter. A nil resolver or +// endpointResolver falls back to a resolver/resolver that finds nothing, so +// an InvokeRouter is always safe to use even before the real stores are +// wired in. +func NewInvokeRouter(resolver CredentialResolver, endpointResolver EndpointResolver) *InvokeRouter { + if resolver == nil { + resolver = NoopCredentialResolver + } + if endpointResolver == nil { + endpointResolver = NewStaticEndpointResolver(nil) + } + return &InvokeRouter{resolver: resolver, endpointResolver: endpointResolver} +} + +// InvokeModel routes modelID to its family adapter via the catalog. Unknown +// modelIds and unresolvable vendors return ResourceNotFoundException; a +// vendor with no resolvable credential returns AccessDeniedException. +func (rt *InvokeRouter) InvokeModel(ctx context.Context, accountID, modelID string, body []byte) ([]byte, string, error) { + entry, ok := lookupCatalogEntry(modelID) + if !ok { + return nil, "", errors.New(awserrors.ErrorResourceNotFoundException) + } + + var a InvokeAdapter + switch { + case entry.Provider == tierSelfHost: + a = newLlamaInvokeAdapter(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) + } + a = newAnthropicInvokeAdapter(key) + default: + return nil, "", errors.New(awserrors.ErrorResourceNotFoundException) + } + default: + return nil, "", errors.New(awserrors.ErrorResourceNotFoundException) + } + + return a.InvokeModel(ctx, modelID, body) +} + +// InvokeModel is the bedrock-runtime InvokeModel entry point used by the +// gateway route table. resolver and endpointResolver may be nil; +// NewInvokeRouter supplies no-op fallbacks. +func InvokeModel(ctx context.Context, accountID, modelID string, body []byte, resolver CredentialResolver, endpointResolver EndpointResolver) ([]byte, string, error) { + return NewInvokeRouter(resolver, endpointResolver).InvokeModel(ctx, accountID, modelID, body) +} diff --git a/spinifex/gateway/bedrock/invoke_anthropic.go b/spinifex/gateway/bedrock/invoke_anthropic.go new file mode 100644 index 000000000..62ebcc6e6 --- /dev/null +++ b/spinifex/gateway/bedrock/invoke_anthropic.go @@ -0,0 +1,102 @@ +package gateway_bedrock + +import ( + "bytes" + "context" + "encoding/json" + "errors" + "io" + "log/slog" + "net/http" + + "github.com/mulgadc/spinifex/spinifex/awserrors" +) + +// anthropicInvokeAdapter forwards Bedrock-native Claude InvokeModel bodies to +// the Anthropic Messages API almost as-is. httpClient and baseURL are +// injectable for tests. +type anthropicInvokeAdapter struct { + httpClient *http.Client + baseURL string +} + +func newAnthropicInvokeAdapterClient() *anthropicInvokeAdapter { + return &anthropicInvokeAdapter{ + httpClient: &http.Client{Timeout: providerHTTPTimeout}, + baseURL: anthropicDefaultBaseURL, + } +} + +// boundAnthropicInvokeAdapter adapts anthropicInvokeAdapter to InvokeAdapter +// by baking in a resolved per-account (or platform-default) API key. +type boundAnthropicInvokeAdapter struct { + inner *anthropicInvokeAdapter + apiKey string +} + +var _ InvokeAdapter = (*boundAnthropicInvokeAdapter)(nil) + +func (b *boundAnthropicInvokeAdapter) InvokeModel(ctx context.Context, modelID string, body []byte) ([]byte, string, error) { + return b.inner.InvokeModel(ctx, modelID, body, b.apiKey) +} + +// newAnthropicInvokeAdapter returns an InvokeAdapter that forwards to the +// Anthropic Messages API with apiKey. +func newAnthropicInvokeAdapter(apiKey string) InvokeAdapter { + return &boundAnthropicInvokeAdapter{inner: newAnthropicInvokeAdapterClient(), apiKey: apiKey} +} + +// InvokeModel rewrites the Bedrock Claude InvokeModel body (drops the +// Bedrock-only anthropic_version field, injects model since Anthropic's API +// carries it in the body rather than the URL) and posts it to the Anthropic +// Messages API, returning the response verbatim. +func (a *anthropicInvokeAdapter) InvokeModel(ctx context.Context, modelID string, body []byte, apiKey string) ([]byte, string, error) { + var fields map[string]json.RawMessage + if err := json.Unmarshal(body, &fields); err != nil { + slog.Error("anthropic invoke: failed to parse request body", "model", modelID, "err", err) + return nil, "", errors.New(awserrors.ErrorValidationException) + } + delete(fields, "anthropic_version") + + modelJSON, err := json.Marshal(anthropicModelID(modelID)) + if err != nil { + slog.Error("anthropic invoke: failed to marshal model id", "model", modelID, "err", err) + return nil, "", errors.New(awserrors.ErrorInternalError) + } + fields["model"] = modelJSON + + reqBody, err := json.Marshal(fields) + if err != nil { + slog.Error("anthropic invoke: failed to marshal request", "model", modelID, "err", err) + return nil, "", errors.New(awserrors.ErrorInternalError) + } + + httpReq, err := http.NewRequestWithContext(ctx, http.MethodPost, a.baseURL+anthropicMessagesPath, bytes.NewReader(reqBody)) + if err != nil { + slog.Error("anthropic invoke: 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") + + resp, err := a.httpClient.Do(httpReq) + if err != nil { + slog.Error("anthropic invoke: 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 invoke: 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 invoke: upstream error", "model", modelID, "status", resp.StatusCode, "body", string(respBody)) + return nil, "", errors.New(mapUpstreamStatus(resp.StatusCode)) + } + + return respBody, "application/json", nil +} diff --git a/spinifex/gateway/bedrock/invoke_anthropic_test.go b/spinifex/gateway/bedrock/invoke_anthropic_test.go new file mode 100644 index 000000000..6c9e452a9 --- /dev/null +++ b/spinifex/gateway/bedrock/invoke_anthropic_test.go @@ -0,0 +1,97 @@ +package gateway_bedrock + +import ( + "context" + "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 TestAnthropicInvokeAdapter_InvokeModel_RewritesRequestAndReturnsBodyVerbatim(t *testing.T) { + canned := `{"id":"msg_123","content":[{"type":"text","text":"Hello there"}],"stop_reason":"end_turn"}` + + var captured map[string]json.RawMessage + var capturedAPIKey, capturedAnthropicVersion, capturedContentType string + + ts := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + capturedAPIKey = r.Header.Get("X-Api-Key") + capturedAnthropicVersion = r.Header.Get("Anthropic-Version") + capturedContentType = r.Header.Get("Content-Type") + require.NoError(t, json.NewDecoder(r.Body).Decode(&captured)) + w.Header().Set("Content-Type", "application/json") + w.WriteHeader(http.StatusOK) + _, _ = w.Write([]byte(canned)) + })) + defer ts.Close() + + a := &anthropicInvokeAdapter{httpClient: ts.Client(), baseURL: ts.URL} + + reqBody := []byte(`{"anthropic_version":"bedrock-2023-05-31","max_tokens":256,"messages":[{"role":"user","content":[{"type":"text","text":"hi"}]}]}`) + + respBody, contentType, err := a.InvokeModel(context.Background(), "anthropic.claude-3-5-sonnet-20240620-v1:0", reqBody, "sk-test-key") + require.NoError(t, err) + + assert.Equal(t, "sk-test-key", capturedAPIKey) + assert.Equal(t, anthropicAPIVersion, capturedAnthropicVersion) + assert.Equal(t, "application/json", capturedContentType) + + _, hasVersion := captured["anthropic_version"] + assert.False(t, hasVersion, "anthropic_version must be dropped from the forwarded request") + + var model string + require.NoError(t, json.Unmarshal(captured["model"], &model)) + assert.Equal(t, "claude-3-5-sonnet-20240620", model) + + require.Contains(t, captured, "max_tokens") + require.Contains(t, captured, "messages") + + assert.Equal(t, "application/json", contentType) + assert.JSONEq(t, canned, string(respBody)) +} + +func TestAnthropicInvokeAdapter_InvokeModel_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() + + a := &anthropicInvokeAdapter{httpClient: ts.Client(), baseURL: ts.URL} + + _, _, err := a.InvokeModel(context.Background(), "anthropic.claude-3-5-sonnet-20240620-v1:0", []byte(`{"anthropic_version":"bedrock-2023-05-31","messages":[]}`), "bad-key") + require.Error(t, err) + assert.Equal(t, awserrors.ErrorAccessDeniedException, err.Error()) +} + +// TestBoundAnthropicInvokeAdapter_InvokeModel exercises newAnthropicInvokeAdapter +// and boundAnthropicInvokeAdapter.InvokeModel (the InvokeAdapter used by +// InvokeRouter for the "provider:anthropic" tier). newAnthropicInvokeAdapter +// hardcodes the real Anthropic base URL, 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 TestBoundAnthropicInvokeAdapter_InvokeModel(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"}`)) + })) + defer ts.Close() + + adapter := newAnthropicInvokeAdapter("sk-bound-test") + bound, ok := adapter.(*boundAnthropicInvokeAdapter) + require.True(t, ok) + require.NotNil(t, bound.inner) + bound.inner.httpClient = ts.Client() + bound.inner.baseURL = ts.URL + + respBody, contentType, err := adapter.InvokeModel(context.Background(), "anthropic.claude-3-5-sonnet-20240620-v1:0", []byte(`{"anthropic_version":"bedrock-2023-05-31","messages":[]}`)) + require.NoError(t, err) + assert.Equal(t, "application/json", contentType) + assert.Contains(t, string(respBody), "bound response") +} diff --git a/spinifex/gateway/bedrock/invoke_llama.go b/spinifex/gateway/bedrock/invoke_llama.go new file mode 100644 index 000000000..ba4313151 --- /dev/null +++ b/spinifex/gateway/bedrock/invoke_llama.go @@ -0,0 +1,166 @@ +package gateway_bedrock + +import ( + "bytes" + "context" + "encoding/json" + "errors" + "io" + "log/slog" + "net/http" + + "github.com/mulgadc/spinifex/spinifex/awserrors" +) + +const llamaCompletionsPath = "/v1/completions" + +// llamaInvokeRequest is the Bedrock-native Llama InvokeModel request body. +type llamaInvokeRequest struct { + Prompt string `json:"prompt"` + MaxGenLen int `json:"max_gen_len,omitempty"` + Temperature *float64 `json:"temperature,omitempty"` + TopP *float64 `json:"top_p,omitempty"` +} + +// llamaInvokeResponse is the Bedrock-native Llama InvokeModel response body. +type llamaInvokeResponse struct { + Generation string `json:"generation"` + PromptTokenCount int `json:"prompt_token_count"` + GenerationTokenCount int `json:"generation_token_count"` + StopReason string `json:"stop_reason"` +} + +// llamaCompletionsRequest is the OpenAI /v1/completions request vLLM serves. +type llamaCompletionsRequest struct { + Model string `json:"model"` + Prompt string `json:"prompt"` + MaxTokens int `json:"max_tokens,omitempty"` + Temperature *float64 `json:"temperature,omitempty"` + TopP *float64 `json:"top_p,omitempty"` +} + +type llamaCompletionsChoice struct { + Text string `json:"text"` + FinishReason string `json:"finish_reason"` +} + +type llamaCompletionsUsage struct { + PromptTokens int `json:"prompt_tokens"` + CompletionTokens int `json:"completion_tokens"` +} + +type llamaCompletionsResponse struct { + Choices []llamaCompletionsChoice `json:"choices"` + Usage llamaCompletionsUsage `json:"usage"` +} + +// llamaFinishReasons maps OpenAI finish_reason values to Bedrock Llama's +// stop_reason values. Unrecognised values fall back to "stop". +var llamaFinishReasons = map[string]string{ + "stop": "stop", + "length": "length", +} + +func mapLlamaStopReason(reason string) string { + if mapped, ok := llamaFinishReasons[reason]; ok { + return mapped + } + return "stop" +} + +// llamaInvokeAdapter serves Bedrock-native Llama InvokeModel requests over a +// self-hosted vLLM OpenAI-completions endpoint. httpClient and +// endpointResolver are injectable for tests. +type llamaInvokeAdapter struct { + endpointResolver EndpointResolver + httpClient *http.Client +} + +var _ InvokeAdapter = (*llamaInvokeAdapter)(nil) + +func newLlamaInvokeAdapter(endpointResolver EndpointResolver) *llamaInvokeAdapter { + return &llamaInvokeAdapter{ + endpointResolver: endpointResolver, + httpClient: &http.Client{Timeout: providerHTTPTimeout}, + } +} + +// InvokeModel resolves modelID's endpoint, translates the Bedrock-native +// Llama request to an OpenAI completions request, calls the endpoint, and +// translates the response back to the Bedrock Llama shape. +func (a *llamaInvokeAdapter) InvokeModel(ctx context.Context, modelID string, body []byte) ([]byte, string, error) { + baseURL, ok, err := a.endpointResolver.Endpoint(ctx, modelID) + if err != nil { + slog.Error("llama invoke: endpoint resolution failed", "model", modelID, "err", err) + return nil, "", errors.New(awserrors.ErrorServiceUnavailableException) + } + if !ok { + return nil, "", errors.New(awserrors.ErrorModelNotReadyException) + } + + var lr llamaInvokeRequest + if err := json.Unmarshal(body, &lr); err != nil { + slog.Error("llama invoke: failed to parse request body", "model", modelID, "err", err) + return nil, "", errors.New(awserrors.ErrorValidationException) + } + + reqBody, err := json.Marshal(llamaCompletionsRequest{ + Model: modelID, + Prompt: lr.Prompt, + MaxTokens: lr.MaxGenLen, + Temperature: lr.Temperature, + TopP: lr.TopP, + }) + if err != nil { + slog.Error("llama invoke: failed to marshal request", "model", modelID, "err", err) + return nil, "", errors.New(awserrors.ErrorInternalError) + } + + httpReq, err := http.NewRequestWithContext(ctx, http.MethodPost, baseURL+llamaCompletionsPath, bytes.NewReader(reqBody)) + if err != nil { + slog.Error("llama invoke: failed to build request", "model", modelID, "err", err) + return nil, "", errors.New(awserrors.ErrorInternalError) + } + httpReq.Header.Set("Content-Type", "application/json") + + resp, err := a.httpClient.Do(httpReq) + if err != nil { + slog.Error("llama invoke: 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("llama invoke: failed to read response body", "model", modelID, "err", err) + return nil, "", errors.New(awserrors.ErrorServiceUnavailableException) + } + + if resp.StatusCode < 200 || resp.StatusCode >= 300 { + slog.Error("llama invoke: upstream error", "model", modelID, "status", resp.StatusCode, "body", string(respBody)) + return nil, "", errors.New(mapUpstreamStatus(resp.StatusCode)) + } + + var cr llamaCompletionsResponse + if err := json.Unmarshal(respBody, &cr); err != nil { + slog.Error("llama invoke: failed to parse response", "model", modelID, "err", err) + return nil, "", errors.New(awserrors.ErrorModelErrorException) + } + if len(cr.Choices) == 0 { + return nil, "", errors.New(awserrors.ErrorModelErrorException) + } + choice := cr.Choices[0] + + out, err := json.Marshal(llamaInvokeResponse{ + Generation: choice.Text, + PromptTokenCount: cr.Usage.PromptTokens, + GenerationTokenCount: cr.Usage.CompletionTokens, + StopReason: mapLlamaStopReason(choice.FinishReason), + }) + if err != nil { + slog.Error("llama invoke: failed to marshal response", "model", modelID, "err", err) + return nil, "", errors.New(awserrors.ErrorInternalError) + } + + return out, "application/json", nil +} diff --git a/spinifex/gateway/bedrock/invoke_llama_test.go b/spinifex/gateway/bedrock/invoke_llama_test.go new file mode 100644 index 000000000..621beb560 --- /dev/null +++ b/spinifex/gateway/bedrock/invoke_llama_test.go @@ -0,0 +1,101 @@ +package gateway_bedrock + +import ( + "context" + "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 TestLlamaInvokeAdapter_InvokeModel_MapsRequestAndResponse(t *testing.T) { + var captured llamaCompletionsRequest + + 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": [{"text": "hi there", "finish_reason": "stop"}], + "usage": {"prompt_tokens": 8, "completion_tokens": 3} + }`)) + })) + defer ts.Close() + + modelID := "meta.llama3-70b-instruct-v1:0" + a := newLlamaInvokeAdapter(NewStaticEndpointResolver(map[string]string{modelID: ts.URL})) + a.httpClient = ts.Client() + + temp := 0.5 + reqBody, err := json.Marshal(llamaInvokeRequest{Prompt: "hello", MaxGenLen: 128, Temperature: &temp}) + require.NoError(t, err) + + respBody, contentType, err := a.InvokeModel(context.Background(), modelID, reqBody) + require.NoError(t, err) + assert.Equal(t, "application/json", contentType) + + assert.Equal(t, modelID, captured.Model) + assert.Equal(t, "hello", captured.Prompt) + assert.Equal(t, 128, captured.MaxTokens) + require.NotNil(t, captured.Temperature) + assert.Equal(t, 0.5, *captured.Temperature) + + var out llamaInvokeResponse + require.NoError(t, json.Unmarshal(respBody, &out)) + assert.Equal(t, "hi there", out.Generation) + assert.Equal(t, 8, out.PromptTokenCount) + assert.Equal(t, 3, out.GenerationTokenCount) + assert.Equal(t, "stop", out.StopReason) +} + +func TestLlamaInvokeAdapter_InvokeModel_MapsLengthFinishReason(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": [{"text": "truncated", "finish_reason": "length"}], + "usage": {"prompt_tokens": 1, "completion_tokens": 1} + }`)) + })) + defer ts.Close() + + modelID := "meta.llama3-70b-instruct-v1:0" + a := newLlamaInvokeAdapter(NewStaticEndpointResolver(map[string]string{modelID: ts.URL})) + a.httpClient = ts.Client() + + respBody, _, err := a.InvokeModel(context.Background(), modelID, []byte(`{"prompt":"hello"}`)) + require.NoError(t, err) + + var out llamaInvokeResponse + require.NoError(t, json.Unmarshal(respBody, &out)) + assert.Equal(t, "length", out.StopReason) +} + +func TestLlamaInvokeAdapter_InvokeModel_EmptyChoicesReturnsModelError(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": [], "usage": {"prompt_tokens": 1, "completion_tokens": 0}}`)) + })) + defer ts.Close() + + modelID := "meta.llama3-70b-instruct-v1:0" + a := newLlamaInvokeAdapter(NewStaticEndpointResolver(map[string]string{modelID: ts.URL})) + a.httpClient = ts.Client() + + _, _, err := a.InvokeModel(context.Background(), modelID, []byte(`{"prompt":"hello"}`)) + require.Error(t, err) + assert.Equal(t, awserrors.ErrorModelErrorException, err.Error()) +} + +func TestLlamaInvokeAdapter_InvokeModel_UnresolvedEndpointReturnsModelNotReady(t *testing.T) { + a := newLlamaInvokeAdapter(NewStaticEndpointResolver(nil)) + + _, _, err := a.InvokeModel(context.Background(), "meta.llama3-70b-instruct-v1:0", []byte(`{"prompt":"hello"}`)) + require.Error(t, err) + assert.Equal(t, awserrors.ErrorModelNotReadyException, err.Error()) +} diff --git a/spinifex/gateway/bedrock/invoke_test.go b/spinifex/gateway/bedrock/invoke_test.go new file mode 100644 index 000000000..31b4b5142 --- /dev/null +++ b/spinifex/gateway/bedrock/invoke_test.go @@ -0,0 +1,84 @@ +package gateway_bedrock + +import ( + "context" + "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 TestInvokeRouter_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": [{"text": "hi", "finish_reason": "stop"}], + "usage": {"prompt_tokens": 1, "completion_tokens": 1} + }`)) + })) + defer ts.Close() + + modelID := "meta.llama3-70b-instruct-v1:0" + rt := NewInvokeRouter(nil, NewStaticEndpointResolver(map[string]string{modelID: ts.URL})) + + respBody, contentType, err := rt.InvokeModel(context.Background(), "000000000001", modelID, []byte(`{"prompt":"hello"}`)) + require.NoError(t, err) + assert.Equal(t, "application/json", contentType) + + var out llamaInvokeResponse + require.NoError(t, json.Unmarshal(respBody, &out)) + assert.Equal(t, "hi", out.Generation) +} + +func TestInvokeRouter_UnknownModelReturnsResourceNotFound(t *testing.T) { + rt := NewInvokeRouter(nil, nil) + _, _, err := rt.InvokeModel(context.Background(), "000000000001", "does.not-exist-v1:0", []byte(`{}`)) + require.Error(t, err) + assert.Equal(t, awserrors.ErrorResourceNotFoundException, err.Error()) +} + +func TestInvokeRouter_AnthropicNoCredentialReturnsAccessDenied(t *testing.T) { + rt := NewInvokeRouter(stubCredentialResolver{ok: false}, nil) + _, _, err := rt.InvokeModel(context.Background(), "000000000001", "anthropic.claude-3-5-sonnet-20240620-v1:0", []byte(`{}`)) + require.Error(t, err) + assert.Equal(t, awserrors.ErrorAccessDeniedException, err.Error()) +} + +func TestInvokeRouter_SelfHostNoEndpointReturnsModelNotReady(t *testing.T) { + rt := NewInvokeRouter(nil, nil) + _, _, err := rt.InvokeModel(context.Background(), "000000000001", "meta.llama3-70b-instruct-v1:0", []byte(`{"prompt":"hello"}`)) + require.Error(t, err) + assert.Equal(t, awserrors.ErrorModelNotReadyException, err.Error()) +} + +func TestInvokeModel_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": [{"text": "via package InvokeModel", "finish_reason": "stop"}], + "usage": {"prompt_tokens": 1, "completion_tokens": 1} + }`)) + })) + defer ts.Close() + + modelID := "meta.llama3-70b-instruct-v1:0" + respBody, contentType, err := InvokeModel(context.Background(), "000000000001", modelID, []byte(`{"prompt":"hello"}`), nil, NewStaticEndpointResolver(map[string]string{modelID: ts.URL})) + require.NoError(t, err) + assert.Equal(t, "application/json", contentType) + + var out llamaInvokeResponse + require.NoError(t, json.Unmarshal(respBody, &out)) + assert.Equal(t, "via package InvokeModel", out.Generation) +} + +func TestInvokeModel_PackageEntryPoint_UnknownModel(t *testing.T) { + _, _, err := InvokeModel(context.Background(), "000000000001", "does.not-exist-v1:0", []byte(`{}`), nil, nil) + require.Error(t, err) + assert.Equal(t, awserrors.ErrorResourceNotFoundException, err.Error()) +} diff --git a/spinifex/gateway/bedrock_request_test.go b/spinifex/gateway/bedrock_request_test.go index ddcf49417..af445691c 100644 --- a/spinifex/gateway/bedrock_request_test.go +++ b/spinifex/gateway/bedrock_request_test.go @@ -35,6 +35,23 @@ func newVLLMStub(t *testing.T) *httptest.Server { return ts } +// newLlamaCompletionsStub stands up an httptest server that answers the +// OpenAI completions wire, mirroring gateway/bedrock/invoke_llama_test.go's +// stub, for the InvokeModel self-host path. +func newLlamaCompletionsStub(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": [{"text": "hi from llama invoke", "finish_reason": "stop"}], + "usage": {"prompt_tokens": 5, "completion_tokens": 3} + }`)) + })) + 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 @@ -154,6 +171,25 @@ func TestBedrockRuntimeRequest_MalformedBodyReturnsValidationException(t *testin assert.Equal(t, awserrors.ErrorValidationException, err.Error()) } +func TestBedrockRuntimeRequest_InvokeModel(t *testing.T) { + ts := newLlamaCompletionsStub(t) + gw := newBedrockRequestGateway(t, ts.URL) + + body := `{"prompt":"hello","max_gen_len":128}` + req := bedrockRequestWithAccount(http.MethodPost, "/model/"+bedrockTestLlamaModelID+"/invoke", body) + w := httptest.NewRecorder() + require.NoError(t, gw.BedrockRuntime_Request(w, req)) + require.Equal(t, http.StatusOK, w.Code) + assert.Equal(t, "application/json", w.Header().Get("Content-Type")) + + assert.JSONEq(t, `{ + "generation": "hi from llama invoke", + "prompt_token_count": 5, + "generation_token_count": 3, + "stop_reason": "stop" + }`, w.Body.String()) +} + func TestBedrockRuntimeRequest_MissingAccountIDReturnsServerInternal(t *testing.T) { ts := newVLLMStub(t) gw := newBedrockRequestGateway(t, ts.URL) diff --git a/spinifex/gateway/bedrockruntime.go b/spinifex/gateway/bedrockruntime.go index e9d0c1bf0..95143fb55 100644 --- a/spinifex/gateway/bedrockruntime.go +++ b/spinifex/gateway/bedrockruntime.go @@ -29,7 +29,9 @@ type bedrockRuntimeRoute struct { // 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. +// bedrockRuntimeRoutes is the dispatch table. InvokeModel has no handler +// function here: BedrockRuntime_Request special-cases its action to bypass +// the JSON-marshaling dispatch below, since its response is raw bytes. var bedrockRuntimeRoutes = []bedrockRuntimeRoute{ {"POST", regexp.MustCompile(`^/model/([^/]+)/converse$`), "Converse", func(ctx context.Context, acct string, p []string, b []byte, resolver gateway_bedrock.CredentialResolver, endpoints gateway_bedrock.EndpointResolver) (any, error) { @@ -41,6 +43,7 @@ var bedrockRuntimeRoutes = []bedrockRuntimeRoute{ } return gateway_bedrock.Converse(ctx, acct, p[0], input, resolver, endpoints) }}, + {"POST", regexp.MustCompile(`^/model/([^/]+)/invoke$`), "InvokeModel", nil}, } // lookupBedrockRuntimeAction matches method+path against bedrockRuntimeRoutes, @@ -103,6 +106,17 @@ func (gw *GatewayConfig) BedrockRuntime_Request(w http.ResponseWriter, r *http.R return errors.New(awserrors.ErrorInvalidParameterValue) } + // InvokeModel returns provider-native bytes, not a struct WriteJSONResponse + // could marshal, so it writes its own response body directly. + if action == "InvokeModel" { + respBody, contentType, err := gateway_bedrock.InvokeModel(r.Context(), accountID, params[0], body, gw.bedrockResolver(), gw.bedrockEndpointResolver()) + if err != nil { + return err + } + gateway_bedrock.WriteRawResponse(w, respBody, contentType) + return nil + } + output, err := handler(r.Context(), accountID, params, body, gw.bedrockResolver(), gw.bedrockEndpointResolver()) if err != nil { return err diff --git a/tests/e2e/bedrock/bedrock_test.go b/tests/e2e/bedrock/bedrock_test.go index 15d57b7e1..ffee8260b 100644 --- a/tests/e2e/bedrock/bedrock_test.go +++ b/tests/e2e/bedrock/bedrock_test.go @@ -3,6 +3,7 @@ package bedrock import ( + "encoding/json" "os" "testing" @@ -118,6 +119,77 @@ func TestConverseAnthropic(t *testing.T) { assert.Greater(t, aws.Int64Value(out.Usage.OutputTokens), int64(0), "OutputTokens must be > 0") } +// TestInvokeModelSelfHost exercises the real vLLM InvokeModel path with the +// Bedrock-native Llama request/response shape. Gated behind +// OCHRE_E2E_SELFHOST=1 — it needs a GPU-backed vLLM endpoint configured on the +// gateway via OCHRE_VLLM_ENDPOINTS at boot. +func TestInvokeModelSelfHost(t *testing.T) { + if os.Getenv("OCHRE_E2E_SELFHOST") != "1" { + t.Skip("OCHRE_E2E_SELFHOST!=1; skipping real self-host InvokeModel") + } + f := requireBedrockFixture(t) + + reqBody, err := json.Marshal(map[string]any{ + "prompt": "Say hello in one word.", + "max_gen_len": 64, + }) + require.NoError(t, err, "marshal invoke-model request") + + out, err := f.AWS.BedrockRuntime.InvokeModel(&bedrockruntime.InvokeModelInput{ + ModelId: aws.String(modelSelfHostLlama), + Body: reqBody, + ContentType: aws.String("application/json"), + }) + require.NoError(t, err, "invoke-model self-host") + assert.Equal(t, "application/json", aws.StringValue(out.ContentType)) + + var respBody struct { + Generation string `json:"generation"` + GenerationTokenCount int `json:"generation_token_count"` + } + require.NoError(t, json.Unmarshal(out.Body, &respBody), "unmarshal invoke-model response") + assert.NotEmpty(t, respBody.Generation, "empty generation") + assert.Greater(t, respBody.GenerationTokenCount, 0, "GenerationTokenCount must be > 0") +} + +// TestInvokeModelAnthropic exercises the real Anthropic-direct InvokeModel +// path with the Bedrock-native Claude Messages request/response shape. 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 TestInvokeModelAnthropic(t *testing.T) { + if os.Getenv("OCHRE_E2E_ANTHROPIC") != "1" { + t.Skip("OCHRE_E2E_ANTHROPIC!=1; skipping real Anthropic InvokeModel") + } + f := requireBedrockFixture(t) + + reqBody, err := json.Marshal(map[string]any{ + "anthropic_version": "bedrock-2023-05-31", + "max_tokens": 64, + "messages": []map[string]any{ + {"role": "user", "content": "Say hello in one word."}, + }, + }) + require.NoError(t, err, "marshal invoke-model request") + + out, err := f.AWS.BedrockRuntime.InvokeModel(&bedrockruntime.InvokeModelInput{ + ModelId: aws.String(modelAnthropic), + Body: reqBody, + ContentType: aws.String("application/json"), + }) + require.NoError(t, err, "invoke-model anthropic") + assert.Equal(t, "application/json", aws.StringValue(out.ContentType)) + + var respBody struct { + Content []struct { + Type string `json:"type"` + Text string `json:"text"` + } `json:"content"` + } + require.NoError(t, json.Unmarshal(out.Body, &respBody), "unmarshal invoke-model response") + require.NotEmpty(t, respBody.Content, "empty content") + assert.NotEmpty(t, respBody.Content[0].Text, "empty assistant text") +} + // converseText concatenates the text content blocks of a Converse output message. func converseText(msg *bedrockruntime.Message) string { var text string