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

Streamline EmbeddingOptions

Browse Source 
* Add model and dimensions to option abstraction
* Use abstraction in Observations directly instead of dedicated implementation
* Clean-up the merge of runtime and default embedding options in OpenAI

Relates to #gh-1148

Signed-off-by: Thomas Vitale <ThomasVitale@users.noreply.github.com>
This commit is contained in:
Thomas Vitale
2024-08-07 15:16:17 +02:00
committed by Christian Tzolov
parent 17ba1fc3ba
commit 08ccc10a16
29 changed files with 352 additions and 249 deletions
Download Patch File Download Diff File Expand all files Collapse all files

7
models/spring-ai-azure-openai/src/main/java/org/springframework/ai/azure/openai/AzureOpenAiEmbeddingOptions.java
View File

@@ -23,6 +23,7 @@ import org.springframework.ai.embedding.EmbeddingOptions;
* The configuration information for the embedding requests.
*
* @author Christian Tzolov
* @author Thomas Vitale
* @since 0.8.0
*/
public class AzureOpenAiEmbeddingOptions implements EmbeddingOptions {
@@ -123,6 +124,11 @@ public class AzureOpenAiEmbeddingOptions implements EmbeddingOptions {
}
@Override
public String getModel() {
return getDeploymentName();
}
public String getUser() {
return this.user;
}
@@ -147,6 +153,7 @@ public class AzureOpenAiEmbeddingOptions implements EmbeddingOptions {
this.inputType = inputType;
}
@Override
public Integer getDimensions() {
return this.dimensions;
}

11
models/spring-ai-bedrock/src/main/java/org/springframework/ai/bedrock/cohere/BedrockCohereEmbeddingOptions.java
View File

@@ -25,6 +25,7 @@ import org.springframework.ai.embedding.EmbeddingOptions;
/**
* @author Christian Tzolov
* @author Thomas Vitale
*/
@JsonInclude(Include.NON_NULL)
public class BedrockCohereEmbeddingOptions implements EmbeddingOptions {
@@ -86,4 +87,14 @@ public class BedrockCohereEmbeddingOptions implements EmbeddingOptions {
this.truncate = truncate;
}
@Override
public String getModel() {
return null;
}
@Override
public Integer getDimensions() {
return null;
}
}

11
models/spring-ai-bedrock/src/main/java/org/springframework/ai/bedrock/titan/BedrockTitanEmbeddingOptions.java
View File

@@ -24,6 +24,7 @@ import org.springframework.util.Assert;
/**
* @author Wei Jiang
* @author Thomas Vitale
*/
@JsonInclude(Include.NON_NULL)
public class BedrockTitanEmbeddingOptions implements EmbeddingOptions {
@@ -62,4 +63,14 @@ public class BedrockTitanEmbeddingOptions implements EmbeddingOptions {
this.inputType = inputType;
}
@Override
public String getModel() {
return null;
}
@Override
public Integer getDimensions() {
return null;
}
}

7
models/spring-ai-minimax/src/main/java/org/springframework/ai/minimax/MiniMaxEmbeddingOptions.java
View File

@@ -24,6 +24,7 @@ import org.springframework.ai.embedding.EmbeddingOptions;
* This class represents the options for MiniMax embedding.
*
* @author Geng Rong
* @author Thomas Vitale
* @since 1.0.0 M1
*/
@JsonInclude(Include.NON_NULL)
@@ -59,6 +60,7 @@ public class MiniMaxEmbeddingOptions implements EmbeddingOptions {
}
@Override
public String getModel() {
return this.model;
}
@@ -67,4 +69,9 @@ public class MiniMaxEmbeddingOptions implements EmbeddingOptions {
this.model = model;
}
@Override
public Integer getDimensions() {
return null;
}
}

7
models/spring-ai-mistral-ai/src/main/java/org/springframework/ai/mistralai/MistralAiEmbeddingOptions.java
View File

@@ -22,6 +22,7 @@ import org.springframework.ai.embedding.EmbeddingOptions;
/**
* @author Ricken Bazolo
* @author Thomas Vitale
* @since 0.8.1
*/
@JsonInclude(Include.NON_NULL)
@@ -41,6 +42,7 @@ public class MistralAiEmbeddingOptions implements EmbeddingOptions {
return new Builder();
}
@Override
public String getModel() {
return this.model;
}
@@ -57,6 +59,11 @@ public class MistralAiEmbeddingOptions implements EmbeddingOptions {
this.encodingFormat = encodingFormat;
}
@Override
public Integer getDimensions() {
return null;
}
public static class Builder {
protected MistralAiEmbeddingOptions options;

6
models/spring-ai-ollama/src/main/java/org/springframework/ai/ollama/api/OllamaOptions.java
View File

@@ -498,6 +498,7 @@ public class OllamaOptions implements FunctionCallingOptions, ChatOptions, Embed
// -------------------
// Getters and Setters
// -------------------
@Override
public String getModel() {
return model;
}
@@ -762,6 +763,11 @@ public class OllamaOptions implements FunctionCallingOptions, ChatOptions, Embed
this.truncate = truncate;
}
@Override
public Integer getDimensions() {
return null;
}
@Override
public List<FunctionCallback> getFunctionCallbacks() {
return this.functionCallbacks;

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

@@ -30,7 +30,6 @@ import org.springframework.ai.embedding.observation.DefaultEmbeddingModelObserva
import org.springframework.ai.embedding.observation.EmbeddingModelObservationConvention;
import org.springframework.ai.embedding.observation.EmbeddingModelObservationDocumentation;
import org.springframework.ai.embedding.observation.EmbeddingModelObservationContext;
import org.springframework.ai.embedding.observation.EmbeddingModelRequestOptions;
import org.springframework.ai.model.ModelOptionsUtils;
import org.springframework.ai.observation.AiOperationMetadata;
import org.springframework.ai.observation.conventions.AiOperationType;
@@ -39,9 +38,9 @@ import org.springframework.ai.openai.api.OpenAiApi;
import org.springframework.ai.openai.api.OpenAiApi.EmbeddingList;
import org.springframework.ai.openai.metadata.OpenAiUsage;
import org.springframework.ai.retry.RetryUtils;
import org.springframework.lang.Nullable;
import org.springframework.retry.support.RetryTemplate;
import org.springframework.util.Assert;
import org.springframework.util.StringUtils;
import java.util.List;
@@ -126,7 +125,7 @@ public class OpenAiEmbeddingModel extends AbstractEmbeddingModel {
*/
public OpenAiEmbeddingModel(OpenAiApi openAiApi, MetadataMode metadataMode, OpenAiEmbeddingOptions options,
RetryTemplate retryTemplate, ObservationRegistry observationRegistry) {
Assert.notNull(openAiApi, "OpenAiService must not be null");
Assert.notNull(openAiApi, "openAiApi must not be null");
Assert.notNull(metadataMode, "metadataMode must not be null");
Assert.notNull(options, "options must not be null");
Assert.notNull(retryTemplate, "retryTemplate must not be null");
@@ -147,12 +146,13 @@ public class OpenAiEmbeddingModel extends AbstractEmbeddingModel {
@Override
public EmbeddingResponse call(EmbeddingRequest request) {
org.springframework.ai.openai.api.OpenAiApi.EmbeddingRequest<List<String>> apiRequest = createRequest(request);
OpenAiEmbeddingOptions requestOptions = mergeOptions(request.getOptions(), this.defaultOptions);
OpenAiApi.EmbeddingRequest<List<String>> apiRequest = createRequest(request, requestOptions);
var observationContext = EmbeddingModelObservationContext.builder()
.embeddingRequest(request)
.operationMetadata(buildOperationMetadata())
.requestOptions(buildRequestOptions(apiRequest))
.requestOptions(requestOptions)
.build();
return EmbeddingModelObservationDocumentation.EMBEDDING_MODEL_OPERATION
@@ -183,21 +183,31 @@ public class OpenAiEmbeddingModel extends AbstractEmbeddingModel {
});
}
@SuppressWarnings("unchecked")
private OpenAiApi.EmbeddingRequest<List<String>> createRequest(EmbeddingRequest request) {
org.springframework.ai.openai.api.OpenAiApi.EmbeddingRequest<List<String>> apiRequest = (this.defaultOptions != null)
? new org.springframework.ai.openai.api.OpenAiApi.EmbeddingRequest<>(request.getInstructions(),
this.defaultOptions.getModel(), this.defaultOptions.getEncodingFormat(),
this.defaultOptions.getDimensions(), this.defaultOptions.getUser())
: new org.springframework.ai.openai.api.OpenAiApi.EmbeddingRequest<>(request.getInstructions(),
OpenAiApi.DEFAULT_EMBEDDING_MODEL);
private OpenAiApi.EmbeddingRequest<List<String>> createRequest(EmbeddingRequest request,
OpenAiEmbeddingOptions requestOptions) {
return new OpenAiApi.EmbeddingRequest<>(request.getInstructions(), requestOptions.getModel(),
requestOptions.getEncodingFormat(), requestOptions.getDimensions(), requestOptions.getUser());
}
if (request.getOptions() != null && !EmbeddingOptions.EMPTY.equals(request.getOptions())) {
apiRequest = ModelOptionsUtils.merge(request.getOptions(), apiRequest,
org.springframework.ai.openai.api.OpenAiApi.EmbeddingRequest.class);
/**
* Merge runtime and default {@link EmbeddingOptions} to compute the final options to
* use in the request.
*/
private OpenAiEmbeddingOptions mergeOptions(@Nullable EmbeddingOptions runtimeOptions,
OpenAiEmbeddingOptions defaultOptions) {
if (runtimeOptions == null) {
return defaultOptions;
}
return apiRequest;
return OpenAiEmbeddingOptions.builder()
// Handle portable embedding options
.withModel(ModelOptionsUtils.mergeOption(runtimeOptions.getModel(), defaultOptions.getModel()))
.withDimensions(
ModelOptionsUtils.mergeOption(runtimeOptions.getDimensions(), defaultOptions.getDimensions()))
// Handle OpenAI specific embedding options
.withEncodingFormat(defaultOptions.getEncodingFormat())
.withUser(defaultOptions.getUser())
.build();
}
private AiOperationMetadata buildOperationMetadata() {
@@ -207,14 +217,6 @@ public class OpenAiEmbeddingModel extends AbstractEmbeddingModel {
.build();
}
private EmbeddingModelRequestOptions buildRequestOptions(OpenAiApi.EmbeddingRequest<List<String>> request) {
return EmbeddingModelRequestOptions.builder()
.model(StringUtils.hasText(request.model()) ? request.model() : "unknown")
.dimensions(request.dimensions())
.encodingFormat(request.encodingFormat())
.build();
}
/**
* Use the provided convention for reporting observation data
* @param observationConvention The provided convention

2
models/spring-ai-openai/src/main/java/org/springframework/ai/openai/OpenAiEmbeddingOptions.java
View File

@@ -85,6 +85,7 @@ public class OpenAiEmbeddingOptions implements EmbeddingOptions {
}
@Override
public String getModel() {
return this.model;
}
@@ -101,6 +102,7 @@ public class OpenAiEmbeddingOptions implements EmbeddingOptions {
this.encodingFormat = encodingFormat;
}
@Override
public Integer getDimensions() {
return this.dimensions;
}

27
models/spring-ai-openai/src/main/java/org/springframework/ai/openai/OpenAiImageOptions.java
View File

@@ -190,7 +190,18 @@ public class OpenAiImageOptions implements ImageOptions {
@Override
public Integer getWidth() {
return this.width;
if (this.width != null) {
return this.width;
}
else if (this.size != null) {
try {
return Integer.parseInt(this.size.split("x")[0]);
}
catch (NumberFormatException ex) {
return null;
}
}
return null;
}
public void setWidth(Integer width) {
@@ -200,7 +211,18 @@ public class OpenAiImageOptions implements ImageOptions {
@Override
public Integer getHeight() {
return this.height;
if (this.height != null) {
return this.height;
}
else if (this.size != null) {
try {
return Integer.parseInt(this.size.split("x")[1]);
}
catch (NumberFormatException ex) {
return null;
}
}
return null;
}
public void setHeight(Integer height) {
@@ -230,7 +252,6 @@ public class OpenAiImageOptions implements ImageOptions {
}
public String getSize() {
if (this.size != null) {
return this.size;
}

76
models/spring-ai-openai/src/test/java/org/springframework/ai/openai/OpenAiImageOptionsTests.java Normal file
View File

@@ -0,0 +1,76 @@
/*
* 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.openai;
import org.junit.jupiter.api.Test;
import static org.assertj.core.api.Assertions.assertThat;
/**
* Unit tests for {@link OpenAiImageOptions}.
*
* @author Thomas Vitale
*/
class OpenAiImageOptionsTests {
@Test
void whenImageDimensionsAreAllUnset() {
OpenAiImageOptions options = new OpenAiImageOptions();
assertThat(options.getHeight()).isEqualTo(null);
assertThat(options.getWidth()).isEqualTo(null);
assertThat(options.getSize()).isEqualTo(null);
}
@Test
void whenSizeIsSet() {
OpenAiImageOptions options = new OpenAiImageOptions();
options.setSize("1920x1080");
assertThat(options.getHeight()).isEqualTo(1080);
assertThat(options.getWidth()).isEqualTo(1920);
assertThat(options.getSize()).isEqualTo("1920x1080");
}
@Test
void whenWidthAndHeightAreSet() {
OpenAiImageOptions options = new OpenAiImageOptions();
options.setWidth(1920);
options.setHeight(1080);
assertThat(options.getHeight()).isEqualTo(1080);
assertThat(options.getWidth()).isEqualTo(1920);
assertThat(options.getSize()).isEqualTo("1920x1080");
}
@Test
void whenWidthIsSet() {
OpenAiImageOptions options = new OpenAiImageOptions();
options.setWidth(1920);
assertThat(options.getHeight()).isEqualTo(null);
assertThat(options.getWidth()).isEqualTo(1920);
// This is because "setWidth()" computes "size" without checking for null values.
assertThat(options.getSize()).isEqualTo("1920xnull");
}
@Test
void whenHeightIsSet() {
OpenAiImageOptions options = new OpenAiImageOptions();
options.setHeight(1080);
assertThat(options.getHeight()).isEqualTo(1080);
assertThat(options.getWidth()).isEqualTo(null);
// This is because "setHeight()" computes "size" without checking for null values.
assertThat(options.getSize()).isEqualTo("nullx1080");
}
}

1
models/spring-ai-openai/src/test/java/org/springframework/ai/openai/embedding/OpenAiEmbeddingModelObservationIT.java
View File

@@ -84,7 +84,6 @@ public class OpenAiEmbeddingModelObservationIT {
OpenAiApi.EmbeddingModel.TEXT_EMBEDDING_3_SMALL.getValue())
.hasLowCardinalityKeyValue(LowCardinalityKeyNames.RESPONSE_MODEL.asString(), responseMetadata.getModel())
.hasHighCardinalityKeyValue(HighCardinalityKeyNames.REQUEST_EMBEDDING_DIMENSIONS.asString(), "1536")
.hasHighCardinalityKeyValue(HighCardinalityKeyNames.REQUEST_EMBEDDING_ENCODING_FORMAT.asString(), "float")
.hasHighCardinalityKeyValue(HighCardinalityKeyNames.USAGE_INPUT_TOKENS.asString(),
String.valueOf(responseMetadata.getUsage().getPromptTokens()))
.hasHighCardinalityKeyValue(HighCardinalityKeyNames.USAGE_TOTAL_TOKENS.asString(),

11
models/spring-ai-postgresml/src/main/java/org/springframework/ai/postgresml/PostgresMlEmbeddingOptions.java
View File

@@ -28,6 +28,7 @@ import org.springframework.ai.postgresml.PostgresMlEmbeddingModel.VectorType;
/**
* @author Christian Tzolov
* @author Thomas Vitale
*/
@JsonInclude(Include.NON_NULL)
public class PostgresMlEmbeddingOptions implements EmbeddingOptions {
@@ -130,4 +131,14 @@ public class PostgresMlEmbeddingOptions implements EmbeddingOptions {
this.metadataMode = metadataMode;
}
@Override
public String getModel() {
return null;
}
@Override
public Integer getDimensions() {
return null;
}
}

7
models/spring-ai-qianfan/src/main/java/org/springframework/ai/qianfan/QianFanEmbeddingOptions.java
View File

@@ -24,6 +24,7 @@ import org.springframework.ai.embedding.EmbeddingOptions;
* This class represents the options for QianFan embedding.
*
* @author Geng Rong
* @author Thomas Vitale
* @since 1.0
*/
@JsonInclude(Include.NON_NULL)
@@ -70,6 +71,7 @@ public class QianFanEmbeddingOptions implements EmbeddingOptions {
}
@Override
public String getModel() {
return this.model;
}
@@ -86,4 +88,9 @@ public class QianFanEmbeddingOptions implements EmbeddingOptions {
this.user = user;
}
@Override
public Integer getDimensions() {
return null;
}
}

2
models/spring-ai-vertex-ai-embedding/src/main/java/org/springframework/ai/vertexai/embedding/multimodal/VertexAiMultimodalEmbeddingOptions.java
View File

@@ -168,6 +168,7 @@ public class VertexAiMultimodalEmbeddingOptions implements EmbeddingOptions {
}
@Override
public String getModel() {
return this.model;
}
@@ -176,6 +177,7 @@ public class VertexAiMultimodalEmbeddingOptions implements EmbeddingOptions {
this.model = model;
}
@Override
public Integer getDimensions() {
return this.dimensions;
}

2
models/spring-ai-vertex-ai-embedding/src/main/java/org/springframework/ai/vertexai/embedding/text/VertexAiTextEmbeddingOptions.java
View File

@@ -183,6 +183,7 @@ public class VertexAiTextEmbeddingOptions implements EmbeddingOptions {
return this;
}
@Override
public String getModel() {
return this.model;
}
@@ -199,6 +200,7 @@ public class VertexAiTextEmbeddingOptions implements EmbeddingOptions {
this.taskType = taskType;
}
@Override
public Integer getDimensions() {
return this.dimensions;
}

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

@@ -24,6 +24,7 @@ import org.springframework.ai.embedding.EmbeddingOptions;
* The ZhiPuAiEmbeddingOptions class represents the options for ZhiPuAI embedding.
*
* @author Geng Rong
* @author Thomas Vitale
* @since 1.0.0 M1
*/
@JsonInclude(Include.NON_NULL)
@@ -59,6 +60,7 @@ public class ZhiPuAiEmbeddingOptions implements EmbeddingOptions {
}
@Override
public String getModel() {
return this.model;
}
@@ -67,4 +69,9 @@ public class ZhiPuAiEmbeddingOptions implements EmbeddingOptions {
this.model = model;
}
@Override
public Integer getDimensions() {
return null;
}
}

5
spring-ai-core/src/main/java/org/springframework/ai/embedding/DocumentEmbeddingRequest.java
View File

@@ -25,6 +25,7 @@ import org.springframework.ai.model.ModelRequest;
* Represents a request to embed a list of documents.
*
* @author Christian Tzolov
* @author Thomas Vitale
* @since 1.0.0
*/
public class DocumentEmbeddingRequest implements ModelRequest<List<Document>> {
@@ -34,11 +35,11 @@ public class DocumentEmbeddingRequest implements ModelRequest<List<Document>> {
private final EmbeddingOptions options;
public DocumentEmbeddingRequest(Document... inputs) {
this(Arrays.asList(inputs), EmbeddingOptions.EMPTY);
this(Arrays.asList(inputs), EmbeddingOptionsBuilder.builder().build());
}
public DocumentEmbeddingRequest(List<Document> inputs) {
this(inputs, EmbeddingOptions.EMPTY);
this(inputs, EmbeddingOptionsBuilder.builder().build());
}
public DocumentEmbeddingRequest(List<Document> inputs, EmbeddingOptions options) {

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

@@ -53,7 +53,7 @@ public interface EmbeddingModel extends Model<EmbeddingRequest, EmbeddingRespons
*/
default List<List<Double>> embed(List<String> texts) {
Assert.notNull(texts, "Texts must not be null");
return this.call(new EmbeddingRequest(texts, EmbeddingOptions.EMPTY))
return this.call(new EmbeddingRequest(texts, EmbeddingOptionsBuilder.builder().build()))
.getResults()
.stream()
.map(Embedding::getOutput)
@@ -67,7 +67,7 @@ public interface EmbeddingModel extends Model<EmbeddingRequest, EmbeddingRespons
*/
default EmbeddingResponse embedForResponse(List<String> texts) {
Assert.notNull(texts, "Texts must not be null");
return this.call(new EmbeddingRequest(texts, EmbeddingOptions.EMPTY));
return this.call(new EmbeddingRequest(texts, EmbeddingOptionsBuilder.builder().build()));
}
/**

14
spring-ai-core/src/main/java/org/springframework/ai/embedding/EmbeddingOptions.java
View File

@@ -16,16 +16,24 @@
package org.springframework.ai.embedding;
import org.springframework.ai.model.ModelOptions;
import org.springframework.lang.Nullable;
/**
* @author Christian Tzolov
* @author Thomas Vitale
*/
public interface EmbeddingOptions extends ModelOptions {
class EmptyEmbeddingOptions implements EmbeddingOptions {
/**
* Use the {@link EmbeddingOptionsBuilder} instead.
*/
@Deprecated(since = "1.0.0", forRemoval = true)
EmbeddingOptions EMPTY = EmbeddingOptionsBuilder.builder().build();
}
@Nullable
String getModel();
EmbeddingOptions EMPTY = new EmptyEmbeddingOptions();
@Nullable
Integer getDimensions();
}

73
spring-ai-core/src/main/java/org/springframework/ai/embedding/EmbeddingOptionsBuilder.java Normal file
View File

@@ -0,0 +1,73 @@
/*
* 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.embedding;
/**
* @author Thomas Vitale
* @since 1.0.0
*/
public class EmbeddingOptionsBuilder {
private final DefaultEmbeddingOptions embeddingOptions = new DefaultEmbeddingOptions();
private EmbeddingOptionsBuilder() {
}
public static EmbeddingOptionsBuilder builder() {
return new EmbeddingOptionsBuilder();
}
public EmbeddingOptionsBuilder withModel(String model) {
embeddingOptions.setModel(model);
return this;
}
public EmbeddingOptionsBuilder withDimensions(Integer dimensions) {
embeddingOptions.setDimensions(dimensions);
return this;
}
public EmbeddingOptions build() {
return embeddingOptions;
}
private static class DefaultEmbeddingOptions implements EmbeddingOptions {
private String model;
private Integer dimensions;
@Override
public String getModel() {
return this.model;
}
public void setModel(String model) {
this.model = model;
}
@Override
public Integer getDimensions() {
return this.dimensions;
}
public void setDimensions(Integer dimensions) {
this.dimensions = dimensions;
}
}
}

35
spring-ai-core/src/main/java/org/springframework/ai/embedding/observation/DefaultEmbeddingModelObservationConvention.java
View File

@@ -27,6 +27,9 @@ import org.springframework.util.StringUtils;
*/
public class DefaultEmbeddingModelObservationConvention implements EmbeddingModelObservationConvention {
private static final KeyValue REQUEST_MODEL_NONE = KeyValue
.of(EmbeddingModelObservationDocumentation.LowCardinalityKeyNames.REQUEST_MODEL, KeyValue.NONE_VALUE);
private static final KeyValue RESPONSE_MODEL_NONE = KeyValue
.of(EmbeddingModelObservationDocumentation.LowCardinalityKeyNames.RESPONSE_MODEL, KeyValue.NONE_VALUE);
@@ -34,10 +37,6 @@ public class DefaultEmbeddingModelObservationConvention implements EmbeddingMode
EmbeddingModelObservationDocumentation.HighCardinalityKeyNames.REQUEST_EMBEDDING_DIMENSIONS,
KeyValue.NONE_VALUE);
private static final KeyValue REQUEST_EMBEDDING_ENCODING_FORMAT_NONE = KeyValue.of(
EmbeddingModelObservationDocumentation.HighCardinalityKeyNames.REQUEST_EMBEDDING_ENCODING_FORMAT,
KeyValue.NONE_VALUE);
private static final KeyValue USAGE_INPUT_TOKENS_NONE = KeyValue
.of(EmbeddingModelObservationDocumentation.HighCardinalityKeyNames.USAGE_INPUT_TOKENS, KeyValue.NONE_VALUE);
@@ -53,8 +52,11 @@ public class DefaultEmbeddingModelObservationConvention implements EmbeddingMode
@Override
public String getContextualName(EmbeddingModelObservationContext context) {
return "%s %s".formatted(context.getOperationMetadata().operationType(),
context.getRequestOptions().getModel());
if (StringUtils.hasText(context.getRequestOptions().getModel())) {
return "%s %s".formatted(context.getOperationMetadata().operationType(),
context.getRequestOptions().getModel());
}
return context.getOperationMetadata().operationType();
}
@Override
@@ -65,7 +67,7 @@ public class DefaultEmbeddingModelObservationConvention implements EmbeddingMode
protected KeyValue aiOperationType(EmbeddingModelObservationContext context) {
return KeyValue.of(EmbeddingModelObservationDocumentation.LowCardinalityKeyNames.AI_OPERATION_TYPE,
context.getOperationMetadata().operationType());
context.getOperationType());
}
protected KeyValue aiProvider(EmbeddingModelObservationContext context) {
@@ -74,8 +76,11 @@ public class DefaultEmbeddingModelObservationConvention implements EmbeddingMode
}
protected KeyValue requestModel(EmbeddingModelObservationContext context) {
return KeyValue.of(EmbeddingModelObservationDocumentation.LowCardinalityKeyNames.REQUEST_MODEL,
context.getRequestOptions().getModel());
if (StringUtils.hasText(context.getRequestOptions().getModel())) {
return KeyValue.of(EmbeddingModelObservationDocumentation.LowCardinalityKeyNames.REQUEST_MODEL,
context.getRequestOptions().getModel());
}
return REQUEST_MODEL_NONE;
}
protected KeyValue responseModel(EmbeddingModelObservationContext context) {
@@ -89,8 +94,7 @@ public class DefaultEmbeddingModelObservationConvention implements EmbeddingMode
@Override
public KeyValues getHighCardinalityKeyValues(EmbeddingModelObservationContext context) {
return KeyValues.of(requestEmbeddingDimension(context), requestEmbeddingFormat(context),
usageInputTokens(context), usageTotalTokens(context));
return KeyValues.of(requestEmbeddingDimension(context), usageInputTokens(context), usageTotalTokens(context));
}
// Request
@@ -104,15 +108,6 @@ public class DefaultEmbeddingModelObservationConvention implements EmbeddingMode
return REQUEST_EMBEDDING_DIMENSION_NONE;
}
protected KeyValue requestEmbeddingFormat(EmbeddingModelObservationContext context) {
if (StringUtils.hasText(context.getRequestOptions().getEncodingFormat())) {
return KeyValue.of(
EmbeddingModelObservationDocumentation.HighCardinalityKeyNames.REQUEST_EMBEDDING_ENCODING_FORMAT,
context.getRequestOptions().getEncodingFormat());
}
return REQUEST_EMBEDDING_ENCODING_FORMAT_NONE;
}
// Response
protected KeyValue usageInputTokens(EmbeddingModelObservationContext context) {

16
spring-ai-core/src/main/java/org/springframework/ai/embedding/observation/EmbeddingModelObservationContext.java
View File

@@ -15,10 +15,12 @@
*/
package org.springframework.ai.embedding.observation;
import org.springframework.ai.embedding.EmbeddingOptions;
import org.springframework.ai.embedding.EmbeddingRequest;
import org.springframework.ai.embedding.EmbeddingResponse;
import org.springframework.ai.model.observation.ModelObservationContext;
import org.springframework.ai.observation.AiOperationMetadata;
import org.springframework.ai.observation.conventions.AiOperationType;
import org.springframework.util.Assert;
/**
@@ -29,19 +31,23 @@ import org.springframework.util.Assert;
*/
public class EmbeddingModelObservationContext extends ModelObservationContext<EmbeddingRequest, EmbeddingResponse> {
private final EmbeddingModelRequestOptions requestOptions;
private final EmbeddingOptions requestOptions;
EmbeddingModelObservationContext(EmbeddingRequest embeddingRequest, AiOperationMetadata operationMetadata,
EmbeddingModelRequestOptions requestOptions) {
EmbeddingOptions requestOptions) {
super(embeddingRequest, operationMetadata);
Assert.notNull(requestOptions, "requestOptions cannot be null");
this.requestOptions = requestOptions;
}
public EmbeddingModelRequestOptions getRequestOptions() {
public EmbeddingOptions getRequestOptions() {
return requestOptions;
}
public String getOperationType() {
return AiOperationType.EMBEDDING.value();
}
public static Builder builder() {
return new Builder();
}
@@ -52,7 +58,7 @@ public class EmbeddingModelObservationContext extends ModelObservationContext<Em
private AiOperationMetadata operationMetadata;
private EmbeddingModelRequestOptions requestOptions;
private EmbeddingOptions requestOptions;
private Builder() {
}
@@ -67,7 +73,7 @@ public class EmbeddingModelObservationContext extends ModelObservationContext<Em
return this;
}
public Builder requestOptions(EmbeddingModelRequestOptions requestOptions) {
public Builder requestOptions(EmbeddingOptions requestOptions) {
this.requestOptions = requestOptions;
return this;
}

10
spring-ai-core/src/main/java/org/springframework/ai/embedding/observation/EmbeddingModelObservationDocumentation.java
View File

@@ -112,16 +112,6 @@ public enum EmbeddingModelObservationDocumentation implements ObservationDocumen
}
},
/**
* The format the embeddings are returned in.
*/
REQUEST_EMBEDDING_ENCODING_FORMAT {
@Override
public String asString() {
return AiObservationAttributes.REQUEST_EMBEDDING_ENCODING_FORMAT.value();
}
},
// Usage
/**

98
spring-ai-core/src/main/java/org/springframework/ai/embedding/observation/EmbeddingModelRequestOptions.java
View File

@@ -1,98 +0,0 @@
/*
* 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.embedding.observation;
import org.springframework.ai.embedding.EmbeddingOptions;
import org.springframework.lang.Nullable;
import org.springframework.util.Assert;
/**
* Represents client-side options for embedding model requests.
*
* @author Thomas Vitale
* @since 1.0.0
*/
public class EmbeddingModelRequestOptions implements EmbeddingOptions {
private final String model;
@Nullable
private final Integer dimensions;
@Nullable
private final String encodingFormat;
EmbeddingModelRequestOptions(Builder builder) {
Assert.hasText(builder.model, "model cannot be null or empty");
this.model = builder.model;
this.dimensions = builder.dimensions;
this.encodingFormat = builder.encodingFormat;
}
public static Builder builder() {
return new Builder();
}
public static class Builder {
private String model;
@Nullable
private Integer dimensions;
@Nullable
private String encodingFormat;
private Builder() {
}
public Builder model(String model) {
this.model = model;
return this;
}
public Builder dimensions(@Nullable Integer dimensions) {
this.dimensions = dimensions;
return this;
}
public Builder encodingFormat(@Nullable String encodingFormat) {
this.encodingFormat = encodingFormat;
return this;
}
public EmbeddingModelRequestOptions build() {
return new EmbeddingModelRequestOptions(this);
}
}
public String getModel() {
return model;
}
@Nullable
public Integer getDimensions() {
return dimensions;
}
@Nullable
public String getEncodingFormat() {
return encodingFormat;
}
}

4
spring-ai-core/src/main/java/org/springframework/ai/observation/conventions/AiObservationAttributes.java
View File

@@ -79,10 +79,6 @@ public enum AiObservationAttributes {
* The number of dimensions the resulting output embeddings have.
*/
REQUEST_EMBEDDING_DIMENSIONS("gen_ai.request.embedding.dimensions"),
/**
* The format the embeddings are returned in.
*/
REQUEST_EMBEDDING_ENCODING_FORMAT("gen_ai.request.embedding.encoding_format"),
/**
* The format in which the generated image is returned.

43
spring-ai-core/src/test/java/org/springframework/ai/embedding/observation/DefaultEmbeddingModelObservationConventionTests.java
View File

@@ -19,7 +19,7 @@ import io.micrometer.common.KeyValue;
import io.micrometer.observation.Observation;
import org.junit.jupiter.api.Test;
import org.springframework.ai.chat.metadata.Usage;
import org.springframework.ai.embedding.EmbeddingOptions;
import org.springframework.ai.embedding.EmbeddingOptionsBuilder;
import org.springframework.ai.embedding.EmbeddingRequest;
import org.springframework.ai.embedding.EmbeddingResponse;
import org.springframework.ai.embedding.EmbeddingResponseMetadata;
@@ -50,32 +50,42 @@ class DefaultEmbeddingModelObservationConventionTests {
}
@Test
void shouldHaveContextualName() {
void contextualNameWhenModelIsDefined() {
EmbeddingModelObservationContext observationContext = EmbeddingModelObservationContext.builder()
.embeddingRequest(generateEmbeddingRequest())
.operationMetadata(generateOperationMetadata())
.requestOptions(EmbeddingModelRequestOptions.builder().model("mistral").build())
.requestOptions(EmbeddingOptionsBuilder.builder().withModel("mistral").build())
.build();
assertThat(this.observationConvention.getContextualName(observationContext)).isEqualTo("embedding mistral");
}
@Test
void contextualNameWhenModelIsNotDefined() {
EmbeddingModelObservationContext observationContext = EmbeddingModelObservationContext.builder()
.embeddingRequest(generateEmbeddingRequest())
.operationMetadata(generateOperationMetadata())
.requestOptions(EmbeddingOptionsBuilder.builder().build())
.build();
assertThat(this.observationConvention.getContextualName(observationContext)).isEqualTo("embedding");
}
@Test
void supportsOnlyEmbeddingModelObservationContext() {
EmbeddingModelObservationContext observationContext = EmbeddingModelObservationContext.builder()
.embeddingRequest(generateEmbeddingRequest())
.operationMetadata(generateOperationMetadata())
.requestOptions(EmbeddingModelRequestOptions.builder().model("supermodel").build())
.requestOptions(EmbeddingOptionsBuilder.builder().withModel("supermodel").build())
.build();
assertThat(this.observationConvention.supportsContext(observationContext)).isTrue();
assertThat(this.observationConvention.supportsContext(new Observation.Context())).isFalse();
}
@Test
void shouldHaveRequiredLowCardinalityKeyValues() {
void shouldHaveLowCardinalityKeyValuesWhenDefined() {
EmbeddingModelObservationContext observationContext = EmbeddingModelObservationContext.builder()
.embeddingRequest(generateEmbeddingRequest())
.operationMetadata(generateOperationMetadata())
.requestOptions(EmbeddingModelRequestOptions.builder().model("mistral").build())
.requestOptions(EmbeddingOptionsBuilder.builder().withModel("mistral").build())
.build();
assertThat(this.observationConvention.getLowCardinalityKeyValues(observationContext)).contains(
KeyValue.of(LowCardinalityKeyNames.AI_OPERATION_TYPE.asString(), "embedding"),
@@ -84,15 +94,11 @@ class DefaultEmbeddingModelObservationConventionTests {
}
@Test
void shouldHaveOptionalKeyValues() {
void shouldHaveLowCardinalityKeyValuesWhenDefinedAndResponse() {
EmbeddingModelObservationContext observationContext = EmbeddingModelObservationContext.builder()
.embeddingRequest(generateEmbeddingRequest())
.operationMetadata(generateOperationMetadata())
.requestOptions(EmbeddingModelRequestOptions.builder()
.model("supermodel")
.dimensions(1492)
.encodingFormat("vector")
.build())
.requestOptions(EmbeddingOptionsBuilder.builder().withModel("mistral").withDimensions(1492).build())
.build();
observationContext.setResponse(new EmbeddingResponse(List.of(),
new EmbeddingResponseMetadata("mistral-42", new TestUsage(), Map.of())));
@@ -100,29 +106,28 @@ class DefaultEmbeddingModelObservationConventionTests {
.contains(KeyValue.of(LowCardinalityKeyNames.RESPONSE_MODEL.asString(), "mistral-42"));
assertThat(this.observationConvention.getHighCardinalityKeyValues(observationContext)).contains(
KeyValue.of(HighCardinalityKeyNames.REQUEST_EMBEDDING_DIMENSIONS.asString(), "1492"),
KeyValue.of(HighCardinalityKeyNames.REQUEST_EMBEDDING_ENCODING_FORMAT.asString(), "vector"),
KeyValue.of(HighCardinalityKeyNames.USAGE_INPUT_TOKENS.asString(), "1000"),
KeyValue.of(HighCardinalityKeyNames.USAGE_TOTAL_TOKENS.asString(), "1000"));
}
@Test
void shouldHaveMissingKeyValues() {
void shouldHaveNoneKeyValuesWhenMissing() {
EmbeddingModelObservationContext observationContext = EmbeddingModelObservationContext.builder()
.embeddingRequest(generateEmbeddingRequest())
.operationMetadata(generateOperationMetadata())
.requestOptions(EmbeddingModelRequestOptions.builder().model("supermodel").build())
.requestOptions(EmbeddingOptionsBuilder.builder().build())
.build();
assertThat(this.observationConvention.getLowCardinalityKeyValues(observationContext))
.contains(KeyValue.of(LowCardinalityKeyNames.RESPONSE_MODEL.asString(), KeyValue.NONE_VALUE));
assertThat(this.observationConvention.getLowCardinalityKeyValues(observationContext)).contains(
KeyValue.of(LowCardinalityKeyNames.REQUEST_MODEL.asString(), KeyValue.NONE_VALUE),
KeyValue.of(LowCardinalityKeyNames.RESPONSE_MODEL.asString(), KeyValue.NONE_VALUE));
assertThat(this.observationConvention.getHighCardinalityKeyValues(observationContext)).contains(
KeyValue.of(HighCardinalityKeyNames.REQUEST_EMBEDDING_DIMENSIONS.asString(), KeyValue.NONE_VALUE),
KeyValue.of(HighCardinalityKeyNames.REQUEST_EMBEDDING_ENCODING_FORMAT.asString(), KeyValue.NONE_VALUE),
KeyValue.of(HighCardinalityKeyNames.USAGE_INPUT_TOKENS.asString(), KeyValue.NONE_VALUE),
KeyValue.of(HighCardinalityKeyNames.USAGE_TOTAL_TOKENS.asString(), KeyValue.NONE_VALUE));
}
private EmbeddingRequest generateEmbeddingRequest() {
return new EmbeddingRequest(List.of(), EmbeddingOptions.EMPTY);
return new EmbeddingRequest(List.of(), EmbeddingOptionsBuilder.builder().build());
}
private AiOperationMetadata generateOperationMetadata() {

6
spring-ai-core/src/test/java/org/springframework/ai/embedding/observation/EmbeddingModelMeterObservationHandlerTests.java
View File

@@ -22,7 +22,7 @@ import io.micrometer.observation.ObservationRegistry;
import org.junit.jupiter.api.BeforeEach;
import org.junit.jupiter.api.Test;
import org.springframework.ai.chat.metadata.Usage;
import org.springframework.ai.embedding.EmbeddingOptions;
import org.springframework.ai.embedding.EmbeddingOptionsBuilder;
import org.springframework.ai.embedding.EmbeddingRequest;
import org.springframework.ai.embedding.EmbeddingResponse;
import org.springframework.ai.embedding.EmbeddingResponseMetadata;
@@ -89,12 +89,12 @@ class EmbeddingModelMeterObservationHandlerTests {
return EmbeddingModelObservationContext.builder()
.embeddingRequest(generateEmbeddingRequest())
.operationMetadata(generateOperationMetadata())
.requestOptions(EmbeddingModelRequestOptions.builder().model("mistral").build())
.requestOptions(EmbeddingOptionsBuilder.builder().withModel("mistral").build())
.build();
}
private EmbeddingRequest generateEmbeddingRequest() {
return new EmbeddingRequest(List.of(), EmbeddingOptions.EMPTY);
return new EmbeddingRequest(List.of(), EmbeddingOptionsBuilder.builder().build());
}
private AiOperationMetadata generateOperationMetadata() {

6
spring-ai-core/src/test/java/org/springframework/ai/embedding/observation/EmbeddingModelObservationContextTests.java
View File

@@ -16,7 +16,7 @@
package org.springframework.ai.embedding.observation;
import org.junit.jupiter.api.Test;
import org.springframework.ai.embedding.EmbeddingOptions;
import org.springframework.ai.embedding.EmbeddingOptionsBuilder;
import org.springframework.ai.embedding.EmbeddingRequest;
import org.springframework.ai.observation.AiOperationMetadata;
import org.springframework.ai.observation.conventions.AiOperationType;
@@ -39,7 +39,7 @@ class EmbeddingModelObservationContextTests {
var observationContext = EmbeddingModelObservationContext.builder()
.embeddingRequest(generateEmbeddingRequest())
.operationMetadata(generateOperationMetadata())
.requestOptions(EmbeddingModelRequestOptions.builder().model("supermodel").build())
.requestOptions(EmbeddingOptionsBuilder.builder().withModel("supermodel").build())
.build();
assertThat(observationContext).isNotNull();
@@ -56,7 +56,7 @@ class EmbeddingModelObservationContextTests {
}
private EmbeddingRequest generateEmbeddingRequest() {
return new EmbeddingRequest(List.of(), EmbeddingOptions.EMPTY);
return new EmbeddingRequest(List.of(), EmbeddingOptionsBuilder.builder().build());
}
private AiOperationMetadata generateOperationMetadata() {

51
spring-ai-core/src/test/java/org/springframework/ai/embedding/observation/EmbeddingModelRequestOptionsTests.java
View File

@@ -1,51 +0,0 @@
/*
* 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.embedding.observation;
import org.junit.jupiter.api.Test;
import static org.assertj.core.api.Assertions.assertThat;
import static org.assertj.core.api.Assertions.assertThatThrownBy;
/**
* Unit tests for {@link EmbeddingModelRequestOptions}.
*
* @author Thomas Vitale
*/
class EmbeddingModelRequestOptionsTests {
@Test
void whenMandatoryRequestOptionsThenReturn() {
var requestOptions = EmbeddingModelRequestOptions.builder().model("rowena").build();
assertThat(requestOptions).isNotNull();
}
@Test
void whenModelIsNullThenThrow() {
assertThatThrownBy(() -> EmbeddingModelRequestOptions.builder().build())
.isInstanceOf(IllegalArgumentException.class)
.hasMessageContaining("model cannot be null or empty");
}
@Test
void whenModelIsEmptyThenThrow() {
assertThatThrownBy(() -> EmbeddingModelRequestOptions.builder().model("").build())
.isInstanceOf(IllegalArgumentException.class)
.hasMessageContaining("model cannot be null or empty");
}
}