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



32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
# File 'lib/legion/llm/hooks/metering.rb', line 32

def extract_metering_data(response, model)
  usage = extract_usage(response)
  resolved_model    = extract_model_id(response) || model.to_s
  resolved_provider = extract_provider(response)

  status = response.is_a?(Hash) && response[:error] ? 'failure' : 'success'

  {
    provider:      resolved_provider,
    model_id:      resolved_model,
    input_tokens:  usage[:input_tokens],
    output_tokens: usage[:output_tokens],
    event_type:    'llm_completion',
    status:        status
  }
end

.extract_model_id(response) ⇒ Object



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

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

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

.extract_provider(response) ⇒ Object



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

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

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

.extract_usage(response) ⇒ Object



71
72
73
74
75
76
77
78
79
# File 'lib/legion/llm/hooks/metering.rb', line 71

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)


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

def gateway_metering?
  defined?(LegionLLMGateway) && LegionLLMGateway.respond_to?(:emit)
end

.installObject



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

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

.metering_available?Boolean

Returns:

  • (Boolean)


49
50
51
# File 'lib/legion/llm/hooks/metering.rb', line 49

def metering_available?
  gateway_metering? || transport_metering?
end

.publish_metering(data) ⇒ Object



61
62
63
64
65
66
67
68
69
# File 'lib/legion/llm/hooks/metering.rb', line 61

def publish_metering(data)
  if gateway_metering?
    LegionLLMGateway.emit(data)
  elsif transport_metering?
    Legion::LLM::Metering.emit(data)
  end
rescue StandardError => e
  handle_exception(e, level: :warn)
end

.record(response, model) ⇒ Object



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

def record(response, model)
  return unless metering_available?

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

  publish_metering(data)
  nil
rescue StandardError => e
  handle_exception(e, level: :warn)
  nil
end

.transport_metering?Boolean

Returns:

  • (Boolean)


57
58
59
# File 'lib/legion/llm/hooks/metering.rb', line 57

def transport_metering?
  defined?(Legion::LLM::Transport::Messages::MeteringEvent)
end