diff --git a/internal/control/catalog/catalog.go b/internal/control/catalog/catalog.go index 35a4552..74452ce 100644 --- a/internal/control/catalog/catalog.go +++ b/internal/control/catalog/catalog.go @@ -87,6 +87,39 @@ type Model struct { // RoutingProfile selects how the router scores endpoints for this model. RoutingProfile string `koanf:"routing_profile"` SLO SLO `koanf:"slo"` + + // ContextLength is how many tokens one request may use, prompt and + // completion together, and MaxOutput bounds the completion alone. Zero + // means unknown, which is what every model was until now. + // + // The gateway knew neither, so an over-long prompt was not refused: it was + // admitted, given a replica, and failed at the engine, and a client sizing + // its prompts from /v1/models had nothing to read. Worse, the number is not + // the one an operator writes in the engine's arguments — llama.cpp divides + // --ctx-size across its parallel slots, so a pool of 8192 across four slots + // is 2048 a request, and nothing said so. + ContextLength int `koanf:"context_length"` + MaxOutput int `koanf:"max_output"` + + // Capabilities is what this model can do beyond plain completion. + Capabilities Capabilities `koanf:"capabilities"` +} + +// Capabilities describes what a client may ask a model for. +// +// Declared by the operator rather than discovered, because there is no way to +// ask an engine what its model can do. That makes these claims, and a wrong +// claim is worse than none — which is why `ctl engine-check` verifies tool +// calling against a running engine rather than trusting this. +// +// The reason it exists: the gateway forwards a tools array to any model, and a +// model that cannot use them answers in prose. Every OpenAI-compatible client +// then behaves as though the model is broken, and the only way to find out was +// to try it. That cost an afternoon. +type Capabilities struct { + Tools bool `koanf:"tools"` + Vision bool `koanf:"vision"` + Embeddings bool `koanf:"embeddings"` } // Catalog is the set of models the cluster should serve. diff --git a/internal/router/context_limit.go b/internal/router/context_limit.go new file mode 100644 index 0000000..40cbb55 --- /dev/null +++ b/internal/router/context_limit.go @@ -0,0 +1,60 @@ +package router + +import ( + "fmt" + + "gitea.mynodes.xyz/ilker/llmgw/internal/control/catalog" +) + +// checkContextFits reports why a request cannot fit in a model's context, or +// "" when it can. +// +// Refused here rather than at the engine. An over-long prompt used to be +// admitted, scheduled, given a replica and a slot, and only then rejected — +// paying the whole cost of the request to learn something knowable from the +// request itself. The engine's error also arrives in whatever shape that engine +// chose, which is not something a client can act on. +// +// The prompt estimate is bytes over four, the same crude stand-in used for +// scoring, so the comparison is approximate in one direction: it can let +// through a request the engine will refuse. That is the right way round. A +// tokenizer would make it exact, and needing one is why plan §7.5 tier 2 is v2. +// To keep an estimate that reads slightly low from refusing valid requests, a +// margin is allowed before anything is rejected. +func checkContextFits(m catalog.Model, promptTokens, maxTokens int) string { + if m.ContextLength <= 0 { + // Unknown, which is every model that has not declared one. Saying + // nothing beats guessing a limit and refusing work over it. + return "" + } + + if m.MaxOutput > 0 && maxTokens > m.MaxOutput { + return fmt.Sprintf("max_tokens is %d, but %q allows at most %d", + maxTokens, m.Name, m.MaxOutput) + } + + // Only the prompt is estimated; max_tokens is exact when given. A request + // that does not say how much it wants back is judged on its prompt alone, + // because assuming the whole remaining window would refuse most requests. + need := promptTokens + maxTokens + allowed := m.ContextLength + contextEstimateMargin(m.ContextLength) + if need <= allowed { + return "" + } + if maxTokens > 0 { + return fmt.Sprintf( + "about %d prompt tokens plus max_tokens %d exceeds the %d-token context of %q", + promptTokens, maxTokens, m.ContextLength, m.Name) + } + return fmt.Sprintf("the prompt is about %d tokens, over the %d-token context of %q", + promptTokens, m.ContextLength, m.Name) +} + +// contextEstimateMargin is the slack allowed for the bytes-over-four estimate. +// +// That estimate is wrong in both directions depending on the text — code and +// non-Latin scripts tokenize very differently from English prose — and being +// wrong by a few percent must not turn a valid request into a refusal. Ten +// percent is enough to absorb that while still catching a prompt that is +// genuinely too big, which is the case worth catching. +func contextEstimateMargin(context int) int { return context / 10 } diff --git a/internal/router/context_limit_test.go b/internal/router/context_limit_test.go new file mode 100644 index 0000000..4295b0a --- /dev/null +++ b/internal/router/context_limit_test.go @@ -0,0 +1,105 @@ +package router + +import ( + "encoding/json" + "strings" + "testing" + + "gitea.mynodes.xyz/ilker/llmgw/internal/control/catalog" +) + +func TestAModelWithNoDeclaredContextRefusesNothing(t *testing.T) { + // Every model was in this state until now, and guessing a limit would turn + // working requests into errors. + m := catalog.Model{Name: "m"} + if why := checkContextFits(m, 1_000_000, 1_000_000); why != "" { + t.Errorf("refused a request against an unknown context: %q", why) + } +} + +func TestAPromptOverTheContextIsRefused(t *testing.T) { + m := catalog.Model{Name: "coder", ContextLength: 2048} + why := checkContextFits(m, 5000, 0) + if why == "" { + t.Fatal("a prompt well over the context was allowed") + } + for _, want := range []string{"5000", "2048", "coder"} { + if !strings.Contains(why, want) { + t.Errorf("the message does not mention %s: %q", want, why) + } + } +} + +func TestPromptPlusMaxTokensMustFitTogether(t *testing.T) { + // The failure people actually hit: a prompt that fits, asking for a + // completion that cannot follow it. + m := catalog.Model{Name: "coder", ContextLength: 2048} + if why := checkContextFits(m, 1800, 1000); why == "" { + t.Error("prompt plus max_tokens over the window was allowed") + } + if why := checkContextFits(m, 1000, 500); why != "" { + t.Errorf("a request that fits was refused: %q", why) + } +} + +func TestTheEstimateIsGivenRoomToBeWrong(t *testing.T) { + // Bytes over four is wrong in both directions — code and non-Latin scripts + // tokenize nothing like English prose — and being a few percent out must + // not turn a valid request into a refusal. + m := catalog.Model{Name: "m", ContextLength: 1000} + if why := checkContextFits(m, 1050, 0); why != "" { + t.Errorf("a request just over the estimate was refused: %q", why) + } + if why := checkContextFits(m, 2000, 0); why == "" { + t.Error("a prompt at twice the context was allowed; the margin is too wide") + } +} + +func TestMaxOutputIsCheckedOnItsOwn(t *testing.T) { + // Asking for more output than the model can produce is worth catching even + // when the total would fit. + m := catalog.Model{Name: "m", ContextLength: 32768, MaxOutput: 4096} + why := checkContextFits(m, 10, 8000) + if why == "" { + t.Fatal("max_tokens over max_output was allowed") + } + if !strings.Contains(why, "4096") { + t.Errorf("the message does not say the limit: %q", why) + } +} + +func TestTheTokenEstimateCoversTheWholePrompt(t *testing.T) { + // It used to be taken from the routing prefix, which stops at prefixBytes, + // so every prompt over 2 KB estimated as 512 tokens however long it was. + // TPM limits therefore undercounted exactly the requests that consume most, + // and no context check could work. + long := strings.Repeat("x", 40_000) + b := &inferenceBody{Messages: []struct { + Role string `json:"role"` + Content json.RawMessage `json:"content"` + }{{Role: "user", Content: json.RawMessage(`"` + long + `"`)}}} + + got := promptBytes(b) + if got < 40_000 { + t.Fatalf("promptBytes = %d for a 40,000 byte prompt", got) + } + if prefix := len(promptPrefix(b)); prefix > prefixBytes { + t.Errorf("the routing prefix grew past its bound: %d", prefix) + } + // The two must stay different things: bounded for routing, whole for limits. + if got <= prefixBytes { + t.Error("the estimate is still bounded by the routing prefix") + } +} + +func TestStructuredContentIsCounted(t *testing.T) { + // Multimodal messages carry an array of parts rather than a string. Counting + // zero for those would let a large request through unmeasured. + b := &inferenceBody{Messages: []struct { + Role string `json:"role"` + Content json.RawMessage `json:"content"` + }{{Role: "user", Content: json.RawMessage(`[{"type":"text","text":"` + strings.Repeat("y", 5000) + `"}]`)}}} + if got := promptBytes(b); got < 5000 { + t.Errorf("structured content counted as %d bytes", got) + } +} diff --git a/internal/router/handlers.go b/internal/router/handlers.go index 4e859c3..05c7f45 100644 --- a/internal/router/handlers.go +++ b/internal/router/handlers.go @@ -195,6 +195,17 @@ func (r *Router) handleInference(path string) http.HandlerFunc { return } + // Before authentication work and before admission: a request that + // cannot fit should not consume a quota reservation or a queue slot on + // its way to being refused. + if _, entry := r.resolve(parsed.Model); entry.ContextLength > 0 { + if why := checkContextFits(entry, parsed.PromptTokens, parsed.MaxTokens); why != "" { + writeError(w, http.StatusBadRequest, "invalid_request_error", why) + r.m.observeRequest(parsed.Model, "", http.StatusBadRequest, time.Since(start), nil, 0) + return + } + } + if p, ok := auth.FromContext(req.Context()); ok { if !r.permits(p, parsed.Model) { auth.WriteForbiddenModel(w, parsed.Model) @@ -373,8 +384,10 @@ func (r *Router) parse(req *http.Request, path string) (*parsedRequest, error) { } // Bytes over four: a crude stand-in for a tokenizer, good enough for - // scoring. It becomes exact only with plan §7.5 tier 2 (see plan §14). - p.PromptTokens = len(p.PromptPrefix) / 4 + // scoring and for limits. It becomes exact only with plan §7.5 tier 2 (see + // plan §14). Measured over the whole prompt, not the routing prefix — + // see promptBytes. + p.PromptTokens = promptBytes(&b) / 4 // Ask the engine for usage the client did not request, so accounting works // regardless of client behaviour; the extra chunk is dropped on the way back. @@ -432,6 +445,40 @@ func promptPrefix(b *inferenceBody) string { return s } +// promptBytes measures the whole prompt, which promptPrefix deliberately does +// not. +// +// The two were the same function for too long, and the estimate was taken from +// the prefix: bytes over four of at most prefixBytes, so every prompt above +// 2 KB counted as 512 tokens however long it really was. Prefix routing wants a +// bounded head of the conversation; a token estimate wants all of it, and using +// one for the other meant TPM limits undercounted exactly the requests that +// consume the most, and nothing could tell whether a prompt fit in a context. +func promptBytes(b *inferenceBody) int { + n := 0 + add := func(raw json.RawMessage) { + if len(raw) == 0 { + return + } + var str string + if err := json.Unmarshal(raw, &str); err == nil { + n += len(str) + return + } + // Not a plain string: structured content, an array of parts, an image. + // The raw JSON overstates the text slightly, which is the safe + // direction for a limit. + n += len(raw) + } + for _, m := range b.Messages { + n += len(m.Role) + 2 + add(m.Content) + } + add(b.Prompt) + add(b.Input) + return n +} + // injectIncludeUsage sets stream_options.include_usage on the outgoing body. // // This is the one place the router rewrites a request, so it is done by @@ -556,9 +603,23 @@ func (r *Router) handleModels(w http.ResponseWriter, _ *http.Request) { if r.models[n].Simulated() { owner = "llmgw-sim" } - data = append(data, map[string]any{ + entry := map[string]any{ "id": n, "object": "model", "created": time.Now().Unix(), "owned_by": owner, - }) + } + // Only when known. A context_length of 0 would read as "no context", + // which is a worse answer than saying nothing. + m := r.models[n] + if m.ContextLength > 0 { + entry["context_length"] = m.ContextLength + } + if m.MaxOutput > 0 { + entry["max_output_tokens"] = m.MaxOutput + } + entry["capabilities"] = map[string]bool{ + "tools": m.Capabilities.Tools, "vision": m.Capabilities.Vision, + "embeddings": m.Capabilities.Embeddings, + } + data = append(data, entry) } writeJSON(w, http.StatusOK, map[string]any{"object": "list", "data": data}) } diff --git a/internal/router/router_test.go b/internal/router/router_test.go index 8627a50..2d51051 100644 --- a/internal/router/router_test.go +++ b/internal/router/router_test.go @@ -415,3 +415,89 @@ func TestTheErrorIsHonestWhenNothingIsServed(t *testing.T) { t.Errorf("an empty fleet is not explained: %q", msg) } } + +func TestAnOverLongRequestIsRefusedBeforeItReachesAnEngine(t *testing.T) { + // The whole point of checking here: an over-long prompt used to be + // admitted, scheduled, given a replica and a slot, and refused by the + // engine — paying the cost of the request to learn something knowable from + // the request itself. + e := newFakeEngine(t, "m") + cfg := baseConfig(config.Endpoint{ID: "a", URL: e.srv.URL, Models: []string{"m"}}) + m := catModel("m", nil) + m.ContextLength = 2048 + cfg.Models = []catalog.Model{m} + rt := newTestRouter(t, cfg) + + before := e.calls.Load() + rec := post(t, rt.Handler(), map[string]any{ + "model": "m", + "messages": []any{map[string]any{"role": "user", "content": strings.Repeat("x", 40_000)}}, + }, nil) + + if rec.Code != http.StatusBadRequest { + t.Fatalf("status %d, expected 400; body %s", rec.Code, rec.Body.String()) + } + if got := e.calls.Load(); got != before { + t.Errorf("the engine was called %d times for a request that could not fit", got-before) + } + if !strings.Contains(rec.Body.String(), "2048") { + t.Errorf("the error does not name the limit: %s", rec.Body.String()) + } +} + +func TestModelsEndpointPublishesContextAndCapabilities(t *testing.T) { + // A client sizing its prompts had nothing to read, and a client deciding + // whether to send tools had nothing either. + e := newFakeEngine(t, "m") + cfg := baseConfig(config.Endpoint{ID: "a", URL: e.srv.URL, Models: []string{"m"}}) + m := catModel("m", nil) + m.ContextLength, m.MaxOutput = 32768, 4096 + m.Capabilities = catalog.Capabilities{Tools: true} + cfg.Models = []catalog.Model{m} + rt := newTestRouter(t, cfg) + + req := httptest.NewRequest(http.MethodGet, "/v1/models", nil) + rec := httptest.NewRecorder() + rt.Handler().ServeHTTP(rec, req) + + var out struct { + Data []struct { + ID string `json:"id"` + Context int `json:"context_length"` + MaxOutput int `json:"max_output_tokens"` + Capabilities map[string]bool `json:"capabilities"` + } `json:"data"` + } + if err := json.Unmarshal(rec.Body.Bytes(), &out); err != nil { + t.Fatal(err) + } + var found bool + for _, d := range out.Data { + if d.ID != "m" { + continue + } + found = true + if d.Context != 32768 || d.MaxOutput != 4096 { + t.Errorf("context %d, max output %d", d.Context, d.MaxOutput) + } + if !d.Capabilities["tools"] || d.Capabilities["vision"] { + t.Errorf("capabilities wrong: %v", d.Capabilities) + } + } + if !found { + t.Fatal("the model was not listed") + } +} + +func TestAModelWithNoContextPublishesNoContextField(t *testing.T) { + // Zero would read as "no context", which is worse than saying nothing. + e := newFakeEngine(t, "m") + rt := newTestRouter(t, baseConfig(config.Endpoint{ID: "a", URL: e.srv.URL, Models: []string{"m"}})) + + req := httptest.NewRequest(http.MethodGet, "/v1/models", nil) + rec := httptest.NewRecorder() + rt.Handler().ServeHTTP(rec, req) + if strings.Contains(rec.Body.String(), "context_length") { + t.Errorf("an unknown context was published as a number: %s", rec.Body.String()) + } +}