1package tokeniser
  2
  3import (
  4	"fmt"
  5	"strings"
  6	"sync"
  7
  8	"github.com/pkoukk/tiktoken-go"
  9)
 10
 11// Encoding names accepted by tiktoken for common 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 cached Tiktoken encoding, initialising it if necessary.
 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 maps a model name to the nearest tiktoken encoding.
 52// Unknown model families fall back to cl100k_base.
 53func EncodingForModel(model string) (string, error) {
 54	model = strings.ToLower(model)
 55
 56	// GPT-4o and newer use o200k_base
 57	if strings.Contains(model, "gpt-4o") ||
 58		strings.Contains(model, "gpt-4.1") ||
 59		strings.Contains(model, "gpt-4.5") ||
 60		strings.HasPrefix(model, "gpt-4o") ||
 61		strings.HasPrefix(model, "gpt-4.1") ||
 62		strings.HasPrefix(model, "gpt-4.5") {
 63		return EncodingO200kBase, nil
 64	}
 65
 66	// GPT-4 family uses cl100k_base
 67	if strings.HasPrefix(model, "gpt-4") {
 68		return EncodingCl100kBase, nil
 69	}
 70
 71	// GPT-3.5-turbo uses cl100k_base
 72	if strings.Contains(model, "gpt-3.5-turbo") {
 73		return EncodingCl100kBase, nil
 74	}
 75
 76	// Embedding models use cl100k_base
 77	if strings.Contains(model, "text-embedding") {
 78		return EncodingCl100kBase, nil
 79	}
 80
 81	// Codex and davinci models use p50k_base
 82	if strings.Contains(model, "code-") ||
 83		strings.Contains(model, "davinci") {
 84		return EncodingP50kBase, nil
 85	}
 86
 87	// Default to cl100k_base (most common for modern models)
 88	return EncodingCl100kBase, nil
 89}
 90
 91// CountTokens returns the number of tokens in text for the given model.
 92func CountTokens(text string, model string) (int, error) {
 93	encodingName, err := EncodingForModel(model)
 94	if err != nil {
 95		return 0, err
 96	}
 97
 98	enc, err := GetEncoding(encodingName)
 99	if err != nil {
100		return 0, fmt.Errorf("failed to get encoding: %w", err)
101	}
102
103	tokens := enc.Encode(text, nil, nil)
104	return len(tokens), nil
105}