Add Bedrock embedding support via InvokeModel API

lib/ruby_llm/protocols/converse.rb

@@ -5,6 +5,7 @@ module Protocols
     # The AWS Bedrock Converse API. Requests are SigV4-signed by the provider.
     class Converse < Protocol
       include Converse::Chat
+      include Converse::Embeddings
       include Converse::Media
       include Converse::Streaming
 
@@ -18,12 +19,17 @@ def self.reasoning_embedded?(model)
       private
 
       def sync_response(payload, additional_headers = {})
+        response = signed_post(completion_url, payload, additional_headers)
+        parse_completion_response(response)
+      end
+
+      def signed_post(url, payload, additional_headers = {})
         body = JSON.generate(payload)
-        response = @connection.post(completion_url, payload) do |req|
-          req.headers.merge!(@provider.sign_headers('POST', completion_url, body))
+
+        @connection.post(url, payload) do |req|
+          req.headers.merge!(@provider.sign_headers('POST', url, body))
           req.headers.merge!(additional_headers) unless additional_headers.empty?
         end
-        parse_completion_response(response)
       end
     end
   end

lib/ruby_llm/protocols/converse/embeddings.rb

@@ -0,0 +1,41 @@
+# frozen_string_literal: true
+
+module RubyLLM
+  module Protocols
+    class Converse
+      # Embeddings methods for Bedrock's InvokeModel API.
+      module Embeddings
+        def embed(text, model:, dimensions:)
+          responses = [text].flatten.map do |value|
+            payload = render_embedding_payload(value, model:, dimensions:)
+            signed_post(embedding_url(model:), payload)
+          end
+
+          parse_embedding_responses(responses, model:, text:)
+        end
+
+        private
+
+        def embedding_url(model:)
+          "/model/#{model}/invoke"
+        end
+
+        def render_embedding_payload(text, model:, dimensions:) # rubocop:disable Lint/UnusedMethodArgument
+          {
+            inputText: text.to_s,
+            dimensions: dimensions,
+            normalize: true
+          }.compact
+        end
+
+        def parse_embedding_responses(responses, model:, text:)
+          vectors = responses.map { |response| response.body['embedding'] }
+          input_tokens = responses.sum { |response| response.body['inputTextTokenCount'] || 0 }
+          vectors = vectors.first unless text.is_a?(Array)
+
+          Embedding.new(vectors:, model:, input_tokens:)
+        end
+      end
+    end
+  end
+end

spec/ruby_llm/embedding_spec.rb

@@ -24,6 +24,7 @@
       it "#{provider}/#{model} can handle a single text with custom dimensions" do
         skip 'Mistral does not support custom dimensions' if provider == :mistral
         skip 'Azure Cohere embeddings do not support custom dimensions' if provider == :azure
+        skip 'Bedrock Titan only supports dimensions of 256, 512, or 1024' if provider == :bedrock
 
         embedding = RubyLLM.embed(test_text, model: model, provider: provider, dimensions: test_dimensions)
         expect(embedding.vectors).to be_an(Array)
@@ -42,6 +43,7 @@
       it "#{provider}/#{model} can handle multiple texts with custom dimensions" do
         skip 'Mistral does not support custom dimensions' if provider == :mistral
         skip 'Azure Cohere embeddings do not support custom dimensions' if provider == :azure
+        skip 'Bedrock Titan only supports dimensions of 256, 512, or 1024' if provider == :bedrock
 
         embeddings = RubyLLM.embed(test_texts, model: model, provider: provider, dimensions: test_dimensions)
         expect(embeddings.vectors).to be_an(Array)

spec/ruby_llm_spec.rb

@@ -8,14 +8,15 @@
     let(:log_file) { double }
     let(:log_level) { double }
 
-    before do
-      described_class.instance_variable_set(:@config, nil)
-      described_class.instance_variable_set(:@logger, nil)
-    end
-
-    after do
+    around do |example|
+      original_config = described_class.instance_variable_get(:@config)
+      original_logger = described_class.instance_variable_get(:@logger)
       described_class.instance_variable_set(:@config, nil)
       described_class.instance_variable_set(:@logger, nil)
+      example.run
+    ensure
+      described_class.instance_variable_set(:@config, original_config)
+      described_class.instance_variable_set(:@logger, original_logger)
     end
 
     context 'with configuration options' do

spec/support/models_to_test.rb

@@ -97,6 +97,7 @@ def filter_local_providers(models)
 
 EMBEDDING_MODELS = [
   { provider: :azure, model: 'Cohere-embed-v3-english' },
+  { provider: :bedrock, model: 'amazon.titan-embed-text-v2:0' },
   { provider: :gemini, model: 'gemini-embedding-001' },
   { provider: :mistral, model: 'mistral-embed' },
   { provider: :openai, model: 'text-embedding-3-small' },