feat(tokens): add text token counting

Assisted-by: Codex:gpt-5

go.mod

@@ -458,7 +458,7 @@ require (
 	github.com/peterbourgon/diskv v2.0.1+incompatible // indirect
 	github.com/pierrec/lz4/v4 v4.1.8 // indirect
 	github.com/pkg/errors v0.9.1
-	github.com/pkoukk/tiktoken-go v0.1.7 // indirect
+	github.com/pkoukk/tiktoken-go v0.1.7
 	github.com/pmezard/go-difflib v1.0.1-0.20181226105442-5d4384ee4fb2 // indirect
 	github.com/polydawn/refmt v0.89.1-0.20231129105047-37766d95467a // indirect
 	github.com/power-devops/perfstat v0.0.0-20240221224432-82ca36839d55 // indirect

pkg/tokens/tokens.go

@@ -0,0 +1,103 @@
+// SPDX-License-Identifier: MIT
+
+package tokens
+
+import (
+	"fmt"
+	"strings"
+
+	"github.com/mudler/LocalAI/core/schema"
+	tiktoken "github.com/pkoukk/tiktoken-go"
+)
+
+const defaultEncoding = tiktoken.MODEL_CL100K_BASE
+
+// EncodingFor resolves the encoding known by tiktoken and uses cl100k_base
+// when the supplied model has no known mapping.
+func EncodingFor(model string) string {
+	if encoding, ok := tiktoken.MODEL_TO_ENCODING[model]; ok {
+		return encoding
+	}
+	for prefix, encoding := range tiktoken.MODEL_PREFIX_TO_ENCODING {
+		if strings.HasPrefix(model, prefix) {
+			return encoding
+		}
+	}
+	return defaultEncoding
+}
+
+// CountText counts text tokens without interpreting literal special-token
+// strings as tokenizer control tokens.
+func CountText(text, model string) (int, error) {
+	encoding, err := tiktoken.GetEncoding(EncodingFor(model))
+	if err != nil {
+		return 0, fmt.Errorf("get token encoding: %w", err)
+	}
+	return len(encoding.EncodeOrdinary(text)), nil
+}
+
+// Count counts only textual message content.
+func Count(messages []schema.Message, model string) (int, error) {
+	total := 0
+	for index, message := range messages {
+		text, err := messageText(message.Content)
+		if err != nil {
+			return 0, fmt.Errorf("message %d: %w", index, err)
+		}
+		count, err := CountText(text, model)
+		if err != nil {
+			return 0, err
+		}
+		total += count
+	}
+	return total, nil
+}
+
+func messageText(content any) (string, error) {
+	switch value := content.(type) {
+	case nil:
+		return "", nil
+	case string:
+		return value, nil
+	case []any:
+		var text strings.Builder
+		for index, rawPart := range value {
+			part, ok := rawPart.(map[string]any)
+			if !ok {
+				return "", fmt.Errorf("unsupported multimodal content part %d of type %T", index, rawPart)
+			}
+
+			if rawText, exists := part["text"]; exists {
+				partText, ok := rawText.(string)
+				if !ok {
+					return "", fmt.Errorf("unsupported multimodal text part %d: text has type %T", index, rawText)
+				}
+				text.WriteString(partText)
+				continue
+			}
+
+			partType, ok := part["type"].(string)
+			if ok && recognizedMediaType(partType) {
+				continue
+			}
+			return "", fmt.Errorf("unsupported multimodal content part %d", index)
+		}
+		return text.String(), nil
+	case map[string]any:
+		if len(value) == 0 {
+			return "", nil
+		}
+		return "", fmt.Errorf("unsupported message content of type %T", content)
+	default:
+		return "", fmt.Errorf("unsupported message content of type %T", content)
+	}
+}
+
+func recognizedMediaType(partType string) bool {
+	switch partType {
+	case "image", "image_url", "audio", "audio_url", "input_audio", "video", "video_url":
+		return true
+	default:
+		return false
+	}
+}

pkg/tokens/tokens_suite_test.go

@@ -0,0 +1,34 @@
+// SPDX-License-Identifier: MIT
+
+package tokens_test
+
+import (
+	"testing"
+
+	. "github.com/onsi/ginkgo/v2"
+	. "github.com/onsi/gomega"
+	tiktoken "github.com/pkoukk/tiktoken-go"
+)
+
+type testBPELoader struct{}
+
+func (testBPELoader) LoadTiktokenBpe(string) (map[string]int, error) {
+	ranks := make(map[string]int, 256)
+	for value := 0; value < 256; value++ {
+		ranks[string([]byte{byte(value)})] = value
+	}
+	for index, token := range []string{
+		"he", "hel", "hell", "hello",
+		"wo", "wor", "worl", "world",
+		" w", " wo", " wor", " worl", " world",
+	} {
+		ranks[token] = 256 + index
+	}
+	return ranks, nil
+}
+
+func TestTokens(t *testing.T) {
+	RegisterFailHandler(Fail)
+	tiktoken.SetBpeLoader(testBPELoader{})
+	RunSpecs(t, "Tokens Suite")
+}

pkg/tokens/tokens_test.go

@@ -0,0 +1,79 @@
+// SPDX-License-Identifier: MIT
+
+package tokens_test
+
+import (
+	"github.com/mudler/LocalAI/core/schema"
+	"github.com/mudler/LocalAI/pkg/tokens"
+
+	. "github.com/onsi/ginkgo/v2"
+	. "github.com/onsi/gomega"
+)
+
+var _ = Describe("Token counting", func() {
+	Describe("EncodingFor", func() {
+		It("uses tiktoken's exact model mapping", func() {
+			Expect(tokens.EncodingFor("text-davinci-003")).To(Equal("p50k_base"))
+		})
+
+		It("uses tiktoken's model prefix mapping", func() {
+			Expect(tokens.EncodingFor("gpt-4o-2024-08-06")).To(Equal("o200k_base"))
+		})
+
+		It("falls back conservatively for unknown models", func() {
+			Expect(tokens.EncodingFor("local-model")).To(Equal("cl100k_base"))
+		})
+	})
+
+	Describe("CountText", func() {
+		It("returns the exact encoded token count", func() {
+			count, err := tokens.CountText("hello world", "gpt-4")
+			Expect(err).NotTo(HaveOccurred())
+			Expect(count).To(Equal(2))
+		})
+
+		It("treats literal special-token strings as ordinary text", func() {
+			var count int
+			Expect(func() {
+				var err error
+				count, err = tokens.CountText("<|endoftext|>", "gpt-4")
+				Expect(err).NotTo(HaveOccurred())
+			}).NotTo(Panic())
+			Expect(count).To(BeNumerically(">", 0))
+		})
+	})
+
+	Describe("Count", func() {
+		It("sums string message content", func() {
+			count, err := tokens.Count([]schema.Message{
+				{Content: "hello"},
+				{Content: "world"},
+			}, "gpt-4")
+			Expect(err).NotTo(HaveOccurred())
+			Expect(count).To(Equal(2))
+		})
+
+		It("extracts text from multimodal content and ignores recognized media", func() {
+			count, err := tokens.Count([]schema.Message{{
+				Content: []any{
+					map[string]any{"type": "text", "text": "hello world"},
+					map[string]any{"type": "image_url", "image_url": map[string]any{"url": "image.png"}},
+				},
+			}}, "gpt-4")
+			Expect(err).NotTo(HaveOccurred())
+			Expect(count).To(Equal(2))
+		})
+
+		DescribeTable("rejects unsupported non-empty content",
+			func(content any) {
+				_, err := tokens.Count([]schema.Message{{Content: content}}, "gpt-4")
+				Expect(err).To(MatchError(ContainSubstring("unsupported")))
+			},
+			Entry("top-level object", map[string]any{"text": "hello"}),
+			Entry("arbitrary multimodal map", []any{map[string]any{"foo": "bar"}}),
+			Entry("malformed text part", []any{map[string]any{"type": "text", "text": 42}}),
+			Entry("unknown typed part", []any{map[string]any{"type": "custom", "custom": "value"}}),
+			Entry("non-object multimodal part", []any{"hello"}),
+		)
+	})
+})