Class: PWN::Plugins::REPL::AIConsoleUsage::Tracker

Inherits:
Object
  • Object
show all
Defined in:
lib/pwn/plugins/repl/ai_console_usage.rb

Overview

Accumulate one observer event per provider completion.

Instance Method Summary collapse

Constructor Details

#initialize ⇒ Tracker

Returns a new instance of Tracker.



10
11
12
13
14
15
16
17
18
19
20
21
# File 'lib/pwn/plugins/repl/ai_console_usage.rb', line 10

def initialize
  @lock = Mutex.new
  @input = 0
  @output = 0
  @cached = 0
  @cache_write = 0
  @calls = 0
  @latest = 0
  @priced = 0.0
  @priced_calls = 0
  @unpriced_calls = 0
end

Instance Method Details

#record(opts = {}) ⇒ Object



23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
# File 'lib/pwn/plugins/repl/ai_console_usage.rb', line 23

def record(opts = {})
  usage = AIConsoleUsage.normalize(opts)
  price = AIConsoleUsage.estimate(opts.merge(usage: usage))
  @lock.synchronize do
    @input += usage[:input_tokens]
    @output += usage[:output_tokens]
    @cached += usage[:cached_tokens]
    @cache_write += usage[:cache_write_tokens]
    @calls += 1
    @latest = usage[:total_tokens]
    if price
      @priced += price
      @priced_calls += 1
    else
      @unpriced_calls += 1
    end
  end
  snapshot
end

#snapshot ⇒ Object



43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
# File 'lib/pwn/plugins/repl/ai_console_usage.rb', line 43

def snapshot
  @lock.synchronize do
    status = if @calls.zero?
               'no provider usage yet'
             elsif @unpriced_calls.zero?
               'estimated from configured USD rates'
             elsif @priced_calls.zero?
               'unavailable: no configured USD rate'
             else
               "partial: #{@unpriced_calls} call(s) have no configured rate"
             end
    {
      input_tokens: @input, output_tokens: @output, cached_tokens: @cached,
      cache_write_tokens: @cache_write, total_tokens: @input + @output,
      latest_context_tokens: @latest, calls: @calls,
      estimated_cost_usd: @priced_calls.zero? ? nil : @priced,
      cost_status: status
    }
  end
end