Spring/spring-ai
1
0
Fork 0
Code Issues Packages Projects Releases Wiki Activity

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

Browse Source
This commit is contained in:
wmz7year
2024-04-22 12:31:25 +08:00
parent 5f9ecdd899
commit a50969ec8b
5 changed files with 113 additions and 7 deletions
Download Patch File Download Diff File Expand all files Collapse all files

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

@@ -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);

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

@@ -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<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 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) {

65
models/spring-ai-bedrock/src/main/java/org/springframework/ai/bedrock/titan/BedrockTitanEmbeddingOptions.java Normal file
View File

@@ -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;
}
}

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

@@ -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));
}

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

@@ -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.