Class: LlmGateway::Agents::InMemorySessionManager

Inherits:
Object
  • Object
show all
Defined in:
lib/llm_gateway/agents/in_memory_session_manager.rb

Direct Known Subclasses

FileSessionManager

Constant Summary collapse

MESSAGE_QUEUED =
:queued
MESSAGE_STARTED =
:started
QUEUES =
[ :steer, :follow_up, :next_turn ].freeze
DRAIN_MODES =
[ :one_at_a_time, :all ].freeze

Instance Attribute Summary collapse

Instance Method Summary collapse

Constructor Details

#initialize(session_id = nil) ⇒ InMemorySessionManager

Returns a new instance of InMemorySessionManager.



16
17
18
19
20
# File 'lib/llm_gateway/agents/in_memory_session_manager.rb', line 16

def initialize(session_id = nil)
  @state = :idle
  @session_id = session_id
  @message_queues = Hash.new { |hash, key| hash[key] = [] }
end

Instance Attribute Details

#session_idObject (readonly)

Returns the value of attribute session_id.



14
15
16
# File 'lib/llm_gateway/agents/in_memory_session_manager.rb', line 14

def session_id
  @session_id
end

#session_startObject (readonly)

Returns the value of attribute session_start.



14
15
16
# File 'lib/llm_gateway/agents/in_memory_session_manager.rb', line 14

def session_start
  @session_start
end

Instance Method Details

#active_messagesObject



97
98
99
# File 'lib/llm_gateway/agents/in_memory_session_manager.rb', line 97

def active_messages
  active_message_events.map { |event| event[:data] }
end

#build_model_input_messagesObject



124
125
126
127
128
# File 'lib/llm_gateway/agents/in_memory_session_manager.rb', line 124

def build_model_input_messages
  return active_messages unless last_compaction_entry

  [ last_compaction_entry[:data], *active_messages ]
end

#busy!Object



22
23
24
# File 'lib/llm_gateway/agents/in_memory_session_manager.rb', line 22

def busy!
  @state = :busy
end

#busy?Boolean

Returns:

  • (Boolean)


44
45
46
# File 'lib/llm_gateway/agents/in_memory_session_manager.rb', line 44

def busy?
  @state == :busy
end

#compaction(adapter) ⇒ Object



140
141
142
143
144
145
146
147
148
149
150
151
152
153
# File 'lib/llm_gateway/agents/in_memory_session_manager.rb', line 140

def compaction(adapter)
  response = adapter.stream(
    active_messages,
    system: "Summarize the conversation so far for future context.",
    tools: []
  )
  message = response.to_h

  push_entry(
    type: "compaction",
    usage: message_usage(message),
    data: message
  )
end

#drain_message_queue(queue = :next_turn, mode: :all) ⇒ Object



30
31
32
33
34
# File 'lib/llm_gateway/agents/in_memory_session_manager.rb', line 30

def drain_message_queue(queue = :next_turn, mode: :all)
  messages = queued_messages(queue, mode)
  messages.each { |message| push_message(message) }
  messages
end

#eventsObject



120
121
122
# File 'lib/llm_gateway/agents/in_memory_session_manager.rb', line 120

def events
  @events ||= [ new_session_event ]
end

#events_until(event_id) ⇒ Object

Raises:

  • (ArgumentError)


113
114
115
116
117
118
# File 'lib/llm_gateway/agents/in_memory_session_manager.rb', line 113

def events_until(event_id)
  index = events.index { |event| event[:id] == event_id }
  raise ArgumentError, "Event not found in session: #{event_id}" unless index

  events[0..index]
end

#idle!Object



26
27
28
# File 'lib/llm_gateway/agents/in_memory_session_manager.rb', line 26

def idle!
  @state = :idle
end

#last_assistant_message_atObject



135
136
137
138
# File 'lib/llm_gateway/agents/in_memory_session_manager.rb', line 135

def last_assistant_message_at
  entry = active_message_events.reverse.find { |event| event.dig(:data, :role) == "assistant" }
  Time.parse(entry[:timestamp]) if entry
end

#last_message_idObject



101
102
103
# File 'lib/llm_gateway/agents/in_memory_session_manager.rb', line 101

def last_message_id
  message_events.last&.dig(:id)
end

#last_model_usedObject



105
106
107
# File 'lib/llm_gateway/agents/in_memory_session_manager.rb', line 105

def last_model_used
  events.reverse.find { |event| event[:type] == "model_change" }&.dig(:model_id)
end

#last_reasoning_level_usedObject



109
110
111
# File 'lib/llm_gateway/agents/in_memory_session_manager.rb', line 109

def last_reasoning_level_used
  events.reverse.find { |event| event[:type] == "reasoning_change" }&.dig(:reasoning)
end

#push_entry(entry) ⇒ Object



84
85
86
87
88
89
90
91
92
93
94
95
# File 'lib/llm_gateway/agents/in_memory_session_manager.rb', line 84

def push_entry(entry)
  id = SecureRandom.uuid
  new_entry = {
    id: id,
    parent_id: parent_id_for_new_entry,
    timestamp: Time.now.iso8601,
    **entry
  }

  persist_entry(new_entry)
  new_entry
end

#push_message(payload) ⇒ Object



74
75
76
77
78
79
80
81
82
# File 'lib/llm_gateway/agents/in_memory_session_manager.rb', line 74

def push_message(payload)
  payload = payload.deep_symbolize_keys

  push_entry(
    type: "message",
    usage: message_usage(payload),
    data: payload,
  )
end

#push_message_to_queue(message, queue = :next_turn) ⇒ Object



40
41
42
# File 'lib/llm_gateway/agents/in_memory_session_manager.rb', line 40

def push_message_to_queue(message, queue = :next_turn)
  @message_queues[validate_queue!(queue)] << message
end

#queued_messages?(queue = :next_turn) ⇒ Boolean

Returns:

  • (Boolean)


36
37
38
# File 'lib/llm_gateway/agents/in_memory_session_manager.rb', line 36

def queued_messages?(queue = :next_turn)
  @message_queues[validate_queue!(queue)].any?
end

#start_or_enqueue_user_message(payload, queue: :next_turn) ⇒ Object



62
63
64
65
66
67
68
69
70
71
72
# File 'lib/llm_gateway/agents/in_memory_session_manager.rb', line 62

def start_or_enqueue_user_message(payload, queue: :next_turn)
  if busy?
    push_message_to_queue(payload, queue)
    MESSAGE_QUEUED
  else
    yield if block_given?
    push_message(payload)
    busy!
    MESSAGE_STARTED
  end
end

#total_tokensObject



130
131
132
133
# File 'lib/llm_gateway/agents/in_memory_session_manager.rb', line 130

def total_tokens
  entry = active_message_events.reverse.find { |event| event.dig(:usage, :total_tokens) }
  entry&.dig(:usage, :total_tokens) || 0
end

#validate_drain_mode!(mode) ⇒ Object

Raises:

  • (ArgumentError)


55
56
57
58
59
60
# File 'lib/llm_gateway/agents/in_memory_session_manager.rb', line 55

def validate_drain_mode!(mode)
  mode = mode.to_sym
  raise ArgumentError, "Invalid queue drain mode: #{mode}" unless DRAIN_MODES.include?(mode)

  mode
end

#validate_queue!(queue) ⇒ Object

Raises:

  • (ArgumentError)


48
49
50
51
52
53
# File 'lib/llm_gateway/agents/in_memory_session_manager.rb', line 48

def validate_queue!(queue)
  queue = queue.to_sym
  raise ArgumentError, "Invalid queue mode: #{queue}" unless QUEUES.include?(queue)

  queue
end