Module: Legion::LLM::TokenTracker

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

Constant Summary collapse

MUTEX = 
Mutex.new

Class Method Summary collapse

Class Method Details

.record(input_tokens:, output_tokens:) ⇒ Hash

Records token usage from a completed LLM call.

Parameters:

  • input_tokens (Integer) —

    number of input tokens consumed

  • output_tokens (Integer) —

    number of output tokens produced

Returns:

  • (Hash) —

    the recorded entry



17
18
19
20
21
22
23
24
25
26
27
 
# File 'lib/legion/llm/token_tracker.rb', line 17

def record(input_tokens:, output_tokens:)
  entry = {
    input_tokens:  input_tokens.to_i,
    output_tokens: output_tokens.to_i,
    recorded_at:   Time.now
  }

  MUTEX.synchronize { store << entry }
  log.debug "[LLM::TokenTracker] recorded #{input_tokens}+#{output_tokens} tokens (total: #{total_tokens})"
  entry
end

.reset! ⇒ Object

Clears all recorded entries (for test isolation).



71
72
73
 
# File 'lib/legion/llm/token_tracker.rb', line 71

def reset!
  MUTEX.synchronize { @store = [] }
end

.session_exceeded? ⇒ Boolean

Returns true when total_tokens >= session_max_tokens (and limit is configured).

Returns:

  • (Boolean) 


53
54
55
56
57
58
 
# File 'lib/legion/llm/token_tracker.rb', line 53

def session_exceeded?
  limit = session_max_tokens
  return false unless limit&.positive?

  total_tokens >= limit
end

.session_warning? ⇒ Boolean

Returns true when total_tokens >= session_warn_tokens (and threshold is configured).

Returns:

  • (Boolean) 


63
64
65
66
67
68
 
# File 'lib/legion/llm/token_tracker.rb', line 63

def session_warning?
  threshold = session_warn_tokens
  return false unless threshold&.positive?

  total_tokens >= threshold
end

.summary ⇒ Hash

Returns a summary hash with totals and configured limits.

Returns:

  • (Hash) 


78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
 
# File 'lib/legion/llm/token_tracker.rb', line 78

def summary
  max   = session_max_tokens
  warn  = session_warn_tokens
  total = total_tokens
  {
    total_tokens:        total,
    total_input_tokens:  total_input_tokens,
    total_output_tokens: total_output_tokens,
    session_max_tokens:  max,
    session_warn_tokens: warn,
    exceeded:            session_exceeded?,
    warning:             session_warning?,
    remaining:           max ? [max - total, 0].max : nil
  }
end

.total_input_tokens ⇒ Integer

Returns total input tokens accumulated.

Returns:

  • (Integer) 


39
40
41
 
# File 'lib/legion/llm/token_tracker.rb', line 39

def total_input_tokens
  MUTEX.synchronize { store.sum { |e| e[:input_tokens] } }
end

.total_output_tokens ⇒ Integer

Returns total output tokens accumulated.

Returns:

  • (Integer) 


46
47
48
 
# File 'lib/legion/llm/token_tracker.rb', line 46

def total_output_tokens
  MUTEX.synchronize { store.sum { |e| e[:output_tokens] } }
end

.total_tokens ⇒ Integer

Returns total tokens accumulated across all recorded calls.

Returns:

  • (Integer) 


32
33
34
 
# File 'lib/legion/llm/token_tracker.rb', line 32

def total_tokens
  MUTEX.synchronize { store.sum { |e| e[:input_tokens] + e[:output_tokens] } }
end