From a50969ec8b029676e26d9753eccfe83420f3a696 Mon Sep 17 00:00:00 2001 From: wmz7year Date: Mon, 22 Apr 2024 12:31:25 +0800 Subject: [PATCH] Bedrock Titan embedding client adds BedrockTitanEmbeddingOptions to support dynamic embedding request types. --- .../ai/bedrock/aot/BedrockRuntimeHints.java | 4 ++ .../titan/BedrockTitanEmbeddingClient.java | 18 ++++- .../titan/BedrockTitanEmbeddingOptions.java | 65 +++++++++++++++++++ .../titan/BedrockTitanEmbeddingClientIT.java | 15 ++++- .../embeddings/bedrock-titan-embedding.adoc | 18 ++++- 5 files changed, 113 insertions(+), 7 deletions(-) create mode 100644 models/spring-ai-bedrock/src/main/java/org/springframework/ai/bedrock/titan/BedrockTitanEmbeddingOptions.java diff --git a/models/spring-ai-bedrock/src/main/java/org/springframework/ai/bedrock/aot/BedrockRuntimeHints.java b/models/spring-ai-bedrock/src/main/java/org/springframework/ai/bedrock/aot/BedrockRuntimeHints.java index edb3c4b13..7db24b3b8 100644 --- a/models/spring-ai-bedrock/src/main/java/org/springframework/ai/bedrock/aot/BedrockRuntimeHints.java +++ b/models/spring-ai-bedrock/src/main/java/org/springframework/ai/bedrock/aot/BedrockRuntimeHints.java @@ -28,6 +28,7 @@ import org.springframework.ai.bedrock.jurassic2.api.Ai21Jurassic2ChatBedrockApi; 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 @@ import static org.springframework.ai.aot.AiRuntimeHints.findJsonAnnotatedClasses * @author Josh Long * @author Christian Tzolov * @author Mark Pollack + * @author Wei Jiang */ public class BedrockRuntimeHints implements RuntimeHintsRegistrar { @@ -72,6 +74,8 @@ public class BedrockRuntimeHints implements RuntimeHintsRegistrar { 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); diff --git a/models/spring-ai-bedrock/src/main/java/org/springframework/ai/bedrock/titan/BedrockTitanEmbeddingClient.java b/models/spring-ai-bedrock/src/main/java/org/springframework/ai/bedrock/titan/BedrockTitanEmbeddingClient.java index d48135f80..1d64f92ef 100644 --- a/models/spring-ai-bedrock/src/main/java/org/springframework/ai/bedrock/titan/BedrockTitanEmbeddingClient.java +++ b/models/spring-ai-bedrock/src/main/java/org/springframework/ai/bedrock/titan/BedrockTitanEmbeddingClient.java @@ -28,6 +28,7 @@ import org.springframework.ai.bedrock.titan.api.TitanEmbeddingBedrockApi.TitanEm 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 @@ import org.springframework.util.Assert; * 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 class BedrockTitanEmbeddingClient extends AbstractEmbeddingClient { List> 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 class BedrockTitanEmbeddingClient extends AbstractEmbeddingClient { 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) { diff --git a/models/spring-ai-bedrock/src/main/java/org/springframework/ai/bedrock/titan/BedrockTitanEmbeddingOptions.java b/models/spring-ai-bedrock/src/main/java/org/springframework/ai/bedrock/titan/BedrockTitanEmbeddingOptions.java new file mode 100644 index 000000000..fd1c609bf --- /dev/null +++ b/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; + } + +} diff --git a/models/spring-ai-bedrock/src/test/java/org/springframework/ai/bedrock/titan/BedrockTitanEmbeddingClientIT.java b/models/spring-ai-bedrock/src/test/java/org/springframework/ai/bedrock/titan/BedrockTitanEmbeddingClientIT.java index dead75901..6400a15f0 100644 --- a/models/spring-ai-bedrock/src/test/java/org/springframework/ai/bedrock/titan/BedrockTitanEmbeddingClientIT.java +++ b/models/spring-ai-bedrock/src/test/java/org/springframework/ai/bedrock/titan/BedrockTitanEmbeddingClientIT.java @@ -22,10 +22,14 @@ import java.util.List; 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; @@ -33,6 +37,8 @@ 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 @@ class BedrockTitanEmbeddingClientIT { .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 @@ class BedrockTitanEmbeddingClientIT { @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)); } diff --git a/spring-ai-docs/src/main/antora/modules/ROOT/pages/api/embeddings/bedrock-titan-embedding.adoc b/spring-ai-docs/src/main/antora/modules/ROOT/pages/api/embeddings/bedrock-titan-embedding.adoc index 6d5d675a0..b7bc8a74e 100644 --- a/spring-ai-docs/src/main/antora/modules/ROOT/pages/api/embeddings/bedrock-titan-embedding.adoc +++ b/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.