ruby-amazon-bedrock 0.1.1

data/Rakefile

@@ -0,0 +1,12 @@
+# frozen_string_literal: true
+
+require "bundler/gem_tasks"
+require "rspec/core/rake_task"
+
+RSpec::Core::RakeTask.new(:spec)
+
+require "rubocop/rake_task"
+
+RuboCop::RakeTask.new
+
+task default: %i[spec rubocop]

data/bin/console

@@ -0,0 +1,12 @@
+#!/usr/bin/env ruby
+# frozen_string_literal: true
+
+require "bundler/setup"
+require "bedrock"
+require "dotenv/load"
+
+# You can add fixtures and/or initialization code here to make experimenting
+# with your gem easier. You can also use a different console, if you like.
+
+require "irb"
+IRB.start(__FILE__)

data/bin/setup

@@ -0,0 +1,8 @@
+#!/usr/bin/env bash
+set -euo pipefail
+IFS=$'\n\t'
+set -vx
+
+bundle install
+
+# Do any other automated setup that you need to do here

data/lib/bedrock/version.rb

@@ -0,0 +1,9 @@
+# frozen_string_literal: true
+
+module Ruby
+  module Amazon
+    module Bedrock
+      VERSION = "0.1.1"
+    end
+  end
+end

data/lib/bedrock.rb

@@ -0,0 +1,10 @@
+# frozen_string_literal: true
+
+require_relative "bedrock/version"
+require_relative "bedrock_runtime/client"
+
+module RubyAmazonBedrock
+  # Raised when an unknown model identifier is provided to PayloadFactory.
+  class UnknownModelError < StandardError; end
+  class Error < StandardError; end
+end

data/lib/bedrock_runtime/client.rb

@@ -0,0 +1,41 @@
+# frozen_string_literal: true
+
+require 'aws-sdk-bedrockruntime'
+
+require 'bedrock_runtime/payload_factory'
+
+module RubyAmazonBedrock
+  # Client for interfacing with the Amazon Bedrock Runtime.
+  #
+  # This class provides methods to initialize a client for AWS BedrockRuntime
+  # and to invoke a model using the client.
+  class Client
+    # Initializes the AWS BedrockRuntime client.
+    #
+    # @note The AWS credentials and region are fetched from the environment variables.
+    def initialize(region:, access_key_id:, secret_access_key:)
+      @client = Aws::BedrockRuntime::Client.new(
+        region: region,
+        access_key_id: access_key_id,
+        secret_access_key: secret_access_key
+      )
+    end
+
+    # Invokes a model using the Bedrock Runtime client.
+    #
+    # @param id [String] The ID of the model to be invoked.
+    # @param input [String] The input string for what needs to be generated.
+    # @param options [Hash] Additional options for the model invocation.
+    # @return [Aws::BedrockRuntime::Types::InvokeModelOutput] The output from invoking the model.
+    # @example Invoke a model with specific ID and input
+    #   client = RubyAmazonBedrock::Client.new
+    #   client.invoke_model(
+    # 		id: 'model_id', input: 'This is what you want to generate', options: { option_key: 'option_value' }
+    # 	)
+    def invoke_model(id:, input:, options: {})
+      payload_builder_class = RubyAmazonBedrock::PayloadFactory.new(id, input, options).create
+      response = @client.invoke_model(payload_builder_class.build)
+      JSON.parse(response.body.read, symbolize_names: true)
+    end
+  end
+end

data/lib/bedrock_runtime/payload_builders/ai_21_labs/base.rb

@@ -0,0 +1,55 @@
+# frozen_string_literal: true
+
+require_relative '../base'
+
+module RubyAmazonBedrock
+  module PayloadBuilders
+    module Ai21Labs
+      # Builds and returns a payload hash suitable for the AI21Labs model processing.
+      # This method constructs a payload with specific parameters like `model_id`,
+      # `content_type`, `accept`, and a `body` that includes various AI-related settings.
+      #
+      # @return [Hash] The constructed payload containing AI model parameters and settings.
+      class Base < RubyAmazonBedrock::PayloadBuilders::Base
+        # Builds and returns a hash representing the payload.
+        # The method prepares the data necessary for processing by an unspecified model.
+        # This includes setting various parameters such as content type, model id, and
+        # specific attributes related to the prompt and its processing characteristics.
+        #
+        # @return [Hash] A hash containing various keys and values required to process the prompt:
+        #   - :model_id [String] The ID of the model to which the payload is intended.
+        #   - :content_type [String] The content type of the payload, typically 'application/json'.
+        #   - :accept [String] Specifies the MIME type that the method response should conform to.
+        #   - :body [String] A JSON string containing:
+        #     - :prompt [String] The input prompt to be processed.
+        #     - :maxTokens [Integer] The maximum number of tokens to generate.
+        #     - :temperature [Integer] Controls randomness in the response generation.
+        #     - :topP [Integer] Controls the diversity of the generated text.
+        #     - :stop_sequences [Array] An array of strings which, when encountered, will end generation.
+        #     - :countPenalty, :presencePenalty, :frequencyPenalty [Hash] Hashes containing a scale key
+        # 			used for penalty configuration.
+        def build
+          {
+            model_id: model_id,
+            content_type: 'application/json',
+            accept: '*/*',
+            body: {
+              prompt: @input,
+              maxTokens: 200,
+              temperature: 0,
+              topP: 250,
+              stop_sequences: [],
+              countPenalty: { scale: 0 },
+              presencePenalty: { scale: 0 },
+              frequencyPenalty: { scale: 0 }
+            }.to_json
+          }
+        end
+
+        def model_id
+          # noop
+        end
+      end
+    end
+  end
+end

data/lib/bedrock_runtime/payload_builders/ai_21_labs/j2_mid_v1.rb

@@ -0,0 +1,22 @@
+# frozen_string_literal: true
+
+require_relative 'base'
+
+module RubyAmazonBedrock
+  module PayloadBuilders
+    module Ai21Labs
+      # J2MidV1 is a subclass of Base. It provides functionalities specific to the AI21 Labs Jurassic-2 Mid model.
+      #
+      # @see https://us-east-1.console.aws.amazon.com/bedrock/home?region=us-east-1#/providers?model=ai21.j2-mid-v1
+      # for more information about the AI21 Labs model.
+      class J2MidV1 < Base
+        # Returns the model ID for the AI21 Labs Jurassic-2 Mid model.
+        #
+        # @return [String] 'ai21labs.j2-mid-v1'
+        def model_id
+          'ai21labs.j2-mid-v1'
+        end
+      end
+    end
+  end
+end

data/lib/bedrock_runtime/payload_builders/ai_21_labs/j2_ultra_v1.rb

@@ -0,0 +1,22 @@
+# frozen_string_literal: true
+
+# require_relative 'base'
+
+module RubyAmazonBedrock
+  module PayloadBuilders
+    module Ai21Labs
+      # J2UltraV1 is a subclass of Base. It provides functionalities specific to the AI21 Labs Jurassic-2 Ultra model.
+      #
+      # @see https://us-east-1.console.aws.amazon.com/bedrock/home?region=us-east-1#/providers?model=ai21.j2-ultra-v1
+      # for more information about the AI21 Labs model.
+      class J2UltraV1 < Base
+        # Returns the model ID for the AI21 Labs Jurassic-2 Ultra model.
+        #
+        # @return [String] 'ai21labs.j2-ultra-v1'
+        def model_id
+          'ai21labs.j2-ultra-v1'
+        end
+      end
+    end
+  end
+end

data/lib/bedrock_runtime/payload_builders/amazon/base.rb

@@ -0,0 +1,51 @@
+# frozen_string_literal: true
+
+module RubyAmazonBedrock
+  module PayloadBuilders
+    module Amazon
+      # Builds and returns a payload hash suitable for the Amazon model processing.
+      # This method constructs a payload with specific parameters like `model_id`,
+      # `content_type`, `accept`, and a `body` that includes various AI-related settings.
+      #
+      # @return [Hash] The constructed payload containing AI model parameters and settings.
+      class Base < RubyAmazonBedrock::PayloadBuilders::Base
+        # Builds and returns a hash representing the payload for a text generation request.
+        # This method prepares the necessary data structure for processing an input text through
+        # a specified model. It includes various configuration parameters for text generation.
+        #
+        # @return [Hash] The method returns a hash with the following key-value pairs:
+        #   - :model_id [String] The identifier for the model that will process the request.
+        #   - :content_type [String] Specifies the content type of the payload, set to 'application/json'.
+        #   - :accept [String] Indicates the MIME type that the response should conform to.
+        #   - :body [String] A JSON string that includes the following details:
+        #     - :inputText [String] The input text to be processed by the model.
+        #     - :textGenerationConfig [Hash] A hash containing configuration parameters for text generation:
+        #       - :maxTokenCount [Integer] The maximum number of tokens to generate.
+        #       - :stopSequences [Array<String>] An array of strings that, when encountered, will signal the end
+        # 				of generation.
+        #       - :temperature [Float] A parameter controlling the randomness in the response generation.
+        #       - :topP [Float] A parameter controlling the diversity of the generated text.
+        def build
+          {
+            model_id: model_id,
+            content_type: 'application/json',
+            accept: '*/*',
+            body: {
+              inputText: @input,
+              textGenerationConfig: {
+                maxTokenCount: 4096,
+                stopSequences: [],
+                temperature: 0,
+                topP: 1
+              }
+            }.to_json
+          }
+        end
+
+        def model_id
+          # noop
+        end
+      end
+    end
+  end
+end

data/lib/bedrock_runtime/payload_builders/amazon/titan_image_generator_v1.rb

@@ -0,0 +1,45 @@
+# frozen_string_literal: true
+
+module RubyAmazonBedrock
+  module PayloadBuilders
+    module Amazon
+      # TitanImageGeneratorV1 is a subclass of Base. It provides functionalities specific to the Amazon
+      # Titan Image Generator G1 - v1 model.
+      #
+      # @see https://us-east-1.console.aws.amazon.com/bedrock/home?region=us-east-1#/providers?model=amazon.titan-image-generator-v1
+      # for more information about the Amazon model.
+      # NOTE: This model is currently in preview release. Changes may be made to this model.
+      class TitanImageGeneratorV1 < RubyAmazonBedrock::PayloadBuilders::Base
+        def build
+          {
+            model_id: model_id,
+            content_type: 'application/json',
+            accept: 'application/json',
+            body: {
+              taskType: "TEXT_IMAGE",
+              textToImageParams: {
+                text: @input,
+                negativeTtext: @options[:negative_text]
+              },
+              imageGenerationConfig: {
+                numberOfImages: @options[:number_of_images] || 1,
+                quality: @options[:quality] || 'standard',
+                height: @options[:height] || 1024,
+                width: @options[:width] || 1024,
+                cfgScale: @options[:cfg_scale] || 8.0,
+                seed: @options[:seed] || 0
+              }
+            }.to_json
+          }
+        end
+
+        # Returns the model ID for the AI21 Labs Titan Image Generator G1 - v1 model.
+        #
+        # @return [String] 'amazon.titan-image-generator-v1'
+        def model_id
+          'amazon.titan-image-generator-v1'
+        end
+      end
+    end
+  end
+end

data/lib/bedrock_runtime/payload_builders/amazon/titan_text_express_v1.rb

@@ -0,0 +1,23 @@
+# frozen_string_literal: true
+
+require_relative 'base'
+
+module RubyAmazonBedrock
+  module PayloadBuilders
+    module Amazon
+      # TitanTextExpressV1 is a subclass of Base. It provides functionalities specific to the Amazon
+      # Titan Text G1 - Express  model.
+      #
+      # @see https://us-east-1.console.aws.amazon.com/bedrock/home?region=us-east-1#/providers?model=amazon.titan-text-express-v1
+      # for more information about the Amazon model.
+      class TitanTextExpressV1 < Base
+        # Returns the model ID for the Amazon Titan Text G1 - Express model.
+        #
+        # @return [String] 'amazon.titan-text-express-v1'
+        def model_id
+          'amazon.titan-text-express-v1'
+        end
+      end
+    end
+  end
+end

data/lib/bedrock_runtime/payload_builders/amazon/titan_text_lite_v1.rb

@@ -0,0 +1,31 @@
+# frozen_string_literal: true
+
+module RubyAmazonBedrock
+  module PayloadBuilders
+    module Amazon
+      # TitanTextLiteV1 is a subclass of Base. It provides functionalities specific to the Amazon
+      # Titan Text G1 - Lite model.
+      #
+      # @see https://us-east-1.console.aws.amazon.com/bedrock/home?region=us-east-1#/providers?model=amazon.titan-text-lite-v1
+      # for more information about the Amazon model.
+      class TitanTextLiteV1 < Base
+        def build
+          {
+            model_id: 'amazon.titan-text-lite-v1',
+            content_type: 'application/json',
+            accept: '*/*',
+            body: {
+              inputText: @input,
+              textGenerationConfig: {
+                maxTokenCount: 4096,
+                stopSequences: [],
+                temperature: 0,
+                topP: 1
+              }
+            }.to_json
+          }
+        end
+      end
+    end
+  end
+end

data/lib/bedrock_runtime/payload_builders/anthropic/base.rb

@@ -0,0 +1,55 @@
+# frozen_string_literal: true
+
+module RubyAmazonBedrock
+  module PayloadBuilders
+    module Anthropic
+      # Builds and returns a payload hash suitable for the Anthropic model processing.
+      # This method constructs a payload with specific parameters like `model_id`,
+      # `content_type`, `accept`, and a `body` that includes various AI-related settings.
+      #
+      # @return [Hash] The constructed payload containing AI model parameters and settings.
+      class Base < RubyAmazonBedrock::PayloadBuilders::Base
+        # Constructs and returns a payload hash for text generation requests.
+        # This method is designed to prepare a structured payload comprising various
+        # parameters that dictate how the input text should be processed by the models.
+        #
+        # @return [Hash] A structured payload hash containing:
+        #   - :model_id [String] Identifies the model ('anthropic.claude-2') to process the request.
+        #   - :content_type [String] Defines the content type of the payload ('application/json').
+        #   - :accept [String] Specifies the MIME type expected in the response.
+        #   - :body [String] A JSON string including the following details:
+        #     - :prompt [String] The input text for the model to process.
+        #     - :max_tokens_to_sample [Integer] The maximum number of tokens the model should generate.
+        #     - :temperature [Float] A parameter influencing the randomness in response generation.
+        #     - :top_k [Integer] A parameter specifying the number of highest probability vocabulary tokens to keep for
+        # 			top-k sampling.
+        #     - :top_p [Float] Controls the diversity of the generated text (nucleus sampling).
+        #     - :stop_sequences [Array<String>] An array containing specific sequences that signal the model to stop
+        # 			text generation.
+        #     - :anthropic_version [String] Specifies the version of the underlying model or API.
+        def build
+          {
+            model_id: model_id,
+            content_type: 'application/json',
+            accept: '*/*',
+            body: {
+              prompt: "\n\nHuman: #{@input}\n\nAssistant:",
+              max_tokens_to_sample: 300,
+              temperature: 0.5,
+              top_k: 250,
+              top_p: 1,
+              stop_sequences: [
+                '\n\nHuman'
+              ],
+              anthropic_version: 'bedrock-2023-05-31'
+            }.to_json
+          }
+        end
+
+        def model_id
+          # noop
+        end
+      end
+    end
+  end
+end

data/lib/bedrock_runtime/payload_builders/anthropic/claude_instant_v1.rb

@@ -0,0 +1,21 @@
+# frozen_string_literal: true
+
+module RubyAmazonBedrock
+  module PayloadBuilders
+    module Anthropic
+      # ClaudeInstantV1 is a subclass of Base. It provides functionalities specific to the Anthropic
+      # Claude Instant 1.2 model.
+      #
+      # @see https://us-east-1.console.aws.amazon.com/bedrock/home?region=us-east-1#/providers?model=anthropic.claude-instant-v1
+      # for more information about the Anthropic model.
+      class ClaudeInstantV1 < Base
+        # Returns the model ID for the Anthropic Claude Instant 1.2 model.
+        #
+        # @return [String] 'amazon.titan-text-express-v1'
+        def model_id
+          'anthropic.claude-instant-v1'
+        end
+      end
+    end
+  end
+end

data/lib/bedrock_runtime/payload_builders/anthropic/claude_v1.rb

@@ -0,0 +1,23 @@
+# frozen_string_literal: true
+
+require_relative 'base'
+
+module RubyAmazonBedrock
+  module PayloadBuilders
+    module Anthropic
+      # ClaudeV1 is a subclass of Base. It provides functionalities specific to the Anthropic
+      # Claude 1.3 model.
+      #
+      # @see https://us-east-1.console.aws.amazon.com/bedrock/home?region=us-east-1#/providers?model=anthropic.claude-v1
+      # for more information about the Anthropic model.
+      class ClaudeV1 < Base
+        # Returns the model ID for the Anthropic Claude 1.3 model.
+        #
+        # @return [String] 'anthropic.claude-v1'
+        def model_id
+          'anthropic.claude-v1'
+        end
+      end
+    end
+  end
+end

data/lib/bedrock_runtime/payload_builders/anthropic/claude_v2.rb

@@ -0,0 +1,21 @@
+# frozen_string_literal: true
+
+module RubyAmazonBedrock
+  module PayloadBuilders
+    module Anthropic
+      # ClaudeV2 is a subclass of Base. It provides functionalities specific to the Anthropic
+      # Claude 2 model.
+      #
+      # @see https://us-east-1.console.aws.amazon.com/bedrock/home?region=us-east-1#/providers?model=anthropic.claude-v2
+      # for more information about the Anthropic model.
+      class ClaudeV2 < Base
+        # Returns the model ID for the Anthropic Claude 2 model.
+        #
+        # @return [String] 'anthropic.claude-v2'
+        def model_id
+          'anthropic.claude-v2'
+        end
+      end
+    end
+  end
+end

data/lib/bedrock_runtime/payload_builders/base.rb

@@ -0,0 +1,33 @@
+# frozen_string_literal: true
+
+module RubyAmazonBedrock
+  module PayloadBuilders
+    # Base class serves as an abstract class for payload builders.
+    # It provides the basic structure and enforces the implementation
+    # of certain methods in derived classes.
+    class Base
+      # Initializes a new instance of the Base class.
+      # @param input [String] The input string for what needs to be generated.
+      # @param options [Hash] optional parameters to customize payload building.
+      # @option options [Any] :key Custom option key-value pairs.
+      def initialize(input, options = {})
+        @input = input
+        @options = options
+      end
+
+      # Abstract method to build the payload.
+      # @raise [NotImplementedError] if the subclass does not implement this method.
+      # @return [Hash] the constructed payload.
+      def build
+        raise NotImplementedError
+      end
+
+      # Abstract method to retrieve the model ID.
+      # @raise [NotImplementedError] if the subclass does not implement this method.
+      # @return [String] the Amazon Bedrock model ID.
+      def model_id
+        raise NotImplementedError
+      end
+    end
+  end
+end

data/lib/bedrock_runtime/payload_builders/cohere/command_base.rb

@@ -0,0 +1,46 @@
+# frozen_string_literal: true
+
+module RubyAmazonBedrock
+  module PayloadBuilders
+    module Cohere
+      # Builds and returns a payload hash suitable for the Cohere command model processing.
+      # This method constructs a payload with specific parameters like `model_id`,
+      # `content_type`, `accept`, and a `body` that includes various AI-related settings.
+      #
+      # @return [Hash] The constructed payload containing AI model parameters and settings.
+      class CommandBase < RubyAmazonBedrock::PayloadBuilders::Base
+        # Builds and returns a hash representing the payload for command generation requests.
+        # This method prepares the data necessary for processing a command input through
+        # a specified AI model. It configures several parameters to control the behavior
+        # and output of the model in response to the command.
+        #
+        # @return [Hash] A hash containing the necessary information for command processing:
+        #   - :model_id [String] The identifier for the model that will process the command.
+        #   - :content_type [String] Specifies the content type of the payload, set to 'application/json'.
+        #   - :accept [String] Indicates the MIME type that the response should conform to.
+        #   - :body [String] A JSON string that includes the following details:
+        #     - :prompt [String] The command input to be processed by the model.
+        #     - :maxTokens [Integer] The maximum number of tokens the model should generate in response.
+        #     - :temperature [Float] A parameter controlling the variability in the model's responses.
+        #     - :return_likelihood [String] Indicates whether the response should include likelihood information.
+        def build
+          {
+            model_id: model_id,
+            content_type: 'application/json',
+            accept: '*/*',
+            body: {
+              prompt: "#{@input}:",
+              max_tokens: 100,
+              temperature: 0.8
+              # return_likelihood: 'GENERATION' NOTE: This was since it was giving an error
+            }.to_json
+          }
+        end
+
+        def model_id
+          # noop
+        end
+      end
+    end
+  end
+end

data/lib/bedrock_runtime/payload_builders/cohere/command_light_text_v14.rb

@@ -0,0 +1,23 @@
+# frozen_string_literal: true
+
+require_relative 'command_base'
+
+module RubyAmazonBedrock
+  module PayloadBuilders
+    module Cohere
+      # CommandLightTextV14 is a subclass of Base. It provides functionalities specific to the Cohere
+      # Command Light model.
+      #
+      # @see https://us-east-1.console.aws.amazon.com/bedrock/home?region=us-east-1#/providers?model=cohere.command-light-text-v14
+      # for more information about the Cohere model.
+      class CommandLightTextV14 < CommandBase
+        # Returns the model ID for the Cohere Command Light model.
+        #
+        # @return [String] 'cohere.command-light-text-v14'
+        def model_id
+          'cohere.command-light-text-v14'
+        end
+      end
+    end
+  end
+end

data/lib/bedrock_runtime/payload_builders/cohere/command_text_v14.rb

@@ -0,0 +1,21 @@
+# frozen_string_literal: true
+
+module RubyAmazonBedrock
+  module PayloadBuilders
+    module Cohere
+      # CommandTextV14 is a subclass of Base. It provides functionalities specific to the Cohere
+      # Command model.
+      #
+      # @see https://us-east-1.console.aws.amazon.com/bedrock/home?region=us-east-1#/providers?model=cohere.command-text-v14
+      # for more information about the Cohere model.
+      class CommandTextV14 < CommandBase
+        # Returns the model ID for the Cohere Command model.
+        #
+        # @return [String] 'cohere.command-text-v14'
+        def model_id
+          'cohere.command-text-v14'
+        end
+      end
+    end
+  end
+end

data/lib/bedrock_runtime/payload_builders/cohere/embed_base.rb

@@ -0,0 +1,41 @@
+# frozen_string_literal: true
+
+module RubyAmazonBedrock
+  module PayloadBuilders
+    module Cohere
+      # Builds and returns a payload hash suitable for the Cohere command model processing.
+      # This method constructs a payload with specific parameters like `model_id`,
+      # `content_type`, `accept`, and a `body` that includes various AI-related settings.
+      #
+      # @return [Hash] The constructed payload containing AI model parameters and settings.
+      class EmbedBase < RubyAmazonBedrock::PayloadBuilders::Base
+        # Builds and returns a payload suitable for text embedding or search tasks.
+        # The method configures the necessary payload structure for processing text input,
+        # specifically for embedding or indexing in an AI-driven search system.
+        #
+        # @return [Hash] A hash representing the structured payload for text embedding or search:
+        #   - :model_id [String] Identifier of the AI model that will process the embedding or search task.
+        #   - :content_type [String] Specifies the content type of the payload, typically 'application/json'.
+        #   - :accept [String] Indicates the MIME type the response should conform to.
+        #   - :body [String] A JSON string that includes the following:
+        #     - :texts [Array<String>] The text(s) to be processed, either a single string or an array of strings.
+        #     - :input_type [String] Specifies the nature of the input, set to 'search_document' for search tasks.
+        def build
+          {
+            model_id: model_id,
+            content_type: 'application/json',
+            accept: '*/*',
+            body: {
+              texts: [@input],
+              input_type: 'search_document'
+            }.to_json
+          }
+        end
+
+        def model_id
+          # noop
+        end
+      end
+    end
+  end
+end