Send an upper-bound prompt estimate with each credit hold

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 {

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)

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
+}

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)
+	}
+}

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)

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
 }