ruby-amazon-bedrock 0.2.2 → 0.2.3

checksums.yaml

@@ -1,7 +1,7 @@
 ---
 SHA256:
-  metadata.gz: e34184a2246b1054a64d6d29bc7bb1144487a508c6d27f717d8f92edb4cb32a2
-  data.tar.gz: ae40e48b1107bf163bbb30e4c058d660784f1bf9cc6fcafa0bfa9843b526dd49
+  metadata.gz: 81db5e2aa4d8b27014d6f2f24bc79083ddae2401a60a2d6d120fbc27e9facd62
+  data.tar.gz: f8f3ae19c0756142df41ef49e9d1d66d33f480013d375b6086cdb70e5365a3f2
 SHA512:
-  metadata.gz: fcdfeb1017d3bf2cefeb1bf77b68a85bd36bd3a7906651b1a56f8ff2544309ba98eee08fc1f04ce80ed22b738302fd38e7811087f56a976c6eeea1f7cfae0010
-  data.tar.gz: f19b95ef5965317c53eda21c96c69e421abdb34ce441cc850731663aa8ed35bbc67000d02a34ebc2f192ce7fbe81e0d18afe3ab59d0dbea38e32c858625fcc89
+  metadata.gz: 4cd2bdd5a70da7ddd8f92dfed38cee2432b3a1800d9d774918fd967f6d623df68afef0b1e250866a90a013a0ff3f4226146cf1e9684c1e28bd8ebdd4dad444bc
+  data.tar.gz: e2b2f8a61ab61967e7d79ce0762002fbee9ce3fdc60086b5a6cb7f0c825caa91835fe07c632380cabcf8d0cadd237f609b104d3335af7bd9ffb08224bdf84d7f

data/.gitignore

@@ -8,7 +8,7 @@
 /tmp/
 .env
 *.gem
-image.jpg
+*.jpg
 
 # rspec failure tracking
 .rspec_status

data/CHANGELOG.md

@@ -1,3 +1,8 @@
+## [0.2.3] - 2024-01-12
+
+- Fix A121 Labs bug for maxTokens parameter
+- Allow Image generator to store images in S3 after being created
+
 ## [0.2.2] - 2024-01-10
 
 - Fix bug for Amazon Titan Text Lite

data/Gemfile.lock

@@ -1,8 +1,9 @@
 PATH
   remote: .
   specs:
-    ruby-amazon-bedrock (0.2.2)
+    ruby-amazon-bedrock (0.2.3)
       aws-sdk-bedrockruntime (~> 1.0)
+      aws-sdk-s3 (~> 1.0)
 
 GEM
   remote: https://rubygems.org/
@@ -20,6 +21,13 @@ GEM
       aws-partitions (~> 1, >= 1.651.0)
       aws-sigv4 (~> 1.8)
       jmespath (~> 1, >= 1.6.1)
+    aws-sdk-kms (1.71.0)
+      aws-sdk-core (~> 3, >= 3.177.0)
+      aws-sigv4 (~> 1.1)
+    aws-sdk-s3 (1.132.0)
+      aws-sdk-core (~> 3, >= 3.179.0)
+      aws-sdk-kms (~> 1)
+      aws-sigv4 (~> 1.6)
     aws-sigv4 (1.8.0)
       aws-eventstream (~> 1, >= 1.0.2)
     bundler-audit (0.9.1)

data/README.md

@@ -14,7 +14,7 @@ Amazon Bedrock is a fully managed service that makes FMs from leading AI startup
 Add the following line to your application's Gemfile:
 
 ```ruby
-gem "ruby-amazon-bedrock", "~> 0.2.2"
+gem "ruby-amazon-bedrock", "~> 0.2.3"
 ```
 
 And then execute:
@@ -28,7 +28,7 @@ $ bundle install
 Or install with:
 
 ```bash
-$ gem install ruby-amazon-bedrock -v 0.2.1
+$ gem install ruby-amazon-bedrock -v 0.2.3
 ```
 
 and require with:

data/lib/amazon_bedrock/version.rb

@@ -1,5 +1,5 @@
 # frozen_string_literal: true
 
 module RubyAmazonBedrock
-  VERSION = "0.2.2"
+  VERSION = "0.2.3"
 end

data/lib/bedrock_runtime/client.rb

@@ -39,8 +39,7 @@ module RubyAmazonBedrock
       payload_builder_class = RubyAmazonBedrock::PayloadFactory.new(id, prompt, options).create
       response = @client.invoke_model(payload_builder_class.build)
 
-      response_builder_class = RubyAmazonBedrock::ResponseFactory.new(payload_builder_class.type, response,
-                                                                      options).create
+      response_builder_class = RubyAmazonBedrock::ResponseFactory.new(id, response, options).create
       response_builder_class.build
     end
   end

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

@@ -35,7 +35,7 @@ module RubyAmazonBedrock
             accept: '*/*',
             body: {
               prompt: @prompt,
-              maxTokenCount: parameters[:maxTokenCount],
+              maxTokens: parameters[:maxTokens],
               temperature: parameters[:temperature],
               topP: parameters[:topP],
               stopSequences: parameters[:stopSequences],
@@ -52,7 +52,7 @@ module RubyAmazonBedrock
 
         def parameters # rubocop:disable Metrics/CyclomaticComplexity
           {
-            maxTokenCount: @options[:max_tokens] || 200,
+            maxTokens: @options[:max_tokens] || 200,
             stopSequences: @options[:stop_sequences] || [],
             temperature: @options[:temperature] || 0,
             topP: @options[:top_p] || 1,

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

@@ -20,7 +20,7 @@ module RubyAmazonBedrock
         #   - :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.
+        #       - :maxTokens [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.

data/lib/bedrock_runtime/response_builders/ai_21_labs.rb

@@ -0,0 +1,32 @@
+# frozen_string_literal: true
+
+module RubyAmazonBedrock
+  module ResponseBuilders
+    # The Ai21Labs class is responsible for parsing and building a structured response from a raw text response.
+    # It converts the HTTP response for an AI21 Labs model into a structured format to make it easier to access
+    # the response data.
+    class Ai21Labs
+      # Initializes a new instance of the Text class.
+      #
+      # @param response [Object] The raw response object with is an HTTP response.
+      def initialize(response, _options = {})
+        @response = response
+      end
+
+      # Builds and returns a structured representation of the raw text response.
+      # This method parses the raw response body as JSON and symbolizes the names
+      # for easier access in Ruby.
+      #
+      # @return [Hash] A hash representing the structured data from the response body.
+      # The keys of the hash are symbolized for convenient access.
+      def build
+        response = JSON.parse(@response.body.read, symbolize_names: true)
+
+        {
+          text: response[:completions]&.first&.dig(:data, :text),
+          full_response: response
+        }
+      end
+    end
+  end
+end

data/lib/bedrock_runtime/response_builders/{image.rb → amazon_image.rb}

@@ -1,13 +1,15 @@
 # frozen_string_literal: true
 
+require 'aws-sdk-s3'
 require 'base64'
+require 'digest'
 
 module RubyAmazonBedrock
   module ResponseBuilders
-    # The Image class is responsible for handling and processing image data received in a response.
+    # The AmazonImage class is responsible for handling and processing image data received in a response.
     # It parses the response, extracts the image data, decodes it from Base64, and saves it as a file.
-    class Image
-      # Initializes a new instance of the Image class.
+    class AmazonImage
+      # Initializes a new instance of the AmazonImage class.
       #
       # @param response [Object] The raw response object which is an HTTP response.
       # The response should contain the image data in Base64 format.
@@ -16,6 +18,7 @@ module RubyAmazonBedrock
       def initialize(response, options = {})
         @response = response
         @file_path = options[:file_path] || 'image.jpg'
+        @upload = options[:upload] || false
       end
 
       # Processes the response to extract and decode image data, then saves it as a file.
@@ -26,15 +29,20 @@ module RubyAmazonBedrock
       #
       # @return [Hash] A hash indicating the result of the operation.
       # If successful, the hash includes :result set to :success and :file_path.
+      # If upload option is provided it will also include :s3_url after saving the image in S3.
       # If failure, the hash includes :result set to :failure and :error with exception details.
       def build
         response_object = JSON.parse(@response.body.read, symbolize_names: true)
         image_data = Base64.decode64(response_object[:artifacts].first[:base64])
         File.binwrite(@file_path, image_data)
+
+        upload_image if can_upload?
+
         {
           result: :success,
-          file_path: @file_path
-        }
+          file_path: @file_path,
+          s3_url: s3_url
+        }.compact
       rescue StandardError => e
         {
           result: :failure,
@@ -42,8 +50,35 @@ module RubyAmazonBedrock
         }
       end
 
-      def model_id
-        # noop
+      private
+
+      def bucket
+        @bucket ||= ENV.fetch('AWS_S3_BUCKET')
+      end
+
+      def key
+        @key ||= Digest::SHA1.hexdigest(Time.now.to_s)[0..11]
+      end
+
+      def can_upload?
+        @upload && bucket && File.exist?(@file_path)
+      end
+
+      def s3_url
+        s3_resource = Aws::S3::Resource.new
+        object = s3_resource.bucket(bucket).object(key)
+        object.presigned_url(:get)
+      end
+
+      def upload_image
+        s3_client = Aws::S3::Client.new
+        s3_client.put_object(
+          bucket: bucket,
+          key: key,
+          body: File.read(@file_path)
+        )
+      rescue StandardError => e
+        Rails.logger.error("Error uploading image to S3: #{e}")
       end
     end
   end

data/lib/bedrock_runtime/response_builders/amazon_text.rb

@@ -0,0 +1,32 @@
+# frozen_string_literal: true
+
+module RubyAmazonBedrock
+  module ResponseBuilders
+    # The AmazonText class is responsible for parsing and building a structured response from a raw text response.
+    # It converts the HTTP response for an Amazon Titan model into a structured format to make it easier to access
+    # the response data.
+    class AmazonText
+      # Initializes a new instance of the Text class.
+      #
+      # @param response [Object] The raw response object with is an HTTP response.
+      def initialize(response, _options = {})
+        @response = response
+      end
+
+      # Builds and returns a structured representation of the raw text response.
+      # This method parses the raw response body as JSON and symbolizes the names
+      # for easier access in Ruby.
+      #
+      # @return [Hash] A hash representing the structured data from the response body.
+      # The keys of the hash are symbolized for convenient access.
+      def build
+        response = JSON.parse(@response.body.read, symbolize_names: true)
+
+        {
+          text: response[:results]&.first&.dig(:outputText),
+          full_response: response
+        }
+      end
+    end
+  end
+end

data/lib/bedrock_runtime/response_builders/anthropic.rb

@@ -0,0 +1,32 @@
+# frozen_string_literal: true
+
+module RubyAmazonBedrock
+  module ResponseBuilders
+    # The Anthropic class is responsible for parsing and building a structured response from a raw text response.
+    # It converts the HTTP response for an Anthropic model into a structured format to make it easier to access
+    # the response data.
+    class Anthropic
+      # Initializes a new instance of the Text class.
+      #
+      # @param response [Object] The raw response object with is an HTTP response.
+      def initialize(response, _options = {})
+        @response = response
+      end
+
+      # Builds and returns a structured representation of the raw text response.
+      # This method parses the raw response body as JSON and symbolizes the names
+      # for easier access in Ruby.
+      #
+      # @return [Hash] A hash representing the structured data from the response body.
+      # The keys of the hash are symbolized for convenient access.
+      def build
+        response = JSON.parse(@response.body.read, symbolize_names: true)
+
+        {
+          text: response[:completion],
+          full_response: response
+        }
+      end
+    end
+  end
+end

data/lib/bedrock_runtime/response_builders/cohere_command.rb

@@ -0,0 +1,32 @@
+# frozen_string_literal: true
+
+module RubyAmazonBedrock
+  module ResponseBuilders
+    # The CohereCommand class is responsible for parsing and building a structured response from a raw text response.
+    # It converts the HTTP response for an Cohere Command model into a structured format to make it easier to access
+    # the response data.
+    class CohereCommand
+      # Initializes a new instance of the Text class.
+      #
+      # @param response [Object] The raw response object with is an HTTP response.
+      def initialize(response, _options = {})
+        @response = response
+      end
+
+      # Builds and returns a structured representation of the raw text response.
+      # This method parses the raw response body as JSON and symbolizes the names
+      # for easier access in Ruby.
+      #
+      # @return [Hash] A hash representing the structured data from the response body.
+      # The keys of the hash are symbolized for convenient access.
+      def build
+        response = JSON.parse(@response.body.read, symbolize_names: true)
+
+        {
+          text: response[:generations]&.first&.[](:text),
+          full_response: response
+        }
+      end
+    end
+  end
+end

data/lib/bedrock_runtime/response_builders/cohere_embed.rb

@@ -0,0 +1,32 @@
+# frozen_string_literal: true
+
+module RubyAmazonBedrock
+  module ResponseBuilders
+    # The CohereEmbed class is responsible for parsing and building a structured response from a raw text response.
+    # It converts the HTTP response for an Cohere Embed model into a structured format to make it easier to access
+    # the response data.
+    class CohereEmbed
+      # Initializes a new instance of the Text class.
+      #
+      # @param response [Object] The raw response object with is an HTTP response.
+      def initialize(response, _options = {})
+        @response = response
+      end
+
+      # Builds and returns a structured representation of the raw text response.
+      # This method parses the raw response body as JSON and symbolizes the names
+      # for easier access in Ruby.
+      #
+      # @return [Hash] A hash representing the structured data from the response body.
+      # The keys of the hash are symbolized for convenient access.
+      def build
+        response = JSON.parse(@response.body.read, symbolize_names: true)
+
+        {
+          text: response[:texts].first,
+          full_response: response
+        }
+      end
+    end
+  end
+end

data/lib/bedrock_runtime/response_builders/{text.rb → meta.rb}

@@ -2,13 +2,14 @@
 
 module RubyAmazonBedrock
   module ResponseBuilders
-    # The Text class is responsible for parsing and building a structured response from a raw text response.
-    # It converts the HTTP response into a structured format to make it easier to access the data.
-    class Text
+    # The Meta class is responsible for parsing and building a structured response from a raw text response.
+    # It converts the HTTP response for an Meta model into a structured format to make it easier to access
+    # the response data.
+    class Meta
       # Initializes a new instance of the Text class.
       #
       # @param response [Object] The raw response object with is an HTTP response.
-      def initialize(response)
+      def initialize(response, _options = {})
         @response = response
       end
 
@@ -19,7 +20,12 @@ module RubyAmazonBedrock
       # @return [Hash] A hash representing the structured data from the response body.
       # The keys of the hash are symbolized for convenient access.
       def build
-        JSON.parse(@response.body.read, symbolize_names: true)
+        response = JSON.parse(@response.body.read, symbolize_names: true)
+
+        {
+          text: response[:generation],
+          full_response: response
+        }
       end
     end
   end

data/lib/bedrock_runtime/response_builders/stability_ai.rb

@@ -0,0 +1,85 @@
+# frozen_string_literal: true
+
+require 'aws-sdk-s3'
+require 'base64'
+require 'digest'
+
+module RubyAmazonBedrock
+  module ResponseBuilders
+    # The StabilityAi class is responsible for handling and processing image data received in a response.
+    # It parses the response, extracts the image data, decodes it from Base64, and saves it as a file.
+    class StabilityAi
+      # Initializes a new instance of the StabilityAi class.
+      #
+      # @param response [Object] The raw response object which is an HTTP response.
+      # The response should contain the image data in Base64 format.
+      # @param options [Hash] Optional parameters, currently supporting :file_path for specifying
+      # the location and name of the file to save the image. Default is 'image.jpg'.
+      def initialize(response, options = {})
+        @response = response
+        @file_path = options[:file_path] || 'image.jpg'
+        @upload = options[:upload] || false
+      end
+
+      # Processes the response to extract and decode image data, then saves it as a file.
+      # This method parses the response body as JSON, extracts the first artifact which is
+      # expected to be an image in Base64 format, decodes it, and writes it to a file.
+      # If the operation is successful, a hash with a success status and file path is returned.
+      # If an error occurs, a hash with a failure status and error details is returned.
+      #
+      # @return [Hash] A hash indicating the result of the operation.
+      # If successful, the hash includes :result set to :success and :file_path.
+      # If upload option is provided it will also include :s3_url after saving the image in S3.
+      # If failure, the hash includes :result set to :failure and :error with exception details.
+      def build
+        response_object = JSON.parse(@response.body.read, symbolize_names: true)
+        image_data = Base64.decode64(response_object[:artifacts].first[:base64])
+        File.binwrite(@file_path, image_data)
+
+        upload_image if can_upload?
+
+        {
+          result: :success,
+          file_path: @file_path,
+          s3_url: s3_url
+        }.compact
+      rescue StandardError => e
+        {
+          result: :failure,
+          error: e
+        }
+      end
+
+      private
+
+      def bucket
+        @bucket ||= ENV.fetch('AWS_S3_BUCKET')
+      end
+
+      def key
+        @key ||= Digest::SHA1.hexdigest(Time.now.to_s)[0..11]
+      end
+
+      def can_upload?
+        @upload && bucket && File.exist?(@file_path)
+      end
+
+      def s3_url
+        s3_resource = Aws::S3::Resource.new
+        object = s3_resource.bucket(bucket).object(key)
+        object.presigned_url(:get)
+      end
+
+      def upload_image
+        s3_client = Aws::S3::Client.new
+        s3_client.put_object(
+          bucket: bucket,
+          key: key,
+          body: File.read(@file_path)
+        )
+      rescue StandardError => e
+        Rails.logger.error("Error uploading image to S3: #{e}")
+      end
+    end
+  end
+end

data/lib/bedrock_runtime/response_factory.rb

@@ -1,19 +1,25 @@
 # frozen_string_literal: true
 
-require_relative 'response_builders/image'
-require_relative 'response_builders/text'
+require_relative 'response_builders/ai_21_labs'
+require_relative 'response_builders/amazon_text'
+require_relative 'response_builders/amazon_image'
+require_relative 'response_builders/anthropic'
+require_relative 'response_builders/cohere_command'
+require_relative 'response_builders/cohere_embed'
+require_relative 'response_builders/meta'
+require_relative 'response_builders/stability_ai'
 
 module RubyAmazonBedrock
   # The ResponseFactory class is a factory for creating different types of response builder objects.
-  # It is designed to instantiate and return either a Text or Image response builder based on the type specified.
+  # It is designed to instantiate and return an object of a specific response builder class
   class ResponseFactory
     # Initializes a new instance of the ResponseFactory class.
     #
-    # @param type [Symbol] The type of response builder to create (:text or :image).
+    # @param model_id [String] The model_id of response builder to create (Amazon Bedrock model id).
     # @param response [Object] The raw response object, typically an HTTP response.
-    # @param optios [Hash] Additional options for the response builder.
-    def initialize(type, response, options = {})
-      @type = type
+    # @param options [Object] optional attributes to customize the response.
+    def initialize(model_id, response, options = {})
+      @model_id = model_id
       @response = response
       @options = options
     end
@@ -22,16 +28,36 @@ module RubyAmazonBedrock
     # This method uses the @type instance variable to determine which type of response
     # builder to instantiate and return.
     #
-    # @return [ResponseBuilders::Text, ResponseBuilders::Image] Depending on the type,
-    # it returns an instance of either ResponseBuilders::Text or ResponseBuilders::Image.
-    # Returns nil if the type does not match any known response builders.
+    # @return [ResponseBuilders::Class] Depending on the model_id,
+    # it returns an instance of any of the response builder classes.
+    # Returns nil if the model_id does not match any known response builders.
     def create
-      case @type
-      when :text
-        ResponseBuilders::Text.new(@response)
-      when :image
-        ResponseBuilders::Image.new(@response, @options)
-      end
+      builder_class = models_to_builders[@model_id]
+      builder_class.new(@response, @options)
+    end
+
+    # Defines a mapping from model identifiers to their respective builder classes.
+    #
+    # @return [Hash] The mapping of model identifiers to builder classes.
+    def models_to_builders
+      {
+        'ai21.j2-mid-v1' => ResponseBuilders::Ai21Labs,
+        'ai21.j2-ultra-v1' => ResponseBuilders::Ai21Labs,
+        'amazon.titan-image-generator-v1' => ResponseBuilders::AmazonImage,
+        'amazon.titan-text-lite-v1' => ResponseBuilders::AmazonText,
+        'amazon.titan-text-express-v1' => ResponseBuilders::AmazonText,
+        'anthropic.claude-instant-v1' => ResponseBuilders::Anthropic,
+        'anthropic.claude-v1' => ResponseBuilders::Anthropic,
+        'anthropic.claude-v2' => ResponseBuilders::Anthropic,
+        'cohere.command-light-text-v14' => ResponseBuilders::CohereCommand,
+        'cohere.command-text-v14' => ResponseBuilders::CohereCommand,
+        'cohere.embed-english-v3' => ResponseBuilders::CohereEmbed,
+        'cohere.embed-multilingual-v3' => ResponseBuilders::CohereEmbed,
+        'meta.llama2-13b-chat-v1' => ResponseBuilders::Meta,
+        'meta.llama2-70b-chat-v1' => ResponseBuilders::Meta,
+        'stability.stable-diffusion-xl-v0' => ResponseBuilders::StabilityAi,
+        'stability.stable-diffusion-xl-v1' => ResponseBuilders::StabilityAi
+      }
     end
   end
 end

data/ruby-amazon-bedrock.gemspec

@@ -36,4 +36,5 @@ Gem::Specification.new do |spec|
   spec.require_paths = ["lib"]
 
   spec.add_dependency "aws-sdk-bedrockruntime", "~> 1.0"
+  spec.add_dependency "aws-sdk-s3", "~> 1.0"
 end

data/spec/bedrock_runtime/payload_builders/ai_21_labs/base_spec.rb

@@ -9,7 +9,7 @@ RSpec.describe RubyAmazonBedrock::PayloadBuilders::Ai21Labs::Base do
   let(:body) do
     {
       prompt: prompt,
-      maxTokenCount: 200,
+      maxTokens: 200,
       temperature: 0,
       topP: 1,
       stopSequences: [],

data/spec/bedrock_runtime/payload_builders/ai_21_labs/j2_mid_v1_spec.rb

@@ -9,7 +9,7 @@ RSpec.describe RubyAmazonBedrock::PayloadBuilders::Ai21Labs::J2MidV1 do
   let(:body) do
     {
       prompt: prompt,
-      maxTokenCount: 200,
+      maxTokens: 200,
       temperature: 0,
       topP: 1,
       stopSequences: [],

data/spec/bedrock_runtime/payload_builders/ai_21_labs/j2_ultra_v1_spec.rb

@@ -9,7 +9,7 @@ RSpec.describe RubyAmazonBedrock::PayloadBuilders::Ai21Labs::J2UltraV1 do
   let(:body) do
     {
       prompt: prompt,
-      maxTokenCount: 200,
+      maxTokens: 200,
       temperature: 0,
       topP: 1,
       stopSequences: [],

data/spec/bedrock_runtime/response_builders/{text_spec.rb → ai_21_labs_spec.rb}

@@ -1,9 +1,9 @@
 # frozen_string_literal: true
 
 require 'spec_helper'
-require 'bedrock_runtime/response_builders/text'
+require 'bedrock_runtime/response_builders/ai_21_labs'
 
-RSpec.describe RubyAmazonBedrock::ResponseBuilders::Text do
+RSpec.describe RubyAmazonBedrock::ResponseBuilders::Ai21Labs do
   let(:response_body) { { key1: 'value1', key2: 'value2' }.to_json }
   let(:response) { double('response', body: StringIO.new(response_body)) }
 
@@ -14,8 +14,7 @@ RSpec.describe RubyAmazonBedrock::ResponseBuilders::Text do
       result = subject.build
       expect(result).to be_a(Hash)
       expect(result.keys).to all(be_a(Symbol))
-      expect(result[:key1]).to eq('value1')
-      expect(result[:key2]).to eq('value2')
+      expect(result[:full_response]).to eq(JSON.parse(response_body, symbolize_names: true))
     end
   end
 

data/spec/bedrock_runtime/response_builders/{image_spec.rb → stability_ai_spec.rb}

@@ -1,9 +1,9 @@
 # frozen_string_literal: true
 
 require 'spec_helper'
-require 'bedrock_runtime/response_builders/image'
+require 'bedrock_runtime/response_builders/stability_ai'
 
-RSpec.describe RubyAmazonBedrock::ResponseBuilders::Image do
+RSpec.describe RubyAmazonBedrock::ResponseBuilders::StabilityAi do
   let(:base64_image) { Base64.encode64(File.binread('spec/fixtures/sample_image.jpg')) }
   let(:response_body) { { artifacts: [{ base64: base64_image }] }.to_json }
   let(:response) { double('response', body: StringIO.new(response_body)) }

data/spec/bedrock_runtime/response_factory_spec.rb

@@ -2,34 +2,94 @@
 
 require 'spec_helper'
 require 'bedrock_runtime/response_factory'
-require 'bedrock_runtime/response_builders/text'
-require 'bedrock_runtime/response_builders/image'
+require 'bedrock_runtime/response_builders/ai_21_labs'
+require 'bedrock_runtime/response_builders/amazon_text'
+require 'bedrock_runtime/response_builders/amazon_image'
+require 'bedrock_runtime/response_builders/anthropic'
+require 'bedrock_runtime/response_builders/cohere_command'
+require 'bedrock_runtime/response_builders/cohere_embed'
+require 'bedrock_runtime/response_builders/meta'
+require 'bedrock_runtime/response_builders/stability_ai'
 
 RSpec.describe RubyAmazonBedrock::ResponseFactory do
   let(:response) { double('response') }
   let(:options) { { some_option: 'value' } }
 
-  context 'when type is :text' do
-    subject { described_class.new(:text, response) }
+  context 'when model_id is from ai_21_labs' do
+    ['ai21.j2-mid-v1', 'ai21.j2-ultra-v1'].each do |model_id|
+      subject { described_class.new(model_id, response) }
 
-    it 'creates a Text response builder' do
-      expect(subject.create).to be_a(RubyAmazonBedrock::ResponseBuilders::Text)
+      it 'creates an Ai21Labs response builder' do
+        expect(subject.create).to be_a(RubyAmazonBedrock::ResponseBuilders::Ai21Labs)
+      end
     end
   end
 
-  context 'when type is :image' do
-    subject { described_class.new(:image, response, options) }
+  context 'when model_id is from Amazon text generators' do
+    ['amazon.titan-text-lite-v1', 'amazon.titan-text-express-v1'].each do |model_id|
+      subject { described_class.new(model_id, response) }
 
-    it 'creates an Image response builder' do
-      expect(subject.create).to be_a(RubyAmazonBedrock::ResponseBuilders::Image)
+      it 'creates an AmazonText response builder' do
+        expect(subject.create).to be_a(RubyAmazonBedrock::ResponseBuilders::AmazonText)
+      end
     end
   end
 
-  context 'when type is invalid' do
-    subject { described_class.new(:invalid_type, response) }
+  context 'when model_id is from Amazon image generators' do
+    subject { described_class.new('amazon.titan-image-generator-v1', response) }
 
-    it 'returns nil' do
-      expect(subject.create).to be_nil
+    it 'creates an AmazonText response builder' do
+      expect(subject.create).to be_a(RubyAmazonBedrock::ResponseBuilders::AmazonImage)
+    end
+  end
+
+  context 'when model_id is from Anthropic generators' do
+    ['anthropic.claude-instant-v1', 'anthropic.claude-v1', 'anthropic.claude-v2'].each do |model_id|
+      subject { described_class.new(model_id, response) }
+
+      it 'creates an AmazonText response builder' do
+        expect(subject.create).to be_a(RubyAmazonBedrock::ResponseBuilders::Anthropic)
+      end
+    end
+  end
+
+  context 'when model_id is from Cohere commnand generators' do
+    ['cohere.command-light-text-v14', 'cohere.command-text-v14'].each do |model_id|
+      subject { described_class.new(model_id, response) }
+
+      it 'creates an AmazonText response builder' do
+        expect(subject.create).to be_a(RubyAmazonBedrock::ResponseBuilders::CohereCommand)
+      end
+    end
+  end
+
+  context 'when model_id is from Cohere embed generators' do
+    ['cohere.embed-english-v3', 'cohere.embed-multilingual-v3'].each do |model_id|
+      subject { described_class.new(model_id, response) }
+
+      it 'creates an AmazonText response builder' do
+        expect(subject.create).to be_a(RubyAmazonBedrock::ResponseBuilders::CohereEmbed)
+      end
+    end
+  end
+
+  context 'when model_id is from Meta generators' do
+    ['meta.llama2-13b-chat-v1', 'meta.llama2-70b-chat-v1'].each do |model_id|
+      subject { described_class.new(model_id, response) }
+
+      it 'creates an AmazonText response builder' do
+        expect(subject.create).to be_a(RubyAmazonBedrock::ResponseBuilders::Meta)
+      end
+    end
+  end
+
+  context 'when model_id is from Stability AI generators' do
+    ['stability.stable-diffusion-xl-v0', 'stability.stable-diffusion-xl-v0'].each do |model_id|
+      subject { described_class.new(model_id, response) }
+
+      it 'creates an AmazonText response builder' do
+        expect(subject.create).to be_a(RubyAmazonBedrock::ResponseBuilders::StabilityAi)
+      end
     end
   end
 end

data/spec/support/helpers.rb

@@ -1,7 +1,5 @@
 # frozen_string_literal: true
 
-require 'pry'
-
 RSpec.shared_examples 'a payload builder' do
   it 'validates the payload structure' do
     payload_builder = described_class.new(prompt, options)
@@ -28,7 +26,7 @@ RSpec.shared_context 'a121 labs parameters' do
   let(:body) do
     {
       prompt: prompt,
-      maxTokenCount: 100,
+      maxTokens: 100,
       temperature: 0.5,
       topP: 2,
       stopSequences: %w[stop1 stop2],

metadata

@@ -1,14 +1,14 @@
 --- !ruby/object:Gem::Specification
 name: ruby-amazon-bedrock
 version: !ruby/object:Gem::Version
-  version: 0.2.2
+  version: 0.2.3
 platform: ruby
 authors:
 - Alvaro Delgado
 autorequire:
 bindir: exe
 cert_chain: []
-date: 2024-01-11 00:00:00.000000000 Z
+date: 2024-01-12 00:00:00.000000000 Z
 dependencies:
 - !ruby/object:Gem::Dependency
   name: aws-sdk-bedrockruntime
@@ -24,6 +24,20 @@ dependencies:
     - - "~>"
       - !ruby/object:Gem::Version
         version: '1.0'
+- !ruby/object:Gem::Dependency
+  name: aws-sdk-s3
+  requirement: !ruby/object:Gem::Requirement
+    requirements:
+    - - "~>"
+      - !ruby/object:Gem::Version
+        version: '1.0'
+  type: :runtime
+  prerelease: false
+  version_requirements: !ruby/object:Gem::Requirement
+    requirements:
+    - - "~>"
+      - !ruby/object:Gem::Version
+        version: '1.0'
 description: The ruby-amazon-bedrock gem offers Ruby developers an efficient and user-friendly
   interface to     Amazon Bedrock, a powerful library for AI-driven text and image
   generation. This gem simplifies the process of     connecting to Amazon Bedrock's
@@ -80,8 +94,14 @@ files:
 - lib/bedrock_runtime/payload_builders/stability_ai/stable_diffusion_xl_v0.rb
 - lib/bedrock_runtime/payload_builders/stability_ai/stable_diffusion_xl_v1.rb
 - lib/bedrock_runtime/payload_factory.rb
-- lib/bedrock_runtime/response_builders/image.rb
-- lib/bedrock_runtime/response_builders/text.rb
+- lib/bedrock_runtime/response_builders/ai_21_labs.rb
+- lib/bedrock_runtime/response_builders/amazon_image.rb
+- lib/bedrock_runtime/response_builders/amazon_text.rb
+- lib/bedrock_runtime/response_builders/anthropic.rb
+- lib/bedrock_runtime/response_builders/cohere_command.rb
+- lib/bedrock_runtime/response_builders/cohere_embed.rb
+- lib/bedrock_runtime/response_builders/meta.rb
+- lib/bedrock_runtime/response_builders/stability_ai.rb
 - lib/bedrock_runtime/response_factory.rb
 - ruby-amazon-bedrock.gemspec
 - sig/ruby_amazon_bedrock.rbs
@@ -112,8 +132,8 @@ files:
 - spec/bedrock_runtime/payload_builders/stability_ai/stable_diffusion_xl_v0_spec.rb
 - spec/bedrock_runtime/payload_builders/stability_ai/stable_diffusion_xl_v1_spec.rb
 - spec/bedrock_runtime/payload_factory_spec.rb
-- spec/bedrock_runtime/response_builders/image_spec.rb
-- spec/bedrock_runtime/response_builders/text_spec.rb
+- spec/bedrock_runtime/response_builders/ai_21_labs_spec.rb
+- spec/bedrock_runtime/response_builders/stability_ai_spec.rb
 - spec/bedrock_runtime/response_factory_spec.rb
 - spec/cassettes/models/amazon_titan-text-express-v1.yml
 - spec/cassettes/models/amazon_titan-text-lite-v1.yml