Class: Llm::MessageContainer

Inherits:
Object
  • Object
show all
Defined in:
lib/llm/message_container.rb

Instance Attribute Summary collapse

Instance Method Summary collapse

Constructor Details

#initialize ⇒ MessageContainer

Returns a new instance of MessageContainer.



4
5
6
7
8
9
10
 
# File 'lib/llm/message_container.rb', line 4

def initialize
  @messages = []
  @metas = []
  @token_encoder = Tiktoken.get_encoding("cl100k_base")

  self.add_default_system_message!
end

Instance Attribute Details

#messages ⇒ Object (readonly)

Returns the value of attribute messages.



3
4
5
 
# File 'lib/llm/message_container.rb', line 3

def messages
  @messages
end

Instance Method Details

#add_default_system_message! ⇒ Object 



58
59
60
 
# File 'lib/llm/message_container.rb', line 58

def add_default_system_message!
  self.add_system_message("Current time is #{Time.now}")
end

#add_raw_message(message) ⇒ Object 



22
23
24
25
26
 
# File 'lib/llm/message_container.rb', line 22

def add_raw_message(message)
  message = message.with_indifferent_access
  @metas << { index: @metas.size, token: @token_encoder.encode(message[:content].to_s).size, role: message[:role] }
  @messages << message
end

#add_raw_messages(messages) ⇒ Object 



32
33
34
35
36
 
# File 'lib/llm/message_container.rb', line 32

def add_raw_messages(messages)
  messages.each do |message|
    self.add_raw_message(message)
  end
end

#add_system_message(content) ⇒ Object 



12
13
14
15
 
# File 'lib/llm/message_container.rb', line 12

def add_system_message(content)
  @metas << { index: @metas.size, token: @token_encoder.encode(content.to_s).size, role: :system }
  @messages << { role: :system, content: }
end

#add_user_message(content) ⇒ Object 



17
18
19
20
 
# File 'lib/llm/message_container.rb', line 17

def add_user_message(content)
  @metas << { index: @metas.size, token: @token_encoder.encode(content.to_s).size, role: :user }
  @messages << { role: :user, content: }
end

#to_capped_messages(token_limit: 28_000) ⇒ Object 



38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
 
# File 'lib/llm/message_container.rb', line 38

def to_capped_messages(token_limit: 28_000)
  return @messages if self.total_token <= token_limit

  # システムメッセージは消さない
  system_metas, not_system_metas = @metas.partition { |message| message[:role] == :system }
  current_token = system_metas.map { |meta| meta[:token] }.sum
  filtered_indexes = system_metas.map { |meta| meta[:index] }

  # 新しいメッセージを優先的に残す
  not_system_metas.reverse.each do |meta|
    break if token_limit < current_token + meta[:token]

    current_token += meta[:token]
    filtered_indexes << meta[:index]
  end

  # システムメッセージが挟み込まれている場合も、並び順を維持すること
  @messages.select.with_index { |message, i| filtered_indexes.include?(i) }
end

#total_token ⇒ Object 



28
29
30
 
# File 'lib/llm/message_container.rb', line 28

def total_token
  @metas.map { |meta| meta[:token] }.sum
end