Module: RubyLLM::Protocol::Streaming

Included in:
RubyLLM::Protocol
Defined in:
lib/ruby_llm/protocol/streaming.rb

Overview

:nodoc: all

Defined Under Namespace

Modules: FaradayHandlers Classes: StreamState

Class Method Summary collapse

Class Method Details

.build_on_data_handler(progress = {}, &handler) ⇒ Object

Parser state lives on the env so every retry attempt starts fresh; ErrorMiddleware clears it per attempt. Faraday v1 passes no env to on_data, so it falls back to one state for the whole request.



66
67
68
69
70
71
72
73
74
75
76
77
78
 
# File 'lib/ruby_llm/protocol/streaming.rb', line 66

def build_on_data_handler(progress = {}, &handler)
  fallback_state = StreamState.new

  FaradayHandlers.build(
    faraday_v1: faraday_1?,
    on_chunk: lambda { |chunk, env|
      process_stream_chunk(chunk, stream_state(env, fallback_state), env, progress, &handler)
    },
    on_failed_response: lambda { |chunk, env|
      handle_failed_response(chunk, stream_state(env, fallback_state).buffer, env)
    }
  )
end

.build_stream_error_response(parsed_data, env, status) ⇒ Object 



181
182
183
184
185
186
187
188
189
 
# File 'lib/ruby_llm/protocol/streaming.rb', line 181

def build_stream_error_response(parsed_data, env, status)
  error_status = failed_http_status(env) || status || 500

  if faraday_1? || env.nil?
    Struct.new(:body, :status).new(parsed_data, error_status)
  else
    env.merge(body: parsed_data, status: error_status)
  end
end

.error_chunk?(chunk) ⇒ Boolean

An error event split across network reads reaches here in pieces, so only a whole one takes this shortcut; the rest go to the parser, which buffers them until the event is complete.

Returns:

  • (Boolean) 


101
102
103
 
# File 'lib/ruby_llm/protocol/streaming.rb', line 101

def error_chunk?(chunk)
  chunk.start_with?('event: error') && chunk.end_with?("\n\n")
end

.failed_http_status(env) ⇒ Object 



191
192
193
194
195
196
 
# File 'lib/ruby_llm/protocol/streaming.rb', line 191

def failed_http_status(env)
  status = env&.status
  return status if status&.>= 400

  nil
end

.faraday_1? ⇒ Boolean

Returns:

  • (Boolean) 


59
60
61
 
# File 'lib/ruby_llm/protocol/streaming.rb', line 59

def faraday_1?
  Faraday::VERSION.start_with?('1')
end

.handle_data(data, env) ⇒ Object 



145
146
147
148
149
150
151
152
 
# File 'lib/ruby_llm/protocol/streaming.rb', line 145

def handle_data(data, env)
  parsed = JSON.parse(data)
  return parsed unless parsed.is_a?(Hash) && parsed.key?('error')

  raise_stream_error(data, parsed, env)
rescue JSON::ParserError => e
  RubyLLM.logger.debug { "Failed to parse data chunk: #{e.message}" }
end

.handle_error_chunk(chunk, env) ⇒ Object 



113
114
115
116
 
# File 'lib/ruby_llm/protocol/streaming.rb', line 113

def handle_error_chunk(chunk, env)
  error_data = chunk.split("\n")[1]&.delete_prefix('data: ')
  parse_error_from_json(error_data, env, 'Failed to parse error chunk') if error_data
end

.handle_error_event(data, env) ⇒ Object 



154
155
156
 
# File 'lib/ruby_llm/protocol/streaming.rb', line 154

def handle_error_event(data, env)
  parse_error_from_json(data, env, 'Failed to parse error event')
end

.handle_failed_response(chunk, buffer, env) ⇒ Object 



118
119
120
121
122
123
124
 
# File 'lib/ruby_llm/protocol/streaming.rb', line 118

def handle_failed_response(chunk, buffer, env)
  buffer << chunk
  error_data = JSON.parse(buffer)
  raise_stream_error(buffer, error_data, env)
rescue JSON::ParserError
  RubyLLM.logger.debug { "Accumulating error chunk: #{chunk}" }
end

.handle_json_error_chunk(chunk, env) ⇒ Object 



109
110
111
 
# File 'lib/ruby_llm/protocol/streaming.rb', line 109

def handle_json_error_chunk(chunk, env)
  parse_error_from_json(chunk, env, 'Failed to parse JSON error chunk')
end

.handle_sse(chunk, parser, env, progress) ⇒ Object

Marking progress just before a chunk reaches the user lets the retry middleware know this stream already delivered, so a mid-stream failure must not be retried into the same accumulator and user block. An error event raises out of handle_data first, leaving the stream retriable.



130
131
132
133
134
135
136
137
138
139
140
141
142
143
 
# File 'lib/ruby_llm/protocol/streaming.rb', line 130

def handle_sse(chunk, parser, env, progress)
  parser.feed(chunk) do |type, data|
    case type.to_sym
    when :error
      handle_error_event(data, env)
    else
      next if data == '[DONE]'

      parsed = handle_data(data, env)
      progress[:started] = true
      yield parsed
    end
  end
end

.handle_stream(&block) ⇒ Object 



51
52
53
54
55
 
# File 'lib/ruby_llm/protocol/streaming.rb', line 51

def handle_stream(&block)
  build_on_data_handler do |data|
    block.call(build_chunk(data)) if data.is_a?(Hash)
  end
end

.json_error_payload?(chunk) ⇒ Boolean

Returns:

  • (Boolean) 


105
106
107
 
# File 'lib/ruby_llm/protocol/streaming.rb', line 105

def json_error_payload?(chunk)
  chunk.lstrip.start_with?('{') && chunk.include?('"error"')
end

.parse_error_from_json(data, env, error_message) ⇒ Object 



174
175
176
177
178
179
 
# File 'lib/ruby_llm/protocol/streaming.rb', line 174

def parse_error_from_json(data, env, error_message)
  parsed_data = JSON.parse(data)
  raise_stream_error(data, parsed_data, env)
rescue JSON::ParserError => e
  RubyLLM.logger.debug { "#{error_message}: #{e.message}" }
end

.parse_streaming_error(data) ⇒ Object 



158
159
160
161
162
163
164
165
 
# File 'lib/ruby_llm/protocol/streaming.rb', line 158

def parse_streaming_error(data)
  error_data = JSON.parse(data)
  message = error_data.is_a?(Hash) ? error_data['message'] : error_data.to_s
  [500, message || 'Unknown streaming error']
rescue JSON::ParserError => e
  RubyLLM.logger.debug { "Failed to parse streaming error: #{e.message}" }
  [500, "Failed to parse error: #{data}"]
end

.process_stream_chunk(chunk, state, env, progress) ⇒ Object 



86
87
88
89
90
91
92
93
94
95
96
 
# File 'lib/ruby_llm/protocol/streaming.rb', line 86

def process_stream_chunk(chunk, state, env, progress, &)
  RubyLLM.logger.debug { "Received chunk: #{chunk}" } if RubyLLM.config.log_stream_debug

  if error_chunk?(chunk)
    handle_error_chunk(chunk, env)
  elsif json_error_payload?(chunk)
    handle_json_error_chunk(chunk, env)
  else
    handle_sse(chunk, state.parser, env, progress, &)
  end
end

.raise_stream_error(raw_data, parsed_data, env) ⇒ Object 



167
168
169
170
171
172
 
# File 'lib/ruby_llm/protocol/streaming.rb', line 167

def raise_stream_error(raw_data, parsed_data, env)
  status, _message = parse_streaming_error(raw_data)
  error_response = build_stream_error_response(parsed_data, env, status)
  env[:streaming_error_response] = error_response if env.respond_to?(:[]=)
  Transport::ErrorMiddleware.parse_error(provider: self, response: error_response)
end

.stream_events(url, payload, additional_headers = {}, &block) ⇒ Object

Posts payload to url as a server-sent event stream, yielding each parsed event Hash. Returns the Faraday response.



34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
 
# File 'lib/ruby_llm/protocol/streaming.rb', line 34

def stream_events(url, payload, additional_headers = {}, &block)
  progress = {}
  on_data = build_on_data_handler(progress) do |data|
    block.call(data) if data.is_a?(Hash)
  end

  @connection.post url, payload, usage: @usage_tracker do |req|
    req.headers = additional_headers.merge(req.headers) unless additional_headers.empty?
    (req.options.context ||= {})[Transport::Connection::STREAM_PROGRESS_KEY] = progress
    if faraday_1?
      req.options[:on_data] = on_data
    else
      req.options.on_data = on_data
    end
  end
end

.stream_response(payload, additional_headers = {}, &block) ⇒ Object 



18
19
20
21
22
23
24
25
26
27
28
29
30
 
# File 'lib/ruby_llm/protocol/streaming.rb', line 18

def stream_response(payload, additional_headers = {}, &block)
  accumulator = StreamAccumulator.new

  response = stream_events(stream_url, payload, additional_headers) do |data|
    chunk = build_chunk(data)
    accumulator.add chunk
    block.call chunk
  end

  message = accumulator.to_message(response)
  RubyLLM.logger.debug { "Stream completed: #{message.content}" }
  message
end

.stream_state(env, fallback_state) ⇒ Object 



80
81
82
83
84
 
# File 'lib/ruby_llm/protocol/streaming.rb', line 80

def stream_state(env, fallback_state)
  return fallback_state unless env.respond_to?(:[]=)

  env[:streaming_state] ||= StreamState.new
end