Class: Crimson::TokenCounter

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

Overview

Token counter using tiktoken when available, with fallback estimation.

Instance Method Summary collapse

Constructor Details

#initialize(model: nil, provider: nil) ⇒ TokenCounter

Returns a new instance of TokenCounter.

Parameters:

  • model (String, nil) (defaults to: nil)

    model name for encoder selection

  • provider (String, Symbol, nil) (defaults to: nil)

    provider name



8
9
10
11
12
13
# File 'lib/crimson/token_counter.rb', line 8

def initialize(model: nil, provider: nil)
  @model = model
  @provider = provider ? provider.to_sym : nil
  @encoder = nil
  @encoder_loaded = false
end

Instance Method Details

#count(text) ⇒ Integer

Count tokens in a text string.

Parameters:

  • text (String, nil)

Returns:

  • (Integer)


18
19
20
21
22
23
24
25
# File 'lib/crimson/token_counter.rb', line 18

def count(text)
  return 0 if text.nil? || text.empty?
  encoder = load_encoder
  return estimate(text) unless encoder
  encoder.encode(text).length
rescue => e
  estimate(text)
end

#count_messages(messages) ⇒ Integer

Count total tokens for an array of messages, including per-message overhead.

Parameters:

Returns:

  • (Integer)


30
31
32
33
34
35
36
37
38
39
40
41
42
43
# File 'lib/crimson/token_counter.rb', line 30

def count_messages(messages)
  total = 0
  messages.each do |msg|
    total += 4
    total += count(msg.content.to_s)
    if msg.respond_to?(:tool_calls) && msg.tool_calls
      msg.tool_calls.each do |tc|
        total += count(tc.name.to_s)
        total += count(tc.arguments.to_s)
      end
    end
  end
  total
end