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

Add model name and token usage to EmbeddingResponseMetadata

Browse Source 
- Moves these from free-text key/value pairs to interface
- Enables programmatic use for evaluation and observability

Signed-off-by: Thomas Vitale <ThomasVitale@users.noreply.github.com>
This commit is contained in:
Thomas Vitale
2024-07-08 18:56:19 +02:00
committed by Mark Pollack
parent 41eab270cb
commit e601f14364
12 changed files with 269 additions and 65 deletions
Download Patch File Download Diff File Expand all files Collapse all files

26
models/spring-ai-azure-openai/src/main/java/org/springframework/ai/azure/openai/AzureOpenAiEmbeddingModel.java
View File

@@ -15,11 +15,13 @@
*/
package org.springframework.ai.azure.openai;
import java.util.ArrayList;
import java.util.List;
import com.azure.ai.openai.OpenAIClient;
import com.azure.ai.openai.models.EmbeddingItem;
import com.azure.ai.openai.models.Embeddings;
import com.azure.ai.openai.models.EmbeddingsOptions;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import org.springframework.ai.azure.openai.metadata.AzureOpenAiEmbeddingUsage;
import org.springframework.ai.document.Document;
import org.springframework.ai.document.MetadataMode;
import org.springframework.ai.embedding.AbstractEmbeddingModel;
@@ -29,11 +31,8 @@ import org.springframework.ai.embedding.EmbeddingResponse;
import org.springframework.ai.embedding.EmbeddingResponseMetadata;
import org.springframework.util.Assert;
import com.azure.ai.openai.OpenAIClient;
import com.azure.ai.openai.models.EmbeddingItem;
import com.azure.ai.openai.models.Embeddings;
import com.azure.ai.openai.models.EmbeddingsOptions;
import com.azure.ai.openai.models.EmbeddingsUsage;
import java.util.ArrayList;
import java.util.List;
public class AzureOpenAiEmbeddingModel extends AbstractEmbeddingModel {
@@ -99,7 +98,8 @@ public class AzureOpenAiEmbeddingModel extends AbstractEmbeddingModel {
private EmbeddingResponse generateEmbeddingResponse(Embeddings embeddings) {
List<Embedding> data = generateEmbeddingList(embeddings.getData());
EmbeddingResponseMetadata metadata = generateMetadata(embeddings.getUsage());
EmbeddingResponseMetadata metadata = new EmbeddingResponseMetadata();
metadata.setUsage(AzureOpenAiEmbeddingUsage.from(embeddings.getUsage()));
return new EmbeddingResponse(data, metadata);
}
@@ -115,14 +115,6 @@ public class AzureOpenAiEmbeddingModel extends AbstractEmbeddingModel {
return data;
}
private EmbeddingResponseMetadata generateMetadata(EmbeddingsUsage embeddingsUsage) {
EmbeddingResponseMetadata metadata = new EmbeddingResponseMetadata();
// metadata.put("model", model);
metadata.put("prompt-tokens", embeddingsUsage.getPromptTokens());
metadata.put("total-tokens", embeddingsUsage.getTotalTokens());
return metadata;
}
public AzureOpenAiEmbeddingOptions getDefaultOptions() {
return this.defaultOptions;
}

66
models/spring-ai-azure-openai/src/main/java/org/springframework/ai/azure/openai/metadata/AzureOpenAiEmbeddingUsage.java Normal file
View File

@@ -0,0 +1,66 @@
/*
* Copyright 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.azure.openai.metadata;
import com.azure.ai.openai.models.EmbeddingsUsage;
import org.springframework.ai.chat.metadata.Usage;
import org.springframework.util.Assert;
/**
* {@link Usage} implementation for {@literal Microsoft Azure OpenAI Service} embedding.
*
* @author Thomas Vitale
* @see EmbeddingsUsage
*/
public class AzureOpenAiEmbeddingUsage implements Usage {
public static AzureOpenAiEmbeddingUsage from(EmbeddingsUsage usage) {
Assert.notNull(usage, "EmbeddingsUsage must not be null");
return new AzureOpenAiEmbeddingUsage(usage);
}
private final EmbeddingsUsage usage;
public AzureOpenAiEmbeddingUsage(EmbeddingsUsage usage) {
Assert.notNull(usage, "EmbeddingsUsage must not be null");
this.usage = usage;
}
protected EmbeddingsUsage getUsage() {
return this.usage;
}
@Override
public Long getPromptTokens() {
return (long) getUsage().getPromptTokens();
}
@Override
public Long getGenerationTokens() {
return 0L;
}
@Override
public Long getTotalTokens() {
return (long) getUsage().getTotalTokens();
}
@Override
public String toString() {
return getUsage().toString();
}
}

2
models/spring-ai-azure-openai/src/main/java/org/springframework/ai/azure/openai/metadata/AzureOpenAiUsage.java
View File

@@ -22,7 +22,7 @@ import org.springframework.ai.chat.metadata.Usage;
import org.springframework.util.Assert;
/**
* {@link Usage} implementation for {@literal Microsoft Azure OpenAI Service}.
* {@link Usage} implementation for {@literal Microsoft Azure OpenAI Service} chat.
*
* @author John Blum
* @see com.azure.ai.openai.models.CompletionsUsage

12
models/spring-ai-minimax/src/main/java/org/springframework/ai/minimax/MiniMaxEmbeddingModel.java
View File

@@ -26,6 +26,7 @@ import org.springframework.ai.embedding.EmbeddingRequest;
import org.springframework.ai.embedding.EmbeddingResponse;
import org.springframework.ai.embedding.EmbeddingResponseMetadata;
import org.springframework.ai.minimax.api.MiniMaxApi;
import org.springframework.ai.minimax.metadata.MiniMaxUsage;
import org.springframework.ai.model.ModelOptionsUtils;
import org.springframework.ai.retry.RetryUtils;
import org.springframework.retry.support.RetryTemplate;
@@ -38,6 +39,7 @@ import java.util.List;
* MiniMax Embedding Model implementation.
*
* @author Geng Rong
* @author Thomas Vitale
* @since 1.0.0 M1
*/
public class MiniMaxEmbeddingModel extends AbstractEmbeddingModel {
@@ -130,7 +132,8 @@ public class MiniMaxEmbeddingModel extends AbstractEmbeddingModel {
return new EmbeddingResponse(List.of());
}
var metadata = generateResponseMetadata(apiEmbeddingResponse.model(), apiEmbeddingResponse.totalTokens());
var metadata = new EmbeddingResponseMetadata(apiEmbeddingResponse.model(),
MiniMaxUsage.from(new MiniMaxApi.Usage(0, 0, apiEmbeddingResponse.totalTokens())));
List<Embedding> embeddings = new ArrayList<>();
for (int i = 0; i < apiEmbeddingResponse.vectors().size(); i++) {
@@ -141,11 +144,4 @@ public class MiniMaxEmbeddingModel extends AbstractEmbeddingModel {
});
}
private EmbeddingResponseMetadata generateResponseMetadata(String model, Integer totalTokens) {
EmbeddingResponseMetadata metadata = new EmbeddingResponseMetadata();
metadata.put("model", model);
metadata.put("total-tokens", totalTokens);
return metadata;
}
}

64
models/spring-ai-minimax/src/main/java/org/springframework/ai/minimax/metadata/MiniMaxUsage.java Normal file
View File

@@ -0,0 +1,64 @@
/*
* Copyright 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.minimax.metadata;
import org.springframework.ai.chat.metadata.Usage;
import org.springframework.ai.minimax.api.MiniMaxApi;
import org.springframework.util.Assert;
/**
* {@link Usage} implementation for {@literal MiniMax}.
*
* @author Thomas Vitale
*/
public class MiniMaxUsage implements Usage {
public static MiniMaxUsage from(MiniMaxApi.Usage usage) {
return new MiniMaxUsage(usage);
}
private final MiniMaxApi.Usage usage;
protected MiniMaxUsage(MiniMaxApi.Usage usage) {
Assert.notNull(usage, "MiniMax Usage must not be null");
this.usage = usage;
}
protected MiniMaxApi.Usage getUsage() {
return this.usage;
}
@Override
public Long getPromptTokens() {
return getUsage().promptTokens().longValue();
}
@Override
public Long getGenerationTokens() {
return getUsage().completionTokens().longValue();
}
@Override
public Long getTotalTokens() {
return getUsage().totalTokens().longValue();
}
@Override
public String toString() {
return getUsage().toString();
}
}

13
models/spring-ai-mistral-ai/src/main/java/org/springframework/ai/mistralai/MistralAiEmbeddingModel.java
View File

@@ -29,6 +29,7 @@ import org.springframework.ai.embedding.EmbeddingRequest;
import org.springframework.ai.embedding.EmbeddingResponse;
import org.springframework.ai.embedding.EmbeddingResponseMetadata;
import org.springframework.ai.mistralai.api.MistralAiApi;
import org.springframework.ai.mistralai.metadata.MistralAiUsage;
import org.springframework.ai.model.ModelOptionsUtils;
import org.springframework.ai.retry.RetryUtils;
import org.springframework.retry.support.RetryTemplate;
@@ -36,6 +37,7 @@ import org.springframework.util.Assert;
/**
* @author Ricken Bazolo
* @author Thomas Vitale
* @since 0.8.1
*/
public class MistralAiEmbeddingModel extends AbstractEmbeddingModel {
@@ -100,7 +102,8 @@ public class MistralAiEmbeddingModel extends AbstractEmbeddingModel {
return new EmbeddingResponse(List.of());
}
var metadata = generateResponseMetadata(apiEmbeddingResponse.model(), apiEmbeddingResponse.usage());
var metadata = new EmbeddingResponseMetadata(apiEmbeddingResponse.model(),
MistralAiUsage.from(apiEmbeddingResponse.usage()));
var embeddings = apiEmbeddingResponse.data()
.stream()
@@ -118,12 +121,4 @@ public class MistralAiEmbeddingModel extends AbstractEmbeddingModel {
return this.embed(document.getFormattedContent(this.metadataMode));
}
private EmbeddingResponseMetadata generateResponseMetadata(String model, MistralAiApi.Usage usage) {
var metadata = new EmbeddingResponseMetadata();
metadata.put("model", model);
metadata.put("prompt-tokens", usage.promptTokens());
metadata.put("total-tokens", usage.totalTokens());
return metadata;
}
}

15
models/spring-ai-openai/src/main/java/org/springframework/ai/openai/OpenAiEmbeddingModel.java
View File

@@ -31,7 +31,7 @@ import org.springframework.ai.embedding.EmbeddingResponseMetadata;
import org.springframework.ai.model.ModelOptionsUtils;
import org.springframework.ai.openai.api.OpenAiApi;
import org.springframework.ai.openai.api.OpenAiApi.EmbeddingList;
import org.springframework.ai.openai.api.OpenAiApi.Usage;
import org.springframework.ai.openai.metadata.OpenAiUsage;
import org.springframework.ai.retry.RetryUtils;
import org.springframework.retry.support.RetryTemplate;
import org.springframework.util.Assert;
@@ -40,6 +40,7 @@ import org.springframework.util.Assert;
* Open AI Embedding Model implementation.
*
* @author Christian Tzolov
* @author Thomas Vitale
*/
public class OpenAiEmbeddingModel extends AbstractEmbeddingModel {
@@ -134,7 +135,8 @@ public class OpenAiEmbeddingModel extends AbstractEmbeddingModel {
return new EmbeddingResponse(List.of());
}
var metadata = generateResponseMetadata(apiEmbeddingResponse.model(), apiEmbeddingResponse.usage());
var metadata = new EmbeddingResponseMetadata(apiEmbeddingResponse.model(),
OpenAiUsage.from(apiEmbeddingResponse.usage()));
List<Embedding> embeddings = apiEmbeddingResponse.data()
.stream()
@@ -146,13 +148,4 @@ public class OpenAiEmbeddingModel extends AbstractEmbeddingModel {
});
}
private EmbeddingResponseMetadata generateResponseMetadata(String model, 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;
}
}

13
models/spring-ai-qianfan/src/main/java/org/springframework/ai/qianfan/QianFanEmbeddingModel.java
View File

@@ -28,6 +28,7 @@ import org.springframework.ai.embedding.EmbeddingResponseMetadata;
import org.springframework.ai.model.ModelOptionsUtils;
import org.springframework.ai.qianfan.api.QianFanApi;
import org.springframework.ai.qianfan.api.QianFanApi.EmbeddingList;
import org.springframework.ai.qianfan.metadata.QianFanUsage;
import org.springframework.ai.retry.RetryUtils;
import org.springframework.retry.support.RetryTemplate;
import org.springframework.util.Assert;
@@ -38,6 +39,7 @@ import java.util.List;
* QianFan Embedding Client implementation.
*
* @author Geng Rong
* @author Thomas Vitale
* @since 1.0
*/
public class QianFanEmbeddingModel extends AbstractEmbeddingModel {
@@ -135,7 +137,8 @@ public class QianFanEmbeddingModel extends AbstractEmbeddingModel {
+ ", message:" + apiEmbeddingResponse.errorNsg());
}
var metadata = generateResponseMetadata(apiEmbeddingResponse.model(), apiEmbeddingResponse.usage());
var metadata = new EmbeddingResponseMetadata(apiEmbeddingResponse.model(),
QianFanUsage.from(apiEmbeddingResponse.usage()));
List<Embedding> embeddings = apiEmbeddingResponse.data()
.stream()
@@ -146,12 +149,4 @@ public class QianFanEmbeddingModel extends AbstractEmbeddingModel {
});
}
private EmbeddingResponseMetadata generateResponseMetadata(String model, QianFanApi.Usage usage) {
EmbeddingResponseMetadata metadata = new EmbeddingResponseMetadata();
metadata.put("model", model);
metadata.put("prompt-tokens", usage.promptTokens());
metadata.put("total-tokens", usage.totalTokens());
return metadata;
}
}

64
models/spring-ai-qianfan/src/main/java/org/springframework/ai/qianfan/metadata/QianFanUsage.java Normal file
View File

@@ -0,0 +1,64 @@
/*
* Copyright 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.qianfan.metadata;
import org.springframework.ai.chat.metadata.Usage;
import org.springframework.ai.qianfan.api.QianFanApi;
import org.springframework.util.Assert;
/**
* {@link Usage} implementation for {@literal QianFan}.
*
* @author Thomas Vitale
*/
public class QianFanUsage implements Usage {
public static QianFanUsage from(QianFanApi.Usage usage) {
return new QianFanUsage(usage);
}
private final QianFanApi.Usage usage;
protected QianFanUsage(QianFanApi.Usage usage) {
Assert.notNull(usage, "QianFan Usage must not be null");
this.usage = usage;
}
protected QianFanApi.Usage getUsage() {
return this.usage;
}
@Override
public Long getPromptTokens() {
return getUsage().promptTokens().longValue();
}
@Override
public Long getGenerationTokens() {
return 0L;
}
@Override
public Long getTotalTokens() {
return getUsage().totalTokens().longValue();
}
@Override
public String toString() {
return getUsage().toString();
}
}

4
spring-ai-core/src/main/java/org/springframework/ai/embedding/EmbeddingResponse.java
View File

@@ -29,12 +29,12 @@ public class EmbeddingResponse implements ModelResponse<Embedding> {
/**
* Embedding data.
*/
private List<Embedding> embeddings;
private final List<Embedding> embeddings;
/**
* Embedding metadata.
*/
private EmbeddingResponseMetadata metadata = new EmbeddingResponseMetadata();
private final EmbeddingResponseMetadata metadata;
/**
* Creates a new {@link EmbeddingResponse} instance with empty metadata.

43
spring-ai-core/src/main/java/org/springframework/ai/embedding/EmbeddingResponseMetadata.java
View File

@@ -15,31 +15,62 @@
*/
package org.springframework.ai.embedding;
import java.io.Serial;
import java.util.HashMap;
import java.util.Map;
import org.springframework.ai.chat.metadata.EmptyUsage;
import org.springframework.ai.chat.metadata.Usage;
import org.springframework.ai.model.ResponseMetadata;
/**
* Common AI provider metadata returned in an embedding response.
*
* @author Christian Tzolov
* @author Thomas Vitale
*/
public class EmbeddingResponseMetadata extends HashMap<String, Object> implements ResponseMetadata {
@Serial
private static final long serialVersionUID = 1L;
private String model;
private Usage usage;
public EmbeddingResponseMetadata() {
}
public EmbeddingResponseMetadata(int initialCapacity) {
super(initialCapacity);
}
public EmbeddingResponseMetadata(int initialCapacity, float loadFactor) {
super(initialCapacity, loadFactor);
public EmbeddingResponseMetadata(String model, Usage usage) {
this.model = model;
this.usage = usage;
}
public EmbeddingResponseMetadata(Map<String, ?> metadata) {
super(metadata);
}
/**
* The model that handled the request.
*/
public String getModel() {
return this.model != null ? this.model : "";
}
public void setModel(String model) {
this.model = model;
}
/**
* The AI provider specific metadata on API usage.
* @see Usage
*/
public Usage getUsage() {
return this.usage != null ? this.usage : new EmptyUsage();
}
public void setUsage(Usage usage) {
this.usage = usage;
}
}

12
spring-ai-core/src/main/java/org/springframework/ai/image/ImageResponseMetadata.java
View File

@@ -15,14 +15,22 @@
*/
package org.springframework.ai.image;
import org.springframework.ai.chat.metadata.ChatResponseMetadata;
import org.springframework.ai.model.ResponseMetadata;
import java.util.HashMap;
/**
* Represents metadata associated with an image response. It provides additional
* information about the generative response from an AI model, including the creation
* timestamp of the generated image.
*
* @author Mark Pollack
* @author Thomas Vitale
* @since 1.0.0
*/
public interface ImageResponseMetadata extends ResponseMetadata {
static class DefaultImageResponseMetadata extends HashMap<String, Object> implements ImageResponseMetadata {
class DefaultImageResponseMetadata extends HashMap<String, Object> implements ImageResponseMetadata {
}