Class: Brute::Middleware::TokenTracking

Inherits:
Base
  • Object
show all
Defined in:
lib/brute/middleware/token_tracking.rb

Overview

Tracks cumulative token usage across all LLM calls in a session.

Runs POST-call: reads usage from the response and accumulates totals in env. Also records per-call usage for the most recent call.

Instance Method Summary collapse

Constructor Details

#initialize(app) ⇒ TokenTracking

Returns a new instance of TokenTracking.



16
17
18
19
20
21
22
 
# File 'lib/brute/middleware/token_tracking.rb', line 16

def initialize(app)
  super(app)
  @total_input = 0
  @total_output = 0
  @total_reasoning = 0
  @call_count = 0
end

Instance Method Details

#call(env) ⇒ Object 



24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
 
# File 'lib/brute/middleware/token_tracking.rb', line 24

def call(env)
  response = @app.call(env)

  if response.respond_to?(:usage) && (usage = response.usage)
    @total_input += usage.input_tokens.to_i
    @total_output += usage.output_tokens.to_i
    @total_reasoning += usage.reasoning_tokens.to_i
    @call_count += 1

    env[:metadata][:tokens] = {
      total_input: @total_input,
      total_output: @total_output,
      total_reasoning: @total_reasoning,
      total: @total_input + @total_output,
      call_count: @call_count,
      last_call: {
        input: usage.input_tokens.to_i,
        output: usage.output_tokens.to_i,
        total: usage.total_tokens.to_i,
      },
    }
  end

  response
end