Refactored Streaming functionality globally, refactored Gemini provider, introduced streaming error parsing for Anthropic

Author: crmne

lib/ruby_llm/error.rb

@@ -19,16 +19,21 @@ def initialize(response = nil, message = nil)
     end
   end
 
-  class ModelNotFoundError < StandardError; end
+  # Error classes for non-HTTP errors
+  class ConfigurationError < StandardError; end
   class InvalidRoleError < StandardError; end
+  class ModelNotFoundError < StandardError; end
   class UnsupportedFunctionsError < StandardError; end
-  class ConfigurationError < StandardError; end
-  class UnauthorizedError < Error; end
-  class PaymentRequiredError < Error; end
-  class ServiceUnavailableError < Error; end
+
+  # Error classes for different HTTP status codes
   class BadRequestError < Error; end
+  class ForbiddenError < Error; end
+  class OverloadedError < Error; end
+  class PaymentRequiredError < Error; end
   class RateLimitError < Error; end
   class ServerError < Error; end
+  class ServiceUnavailableError < Error; end
+  class UnauthorizedError < Error; end
 
   # Faraday middleware that maps provider-specific API errors to RubyLLM errors.
   # Uses provider's parse_error method to extract meaningful error messages.
@@ -57,12 +62,17 @@ def parse_error(provider:, response:) # rubocop:disable Metrics/CyclomaticComple
           raise UnauthorizedError.new(response, message || 'Invalid API key - check your credentials')
         when 402
           raise PaymentRequiredError.new(response, message || 'Payment required - please top up your account')
+        when 403
+          raise ForbiddenError.new(response,
+                                   message || 'Forbidden - you do not have permission to access this resource')
         when 429
           raise RateLimitError.new(response, message || 'Rate limit exceeded - please wait a moment')
         when 500
           raise ServerError.new(response, message || 'API server error - please try again')
         when 502..503
           raise ServiceUnavailableError.new(response, message || 'API server unavailable - please try again later')
+        when 529
+          raise OverloadedError.new(response, message || 'Service overloaded - please try again later')
         else
           raise Error.new(response, message || 'An unknown error occurred')
         end

lib/ruby_llm/provider.rb

@@ -7,7 +7,9 @@ module RubyLLM
   module Provider
     # Common functionality for all LLM providers. Implements the core provider
     # interface so specific providers only need to implement a few key methods.
-    module Methods # rubocop:disable Metrics/ModuleLength
+    module Methods
+      extend Streaming
+
       def complete(messages, tools:, temperature:, model:, &block) # rubocop:disable Metrics/MethodLength
         normalized_temperature = if capabilities.respond_to?(:normalize_temperature)
                                    capabilities.normalize_temperature(temperature, model)
@@ -80,19 +82,6 @@ def sync_response(payload)
         parse_completion_response response
       end
 
-      def stream_response(payload, &block)
-        accumulator = StreamAccumulator.new
-
-        post stream_url, payload do |req|
-          req.options.on_data = handle_stream do |chunk|
-            accumulator.add chunk
-            block.call chunk
-          end
-        end
-
-        accumulator.to_message
-      end
-
       def post(url, payload)
         connection.post url, payload do |req|
           req.headers.merge! headers
@@ -141,33 +130,6 @@ def connection # rubocop:disable Metrics/MethodLength,Metrics/AbcSize
           f.use :llm_errors, provider: self
         end
       end
-
-      def to_json_stream(&block) # rubocop:disable Metrics/MethodLength
-        buffer = String.new
-        parser = EventStreamParser::Parser.new
-
-        proc do |chunk, _bytes, env|
-          if env && env.status != 200
-            # Accumulate error chunks
-            buffer << chunk
-            begin
-              error_data = JSON.parse(buffer)
-              error_response = env.merge(body: error_data)
-              ErrorMiddleware.parse_error(provider: self, response: error_response)
-            rescue JSON::ParserError
-              # Keep accumulating if we don't have complete JSON yet
-              RubyLLM.logger.debug "Accumulating error chunk: #{chunk}"
-            end
-          else
-            parser.feed(chunk) do |_type, data|
-              unless data == '[DONE]'
-                parsed_data = JSON.parse(data)
-                block.call(parsed_data)
-              end
-            end
-          end
-        end
-      end
     end
 
     def try_parse_json(maybe_json)
@@ -207,6 +169,7 @@ def parse_data_uri(uri)
     class << self
       def extended(base)
         base.extend(Methods)
+        base.extend(Streaming)
       end
 
       def register(name, provider_module)

lib/ruby_llm/providers/anthropic/streaming.rb

@@ -11,12 +11,6 @@ def stream_url
           completion_url
         end
 
-        def handle_stream(&block)
-          to_json_stream do |data|
-            block.call(build_chunk(data))
-          end
-        end
-
         def build_chunk(data)
           Chunk.new(
             role: :assistant,
@@ -31,6 +25,18 @@ def build_chunk(data)
         def json_delta?(data)
           data['type'] == 'content_block_delta' && data.dig('delta', 'type') == 'input_json_delta'
         end
+
+        def parse_streaming_error(data)
+          error_data = JSON.parse(data)
+          return unless error_data['type'] == 'error'
+
+          case error_data.dig('error', 'type')
+          when 'overloaded_error'
+            [529, error_data['error']['message']]
+          else
+            [500, error_data['error']['message']]
+          end
+        end
       end
     end
   end

lib/ruby_llm/providers/gemini/chat.rb

@@ -4,9 +4,13 @@ module RubyLLM
   module Providers
     module Gemini
       # Chat methods for the Gemini API implementation
-      module Chat # rubocop:disable Metrics/ModuleLength
-        # Must be public for Provider to use
+      module Chat
+        def completion_url
+          "models/#{@model}:generateContent"
+        end
+
         def complete(messages, tools:, temperature:, model:, &block) # rubocop:disable Metrics/MethodLength
+          @model = model
           payload = {
             contents: format_messages(messages),
             generationConfig: {
@@ -20,26 +24,15 @@ def complete(messages, tools:, temperature:, model:, &block) # rubocop:disable M
           @tools = tools
 
           if block_given?
-            stream_completion(model, payload, &block)
+            stream_response payload, &block
           else
-            generate_completion(model, payload)
+            sync_response payload
           end
         end
 
         # Format methods can be private
         private
 
-        def generate_completion(model, payload)
-          url = "models/#{model}:generateContent"
-          response = post(url, payload)
-          result = parse_completion_response(response)
-
-          # If this contains a tool call, log it
-          result.tool_calls.values.first if result.tool_call?
-
-          result
-        end
-
         def format_messages(messages)
           messages.map do |msg|
             {

lib/ruby_llm/providers/gemini/images.rb

@@ -5,18 +5,13 @@ module Providers
     module Gemini
       # Image generation methods for the Gemini API implementation
       module Images
-        def images_url(model:)
-          "models/#{model}:predict"
+        def images_url
+          "models/#{@model}:predict"
         end
 
-        def paint(prompt, model:, size:) # rubocop:disable Lint/UnusedMethodArgument
-          payload = render_image_payload(prompt)
-
-          response = post(images_url(model:), payload)
-          parse_image_response(response)
-        end
-
-        def render_image_payload(prompt)
+        def render_image_payload(prompt, model:, size:) # rubocop:disable Metrics/MethodLength
+          RubyLLM.logger.debug "Ignoring size #{size}. Gemini does not support image size customization."
+          @model = model
           {
             instances: [
               {

lib/ruby_llm/providers/gemini/streaming.rb

@@ -5,93 +5,52 @@ module Providers
     module Gemini
       # Streaming methods for the Gemini API implementation
       module Streaming
-        # Need to make stream_completion public for chat.rb to access
-        def stream_completion(model, payload, &block) # rubocop:disable Metrics/AbcSize,Metrics/MethodLength
-          url = "models/#{model}:streamGenerateContent?alt=sse"
-          accumulator = StreamAccumulator.new
-
-          post(url, payload) do |req|
-            req.options.on_data = stream_handler(accumulator, &block)
-          end
-
-          # If this is a tool call, immediately execute it and include the result
-          message = accumulator.to_message
-          if message.tool_call? && message.content.to_s.empty? && @tools && !@tools.empty?
-            tool_call = message.tool_calls.values.first
-            tool = @tools[tool_call.name.to_sym]
-
-            if tool
-              tool_result = tool.call(tool_call.arguments)
-              # Create a new chunk with the result
-              result_chunk = Chunk.new(
-                role: :assistant,
-                content: "The result is #{tool_result}",
-                model_id: message.model_id,
-                input_tokens: message.input_tokens,
-                output_tokens: message.output_tokens,
-                tool_calls: message.tool_calls
-              )
-
-              # Add to accumulator and call the block
-              accumulator.add(result_chunk)
-              block.call(result_chunk)
-            end
-          end
+        def stream_url
+          "models/#{@model}:streamGenerateContent?alt=sse"
+        end
 
-          accumulator.to_message
+        def build_chunk(data)
+          Chunk.new(
+            role: :assistant,
+            model_id: extract_model_id(data),
+            content: extract_content(data),
+            input_tokens: extract_input_tokens(data),
+            output_tokens: extract_output_tokens(data),
+            tool_calls: extract_tool_calls(data)
+          )
         end
 
         private
 
-        # Handle streaming
-        def stream_handler(accumulator, &block) # rubocop:disable Metrics/AbcSize,Metrics/CyclomaticComplexity,Metrics/MethodLength,Metrics/PerceivedComplexity
-          to_json_stream do |data| # rubocop:disable Metrics/BlockLength
-            next unless data['candidates']&.any?
-
-            candidate = data['candidates'][0]
-            parts = candidate.dig('content', 'parts')
-            model_id = data['modelVersion']
+        def extract_model_id(data)
+          data['modelVersion']
+        end
 
-            # First attempt to extract tool calls
-            tool_calls = nil
+        def extract_content(data)
+          return nil unless data['candidates']&.any?
 
-            # Check if any part contains a functionCall
-            if parts&.any? { |p| p['functionCall'] }
-              function_part = parts.find { |p| p['functionCall'] }
-              function_data = function_part['functionCall']
+          candidate = data['candidates'][0]
+          parts = candidate.dig('content', 'parts')
+          return nil unless parts
 
-              if function_data && function_data['name']
-                # Create a tool call with proper structure - convert args to JSON string
-                id = SecureRandom.uuid
-                tool_calls = {
-                  id => ToolCall.new(
-                    id: id,
-                    name: function_data['name'],
-                    arguments: JSON.generate(function_data['args']) # Convert Hash to JSON string
-                  )
-                }
-              end
-            end
+          text_parts = parts.select { |p| p['text'] }
+          text_parts.map { |p| p['text'] }.join if text_parts.any?
+        end
 
-            # Extract text content (if any)
-            text = nil
-            if parts
-              text_parts = parts.select { |p| p['text'] }
-              text = text_parts.map { |p| p['text'] }.join if text_parts.any?
-            end
+        def extract_input_tokens(data)
+          data.dig('usageMetadata', 'promptTokenCount')
+        end
 
-            chunk = Chunk.new(
-              role: :assistant,
-              content: text,
-              model_id: model_id,
-              input_tokens: data.dig('usageMetadata', 'promptTokenCount'),
-              output_tokens: data.dig('usageMetadata', 'candidatesTokenCount'),
-              tool_calls: tool_calls
-            )
+        def extract_output_tokens(data)
+          data.dig('usageMetadata', 'candidatesTokenCount')
+        end
 
-            accumulator.add(chunk)
-            block.call(chunk)
-          end
+        def parse_streaming_error(data)
+          error_data = JSON.parse(data)
+          [error_data['error']['code'], error_data['error']['message']]
+        rescue JSON::ParserError => e
+          RubyLLM.logger.debug "Failed to parse streaming error: #{e.message}"
+          [500, "Failed to parse error: #{data}"]
         end
       end
     end

lib/ruby_llm/providers/openai/streaming.rb

@@ -11,19 +11,15 @@ def stream_url
           completion_url
         end
 
-        def handle_stream(&block) # rubocop:disable Metrics/MethodLength
-          to_json_stream do |data|
-            block.call(
-              Chunk.new(
-                role: :assistant,
-                model_id: data['model'],
-                content: data.dig('choices', 0, 'delta', 'content'),
-                tool_calls: parse_tool_calls(data.dig('choices', 0, 'delta', 'tool_calls'), parse_arguments: false),
-                input_tokens: data.dig('usage', 'prompt_tokens'),
-                output_tokens: data.dig('usage', 'completion_tokens')
-              )
-            )
-          end
+        def build_chunk(data)
+          Chunk.new(
+            role: :assistant,
+            model_id: data['model'],
+            content: data.dig('choices', 0, 'delta', 'content'),
+            tool_calls: parse_tool_calls(data.dig('choices', 0, 'delta', 'tool_calls'), parse_arguments: false),
+            input_tokens: data.dig('usage', 'prompt_tokens'),
+            output_tokens: data.dig('usage', 'completion_tokens')
+          )
         end
       end
     end