ruby-amazon-bedrock 0.2.0 → 0.2.2

data/spec/bedrock_runtime/payload_builders/cohere/command_base_spec.rb

@@ -4,24 +4,29 @@ require 'spec_helper'
 require 'bedrock_runtime/payload_builders/cohere/command_base'
 
 RSpec.describe RubyAmazonBedrock::PayloadBuilders::Cohere::CommandBase do
-  let(:input) { 'example_input' }
-  let(:options) { { key: 'value' } }
+  let(:prompt) { 'example_prompt' }
+  let(:options) { {} }
   let(:body) do
     {
-      prompt: "#{input}:",
-      max_tokens: 100,
-      temperature: 0.8
+      prompt: "#{prompt}:",
+      temperature: 0.9,
+      p: 0.75,
+      k: 0,
+      max_tokens: 20,
+      num_generations: 1,
+      return_likelihoods: 'GENERATION',
+      stop_sequences: [],
+      stream: false,
+      truncate: 'NONE'
     }.to_json
   end
 
   describe '#build' do
-    it 'returns a hash with the expected structure' do
-      payload_builder = described_class.new(input, options)
-      payload = payload_builder.build
+    it_behaves_like 'a payload builder'
 
-      expect(payload[:content_type]).to eq('application/json')
-      expect(payload[:accept]).to eq('*/*')
-      expect(payload[:body]).to eq(body)
+    context 'with custom parameters' do
+      include_context 'cohere command parameters'
+      it_should_behave_like 'a payload builder'
     end
   end
 end

data/spec/bedrock_runtime/payload_builders/cohere/command_light_text_v14_spec.rb

@@ -4,25 +4,29 @@ require 'spec_helper'
 require 'bedrock_runtime/payload_builders/cohere/command_light_text_v14'
 
 RSpec.describe RubyAmazonBedrock::PayloadBuilders::Cohere::CommandLightTextV14 do
-  let(:input) { 'example_input' }
-  let(:options) { { key: 'value' } }
+  let(:prompt) { 'example_input' }
+  let(:options) { {} }
   let(:body) do
     {
-      prompt: "#{input}:",
-      max_tokens: 100,
-      temperature: 0.8
+      prompt: "#{prompt}:",
+      temperature: 0.9,
+      p: 0.75,
+      k: 0,
+      max_tokens: 20,
+      num_generations: 1,
+      return_likelihoods: 'GENERATION',
+      stop_sequences: [],
+      stream: false,
+      truncate: 'NONE'
     }.to_json
   end
 
   describe '#build' do
-    it 'returns a hash with the expected structure' do
-      payload_builder = described_class.new(input, options)
-      payload = payload_builder.build
+    it_behaves_like 'a payload builder'
 
-      expect(payload[:model_id]).to eq('cohere.command-light-text-v14')
-      expect(payload[:content_type]).to eq('application/json')
-      expect(payload[:accept]).to eq('*/*')
-      expect(payload[:body]).to eq(body)
+    context 'with custom parameters' do
+      include_context 'cohere command parameters'
+      it_should_behave_like 'a payload builder'
     end
   end
 end

data/spec/bedrock_runtime/payload_builders/cohere/command_text_v14_spec.rb

@@ -4,25 +4,29 @@ require 'spec_helper'
 require 'bedrock_runtime/payload_builders/cohere/command_text_v14'
 
 RSpec.describe RubyAmazonBedrock::PayloadBuilders::Cohere::CommandTextV14 do
-  let(:input) { 'example_input' }
-  let(:options) { { key: 'value' } }
+  let(:prompt) { 'example_prompt' }
+  let(:options) { {} }
   let(:body) do
     {
-      prompt: "#{input}:",
-      max_tokens: 100,
-      temperature: 0.8
+      prompt: "#{prompt}:",
+      temperature: 0.9,
+      p: 0.75,
+      k: 0,
+      max_tokens: 20,
+      num_generations: 1,
+      return_likelihoods: 'GENERATION',
+      stop_sequences: [],
+      stream: false,
+      truncate: 'NONE'
     }.to_json
   end
 
   describe '#build' do
-    it 'returns a hash with the expected structure' do
-      payload_builder = described_class.new(input, options)
-      payload = payload_builder.build
+    it_behaves_like 'a payload builder'
 
-      expect(payload[:model_id]).to eq('cohere.command-text-v14')
-      expect(payload[:content_type]).to eq('application/json')
-      expect(payload[:accept]).to eq('*/*')
-      expect(payload[:body]).to eq(body)
+    context 'with custom parameters' do
+      include_context 'cohere command parameters'
+      it_should_behave_like 'a payload builder'
     end
   end
 end

data/spec/bedrock_runtime/payload_builders/cohere/embed_base_spec.rb

@@ -4,23 +4,21 @@ require 'spec_helper'
 require 'bedrock_runtime/payload_builders/cohere/embed_base'
 
 RSpec.describe RubyAmazonBedrock::PayloadBuilders::Cohere::EmbedBase do
-  let(:input) { 'example_input' }
-  let(:options) { { key: 'value' } }
+  let(:prompt) { 'example_prompt' }
+  let(:options) { {} }
   let(:body) do
     {
-      texts: [input],
-      input_type: 'search_document'
+      texts: [prompt],
+      input_type: 'search_document',
+      truncate: 'NONE'
     }.to_json
   end
 
   describe '#build' do
-    it 'returns a hash with the expected structure' do
-      payload_builder = described_class.new(input, options)
-      payload = payload_builder.build
+    it_behaves_like 'a payload builder'
 
-      expect(payload[:content_type]).to eq('application/json')
-      expect(payload[:accept]).to eq('*/*')
-      expect(payload[:body]).to eq(body)
+    context 'with custom parameters' do
+      it_should_behave_like 'a payload builder'
     end
   end
 end

data/spec/bedrock_runtime/payload_builders/cohere/embed_english_v3_spec.rb

@@ -4,24 +4,22 @@ require 'spec_helper'
 require 'bedrock_runtime/payload_builders/cohere/embed_english_v3'
 
 RSpec.describe RubyAmazonBedrock::PayloadBuilders::Cohere::EmbedEnglishV3 do
-  let(:input) { 'example_input' }
-  let(:options) { { key: 'value' } }
+  let(:prompt) { 'example_prompt' }
+  let(:options) { {} }
   let(:body) do
     {
-      texts: [input],
-      input_type: 'search_document'
+      texts: [prompt],
+      input_type: 'search_document',
+      truncate: 'NONE'
     }.to_json
   end
 
   describe '#build' do
-    it 'returns a hash with the expected structure' do
-      payload_builder = described_class.new(input, options)
-      payload = payload_builder.build
+    it_behaves_like 'a payload builder'
 
-      expect(payload[:model_id]).to eq('cohere.embed-english-v3')
-      expect(payload[:content_type]).to eq('application/json')
-      expect(payload[:accept]).to eq('*/*')
-      expect(payload[:body]).to eq(body)
+    context 'with custom parameters' do
+      include_context 'cohere embed parameters'
+      it_should_behave_like 'a payload builder'
     end
   end
 end

data/spec/bedrock_runtime/payload_builders/cohere/embed_multilingual_v3_spec.rb

@@ -4,24 +4,22 @@ require 'spec_helper'
 require 'bedrock_runtime/payload_builders/cohere/embed_multilingual_v3'
 
 RSpec.describe RubyAmazonBedrock::PayloadBuilders::Cohere::EmbedMultilingualV3 do
-  let(:input) { 'example_input' }
-  let(:options) { { key: 'value' } }
+  let(:prompt) { 'example_prompt' }
+  let(:options) { {} }
   let(:body) do
     {
-      texts: [input],
-      input_type: 'search_document'
+      texts: [prompt],
+      input_type: 'search_document',
+      truncate: 'NONE'
     }.to_json
   end
 
   describe '#build' do
-    it 'returns a hash with the expected structure' do
-      payload_builder = described_class.new(input, options)
-      payload = payload_builder.build
+    it_behaves_like 'a payload builder'
 
-      expect(payload[:model_id]).to eq('cohere.embed-multilingual-v3')
-      expect(payload[:content_type]).to eq('application/json')
-      expect(payload[:accept]).to eq('*/*')
-      expect(payload[:body]).to eq(body)
+    context 'with custom parameters' do
+      include_context 'cohere embed parameters'
+      it_should_behave_like 'a payload builder'
     end
   end
 end

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

@@ -4,25 +4,23 @@ require 'spec_helper'
 require 'bedrock_runtime/payload_builders/meta/base'
 
 RSpec.describe RubyAmazonBedrock::PayloadBuilders::Meta::Base do
-  let(:input) { 'example_input' }
-  let(:options) { { key: 'value' } }
+  let(:prompt) { 'example_prompt' }
+  let(:options) { {} }
   let(:body) do
     {
-      prompt: input,
+      prompt: prompt,
       max_gen_len: 512,
-      temperature: 0.2,
+      temperature: 0.5,
       top_p: 0.9
     }.to_json
   end
 
   describe '#build' do
-    it 'returns a hash with the expected structure' do
-      payload_builder = described_class.new(input, options)
-      payload = payload_builder.build
+    it_should_behave_like 'a payload builder'
 
-      expect(payload[:content_type]).to eq('application/json')
-      expect(payload[:accept]).to eq('*/*')
-      expect(payload[:body]).to eq(body)
+    context 'with custom parameters' do
+      include_context 'meta parameters'
+      it_should_behave_like 'a payload builder'
     end
   end
 end

data/spec/bedrock_runtime/payload_builders/meta/llama213b_chat_v1_spec.rb

@@ -4,26 +4,23 @@ require 'spec_helper'
 require 'bedrock_runtime/payload_builders/meta/llama213b_chat_v1'
 
 RSpec.describe RubyAmazonBedrock::PayloadBuilders::Meta::Llama213bChatV1 do
-  let(:input) { 'example_input' }
-  let(:options) { { key: 'value' } }
+  let(:prompt) { 'example_prompt' }
+  let(:options) { {} }
   let(:body) do
     {
-      prompt: input,
+      prompt: prompt,
       max_gen_len: 512,
-      temperature: 0.2,
+      temperature: 0.5,
       top_p: 0.9
     }.to_json
   end
 
   describe '#build' do
-    it 'returns a hash with the expected structure' do
-      payload_builder = described_class.new(input, options)
-      payload = payload_builder.build
+    it_should_behave_like 'a payload builder'
 
-      expect(payload[:model_id]).to eq('meta.llama2-13b-chat-v1')
-      expect(payload[:content_type]).to eq('application/json')
-      expect(payload[:accept]).to eq('*/*')
-      expect(payload[:body]).to eq(body)
+    context 'with custom parameters' do
+      include_context 'meta parameters'
+      it_should_behave_like 'a payload builder'
     end
   end
 end

data/spec/bedrock_runtime/payload_builders/meta/llama270b_chat_v1_spec.rb

@@ -4,26 +4,23 @@ require 'spec_helper'
 require 'bedrock_runtime/payload_builders/meta/llama270b_chat_v1'
 
 RSpec.describe RubyAmazonBedrock::PayloadBuilders::Meta::Llama270bChatV1 do
-  let(:input) { 'example_input' }
-  let(:options) { { key: 'value' } }
+  let(:prompt) { 'example_prompt' }
+  let(:options) { {} }
   let(:body) do
     {
-      prompt: input,
+      prompt: prompt,
       max_gen_len: 512,
-      temperature: 0.2,
+      temperature: 0.5,
       top_p: 0.9
     }.to_json
   end
 
   describe '#build' do
-    it 'returns a hash with the expected structure' do
-      payload_builder = described_class.new(input, options)
-      payload = payload_builder.build
+    it_should_behave_like 'a payload builder'
 
-      expect(payload[:model_id]).to eq('meta.llama2-70b-chat-v1')
-      expect(payload[:content_type]).to eq('application/json')
-      expect(payload[:accept]).to eq('*/*')
-      expect(payload[:body]).to eq(body)
+    context 'with custom parameters' do
+      include_context 'meta parameters'
+      it_should_behave_like 'a payload builder'
     end
   end
 end

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

@@ -4,27 +4,25 @@ require 'spec_helper'
 require 'bedrock_runtime/payload_builders/stability_ai/base'
 
 RSpec.describe RubyAmazonBedrock::PayloadBuilders::StabilityAi::Base do
-  let(:input) { 'example_input' }
-  let(:options) { { key: 'value' } }
+  let(:prompt) { 'example_prompt' }
+  let(:options) { {} }
   let(:body) do
     {
       text_prompts: [
-        { text: input }
+        { text: prompt }
       ],
       cfg_scale: 10,
       seed: 0,
-      steps: 50
+      steps: 30
     }.to_json
   end
 
   describe '#build' do
-    it 'returns a hash with the expected structure' do
-      payload_builder = described_class.new(input, options)
-      payload = payload_builder.build
+    it_should_behave_like 'a payload builder'
 
-      expect(payload[:content_type]).to eq('application/json')
-      expect(payload[:accept]).to eq('*/*')
-      expect(payload[:body]).to eq(body)
+    context 'with custom parameters' do
+      include_context 'stability ai parameters'
+      it_should_behave_like 'a payload builder'
     end
   end
 end

data/spec/bedrock_runtime/payload_builders/stability_ai/stable_diffusion_xl_v0_spec.rb

@@ -4,28 +4,25 @@ require 'spec_helper'
 require 'bedrock_runtime/payload_builders/stability_ai/stable_diffusion_xl_v0'
 
 RSpec.describe RubyAmazonBedrock::PayloadBuilders::StabilityAi::StableDiffusionXlV0 do
-  let(:input) { 'example_input' }
-  let(:options) { { key: 'value' } }
+  let(:prompt) { 'example_prompt' }
+  let(:options) { {} }
   let(:body) do
     {
       text_prompts: [
-        { text: input }
+        { text: prompt }
       ],
       cfg_scale: 10,
       seed: 0,
-      steps: 50
+      steps: 30
     }.to_json
   end
 
   describe '#build' do
-    it 'returns a hash with the expected structure' do
-      payload_builder = described_class.new(input, options)
-      payload = payload_builder.build
+    it_should_behave_like 'a payload builder'
 
-      expect(payload[:model_id]).to eq('stability.stable-diffusion-xl-v0')
-      expect(payload[:content_type]).to eq('application/json')
-      expect(payload[:accept]).to eq('*/*')
-      expect(payload[:body]).to eq(body)
+    context 'with custom parameters' do
+      include_context 'stability ai parameters'
+      it_should_behave_like 'a payload builder'
     end
   end
 end

data/spec/bedrock_runtime/payload_builders/stability_ai/stable_diffusion_xl_v1_spec.rb

@@ -4,28 +4,25 @@ require 'spec_helper'
 require 'bedrock_runtime/payload_builders/stability_ai/stable_diffusion_xl_v1'
 
 RSpec.describe RubyAmazonBedrock::PayloadBuilders::StabilityAi::StableDiffusionXlV1 do
-  let(:input) { 'example_input' }
-  let(:options) { { key: 'value' } }
+  let(:prompt) { 'example_prompt' }
+  let(:options) { {} }
   let(:body) do
     {
       text_prompts: [
-        { text: input }
+        { text: prompt }
       ],
       cfg_scale: 10,
       seed: 0,
-      steps: 50
+      steps: 30
     }.to_json
   end
 
   describe '#build' do
-    it 'returns a hash with the expected structure' do
-      payload_builder = described_class.new(input, options)
-      payload = payload_builder.build
+    it_should_behave_like 'a payload builder'
 
-      expect(payload[:model_id]).to eq('stability.stable-diffusion-xl-v1')
-      expect(payload[:content_type]).to eq('application/json')
-      expect(payload[:accept]).to eq('*/*')
-      expect(payload[:body]).to eq(body)
+    context 'with custom parameters' do
+      include_context 'stability ai parameters'
+      it_should_behave_like 'a payload builder'
     end
   end
 end

data/spec/bedrock_runtime/payload_factory_spec.rb

@@ -5,13 +5,13 @@ require 'bedrock_runtime/payload_factory'
 
 RSpec.describe RubyAmazonBedrock::PayloadFactory do
   let(:bedrock_models) do
-    ["ai21labs.j2-mid-v1",
-     "ai21labs.j2-ultra-v1",
+    ["ai21.j2-mid-v1",
+     "ai21.j2-ultra-v1",
      "amazon.titan-image-generator-v1",
      "amazon.titan-text-lite-v1",
      "amazon.titan-text-express-v1",
-     "anthropic.claude-v1",
      "anthropic.claude-instant-v1",
+     "anthropic.claude-v1",
      "anthropic.claude-v2",
      "cohere.command-light-text-v14",
      "cohere.command-text-v14",

data/spec/spec_helper.rb

@@ -1,14 +1,13 @@
 # frozen_string_literal: true
 
 require 'dotenv/load'
-require 'simplecov'
 require 'vcr'
 
 require 'amazon_bedrock'
 require 'bedrock_runtime/client'
 require 'bedrock_runtime/payload_factory'
 
-SimpleCov.start
+require 'support/helpers'
 
 RSpec.configure do |config|
   # Enable flags like --only-failures and --next-failure

data/spec/support/helpers.rb

@@ -0,0 +1,169 @@
+# 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)
+    payload = payload_builder.build
+
+    expect(payload[:content_type]).to eq('application/json')
+    expect(payload[:accept]).to eq('*/*')
+    expect(payload[:body]).to eq(body)
+  end
+end
+
+RSpec.shared_context 'a121 labs parameters' do
+  let(:options) do
+    {
+      max_tokens: 100,
+      temperature: 0.5,
+      top_p: 2,
+      stop_sequences: %w[stop1 stop2],
+      count_penalty: 1,
+      presence_penalty: 1,
+      frequency_penalty: 1
+    }
+  end
+  let(:body) do
+    {
+      prompt: prompt,
+      maxTokenCount: 100,
+      temperature: 0.5,
+      topP: 2,
+      stopSequences: %w[stop1 stop2],
+      countPenalty: { scale: 1 },
+      presencePenalty: { scale: 1 },
+      frequencyPenalty: { scale: 1 }
+    }.to_json
+  end
+end
+
+RSpec.shared_context 'amazon titan parameters' do
+  let(:options) do
+    {
+      max_tokens: 100,
+      stop_sequences: %w[stop1 stop2],
+      temperature: 0.5,
+      top_p: 2
+    }
+  end
+  let(:body) do
+    {
+      inputText: prompt,
+      textGenerationConfig: {
+        maxTokenCount: 100,
+        stopSequences: %w[stop1 stop2],
+        temperature: 0.5,
+        topP: 2
+      }
+    }.to_json
+  end
+end
+
+RSpec.shared_context 'anthropic parameters' do
+  let(:options) do
+    {
+      max_tokens: 100,
+      stop_sequences: [],
+      temperature: 0.1,
+      top_k: 150,
+      top_p: 2
+    }
+  end
+  let(:body) do
+    {
+      prompt: "\n\nHuman: #{prompt}\n\nAssistant:",
+      max_tokens_to_sample: 100,
+      temperature: 0.1,
+      top_k: 150,
+      top_p: 2,
+      stop_sequences: [],
+      anthropic_version: 'bedrock-2023-05-31'
+    }.to_json
+  end
+end
+
+RSpec.shared_context 'cohere command parameters' do
+  let(:options) do
+    {
+      temperature: 0.1,
+      top_p: 0.2,
+      top_k: 0.3,
+      max_tokens: 100,
+      num_generations: 2,
+      return_likelihoods: 'ALL',
+      stop_sequences: %w[stop1 stop2],
+      stream: true,
+      truncate: 'LINE'
+    }
+  end
+  let(:body) do
+    {
+      prompt: "#{prompt}:",
+      temperature: 0.1,
+      p: 0.2,
+      k: 0.3,
+      max_tokens: 100,
+      num_generations: 2,
+      return_likelihoods: 'ALL',
+      stop_sequences: %w[stop1 stop2],
+      stream: true,
+      truncate: 'LINE'
+    }.to_json
+  end
+end
+
+RSpec.shared_context 'cohere embed parameters' do
+  let(:options) do
+    {
+      input_type: 'classification',
+      truncate: 'LEFT'
+    }
+  end
+  let(:body) do
+    {
+      texts: [prompt],
+      input_type: 'classification',
+      truncate: 'LEFT'
+    }.to_json
+  end
+end
+
+RSpec.shared_context 'meta parameters' do
+  let(:options) do
+    {
+      max_tokens: 256,
+      temperature: 0.1,
+      top_p: 0.2
+    }
+  end
+  let(:body) do
+    {
+      prompt: prompt,
+      max_gen_len: 256,
+      temperature: 0.1,
+      top_p: 0.2
+    }.to_json
+  end
+end
+
+RSpec.shared_context 'stability ai parameters' do
+  let(:options) do
+    {
+      cfg_scale: 20,
+      seed: 1,
+      steps: 40
+    }
+  end
+  let(:body) do
+    {
+      text_prompts: [
+        { text: prompt }
+      ],
+      cfg_scale: 20,
+      seed: 1,
+      steps: 40
+    }.to_json
+  end
+end