Module: LlmCostTracker::Integrations::Base

Included in:
Anthropic, Openai, RubyLlm
Defined in:
lib/llm_cost_tracker/integrations/base.rb

Instance Method Summary collapse

Instance Method Details

#active? ⇒ Boolean

Returns:

  • (Boolean) 


24
25
26
 
# File 'lib/llm_cost_tracker/integrations/base.rb', line 24

def active?
  LlmCostTracker.configuration.instrumented?(integration_name)
end

#enforce_budget!(request:, provider: self.provider) ⇒ Object 



51
52
53
54
55
56
57
58
59
 
# File 'lib/llm_cost_tracker/integrations/base.rb', line 51

def enforce_budget!(request:, provider: self.provider)
  return unless active?

  LlmCostTracker::Budget.enforce!(
    provider: provider,
    model: request[:model],
    request: request
  )
end

#gem_version ⇒ Object 



127
128
129
 
# File 'lib/llm_cost_tracker/integrations/base.rb', line 127

def gem_version
  Gem.loaded_specs[integration_name.to_s]&.version
end

#install ⇒ Object 



28
29
30
31
32
33
34
 
# File 'lib/llm_cost_tracker/integrations/base.rb', line 28

def install
  validate_contract!
  patch_targets.each do |target|
    target_class = target.fetch(:constant_name).to_s.safe_constantize
    install_patch(target_class, target.fetch(:patch)) if target_class
  end
end

#integration_name ⇒ Object 



15
16
17
 
# File 'lib/llm_cost_tracker/integrations/base.rb', line 15

def integration_name
  @integration_name ||= name.demodulize.underscore.to_sym
end

#minimum_version(value = nil) ⇒ Object 



122
123
124
125
 
# File 'lib/llm_cost_tracker/integrations/base.rb', line 122

def minimum_version(value = nil)
  @minimum_version = value if value
  @minimum_version
end

#patch_target(constant_name, with:, optional: false, skip_when_methods_missing: false) ⇒ Object 



133
134
135
136
137
138
139
140
141
 
# File 'lib/llm_cost_tracker/integrations/base.rb', line 133

def patch_target(constant_name, with:, optional: false, skip_when_methods_missing: false)
  {
    constant_name: constant_name,
    patch: with,
    method_names: with.instance_methods,
    optional: optional,
    skip_when_methods_missing: skip_when_methods_missing
  }
end

#patch_targets ⇒ Object 



131
 
# File 'lib/llm_cost_tracker/integrations/base.rb', line 131

def patch_targets = []

#provider(value = nil) ⇒ Object 



19
20
21
22
 
# File 'lib/llm_cost_tracker/integrations/base.rb', line 19

def provider(value = nil)
  @provider = value.to_s if value
  @provider ||= integration_name.to_s
end

#record_safely ⇒ Object 



61
62
63
64
65
66
67
 
# File 'lib/llm_cost_tracker/integrations/base.rb', line 61

def record_safely
  yield
rescue LlmCostTracker::Error
  raise
rescue StandardError => e
  Logging.warn("#{integration_name} integration failed to record usage: #{e.class}: #{e.message}")
end

#request_params(args, kwargs) ⇒ Object 



69
70
71
72
73
74
75
76
77
78
79
 
# File 'lib/llm_cost_tracker/integrations/base.rb', line 69

def request_params(args, kwargs)
  params =
    case args.first
    when Hash then args.first
    when nil then {}
    else args.first.to_h
    end
  params.merge(kwargs).with_indifferent_access
rescue StandardError
  kwargs.to_h.with_indifferent_access
end

#status ⇒ Object 



36
37
38
39
40
41
42
43
44
45
46
47
48
49
 
# File 'lib/llm_cost_tracker/integrations/base.rb', line 36

def status
  name = integration_name.to_s
  problems = version_problems + target_problems
  if problems.any?
    return Check.new(:warn, name, "#{name} integration cannot be installed: #{problems.join('; ')}")
  end

  installed = patch_targets.reject { |target| target.fetch(:optional) }.all? do |target|
    target.fetch(:constant_name).to_s.safe_constantize&.ancestors&.include?(target.fetch(:patch))
  end
  return Check.new(:ok, name, "#{name} integration installed") if installed

  Check.new(:warn, name, "#{name} integration is enabled but not installed")
end

#stream_collector(request, provider: self.provider) ⇒ Object 



109
110
111
112
113
114
115
116
 
# File 'lib/llm_cost_tracker/integrations/base.rb', line 109

def stream_collector(request, provider: self.provider)
  LlmCostTracker::Capture::StreamCollector.new(
    provider: provider,
    model: request[:model],
    pricing_mode: stream_pricing_mode(request),
    request: request
  )
end

#stream_pricing_mode(_request) ⇒ Object 



118
119
120
 
# File 'lib/llm_cost_tracker/integrations/base.rb', line 118

def stream_pricing_mode(_request)
  nil
end

#track_stream(stream, collector:) ⇒ Object 



98
99
100
101
102
103
104
105
106
107
 
# File 'lib/llm_cost_tracker/integrations/base.rb', line 98

def track_stream(stream, collector:)
  return stream unless active?

  LlmCostTracker::Capture::StreamTracker.new(
    stream: stream,
    collector: collector,
    active: -> { active? },
    finish: ->(errored) { record_safely { collector.finish!(errored: errored) } }
  ).wrap
end

#wrap_blocking(args, kwargs, record:, provider: self.provider) ⇒ Object 



81
82
83
84
85
86
87
88
 
# File 'lib/llm_cost_tracker/integrations/base.rb', line 81

def wrap_blocking(args, kwargs, record:, provider: self.provider)
  request = request_params(args, kwargs)
  enforce_budget!(request: request, provider: provider)
  started_at = LlmCostTracker::Timing.now_monotonic
  response = yield
  record.call(response, request, LlmCostTracker::Timing.elapsed_ms(started_at))
  response
end

#wrap_stream(args, kwargs, collector:, provider: self.provider) ⇒ Object 



90
91
92
93
94
95
96
 
# File 'lib/llm_cost_tracker/integrations/base.rb', line 90

def wrap_stream(args, kwargs, collector:, provider: self.provider)
  request = request_params(args, kwargs)
  enforce_budget!(request: request, provider: provider)
  built = collector.call(request)
  stream = yield(built)
  track_stream(stream, collector: built)
end