Replace the Embedding format from List<Double> to float[]

 - Adjust all affected classes including the Document.
 - Update docs.

Related to #405

Author: tzolov

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

@@ -29,11 +29,21 @@
 import org.springframework.ai.embedding.EmbeddingRequest;
 import org.springframework.ai.embedding.EmbeddingResponse;
 import org.springframework.ai.embedding.EmbeddingResponseMetadata;
+import org.springframework.ai.model.EmbeddingUtils;
 import org.springframework.util.Assert;
+import org.springframework.util.CollectionUtils;
 
 import java.util.ArrayList;
 import java.util.List;
 
+/**
+ * Azure Open AI Embedding Model implementation.
+ *
+ * @author Mark Pollack
+ * @author Christian Tzolov
+ * @author Thomas Vitale
+ * @since 1.0.0
+ */
 public class AzureOpenAiEmbeddingModel extends AbstractEmbeddingModel {
 
 	private static final Logger logger = LoggerFactory.getLogger(AzureOpenAiEmbeddingModel.class);
@@ -64,13 +74,17 @@ public AzureOpenAiEmbeddingModel(OpenAIClient azureOpenAiClient, MetadataMode me
 	}
 
 	@Override
-	public List<Double> embed(Document document) {
+	public float[] embed(Document document) {
 		logger.debug("Retrieving embeddings");
 
 		EmbeddingResponse response = this
 			.call(new EmbeddingRequest(List.of(document.getFormattedContent(this.metadataMode)), null));
 		logger.debug("Embeddings retrieved");
-		return response.getResults().stream().map(embedding -> embedding.getOutput()).flatMap(List::stream).toList();
+
+		if (CollectionUtils.isEmpty(response.getResults())) {
+			return new float[0];
+		}
+		return response.getResults().get(0).getOutput();
 	}
 
 	@Override
@@ -108,8 +122,7 @@ private List<Embedding> generateEmbeddingList(List<EmbeddingItem> nativeData) {
 		for (EmbeddingItem nativeDatum : nativeData) {
 			List<Float> nativeDatumEmbedding = nativeDatum.getEmbedding();
 			int nativeIndex = nativeDatum.getPromptIndex();
-			Embedding embedding = new Embedding(nativeDatumEmbedding.stream().map(f -> f.doubleValue()).toList(),
-					nativeIndex);
+			Embedding embedding = new Embedding(EmbeddingUtils.toPrimitive(nativeDatumEmbedding), nativeIndex);
 			data.add(embedding);
 		}
 		return data;

models/spring-ai-bedrock/src/main/java/org/springframework/ai/bedrock/cohere/BedrockCohereEmbeddingModel.java

@@ -66,33 +66,8 @@ public BedrockCohereEmbeddingModel(CohereEmbeddingBedrockApi cohereEmbeddingBedr
 		this.defaultOptions = options;
 	}
 
-	// /**
-	// * Cohere Embedding API input types.
-	// * @param inputType the input type to use.
-	// * @return this client.
-	// */
-	// public BedrockCohereEmbeddingModel withInputType(CohereEmbeddingRequest.InputType
-	// inputType) {
-	// this.inputType = inputType;
-	// return this;
-	// }
-
-	// /**
-	// * Specifies how the API handles inputs longer than the maximum token length. If you
-	// specify LEFT or RIGHT, the
-	// * model discards the input until the remaining input is exactly the maximum input
-	// token length for the model.
-	// * @param truncate the truncate option to use.
-	// * @return this client.
-	// */
-	// public BedrockCohereEmbeddingModel withTruncate(CohereEmbeddingRequest.Truncate
-	// truncate) {
-	// this.truncate = truncate;
-	// return this;
-	// }
-
 	@Override
-	public List<Double> embed(Document document) {
+	public float[] embed(Document document) {
 		return embed(document.getContent());
 	}
 

models/spring-ai-bedrock/src/main/java/org/springframework/ai/bedrock/cohere/api/CohereEmbeddingBedrockApi.java

@@ -183,7 +183,7 @@ public enum Truncate {
 	@JsonInclude(Include.NON_NULL)
 	public record CohereEmbeddingResponse(
 			@JsonProperty("id") String id,
-			@JsonProperty("embeddings") List<List<Double>> embeddings,
+			@JsonProperty("embeddings") List<float[]> embeddings,
 			@JsonProperty("texts") List<String> texts,
 			@JsonProperty("response_type") String responseType,
 			// For future use: Currently bedrock doesn't return invocationMetrics for the cohere embedding model.

models/spring-ai-bedrock/src/main/java/org/springframework/ai/bedrock/titan/BedrockTitanEmbeddingModel.java

@@ -75,7 +75,7 @@ public BedrockTitanEmbeddingModel withInputType(InputType inputType) {
 	}
 
 	@Override
-	public List<Double> embed(Document document) {
+	public float[] embed(Document document) {
 		return embed(document.getContent());
 	}
 
@@ -87,16 +87,13 @@ public EmbeddingResponse call(EmbeddingRequest request) {
 					"Titan Embedding does not support batch embedding. Will make multiple API calls to embed(Document)");
 		}
 
-		List<List<Double>> embeddingList = new ArrayList<>();
+		List<Embedding> embeddings = new ArrayList<>();
+		var indexCounter = new AtomicInteger(0);
 		for (String inputContent : request.getInstructions()) {
 			var apiRequest = createTitanEmbeddingRequest(inputContent, request.getOptions());
 			TitanEmbeddingResponse response = this.embeddingApi.embedding(apiRequest);
-			embeddingList.add(response.embedding());
+			embeddings.add(new Embedding(response.embedding(), indexCounter.getAndIncrement()));
 		}
-		var indexCounter = new AtomicInteger(0);
-		List<Embedding> embeddings = embeddingList.stream()
-			.map(e -> new Embedding(e, indexCounter.getAndIncrement()))
-			.toList();
 		return new EmbeddingResponse(embeddings);
 	}
 

models/spring-ai-bedrock/src/main/java/org/springframework/ai/bedrock/titan/api/TitanEmbeddingBedrockApi.java

@@ -137,7 +137,7 @@ public TitanEmbeddingRequest build() {
 	 */
 	@JsonInclude(Include.NON_NULL)
 	public record TitanEmbeddingResponse(
-			@JsonProperty("embedding") List<Double> embedding,
+			@JsonProperty("embedding") float[] embedding,
 			@JsonProperty("inputTextTokenCount") Integer inputTextTokenCount,
 			@JsonProperty("message") Object message) {
 	}

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

@@ -105,7 +105,7 @@ public MiniMaxEmbeddingModel(MiniMaxApi miniMaxApi, MetadataMode metadataMode, M
 	}
 
 	@Override
-	public List<Double> embed(Document document) {
+	public float[] embed(Document document) {
 		Assert.notNull(document, "Document must not be null");
 		return this.embed(document.getFormattedContent(this.metadataMode));
 	}
@@ -137,7 +137,7 @@ public EmbeddingResponse call(EmbeddingRequest request) {
 
 			List<Embedding> embeddings = new ArrayList<>();
 			for (int i = 0; i < apiEmbeddingResponse.vectors().size(); i++) {
-				List<Double> vector = apiEmbeddingResponse.vectors().get(i);
+				float[] vector = apiEmbeddingResponse.vectors().get(i);
 				embeddings.add(new Embedding(vector, i));
 			}
 			return new EmbeddingResponse(embeddings, metadata);

models/spring-ai-minimax/src/main/java/org/springframework/ai/minimax/api/MiniMaxApi.java

@@ -865,7 +865,7 @@ public EmbeddingRequest(List<String> texts, EmbeddingType type) {
 	 */
 	@JsonInclude(Include.NON_NULL)
 	public record EmbeddingList(
-			@JsonProperty("vectors") List<List<Double>> vectors,
+			@JsonProperty("vectors") List<float[]> vectors,
 			@JsonProperty("model") String model,
 			@JsonProperty("total_tokens") Integer totalTokens) {
 	}

models/spring-ai-minimax/src/test/java/org/springframework/ai/minimax/api/MiniMaxRetryTests.java

@@ -157,7 +157,7 @@ public void miniMaxChatStreamNonTransientError() {
 	@Test
 	public void miniMaxEmbeddingTransientError() {
 
-		EmbeddingList expectedEmbeddings = new EmbeddingList(List.of(List.of(9.9, 8.8)), "model", 10);
+		EmbeddingList expectedEmbeddings = new EmbeddingList(List.of(new float[] { 9.9f, 8.8f }), "model", 10);
 
 		when(miniMaxApi.embeddings(isA(EmbeddingRequest.class)))
 			.thenThrow(new TransientAiException("Transient Error 1"))
@@ -168,7 +168,7 @@ public void miniMaxEmbeddingTransientError() {
 			.call(new org.springframework.ai.embedding.EmbeddingRequest(List.of("text1", "text2"), null));
 
 		assertThat(result).isNotNull();
-		assertThat(result.getResult().getOutput()).isEqualTo(List.of(9.9, 8.8));
+		assertThat(result.getResult().getOutput()).isEqualTo(new float[] { 9.9f, 8.8f });
 		assertThat(retryListener.onSuccessRetryCount).isEqualTo(2);
 		assertThat(retryListener.onErrorRetryCount).isEqualTo(2);
 	}

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

@@ -116,7 +116,7 @@ public EmbeddingResponse call(EmbeddingRequest request) {
 	}
 
 	@Override
-	public List<Double> embed(Document document) {
+	public float[] embed(Document document) {
 		Assert.notNull(document, "Document must not be null");
 		return this.embed(document.getFormattedContent(this.metadataMode));
 	}

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

@@ -196,7 +196,7 @@ public record Usage(
 	public record Embedding(
 	// @formatter:off
 		 @JsonProperty("index") Integer index,
-		 @JsonProperty("embedding") List<Double> embedding,
+		 @JsonProperty("embedding") float[] embedding,
 		 @JsonProperty("object") String object) {
 		 // @formatter:on
 
@@ -207,7 +207,7 @@ public record Embedding(
 		 * @param embedding The embedding vector, which is a list of floats. The length of
 		 * vector depends on the model.
 		 */
-		public Embedding(Integer index, List<Double> embedding) {
+		public Embedding(Integer index, float[] embedding) {
 			this(index, embedding, "embedding");
 		}
 	}