Module: Legion::Extensions::Llm::Streaming

Extended by:
Logging::Helper
Includes:
Logging::Helper
Included in:
Provider
Defined in:
lib/legion/extensions/llm/streaming.rb

Overview

Handles streaming responses from AI providers. The provider's build_chunk yields Canonical::Chunk objects; the caller's block receives Canonical::Chunk objects and the sequence ends in exactly one done chunk (or an error chunk followed by the raise, 05 O5). FaradayHandlers / SSE parsing / error paths are unchanged plumbing; the status→error mapping delegates to the ONE mapper (ErrorMiddleware, 10 U8).

Defined Under Namespace

Modules: FaradayHandlers

Class Method Summary collapse

Class Method Details

.build_on_data_handler ⇒ Object 



76
77
78
79
80
81
82
83
84
85
 
# File 'lib/legion/extensions/llm/streaming.rb', line 76

def build_on_data_handler(&)
  buffer = +''
  parser = EventStreamParser::Parser.new

  FaradayHandlers.build(
    faraday_v1: faraday_1?,
    on_chunk: ->(chunk, env) { process_stream_chunk(chunk, parser, env, &) },
    on_failed_response: ->(chunk, env) { handle_failed_response(chunk, buffer, env) }
  )
end

.build_stream_callback(accumulator, block) ⇒ Object 



47
48
49
50
51
52
53
 
# File 'lib/legion/extensions/llm/streaming.rb', line 47

def build_stream_callback(accumulator, block)
  proc do |chunk|
    next unless chunk

    accumulator.add(chunk).each { |emitted| block&.call(emitted) }
  end
end

.build_stream_error_response(parsed_data, env, status) ⇒ Object 



237
238
239
240
241
242
243
244
245
 
# File 'lib/legion/extensions/llm/streaming.rb', line 237

def build_stream_error_response(parsed_data, env, status)
  error_status = 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

Returns:

  • (Boolean) 


99
100
101
 
# File 'lib/legion/extensions/llm/streaming.rb', line 99

def error_chunk?(chunk)
  chunk.start_with?('event: error')
end

.faraday_1? ⇒ Boolean

Returns:

  • (Boolean) 


72
73
74
 
# File 'lib/legion/extensions/llm/streaming.rb', line 72

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

.handle_data(data, env) ⇒ Object 



182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
 
# File 'lib/legion/extensions/llm/streaming.rb', line 182

def handle_data(data, env)
  # An empty data frame carries nothing — it is not a parse failure.
  return if data.to_s.strip.empty?

  parsed = Legion::JSON.parse(data, symbolize_names: false)
  return parsed unless parsed.is_a?(Hash) && parsed.key?('error')

  handle_parsed_error(parsed, env)
rescue Legion::JSON::ParseError => e
  # M1: an unparseable mid-stream data frame is still a provider
  # error — a classified failure through the one mapper, never a
  # silent drop (the stream must not complete "successfully").
  handle_exception(e, level: :warn, handled: true, operation: 'llm.streaming.handle_data')
  raise_unparseable_streaming_error(env, data, 'data frame')
end

.handle_error_chunk(chunk, env) ⇒ Object 



111
112
113
114
 
# File 'lib/legion/extensions/llm/streaming.rb', line 111

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')
end

.handle_error_event(data, env) ⇒ Object 



198
199
200
 
# File 'lib/legion/extensions/llm/streaming.rb', line 198

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 



116
117
118
119
120
121
122
123
124
125
 
# File 'lib/legion/extensions/llm/streaming.rb', line 116

def handle_failed_response(chunk, buffer, env)
  buffer << chunk
  body_persisted = persist_failed_response_body(buffer, env)
  error_data = Legion::JSON.parse(buffer, symbolize_names: false)
  handle_parsed_error(error_data, env)
rescue Legion::JSON::ParseError
  return if body_persisted

  raise_partial_streaming_error(buffer, env)
end

.handle_json_error_chunk(chunk, env) ⇒ Object 



107
108
109
 
# File 'lib/legion/extensions/llm/streaming.rb', line 107

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

.handle_parsed_error(parsed_data, env) ⇒ Object 



210
211
212
213
214
 
# File 'lib/legion/extensions/llm/streaming.rb', line 210

def handle_parsed_error(parsed_data, env)
  status, _message = parse_streaming_error(parsed_data.to_json)
  error_response = build_stream_error_response(parsed_data, env, status)
  ErrorMiddleware.parse_error(provider: self, response: error_response)
end

.handle_sse(chunk, parser, env) ⇒ Object 



171
172
173
174
175
176
177
178
179
180
 
# File 'lib/legion/extensions/llm/streaming.rb', line 171

def handle_sse(chunk, parser, env, &)
  parser.feed(chunk) do |type, data|
    case type.to_sym
    when :error
      handle_error_event(data, env)
    else
      yield handle_data(data, env, &) unless data == '[DONE]'
    end
  end
end

.handle_stream(&block) ⇒ Object 



55
56
57
58
59
60
61
62
63
64
65
66
67
68
 
# File 'lib/legion/extensions/llm/streaming.rb', line 55

def handle_stream(&block)
  build_on_data_handler do |data|
    next unless data.is_a?(Hash)

    result = build_chunk(data)
    next unless result

    if result.is_a?(Array)
      result.each { |chunk| block.call(chunk) if chunk }
    else
      block.call(result)
    end
  end
end

.json_error_payload?(chunk) ⇒ Boolean

Returns:

  • (Boolean) 


103
104
105
 
# File 'lib/legion/extensions/llm/streaming.rb', line 103

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

.parse_error_from_json(data, env, _error_message) ⇒ Object 



216
217
218
219
220
221
222
223
224
225
226
 
# File 'lib/legion/extensions/llm/streaming.rb', line 216

def parse_error_from_json(data, env, _error_message)
  parsed_data = Legion::JSON.parse(data, symbolize_names: false)
  handle_parsed_error(parsed_data, env)
rescue Legion::JSON::ParseError => e
  # M1: an error event that cannot be parsed is STILL an error event
  # — a classified failure, never a silent nil (the pre-M1 fail-open
  # completed the stream "successfully" on the provider's explicit
  # error signal).
  handle_exception(e, level: :warn, handled: true, operation: 'llm.streaming.parse_error_from_json')
  raise_unparseable_streaming_error(env, data, 'error event')
end

.parse_streaming_error(data) ⇒ Object 



202
203
204
205
206
207
208
 
# File 'lib/legion/extensions/llm/streaming.rb', line 202

def parse_streaming_error(data)
  error_data = Legion::JSON.parse(data, symbolize_names: false)
  [500, error_data['message'] || 'Unknown streaming error']
rescue Legion::JSON::ParseError => e
  handle_exception(e, level: :warn, handled: true, operation: 'llm.streaming.parse_streaming_error')
  [500, "Failed to parse error: #{data}"]
end

.persist_failed_response_body(buffer, env) ⇒ Object 



127
128
129
130
131
 
# File 'lib/legion/extensions/llm/streaming.rb', line 127

def persist_failed_response_body(buffer, env)
  custom_persisted = persist_failed_response_custom_body?(buffer, env)
  body_persisted = persist_failed_response_env_body?(buffer, env)
  custom_persisted || body_persisted
end

.persist_failed_response_custom_body?(buffer, env) ⇒ Boolean

Returns:

  • (Boolean) 


140
141
142
143
144
145
146
147
 
# File 'lib/legion/extensions/llm/streaming.rb', line 140

def persist_failed_response_custom_body?(buffer, env)
  return false unless env.respond_to?(:[]=)

  env[ErrorMiddleware::STREAM_ERROR_BODY_KEY] = buffer.dup
  true
rescue StandardError
  false
end

.persist_failed_response_env_body?(buffer, env) ⇒ Boolean

Returns:

  • (Boolean) 


133
134
135
136
137
138
 
# File 'lib/legion/extensions/llm/streaming.rb', line 133

def persist_failed_response_env_body?(buffer, env)
  return false unless env.respond_to?(:body=)

  env.body = buffer.dup
  true
end

.process_stream_chunk(chunk, parser, env) ⇒ Object 



87
88
89
90
91
92
93
94
95
96
97
 
# File 'lib/legion/extensions/llm/streaming.rb', line 87

def process_stream_chunk(chunk, parser, env, &)
  log.debug { "Received chunk: #{chunk}" } if Legion::Extensions::Llm.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
    yield handle_sse(chunk, parser, env, &)
  end
end

.raise_partial_streaming_error(buffer, env) ⇒ Object 



149
150
151
152
153
154
155
156
157
158
159
160
 
# File 'lib/legion/extensions/llm/streaming.rb', line 149

def raise_partial_streaming_error(buffer, env)
  partial = buffer[/"message"\s*:\s*"([^"]{1,200})/, 1]
  status  = env&.status || 0
  msg     = if partial
              "Provider error (status #{status}): #{partial}"
            else
              "Provider error (status #{status}) - response body incomplete"
            end
  log.warn "[llm][streaming] action=handle_failed_response status=#{status} " \
           "partial_body=#{buffer.length}b msg=#{partial.inspect}"
  raise_streaming_status_error(status, msg)
end

.raise_streaming_status_error(status, message) ⇒ Object

10 U8: the streaming path delegates to the ONE status→error mapper (ErrorMiddleware.parse_error) — the inlined case table is deleted. The provider's parse_error extracts the message from the synthesized body.



166
167
168
169
 
# File 'lib/legion/extensions/llm/streaming.rb', line 166

def raise_streaming_status_error(status, message)
  response = Struct.new(:body, :status).new({ 'error' => { 'message' => message } }, status)
  ErrorMiddleware.parse_error(provider: respond_to?(:parse_error) ? self : nil, response:)
end

.raise_unparseable_streaming_error(_env, data, kind) ⇒ Object

One classified failure for unparseable stream error content — the same 500 classification the PARSEABLE in-band error path uses (parse_streaming_error): an in-band provider error event is a failure regardless of the stream's HTTP status. The raw payload is bounded (200 chars) — never the full body.



233
234
235
 
# File 'lib/legion/extensions/llm/streaming.rb', line 233

def raise_unparseable_streaming_error(_env, data, kind)
  raise_streaming_status_error(500, "Provider error: unparseable #{kind} (#{data.to_s[0, 200].inspect})")
end

.stream_response(connection, payload, additional_headers = {}, model: nil, &block) ⇒ Object 



18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
 
# File 'lib/legion/extensions/llm/streaming.rb', line 18

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

  begin
    connection.post stream_url, payload do |req|
      req.headers = additional_headers.merge(req.headers) unless additional_headers.empty?
      on_chunk = build_stream_callback(accumulator, block)
      log.debug { "Stream callback prepared: #{on_chunk.inspect}" } if Legion::Extensions::Llm.config.log_stream_debug
      if faraday_1?
        req.options[:on_data] = handle_stream(&on_chunk)
      else
        req.options.on_data = handle_stream(&on_chunk)
      end
    end
  rescue StandardError => e
    block&.call(Canonical::Chunk.error_chunk(error: e, request_id: nil))
    raise
  end

  # Release any text held by the untagged-preamble heuristic so short
  # responses still stream at least one delta to the caller.
  accumulator.flush_pending_chunk.each { |chunk| block&.call(chunk) }

  message = accumulator.to_response(model:)
  log.debug { "Stream completed: #{message.text}" }
  block&.call(Canonical::Chunk.done(request_id: nil, usage: message.usage, stop_reason: message.stop_reason))
  message
end