Class: DSPy::DeepSearch::TokenBudget

Inherits:
Object
  • Object
show all
Extended by:
T::Sig
Defined in:
lib/dspy/deep_search/token_budget.rb

Defined Under Namespace

Classes: Exceeded

Instance Attribute Summary collapse

Instance Method Summary collapse

Constructor Details

#initialize(limit:) ⇒ TokenBudget 



17
18
19
20
 
# File 'lib/dspy/deep_search/token_budget.rb', line 17

def initialize(limit:)
  @limit = limit
  @total_tokens = T.let(0, Integer)
end

Instance Attribute Details

#limit ⇒ Object (readonly)

Returns the value of attribute limit.



11
12
13
 
# File 'lib/dspy/deep_search/token_budget.rb', line 11

def limit
  @limit
end

#total_tokens ⇒ Object (readonly)

Returns the value of attribute total_tokens.



14
15
16
 
# File 'lib/dspy/deep_search/token_budget.rb', line 14

def total_tokens
  @total_tokens
end

Instance Method Details

#track!(prompt_tokens:, completion_tokens:) ⇒ Object 



28
29
30
31
32
33
34
35
36
37
38
39
40
 
# File 'lib/dspy/deep_search/token_budget.rb', line 28

def track!(prompt_tokens:, completion_tokens:)
  prompt = T.must(prompt_tokens)
  completion = T.must(completion_tokens)

  increment = prompt + completion
  new_total = @total_tokens + increment

  if new_total >= limit
    raise Exceeded, "Token budget exceeded: #{new_total}/#{limit}"
  end

  @total_tokens = new_total
end