Skip to content

Commit 62ce753

Browse files
Assassinxcilayaperumalg
authored andcommitted
Add dimensions option for ZhiPuAi embedding model
- Add dimentions options - Add doc for ZhiPuAi embedding dimensions option Signed-off-by: Assassinxc <[email protected]>
1 parent 14ec6b2 commit 62ce753

File tree

5 files changed

+53
-8
lines changed

5 files changed

+53
-8
lines changed

models/spring-ai-zhipuai/src/main/java/org/springframework/ai/zhipuai/ZhiPuAiEmbeddingModel.java

Lines changed: 3 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -224,11 +224,13 @@ private ZhiPuAiEmbeddingOptions mergeOptions(@Nullable EmbeddingOptions runtimeO
224224

225225
return ZhiPuAiEmbeddingOptions.builder()
226226
.model(ModelOptionsUtils.mergeOption(runtimeOptionsForProvider.getModel(), defaultOptions.getModel()))
227+
.dimensions(ModelOptionsUtils.mergeOption(runtimeOptionsForProvider.getDimensions(),
228+
defaultOptions.getDimensions()))
227229
.build();
228230
}
229231

230232
private ZhiPuAiApi.EmbeddingRequest<String> createEmbeddingRequest(String text, EmbeddingOptions requestOptions) {
231-
return new ZhiPuAiApi.EmbeddingRequest<>(text, requestOptions.getModel());
233+
return new ZhiPuAiApi.EmbeddingRequest<>(text, requestOptions.getModel(), requestOptions.getDimensions());
232234
}
233235

234236
public void setObservationConvention(EmbeddingModelObservationConvention observationConvention) {

models/spring-ai-zhipuai/src/main/java/org/springframework/ai/zhipuai/ZhiPuAiEmbeddingOptions.java

Lines changed: 13 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -39,6 +39,10 @@ public class ZhiPuAiEmbeddingOptions implements EmbeddingOptions {
3939
* ID of the model to use.
4040
*/
4141
private @JsonProperty("model") String model;
42+
/**
43+
* Dimension value of the model to use.
44+
*/
45+
private @JsonProperty("dimensions") Integer dimensions;
4246
// @formatter:on
4347

4448
public static Builder builder() {
@@ -54,6 +58,10 @@ public void setModel(String model) {
5458
this.model = model;
5559
}
5660

61+
public void setDimensions(Integer dimensions) {
62+
this.dimensions = dimensions;
63+
}
64+
5765
@Override
5866
@JsonIgnore
5967
public Integer getDimensions() {
@@ -73,6 +81,11 @@ public Builder model(String model) {
7381
return this;
7482
}
7583

84+
public Builder dimensions(Integer dimensions) {
85+
this.options.setDimensions(dimensions);
86+
return this;
87+
}
88+
7689
public ZhiPuAiEmbeddingOptions build() {
7790
return this.options;
7891
}

models/spring-ai-zhipuai/src/main/java/org/springframework/ai/zhipuai/api/ZhiPuAiApi.java

Lines changed: 23 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -295,7 +295,12 @@ public enum EmbeddingModel {
295295
/**
296296
* DIMENSION: 1024
297297
*/
298-
Embedding_2("Embedding-2");
298+
Embedding_2("Embedding-2"),
299+
300+
/**
301+
* DIMENSION: up to 2048
302+
*/
303+
Embedding_3("Embedding-3");
299304

300305
public final String value;
301306

@@ -956,15 +961,27 @@ public String toString() {
956961
@JsonInclude(Include.NON_NULL)
957962
public record EmbeddingRequest<T>(
958963
@JsonProperty("input") T input,
959-
@JsonProperty("model") String model) {
964+
@JsonProperty("model") String model,
965+
@JsonProperty("dimensions") Integer dimensions) {
960966

961967

962968
/**
963-
* Create an embedding request with the given input. Encoding model is set to 'embedding-2'.
964-
* @param input Input text to embed.
965-
*/
969+
* Create an embedding request with the given input. Encoding model is set to 'embedding-2'.
970+
*
971+
* @param input Input text to embed.
972+
*/
966973
public EmbeddingRequest(T input) {
967-
this(input, DEFAULT_EMBEDDING_MODEL);
974+
this(input, DEFAULT_EMBEDDING_MODEL,null);
975+
}
976+
977+
/**
978+
* Create an embedding request with the given input and model.
979+
*
980+
* @param input
981+
* @param model
982+
*/
983+
public EmbeddingRequest(T input, String model) {
984+
this(input, model,null);
968985
}
969986
}
970987

models/spring-ai-zhipuai/src/test/java/org/springframework/ai/zhipuai/api/ZhiPuAiApiIT.java

Lines changed: 11 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -83,4 +83,15 @@ void embeddings() {
8383
assertThat(response.getBody().data().get(0).embedding()).hasSize(1024);
8484
}
8585

86+
@Test
87+
void embeddingsWithDimensions() {
88+
ResponseEntity<EmbeddingList<Embedding>> response = this.zhiPuAiApi
89+
.embeddings(new ZhiPuAiApi.EmbeddingRequest<>("Hello world",
90+
ZhiPuAiApi.EmbeddingModel.Embedding_3.getValue(), 1536));
91+
92+
assertThat(response).isNotNull();
93+
assertThat(Objects.requireNonNull(response.getBody()).data()).hasSize(1);
94+
assertThat(response.getBody().data().get(0).embedding()).hasSize(1536);
95+
}
96+
8697
}

spring-ai-docs/src/main/antora/modules/ROOT/pages/api/embeddings/zhipuai-embeddings.adoc

Lines changed: 3 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -91,6 +91,7 @@ The prefix `spring.ai.zhipuai.embedding` is property prefix that configures the
9191
| spring.ai.zhipuai.embedding.base-url | Optional overrides the spring.ai.zhipuai.base-url to provide embedding specific url | -
9292
| spring.ai.zhipuai.embedding.api-key | Optional overrides the spring.ai.zhipuai.api-key to provide embedding specific api-key | -
9393
| spring.ai.zhipuai.embedding.options.model | The model to use | embedding-2
94+
| spring.ai.zhipuai.embedding.options.dimensions | The number of dimensions, the default value is 2048 when the model is embedding-3 | -
9495
|====
9596

9697
NOTE: You can override the common `spring.ai.zhipuai.base-url` and `spring.ai.zhipuai.api-key` for the `ChatModel` and `EmbeddingModel` implementations.
@@ -185,7 +186,8 @@ var zhiPuAiApi = new ZhiPuAiApi(System.getenv("ZHIPU_AI_API_KEY"));
185186
186187
var embeddingModel = new ZhiPuAiEmbeddingModel(api, MetadataMode.EMBED,
187188
ZhiPuAiEmbeddingOptions.builder()
188-
.model("embedding-2")
189+
.model("embedding-3")
190+
.dimensions(1536)
189191
.build());
190192
191193
EmbeddingResponse embeddingResponse = this.embeddingModel

0 commit comments

Comments
 (0)