Bedrock Titan embedding client adds BedrockTitanEmbeddingOptions to support dynamic embedding request types.

Author: wmz7year

models/spring-ai-bedrock/src/main/java/org/springframework/ai/bedrock/aot/BedrockRuntimeHints.java

@@ -28,6 +28,7 @@
 import org.springframework.ai.bedrock.llama.BedrockLlamaChatOptions;
 import org.springframework.ai.bedrock.llama.api.LlamaChatBedrockApi;
 import org.springframework.ai.bedrock.titan.BedrockTitanChatOptions;
+import org.springframework.ai.bedrock.titan.BedrockTitanEmbeddingOptions;
 import org.springframework.ai.bedrock.titan.api.TitanChatBedrockApi;
 import org.springframework.ai.bedrock.titan.api.TitanEmbeddingBedrockApi;
 import org.springframework.aot.hint.MemberCategory;
@@ -43,6 +44,7 @@
  * @author Josh Long
  * @author Christian Tzolov
  * @author Mark Pollack
+ * @author Wei Jiang
  */
 public class BedrockRuntimeHints implements RuntimeHintsRegistrar {
 
@@ -72,6 +74,8 @@ public void registerHints(RuntimeHints hints, ClassLoader classLoader) {
 			hints.reflection().registerType(tr, mcs);
 		for (var tr : findJsonAnnotatedClassesInPackage(BedrockTitanChatOptions.class))
 			hints.reflection().registerType(tr, mcs);
+		for (var tr : findJsonAnnotatedClassesInPackage(BedrockTitanEmbeddingOptions.class))
+			hints.reflection().registerType(tr, mcs);
 		for (var tr : findJsonAnnotatedClassesInPackage(TitanEmbeddingBedrockApi.class))
 			hints.reflection().registerType(tr, mcs);
 

models/spring-ai-bedrock/src/main/java/org/springframework/ai/bedrock/titan/BedrockTitanEmbeddingClient.java

@@ -28,6 +28,7 @@
 import org.springframework.ai.document.Document;
 import org.springframework.ai.embedding.AbstractEmbeddingClient;
 import org.springframework.ai.embedding.Embedding;
+import org.springframework.ai.embedding.EmbeddingOptions;
 import org.springframework.ai.embedding.EmbeddingRequest;
 import org.springframework.ai.embedding.EmbeddingResponse;
 import org.springframework.util.Assert;
@@ -40,6 +41,7 @@
  * Note: Titan Embedding does not support batch embedding.
  *
  * @author Christian Tzolov
+ * @author Wei Jiang
  * @since 0.8.0
  */
 public class BedrockTitanEmbeddingClient extends AbstractEmbeddingClient {
@@ -87,9 +89,7 @@ public EmbeddingResponse call(EmbeddingRequest request) {
 
 		List<List<Double>> embeddingList = new ArrayList<>();
 		for (String inputContent : request.getInstructions()) {
-			var apiRequest = (this.inputType == InputType.IMAGE)
-					? new TitanEmbeddingRequest.Builder().withInputImage(inputContent).build()
-					: new TitanEmbeddingRequest.Builder().withInputText(inputContent).build();
+			var apiRequest = createTitanEmbeddingRequest(inputContent, request.getOptions());
 			TitanEmbeddingResponse response = this.embeddingApi.embedding(apiRequest);
 			embeddingList.add(response.embedding());
 		}
@@ -100,6 +100,18 @@ public EmbeddingResponse call(EmbeddingRequest request) {
 		return new EmbeddingResponse(embeddings);
 	}
 
+	private TitanEmbeddingRequest createTitanEmbeddingRequest(String inputContent, EmbeddingOptions requestOptions) {
+		InputType inputType = this.inputType;
+
+		if (requestOptions != null
+				&& requestOptions instanceof BedrockTitanEmbeddingOptions bedrockTitanEmbeddingOptions) {
+			inputType = bedrockTitanEmbeddingOptions.getInputType();
+		}
+
+		return (inputType == InputType.IMAGE) ? new TitanEmbeddingRequest.Builder().withInputImage(inputContent).build()
+				: new TitanEmbeddingRequest.Builder().withInputText(inputContent).build();
+	}
+
 	@Override
 	public int dimensions() {
 		if (this.inputType == InputType.IMAGE) {

models/spring-ai-bedrock/src/main/java/org/springframework/ai/bedrock/titan/BedrockTitanEmbeddingOptions.java

@@ -0,0 +1,65 @@
+/*
+ * Copyright 2023 - 2024 the original author or authors.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ * https://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+package org.springframework.ai.bedrock.titan;
+
+import com.fasterxml.jackson.annotation.JsonInclude;
+import com.fasterxml.jackson.annotation.JsonInclude.Include;
+
+import org.springframework.ai.bedrock.titan.BedrockTitanEmbeddingClient.InputType;
+import org.springframework.ai.embedding.EmbeddingOptions;
+import org.springframework.util.Assert;
+
+/**
+ * @author Wei Jiang
+ */
+@JsonInclude(Include.NON_NULL)
+public class BedrockTitanEmbeddingOptions implements EmbeddingOptions {
+
+	/**
+	 * Titan Embedding API input types. Could be either text or image (encoded in base64).
+	 */
+	private InputType inputType;
+
+	public static Builder builder() {
+		return new Builder();
+	}
+
+	public static class Builder {
+
+		private BedrockTitanEmbeddingOptions options = new BedrockTitanEmbeddingOptions();
+
+		public Builder withInputType(InputType inputType) {
+			Assert.notNull(inputType, "input type can not be null.");
+
+			this.options.setInputType(inputType);
+			return this;
+		}
+
+		public BedrockTitanEmbeddingOptions build() {
+			return this.options;
+		}
+
+	}
+
+	public InputType getInputType() {
+		return this.inputType;
+	}
+
+	public void setInputType(InputType inputType) {
+		this.inputType = inputType;
+	}
+
+}

models/spring-ai-bedrock/src/test/java/org/springframework/ai/bedrock/titan/BedrockTitanEmbeddingClientIT.java

@@ -22,17 +22,23 @@
 
 import org.junit.jupiter.api.Test;
 import org.junit.jupiter.api.condition.EnabledIfEnvironmentVariable;
+
+import software.amazon.awssdk.auth.credentials.EnvironmentVariableCredentialsProvider;
 import software.amazon.awssdk.regions.Region;
 
+import org.springframework.ai.bedrock.titan.BedrockTitanEmbeddingClient.InputType;
 import org.springframework.ai.bedrock.titan.api.TitanEmbeddingBedrockApi;
 import org.springframework.ai.bedrock.titan.api.TitanEmbeddingBedrockApi.TitanEmbeddingModel;
+import org.springframework.ai.embedding.EmbeddingRequest;
 import org.springframework.ai.embedding.EmbeddingResponse;
 import org.springframework.beans.factory.annotation.Autowired;
 import org.springframework.boot.SpringBootConfiguration;
 import org.springframework.boot.test.context.SpringBootTest;
 import org.springframework.context.annotation.Bean;
 import org.springframework.core.io.DefaultResourceLoader;
 
+import com.fasterxml.jackson.databind.ObjectMapper;
+
 import static org.assertj.core.api.Assertions.assertThat;
 
 @SpringBootTest
@@ -46,7 +52,8 @@ class BedrockTitanEmbeddingClientIT {
 	@Test
 	void singleEmbedding() {
 		assertThat(embeddingClient).isNotNull();
-		EmbeddingResponse embeddingResponse = embeddingClient.embedForResponse(List.of("Hello World"));
+		EmbeddingResponse embeddingResponse = embeddingClient.call(new EmbeddingRequest(List.of("Hello World"),
+				BedrockTitanEmbeddingOptions.builder().withInputType(InputType.TEXT).build()));
 		assertThat(embeddingResponse.getResults()).hasSize(1);
 		assertThat(embeddingResponse.getResults().get(0).getOutput()).isNotEmpty();
 		assertThat(embeddingClient.dimensions()).isEqualTo(1024);
@@ -59,7 +66,8 @@ void imageEmbedding() throws IOException {
 			.getContentAsByteArray();
 
 		EmbeddingResponse embeddingResponse = embeddingClient
-			.embedForResponse(List.of(Base64.getEncoder().encodeToString(image)));
+			.call(new EmbeddingRequest(List.of(Base64.getEncoder().encodeToString(image)),
+					BedrockTitanEmbeddingOptions.builder().withInputType(InputType.IMAGE).build()));
 		assertThat(embeddingResponse.getResults()).hasSize(1);
 		assertThat(embeddingResponse.getResults().get(0).getOutput()).isNotEmpty();
 		assertThat(embeddingClient.dimensions()).isEqualTo(1024);
@@ -70,7 +78,8 @@ public static class TestConfiguration {
 
 		@Bean
 		public TitanEmbeddingBedrockApi titanEmbeddingApi() {
-			return new TitanEmbeddingBedrockApi(TitanEmbeddingModel.TITAN_EMBED_IMAGE_V1.id(), Region.US_EAST_1.id(),
+			return new TitanEmbeddingBedrockApi(TitanEmbeddingModel.TITAN_EMBED_IMAGE_V1.id(),
+					EnvironmentVariableCredentialsProvider.create(), Region.US_EAST_1.id(), new ObjectMapper(),
 					Duration.ofMinutes(2));
 		}
 

spring-ai-docs/src/main/antora/modules/ROOT/pages/api/embeddings/bedrock-titan-embedding.adoc

@@ -81,6 +81,22 @@ The prefix `spring.ai.bedrock.titan.embedding` (defined in `BedrockTitanEmbeddin
 Supported values are: `amazon.titan-embed-image-v1` and `amazon.titan-embed-text-v1`.
 Model ID values can also be found in the https://docs.aws.amazon.com/bedrock/latest/userguide/model-ids-arns.html[AWS Bedrock documentation for base model IDs].
 
+== Runtime Options [[embedding-options]]
+
+The https://github.com/spring-projects/spring-ai/blob/main/models/spring-ai-bedrock/src/main/java/org/springframework/ai/bedrock/titan/BedrockTitanEmbeddingOptions.java[BedrockTitanEmbeddingOptions.java] provides model configurations, such as `input-type`.
+On start-up, the default options can be configured with the `BedrockTitanEmbeddingClient(api).withInputType(type)` method or the `spring.ai.bedrock.titan.embedding.input-type` properties.
+
+At run-time you can override the default options by adding new, request specific, options to the `EmbeddingRequest` call.
+For example to override the default temperature for a specific request:
+
+[source,java]
+----
+EmbeddingResponse embeddingResponse = embeddingClient.call(
+    new EmbeddingRequest(List.of("Hello World", "World is big and salvation is near"),
+        BedrockTitanEmbeddingOptions.builder()
+        	.withInputType(InputType.TEXT)
+        .build()));
+----
 
 == Sample Controller
 
@@ -154,7 +170,7 @@ Next, create an https://github.com/spring-projects/spring-ai/blob/main/models/sp
 var titanEmbeddingApi = new TitanEmbeddingBedrockApi(
 	TitanEmbeddingModel.TITAN_EMBED_IMAGE_V1.id(), Region.US_EAST_1.id());
 
-var embeddingClient  new BedrockTitanEmbeddingClient(titanEmbeddingApi);
+var embeddingClient = new BedrockTitanEmbeddingClient(titanEmbeddingApi);
 
 EmbeddingResponse embeddingResponse = embeddingClient
 	.embedForResponse(List.of("Hello World")); // NOTE titan does not support batch embedding.