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 }