Add Mistral AI Codestral Embed model

Signed-off-by: Nicolas Krier <[email protected]>

Author: nicolaskrier

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

@@ -17,6 +17,7 @@
 package org.springframework.ai.mistralai;
 
 import java.util.List;
+import java.util.Map;
 
 import io.micrometer.observation.ObservationRegistry;
 import org.slf4j.Logger;
@@ -41,6 +42,9 @@
 import org.springframework.retry.support.RetryTemplate;
 import org.springframework.util.Assert;
 
+import static org.springframework.ai.mistralai.api.MistralAiApi.EmbeddingModel.CODESTRAL_EMBED;
+import static org.springframework.ai.mistralai.api.MistralAiApi.EmbeddingModel.EMBED;
+
 /**
  * Provides the Mistral AI Embedding Model.
  *
@@ -53,6 +57,9 @@ public class MistralAiEmbeddingModel extends AbstractEmbeddingModel {
 
 	private static final Logger logger = LoggerFactory.getLogger(MistralAiEmbeddingModel.class);
 
+	private static final Map<String, Integer> KNOWN_EMBEDDING_DIMENSIONS = Map.of(EMBED.getValue(), 1024,
+			CODESTRAL_EMBED.getValue(), 1536);
+
 	private static final EmbeddingModelObservationConvention DEFAULT_OBSERVATION_CONVENTION = new DefaultEmbeddingModelObservationConvention();
 
 	private final MistralAiEmbeddingOptions defaultOptions;
@@ -78,8 +85,7 @@ public MistralAiEmbeddingModel(MistralAiApi mistralAiApi) {
 	}
 
 	public MistralAiEmbeddingModel(MistralAiApi mistralAiApi, MetadataMode metadataMode) {
-		this(mistralAiApi, metadataMode,
-				MistralAiEmbeddingOptions.builder().withModel(MistralAiApi.EmbeddingModel.EMBED.getValue()).build(),
+		this(mistralAiApi, metadataMode, MistralAiEmbeddingOptions.builder().withModel(EMBED.getValue()).build(),
 				RetryUtils.DEFAULT_RETRY_TEMPLATE);
 	}
 
@@ -179,6 +185,11 @@ public float[] embed(Document document) {
 		return this.embed(document.getFormattedContent(this.metadataMode));
 	}
 
+	@Override
+	public int dimensions() {
+		return KNOWN_EMBEDDING_DIMENSIONS.getOrDefault(this.defaultOptions.getModel(), super.dimensions());
+	}
+
 	/**
 	 * Use the provided convention for reporting observation data
 	 * @param observationConvention The provided convention

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

@@ -321,7 +321,8 @@ public String getName() {
 	public enum EmbeddingModel {
 
 		// @formatter:off
-		EMBED("mistral-embed");
+		EMBED("mistral-embed"),
+		CODESTRAL_EMBED("codestral-embed");
 		 // @formatter:on
 
 		private final String value;

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

@@ -16,50 +16,53 @@
 
 package org.springframework.ai.mistralai;
 
-import java.util.List;
-
 import org.junit.jupiter.api.Test;
 import org.junit.jupiter.api.condition.EnabledIfEnvironmentVariable;
-
+import org.junit.jupiter.params.ParameterizedTest;
+import org.junit.jupiter.params.provider.CsvSource;
 import org.springframework.ai.embedding.EmbeddingRequest;
+import org.springframework.ai.mistralai.api.MistralAiApi;
 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;
 
 @SpringBootTest(classes = MistralAiTestConfiguration.class)
 @EnabledIfEnvironmentVariable(named = "MISTRAL_AI_API_KEY", matches = ".+")
 class MistralAiEmbeddingIT {
 
-	@Autowired
-	private MistralAiEmbeddingModel mistralAiEmbeddingModel;
-
 	@Test
-	void defaultEmbedding() {
-		assertThat(this.mistralAiEmbeddingModel).isNotNull();
-		var embeddingResponse = this.mistralAiEmbeddingModel.embedForResponse(List.of("Hello World"));
+	void defaultEmbedding(@Autowired MistralAiEmbeddingModel mistralAiEmbeddingModel) {
+		assertThat(mistralAiEmbeddingModel).isNotNull();
+		var embeddingResponse = mistralAiEmbeddingModel.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(embeddingResponse.getMetadata().getModel()).isEqualTo("mistral-embed");
 		assertThat(embeddingResponse.getMetadata().getUsage().getTotalTokens()).isEqualTo(4);
 		assertThat(embeddingResponse.getMetadata().getUsage().getPromptTokens()).isEqualTo(4);
-		assertThat(this.mistralAiEmbeddingModel.dimensions()).isEqualTo(1024);
+		assertThat(mistralAiEmbeddingModel.dimensions()).isEqualTo(1024);
 	}
 
-	@Test
-	void embeddingTest() {
-		assertThat(this.mistralAiEmbeddingModel).isNotNull();
-		var embeddingResponse = this.mistralAiEmbeddingModel.call(new EmbeddingRequest(
-				List.of("Hello World", "World is big"),
-				MistralAiEmbeddingOptions.builder().withModel("mistral-embed").withEncodingFormat("float").build()));
+	@CsvSource({ "mistral-embed, 1024", "codestral-embed, 1536" })
+	@ParameterizedTest
+	void embeddingTest(String model, int dimensions, @Autowired MistralAiApi mistralAiApi) {
+		var mistralAiEmbeddingOptions = MistralAiEmbeddingOptions.builder().withModel(model).build();
+		var mistralAiEmbeddingModel = new MistralAiEmbeddingModel(mistralAiApi, mistralAiEmbeddingOptions);
+		var embeddingOptions = MistralAiEmbeddingOptions.builder().withModel(model).withEncodingFormat("float").build();
+		var embeddingRequest = new EmbeddingRequest(List.of("Hello World", "World is big"), embeddingOptions);
+		var embeddingResponse = mistralAiEmbeddingModel.call(embeddingRequest);
 		assertThat(embeddingResponse.getResults()).hasSize(2);
 		assertThat(embeddingResponse.getResults().get(0)).isNotNull();
-		assertThat(embeddingResponse.getResults().get(0).getOutput()).hasSize(1024);
-		assertThat(embeddingResponse.getMetadata().getModel()).isEqualTo("mistral-embed");
+		assertThat(embeddingResponse.getResults().get(0).getOutput()).hasSize(dimensions);
+		assertThat(embeddingResponse.getResults().get(1)).isNotNull();
+		assertThat(embeddingResponse.getResults().get(1).getOutput()).hasSize(dimensions);
+		assertThat(embeddingResponse.getMetadata().getModel()).isEqualTo(model);
 		assertThat(embeddingResponse.getMetadata().getUsage().getTotalTokens()).isEqualTo(9);
 		assertThat(embeddingResponse.getMetadata().getUsage().getPromptTokens()).isEqualTo(9);
-		assertThat(this.mistralAiEmbeddingModel.dimensions()).isEqualTo(1024);
+		assertThat(mistralAiEmbeddingModel.dimensions()).isEqualTo(dimensions);
 	}
 
 }