Andy Dixon
f165adc5f2
Kind `bedrock` speaks the Converse API for chat roles and InvokeModel for Titan and Cohere embeddings. Requests are signed with AWS Signature Version 4 implemented with crypto/hmac and crypto/sha256 only; the signer is pinned by tests to signatures botocore produces for the same requests, with and without a session token, so drift fails the build. The credential is one secret reference, ACCESS_KEY_ID:SECRET_ACCESS_KEY [:SESSION_TOKEN], encrypted together like every other provider key. The region is read from the endpoint host, which must be a Bedrock runtime host over HTTPS unless private endpoints are explicitly allowed; other hosts are refused. Provider errors are scrubbed of the secret and token. Migration 0024 widens the provider kind constraint; existing rows are untouched and the down migration refuses while bedrock rows exist rather than dropping configuration. UI kind selects, docs/llm-providers.md, ADR-0010, ARCHITECTURE and the changelog are updated.
197 lines
10 KiB
Go
197 lines
10 KiB
Go
package providers
|
|
|
|
import (
|
|
"bytes"
|
|
"context"
|
|
"encoding/json"
|
|
"errors"
|
|
"fmt"
|
|
"net/http"
|
|
"net/http/httptest"
|
|
"strings"
|
|
"testing"
|
|
"time"
|
|
|
|
"github.com/andydixon/noema/internal/llm"
|
|
)
|
|
|
|
func TestOpenAIWireFormat(t *testing.T) {
|
|
var got map[string]any
|
|
var auth string
|
|
srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
|
auth = r.Header.Get("Authorization")
|
|
if r.URL.Path != "/v1/chat/completions" {
|
|
http.NotFound(w, r)
|
|
return
|
|
}
|
|
json.NewDecoder(r.Body).Decode(&got) //nolint:errcheck
|
|
w.Write([]byte("{\"model\":\"gpt-x\",\"choices\":[{\"message\":{\"content\":\"```json\\n{\\\"intent\\\":\\\"question\\\"}\\n```\"},\"finish_reason\":\"stop\"}],\"usage\":{\"prompt_tokens\":12,\"completion_tokens\":5,\"total_tokens\":17}}")) //nolint:errcheck
|
|
}))
|
|
defer srv.Close()
|
|
p := NewOpenAI(Config{Name: "test", Endpoint: srv.URL + "/v1", APIKey: "sk-test-123", Model: "gpt-x", AllowPrivate: true})
|
|
seed := int64(7)
|
|
resp, err := p.Complete(context.Background(), llm.Request{System: "sys", Messages: []llm.Message{{Role: "user", Content: "hi"}}, JSONSchema: json.RawMessage(`{"type":"object"}`), Seed: &seed, MaxTokens: 50})
|
|
if err != nil {
|
|
t.Fatal(err)
|
|
}
|
|
if auth != "Bearer sk-test-123" || got["model"] != "gpt-x" || got["seed"] != float64(7) || got["max_tokens"] != float64(50) {
|
|
t.Fatalf("request wrong: auth=%q body=%v", auth, got)
|
|
}
|
|
if msgs := got["messages"].([]any); len(msgs) != 2 || msgs[0].(map[string]any)["role"] != "system" {
|
|
t.Fatalf("system message not prepended: %v", msgs)
|
|
}
|
|
if string(resp.JSON) != `{"intent":"question"}` || resp.Usage.TotalTokens != 17 || resp.Model != "gpt-x" {
|
|
t.Fatalf("response wrong: %+v", resp)
|
|
}
|
|
}
|
|
|
|
func TestAnthropicWireFormat(t *testing.T) {
|
|
var got map[string]any
|
|
var key string
|
|
srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
|
key = r.Header.Get("x-api-key")
|
|
json.NewDecoder(r.Body).Decode(&got) //nolint:errcheck
|
|
w.Write([]byte(`{"model":"claude-x","content":[{"type":"text","text":"Hello"}],"stop_reason":"end_turn","usage":{"input_tokens":3,"output_tokens":2}}`)) //nolint:errcheck
|
|
}))
|
|
defer srv.Close()
|
|
p := NewAnthropic(Config{Name: "anth", Endpoint: srv.URL, APIKey: "ak", Model: "claude-x", AllowPrivate: true})
|
|
resp, err := p.Complete(context.Background(), llm.Request{System: "be brief", Messages: []llm.Message{{Role: "user", Content: "hi"}}})
|
|
if err != nil || resp.Text != "Hello" || resp.Usage.TotalTokens != 5 || key != "ak" || got["system"] != "be brief" {
|
|
t.Fatalf("anthropic wrong: %v %+v %v", err, resp, got)
|
|
}
|
|
}
|
|
|
|
func TestOllamaWireFormatAndErrors(t *testing.T) {
|
|
srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
|
switch r.URL.Path {
|
|
case "/api/chat":
|
|
var got map[string]any
|
|
json.NewDecoder(r.Body).Decode(&got) //nolint:errcheck
|
|
if got["format"] != "json" {
|
|
w.WriteHeader(500)
|
|
w.Write([]byte(`{"error":"expected json format, key sk-leak"}`)) //nolint:errcheck
|
|
return
|
|
}
|
|
w.Write([]byte(`{"model":"llama","message":{"role":"assistant","content":"not json at all"},"done_reason":"stop","prompt_eval_count":4,"eval_count":6}`)) //nolint:errcheck
|
|
case "/api/embed":
|
|
w.Write([]byte(`{"embeddings":[[0.1,0.2,0.3]]}`)) //nolint:errcheck
|
|
}
|
|
}))
|
|
defer srv.Close()
|
|
p := NewOllama(Config{Name: "local", Endpoint: srv.URL, Model: "llama", AllowPrivate: true})
|
|
_, err := p.Complete(context.Background(), llm.Request{Messages: []llm.Message{{Role: "user", Content: "x"}}})
|
|
if !errors.Is(err, llm.ErrProvider) || strings.Contains(err.Error(), "sk-leak") || !strings.Contains(err.Error(), "[redacted]") {
|
|
t.Fatalf("error should be a scrubbed provider error: %v", err)
|
|
}
|
|
_, err = p.Complete(context.Background(), llm.Request{Messages: []llm.Message{{Role: "user", Content: "x"}}, JSONSchema: json.RawMessage(`{}`)})
|
|
if !errors.Is(err, llm.ErrBadJSON) {
|
|
t.Fatalf("expected bad json, got %v", err)
|
|
}
|
|
e := NewOllamaEmbedder(Config{Name: "local", Endpoint: srv.URL, Model: "nomic", AllowPrivate: true}, 0)
|
|
vecs, err := e.Embed(context.Background(), []string{"a"})
|
|
if err != nil || len(vecs) != 1 || len(vecs[0]) != 3 || e.Dims() != 3 || e.Model() != "local/nomic" {
|
|
t.Fatalf("embed wrong: %v %v", err, vecs)
|
|
}
|
|
}
|
|
|
|
func TestPrivateEndpointRefusedByDefault(t *testing.T) {
|
|
srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { w.Write([]byte(`{}`)) })) //nolint:errcheck
|
|
defer srv.Close()
|
|
p := NewOpenAI(Config{Name: "x", Endpoint: srv.URL, Model: "m"})
|
|
if _, err := p.Complete(context.Background(), llm.Request{Messages: []llm.Message{{Role: "user", Content: "x"}}}); err == nil || !strings.Contains(err.Error(), "private") {
|
|
t.Fatalf("loopback endpoint should need allow_private: %v", err)
|
|
}
|
|
}
|
|
|
|
// TestSigV4MatchesReferenceSigner pins the signer to signatures computed by
|
|
// botocore for the same request, with and without a session token.
|
|
func TestSigV4MatchesReferenceSigner(t *testing.T) {
|
|
body := []byte(`{"messages":[{"role":"user","content":[{"text":"hi"}]}]}`)
|
|
at := time.Date(2026, 9, 15, 10, 0, 0, 0, time.UTC)
|
|
mk := func() *http.Request {
|
|
req, _ := http.NewRequest(http.MethodPost, "https://bedrock-runtime.eu-west-2.amazonaws.com/model/"+awsPathSegment("anthropic.claude-3-haiku-20240307-v1:0")+"/converse", bytes.NewReader(body))
|
|
req.Header.Set("Content-Type", "application/json")
|
|
return req
|
|
}
|
|
req := mk()
|
|
signV4(req, body, awsCredentials{AccessKeyID: "AKIDEXAMPLE", SecretAccessKey: "wJalrXUtnFEMI/K7MDENG+bPxRfiCYEXAMPLEKEY"}, "eu-west-2", "bedrock", at)
|
|
want := "AWS4-HMAC-SHA256 Credential=AKIDEXAMPLE/20260915/eu-west-2/bedrock/aws4_request, SignedHeaders=content-type;host;x-amz-date, Signature=91ab2cb0d4ac1c8ff730487209ce46355c4d033b4d693361ca422142f0f66c9f"
|
|
if got := req.Header.Get("Authorization"); got != want {
|
|
t.Fatalf("signature\n got %s\nwant %s", got, want)
|
|
}
|
|
if req.Header.Get("X-Amz-Date") != "20260915T100000Z" {
|
|
t.Fatal(req.Header.Get("X-Amz-Date"))
|
|
}
|
|
req = mk()
|
|
signV4(req, body, awsCredentials{AccessKeyID: "AKIDEXAMPLE", SecretAccessKey: "wJalrXUtnFEMI/K7MDENG+bPxRfiCYEXAMPLEKEY", SessionToken: "SESSIONTOKENEXAMPLE"}, "eu-west-2", "bedrock", at)
|
|
want = "AWS4-HMAC-SHA256 Credential=AKIDEXAMPLE/20260915/eu-west-2/bedrock/aws4_request, SignedHeaders=content-type;host;x-amz-date;x-amz-security-token, Signature=72345f5e4ff70084eb6c87e22a2b47c7234de67f71fe2faca1e7840a74c66a2b"
|
|
if got := req.Header.Get("Authorization"); got != want || req.Header.Get("X-Amz-Security-Token") != "SESSIONTOKENEXAMPLE" {
|
|
t.Fatalf("session signature\n got %s\nwant %s", got, want)
|
|
}
|
|
if c, ok := parseAWSCredentials("AKID:secret:tok"); !ok || c.SessionToken != "tok" {
|
|
t.Fatal("credential parsing")
|
|
}
|
|
if _, ok := parseAWSCredentials("just-a-key"); ok {
|
|
t.Fatal("a bare key must be rejected")
|
|
}
|
|
}
|
|
|
|
// TestBedrockWireFormat: Converse request shape, signed headers, response
|
|
// parsing, JSON-mode instruction, and embeddings for Titan and Cohere.
|
|
func TestBedrockWireFormat(t *testing.T) {
|
|
var got map[string]any
|
|
var auth, date, path string
|
|
srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
|
auth, date, path = r.Header.Get("Authorization"), r.Header.Get("X-Amz-Date"), r.URL.RawPath
|
|
if path == "" {
|
|
path = r.URL.Path
|
|
}
|
|
json.NewDecoder(r.Body).Decode(&got) //nolint:errcheck
|
|
switch {
|
|
case strings.HasSuffix(r.URL.Path, "/converse"):
|
|
w.Write([]byte(`{"output":{"message":{"role":"assistant","content":[{"text":"{\"intent\":\"statement\"}"}]}},"stopReason":"end_turn","usage":{"inputTokens":7,"outputTokens":3,"totalTokens":10}}`)) //nolint:errcheck
|
|
case strings.Contains(r.URL.Path, "cohere"):
|
|
w.Write([]byte(`{"embeddings":[[0.1,0.2],[0.3,0.4]]}`)) //nolint:errcheck
|
|
default:
|
|
w.Write([]byte(`{"embedding":[0.5,0.6]}`)) //nolint:errcheck
|
|
}
|
|
}))
|
|
defer srv.Close()
|
|
if _, err := NewBedrock(Config{Name: "b", Endpoint: "https://bedrock-runtime.eu-west-2.amazonaws.com", APIKey: "nokey", Model: "m"}); err == nil {
|
|
t.Fatal("malformed credential accepted")
|
|
}
|
|
if _, err := NewBedrock(Config{Name: "b", Endpoint: "https://example.com", APIKey: "a:b", Model: "m"}); err == nil {
|
|
t.Fatal("non-Bedrock endpoint accepted")
|
|
}
|
|
p, err := NewBedrock(Config{Name: "b", Endpoint: srv.URL, APIKey: "AKID:SECRET:TOKEN", Model: "anthropic.claude-3-haiku-20240307-v1:0", AllowPrivate: true, Headers: map[string]string{"X-Noema-Region": "eu-west-2"}})
|
|
if err != nil {
|
|
t.Fatal(err)
|
|
}
|
|
p.now = func() time.Time { return time.Date(2026, 9, 15, 10, 0, 0, 0, time.UTC) }
|
|
resp, err := p.Complete(context.Background(), llm.Request{System: "be brief", Messages: []llm.Message{{Role: "user", Content: "hi"}}, MaxTokens: 50, JSONSchema: json.RawMessage(`{"type":"object"}`)})
|
|
if err != nil || resp.Text == "" || resp.Usage.TotalTokens != 10 || string(resp.JSON) != `{"intent":"statement"}` || resp.FinishReason != "end_turn" {
|
|
t.Fatalf("converse: %v %+v", err, resp)
|
|
}
|
|
if !strings.HasPrefix(auth, "AWS4-HMAC-SHA256 Credential=AKID/20260915/eu-west-2/bedrock/aws4_request, SignedHeaders=content-type;host;x-amz-date;x-amz-security-token, Signature=") || date != "20260915T100000Z" {
|
|
t.Fatalf("signing headers: %q %q", auth, date)
|
|
}
|
|
if path != "/model/anthropic.claude-3-haiku-20240307-v1%3A0/converse" {
|
|
t.Fatalf("model id must be encoded once in the path: %s", path)
|
|
}
|
|
sys := got["system"].([]any)[0].(map[string]any)["text"].(string)
|
|
if !strings.HasPrefix(sys, "be brief") || !strings.Contains(sys, "single JSON object") || got["inferenceConfig"].(map[string]any)["maxTokens"] != float64(50) {
|
|
t.Fatalf("request body: %+v", got)
|
|
}
|
|
if strings.Contains(fmt.Sprint(got), "SECRET") {
|
|
t.Fatal("secret leaked into the body")
|
|
}
|
|
titan, _ := NewBedrockEmbedder(Config{Name: "b", Endpoint: srv.URL, APIKey: "AKID:SECRET", Model: "amazon.titan-embed-text-v2:0", AllowPrivate: true}, 2)
|
|
if v, err := titan.Embed(context.Background(), []string{"a", "b"}); err != nil || len(v) != 2 || v[1][0] != 0.5 || titan.Model() != "bedrock/amazon.titan-embed-text-v2:0" {
|
|
t.Fatalf("titan: %v %v", err, v)
|
|
}
|
|
cohere, _ := NewBedrockEmbedder(Config{Name: "b", Endpoint: srv.URL, APIKey: "AKID:SECRET", Model: "cohere.embed-english-v3", AllowPrivate: true}, 2)
|
|
if v, err := cohere.Embed(context.Background(), []string{"a", "b"}); err != nil || len(v) != 2 || v[1][1] != 0.4 || got["input_type"] != "search_document" {
|
|
t.Fatalf("cohere: %v %v", err, v)
|
|
}
|
|
}
|