diff --git a/.lorx/actions/ci.yml b/.lorx/actions/ci.yml new file mode 100644 index 0000000..dbe6ed2 --- /dev/null +++ b/.lorx/actions/ci.yml @@ -0,0 +1,37 @@ +# CI run by Lorx's engine: `bit run --all` here, or a Lorx runner on push. +# GitHub Actions on this account cannot start jobs (HAL-1712), so this is the merge gate. +# Jobs run in container images so a runner needs only docker or podman, not a Go toolchain. +# GOFLAGS=-buildvcs=false: the checkout is owned by the runner's user, not the container's, +# so git refuses it and Go's VCS stamping fails the build. + +fmt: + image: golang:1.25 + env: { GOFLAGS: -buildvcs=false } + run: 'test -z "$(gofmt -l . | grep -v ^vendor/)"' + +vet: + image: golang:1.25 + env: { GOFLAGS: -buildvcs=false } + run: go vet ./... + timeout: 20m + +build: + needs: [fmt] + image: golang:1.25 + env: { GOFLAGS: -buildvcs=false } + run: go build ./... + timeout: 20m + +lint: + needs: [build] + image: golangci/golangci-lint:v2.14.0 + env: { GOFLAGS: -buildvcs=false } + run: golangci-lint run ./... + timeout: 20m + +test: + needs: [vet, build] + image: golang:1.25 + env: { GOFLAGS: -buildvcs=false } + run: go test ./... + timeout: 30m diff --git a/examples/judge/main.go b/examples/judge/main.go index d2a5e40..6804c4a 100644 --- a/examples/judge/main.go +++ b/examples/judge/main.go @@ -97,7 +97,8 @@ func main() { }) if err != nil { fmt.Fprintln(os.Stderr, "judge:", err) - os.Exit(1) + cancel() + os.Exit(1) //nolint:gocritic // cancel already ran; exiting non-zero is the point } fmt.Printf("model %s\n", result.Model) diff --git a/judge/typesafe/count_internal_test.go b/judge/typesafe/count_internal_test.go new file mode 100644 index 0000000..16aa065 --- /dev/null +++ b/judge/typesafe/count_internal_test.go @@ -0,0 +1,67 @@ +package typesafe + +import ( + "encoding/json" + "fmt" + "math/rand" + "strings" + "testing" +) + +// The chunked count must equal the serial count exactly, not +// approximately: the guard compares it against hard limits. +func TestChunkedCountEqualsSerialCount(t *testing.T) { + enc, err := cl100k() + if err != nil { + t.Fatal(err) + } + rng := rand.New(rand.NewSource(1708)) + words := []string{ + "Revenue", "was", "$12,345", "million,", "up", "4.5%", "(2021:", "11,812)", "—", "Note", "21.", + "the", "Company's", "don't", "it'll", "\n\n", "\t", " ", "café", "naïve", "東京", "£1.2bn", "10-K", + "\"quoted\"", "{json}", "a", "I", "x2", "2022", "1234567", "...", "\r\n", " indent", "émigré", + } + var texts []string + for n := 0; n < 60; n++ { + var b strings.Builder + for b.Len() < 20_000+rng.Intn(60_000) { + b.WriteString(words[rng.Intn(len(words))]) + if rng.Intn(5) > 0 { + b.WriteByte(' ') + } + } + texts = append(texts, b.String()) + } + // And the shape the guard actually counts: a marshalled request state. + state := map[string]any{"question": "What was FY2022 revenue?"} + for i := 0; i < 16; i++ { + state[fmt.Sprintf("p_%d", i)] = map[string]any{"page": i, "text": texts[i][:4500]} + } + sj, _ := json.Marshal(state) + texts = append(texts, string(sj)) + + for i, text := range texts { + chunks := splitForCount(text, parallelCountBytes) + if strings.Join(chunks, "") != text { + t.Fatalf("text %d: chunks do not reassemble the text", i) + } + want := len(enc.Encode(text, nil, nil)) + got, err := EstimateTokens(text) + if err != nil { + t.Fatal(err) + } + if got != want { + t.Errorf("text %d (%d bytes, %d chunks): chunked count %d, serial %d", i, len(text), len(chunks), got, want) + } + } +} + +func TestSplitForCountLeavesShortTextWhole(t *testing.T) { + if got := splitForCount("short text", parallelCountBytes); len(got) != 1 { + t.Errorf("got %d chunks", len(got)) + } + dense := strings.Repeat("9", 3*parallelCountBytes) + if got := splitForCount(dense, parallelCountBytes); len(got) != 1 || got[0] != dense { + t.Errorf("text with no safe cut must come back whole, got %d chunks", len(got)) + } +} diff --git a/judge/typesafe/guard_test.go b/judge/typesafe/guard_test.go new file mode 100644 index 0000000..a6fa240 --- /dev/null +++ b/judge/typesafe/guard_test.go @@ -0,0 +1,312 @@ +package typesafe_test + +import ( + "context" + "fmt" + "net/http" + "net/http/httptest" + "strings" + "sync" + "sync/atomic" + "testing" + "time" + + "github.com/hallelx2/llmgate" + "github.com/hallelx2/llmgate/judge/typesafe" +) + +// These tests run the guard with its real estimator, not the byte-counting +// stub newTestJudge installs: the property under test is the estimator's +// cost on the request path (HAL-1708). + +// newRealGuardJudge wires a Judge with the default estimator at srv, and +// records every RequestTrace. +func newRealGuardJudge(t testing.TB, h http.HandlerFunc) (*typesafe.Judge, *[]typesafe.RequestTrace) { + t.Helper() + srv := httptest.NewServer(h) + t.Cleanup(srv.Close) + var mu sync.Mutex + traces := &[]typesafe.RequestTrace{} + j, err := typesafe.New(typesafe.Config{ + APIKey: "test-key", + BaseURL: srv.URL, + OnRequest: func(tr typesafe.RequestTrace) { + mu.Lock() + *traces = append(*traces, tr) + mu.Unlock() + }, + }) + if err != nil { + t.Fatalf("New: %v", err) + } + return j, traces +} + +// answerAll answers every question the request asked with a Noul. +func answerAll(n int) string { + var b strings.Builder + b.WriteString("{") + for i := 0; i < n; i++ { + if i > 0 { + b.WriteString(",") + } + fmt.Fprintf(&b, `"p_%d":{"type":"noul","noul":0.5}`, i) + } + b.WriteString("}") + return b.String() +} + +// filingPage is dense 10-K prose with figures, about 6,000 characters. +var filingPage = strings.Repeat("Revenue for fiscal 2022 was $12,345 million, up 4.5% year over year, driven by services. ", 66) + +// fullBatch is the shape retrieval's full-page read sends: sixteen pages +// in the state, one Noul per page. +func fullBatch(pages int) llmgate.JudgeRequest { + state := map[string]any{"question": "What was FY2022 revenue?"} + qs := map[string]llmgate.Question{} + for i := 0; i < pages; i++ { + k := fmt.Sprintf("p_%d", i) + state[k] = map[string]any{"page": i + 1, "text": filingPage[:len(filingPage)*3/4]} + qs[k] = llmgate.Noul{ + Instructions: fmt.Sprintf("`question` is a question about a document. `%s.text` is the text of page `%s.page`. Does this page contain the specific facts or figures needed to answer the question?", k, k), + Criteria: &llmgate.NoulCriteria{True: "The answer is on this page", False: "It is not"}, + } + } + return llmgate.JudgeRequest{State: state, Questions: qs} +} + +// headBatch is the coarse skim: 120 page heads, 700 characters each. +func headBatch(pages int) llmgate.JudgeRequest { + state := map[string]any{"question": "What was FY2022 revenue?"} + qs := map[string]llmgate.Question{} + for i := 0; i < pages; i++ { + k := fmt.Sprintf("p_%d", i) + state[k] = map[string]any{"page": i + 1, "text": filingPage[:700]} + qs[k] = llmgate.Noul{Instructions: fmt.Sprintf("Could the answer be on page `%s.page`, judging from `%s.text`?", k, k)} + } + return llmgate.JudgeRequest{State: state, Questions: qs} +} + +func TestGuardOnAFullBatchIsCheap(t *testing.T) { + req := fullBatch(16) + j, traces := newRealGuardJudge(t, func(w http.ResponseWriter, r *http.Request) { + writeAnswers(w, answerAll(len(req.Questions))) + }) + // Build the encoder outside the measurement: a process pays that once. + if _, err := typesafe.EstimateTokens("warm"); err != nil { + t.Fatal(err) + } + if _, err := j.Judge(context.Background(), req); err != nil { + t.Fatalf("Judge: %v", err) + } + tr := (*traces)[0] + if !tr.GuardCounted { + t.Fatalf("a ~%d-byte state should be past the byte bound and tokenised", tr.StateBytes) + } + // Before HAL-1708 this was 17 encoder builds at ~256 ms each: over 4 s. + // The ceiling is loose so -race and a busy CI box still pass it. + if tr.Prepare > time.Second { + t.Errorf("prepare took %v for a 16-page batch; the guard is on the request path and must not dominate it", tr.Prepare) + } +} + +func TestGuardSkipsTheTokenizerWhenBytesProveItFits(t *testing.T) { + req := headBatch(40) + j, traces := newRealGuardJudge(t, func(w http.ResponseWriter, r *http.Request) { + writeAnswers(w, answerAll(len(req.Questions))) + }) + if _, err := j.Judge(context.Background(), req); err != nil { + t.Fatalf("Judge: %v", err) + } + if tr := (*traces)[0]; tr.GuardCounted { + t.Errorf("state %d bytes: every token covers a byte, so this provably fits and should not be tokenised", tr.StateBytes) + } +} + +func TestRealGuardStillRejectsAnOversizedState(t *testing.T) { + var called atomic.Bool + j, _ := newRealGuardJudge(t, func(w http.ResponseWriter, r *http.Request) { + called.Store(true) + writeAnswers(w, `{"q":{"type":"noul","noul":0.5}}`) + }) + // ~80k tokens of ordinary words: over the 64k total by any tokenizer. + _, err := j.Judge(context.Background(), llmgate.JudgeRequest{ + State: strings.Repeat("hello world ", 40_000), + Questions: map[string]llmgate.Question{"q": llmgate.Noul{Instructions: "urgent?"}}, + }) + if got := llmgate.Classify(err); got != llmgate.ErrClassContextLength { + t.Fatalf("Classify(%v) = %v, want context_length", err, got) + } + if called.Load() { + t.Error("the oversized request reached the server") + } +} + +func TestRealGuardStillRejectsOneOversizedQuestion(t *testing.T) { + j, _ := newRealGuardJudge(t, func(w http.ResponseWriter, r *http.Request) { + writeAnswers(w, `{"big":{"type":"noul","noul":0.5}}`) + }) + // A 20k-token state is fine alone; a 20k-token question beside it is not. + _, err := j.Judge(context.Background(), llmgate.JudgeRequest{ + State: strings.Repeat("hello world ", 10_000), + Questions: map[string]llmgate.Question{ + "big": llmgate.Noul{Instructions: strings.Repeat("hello world ", 10_000)}, + }, + }) + if got := llmgate.Classify(err); got != llmgate.ErrClassContextLength { + t.Fatalf("Classify(%v) = %v, want context_length", err, got) + } + if !strings.Contains(err.Error(), "big") { + t.Errorf("error %q should name the offending question", err) + } +} + +// Near the ceiling bytes alone cannot decide and the state must be counted: +// ~31k bytes that tokenise to ~5k, beside a question just over 1k bytes, +// fit in tokens though not in bytes. +func TestGuardCountsWhenBytesCannotDecide(t *testing.T) { + j, traces := newRealGuardJudge(t, func(w http.ResponseWriter, r *http.Request) { + writeAnswers(w, `{"q":{"type":"noul","noul":0.5}}`) + }) + _, err := j.Judge(context.Background(), llmgate.JudgeRequest{ + State: strings.Repeat("hello world ", 2_600), + Questions: map[string]llmgate.Question{"q": llmgate.Noul{Instructions: strings.Repeat("is it so? ", 120)}}, + }) + if err != nil { + t.Fatalf("a request that fits in tokens was refused: %v", err) + } + if !(*traces)[0].GuardCounted { + t.Error("the byte bound alone cannot prove this fits; the guard should have counted") + } +} + +func TestOnRequestReportsWhereTheTimeWent(t *testing.T) { + j, traces := newRealGuardJudge(t, func(w http.ResponseWriter, r *http.Request) { + time.Sleep(20 * time.Millisecond) + w.Header().Set("Content-Type", "application/json") + fmt.Fprint(w, `{"model":"jev-1.13","answers":{"p_0":{"type":"noul","noul":0.9}},"usage":{"input_tokens":321,"output_tokens":0}}`) + }) + if _, err := j.Judge(context.Background(), headBatch(1)); err != nil { + t.Fatalf("Judge: %v", err) + } + tr := (*traces)[0] + if tr.Questions != 1 || tr.StatusCode != 200 || tr.InputTokens != 321 || tr.Err != nil { + t.Errorf("trace = %+v", tr) + } + if tr.FirstByte < 20*time.Millisecond { + t.Errorf("FirstByte = %v, want at least the server's 20ms", tr.FirstByte) + } + if tr.Total < tr.FirstByte || tr.Total < tr.Prepare { + t.Errorf("Total %v must cover Prepare %v and FirstByte %v", tr.Total, tr.Prepare, tr.FirstByte) + } + if tr.StateBytes == 0 || tr.BodyBytes <= tr.StateBytes { + t.Errorf("StateBytes %d, BodyBytes %d", tr.StateBytes, tr.BodyBytes) + } + if tr.Model != typesafe.DefaultModel { + t.Errorf("Model = %q", tr.Model) + } +} + +func TestOnRequestSeesFailures(t *testing.T) { + j, traces := newRealGuardJudge(t, func(w http.ResponseWriter, r *http.Request) { + http.Error(w, `{"message":"overloaded"}`, 529) + }) + _, err := j.Judge(context.Background(), headBatch(1)) + if err == nil { + t.Fatal("want an error") + } + tr := (*traces)[0] + if tr.Err == nil || tr.StatusCode != 529 { + t.Errorf("trace = %+v, want the 529 and its error", tr) + } +} + +func TestWarmLeavesAConnectionForTheFirstCall(t *testing.T) { + var gets atomic.Int32 + j, traces := newRealGuardJudge(t, func(w http.ResponseWriter, r *http.Request) { + if r.Method == http.MethodGet { + gets.Add(1) + http.NotFound(w, r) + return + } + writeAnswers(w, answerAll(1)) + }) + if err := j.Warm(context.Background()); err != nil { + t.Fatalf("Warm: %v", err) + } + if gets.Load() != 1 { + t.Fatalf("Warm sent %d GETs, want 1", gets.Load()) + } + if _, err := j.Judge(context.Background(), headBatch(1)); err != nil { + t.Fatalf("Judge: %v", err) + } + if !(*traces)[0].ConnReused { + t.Error("the first judgment opened a new connection; Warm should have left one in the pool") + } +} + +// The encoder is now shared across goroutines; retrieval sends a level's +// batches concurrently. +func TestEstimateTokensIsSafeConcurrently(t *testing.T) { + want, err := typesafe.EstimateTokens(filingPage) + if err != nil { + t.Fatal(err) + } + var wg sync.WaitGroup + errs := make(chan error, 32) + for i := 0; i < 32; i++ { + wg.Add(1) + go func() { + defer wg.Done() + got, err := typesafe.EstimateTokens(filingPage) + if err != nil || got != want { + errs <- fmt.Errorf("got %d, %v; want %d", got, err, want) + } + }() + } + wg.Wait() + close(errs) + for err := range errs { + t.Error(err) + } +} + +// Before HAL-1708: ~256 ms/op. Acceptance: under 1 ms. +func BenchmarkEstimateTokensShortQuestion(b *testing.B) { + q := `{"type":"noul","instructions":"Does this page contain the specific facts or figures needed to answer the question?"}` + typesafe.EstimateTokens(q) + b.ResetTimer() + for i := 0; i < b.N; i++ { + typesafe.EstimateTokens(q) + } +} + +// The coarse skim's request, end to end against a local stub: guard, +// marshal, one loopback round trip. Acceptance: the guard under 50 ms. +func BenchmarkJudgeHeadBatch120(b *testing.B) { + req := headBatch(120) + j, _ := newRealGuardJudge(b, func(w http.ResponseWriter, r *http.Request) { + writeAnswers(w, answerAll(120)) + }) + typesafe.EstimateTokens("warm") + b.ResetTimer() + for i := 0; i < b.N; i++ { + if _, err := j.Judge(context.Background(), req); err != nil { + b.Fatal(err) + } + } +} + +func BenchmarkJudgeFullBatch16(b *testing.B) { + req := fullBatch(16) + j, _ := newRealGuardJudge(b, func(w http.ResponseWriter, r *http.Request) { + writeAnswers(w, answerAll(16)) + }) + typesafe.EstimateTokens("warm") + b.ResetTimer() + for i := 0; i < b.N; i++ { + if _, err := j.Judge(context.Background(), req); err != nil { + b.Fatal(err) + } + } +} diff --git a/judge/typesafe/typesafe.go b/judge/typesafe/typesafe.go index bdb6534..bfef69d 100644 --- a/judge/typesafe/typesafe.go +++ b/judge/typesafe/typesafe.go @@ -28,11 +28,13 @@ import ( "io" "net" "net/http" + "net/http/httptrace" "net/url" "os" "sort" "strconv" "strings" + "sync" "time" "github.com/pkoukk/tiktoken-go" @@ -110,6 +112,47 @@ type Config struct { // approximate guard, and it is better to lose the warning than to lose // the call. SkipContextGuard bool + + // OnRequest, when set, is called once per Judge call — success or + // failure — with where the time went. It runs on the calling + // goroutine after the call completes, so it must be cheap and must + // not block. + // + // It exists because a round-trip number alone cannot tell local CPU + // from provider latency: the 2026-09-25 query evaluation measured + // 28.6 s against a predicted ~9 s, and the difference was the guard + // rebuilding its tokenizer on every count (HAL-1708), inside the + // concurrency limiter where it looked like the provider being slow. + OnRequest func(RequestTrace) +} + +// RequestTrace is where one Judge call's time went. +type RequestTrace struct { + Model string + Questions int + StateBytes int // the marshalled state + BodyBytes int // the whole request body + + // Prepare covers validation, marshalling and the context guard — + // everything local before the request goes on the wire. + Prepare time.Duration + // GuardCounted is true when the guard ran the tokenizer, false when + // the byte bound proved the request fits (or the guard was skipped). + GuardCounted bool + + // Connect is DNS + TCP + TLS for a new connection; zero when an idle + // connection (or an HTTP/2 stream on one) was reused. + Connect time.Duration + ConnReused bool + // FirstByte runs from the request going out to the first response + // byte: the provider's own time plus one network round trip. + FirstByte time.Duration + // Total is the whole call, Prepare included. + Total time.Duration + + StatusCode int + InputTokens int + Err error } // Judge is a System One client implementing llmgate.Judge. @@ -119,7 +162,14 @@ type Judge struct { model string httpClient *http.Client countFn func(string) (int, error) - skipGuard bool + // byteBound is true when countFn is the built-in estimator, whose + // tokens each cover at least one byte — so a byte count is a proven + // upper bound on the token count and the guard can skip counting + // whenever bytes alone fit. A caller-supplied counter makes no such + // promise, so it is always consulted. + byteBound bool + skipGuard bool + onRequest func(RequestTrace) } // Compile-time check that the transport satisfies the seam. @@ -159,8 +209,10 @@ func New(cfg Config) (*Judge, error) { } countFn := cfg.CountTokens + byteBound := false if countFn == nil { countFn = EstimateTokens + byteBound = true } return &Judge{ @@ -169,7 +221,9 @@ func New(cfg Config) (*Judge, error) { model: model, httpClient: httpClient, countFn: countFn, + byteBound: byteBound, skipGuard: cfg.SkipContextGuard, + onRequest: cfg.OnRequest, }, nil } @@ -248,7 +302,20 @@ type wireResponse struct { } // Judge evaluates req.State against every question in one HTTP request. -func (j *Judge) Judge(ctx context.Context, req llmgate.JudgeRequest) (*llmgate.Judgment, error) { +func (j *Judge) Judge(ctx context.Context, req llmgate.JudgeRequest) (res *llmgate.Judgment, err error) { + start := time.Now() + tr := RequestTrace{Questions: len(req.Questions)} + if j.onRequest != nil { + defer func() { + tr.Total = time.Since(start) + tr.Err = err + if res != nil { + tr.InputTokens = res.Usage.InputTokens + } + j.onRequest(tr) + }() + } + if err := req.Validate(); err != nil { return nil, err } @@ -257,23 +324,81 @@ func (j *Judge) Judge(ctx context.Context, req llmgate.JudgeRequest) (*llmgate.J if model == "" { model = j.model } + tr.Model = model + + // The state is marshalled once: the body embeds it verbatim and the + // guard measures it, rather than each serialising it separately. + stateJSON, err := json.Marshal(req.State) + if err != nil { + return nil, fmt.Errorf("typesafe: marshal request: %w", err) + } + tr.StateBytes = len(stateJSON) body, err := json.Marshal(wireRequest{ - State: req.State, + State: json.RawMessage(stateJSON), Model: model, Questions: req.Questions, }) if err != nil { return nil, fmt.Errorf("typesafe: marshal request: %w", err) } + tr.BodyBytes = len(body) if !j.skipGuard { - if err := j.checkContextBudget(req); err != nil { + counted, err := j.checkContextBudget(stateJSON, req.Questions) + tr.GuardCounted = counted + if err != nil { return nil, err } } - - httpReq, err := http.NewRequestWithContext(ctx, http.MethodPost, j.baseURL+evaluatePath, bytes.NewReader(body)) + tr.Prepare = time.Since(start) + + // The transport calls these hooks on its own goroutines, and can still + // be running one when Do returns with an error; the timings are held + // under a lock and copied into the trace when the call ends. + var ( + wireMu sync.Mutex + connStart time.Time + sent time.Time + connect time.Duration + reused bool + firstByte time.Duration + ) + if j.onRequest != nil { + defer func() { + wireMu.Lock() + tr.Connect, tr.ConnReused, tr.FirstByte = connect, reused, firstByte + wireMu.Unlock() + }() + } + trace := &httptrace.ClientTrace{ + GetConn: func(string) { + wireMu.Lock() + connStart = time.Now() + wireMu.Unlock() + }, + GotConn: func(info httptrace.GotConnInfo) { + wireMu.Lock() + reused = info.Reused + if !info.Reused && !connStart.IsZero() { + connect = time.Since(connStart) + } + wireMu.Unlock() + }, + WroteRequest: func(httptrace.WroteRequestInfo) { + wireMu.Lock() + sent = time.Now() + wireMu.Unlock() + }, + GotFirstResponseByte: func() { + wireMu.Lock() + if !sent.IsZero() { + firstByte = time.Since(sent) + } + wireMu.Unlock() + }, + } + httpReq, err := http.NewRequestWithContext(httptrace.WithClientTrace(ctx, trace), http.MethodPost, j.baseURL+evaluatePath, bytes.NewReader(body)) if err != nil { return nil, fmt.Errorf("typesafe: build request: %w", err) } @@ -292,7 +417,8 @@ func (j *Judge) Judge(ctx context.Context, req llmgate.JudgeRequest) (*llmgate.J Cause: err, } } - defer resp.Body.Close() + defer func() { _ = resp.Body.Close() }() + tr.StatusCode = resp.StatusCode // Bounded: one extra byte past the cap tells us it overflowed rather // than merely reached the limit exactly. @@ -711,6 +837,7 @@ func retryAfter(h http.Header) time.Duration { } // checkContextBudget rejects a request that cannot fit, before it is sent. +// counted reports whether the tokenizer ran. // // Two limits, because the model reads the state once and then runs every // question against it: a total budget, and a per-question budget the @@ -720,30 +847,52 @@ func retryAfter(h http.Header) time.Duration { // uses cl100k_base. It is deliberately a guard and not an authority: the // server rules on the real thing, and Config.SkipContextGuard exists for // the case where an approximation rejects something the API would accept. -func (j *Judge) checkContextBudget(req llmgate.JudgeRequest) error { - stateJSON, err := json.Marshal(req.State) - if err != nil { - return fmt.Errorf("typesafe: marshal state for budget check: %w", err) +// +// It is also on every request's critical path, so it does the least work +// that can decide. With the built-in estimator every token covers at least +// one byte, so byte lengths are an upper bound on token counts: when the +// bytes fit, the tokens fit, and nothing is tokenised. Past that, the state +// is counted exactly once and each question is still bounded by its bytes; +// questions are tokenised only if the bound is not enough to decide. +func (j *Judge) checkContextBudget(stateJSON []byte, questions map[string]llmgate.Question) (counted bool, err error) { + ids := sortedQuestionIDs(questions) + qJSON := make([][]byte, len(ids)) + qBytes, longestBytes := 0, 0 + for i, id := range ids { + b, err := json.Marshal(questions[id]) + if err != nil { + return false, fmt.Errorf("typesafe: marshal question %q for budget check: %w", id, err) + } + qJSON[i] = b + qBytes += len(b) + longestBytes = max(longestBytes, len(b)) + } + + if j.byteBound { + if len(stateJSON)+qBytes <= MaxTotalTokens && len(stateJSON)+longestBytes <= MaxStateAndQuestionTokens { + return false, nil + } } + stateTokens, err := j.countFn(string(stateJSON)) if err != nil { // A broken estimator must not block a request that may well be // fine. Skip the guard and let the server decide. - return nil + return true, nil + } + if j.byteBound { + if stateTokens+qBytes <= MaxTotalTokens && stateTokens+longestBytes <= MaxStateAndQuestionTokens { + return true, nil + } } total := stateTokens longest := 0 longestID := "" - - for _, id := range sortedQuestionIDs(req.Questions) { - qJSON, err := json.Marshal(req.Questions[id]) + for i, id := range ids { + n, err := j.countFn(string(qJSON[i])) if err != nil { - return fmt.Errorf("typesafe: marshal question %q for budget check: %w", id, err) - } - n, err := j.countFn(string(qJSON)) - if err != nil { - return nil + return true, nil } total += n if n > longest { @@ -752,17 +901,17 @@ func (j *Judge) checkContextBudget(req llmgate.JudgeRequest) error { } if total > MaxTotalTokens { - return &llmgate.LLMError{ + return true, &llmgate.LLMError{ Class: llmgate.ErrClassContextLength, Provider: llmgate.ProviderTypeSafe, Message: fmt.Sprintf( "state plus %d questions is about %d tokens, over the %d limit — send fewer questions per call, or trim the state", - len(req.Questions), total, MaxTotalTokens), + len(questions), total, MaxTotalTokens), } } if stateTokens+longest > MaxStateAndQuestionTokens { - return &llmgate.LLMError{ + return true, &llmgate.LLMError{ Class: llmgate.ErrClassContextLength, Provider: llmgate.ProviderTypeSafe, Message: fmt.Sprintf( @@ -770,7 +919,7 @@ func (j *Judge) checkContextBudget(req llmgate.JudgeRequest) error { longestID, stateTokens+longest, MaxStateAndQuestionTokens), } } - return nil + return true, nil } // sortedQuestionIDs keeps budget accounting and its error messages @@ -791,12 +940,124 @@ func sortedQuestionIDs(q map[string]llmgate.Question) []string { // neighbourhood rather than exact. It exists to catch a request that is // obviously too large before it costs a round-trip, not to predict the // bill. +// +// The encoder is built once per process. tiktoken-go caches the rank +// table but builds a fresh BPE — a 100k-entry map and a regex compile — +// on every GetEncoding, which measured 256 ms per call; the guard used to +// pay that once per question (HAL-1708). +// +// Past that, encoding itself is the cost: tiktoken-go's pre-tokenizer is a +// backtracking regex, about 4 µs a byte, so a 16-page request state took +// 290 ms to count — on every request, inside the concurrency limiter's +// slot. Long text is therefore cut into pieces and counted in parallel. +// The count is exact, not approximate: see splitForCount. func EstimateTokens(text string) (int, error) { + enc, err := cl100k() + if err != nil { + return 0, err + } + chunks := splitForCount(text, parallelCountBytes) + if len(chunks) == 1 { + return len(enc.Encode(text, nil, nil)), nil + } + counts := make([]int, len(chunks)) + var wg sync.WaitGroup + for i, c := range chunks { + wg.Add(1) + go func() { + defer wg.Done() + counts[i] = len(enc.Encode(c, nil, nil)) + }() + } + wg.Wait() + n := 0 + for _, c := range counts { + n += c + } + return n, nil +} + +// parallelCountBytes is the chunk size above which counting fans out. +// Below it a goroutine costs more than it saves. +const parallelCountBytes = 8 << 10 + +// splitForCount cuts text into chunks of roughly size bytes, only at +// positions where the count of the parts provably equals the count of the +// whole. +// +// cl100k's BPE never merges across a pre-tokenizer piece, so a cut that +// falls on a piece boundary changes nothing. One such boundary is easy to +// recognise without running the regex: a single space, preceded by a +// non-space and followed by a letter. The pattern gives that space to the +// word after it (` ?\p{L}+`), and the piece before it ends at the +// non-space; no alternative in the pattern spans a lone space between a +// non-space and a letter. Cutting just before the space keeps " word" +// whole in the next chunk. +// +// Text with no such position — a single enormous token-dense run — comes +// back as one chunk and is counted serially, which is merely slower. +func splitForCount(text string, size int) []string { + if len(text) <= size { + return []string{text} + } + var out []string + for len(text) > size { + cut := -1 + for i := size; i < len(text)-1 && i < size*2; i++ { + if text[i] == ' ' && !isSpaceByte(text[i-1]) && isASCIILetter(text[i+1]) { + cut = i + break + } + } + if cut < 0 { + break + } + out = append(out, text[:cut]) + text = text[cut:] + } + return append(out, text) +} + +func isSpaceByte(b byte) bool { + return b == ' ' || b == '\t' || b == '\n' || b == '\r' || b == '\v' || b == '\f' +} + +func isASCIILetter(b byte) bool { + return (b >= 'a' && b <= 'z') || (b >= 'A' && b <= 'Z') +} + +var cl100k = sync.OnceValues(func() (*tiktoken.Tiktoken, error) { enc, err := tiktoken.GetEncoding("cl100k_base") if err != nil { - return 0, fmt.Errorf("typesafe: load tokenizer: %w", err) + return nil, fmt.Errorf("typesafe: load tokenizer: %w", err) + } + return enc, nil +}) + +// Warm opens a connection to the API ahead of the first judgment, so the +// first real call does not pay DNS, TCP and TLS on the critical path, and +// builds the guard's tokenizer so the first oversized request does not pay +// for that either. +// +// It sends an unauthenticated GET to the API root; the status is ignored — +// any response means a connection now sits in the pool. HTTP/2 multiplexes +// later calls onto it. Call it at startup or when a session begins; it is +// never required. +func (j *Judge) Warm(ctx context.Context) error { + if j.byteBound && !j.skipGuard { + go func() { _, _ = cl100k() }() // an error resurfaces on first use + } + req, err := http.NewRequestWithContext(ctx, http.MethodGet, j.baseURL+"/", nil) + if err != nil { + return fmt.Errorf("typesafe: warm: %w", err) + } + resp, err := j.httpClient.Do(req) + if err != nil { + return fmt.Errorf("typesafe: warm: %w", err) } - return len(enc.Encode(text, nil, nil)), nil + // Drained so the connection returns to the pool. + _, _ = io.Copy(io.Discard, io.LimitReader(resp.Body, 64<<10)) + return resp.Body.Close() } // truncate shortens s to at most n characters, marking that it was cut. diff --git a/middleware/limit/limit.go b/middleware/limit/limit.go index 97aa2ae..45ea2ec 100644 --- a/middleware/limit/limit.go +++ b/middleware/limit/limit.go @@ -257,6 +257,8 @@ type limitedClient struct { l *Limiter } +// Complete waits for a slot from the limiter, then forwards the request, +// feeding the outcome back so the limit adapts. func (c *limitedClient) Complete(ctx context.Context, req llmgate.Request) (*llmgate.Response, error) { release, err := c.l.Acquire(ctx) if err != nil { @@ -287,6 +289,8 @@ type limitedJudge struct { l *Limiter } +// Judge waits for a slot from the limiter, then forwards the batch, +// feeding the outcome back so the limit adapts. func (j *limitedJudge) Judge(ctx context.Context, req llmgate.JudgeRequest) (*llmgate.Judgment, error) { release, err := j.l.Acquire(ctx) if err != nil { diff --git a/middleware/limit/limit_test.go b/middleware/limit/limit_test.go index 7b370a5..93b9705 100644 --- a/middleware/limit/limit_test.go +++ b/middleware/limit/limit_test.go @@ -86,9 +86,10 @@ func TestLimiterConvergesOnTheProvidersCapacityAndLosesNothing(t *testing.T) { defer emu.Unlock() var dec, inc int for _, e := range events { - if e.Cause == "rate_limited" { + switch e.Cause { + case "rate_limited": dec++ - } else if e.Cause == "success" { + case "success": inc++ } }