Class: SmartPrompt::TokenCounter

Inherits:
Object
  • Object
show all
Defined in:
lib/smart_prompt/token_counter.rb

Overview

TokenCounter provides token counting functionality with caching Uses tiktoken for accurate token counting compatible with OpenAI models

Instance Method Summary collapse

Constructor Details

#initialize(model: "gpt-3.5-turbo") ⇒ TokenCounter

Returns a new instance of TokenCounter.



5
6
7
8
9
10
 
# File 'lib/smart_prompt/token_counter.rb', line 5

def initialize(model: "gpt-3.5-turbo")
  @cache = {}
  @model = model
  @encoding = nil
  @use_tiktoken = load_tiktoken
end

Instance Method Details

#cache_size ⇒ Object

Get cache size



42
43
44
 
# File 'lib/smart_prompt/token_counter.rb', line 42

def cache_size
  @cache.size
end

#clear_cache ⇒ Object

Clear the cache



37
38
39
 
# File 'lib/smart_prompt/token_counter.rb', line 37

def clear_cache
  @cache.clear
end

#count(text) ⇒ Object

Count tokens in text with caching



13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
 
# File 'lib/smart_prompt/token_counter.rb', line 13

def count(text)
  return 0 if text.nil? || text.empty?
  
  # Return cached result if available
  return @cache[text] if @cache.key?(text)
  
  # Calculate and cache the result
  token_count = if @use_tiktoken
    count_with_tiktoken(text)
  else
    count_with_fallback(text)
  end
  
  @cache[text] = token_count
end

#count_messages(messages) ⇒ Object

Count tokens across multiple messages



30
31
32
33
34
 
# File 'lib/smart_prompt/token_counter.rb', line 30

def count_messages(messages)
  return 0 if messages.nil? || messages.empty?
  
  messages.sum { |msg| count(msg.content) }
end