diff --git a/CHANGELOG.md b/CHANGELOG.md index 67881d3..f34fb98 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -10,6 +10,17 @@ surface is governed by [`COMPATIBILITY.md`](COMPATIBILITY.md). ## [Unreleased] ### Added +- **Native AWS Bedrock provider.** Run Bedrock-hosted models through RiskKernel with + full budgets / approvals / audit / OTel — set the standard AWS env vars + (`AWS_ACCESS_KEY_ID`, `AWS_SECRET_ACCESS_KEY`, optional `AWS_SESSION_TOKEN`, + `AWS_REGION`) and `RISKKERNEL_DEFAULT_PROVIDER=bedrock`, then call with a Bedrock + model id (e.g. `anthropic.claude-3-5-sonnet-20240620-v1:0`). Implemented against + the Bedrock Runtime **Converse** API and signed with hand-rolled AWS SigV4 — **no + AWS SDK dependency** (the signer is verified against AWS's published SigV4 + known-answer test vector). Registered only when AWS credentials are present; + `RISKKERNEL_BEDROCK_BASE_URL` overrides the endpoint for a VPC/PrivateLink setup. + Long-tail Bedrock model ids meter `priced:false` until a `RISKKERNEL_PRICING_FILE` + rate is added (token budget yes, dollar budget no). See [`docs/PROVIDERS.md`](docs/PROVIDERS.md#aws-bedrock-native). - **Importable Datadog dashboard.** [`examples/otel/datadog`](examples/otel/datadog) ships a ready-made Datadog dashboard (spend per run, budget halts by reason, tool-call outcomes, latency and token burn by model) built from the same diff --git a/README.md b/README.md index d9a61db..11eb356 100644 --- a/README.md +++ b/README.md @@ -45,7 +45,7 @@ It is **not** another gateway (LiteLLM/Portkey own routing), **not** another obs ## Three ways to adopt — pick the one that fits -1. **Proxy (zero code).** Set one env var: `OPENAI_BASE_URL=http://localhost:7070/v1` (or `ANTHROPIC_BASE_URL` for `/v1/messages`). Every call — streaming or not — is intercepted, budgeted, logged, checkpointed, and forwarded to the real provider with your key. Native providers: Anthropic, OpenAI, and Ollama (local); front the long tail (Gemini, Cohere, Mistral, Bedrock, …) with [LiteLLM upstream](docs/PROVIDERS.md). +1. **Proxy (zero code).** Set one env var: `OPENAI_BASE_URL=http://localhost:7070/v1` (or `ANTHROPIC_BASE_URL` for `/v1/messages`). Every call — streaming or not — is intercepted, budgeted, logged, checkpointed, and forwarded to the real provider with your key. Native providers: Anthropic, OpenAI, Ollama (local), and AWS Bedrock; front the long tail (Gemini, Cohere, Mistral, …) with [LiteLLM upstream](docs/PROVIDERS.md). 2. **SDK (deep control).** `pip install riskkernel` (Python) or `npm install @riskkernel/sdk` (TypeScript), then governed runs, per-step loop/time budgets, checkpoints, and approval gates. Framework adapters for the Claude Agent SDK, OpenAI Agents SDK, LangChain, LlamaIndex, CrewAI, AutoGen, and PydanticAI (Python), and the Vercel AI SDK (TypeScript). 3. **OpenTelemetry (universal).** RiskKernel is an OTLP endpoint *and* emitter — ingest GenAI spans (`POST /v1/traces`) to meter apps already instrumented with OpenLLMetry / the OpenAI Agents SDK / the Vercel AI SDK, and export cost/halt/tool spans to the backend you already run. diff --git a/ROADMAP.md b/ROADMAP.md index 4a1af38..1b6d909 100644 --- a/ROADMAP.md +++ b/ROADMAP.md @@ -18,7 +18,8 @@ The core runtime is built and released: ([`docs/RESUME.md`](docs/RESUME.md)). - **OpenAI- and Anthropic-compatible proxy** — point one env var at RiskKernel and every call (streaming or not) is metered, priced, and budget-enforced (BYO key). - Native providers: Anthropic, OpenAI, and Ollama (local, key-free). + Native providers: Anthropic, OpenAI, Ollama (local, key-free), and AWS Bedrock; the + long tail is fronted via LiteLLM upstream ([`docs/PROVIDERS.md`](docs/PROVIDERS.md)). - **Human-in-the-loop approval** — gate side-effecting tools; resolve from the CLI, a local web page, a webhook, or **Slack** ([`docs/APPROVALS_SLACK.md`](docs/APPROVALS_SLACK.md)). - **Policy-as-code, enforced per-run** — named policy bundles via `POST /v1/policies` @@ -26,7 +27,8 @@ The core runtime is built and released: under a bundle is governed by its tool allowlist and approval rules, not just its budget ([`docs/POLICY.md`](docs/POLICY.md)). - **OpenTelemetry GenAI — export and ingress** — emit cost/halt/tool spans into your - existing backend (ready-made **Grafana + Tempo** and **SigNoz** dashboards), *and* + existing backend (ready-made **Grafana + Tempo**, **SigNoz**, and **Datadog** + dashboards), *and* ingest GenAI spans (`POST /v1/traces`) to meter apps RiskKernel never proxied ([`docs/OTLP_INGRESS.md`](docs/OTLP_INGRESS.md)). - **Spend attribution** — roll cost up across runs by team/user/feature @@ -51,10 +53,6 @@ The core runtime is built and released: Where the work is heading near-term: -- **More native providers** — AWS Bedrock ([#24](https://github.com/prashar32/riskkernel/issues/24)); - the long tail via LiteLLM-as-upstream. -- **More backend dashboards** — a Datadog dashboard to join the Grafana and SigNoz - examples. - **Easier install** — a Homebrew tap for `brew install riskkernel` ([#97](https://github.com/prashar32/riskkernel/issues/97)). diff --git a/docs/PROVIDERS.md b/docs/PROVIDERS.md index a28767b..5ba47cc 100644 --- a/docs/PROVIDERS.md +++ b/docs/PROVIDERS.md @@ -1,11 +1,11 @@ # Providers — native, and the long tail via LiteLLM -RiskKernel implements the top providers **natively** in Go: Anthropic, OpenAI, and -Ollama (local). For those, point your app at the proxy and you're done — no extra -moving parts. +RiskKernel implements the top providers **natively** in Go: Anthropic, OpenAI, +Ollama (local), and AWS Bedrock. For those, point your app at the proxy and you're +done — no extra moving parts. (Bedrock setup is in its own section, [below](#aws-bedrock-native).) The other 100+ providers (Google Gemini, Cohere, Mistral, Groq, Together, Azure -OpenAI, AWS Bedrock, OpenRouter, …) are not reimplemented inside RiskKernel — that +OpenAI, OpenRouter, …) are not reimplemented inside RiskKernel — that isn't the product. Instead you front them with **[LiteLLM](https://github.com/BerriAI/litellm)**, which already speaks all of them through one OpenAI-compatible endpoint, and put RiskKernel **in front of LiteLLM**. RiskKernel governs every call; LiteLLM does the @@ -29,7 +29,8 @@ upstream. LiteLLM is just the real-world version of that upstream. | Anthropic (`claude-*`) | Native. Set `ANTHROPIC_API_KEY` on the daemon; use a `claude-*` model. | | OpenAI (`gpt-*`, `o1`, `o3`) | Native. Set `OPENAI_API_KEY` on the daemon; use a `gpt-*`/`o1`/`o3` model. | | A local Ollama model | Native. Set `RISKKERNEL_OLLAMA_BASE_URL`; use a model the routing sends to Ollama. | -| Anything else (Gemini, Cohere, Mistral, Groq, Bedrock, …) | Front it with **LiteLLM** as described below. | +| AWS Bedrock | Native. Set AWS credentials + `RISKKERNEL_DEFAULT_PROVIDER=bedrock`; see [below](#aws-bedrock-native). | +| Anything else (Gemini, Cohere, Mistral, Groq, …) | Front it with **LiteLLM** as described below. | You only need LiteLLM for the long tail. If your stack is purely Anthropic/OpenAI/ Ollama, skip this page. @@ -203,3 +204,40 @@ section of [`docs/BUDGETS.md`](BUDGETS.md) for the full format and stability pro Anthropic/OpenAI/Ollama stacks, native providers need no LiteLLM at all. - **Streaming** works end-to-end (LiteLLM is OpenAI-compatible SSE), and the call is metered from the stream's final usage chunk — same as a native streamed call. + +## AWS Bedrock (native) + +Bedrock is implemented natively against the Bedrock Runtime **Converse** API, signed +with hand-rolled AWS SigV4 — **no AWS SDK dependency**. BYO AWS credentials, read +from the standard AWS environment variables and never stored. + +```bash +export AWS_ACCESS_KEY_ID=AKIA... +export AWS_SECRET_ACCESS_KEY=... +export AWS_SESSION_TOKEN=... # optional, for STS temporary credentials +export AWS_REGION=us-east-1 # or AWS_DEFAULT_REGION +export RISKKERNEL_DEFAULT_PROVIDER=bedrock +riskkernel serve +``` + +Then call with a **Bedrock model id** (a foundation-model id or an inference-profile +id), e.g. `anthropic.claude-3-5-sonnet-20240620-v1:0` or +`us.anthropic.claude-3-5-sonnet-20241022-v2:0`. + +**Routing.** RiskKernel routes by model-name prefix: `claude-*` → native Anthropic, +`gpt-*`/`o1`/`o3` → native OpenAI, everything else → the **default provider**. A +Bedrock model id starts with a vendor namespace (`anthropic.`, `amazon.`, `meta.`, +`us.`, …), so it doesn't match those prefixes and falls to the default — which is why +you set `RISKKERNEL_DEFAULT_PROVIDER=bedrock`. (Bare `claude-*`/`gpt-*` names still +route to the native Anthropic/OpenAI providers if those are also configured.) + +Bedrock is registered only when `AWS_ACCESS_KEY_ID` and `AWS_SECRET_ACCESS_KEY` are +set; with `RISKKERNEL_DEFAULT_PROVIDER=bedrock` and no credentials, the daemon +refuses to start with a clear error. For a VPC/PrivateLink endpoint, override the +runtime URL with `RISKKERNEL_BEDROCK_BASE_URL`. + +**Cost caveat (same as the long tail).** The built-in pricing table only knows +`claude-*` / `gpt-*` names, so a Bedrock model id meters `priced: false` (cost `$0`) +until you add a rate for it via `RISKKERNEL_PRICING_FILE` — its tokens still count +toward the **token** budget, just not the **dollar** budget. Cost is metered from +Bedrock's own returned token usage; RiskKernel only needs the per-token rate. diff --git a/internal/app/bootstrap.go b/internal/app/bootstrap.go index 8d43059..05ff4e0 100644 --- a/internal/app/bootstrap.go +++ b/internal/app/bootstrap.go @@ -198,22 +198,25 @@ func OpenStore(cfg *config.Config, log *slog.Logger) (storage.Store, error) { } // BuildRegistry constructs the provider registry from config. Anthropic, OpenAI, -// and Ollama are implemented natively; Bedrock is a stub config can reference. -// The default provider must be usable. +// Ollama, and AWS Bedrock are implemented natively. The default provider must be +// usable. func BuildRegistry(cfg *config.Config) (*provider.Registry, error) { // Anthropic is always registered (native provider). When the key is absent the // daemon still boots — health and routing work — and only an actual Chat call // returns a clear "missing API key" error. Ollama is native and key-free - // (local models); OpenAI is registered (native) when a key is present; Bedrock - // is a stub config can name before it's built out. + // (local models); OpenAI and Bedrock are registered (native) when their + // credentials are present. ps := []provider.Provider{ provider.NewAnthropic(cfg.AnthropicAPIKey).WithBaseURL(cfg.AnthropicBaseURL), - provider.NewBedrock(), provider.NewOllama(cfg.OllamaBaseURL), // empty → local default } if cfg.OpenAIAPIKey != "" { ps = append(ps, provider.NewOpenAI(cfg.OpenAIAPIKey).WithBaseURL(cfg.OpenAIBaseURL)) } + if cfg.AWSAccessKeyID != "" && cfg.AWSSecretAccessKey != "" { + ps = append(ps, provider.NewBedrock(cfg.AWSAccessKeyID, cfg.AWSSecretAccessKey, + cfg.AWSSessionToken, cfg.AWSRegion).WithBaseURL(cfg.BedrockBaseURL)) + } return provider.NewRegistry(cfg.DefaultProvider, ps...) } diff --git a/internal/config/config.go b/internal/config/config.go index 0d7179b..cbaec4a 100644 --- a/internal/config/config.go +++ b/internal/config/config.go @@ -63,6 +63,19 @@ type Config struct { // OllamaBaseURL points the native Ollama provider at a server. Empty uses the // local default (http://localhost:11434). Read from RISKKERNEL_OLLAMA_BASE_URL. OllamaBaseURL string + // BedrockBaseURL overrides the Bedrock runtime endpoint (a VPC/PrivateLink + // endpoint, or a mock). Empty uses the regional default. Read from + // RISKKERNEL_BEDROCK_BASE_URL. + BedrockBaseURL string + + // AWS credentials + region for the native Bedrock provider, read from the + // standard AWS env vars so existing setups need no new config. Bedrock is + // registered only when an access key and secret are present. Never stored or + // logged. AWSSessionToken is optional (STS temporary credentials). + AWSAccessKeyID string // AWS_ACCESS_KEY_ID + AWSSecretAccessKey string // AWS_SECRET_ACCESS_KEY + AWSSessionToken string // AWS_SESSION_TOKEN + AWSRegion string // AWS_REGION, then AWS_DEFAULT_REGION // DefaultBudget is applied to runs created without an explicit budget — e.g. // proxy calls that supply only a run-id. Any zero field is unlimited. When no @@ -216,20 +229,25 @@ func Load() (*Config, error) { } cfg := &Config{ - Port: port, - DataDir: getenvDefault("RISKKERNEL_DATA_DIR", "./data"), - DatabaseURL: os.Getenv("RISKKERNEL_DATABASE_URL"), - APIToken: os.Getenv("RISKKERNEL_API_TOKEN"), - DefaultProvider: getenvDefault("RISKKERNEL_DEFAULT_PROVIDER", "anthropic"), - AnthropicAPIKey: os.Getenv("ANTHROPIC_API_KEY"), - OpenAIAPIKey: os.Getenv("OPENAI_API_KEY"), - AnthropicBaseURL: os.Getenv("RISKKERNEL_ANTHROPIC_BASE_URL"), - OpenAIBaseURL: os.Getenv("RISKKERNEL_OPENAI_BASE_URL"), - OllamaBaseURL: os.Getenv("RISKKERNEL_OLLAMA_BASE_URL"), - DefaultBudget: budget, - PricingFile: os.Getenv("RISKKERNEL_PRICING_FILE"), - PolicyFile: os.Getenv("RISKKERNEL_POLICY_FILE"), - OTel: loadOTel(), + Port: port, + DataDir: getenvDefault("RISKKERNEL_DATA_DIR", "./data"), + DatabaseURL: os.Getenv("RISKKERNEL_DATABASE_URL"), + APIToken: os.Getenv("RISKKERNEL_API_TOKEN"), + DefaultProvider: getenvDefault("RISKKERNEL_DEFAULT_PROVIDER", "anthropic"), + AnthropicAPIKey: os.Getenv("ANTHROPIC_API_KEY"), + OpenAIAPIKey: os.Getenv("OPENAI_API_KEY"), + AnthropicBaseURL: os.Getenv("RISKKERNEL_ANTHROPIC_BASE_URL"), + OpenAIBaseURL: os.Getenv("RISKKERNEL_OPENAI_BASE_URL"), + OllamaBaseURL: os.Getenv("RISKKERNEL_OLLAMA_BASE_URL"), + BedrockBaseURL: os.Getenv("RISKKERNEL_BEDROCK_BASE_URL"), + AWSAccessKeyID: os.Getenv("AWS_ACCESS_KEY_ID"), + AWSSecretAccessKey: os.Getenv("AWS_SECRET_ACCESS_KEY"), + AWSSessionToken: os.Getenv("AWS_SESSION_TOKEN"), + AWSRegion: getenvDefault("AWS_REGION", os.Getenv("AWS_DEFAULT_REGION")), + DefaultBudget: budget, + PricingFile: os.Getenv("RISKKERNEL_PRICING_FILE"), + PolicyFile: os.Getenv("RISKKERNEL_POLICY_FILE"), + OTel: loadOTel(), Approval: ApprovalConfig{ DefaultSafe: envBoolDefault("RISKKERNEL_APPROVAL_DEFAULT_SAFE", true), WebhookURL: os.Getenv("RISKKERNEL_APPROVAL_WEBHOOK"), diff --git a/internal/provider/bedrock.go b/internal/provider/bedrock.go new file mode 100644 index 0000000..c009a07 --- /dev/null +++ b/internal/provider/bedrock.go @@ -0,0 +1,189 @@ +package provider + +import ( + "bytes" + "context" + "encoding/json" + "fmt" + "io" + "net/http" + "net/url" + "strings" + "time" +) + +// defaultBedrockMaxTokens is used when a Request omits MaxTokens. +const defaultBedrockMaxTokens = 1024 + +// Bedrock implements Provider against the AWS Bedrock Runtime **Converse** API, +// signed with hand-rolled SigV4 (no AWS SDK dependency — see sigv4.go). The unified +// Converse API works across Bedrock's hosted models and returns token usage, which +// the governor meters and the ledger prices. BYO AWS credentials, read from the +// standard AWS env vars and never stored. +type Bedrock struct { + creds awsCreds + region string + baseURL string // override; empty → the regional bedrock-runtime endpoint + http *http.Client +} + +// NewBedrock constructs a Bedrock provider for the given region and credentials. +func NewBedrock(accessKey, secretKey, sessionToken, region string) *Bedrock { + return &Bedrock{ + creds: awsCreds{accessKey: accessKey, secretKey: secretKey, sessionToken: sessionToken}, + region: region, + http: &http.Client{Timeout: 120 * time.Second}, + } +} + +// WithBaseURL overrides the runtime endpoint — for a VPC/PrivateLink endpoint or a +// test mock. Empty keeps the regional default. Returns the provider for chaining. +func (b *Bedrock) WithBaseURL(u string) *Bedrock { + if u != "" { + b.baseURL = strings.TrimRight(u, "/") + } + return b +} + +// Name returns the stable provider identifier. +func (b *Bedrock) Name() string { return "bedrock" } + +func (b *Bedrock) endpoint() string { + if b.baseURL != "" { + return b.baseURL + } + return "https://bedrock-runtime." + b.region + ".amazonaws.com" +} + +// --- Converse wire types --- + +type bedrockTextBlock struct { + Text string `json:"text"` +} + +type bedrockMessage struct { + Role string `json:"role"` + Content []bedrockTextBlock `json:"content"` +} + +type bedrockReq struct { + Messages []bedrockMessage `json:"messages"` + System []bedrockTextBlock `json:"system,omitempty"` + InferenceConfig *bedrockInference `json:"inferenceConfig,omitempty"` +} + +type bedrockInference struct { + MaxTokens int `json:"maxTokens,omitempty"` + Temperature *float64 `json:"temperature,omitempty"` +} + +type bedrockResp struct { + Output struct { + Message bedrockMessage `json:"message"` + } `json:"output"` + StopReason string `json:"stopReason"` + Usage struct { + InputTokens int64 `json:"inputTokens"` + OutputTokens int64 `json:"outputTokens"` + } `json:"usage"` +} + +// Chat performs one chat completion against the Bedrock Converse API. +func (b *Bedrock) Chat(ctx context.Context, req Request) (*Response, error) { + if b.creds.accessKey == "" || b.creds.secretKey == "" { + return nil, fmt.Errorf("bedrock: missing AWS credentials (set AWS_ACCESS_KEY_ID and AWS_SECRET_ACCESS_KEY)") + } + if b.region == "" { + return nil, fmt.Errorf("bedrock: missing AWS region (set AWS_REGION)") + } + + // Bedrock takes the system prompt as a separate field. Lift any system message + // out of Messages and map the rest into Converse content blocks. + var system []bedrockTextBlock + if req.System != "" { + system = append(system, bedrockTextBlock{Text: req.System}) + } + msgs := make([]bedrockMessage, 0, len(req.Messages)) + for _, m := range req.Messages { + if m.Role == RoleSystem { + system = append(system, bedrockTextBlock{Text: m.Content}) + continue + } + msgs = append(msgs, bedrockMessage{Role: string(m.Role), Content: []bedrockTextBlock{{Text: m.Content}}}) + } + + maxTokens := req.MaxTokens + if maxTokens <= 0 { + maxTokens = defaultBedrockMaxTokens + } + body, err := json.Marshal(bedrockReq{ + Messages: msgs, + System: system, + InferenceConfig: &bedrockInference{MaxTokens: maxTokens, Temperature: req.Temperature}, + }) + if err != nil { + return nil, fmt.Errorf("bedrock: marshaling request: %w", err) + } + + // Build the URL so the SigV4 canonical path matches what's sent: the model id + // can contain ':' (e.g. ...-v1:0), which AWS encodes as %3A — set Path (decoded) + // and RawPath (AWS-encoded) so EscapedPath() is identical to the wire path. + u, err := url.Parse(b.endpoint()) + if err != nil { + return nil, fmt.Errorf("bedrock: bad endpoint %q: %w", b.endpoint(), err) + } + u.Path = "/model/" + req.Model + "/converse" + u.RawPath = "/model/" + awsURIEncodeSegment(req.Model) + "/converse" + + httpReq, err := http.NewRequestWithContext(ctx, http.MethodPost, b.endpoint(), bytes.NewReader(body)) + if err != nil { + return nil, fmt.Errorf("bedrock: building request: %w", err) + } + httpReq.URL = u + httpReq.Header.Set("content-type", "application/json") + signV4(httpReq, body, b.creds, b.region, "bedrock", time.Now()) + + resp, err := b.http.Do(httpReq) + if err != nil { + // Propagate context errors verbatim so the governor can distinguish a + // kill-switch/timeout cancellation from a transport failure. + if ctxErr := ctx.Err(); ctxErr != nil { + return nil, ctxErr + } + return nil, fmt.Errorf("bedrock: request failed: %w", err) + } + defer resp.Body.Close() + + raw, err := io.ReadAll(resp.Body) + if err != nil { + return nil, fmt.Errorf("bedrock: reading response: %w", err) + } + if resp.StatusCode != http.StatusOK { + var apiErr struct { + Message string `json:"message"` + } + if json.Unmarshal(raw, &apiErr) == nil && apiErr.Message != "" { + return nil, fmt.Errorf("bedrock: %s (http %d)", apiErr.Message, resp.StatusCode) + } + return nil, fmt.Errorf("bedrock: http %d: %s", resp.StatusCode, strings.TrimSpace(string(raw))) + } + + var out bedrockResp + if err := json.Unmarshal(raw, &out); err != nil { + return nil, fmt.Errorf("bedrock: decoding response: %w", err) + } + var sb strings.Builder + for _, c := range out.Output.Message.Content { + sb.WriteString(c.Text) + } + return &Response{ + ID: resp.Header.Get("x-amzn-RequestId"), + Model: req.Model, + Content: sb.String(), + FinishReason: out.StopReason, + Usage: Usage{ + PromptTokens: out.Usage.InputTokens, + CompletionTokens: out.Usage.OutputTokens, + }, + }, nil +} diff --git a/internal/provider/bedrock_test.go b/internal/provider/bedrock_test.go new file mode 100644 index 0000000..8bc931b --- /dev/null +++ b/internal/provider/bedrock_test.go @@ -0,0 +1,143 @@ +package provider + +import ( + "context" + "encoding/json" + "net/http" + "net/http/httptest" + "strings" + "testing" +) + +const bedrockModel = "anthropic.claude-3-5-sonnet-20240620-v1:0" + +func TestBedrockChat_Success(t *testing.T) { + var gotPath, gotAuth, gotDate, gotCT string + var gotBody bedrockReq + srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + gotPath = r.RequestURI + gotAuth = r.Header.Get("Authorization") + gotDate = r.Header.Get("X-Amz-Date") + gotCT = r.Header.Get("Content-Type") + _ = json.NewDecoder(r.Body).Decode(&gotBody) + w.Header().Set("x-amzn-RequestId", "req-123") + w.Header().Set("content-type", "application/json") + _, _ = w.Write([]byte(`{ + "output":{"message":{"role":"assistant","content":[{"text":"hello "},{"text":"world"}]}}, + "stopReason":"end_turn", + "usage":{"inputTokens":11,"outputTokens":4,"totalTokens":15} + }`)) + })) + defer srv.Close() + + b := NewBedrock("AKID", "secret", "", "us-east-1").WithBaseURL(srv.URL) + resp, err := b.Chat(context.Background(), Request{ + Model: bedrockModel, + System: "be terse", + MaxTokens: 256, + Messages: []Message{{Role: RoleUser, Content: "hi"}}, + }) + if err != nil { + t.Fatalf("Chat: %v", err) + } + + // The model id's ':' is percent-encoded in the path that's actually sent (so it + // matches the signed canonical URI). + if gotPath != "/model/anthropic.claude-3-5-sonnet-20240620-v1%3A0/converse" { + t.Errorf("request path = %q", gotPath) + } + if gotCT != "application/json" || gotDate == "" { + t.Errorf("headers: content-type=%q x-amz-date=%q", gotCT, gotDate) + } + // A well-formed SigV4 Authorization scoped to the bedrock service, signing the + // content-type, host, and date headers. + if !strings.HasPrefix(gotAuth, "AWS4-HMAC-SHA256 Credential=AKID/") || + !strings.Contains(gotAuth, "/us-east-1/bedrock/aws4_request") || + !strings.Contains(gotAuth, "SignedHeaders=content-type;host;x-amz-date") || + !strings.Contains(gotAuth, "Signature=") { + t.Errorf("Authorization not well-formed: %q", gotAuth) + } + // System lifted out; the user message mapped to a Converse content block. + if len(gotBody.System) != 1 || gotBody.System[0].Text != "be terse" { + t.Errorf("system = %+v", gotBody.System) + } + if len(gotBody.Messages) != 1 || gotBody.Messages[0].Role != "user" || + len(gotBody.Messages[0].Content) != 1 || gotBody.Messages[0].Content[0].Text != "hi" { + t.Errorf("messages = %+v", gotBody.Messages) + } + if gotBody.InferenceConfig == nil || gotBody.InferenceConfig.MaxTokens != 256 { + t.Errorf("inferenceConfig = %+v", gotBody.InferenceConfig) + } + + if resp.Content != "hello world" || resp.FinishReason != "end_turn" || resp.Model != bedrockModel { + t.Errorf("resp = %+v", resp) + } + if resp.ID != "req-123" { + t.Errorf("id = %q, want req-123 (from x-amzn-RequestId)", resp.ID) + } + if resp.Usage.PromptTokens != 11 || resp.Usage.CompletionTokens != 4 || resp.Usage.Total() != 15 { + t.Errorf("usage = %+v", resp.Usage) + } +} + +func TestBedrockChat_SessionToken(t *testing.T) { + var gotToken, gotAuth string + srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + gotToken = r.Header.Get("X-Amz-Security-Token") + gotAuth = r.Header.Get("Authorization") + _, _ = w.Write([]byte(`{"output":{"message":{"content":[{"text":"ok"}]}},"stopReason":"end_turn","usage":{"inputTokens":1,"outputTokens":1}}`)) + })) + defer srv.Close() + + b := NewBedrock("AKID", "secret", "session-tok", "us-west-2").WithBaseURL(srv.URL) + if _, err := b.Chat(context.Background(), Request{Model: bedrockModel, Messages: []Message{{Role: RoleUser, Content: "hi"}}}); err != nil { + t.Fatalf("Chat: %v", err) + } + if gotToken != "session-tok" { + t.Errorf("x-amz-security-token = %q", gotToken) + } + // The session token must be part of the signed headers (some services require it). + if !strings.Contains(gotAuth, "x-amz-security-token") { + t.Errorf("session token not in SignedHeaders: %q", gotAuth) + } +} + +func TestBedrockChat_APIError(t *testing.T) { + srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, _ *http.Request) { + w.WriteHeader(http.StatusBadRequest) + _, _ = w.Write([]byte(`{"message":"The provided model identifier is invalid."}`)) + })) + defer srv.Close() + + b := NewBedrock("AKID", "secret", "", "us-east-1").WithBaseURL(srv.URL) + _, err := b.Chat(context.Background(), Request{Model: "bad", Messages: []Message{{Role: RoleUser, Content: "hi"}}}) + if err == nil || !strings.Contains(err.Error(), "provided model identifier is invalid") { + t.Fatalf("expected API error surfaced, got %v", err) + } +} + +func TestBedrockChat_MissingCreds(t *testing.T) { + b := NewBedrock("", "", "", "us-east-1") + _, err := b.Chat(context.Background(), Request{Model: bedrockModel, Messages: []Message{{Role: RoleUser, Content: "hi"}}}) + if err == nil || !strings.Contains(err.Error(), "missing AWS credentials") { + t.Fatalf("expected missing-credentials error, got %v", err) + } +} + +func TestBedrockChat_MissingRegion(t *testing.T) { + b := NewBedrock("AKID", "secret", "", "") + _, err := b.Chat(context.Background(), Request{Model: bedrockModel, Messages: []Message{{Role: RoleUser, Content: "hi"}}}) + if err == nil || !strings.Contains(err.Error(), "missing AWS region") { + t.Fatalf("expected missing-region error, got %v", err) + } +} + +func TestBedrockChat_ContextCancel(t *testing.T) { + b := NewBedrock("AKID", "secret", "", "us-east-1").WithBaseURL("http://127.0.0.1:0") + ctx, cancel := context.WithCancel(context.Background()) + cancel() + _, err := b.Chat(ctx, Request{Model: bedrockModel, Messages: []Message{{Role: RoleUser, Content: "hi"}}}) + if err != context.Canceled { + t.Fatalf("expected context.Canceled, got %v", err) + } +} diff --git a/internal/provider/provider.go b/internal/provider/provider.go index f585940..2d78b82 100644 --- a/internal/provider/provider.go +++ b/internal/provider/provider.go @@ -3,19 +3,15 @@ // outbound network calls. Each Chat call returns token Usage so the deterministic // governor and cost ledger can attribute spend to a run. // -// v0.1 implements Anthropic natively; OpenAI, Bedrock, and Ollama are stubs that -// return ErrNotImplemented. The interface is intentionally provider-neutral so -// the gateway and governor never special-case a vendor. +// Anthropic, OpenAI, Ollama, and AWS Bedrock are implemented natively. The +// interface is intentionally provider-neutral so the gateway and governor never +// special-case a vendor; the long tail is fronted via LiteLLM (see docs). package provider import ( "context" - "errors" ) -// ErrNotImplemented is returned by stub providers that are not yet wired up. -var ErrNotImplemented = errors.New("provider: not implemented in v0.1") - // Role identifies the author of a message. type Role string diff --git a/internal/provider/sigv4.go b/internal/provider/sigv4.go new file mode 100644 index 0000000..4592404 --- /dev/null +++ b/internal/provider/sigv4.go @@ -0,0 +1,131 @@ +package provider + +import ( + "crypto/hmac" + "crypto/sha256" + "encoding/hex" + "fmt" + "net/http" + "sort" + "strings" + "time" +) + +// AWS Signature Version 4 signing, hand-rolled so the Bedrock provider needs no +// AWS SDK dependency. Scope is exactly what Bedrock's runtime needs: a signed POST +// (or GET) with host + x-amz-date, plus content-type and x-amz-security-token when +// present. Verified against the AWS SigV4 "get-vanilla" known-answer test vector +// (see sigv4_test.go). + +// awsCreds are the AWS credentials used to sign a request. sessionToken is set only +// for temporary (STS) credentials. +type awsCreds struct { + accessKey string + secretKey string + sessionToken string +} + +// signV4 signs req in place with AWS SigV4, setting the X-Amz-Date, +// X-Amz-Security-Token (when temporary creds), and Authorization headers. body is +// the exact request body that will be sent (nil/empty for a GET). The caller must +// have set req.URL such that EscapedPath() equals the path that will be sent on the +// wire (so the signed canonical URI matches the request) — see bedrock.go. +func signV4(req *http.Request, body []byte, c awsCreds, region, service string, t time.Time) { + t = t.UTC() + amzDate := t.Format("20060102T150405Z") + dateStamp := t.Format("20060102") + + req.Header.Set("X-Amz-Date", amzDate) + if c.sessionToken != "" { + req.Header.Set("X-Amz-Security-Token", c.sessionToken) + } + + // The headers we sign: host and x-amz-date always; content-type and the + // session token when present. AWS requires lowercase names, sorted. + headers := map[string]string{ + "host": req.URL.Host, + "x-amz-date": amzDate, + } + if ct := req.Header.Get("Content-Type"); ct != "" { + headers["content-type"] = ct + } + if c.sessionToken != "" { + headers["x-amz-security-token"] = c.sessionToken + } + names := make([]string, 0, len(headers)) + for n := range headers { + names = append(names, n) + } + sort.Strings(names) + + var canonHeaders strings.Builder + for _, n := range names { + canonHeaders.WriteString(n) + canonHeaders.WriteByte(':') + canonHeaders.WriteString(strings.TrimSpace(headers[n])) + canonHeaders.WriteByte('\n') + } + signedHeaders := strings.Join(names, ";") + + canonicalURI := req.URL.EscapedPath() + if canonicalURI == "" { + canonicalURI = "/" + } + + canonicalRequest := strings.Join([]string{ + req.Method, + canonicalURI, + req.URL.RawQuery, // empty for our calls; assumed already canonical otherwise + canonHeaders.String(), + signedHeaders, + hexSHA256(body), + }, "\n") + + scope := dateStamp + "/" + region + "/" + service + "/aws4_request" + stringToSign := strings.Join([]string{ + "AWS4-HMAC-SHA256", + amzDate, + scope, + hexSHA256([]byte(canonicalRequest)), + }, "\n") + + kDate := hmacSHA256([]byte("AWS4"+c.secretKey), dateStamp) + kRegion := hmacSHA256(kDate, region) + kService := hmacSHA256(kRegion, service) + kSigning := hmacSHA256(kService, "aws4_request") + signature := hex.EncodeToString(hmacSHA256(kSigning, stringToSign)) + + req.Header.Set("Authorization", "AWS4-HMAC-SHA256 "+ + "Credential="+c.accessKey+"/"+scope+", "+ + "SignedHeaders="+signedHeaders+", "+ + "Signature="+signature) +} + +func hmacSHA256(key []byte, data string) []byte { + h := hmac.New(sha256.New, key) + h.Write([]byte(data)) + return h.Sum(nil) +} + +func hexSHA256(b []byte) string { + sum := sha256.Sum256(b) + return hex.EncodeToString(sum[:]) +} + +// awsURIEncodeSegment percent-encodes a single path segment per AWS's UriEncode +// rules: every byte except the unreserved set (A-Z a-z 0-9 - _ . ~) is encoded. +// This is stricter than Go's url.PathEscape (which leaves ':' alone), and Bedrock +// model ids contain ':' (e.g. ...-v1:0) — so the signed path matches the wire path. +func awsURIEncodeSegment(s string) string { + const unreserved = "ABCDEFGHIJKLMNOPQRSTUVWXYZabcdefghijklmnopqrstuvwxyz0123456789-_.~" + var b strings.Builder + for i := 0; i < len(s); i++ { + ch := s[i] + if strings.IndexByte(unreserved, ch) >= 0 { + b.WriteByte(ch) + } else { + fmt.Fprintf(&b, "%%%02X", ch) + } + } + return b.String() +} diff --git a/internal/provider/sigv4_test.go b/internal/provider/sigv4_test.go new file mode 100644 index 0000000..793c014 --- /dev/null +++ b/internal/provider/sigv4_test.go @@ -0,0 +1,52 @@ +package provider + +import ( + "net/http" + "testing" + "time" +) + +// TestSignV4_GetVanilla checks the signer against AWS's canonical SigV4 +// "get-vanilla" known-answer test vector. The expected signature is AWS's +// published value (independently reproduced with openssl), so this catches any +// drift in the canonical request, string-to-sign, or key-derivation steps. +func TestSignV4_GetVanilla(t *testing.T) { + req, err := http.NewRequest(http.MethodGet, "https://example.amazonaws.com/", nil) + if err != nil { + t.Fatal(err) + } + creds := awsCreds{ + accessKey: "AKIDEXAMPLE", + secretKey: "wJalrXUtnFEMI/K7MDENG+bPxRfiCYEXAMPLEKEY", + } + when := time.Date(2015, time.August, 30, 12, 36, 0, 0, time.UTC) + + signV4(req, nil, creds, "us-east-1", "service", when) + + if got := req.Header.Get("X-Amz-Date"); got != "20150830T123600Z" { + t.Errorf("X-Amz-Date = %q", got) + } + const want = "AWS4-HMAC-SHA256 " + + "Credential=AKIDEXAMPLE/20150830/us-east-1/service/aws4_request, " + + "SignedHeaders=host;x-amz-date, " + + "Signature=5fa00fa31553b73ebf1942676e86291e8372ff2a2260956d9b8aae1d763fbf31" + if got := req.Header.Get("Authorization"); got != want { + t.Errorf("Authorization mismatch:\n got %q\nwant %q", got, want) + } +} + +func TestAWSURIEncodeSegment(t *testing.T) { + // Unreserved chars pass through; everything else (notably ':' in Bedrock model + // ids) is percent-encoded uppercase. + cases := map[string]string{ + "anthropic.claude-3-5-sonnet-20240620-v1:0": "anthropic.claude-3-5-sonnet-20240620-v1%3A0", + "a_b.c-d~e": "a_b.c-d~e", + "x/y": "x%2Fy", + "a b": "a%20b", + } + for in, want := range cases { + if got := awsURIEncodeSegment(in); got != want { + t.Errorf("awsURIEncodeSegment(%q) = %q, want %q", in, got, want) + } + } +} diff --git a/internal/provider/stream.go b/internal/provider/stream.go index 262880a..6fe423c 100644 --- a/internal/provider/stream.go +++ b/internal/provider/stream.go @@ -4,8 +4,8 @@ import "context" // Streamer is the optional interface a provider implements to support streaming // chat. The gateway type-asserts for it; a provider that doesn't implement it -// (Anthropic, Ollama, the stubs, for now) makes a streaming request fall back to a -// clear "unsupported" error rather than silently buffering. +// (Ollama and Bedrock, for now) makes a streaming request fall back to a clear +// "unsupported" error rather than silently buffering. type Streamer interface { // ChatStream starts a streaming completion. The returned ChatStream yields the // provider's raw SSE chunks verbatim (so the client receives authentic, diff --git a/internal/provider/stubs.go b/internal/provider/stubs.go deleted file mode 100644 index e36a15e..0000000 --- a/internal/provider/stubs.go +++ /dev/null @@ -1,18 +0,0 @@ -package provider - -import "context" - -// The provider below is a stub. It satisfies the Provider interface so routing and -// config can reference it, but Chat returns ErrNotImplemented until it is built -// out. For the long tail of providers, the documented path is to front RiskKernel -// with LiteLLM rather than reimplement 100+ vendors here. (Anthropic, OpenAI, and -// Ollama are implemented natively — see anthropic.go, openai.go, ollama.go.) - -// Bedrock is a stub. Native implementation is planned. -type Bedrock struct{} - -func NewBedrock() *Bedrock { return &Bedrock{} } -func (b *Bedrock) Name() string { return "bedrock" } -func (b *Bedrock) Chat(context.Context, Request) (*Response, error) { - return nil, ErrNotImplemented -}