llmgw/internal/router/context_limit_test.go

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

func TestAPictureIsCountedAsWhatAModelMakesOfItNotItsBase64(t *testing.T) {
	// A 1280x720 PNG was counted as 138 thousand tokens of base64 and an 8k
	// vision model refused it; its projector makes about a thousand of it.
	pic := `{"type":"image_url","image_url":{"url":"data:image/png;base64,` + strings.Repeat("A", 550_000) + `"}}`
	b := &inferenceBody{Messages: []struct {
		Role    string          `json:"role"`
		Content json.RawMessage `json:"content"`
	}{{Role: "user", Content: json.RawMessage(`[{"type":"text","text":"check it"},` + pic + `,` + pic + `]`)}}}
	got := promptBytes(b) / 4
	if got < 2000 || got > 2200 {
		t.Errorf("two pictures and a line estimated as %d tokens", got)
	}
	m := catalog.Model{Name: "vlm", ContextLength: 8192}
	if why := checkContextFits(m, got, 300); why != "" {
		t.Errorf("refused: %s", why)
	}
}