1package tokeniser
  2
  3import (
  4	"fmt"
  5	"strings"
  6	"sync"
  7
  8	"github.com/pkoukk/tiktoken-go"
  9)
 10
 11// Encoding names for different model families
 12const (
 13	EncodingCl100kBase = "cl100k_base" // GPT-4, GPT-3.5-turbo, embeddings
 14	EncodingO200kBase  = "o200k_base"  // GPT-4o, GPT-4.1, GPT-4.5
 15	EncodingP50kBase   = "p50k_base"   // Codex models, davinci-002, davinci-003
 16	EncodingR50kBase   = "r50k_base"   // GPT-3 models
 17)
 18
 19// tokenizerCache caches encodings to avoid repeated initialization overhead
 20type tokenizerCache struct {
 21	mu        sync.RWMutex
 22	encodings map[string]*tiktoken.Tiktoken
 23}
 24
 25var cache = &tokenizerCache{encodings: make(map[string]*tiktoken.Tiktoken)}
 26
 27// GetEncoding returns a Tiktoken encoding for the given encoding name.
 28func GetEncoding(encodingName string) (*tiktoken.Tiktoken, error) {
 29	// Check cache first
 30	cache.mu.RLock()
 31	if enc, ok := cache.encodings[encodingName]; ok {
 32		cache.mu.RUnlock()
 33		return enc, nil
 34	}
 35	cache.mu.RUnlock()
 36
 37	// Initialise encoding
 38	enc, err := tiktoken.GetEncoding(encodingName)
 39	if err != nil {
 40		return nil, fmt.Errorf("failed to get encoding %s: %w", encodingName, err)
 41	}
 42
 43	// Cache it
 44	cache.mu.Lock()
 45	cache.encodings[encodingName] = enc
 46	cache.mu.Unlock()
 47
 48	return enc, nil
 49}
 50
 51// EncodingForModel returns the appropriate encoding name for a given model.
 52func EncodingForModel(model string) (string, error) {
 53	model = strings.ToLower(model)
 54
 55	// GPT-4o and newer use o200k_base
 56	if strings.Contains(model, "gpt-4o") ||
 57		strings.Contains(model, "gpt-4.1") ||
 58		strings.Contains(model, "gpt-4.5") ||
 59		strings.HasPrefix(model, "gpt-4o") ||
 60		strings.HasPrefix(model, "gpt-4.1") ||
 61		strings.HasPrefix(model, "gpt-4.5") {
 62		return EncodingO200kBase, nil
 63	}
 64
 65	// GPT-4 family uses cl100k_base
 66	if strings.HasPrefix(model, "gpt-4") {
 67		return EncodingCl100kBase, nil
 68	}
 69
 70	// GPT-3.5-turbo uses cl100k_base
 71	if strings.Contains(model, "gpt-3.5-turbo") {
 72		return EncodingCl100kBase, nil
 73	}
 74
 75	// Embedding models use cl100k_base
 76	if strings.Contains(model, "text-embedding") {
 77		return EncodingCl100kBase, nil
 78	}
 79
 80	// Codex and davinci models use p50k_base
 81	if strings.Contains(model, "code-") ||
 82		strings.Contains(model, "davinci") {
 83		return EncodingP50kBase, nil
 84	}
 85
 86	// Default to cl100k_base (most common for modern models)
 87	return EncodingCl100kBase, nil
 88}
 89
 90// CountTokens counts the tokens in a text string using the appropriate encoding.
 91func CountTokens(text string, model string) (int, error) {
 92	encodingName, err := EncodingForModel(model)
 93	if err != nil {
 94		return 0, err
 95	}
 96
 97	enc, err := GetEncoding(encodingName)
 98	if err != nil {
 99		return 0, fmt.Errorf("failed to get encoding: %w", err)
100	}
101
102	tokens := enc.Encode(text, nil, nil)
103	return len(tokens), nil
104}