diff --git a/spring-ai-bedrock/README.md b/spring-ai-bedrock/README.md index 492b62b03..c61e82c0d 100644 --- a/spring-ai-bedrock/README.md +++ b/spring-ai-bedrock/README.md @@ -79,5 +79,5 @@ Next you can use the `spring.ai.bedrock...*` properties t * [Spring AI Bedrock Cohere Embedding](./README_COHERE_EMBEDDING.md) - `spring.ai.bedrock.cohere.embedding=true` * [Spring AI Bedrock Anthropic Chat](./README_ANTHROPIC_CHAT.md) - `spring.ai.bedrock.anthropic.chat=true` * (WIP) [Spring AI Bedrock Titan Chat](./README_TITAN_CHAT.md) - `spring.ai.bedrock.titan.chat=true` -* (WIP) [Spring AI Bedrock Titan Embedding](./README_TITAN_EMBEDING.md) - `spring.ai.bedrock.titan.embedding=true` +* [Spring AI Bedrock Titan Embedding](./README_TITAN_EMBEDDING.md) - `spring.ai.bedrock.titan.embedding=true` * (WIP) [Spring AI Bedrock Ai21 Jurassic2 Chat](./README_JURASSIC2_CHAT.md) - `spring.ai.bedrock.jurassic2.chat=true` diff --git a/spring-ai-bedrock/README_ANTHROPIC_CHAT.md b/spring-ai-bedrock/README_ANTHROPIC_CHAT.md index a8e6b5fcd..c87de34e7 100644 --- a/spring-ai-bedrock/README_ANTHROPIC_CHAT.md +++ b/spring-ai-bedrock/README_ANTHROPIC_CHAT.md @@ -82,8 +82,10 @@ Use the `BedrockAnthropicChatProperties` to configure the Bedrock Llama2 Chat cl | Property | Description | Default | | ------------- | ------------- | ------------- | +| spring.ai.bedrock.embedding.aws.region | AWS region to use. | us-east-1 | +| spring.ai.bedrock.embedding.aws.accessKey | AWS credentials access key. | | +| spring.ai.bedrock.embedding.aws.secretKey | AWS credentials secret key. | | | spring.ai.bedrock.anthropic.chat.enable | Enable Bedrock Llama2 chat client. Disabled by default | false | -| spring.ai.bedrock.anthropic.chat.awsRegion | AWS region to use. | us-east-1 | | spring.ai.bedrock.anthropic.chat.temperature | Controls the randomness of the output. Values can range over [0.0,1.0] | 0.8 | | spring.ai.bedrock.anthropic.chat.topP | The maximum cumulative probability of tokens to consider when sampling. | AWS Bedrock default | | spring.ai.bedrock.anthropic.chat.maxGenLen | Specify the maximum number of tokens to use in the generated response. | 300 | diff --git a/spring-ai-bedrock/README_COHERE_CHAT.md b/spring-ai-bedrock/README_COHERE_CHAT.md index 37e64a506..75b9d1466 100644 --- a/spring-ai-bedrock/README_COHERE_CHAT.md +++ b/spring-ai-bedrock/README_COHERE_CHAT.md @@ -92,8 +92,10 @@ Use the `BedrockCohereChatProperties` to configure the Bedrock Cohere Chat clien | Property | Description | Default | | ------------- | ------------- | ------------- | +| spring.ai.bedrock.embedding.aws.region | AWS region to use. | us-east-1 | +| spring.ai.bedrock.embedding.aws.accessKey | AWS credentials access key. | | +| spring.ai.bedrock.embedding.aws.secretKey | AWS credentials secret key. | | | spring.ai.bedrock.cohere.chat.enable | Enable Bedrock Cohere chat client. Disabled by default | false | -| spring.ai.bedrock.cohere.chat.awsRegion | AWS region to use. | us-east-1 | | spring.ai.bedrock.cohere.chat.model | The model id to use. See the `CohereChatModel` for the supported models. | cohere.command-text-v14 | | spring.ai.bedrock.cohere.chat.temperature | Controls the randomness of the output. Values can range over [0.0,1.0] | 0.7 | | spring.ai.bedrock.cohere.chat.topP | The maximum cumulative probability of tokens to consider when sampling. | AWS Bedrock default | diff --git a/spring-ai-bedrock/README_COHERE_EMBEDDING.md b/spring-ai-bedrock/README_COHERE_EMBEDDING.md index 3335c040f..3c076f024 100644 --- a/spring-ai-bedrock/README_COHERE_EMBEDDING.md +++ b/spring-ai-bedrock/README_COHERE_EMBEDDING.md @@ -58,15 +58,17 @@ or you can leverage the `spring-ai-bedrock-ai-spring-boot-starter` Boot starter. ``` -**NOTE:** You have to enable the Bedrock Cohere chat client with `spring.ai.bedrock.cohere.embedding.enabled=true`. +**NOTE:** You have to enable the Bedrock Cohere embedding client with `spring.ai.bedrock.cohere.embedding.enabled=true`. By default the client is disabled. Use the `BedrockCohereEmbeddingProperties` to configure the Bedrock Cohere Chat client: | Property | Description | Default | | ------------- | ------------- | ------------- | -| spring.ai.bedrock.cohere.embedding.enable | Enable Bedrock Cohere chat client. Disabled by default | false | -| spring.ai.bedrock.cohere.embedding.awsRegion | AWS region to use. | us-east-1 | +| spring.ai.bedrock.embedding.aws.region | AWS region to use. | us-east-1 | +| spring.ai.bedrock.embedding.aws.accessKey | AWS credentials access key. | | +| spring.ai.bedrock.embedding.aws.secretKey | AWS credentials secret key. | | +| spring.ai.bedrock.cohere.embedding.enable | Enable Bedrock Cohere embedding client. Disabled by default | false | | spring.ai.bedrock.cohere.embedding.model | The model id to use. See the `CohereEmbeddingModel` for the supported models. | cohere.embed-multilingual-v3 | | spring.ai.bedrock.cohere.embedding.inputType | Prepends special tokens to differentiate each type from one another. You should not mix different types together, except when mixing types for for search and retrieval. In this case, embed your corpus with the search_document type and embedded queries with type search_query type. | search_document | | spring.ai.bedrock.cohere.embedding.truncate | Specifies how the API handles inputs longer than the maximum token length. | NONE | diff --git a/spring-ai-bedrock/README_LLAMA2_CHAT.md b/spring-ai-bedrock/README_LLAMA2_CHAT.md index 9c448bac9..d7234a644 100644 --- a/spring-ai-bedrock/README_LLAMA2_CHAT.md +++ b/spring-ai-bedrock/README_LLAMA2_CHAT.md @@ -80,8 +80,10 @@ Use the `BedrockLlama2ChatProperties` to configure the Bedrock Llama2 Chat clien | Property | Description | Default | | ------------- | ------------- | ------------- | +| spring.ai.bedrock.embedding.aws.region | AWS region to use. | us-east-1 | +| spring.ai.bedrock.embedding.aws.accessKey | AWS credentials access key. | | +| spring.ai.bedrock.embedding.aws.secretKey | AWS credentials secret key. | | | spring.ai.bedrock.llama2.chat.enable | Enable Bedrock Llama2 chat client. Disabled by default | false | -| spring.ai.bedrock.llama2.chat.awsRegion | AWS region to use. | us-east-1 | | spring.ai.bedrock.llama2.chat.temperature | Controls the randomness of the output. Values can range over [0.0,1.0] | 0.7 | | spring.ai.bedrock.llama2.chat.topP | The maximum cumulative probability of tokens to consider when sampling. | AWS Bedrock default | | spring.ai.bedrock.llama2.chat.maxGenLen | Specify the maximum number of tokens to use in the generated response. | 300 | diff --git a/spring-ai-bedrock/README_TITAN_EMBEDDING.md b/spring-ai-bedrock/README_TITAN_EMBEDDING.md index 0c41c4049..ca11f14dd 100644 --- a/spring-ai-bedrock/README_TITAN_EMBEDDING.md +++ b/spring-ai-bedrock/README_TITAN_EMBEDDING.md @@ -1,6 +1,95 @@ -# (WIP) Bedrock Titan Embedding +# 1. Bedrock Titan Embedding Use the [TitanEmbeddingBedrockApi.java](src/main/java/org/springframework/ai/bedrock/titan/api/TitanEmbeddingBedrockApi.java) Bedrock Embedding client to implement `EmbeddingClient`. Consult the the existing Cohere embedding client implementation. Mind that Titan doesn't support batch embedding. You have to either emulate it (could be very expensive) or throw a not supported exception. + +## 1.1 TitanEmbeddingBedrockApi + +[TitanEmbeddingBedrockApi](./src/main/java/org/springframework/ai/bedrock/titan/api/TitanEmbeddingBedrockApi.java) provides is lightweight Java client on top of AWS Bedrock [Titan Embedding models](https://docs.aws.amazon.com/bedrock/latest/userguide/titan-embedding-models.html). + +> TitanEmbeddingBedrockApi supports Text and Image embedding. + + +> **NOTE:** TitanEmbeddingBedrockApi does NOT support batch embedding. + +Following class diagram illustrates the Llama2ChatBedrockApi interface and building blocks: + +![TitanEmbeddingBedrockApi Class Diagram](./src/test/resources/doc/Bedrock%20Titan%20Embedding%20API.jpg) + +The CohereEmbeddingBedrockApi supports the `amazon.titan-embed-image-v1` and `amazon.titan-embed-image-v1` models for single and batch embedding computation. + +Here is a simple snippet how to use the api programmatically: + +```java +TitanEmbeddingBedrockApi titanEmbedApi = new TitanEmbeddingBedrockApi( + TitanEmbeddingModel.TITAN_EMBED_TEXT_V1.id(), Region.US_EAST_1.id()); + +TitanEmbeddingRequest request = TitanEmbeddingRequest.builder() + .withInputText("I like to eat apples.") + .build(); + +TitanEmbeddingResponse response = titanEmbedApi.embedding(request); +``` + +To embed an image you need to convert it into base64 format: + +```json +TitanEmbeddingBedrockApi titanEmbedApi = new TitanEmbeddingBedrockApi( + TitanEmbeddingModel.TITAN_EMBED_IMAGE_V1.id(), Region.US_EAST_1.id()); + +byte[] image = new DefaultResourceLoader() + .getResource("classpath:/spring_framework.png") + .getContentAsByteArray(); + + +TitanEmbeddingRequest request = TitanEmbeddingRequest.builder() + .withInputImage(Base64.getEncoder().encodeToString(image)) + .build(); + +TitanEmbeddingResponse response = titanEmbedApi.embedding(request); +``` + +## 1.2 BedrockTitanEmbeddingClient + +[BedrockTitanEmbeddingClient](./src/main/java/org/springframework/ai/bedrock/titan/BedrockTitanEmbeddingClient.java) implements the Spring-Ai `EmbeddingClient` on top of the `TitanEmbeddingBedrockApi`. + +You can use like this: + +```java +@Bean +public TitanEmbeddingBedrockApi titanEmbeddingApi() { + return new TitanEmbeddingBedrockApi( + TitanEmbeddingModel.TITAN_EMBED_IMAGE_V1.id(), Region.US_EAST_1.id()); +} + +@Bean +public BedrockTitanEmbeddingClient titanEmbedding(TitanEmbeddingBedrockApi titanEmbeddingApi) { + return new BedrockTitanEmbeddingClient(titanEmbeddingApi); +} +``` + +or you can leverage the `spring-ai-bedrock-ai-spring-boot-starter` Boot starter. For this add the following dependency: + +```xml + + spring-ai-bedrock-ai-spring-boot-starter + org.springframework.ai + 0.8.0-SNAPSHOT + +``` + +**NOTE:** You have to enable the Bedrock Titan embedding client with `spring.ai.bedrock.titan.embedding.enabled=true`. +By default the client is disabled. + +Use the `BedrockTitanEmbeddingProperties` to configure the Bedrock Titan embedding client: + +| Property | Description | Default | +| ------------- | ------------- | ------------- | +| spring.ai.bedrock.embedding.aws.region | AWS region to use. | us-east-1 | +| spring.ai.bedrock.embedding.aws.accessKey | AWS credentials access key. | | +| spring.ai.bedrock.embedding.aws.secretKey | AWS credentials secret key. | | +| spring.ai.bedrock.titan.embedding.enable | Enable Bedrock Titan embedding client. Disabled by default | false | +| spring.ai.bedrock.titan.embedding.model | The model id to use. See the `TitanEmbeddingModel` for the supported models. | amazon.titan-embed-image-v1 | +| spring.ai.bedrock.titan.embedding.inputType | Titan Embedding API input types. Could be either text or image (encoded in base64). | TEXT | diff --git a/spring-ai-bedrock/src/main/java/org/springframework/ai/bedrock/api/AbstractBedrockApi.java b/spring-ai-bedrock/src/main/java/org/springframework/ai/bedrock/api/AbstractBedrockApi.java index a48baab8a..8db590aa8 100644 --- a/spring-ai-bedrock/src/main/java/org/springframework/ai/bedrock/api/AbstractBedrockApi.java +++ b/spring-ai-bedrock/src/main/java/org/springframework/ai/bedrock/api/AbstractBedrockApi.java @@ -109,6 +109,20 @@ public abstract class AbstractBedrockApi { .build(); } + /** + * @return The model id. + */ + public String getModelId() { + return modelId; + } + + /** + * @return The AWS region. + */ + public String getRegion() { + return region; + } + /** * Encapsulates the metrics about the model invocation. * https://docs.aws.amazon.com/bedrock/latest/userguide/model-parameters-claude.html diff --git a/spring-ai-bedrock/src/main/java/org/springframework/ai/bedrock/titan/BedrockTitanEmbeddingClient.java b/spring-ai-bedrock/src/main/java/org/springframework/ai/bedrock/titan/BedrockTitanEmbeddingClient.java new file mode 100644 index 000000000..58038fa31 --- /dev/null +++ b/spring-ai-bedrock/src/main/java/org/springframework/ai/bedrock/titan/BedrockTitanEmbeddingClient.java @@ -0,0 +1,119 @@ +/* + * Copyright 2023-2023 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 java.util.List; +import java.util.concurrent.atomic.AtomicInteger; + +import org.springframework.ai.bedrock.titan.api.TitanEmbeddingBedrockApi; +import org.springframework.ai.bedrock.titan.api.TitanEmbeddingBedrockApi.TitanEmbeddingRequest; +import org.springframework.ai.bedrock.titan.api.TitanEmbeddingBedrockApi.TitanEmbeddingResponse; +import org.springframework.ai.document.Document; +import org.springframework.ai.embedding.AbstractEmbeddingClient; +import org.springframework.ai.embedding.Embedding; +import org.springframework.ai.embedding.EmbeddingResponse; +import org.springframework.ai.embedding.EmbeddingUtil; +import org.springframework.util.Assert; + +/** + * {@link org.springframework.ai.embedding.EmbeddingClient} implementation that uses the + * Bedrock Titan Embedding API. Titan Embedding supports text and image (encoded in + * base64) inputs. + * + * Note: Titan Embedding does not support batch embedding. + * + * @author Christian Tzolov + * @since 0.8.0 + */ +public class BedrockTitanEmbeddingClient extends AbstractEmbeddingClient { + + private final TitanEmbeddingBedrockApi embeddingApi; + + public enum InputType { + + TEXT, IMAGE + + } + + /** + * Titan Embedding API input types. Could be either text or image (encoded in base64). + */ + private InputType inputType = InputType.TEXT; + + public BedrockTitanEmbeddingClient(TitanEmbeddingBedrockApi titanEmbeddingBedrockApi) { + this.embeddingApi = titanEmbeddingBedrockApi; + } + + /** + * Titan Embedding API input types. Could be either text or image (encoded in base64). + * @param inputType the input type to use. + */ + public BedrockTitanEmbeddingClient withInputType(InputType inputType) { + this.inputType = inputType; + return this; + } + + @Override + public List embed(String inputContent) { + return this.embed(List.of(inputContent)).iterator().next(); + } + + @Override + public List embed(Document document) { + return embed(document.getContent()); + } + + @Override + public EmbeddingResponse embedForResponse(List texts) { + var indexCounter = new AtomicInteger(0); + List embeddings = this.embed(texts) + .stream() + .map(e -> new Embedding(e, indexCounter.getAndIncrement())) + .toList(); + return new EmbeddingResponse(embeddings); + } + + @Override + public List> embed(List inputContents) { + Assert.notEmpty(inputContents, "At least one text is required!"); + Assert.isTrue(inputContents.size() == 1, "Titan Embedding does not support batch embedding!"); + + String inputContent = inputContents.iterator().next(); + + var request = (this.inputType == InputType.IMAGE) + ? new TitanEmbeddingRequest.Builder().withInputImage(inputContent).build() + : new TitanEmbeddingRequest.Builder().withInputText(inputContent).build(); + + TitanEmbeddingResponse response = this.embeddingApi.embedding(request); + + return List.of(response.embedding()); + } + + @Override + public int dimensions() { + if (this.inputType == InputType.IMAGE) { + if (this.embeddingDimensions.get() < 0) { + this.embeddingDimensions.set(EmbeddingUtil.dimensions(this, embeddingApi.getModelId(), + // small base64 encoded image + "iVBORw0KGgoAAAANSUhEUgAAAAEAAAABCAQAAAC1HAwCAAAAC0lEQVR42mNk+A8AAQUBAScY42YAAAAASUVORK5CYII=")); + } + } + return super.dimensions(); + + } + +} diff --git a/spring-ai-bedrock/src/test/java/org/springframework/ai/bedrock/titan/BedrockTitanEmbeddingClientIT.java b/spring-ai-bedrock/src/test/java/org/springframework/ai/bedrock/titan/BedrockTitanEmbeddingClientIT.java new file mode 100644 index 000000000..267d0eabe --- /dev/null +++ b/spring-ai-bedrock/src/test/java/org/springframework/ai/bedrock/titan/BedrockTitanEmbeddingClientIT.java @@ -0,0 +1,76 @@ +package org.springframework.ai.bedrock.titan; + +import java.io.IOException; +import java.util.Base64; +import java.util.List; + +import org.junit.jupiter.api.Test; +import org.junit.jupiter.api.condition.EnabledIfEnvironmentVariable; +import software.amazon.awssdk.regions.Region; + +import org.springframework.ai.bedrock.titan.api.TitanEmbeddingBedrockApi; +import org.springframework.ai.bedrock.titan.api.TitanEmbeddingBedrockApi.TitanEmbeddingModel; +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 static org.assertj.core.api.Assertions.assertThat; +import static org.assertj.core.api.Assertions.assertThatThrownBy; + +@SpringBootTest +@EnabledIfEnvironmentVariable(named = "AWS_ACCESS_KEY_ID", matches = ".*") +@EnabledIfEnvironmentVariable(named = "AWS_SECRET_ACCESS_KEY", matches = ".*") +class BedrockTitanEmbeddingClientIT { + + @Autowired + private BedrockTitanEmbeddingClient embeddingClient; + + @Test + void singleEmbedding() { + assertThat(embeddingClient).isNotNull(); + EmbeddingResponse embeddingResponse = embeddingClient.embedForResponse(List.of("Hello World")); + assertThat(embeddingResponse.getData()).hasSize(1); + assertThat(embeddingResponse.getData().get(0).getEmbedding()).isNotEmpty(); + assertThat(embeddingClient.dimensions()).isEqualTo(1024); + } + + @Test + void batchEmbedding() { + assertThatThrownBy( + () -> embeddingClient.embedForResponse(List.of("Hello World", "World is big and salvation is near"))) + .isInstanceOf(IllegalArgumentException.class) + .hasMessageContaining("Titan Embedding does not support batch embedding!"); + } + + @Test + void imageEmbedding() throws IOException { + + byte[] image = new DefaultResourceLoader().getResource("classpath:/spring_framework.png") + .getContentAsByteArray(); + + EmbeddingResponse embeddingResponse = embeddingClient + .embedForResponse(List.of(Base64.getEncoder().encodeToString(image))); + assertThat(embeddingResponse.getData()).hasSize(1); + assertThat(embeddingResponse.getData().get(0).getEmbedding()).isNotEmpty(); + assertThat(embeddingClient.dimensions()).isEqualTo(1024); + } + + @SpringBootConfiguration + public static class TestConfiguration { + + @Bean + public TitanEmbeddingBedrockApi titanEmbeddingApi() { + return new TitanEmbeddingBedrockApi(TitanEmbeddingModel.TITAN_EMBED_IMAGE_V1.id(), Region.US_EAST_1.id()); + } + + @Bean + public BedrockTitanEmbeddingClient titanEmbedding(TitanEmbeddingBedrockApi titanEmbeddingApi) { + return new BedrockTitanEmbeddingClient(titanEmbeddingApi); + } + + } + +} diff --git a/spring-ai-bedrock/src/test/resources/doc/Bedrock Titan Embedding API.jpg b/spring-ai-bedrock/src/test/resources/doc/Bedrock Titan Embedding API.jpg new file mode 100644 index 000000000..1fb9a7565 Binary files /dev/null and b/spring-ai-bedrock/src/test/resources/doc/Bedrock Titan Embedding API.jpg differ diff --git a/spring-ai-bedrock/src/test/resources/doc/Bedrock-Anthropic-Chat-API.jpg b/spring-ai-bedrock/src/test/resources/doc/Bedrock-Anthropic-Chat-API.jpg new file mode 100644 index 000000000..acd5dddde Binary files /dev/null and b/spring-ai-bedrock/src/test/resources/doc/Bedrock-Anthropic-Chat-API.jpg differ diff --git a/spring-ai-core/src/main/java/org/springframework/ai/embedding/AbstractEmbeddingClient.java b/spring-ai-core/src/main/java/org/springframework/ai/embedding/AbstractEmbeddingClient.java index c1bfb3698..63ee89021 100644 --- a/spring-ai-core/src/main/java/org/springframework/ai/embedding/AbstractEmbeddingClient.java +++ b/spring-ai-core/src/main/java/org/springframework/ai/embedding/AbstractEmbeddingClient.java @@ -26,12 +26,12 @@ import java.util.concurrent.atomic.AtomicInteger; */ public abstract class AbstractEmbeddingClient implements EmbeddingClient { - private final AtomicInteger embeddingDimensions = new AtomicInteger(-1); + protected final AtomicInteger embeddingDimensions = new AtomicInteger(-1); @Override public int dimensions() { if (this.embeddingDimensions.get() < 0) { - this.embeddingDimensions.set(EmbeddingUtil.dimensions(this, "Test")); + this.embeddingDimensions.set(EmbeddingUtil.dimensions(this, "Test", "Hello World")); } return this.embeddingDimensions.get(); } diff --git a/spring-ai-core/src/main/java/org/springframework/ai/embedding/EmbeddingUtil.java b/spring-ai-core/src/main/java/org/springframework/ai/embedding/EmbeddingUtil.java index 7cd1e7582..451cea484 100644 --- a/spring-ai-core/src/main/java/org/springframework/ai/embedding/EmbeddingUtil.java +++ b/spring-ai-core/src/main/java/org/springframework/ai/embedding/EmbeddingUtil.java @@ -36,9 +36,10 @@ public class EmbeddingUtil { * the response dimensions. * @param embeddingClient Fall-back client to determine, empirically the dimensions. * @param modelName Embedding model name to retrieve the dimensions for. + * @param dummyContent Dummy content to use for the empirical dimension calculation. * @return Returns the embedding dimensions for the modelName. */ - public static int dimensions(EmbeddingClient embeddingClient, String modelName) { + public static int dimensions(EmbeddingClient embeddingClient, String modelName, String dummyContent) { if (KNOWN_EMBEDDING_DIMENSIONS.containsKey(modelName)) { // Retrieve the dimension from a pre-configured file. @@ -47,7 +48,7 @@ public class EmbeddingUtil { else { // Determine the dimensions empirically. // Generate an embedding and count the dimension size; - return embeddingClient.embed("Test String").size(); + return embeddingClient.embed(dummyContent).size(); } } diff --git a/spring-ai-core/src/test/java/org/springframework/ai/embedding/EmbeddingUtilTests.java b/spring-ai-core/src/test/java/org/springframework/ai/embedding/EmbeddingUtilTests.java index 92cea53ea..bfd52b167 100644 --- a/spring-ai-core/src/test/java/org/springframework/ai/embedding/EmbeddingUtilTests.java +++ b/spring-ai-core/src/test/java/org/springframework/ai/embedding/EmbeddingUtilTests.java @@ -75,7 +75,8 @@ public class EmbeddingUtilTests { @ParameterizedTest @CsvFileSource(resources = "/embedding/embedding-model-dimensions.properties", numLinesToSkip = 1, delimiter = '=') public void testKnownEmbeddingModelDimensions(String model, String dimension) { - assertThat(EmbeddingUtil.dimensions(embeddingClient, model)).isEqualTo(Integer.valueOf(dimension)); + assertThat(EmbeddingUtil.dimensions(embeddingClient, model, "Hello world!")) + .isEqualTo(Integer.valueOf(dimension)); verify(embeddingClient, never()).embed(any(String.class)); verify(embeddingClient, never()).embed(any(Document.class)); } @@ -83,7 +84,7 @@ public class EmbeddingUtilTests { @Test public void testUnknownModelDimension() { when(embeddingClient.embed(eq("Test String"))).thenReturn(List.of(0.1, 0.1, 0.1)); - assertThat(EmbeddingUtil.dimensions(embeddingClient, "unknown_model")).isEqualTo(3); + assertThat(EmbeddingUtil.dimensions(embeddingClient, "unknown_model", "Hello world!")).isEqualTo(3); } } diff --git a/spring-ai-spring-boot-autoconfigure/src/main/java/org/springframework/ai/autoconfigure/bedrock/titan/BedrockTitanEmbeddingAutoConfiguration.java b/spring-ai-spring-boot-autoconfigure/src/main/java/org/springframework/ai/autoconfigure/bedrock/titan/BedrockTitanEmbeddingAutoConfiguration.java new file mode 100644 index 000000000..c7f290cb1 --- /dev/null +++ b/spring-ai-spring-boot-autoconfigure/src/main/java/org/springframework/ai/autoconfigure/bedrock/titan/BedrockTitanEmbeddingAutoConfiguration.java @@ -0,0 +1,66 @@ +/* + * Copyright 2023-2023 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.autoconfigure.bedrock.titan; + +import com.fasterxml.jackson.databind.ObjectMapper; +import software.amazon.awssdk.auth.credentials.AwsCredentialsProvider; + +import org.springframework.ai.autoconfigure.NativeHints; +import org.springframework.ai.autoconfigure.bedrock.BedrockAwsConnectionConfiguration; +import org.springframework.ai.autoconfigure.bedrock.BedrockAwsConnectionProperties; +import org.springframework.ai.bedrock.titan.BedrockTitanEmbeddingClient; +import org.springframework.ai.bedrock.titan.api.TitanEmbeddingBedrockApi; +import org.springframework.boot.autoconfigure.AutoConfiguration; +import org.springframework.boot.autoconfigure.condition.ConditionalOnClass; +import org.springframework.boot.autoconfigure.condition.ConditionalOnMissingBean; +import org.springframework.boot.autoconfigure.condition.ConditionalOnProperty; +import org.springframework.boot.context.properties.EnableConfigurationProperties; +import org.springframework.context.annotation.Bean; +import org.springframework.context.annotation.Import; +import org.springframework.context.annotation.ImportRuntimeHints; + +/** + * {@link AutoConfiguration Auto-configuration} for Bedrock Titan Embedding Client. + * + * @author Christian Tzolov + * @since 0.8.0 + */ +@AutoConfiguration +@ConditionalOnClass(TitanEmbeddingBedrockApi.class) +@EnableConfigurationProperties({ BedrockTitanEmbeddingProperties.class, BedrockAwsConnectionProperties.class }) +@ConditionalOnProperty(prefix = BedrockTitanEmbeddingProperties.CONFIG_PREFIX, name = "enabled", havingValue = "true") +@Import(BedrockAwsConnectionConfiguration.class) +@ImportRuntimeHints(NativeHints.class) +public class BedrockTitanEmbeddingAutoConfiguration { + + @Bean + @ConditionalOnMissingBean + public TitanEmbeddingBedrockApi titanApi(AwsCredentialsProvider credentialsProvider, + BedrockTitanEmbeddingProperties properties, BedrockAwsConnectionProperties awsProperties) { + return new TitanEmbeddingBedrockApi(properties.getModel(), credentialsProvider, awsProperties.getRegion(), + new ObjectMapper()); + } + + @Bean + @ConditionalOnMissingBean + public BedrockTitanEmbeddingClient cohereEmbeddingClient(TitanEmbeddingBedrockApi titanEmbeddingApi, + BedrockTitanEmbeddingProperties properties) { + + return new BedrockTitanEmbeddingClient(titanEmbeddingApi).withInputType(properties.getInputType()); + } + +} diff --git a/spring-ai-spring-boot-autoconfigure/src/main/java/org/springframework/ai/autoconfigure/bedrock/titan/BedrockTitanEmbeddingProperties.java b/spring-ai-spring-boot-autoconfigure/src/main/java/org/springframework/ai/autoconfigure/bedrock/titan/BedrockTitanEmbeddingProperties.java new file mode 100644 index 000000000..f74026c81 --- /dev/null +++ b/spring-ai-spring-boot-autoconfigure/src/main/java/org/springframework/ai/autoconfigure/bedrock/titan/BedrockTitanEmbeddingProperties.java @@ -0,0 +1,78 @@ +/* + * Copyright 2023-2023 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.autoconfigure.bedrock.titan; + +import org.springframework.ai.bedrock.titan.BedrockTitanEmbeddingClient.InputType; +import org.springframework.ai.bedrock.titan.api.TitanEmbeddingBedrockApi.TitanEmbeddingModel; +import org.springframework.boot.context.properties.ConfigurationProperties; + +/** + * Bedrock Titan Embedding autoconfiguration properties. + * + * @author Christian Tzolov + * @since 0.8.0 + */ +@ConfigurationProperties(BedrockTitanEmbeddingProperties.CONFIG_PREFIX) +public class BedrockTitanEmbeddingProperties { + + public static final String CONFIG_PREFIX = "spring.ai.bedrock.titan.embedding"; + + /** + * Enable Bedrock Titan Embedding Client. False by default. + */ + private boolean enabled = false; + + /** + * Bedrock Titan Embedding model name. Defaults to 'amazon.titan-embed-image-v1'. + */ + private String model = TitanEmbeddingModel.TITAN_EMBED_IMAGE_V1.id(); + + /** + * Titan Embedding API input types. Could be either text or image (encoded in base64). + * Defaults to {@link InputType#IMAGE}. + */ + private InputType inputType = InputType.IMAGE; + + public boolean isEnabled() { + return enabled; + } + + public void setEnabled(boolean enabled) { + this.enabled = enabled; + } + + public String getModel() { + return model; + } + + public void setModel(String model) { + this.model = model; + } + + public static String getConfigPrefix() { + return CONFIG_PREFIX; + } + + public void setInputType(InputType inputType) { + this.inputType = inputType; + } + + public InputType getInputType() { + return inputType; + } + +} diff --git a/spring-ai-spring-boot-autoconfigure/src/test/java/org/springframework/ai/autoconfigure/bedrock/cohere/BedrockCohereEmbeddingAutoConfigurationIT.java b/spring-ai-spring-boot-autoconfigure/src/test/java/org/springframework/ai/autoconfigure/bedrock/cohere/BedrockCohereEmbeddingAutoConfigurationIT.java index afd747ddb..596487a77 100644 --- a/spring-ai-spring-boot-autoconfigure/src/test/java/org/springframework/ai/autoconfigure/bedrock/cohere/BedrockCohereEmbeddingAutoConfigurationIT.java +++ b/spring-ai-spring-boot-autoconfigure/src/test/java/org/springframework/ai/autoconfigure/bedrock/cohere/BedrockCohereEmbeddingAutoConfigurationIT.java @@ -47,8 +47,8 @@ public class BedrockCohereEmbeddingAutoConfigurationIT { "spring.ai.bedrock.aws.secret-key=" + System.getenv("AWS_SECRET_ACCESS_KEY"), "spring.ai.bedrock.aws.region=" + Region.US_EAST_1.id(), "spring.ai.bedrock.cohere.embedding.model=" + CohereEmbeddingModel.COHERE_EMBED_MULTILINGUAL_V1.id(), - "spring.ai.bedrock.cohere.chat.inputType=search_document", - "spring.ai.bedrock.cohere.chat.truncate=NONE") + "spring.ai.bedrock.cohere.embedding.inputType=search_document", + "spring.ai.bedrock.cohere.embedding.truncate=NONE") .withConfiguration(AutoConfigurations.of(BedrockCohereEmbeddingAutoConfiguration.class)); @Test @@ -95,15 +95,15 @@ public class BedrockCohereEmbeddingAutoConfigurationIT { "spring.ai.bedrock.cohere.embedding.truncate=RIGHT") .withConfiguration(AutoConfigurations.of(BedrockCohereEmbeddingAutoConfiguration.class)) .run(context -> { - var chatProperties = context.getBean(BedrockCohereEmbeddingProperties.class); + var properties = context.getBean(BedrockCohereEmbeddingProperties.class); var awsProperties = context.getBean(BedrockAwsConnectionProperties.class); - assertThat(chatProperties.isEnabled()).isTrue(); + assertThat(properties.isEnabled()).isTrue(); assertThat(awsProperties.getRegion()).isEqualTo(Region.EU_CENTRAL_1.id()); - assertThat(chatProperties.getModel()).isEqualTo("MODEL_XYZ"); + assertThat(properties.getModel()).isEqualTo("MODEL_XYZ"); - assertThat(chatProperties.getInputType()).isEqualTo(InputType.classification); - assertThat(chatProperties.getTruncate()).isEqualTo(CohereEmbeddingRequest.Truncate.RIGHT); + assertThat(properties.getInputType()).isEqualTo(InputType.classification); + assertThat(properties.getTruncate()).isEqualTo(CohereEmbeddingRequest.Truncate.RIGHT); assertThat(awsProperties.getAccessKey()).isEqualTo("ACCESS_KEY"); assertThat(awsProperties.getSecretKey()).isEqualTo("SECRET_KEY"); @@ -111,7 +111,7 @@ public class BedrockCohereEmbeddingAutoConfigurationIT { } @Test - public void chatCompletionDisabled() { + public void embeddingDisabled() { // It is disabled by default new ApplicationContextRunner() @@ -121,7 +121,7 @@ public class BedrockCohereEmbeddingAutoConfigurationIT { assertThat(context.getBeansOfType(BedrockCohereEmbeddingClient.class)).isEmpty(); }); - // Explicitly enable the chat auto-configuration. + // Explicitly enable the embedding auto-configuration. new ApplicationContextRunner().withPropertyValues("spring.ai.bedrock.cohere.embedding.enabled=true") .withConfiguration(AutoConfigurations.of(BedrockCohereEmbeddingAutoConfiguration.class)) .run(context -> { @@ -129,7 +129,7 @@ public class BedrockCohereEmbeddingAutoConfigurationIT { assertThat(context.getBeansOfType(BedrockCohereEmbeddingClient.class)).isNotEmpty(); }); - // Explicitly disable the chat auto-configuration. + // Explicitly disable the embedding auto-configuration. new ApplicationContextRunner().withPropertyValues("spring.ai.bedrock.cohere.embedding.enabled=false") .withConfiguration(AutoConfigurations.of(BedrockCohereEmbeddingAutoConfiguration.class)) .run(context -> { diff --git a/spring-ai-spring-boot-autoconfigure/src/test/java/org/springframework/ai/autoconfigure/bedrock/titan/BedrockTitanEmbeddingAutoConfigurationIT.java b/spring-ai-spring-boot-autoconfigure/src/test/java/org/springframework/ai/autoconfigure/bedrock/titan/BedrockTitanEmbeddingAutoConfigurationIT.java new file mode 100644 index 000000000..591699905 --- /dev/null +++ b/spring-ai-spring-boot-autoconfigure/src/test/java/org/springframework/ai/autoconfigure/bedrock/titan/BedrockTitanEmbeddingAutoConfigurationIT.java @@ -0,0 +1,135 @@ +/* + * Copyright 2023-2023 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.autoconfigure.bedrock.titan; + +import java.util.Base64; +import java.util.List; + +import org.junit.jupiter.api.Test; +import org.junit.jupiter.api.condition.EnabledIfEnvironmentVariable; +import software.amazon.awssdk.regions.Region; + +import org.springframework.ai.autoconfigure.bedrock.BedrockAwsConnectionProperties; +import org.springframework.ai.bedrock.titan.BedrockTitanEmbeddingClient; +import org.springframework.ai.bedrock.titan.BedrockTitanEmbeddingClient.InputType; +import org.springframework.ai.bedrock.titan.api.TitanEmbeddingBedrockApi.TitanEmbeddingModel; +import org.springframework.ai.embedding.EmbeddingResponse; +import org.springframework.boot.autoconfigure.AutoConfigurations; +import org.springframework.boot.test.context.runner.ApplicationContextRunner; +import org.springframework.core.io.DefaultResourceLoader; + +import static org.assertj.core.api.Assertions.assertThat; + +/** + * @author Christian Tzolov + * @since 0.8.0 + */ +@EnabledIfEnvironmentVariable(named = "AWS_ACCESS_KEY_ID", matches = ".*") +@EnabledIfEnvironmentVariable(named = "AWS_SECRET_ACCESS_KEY", matches = ".*") +public class BedrockTitanEmbeddingAutoConfigurationIT { + + private final ApplicationContextRunner contextRunner = new ApplicationContextRunner() + .withPropertyValues("spring.ai.bedrock.titan.embedding.enabled=true", + "spring.ai.bedrock.aws.access-key=" + System.getenv("AWS_ACCESS_KEY_ID"), + "spring.ai.bedrock.aws.secret-key=" + System.getenv("AWS_SECRET_ACCESS_KEY"), + "spring.ai.bedrock.aws.region=" + Region.US_EAST_1.id(), + "spring.ai.bedrock.titan.embedding.model=" + TitanEmbeddingModel.TITAN_EMBED_IMAGE_V1.id()) + .withConfiguration(AutoConfigurations.of(BedrockTitanEmbeddingAutoConfiguration.class)); + + @Test + public void singleTextEmbedding() { + contextRunner.withPropertyValues("spring.ai.bedrock.titan.embedding.inputType=TEXT").run(context -> { + BedrockTitanEmbeddingClient embeddingClient = context.getBean(BedrockTitanEmbeddingClient.class); + assertThat(embeddingClient).isNotNull(); + EmbeddingResponse embeddingResponse = embeddingClient.embedForResponse(List.of("Hello World")); + assertThat(embeddingResponse.getData()).hasSize(1); + assertThat(embeddingResponse.getData().get(0).getEmbedding()).isNotEmpty(); + assertThat(embeddingClient.dimensions()).isEqualTo(1024); + }); + } + + @Test + public void singleImageEmbedding() { + contextRunner.withPropertyValues("spring.ai.bedrock.titan.embedding.inputType=IMAGE").run(context -> { + BedrockTitanEmbeddingClient embeddingClient = context.getBean(BedrockTitanEmbeddingClient.class); + assertThat(embeddingClient).isNotNull(); + + byte[] image = new DefaultResourceLoader().getResource("classpath:/spring_framework.png") + .getContentAsByteArray(); + + var base64Image = Base64.getEncoder().encodeToString(image); + + EmbeddingResponse embeddingResponse = embeddingClient.embedForResponse(List.of(base64Image)); + + assertThat(embeddingResponse.getData()).hasSize(1); + assertThat(embeddingResponse.getData().get(0).getEmbedding()).isNotEmpty(); + assertThat(embeddingClient.dimensions()).isEqualTo(1024); + }); + } + + @Test + public void propertiesTest() { + + new ApplicationContextRunner().withPropertyValues("spring.ai.bedrock.titan.embedding.enabled=true", + "spring.ai.bedrock.aws.access-key=ACCESS_KEY", "spring.ai.bedrock.aws.secret-key=SECRET_KEY", + "spring.ai.bedrock.aws.region=" + Region.EU_CENTRAL_1.id(), + "spring.ai.bedrock.titan.embedding.model=MODEL_XYZ", "spring.ai.bedrock.titan.embedding.inputType=TEXT") + .withConfiguration(AutoConfigurations.of(BedrockTitanEmbeddingAutoConfiguration.class)) + .run(context -> { + var properties = context.getBean(BedrockTitanEmbeddingProperties.class); + var awsProperties = context.getBean(BedrockAwsConnectionProperties.class); + + assertThat(properties.isEnabled()).isTrue(); + assertThat(awsProperties.getRegion()).isEqualTo(Region.EU_CENTRAL_1.id()); + assertThat(properties.getModel()).isEqualTo("MODEL_XYZ"); + + assertThat(properties.getInputType()).isEqualTo(InputType.TEXT); + + assertThat(awsProperties.getAccessKey()).isEqualTo("ACCESS_KEY"); + assertThat(awsProperties.getSecretKey()).isEqualTo("SECRET_KEY"); + }); + } + + @Test + public void embeddingDisabled() { + + // It is disabled by default + new ApplicationContextRunner() + .withConfiguration(AutoConfigurations.of(BedrockTitanEmbeddingAutoConfiguration.class)) + .run(context -> { + assertThat(context.getBeansOfType(BedrockTitanEmbeddingProperties.class)).isEmpty(); + assertThat(context.getBeansOfType(BedrockTitanEmbeddingClient.class)).isEmpty(); + }); + + // Explicitly enable the embedding auto-configuration. + new ApplicationContextRunner().withPropertyValues("spring.ai.bedrock.titan.embedding.enabled=true") + .withConfiguration(AutoConfigurations.of(BedrockTitanEmbeddingAutoConfiguration.class)) + .run(context -> { + assertThat(context.getBeansOfType(BedrockTitanEmbeddingProperties.class)).isNotEmpty(); + assertThat(context.getBeansOfType(BedrockTitanEmbeddingClient.class)).isNotEmpty(); + }); + + // Explicitly disable the embedding auto-configuration. + new ApplicationContextRunner().withPropertyValues("spring.ai.bedrock.titan.embedding.enabled=false") + .withConfiguration(AutoConfigurations.of(BedrockTitanEmbeddingAutoConfiguration.class)) + .run(context -> { + assertThat(context.getBeansOfType(BedrockTitanEmbeddingProperties.class)).isEmpty(); + assertThat(context.getBeansOfType(BedrockTitanEmbeddingClient.class)).isEmpty(); + }); + } + +} diff --git a/spring-ai-spring-boot-autoconfigure/src/test/resources/spring_framework.png b/spring-ai-spring-boot-autoconfigure/src/test/resources/spring_framework.png new file mode 100644 index 000000000..735057a6b Binary files /dev/null and b/spring-ai-spring-boot-autoconfigure/src/test/resources/spring_framework.png differ