Class: Smith::Tool::CallAllowance

Inherits:
Object
  • Object
show all
Defined in:
lib/smith/tool/call_allowance.rb

Constant Summary collapse

EXHAUSTION_POLICIES = 
%i[raise complete].freeze

Class Method Summary collapse

Instance Method Summary collapse

Constructor Details

#initialize(remaining, on_exhaustion: :raise) ⇒ CallAllowance

Returns a new instance of CallAllowance.

Raises:

  • (ArgumentError) 


12
13
14
15
16
17
18
 
# File 'lib/smith/tool/call_allowance.rb', line 12

def initialize(remaining, on_exhaustion: :raise)
  invalid = (remaining.is_a?(Integer) && remaining.negative?) ||
            (!remaining.is_a?(Integer) && !remaining.is_a?(CallBudget))
  raise ArgumentError, "tool call allowance must be a non-negative integer" if invalid

  initialize_scope(CallBudget.coerce(remaining), on_exhaustion:, parent: nil)
end

Class Method Details

.charge_legacy!(allowance) ⇒ Object 



8
9
10
 
# File 'lib/smith/tool/call_allowance.rb', line 8

def self.charge_legacy!(allowance)
  LegacyCallAllowance.charge!(allowance) { yield if block_given? }
end

Instance Method Details

#[](key) ⇒ Object 



88
89
90
 
# File 'lib/smith/tool/call_allowance.rb', line 88

def [](key)
  remaining if key == :remaining
end

#charge!(tool_name = nil) ⇒ Object 



26
27
28
29
30
31
32
33
34
35
36
37
 
# File 'lib/smith/tool/call_allowance.rb', line 26

def charge!(tool_name = nil)
  batch = CallBatch.coerce(exact? ? [tool_name] : 1, exact: exact?)

  synchronize do
    Thread.handle_interrupt(Object => :never) do
      raise BudgetExceeded, "agent tool_calls budget exceeded" unless reservable?(batch)

      yield if block_given?
      consume!(batch)
    end
  end
end

#complete_on_exhaustion? ⇒ Boolean

Returns:

  • (Boolean) 


76
77
78
 
# File 'lib/smith/tool/call_allowance.rb', line 76

def complete_on_exhaustion?
  @on_exhaustion == :complete
end

#exact? ⇒ Boolean

Returns:

  • (Boolean) 


80
81
82
 
# File 'lib/smith/tool/call_allowance.rb', line 80

def exact?
  lineage.any? { _1.__send__(:budget).exact? }
end

#remaining ⇒ Object 



39
40
41
 
# File 'lib/smith/tool/call_allowance.rb', line 39

def remaining
  synchronize { lineage.map { _1.__send__(:counter).remaining }.min }
end

#remaining_for(tool_name) ⇒ Object 



43
44
45
46
47
48
49
50
51
 
# File 'lib/smith/tool/call_allowance.rb', line 43

def remaining_for(tool_name)
  name = canonical_tool_name(tool_name)
  synchronize do
    exact_scopes = lineage.select { _1.__send__(:budget).exact? }
    return nil if exact_scopes.empty?

    exact_scopes.map { _1.__send__(:counter).remaining_for(name) }.min
  end
end

#reserve_batch(tool_names_or_size, ledger: nil) ⇒ Object 



53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
 
# File 'lib/smith/tool/call_allowance.rb', line 53

def reserve_batch(tool_names_or_size, ledger: nil)
  batch = CallBatch.coerce(tool_names_or_size, exact: exact?)

  synchronize do
    Thread.handle_interrupt(Object => :never) do
      return unless reservable?(batch)

      workflow_ledger = ledger if ledger&.limits&.key?(:tool_calls)
      ledger_reservation = workflow_ledger&.reserve!(:tool_calls, batch.size)
      reservation = CallReservation.new(
        limit: batch.size,
        ledger: workflow_ledger,
        ledger_reservation: ledger_reservation
      )

      consume!(batch)
      reservation
    end
  end
rescue BudgetExceeded
  nil
end

#scope(budget, on_exhaustion: @on_exhaustion) ⇒ Object 



20
21
22
23
24
 
# File 'lib/smith/tool/call_allowance.rb', line 20

def scope(budget, on_exhaustion: @on_exhaustion)
  self.class.allocate.tap do |allowance|
    allowance.__send__(:initialize_scope, CallBudget.coerce(budget), on_exhaustion:, parent: self)
  end
end

#used? ⇒ Boolean

Returns:

  • (Boolean) 


84
85
86
 
# File 'lib/smith/tool/call_allowance.rb', line 84

def used?
  synchronize { counter.used? }
end