Module: Legion::LLM::Hooks::CostTracking

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

Class Method Summary collapse

Class Method Details

.extract_model(response) ⇒ Object 



49
50
51
52
53
 
# File 'lib/legion/llm/hooks/cost_tracking.rb', line 49

def extract_model(response)
  return nil unless response.is_a?(Hash)

  response.dig(:meta, :model) || response[:model]
end

.extract_provider(response) ⇒ Object 



43
44
45
46
47
 
# File 'lib/legion/llm/hooks/cost_tracking.rb', line 43

def extract_provider(response)
  return nil unless response.is_a?(Hash)

  response.dig(:meta, :provider) || response[:provider]
end

.extract_usage(response) ⇒ Object 



33
34
35
36
37
38
39
40
41
 
# File 'lib/legion/llm/hooks/cost_tracking.rb', line 33

def extract_usage(response)
  return { input_tokens: 0, output_tokens: 0 } unless response.is_a?(Hash)

  usage = response[:usage] || {}
  {
    input_tokens:  usage[:input_tokens] || usage[:prompt_tokens] || 0,
    output_tokens: usage[:output_tokens] || usage[:completion_tokens] || 0
  }
end

.install ⇒ Object 



12
13
14
15
16
17
 
# File 'lib/legion/llm/hooks/cost_tracking.rb', line 12

def install
  Legion::LLM::Hooks.after_chat do |response:, model:, **|
    track(response, model)
    nil
  end
end

.track(response, model) ⇒ Object 



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

def track(response, model)
  usage = extract_usage(response)
  return if usage[:input_tokens].zero? && usage[:output_tokens].zero?

  CostTracker.record(
    model:         (extract_model(response) || model).to_s,
    input_tokens:  usage[:input_tokens],
    output_tokens: usage[:output_tokens],
    provider:      extract_provider(response)
  )
rescue StandardError => e
  handle_exception(e, level: :debug)
end