Skip to content

Commit 35ce0d1

Browse files
tzolovtzolov
committed
Add Bedrock Cohere Embedding Options support
 - Realign the Embeddings API docs strcture
1 parent 056b95a commit 35ce0d1

File tree

22 files changed

+507
-248
lines changed

22 files changed

+507
-248
lines changed

‎models/spring-ai-bedrock/README_COHERE_EMBEDDING.md

Lines changed: 2 additions & 72 deletions
Original file line numberDiff line numberDiff line change
@@ -1,74 +1,4 @@
1-
# 1. Bedrock Cohere Embedding
1+
# Bedrock Cohere Embedding
22

3-
## 1.1 CohereEmbeddingBedrockApi
3+
Visit the Spring AI [Bedrock Cohere Embedding Documentation](https://docs.spring.io/spring-ai/reference/api/embeddings/bedrock-cohere-embedding.html).
44

5-
[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).
6-
7-
Following class diagram illustrates the Llama2ChatBedrockApi interface and building blocks:
8-
9-
![CohereEmbeddingBedrockApi Class Diagram](./src/test/resources/doc/Bedrock%20Cohere%20Embedding%20API.jpg)
10-
11-
The CohereEmbeddingBedrockApi supports the `cohere.embed-english-v3` and `cohere.embed-multilingual-v3` models for single and batch embedding computation.
12-
13-
Here is a simple snippet how to use the api programmatically:
14-
15-
```java
16-
CohereEmbeddingBedrockApi api = new CohereEmbeddingBedrockApi(
17-
CohereEmbeddingModel.COHERE_EMBED_MULTILINGUAL_V1.id(),
18-
EnvironmentVariableCredentialsProvider.create(),
19-
Region.US_EAST_1.id(), new ObjectMapper());
20-
21-
CohereEmbeddingRequest request = new CohereEmbeddingRequest(
22-
List.of("I like to eat apples", "I like to eat oranges"),
23-
CohereEmbeddingRequest.InputType.search_document,
24-
CohereEmbeddingRequest.Truncate.NONE);
25-
26-
CohereEmbeddingResponse response = api.embedding(request);
27-
28-
assertThat(response.embeddings()).hasSize(2);
29-
assertThat(response.embeddings().get(0)).hasSize(1024);
30-
```
31-
32-
## 1.2 BedrockCohereEmbeddingClient
33-
34-
[BedrockCohereEmbeddingClient](./src/main/java/org/springframework/ai/bedrock/cohere/BedrockCohereEmbeddingClient.java) implements the Spring-Ai `EmbeddingClient` on top of the `CohereEmbeddingBedrockApi`.
35-
36-
You can use like this:
37-
38-
```java
39-
@Bean
40-
public CohereEmbeddingBedrockApi cohereEmbeddingApi() {
41-
return new CohereEmbeddingBedrockApi(CohereEmbeddingModel.COHERE_EMBED_MULTILINGUAL_V1.id(),
42-
EnvironmentVariableCredentialsProvider.create(), Region.US_EAST_1.id(), new ObjectMapper());
43-
}
44-
45-
@Bean
46-
public BedrockCohereEmbeddingClient cohereAiEmbedding(CohereEmbeddingBedrockApi cohereEmbeddingApi) {
47-
return new BedrockCohereEmbeddingClient(cohereEmbeddingApi);
48-
}
49-
```
50-
51-
or you can leverage the `spring-ai-bedrock-ai-spring-boot-starter` Boot starter. For this add the following dependency:
52-
53-
```xml
54-
<dependency>
55-
<artifactId>spring-ai-bedrock-ai-spring-boot-starter</artifactId>
56-
<groupId>org.springframework.ai</groupId>
57-
<version>0.8.0-SNAPSHOT</version>
58-
</dependency>
59-
```
60-
61-
**NOTE:** You have to enable the Bedrock Cohere embedding client with `spring.ai.bedrock.cohere.embedding.enabled=true`.
62-
By default the client is disabled.
63-
64-
Use the `BedrockCohereEmbeddingProperties` to configure the Bedrock Cohere Chat client:
65-
66-
| Property | Description | Default |
67-
| ------------- | ------------- | ------------- |
68-
| spring.ai.bedrock.aws.region | AWS region to use. | us-east-1 |
69-
| spring.ai.bedrock.aws.accessKey | AWS credentials access key. | |
70-
| spring.ai.bedrock.aws.secretKey | AWS credentials secret key. | |
71-
| spring.ai.bedrock.cohere.embedding.enable | Enable Bedrock Cohere embedding client. Disabled by default | false |
72-
| spring.ai.bedrock.cohere.embedding.model | The model id to use. See the `CohereEmbeddingModel` for the supported models. | cohere.embed-multilingual-v3 |
73-
| 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 |
74-
| spring.ai.bedrock.cohere.embedding.truncate | Specifies how the API handles inputs longer than the maximum token length. | NONE |

‎models/spring-ai-bedrock/src/main/java/org/springframework/ai/bedrock/cohere/BedrockCohereEmbeddingClient.java

Lines changed: 67 additions & 23 deletions
Original file line numberDiff line numberDiff line change
@@ -25,8 +25,10 @@
2525
import org.springframework.ai.document.Document;
2626
import org.springframework.ai.embedding.AbstractEmbeddingClient;
2727
import org.springframework.ai.embedding.Embedding;
28+
import org.springframework.ai.embedding.EmbeddingOptions;
2829
import org.springframework.ai.embedding.EmbeddingRequest;
2930
import org.springframework.ai.embedding.EmbeddingResponse;
31+
import org.springframework.ai.model.ModelOptionsUtils;
3032
import org.springframework.util.Assert;
3133

3234
/**
@@ -41,35 +43,54 @@ public class BedrockCohereEmbeddingClient extends AbstractEmbeddingClient {
4143

4244
private final CohereEmbeddingBedrockApi embeddingApi;
4345

44-
private CohereEmbeddingRequest.InputType inputType = CohereEmbeddingRequest.InputType.SEARCH_DOCUMENT;
46+
private final BedrockCohereEmbeddingOptions defaultOptions;
4547

46-
private CohereEmbeddingRequest.Truncate truncate = CohereEmbeddingRequest.Truncate.NONE;
48+
// private CohereEmbeddingRequest.InputType inputType =
49+
// CohereEmbeddingRequest.InputType.SEARCH_DOCUMENT;
50+
51+
// private CohereEmbeddingRequest.Truncate truncate =
52+
// CohereEmbeddingRequest.Truncate.NONE;
4753

4854
public BedrockCohereEmbeddingClient(CohereEmbeddingBedrockApi cohereEmbeddingBedrockApi) {
49-
this.embeddingApi = cohereEmbeddingBedrockApi;
55+
this(cohereEmbeddingBedrockApi,
56+
BedrockCohereEmbeddingOptions.builder()
57+
.withInputType(CohereEmbeddingRequest.InputType.SEARCH_DOCUMENT)
58+
.withTruncate(CohereEmbeddingRequest.Truncate.NONE)
59+
.build());
5060
}
5161

52-
/**
53-
* Cohere Embedding API input types.
54-
* @param inputType the input type to use.
55-
* @return this client.
56-
*/
57-
public BedrockCohereEmbeddingClient withInputType(CohereEmbeddingRequest.InputType inputType) {
58-
this.inputType = inputType;
59-
return this;
62+
public BedrockCohereEmbeddingClient(CohereEmbeddingBedrockApi cohereEmbeddingBedrockApi,
63+
BedrockCohereEmbeddingOptions options) {
64+
Assert.notNull(cohereEmbeddingBedrockApi, "CohereEmbeddingBedrockApi must not be null");
65+
Assert.notNull(options, "BedrockCohereEmbeddingOptions must not be null");
66+
this.embeddingApi = cohereEmbeddingBedrockApi;
67+
this.defaultOptions = options;
6068
}
6169

62-
/**
63-
* Specifies how the API handles inputs longer than the maximum token length. If you
64-
* specify LEFT or RIGHT, the model discards the input until the remaining input is
65-
* exactly the maximum input token length for the model.
66-
* @param truncate the truncate option to use.
67-
* @return this client.
68-
*/
69-
public BedrockCohereEmbeddingClient withTruncate(CohereEmbeddingRequest.Truncate truncate) {
70-
this.truncate = truncate;
71-
return this;
72-
}
70+
// /**
71+
// * Cohere Embedding API input types.
72+
// * @param inputType the input type to use.
73+
// * @return this client.
74+
// */
75+
// public BedrockCohereEmbeddingClient withInputType(CohereEmbeddingRequest.InputType
76+
// inputType) {
77+
// this.inputType = inputType;
78+
// return this;
79+
// }
80+
81+
// /**
82+
// * Specifies how the API handles inputs longer than the maximum token length. If you
83+
// specify LEFT or RIGHT, the
84+
// * model discards the input until the remaining input is exactly the maximum input
85+
// token length for the model.
86+
// * @param truncate the truncate option to use.
87+
// * @return this client.
88+
// */
89+
// public BedrockCohereEmbeddingClient withTruncate(CohereEmbeddingRequest.Truncate
90+
// truncate) {
91+
// this.truncate = truncate;
92+
// return this;
93+
// }
7394

7495
@Override
7596
public List<Double> embed(Document document) {
@@ -80,7 +101,10 @@ public List<Double> embed(Document document) {
80101
public EmbeddingResponse call(EmbeddingRequest request) {
81102
Assert.notEmpty(request.getInstructions(), "At least one text is required!");
82103

83-
var apiRequest = new CohereEmbeddingRequest(request.getInstructions(), this.inputType, this.truncate);
104+
final BedrockCohereEmbeddingOptions optionsToUse = this.mergeOptions(request.getOptions());
105+
106+
var apiRequest = new CohereEmbeddingRequest(request.getInstructions(), optionsToUse.getInputType(),
107+
optionsToUse.getTruncate());
84108
CohereEmbeddingResponse apiResponse = this.embeddingApi.embedding(apiRequest);
85109
var indexCounter = new AtomicInteger(0);
86110
List<Embedding> embeddings = apiResponse.embeddings()
@@ -90,4 +114,24 @@ public EmbeddingResponse call(EmbeddingRequest request) {
90114
return new EmbeddingResponse(embeddings);
91115
}
92116

117+
/**
118+
* Merge the default and request options.
119+
* @param requestOptions request options to merge.
120+
* @return the merged options.
121+
*/
122+
BedrockCohereEmbeddingOptions mergeOptions(EmbeddingOptions requestOptions) {
123+
124+
BedrockCohereEmbeddingOptions options = (this.defaultOptions != null) ? this.defaultOptions
125+
: BedrockCohereEmbeddingOptions.builder()
126+
.withInputType(CohereEmbeddingRequest.InputType.SEARCH_DOCUMENT)
127+
.withTruncate(CohereEmbeddingRequest.Truncate.NONE)
128+
.build();
129+
130+
if (requestOptions != null && !EmbeddingOptions.EMPTY.equals(requestOptions)) {
131+
options = ModelOptionsUtils.merge(requestOptions, options, BedrockCohereEmbeddingOptions.class);
132+
}
133+
134+
return options;
135+
}
136+
93137
}

‎models/spring-ai-bedrock/src/main/java/org/springframework/ai/bedrock/cohere/BedrockCohereEmbeddingOptions.java

Lines changed: 89 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,89 @@
1+
/*
2+
* Copyright 2024-2024 the original author or authors.
3+
*
4+
* Licensed under the Apache License, Version 2.0 (the "License");
5+
* you may not use this file except in compliance with the License.
6+
* You may obtain a copy of the License at
7+
*
8+
* https://www.apache.org/licenses/LICENSE-2.0
9+
*
10+
* Unless required by applicable law or agreed to in writing, software
11+
* distributed under the License is distributed on an "AS IS" BASIS,
12+
* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
13+
* See the License for the specific language governing permissions and
14+
* limitations under the License.
15+
*/
16+
17+
package org.springframework.ai.bedrock.cohere;
18+
19+
import com.fasterxml.jackson.annotation.JsonInclude;
20+
import com.fasterxml.jackson.annotation.JsonInclude.Include;
21+
import com.fasterxml.jackson.annotation.JsonProperty;
22+
23+
import org.springframework.ai.bedrock.cohere.api.CohereEmbeddingBedrockApi.CohereEmbeddingRequest.InputType;
24+
import org.springframework.ai.bedrock.cohere.api.CohereEmbeddingBedrockApi.CohereEmbeddingRequest.Truncate;
25+
26+
/**
27+
* @author Christian Tzolov
28+
*/
29+
@JsonInclude(Include.NON_NULL)
30+
public class BedrockCohereEmbeddingOptions {
31+
32+
// @formatter:off
33+
/**
34+
* Prepends special tokens to differentiate each type from one another. You should not mix
35+
* different types together, except when mixing types for for search and retrieval.
36+
* In this case, embed your corpus with the search_document type and embedded queries with
37+
* type search_query type.
38+
*/
39+
private @JsonProperty("input_type") InputType inputType;
40+
41+
/**
42+
* Specifies how the API handles inputs longer than the maximum token length. If you specify LEFT or
43+
* RIGHT, the model discards the input until the remaining input is exactly the maximum input token length for the
44+
* model.
45+
*/
46+
private @JsonProperty("truncate") Truncate truncate;
47+
// @formatter:on
48+
49+
public static Builder builder() {
50+
return new Builder();
51+
}
52+
53+
public static class Builder {
54+
55+
private BedrockCohereEmbeddingOptions options = new BedrockCohereEmbeddingOptions();
56+
57+
public Builder withInputType(InputType inputType) {
58+
this.options.setInputType(inputType);
59+
return this;
60+
}
61+
62+
public Builder withTruncate(Truncate truncate) {
63+
this.options.setTruncate(truncate);
64+
return this;
65+
}
66+
67+
public BedrockCohereEmbeddingOptions build() {
68+
return this.options;
69+
}
70+
71+
}
72+
73+
public InputType getInputType() {
74+
return this.inputType;
75+
}
76+
77+
public void setInputType(InputType inputType) {
78+
this.inputType = inputType;
79+
}
80+
81+
public Truncate getTruncate() {
82+
return this.truncate;
83+
}
84+
85+
public void setTruncate(Truncate truncate) {
86+
this.truncate = truncate;
87+
}
88+
89+
}

‎models/spring-ai-bedrock/src/main/java/org/springframework/ai/bedrock/cohere/api/CohereEmbeddingBedrockApi.java

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -163,7 +163,7 @@ public enum CohereEmbeddingModel {
163163
* @return The model id.
164164
*/
165165
public String id() {
166-
return id;
166+
return this.id;
167167
}
168168

169169
CohereEmbeddingModel(String value) {

‎spring-ai-docs/src/main/antora/modules/ROOT/images/bedrock/bedrock-cohere-chat-api.jpg

-164 KB
Binary file not shown.

‎spring-ai-docs/src/main/antora/modules/ROOT/images/bedrock/bedrock-cohere-chat-low-level-api.jpg

442 KB
Loading

‎spring-ai-docs/src/main/antora/modules/ROOT/images/bedrock/bedrock-cohere-embedding-low-level-api.jpg

143 KB
Loading

‎spring-ai-docs/src/main/antora/modules/ROOT/nav.adoc

Lines changed: 3 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -8,10 +8,12 @@
88
*** xref:api/embeddings/ollama-embeddings.adoc[]
99
*** xref:api/embeddings/azure-openai-embeddings.adoc[]
1010
*** xref:api/embeddings/postgresml-embeddings.adoc[]
11+
*** xref:api/bedrock.adoc[Amazon Bedrock Embedding]
12+
**** xref:api/embeddings/bedrock-cohere-embedding.adoc[]
1113
** xref:api/chatclient.adoc[]
1214
*** xref:api/clients/openai-chat.adoc[]
1315
*** xref:api/clients/azure-openai-chat.adoc[]
14-
*** xref:api/clients/bedrock.adoc[]
16+
*** xref:api/bedrock.adoc[Amazon Bedrock Chat]
1517
**** xref:api/clients/bedrock/bedrock-anthropic.adoc[]
1618
**** xref:api/clients/bedrock/bedrock-llama2.adoc[]
1719
**** xref:api/clients/bedrock/bedrock-cohere.adoc[]

‎spring-ai-docs/src/main/antora/modules/ROOT/pages/api/clients/bedrock.adoc renamed to ‎spring-ai-docs/src/main/antora/modules/ROOT/pages/api/bedrock.adoc

Lines changed: 3 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -1,4 +1,4 @@
1-
= Amazon Bedrock Chat
1+
= Amazon Bedrock
22

33
link:https://docs.aws.amazon.com/bedrock/latest/userguide/what-is-bedrock.html[Amazon Bedrock] is a managed service that provides foundation models from various AI providers, available through a unified API.
44

@@ -87,10 +87,5 @@ For more information, refer to the documentation below for each supported model.
8787
* xref:api/clients/bedrock/bedrock-anthropic.adoc[Spring AI Bedrock Anthropic Chat]: `spring.ai.bedrock.anthropic.chat.enabled=true`
8888
* xref:api/clients/bedrock/bedrock-llama2.adoc[Spring AI Bedrock Llama2 Chat]: `spring.ai.bedrock.llama2.chat.enabled=true`
8989
* xref:api/clients/bedrock/bedrock-cohere.adoc[Spring AI Bedrock Cohere Chat]: `spring.ai.bedrock.cohere.chat.enabled=true`
90-
91-
92-
// * [Spring AI Bedrock Cohere Chat](./README_COHERE_CHAT.md) - `spring.ai.bedrock.cohere.chat.enabled=true`
93-
// * [Spring AI Bedrock Cohere Embedding](./README_COHERE_EMBEDDING.md) - `spring.ai.bedrock.cohere.embedding.enabled=true`
94-
// * [Spring AI Bedrock Titan Chat](./README_TITAN_CHAT.md) - `spring.ai.bedrock.titan.chat.enabled=true`
95-
// * [Spring AI Bedrock Titan Embedding](./README_TITAN_EMBEDDING.md) - `spring.ai.bedrock.titan.embedding.enabled=true`
96-
// * (WIP) [Spring AI Bedrock Ai21 Jurassic2 Chat](./README_JURASSIC2_CHAT.md) - `spring.ai.bedrock.jurassic2.chat.enabled=true`
90+
* xref:api/embeddings/bedrock-cohere-embedding.adoc[Spring AI Bedrock Cohere Embeddings]: `spring.ai.bedrock.cohere.embedding.enabled=true`
91+
// * xref:api/clients/bedrock/bedrock-jurassic2-chat.adoc[(WIP)Spring AI Bedrock Jurassic Chat]: `spring.ai.bedrock.jurassic2.chat.enabled=true`

‎spring-ai-docs/src/main/antora/modules/ROOT/pages/api/clients/bedrock/bedrock-anthropic.adoc

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -11,7 +11,7 @@ The https://aws.amazon.com/bedrock/claude[AWS Bedrock Anthropic Model Page] and
1111

1212
== Prerequisites
1313

14-
Refer to the xref:api/clients/bedrock.adoc[Spring AI documentation on Amazon Bedrock] for setting up API access.
14+
Refer to the xref:api/bedrock.adoc[Spring AI documentation on Amazon Bedrock] for setting up API access.
1515

1616
== Auto-configuration
1717

0 commit comments

Comments
 (0)