Class: Brute::Middleware::OTel::TokenUsage

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

Overview

Records token usage from the LLM response as span attributes.

Runs POST-call: reads token counts from the response usage object and sets them as attributes on the span.

Instance Method Summary collapse

Methods inherited from Base

#initialize

Constructor Details

This class inherits a constructor from Brute::Middleware::Base

Instance Method Details

#call(env) ⇒ Object 



17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
 
# File 'lib/brute/middleware/otel/token_usage.rb', line 17

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

  span = env[:span]
  if span && response.respond_to?(:usage) && (usage = response.usage)
    span.set_attribute("gen_ai.usage.input_tokens", usage.input_tokens.to_i)
    span.set_attribute("gen_ai.usage.output_tokens", usage.output_tokens.to_i)
    span.set_attribute("gen_ai.usage.total_tokens", usage.total_tokens.to_i)

    reasoning = usage.reasoning_tokens.to_i
    span.set_attribute("gen_ai.usage.reasoning_tokens", reasoning) if reasoning > 0
  end

  response
end