Class: TurnKit::Budget

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

Instance Attribute Summary collapse

Instance Method Summary collapse

Constructor Details

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

Returns a new instance of Budget.



7
8
9
10
11
12
13
14
15
16
17
18
19
20
 
# File 'lib/turnkit/budget.rb', line 7

def initialize(max_iterations:, timeout:, max_depth:, max_tool_executions:, max_tool_executions_by_name: {}, cost_limit: 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)
  @cost_limit = cost_limit
  @iterations = 0
  @tool_executions = 0
  @tool_executions_by_name = Hash.new(0)
  @cost = 0
  @mutex = Mutex.new
end

Instance Attribute Details

#cost_limit ⇒ Object (readonly)

Returns the value of attribute cost_limit.



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

def cost_limit
  @cost_limit
end

#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_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

Instance Method Details

#add_cost!(cost) ⇒ Object 



46
47
48
49
50
51
52
53
 
# File 'lib/turnkit/budget.rb', line 46

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

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

#add_usage!(usage) ⇒ Object 



42
43
44
 
# File 'lib/turnkit/budget.rb', line 42

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

#check!(depth:) ⇒ Object

Raises:



55
56
57
58
 
# File 'lib/turnkit/budget.rb', line 55

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 



22
23
24
25
26
27
28
 
# File 'lib/turnkit/budget.rb', line 22

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 



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

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