Class: Phronomy::Agent::ContextParts::Budget::TokenBudgetPacker

Inherits:
Object
  • Object
show all
Defined in:
lib/phronomy/agent/context_parts/budget/token_budget_packer.rb

Instance Method Summary collapse

Instance Method Details

#pack(request:, units:) ⇒ Object



8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
# File 'lib/phronomy/agent/context_parts/budget/token_budget_packer.rb', line 8

def pack(request:, units:)
  return Array(units).freeze unless request.token_budget

  candidate_index = request.candidates.to_h { |candidate| [candidate.candidate_id, candidate] }
  mandatory = Integer(request.["mandatory_token_estimate"] || 0)
  limit = request.token_budget.effective_input_limit
  if mandatory > limit
    raise Phronomy::ContextBudgetExceededError,
      "Mandatory content (estimated #{mandatory} tokens) exceeds " \
      "available input budget (#{limit} tokens)"
  end

  remaining = limit - mandatory
  required, optional = Array(units).partition { |unit| unit.requirement != :optional }
  required_cost = required.sum { |unit| unit_cost(unit, candidate_index) }
  if required_cost > remaining
    raise Phronomy::ContextBudgetExceededError,
      "Required Context (estimated #{required_cost} tokens) exceeds " \
      "remaining input budget (#{remaining} tokens after mandatory content)"
  end

  selected = required.dup
  remaining -= required_cost
  optional.each do |unit|
    cost = unit_cost(unit, candidate_index)
    next if cost > remaining

    selected << unit
    remaining -= cost
  end
  selected.freeze
end