Class: Legion::TTY::Components::TokenTracker

Inherits:
Object
  • Object
show all
Defined in:
lib/legion/tty/components/token_tracker.rb

Constant Summary collapse

PRICING = 
{
  'claude' => { input: 0.003, output: 0.015 },
  'openai' => { input: 0.005, output: 0.015 },
  'gemini' => { input: 0.001, output: 0.002 },
  'azure' => { input: 0.005, output: 0.015 },
  'local' => { input: 0.0, output: 0.0 }
}.freeze

Instance Attribute Summary collapse

Instance Method Summary collapse

Constructor Details

#initialize(provider: 'claude') ⇒ TokenTracker

Returns a new instance of TokenTracker.



17
18
19
20
21
22
 
# File 'lib/legion/tty/components/token_tracker.rb', line 17

def initialize(provider: 'claude')
  @provider = provider
  @total_input_tokens = 0
  @total_output_tokens = 0
  @total_cost = 0.0
end

Instance Attribute Details

#total_cost ⇒ Object (readonly)

Returns the value of attribute total_cost.



15
16
17
 
# File 'lib/legion/tty/components/token_tracker.rb', line 15

def total_cost
  @total_cost
end

#total_input_tokens ⇒ Object (readonly)

Returns the value of attribute total_input_tokens.



15
16
17
 
# File 'lib/legion/tty/components/token_tracker.rb', line 15

def total_input_tokens
  @total_input_tokens
end

#total_output_tokens ⇒ Object (readonly)

Returns the value of attribute total_output_tokens.



15
16
17
 
# File 'lib/legion/tty/components/token_tracker.rb', line 15

def total_output_tokens
  @total_output_tokens
end

Instance Method Details

#summary ⇒ Object 



31
32
33
34
35
36
 
# File 'lib/legion/tty/components/token_tracker.rb', line 31

def summary
  input = format_number(@total_input_tokens)
  output = format_number(@total_output_tokens)
  cost = format('%.4f', @total_cost)
  "Tokens: #{input} in / #{output} out | Cost: $#{cost}"
end

#track(input_tokens:, output_tokens:) ⇒ Object 



24
25
26
27
28
29
 
# File 'lib/legion/tty/components/token_tracker.rb', line 24

def track(input_tokens:, output_tokens:)
  @total_input_tokens += input_tokens.to_i
  @total_output_tokens += output_tokens.to_i
  rates = PRICING[@provider] || PRICING['claude']
  @total_cost += (input_tokens.to_i * rates[:input] / 1000.0) + (output_tokens.to_i * rates[:output] / 1000.0)
end