Class: RubyLLM::Transport::WebsocketConnection

Inherits:
Object
  • Object
show all
Defined in:
lib/ruby_llm/transport/websocket_connection.rb

Overview

:nodoc: all

Constant Summary collapse

MAX_FRAME_BYTES = 
16 * 1024 * 1024

Instance Attribute Summary collapse

Class Method Summary collapse

Instance Method Summary collapse

Constructor Details

#initialize(url, headers:, config:) ⇒ WebsocketConnection

Returns a new instance of WebsocketConnection.

Raises:

  • (ArgumentError) 


22
23
24
25
26
27
28
29
30
31
32
33
34
35
 
# File 'lib/ruby_llm/transport/websocket_connection.rb', line 22

def initialize(url, headers:, config:)
  @url = url
  @uri = URI.parse(url)
  raise ArgumentError, 'WebSocket URL must use ws or wss' unless %w[ws wss].include?(@uri.scheme)
  if @uri.userinfo || @uri.fragment
    raise ArgumentError, 'WebSocket URL must not contain credentials or a fragment'
  end
  raise ArgumentError, 'WebSocket connections do not support HTTP proxies' if config.http_proxy

  @headers = headers
  @timeout = config.request_timeout
  @messages = []
  @write_lock = Mutex.new
end

Instance Attribute Details

#url ⇒ Object (readonly)

Returns the value of attribute url.



12
13
14
 
# File 'lib/ruby_llm/transport/websocket_connection.rb', line 12

def url
  @url
end

Class Method Details

.open(url, headers:, config:) ⇒ Object 



14
15
16
17
18
19
20
 
# File 'lib/ruby_llm/transport/websocket_connection.rb', line 14

def self.open(url, headers:, config:)
  connection = new(url, headers:, config:)
  connection.connect
  yield connection
ensure
  connection&.close
end

Instance Method Details

#close ⇒ Object 



110
111
112
113
114
115
116
117
 
# File 'lib/ruby_llm/transport/websocket_connection.rb', line 110

def close
  @driver.close if @driver && @opened && !@closed
rescue IOError, SystemCallError, OpenSSL::SSL::SSLError, Error
  nil
ensure
  @closed = true
  close_socket
end

#connect ⇒ Object 



37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
 
# File 'lib/ruby_llm/transport/websocket_connection.rb', line 37

def connect
  load_driver
  deadline = monotonic_time + @timeout
  open_socket(deadline)
  configure_driver
  @driver.start
  read_frame(deadline) until @opened || @closed || @error
  raise @error if @error
  raise Error, 'WebSocket closed before opening' unless @opened

  self
rescue StandardError
  close
  raise
end

#each_message(write:) ⇒ Object 



88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
 
# File 'lib/ruby_llm/transport/websocket_connection.rb', line 88

def each_message(write:)
  writer_error = nil
  writer = Thread.new do
    write.call(self)
  rescue StandardError => e
    writer_error = e
    close
  end
  writer.report_on_exception = false

  while (message = read)
    yield message
  end
  finish_writer(writer)
  raise writer_error if writer_error
rescue StandardError => e
  raise writer_error || e
ensure
  close
  writer.kill.join if writer&.alive?
end

#read(timeout: @timeout) ⇒ Object 



63
64
65
66
67
68
69
70
71
 
# File 'lib/ruby_llm/transport/websocket_connection.rb', line 63

def read(timeout: @timeout)
  deadline = monotonic_time + timeout
  read_frame(deadline) while @messages.empty? && !@closed && !@error
  raise @error if @error

  @messages.shift
rescue IOError, SystemCallError, OpenSSL::SSL::SSLError
  raise unless @closed
end

#read_available ⇒ Object 



73
74
75
76
77
78
79
80
81
82
83
84
85
86
 
# File 'lib/ruby_llm/transport/websocket_connection.rb', line 73

def read_available
  ensure_open
  until @messages.any? || @closed || @error
    break unless read_available_frame
  end
  ensure_open

  @messages.shift
rescue EOFError, Errno::ECONNRESET
  @closed = true
  raise Error, 'WebSocket connection ended without a close frame'
rescue IOError, SystemCallError, OpenSSL::SSL::SSLError
  raise unless @closed
end

#send_binary(data) ⇒ Object 



58
59
60
61
 
# File 'lib/ruby_llm/transport/websocket_connection.rb', line 58

def send_binary(data)
  ensure_open
  @driver.binary(data.b)
end

#send_text(text) ⇒ Object 



53
54
55
56
 
# File 'lib/ruby_llm/transport/websocket_connection.rb', line 53

def send_text(text)
  ensure_open
  @driver.text(text)
end

#write(data) ⇒ Object 



119
120
121
122
123
124
125
126
127
128
129
 
# File 'lib/ruby_llm/transport/websocket_connection.rb', line 119

def write(data)
  @write_lock.synchronize do
    deadline = monotonic_time + @timeout
    offset = 0
    while offset < data.bytesize
      written = socket_operation(deadline) { @socket.write_nonblock(data.byteslice(offset..)) }
      offset += written
    end
  end
  data.bytesize
end