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

Add Bedrock Cohere Embedding Options support

Browse Source 
- Realign the Embeddings API docs strcture
This commit is contained in:
Christian Tzolov
2024-02-11 12:25:21 +01:00
parent 056b95a00e
commit 35ce0d1a4b
22 changed files with 507 additions and 248 deletions
Download Patch File Download Diff File Expand all files Collapse all files

74
models/spring-ai-bedrock/README_COHERE_EMBEDDING.md
View File

@@ -1,74 +1,4 @@
# 1. Bedrock Cohere Embedding
# Bedrock Cohere Embedding
## 1.1 CohereEmbeddingBedrockApi
Visit the Spring AI [Bedrock Cohere Embedding Documentation](https://docs.spring.io/spring-ai/reference/api/embeddings/bedrock-cohere-embedding.html).
[CohereEmbeddingBedrockApi](./src/main/java/org/springframework/ai/bedrock/cohere/api/CohereEmbeddingBedrockApi.java) provides is lightweight Java client on top of AWS Bedrock [Cohere Embed models](https://docs.aws.amazon.com/bedrock/latest/userguide/model-parameters-embed.html).
Following class diagram illustrates the Llama2ChatBedrockApi interface and building blocks:
![CohereEmbeddingBedrockApi Class Diagram](./src/test/resources/doc/Bedrock%20Cohere%20Embedding%20API.jpg)
The CohereEmbeddingBedrockApi supports the `cohere.embed-english-v3` and `cohere.embed-multilingual-v3` models for single and batch embedding computation.
Here is a simple snippet how to use the api programmatically:
```java
CohereEmbeddingBedrockApi api = new CohereEmbeddingBedrockApi(
CohereEmbeddingModel.COHERE_EMBED_MULTILINGUAL_V1.id(),
EnvironmentVariableCredentialsProvider.create(),
Region.US_EAST_1.id(), new ObjectMapper());
CohereEmbeddingRequest request = new CohereEmbeddingRequest(
List.of("I like to eat apples", "I like to eat oranges"),
CohereEmbeddingRequest.InputType.search_document,
CohereEmbeddingRequest.Truncate.NONE);
CohereEmbeddingResponse response = api.embedding(request);
assertThat(response.embeddings()).hasSize(2);
assertThat(response.embeddings().get(0)).hasSize(1024);
```
## 1.2 BedrockCohereEmbeddingClient
[BedrockCohereEmbeddingClient](./src/main/java/org/springframework/ai/bedrock/cohere/BedrockCohereEmbeddingClient.java) implements the Spring-Ai `EmbeddingClient` on top of the `CohereEmbeddingBedrockApi`.
You can use like this:
```java
@Bean
public CohereEmbeddingBedrockApi cohereEmbeddingApi() {
return new CohereEmbeddingBedrockApi(CohereEmbeddingModel.COHERE_EMBED_MULTILINGUAL_V1.id(),
EnvironmentVariableCredentialsProvider.create(), Region.US_EAST_1.id(), new ObjectMapper());
}
@Bean
public BedrockCohereEmbeddingClient cohereAiEmbedding(CohereEmbeddingBedrockApi cohereEmbeddingApi) {
return new BedrockCohereEmbeddingClient(cohereEmbeddingApi);
}
```
or you can leverage the `spring-ai-bedrock-ai-spring-boot-starter` Boot starter. For this add the following dependency:
```xml
<dependency>
<artifactId>spring-ai-bedrock-ai-spring-boot-starter</artifactId>
<groupId>org.springframework.ai</groupId>
<version>0.8.0-SNAPSHOT</version>
</dependency>
```
**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.aws.region | AWS region to use. | us-east-1 |
| spring.ai.bedrock.aws.accessKey | AWS credentials access key. | |
| spring.ai.bedrock.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 |

90
models/spring-ai-bedrock/src/main/java/org/springframework/ai/bedrock/cohere/BedrockCohereEmbeddingClient.java
View File

@@ -25,8 +25,10 @@ import org.springframework.ai.bedrock.cohere.api.CohereEmbeddingBedrockApi.Coher
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.ai.model.ModelOptionsUtils;
import org.springframework.util.Assert;
/**
@@ -41,35 +43,54 @@ public class BedrockCohereEmbeddingClient extends AbstractEmbeddingClient {
private final CohereEmbeddingBedrockApi embeddingApi;
private CohereEmbeddingRequest.InputType inputType = CohereEmbeddingRequest.InputType.SEARCH_DOCUMENT;
private final BedrockCohereEmbeddingOptions defaultOptions;
private CohereEmbeddingRequest.Truncate truncate = CohereEmbeddingRequest.Truncate.NONE;
// private CohereEmbeddingRequest.InputType inputType =
// CohereEmbeddingRequest.InputType.SEARCH_DOCUMENT;
// private CohereEmbeddingRequest.Truncate truncate =
// CohereEmbeddingRequest.Truncate.NONE;
public BedrockCohereEmbeddingClient(CohereEmbeddingBedrockApi cohereEmbeddingBedrockApi) {
this(cohereEmbeddingBedrockApi,
BedrockCohereEmbeddingOptions.builder()
.withInputType(CohereEmbeddingRequest.InputType.SEARCH_DOCUMENT)
.withTruncate(CohereEmbeddingRequest.Truncate.NONE)
.build());
}
public BedrockCohereEmbeddingClient(CohereEmbeddingBedrockApi cohereEmbeddingBedrockApi,
BedrockCohereEmbeddingOptions options) {
Assert.notNull(cohereEmbeddingBedrockApi, "CohereEmbeddingBedrockApi must not be null");
Assert.notNull(options, "BedrockCohereEmbeddingOptions must not be null");
this.embeddingApi = cohereEmbeddingBedrockApi;
this.defaultOptions = options;
}
/**
* Cohere Embedding API input types.
* @param inputType the input type to use.
* @return this client.
*/
public BedrockCohereEmbeddingClient withInputType(CohereEmbeddingRequest.InputType inputType) {
this.inputType = inputType;
return this;
}
// /**
// * Cohere Embedding API input types.
// * @param inputType the input type to use.
// * @return this client.
// */
// public BedrockCohereEmbeddingClient withInputType(CohereEmbeddingRequest.InputType
// inputType) {
// this.inputType = inputType;
// return this;
// }
/**
* Specifies how the API handles inputs longer than the maximum token length. If you
* specify LEFT or RIGHT, the model discards the input until the remaining input is
* exactly the maximum input token length for the model.
* @param truncate the truncate option to use.
* @return this client.
*/
public BedrockCohereEmbeddingClient withTruncate(CohereEmbeddingRequest.Truncate truncate) {
this.truncate = truncate;
return this;
}
// /**
// * Specifies how the API handles inputs longer than the maximum token length. If you
// specify LEFT or RIGHT, the
// * model discards the input until the remaining input is exactly the maximum input
// token length for the model.
// * @param truncate the truncate option to use.
// * @return this client.
// */
// public BedrockCohereEmbeddingClient withTruncate(CohereEmbeddingRequest.Truncate
// truncate) {
// this.truncate = truncate;
// return this;
// }
@Override
public List<Double> embed(Document document) {
@@ -80,7 +101,10 @@ public class BedrockCohereEmbeddingClient extends AbstractEmbeddingClient {
public EmbeddingResponse call(EmbeddingRequest request) {
Assert.notEmpty(request.getInstructions(), "At least one text is required!");
var apiRequest = new CohereEmbeddingRequest(request.getInstructions(), this.inputType, this.truncate);
final BedrockCohereEmbeddingOptions optionsToUse = this.mergeOptions(request.getOptions());
var apiRequest = new CohereEmbeddingRequest(request.getInstructions(), optionsToUse.getInputType(),
optionsToUse.getTruncate());
CohereEmbeddingResponse apiResponse = this.embeddingApi.embedding(apiRequest);
var indexCounter = new AtomicInteger(0);
List<Embedding> embeddings = apiResponse.embeddings()
@@ -90,4 +114,24 @@ public class BedrockCohereEmbeddingClient extends AbstractEmbeddingClient {
return new EmbeddingResponse(embeddings);
}
/**
* Merge the default and request options.
* @param requestOptions request options to merge.
* @return the merged options.
*/
BedrockCohereEmbeddingOptions mergeOptions(EmbeddingOptions requestOptions) {
BedrockCohereEmbeddingOptions options = (this.defaultOptions != null) ? this.defaultOptions
: BedrockCohereEmbeddingOptions.builder()
.withInputType(CohereEmbeddingRequest.InputType.SEARCH_DOCUMENT)
.withTruncate(CohereEmbeddingRequest.Truncate.NONE)
.build();
if (requestOptions != null && !EmbeddingOptions.EMPTY.equals(requestOptions)) {
options = ModelOptionsUtils.merge(requestOptions, options, BedrockCohereEmbeddingOptions.class);
}
return options;
}
}

89
models/spring-ai-bedrock/src/main/java/org/springframework/ai/bedrock/cohere/BedrockCohereEmbeddingOptions.java Normal file
View File

@@ -0,0 +1,89 @@
/*
* Copyright 2024-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.cohere;
import com.fasterxml.jackson.annotation.JsonInclude;
import com.fasterxml.jackson.annotation.JsonInclude.Include;
import com.fasterxml.jackson.annotation.JsonProperty;
import org.springframework.ai.bedrock.cohere.api.CohereEmbeddingBedrockApi.CohereEmbeddingRequest.InputType;
import org.springframework.ai.bedrock.cohere.api.CohereEmbeddingBedrockApi.CohereEmbeddingRequest.Truncate;
/**
* @author Christian Tzolov
*/
@JsonInclude(Include.NON_NULL)
public class BedrockCohereEmbeddingOptions {
// @formatter:off
/**
* 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.
*/
private @JsonProperty("input_type") InputType inputType;
/**
* Specifies how the API handles inputs longer than the maximum token length. If you specify LEFT or
* RIGHT, the model discards the input until the remaining input is exactly the maximum input token length for the
* model.
*/
private @JsonProperty("truncate") Truncate truncate;
// @formatter:on
public static Builder builder() {
return new Builder();
}
public static class Builder {
private BedrockCohereEmbeddingOptions options = new BedrockCohereEmbeddingOptions();
public Builder withInputType(InputType inputType) {
this.options.setInputType(inputType);
return this;
}
public Builder withTruncate(Truncate truncate) {
this.options.setTruncate(truncate);
return this;
}
public BedrockCohereEmbeddingOptions build() {
return this.options;
}
}
public InputType getInputType() {
return this.inputType;
}
public void setInputType(InputType inputType) {
this.inputType = inputType;
}
public Truncate getTruncate() {
return this.truncate;
}
public void setTruncate(Truncate truncate) {
this.truncate = truncate;
}
}

2
models/spring-ai-bedrock/src/main/java/org/springframework/ai/bedrock/cohere/api/CohereEmbeddingBedrockApi.java
View File

@@ -163,7 +163,7 @@ public class CohereEmbeddingBedrockApi extends
* @return The model id.
*/
public String id() {
return id;
return this.id;
}
CohereEmbeddingModel(String value) {