From c2b3cab8db7caf4d9af2637690a3fc2a9f93bd13 Mon Sep 17 00:00:00 2001 From: Marketen Date: Fri, 30 Oct 2026 09:15:36 +0100 Subject: [PATCH] feat: send an upper-bound prompt estimate with each credit hold Metering used to hold the model's whole context window for every prompt. The gateway now sends an upper bound of the prompt's tokens: the bytes of the body sent upstream, plus 1,024 for what providers add, with each image counted as 16,384 tokens instead of its bytes. Tokenizers never produce more tokens than UTF-8 bytes, so the bound holds for text. Requests with audio, video, files or unknown content parts, and SystemOne requests, send no estimate, and Metering holds the full context window as before. Co-Authored-By: Claude Opus 5.5 --- .../internal/adapters/metering/client.go | 2 + .../internal/adapters/metering/client_test.go | 31 ++++ .../application/services/input_estimate.go | 79 ++++++++++ .../services/input_estimate_test.go | 146 ++++++++++++++++++ .../internal/application/services/proxy.go | 4 + pkg/domain/generation.go | 11 +- 6 files changed, 269 insertions(+), 4 deletions(-) create mode 100644 apps/gateway/internal/application/services/input_estimate.go create mode 100644 apps/gateway/internal/application/services/input_estimate_test.go diff --git a/apps/gateway/internal/adapters/metering/client.go b/apps/gateway/internal/adapters/metering/client.go index f95852e..6407457 100644 --- a/apps/gateway/internal/adapters/metering/client.go +++ b/apps/gateway/internal/adapters/metering/client.go @@ -58,6 +58,7 @@ func (c *Client) Reserve(ctx context.Context, auth domain.AuthContext, endpoint FallbackUsed: req.FallbackUsed, Stream: req.Stream, MaxOutputTokens: req.MaxOutputTokens, + EstimatedInputTokens: req.EstimatedInputTokens, } var res reserveResponse sent := time.Now() // A lease measured from here never outlasts Metering's. @@ -420,6 +421,7 @@ type reserveRequest struct { FallbackUsed *bool `json:"fallback_used,omitempty"` Stream bool `json:"stream"` MaxOutputTokens *int `json:"max_output_tokens,omitempty"` + EstimatedInputTokens *int64 `json:"estimated_input_tokens,omitempty"` } type reserveResponse struct { diff --git a/apps/gateway/internal/adapters/metering/client_test.go b/apps/gateway/internal/adapters/metering/client_test.go index c7c30fa..bd19297 100644 --- a/apps/gateway/internal/adapters/metering/client_test.go +++ b/apps/gateway/internal/adapters/metering/client_test.go @@ -38,6 +38,37 @@ func TestClientReserve_MapsInsufficientBalance(t *testing.T) { } } +func TestClientReserve_SendsHoldSizing(t *testing.T) { + var payload map[string]any + server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + if err := json.NewDecoder(r.Body).Decode(&payload); err != nil { + t.Fatalf("decode reserve payload: %v", err) + } + _, _ = w.Write([]byte(`{"reservation_id":"r-1"}`)) + })) + defer server.Close() + + estimate, maxTokens := int64(9_216), 512 + client := NewClient(server.URL, "token", time.Second) + req := domain.GenerateRequest{PublicModelID: "model", EstimatedInputTokens: &estimate, MaxOutputTokens: &maxTokens} + if _, err := client.Reserve(context.Background(), testAuth(), domain.EndpointChatCompletions, req, testModel(), "gw-1"); err != nil { + t.Fatalf("Reserve: %v", err) + } + if payload["estimated_input_tokens"] != float64(9_216) || payload["max_output_tokens"] != float64(512) { + t.Fatalf("hold sizing = estimate %v/max output %v, want 9216/512", payload["estimated_input_tokens"], payload["max_output_tokens"]) + } + + // Without an estimate the field is omitted, so Metering holds the whole + // context window. + payload = nil + if _, err := client.Reserve(context.Background(), testAuth(), domain.EndpointChatCompletions, domain.GenerateRequest{PublicModelID: "model"}, testModel(), "gw-2"); err != nil { + t.Fatalf("Reserve without estimate: %v", err) + } + if _, present := payload["estimated_input_tokens"]; present { + t.Fatalf("payload = %v, want no estimated_input_tokens", payload) + } +} + func TestClientReserve_MapsInternalContractErrorsToInternalError(t *testing.T) { server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, _ *http.Request) { w.WriteHeader(http.StatusConflict) diff --git a/apps/gateway/internal/application/services/input_estimate.go b/apps/gateway/internal/application/services/input_estimate.go new file mode 100644 index 0000000..a4e7caf --- /dev/null +++ b/apps/gateway/internal/application/services/input_estimate.go @@ -0,0 +1,79 @@ +package services + +import ( + "bytes" + "encoding/json" + + "github.com/dappnode/dappnode-nexus-gateway/pkg/domain" +) + +// Metering holds credit for a request's prompt before the provider runs it. +// The hold needs an upper bound, not a guess: the prompt is charged at its +// real size when the request settles, and an estimate that is too low lets a +// request start that the account may not be able to pay for. +// +// The bound is the body's size in bytes. The tokenizers behind the models +// Nexus serves (byte-level BPE, and SentencePiece with byte fallback) never +// produce more tokens than the UTF-8 bytes of the text they encode. The JSON +// body carries every prompt string at least once, plus quotes, keys and +// escapes, so its length also covers the role markers the chat template adds +// per message. promptOverheadTokens covers what a provider adds that is not in +// the body: rendered tool instructions, a default system prompt, the upstream +// model name. +// +// Media is not text. An image part counts imageTokens instead of its bytes: an +// inline base64 image is far longer in bytes than in tokens, and a linked one +// far shorter. Audio, video, files, and part types the gateway does not know +// have no bound it can compute, so those requests get no estimate and the hold +// covers the model's whole context window, as it did before. +const ( + // imageTokens is above what any model Nexus serves spends on one image. + imageTokens = 16_384 + promptOverheadTokens = 1_024 +) + +// estimateInputTokens returns an upper bound on the prompt tokens of body, +// the request exactly as it goes upstream. ok is false when there is no bound. +func estimateInputTokens(endpoint string, body []byte) (tokens int64, ok bool) { + // SystemOne's provider builds its own prompt from the request, so its + // size is unknown here. + if endpoint != domain.EndpointChatCompletions { + return 0, false + } + var req struct { + Messages []struct { + Content json.RawMessage `json:"content"` + } `json:"messages"` + } + if err := json.Unmarshal(body, &req); err != nil { + return 0, false + } + tokens = int64(len(body)) + promptOverheadTokens + for _, message := range req.Messages { + content := bytes.TrimSpace(message.Content) + if len(content) == 0 || content[0] != '[' { + continue // A string or null: text, already counted. + } + var parts []json.RawMessage + if err := json.Unmarshal(content, &parts); err != nil { + return 0, false + } + for _, part := range parts { + var p struct { + Type string `json:"type"` + } + if err := json.Unmarshal(part, &p); err != nil { + return 0, false + } + switch p.Type { + case "text", "input_text", "output_text", "refusal": + // Text, already counted in the body's bytes. + case "image_url", "image", "input_image": + tokens += imageTokens - int64(len(part)) + default: + return 0, false + } + } + } + return max(tokens, 1), true +} diff --git a/apps/gateway/internal/application/services/input_estimate_test.go b/apps/gateway/internal/application/services/input_estimate_test.go new file mode 100644 index 0000000..f5298da --- /dev/null +++ b/apps/gateway/internal/application/services/input_estimate_test.go @@ -0,0 +1,146 @@ +package services + +import ( + "context" + "strings" + "testing" + + "github.com/dappnode/dappnode-nexus-gateway/pkg/domain" +) + +func TestEstimateInputTokens(t *testing.T) { + inlineImage := `{"type":"image_url","image_url":{"url":"data:image/png;base64,` + strings.Repeat("A", 200_000) + `"}}` + linkedImage := `{"type":"image_url","image_url":{"url":"https://example.com/cat.png"}}` + text := `{"type":"text","text":"what is this?"}` + + cases := []struct { + name string + body string + // extra is what the estimate adds to the body's bytes; ok=false means + // no bound. + extra func(body string) int64 + ok bool + }{ + { + name: "string content counts every byte", + body: `{"model":"m","messages":[{"role":"system","content":"be brief"},{"role":"user","content":"hello"}]}`, + extra: func(string) int64 { return 0 }, + ok: true, + }, + { + name: "text parts count every byte", + body: `{"model":"m","messages":[{"role":"user","content":[` + text + `,{"type":"input_text","text":"x"}]}]}`, + extra: func(string) int64 { return 0 }, + ok: true, + }, + { + name: "null content and tool calls", + body: `{"model":"m","messages":[{"role":"assistant","content":null,"tool_calls":[{"id":"1","type":"function","function":{"name":"f","arguments":"{}"}}]},{"role":"tool","tool_call_id":"1","content":"ok"}]}`, + extra: func(string) int64 { return 0 }, + ok: true, + }, + { + name: "inline image counts a fixed allowance instead of its bytes", + body: `{"model":"m","messages":[{"role":"user","content":[` + text + `,` + inlineImage + `]}]}`, + extra: func(string) int64 { return imageTokens - int64(len(inlineImage)) }, + ok: true, + }, + { + name: "linked image counts the allowance too", + body: `{"model":"m","messages":[{"role":"user","content":[` + linkedImage + `]}]}`, + extra: func(string) int64 { return imageTokens - int64(len(linkedImage)) }, + ok: true, + }, + { + name: "every image counts", + body: `{"model":"m","messages":[{"role":"user","content":[` + linkedImage + `]},{"role":"tool","tool_call_id":"1","content":[` + linkedImage + `]}]}`, + extra: func(string) int64 { return 2 * (imageTokens - int64(len(linkedImage))) }, + ok: true, + }, + { + name: "audio has no bound", + body: `{"model":"m","messages":[{"role":"user","content":[{"type":"input_audio","input_audio":{"data":"AAAA","format":"wav"}}]}]}`, + }, + { + name: "video has no bound", + body: `{"model":"m","messages":[{"role":"user","content":[{"type":"video_url","video_url":{"url":"https://example.com/a.mp4"}}]}]}`, + }, + { + name: "files have no bound", + body: `{"model":"m","messages":[{"role":"user","content":[{"type":"file","file":{"file_id":"f-1"}}]}]}`, + }, + { + name: "an unknown part type has no bound", + body: `{"model":"m","messages":[{"role":"user","content":[{"type":"hologram"}]}]}`, + }, + { + name: "a part without a type has no bound", + body: `{"model":"m","messages":[{"role":"user","content":[{"text":"hi"}]}]}`, + }, + { + name: "a part that is not an object has no bound", + body: `{"model":"m","messages":[{"role":"user","content":["hi"]}]}`, + }, + { + name: "invalid JSON has no bound", + body: `{"model":"m","messages":[`, + }, + } + for _, tc := range cases { + t.Run(tc.name, func(t *testing.T) { + got, ok := estimateInputTokens(domain.EndpointChatCompletions, []byte(tc.body)) + if ok != tc.ok { + t.Fatalf("ok = %v, want %v", ok, tc.ok) + } + if !ok { + return + } + want := int64(len(tc.body)) + promptOverheadTokens + tc.extra(tc.body) + if got != want { + t.Fatalf("estimate = %d, want %d", got, want) + } + }) + } +} + +func TestEstimateInputTokens_InlineImageDoesNotHoldItsBytes(t *testing.T) { + image := `{"type":"image_url","image_url":{"url":"data:image/jpeg;base64,` + strings.Repeat("A", 2_000_000) + `"}}` + body := `{"model":"m","messages":[{"role":"user","content":[` + image + `]}]}` + got, ok := estimateInputTokens(domain.EndpointChatCompletions, []byte(body)) + if !ok || got > imageTokens+promptOverheadTokens+100 { + t.Fatalf("estimate = %d (ok=%v), want about one image allowance, not 2M bytes", got, ok) + } +} + +func TestEstimateInputTokens_SystemOneHasNoBound(t *testing.T) { + if _, ok := estimateInputTokens(domain.EndpointSystemOne, []byte(`{"model":"m","input":"hi"}`)); ok { + t.Fatal("systemone builds its own prompt, so it must not get an estimate") + } +} + +func TestProxy_SendsThePromptEstimateWithTheHold(t *testing.T) { + meter := &stubUsageMeter{} + svc := newDirectModelGenerateService(meter, &stubProvider{}) + body := `{"model":"model","messages":[{"role":"user","content":"hello"}]}` + + if _, err := svc.Proxy(context.Background(), domain.EndpointChatCompletions, []byte(body), domain.GenerateRequest{PublicModelID: "model"}, "sk-test"); err != nil { + t.Fatalf("Proxy: %v", err) + } + estimate := meter.lastReserveReq.EstimatedInputTokens + if want := int64(len(body)) + promptOverheadTokens; estimate == nil || *estimate != want { + t.Fatalf("estimate sent with the hold = %v, want %d", estimate, want) + } +} + +func TestProxy_SendsNoEstimateWhenThePromptHasNoBound(t *testing.T) { + meter := &stubUsageMeter{} + svc := newDirectModelGenerateService(meter, &stubProvider{}) + body := `{"model":"model","messages":[{"role":"user","content":[{"type":"input_audio","input_audio":{"data":"AAAA","format":"wav"}}]}]}` + + if _, err := svc.Proxy(context.Background(), domain.EndpointChatCompletions, []byte(body), domain.GenerateRequest{PublicModelID: "model"}, "sk-test"); err != nil { + t.Fatalf("Proxy: %v", err) + } + if meter.lastReserveReq.EstimatedInputTokens != nil { + t.Fatalf("estimate = %d, want none so the hold covers the context window", *meter.lastReserveReq.EstimatedInputTokens) + } +} diff --git a/apps/gateway/internal/application/services/proxy.go b/apps/gateway/internal/application/services/proxy.go index 5dde257..b7d5982 100644 --- a/apps/gateway/internal/application/services/proxy.go +++ b/apps/gateway/internal/application/services/proxy.go @@ -83,6 +83,10 @@ func (s *GenerateService) Proxy(ctx context.Context, endpoint string, raw []byte s.recordFailure(ctx, nil, &authCtx, endpoint, &execReq, &model, err, nil, time.Since(start).Milliseconds()) return nil, reject(err, &model) } + // The hold covers the prompt as sent upstream, after masking. + if tokens, ok := estimateInputTokens(endpoint, masked); ok { + execReq.EstimatedInputTokens = &tokens + } reservationID, err := s.metering.Reserve(ctx, authCtx, endpoint, execReq, model, uuid.NewString()) if err != nil { err = clientGone(ctx, err) diff --git a/pkg/domain/generation.go b/pkg/domain/generation.go index 9daabf6..8a22524 100644 --- a/pkg/domain/generation.go +++ b/pkg/domain/generation.go @@ -27,10 +27,13 @@ type GenerateRequest struct { Input []InputItem Tools []ToolDefinition - Stream bool - MaxOutputTokens *int - ServiceTier *string - ParallelToolCalls *bool + Stream bool + MaxOutputTokens *int + // EstimatedInputTokens is an upper bound on the prompt's tokens, sized + // from the body sent upstream; nil when no bound is known. + EstimatedInputTokens *int64 + ServiceTier *string + ParallelToolCalls *bool // StructuredOutput is set when the client asks for JSON output. StructuredOutput bool }