Class: RubyLLM::Accounting::Usage::Tracker

Inherits:
Object
  • Object
show all
Defined in:
lib/ruby_llm/accounting/usage.rb

Overview

Tracks the physical transport attempts belonging to one operation.

Constant Summary collapse

CATEGORY_BY_OPERATION =

:nodoc:

{
  chat: :text_tokens,
  embedding: :embeddings,
  moderation: :text_tokens,
  image: :images,
  speech: :audio_tokens,
  transcription: :audio_tokens,
  ocr: :text_tokens,
  rerank: :embeddings
}.freeze

Instance Attribute Summary collapse

Instance Method Summary collapse

Constructor Details

#initialize(operation:, provider:, model:, config:, on_finish: nil) ⇒ Tracker

Returns a new instance of Tracker.



123
124
125
126
127
128
129
130
131
 
# File 'lib/ruby_llm/accounting/usage.rb', line 123

def initialize(operation:, provider:, model:, config:, on_finish: nil)
  @operation = operation.to_sym
  @provider = provider
  @model_info = model
  @config = config
  @on_finish = on_finish
  @entries = []
  @pending = []
end

Instance Attribute Details

#entries ⇒ Object (readonly)

Returns the value of attribute entries.



121
122
123
 
# File 'lib/ruby_llm/accounting/usage.rb', line 121

def entries
  @entries
end

Instance Method Details

#fail_attempt(entry, error) ⇒ Object 



151
152
153
154
155
156
 
# File 'lib/ruby_llm/accounting/usage.rb', line 151

def fail_attempt(entry, error)
  return unless entry&.pending?

  status = error.is_a?(CancelledError) ? :cancelled : :failed
  finish(entry, status:, tokens: failure_tokens(entry, error))
end

#fail_pending(error) ⇒ Object 



158
159
160
 
# File 'lib/ruby_llm/accounting/usage.rb', line 158

def fail_pending(error)
  @pending.dup.each { |entry| fail_attempt(entry, error) }
end

#observe(chunk) ⇒ Object 



140
141
142
143
144
145
146
147
148
149
 
# File 'lib/ruby_llm/accounting/usage.rb', line 140

def observe(chunk)
  entry = @pending.last
  return unless entry
  return unless chunk.respond_to?(:tokens)

  entry.finish(
    status: :pending,
    tokens: merge_stream_tokens(entry.tokens, chunk.tokens)
  )
end

#start ⇒ Object 



133
134
135
136
137
138
 
# File 'lib/ruby_llm/accounting/usage.rb', line 133

def start
  entry = Entry.new(operation: @operation, provider: @provider.slug, model: @model_info&.id)
  @entries << entry
  @pending << entry
  entry
end

#succeed(result) ⇒ Object 



162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
 
# File 'lib/ruby_llm/accounting/usage.rb', line 162

def succeed(result)
  # A request that produced several results, like a multi-image
  # generation, is billed once: the first result carries the call.
  billed = result.is_a?(Array) ? result.first : result
  billed.model_info = message_model(billed) if billed.is_a?(Message)
  pending = @pending.dup
  if pending.empty?
    attach_to_result(billed)
    return result
  end

  tokens = billed.respond_to?(:tokens) ? billed.tokens : Tokens.new
  cost = billed.cost if billed.respond_to?(:cost)
  pending[0...-1].each do |entry|
    finish(entry, status: :succeeded, tokens: Tokens.new)
  end
  finish(pending.last, status: :succeeded, tokens:, cost:)
  attach_to_result(billed)
  result
end

#succeed_attempts(tokens:) ⇒ Object 



183
184
185
186
187
 
# File 'lib/ruby_llm/accounting/usage.rb', line 183

def succeed_attempts(tokens:)
  @pending.dup.zip(tokens).each do |entry, usage|
    finish(entry, status: :succeeded, tokens: usage)
  end
end