ruby_llm_community 0.0.1 → 0.0.2

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

@@ -0,0 +1,63 @@
+# frozen_string_literal: true
+
+module RubyLLM
+  module Providers
+    class 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 stream_response(connection, payload, additional_headers = {}, &block)
+            signature = sign_request("#{connection.connection.url_prefix}#{stream_url}", payload:)
+            accumulator = StreamAccumulator.new
+
+            response = connection.post stream_url, payload do |req|
+              req.headers.merge! build_headers(signature.headers, streaming: block_given?)
+              # Merge additional headers, with existing headers taking precedence
+              req.headers = additional_headers.merge(req.headers) unless additional_headers.empty?
+              req.options.on_data = handle_stream do |chunk|
+                accumulator.add chunk
+                block.call chunk
+              end
+            end
+
+            accumulator.to_message(response)
+          end
+
+          def handle_stream(&block)
+            buffer = +''
+            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
+    class 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
+    class 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
+    class 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
+    class 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)
+            valid_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)
+            valid_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 valid_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
+    class 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.
+    class Bedrock < Provider
+      include Bedrock::Chat
+      include Bedrock::Streaming
+      include Bedrock::Models
+      include Bedrock::Signing
+      include Bedrock::Media
+      include Anthropic::Tools
+
+      def api_base
+        "https://bedrock-runtime.#{@config.bedrock_region}.amazonaws.com"
+      end
+
+      def parse_error(response)
+        return if response.body.empty?
+
+        body = try_parse_json(response.body)
+        case body
+        when Hash
+          body['message']
+        when Array
+          body.map do |part|
+            part['message']
+          end.join('. ')
+        else
+          body
+        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: @config.bedrock_api_key,
+                              secret_access_key: @config.bedrock_secret_key,
+                              session_token: @config.bedrock_session_token,
+                              region: @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
+
+      class << self
+        def capabilities
+          Bedrock::Capabilities
+        end
+
+        def configuration_requirements
+          %i[bedrock_api_key bedrock_secret_key bedrock_region]
+        end
+      end
+    end
+  end
+end

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

@@ -0,0 +1,131 @@
+# frozen_string_literal: true
+
+module RubyLLM
+  module Providers
+    class DeepSeek
+      # Determines capabilities and pricing for DeepSeek models
+      module Capabilities
+        module_function
+
+        def context_window_for(model_id)
+          case model_id
+          when /deepseek-(?:chat|reasoner)/ then 64_000
+          else 32_768 # Sensible default
+          end
+        end
+
+        def max_tokens_for(model_id)
+          case model_id
+          when /deepseek-(?:chat|reasoner)/ then 8_192
+          else 4_096
+          end
+        end
+
+        def input_price_for(model_id)
+          PRICES.dig(model_family(model_id), :input_miss) || default_input_price
+        end
+
+        def output_price_for(model_id)
+          PRICES.dig(model_family(model_id), :output) || default_output_price
+        end
+
+        def cache_hit_price_for(model_id)
+          PRICES.dig(model_family(model_id), :input_hit) || default_cache_hit_price
+        end
+
+        def supports_vision?(_model_id)
+          false
+        end
+
+        def supports_functions?(model_id)
+          model_id.match?(/deepseek-chat/)
+        end
+
+        def supports_json_mode?(_model_id)
+          false
+        end
+
+        def format_display_name(model_id)
+          case model_id
+          when 'deepseek-chat' then 'DeepSeek V3'
+          when 'deepseek-reasoner' then 'DeepSeek R1'
+          else
+            model_id.split('-')
+                    .map(&:capitalize)
+                    .join(' ')
+          end
+        end
+
+        def model_type(_model_id)
+          'chat'
+        end
+
+        def model_family(model_id)
+          case model_id
+          when /deepseek-reasoner/ then :reasoner
+          else :chat
+          end
+        end
+
+        # Pricing information for DeepSeek models (USD per 1M tokens)
+        PRICES = {
+          chat: {
+            input_hit: 0.07,
+            input_miss: 0.27,
+            output: 1.10
+          },
+          reasoner: {
+            input_hit: 0.14,
+            input_miss: 0.55,
+            output: 2.19
+          }
+        }.freeze
+
+        def default_input_price
+          0.27
+        end
+
+        def default_output_price
+          1.10
+        end
+
+        def default_cache_hit_price
+          0.07
+        end
+
+        def modalities_for(_model_id)
+          {
+            input: ['text'],
+            output: ['text']
+          }
+        end
+
+        def capabilities_for(model_id)
+          capabilities = ['streaming']
+
+          capabilities << 'function_calling' if model_id.match?(/deepseek-chat/)
+
+          capabilities
+        end
+
+        def pricing_for(model_id)
+          family = model_family(model_id)
+          prices = PRICES.fetch(family, { input_miss: default_input_price, output: default_output_price })
+
+          standard_pricing = {
+            input_per_million: prices[:input_miss],
+            output_per_million: prices[:output]
+          }
+
+          standard_pricing[:cached_input_per_million] = prices[:input_hit] if prices[:input_hit]
+
+          {
+            text_tokens: {
+              standard: standard_pricing
+            }
+          }
+        end
+      end
+    end
+  end
+end

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

@@ -0,0 +1,17 @@
+# frozen_string_literal: true
+
+module RubyLLM
+  module Providers
+    class 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

@@ -0,0 +1,30 @@
+# frozen_string_literal: true
+
+module RubyLLM
+  module Providers
+    # DeepSeek API integration.
+    class DeepSeek < OpenAIBase
+      include DeepSeek::Chat
+
+      def api_base
+        'https://api.deepseek.com'
+      end
+
+      def headers
+        {
+          'Authorization' => "Bearer #{@config.deepseek_api_key}"
+        }
+      end
+
+      class << self
+        def capabilities
+          DeepSeek::Capabilities
+        end
+
+        def configuration_requirements
+          %i[deepseek_api_key]
+        end
+      end
+    end
+  end
+end