Class: RubynCode::Observability::BudgetEnforcer

Inherits:
Object
  • Object
show all
Defined in:
lib/rubyn_code/observability/budget_enforcer.rb

Overview

Tracks API spend and halts the agent when session or daily budgets are exceeded. Cost records are persisted to SQLite so budgets survive restarts.

Constant Summary collapse

DEFAULT_SESSION_LIMIT = 
5.00
DEFAULT_DAILY_LIMIT = 
10.00
TABLE_NAME = 
'cost_records'

Instance Method Summary collapse

Constructor Details

#initialize(db, session_id:, session_limit: DEFAULT_SESSION_LIMIT, daily_limit: DEFAULT_DAILY_LIMIT) ⇒ BudgetEnforcer

Returns a new instance of BudgetEnforcer.

Parameters:

  • db (DB::Connection) —

    database connection

  • session_id (String) —

    current session identifier

  • session_limit (Float) (defaults to: DEFAULT_SESSION_LIMIT) —

    maximum USD spend per session

  • daily_limit (Float) (defaults to: DEFAULT_DAILY_LIMIT) —

    maximum USD spend per calendar day



22
23
24
25
26
27
28
29
 
# File 'lib/rubyn_code/observability/budget_enforcer.rb', line 22

def initialize(db, session_id:, session_limit: DEFAULT_SESSION_LIMIT, daily_limit: DEFAULT_DAILY_LIMIT)
  @db            = db
  @session_id    = session_id
  @session_limit = session_limit.to_f
  @daily_limit   = daily_limit.to_f

  ensure_table_exists
end

Instance Method Details

#check! ⇒ void

This method returns an undefined value.

Raises BudgetExceededError if either the session or daily budget is exceeded.

Raises:



57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
 
# File 'lib/rubyn_code/observability/budget_enforcer.rb', line 57

def check!
  sc = session_cost
  if sc >= @session_limit
    raise BudgetExceededError,
          format('Session budget exceeded: $%<cost>.4f >= $%<limit>.2f limit',
                 cost: sc, limit: @session_limit)
  end

  dc = daily_cost
  return unless dc >= @daily_limit

  raise BudgetExceededError,
        format('Daily budget exceeded: $%<cost>.4f >= $%<limit>.2f limit',
               cost: dc, limit: @daily_limit)
end

#daily_cost ⇒ Float

Returns the total cost accumulated today (UTC).

Returns:

  • (Float) —

    total daily cost in USD



87
88
89
90
91
92
93
94
 
# File 'lib/rubyn_code/observability/budget_enforcer.rb', line 87

def daily_cost
  today = Time.now.utc.strftime('%Y-%m-%d')
  rows = @db.query(
    "SELECT COALESCE(SUM(cost_usd), 0.0) AS total FROM #{TABLE_NAME} WHERE created_at >= ?",
    ["#{today}T00:00:00Z"]
  ).to_a
  extract_total(rows)
end

#record!(model:, input_tokens:, output_tokens:, **opts) ⇒ CostRecord

Records a cost entry for an API call and persists it to the database.

Parameters:

  • model (String) —

    the model identifier

  • input_tokens (Integer) —

    input token count

  • output_tokens (Integer) —

    output token count

  • cache_read_tokens (Integer) —

    cache-read token count

  • cache_write_tokens (Integer) —

    cache-write token count

  • request_type (String) —

    the type of request (e.g., “chat”, “compact”)

Returns:



40
41
42
43
44
45
46
47
48
49
50
51
 
# File 'lib/rubyn_code/observability/budget_enforcer.rb', line 40

def record!(model:, input_tokens:, output_tokens:, **opts)
  cache_read  = opts.fetch(:cache_read_tokens, 0)
  cache_write = opts.fetch(:cache_write_tokens, 0)
  req_type    = opts.fetch(:request_type, 'chat')

  cost = CostCalculator.calculate(
    model: model, input_tokens: input_tokens, output_tokens: output_tokens,
    cache_read_tokens: cache_read, cache_write_tokens: cache_write
  )

  persist_cost_record(model, input_tokens, output_tokens, cache_read, cache_write, cost, req_type)
end

#remaining_budget ⇒ Float

Returns the smaller of the session and daily remaining budgets.

Returns:

  • (Float) —

    remaining budget in USD



99
100
101
102
103
 
# File 'lib/rubyn_code/observability/budget_enforcer.rb', line 99

def remaining_budget
  session_remaining = @session_limit - session_cost
  daily_remaining   = @daily_limit - daily_cost
  [session_remaining, daily_remaining].min
end

#session_cost ⇒ Float

Returns the total cost accumulated in the current session.

Returns:

  • (Float) —

    total session cost in USD



76
77
78
79
80
81
82
 
# File 'lib/rubyn_code/observability/budget_enforcer.rb', line 76

def session_cost
  rows = @db.query(
    "SELECT COALESCE(SUM(cost_usd), 0.0) AS total FROM #{TABLE_NAME} WHERE session_id = ?",
    [@session_id]
  ).to_a
  extract_total(rows)
end