Module: Ask::TokenUsage::Instrumentation

Defined in:
lib/ask/token_usage/instrumentation.rb

Overview

Optional bridge to ask-instrumentation events. When installed, every LLM chat completion (chat.ask / chat.stream.ask) automatically deducts the measured token cost from the caller's wallet.

require "ask/token_usage/instrumentation"

Ask::TokenUsage::Instrumentation.install do |event|
# Return the owner (wallet key) for this LLM call, or nil to skip.
event.payload[:workspace]   # whatever object/identifier you key wallets on
end

The block receives every matching instrumentation event and must return either the wallet owner or nil. Token cost is computed from the +input_tokens+/+output_tokens+ reported in the event payload.

Requires ask-instrumentation as a runtime dependency. The gem works without it; this file is only loaded when explicitly required.

Constant Summary collapse

PATTERNS = 
%w[chat.ask chat.stream.ask].freeze

Class Method Summary collapse

Class Method Details

.handle_event(payload) ⇒ Object 



54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
 
# File 'lib/ask/token_usage/instrumentation.rb', line 54

def handle_event(payload)
  owner = @owner_resolver.call(payload)
  return unless owner

  wallet = Ask::TokenUsage.wallet_for(owner)
  input  = payload[:usage] ? payload[:usage][:input_tokens].to_i : payload[:input_tokens].to_i
  output = payload[:usage] ? payload[:usage][:output_tokens].to_i : payload[:output_tokens].to_i
  tokens = input + output
  return if tokens <= 0

  model    = payload[:model] || payload[:model_id]
  provider = payload[:provider]

  wallet.deduct!(tokens,
    reason: :llm_call,
    metadata: {
      input_tokens: input,
      output_tokens: output,
      model_id: model,
      provider: provider
    })
rescue InsufficientTokens
  # let the caller handle insufficient — this module only auto-deducts
  # when the wallet can cover the cost.
rescue StandardError => e
  warn "[ask-token-usage instrumentation] failed to deduct: #{e.class}: #{e.message}"
end

.install(&block) ⇒ Object

Subscribe to ask-instrumentation chat events. &block receives each event and must return the wallet owner (for the wallet) or nil to skip deduction. Idempotent — calling install twice is safe.

Raises:

  • (ArgumentError) 


31
32
33
34
35
36
37
38
39
40
41
42
 
# File 'lib/ask/token_usage/instrumentation.rb', line 31

def install(&block)
  return if @installed

  raise ArgumentError, "block is required — it maps each event to a wallet owner" unless block

  @owner_resolver = block
  @installed = true

  ActiveSupport::Notifications.subscribe(/chat(\.stream)?\.ask/) do |_name, _start, _finish, _id, payload|
    handle_event(payload)
  end
end

.installed? ⇒ Boolean

Returns:

  • (Boolean) 


50
51
52
 
# File 'lib/ask/token_usage/instrumentation.rb', line 50

def installed?
  @installed
end

.uninstall ⇒ Object 



44
45
46
47
48
 
# File 'lib/ask/token_usage/instrumentation.rb', line 44

def uninstall
  ActiveSupport::Notifications.unsubscribe(/chat(\.stream)?\.ask/)
  @owner_resolver = nil
  @installed = false
end