Module: Legion::LLM::Hooks::Metering

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

Class Method Summary collapse

Class Method Details

.extract_metering_data(response, model) ⇒ Object 



30
31
32
33
34
35
36
37
38
39
40
 
# File 'lib/legion/llm/hooks/metering.rb', line 30

def extract_metering_data(response, model)
  usage = extract_usage(response)
  {
    provider:      extract_provider(response),
    model_id:      (extract_model(response) || model).to_s,
    input_tokens:  usage[:input_tokens],
    output_tokens: usage[:output_tokens],
    event_type:    'llm_completion',
    status:        response.is_a?(Hash) && response[:error] ? 'failure' : 'success'
  }
end

.extract_model(response) ⇒ Object 



58
59
60
61
62
 
# File 'lib/legion/llm/hooks/metering.rb', line 58

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

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

.extract_provider(response) ⇒ Object 



52
53
54
55
56
 
# File 'lib/legion/llm/hooks/metering.rb', line 52

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

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

.extract_usage(response) ⇒ Object 



42
43
44
45
46
47
48
49
50
 
# File 'lib/legion/llm/hooks/metering.rb', line 42

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

.gateway_metering? ⇒ Boolean

Returns:

  • (Boolean) 


75
76
77
 
# File 'lib/legion/llm/hooks/metering.rb', line 75

def gateway_metering?
  defined?(Legion::Extensions::LLM::Gateway::Runners::MeteringWriter)
end

.install ⇒ Object 



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

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

.metering_available? ⇒ Boolean

Returns:

  • (Boolean) 


87
88
89
 
# File 'lib/legion/llm/hooks/metering.rb', line 87

def metering_available?
  gateway_metering? || transport_metering?
end

.publish_metering(payload) ⇒ Object 



64
65
66
67
68
69
70
71
72
73
 
# File 'lib/legion/llm/hooks/metering.rb', line 64

def publish_metering(payload)
  if gateway_metering?
    Legion::Extensions::LLM::Gateway::Runners::MeteringWriter.write_metering_record(payload)
  elsif transport_metering?
    Legion::Transport.publish(
      'lex.metering.record',
      Legion::JSON.dump(payload)
    )
  end
end

.record(response, model) ⇒ Object 



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

def record(response, model)
  return unless metering_available?

  payload = extract_metering_data(response, model)
  return if payload[:input_tokens].zero? && payload[:output_tokens].zero?

  publish_metering(payload)
rescue StandardError => e
  handle_exception(e, level: :debug)
end

.transport_metering? ⇒ Boolean

Returns:

  • (Boolean) 


79
80
81
82
83
84
85
 
# File 'lib/legion/llm/hooks/metering.rb', line 79

def transport_metering?
  defined?(Legion::Settings) &&
    Legion::Settings[:transport][:connected] == true
rescue StandardError => e
  handle_exception(e, level: :debug)
  false
end