Optimize MistralAiEmbeddingModel dimensions method

- Calculate and cache values for unknown models only if necessary
- Make known embedding dimensions a mutable map attribute
- Polish MistralAiEmbeddingModelTests

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

@@ -16,6 +16,7 @@
 
 package org.springframework.ai.mistralai;
 
+import java.util.HashMap;
 import java.util.List;
 import java.util.Map;
 
@@ -56,16 +57,14 @@ public class MistralAiEmbeddingModel extends AbstractEmbeddingModel {
 
 	private static final Logger logger = LoggerFactory.getLogger(MistralAiEmbeddingModel.class);
 
+	private static final EmbeddingModelObservationConvention DEFAULT_OBSERVATION_CONVENTION = new DefaultEmbeddingModelObservationConvention();
+
 	/**
 	 * Known embedding dimensions for Mistral AI models. Maps model names to their
 	 * respective embedding vector dimensions. This allows the dimensions() method to
 	 * return the correct value without making an API call.
 	 */
-	private static final Map<String, Integer> KNOWN_EMBEDDING_DIMENSIONS = Map.of(
-			MistralAiApi.EmbeddingModel.EMBED.getValue(), 1024, MistralAiApi.EmbeddingModel.CODESTRAL_EMBED.getValue(),
-			1536);
-
-	private static final EmbeddingModelObservationConvention DEFAULT_OBSERVATION_CONVENTION = new DefaultEmbeddingModelObservationConvention();
+	private final Map<String, Integer> knownEmbeddingDimensions = createKnownEmbeddingDimensions();
 
 	private final MistralAiEmbeddingOptions defaultOptions;
 
@@ -85,6 +84,14 @@ public class MistralAiEmbeddingModel extends AbstractEmbeddingModel {
 	 */
 	private EmbeddingModelObservationConvention observationConvention = DEFAULT_OBSERVATION_CONVENTION;
 
+	private static Map<String, Integer> createKnownEmbeddingDimensions() {
+		Map<String, Integer> knownEmbeddingDimensions = new HashMap<>();
+		knownEmbeddingDimensions.put(MistralAiApi.EmbeddingModel.EMBED.getValue(), 1024);
+		knownEmbeddingDimensions.put(MistralAiApi.EmbeddingModel.CODESTRAL_EMBED.getValue(), 1536);
+
+		return knownEmbeddingDimensions;
+	}
+
 	public MistralAiEmbeddingModel(MistralAiApi mistralAiApi, MetadataMode metadataMode,
 			MistralAiEmbeddingOptions options, RetryTemplate retryTemplate, ObservationRegistry observationRegistry) {
 		Assert.notNull(mistralAiApi, "mistralAiApi must not be null");
@@ -174,7 +181,8 @@ public float[] embed(Document document) {
 
 	@Override
 	public int dimensions() {
-		return KNOWN_EMBEDDING_DIMENSIONS.getOrDefault(this.defaultOptions.getModel(), super.dimensions());
+		return this.knownEmbeddingDimensions.computeIfAbsent(this.defaultOptions.getModel(),
+				model -> super.dimensions());
 	}
 
 	/**

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

@@ -16,6 +16,7 @@
 
 package org.springframework.ai.mistralai;
 
+import java.util.Arrays;
 import java.util.List;
 
 import org.junit.jupiter.api.Test;
@@ -94,7 +95,7 @@ void testDimensionsFallbackForUnknownModel() {
 
 	@Test
 	void testAllEmbeddingModelsHaveDimensionMapping() {
-		// This test ensures that KNOWN_EMBEDDING_DIMENSIONS map stays in sync with the
+		// This test ensures that knownEmbeddingDimensions map stays in sync with the
 		// EmbeddingModel enum
 		// If a new model is added to the enum but not to the dimensions map, this test
 		// will help catch it
@@ -138,16 +139,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));