diff --git a/models/spring-ai-azure-openai/src/main/java/org/springframework/ai/azure/openai/AzureOpenAiEmbeddingModel.java b/models/spring-ai-azure-openai/src/main/java/org/springframework/ai/azure/openai/AzureOpenAiEmbeddingModel.java index 625c8e02b..546c4ed29 100644 --- a/models/spring-ai-azure-openai/src/main/java/org/springframework/ai/azure/openai/AzureOpenAiEmbeddingModel.java +++ b/models/spring-ai-azure-openai/src/main/java/org/springframework/ai/azure/openai/AzureOpenAiEmbeddingModel.java @@ -29,11 +29,21 @@ import org.springframework.ai.embedding.Embedding; 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 class AzureOpenAiEmbeddingModel extends AbstractEmbeddingModel { } @Override - public List 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 @@ public class AzureOpenAiEmbeddingModel extends AbstractEmbeddingModel { for (EmbeddingItem nativeDatum : nativeData) { List 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; diff --git a/models/spring-ai-bedrock/src/main/java/org/springframework/ai/bedrock/cohere/BedrockCohereEmbeddingModel.java b/models/spring-ai-bedrock/src/main/java/org/springframework/ai/bedrock/cohere/BedrockCohereEmbeddingModel.java index 25e3f35b4..c070f2506 100644 --- a/models/spring-ai-bedrock/src/main/java/org/springframework/ai/bedrock/cohere/BedrockCohereEmbeddingModel.java +++ b/models/spring-ai-bedrock/src/main/java/org/springframework/ai/bedrock/cohere/BedrockCohereEmbeddingModel.java @@ -66,33 +66,8 @@ public class BedrockCohereEmbeddingModel extends AbstractEmbeddingModel { 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 embed(Document document) { + public float[] embed(Document document) { return embed(document.getContent()); } diff --git a/models/spring-ai-bedrock/src/main/java/org/springframework/ai/bedrock/cohere/api/CohereEmbeddingBedrockApi.java b/models/spring-ai-bedrock/src/main/java/org/springframework/ai/bedrock/cohere/api/CohereEmbeddingBedrockApi.java index 1ae424227..30938fe9a 100644 --- a/models/spring-ai-bedrock/src/main/java/org/springframework/ai/bedrock/cohere/api/CohereEmbeddingBedrockApi.java +++ b/models/spring-ai-bedrock/src/main/java/org/springframework/ai/bedrock/cohere/api/CohereEmbeddingBedrockApi.java @@ -183,7 +183,7 @@ public class CohereEmbeddingBedrockApi extends @JsonInclude(Include.NON_NULL) public record CohereEmbeddingResponse( @JsonProperty("id") String id, - @JsonProperty("embeddings") List> embeddings, + @JsonProperty("embeddings") List embeddings, @JsonProperty("texts") List texts, @JsonProperty("response_type") String responseType, // For future use: Currently bedrock doesn't return invocationMetrics for the cohere embedding model. diff --git a/models/spring-ai-bedrock/src/main/java/org/springframework/ai/bedrock/titan/BedrockTitanEmbeddingModel.java b/models/spring-ai-bedrock/src/main/java/org/springframework/ai/bedrock/titan/BedrockTitanEmbeddingModel.java index e3089eec5..84a646b84 100644 --- a/models/spring-ai-bedrock/src/main/java/org/springframework/ai/bedrock/titan/BedrockTitanEmbeddingModel.java +++ b/models/spring-ai-bedrock/src/main/java/org/springframework/ai/bedrock/titan/BedrockTitanEmbeddingModel.java @@ -75,7 +75,7 @@ public class BedrockTitanEmbeddingModel extends AbstractEmbeddingModel { } @Override - public List embed(Document document) { + public float[] embed(Document document) { return embed(document.getContent()); } @@ -87,16 +87,13 @@ public class BedrockTitanEmbeddingModel extends AbstractEmbeddingModel { "Titan Embedding does not support batch embedding. Will make multiple API calls to embed(Document)"); } - List> embeddingList = new ArrayList<>(); + List 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 embeddings = embeddingList.stream() - .map(e -> new Embedding(e, indexCounter.getAndIncrement())) - .toList(); return new EmbeddingResponse(embeddings); } diff --git a/models/spring-ai-bedrock/src/main/java/org/springframework/ai/bedrock/titan/api/TitanEmbeddingBedrockApi.java b/models/spring-ai-bedrock/src/main/java/org/springframework/ai/bedrock/titan/api/TitanEmbeddingBedrockApi.java index 016ec4306..9b725e325 100644 --- a/models/spring-ai-bedrock/src/main/java/org/springframework/ai/bedrock/titan/api/TitanEmbeddingBedrockApi.java +++ b/models/spring-ai-bedrock/src/main/java/org/springframework/ai/bedrock/titan/api/TitanEmbeddingBedrockApi.java @@ -137,7 +137,7 @@ public class TitanEmbeddingBedrockApi extends */ @JsonInclude(Include.NON_NULL) public record TitanEmbeddingResponse( - @JsonProperty("embedding") List embedding, + @JsonProperty("embedding") float[] embedding, @JsonProperty("inputTextTokenCount") Integer inputTextTokenCount, @JsonProperty("message") Object message) { } diff --git a/models/spring-ai-minimax/src/main/java/org/springframework/ai/minimax/MiniMaxEmbeddingModel.java b/models/spring-ai-minimax/src/main/java/org/springframework/ai/minimax/MiniMaxEmbeddingModel.java index 67477a308..cef854e5f 100644 --- a/models/spring-ai-minimax/src/main/java/org/springframework/ai/minimax/MiniMaxEmbeddingModel.java +++ b/models/spring-ai-minimax/src/main/java/org/springframework/ai/minimax/MiniMaxEmbeddingModel.java @@ -105,7 +105,7 @@ public class MiniMaxEmbeddingModel extends AbstractEmbeddingModel { } @Override - public List 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 class MiniMaxEmbeddingModel extends AbstractEmbeddingModel { List embeddings = new ArrayList<>(); for (int i = 0; i < apiEmbeddingResponse.vectors().size(); i++) { - List vector = apiEmbeddingResponse.vectors().get(i); + float[] vector = apiEmbeddingResponse.vectors().get(i); embeddings.add(new Embedding(vector, i)); } return new EmbeddingResponse(embeddings, metadata); diff --git a/models/spring-ai-minimax/src/main/java/org/springframework/ai/minimax/api/MiniMaxApi.java b/models/spring-ai-minimax/src/main/java/org/springframework/ai/minimax/api/MiniMaxApi.java index 725bc9620..9218cd128 100644 --- a/models/spring-ai-minimax/src/main/java/org/springframework/ai/minimax/api/MiniMaxApi.java +++ b/models/spring-ai-minimax/src/main/java/org/springframework/ai/minimax/api/MiniMaxApi.java @@ -865,7 +865,7 @@ public class MiniMaxApi { */ @JsonInclude(Include.NON_NULL) public record EmbeddingList( - @JsonProperty("vectors") List> vectors, + @JsonProperty("vectors") List vectors, @JsonProperty("model") String model, @JsonProperty("total_tokens") Integer totalTokens) { } diff --git a/models/spring-ai-minimax/src/test/java/org/springframework/ai/minimax/api/MiniMaxRetryTests.java b/models/spring-ai-minimax/src/test/java/org/springframework/ai/minimax/api/MiniMaxRetryTests.java index 157f7d063..493874de5 100644 --- a/models/spring-ai-minimax/src/test/java/org/springframework/ai/minimax/api/MiniMaxRetryTests.java +++ b/models/spring-ai-minimax/src/test/java/org/springframework/ai/minimax/api/MiniMaxRetryTests.java @@ -157,7 +157,7 @@ public class MiniMaxRetryTests { @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 class MiniMaxRetryTests { .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); } diff --git a/models/spring-ai-mistral-ai/src/main/java/org/springframework/ai/mistralai/MistralAiEmbeddingModel.java b/models/spring-ai-mistral-ai/src/main/java/org/springframework/ai/mistralai/MistralAiEmbeddingModel.java index 5a79f9907..baf343222 100644 --- a/models/spring-ai-mistral-ai/src/main/java/org/springframework/ai/mistralai/MistralAiEmbeddingModel.java +++ b/models/spring-ai-mistral-ai/src/main/java/org/springframework/ai/mistralai/MistralAiEmbeddingModel.java @@ -116,7 +116,7 @@ public class MistralAiEmbeddingModel extends AbstractEmbeddingModel { } @Override - public List embed(Document document) { + public float[] embed(Document document) { Assert.notNull(document, "Document must not be null"); return this.embed(document.getFormattedContent(this.metadataMode)); } diff --git a/models/spring-ai-mistral-ai/src/main/java/org/springframework/ai/mistralai/api/MistralAiApi.java b/models/spring-ai-mistral-ai/src/main/java/org/springframework/ai/mistralai/api/MistralAiApi.java index c193e856b..7796b58e0 100644 --- a/models/spring-ai-mistral-ai/src/main/java/org/springframework/ai/mistralai/api/MistralAiApi.java +++ b/models/spring-ai-mistral-ai/src/main/java/org/springframework/ai/mistralai/api/MistralAiApi.java @@ -196,7 +196,7 @@ public class MistralAiApi { public record Embedding( // @formatter:off @JsonProperty("index") Integer index, - @JsonProperty("embedding") List embedding, + @JsonProperty("embedding") float[] embedding, @JsonProperty("object") String object) { // @formatter:on @@ -207,7 +207,7 @@ public class MistralAiApi { * @param embedding The embedding vector, which is a list of floats. The length of * vector depends on the model. */ - public Embedding(Integer index, List embedding) { + public Embedding(Integer index, float[] embedding) { this(index, embedding, "embedding"); } } diff --git a/models/spring-ai-mistral-ai/src/test/java/org/springframework/ai/mistralai/MistralAiRetryTests.java b/models/spring-ai-mistral-ai/src/test/java/org/springframework/ai/mistralai/MistralAiRetryTests.java index c75871222..d0523397e 100644 --- a/models/spring-ai-mistral-ai/src/test/java/org/springframework/ai/mistralai/MistralAiRetryTests.java +++ b/models/spring-ai-mistral-ai/src/test/java/org/springframework/ai/mistralai/MistralAiRetryTests.java @@ -15,6 +15,11 @@ */ package org.springframework.ai.mistralai; +import static org.assertj.core.api.Assertions.assertThat; +import static org.junit.jupiter.api.Assertions.assertThrows; +import static org.mockito.ArgumentMatchers.isA; +import static org.mockito.Mockito.when; + import java.util.List; import java.util.Optional; @@ -23,8 +28,6 @@ import org.junit.jupiter.api.Test; import org.junit.jupiter.api.extension.ExtendWith; import org.mockito.Mock; import org.mockito.junit.jupiter.MockitoExtension; -import reactor.core.publisher.Flux; - import org.springframework.ai.chat.prompt.Prompt; import org.springframework.ai.document.MetadataMode; import org.springframework.ai.mistralai.api.MistralAiApi; @@ -45,10 +48,7 @@ import org.springframework.retry.RetryContext; import org.springframework.retry.RetryListener; import org.springframework.retry.support.RetryTemplate; -import static org.assertj.core.api.Assertions.assertThat; -import static org.junit.jupiter.api.Assertions.assertThrows; -import static org.mockito.ArgumentMatchers.isA; -import static org.mockito.Mockito.when; +import reactor.core.publisher.Flux; /** * @author Christian Tzolov @@ -166,7 +166,7 @@ public class MistralAiRetryTests { public void mistralAiEmbeddingTransientError() { EmbeddingList expectedEmbeddings = new EmbeddingList<>("list", - List.of(new Embedding(0, List.of(9.9, 8.8))), "model", new MistralAiApi.Usage(10, 10, 10)); + List.of(new Embedding(0, new float[] { 9.9f, 8.8f })), "model", new MistralAiApi.Usage(10, 10, 10)); when(mistralAiApi.embeddings(isA(EmbeddingRequest.class))) .thenThrow(new TransientAiException("Transient Error 1")) @@ -177,7 +177,7 @@ public class MistralAiRetryTests { .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); } diff --git a/models/spring-ai-ollama/src/main/java/org/springframework/ai/ollama/OllamaEmbeddingModel.java b/models/spring-ai-ollama/src/main/java/org/springframework/ai/ollama/OllamaEmbeddingModel.java index 439e8a425..4b1777960 100644 --- a/models/spring-ai-ollama/src/main/java/org/springframework/ai/ollama/OllamaEmbeddingModel.java +++ b/models/spring-ai-ollama/src/main/java/org/springframework/ai/ollama/OllamaEmbeddingModel.java @@ -76,7 +76,7 @@ public class OllamaEmbeddingModel extends AbstractEmbeddingModel { } @Override - public List embed(Document document) { + public float[] embed(Document document) { return embed(document.getContent()); } diff --git a/models/spring-ai-ollama/src/main/java/org/springframework/ai/ollama/api/OllamaApi.java b/models/spring-ai-ollama/src/main/java/org/springframework/ai/ollama/api/OllamaApi.java index 1ad37f0a8..eeb003d53 100644 --- a/models/spring-ai-ollama/src/main/java/org/springframework/ai/ollama/api/OllamaApi.java +++ b/models/spring-ai-ollama/src/main/java/org/springframework/ai/ollama/api/OllamaApi.java @@ -751,7 +751,7 @@ public class OllamaApi { @Deprecated(since = "1.0.0-M2", forRemoval = true) @JsonInclude(Include.NON_NULL) public record EmbeddingResponse( - @JsonProperty("embedding") List embedding) { + @JsonProperty("embedding") List embedding) { } @@ -764,7 +764,7 @@ public class OllamaApi { @JsonInclude(Include.NON_NULL) public record EmbeddingsResponse( @JsonProperty("model") String model, - @JsonProperty("embeddings") List> embeddings) { + @JsonProperty("embeddings") List embeddings) { } /** diff --git a/models/spring-ai-ollama/src/test/java/org/springframework/ai/ollama/OllamaEmbeddingModelTests.java b/models/spring-ai-ollama/src/test/java/org/springframework/ai/ollama/OllamaEmbeddingModelTests.java index 975078ea4..a6ed94d41 100644 --- a/models/spring-ai-ollama/src/test/java/org/springframework/ai/ollama/OllamaEmbeddingModelTests.java +++ b/models/spring-ai-ollama/src/test/java/org/springframework/ai/ollama/OllamaEmbeddingModelTests.java @@ -55,9 +55,9 @@ public class OllamaEmbeddingModelTests { when(ollamaApi.embed(embeddingsRequestCaptor.capture())) .thenReturn( - new EmbeddingsResponse("RESPONSE_MODEL_NAME", List.of(List.of(1d, 2d, 3d), List.of(4d, 5d, 6d)))) + new EmbeddingsResponse("RESPONSE_MODEL_NAME", List.of(new float[]{1f, 2f, 3f}, new float[]{4f, 5f, 6f}))) .thenReturn(new EmbeddingsResponse("RESPONSE_MODEL_NAME2", - List.of(List.of(7d, 8d, 9d), List.of(10d, 11d, 12d)))); + List.of(new float[]{7f, 8f, 9f}, new float[]{10f, 11f, 12f}))); // Tests default options var defaultOptions = OllamaOptions.builder().withModel("DEFAULT_MODEL").build(); @@ -69,10 +69,10 @@ public class OllamaEmbeddingModelTests { assertThat(response.getResults()).hasSize(2); assertThat(response.getResults().get(0).getIndex()).isEqualTo(0); - assertThat(response.getResults().get(0).getOutput()).isEqualTo(List.of(1d, 2d, 3d)); + assertThat(response.getResults().get(0).getOutput()).isEqualTo(new float[]{1f, 2f, 3f}); assertThat(response.getResults().get(0).getMetadata()).isEqualTo(EmbeddingResultMetadata.EMPTY); assertThat(response.getResults().get(1).getIndex()).isEqualTo(1); - assertThat(response.getResults().get(1).getOutput()).isEqualTo(List.of(4d, 5d, 6d)); + assertThat(response.getResults().get(1).getOutput()).isEqualTo(new float[]{4f, 5f, 6f}); assertThat(response.getResults().get(1).getMetadata()).isEqualTo(EmbeddingResultMetadata.EMPTY); assertThat(response.getMetadata().getModel()).isEqualTo("RESPONSE_MODEL_NAME"); @@ -94,10 +94,10 @@ public class OllamaEmbeddingModelTests { assertThat(response.getResults()).hasSize(2); assertThat(response.getResults().get(0).getIndex()).isEqualTo(0); - assertThat(response.getResults().get(0).getOutput()).isEqualTo(List.of(7d, 8d, 9d)); + assertThat(response.getResults().get(0).getOutput()).isEqualTo(new float[]{7f, 8f, 9f}); assertThat(response.getResults().get(0).getMetadata()).isEqualTo(EmbeddingResultMetadata.EMPTY); assertThat(response.getResults().get(1).getIndex()).isEqualTo(1); - assertThat(response.getResults().get(1).getOutput()).isEqualTo(List.of(10d, 11d, 12d)); + assertThat(response.getResults().get(1).getOutput()).isEqualTo(new float[]{10f, 11f, 12f}); assertThat(response.getResults().get(1).getMetadata()).isEqualTo(EmbeddingResultMetadata.EMPTY); assertThat(response.getMetadata().getModel()).isEqualTo("RESPONSE_MODEL_NAME2"); diff --git a/models/spring-ai-openai/src/main/java/org/springframework/ai/openai/OpenAiEmbeddingModel.java b/models/spring-ai-openai/src/main/java/org/springframework/ai/openai/OpenAiEmbeddingModel.java index 4265589ae..89f61224a 100644 --- a/models/spring-ai-openai/src/main/java/org/springframework/ai/openai/OpenAiEmbeddingModel.java +++ b/models/spring-ai-openai/src/main/java/org/springframework/ai/openai/OpenAiEmbeddingModel.java @@ -137,7 +137,7 @@ public class OpenAiEmbeddingModel extends AbstractEmbeddingModel { } @Override - public List embed(Document document) { + public float[] embed(Document document) { Assert.notNull(document, "Document must not be null"); return this.embed(document.getFormattedContent(this.metadataMode)); } diff --git a/models/spring-ai-openai/src/main/java/org/springframework/ai/openai/api/OpenAiApi.java b/models/spring-ai-openai/src/main/java/org/springframework/ai/openai/api/OpenAiApi.java index bb289bff9..73149dd75 100644 --- a/models/spring-ai-openai/src/main/java/org/springframework/ai/openai/api/OpenAiApi.java +++ b/models/spring-ai-openai/src/main/java/org/springframework/ai/openai/api/OpenAiApi.java @@ -1102,7 +1102,7 @@ public class OpenAiApi { @JsonInclude(Include.NON_NULL) public record Embedding(// @formatter:off @JsonProperty("index") Integer index, - @JsonProperty("embedding") List embedding, + @JsonProperty("embedding") float[] embedding, @JsonProperty("object") String object) {// @formatter:on /** @@ -1112,7 +1112,7 @@ public class OpenAiApi { * @param embedding The embedding vector, which is a list of floats. The length of * vector depends on the model. */ - public Embedding(Integer index, List embedding) { + public Embedding(Integer index, float[] embedding) { this(index, embedding, "embedding"); } } diff --git a/models/spring-ai-openai/src/test/java/org/springframework/ai/openai/chat/OpenAiRetryTests.java b/models/spring-ai-openai/src/test/java/org/springframework/ai/openai/chat/OpenAiRetryTests.java index 0f486d4c5..bae33e60c 100644 --- a/models/spring-ai-openai/src/test/java/org/springframework/ai/openai/chat/OpenAiRetryTests.java +++ b/models/spring-ai-openai/src/test/java/org/springframework/ai/openai/chat/OpenAiRetryTests.java @@ -197,7 +197,7 @@ public class OpenAiRetryTests { public void openAiEmbeddingTransientError() { EmbeddingList expectedEmbeddings = new EmbeddingList<>("list", - List.of(new Embedding(0, List.of(9.9, 8.8))), "model", new OpenAiApi.Usage(10, 10, 10)); + List.of(new Embedding(0, new float[] { 9.9f, 8.8f })), "model", new OpenAiApi.Usage(10, 10, 10)); when(openAiApi.embeddings(isA(EmbeddingRequest.class))).thenThrow(new TransientAiException("Transient Error 1")) .thenThrow(new TransientAiException("Transient Error 2")) @@ -207,7 +207,7 @@ public class OpenAiRetryTests { .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); } diff --git a/models/spring-ai-postgresml/src/main/java/org/springframework/ai/postgresml/PostgresMlEmbeddingModel.java b/models/spring-ai-postgresml/src/main/java/org/springframework/ai/postgresml/PostgresMlEmbeddingModel.java index 3ba6ab13a..80ddb93f2 100644 --- a/models/spring-ai-postgresml/src/main/java/org/springframework/ai/postgresml/PostgresMlEmbeddingModel.java +++ b/models/spring-ai-postgresml/src/main/java/org/springframework/ai/postgresml/PostgresMlEmbeddingModel.java @@ -31,6 +31,7 @@ import org.springframework.ai.embedding.EmbeddingOptions; 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.ai.model.ModelOptionsUtils; import org.springframework.beans.factory.InitializingBean; import org.springframework.jdbc.core.JdbcTemplate; @@ -57,22 +58,23 @@ public class PostgresMlEmbeddingModel extends AbstractEmbeddingModel implements PG_ARRAY("", null, (rs, i) -> { Array embedding = rs.getArray("embedding"); - return Arrays.stream((Float[]) embedding.getArray()).map(Float::doubleValue).toList(); + return EmbeddingUtils.toPrimitive((Float[]) embedding.getArray()); + }), PG_VECTOR("::vector", "vector", (rs, i) -> { String embedding = rs.getString("embedding"); - return Arrays.stream((embedding.substring(1, embedding.length() - 1) - /* remove leading '[' and trailing ']' */.split(","))).map(Double::parseDouble).toList(); + return EmbeddingUtils.toPrimitive(Arrays.stream((embedding.substring(1, embedding.length() - 1) + /* remove leading '[' and trailing ']' */.split(","))).map(Float::parseFloat).toList()); }); private final String cast; private final String extensionName; - private final RowMapper> rowMapper; + private final RowMapper rowMapper; - VectorType(String cast, String extensionName, RowMapper> rowMapper) { + VectorType(String cast, String extensionName, RowMapper rowMapper) { this.cast = cast; this.extensionName = extensionName; this.rowMapper = rowMapper; @@ -156,7 +158,7 @@ public class PostgresMlEmbeddingModel extends AbstractEmbeddingModel implements @SuppressWarnings("null") @Override - public List embed(String text) { + public float[] embed(String text) { return this.jdbcTemplate.queryForObject( "SELECT pgml.embed(?, ?, ?::JSONB)" + this.defaultOptions.getVectorType().cast + " AS embedding", this.defaultOptions.getVectorType().rowMapper, this.defaultOptions.getTransformer(), text, @@ -164,7 +166,7 @@ public class PostgresMlEmbeddingModel extends AbstractEmbeddingModel implements } @Override - public List embed(Document document) { + public float[] embed(Document document) { return this.embed(document.getFormattedContent(this.defaultOptions.getMetadataMode())); } @@ -175,7 +177,7 @@ public class PostgresMlEmbeddingModel extends AbstractEmbeddingModel implements final PostgresMlEmbeddingOptions optionsToUse = this.mergeOptions(request.getOptions()); List data = new ArrayList<>(); - List> embed = List.of(); + List embed = List.of(); List texts = request.getInstructions(); if (!CollectionUtils.isEmpty(texts)) { @@ -187,7 +189,7 @@ public class PostgresMlEmbeddingModel extends AbstractEmbeddingModel implements preparedStatement.setArray(3, connection.createArrayOf("TEXT", texts.toArray(Object[]::new))); return preparedStatement; }, rs -> { - List> result = new ArrayList<>(); + List result = new ArrayList<>(); while (rs.next()) { result.add(optionsToUse.getVectorType().rowMapper.mapRow(rs, -1)); } diff --git a/models/spring-ai-postgresml/src/test/java/org/springframework/ai/postgresml/PostgresMlEmbeddingModelIT.java b/models/spring-ai-postgresml/src/test/java/org/springframework/ai/postgresml/PostgresMlEmbeddingModelIT.java index 323554fe9..23697f934 100644 --- a/models/spring-ai-postgresml/src/test/java/org/springframework/ai/postgresml/PostgresMlEmbeddingModelIT.java +++ b/models/spring-ai-postgresml/src/test/java/org/springframework/ai/postgresml/PostgresMlEmbeddingModelIT.java @@ -84,7 +84,7 @@ class PostgresMlEmbeddingModelIT { PostgresMlEmbeddingModel embeddingModel = new PostgresMlEmbeddingModel(this.jdbcTemplate); embeddingModel.afterPropertiesSet(); - List embed = embeddingModel.embed("Hello World!"); + float[] embed = embeddingModel.embed("Hello World!"); assertThat(embed).hasSize(768); } @@ -98,7 +98,7 @@ class PostgresMlEmbeddingModelIT { .build()); embeddingModel.afterPropertiesSet(); - List embed = embeddingModel.embed(new Document("Hello World!")); + float[] embed = embeddingModel.embed(new Document("Hello World!")); assertThat(embed).hasSize(768); } @@ -109,7 +109,7 @@ class PostgresMlEmbeddingModelIT { PostgresMlEmbeddingOptions.builder().withTransformer("intfloat/e5-small").build()); embeddingModel.afterPropertiesSet(); - List embed = embeddingModel.embed(new Document("Hello World!")); + float[] embed = embeddingModel.embed(new Document("Hello World!")); assertThat(embed).hasSize(384); } @@ -125,7 +125,7 @@ class PostgresMlEmbeddingModelIT { .build()); embeddingModel.afterPropertiesSet(); - List embed = embeddingModel.embed(new Document("Hello World!")); + float[] embed = embeddingModel.embed(new Document("Hello World!")); assertThat(embed).hasSize(768); } diff --git a/models/spring-ai-qianfan/src/main/java/org/springframework/ai/qianfan/QianFanEmbeddingModel.java b/models/spring-ai-qianfan/src/main/java/org/springframework/ai/qianfan/QianFanEmbeddingModel.java index f8a70a983..1dc08be1a 100644 --- a/models/spring-ai-qianfan/src/main/java/org/springframework/ai/qianfan/QianFanEmbeddingModel.java +++ b/models/spring-ai-qianfan/src/main/java/org/springframework/ai/qianfan/QianFanEmbeddingModel.java @@ -105,7 +105,7 @@ public class QianFanEmbeddingModel extends AbstractEmbeddingModel { } @Override - public List embed(Document document) { + public float[] embed(Document document) { Assert.notNull(document, "Document must not be null"); return this.embed(document.getFormattedContent(this.metadataMode)); } diff --git a/models/spring-ai-qianfan/src/main/java/org/springframework/ai/qianfan/api/QianFanApi.java b/models/spring-ai-qianfan/src/main/java/org/springframework/ai/qianfan/api/QianFanApi.java index f0afd63c6..36eabe441 100644 --- a/models/spring-ai-qianfan/src/main/java/org/springframework/ai/qianfan/api/QianFanApi.java +++ b/models/spring-ai-qianfan/src/main/java/org/springframework/ai/qianfan/api/QianFanApi.java @@ -484,7 +484,7 @@ public class QianFanApi extends AuthApi { public record Embedding( // @formatter:off @JsonProperty("index") Integer index, - @JsonProperty("embedding") List embedding, + @JsonProperty("embedding") float[] embedding, @JsonProperty("object") String object) { // @formatter:on @@ -495,7 +495,7 @@ public class QianFanApi extends AuthApi { * @param embedding The embedding vector, which is a list of floats. The length of * vector depends on the model. */ - public Embedding(Integer index, List embedding) { + public Embedding(Integer index, float[] embedding) { this(index, embedding, "embedding"); } } diff --git a/models/spring-ai-qianfan/src/test/java/org/springframework/ai/qianfan/api/QianFanRetryTests.java b/models/spring-ai-qianfan/src/test/java/org/springframework/ai/qianfan/api/QianFanRetryTests.java index bc23ba3f7..90f95c620 100644 --- a/models/spring-ai-qianfan/src/test/java/org/springframework/ai/qianfan/api/QianFanRetryTests.java +++ b/models/spring-ai-qianfan/src/test/java/org/springframework/ai/qianfan/api/QianFanRetryTests.java @@ -125,10 +125,11 @@ public class QianFanRetryTests { } @Test - public void qianFanChatNonTransientError() { - when(qianFanApi.chatCompletionEntity(isA(ChatCompletionRequest.class))).thenThrow(new RuntimeException("Non Transient Error")); - assertThrows(RuntimeException.class, () -> chatClient.call(new Prompt("text"))); - } + public void qianFanChatNonTransientError() { + when(qianFanApi.chatCompletionEntity(isA(ChatCompletionRequest.class))) + .thenThrow(new RuntimeException("Non Transient Error")); + assertThrows(RuntimeException.class, () -> chatClient.call(new Prompt("text"))); + } @Test public void qianFanChatStreamTransientError() { @@ -150,14 +151,15 @@ public class QianFanRetryTests { } @Test - public void qianFanChatStreamNonTransientError() { - when(qianFanApi.chatCompletionStream(isA(ChatCompletionRequest.class))).thenThrow(new RuntimeException("Non Transient Error")); - assertThrows(RuntimeException.class, () -> chatClient.stream(new Prompt("text"))); - } + public void qianFanChatStreamNonTransientError() { + when(qianFanApi.chatCompletionStream(isA(ChatCompletionRequest.class))) + .thenThrow(new RuntimeException("Non Transient Error")); + assertThrows(RuntimeException.class, () -> chatClient.stream(new Prompt("text"))); + } @Test public void qianFanEmbeddingTransientError() { - QianFanApi.Embedding embedding = new QianFanApi.Embedding(1, List.of(9.9, 8.8)); + QianFanApi.Embedding embedding = new QianFanApi.Embedding(1, new float[] { 9.9f, 8.8f }); EmbeddingList expectedEmbeddings = new EmbeddingList("embedding_list", List.of(embedding), "model", null, null, new Usage(10, 10)); @@ -170,16 +172,17 @@ public class QianFanRetryTests { .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); } @Test - public void qianFanEmbeddingNonTransientError() { - when(qianFanApi.embeddings(isA(EmbeddingRequest.class))).thenThrow(new RuntimeException("Non Transient Error")); - assertThrows(RuntimeException.class, () -> embeddingClient.call(new org.springframework.ai.embedding.EmbeddingRequest(List.of("text1", "text2"), null))); - } + public void qianFanEmbeddingNonTransientError() { + when(qianFanApi.embeddings(isA(EmbeddingRequest.class))).thenThrow(new RuntimeException("Non Transient Error")); + assertThrows(RuntimeException.class, () -> embeddingClient + .call(new org.springframework.ai.embedding.EmbeddingRequest(List.of("text1", "text2"), null))); + } @Test public void qianFanImageTransientError() { @@ -200,11 +203,11 @@ public class QianFanRetryTests { } @Test - public void qianFanImageNonTransientError() { - when(qianFanImageApi.createImage(isA(QianFanImageRequest.class))) - .thenThrow(new RuntimeException("Transient Error 1")); - assertThrows(RuntimeException.class, - () -> imageModel.call(new ImagePrompt(List.of(new ImageMessage("Image Message"))))); - } + public void qianFanImageNonTransientError() { + when(qianFanImageApi.createImage(isA(QianFanImageRequest.class))) + .thenThrow(new RuntimeException("Transient Error 1")); + assertThrows(RuntimeException.class, + () -> imageModel.call(new ImagePrompt(List.of(new ImageMessage("Image Message"))))); + } } diff --git a/models/spring-ai-transformers/src/main/java/org/springframework/ai/transformers/TransformersEmbeddingModel.java b/models/spring-ai-transformers/src/main/java/org/springframework/ai/transformers/TransformersEmbeddingModel.java index 1f6e1b917..e666fbdf2 100644 --- a/models/spring-ai-transformers/src/main/java/org/springframework/ai/transformers/TransformersEmbeddingModel.java +++ b/models/spring-ai-transformers/src/main/java/org/springframework/ai/transformers/TransformersEmbeddingModel.java @@ -210,19 +210,19 @@ public class TransformersEmbeddingModel extends AbstractEmbeddingModel implement } @Override - public List embed(String text) { + public float[] embed(String text) { return embed(List.of(text)).get(0); } @Override - public List embed(Document document) { + public float[] embed(Document document) { return this.embed(document.getFormattedContent(this.metadataMode)); } @Override public EmbeddingResponse embedForResponse(List texts) { List data = new ArrayList<>(); - List> embed = this.embed(texts); + List embed = this.embed(texts); for (int i = 0; i < embed.size(); i++) { data.add(new Embedding(embed.get(i), i)); } @@ -230,7 +230,7 @@ public class TransformersEmbeddingModel extends AbstractEmbeddingModel implement } @Override - public List> embed(List texts) { + public List embed(List texts) { return this.call(new EmbeddingRequest(texts, EmbeddingOptions.EMPTY)) .getResults() .stream() @@ -241,7 +241,7 @@ public class TransformersEmbeddingModel extends AbstractEmbeddingModel implement @Override public EmbeddingResponse call(EmbeddingRequest request) { - List> resultEmbeddings = new ArrayList<>(); + List resultEmbeddings = new ArrayList<>(); try { @@ -286,7 +286,7 @@ public class TransformersEmbeddingModel extends AbstractEmbeddingModel implement NDArray embedding = meanPooling(ndTokenEmbeddings, ndAttentionMask); for (int i = 0; i < embedding.size(0); i++) { - resultEmbeddings.add(toDoubleList(embedding.get(i).toFloatArray())); + resultEmbeddings.add(embedding.get(i).toFloatArray()); } } } @@ -343,16 +343,6 @@ public class TransformersEmbeddingModel extends AbstractEmbeddingModel implement return sumEmbeddings.div(sumMask); } - private List toDoubleList(float[] floats) { - List result = new ArrayList<>(); - if (floats != null && floats.length > 0) { - for (int i = 0; i < floats.length; i++) { - result.add((double) floats[i]); - } - } - return result; - } - private static Resource toResource(String uri) { return new DefaultResourceLoader().getResource(uri); } diff --git a/models/spring-ai-transformers/src/test/java/org/springframework/ai/transformers/TransformersEmbeddingModelTests.java b/models/spring-ai-transformers/src/test/java/org/springframework/ai/transformers/TransformersEmbeddingModelTests.java index 57ee908b3..8c2fb3f01 100644 --- a/models/spring-ai-transformers/src/test/java/org/springframework/ai/transformers/TransformersEmbeddingModelTests.java +++ b/models/spring-ai-transformers/src/test/java/org/springframework/ai/transformers/TransformersEmbeddingModelTests.java @@ -38,35 +38,35 @@ public class TransformersEmbeddingModelTests { TransformersEmbeddingModel embeddingModel = new TransformersEmbeddingModel(); embeddingModel.afterPropertiesSet(); - List embed = embeddingModel.embed("Hello world"); + float[] embed = embeddingModel.embed("Hello world"); assertThat(embed).hasSize(384); - assertThat(DF.format(embed.get(0))).isEqualTo(DF.format(-0.19744634628295898)); - assertThat(DF.format(embed.get(383))).isEqualTo(DF.format(0.17298996448516846)); + assertThat(DF.format(embed[0])).isEqualTo(DF.format(-0.19744634628295898)); + assertThat(DF.format(embed[383])).isEqualTo(DF.format(0.17298996448516846)); } @Test void embedDocument() throws Exception { TransformersEmbeddingModel embeddingModel = new TransformersEmbeddingModel(); embeddingModel.afterPropertiesSet(); - List embed = embeddingModel.embed(new Document("Hello world")); + float[] embed = embeddingModel.embed(new Document("Hello world")); assertThat(embed).hasSize(384); - assertThat(DF.format(embed.get(0))).isEqualTo(DF.format(-0.19744634628295898)); - assertThat(DF.format(embed.get(383))).isEqualTo(DF.format(0.17298996448516846)); + assertThat(DF.format(embed[0])).isEqualTo(DF.format(-0.19744634628295898)); + assertThat(DF.format(embed[383])).isEqualTo(DF.format(0.17298996448516846)); } @Test void embedList() throws Exception { TransformersEmbeddingModel embeddingModel = new TransformersEmbeddingModel(); embeddingModel.afterPropertiesSet(); - List> embed = embeddingModel.embed(List.of("Hello world", "World is big")); + List embed = embeddingModel.embed(List.of("Hello world", "World is big")); assertThat(embed).hasSize(2); assertThat(embed.get(0)).hasSize(384); - assertThat(DF.format(embed.get(0).get(0))).isEqualTo(DF.format(-0.19744634628295898)); - assertThat(DF.format(embed.get(0).get(383))).isEqualTo(DF.format(0.17298996448516846)); + assertThat(DF.format(embed.get(0)[0])).isEqualTo(DF.format(-0.19744634628295898)); + assertThat(DF.format(embed.get(0)[383])).isEqualTo(DF.format(0.17298996448516846)); assertThat(embed.get(1)).hasSize(384); - assertThat(DF.format(embed.get(1).get(0))).isEqualTo(DF.format(0.4293745160102844)); - assertThat(DF.format(embed.get(1).get(383))).isEqualTo(DF.format(0.05501303821802139)); + assertThat(DF.format(embed.get(1)[0])).isEqualTo(DF.format(0.4293745160102844)); + assertThat(DF.format(embed.get(1)[383])).isEqualTo(DF.format(0.05501303821802139)); assertThat(embed.get(0)).isNotEqualTo(embed.get(1)); } @@ -80,12 +80,12 @@ public class TransformersEmbeddingModelTests { assertTrue(embed.getMetadata().isEmpty(), "Expected embed metadata to be empty, but it was not."); assertThat(embed.getResults().get(0).getOutput()).hasSize(384); - assertThat(DF.format(embed.getResults().get(0).getOutput().get(0))).isEqualTo(DF.format(-0.19744634628295898)); - assertThat(DF.format(embed.getResults().get(0).getOutput().get(383))).isEqualTo(DF.format(0.17298996448516846)); + assertThat(DF.format(embed.getResults().get(0).getOutput()[0])).isEqualTo(DF.format(-0.19744634628295898)); + assertThat(DF.format(embed.getResults().get(0).getOutput()[383])).isEqualTo(DF.format(0.17298996448516846)); assertThat(embed.getResults().get(1).getOutput()).hasSize(384); - assertThat(DF.format(embed.getResults().get(1).getOutput().get(0))).isEqualTo(DF.format(0.4293745160102844)); - assertThat(DF.format(embed.getResults().get(1).getOutput().get(383))).isEqualTo(DF.format(0.05501303821802139)); + assertThat(DF.format(embed.getResults().get(1).getOutput()[0])).isEqualTo(DF.format(0.4293745160102844)); + assertThat(DF.format(embed.getResults().get(1).getOutput()[383])).isEqualTo(DF.format(0.05501303821802139)); } @Test diff --git a/models/spring-ai-vertex-ai-embedding/src/main/java/org/springframework/ai/vertexai/embedding/VertexAiEmbeddingUtils.java b/models/spring-ai-vertex-ai-embedding/src/main/java/org/springframework/ai/vertexai/embedding/VertexAiEmbeddingUtils.java index b2122f7bc..a160baeda 100644 --- a/models/spring-ai-vertex-ai-embedding/src/main/java/org/springframework/ai/vertexai/embedding/VertexAiEmbeddingUtils.java +++ b/models/spring-ai-vertex-ai-embedding/src/main/java/org/springframework/ai/vertexai/embedding/VertexAiEmbeddingUtils.java @@ -427,13 +427,14 @@ public abstract class VertexAiEmbeddingUtils { return builder.build(); } - public static List toVector(Value value) { - return value.getListValue() - .getValuesList() - .stream() - .map(Value::getNumberValue) - // .map(Double::floatValue) - .toList(); + public static float[] toVector(Value value) { + float[] floats = new float[value.getListValue().getValuesList().size()]; + int index = 0; + for (Value v : value.getListValue().getValuesList()) { + double d = v.getNumberValue(); + floats[index++] = Double.valueOf(d).floatValue(); + } + return floats; } } diff --git a/models/spring-ai-vertex-ai-embedding/src/main/java/org/springframework/ai/vertexai/embedding/multimodal/VertexAiMultimodalEmbeddingModel.java b/models/spring-ai-vertex-ai-embedding/src/main/java/org/springframework/ai/vertexai/embedding/multimodal/VertexAiMultimodalEmbeddingModel.java index ac787d667..4d3a8066e 100644 --- a/models/spring-ai-vertex-ai-embedding/src/main/java/org/springframework/ai/vertexai/embedding/multimodal/VertexAiMultimodalEmbeddingModel.java +++ b/models/spring-ai-vertex-ai-embedding/src/main/java/org/springframework/ai/vertexai/embedding/multimodal/VertexAiMultimodalEmbeddingModel.java @@ -198,7 +198,7 @@ public class VertexAiMultimodalEmbeddingModel implements DocumentEmbeddingModel for (Value prediction : embeddingResponse.getPredictionsList()) { if (prediction.getStructValue().containsFields("textEmbedding")) { Value textEmbedding = prediction.getStructValue().getFieldsOrThrow("textEmbedding"); - List textVector = VertexAiEmbeddingUtils.toVector(textEmbedding); + float[] textVector = VertexAiEmbeddingUtils.toVector(textEmbedding); var docMetadata = documentMetadata.get(ModalityType.TEXT); embeddingList.add(new Embedding(textVector, index++, new EmbeddingResultMetadata(docMetadata.documentId, @@ -206,7 +206,7 @@ public class VertexAiMultimodalEmbeddingModel implements DocumentEmbeddingModel } if (prediction.getStructValue().containsFields("imageEmbedding")) { Value imageEmbedding = prediction.getStructValue().getFieldsOrThrow("imageEmbedding"); - List imageVector = VertexAiEmbeddingUtils.toVector(imageEmbedding); + float[] imageVector = VertexAiEmbeddingUtils.toVector(imageEmbedding); var docMetadata = documentMetadata.get(ModalityType.IMAGE); embeddingList @@ -220,7 +220,7 @@ public class VertexAiMultimodalEmbeddingModel implements DocumentEmbeddingModel .getValues(0) .getStructValue() .getFieldsOrThrow("embedding"); - List videoVector = VertexAiEmbeddingUtils.toVector(embeddings); + float[] videoVector = VertexAiEmbeddingUtils.toVector(embeddings); var docMetadata = documentMetadata.get(ModalityType.VIDEO); embeddingList diff --git a/models/spring-ai-vertex-ai-embedding/src/main/java/org/springframework/ai/vertexai/embedding/text/VertexAiTextEmbeddingModel.java b/models/spring-ai-vertex-ai-embedding/src/main/java/org/springframework/ai/vertexai/embedding/text/VertexAiTextEmbeddingModel.java index aba0dc2ee..7f2654db5 100644 --- a/models/spring-ai-vertex-ai-embedding/src/main/java/org/springframework/ai/vertexai/embedding/text/VertexAiTextEmbeddingModel.java +++ b/models/spring-ai-vertex-ai-embedding/src/main/java/org/springframework/ai/vertexai/embedding/text/VertexAiTextEmbeddingModel.java @@ -66,7 +66,7 @@ public class VertexAiTextEmbeddingModel extends AbstractEmbeddingModel { } @Override - public List embed(Document document) { + public float[] embed(Document document) { Assert.notNull(document, "Document must not be null"); return this.embed(document.getFormattedContent()); } @@ -125,7 +125,7 @@ public class VertexAiTextEmbeddingModel extends AbstractEmbeddingModel { Value values = embeddings.getStructValue().getFieldsOrThrow("values"); - List vectorValues = VertexAiEmbeddingUtils.toVector(values); + float[] vectorValues = VertexAiEmbeddingUtils.toVector(values); embeddingList.add(new Embedding(vectorValues, index++)); } diff --git a/models/spring-ai-vertex-ai-embedding/src/test/java/org/springframework/ai/vertexai/embedding/multimodal/VertexAiMultimodalEmbeddingModelIT.java b/models/spring-ai-vertex-ai-embedding/src/test/java/org/springframework/ai/vertexai/embedding/multimodal/VertexAiMultimodalEmbeddingModelIT.java index 8c6777149..7f3f1a616 100644 --- a/models/spring-ai-vertex-ai-embedding/src/test/java/org/springframework/ai/vertexai/embedding/multimodal/VertexAiMultimodalEmbeddingModelIT.java +++ b/models/spring-ai-vertex-ai-embedding/src/test/java/org/springframework/ai/vertexai/embedding/multimodal/VertexAiMultimodalEmbeddingModelIT.java @@ -74,7 +74,7 @@ class VertexAiMultimodalEmbeddingModelIT { assertThat(embeddingResponse.getMetadata().getUsage().getTotalTokens()) .as("Total tokens in metadata should be 0") - .isEqualTo("0"); + .isEqualTo(0L); assertThat(multiModelEmbeddingModel.dimensions()).isEqualTo(1408); } diff --git a/models/spring-ai-vertex-ai-palm2/src/main/java/org/springframework/ai/vertexai/palm2/VertexAiPaLm2EmbeddingModel.java b/models/spring-ai-vertex-ai-palm2/src/main/java/org/springframework/ai/vertexai/palm2/VertexAiPaLm2EmbeddingModel.java index 59ee71cd1..3fcdb9357 100644 --- a/models/spring-ai-vertex-ai-palm2/src/main/java/org/springframework/ai/vertexai/palm2/VertexAiPaLm2EmbeddingModel.java +++ b/models/spring-ai-vertex-ai-palm2/src/main/java/org/springframework/ai/vertexai/palm2/VertexAiPaLm2EmbeddingModel.java @@ -37,7 +37,7 @@ public class VertexAiPaLm2EmbeddingModel extends AbstractEmbeddingModel { } @Override - public List embed(Document document) { + public float[] embed(Document document) { return embed(document.getContent()); } diff --git a/models/spring-ai-vertex-ai-palm2/src/main/java/org/springframework/ai/vertexai/palm2/api/VertexAiPaLm2Api.java b/models/spring-ai-vertex-ai-palm2/src/main/java/org/springframework/ai/vertexai/palm2/api/VertexAiPaLm2Api.java index 94e7c3295..cf47be62c 100644 --- a/models/spring-ai-vertex-ai-palm2/src/main/java/org/springframework/ai/vertexai/palm2/api/VertexAiPaLm2Api.java +++ b/models/spring-ai-vertex-ai-palm2/src/main/java/org/springframework/ai/vertexai/palm2/api/VertexAiPaLm2Api.java @@ -16,6 +16,7 @@ package org.springframework.ai.vertexai.palm2.api; import java.io.IOException; +import java.util.Arrays; import java.util.List; import java.util.Map; import java.util.function.Consumer; @@ -364,8 +365,17 @@ public class VertexAiPaLm2Api { */ @JsonInclude(Include.NON_NULL) public record Embedding( - @JsonProperty("value") List value) { + @JsonProperty("value") float[] value) { + @Override + public final int hashCode() { + return Arrays.hashCode(value); + } + + @Override + public final boolean equals(Object arg0) { + return Arrays.equals(value,((Embedding) arg0).value); + } } /** diff --git a/models/spring-ai-vertex-ai-palm2/src/test/java/org/springframework/ai/vertexai/palm2/api/VertexAiPaLm2ApiTests.java b/models/spring-ai-vertex-ai-palm2/src/test/java/org/springframework/ai/vertexai/palm2/api/VertexAiPaLm2ApiTests.java index aedce410c..2f58d9dd6 100644 --- a/models/spring-ai-vertex-ai-palm2/src/test/java/org/springframework/ai/vertexai/palm2/api/VertexAiPaLm2ApiTests.java +++ b/models/spring-ai-vertex-ai-palm2/src/test/java/org/springframework/ai/vertexai/palm2/api/VertexAiPaLm2ApiTests.java @@ -95,7 +95,7 @@ public class VertexAiPaLm2ApiTests { String text = "Hello, how are you?"; - Embedding expectedEmbedding = new Embedding(List.of(0.1, 0.2, 0.3)); + Embedding expectedEmbedding = new Embedding(new float[] { 0.1f, 0.2f, 0.3f }); server .expect(requestToUriTemplate("/models/{generative}:embedText?key={apiKey}", @@ -117,8 +117,8 @@ public class VertexAiPaLm2ApiTests { List texts = List.of("Hello, how are you?", "I'm fine, thank you."); - List expectedEmbeddings = List.of(new Embedding(List.of(0.1, 0.2, 0.3)), - new Embedding(List.of(0.4, 0.5, 0.6))); + List expectedEmbeddings = List.of(new Embedding(new float[] { 0.1f, 0.2f, 0.3f }), + new Embedding(new float[] { 0.4f, 0.5f, 0.6f })); server .expect(requestToUriTemplate("/models/{generative}:batchEmbedText?key={apiKey}", diff --git a/models/spring-ai-zhipuai/src/main/java/org/springframework/ai/zhipuai/ZhiPuAiEmbeddingModel.java b/models/spring-ai-zhipuai/src/main/java/org/springframework/ai/zhipuai/ZhiPuAiEmbeddingModel.java index 3ee5db59f..573b4bacc 100644 --- a/models/spring-ai-zhipuai/src/main/java/org/springframework/ai/zhipuai/ZhiPuAiEmbeddingModel.java +++ b/models/spring-ai-zhipuai/src/main/java/org/springframework/ai/zhipuai/ZhiPuAiEmbeddingModel.java @@ -103,7 +103,7 @@ public class ZhiPuAiEmbeddingModel extends AbstractEmbeddingModel { } @Override - public List embed(Document document) { + public float[] embed(Document document) { Assert.notNull(document, "Document must not be null"); return this.embed(document.getFormattedContent(this.metadataMode)); } @@ -119,14 +119,14 @@ public class ZhiPuAiEmbeddingModel extends AbstractEmbeddingModel { "ZhiPu Embedding does not support batch embedding. Will make multiple API calls to embed(Document)"); } - List> embeddingList = new ArrayList<>(); + List embeddingList = new ArrayList<>(); for (String inputContent : request.getInstructions()) { var apiRequest = createZhiPuEmbeddingRequest(inputContent, request.getOptions()); ZhiPuAiApi.EmbeddingList response = this.zhiPuAiApi.embeddings(apiRequest) .getBody(); if (response == null || response.data() == null || response.data().isEmpty()) { logger.warn("No embeddings returned for input: {}", inputContent); - embeddingList.add(List.of()); + embeddingList.add(new float[0]); } else { embeddingList.add(response.data().get(0).embedding()); diff --git a/models/spring-ai-zhipuai/src/main/java/org/springframework/ai/zhipuai/api/ZhiPuAiApi.java b/models/spring-ai-zhipuai/src/main/java/org/springframework/ai/zhipuai/api/ZhiPuAiApi.java index cae8a6cc9..a7fc4c9c9 100644 --- a/models/spring-ai-zhipuai/src/main/java/org/springframework/ai/zhipuai/api/ZhiPuAiApi.java +++ b/models/spring-ai-zhipuai/src/main/java/org/springframework/ai/zhipuai/api/ZhiPuAiApi.java @@ -738,7 +738,7 @@ public class ZhiPuAiApi { @JsonInclude(Include.NON_NULL) public record Embedding( @JsonProperty("index") Integer index, - @JsonProperty("embedding") List embedding, + @JsonProperty("embedding") float[] embedding, @JsonProperty("object") String object) { /** @@ -747,7 +747,7 @@ public class ZhiPuAiApi { * @param index The index of the embedding in the list of embeddings. * @param embedding The embedding vector, which is a list of floats. The length of vector depends on the model. */ - public Embedding(Integer index, List embedding) { + public Embedding(Integer index, float[] embedding) { this(index, embedding, "embedding"); } } diff --git a/models/spring-ai-zhipuai/src/test/java/org/springframework/ai/zhipuai/api/ZhiPuAiRetryTests.java b/models/spring-ai-zhipuai/src/test/java/org/springframework/ai/zhipuai/api/ZhiPuAiRetryTests.java index 6feb7747c..8b81ce131 100644 --- a/models/spring-ai-zhipuai/src/test/java/org/springframework/ai/zhipuai/api/ZhiPuAiRetryTests.java +++ b/models/spring-ai-zhipuai/src/test/java/org/springframework/ai/zhipuai/api/ZhiPuAiRetryTests.java @@ -171,7 +171,7 @@ public class ZhiPuAiRetryTests { public void zhiPuAiEmbeddingTransientError() { EmbeddingList expectedEmbeddings = new EmbeddingList<>("list", - List.of(new Embedding(0, List.of(9.9, 8.8))), "model", new ZhiPuAiApi.Usage(10, 10, 10)); + List.of(new Embedding(0, new float[] { 9.9f, 8.8f })), "model", new ZhiPuAiApi.Usage(10, 10, 10)); when(zhiPuAiApi.embeddings(isA(EmbeddingRequest.class))) .thenThrow(new TransientAiException("Transient Error 1")) @@ -182,7 +182,7 @@ public class ZhiPuAiRetryTests { .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); } diff --git a/spring-ai-core/src/main/java/org/springframework/ai/document/Document.java b/spring-ai-core/src/main/java/org/springframework/ai/document/Document.java index 7b65ad95b..7722a4dfe 100644 --- a/spring-ai-core/src/main/java/org/springframework/ai/document/Document.java +++ b/spring-ai-core/src/main/java/org/springframework/ai/document/Document.java @@ -65,7 +65,7 @@ public class Document implements MediaContent { * Embedding of the document. Note: ephemeral field. */ @JsonProperty(index = 100) - private List embedding = new ArrayList<>(); + private float[] embedding = new float[0]; /** * Mutable, ephemeral, content to text formatter. Defaults to Document text. @@ -206,7 +206,7 @@ public class Document implements MediaContent { return formatter.format(this, metadataMode); } - public void setEmbedding(List embedding) { + public void setEmbedding(float[] embedding) { Assert.notNull(embedding, "embedding must not be null"); this.embedding = embedding; } @@ -224,7 +224,7 @@ public class Document implements MediaContent { return this.metadata; } - public List getEmbedding() { + public float[] getEmbedding() { return this.embedding; } diff --git a/spring-ai-core/src/main/java/org/springframework/ai/embedding/AbstractEmbeddingModel.java b/spring-ai-core/src/main/java/org/springframework/ai/embedding/AbstractEmbeddingModel.java index 6d857d3a3..1c6c1c374 100644 --- a/spring-ai-core/src/main/java/org/springframework/ai/embedding/AbstractEmbeddingModel.java +++ b/spring-ai-core/src/main/java/org/springframework/ai/embedding/AbstractEmbeddingModel.java @@ -53,7 +53,7 @@ public abstract class AbstractEmbeddingModel implements EmbeddingModel { else { // Determine the dimensions empirically. // Generate an embedding and count the dimension size; - return embeddingModel.embed(dummyContent).size(); + return embeddingModel.embed(dummyContent).length; } } diff --git a/spring-ai-core/src/main/java/org/springframework/ai/embedding/Embedding.java b/spring-ai-core/src/main/java/org/springframework/ai/embedding/Embedding.java index c5eefe0a7..daaa20d0e 100644 --- a/spring-ai-core/src/main/java/org/springframework/ai/embedding/Embedding.java +++ b/spring-ai-core/src/main/java/org/springframework/ai/embedding/Embedding.java @@ -15,7 +15,6 @@ */ package org.springframework.ai.embedding; -import java.util.List; import java.util.Objects; import org.springframework.ai.model.ModelResult; @@ -23,9 +22,9 @@ import org.springframework.ai.model.ModelResult; /** * Represents a single embedding vector. */ -public class Embedding implements ModelResult> { +public class Embedding implements ModelResult { - private List embedding; + private float[] embedding; private Integer index; @@ -36,7 +35,7 @@ public class Embedding implements ModelResult> { * @param embedding the embedding vector values. * @param index the embedding index in a list of embeddings. */ - public Embedding(List embedding, Integer index) { + public Embedding(float[] embedding, Integer index) { this(embedding, index, EmbeddingResultMetadata.EMPTY); } @@ -46,7 +45,7 @@ public class Embedding implements ModelResult> { * @param index the embedding index in a list of embeddings. * @param metadata the metadata associated with the embedding. */ - public Embedding(List embedding, Integer index, EmbeddingResultMetadata metadata) { + public Embedding(float[] embedding, Integer index, EmbeddingResultMetadata metadata) { this.embedding = embedding; this.index = index; this.metadata = metadata; @@ -56,8 +55,8 @@ public class Embedding implements ModelResult> { * @return Get the embedding vector values. */ @Override - public List getOutput() { - return this.embedding; + public float[] getOutput() { + return embedding; } /** @@ -91,7 +90,7 @@ public class Embedding implements ModelResult> { @Override public String toString() { - String message = this.embedding.isEmpty() ? "" : ""; + String message = this.embedding.length == 0 ? "" : ""; return "Embedding{" + "embedding=" + message + ", index=" + this.index + '}'; } diff --git a/spring-ai-core/src/main/java/org/springframework/ai/embedding/EmbeddingModel.java b/spring-ai-core/src/main/java/org/springframework/ai/embedding/EmbeddingModel.java index a40ee22a6..aefc30b4e 100644 --- a/spring-ai-core/src/main/java/org/springframework/ai/embedding/EmbeddingModel.java +++ b/spring-ai-core/src/main/java/org/springframework/ai/embedding/EmbeddingModel.java @@ -34,9 +34,10 @@ public interface EmbeddingModel extends Model embed(String text) { + default float[] embed(String text) { Assert.notNull(text, "Text must not be null"); - return this.embed(List.of(text)).iterator().next(); + List response = this.embed(List.of(text)); + return response.iterator().next(); } /** @@ -44,14 +45,14 @@ public interface EmbeddingModel extends Model embed(Document document); + float[] embed(Document document); /** * Embeds a batch of texts into vectors. * @param texts list of texts to embed. * @return list of list of embedded vectors. */ - default List> embed(List texts) { + default List embed(List texts) { Assert.notNull(texts, "Texts must not be null"); return this.call(new EmbeddingRequest(texts, EmbeddingOptionsBuilder.builder().build())) .getResults() @@ -75,7 +76,7 @@ public interface EmbeddingModel extends Model doubleToFloat(final List doubles) { + return doubles.stream().map(f -> f.floatValue()).toList(); + } + + public static float[] toPrimitive(List floats) { + return toPrimitive(floats.toArray(new Float[floats.size()])); + } + + public static float[] toPrimitive(final Float[] array) { + if (array == null) { + return null; + } + if (array.length == 0) { + return EMPTY_FLOAT_ARRAY; + } + final float[] result = new float[array.length]; + for (int i = 0; i < array.length; i++) { + result[i] = array[i].floatValue(); + } + return result; + } + + public static Float[] toFloatArray(final float[] array) { + if (array == null) { + return null; + } + if (array.length == 0) { + return new Float[0]; + } + final Float[] result = new Float[array.length]; + for (int i = 0; i < array.length; i++) { + result[i] = array[i]; + } + return result; + } + + public static List toList(float[] floats) { + + List output = new ArrayList(); + for (float value : floats) { + output.add(value); + } + return output; + } + +} diff --git a/spring-ai-core/src/main/java/org/springframework/ai/vectorstore/SimpleVectorStore.java b/spring-ai-core/src/main/java/org/springframework/ai/vectorstore/SimpleVectorStore.java index 7c681b678..7026fa867 100644 --- a/spring-ai-core/src/main/java/org/springframework/ai/vectorstore/SimpleVectorStore.java +++ b/spring-ai-core/src/main/java/org/springframework/ai/vectorstore/SimpleVectorStore.java @@ -72,7 +72,7 @@ public class SimpleVectorStore implements VectorStore { public void add(List documents) { for (Document document : documents) { logger.info("Calling EmbeddingModel for document id = {}", document.getId()); - List embedding = this.embeddingModel.embed(document); + float[] embedding = this.embeddingModel.embed(document); document.setEmbedding(embedding); this.store.put(document.getId(), document); } @@ -93,7 +93,7 @@ public class SimpleVectorStore implements VectorStore { "The [" + this.getClass() + "] doesn't support metadata filtering!"); } - List userQueryEmbedding = getUserQueryEmbedding(request.getQuery()); + float[] userQueryEmbedding = getUserQueryEmbedding(request.getQuery()); return this.store.values() .stream() .map(entry -> new Similarity(entry.getId(), @@ -186,7 +186,7 @@ public class SimpleVectorStore implements VectorStore { return json; } - private List getUserQueryEmbedding(String query) { + private float[] getUserQueryEmbedding(String query) { return this.embeddingModel.embed(query); } @@ -209,17 +209,17 @@ public class SimpleVectorStore implements VectorStore { throw new UnsupportedOperationException("This is a utility class and cannot be instantiated"); } - public static double cosineSimilarity(List vectorX, List vectorY) { + public static double cosineSimilarity(float[] vectorX, float[] vectorY) { if (vectorX == null || vectorY == null) { throw new RuntimeException("Vectors must not be null"); } - if (vectorX.size() != vectorY.size()) { + if (vectorX.length != vectorY.length) { throw new IllegalArgumentException("Vectors lengths must be equal"); } - double dotProduct = dotProduct(vectorX, vectorY); - double normX = norm(vectorX); - double normY = norm(vectorY); + float dotProduct = dotProduct(vectorX, vectorY); + float normX = norm(vectorX); + float normY = norm(vectorY); if (normX == 0 || normY == 0) { throw new IllegalArgumentException("Vectors cannot have zero norm"); @@ -228,20 +228,20 @@ public class SimpleVectorStore implements VectorStore { return dotProduct / (Math.sqrt(normX) * Math.sqrt(normY)); } - public static double dotProduct(List vectorX, List vectorY) { - if (vectorX.size() != vectorY.size()) { + public static float dotProduct(float[] vectorX, float[] vectorY) { + if (vectorX.length != vectorY.length) { throw new IllegalArgumentException("Vectors lengths must be equal"); } - double result = 0; - for (int i = 0; i < vectorX.size(); ++i) { - result += vectorX.get(i) * vectorY.get(i); + float result = 0; + for (int i = 0; i < vectorX.length; ++i) { + result += vectorX[i] * vectorY[i]; } return result; } - public static double norm(List vector) { + public static float norm(float[] vector) { return dotProduct(vector, vector); } diff --git a/spring-ai-core/src/test/java/org/springframework/ai/embedding/AbstractEmbeddingModelTests.java b/spring-ai-core/src/test/java/org/springframework/ai/embedding/AbstractEmbeddingModelTests.java index d6d51454b..327732c9e 100644 --- a/spring-ai-core/src/test/java/org/springframework/ai/embedding/AbstractEmbeddingModelTests.java +++ b/spring-ai-core/src/test/java/org/springframework/ai/embedding/AbstractEmbeddingModelTests.java @@ -48,17 +48,17 @@ public class AbstractEmbeddingModelTests { EmbeddingModel dummy = new EmbeddingModel() { @Override - public List embed(String text) { - return List.of(0.1, 0.1, 0.1); + public float[] embed(String text) { + return new float[] { 0.1f, 0.1f, 0.1f }; } @Override - public List embed(Document document) { + public float[] embed(Document document) { throw new UnsupportedOperationException("Unimplemented method 'embed'"); } @Override - public List> embed(List texts) { + public List embed(List texts) { throw new UnsupportedOperationException("Unimplemented method 'embed'"); } @@ -87,7 +87,7 @@ public class AbstractEmbeddingModelTests { @Test public void testUnknownModelDimension() { - when(embeddingModel.embed(eq("Hello world!"))).thenReturn(List.of(0.1, 0.1, 0.1)); + when(embeddingModel.embed(eq("Hello world!"))).thenReturn(new float[]{0.1f, 0.1f, 0.1f}); assertThat(AbstractEmbeddingModel.dimensions(embeddingModel, "unknown_model", "Hello world!")).isEqualTo(3); } diff --git a/spring-ai-docs/src/main/antora/modules/ROOT/pages/api/embeddings.adoc b/spring-ai-docs/src/main/antora/modules/ROOT/pages/api/embeddings.adoc index c98c246f8..f796b9a6f 100644 --- a/spring-ai-docs/src/main/antora/modules/ROOT/pages/api/embeddings.adoc +++ b/spring-ai-docs/src/main/antora/modules/ROOT/pages/api/embeddings.adoc @@ -50,14 +50,14 @@ public interface EmbeddingModel extends Model embed(Document document); + float[] embed(Document document); /** * Embeds the given text into a vector. * @param text the text to embed. * @return the embedded vector. */ - default List embed(String text) { + default float[] embed(String text) { Assert.notNull(text, "Text must not be null"); return this.embed(List.of(text)).iterator().next(); } @@ -67,7 +67,7 @@ public interface EmbeddingModel extends Model> embed(List texts) { + default List embed(List texts) { Assert.notNull(texts, "Texts must not be null"); return this.call(new EmbeddingRequest(texts, EmbeddingOptions.EMPTY)) .getResults() @@ -102,7 +102,7 @@ The embed methods offer various options for converting text into embeddings, acc Multiple shortcut methods are provided for embedding text, including the `embed(String text)` method, which takes a single string and returns the corresponding embedding vector. All shortcuts are implemented around the `call` method, which is the primary method for invoking the embedding model. -Typically the embedding returns a lists of doubles, representing the embeddings in a numerical vector format. +Typically the embedding returns a lists of floats, representing the embeddings in a numerical vector format. The `embedForResponse` method provides a more comprehensive output, potentially including additional information about the embeddings. @@ -146,8 +146,8 @@ The `Embedding` represents a single embedding vector. [source,java] ---- -public class Embedding implements ModelResult> { - private List embedding; +public class Embedding implements ModelResult { + private float[] embedding; private Integer index; private EmbeddingResultMetadata metadata; // other methods omitted diff --git a/spring-ai-docs/src/main/antora/modules/ROOT/pages/api/vectordbs.adoc b/spring-ai-docs/src/main/antora/modules/ROOT/pages/api/vectordbs.adoc index dc58048c1..6e7e0e669 100644 --- a/spring-ai-docs/src/main/antora/modules/ROOT/pages/api/vectordbs.adoc +++ b/spring-ai-docs/src/main/antora/modules/ROOT/pages/api/vectordbs.adoc @@ -71,7 +71,7 @@ To insert data into the vector database, encapsulate it within a `Document` obje The `Document` class encapsulates content from a data source, such as a PDF or Word document, and includes text represented as a string. It also contains metadata in the form of key-value pairs, including details such as the filename. -Upon insertion into the vector database, the text content is transformed into a numerical array, or a `List`, known as vector embeddings, using an embedding model. Embedding models, such as https://en.wikipedia.org/wiki/Word2vec[Word2Vec], https://en.wikipedia.org/wiki/GloVe_(machine_learning)[GLoVE], and https://en.wikipedia.org/wiki/BERT_(language_model)[BERT], or OpenAI's `text-embedding-ada-002`, are used to convert words, sentences, or paragraphs into these vector embeddings. +Upon insertion into the vector database, the text content is transformed into a numerical array, or a `float[]`, known as vector embeddings, using an embedding model. Embedding models, such as https://en.wikipedia.org/wiki/Word2vec[Word2Vec], https://en.wikipedia.org/wiki/GloVe_(machine_learning)[GLoVE], and https://en.wikipedia.org/wiki/BERT_(language_model)[BERT], or OpenAI's `text-embedding-ada-002`, are used to convert words, sentences, or paragraphs into these vector embeddings. The vector database's role is to store and facilitate similarity searches for these embeddings. It does not generate the embeddings itself. For creating vector embeddings, the `EmbeddingModel` should be utilized. diff --git a/spring-ai-spring-boot-autoconfigure/src/test/java/org/springframework/ai/autoconfigure/transformers/TransformersEmbeddingModelAutoConfigurationIT.java b/spring-ai-spring-boot-autoconfigure/src/test/java/org/springframework/ai/autoconfigure/transformers/TransformersEmbeddingModelAutoConfigurationIT.java index 2cb96123e..6a7cacabb 100644 --- a/spring-ai-spring-boot-autoconfigure/src/test/java/org/springframework/ai/autoconfigure/transformers/TransformersEmbeddingModelAutoConfigurationIT.java +++ b/spring-ai-spring-boot-autoconfigure/src/test/java/org/springframework/ai/autoconfigure/transformers/TransformersEmbeddingModelAutoConfigurationIT.java @@ -50,10 +50,10 @@ public class TransformersEmbeddingModelAutoConfigurationIT { EmbeddingModel embeddingModel = context.getBean(EmbeddingModel.class); assertThat(embeddingModel).isInstanceOf(TransformersEmbeddingModel.class); - List> embeddings = embeddingModel.embed(List.of("Spring Framework", "Spring AI")); + List embeddings = embeddingModel.embed(List.of("Spring Framework", "Spring AI")); assertThat(embeddings.size()).isEqualTo(2); // batch size - assertThat(embeddings.get(0).size()).isEqualTo(embeddingModel.dimensions()); // dimensions + assertThat(embeddings.get(0).length).isEqualTo(embeddingModel.dimensions()); // dimensions // size }); } @@ -80,10 +80,10 @@ public class TransformersEmbeddingModelAutoConfigurationIT { assertThat(embeddingModel.dimensions()).isEqualTo(384); - List> embeddings = embeddingModel.embed(List.of("Spring Framework", "Spring AI")); + List embeddings = embeddingModel.embed(List.of("Spring Framework", "Spring AI")); assertThat(embeddings.size()).isEqualTo(2); // batch size - assertThat(embeddings.get(0).size()).isEqualTo(embeddingModel.dimensions()); // dimensions + assertThat(embeddings.get(0).length).isEqualTo(embeddingModel.dimensions()); // dimensions // size }); } diff --git a/spring-ai-spring-boot-autoconfigure/src/test/java/org/springframework/ai/autoconfigure/vertexai/embedding/VertexAiTextEmbeddingModelAutoConfigurationIT.java b/spring-ai-spring-boot-autoconfigure/src/test/java/org/springframework/ai/autoconfigure/vertexai/embedding/VertexAiTextEmbeddingModelAutoConfigurationIT.java index b4385d64f..3cb3481a2 100644 --- a/spring-ai-spring-boot-autoconfigure/src/test/java/org/springframework/ai/autoconfigure/vertexai/embedding/VertexAiTextEmbeddingModelAutoConfigurationIT.java +++ b/spring-ai-spring-boot-autoconfigure/src/test/java/org/springframework/ai/autoconfigure/vertexai/embedding/VertexAiTextEmbeddingModelAutoConfigurationIT.java @@ -60,10 +60,10 @@ public class VertexAiTextEmbeddingModelAutoConfigurationIT { VertexAiTextEmbeddingModel embeddingModel = context.getBean(VertexAiTextEmbeddingModel.class); assertThat(embeddingModel).isInstanceOf(VertexAiTextEmbeddingModel.class); - List> embeddings = embeddingModel.embed(List.of("Spring Framework", "Spring AI")); + List embeddings = embeddingModel.embed(List.of("Spring Framework", "Spring AI")); assertThat(embeddings.size()).isEqualTo(2); // batch size - assertThat(embeddings.get(0).size()).isEqualTo(embeddingModel.dimensions()); + assertThat(embeddings.get(0).length).isEqualTo(embeddingModel.dimensions()); }); } @@ -113,7 +113,7 @@ public class VertexAiTextEmbeddingModelAutoConfigurationIT { assertThat(embeddingResponse.getResults().get(0).getOutput()).hasSize(1408); assertThat(embeddingResponse.getMetadata().getModel()).isEqualTo("multimodalembedding@001"); - assertThat(embeddingResponse.getMetadata().getUsage()).isEqualTo(0); + assertThat(embeddingResponse.getMetadata().getUsage().getPromptTokens()).isEqualTo(0); assertThat(multiModelEmbeddingModel.dimensions()).isEqualTo(1408); diff --git a/vector-stores/spring-ai-azure-store/src/main/java/org/springframework/ai/vectorstore/azure/AzureVectorStore.java b/vector-stores/spring-ai-azure-store/src/main/java/org/springframework/ai/vectorstore/azure/AzureVectorStore.java index d09b8dc71..b51f292fc 100644 --- a/vector-stores/spring-ai-azure-store/src/main/java/org/springframework/ai/vectorstore/azure/AzureVectorStore.java +++ b/vector-stores/spring-ai-azure-store/src/main/java/org/springframework/ai/vectorstore/azure/AzureVectorStore.java @@ -38,6 +38,7 @@ import org.slf4j.Logger; import org.slf4j.LoggerFactory; import org.springframework.ai.document.Document; import org.springframework.ai.embedding.EmbeddingModel; +import org.springframework.ai.model.EmbeddingUtils; import org.springframework.ai.vectorstore.SearchRequest; import org.springframework.ai.vectorstore.VectorStore; import org.springframework.ai.vectorstore.filter.FilterExpressionConverter; @@ -281,9 +282,10 @@ public class AzureVectorStore implements VectorStore, InitializingBean { Assert.notNull(request, "The search request must not be null."); - var searchEmbedding = toFloatList(embeddingModel.embed(request.getQuery())); + var searchEmbedding = embeddingModel.embed(request.getQuery()); - final var vectorQuery = new VectorizedQuery(searchEmbedding).setKNearestNeighborsCount(request.getTopK()) + final var vectorQuery = new VectorizedQuery(EmbeddingUtils.toList(searchEmbedding)) + .setKNearestNeighborsCount(request.getTopK()) // Set the fields to compare the vector against. This is a comma-delimited // list of field names. .setFields(EMBEDDING_FIELD_NAME); @@ -311,7 +313,7 @@ public class AzureVectorStore implements VectorStore, InitializingBean { metadata.put(DISTANCE_METADATA_FIELD_NAME, 1 - (float) result.getScore()); final Document doc = new Document(entry.id(), entry.content(), metadata); - doc.setEmbedding(entry.embedding()); + doc.setEmbedding(EmbeddingUtils.toPrimitive(entry.embedding())); return doc; @@ -319,14 +321,10 @@ public class AzureVectorStore implements VectorStore, InitializingBean { .collect(Collectors.toList()); } - private List toFloatList(List doubleList) { - return doubleList.stream().map(Double::floatValue).toList(); - } - /** * Internal data structure for retrieving and storing documents. */ - private record AzureSearchDocument(String id, String content, List embedding, String metadata) { + private record AzureSearchDocument(String id, String content, List embedding, String metadata) { } @Override diff --git a/vector-stores/spring-ai-cassandra-store/src/main/java/org/springframework/ai/vectorstore/CassandraVectorStore.java b/vector-stores/spring-ai-cassandra-store/src/main/java/org/springframework/ai/vectorstore/CassandraVectorStore.java index 09cce8a67..00943cca2 100644 --- a/vector-stores/spring-ai-cassandra-store/src/main/java/org/springframework/ai/vectorstore/CassandraVectorStore.java +++ b/vector-stores/spring-ai-cassandra-store/src/main/java/org/springframework/ai/vectorstore/CassandraVectorStore.java @@ -34,6 +34,7 @@ import org.slf4j.LoggerFactory; import org.springframework.ai.document.Document; import org.springframework.ai.embedding.EmbeddingModel; +import org.springframework.ai.model.EmbeddingUtils; import org.springframework.ai.vectorstore.CassandraVectorStoreConfig.SchemaColumn; import org.springframework.ai.vectorstore.filter.FilterExpressionConverter; @@ -163,7 +164,7 @@ public class CassandraVectorStore implements VectorStore, AutoCloseable { futures[i++] = CompletableFuture.runAsync(() -> { List primaryKeyValues = this.conf.documentIdTranslator.apply(d.getId()); - if (null == d.getEmbedding() || d.getEmbedding().isEmpty()) { + if (null == d.getEmbedding() || d.getEmbedding().length == 0) { d.setEmbedding(this.embeddingModel.embed(d)); } @@ -175,8 +176,7 @@ public class CassandraVectorStore implements VectorStore, AutoCloseable { builder = builder.setString(this.conf.schema.content(), d.getContent()) .setVector(this.conf.schema.embedding(), - CqlVector.newInstance(d.getEmbedding().stream().map(Double::floatValue).toList()), - Float.class); + CqlVector.newInstance(EmbeddingUtils.toList(d.getEmbedding())), Float.class); for (var metadataColumn : this.conf.schema.metadataColumns() .stream() @@ -241,10 +241,8 @@ public class CassandraVectorStore implements VectorStore, AutoCloseable { Document doc = new Document(getDocumentId(row), row.getString(this.conf.schema.content()), docFields); if (this.conf.returnEmbeddings) { - doc.setEmbedding(row.getVector(this.conf.schema.embedding(), Float.class) - .stream() - .map(Float::doubleValue) - .toList()); + doc.setEmbedding(EmbeddingUtils + .toPrimitive(row.getVector(this.conf.schema.embedding(), Float.class).stream().toList())); } documents.add(doc); } @@ -359,10 +357,10 @@ public class CassandraVectorStore implements VectorStore, AutoCloseable { return this.conf.primaryKeyTranslator.apply(primaryKeyValues); } - private static Float[] toFloatArray(List embeddingDouble) { - Float[] embeddingFloat = new Float[embeddingDouble.size()]; + private static Float[] toFloatArray(float[] embedding) { + Float[] embeddingFloat = new Float[embedding.length]; int i = 0; - for (Double d : embeddingDouble) { + for (Float d : embedding) { embeddingFloat[i++] = d.floatValue(); } return embeddingFloat; diff --git a/vector-stores/spring-ai-chroma-store/src/main/java/org/springframework/ai/chroma/ChromaApi.java b/vector-stores/spring-ai-chroma-store/src/main/java/org/springframework/ai/chroma/ChromaApi.java index 343489ffe..6ce419d91 100644 --- a/vector-stores/spring-ai-chroma-store/src/main/java/org/springframework/ai/chroma/ChromaApi.java +++ b/vector-stores/spring-ai-chroma-store/src/main/java/org/springframework/ai/chroma/ChromaApi.java @@ -182,7 +182,7 @@ public class ChromaApi { * @param documents List of document contents. One for each returned document. * @param metadata List of document metadata. One for each returned document. */ - public record GetEmbeddingResponse(List ids, List> embeddings, List documents, + public record GetEmbeddingResponse(List ids, List embeddings, List documents, @JsonProperty("metadatas") List> metadata) { } @@ -198,7 +198,7 @@ public class ChromaApi { * "metadatas", "documents", "distances". Ids are always included. Defaults to * [metadatas, documents, distances]. */ - public record QueryRequest(@JsonProperty("query_embeddings") List> queryEmbeddings, + public record QueryRequest(@JsonProperty("query_embeddings") List queryEmbeddings, @JsonProperty("n_results") int nResults, Map where, List include) { public enum Include { @@ -222,11 +222,11 @@ public class ChromaApi { /** * Convenience to query for a single embedding instead of a batch of embeddings. */ - public QueryRequest(List queryEmbedding, int nResults) { + public QueryRequest(float[] queryEmbedding, int nResults) { this(List.of(queryEmbedding), nResults, Map.of(), Include.all); } - public QueryRequest(List queryEmbedding, int nResults, Map where) { + public QueryRequest(float[] queryEmbedding, int nResults, Map where) { this(List.of(queryEmbedding), nResults, where, Include.all); } } @@ -241,15 +241,14 @@ public class ChromaApi { * @param metadata List of list of document metadata. One for each returned document. * @param distances List of list of search distances. One for each returned document. */ - public record QueryResponse(List> ids, List>> embeddings, - List> documents, @JsonProperty("metadatas") List>> metadata, - List> distances) { + public record QueryResponse(List> ids, List> embeddings, List> documents, + @JsonProperty("metadatas") List>> metadata, List> distances) { } /** * Single query embedding response. */ - public record Embedding(String id, List embedding, String document, Map metadata, + public record Embedding(String id, float[] embedding, String document, Map metadata, Double distances) { } diff --git a/vector-stores/spring-ai-chroma-store/src/main/java/org/springframework/ai/vectorstore/ChromaVectorStore.java b/vector-stores/spring-ai-chroma-store/src/main/java/org/springframework/ai/vectorstore/ChromaVectorStore.java index 9fe7bacab..22e910899 100644 --- a/vector-stores/spring-ai-chroma-store/src/main/java/org/springframework/ai/vectorstore/ChromaVectorStore.java +++ b/vector-stores/spring-ai-chroma-store/src/main/java/org/springframework/ai/vectorstore/ChromaVectorStore.java @@ -97,7 +97,7 @@ public class ChromaVectorStore implements VectorStore, InitializingBean { metadatas.add(document.getMetadata()); contents.add(document.getContent()); document.setEmbedding(this.embeddingModel.embed(document)); - embeddings.add(JsonUtils.toFloatArray(document.getEmbedding())); + embeddings.add(document.getEmbedding()); } this.chromaApi.upsertEmbeddings(this.collectionId, @@ -121,10 +121,10 @@ public class ChromaVectorStore implements VectorStore, InitializingBean { String query = request.getQuery(); Assert.notNull(query, "Query string must not be null"); - List embedding = this.embeddingModel.embed(query); + float[] embedding = this.embeddingModel.embed(query); Map where = (StringUtils.hasText(nativeFilterExpression)) ? JsonUtils.jsonToMap(nativeFilterExpression) : Map.of(); - var queryRequest = new ChromaApi.QueryRequest(JsonUtils.toFloatList(embedding), request.getTopK(), where); + var queryRequest = new ChromaApi.QueryRequest(embedding, request.getTopK(), where); var queryResponse = this.chromaApi.queryCollection(this.collectionId, queryRequest); var embeddings = this.chromaApi.toEmbeddingResponseList(queryResponse); @@ -141,7 +141,7 @@ public class ChromaVectorStore implements VectorStore, InitializingBean { } metadata.put(DISTANCE_FIELD_NAME, distance); Document document = new Document(id, content, metadata); - document.setEmbedding(JsonUtils.toDouble(chromaEmbedding.embedding())); + document.setEmbedding(chromaEmbedding.embedding()); responseDocuments.add(document); } } diff --git a/vector-stores/spring-ai-chroma-store/src/main/java/org/springframework/ai/vectorstore/JsonUtils.java b/vector-stores/spring-ai-chroma-store/src/main/java/org/springframework/ai/vectorstore/JsonUtils.java index 41f891b49..81860ccd8 100644 --- a/vector-stores/spring-ai-chroma-store/src/main/java/org/springframework/ai/vectorstore/JsonUtils.java +++ b/vector-stores/spring-ai-chroma-store/src/main/java/org/springframework/ai/vectorstore/JsonUtils.java @@ -44,24 +44,24 @@ public class JsonUtils { } } - /** - * Converts a list of doubles to a list of floats. - * @param embeddingDouble the list of doubles to convert - * @return the list of floats - */ - public static List toFloatList(List embeddingDouble) { - return embeddingDouble.stream().map(Number::floatValue).toList(); - } + // /** + // * Converts a list of doubles to a list of floats. + // * @param embeddingDouble the list of doubles to convert + // * @return the list of floats + // */ + // public static List toFloatList(List embeddingDouble) { + // return embeddingDouble.stream().map(Number::floatValue).toList(); + // } /** * Converts a list of doubles to a float array. - * @param embeddingDouble the list of doubles to convert + * @param embedding the list of doubles to convert * @return the float array */ - public static float[] toFloatArray(List embeddingDouble) { - float[] embeddingFloat = new float[embeddingDouble.size()]; + public static float[] toFloatArray(List embedding) { + float[] embeddingFloat = new float[embedding.size()]; int i = 0; - for (Double d : embeddingDouble) { + for (Float d : embedding) { embeddingFloat[i++] = d.floatValue(); } return embeddingFloat; diff --git a/vector-stores/spring-ai-chroma-store/src/test/java/org/springframework/ai/chroma/ChromaApiIT.java b/vector-stores/spring-ai-chroma-store/src/test/java/org/springframework/ai/chroma/ChromaApiIT.java index 638c2a1f2..8a516a1ae 100644 --- a/vector-stores/spring-ai-chroma-store/src/test/java/org/springframework/ai/chroma/ChromaApiIT.java +++ b/vector-stores/spring-ai-chroma-store/src/test/java/org/springframework/ai/chroma/ChromaApiIT.java @@ -92,7 +92,7 @@ public class ChromaApiIT { assertThat(chroma.countEmbeddings(newCollection.id())).isEqualTo(3); var queryResult = chroma.queryCollection(newCollection.id(), - new QueryRequest(List.of(1f, 1f, 1f), 3, chroma.where(""" + new QueryRequest(new float[] { 1f, 1f, 1f }, 3, chroma.where(""" { "key2" : { "$eq": true } } @@ -108,7 +108,7 @@ public class ChromaApiIT { assertThat(result.ids().get(0)).isEqualTo("id2"); queryResult = chroma.queryCollection(newCollection.id(), - new QueryRequest(List.of(1f, 1f, 1f), 3, chroma.where(""" + new QueryRequest(new float[] { 1f, 1f, 1f }, 3, chroma.where(""" { "key2" : { "$eq": true } } @@ -139,7 +139,7 @@ public class ChromaApiIT { assertThat(chroma.countEmbeddings(collection.id())).isEqualTo(3); - var queryResult = chroma.queryCollection(collection.id(), new QueryRequest(List.of(1f, 1f, 1f), 3)); + var queryResult = chroma.queryCollection(collection.id(), new QueryRequest(new float[] { 1f, 1f, 1f }, 3)); assertThat(queryResult.ids().get(0)).hasSize(3); assertThat(queryResult.ids().get(0)).containsExactlyInAnyOrder("id1", "id2", "id3"); @@ -149,26 +149,28 @@ public class ChromaApiIT { assertThat(chromaEmbeddings).hasSize(3); assertThat(chromaEmbeddings).hasSize(3); - queryResult = chroma.queryCollection(collection.id(), new QueryRequest(List.of(1f, 1f, 1f), 3, chroma.where(""" - { - "$and" : [ - {"country" : { "$eq": "BG"}}, - {"year" : { "$gte": 2020}} - ] - } - """))); + queryResult = chroma.queryCollection(collection.id(), + new QueryRequest(new float[] { 1f, 1f, 1f }, 3, chroma.where(""" + { + "$and" : [ + {"country" : { "$eq": "BG"}}, + {"year" : { "$gte": 2020}} + ] + } + """))); assertThat(queryResult.ids().get(0)).hasSize(2); assertThat(queryResult.ids().get(0)).containsExactlyInAnyOrder("id1", "id3"); - queryResult = chroma.queryCollection(collection.id(), new QueryRequest(List.of(1f, 1f, 1f), 3, chroma.where(""" - { - "$and" : [ - {"country" : { "$eq": "BG"}}, - {"year" : { "$gte": 2020}}, - {"active" : { "$eq": true}} - ] - } - """))); + queryResult = chroma.queryCollection(collection.id(), + new QueryRequest(new float[] { 1f, 1f, 1f }, 3, chroma.where(""" + { + "$and" : [ + {"country" : { "$eq": "BG"}}, + {"year" : { "$gte": 2020}}, + {"active" : { "$eq": true}} + ] + } + """))); assertThat(queryResult.ids().get(0)).hasSize(1); assertThat(queryResult.ids().get(0)).containsExactlyInAnyOrder("id1"); } diff --git a/vector-stores/spring-ai-elasticsearch-store/src/main/java/org/springframework/ai/vectorstore/ElasticsearchVectorStore.java b/vector-stores/spring-ai-elasticsearch-store/src/main/java/org/springframework/ai/vectorstore/ElasticsearchVectorStore.java index 2606fe662..80aa78d9a 100644 --- a/vector-stores/spring-ai-elasticsearch-store/src/main/java/org/springframework/ai/vectorstore/ElasticsearchVectorStore.java +++ b/vector-stores/spring-ai-elasticsearch-store/src/main/java/org/springframework/ai/vectorstore/ElasticsearchVectorStore.java @@ -30,6 +30,7 @@ import org.slf4j.Logger; import org.slf4j.LoggerFactory; import org.springframework.ai.document.Document; import org.springframework.ai.embedding.EmbeddingModel; +import org.springframework.ai.model.EmbeddingUtils; import org.springframework.ai.vectorstore.filter.Filter; import org.springframework.ai.vectorstore.filter.FilterExpressionConverter; import org.springframework.beans.factory.InitializingBean; @@ -94,7 +95,7 @@ public class ElasticsearchVectorStore implements VectorStore, InitializingBean { BulkRequest.Builder bulkRequestBuilder = new BulkRequest.Builder(); for (Document document : documents) { - if (Objects.isNull(document.getEmbedding()) || document.getEmbedding().isEmpty()) { + if (Objects.isNull(document.getEmbedding()) || document.getEmbedding().length == 0) { logger.debug("Calling EmbeddingModel for document id = " + document.getId()); document.setEmbedding(this.embeddingModel.embed(document)); } @@ -150,14 +151,11 @@ public class ElasticsearchVectorStore implements VectorStore, InitializingBean { threshold = 1 - threshold; } final float finalThreshold = threshold; - List vectors = this.embeddingModel.embed(searchRequest.getQuery()) - .stream() - .map(Double::floatValue) - .toList(); + float[] vectors = this.embeddingModel.embed(searchRequest.getQuery()); SearchResponse res = elasticsearchClient.search( sr -> sr.index(options.getIndexName()) - .knn(knn -> knn.queryVector(vectors) + .knn(knn -> knn.queryVector(EmbeddingUtils.toList(vectors)) .similarity(finalThreshold) .k((long) searchRequest.getTopK()) .field("embedding") diff --git a/vector-stores/spring-ai-gemfire-store/src/main/java/org/springframework/ai/vectorstore/GemFireVectorStore.java b/vector-stores/spring-ai-gemfire-store/src/main/java/org/springframework/ai/vectorstore/GemFireVectorStore.java index 1452eb647..3396aaa21 100644 --- a/vector-stores/spring-ai-gemfire-store/src/main/java/org/springframework/ai/vectorstore/GemFireVectorStore.java +++ b/vector-stores/spring-ai-gemfire-store/src/main/java/org/springframework/ai/vectorstore/GemFireVectorStore.java @@ -256,12 +256,12 @@ public class GemFireVectorStore implements VectorStore, InitializingBean { private final String key; - private List vector; + private float[] vector; @JsonInclude(JsonInclude.Include.NON_NULL) private Map metadata; - public Embedding(@JsonProperty("key") String key, @JsonProperty("vector") List vector, + public Embedding(@JsonProperty("key") String key, @JsonProperty("vector") float[] vector, String contentName, String content, @JsonProperty("metadata") Map metadata) { this.key = key; this.vector = vector; @@ -273,7 +273,7 @@ public class GemFireVectorStore implements VectorStore, InitializingBean { return key; } - public List getVector() { + public float[] getVector() { return vector; } @@ -289,7 +289,7 @@ public class GemFireVectorStore implements VectorStore, InitializingBean { @JsonProperty("vector") @NonNull - private final List vector; + private final float[] vector; @JsonProperty("top-k") private final int k; @@ -300,14 +300,14 @@ public class GemFireVectorStore implements VectorStore, InitializingBean { @JsonProperty("include-metadata") private final boolean includeMetadata; - public QueryRequest(List vector, int k, int kPerBucket, boolean includeMetadata) { + public QueryRequest(float[] vector, int k, int kPerBucket, boolean includeMetadata) { this.vector = vector; this.k = k; this.kPerBucket = kPerBucket; this.includeMetadata = includeMetadata; } - public List getVector() { + public float[] getVector() { return vector; } @@ -378,7 +378,7 @@ public class GemFireVectorStore implements VectorStore, InitializingBean { UploadRequest upload = new UploadRequest(documents.stream().map(document -> { // Compute and assign an embedding to the document. document.setEmbedding(this.embeddingModel.embed(document)); - List floatVector = document.getEmbedding().stream().map(Double::floatValue).toList(); + float[] floatVector = document.getEmbedding(); return new UploadRequest.Embedding(document.getId(), floatVector, DOCUMENT_FIELD, document.getContent(), document.getMetadata()); }).toList()); @@ -425,8 +425,7 @@ public class GemFireVectorStore implements VectorStore, InitializingBean { if (request.hasFilterExpression()) { throw new UnsupportedOperationException("GemFire currently does not support metadata filter expressions."); } - List vector = this.embeddingModel.embed(request.getQuery()); - List floatVector = vector.stream().map(Double::floatValue).toList(); + float[] floatVector = this.embeddingModel.embed(request.getQuery()); return client.post() .uri("/" + indexName + QUERY) .contentType(MediaType.APPLICATION_JSON) diff --git a/vector-stores/spring-ai-hanadb-store/src/main/java/org/springframework/ai/vectorstore/HanaCloudVectorStore.java b/vector-stores/spring-ai-hanadb-store/src/main/java/org/springframework/ai/vectorstore/HanaCloudVectorStore.java index db8d44e01..ece41a15a 100644 --- a/vector-stores/spring-ai-hanadb-store/src/main/java/org/springframework/ai/vectorstore/HanaCloudVectorStore.java +++ b/vector-stores/spring-ai-hanadb-store/src/main/java/org/springframework/ai/vectorstore/HanaCloudVectorStore.java @@ -20,6 +20,7 @@ import org.slf4j.Logger; import org.slf4j.LoggerFactory; import org.springframework.ai.document.Document; import org.springframework.ai.embedding.EmbeddingModel; +import org.springframework.ai.model.EmbeddingUtils; import java.util.Collections; import java.util.List; @@ -130,15 +131,17 @@ public class HanaCloudVectorStore implements VectorStore { } private String getEmbedding(SearchRequest searchRequest) { - return "[" + this.embeddingModel.embed(searchRequest.getQuery()) + return "[" + EmbeddingUtils.toList(this.embeddingModel.embed(searchRequest.getQuery())) .stream() .map(String::valueOf) .collect(Collectors.joining(", ")) + "]"; } private String getEmbedding(Document document) { - return "[" + this.embeddingModel.embed(document).stream().map(String::valueOf).collect(Collectors.joining(", ")) - + "]"; + return "[" + EmbeddingUtils.toList(this.embeddingModel.embed(document)) + .stream() + .map(String::valueOf) + .collect(Collectors.joining(", ")) + "]"; } } diff --git a/vector-stores/spring-ai-milvus-store/src/main/java/org/springframework/ai/vectorstore/MilvusVectorStore.java b/vector-stores/spring-ai-milvus-store/src/main/java/org/springframework/ai/vectorstore/MilvusVectorStore.java index d871d8e7d..c51f48bf3 100644 --- a/vector-stores/spring-ai-milvus-store/src/main/java/org/springframework/ai/vectorstore/MilvusVectorStore.java +++ b/vector-stores/spring-ai-milvus-store/src/main/java/org/springframework/ai/vectorstore/MilvusVectorStore.java @@ -24,6 +24,7 @@ import org.slf4j.Logger; import org.slf4j.LoggerFactory; import org.springframework.ai.document.Document; import org.springframework.ai.embedding.EmbeddingModel; +import org.springframework.ai.model.EmbeddingUtils; import org.springframework.ai.vectorstore.filter.FilterExpressionConverter; import org.springframework.beans.factory.InitializingBean; import org.springframework.util.Assert; @@ -45,7 +46,6 @@ import io.milvus.param.RpcStatus; import io.milvus.param.collection.CreateCollectionParam; import io.milvus.param.collection.DropCollectionParam; import io.milvus.param.collection.FieldType; -import io.milvus.param.collection.FlushParam; import io.milvus.param.collection.HasCollectionParam; import io.milvus.param.collection.LoadCollectionParam; import io.milvus.param.collection.ReleaseCollectionParam; @@ -271,14 +271,14 @@ public class MilvusVectorStore implements VectorStore, InitializingBean { List> embeddingArray = new ArrayList<>(); for (Document document : documents) { - List embedding = this.embeddingModel.embed(document); + float[] embedding = this.embeddingModel.embed(document); document.setEmbedding(embedding); docIdArray.add(document.getId()); // Use a (future) DocumentTextLayoutFormatter instance to extract // the content used to compute the embeddings contentArray.add(document.getContent()); metadataArray.add(new JSONObject(document.getMetadata())); - embeddingArray.add(toFloatList(embedding)); + embeddingArray.add(EmbeddingUtils.toList(embedding)); } List fields = new ArrayList<>(); @@ -327,7 +327,7 @@ public class MilvusVectorStore implements VectorStore, InitializingBean { Assert.notNull(request.getQuery(), "Query string must not be null"); - List embedding = this.embeddingModel.embed(request.getQuery()); + float[] embedding = this.embeddingModel.embed(request.getQuery()); var searchParamBuilder = SearchParam.newBuilder() .withCollectionName(this.config.collectionName) @@ -335,7 +335,7 @@ public class MilvusVectorStore implements VectorStore, InitializingBean { .withMetricType(this.config.metricType) .withOutFields(SEARCH_OUTPUT_FIELDS) .withTopK(request.getTopK()) - .withVectors(List.of(toFloatList(embedding))) + .withVectors(List.of(EmbeddingUtils.toList(embedding))) .withVectorFieldName(EMBEDDING_FIELD_NAME); if (StringUtils.hasText(nativeFilterExpressions)) { @@ -370,10 +370,6 @@ public class MilvusVectorStore implements VectorStore, InitializingBean { : (1 - distance); } - private List toFloatList(List embeddingDouble) { - return embeddingDouble.stream().map(Number::floatValue).toList(); - } - // --------------------------------------------------------------------------------- // Initialization // --------------------------------------------------------------------------------- diff --git a/vector-stores/spring-ai-mongodb-atlas-store/src/main/java/org/springframework/ai/vectorstore/MongoDBAtlasVectorStore.java b/vector-stores/spring-ai-mongodb-atlas-store/src/main/java/org/springframework/ai/vectorstore/MongoDBAtlasVectorStore.java index e07d88474..2411ea63b 100644 --- a/vector-stores/spring-ai-mongodb-atlas-store/src/main/java/org/springframework/ai/vectorstore/MongoDBAtlasVectorStore.java +++ b/vector-stores/spring-ai-mongodb-atlas-store/src/main/java/org/springframework/ai/vectorstore/MongoDBAtlasVectorStore.java @@ -24,6 +24,7 @@ import java.util.Optional; import com.mongodb.MongoCommandException; import org.springframework.ai.document.Document; import org.springframework.ai.embedding.EmbeddingModel; +import org.springframework.ai.model.EmbeddingUtils; import org.springframework.beans.factory.InitializingBean; import org.springframework.data.mongodb.UncategorizedMongoDbException; import org.springframework.data.mongodb.core.MongoTemplate; @@ -146,10 +147,10 @@ public class MongoDBAtlasVectorStore implements VectorStore, InitializingBean { String id = mongoDocument.getString(ID_FIELD_NAME); String content = mongoDocument.getString(CONTENT_FIELD_NAME); Map metadata = mongoDocument.get(METADATA_FIELD_NAME, org.bson.Document.class); - List embedding = mongoDocument.getList(this.config.pathName, Double.class); + List embedding = mongoDocument.getList(this.config.pathName, Float.class); Document document = new Document(id, content, metadata); - document.setEmbedding(embedding); + document.setEmbedding(EmbeddingUtils.toPrimitive(embedding)); return document; } @@ -157,7 +158,7 @@ public class MongoDBAtlasVectorStore implements VectorStore, InitializingBean { @Override public void add(List documents) { for (Document document : documents) { - List embedding = this.embeddingModel.embed(document); + float[] embedding = this.embeddingModel.embed(document); document.setEmbedding(embedding); this.mongoTemplate.save(document, this.config.collectionName); } @@ -184,9 +185,9 @@ public class MongoDBAtlasVectorStore implements VectorStore, InitializingBean { String nativeFilterExpressions = (request.getFilterExpression() != null) ? this.filterExpressionConverter.convertExpression(request.getFilterExpression()) : ""; - List queryEmbedding = this.embeddingModel.embed(request.getQuery()); - var vectorSearch = new VectorSearchAggregation(queryEmbedding, this.config.pathName, this.config.numCandidates, - this.config.vectorIndexName, request.getTopK(), nativeFilterExpressions); + float[] queryEmbedding = this.embeddingModel.embed(request.getQuery()); + var vectorSearch = new VectorSearchAggregation(EmbeddingUtils.toList(queryEmbedding), this.config.pathName, + this.config.numCandidates, this.config.vectorIndexName, request.getTopK(), nativeFilterExpressions); Aggregation aggregation = Aggregation.newAggregation(vectorSearch, Aggregation.addFields() diff --git a/vector-stores/spring-ai-mongodb-atlas-store/src/main/java/org/springframework/ai/vectorstore/VectorSearchAggregation.java b/vector-stores/spring-ai-mongodb-atlas-store/src/main/java/org/springframework/ai/vectorstore/VectorSearchAggregation.java index 572749af4..6888b9e11 100644 --- a/vector-stores/spring-ai-mongodb-atlas-store/src/main/java/org/springframework/ai/vectorstore/VectorSearchAggregation.java +++ b/vector-stores/spring-ai-mongodb-atlas-store/src/main/java/org/springframework/ai/vectorstore/VectorSearchAggregation.java @@ -15,14 +15,14 @@ */ package org.springframework.ai.vectorstore; +import java.util.List; + import org.bson.Document; import org.springframework.data.mongodb.core.aggregation.AggregationOperation; import org.springframework.data.mongodb.core.aggregation.AggregationOperationContext; import org.springframework.lang.NonNull; -import java.util.List; - -record VectorSearchAggregation(List embeddings, String path, int numCandidates, String index, int count, +record VectorSearchAggregation(List embeddings, String path, int numCandidates, String index, int count, String filter) implements AggregationOperation { @SuppressWarnings("null") diff --git a/vector-stores/spring-ai-mongodb-atlas-store/src/test/java/org/springframework/ai/vectorstore/VectorSearchAggregationTest.java b/vector-stores/spring-ai-mongodb-atlas-store/src/test/java/org/springframework/ai/vectorstore/VectorSearchAggregationTest.java index 4caf4f8ca..d217eb54b 100644 --- a/vector-stores/spring-ai-mongodb-atlas-store/src/test/java/org/springframework/ai/vectorstore/VectorSearchAggregationTest.java +++ b/vector-stores/spring-ai-mongodb-atlas-store/src/test/java/org/springframework/ai/vectorstore/VectorSearchAggregationTest.java @@ -27,13 +27,13 @@ class VectorSearchAggregationTest { @Test void toDocumentNoFilter() { - var vectorSearchAggregation = new VectorSearchAggregation(List.of(1.0, 2.0, 3.0), "embedding", 10, + var vectorSearchAggregation = new VectorSearchAggregation(List.of(1.0f, 2.0f, 3.0f), "embedding", 10, "vector_store", 10, ""); var aggregation = Aggregation.newAggregation(vectorSearchAggregation); var document = aggregation.toDocument("vector_store", Aggregation.DEFAULT_CONTEXT); var vectorSearchDocument = new Document("$vectorSearch", - new Document("queryVector", List.of(1.0, 2.0, 3.0)).append("path", "embedding") + new Document("queryVector", List.of(1.0f, 2.0f, 3.0f)).append("path", "embedding") .append("numCandidates", 10) .append("index", "vector_store") .append("limit", 10)); @@ -44,13 +44,13 @@ class VectorSearchAggregationTest { @Test void toDocumentWithFilter() { - var vectorSearchAggregation = new VectorSearchAggregation(List.of(1.0, 2.0, 3.0), "embedding", 10, + var vectorSearchAggregation = new VectorSearchAggregation(List.of(1.0f, 2.0f, 3.0f), "embedding", 10, "vector_store", 10, "{\"metadata.country\":{$eq:\"BG\"}}"); var aggregation = Aggregation.newAggregation(vectorSearchAggregation); var document = aggregation.toDocument("vector_store", Aggregation.DEFAULT_CONTEXT); var vectorSearchDocument = new Document("$vectorSearch", - new Document("queryVector", List.of(1.0, 2.0, 3.0)).append("path", "embedding") + new Document("queryVector", List.of(1.0f, 2.0f, 3.0f)).append("path", "embedding") .append("numCandidates", 10) .append("index", "vector_store") .append("filter", new Document("metadata.country", new Document().append("$eq", "BG"))) diff --git a/vector-stores/spring-ai-neo4j-store/src/main/java/org/springframework/ai/vectorstore/Neo4jVectorStore.java b/vector-stores/spring-ai-neo4j-store/src/main/java/org/springframework/ai/vectorstore/Neo4jVectorStore.java index c5e546e6f..e3d684a2b 100644 --- a/vector-stores/spring-ai-neo4j-store/src/main/java/org/springframework/ai/vectorstore/Neo4jVectorStore.java +++ b/vector-stores/spring-ai-neo4j-store/src/main/java/org/springframework/ai/vectorstore/Neo4jVectorStore.java @@ -15,6 +15,12 @@ */ package org.springframework.ai.vectorstore; +import java.util.HashMap; +import java.util.List; +import java.util.Map; +import java.util.Optional; +import java.util.function.Predicate; + import org.neo4j.cypherdsl.support.schema_name.SchemaNames; import org.neo4j.driver.Driver; import org.neo4j.driver.SessionConfig; @@ -25,12 +31,6 @@ import org.springframework.ai.vectorstore.filter.Neo4jVectorFilterExpressionConv import org.springframework.beans.factory.InitializingBean; import org.springframework.util.Assert; -import java.util.HashMap; -import java.util.List; -import java.util.Map; -import java.util.Optional; -import java.util.function.Predicate; - /** * @author Gerrit Meier * @author Michael Simons @@ -332,7 +332,7 @@ public class Neo4jVectorStore implements VectorStore, InitializingBean { Assert.isTrue(request.getSimilarityThreshold() >= 0 && request.getSimilarityThreshold() <= 1, "The similarity score is bounded between 0 and 1; least to most similar respectively."); - var embedding = Values.value(toFloatArray(this.embeddingModel.embed(request.getQuery()))); + var embedding = Values.value(this.embeddingModel.embed(request.getQuery())); try (var session = this.driver.session(this.config.sessionConfig)) { StringBuilder condition = new StringBuilder("score >= $threshold"); if (request.hasFilterExpression()) { @@ -393,19 +393,10 @@ public class Neo4jVectorStore implements VectorStore, InitializingBean { document.getMetadata().forEach((k, v) -> properties.put("metadata." + k, Values.value(v))); row.put("properties", properties); - row.put(this.config.embeddingProperty, Values.value(toFloatArray(embedding))); + row.put(this.config.embeddingProperty, Values.value(embedding)); return row; } - private static float[] toFloatArray(List embeddingDouble) { - float[] embeddingFloat = new float[embeddingDouble.size()]; - int i = 0; - for (Double d : embeddingDouble) { - embeddingFloat[i++] = d.floatValue(); - } - return embeddingFloat; - } - private Document recordToDocument(org.neo4j.driver.Record neoRecord) { var node = neoRecord.get("node").asNode(); var score = neoRecord.get("score").asFloat(); diff --git a/vector-stores/spring-ai-opensearch-store/src/main/java/org/springframework/ai/vectorstore/OpenSearchVectorStore.java b/vector-stores/spring-ai-opensearch-store/src/main/java/org/springframework/ai/vectorstore/OpenSearchVectorStore.java index 655c48fbd..9522ccf1c 100644 --- a/vector-stores/spring-ai-opensearch-store/src/main/java/org/springframework/ai/vectorstore/OpenSearchVectorStore.java +++ b/vector-stores/spring-ai-opensearch-store/src/main/java/org/springframework/ai/vectorstore/OpenSearchVectorStore.java @@ -116,7 +116,7 @@ public class OpenSearchVectorStore implements VectorStore, InitializingBean { public void add(List documents) { BulkRequest.Builder bulkRequestBuilder = new BulkRequest.Builder(); for (Document document : documents) { - if (Objects.isNull(document.getEmbedding()) || document.getEmbedding().isEmpty()) { + if (Objects.isNull(document.getEmbedding()) || document.getEmbedding().length == 0) { logger.debug("Calling EmbeddingModel for document id = " + document.getId()); document.setEmbedding(this.embeddingModel.embed(document)); } @@ -150,7 +150,7 @@ public class OpenSearchVectorStore implements VectorStore, InitializingBean { searchRequest.getSimilarityThreshold(), searchRequest.getFilterExpression()); } - public List similaritySearch(List embedding, int topK, double similarityThreshold, + public List similaritySearch(float[] embedding, int topK, double similarityThreshold, Filter.Expression filterExpression) { return similaritySearch(new org.opensearch.client.opensearch.core.SearchRequest.Builder() .query(getOpenSearchSimilarityQuery(embedding, filterExpression)) @@ -161,7 +161,7 @@ public class OpenSearchVectorStore implements VectorStore, InitializingBean { .build()); } - private Query getOpenSearchSimilarityQuery(List embedding, Filter.Expression filterExpression) { + private Query getOpenSearchSimilarityQuery(float[] embedding, Filter.Expression filterExpression) { return Query.of(queryBuilder -> queryBuilder.scriptScore(scriptScoreQueryBuilder -> { scriptScoreQueryBuilder .query(queryBuilder2 -> queryBuilder2.queryString(queryStringQuerybuilder -> queryStringQuerybuilder diff --git a/vector-stores/spring-ai-oracle-store/src/main/java/org/springframework/ai/vectorstore/OracleVectorStore.java b/vector-stores/spring-ai-oracle-store/src/main/java/org/springframework/ai/vectorstore/OracleVectorStore.java index 0c3610c2f..d3b8bfd46 100644 --- a/vector-stores/spring-ai-oracle-store/src/main/java/org/springframework/ai/vectorstore/OracleVectorStore.java +++ b/vector-stores/spring-ai-oracle-store/src/main/java/org/springframework/ai/vectorstore/OracleVectorStore.java @@ -322,14 +322,14 @@ public class OracleVectorStore implements VectorStore, InitializingBean { /** * Converts a list of Double values into an Oracle VECTOR object ready to be inserted. * Optionally normalize the vector beforehand (see forcedNormalization). - * @param doubleList + * @param floatList * @return * @throws SQLException */ - private VECTOR toVECTOR(final List doubleList) throws SQLException { - final double[] doubles = new double[doubleList.size()]; + private VECTOR toVECTOR(final float[] floatList) throws SQLException { + final double[] doubles = new double[floatList.length]; int i = 0; - for (double d : doubleList) { + for (double d : floatList) { doubles[i++] = d; } @@ -400,8 +400,8 @@ public class OracleVectorStore implements VectorStore, InitializingBean { metadata.put("distance", rs.getDouble(5)); final Document document = new Document(rs.getString(1), rs.getString(2), metadata); - final double[] embedding = rs.getObject(4, double[].class); - document.setEmbedding(toDoubleList(embedding)); + final float[] embedding = rs.getObject(4, float[].class); + document.setEmbedding(embedding); return document; } @@ -418,9 +418,9 @@ public class OracleVectorStore implements VectorStore, InitializingBean { return result; } - private List toDoubleList(final double[] embeddings) { - final List result = new ArrayList<>(embeddings.length); - for (double v : embeddings) { + private List toFloatList(final float[] embeddings) { + final List result = new ArrayList<>(embeddings.length); + for (float v : embeddings) { result.add(v); } return result; diff --git a/vector-stores/spring-ai-pgvector-store/src/main/java/org/springframework/ai/vectorstore/PgVectorStore.java b/vector-stores/spring-ai-pgvector-store/src/main/java/org/springframework/ai/vectorstore/PgVectorStore.java index bf2c662f7..aa7966e76 100644 --- a/vector-stores/spring-ai-pgvector-store/src/main/java/org/springframework/ai/vectorstore/PgVectorStore.java +++ b/vector-stores/spring-ai-pgvector-store/src/main/java/org/springframework/ai/vectorstore/PgVectorStore.java @@ -22,7 +22,6 @@ import java.util.List; import java.util.Map; import java.util.Optional; import java.util.UUID; -import java.util.stream.IntStream; import org.postgresql.util.PGobject; import org.slf4j.Logger; @@ -168,7 +167,7 @@ public class PgVectorStore implements VectorStore, InitializingBean { var json = toJson(document.getMetadata()); var embedding = embeddingModel.embed(document); document.setEmbedding(embedding); - var pGvector = new PGvector(toFloatArray(embedding)); + var pGvector = new PGvector(embedding); StatementCreatorUtils.setParameterValue(ps, 1, SqlTypeValue.TYPE_UNKNOWN, UUID.fromString(document.getId())); @@ -196,10 +195,10 @@ public class PgVectorStore implements VectorStore, InitializingBean { } } - private float[] toFloatArray(List embeddingDouble) { - float[] embeddingFloat = new float[embeddingDouble.size()]; + private float[] toFloatArray(List embedding) { + float[] embeddingFloat = new float[embedding.size()]; int i = 0; - for (Double d : embeddingDouble) { + for (Float d : embedding) { embeddingFloat[i++] = d.floatValue(); } return embeddingFloat; @@ -253,8 +252,8 @@ public class PgVectorStore implements VectorStore, InitializingBean { } private PGvector getQueryEmbedding(String query) { - List embedding = this.embeddingModel.embed(query); - return new PGvector(toFloatArray(embedding)); + float[] embedding = this.embeddingModel.embed(query); + return new PGvector(embedding); } private String comparisonOperator() { @@ -441,14 +440,13 @@ public class PgVectorStore implements VectorStore, InitializingBean { metadata.put(COLUMN_DISTANCE, distance); Document document = new Document(id, content, metadata); - document.setEmbedding(toDoubleList(embedding)); + document.setEmbedding(toFloatArray(embedding)); return document; } - private List toDoubleList(PGobject embedding) throws SQLException { - float[] floatArray = new PGvector(embedding.getValue()).toArray(); - return IntStream.range(0, floatArray.length).mapToDouble(i -> floatArray[i]).boxed().toList(); + private float[] toFloatArray(PGobject embedding) throws SQLException { + return new PGvector(embedding.getValue()).toArray(); } private Map toMap(PGobject pgObject) { diff --git a/vector-stores/spring-ai-pinecone-store/src/main/java/org/springframework/ai/vectorstore/PineconeVectorStore.java b/vector-stores/spring-ai-pinecone-store/src/main/java/org/springframework/ai/vectorstore/PineconeVectorStore.java index b5b5def73..8515b7fb1 100644 --- a/vector-stores/spring-ai-pinecone-store/src/main/java/org/springframework/ai/vectorstore/PineconeVectorStore.java +++ b/vector-stores/spring-ai-pinecone-store/src/main/java/org/springframework/ai/vectorstore/PineconeVectorStore.java @@ -38,6 +38,7 @@ import io.pinecone.proto.Vector; import org.springframework.ai.document.Document; import org.springframework.ai.embedding.EmbeddingModel; +import org.springframework.ai.model.EmbeddingUtils; import org.springframework.ai.vectorstore.filter.FilterExpressionConverter; import org.springframework.ai.vectorstore.filter.converter.PineconeFilterExpressionConverter; import org.springframework.util.Assert; @@ -275,7 +276,7 @@ public class PineconeVectorStore implements VectorStore { return Vector.newBuilder() .setId(document.getId()) - .addAllValues(toFloatList(document.getEmbedding())) + .addAllValues(EmbeddingUtils.toList(document.getEmbedding())) .setMetadata(metadataToStruct(document)) .build(); }).toList(); @@ -360,10 +361,10 @@ public class PineconeVectorStore implements VectorStore { String nativeExpressionFilters = (request.getFilterExpression() != null) ? this.filterExpressionConverter.convertExpression(request.getFilterExpression()) : ""; - List queryEmbedding = this.embeddingModel.embed(request.getQuery()); + float[] queryEmbedding = this.embeddingModel.embed(request.getQuery()); var queryRequestBuilder = QueryRequest.newBuilder() - .addAllVector(toFloatList(queryEmbedding)) + .addAllVector(EmbeddingUtils.toList(queryEmbedding)) .setTopK(request.getTopK()) .setIncludeMetadata(true) .setNamespace(namespace); @@ -423,13 +424,4 @@ public class PineconeVectorStore implements VectorStore { } } - /** - * Converts a list of doubles to a list of floats. - * @param doubleList The list of doubles. - * @return The converted list of floats. - */ - private List toFloatList(List doubleList) { - return doubleList.stream().map(d -> d.floatValue()).toList(); - } - } diff --git a/vector-stores/spring-ai-qdrant-store/src/main/java/org/springframework/ai/vectorstore/qdrant/QdrantVectorStore.java b/vector-stores/spring-ai-qdrant-store/src/main/java/org/springframework/ai/vectorstore/qdrant/QdrantVectorStore.java index 73f0b29b1..184af9dd7 100644 --- a/vector-stores/spring-ai-qdrant-store/src/main/java/org/springframework/ai/vectorstore/qdrant/QdrantVectorStore.java +++ b/vector-stores/spring-ai-qdrant-store/src/main/java/org/springframework/ai/vectorstore/qdrant/QdrantVectorStore.java @@ -15,12 +15,25 @@ */ package org.springframework.ai.vectorstore.qdrant; +import static io.qdrant.client.PointIdFactory.id; +import static io.qdrant.client.ValueFactory.value; +import static io.qdrant.client.VectorsFactory.vectors; +import static io.qdrant.client.WithPayloadSelectorFactory.enable; + import java.util.List; import java.util.Map; import java.util.Optional; import java.util.UUID; import java.util.concurrent.ExecutionException; +import org.springframework.ai.document.Document; +import org.springframework.ai.embedding.EmbeddingModel; +import org.springframework.ai.model.EmbeddingUtils; +import org.springframework.ai.vectorstore.SearchRequest; +import org.springframework.ai.vectorstore.VectorStore; +import org.springframework.beans.factory.InitializingBean; +import org.springframework.util.Assert; + import io.qdrant.client.QdrantClient; import io.qdrant.client.grpc.Collections.Distance; import io.qdrant.client.grpc.Collections.VectorParams; @@ -31,17 +44,6 @@ import io.qdrant.client.grpc.Points.PointStruct; import io.qdrant.client.grpc.Points.ScoredPoint; import io.qdrant.client.grpc.Points.SearchPoints; import io.qdrant.client.grpc.Points.UpdateStatus; -import org.springframework.ai.document.Document; -import org.springframework.ai.embedding.EmbeddingModel; -import org.springframework.ai.vectorstore.SearchRequest; -import org.springframework.ai.vectorstore.VectorStore; -import org.springframework.beans.factory.InitializingBean; -import org.springframework.util.Assert; - -import static io.qdrant.client.PointIdFactory.id; -import static io.qdrant.client.ValueFactory.value; -import static io.qdrant.client.VectorsFactory.vectors; -import static io.qdrant.client.WithPayloadSelectorFactory.enable; /** * Qdrant vectorStore implementation. This store supports creating, updating, deleting, @@ -176,7 +178,7 @@ public class QdrantVectorStore implements VectorStore, InitializingBean { return PointStruct.newBuilder() .setId(id(UUID.fromString(document.getId()))) - .setVectors(vectors(toFloatList(document.getEmbedding()))) + .setVectors(vectors(document.getEmbedding())) .putAllPayload(toPayload(document)) .build(); }).toList(); @@ -220,13 +222,13 @@ public class QdrantVectorStore implements VectorStore, InitializingBean { ? this.filterExpressionConverter.convertExpression(request.getFilterExpression()) : Filter.getDefaultInstance(); - List queryEmbedding = this.embeddingModel.embed(request.getQuery()); + float[] queryEmbedding = this.embeddingModel.embed(request.getQuery()); var searchPoints = SearchPoints.newBuilder() .setCollectionName(this.collectionName) .setLimit(request.getTopK()) .setWithPayload(enable(true)) - .addAllVector(toFloatList(queryEmbedding)) + .addAllVector(EmbeddingUtils.toList(queryEmbedding)) .setFilter(filter) .setScoreThreshold((float) request.getSimilarityThreshold()) .build(); @@ -280,15 +282,6 @@ public class QdrantVectorStore implements VectorStore, InitializingBean { } } - /** - * Converts a list of doubles to a list of floats. - * @param doubleList The list of doubles. - * @return The converted list of floats. - */ - private List toFloatList(List doubleList) { - return doubleList.stream().map(d -> d.floatValue()).toList(); - } - @Override public void afterPropertiesSet() throws Exception { diff --git a/vector-stores/spring-ai-redis-store/src/main/java/org/springframework/ai/vectorstore/RedisVectorStore.java b/vector-stores/spring-ai-redis-store/src/main/java/org/springframework/ai/vectorstore/RedisVectorStore.java index 990bee6f5..c4498a82b 100644 --- a/vector-stores/spring-ai-redis-store/src/main/java/org/springframework/ai/vectorstore/RedisVectorStore.java +++ b/vector-stores/spring-ai-redis-store/src/main/java/org/springframework/ai/vectorstore/RedisVectorStore.java @@ -352,7 +352,7 @@ public class RedisVectorStore implements VectorStore, InitializingBean { returnFields.add(this.config.embeddingFieldName); returnFields.add(this.config.contentFieldName); returnFields.add(DISTANCE_FIELD_NAME); - var embedding = toFloatArray(this.embeddingModel.embed(request.getQuery())); + var embedding = this.embeddingModel.embed(request.getQuery()); Query query = new Query(queryString).addParam(EMBEDDING_PARAM_NAME, RediSearchUtil.toByteArray(embedding)) .returnFields(returnFields.toArray(new String[0])) .setSortBy(DISTANCE_FIELD_NAME, true) @@ -457,10 +457,10 @@ public class RedisVectorStore implements VectorStore, InitializingBean { return JSON_PATH_PREFIX + field; } - private static float[] toFloatArray(List embeddingDouble) { - float[] embeddingFloat = new float[embeddingDouble.size()]; + private static float[] toFloatArray(List embedding) { + float[] embeddingFloat = new float[embedding.size()]; int i = 0; - for (Double d : embeddingDouble) { + for (Float d : embedding) { embeddingFloat[i++] = d.floatValue(); } return embeddingFloat; diff --git a/vector-stores/spring-ai-typesense-store/src/main/java/org/springframework/ai/vectorstore/TypesenseVectorStore.java b/vector-stores/spring-ai-typesense-store/src/main/java/org/springframework/ai/vectorstore/TypesenseVectorStore.java index 2e032337f..885cc7f5a 100644 --- a/vector-stores/spring-ai-typesense-store/src/main/java/org/springframework/ai/vectorstore/TypesenseVectorStore.java +++ b/vector-stores/spring-ai-typesense-store/src/main/java/org/springframework/ai/vectorstore/TypesenseVectorStore.java @@ -16,10 +16,13 @@ package org.springframework.ai.vectorstore; +import java.util.Arrays; import java.util.HashMap; import java.util.List; import java.util.Map; import java.util.Optional; +import java.util.stream.IntStream; +import java.util.stream.Stream; import org.slf4j.Logger; import org.slf4j.LoggerFactory; @@ -169,7 +172,7 @@ public class TypesenseVectorStore implements VectorStore, InitializingBean { typesenseDoc.put(DOC_ID_FIELD_NAME, document.getId()); typesenseDoc.put(CONTENT_FIELD_NAME, document.getContent()); typesenseDoc.put(METADATA_FIELD_NAME, document.getMetadata()); - List embedding = this.embeddingModel.embed(document.getContent()); + float[] embedding = this.embeddingModel.embed(document.getContent()); typesenseDoc.put(EMBEDDING_FIELD_NAME, embedding); return typesenseDoc; @@ -222,16 +225,17 @@ public class TypesenseVectorStore implements VectorStore, InitializingBean { logger.info("Filter expression: {}", nativeFilterExpressions); - List embedding = this.embeddingModel.embed(request.getQuery()); + float[] embedding = this.embeddingModel.embed(request.getQuery()); MultiSearchCollectionParameters multiSearchCollectionParameters = new MultiSearchCollectionParameters(); multiSearchCollectionParameters.collection(this.config.collectionName); multiSearchCollectionParameters.q("*"); - // typesnese uses only cosine similarity + Stream floatStream = IntStream.range(0, embedding.length).mapToObj(i -> embedding[i]); + // typesense uses only cosine similarity String vectorQuery = EMBEDDING_FIELD_NAME + ":(" + "[" - + String.join(",", embedding.stream().map(String::valueOf).toList()) + "], " + "k: " + request.getTopK() - + ", " + "distance_threshold: " + (1 - request.getSimilarityThreshold()) + ")"; + + String.join(",", floatStream.map(String::valueOf).toList()) + "], " + "k: " + request.getTopK() + ", " + + "distance_threshold: " + (1 - request.getSimilarityThreshold()) + ")"; multiSearchCollectionParameters.vectorQuery(vectorQuery); multiSearchCollectionParameters.filterBy(nativeFilterExpressions); diff --git a/vector-stores/spring-ai-weaviate-store/src/main/java/org/springframework/ai/vectorstore/WeaviateVectorStore.java b/vector-stores/spring-ai-weaviate-store/src/main/java/org/springframework/ai/vectorstore/WeaviateVectorStore.java index 8f8fc76d8..ad1b79136 100644 --- a/vector-stores/spring-ai-weaviate-store/src/main/java/org/springframework/ai/vectorstore/WeaviateVectorStore.java +++ b/vector-stores/spring-ai-weaviate-store/src/main/java/org/springframework/ai/vectorstore/WeaviateVectorStore.java @@ -44,6 +44,7 @@ import io.weaviate.client.v1.graphql.query.fields.Fields; import org.springframework.ai.document.Document; import org.springframework.ai.embedding.EmbeddingModel; +import org.springframework.ai.model.EmbeddingUtils; import org.springframework.ai.vectorstore.WeaviateVectorStore.WeaviateVectorStoreConfig.ConsistentLevel; import org.springframework.ai.vectorstore.WeaviateVectorStore.WeaviateVectorStoreConfig.MetadataField; import org.springframework.beans.factory.InitializingBean; @@ -361,8 +362,8 @@ public class WeaviateVectorStore implements VectorStore { private WeaviateObject toWeaviateObject(Document document) { - if (CollectionUtils.isEmpty(document.getEmbedding())) { - List embedding = this.embeddingModel.embed(document); + if (document.getEmbedding() == null || document.getEmbedding().length == 0) { + float[] embedding = this.embeddingModel.embed(document); document.setEmbedding(embedding); } @@ -388,7 +389,7 @@ public class WeaviateVectorStore implements VectorStore { return WeaviateObject.builder() .className(this.weaviateObjectClass) .id(document.getId()) - .vector(toFloatArray(document.getEmbedding())) + .vector(EmbeddingUtils.toFloatArray(document.getEmbedding())) .properties(fields) .build(); } @@ -422,13 +423,13 @@ public class WeaviateVectorStore implements VectorStore { @Override public List similaritySearch(SearchRequest request) { - Float[] embedding = toFloatArray(this.embeddingModel.embed(request.getQuery())); + float[] embedding = this.embeddingModel.embed(request.getQuery()); GetBuilder.GetBuilderBuilder builder = GetBuilder.builder(); GetBuilderBuilder queryBuilder = builder.className(this.weaviateObjectClass) .withNearVectorFilter(NearVectorArgument.builder() - .vector(embedding) + .vector(EmbeddingUtils.toFloatArray(embedding)) .certainty((float) request.getSimilarityThreshold()) .build()) .limit(request.getTopK()) @@ -512,18 +513,9 @@ public class WeaviateVectorStore implements VectorStore { String content = (String) item.get(CONTENT_FIELD_NAME); var document = new Document(id, content, metadata); - document.setEmbedding(embedding); + document.setEmbedding(EmbeddingUtils.toPrimitive(EmbeddingUtils.doubleToFloat(embedding))); return document; } - /** - * Converts a list of doubles to an array of floats. - * @param doubleList The list of doubles. - * @return The converted array of floats. - */ - private Float[] toFloatArray(List doubleList) { - return doubleList.stream().map(Number::floatValue).toList().toArray(new Float[0]); - } - } \ No newline at end of file