Optimize MistralAiEmbeddingModel dimensions method

- Calculate default value only if necessary
- Fix warnings in MistralAiEmbeddingModelTests unit tests

Signed-off-by: Nicolas Krier <7557886+nicolaskrier@users.noreply.github.com>

Author: nicolaskrier

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

@@ -197,7 +197,13 @@ public float[] embed(Document document) {
 
 	@Override
 	public int dimensions() {
-		return KNOWN_EMBEDDING_DIMENSIONS.getOrDefault(this.defaultOptions.getModel(), super.dimensions());
+		var dimensions = KNOWN_EMBEDDING_DIMENSIONS.get(this.defaultOptions.getModel());
+
+		if (dimensions == null) {
+			dimensions = super.dimensions();
+		}
+
+		return dimensions;
 	}
 
 	/**

models/spring-ai-mistral-ai/src/test/java/org/springframework/ai/mistralai/MistralAiEmbeddingModelTests.java

@@ -16,8 +16,10 @@
 
 package org.springframework.ai.mistralai;
 
+import java.util.Arrays;
 import java.util.List;
 
+import io.micrometer.observation.ObservationRegistry;
 import org.junit.jupiter.api.Test;
 import org.mockito.Mockito;
 
@@ -46,7 +48,7 @@ void testDimensionsForMistralEmbedModel() {
 			.build();
 
 		MistralAiEmbeddingModel model = new MistralAiEmbeddingModel(mockApi, MetadataMode.EMBED, options,
-				RetryUtils.DEFAULT_RETRY_TEMPLATE);
+				RetryUtils.DEFAULT_RETRY_TEMPLATE, ObservationRegistry.NOOP);
 
 		assertThat(model.dimensions()).isEqualTo(1024);
 	}
@@ -60,7 +62,7 @@ void testDimensionsForCodestralEmbedModel() {
 			.build();
 
 		MistralAiEmbeddingModel model = new MistralAiEmbeddingModel(mockApi, MetadataMode.EMBED, options,
-				RetryUtils.DEFAULT_RETRY_TEMPLATE);
+				RetryUtils.DEFAULT_RETRY_TEMPLATE, ObservationRegistry.NOOP);
 
 		assertThat(model.dimensions()).isEqualTo(1536);
 	}
@@ -73,7 +75,7 @@ void testDimensionsFallbackForUnknownModel() {
 		MistralAiEmbeddingOptions options = MistralAiEmbeddingOptions.builder().withModel("unknown-model").build();
 
 		MistralAiEmbeddingModel model = new MistralAiEmbeddingModel(mockApi, MetadataMode.EMBED, options,
-				RetryUtils.DEFAULT_RETRY_TEMPLATE);
+				RetryUtils.DEFAULT_RETRY_TEMPLATE, ObservationRegistry.NOOP);
 
 		// Should fall back to super.dimensions() which detects dimensions from the API
 		// response
@@ -94,7 +96,7 @@ void testAllEmbeddingModelsHaveDimensionMapping() {
 				.build();
 
 			MistralAiEmbeddingModel model = new MistralAiEmbeddingModel(mockApi, MetadataMode.EMBED, options,
-					RetryUtils.DEFAULT_RETRY_TEMPLATE);
+					RetryUtils.DEFAULT_RETRY_TEMPLATE, ObservationRegistry.NOOP);
 
 			// Each model should have a valid dimension (not the fallback -1)
 			assertThat(model.dimensions()).as("Model %s should have a dimension mapping", embeddingModel.getValue())
@@ -122,16 +124,13 @@ private MistralAiApi createMockApiWithEmbeddingResponse(int dimensions) {
 
 		// Create a mock embedding response with the specified dimensions
 		float[] embedding = new float[dimensions];
-		for (int i = 0; i < dimensions; i++) {
-			embedding[i] = 0.1f;
-		}
+		Arrays.fill(embedding, 0.1f);
 
 		MistralAiApi.Embedding embeddingData = new MistralAiApi.Embedding(0, embedding, "embedding");
 
 		MistralAiApi.Usage usage = new MistralAiApi.Usage(10, 0, 10);
 
-		MistralAiApi.EmbeddingList embeddingList = new MistralAiApi.EmbeddingList("object", List.of(embeddingData),
-				"model", usage);
+		var embeddingList = new MistralAiApi.EmbeddingList<>("object", List.of(embeddingData), "model", usage);
 
 		when(mockApi.embeddings(any())).thenReturn(ResponseEntity.ok(embeddingList));