Class: LlamaBotRails::ChatChannel

Inherits:
ApplicationCable::Channel show all
Defined in:
app/channels/llama_bot_rails/chat_channel.rb

Instance Method Summary collapse

Instance Method Details

#receive(data) ⇒ Object

Receive messages from _chat.html.erb frontend and send to llamabot FastAPI backend, frontend comes from the llamabot/_chat.html.erb chatbot, sent through external websocket to FastAPI/Python backend.



125
126
127
128
129
130
131
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
# File 'app/channels/llama_bot_rails/chat_channel.rb', line 125

def receive(data)
  begin
    #used to validate the message before it's sent to the llamabot-backend.

    # Get the currently logged in user from the environment.
    current_user = LlamaBotRails.current_user_resolver.call(connection.env)

    @api_token = Rails.application.message_verifier(:llamabot_ws).generate(
      { session_id: SecureRandom.uuid, user_id: current_user&.id},
      expires_in: 30.minutes
    )

    #This could be an example of how we might implement hooks & filters in the future.
    validate_message(data) #Placeholder for now, we are using this to mock errors being thrown. In the future, we can add actual validation logic.        
    # Forward the processed data to the LlamaBot Backend Socket
    message = data["message"]

    builder = state_builder_class.new(
      params: data,
      context: { api_token: @api_token }.with_indifferent_access
    )

    # 2. Construct the LangGraph-ready state
    state_payload = builder.build

    # 3. Ship it over the existing WebSocket
    send_to_external_application(state_payload)

    # Log the incoming WebSocket data
    Rails.logger.info "[LlamaBot] Got message from Javascript LlamaBot Frontend: #{data.inspect}"
  rescue => e
    Rails.logger.error "[LlamaBot] Error in receive method: #{e.message}"
    Rails.logger.error "[LlamaBot] Backtrace: #{e.backtrace.join("\n")}"
    send_message_to_frontend("error", e.message)
  end
end

#send_message_to_frontend(type, message, trace_info = nil) ⇒ Object



162
163
164
165
166
167
168
169
170
171
172
173
174
175
# File 'app/channels/llama_bot_rails/chat_channel.rb', line 162

def send_message_to_frontend(type, message, trace_info = nil)
  
  # Log trace info for debugging
  Rails.logger.info "[LlamaBot] TRACE INFO DEBUG: Type: #{type}, Has trace info: #{trace_info.present?}"

  message_data = {
    type: type,
    content: message
  }
  
  formatted_message = { message: message_data.to_json }.to_json
  
  ActionCable.server.broadcast "chat_channel_#{params[:session_id]}", formatted_message
end

#subscribedObject

_chat.html.erb front-end subscribes to this channel in _websocket.html.erb.



38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
# File 'app/channels/llama_bot_rails/chat_channel.rb', line 38

def subscribed
  begin
    stream_from "chat_channel_#{params[:session_id]}" # Public stream for session-based messages <- this is the channel we're subscribing to in _websocket.html.erb
    Rails.logger.info "[LlamaBot] Subscribed to chat channel with session ID: #{params[:session_id]}"
    @agent_state_builder_class = params[:agent_state_builder_class]
    if @agent_state_builder_class.blank? #defaults to the config file if not provided
      @agent_state_builder_class = LlamaBotRails.config.state_builder_class || 'LlamaBotRails::AgentStateBuilder'
    end

    @connection_id = SecureRandom.uuid
    Rails.logger.info "[LlamaBot] Created new connection with ID: #{@connection_id}"
    Rails.logger.info "[LlamaBot] Secure API token generated."

    # Use a begin/rescue block to catch thread creation errors
  begin
    @worker = Thread.new do
        Thread.current[:connection_id] = @connection_id
      Thread.current.abort_on_exception = true  # This will help surface errors
      setup_external_websocket(@connection_id)
    end
  rescue => e
    Rails.logger.error "[LlamaBot] Error in WebSocket subscription: #{e.message}"
    Rails.logger.error e.backtrace.join("\n")
    
    # Send error message to frontend before rejecting
    begin
      send_message_to_frontend("error", "Failed to establish chat connection: #{e.message}")
    rescue => send_error
      Rails.logger.error "[LlamaBot] Could not send error to frontend: #{send_error.message}"
    end
    
    reject # Reject the connection if there's an error
    end
  rescue ThreadError => e
    Rails.logger.error "[LlamaBot] Failed to allocate thread: #{e.message}"
    # Handle the error gracefully - potentially notify the client
    send_message_to_frontend("error", "Failed to establish connection: #{e.message}")
  end
end

#unsubscribedObject



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
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
# File 'app/channels/llama_bot_rails/chat_channel.rb', line 78

def unsubscribed
  connection_id = @connection_id
  Rails.logger.info "[LlamaBot] Unsubscribing connection: #{connection_id}"
  
  begin
    # Only kill the worker if it belongs to this connection
    if @worker && @worker[:connection_id] == connection_id
      begin
        @worker.kill
        @worker = nil
        Rails.logger.info "[LlamaBot] Killed worker thread for connection: #{connection_id}"
      rescue => e
        Rails.logger.error "[LlamaBot] Error killing worker thread: #{e.message}"
      end
    end

    # Close the external WebSocket connection BEFORE stopping async tasks
    if @external_ws_connection
      begin
        @external_ws_connection.close
        Rails.logger.info "👋 [LlamaBot] Gracefully closed external WebSocket connection for: #{connection_id}"
      rescue => e
        Rails.logger.warn "❌ [LlamaBot] Could not close WebSocket connection: #{e.message}"
      end
    end

    # Clean up async tasks with better error handling
    begin
      @listener_task&.stop rescue nil
      @keepalive_task&.stop rescue nil
      @external_ws_task&.stop rescue nil
    rescue => e
      Rails.logger.error "[LlamaBot] Error stopping async tasks: #{e.message}"
    end
    
    # Force garbage collection in development/test environments to help clean up
    if !Rails.env.production?
      GC.start
    end
  rescue => e
    Rails.logger.error "[LlamaBot] Fatal error during channel unsubscription: #{e.message}"
    Rails.logger.error e.backtrace.join("\n")
  end
end