diff --git a/models/spring-ai-zhipuai/src/main/java/org/springframework/ai/zhipuai/ZhiPuAiEmbeddingModel.java b/models/spring-ai-zhipuai/src/main/java/org/springframework/ai/zhipuai/ZhiPuAiEmbeddingModel.java index 8ec16a109..9dc29fc6d 100644 --- a/models/spring-ai-zhipuai/src/main/java/org/springframework/ai/zhipuai/ZhiPuAiEmbeddingModel.java +++ b/models/spring-ai-zhipuai/src/main/java/org/springframework/ai/zhipuai/ZhiPuAiEmbeddingModel.java @@ -24,14 +24,15 @@ 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.embedding.EmbeddingResponseMetadata; import org.springframework.ai.model.ModelOptionsUtils; import org.springframework.ai.retry.RetryUtils; import org.springframework.ai.zhipuai.api.ZhiPuAiApi; import org.springframework.retry.support.RetryTemplate; import org.springframework.util.Assert; +import java.util.ArrayList; import java.util.List; +import java.util.concurrent.atomic.AtomicInteger; /** * ZhiPuAI Embedding Client implementation. @@ -112,43 +113,43 @@ public class ZhiPuAiEmbeddingModel extends AbstractEmbeddingModel { public EmbeddingResponse call(EmbeddingRequest request) { return this.retryTemplate.execute(ctx -> { - - ZhiPuAiApi.EmbeddingRequest> apiRequest = (this.defaultOptions != null) - ? new ZhiPuAiApi.EmbeddingRequest<>(request.getInstructions(), this.defaultOptions.getModel()) - : new ZhiPuAiApi.EmbeddingRequest<>(request.getInstructions(), ZhiPuAiApi.DEFAULT_EMBEDDING_MODEL); - - if (request.getOptions() != null && !EmbeddingOptions.EMPTY.equals(request.getOptions())) { - apiRequest = ModelOptionsUtils.merge(request.getOptions(), apiRequest, - ZhiPuAiApi.EmbeddingRequest.class); + Assert.notEmpty(request.getInstructions(), "At least one text is required!"); + if (request.getInstructions().size() != 1) { + logger.warn( + "ZhiPu Embedding does not support batch embedding. Will make multiple API calls to embed(Document)"); } - ZhiPuAiApi.EmbeddingList apiEmbeddingResponse = this.zhiPuAiApi.embeddings(apiRequest) - .getBody(); - - if (apiEmbeddingResponse == null) { - logger.warn("No embeddings returned for request: {}", request); - return new EmbeddingResponse(List.of()); + List> embeddingList = new ArrayList<>(); + for (String inputContent : request.getInstructions()) { + var apiRequest = createZhiPuEmbeddingRequest(inputContent, request.getOptions()); + ZhiPuAiApi.EmbeddingList response = this.zhiPuAiApi.embeddings(apiRequest) + .getBody(); + if (response == null || response.data() == null || response.data().isEmpty()) { + logger.warn("No embeddings returned for input: {}", inputContent); + embeddingList.add(List.of()); + } + else { + embeddingList.add(response.data().get(0).embedding()); + } } - - var metadata = generateResponseMetadata(apiEmbeddingResponse.model(), apiEmbeddingResponse.usage()); - - List embeddings = apiEmbeddingResponse.data() - .stream() - .map(e -> new Embedding(e.embedding(), e.index())) + var indexCounter = new AtomicInteger(0); + List embeddings = embeddingList.stream() + .map(e -> new Embedding(e, indexCounter.getAndIncrement())) .toList(); - - return new EmbeddingResponse(embeddings, metadata); - + return new EmbeddingResponse(embeddings); }); } - private EmbeddingResponseMetadata generateResponseMetadata(String model, ZhiPuAiApi.Usage usage) { - EmbeddingResponseMetadata metadata = new EmbeddingResponseMetadata(); - metadata.put("model", model); - metadata.put("prompt-tokens", usage.promptTokens()); - metadata.put("completion-tokens", usage.completionTokens()); - metadata.put("total-tokens", usage.totalTokens()); - return metadata; + private ZhiPuAiApi.EmbeddingRequest createZhiPuEmbeddingRequest(String text, + EmbeddingOptions requestOptions) { + ZhiPuAiApi.EmbeddingRequest apiRequest = (this.defaultOptions != null) + ? new ZhiPuAiApi.EmbeddingRequest<>(text, this.defaultOptions.getModel()) + : new ZhiPuAiApi.EmbeddingRequest<>(text, ZhiPuAiApi.DEFAULT_EMBEDDING_MODEL); + + if (requestOptions != null && !EmbeddingOptions.EMPTY.equals(requestOptions)) { + apiRequest = ModelOptionsUtils.merge(requestOptions, apiRequest, ZhiPuAiApi.EmbeddingRequest.class); + } + return apiRequest; } } diff --git a/models/spring-ai-zhipuai/src/test/java/org/springframework/ai/zhipuai/embedding/EmbeddingIT.java b/models/spring-ai-zhipuai/src/test/java/org/springframework/ai/zhipuai/embedding/EmbeddingIT.java new file mode 100644 index 000000000..7955b963e --- /dev/null +++ b/models/spring-ai-zhipuai/src/test/java/org/springframework/ai/zhipuai/embedding/EmbeddingIT.java @@ -0,0 +1,68 @@ +/* + * Copyright 2023 - 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.zhipuai.embedding; + +import org.junit.jupiter.api.Test; +import org.springframework.ai.embedding.EmbeddingResponse; +import org.springframework.ai.zhipuai.ZhiPuAiEmbeddingModel; +import org.springframework.beans.factory.annotation.Autowired; +import org.springframework.boot.test.context.SpringBootTest; + +import java.util.List; + +import static org.assertj.core.api.Assertions.assertThat; + +/** + * @author Geng Rong + */ +@SpringBootTest +class EmbeddingIT { + + @Autowired + private ZhiPuAiEmbeddingModel embeddingModel; + + @Test + void defaultEmbedding() { + assertThat(embeddingModel).isNotNull(); + + EmbeddingResponse embeddingResponse = embeddingModel.embedForResponse(List.of("Hello World")); + + assertThat(embeddingResponse.getResults()).hasSize(1); + + assertThat(embeddingResponse.getResults().get(0)).isNotNull(); + assertThat(embeddingResponse.getResults().get(0).getOutput()).hasSize(1024); + + assertThat(embeddingModel.dimensions()).isEqualTo(1024); + } + + @Test + void batchEmbedding() { + assertThat(embeddingModel).isNotNull(); + + EmbeddingResponse embeddingResponse = embeddingModel.embedForResponse(List.of("Hello World", "HI")); + + assertThat(embeddingResponse.getResults()).hasSize(2); + + assertThat(embeddingResponse.getResults().get(0)).isNotNull(); + assertThat(embeddingResponse.getResults().get(0).getOutput()).hasSize(1024); + + assertThat(embeddingResponse.getResults().get(1)).isNotNull(); + assertThat(embeddingResponse.getResults().get(1).getOutput()).hasSize(1024); + + assertThat(embeddingModel.dimensions()).isEqualTo(1024); + } + +}