Module: RubyLLM::Agents::BudgetTracker

Extended by:
CacheHelper
Defined in:
lib/ruby_llm/agents/infrastructure/budget_tracker.rb

Overview

Cache-based budget tracking for cost governance

Tracks spending against configured budget limits using cache counters. Supports daily and monthly budgets at both global and per-agent levels. In multi-tenant mode, budgets are tracked separately per tenant.

Note: Uses best-effort enforcement with cache counters. In high-concurrency scenarios, slight overruns may occur due to race conditions. This is an acceptable trade-off for performance.

Examples:

Checking budget before execution

BudgetTracker.check_budget!("MyAgent")  # raises BudgetExceededError if over limit

Recording spend after execution

BudgetTracker.record_spend!("MyAgent", 0.05)

Multi-tenant usage

BudgetTracker.check_budget!("MyAgent", tenant_id: "acme_corp")
BudgetTracker.record_spend!("MyAgent", 0.05, tenant_id: "acme_corp")

See Also:

Constant Summary

Constants included from CacheHelper

CacheHelper::NAMESPACE

Class Method Summary collapse

Methods included from CacheHelper

cache_delete, cache_exist?, cache_increment, cache_key, cache_read, cache_store, cache_write

Class Method Details

.calculate_forecast(tenant_id: nil) ⇒ Hash?

Calculates budget forecasts based on current spending trends

Parameters:

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

    Optional tenant identifier (uses resolver if not provided)

Returns:

  • (Hash, nil) —

    Forecast information



169
170
171
172
173
174
 
# File 'lib/ruby_llm/agents/infrastructure/budget_tracker.rb', line 169

def calculate_forecast(tenant_id: nil)
  tenant_id = resolve_tid(tenant_id)
  budget_config = Budget::ConfigResolver.resolve_budget_config(tenant_id)

  Budget::Forecaster.calculate_forecast(tenant_id: tenant_id, budget_config: budget_config)
end

.check_budget!(agent_type, tenant_id: nil, tenant_config: nil) ⇒ void

This method returns an undefined value.

Checks if the current spend exceeds budget limits

Parameters:

  • agent_type (String) —

    The agent class name

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

    Optional tenant identifier (uses resolver if not provided)

  • tenant_config (Hash, nil) (defaults to: nil) —

    Optional runtime tenant config (takes priority over resolver/DB)

Raises:



46
47
48
49
50
51
52
53
54
 
# File 'lib/ruby_llm/agents/infrastructure/budget_tracker.rb', line 46

def check_budget!(agent_type, tenant_id: nil, tenant_config: nil)
  tenant_id = resolve_tid(tenant_id)
  budget_config = Budget::ConfigResolver.resolve_budget_config(tenant_id, runtime_config: tenant_config)

  return unless budget_config[:enabled]
  return unless budget_config[:enforcement] == :hard

  check_budget_limits!(agent_type, tenant_id, budget_config)
end

.check_token_budget!(agent_type, tenant_id: nil, tenant_config: nil) ⇒ void

This method returns an undefined value.

Checks if the current token usage exceeds budget limits

Parameters:

  • agent_type (String) —

    The agent class name

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

    Optional tenant identifier (uses resolver if not provided)

  • tenant_config (Hash, nil) (defaults to: nil) —

    Optional runtime tenant config (takes priority over resolver/DB)

Raises:



63
64
65
66
67
68
69
70
71
 
# File 'lib/ruby_llm/agents/infrastructure/budget_tracker.rb', line 63

def check_token_budget!(agent_type, tenant_id: nil, tenant_config: nil)
  tenant_id = resolve_tid(tenant_id)
  budget_config = Budget::ConfigResolver.resolve_budget_config(tenant_id, runtime_config: tenant_config)

  return unless budget_config[:enabled]
  return unless budget_config[:enforcement] == :hard

  check_token_limits!(agent_type, tenant_id, budget_config)
end

.current_spend(scope, period, agent_type: nil, tenant_id: nil) ⇒ Float

Returns the current spend for a scope and period

Parameters:

  • scope (Symbol) —

    :global or :agent

  • period (Symbol) —

    :daily or :monthly

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

    Required when scope is :agent

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

    Optional tenant identifier (uses resolver if not provided)

Returns:

  • (Float) —

    Current spend in USD



112
113
114
115
 
# File 'lib/ruby_llm/agents/infrastructure/budget_tracker.rb', line 112

def current_spend(scope, period, agent_type: nil, tenant_id: nil)
  tenant_id = resolve_tid(tenant_id)
  Budget::BudgetQuery.current_spend(scope, period, agent_type: agent_type, tenant_id: tenant_id)
end

.current_tokens(period, tenant_id: nil) ⇒ Integer

Returns the current token usage for a period (global only)

Parameters:

  • period (Symbol) —

    :daily or :monthly

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

    Optional tenant identifier (uses resolver if not provided)

Returns:

  • (Integer) —

    Current token usage



122
123
124
125
 
# File 'lib/ruby_llm/agents/infrastructure/budget_tracker.rb', line 122

def current_tokens(period, tenant_id: nil)
  tenant_id = resolve_tid(tenant_id)
  Budget::BudgetQuery.current_tokens(period, tenant_id: tenant_id)
end

.record_spend!(agent_type, amount, tenant_id: nil, tenant_config: nil) ⇒ void

This method returns an undefined value.

Records spend and checks for soft cap alerts

Parameters:

  • agent_type (String) —

    The agent class name

  • amount (Float) —

    The amount spent in USD

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

    Optional tenant identifier (uses resolver if not provided)

  • tenant_config (Hash, nil) (defaults to: nil) —

    Optional runtime tenant config (takes priority over resolver/DB)



80
81
82
83
84
85
86
87
 
# File 'lib/ruby_llm/agents/infrastructure/budget_tracker.rb', line 80

def record_spend!(agent_type, amount, tenant_id: nil, tenant_config: nil)
  return if amount.nil? || amount <= 0

  tenant_id = resolve_tid(tenant_id)
  budget_config = Budget::ConfigResolver.resolve_budget_config(tenant_id, runtime_config: tenant_config)

  Budget::SpendRecorder.record_spend!(agent_type, amount, tenant_id: tenant_id, budget_config: budget_config)
end

.record_tokens!(agent_type, tokens, tenant_id: nil, tenant_config: nil) ⇒ void

This method returns an undefined value.

Records token usage and checks for soft cap alerts

Parameters:

  • agent_type (String) —

    The agent class name

  • tokens (Integer) —

    The number of tokens used

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

    Optional tenant identifier (uses resolver if not provided)

  • tenant_config (Hash, nil) (defaults to: nil) —

    Optional runtime tenant config (takes priority over resolver/DB)



96
97
98
99
100
101
102
103
 
# File 'lib/ruby_llm/agents/infrastructure/budget_tracker.rb', line 96

def record_tokens!(agent_type, tokens, tenant_id: nil, tenant_config: nil)
  return if tokens.nil? || tokens <= 0

  tenant_id = resolve_tid(tenant_id)
  budget_config = Budget::ConfigResolver.resolve_budget_config(tenant_id, runtime_config: tenant_config)

  Budget::SpendRecorder.record_tokens!(agent_type, tokens, tenant_id: tenant_id, budget_config: budget_config)
end

.remaining_budget(scope, period, agent_type: nil, tenant_id: nil) ⇒ Float?

Returns the remaining budget for a scope and period

Parameters:

  • scope (Symbol) —

    :global or :agent

  • period (Symbol) —

    :daily or :monthly

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

    Required when scope is :agent

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

    Optional tenant identifier (uses resolver if not provided)

Returns:

  • (Float, nil) —

    Remaining budget in USD, or nil if no limit configured



134
135
136
137
138
139
 
# File 'lib/ruby_llm/agents/infrastructure/budget_tracker.rb', line 134

def remaining_budget(scope, period, agent_type: nil, tenant_id: nil)
  tenant_id = resolve_tid(tenant_id)
  budget_config = Budget::ConfigResolver.resolve_budget_config(tenant_id)

  Budget::BudgetQuery.remaining_budget(scope, period, agent_type: agent_type, tenant_id: tenant_id, budget_config: budget_config)
end

.remaining_token_budget(period, tenant_id: nil) ⇒ Integer?

Returns the remaining token budget for a period (global only)

Parameters:

  • period (Symbol) —

    :daily or :monthly

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

    Optional tenant identifier (uses resolver if not provided)

Returns:

  • (Integer, nil) —

    Remaining token budget, or nil if no limit configured



146
147
148
149
150
151
 
# File 'lib/ruby_llm/agents/infrastructure/budget_tracker.rb', line 146

def remaining_token_budget(period, tenant_id: nil)
  tenant_id = resolve_tid(tenant_id)
  budget_config = Budget::ConfigResolver.resolve_budget_config(tenant_id)

  Budget::BudgetQuery.remaining_token_budget(period, tenant_id: tenant_id, budget_config: budget_config)
end

.reset!(tenant_id: nil) ⇒ void

This method returns an undefined value.

Resets all budget counters (useful for testing)

Parameters:

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

    Optional tenant identifier to reset only that tenant’s counters



180
181
182
183
184
185
186
187
188
189
190
191
 
# File 'lib/ruby_llm/agents/infrastructure/budget_tracker.rb', line 180

def reset!(tenant_id: nil)
  tenant_id = resolve_tid(tenant_id)
  tenant_part = Budget::SpendRecorder.tenant_key_part(tenant_id)
  today = Budget::SpendRecorder.date_key_part(:daily)
  month = Budget::SpendRecorder.date_key_part(:monthly)

  BudgetTracker.cache_delete(BudgetTracker.cache_key("budget", tenant_part, today))
  BudgetTracker.cache_delete(BudgetTracker.cache_key("budget", tenant_part, month))

  # Reset memoized table existence check (useful for testing)
  Budget::ConfigResolver.reset_tenant_budget_table_check!
end

.status(agent_type: nil, tenant_id: nil) ⇒ Hash

Returns a summary of all budget statuses

Parameters:

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

    Optional agent type for per-agent budgets

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

    Optional tenant identifier (uses resolver if not provided)

Returns:

  • (Hash) —

    Budget status information



158
159
160
161
162
163
 
# File 'lib/ruby_llm/agents/infrastructure/budget_tracker.rb', line 158

def status(agent_type: nil, tenant_id: nil)
  tenant_id = resolve_tid(tenant_id)
  budget_config = Budget::ConfigResolver.resolve_budget_config(tenant_id)

  Budget::BudgetQuery.status(agent_type: agent_type, tenant_id: tenant_id, budget_config: budget_config)
end