Class: Roast::Workflow::ContextManager

Inherits:
Object
  • Object
show all
Defined in:
lib/roast/workflow/context_manager.rb

Instance Attribute Summary collapse

Instance Method Summary collapse

Constructor Details

#initialize(token_counter: nil, threshold_checker: nil) ⇒ ContextManager

Returns a new instance of ContextManager.



8
9
10
11
12
13
14
15
16
 
# File 'lib/roast/workflow/context_manager.rb', line 8

def initialize(token_counter: nil, threshold_checker: nil)
  @token_counter = token_counter || Services::TokenCountingService.new
  @threshold_checker = threshold_checker || Services::ContextThresholdChecker.new
  @total_tokens = 0
  @message_count = 0
  @config = default_config
  @last_actual_update = nil
  @estimated_tokens_since_update = 0
end

Instance Attribute Details

#total_tokens ⇒ Object (readonly)

Returns the value of attribute total_tokens.



6
7
8
 
# File 'lib/roast/workflow/context_manager.rb', line 6

def total_tokens
  @total_tokens
end

Instance Method Details

#check_warnings(token_count = @total_tokens) ⇒ Object 



43
44
45
46
47
48
49
50
51
52
53
54
55
 
# File 'lib/roast/workflow/context_manager.rb', line 43

def check_warnings(token_count = @total_tokens)
  return unless @config[:enabled]

  warning = @threshold_checker.check_warning_threshold(
    token_count,
    @config[:threshold],
    @config[:max_tokens],
  )

  if warning
    ActiveSupport::Notifications.instrument("roast.context_warning", warning)
  end
end

#configure(config) ⇒ Object 



18
19
20
 
# File 'lib/roast/workflow/context_manager.rb', line 18

def configure(config)
  @config = default_config.merge(config)
end

#reset ⇒ Object 



57
58
59
60
 
# File 'lib/roast/workflow/context_manager.rb', line 57

def reset
  @total_tokens = 0
  @message_count = 0
end

#should_compact?(token_count = @total_tokens) ⇒ Boolean

Returns:

  • (Boolean) 


33
34
35
36
37
38
39
40
41
 
# File 'lib/roast/workflow/context_manager.rb', line 33

def should_compact?(token_count = @total_tokens)
  return false unless @config[:enabled]

  @threshold_checker.should_compact?(
    token_count,
    @config[:threshold],
    @config[:max_tokens],
  )
end

#statistics ⇒ Object 



62
63
64
65
66
67
68
 
# File 'lib/roast/workflow/context_manager.rb', line 62

def statistics
  {
    total_tokens: @total_tokens,
    message_count: @message_count,
    average_tokens_per_message: @message_count > 0 ? @total_tokens / @message_count : 0,
  }
end

#track_usage(messages) ⇒ Object 



22
23
24
25
26
27
28
29
30
31
 
# File 'lib/roast/workflow/context_manager.rb', line 22

def track_usage(messages)
  current_tokens = @token_counter.count_messages(messages)
  @total_tokens += current_tokens
  @message_count += messages.size

  {
    current_tokens: current_tokens,
    total_tokens: @total_tokens,
  }
end

#update_with_actual_usage(actual_total) ⇒ Object 



70
71
72
73
74
75
76
 
# File 'lib/roast/workflow/context_manager.rb', line 70

def update_with_actual_usage(actual_total)
  return unless actual_total && actual_total > 0

  @total_tokens = actual_total
  @last_actual_update = Time.now
  @estimated_tokens_since_update = 0
end