Add class for embedding with Titan

Adds a `TextToEmbedding::Titan` class that calls out to Titan via
Bedrock to embed a string or array of strings of text.

Works in exactly the same way that the current OpenAI one does, and the
wrapper `TextToEmbedding` class will call this class if the env var is
set to use Titan as the embedding provider.

Author: jackbot

lib/bedrock_models.rb

@@ -1,3 +1,4 @@
 module BedrockModels
   CLAUDE_3_7_SONNET = "eu.anthropic.claude-3-7-sonnet-20250219-v1:0".freeze
+  TITAN_EMBED_V2 = "amazon.titan-embed-text-v2:0".freeze
 end

lib/search/text_to_embedding.rb

@@ -4,6 +4,8 @@ def self.call(single_or_collection_of_text, llm_provider:)
       case llm_provider.to_sym
       when :openai
         Search::TextToEmbedding::OpenAI.call(single_or_collection_of_text)
+      when :titan
+        Search::TextToEmbedding::Titan.call(single_or_collection_of_text)
       else
         raise "Unknown provider: #{llm_provider}"
       end

lib/search/text_to_embedding/titan.rb

@@ -0,0 +1,51 @@
+class Search::TextToEmbedding
+  class Titan
+    INPUT_TEXT_LENGTH_LIMIT = 50_000
+
+    def self.call(...) = new(...).call
+
+    def initialize(single_or_collection_of_text)
+      @string_input = single_or_collection_of_text.is_a?(String)
+      @text_collection = Array(single_or_collection_of_text)
+    end
+
+    def call
+      # For strings longer than the embedding model limit we have to truncate
+      # the text.
+      # This is done silently and in future we may want to log this in a
+      # database or log a warning.
+      to_embed = text_collection.map(&method(:keep_input_within_text_limit))
+
+      embeddings = convert_text_to_embeddings(to_embed)
+
+      # return just first embedding rather than an array of embeddings if we
+      # weren't given an array input
+      string_input ? embeddings.first : embeddings
+    end
+
+  private
+
+    attr_reader :string_input, :text_collection
+
+    def bedrock_client
+      @bedrock_client ||= Aws::BedrockRuntime::Client.new
+    end
+
+    def keep_input_within_text_limit(text)
+      text[0...INPUT_TEXT_LENGTH_LIMIT]
+    end
+
+    def convert_text_to_embeddings(to_embed_collection)
+      to_embed_collection.map do |text|
+        response = bedrock_client.invoke_model(
+          model_id: BedrockModels::TITAN_EMBED_V2,
+          body: {
+            inputText: text,
+          }.to_json,
+        )
+
+        JSON.parse(response.body.read)["embedding"]
+      end
+    end
+  end
+end

spec/lib/search/text_to_embedding/titan_spec.rb

@@ -0,0 +1,46 @@
+RSpec.describe Search::TextToEmbedding::Titan do
+  describe ".call" do
+    it "returns a single embedding array for a string input" do
+      client = stub_bedrock_invoke_model(
+        bedrock_titan_embedding_response([1.0, 2.0, 3.0]),
+      )
+
+      embedding = described_class.call("text")
+
+      expect(client.api_requests.size).to eq(1)
+
+      expect(embedding).to eq([1.0, 2.0, 3.0])
+    end
+
+    it "returns an array of embedding arrays for an array input" do
+      client = stub_bedrock_invoke_model(
+        bedrock_titan_embedding_response([1.0, 2.0, 3.0]),
+        bedrock_titan_embedding_response([4.0, 5.0, 6.0]),
+      )
+
+      embedding = described_class.call(["Embed this", "Embed that"])
+
+      expect(client.api_requests.size).to eq(2)
+
+      expect(embedding).to eq([[1.0, 2.0, 3.0], [4.0, 5.0, 6.0]])
+    end
+
+    it "truncates input text to the length limit" do
+      client = stub_bedrock_invoke_model(
+        bedrock_titan_embedding_response([1.0, 2.0, 3.0]),
+      )
+
+      long_text = "a" * (described_class::INPUT_TEXT_LENGTH_LIMIT + 1)
+      described_class.call(long_text)
+
+      expect(client.api_requests.size).to eq(1)
+
+      request_body = JSON.parse(
+        client.api_requests.first.dig(:params, :body),
+      )
+
+      expect(request_body["inputText"].length)
+        .to eq(described_class::INPUT_TEXT_LENGTH_LIMIT)
+    end
+  end
+end

spec/lib/search/text_to_embedding_spec.rb

@@ -1,19 +1,21 @@
 RSpec.describe Search::TextToEmbedding do
   describe ".call" do
     let(:text) { "The text" }
-    let(:provider) { :openai }
 
     it "calls the OpenAI embedding provider" do
       expect(Search::TextToEmbedding::OpenAI).to receive(:call).with(text)
-      described_class.call(text, llm_provider: provider)
+      described_class.call(text, llm_provider: :openai)
     end
 
-    context "when an unknown provider is specified" do
-      let(:provider) { :unknown_provider }
+    it "calls the Titan embedding provider" do
+      expect(Search::TextToEmbedding::Titan).to receive(:call).with(text)
+      described_class.call(text, llm_provider: :titan)
+    end
 
+    context "when an unknown provider is specified" do
       it "raises an error" do
-        expect { described_class.call(text, llm_provider: provider) }
-          .to raise_error(RuntimeError, "Unknown provider: #{provider}")
+        expect { described_class.call(text, llm_provider: "notreal") }
+          .to raise_error(RuntimeError, "Unknown provider: notreal")
       end
     end
   end

spec/support/stub_bedrock.rb

@@ -28,6 +28,13 @@ def stub_bedrock_converse(*responses)
     bedrock_client
   end
 
+  def stub_bedrock_invoke_model(*responses)
+    bedrock_client = Aws::BedrockRuntime::Client.new(stub_responses: true)
+    allow(Aws::BedrockRuntime::Client).to receive(:new).and_return(bedrock_client)
+    bedrock_client.stub_responses(:invoke_model, responses)
+    bedrock_client
+  end
+
   def bedrock_claude_structured_answer_response(question, answer, answered: true)
     lambda do |context|
       given_question = context.params.dig(:messages, -1, :content, 0, :text)
@@ -88,6 +95,15 @@ def bedrock_claude_guardrail_response(triggered: false, triggered_guardrails: []
     end
   end
 
+  def bedrock_titan_embedding_response(embedding_array)
+    {
+      content_type: "application/json",
+      body: {
+        embedding: embedding_array,
+      }.to_json,
+    }
+  end
+
   def bedrock_claude_text_response(response_text,
                                    user_message: nil,
                                    input_tokens: 10,