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

fix: ZhiPu embedding client individual calls to embed document collection

Browse Source 
fixes #771
This commit is contained in:
GR
2024-05-27 08:47:43 +08:00
committed by Mark Pollack
parent c4784a070f
commit 23e448cf2a
2 changed files with 100 additions and 31 deletions
Download Patch File Download Diff File Expand all files Collapse all files

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

@@ -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<List<String>> 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<ZhiPuAiApi.Embedding> apiEmbeddingResponse = this.zhiPuAiApi.embeddings(apiRequest)
.getBody();
if (apiEmbeddingResponse == null) {
logger.warn("No embeddings returned for request: {}", request);
return new EmbeddingResponse(List.of());
List<List<Double>> embeddingList = new ArrayList<>();
for (String inputContent : request.getInstructions()) {
var apiRequest = createZhiPuEmbeddingRequest(inputContent, request.getOptions());
ZhiPuAiApi.EmbeddingList<ZhiPuAiApi.Embedding> 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<Embedding> embeddings = apiEmbeddingResponse.data()
.stream()
.map(e -> new Embedding(e.embedding(), e.index()))
var indexCounter = new AtomicInteger(0);
List<Embedding> 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<String> createZhiPuEmbeddingRequest(String text,
EmbeddingOptions requestOptions) {
ZhiPuAiApi.EmbeddingRequest<String> 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;
}
}

68
models/spring-ai-zhipuai/src/test/java/org/springframework/ai/zhipuai/embedding/EmbeddingIT.java Normal file
View File

@@ -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);
}
}