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 (
 26	cache     = &tokenizerCache{encodings: make(map[string]*tiktoken.Tiktoken)}
 27	modelLock sync.RWMutex
 28)
 29
 30// GetEncoding returns a Tiktoken encoding for the given encoding name.
 31func GetEncoding(encodingName string) (*tiktoken.Tiktoken, error) {
 32	// Check cache first
 33	cache.mu.RLock()
 34	if enc, ok := cache.encodings[encodingName]; ok {
 35		cache.mu.RUnlock()
 36		return enc, nil
 37	}
 38	cache.mu.RUnlock()
 39
 40	// Initialise encoding
 41	enc, err := tiktoken.GetEncoding(encodingName)
 42	if err != nil {
 43		return nil, fmt.Errorf("failed to get encoding %s: %w", encodingName, err)
 44	}
 45
 46	// Cache it
 47	cache.mu.Lock()
 48	cache.encodings[encodingName] = enc
 49	cache.mu.Unlock()
 50
 51	return enc, nil
 52}
 53
 54// EncodingForModel returns the appropriate encoding name for a given model.
 55func EncodingForModel(model string) (string, error) {
 56	model = strings.ToLower(model)
 57
 58	// GPT-4o and newer use o200k_base
 59	if strings.Contains(model, "gpt-4o") ||
 60		strings.Contains(model, "gpt-4.1") ||
 61		strings.Contains(model, "gpt-4.5") ||
 62		strings.HasPrefix(model, "gpt-4o") ||
 63		strings.HasPrefix(model, "gpt-4.1") ||
 64		strings.HasPrefix(model, "gpt-4.5") {
 65		return EncodingO200kBase, nil
 66	}
 67
 68	// GPT-4 family uses cl100k_base
 69	if strings.HasPrefix(model, "gpt-4") {
 70		return EncodingCl100kBase, nil
 71	}
 72
 73	// GPT-3.5-turbo uses cl100k_base
 74	if strings.Contains(model, "gpt-3.5-turbo") {
 75		return EncodingCl100kBase, nil
 76	}
 77
 78	// Embedding models use cl100k_base
 79	if strings.Contains(model, "text-embedding") {
 80		return EncodingCl100kBase, nil
 81	}
 82
 83	// Codex and davinci models use p50k_base
 84	if strings.Contains(model, "code-") ||
 85		strings.Contains(model, "davinci") {
 86		return EncodingP50kBase, nil
 87	}
 88
 89	// Default to cl100k_base (most common for modern models)
 90	return EncodingCl100kBase, nil
 91}
 92
 93// CountTokens counts the tokens in a text string using the appropriate encoding.
 94func CountTokens(text string, model string) (int, error) {
 95	encodingName, err := EncodingForModel(model)
 96	if err != nil {
 97		return 0, err
 98	}
 99
100	enc, err := GetEncoding(encodingName)
101	if err != nil {
102		return 0, fmt.Errorf("failed to get encoding: %w", err)
103	}
104
105	tokens := enc.Encode(text, nil, nil)
106	return len(tokens), nil
107}
108
109// CountMessagesTokens counts tokens for chat API messages.
110// Based on OpenAI's official token counting method.
111func CountMessagesTokens(messages []Message, model string) (int, error) {
112	encodingName, err := EncodingForModel(model)
113	if err != nil {
114		return 0, err
115	}
116
117	enc, err := GetEncoding(encodingName)
118	if err != nil {
119		return 0, fmt.Errorf("failed to get encoding: %w", err)
120	}
121
122	var tokensPerMessage, tokensPerName int
123	switch {
124	case strings.Contains(model, "gpt-3.5-turbo-0613"),
125		strings.Contains(model, "gpt-3.5-turbo-16k-0613"),
126		strings.Contains(model, "gpt-4-0314"),
127		strings.Contains(model, "gpt-4-32k-0314"),
128		strings.Contains(model, "gpt-4-0613"),
129		strings.Contains(model, "gpt-4-32k-0613"):
130		tokensPerMessage = 3
131		tokensPerName = 1
132	case strings.Contains(model, "gpt-3.5-turbo-0301"):
133		tokensPerMessage = 4
134		tokensPerName = -1
135	default:
136		// Use default values for other models
137		tokensPerMessage = 3
138		tokensPerName = 1
139	}
140
141	numTokens := 0
142	for _, msg := range messages {
143		numTokens += tokensPerMessage
144		numTokens += len(enc.Encode(msg.Content, nil, nil))
145		numTokens += len(enc.Encode(msg.Role, nil, nil))
146		numTokens += len(enc.Encode(msg.Name, nil, nil))
147		if msg.Name != "" {
148			numTokens += tokensPerName
149		}
150	}
151	numTokens += 3 // Every reply is primed with <|start|>assistant<|message|>
152
153	return numTokens, nil
154}
155
156// Message represents a chat message for token counting.
157type Message struct {
158	Role    string
159	Content string
160	Name    string
161}