Class: AgentHarness::TokenTracker

Inherits:
Object
  • Object
show all
Defined in:
lib/agent_harness/token_tracker.rb

Overview

Tracks token usage across provider interactions

Provides in-memory tracking of token usage with support for callbacks when tokens are used. Consumers can register callbacks to persist usage data externally.

Examples:

Basic usage

tracker = AgentHarness::TokenTracker.new
tracker.record(provider: :claude, input_tokens: 100, output_tokens: 50)
puts tracker.summary

With callback

tracker.on_tokens_used do |event|
  MyDatabase.save_usage(event)
end

Defined Under Namespace

Classes: TokenEvent

Instance Method Summary collapse

Constructor Details

#initialize ⇒ TokenTracker

Returns a new instance of TokenTracker.



29
30
31
32
33
 
# File 'lib/agent_harness/token_tracker.rb', line 29

def initialize
  @events = []
  @callbacks = []
  @mutex = Mutex.new
end

Instance Method Details

#clear! ⇒ void

This method returns an undefined value.

Clear all recorded events



106
107
108
109
110
 
# File 'lib/agent_harness/token_tracker.rb', line 106

def clear!
  @mutex.synchronize do
    @events.clear
  end
end

#event_count ⇒ Integer

Get event count

Returns:

  • (Integer) —

    number of recorded events



123
124
125
126
127
 
# File 'lib/agent_harness/token_tracker.rb', line 123

def event_count
  @mutex.synchronize do
    @events.size
  end
end

#on_tokens_used {|TokenEvent| ... } ⇒ void

This method returns an undefined value.

Register callback for token events

Yields:

  • (TokenEvent) —

    called when tokens are recorded



99
100
101
 
# File 'lib/agent_harness/token_tracker.rb', line 99

def on_tokens_used(&block)
  @callbacks << block
end

#recent_events(limit: 100) ⇒ Array<TokenEvent>

Get recent events

Parameters:

  • limit (Integer) (defaults to: 100) —

    maximum number of events to return

Returns:



89
90
91
92
93
 
# File 'lib/agent_harness/token_tracker.rb', line 89

def recent_events(limit: 100)
  @mutex.synchronize do
    @events.last(limit)
  end
end

#record(provider:, model: nil, input_tokens: 0, output_tokens: 0, total_tokens: nil, request_id: nil) ⇒ TokenEvent

Record token usage

Parameters:

  • provider (Symbol, String) —

    the provider name

  • model (String, nil) (defaults to: nil) —

    the model used

  • input_tokens (Integer) (defaults to: 0) —

    input tokens used

  • output_tokens (Integer) (defaults to: 0) —

    output tokens used

  • total_tokens (Integer, nil) (defaults to: nil) —

    total tokens (calculated if nil)

  • request_id (String, nil) (defaults to: nil) —

    unique request ID (generated if nil)

Returns:



44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
 
# File 'lib/agent_harness/token_tracker.rb', line 44

def record(provider:, model: nil, input_tokens: 0, output_tokens: 0, total_tokens: nil, request_id: nil)
  total = total_tokens || (input_tokens + output_tokens)

  event = TokenEvent.new(
    provider: provider.to_sym,
    model: model,
    input_tokens: input_tokens,
    output_tokens: output_tokens,
    total_tokens: total,
    timestamp: Time.now,
    request_id: request_id || SecureRandom.uuid
  )

  @mutex.synchronize do
    @events << event
  end

  # Notify callbacks
  notify_callbacks(event)

  event
end

#summary(since: nil, provider: nil) ⇒ Hash

Get usage summary

Parameters:

  • since (Time, nil) (defaults to: nil) —

    only include events after this time

  • provider (Symbol, String, nil) (defaults to: nil) —

    filter by provider

Returns:

  • (Hash) —

    usage summary



72
73
74
75
76
77
78
79
80
81
82
83
 
# File 'lib/agent_harness/token_tracker.rb', line 72

def summary(since: nil, provider: nil)
  events = filtered_events(since: since, provider: provider)

  {
    total_requests: events.size,
    total_input_tokens: events.sum(&:input_tokens),
    total_output_tokens: events.sum(&:output_tokens),
    total_tokens: events.sum(&:total_tokens),
    by_provider: group_by_provider(events),
    by_model: group_by_model(events)
  }
end

#total_tokens(since: nil) ⇒ Integer

Get total token count

Parameters:

  • since (Time, nil) (defaults to: nil) —

    only include events after this time

Returns:

  • (Integer) —

    total tokens used



116
117
118
 
# File 'lib/agent_harness/token_tracker.rb', line 116

def total_tokens(since: nil)
  filtered_events(since: since).sum(&:total_tokens)
end