ruby_llm_community 0.0.1 → 0.0.2

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

@@ -0,0 +1,167 @@
+# frozen_string_literal: true
+
+module RubyLLM
+  module Providers
+    class Bedrock
+      # Determines capabilities and pricing for AWS Bedrock models
+      module Capabilities
+        module_function
+
+        def context_window_for(model_id)
+          case model_id
+          when /anthropic\.claude-2/ then 100_000
+          else 200_000
+          end
+        end
+
+        def max_tokens_for(_model_id)
+          4_096
+        end
+
+        def input_price_for(model_id)
+          PRICES.dig(model_family(model_id), :input) || default_input_price
+        end
+
+        def output_price_for(model_id)
+          PRICES.dig(model_family(model_id), :output) || default_output_price
+        end
+
+        def supports_chat?(model_id)
+          model_id.match?(/anthropic\.claude/)
+        end
+
+        def supports_streaming?(model_id)
+          model_id.match?(/anthropic\.claude/)
+        end
+
+        def supports_images?(model_id)
+          model_id.match?(/anthropic\.claude/)
+        end
+
+        def supports_vision?(model_id)
+          model_id.match?(/anthropic\.claude/)
+        end
+
+        def supports_functions?(model_id)
+          model_id.match?(/anthropic\.claude/)
+        end
+
+        def supports_audio?(_model_id)
+          false
+        end
+
+        def supports_json_mode?(model_id)
+          model_id.match?(/anthropic\.claude/)
+        end
+
+        def format_display_name(model_id)
+          model_id.then { |id| humanize(id) }
+        end
+
+        def model_type(_model_id)
+          'chat'
+        end
+
+        def supports_structured_output?(_model_id)
+          false
+        end
+
+        # Model family patterns for capability lookup
+        MODEL_FAMILIES = {
+          /anthropic\.claude-3-opus/ => :claude3_opus,
+          /anthropic\.claude-3-sonnet/ => :claude3_sonnet,
+          /anthropic\.claude-3-5-sonnet/ => :claude3_sonnet,
+          /anthropic\.claude-3-7-sonnet/ => :claude3_sonnet,
+          /anthropic\.claude-3-haiku/ => :claude3_haiku,
+          /anthropic\.claude-3-5-haiku/ => :claude3_5_haiku,
+          /anthropic\.claude-v2/ => :claude2,
+          /anthropic\.claude-instant/ => :claude_instant
+        }.freeze
+
+        def model_family(model_id)
+          MODEL_FAMILIES.find { |pattern, _family| model_id.match?(pattern) }&.last || :other
+        end
+
+        # Pricing information for Bedrock models (per million tokens)
+        PRICES = {
+          claude3_opus: { input: 15.0, output: 75.0 },
+          claude3_sonnet: { input: 3.0, output: 15.0 },
+          claude3_haiku: { input: 0.25, output: 1.25 },
+          claude3_5_haiku: { input: 0.8, output: 4.0 },
+          claude2: { input: 8.0, output: 24.0 },
+          claude_instant: { input: 0.8, output: 2.4 }
+        }.freeze
+
+        def default_input_price
+          0.1
+        end
+
+        def default_output_price
+          0.2
+        end
+
+        def humanize(id)
+          id.tr('-', ' ')
+            .split('.')
+            .last
+            .split
+            .map(&:capitalize)
+            .join(' ')
+        end
+
+        def modalities_for(model_id)
+          modalities = {
+            input: ['text'],
+            output: ['text']
+          }
+
+          if model_id.match?(/anthropic\.claude/) && supports_vision?(model_id)
+            modalities[:input] << 'image'
+            modalities[:input] << 'pdf'
+          end
+
+          modalities
+        end
+
+        def capabilities_for(model_id)
+          capabilities = []
+
+          capabilities << 'streaming' if model_id.match?(/anthropic\.claude/)
+
+          capabilities << 'function_calling' if supports_functions?(model_id)
+
+          capabilities << 'reasoning' if model_id.match?(/claude-3-7/)
+
+          if model_id.match?(/claude-3\.5|claude-3-7/)
+            capabilities << 'batch'
+            capabilities << 'citations'
+          end
+
+          capabilities
+        end
+
+        def pricing_for(model_id)
+          family = model_family(model_id)
+          prices = PRICES.fetch(family, { input: default_input_price, output: default_output_price })
+
+          standard_pricing = {
+            input_per_million: prices[:input],
+            output_per_million: prices[:output]
+          }
+
+          batch_pricing = {
+            input_per_million: prices[:input] * 0.5,
+            output_per_million: prices[:output] * 0.5
+          }
+
+          {
+            text_tokens: {
+              standard: standard_pricing,
+              batch: batch_pricing
+            }
+          }
+        end
+      end
+    end
+  end
+end

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

@@ -0,0 +1,65 @@
+# frozen_string_literal: true
+
+module RubyLLM
+  module Providers
+    class Bedrock
+      # Chat methods for the AWS Bedrock API implementation
+      module Chat
+        module_function
+
+        def sync_response(connection, payload, additional_headers = {})
+          signature = sign_request("#{connection.connection.url_prefix}#{completion_url}", payload:)
+          response = connection.post completion_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?
+          end
+          Anthropic::Chat.parse_completion_response response
+        end
+
+        def format_message(msg)
+          if msg.tool_call?
+            Anthropic::Tools.format_tool_call(msg)
+          elsif msg.tool_result?
+            Anthropic::Tools.format_tool_result(msg)
+          else
+            format_basic_message(msg)
+          end
+        end
+
+        def format_basic_message(msg)
+          {
+            role: Anthropic::Chat.convert_role(msg.role),
+            content: Media.format_content(msg.content)
+          }
+        end
+
+        private
+
+        def completion_url
+          "model/#{@model_id}/invoke"
+        end
+
+        def render_payload(messages, tools:, temperature:, model:, stream: false, schema: nil) # rubocop:disable Lint/UnusedMethodArgument,Metrics/ParameterLists
+          # Hold model_id in instance variable for use in completion_url and stream_url
+          @model_id = model
+
+          system_messages, chat_messages = Anthropic::Chat.separate_messages(messages)
+          system_content = Anthropic::Chat.build_system_content(system_messages)
+
+          build_base_payload(chat_messages, model).tap do |payload|
+            Anthropic::Chat.add_optional_fields(payload, system_content:, tools:, temperature:)
+          end
+        end
+
+        def build_base_payload(chat_messages, model)
+          {
+            anthropic_version: 'bedrock-2023-05-31',
+            messages: chat_messages.map { |msg| format_message(msg) },
+            max_tokens: RubyLLM.models.find(model)&.max_tokens || 4096
+          }
+        end
+      end
+    end
+  end
+end

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

@@ -0,0 +1,61 @@
+# frozen_string_literal: true
+
+module RubyLLM
+  module Providers
+    class Bedrock
+      # Media handling methods for the Bedrock API integration
+      # NOTE: Bedrock does not support url attachments
+      module Media
+        extend Anthropic::Media
+
+        module_function
+
+        def format_content(content)
+          # Convert Hash/Array back to JSON string for API
+          return [Anthropic::Media.format_text(content.to_json)] if content.is_a?(Hash) || content.is_a?(Array)
+          return [Anthropic::Media.format_text(content)] unless content.is_a?(Content)
+
+          parts = []
+          parts << Anthropic::Media.format_text(content.text) if content.text
+
+          content.attachments.each do |attachment|
+            case attachment.type
+            when :image
+              parts << format_image(attachment)
+            when :pdf
+              parts << format_pdf(attachment)
+            when :text
+              parts << Anthropic::Media.format_text_file(attachment)
+            else
+              raise UnsupportedAttachmentError, attachment.type
+            end
+          end
+
+          parts
+        end
+
+        def format_image(image)
+          {
+            type: 'image',
+            source: {
+              type: 'base64',
+              media_type: image.mime_type,
+              data: image.encoded
+            }
+          }
+        end
+
+        def format_pdf(pdf)
+          {
+            type: 'document',
+            source: {
+              type: 'base64',
+              media_type: pdf.mime_type,
+              data: pdf.encoded
+            }
+          }
+        end
+      end
+    end
+  end
+end

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

@@ -0,0 +1,82 @@
+# frozen_string_literal: true
+
+module RubyLLM
+  module Providers
+    class Bedrock
+      # Models methods for the AWS Bedrock API implementation
+      module Models
+        def list_models
+          mgmt_api_base = "https://bedrock.#{@config.bedrock_region}.amazonaws.com"
+          full_models_url = "#{mgmt_api_base}/#{models_url}"
+          signature = sign_request(full_models_url, method: :get)
+          response = @connection.get(full_models_url) do |req|
+            req.headers.merge! signature.headers
+          end
+
+          parse_list_models_response(response, slug, capabilities)
+        end
+
+        module_function
+
+        def models_url
+          'foundation-models'
+        end
+
+        def parse_list_models_response(response, slug, capabilities)
+          models = Array(response.body['modelSummaries'])
+
+          # Filter to include only models we care about
+          models.select { |m| m['modelId'].include?('claude') }.map do |model_data|
+            model_id = model_data['modelId']
+
+            Model::Info.new(
+              id: model_id_with_region(model_id, model_data),
+              name: model_data['modelName'] || capabilities.format_display_name(model_id),
+              provider: slug,
+              family: capabilities.model_family(model_id),
+              created_at: nil,
+              context_window: capabilities.context_window_for(model_id),
+              max_output_tokens: capabilities.max_tokens_for(model_id),
+              modalities: capabilities.modalities_for(model_id),
+              capabilities: capabilities.capabilities_for(model_id),
+              pricing: capabilities.pricing_for(model_id),
+              metadata: {
+                provider_name: model_data['providerName'],
+                inference_types: model_data['inferenceTypesSupported'] || [],
+                streaming_supported: model_data['responseStreamingSupported'] || false,
+                input_modalities: model_data['inputModalities'] || [],
+                output_modalities: model_data['outputModalities'] || []
+              }
+            )
+          end
+        end
+
+        # Simple test-friendly method that only sets the ID
+        def create_model_info(model_data, slug, _capabilities)
+          model_id = model_data['modelId']
+
+          Model::Info.new(
+            id: model_id_with_region(model_id, model_data),
+            name: model_data['modelName'] || model_id,
+            provider: slug,
+            family: 'claude',
+            created_at: nil,
+            context_window: 200_000,
+            max_output_tokens: 4096,
+            modalities: { input: ['text'], output: ['text'] },
+            capabilities: [],
+            pricing: {},
+            metadata: {}
+          )
+        end
+
+        def model_id_with_region(model_id, model_data)
+          return model_id unless model_data['inferenceTypesSupported']&.include?('INFERENCE_PROFILE')
+          return model_id if model_data['inferenceTypesSupported']&.include?('ON_DEMAND')
+
+          "us.#{model_id}"
+        end
+      end
+    end
+  end
+end