Module: Legion::LLM::Inference::Steps::TokenBudget

Includes:

Logging, Legion::Logging::Helper

Included in:

Executor

Defined in:

lib/legion/llm/inference/steps/token_budget.rb

Instance Method Summary

• #step_token_budget ⇒ Object

Instance Method Details

#step_token_budget ⇒ Object

# File 'lib/legion/llm/inference/steps/token_budget.rb', line 13

def step_token_budget
  max_input = @request.extra&.dig(:max_input_tokens)
  session_total = Legion::LLM::Metering::Tokens.total_tokens
  session_limit = Legion::LLM::Metering::Tokens.summary[:session_max_tokens]
  log_step_debug(:token_budget, :checking, max_input_tokens: max_input || 'none',
                                           session_total: session_total, session_limit: session_limit || 'none')
  check_input_cap(max_input) if max_input&.positive?
  check_session_budget
  log_step_debug(:token_budget, :passed, max_input_tokens: max_input || 'none')
rescue Legion::LLM::TokenBudgetExceeded
  log_step_info(:token_budget, :blocked)
  raise
rescue StandardError => e
  @warnings << { type: :token_budget_check_failed, message: e.message }
  handle_exception(e, level: :warn, operation: 'llm.pipeline.steps.token_budget')
end