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 



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

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

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

.extract_provider(response) ⇒ Object 



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

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

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

.extract_usage(response) ⇒ Object 



38
39
40
41
42
43
44
45
46
 
# File 'lib/legion/llm/hooks/cost_tracking.rb', line 38

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 



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

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

.track(response, model) ⇒ Object 



19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
 
# 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?

  resolved_model    = extract_model(response) || model.to_s
  resolved_provider = extract_provider(response)

  Legion::LLM::Metering::Recorder.record(
    model:         resolved_model,
    input_tokens:  usage[:input_tokens],
    output_tokens: usage[:output_tokens],
    provider:      resolved_provider
  )
  nil
rescue StandardError => e
  handle_exception(e, level: :warn)
  nil
end