ruby_llm 1.16.0 → 2.0.0.rc2

data/lib/ruby_llm/providers/bedrock/streaming.rb

@@ -1,328 +0,0 @@
-# frozen_string_literal: true
-
-require 'base64'
-require 'faraday'
-require 'json'
-
-module RubyLLM
-  module Providers
-    class Bedrock
-      # Streaming implementation for Bedrock ConverseStream (AWS Event Stream).
-      module Streaming
-        private
-
-        def stream_url
-          "/model/#{@model.id}/converse-stream"
-        end
-
-        def stream_response(connection, payload, additional_headers = {}, &block)
-          accumulator = StreamAccumulator.new
-          decoder = event_stream_decoder
-          request_payload = api_payload(payload)
-          body = JSON.generate(request_payload)
-
-          response = connection.post(stream_url, request_payload) do |req|
-            req.headers.merge!(sign_headers('POST', stream_url, body))
-            req.headers.merge!(additional_headers) unless additional_headers.empty?
-            req.headers['Accept'] = 'application/vnd.amazon.eventstream'
-
-            if Faraday::VERSION.start_with?('1')
-              req.options[:on_data] = proc do |chunk, _size|
-                parse_stream_chunk(decoder, chunk, accumulator, &block)
-              end
-            else
-              req.options.on_data = proc do |chunk, _bytes, env|
-                if env&.status == 200
-                  parse_stream_chunk(decoder, chunk, accumulator, &block)
-                else
-                  handle_failed_stream(chunk, env)
-                end
-              end
-            end
-          end
-
-          message = accumulator.to_message(response)
-          RubyLLM.logger.debug { "Stream completed: #{message.content}" }
-          message
-        end
-
-        def event_stream_decoder
-          require 'aws-eventstream'
-          Aws::EventStream::Decoder.new
-        rescue LoadError
-          raise Error,
-                'The aws-eventstream gem is required for Bedrock streaming. ' \
-                'Please add it to your Gemfile: gem "aws-eventstream"'
-        end
-
-        def handle_failed_stream(chunk, env)
-          data = JSON.parse(chunk)
-          error_response = env.merge(body: data)
-          ErrorMiddleware.parse_error(provider: self, response: error_response)
-        rescue JSON::ParserError
-          RubyLLM.logger.debug { "Failed Bedrock stream error chunk: #{chunk}" }
-        end
-
-        def parse_stream_chunk(decoder, raw_chunk, accumulator)
-          handle_non_eventstream_error_chunk(raw_chunk)
-
-          decode_events(decoder, raw_chunk).each do |event|
-            chunk = build_chunk(event)
-            next unless chunk
-
-            accumulator.add(chunk)
-            yield chunk
-          end
-        end
-
-        def handle_non_eventstream_error_chunk(raw_chunk)
-          text = raw_chunk.to_s
-
-          if text.start_with?('event: error')
-            payload = text.lines.find { |line| line.start_with?('data:') }&.delete_prefix('data:')&.strip
-            raise_streaming_chunk_error(payload) if payload
-            return
-          end
-
-          return unless text.lstrip.start_with?('{') && text.include?('"error"')
-
-          raise_streaming_chunk_error(text)
-        end
-
-        def raise_streaming_chunk_error(payload)
-          parsed = JSON.parse(payload)
-          message = parsed.dig('error', 'message') || parsed['message'] || 'Bedrock streaming error'
-          response = Struct.new(:body, :status).new({ 'message' => message }, 500)
-          ErrorMiddleware.parse_error(provider: self, response: response)
-        rescue JSON::ParserError
-          nil
-        end
-
-        def decode_events(decoder, raw_chunk)
-          events = []
-          message, eof = decoder.decode_chunk(raw_chunk)
-
-          while message
-            event = decode_event_payload(message.payload.read)
-            if event && RubyLLM.config.log_stream_debug
-              RubyLLM.logger.debug do
-                "Bedrock stream event keys: #{event.keys}"
-              end
-            end
-            events << event if event
-            break if eof
-
-            message, eof = decoder.decode_chunk
-          end
-
-          events
-        end
-
-        def decode_event_payload(payload)
-          outer = JSON.parse(payload)
-
-          if outer['bytes'].is_a?(String)
-            JSON.parse(Base64.decode64(outer['bytes']))
-          else
-            outer
-          end
-        rescue JSON::ParserError => e
-          RubyLLM.logger.debug { "Failed to decode Bedrock stream event payload: #{e.message}" }
-          nil
-        end
-
-        def build_chunk(event)
-          raise_stream_error(event) if stream_error_event?(event)
-
-          metadata_usage, usage, message_usage = event_usage(event)
-
-          Chunk.new(
-            role: :assistant,
-            model_id: event['modelId'] || event.dig('message', 'model') || @model&.id,
-            content: extract_content_delta(event),
-            thinking: Thinking.build(
-              text: extract_thinking_delta(event),
-              signature: extract_thinking_signature(event)
-            ),
-            tool_calls: extract_tool_calls(event),
-            input_tokens: extract_input_tokens(metadata_usage, usage, message_usage),
-            output_tokens: extract_output_tokens(metadata_usage, usage),
-            cached_tokens: extract_cached_tokens(metadata_usage, usage),
-            cache_creation_tokens: extract_cache_creation_tokens(metadata_usage, usage),
-            thinking_tokens: extract_reasoning_tokens(metadata_usage, usage)
-          )
-        end
-
-        def event_usage(event)
-          [
-            event.dig('metadata', 'usage') || {},
-            event['usage'] || {},
-            event.dig('message', 'usage') || {}
-          ]
-        end
-
-        def extract_input_tokens(metadata_usage, usage, message_usage)
-          bedrock_usage = metadata_usage['inputTokens'] ? metadata_usage : usage
-          return Bedrock::Chat.input_tokens(bedrock_usage) if bedrock_usage['inputTokens']
-
-          message_usage['input_tokens']
-        end
-
-        def extract_output_tokens(metadata_usage, usage)
-          metadata_usage['outputTokens'] || usage['outputTokens'] || usage['output_tokens']
-        end
-
-        def extract_cached_tokens(metadata_usage, usage)
-          metadata_usage['cacheReadInputTokens'] || usage['cacheReadInputTokens'] || usage['cache_read_input_tokens']
-        end
-
-        def extract_cache_creation_tokens(metadata_usage, usage)
-          metadata_usage['cacheWriteInputTokens'] || usage['cacheWriteInputTokens'] ||
-            usage['cache_creation_input_tokens']
-        end
-
-        def extract_reasoning_tokens(metadata_usage, usage)
-          metadata_usage['reasoningTokens'] || usage['reasoningTokens'] ||
-            usage.dig('output_tokens_details', 'thinking_tokens')
-        end
-
-        def stream_error_event?(event)
-          event.keys.any? { |key| key.end_with?('Exception') } || event['type'] == 'error'
-        end
-
-        def raise_stream_error(event)
-          if event['type'] == 'error'
-            message = event.dig('error', 'message') || 'Bedrock streaming error'
-            response = Struct.new(:body, :status).new({ 'message' => message }, 500)
-            ErrorMiddleware.parse_error(provider: self, response: response)
-            return
-          end
-
-          key = event.keys.find { |candidate| candidate.end_with?('Exception') }
-          payload = event[key]
-          message = payload['message'] || key
-          status = case key
-                   when 'throttlingException' then 429
-                   when 'validationException' then 400
-                   when 'accessDeniedException', 'unrecognizedClientException' then 401
-                   when 'serviceUnavailableException' then 503
-                   else 500
-                   end
-
-          response = Struct.new(:body, :status).new({ 'message' => message }, status)
-          ErrorMiddleware.parse_error(provider: self, response: response)
-        end
-
-        def extract_content_delta(event)
-          delta = normalized_delta(event)
-          return delta['text'] if delta['text']
-
-          return event.dig('delta', 'text') if event.dig('delta', 'type') == 'text_delta'
-
-          nil
-        end
-
-        def extract_thinking_delta(event)
-          delta = normalized_delta(event)
-          reasoning_content = delta['reasoningContent'] || {}
-
-          reasoning_text = reasoning_content['reasoningText'] || {}
-          return reasoning_text['text'] if reasoning_text['text']
-          return reasoning_content['text'] if reasoning_content['text']
-          return event.dig('delta', 'thinking') if event.dig('delta', 'type') == 'thinking_delta'
-
-          nil
-        end
-
-        def extract_thinking_signature(event)
-          signature = extract_signature_from_delta(event)
-          return signature if signature
-
-          signature = extract_signature_from_start(event)
-          return signature if signature
-
-          nil
-        end
-
-        def extract_signature_from_delta(event)
-          delta = normalized_delta(event)
-          reasoning_content = delta['reasoningContent'] || {}
-          reasoning_text = reasoning_content['reasoningText'] || {}
-          return reasoning_text['signature'] if reasoning_text['signature']
-          return reasoning_content['signature'] if reasoning_content['signature']
-          return event.dig('delta', 'signature') if event.dig('delta', 'type') == 'signature_delta'
-
-          nil
-        end
-
-        def extract_signature_from_start(event)
-          start = event.dig('contentBlockStart', 'start', 'reasoningContent')
-          return nil unless start
-
-          reasoning_text = start['reasoningText'] || {}
-          return reasoning_text['signature'] if reasoning_text['signature']
-          return start['redactedContent'] if start['redactedContent']
-
-          nil
-        end
-
-        def extract_tool_calls(event)
-          return extract_tool_call_start(event) if tool_call_start_event?(event)
-          return extract_tool_call_delta(event) if tool_call_delta_event?(event)
-
-          nil
-        end
-
-        def tool_call_start_event?(event)
-          event['contentBlockStart'] || event['start'] || event.dig('content_block', 'tool_use')
-        end
-
-        def tool_call_delta_event?(event)
-          event['contentBlockDelta'] || event.dig('delta', 'toolUse') || event.dig('delta', 'tool_use') ||
-            event.dig('delta', 'partial_json')
-        end
-
-        def extract_tool_call_start(event)
-          tool_use = event.dig('contentBlockStart', 'start', 'toolUse')
-          tool_use ||= event.dig('start', 'toolUse')
-          tool_use ||= event.dig('content_block', 'tool_use') if event['type'] == 'content_block_start'
-          return nil unless tool_use
-
-          tool_use_id = tool_use['toolUseId'] || tool_use['id']
-          tool_name = tool_use['name']
-          tool_input = tool_use['input'] || {}
-
-          {
-            tool_use_id => ToolCall.new(
-              id: tool_use_id,
-              name: tool_name,
-              arguments: tool_input
-            )
-          }
-        end
-
-        def extract_tool_call_delta(event)
-          input = normalized_delta(event).dig('toolUse', 'input')
-          input ||= normalized_delta(event).dig('tool_use', 'input')
-          input ||= event.dig('delta', 'partial_json') if event.dig('delta', 'type') == 'input_json_delta'
-          return nil unless input
-
-          { nil => ToolCall.new(id: nil, name: nil, arguments: input) }
-        end
-
-        def normalized_delta(event)
-          delta = event.dig('contentBlockDelta', 'delta') || event['delta'] || {}
-          return delta if delta.is_a?(Hash)
-
-          if delta.is_a?(String) && !delta.empty?
-            JSON.parse(delta)
-          else
-            {}
-          end
-        rescue JSON::ParserError
-          {}
-        end
-      end
-    end
-  end
-end