llmgw/internal/router/kvindex/tokenize.go

package kvindex

import (
	"bytes"
	"context"
	"encoding/json"
	"errors"
	"fmt"
	"io"
	"net/http"
)

// Why tokenizing goes to the engine rather than happening here.
//
// Block hashes are over token ids, so the router needs the exact ids the
// engine will compute. For a chat request those come from the model's chat
// template — a Jinja program shipped with the tokenizer, which differs per
// model and takes the tools and template arguments as input — and only then
// from the tokenizer. A Go tokenizer for tokenizer.json would get the second
// half right and still need a Jinja interpreter faithful to transformers' for
// the first; a mismatch anywhere is an index that never matches, silently.
// vLLM's /tokenize runs the same code the request will, so it agrees by
// construction.
//
// The price is a round trip to a replica on the management network before
// the request is placed, and the template and tokenizer run twice on the
// engine — once here, once for the request. Against a fake engine on
// loopback the router's side of it — the call, decoding 2048 token ids and
// hashing 128 blocks — measured about 0.2ms (BenchmarkTokenize in the router
// package). Against a real vLLM add its templating and tokenizing, which has
// not been measured here; it is expected to be milliseconds for prompts of a
// few thousand tokens. TokenizeTimeout
// bounds it: a slow answer means routing as tier 1 would, not a slow request.

// ErrUnsupported says a request cannot be indexed: its blocks would be hashed
// with more than their tokens, or it is not a shape /tokenize takes.
var ErrUnsupported = errors.New("request not indexable")

// TokenizeBody builds the /tokenize request for an inference request, or
// says why there is none.
//
// Only what changes the token ids is carried. A request the engine would
// hash with extra keys — a cache salt, an image — gets ErrUnsupported: its
// blocks hash differently from its tokens alone, and asking the index about
// them would only ever miss.
func TokenizeBody(path string, body []byte, model string) ([]byte, error) {
	var fields map[string]json.RawMessage
	if err := json.Unmarshal(body, &fields); err != nil {
		return nil, fmt.Errorf("%w: %v", ErrUnsupported, err)
	}
	if _, ok := fields["cache_salt"]; ok {
		return nil, fmt.Errorf("%w: cache_salt", ErrUnsupported)
	}
	out := map[string]json.RawMessage{}
	modelJSON, _ := json.Marshal(model)
	out["model"] = modelJSON

	switch path {
	case "/v1/chat/completions":
		msgs, ok := fields["messages"]
		if !ok {
			return nil, fmt.Errorf("%w: no messages", ErrUnsupported)
		}
		if !textOnly(msgs) {
			return nil, fmt.Errorf("%w: content other than text", ErrUnsupported)
		}
		out["messages"] = msgs
		for _, k := range []string{"tools", "add_generation_prompt", "continue_final_message",
			"chat_template", "chat_template_kwargs", "add_special_tokens"} {
			if v, ok := fields[k]; ok && string(v) != "null" {
				out[k] = v
			}
		}
	case "/v1/completions":
		var s string
		if err := json.Unmarshal(fields["prompt"], &s); err != nil {
			// A list of prompts, or of token ids: a batch the index would have
			// to place as one, which it cannot.
			return nil, fmt.Errorf("%w: prompt is not a single string", ErrUnsupported)
		}
		out["prompt"] = fields["prompt"]
		if v, ok := fields["add_special_tokens"]; ok {
			out["add_special_tokens"] = v
		}
	default:
		return nil, fmt.Errorf("%w: %s", ErrUnsupported, path)
	}
	return json.Marshal(out)
}

// textOnly reports whether every message's content is a string or a list of
// text parts. An image or audio part is hashed by the engine with its
// identity as an extra key.
func textOnly(raw json.RawMessage) bool {
	var msgs []struct {
		Content json.RawMessage `json:"content"`
	}
	if err := json.Unmarshal(raw, &msgs); err != nil {
		return false
	}
	for _, m := range msgs {
		if len(m.Content) == 0 || m.Content[0] != '[' {
			continue
		}
		var parts []struct {
			Type string `json:"type"`
		}
		if err := json.Unmarshal(m.Content, &parts); err != nil {
			return false
		}
		for _, p := range parts {
			if p.Type != "text" {
				return false
			}
		}
	}
	return true
}

// Tokenize asks an engine for a request's token ids.
func Tokenize(ctx context.Context, client *http.Client, baseURL, apiKey string, body []byte) ([]int, error) {
	req, err := http.NewRequestWithContext(ctx, http.MethodPost, baseURL+"/tokenize", bytes.NewReader(body))
	if err != nil {
		return nil, err
	}
	req.Header.Set("Content-Type", "application/json")
	if apiKey != "" {
		req.Header.Set("Authorization", "Bearer "+apiKey)
	}
	resp, err := client.Do(req)
	if err != nil {
		return nil, err
	}
	defer func() { _ = resp.Body.Close() }()
	if resp.StatusCode != http.StatusOK {
		msg, _ := io.ReadAll(io.LimitReader(resp.Body, 512))
		return nil, fmt.Errorf("tokenize: %s: %s", resp.Status, bytes.TrimSpace(msg))
	}
	var out struct {
		Tokens []int `json:"tokens"`
	}
	if err := json.NewDecoder(io.LimitReader(resp.Body, 64<<20)).Decode(&out); err != nil {
		return nil, fmt.Errorf("tokenize: %w", err)
	}
	return out.Tokens, nil
}