Module: Legion::LLM::TokenEstimation

Extended by:
Legion::Logging::Helper
Defined in:
lib/legion/llm/token_estimation.rb

Constant Summary collapse

CHARS_PER_TOKEN =
4
OVERHEAD_PER_MESSAGE =
4

Class Method Summary collapse

Class Method Details

.content_length(content) ⇒ Object



33
34
35
36
37
38
39
40
41
42
43
44
45
# File 'lib/legion/llm/token_estimation.rb', line 33

def content_length(content)
  case content
  when String then content.length
  when Array
    content.sum do |block|
      b = block.respond_to?(:transform_keys) ? block.transform_keys(&:to_sym) : block
      (b[:text] || b['text']).to_s.length
    end
  when nil then 0
  else
    content.to_s.length
  end
end

.estimate(messages:, model:, system: nil, tools: nil) ⇒ Object



16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
# File 'lib/legion/llm/token_estimation.rb', line 16

def estimate(messages:, model:, system: nil, tools: nil)
  log.debug("[llm][token_estimation] action=estimate model=#{model} messages=#{messages.size}")

  total_chars = 0
  total_chars += system.to_s.length if system

  messages.each do |msg|
    content = msg[:content] || msg['content']
    total_chars += content_length(content)
    total_chars += OVERHEAD_PER_MESSAGE
  end

  tools.each { |tool| total_chars += Legion::JSON.dump(tool).length } if tools.is_a?(Array) && tools.any?

  { input_tokens: (total_chars.to_f / CHARS_PER_TOKEN).ceil }
end