Add dimensions option for ZhiPuAi embedding model

Signed-off-by: Assassinxc <[email protected]>
spring-projects · Jan 21, 2025 · 85c1420 · 85c1420
1 parent 6a53268
commit 85c1420
Show file tree

Hide file tree

Showing 5 changed files with 53 additions and 8 deletions.
diff --git a/...spring-ai-zhipuai/src/main/java/org/springframework/ai/zhipuai/ZhiPuAiEmbeddingModel.java b/...spring-ai-zhipuai/src/main/java/org/springframework/ai/zhipuai/ZhiPuAiEmbeddingModel.java
@@ -221,11 +221,13 @@ private ZhiPuAiEmbeddingOptions mergeOptions(@Nullable EmbeddingOptions runtimeO
 
 		return ZhiPuAiEmbeddingOptions.builder()
 			.model(ModelOptionsUtils.mergeOption(runtimeOptionsForProvider.getModel(), defaultOptions.getModel()))
+			.dimensions(ModelOptionsUtils.mergeOption(runtimeOptionsForProvider.getDimensions(),
+					defaultOptions.getDimensions()))
 			.build();
 	}
 
 	private ZhiPuAiApi.EmbeddingRequest<String> createEmbeddingRequest(String text, EmbeddingOptions requestOptions) {
-		return new ZhiPuAiApi.EmbeddingRequest<>(text, requestOptions.getModel());
+		return new ZhiPuAiApi.EmbeddingRequest<>(text, requestOptions.getModel(), requestOptions.getDimensions());
 	}
 
 	public void setObservationConvention(EmbeddingModelObservationConvention observationConvention) {

diff --git a/...ring-ai-zhipuai/src/main/java/org/springframework/ai/zhipuai/ZhiPuAiEmbeddingOptions.java b/...ring-ai-zhipuai/src/main/java/org/springframework/ai/zhipuai/ZhiPuAiEmbeddingOptions.java
@@ -39,6 +39,10 @@ public class ZhiPuAiEmbeddingOptions implements EmbeddingOptions {
 	 * ID of the model to use.
 	 */
 	private @JsonProperty("model") String model;
+	/**
+	 * Dimension value of the model to use.
+	 */
+	private @JsonProperty("dimensions") Integer dimensions;
 	// @formatter:on
 
 	public static Builder builder() {
@@ -54,6 +58,10 @@ public void setModel(String model) {
 		this.model = model;
 	}
 
+	public void setDimensions(Integer dimensions) {
+		this.dimensions = dimensions;
+	}
+
 	@Override
 	@JsonIgnore
 	public Integer getDimensions() {
@@ -73,6 +81,11 @@ public Builder model(String model) {
 			return this;
 		}
 
+		public Builder dimensions(Integer dimensions) {
+			this.options.setDimensions(dimensions);
+			return this;
+		}
+
 		public ZhiPuAiEmbeddingOptions build() {
 			return this.options;
 		}

diff --git a/models/spring-ai-zhipuai/src/main/java/org/springframework/ai/zhipuai/api/ZhiPuAiApi.java b/models/spring-ai-zhipuai/src/main/java/org/springframework/ai/zhipuai/api/ZhiPuAiApi.java
@@ -295,7 +295,12 @@ public enum EmbeddingModel {
 		/**
 		 * DIMENSION: 1024
 		 */
-		Embedding_2("Embedding-2");
+		Embedding_2("Embedding-2"),
+
+		/**
+		 * DIMENSION: up to 2048
+		 */
+		Embedding_3("Embedding-3");
 
 		public final String  value;
 
@@ -956,15 +961,27 @@ public String toString() {
 	@JsonInclude(Include.NON_NULL)
 	public record EmbeddingRequest<T>(
 			@JsonProperty("input") T input,
-			@JsonProperty("model") String model) {
+			@JsonProperty("model") String model,
+			@JsonProperty("dimensions") Integer dimensions) {
 
 
 		/**
-		 * Create an embedding request with the given input. Encoding model is set to 'embedding-2'.
-		 * @param input Input text to embed.
-		 */
+         * Create an embedding request with the given input. Encoding model is set to 'embedding-2'.
+         *
+         * @param input Input text to embed.
+         */
 		public EmbeddingRequest(T input) {
-			this(input, DEFAULT_EMBEDDING_MODEL);
+			this(input, DEFAULT_EMBEDDING_MODEL,null);
+		}
+
+		/**
+		 * Create an embedding request with the given input and model.
+		 *
+		 * @param input
+		 * @param model
+		 */
+		public EmbeddingRequest(T input, String model) {
+			this(input, model,null);
 		}
 	}
 

diff --git a/models/spring-ai-zhipuai/src/test/java/org/springframework/ai/zhipuai/api/ZhiPuAiApiIT.java b/models/spring-ai-zhipuai/src/test/java/org/springframework/ai/zhipuai/api/ZhiPuAiApiIT.java
@@ -83,4 +83,15 @@ void embeddings() {
 		assertThat(response.getBody().data().get(0).embedding()).hasSize(1024);
 	}
 
+	@Test
+	void embeddingsWithDimensions() {
+		ResponseEntity<EmbeddingList<Embedding>> response = this.zhiPuAiApi
+			.embeddings(new ZhiPuAiApi.EmbeddingRequest<>("Hello world",
+					ZhiPuAiApi.EmbeddingModel.Embedding_3.getValue(), 1536));
+
+		assertThat(response).isNotNull();
+		assertThat(Objects.requireNonNull(response.getBody()).data()).hasSize(1);
+		assertThat(response.getBody().data().get(0).embedding()).hasSize(1536);
+	}
+
 }
diff --git a/...-docs/src/main/antora/modules/ROOT/pages/api/embeddings/zhipuai-embeddings.adoc b/...-docs/src/main/antora/modules/ROOT/pages/api/embeddings/zhipuai-embeddings.adoc
@@ -91,6 +91,7 @@ The prefix `spring.ai.zhipuai.embedding` is property prefix that configures the
 | spring.ai.zhipuai.embedding.base-url   | Optional overrides the spring.ai.zhipuai.base-url to provide embedding specific url | -
 | spring.ai.zhipuai.embedding.api-key    | Optional overrides the spring.ai.zhipuai.api-key to provide embedding specific api-key  | -
 | spring.ai.zhipuai.embedding.options.model      | The model to use      | embedding-2
+| spring.ai.zhipuai.embedding.options.dimensions      | The number of dimensions      | -
 |====
 
 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"));
 
 var embeddingModel = new ZhiPuAiEmbeddingModel(api, MetadataMode.EMBED,
 				ZhiPuAiEmbeddingOptions.builder()
-						.model("embedding-2")
+						.model("embedding-3")
+						.dimensions(1536)
 						.build());
 
 EmbeddingResponse embeddingResponse = this.embeddingModel