Class: TurnKit::Budget

Inherits:
Object
  • Object
show all
Defined in:
lib/turnkit/budget.rb

Instance Attribute Summary collapse

Class Method Summary collapse

Instance Method Summary collapse

Constructor Details

#initialize(max_iterations:, timeout:, max_depth:, max_tool_executions:, max_tool_executions_by_name: {}, max_spend: nil, root_started_at: Clock.now) ⇒ Budget

Returns a new instance of Budget.



15
16
17
18
19
20
21
22
23
24
25
26
27
28
 
# File 'lib/turnkit/budget.rb', line 15

def initialize(max_iterations:, timeout:, max_depth:, max_tool_executions:, max_tool_executions_by_name: {}, max_spend: nil, root_started_at: Clock.now)
  @root_started_at = root_started_at
  @max_iterations = max_iterations
  @timeout = timeout
  @max_depth = max_depth
  @max_tool_executions = max_tool_executions
  @max_tool_executions_by_name = normalize_tool_limits(max_tool_executions_by_name)
  @max_spend = max_spend
  @iterations = 0
  @tool_executions = 0
  @tool_executions_by_name = Hash.new(0)
  @cost = 0
  @mutex = Mutex.new
end

Instance Attribute Details

#max_depth ⇒ Object (readonly)

Returns the value of attribute max_depth.



5
6
7
 
# File 'lib/turnkit/budget.rb', line 5

def max_depth
  @max_depth
end

#max_iterations ⇒ Object (readonly)

Returns the value of attribute max_iterations.



5
6
7
 
# File 'lib/turnkit/budget.rb', line 5

def max_iterations
  @max_iterations
end

#max_spend ⇒ Object (readonly)

Returns the value of attribute max_spend.



5
6
7
 
# File 'lib/turnkit/budget.rb', line 5

def max_spend
  @max_spend
end

#max_tool_executions ⇒ Object (readonly)

Returns the value of attribute max_tool_executions.



5
6
7
 
# File 'lib/turnkit/budget.rb', line 5

def max_tool_executions
  @max_tool_executions
end

#max_tool_executions_by_name ⇒ Object (readonly)

Returns the value of attribute max_tool_executions_by_name.



5
6
7
 
# File 'lib/turnkit/budget.rb', line 5

def max_tool_executions_by_name
  @max_tool_executions_by_name
end

#root_started_at ⇒ Object (readonly)

Returns the value of attribute root_started_at.



5
6
7
 
# File 'lib/turnkit/budget.rb', line 5

def root_started_at
  @root_started_at
end

#timeout ⇒ Object (readonly)

Returns the value of attribute timeout.



5
6
7
 
# File 'lib/turnkit/budget.rb', line 5

def timeout
  @timeout
end

Class Method Details

.resume(store:, root_turn_id:, limits: {}) ⇒ Object 



7
8
9
10
11
12
13
 
# File 'lib/turnkit/budget.rb', line 7

def self.resume(store:, root_turn_id:, limits: {})
  turns = store.list_turns(root_turn_id: root_turn_id)
  root = turns.find { |turn| turn.fetch("id") == root_turn_id } || turns.first || {}
  budget = new(**limits.merge(root_started_at: root["started_at"] || Clock.now))
  budget.seed!(turns: turns, tool_executions: turns.flat_map { |turn| store.list_tool_executions(turn_id: turn.fetch("id")) })
  budget
end

Instance Method Details

#add_cost!(cost) ⇒ Object 



65
66
67
68
69
70
71
72
 
# File 'lib/turnkit/budget.rb', line 65

def add_cost!(cost)
  return unless cost && max_spend

  @mutex.synchronize do
    @cost += cost.to_f
    raise BudgetError, "cost limit reached" if @cost > max_spend
  end
end

#add_usage!(usage) ⇒ Object 



61
62
63
 
# File 'lib/turnkit/budget.rb', line 61

def add_usage!(usage)
  add_cost!(usage&.cost)
end

#check!(depth:) ⇒ Object

Raises:



74
75
76
77
 
# File 'lib/turnkit/budget.rb', line 74

def check!(depth:)
  raise BudgetError, "maximum sub-agent depth reached" if max_depth && depth > max_depth
  raise BudgetError, "turn timed out" if timeout && Clock.now >= root_started_at + timeout
end

#count_iteration! ⇒ Object 



41
42
43
44
45
46
47
 
# File 'lib/turnkit/budget.rb', line 41

def count_iteration!
  @mutex.synchronize do
    raise BudgetError, "maximum iterations reached" if max_iterations && @iterations >= max_iterations

    @iterations += 1
  end
end

#count_tool_execution!(name = nil) ⇒ Object 



49
50
51
52
53
54
55
56
57
58
59
 
# File 'lib/turnkit/budget.rb', line 49

def count_tool_execution!(name = nil)
  @mutex.synchronize do
    key = name.to_s if name
    limit = max_tool_executions_by_name[key] if key
    raise BudgetError, "maximum tool executions reached" if max_tool_executions && @tool_executions >= max_tool_executions
    raise BudgetError, "maximum executions reached for tool #{key}" if limit && @tool_executions_by_name[key] >= limit

    @tool_executions += 1
    @tool_executions_by_name[key] += 1 if key
  end
end

#seed!(turns:, tool_executions:) ⇒ Object 



30
31
32
33
34
35
36
37
38
39
 
# File 'lib/turnkit/budget.rb', line 30

def seed!(turns:, tool_executions:)
  @mutex.synchronize do
    @iterations = Array(turns).sum { |turn| (turn["options"] || {})["iterations"].to_i }
    completed = Array(tool_executions).select { |execution| %w[completed failed].include?(execution["status"]) && !execution.dig("error", "details", "budget_denied") }
    @tool_executions = completed.length
    completed.each { |execution| @tool_executions_by_name[execution.fetch("tool_name").to_s] += 1 }
    @cost = Array(turns).sum { |turn| turn["cost"].to_f }
  end
  self
end