ruby-pi 0.1.8 → 0.1.10

data/lib/ruby_pi/llm/openai.rb

@@ -201,11 +201,10 @@ module RubyPi
                               begin
                                 JSON.parse(tc_args)
                                 tc_args
-                              rescue JSON::ParserError => e
+                              rescue JSON::ParserError
                                 raise RubyPi::ProviderError.new(
                                   "Invalid JSON in assistant tool_call.arguments " \
-                                  "for tool '#{tc_name || "unknown"}': #{e.message} " \
-                                  "(raw: #{tc_args.inspect})",
+                                  "for tool '#{tc_name || "unknown"}'",
                                   provider: :openai
                                 )
                               end
@@ -291,7 +290,7 @@ module RubyPi
         end
 
         handle_error_response(response) unless response.success?
-        parse_response(JSON.parse(response.body))
+        parse_response(parse_json_response(response.body))
       end
 
       # Executes a streaming request to the OpenAI API, yielding events.
@@ -315,20 +314,63 @@ module RubyPi
         # Issue #21: Accumulate usage data from the final SSE chunk
         streaming_usage = {}
 
-        # Buffer for incomplete SSE lines across on_data chunks. Faraday's
-        # on_data callback delivers raw bytes as they arrive from the network,
-        # which may split SSE events mid-line. We accumulate a line buffer and
-        # process complete lines incrementally so that deltas reach the caller
-        # as soon as each SSE event is fully received.
-        # BINARY buffer: chunks arrive as ASCII-8BIT and may end mid-way
-        # through a multi-byte UTF-8 character; appending such a chunk to a
-        # UTF-8 buffer holding non-ASCII text raises
-        # Encoding::CompatibilityError. Complete lines are re-encoded to
-        # UTF-8 (and scrubbed) before parsing.
-        sse_buffer = (+"").force_encoding(Encoding::BINARY)
+        sse_parser = SSEParser.new(provider: provider_name)
+        stream_completed = false
         response_status = nil
         error_body = (+"").force_encoding(Encoding::BINARY)
 
+        process_payload = proc do |data_str|
+          if data_str == "[DONE]"
+            stream_completed = true
+            next
+          end
+
+          data = sse_parser.parse_json(data_str)
+          if data.key?("usage") && data["usage"]
+            usage_info = data["usage"]
+            streaming_usage = {
+              prompt_tokens: usage_info["prompt_tokens"],
+              completion_tokens: usage_info["completion_tokens"],
+              total_tokens: usage_info["total_tokens"]
+            }
+          end
+
+          choice = (data["choices"] || []).first
+          next unless choice
+
+          delta = choice["delta"] || {}
+          finish_reason = choice["finish_reason"] if choice["finish_reason"]
+
+          if delta.key?("content") && delta["content"]
+            text = delta["content"]
+            append_stream_data!(accumulated_text, text, limit: MAX_STREAM_OUTPUT_BYTES, label: "stream output")
+            block.call(StreamEvent.new(type: :text_delta, data: text))
+          end
+
+          (delta["tool_calls"] || []).each do |tc_delta|
+            index = tc_delta["index"] || 0
+            tool_call_accumulators[index] ||= { id: nil, name: +"", arguments: +"" }
+            enforce_stream_tool_call_limit!(tool_call_accumulators.size)
+            acc = tool_call_accumulators[index]
+            acc[:id] = tc_delta["id"] if tc_delta["id"]
+            acc[:name] << tc_delta["function"]["name"] if tc_delta.dig("function", "name")
+
+            if tc_delta.dig("function", "arguments")
+              append_stream_data!(
+                acc[:arguments], tc_delta["function"]["arguments"],
+                limit: MAX_STREAM_TOOL_ARGUMENT_BYTES, label: "tool arguments"
+              )
+            end
+
+            block.call(StreamEvent.new(type: :tool_call_delta, data: {
+              index: index,
+              id: acc[:id],
+              name: acc[:name].dup.freeze,
+              arguments_fragment: tc_delta.dig("function", "arguments") || ""
+            }))
+          end
+        end
+
         response = with_transport_errors do
           conn.post("/v1/chat/completions") do |req|
             req.headers["Content-Type"] = "application/json"
@@ -344,83 +386,11 @@ module RubyPi
             # If the HTTP status indicates an error, accumulate the body for
             # the error handler instead of parsing it as SSE events.
             if response_status && response_status >= 400
-              error_body << chunk.b
+              append_error_body(error_body, chunk)
               next
             end
 
-            sse_buffer << chunk.b
-            # Process all complete lines in the buffer. A complete line holds
-            # complete UTF-8 sequences (multi-byte characters split across
-            # chunks are repaired by the buffering), so re-encode it to UTF-8
-            # here; scrub guards against a server sending invalid bytes.
-            while (line_end = sse_buffer.index("\n"))
-              line = sse_buffer.slice!(0, line_end + 1).force_encoding(Encoding::UTF_8).scrub.strip
-              next if line.empty?
-              next unless line.start_with?("data: ")
-
-              data_str = line.sub(/\Adata: /, "")
-              next if data_str == "[DONE]"
-
-              begin
-                data = JSON.parse(data_str)
-              rescue JSON::ParserError
-                next
-              end
-
-              # Issue #21: Capture usage data from the final SSE chunk.
-              # OpenAI sends usage in a dedicated chunk when include_usage is true.
-              if data.key?("usage") && data["usage"]
-                usage_info = data["usage"]
-                streaming_usage = {
-                  prompt_tokens: usage_info["prompt_tokens"],
-                  completion_tokens: usage_info["completion_tokens"],
-                  total_tokens: usage_info["total_tokens"]
-                }
-              end
-
-              # Process this SSE event
-              choices = data["choices"] || []
-              choice = choices.first
-              next unless choice
-
-              delta = choice["delta"] || {}
-              finish_reason = choice["finish_reason"] if choice["finish_reason"]
-
-              # Handle text content deltas
-              if delta.key?("content") && delta["content"]
-                text = delta["content"]
-                accumulated_text << text
-                block.call(StreamEvent.new(type: :text_delta, data: text))
-              end
-
-              # Handle tool call deltas
-              if delta.key?("tool_calls")
-                delta["tool_calls"].each do |tc_delta|
-                  index = tc_delta["index"] || 0
-
-                  # Initialize accumulator for this tool call
-                  tool_call_accumulators[index] ||= { id: nil, name: +"", arguments: +"" }
-                  acc = tool_call_accumulators[index]
-
-                  acc[:id] = tc_delta["id"] if tc_delta["id"]
-
-                  if tc_delta.dig("function", "name")
-                    acc[:name] << tc_delta["function"]["name"]
-                  end
-
-                  if tc_delta.dig("function", "arguments")
-                    acc[:arguments] << tc_delta["function"]["arguments"]
-                  end
-
-                  block.call(StreamEvent.new(type: :tool_call_delta, data: {
-                    index: index,
-                    id: acc[:id],
-                    name: acc[:name],
-                    arguments_fragment: tc_delta.dig("function", "arguments") || ""
-                  }))
-                end
-              end
-            end
+            sse_parser.feed(chunk, &process_payload)
           end
           end # conn.post
         end # with_transport_errors
@@ -433,6 +403,15 @@ module RubyPi
           handle_error_response(response, override_body: error_body_str)
         end
 
+        sse_parser.finish(&process_payload)
+
+        unless stream_completed
+          raise RubyPi::StreamingProtocolError.new(
+            "openai stream ended before data: [DONE]",
+            status_code: nil
+          )
+        end
+
         # Build final tool calls from accumulators
         # Issue #12: Guard JSON.parse against empty strings. An empty string
         # is truthy in Ruby, so the previous `empty? ? {} : JSON.parse(...)` check
@@ -543,9 +522,9 @@ module RubyPi
 
         # Attempt to parse the JSON string
         JSON.parse(raw_args)
-      rescue JSON::ParserError => e
+      rescue JSON::ParserError
         raise RubyPi::ProviderError.new(
-          "Failed to parse tool call arguments from OpenAI: #{e.message} (raw: #{raw_args.inspect})",
+          "Failed to parse tool call arguments from OpenAI",
           provider: :openai
         )
       end

data/lib/ruby_pi/llm/sse_parser.rb

@@ -0,0 +1,125 @@
+# frozen_string_literal: true
+
+require "json"
+
+module RubyPi
+  module LLM
+    # Incremental, bounded Server-Sent Events decoder.
+    #
+    # Faraday delivers arbitrary byte chunks, including chunks split inside a
+    # UTF-8 character or line ending. This decoder keeps bytes in BINARY form
+    # until a complete SSE event is available, then validates UTF-8 and yields
+    # the joined value of its data fields.
+    class SSEParser
+      DEFAULT_MAX_EVENT_BYTES = 1_048_576
+
+      def initialize(provider:, max_event_bytes: DEFAULT_MAX_EVENT_BYTES)
+        @provider = provider
+        @max_event_bytes = max_event_bytes
+        @buffer = (+"").force_encoding(Encoding::BINARY)
+        @data_lines = []
+        @event_bytes = 0
+      end
+
+      def feed(chunk, &block)
+        @buffer << chunk.to_s.b
+        extract_lines(finishing: false, &block)
+        enforce_size!(@buffer.bytesize + @event_bytes)
+      end
+
+      def finish(&block)
+        extract_lines(finishing: true, &block)
+
+        unless @buffer.empty?
+          line = @buffer.slice!(0, @buffer.bytesize)
+          process_line(line, &block)
+        end
+
+        dispatch(&block)
+      end
+
+      def parse_json(payload)
+        JSON.parse(payload)
+      rescue JSON::ParserError => e
+        raise RubyPi::StreamingProtocolError.new(
+          "#{@provider} stream contained malformed JSON: #{e.message}",
+          status_code: nil
+        )
+      end
+
+      private
+
+      def extract_lines(finishing:, &block)
+        loop do
+          line_end = next_line_ending
+          break unless line_end
+
+          if @buffer.getbyte(line_end) == 13 && line_end == @buffer.bytesize - 1 && !finishing
+            break
+          end
+
+          terminator_size = if @buffer.getbyte(line_end) == 13 && @buffer.getbyte(line_end + 1) == 10
+                              2
+                            else
+                              1
+                            end
+          line = @buffer.slice!(0, line_end + terminator_size)
+          line = line.byteslice(0, line.bytesize - terminator_size)
+          process_line(line, &block)
+        end
+      end
+
+      def next_line_ending
+        cr = @buffer.index("\r")
+        lf = @buffer.index("\n")
+        [cr, lf].compact.min
+      end
+
+      def process_line(raw_line, &block)
+        line = raw_line.dup.force_encoding(Encoding::UTF_8)
+        unless line.valid_encoding?
+          raise RubyPi::StreamingProtocolError.new(
+            "#{@provider} stream contained invalid UTF-8",
+            status_code: nil
+          )
+        end
+
+        if line.empty?
+          dispatch(&block)
+          return
+        end
+
+        return if line.start_with?(":")
+
+        field, separator, value = line.partition(":")
+        return unless field == "data"
+
+        value = "" unless separator == ":"
+        value = value.byteslice(1..) if value.start_with?(" ")
+        separator_bytes = @data_lines.empty? ? 0 : 1
+        prospective_bytes = @event_bytes + separator_bytes + value.bytesize
+        enforce_size!(prospective_bytes)
+        @data_lines << value
+        @event_bytes = prospective_bytes
+      end
+
+      def dispatch
+        return if @data_lines.empty?
+
+        payload = @data_lines.join("\n")
+        @data_lines.clear
+        @event_bytes = 0
+        yield payload
+      end
+
+      def enforce_size!(bytes)
+        return if bytes <= @max_event_bytes
+
+        raise RubyPi::StreamingProtocolError.new(
+          "#{@provider} stream event exceeded #{@max_event_bytes} bytes",
+          status_code: nil
+        )
+      end
+    end
+  end
+end

data/lib/ruby_pi/llm/stream_event.rb

@@ -25,10 +25,10 @@ module RubyPi
     #   end
     class StreamEvent
       # Valid event types for stream events.
-      VALID_TYPES = %i[text_delta tool_call_delta done fallback_start].freeze
+      VALID_TYPES = %i[text_delta tool_call_delta done retry_start fallback_start].freeze
 
       # @return [Symbol] the type of stream event — one of :text_delta,
-      #   :tool_call_delta, :done, or :fallback_start
+      #   :tool_call_delta, :done, :retry_start, or :fallback_start
       attr_reader :type
 
       # @return [Object] the event payload. For :text_delta this is a String
@@ -38,7 +38,8 @@ module RubyPi
 
       # Creates a new StreamEvent instance.
       #
-      # @param type [Symbol] event type (:text_delta, :tool_call_delta, :done, :fallback_start)
+      # @param type [Symbol] event type (:text_delta, :tool_call_delta, :done,
+      #   :retry_start, :fallback_start)
       # @param data [Object] event payload
       # @raise [ArgumentError] if the type is not recognized
       def initialize(type:, data: nil)
@@ -71,6 +72,12 @@ module RubyPi
         @type == :done
       end
 
+      # Returns true when a failed streaming attempt is being discarded and
+      # the same provider is about to retry from the beginning.
+      def retry_start?
+        @type == :retry_start
+      end
+
       # Returns true if this is a fallback_start event, signaling that the
       # primary provider failed mid-stream and the fallback provider is
       # taking over. Consumers should clear any partial output rendered