Module: AIA::Directives::Checkpoint

Defined in:
lib/aia/directives/checkpoint.rb

Class Attribute Summary collapse

Class Method Summary collapse

Class Attribute Details

.checkpoint_counter ⇒ Object

Returns the value of attribute checkpoint_counter.



17
18
19
 
# File 'lib/aia/directives/checkpoint.rb', line 17

def checkpoint_counter
  @checkpoint_counter
end

.checkpoint_store ⇒ Object

Returns the value of attribute checkpoint_store.



17
18
19
 
# File 'lib/aia/directives/checkpoint.rb', line 17

def checkpoint_store
  @checkpoint_store
end

.last_checkpoint_name ⇒ Object

Returns the value of attribute last_checkpoint_name.



17
18
19
 
# File 'lib/aia/directives/checkpoint.rb', line 17

def last_checkpoint_name
  @last_checkpoint_name
end

Class Method Details

.checkpoint(args, _unused = nil) ⇒ Object Also known as: ckp, cp

//checkpoint [name] Creates a named checkpoint of the current conversation state. If no name is provided, uses an auto-incrementing number.



30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
 
# File 'lib/aia/directives/checkpoint.rb', line 30

def self.checkpoint(args, _unused = nil)
  name = args.empty? ? nil : args.join(' ').strip

  if name.nil? || name.empty?
    self.checkpoint_counter += 1
    name = checkpoint_counter.to_s
  end

  chats = get_chats
  return "Error: No active chat sessions found." if chats.nil? || chats.empty?

  # Deep copy messages from all chats
  first_chat_messages = chats.values.first&.messages || []
  checkpoint_store[name] = {
    messages: chats.transform_values { |chat|
      chat.messages.map { |msg| deep_copy_message(msg) }
    },
    position: first_chat_messages.size,
    created_at: Time.now,
    topic_preview: extract_last_user_message(first_chat_messages)
  }
  self.last_checkpoint_name = name

  puts "Checkpoint '#{name}' created at position #{checkpoint_store[name][:position]}."
  ""
end

.checkpoint_names ⇒ Object

Helper methods



193
194
195
 
# File 'lib/aia/directives/checkpoint.rb', line 193

def self.checkpoint_names
  checkpoint_store.keys
end

.checkpoint_positions ⇒ Object 



197
198
199
200
201
202
203
204
205
 
# File 'lib/aia/directives/checkpoint.rb', line 197

def self.checkpoint_positions
  positions = {}
  checkpoint_store.each do |name, data|
    pos = data[:position]
    positions[pos] ||= []
    positions[pos] << name
  end
  positions
end

.checkpoints_list(args, _unused = nil) ⇒ Object Also known as: checkpoints

//checkpoints Lists all available checkpoints with their details.



174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
 
# File 'lib/aia/directives/checkpoint.rb', line 174

def self.checkpoints_list(args, _unused = nil)
  if checkpoint_store.empty?
    puts "No checkpoints available."
    return ""
  end

  puts "\n=== Available Checkpoints ==="
  checkpoint_store.each do |name, data|
    created = data[:created_at]&.strftime('%H:%M:%S') || 'unknown'
    puts "  #{name}: position #{data[:position]}, created #{created}"
    if data[:topic_preview] && !data[:topic_preview].empty?
      puts "    → \"#{data[:topic_preview]}\""
    end
  end
  puts "=== End of Checkpoints ==="
  ""
end

.clear(args, _unused = nil) ⇒ Object

//clear Clears the conversation context, optionally keeping the system prompt.



106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
 
# File 'lib/aia/directives/checkpoint.rb', line 106

def self.clear(args, _unused = nil)
  keep_system = !args.include?('--all')

  chats = get_chats
  return "Error: No active chat sessions found." if chats.nil? || chats.empty?

  chats.each do |_model_id, chat|
    if keep_system
      system_msg = chat.messages.find { |m| m.role == :system }
      chat.instance_variable_set(:@messages, [])
      chat.add_message(system_msg) if system_msg
    else
      chat.instance_variable_set(:@messages, [])
    end
  end

  # Clear all checkpoints
  checkpoint_store.clear
  self.checkpoint_counter = 0
  self.last_checkpoint_name = nil

  "Chat context cleared."
end

.deep_copy_message(msg) ⇒ Object 



233
234
235
236
237
238
239
240
241
242
243
 
# File 'lib/aia/directives/checkpoint.rb', line 233

def self.deep_copy_message(msg)
  RubyLLM::Message.new(
    role: msg.role,
    content: msg.content,
    tool_calls: msg.tool_calls&.transform_values { |tc| tc.dup rescue tc },
    tool_call_id: msg.tool_call_id,
    input_tokens: msg.input_tokens,
    output_tokens: msg.output_tokens,
    model_id: msg.model_id
  )
end

.extract_last_user_message(messages, max_length: 70) ⇒ Object

Extract a preview of the last user message for checkpoint context



259
260
261
262
263
264
265
266
267
268
269
270
 
# File 'lib/aia/directives/checkpoint.rb', line 259

def self.extract_last_user_message(messages, max_length: 70)
  return "" if messages.nil? || messages.empty?

  # Find the last user message (not system, not assistant, not tool)
  last_user_msg = messages.reverse.find { |msg| msg.role == :user }
  return "" unless last_user_msg

  content = last_user_msg.content.to_s.strip
  # Collapse whitespace and truncate
  content = content.gsub(/\s+/, ' ')
  content.length > max_length ? "#{content[0..max_length - 4]}..." : content
end

.find_previous_checkpoint ⇒ Object

Find the previous checkpoint (second-to-last by position) This is used when //restore is called without a name



217
218
219
220
221
222
223
 
# File 'lib/aia/directives/checkpoint.rb', line 217

def self.find_previous_checkpoint
  return nil if checkpoint_store.size < 2

  # Sort checkpoints by position descending, return the second one
  sorted = checkpoint_store.sort_by { |_name, data| -data[:position] }
  sorted[1]&.first  # Return the name of the second-to-last checkpoint
end

.format_message_content(msg) ⇒ Object 



245
246
247
248
249
250
251
252
253
254
255
256
 
# File 'lib/aia/directives/checkpoint.rb', line 245

def self.format_message_content(msg)
  if msg.tool_call?
    tool_names = msg.tool_calls.values.map(&:name).join(', ')
    "[Tool calls: #{tool_names}]"
  elsif msg.tool_result?
    result_preview = msg.content.to_s[0..50]
    "[Tool result for: #{msg.tool_call_id}] #{result_preview}..."
  else
    content = msg.content.to_s
    content.length > 150 ? "#{content[0..147]}..." : content
  end
end

.get_chats ⇒ Object 



227
228
229
230
231
 
# File 'lib/aia/directives/checkpoint.rb', line 227

def self.get_chats
  return nil unless AIA.client.respond_to?(:chats)

  AIA.client.chats
end

.remove_invalid_checkpoints(max_position) ⇒ Object

Remove checkpoints with position > the given position Returns the count of removed checkpoints



209
210
211
212
213
 
# File 'lib/aia/directives/checkpoint.rb', line 209

def self.remove_invalid_checkpoints(max_position)
  invalid_names = checkpoint_store.select { |_name, data| data[:position] > max_position }.keys
  invalid_names.each { |name| checkpoint_store.delete(name) }
  invalid_names.size
end

.reset! ⇒ Object

Reset all checkpoint state (useful for testing)



20
21
22
23
24
 
# File 'lib/aia/directives/checkpoint.rb', line 20

def reset!
  @checkpoint_store = {}
  @checkpoint_counter = 0
  @last_checkpoint_name = nil
end

.restore(args, _unused = nil) ⇒ Object

//restore [name] Restores the conversation state to a previously saved checkpoint. If no name is provided, restores to the previous checkpoint (one step back).



60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
 
# File 'lib/aia/directives/checkpoint.rb', line 60

def self.restore(args, _unused = nil)
  name = args.empty? ? nil : args.join(' ').strip

  if name.nil? || name.empty?
    # Find the previous checkpoint (second-to-last by position)
    name = find_previous_checkpoint
    if name.nil?
      return "Error: No previous checkpoint to restore to."
    end
  end

  unless checkpoint_store.key?(name)
    available = checkpoint_names.empty? ? "none" : checkpoint_names.join(', ')
    return "Error: Checkpoint '#{name}' not found. Available: #{available}"
  end

  checkpoint_data = checkpoint_store[name]
  chats = get_chats

  return "Error: No active chat sessions found." if chats.nil? || chats.empty?

  # Restore messages to each chat
  checkpoint_data[:messages].each do |model_id, saved_messages|
    chat = chats[model_id]
    next unless chat

    # Replace the chat's messages with the saved ones
    restored_messages = saved_messages.map { |msg| deep_copy_message(msg) }
    chat.instance_variable_set(:@messages, restored_messages)
  end

  # Remove checkpoints that are now invalid (position > restored position)
  # because they reference conversation states that no longer exist
  restored_position = checkpoint_data[:position]
  removed_count = remove_invalid_checkpoints(restored_position)

  # Update last_checkpoint_name to this checkpoint
  self.last_checkpoint_name = name

  msg = "Context restored to checkpoint '#{name}' (position #{restored_position})."
  msg += " Removed #{removed_count} checkpoint(s) that were beyond this position." if removed_count > 0
  msg
end

.review(args, _unused = nil) ⇒ Object Also known as: context

//review Displays the current conversation context with checkpoint markers.



132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
 
# File 'lib/aia/directives/checkpoint.rb', line 132

def self.review(args, _unused = nil)
  chats = get_chats
  return "Error: No active chat sessions found." if chats.nil? || chats.empty?

  # For multi-model, show first chat's messages (they should be similar for user messages)
  first_chat = chats.values.first
  messages = first_chat&.messages || []

  puts "\n=== Chat Context (RubyLLM) ==="
  puts "Total messages: #{messages.size}"
  puts "Models: #{chats.keys.join(', ')}"
  puts "Checkpoints: #{checkpoint_names.join(', ')}" if checkpoint_names.any?
  puts

  positions = checkpoint_positions

  messages.each_with_index do |msg, index|
    # Show checkpoint marker if one exists at this position
    if positions[index]
      puts "📍 [Checkpoint: #{positions[index].join(', ')}]"
      puts "-" * 40
    end

    role = msg.role.to_s.capitalize
    content = format_message_content(msg)

    puts "#{index + 1}. [#{role}]: #{content}"
    puts
  end

  # Check for checkpoint at the end
  if positions[messages.size]
    puts "📍 [Checkpoint: #{positions[messages.size].join(', ')}]"
    puts "-" * 40
  end

  puts "=== End of Context ==="
  ""
end