ruby_llm 1.0.1 → 1.1.0rc1

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

@@ -0,0 +1,46 @@
+# frozen_string_literal: true
+
+module RubyLLM
+  module Providers
+    module Bedrock
+      module Streaming
+        # Base module for AWS Bedrock streaming functionality.
+        # Serves as the core module that includes all other streaming-related modules
+        # and provides fundamental streaming operations.
+        #
+        # Responsibilities:
+        # - Stream URL management
+        # - Stream handling and error processing
+        # - Coordinating the functionality of other streaming modules
+        #
+        # @example
+        #   module MyStreamingImplementation
+        #     include RubyLLM::Providers::Bedrock::Streaming::Base
+        #   end
+        module Base
+          def self.included(base)
+            base.include ContentExtraction
+            base.include MessageProcessing
+            base.include PayloadProcessing
+            base.include PreludeHandling
+          end
+
+          def stream_url
+            "model/#{@model_id}/invoke-with-response-stream"
+          end
+
+          def handle_stream(&block)
+            buffer = String.new
+            proc do |chunk, _bytes, env|
+              if env && env.status != 200
+                handle_failed_response(chunk, buffer, env)
+              else
+                process_chunk(chunk, &block)
+              end
+            end
+          end
+        end
+      end
+    end
+  end
+end

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

@@ -0,0 +1,63 @@
+# frozen_string_literal: true
+
+module RubyLLM
+  module Providers
+    module Bedrock
+      module Streaming
+        # Module for handling content extraction from AWS Bedrock streaming responses.
+        # Provides methods to extract and process various types of content from the response data.
+        #
+        # Responsibilities:
+        # - Extracting content from different response formats
+        # - Processing JSON deltas and content blocks
+        # - Extracting metadata (tokens, model IDs, tool calls)
+        # - Handling different content structures (arrays, blocks, completions)
+        module ContentExtraction
+          def json_delta?(data)
+            data['type'] == 'content_block_delta' && data.dig('delta', 'type') == 'input_json_delta'
+          end
+
+          def extract_streaming_content(data)
+            return '' unless data.is_a?(Hash)
+
+            extract_content_by_type(data)
+          end
+
+          def extract_tool_calls(data)
+            data.dig('message', 'tool_calls') || data['tool_calls']
+          end
+
+          def extract_model_id(data)
+            data.dig('message', 'model') || @model_id
+          end
+
+          def extract_input_tokens(data)
+            data.dig('message', 'usage', 'input_tokens')
+          end
+
+          def extract_output_tokens(data)
+            data.dig('message', 'usage', 'output_tokens') || data.dig('usage', 'output_tokens')
+          end
+
+          private
+
+          def extract_content_by_type(data)
+            case data['type']
+            when 'content_block_start' then extract_block_start_content(data)
+            when 'content_block_delta' then extract_delta_content(data)
+            else ''
+            end
+          end
+
+          def extract_block_start_content(data)
+            data.dig('content_block', 'text').to_s
+          end
+
+          def extract_delta_content(data)
+            data.dig('delta', 'text').to_s
+          end
+        end
+      end
+    end
+  end
+end

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

@@ -0,0 +1,79 @@
+# frozen_string_literal: true
+
+module RubyLLM
+  module Providers
+    module Bedrock
+      module Streaming
+        # Module for processing streaming messages from AWS Bedrock.
+        # Handles the core message processing logic, including validation and chunking.
+        #
+        # Responsibilities:
+        # - Processing incoming message chunks
+        # - Validating message structure and content
+        # - Managing message offsets and boundaries
+        # - Error handling during message processing
+        #
+        # @example Processing a message chunk
+        #   offset = process_message(chunk, current_offset) do |processed_chunk|
+        #     handle_processed_chunk(processed_chunk)
+        #   end
+        module MessageProcessing
+          def process_chunk(chunk, &)
+            offset = 0
+            offset = process_message(chunk, offset, &) while offset < chunk.bytesize
+          rescue StandardError => e
+            RubyLLM.logger.debug "Error processing chunk: #{e.message}"
+            RubyLLM.logger.debug "Chunk size: #{chunk.bytesize}"
+          end
+
+          def process_message(chunk, offset, &)
+            return chunk.bytesize unless can_read_prelude?(chunk, offset)
+
+            message_info = extract_message_info(chunk, offset)
+            return find_next_message(chunk, offset) unless message_info
+
+            process_valid_message(chunk, offset, message_info, &)
+          end
+
+          def process_valid_message(chunk, offset, message_info, &)
+            payload = extract_payload(chunk, message_info[:headers_end], message_info[:payload_end])
+            return find_next_message(chunk, offset) unless valid_payload?(payload)
+
+            process_payload(payload, &)
+            offset + message_info[:total_length]
+          end
+
+          private
+
+          def extract_message_info(chunk, offset)
+            total_length, headers_length = read_prelude(chunk, offset)
+            return unless valid_lengths?(total_length, headers_length)
+
+            message_end = offset + total_length
+            return unless chunk.bytesize >= message_end
+
+            headers_end, payload_end = calculate_positions(offset, total_length, headers_length)
+            return unless valid_positions?(headers_end, payload_end, chunk.bytesize)
+
+            { total_length:, headers_length:, headers_end:, payload_end: }
+          end
+
+          def extract_payload(chunk, headers_end, payload_end)
+            chunk[headers_end...payload_end]
+          end
+
+          def valid_payload?(payload)
+            return false if payload.nil? || payload.empty?
+
+            json_start = payload.index('{')
+            json_end = payload.rindex('}')
+
+            return false if json_start.nil? || json_end.nil? || json_start >= json_end
+
+            true
+          end
+        end
+      end
+    end
+  end
+end

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

@@ -0,0 +1,90 @@
+# frozen_string_literal: true
+
+require 'base64'
+
+module RubyLLM
+  module Providers
+    module Bedrock
+      module Streaming
+        # Module for processing payloads from AWS Bedrock streaming responses.
+        # Handles JSON payload extraction, decoding, and chunk creation.
+        #
+        # Responsibilities:
+        # - Extracting and validating JSON payloads
+        # - Decoding Base64-encoded response data
+        # - Creating response chunks from processed data
+        # - Error handling for JSON parsing and processing
+        #
+        # @example Processing a payload
+        #   process_payload(raw_payload) do |chunk|
+        #     yield_chunk_to_client(chunk)
+        #   end
+        module PayloadProcessing
+          def process_payload(payload, &)
+            json_payload = extract_json_payload(payload)
+            parse_and_process_json(json_payload, &)
+          rescue JSON::ParserError => e
+            log_json_parse_error(e, json_payload)
+          rescue StandardError => e
+            log_general_error(e)
+          end
+
+          private
+
+          def extract_json_payload(payload)
+            json_start = payload.index('{')
+            json_end = payload.rindex('}')
+            payload[json_start..json_end]
+          end
+
+          def parse_and_process_json(json_payload, &)
+            json_data = JSON.parse(json_payload)
+            process_json_data(json_data, &)
+          end
+
+          def process_json_data(json_data, &)
+            return unless json_data['bytes']
+
+            data = decode_and_parse_data(json_data)
+            create_and_yield_chunk(data, &)
+          end
+
+          def decode_and_parse_data(json_data)
+            decoded_bytes = Base64.strict_decode64(json_data['bytes'])
+            JSON.parse(decoded_bytes)
+          end
+
+          def create_and_yield_chunk(data, &block)
+            block.call(build_chunk(data))
+          end
+
+          def build_chunk(data)
+            Chunk.new(
+              **extract_chunk_attributes(data)
+            )
+          end
+
+          def extract_chunk_attributes(data)
+            {
+              role: :assistant,
+              model_id: extract_model_id(data),
+              content: extract_streaming_content(data),
+              input_tokens: extract_input_tokens(data),
+              output_tokens: extract_output_tokens(data),
+              tool_calls: extract_tool_calls(data)
+            }
+          end
+
+          def log_json_parse_error(error, json_payload)
+            RubyLLM.logger.debug "Failed to parse payload as JSON: #{error.message}"
+            RubyLLM.logger.debug "Attempted JSON payload: #{json_payload.inspect}"
+          end
+
+          def log_general_error(error)
+            RubyLLM.logger.debug "Error processing payload: #{error.message}"
+          end
+        end
+      end
+    end
+  end
+end

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

@@ -0,0 +1,91 @@
+# frozen_string_literal: true
+
+module RubyLLM
+  module Providers
+    module Bedrock
+      module Streaming
+        # Module for handling message preludes in AWS Bedrock streaming responses.
+        # Manages the parsing and validation of message headers and prelude data.
+        #
+        # Responsibilities:
+        # - Reading and validating message preludes
+        # - Calculating message positions and boundaries
+        # - Finding and validating prelude positions in chunks
+        # - Ensuring message integrity through length validation
+        #
+        # @example Reading a prelude
+        #   if can_read_prelude?(chunk, offset)
+        #     total_length, headers_length = read_prelude(chunk, offset)
+        #     process_message_with_lengths(total_length, headers_length)
+        #   end
+        module PreludeHandling
+          def can_read_prelude?(chunk, offset)
+            chunk.bytesize - offset >= 12
+          end
+
+          def read_prelude(chunk, offset)
+            total_length = chunk[offset...offset + 4].unpack1('N')
+            headers_length = chunk[offset + 4...offset + 8].unpack1('N')
+            [total_length, headers_length]
+          end
+
+          def valid_lengths?(total_length, headers_length)
+            validate_length_constraints(total_length, headers_length)
+          end
+
+          def calculate_positions(offset, total_length, headers_length)
+            headers_end = offset + 12 + headers_length
+            payload_end = offset + total_length - 4 # Subtract 4 bytes for message CRC
+            [headers_end, payload_end]
+          end
+
+          def valid_positions?(headers_end, payload_end, chunk_size)
+            return false if headers_end >= payload_end
+            return false if headers_end >= chunk_size
+            return false if payload_end > chunk_size
+
+            true
+          end
+
+          def find_next_message(chunk, offset)
+            next_prelude = find_next_prelude(chunk, offset + 4)
+            next_prelude || chunk.bytesize
+          end
+
+          def find_next_prelude(chunk, start_offset)
+            scan_range(chunk, start_offset).each do |pos|
+              return pos if valid_prelude_at_position?(chunk, pos)
+            end
+            nil
+          end
+
+          private
+
+          def scan_range(chunk, start_offset)
+            (start_offset...(chunk.bytesize - 8))
+          end
+
+          def valid_prelude_at_position?(chunk, pos)
+            lengths = extract_potential_lengths(chunk, pos)
+            validate_length_constraints(*lengths)
+          end
+
+          def extract_potential_lengths(chunk, pos)
+            [
+              chunk[pos...pos + 4].unpack1('N'),
+              chunk[pos + 4...pos + 8].unpack1('N')
+            ]
+          end
+
+          def validate_length_constraints(total_length, headers_length)
+            return false if total_length.nil? || headers_length.nil?
+            return false if total_length <= 0 || total_length > 1_000_000
+            return false if headers_length <= 0 || headers_length >= total_length
+
+            true
+          end
+        end
+      end
+    end
+  end
+end

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

@@ -0,0 +1,36 @@
+# frozen_string_literal: true
+
+require_relative 'streaming/base'
+require_relative 'streaming/content_extraction'
+require_relative 'streaming/message_processing'
+require_relative 'streaming/payload_processing'
+require_relative 'streaming/prelude_handling'
+
+module RubyLLM
+  module Providers
+    module Bedrock
+      # Streaming implementation for the AWS Bedrock API.
+      # This module provides functionality for handling streaming responses from AWS Bedrock,
+      # including message processing, content extraction, and error handling.
+      #
+      # The implementation is split into several focused modules:
+      # - Base: Core streaming functionality and module coordination
+      # - ContentExtraction: Extracting content from response data
+      # - MessageProcessing: Processing streaming message chunks
+      # - PayloadProcessing: Handling JSON payloads and chunk creation
+      # - PreludeHandling: Managing message preludes and headers
+      #
+      # @example Using the streaming module
+      #   class BedrockClient
+      #     include RubyLLM::Providers::Bedrock::Streaming
+      #
+      #     def stream_response(&block)
+      #       handle_stream(&block)
+      #     end
+      #   end
+      module Streaming
+        include Base
+      end
+    end
+  end
+end

data/lib/ruby_llm/providers/bedrock.rb

@@ -0,0 +1,83 @@
+# frozen_string_literal: true
+
+require 'openssl'
+require 'time'
+
+module RubyLLM
+  module Providers
+    # AWS Bedrock API integration. Handles chat completion and streaming
+    # for Claude models.
+    module Bedrock
+      extend Provider
+      extend Bedrock::Chat
+      extend Bedrock::Streaming
+      extend Bedrock::Models
+      extend Bedrock::Signing
+
+      # This provider currently only supports Anthropic models, so the tools/media implementation is shared
+      extend Anthropic::Media
+      extend Anthropic::Tools
+
+      module_function
+
+      def api_base
+        @api_base ||= "https://bedrock-runtime.#{RubyLLM.config.bedrock_region}.amazonaws.com"
+      end
+
+      def post(url, payload)
+        signature = sign_request("#{connection.url_prefix}#{url}", payload:)
+        connection.post url, payload do |req|
+          req.headers.merge! build_headers(signature.headers, streaming: block_given?)
+
+          yield req if block_given?
+        end
+      end
+
+      def sign_request(url, method: :post, payload: nil)
+        signer = create_signer
+        request = build_request(url, method:, payload:)
+        signer.sign_request(request)
+      end
+
+      def create_signer
+        Signing::Signer.new({
+                              access_key_id: RubyLLM.config.bedrock_api_key,
+                              secret_access_key: RubyLLM.config.bedrock_secret_key,
+                              session_token: RubyLLM.config.bedrock_session_token,
+                              region: RubyLLM.config.bedrock_region,
+                              service: 'bedrock'
+                            })
+      end
+
+      def build_request(url, method: :post, payload: nil)
+        {
+          connection: connection,
+          http_method: method,
+          url: url || completion_url,
+          body: payload ? JSON.generate(payload, ascii_only: false) : nil
+        }
+      end
+
+      def build_headers(signature_headers, streaming: false)
+        accept_header = streaming ? 'application/vnd.amazon.eventstream' : 'application/json'
+
+        signature_headers.merge(
+          'Content-Type' => 'application/json',
+          'Accept' => accept_header
+        )
+      end
+
+      def capabilities
+        Bedrock::Capabilities
+      end
+
+      def slug
+        'bedrock'
+      end
+
+      def configuration_requirements
+        %i[bedrock_api_key bedrock_secret_key bedrock_region]
+      end
+    end
+  end
+end

data/lib/ruby_llm/providers/deepseek/chat.rb

@@ -0,0 +1,17 @@
+# frozen_string_literal: true
+
+module RubyLLM
+  module Providers
+    module DeepSeek
+      # Chat methods of the DeepSeek API integration
+      module Chat
+        module_function
+
+        def format_role(role)
+          # DeepSeek doesn't use the new OpenAI convention for system prompts
+          role.to_s
+        end
+      end
+    end
+  end
+end

data/lib/ruby_llm/providers/deepseek.rb

@@ -5,6 +5,7 @@ module RubyLLM
     # DeepSeek API integration.
     module DeepSeek
       extend OpenAI
+      extend DeepSeek::Chat
 
       module_function
 
@@ -25,6 +26,10 @@ module RubyLLM
       def slug
         'deepseek'
       end
+
+      def configuration_requirements
+        %i[deepseek_api_key]
+      end
     end
   end
 end