ruby-amazon-bedrock 0.2.1 → 0.2.2

data/lib/bedrock_runtime/payload_factory.rb

@@ -29,7 +29,7 @@ module RubyAmazonBedrock
   class PayloadFactory
     def initialize(model_id, input, options = {})
       @model_id = model_id
-      @input = input
+      @prompt = input
       @options = options
     end
 
@@ -42,7 +42,7 @@ module RubyAmazonBedrock
 
       raise UnknownModelError, "Unknown modelId: #{@model_id}" unless builder_class
 
-      builder_class.new(@input, @options)
+      builder_class.new(@prompt, @options)
     end
 
     # Defines a mapping from model identifiers to their respective builder classes.
@@ -55,8 +55,8 @@ module RubyAmazonBedrock
         'amazon.titan-image-generator-v1' => PayloadBuilders::Amazon::TitanImageGeneratorV1,
         'amazon.titan-text-lite-v1' => PayloadBuilders::Amazon::TitanTextLiteV1,
         'amazon.titan-text-express-v1' => PayloadBuilders::Amazon::TitanTextExpressV1,
-        'anthropic.claude-v1' => PayloadBuilders::Anthropic::ClaudeV1,
         'anthropic.claude-instant-v1' => PayloadBuilders::Anthropic::ClaudeInstantV1,
+        'anthropic.claude-v1' => PayloadBuilders::Anthropic::ClaudeV1,
         'anthropic.claude-v2' => PayloadBuilders::Anthropic::ClaudeV2,
         'cohere.command-light-text-v14' => PayloadBuilders::Cohere::CommandLightTextV14,
         'cohere.command-text-v14' => PayloadBuilders::Cohere::CommandTextV14,

data/lib/bedrock_runtime/response_builders/image.rb

@@ -41,6 +41,10 @@ module RubyAmazonBedrock
           error: e
         }
       end
+
+      def model_id
+        # noop
+      end
     end
   end
 end

data/spec/bedrock_runtime/client_spec.rb

@@ -33,7 +33,7 @@ RSpec.describe RubyAmazonBedrock::Client do
     models.each do |model|
       it "invokes #{model} model with the given payload" do
         VCR.use_cassette("models/#{model}") do
-          response = client.invoke_model(id: model, input: 'Give me a short list of steps for creating a Ruby gem', options: {})
+          response = client.invoke_model(id: model, prompt: 'Give me a short list of steps for creating a Ruby gem', options: {})
           expect(response).to be_a(Hash)
         end
       end
@@ -42,7 +42,7 @@ RSpec.describe RubyAmazonBedrock::Client do
     context 'when the model is not found' do
       it 'raises an UnkownModelError' do
         VCR.use_cassette('models/not_found') do
-          expect { client.invoke_model(id: 'unknown-model', input: 'Hello World', options: {}) }.to raise_error(RubyAmazonBedrock::UnknownModelError)
+          expect { client.invoke_model(id: 'unknown-model', prompt: 'Hello World', options: {}) }.to raise_error(RubyAmazonBedrock::UnknownModelError)
         end
       end
     end

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

@@ -4,14 +4,14 @@ require 'spec_helper'
 require 'bedrock_runtime/payload_builders/ai_21_labs/base'
 
 RSpec.describe RubyAmazonBedrock::PayloadBuilders::Ai21Labs::Base do
-  let(:input) { 'example_input' }
-  let(:options) { { key: 'value' } }
+  let(:prompt) { 'example_prompt' }
+  let(:options) { {} }
   let(:body) do
     {
-      prompt: input,
-      maxTokens: 200,
+      prompt: prompt,
+      maxTokenCount: 200,
       temperature: 0,
-      topP: 1.0,
+      topP: 1,
       stopSequences: [],
       countPenalty: { scale: 0 },
       presencePenalty: { scale: 0 },
@@ -20,13 +20,11 @@ RSpec.describe RubyAmazonBedrock::PayloadBuilders::Ai21Labs::Base do
   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 'a121 labs parameters'
+      it_should_behave_like 'a payload builder'
     end
   end
 end

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

@@ -4,14 +4,14 @@ require 'spec_helper'
 require 'bedrock_runtime/payload_builders/ai_21_labs/j2_mid_v1'
 
 RSpec.describe RubyAmazonBedrock::PayloadBuilders::Ai21Labs::J2MidV1 do
-  let(:input) { 'example_input' }
-  let(:options) { { key: 'value' } }
+  let(:prompt) { 'example_prompt' }
+  let(:options) { {} }
   let(:body) do
     {
-      prompt: input,
-      maxTokens: 200,
+      prompt: prompt,
+      maxTokenCount: 200,
       temperature: 0,
-      topP: 1.0,
+      topP: 1,
       stopSequences: [],
       countPenalty: { scale: 0 },
       presencePenalty: { scale: 0 },
@@ -20,14 +20,11 @@ RSpec.describe RubyAmazonBedrock::PayloadBuilders::Ai21Labs::J2MidV1 do
   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('ai21.j2-mid-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 'a121 labs parameters'
+      it_should_behave_like 'a payload builder'
     end
   end
 end

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

@@ -4,14 +4,14 @@ require 'spec_helper'
 require 'bedrock_runtime/payload_builders/ai_21_labs/j2_ultra_v1'
 
 RSpec.describe RubyAmazonBedrock::PayloadBuilders::Ai21Labs::J2UltraV1 do
-  let(:input) { 'example_input' }
-  let(:options) { { key: 'value' } }
+  let(:prompt) { 'example_prompt' }
+  let(:options) { {} }
   let(:body) do
     {
-      prompt: input,
-      maxTokens: 200,
+      prompt: prompt,
+      maxTokenCount: 200,
       temperature: 0,
-      topP: 1.0,
+      topP: 1,
       stopSequences: [],
       countPenalty: { scale: 0 },
       presencePenalty: { scale: 0 },
@@ -20,14 +20,11 @@ RSpec.describe RubyAmazonBedrock::PayloadBuilders::Ai21Labs::J2UltraV1 do
   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'
+  end
 
-      expect(payload[:model_id]).to eq('ai21.j2-ultra-v1')
-      expect(payload[:content_type]).to eq('application/json')
-      expect(payload[:accept]).to eq('*/*')
-      expect(payload[:body]).to eq(body)
-    end
+  context 'with custom parameters' do
+    include_context 'a121 labs parameters'
+    it_should_behave_like 'a payload builder'
   end
 end

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

@@ -4,11 +4,11 @@ require 'spec_helper'
 require 'bedrock_runtime/payload_builders/amazon/base'
 
 RSpec.describe RubyAmazonBedrock::PayloadBuilders::Amazon::Base do
-  let(:input) { 'example_input' }
-  let(:options) { { key: 'value' } }
+  let(:prompt) { 'example_prompt' }
+  let(:options) { {} }
   let(:body) do
     {
-      inputText: input,
+      inputText: prompt,
       textGenerationConfig: {
         maxTokenCount: 4096,
         stopSequences: [],
@@ -19,13 +19,11 @@ RSpec.describe RubyAmazonBedrock::PayloadBuilders::Amazon::Base do
   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 'amazon titan parameters'
+      it_should_behave_like 'a payload builder'
     end
   end
 end

data/spec/bedrock_runtime/payload_builders/amazon/titan_text_express_v1_spec.rb

@@ -4,11 +4,11 @@ require 'spec_helper'
 require 'bedrock_runtime/payload_builders/amazon/titan_text_express_v1'
 
 RSpec.describe RubyAmazonBedrock::PayloadBuilders::Amazon::TitanTextExpressV1 do
-  let(:input) { 'example_input' }
-  let(:options) { { key: 'value' } }
+  let(:prompt) { 'example_prompt' }
+  let(:options) { {} }
   let(:body) do
     {
-      inputText: input,
+      inputText: prompt,
       textGenerationConfig: {
         maxTokenCount: 4096,
         stopSequences: [],
@@ -19,14 +19,11 @@ RSpec.describe RubyAmazonBedrock::PayloadBuilders::Amazon::TitanTextExpressV1 do
   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('amazon.titan-text-express-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 'amazon titan parameters'
+      it_should_behave_like 'a payload builder'
     end
   end
 end

data/spec/bedrock_runtime/payload_builders/amazon/titan_text_lite_v1_spec.rb

@@ -4,11 +4,11 @@ require 'spec_helper'
 require 'bedrock_runtime/payload_builders/amazon/titan_text_lite_v1'
 
 RSpec.describe RubyAmazonBedrock::PayloadBuilders::Amazon::TitanTextLiteV1 do
-  let(:input) { 'example_input' }
-  let(:options) { { key: 'value' } }
+  let(:prompt) { 'example_prompt' }
+  let(:options) { {} }
   let(:body) do
     {
-      inputText: input,
+      inputText: prompt,
       textGenerationConfig: {
         maxTokenCount: 4096,
         stopSequences: [],
@@ -19,14 +19,11 @@ RSpec.describe RubyAmazonBedrock::PayloadBuilders::Amazon::TitanTextLiteV1 do
   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('amazon.titan-text-lite-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 'amazon titan parameters'
+      it_should_behave_like 'a payload builder'
     end
   end
 end

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

@@ -4,12 +4,12 @@ require 'spec_helper'
 require 'bedrock_runtime/payload_builders/anthropic/base'
 
 RSpec.describe RubyAmazonBedrock::PayloadBuilders::Anthropic::Base do
-  let(:input) { 'example_input' }
-  let(:options) { { key: 'value' } }
+  let(:prompt) { 'example_prompt' }
+  let(:options) { {} }
   let(:body) do
     {
-      prompt: "\n\nHuman: #{input}\n\nAssistant:",
-      max_tokens_to_sample: 300,
+      prompt: "\n\nHuman: #{prompt}\n\nAssistant:",
+      max_tokens_to_sample: 200,
       temperature: 0.5,
       top_k: 250,
       top_p: 1,
@@ -21,13 +21,11 @@ RSpec.describe RubyAmazonBedrock::PayloadBuilders::Anthropic::Base do
   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 'anthropic parameters'
+      it_should_behave_like 'a payload builder'
     end
   end
 end

data/spec/bedrock_runtime/payload_builders/anthropic/claude_instant_v1_spec.rb

@@ -4,12 +4,12 @@ require 'spec_helper'
 require 'bedrock_runtime/payload_builders/anthropic/claude_instant_v1'
 
 RSpec.describe RubyAmazonBedrock::PayloadBuilders::Anthropic::ClaudeInstantV1 do
-  let(:input) { 'example_input' }
-  let(:options) { { key: 'value' } }
+  let(:prompt) { 'example_prompt' }
+  let(:options) { {} }
   let(:body) do
     {
-      prompt: "\n\nHuman: #{input}\n\nAssistant:",
-      max_tokens_to_sample: 300,
+      prompt: "\n\nHuman: #{prompt}\n\nAssistant:",
+      max_tokens_to_sample: 200,
       temperature: 0.5,
       top_k: 250,
       top_p: 1,
@@ -21,14 +21,11 @@ RSpec.describe RubyAmazonBedrock::PayloadBuilders::Anthropic::ClaudeInstantV1 do
   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('anthropic.claude-instant-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 'anthropic parameters'
+      it_should_behave_like 'a payload builder'
     end
   end
 end

data/spec/bedrock_runtime/payload_builders/anthropic/claude_v1_spec.rb

@@ -4,12 +4,12 @@ require 'spec_helper'
 require 'bedrock_runtime/payload_builders/anthropic/claude_v1'
 
 RSpec.describe RubyAmazonBedrock::PayloadBuilders::Anthropic::ClaudeV1 do
-  let(:input) { 'example_input' }
-  let(:options) { { key: 'value' } }
+  let(:prompt) { 'example_prompt' }
+  let(:options) { {} }
   let(:body) do
     {
-      prompt: "\n\nHuman: #{input}\n\nAssistant:",
-      max_tokens_to_sample: 300,
+      prompt: "\n\nHuman: #{prompt}\n\nAssistant:",
+      max_tokens_to_sample: 200,
       temperature: 0.5,
       top_k: 250,
       top_p: 1,
@@ -21,14 +21,11 @@ RSpec.describe RubyAmazonBedrock::PayloadBuilders::Anthropic::ClaudeV1 do
   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('anthropic.claude-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 'anthropic parameters'
+      it_should_behave_like 'a payload builder'
     end
   end
 end

data/spec/bedrock_runtime/payload_builders/anthropic/claude_v2_spec.rb

@@ -4,12 +4,12 @@ require 'spec_helper'
 require 'bedrock_runtime/payload_builders/anthropic/claude_v2'
 
 RSpec.describe RubyAmazonBedrock::PayloadBuilders::Anthropic::ClaudeV2 do
-  let(:input) { 'example_input' }
-  let(:options) { { key: 'value' } }
+  let(:prompt) { 'example_prompt' }
+  let(:options) { {} }
   let(:body) do
     {
-      prompt: "\n\nHuman: #{input}\n\nAssistant:",
-      max_tokens_to_sample: 300,
+      prompt: "\n\nHuman: #{prompt}\n\nAssistant:",
+      max_tokens_to_sample: 200,
       temperature: 0.5,
       top_k: 250,
       top_p: 1,
@@ -21,14 +21,11 @@ RSpec.describe RubyAmazonBedrock::PayloadBuilders::Anthropic::ClaudeV2 do
   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('anthropic.claude-v2')
-      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 'anthropic parameters'
+      it_should_behave_like 'a payload builder'
     end
   end
 end

data/spec/bedrock_runtime/payload_builders/base_spec.rb

@@ -11,7 +11,7 @@ RSpec.describe RubyAmazonBedrock::PayloadBuilders::Base do
 
       payload_builder = described_class.new(input, options)
 
-      expect(payload_builder.instance_variable_get(:@input)).to eq(input)
+      expect(payload_builder.instance_variable_get(:@prompt)).to eq(input)
       expect(payload_builder.instance_variable_get(:@options)).to eq(options)
     end
   end

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