Spring/spring-ai
1
0
Fork 0
Code Issues Packages Projects Releases Wiki Activity

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

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

Related to #405
This commit is contained in:
Christian Tzolov
2024-07-03 14:42:27 +02:00
committed by Mark Pollack
parent 656fa8b4fe
commit d538e00643
67 changed files with 442 additions and 412 deletions
Download Patch File Download Diff File Expand all files Collapse all files

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

@@ -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<Double> embed(Document document) {
public float[] embed(Document document) {
logger.debug("Retrieving embeddings");
EmbeddingResponse response = this
.call(new EmbeddingRequest(List.of(document.getFormattedContent(this.metadataMode)), null));
logger.debug("Embeddings retrieved");
return response.getResults().stream().map(embedding -> embedding.getOutput()).flatMap(List::stream).toList();
if (CollectionUtils.isEmpty(response.getResults())) {
return new float[0];
}
return response.getResults().get(0).getOutput();
}
@Override
@@ -108,8 +122,7 @@ public class AzureOpenAiEmbeddingModel extends AbstractEmbeddingModel {
for (EmbeddingItem nativeDatum : nativeData) {
List<Float> nativeDatumEmbedding = nativeDatum.getEmbedding();
int nativeIndex = nativeDatum.getPromptIndex();
Embedding embedding = new Embedding(nativeDatumEmbedding.stream().map(f -> f.doubleValue()).toList(),
nativeIndex);
Embedding embedding = new Embedding(EmbeddingUtils.toPrimitive(nativeDatumEmbedding), nativeIndex);
data.add(embedding);
}
return data;

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

@@ -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<Double> embed(Document document) {
public float[] embed(Document document) {
return embed(document.getContent());
}

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

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

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

@@ -75,7 +75,7 @@ public class BedrockTitanEmbeddingModel extends AbstractEmbeddingModel {
}
@Override
public List<Double> 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<List<Double>> embeddingList = new ArrayList<>();
List<Embedding> embeddings = new ArrayList<>();
var indexCounter = new AtomicInteger(0);
for (String inputContent : request.getInstructions()) {
var apiRequest = createTitanEmbeddingRequest(inputContent, request.getOptions());
TitanEmbeddingResponse response = this.embeddingApi.embedding(apiRequest);
embeddingList.add(response.embedding());
embeddings.add(new Embedding(response.embedding(), indexCounter.getAndIncrement()));
}
var indexCounter = new AtomicInteger(0);
List<Embedding> embeddings = embeddingList.stream()
.map(e -> new Embedding(e, indexCounter.getAndIncrement()))
.toList();
return new EmbeddingResponse(embeddings);
}

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

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

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

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

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

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

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

@@ -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);
}

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

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

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

@@ -196,7 +196,7 @@ public class MistralAiApi {
public record Embedding(
// @formatter:off
@JsonProperty("index") Integer index,
@JsonProperty("embedding") List<Double> embedding,
@JsonProperty("embedding") float[] embedding,
@JsonProperty("object") String object) {
// @formatter:on
@@ -207,7 +207,7 @@ public 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<Double> embedding) {
public Embedding(Integer index, float[] embedding) {
this(index, embedding, "embedding");
}
}

16
models/spring-ai-mistral-ai/src/test/java/org/springframework/ai/mistralai/MistralAiRetryTests.java
View File

@@ -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<Embedding> 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);
}

2
models/spring-ai-ollama/src/main/java/org/springframework/ai/ollama/OllamaEmbeddingModel.java
View File

@@ -76,7 +76,7 @@ public class OllamaEmbeddingModel extends AbstractEmbeddingModel {
}
@Override
public List<Double> embed(Document document) {
public float[] embed(Document document) {
return embed(document.getContent());
}

4
models/spring-ai-ollama/src/main/java/org/springframework/ai/ollama/api/OllamaApi.java
View File

@@ -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<Double> embedding) {
@JsonProperty("embedding") List<Float> embedding) {
}
@@ -764,7 +764,7 @@ public class OllamaApi {
@JsonInclude(Include.NON_NULL)
public record EmbeddingsResponse(
@JsonProperty("model") String model,
@JsonProperty("embeddings") List<List<Double>> embeddings) {
@JsonProperty("embeddings") List<float[]> embeddings) {
}
/**

12
models/spring-ai-ollama/src/test/java/org/springframework/ai/ollama/OllamaEmbeddingModelTests.java
View File

@@ -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");

2
models/spring-ai-openai/src/main/java/org/springframework/ai/openai/OpenAiEmbeddingModel.java
View File

@@ -137,7 +137,7 @@ public class OpenAiEmbeddingModel extends AbstractEmbeddingModel {
}
@Override
public List<Double> embed(Document document) {
public float[] embed(Document document) {
Assert.notNull(document, "Document must not be null");
return this.embed(document.getFormattedContent(this.metadataMode));
}

4
models/spring-ai-openai/src/main/java/org/springframework/ai/openai/api/OpenAiApi.java
View File

@@ -1102,7 +1102,7 @@ public class OpenAiApi {
@JsonInclude(Include.NON_NULL)
public record Embedding(// @formatter:off
@JsonProperty("index") Integer index,
@JsonProperty("embedding") List<Double> 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<Double> embedding) {
public Embedding(Integer index, float[] embedding) {
this(index, embedding, "embedding");
}
}

4
models/spring-ai-openai/src/test/java/org/springframework/ai/openai/chat/OpenAiRetryTests.java
View File

@@ -197,7 +197,7 @@ public class OpenAiRetryTests {
public void openAiEmbeddingTransientError() {
EmbeddingList<Embedding> 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);
}

20
models/spring-ai-postgresml/src/main/java/org/springframework/ai/postgresml/PostgresMlEmbeddingModel.java
View File

@@ -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<List<Double>> rowMapper;
private final RowMapper<float[]> rowMapper;
VectorType(String cast, String extensionName, RowMapper<List<Double>> rowMapper) {
VectorType(String cast, String extensionName, RowMapper<float[]> 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<Double> 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<Double> 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<Embedding> data = new ArrayList<>();
List<List<Double>> embed = List.of();
List<float[]> embed = List.of();
List<String> 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<List<Double>> result = new ArrayList<>();
List<float[]> result = new ArrayList<>();
while (rs.next()) {
result.add(optionsToUse.getVectorType().rowMapper.mapRow(rs, -1));
}

8
models/spring-ai-postgresml/src/test/java/org/springframework/ai/postgresml/PostgresMlEmbeddingModelIT.java
View File

@@ -84,7 +84,7 @@ class PostgresMlEmbeddingModelIT {
PostgresMlEmbeddingModel embeddingModel = new PostgresMlEmbeddingModel(this.jdbcTemplate);
embeddingModel.afterPropertiesSet();
List<Double> 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<Double> 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<Double> 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<Double> embed = embeddingModel.embed(new Document("Hello World!"));
float[] embed = embeddingModel.embed(new Document("Hello World!"));
assertThat(embed).hasSize(768);
}

2
models/spring-ai-qianfan/src/main/java/org/springframework/ai/qianfan/QianFanEmbeddingModel.java
View File

@@ -105,7 +105,7 @@ public class QianFanEmbeddingModel extends AbstractEmbeddingModel {
}
@Override
public List<Double> embed(Document document) {
public float[] embed(Document document) {
Assert.notNull(document, "Document must not be null");
return this.embed(document.getFormattedContent(this.metadataMode));
}

4
models/spring-ai-qianfan/src/main/java/org/springframework/ai/qianfan/api/QianFanApi.java
View File

@@ -484,7 +484,7 @@ public class QianFanApi extends AuthApi {
public record Embedding(
// @formatter:off
@JsonProperty("index") Integer index,
@JsonProperty("embedding") List<Double> 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<Double> embedding) {
public Embedding(Integer index, float[] embedding) {
this(index, embedding, "embedding");
}
}

43
models/spring-ai-qianfan/src/test/java/org/springframework/ai/qianfan/api/QianFanRetryTests.java
View File

@@ -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")))));
}
}

22
models/spring-ai-transformers/src/main/java/org/springframework/ai/transformers/TransformersEmbeddingModel.java
View File

@@ -210,19 +210,19 @@ public class TransformersEmbeddingModel extends AbstractEmbeddingModel implement
}
@Override
public List<Double> embed(String text) {
public float[] embed(String text) {
return embed(List.of(text)).get(0);
}
@Override
public List<Double> embed(Document document) {
public float[] embed(Document document) {
return this.embed(document.getFormattedContent(this.metadataMode));
}
@Override
public EmbeddingResponse embedForResponse(List<String> texts) {
List<Embedding> data = new ArrayList<>();
List<List<Double>> embed = this.embed(texts);
List<float[]> 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<List<Double>> embed(List<String> texts) {
public List<float[]> embed(List<String> 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<List<Double>> resultEmbeddings = new ArrayList<>();
List<float[]> 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<Double> toDoubleList(float[] floats) {
List<Double> 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);
}

30
models/spring-ai-transformers/src/test/java/org/springframework/ai/transformers/TransformersEmbeddingModelTests.java
View File

@@ -38,35 +38,35 @@ public class TransformersEmbeddingModelTests {
TransformersEmbeddingModel embeddingModel = new TransformersEmbeddingModel();
embeddingModel.afterPropertiesSet();
List<Double> 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<Double> 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<List<Double>> embed = embeddingModel.embed(List.of("Hello world", "World is big"));
List<float[]> 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

15
models/spring-ai-vertex-ai-embedding/src/main/java/org/springframework/ai/vertexai/embedding/VertexAiEmbeddingUtils.java
View File

@@ -427,13 +427,14 @@ public abstract class VertexAiEmbeddingUtils {
return builder.build();
}
public static List<Double> 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;
}
}

6
models/spring-ai-vertex-ai-embedding/src/main/java/org/springframework/ai/vertexai/embedding/multimodal/VertexAiMultimodalEmbeddingModel.java
View File

@@ -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<Double> 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<Double> 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<Double> videoVector = VertexAiEmbeddingUtils.toVector(embeddings);
float[] videoVector = VertexAiEmbeddingUtils.toVector(embeddings);
var docMetadata = documentMetadata.get(ModalityType.VIDEO);
embeddingList

4
models/spring-ai-vertex-ai-embedding/src/main/java/org/springframework/ai/vertexai/embedding/text/VertexAiTextEmbeddingModel.java
View File

@@ -66,7 +66,7 @@ public class VertexAiTextEmbeddingModel extends AbstractEmbeddingModel {
}
@Override
public List<Double> embed(Document document) {
public float[] embed(Document document) {
Assert.notNull(document, "Document must not be null");
return this.embed(document.getFormattedContent());
}
@@ -125,7 +125,7 @@ public class VertexAiTextEmbeddingModel extends AbstractEmbeddingModel {
Value values = embeddings.getStructValue().getFieldsOrThrow("values");
List<Double> vectorValues = VertexAiEmbeddingUtils.toVector(values);
float[] vectorValues = VertexAiEmbeddingUtils.toVector(values);
embeddingList.add(new Embedding(vectorValues, index++));
}

2
models/spring-ai-vertex-ai-embedding/src/test/java/org/springframework/ai/vertexai/embedding/multimodal/VertexAiMultimodalEmbeddingModelIT.java
View File

@@ -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);
}

2
models/spring-ai-vertex-ai-palm2/src/main/java/org/springframework/ai/vertexai/palm2/VertexAiPaLm2EmbeddingModel.java
View File

@@ -37,7 +37,7 @@ public class VertexAiPaLm2EmbeddingModel extends AbstractEmbeddingModel {
}
@Override
public List<Double> embed(Document document) {
public float[] embed(Document document) {
return embed(document.getContent());
}

12
models/spring-ai-vertex-ai-palm2/src/main/java/org/springframework/ai/vertexai/palm2/api/VertexAiPaLm2Api.java
View File

@@ -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<Double> 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);
}
}
/**

6
models/spring-ai-vertex-ai-palm2/src/test/java/org/springframework/ai/vertexai/palm2/api/VertexAiPaLm2ApiTests.java
View File

@@ -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<String> texts = List.of("Hello, how are you?", "I'm fine, thank you.");
List<Embedding> expectedEmbeddings = List.of(new Embedding(List.of(0.1, 0.2, 0.3)),
new Embedding(List.of(0.4, 0.5, 0.6)));
List<Embedding> 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}",

6
models/spring-ai-zhipuai/src/main/java/org/springframework/ai/zhipuai/ZhiPuAiEmbeddingModel.java
View File

@@ -103,7 +103,7 @@ public class ZhiPuAiEmbeddingModel extends AbstractEmbeddingModel {
}
@Override
public List<Double> embed(Document document) {
public float[] embed(Document document) {
Assert.notNull(document, "Document must not be null");
return this.embed(document.getFormattedContent(this.metadataMode));
}
@@ -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<List<Double>> embeddingList = new ArrayList<>();
List<float[]> embeddingList = new ArrayList<>();
for (String inputContent : request.getInstructions()) {
var apiRequest = createZhiPuEmbeddingRequest(inputContent, request.getOptions());
ZhiPuAiApi.EmbeddingList<ZhiPuAiApi.Embedding> 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());

4
models/spring-ai-zhipuai/src/main/java/org/springframework/ai/zhipuai/api/ZhiPuAiApi.java
View File

@@ -738,7 +738,7 @@ public class ZhiPuAiApi {
@JsonInclude(Include.NON_NULL)
public record Embedding(
@JsonProperty("index") Integer index,
@JsonProperty("embedding") List<Double> 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<Double> embedding) {
public Embedding(Integer index, float[] embedding) {
this(index, embedding, "embedding");
}
}

4
models/spring-ai-zhipuai/src/test/java/org/springframework/ai/zhipuai/api/ZhiPuAiRetryTests.java
View File

@@ -171,7 +171,7 @@ public class ZhiPuAiRetryTests {
public void zhiPuAiEmbeddingTransientError() {
EmbeddingList<Embedding> 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);
}

6
spring-ai-core/src/main/java/org/springframework/ai/document/Document.java
View File

@@ -65,7 +65,7 @@ public class Document implements MediaContent {
* Embedding of the document. Note: ephemeral field.
*/
@JsonProperty(index = 100)
private List<Double> 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<Double> 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<Double> getEmbedding() {
public float[] getEmbedding() {
return this.embedding;
}

2
spring-ai-core/src/main/java/org/springframework/ai/embedding/AbstractEmbeddingModel.java
View File

@@ -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;
}
}

15
spring-ai-core/src/main/java/org/springframework/ai/embedding/Embedding.java
View File

@@ -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<List<Double>> {
public class Embedding implements ModelResult<float[]> {
private List<Double> embedding;
private float[] embedding;
private Integer index;
@@ -36,7 +35,7 @@ public class Embedding implements ModelResult<List<Double>> {
* @param embedding the embedding vector values.
* @param index the embedding index in a list of embeddings.
*/
public Embedding(List<Double> embedding, Integer index) {
public Embedding(float[] embedding, Integer index) {
this(embedding, index, EmbeddingResultMetadata.EMPTY);
}
@@ -46,7 +45,7 @@ public class Embedding implements ModelResult<List<Double>> {
* @param index the embedding index in a list of embeddings.
* @param metadata the metadata associated with the embedding.
*/
public Embedding(List<Double> 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<List<Double>> {
* @return Get the embedding vector values.
*/
@Override
public List<Double> getOutput() {
return this.embedding;
public float[] getOutput() {
return embedding;
}
/**
@@ -91,7 +90,7 @@ public class Embedding implements ModelResult<List<Double>> {
@Override
public String toString() {
String message = this.embedding.isEmpty() ? "<empty>" : "<has data>";
String message = this.embedding.length == 0 ? "<empty>" : "<has data>";
return "Embedding{" + "embedding=" + message + ", index=" + this.index + '}';
}

11
spring-ai-core/src/main/java/org/springframework/ai/embedding/EmbeddingModel.java
View File

@@ -34,9 +34,10 @@ public interface EmbeddingModel extends Model<EmbeddingRequest, EmbeddingRespons
* @param text the text to embed.
* @return the embedded vector.
*/
default List<Double> 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<float[]> response = this.embed(List.of(text));
return response.iterator().next();
}
/**
@@ -44,14 +45,14 @@ public interface EmbeddingModel extends Model<EmbeddingRequest, EmbeddingRespons
* @param document the document to embed.
* @return the embedded vector.
*/
List<Double> 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<List<Double>> embed(List<String> texts) {
default List<float[]> embed(List<String> 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<EmbeddingRequest, EmbeddingRespons
* specific.
*/
default int dimensions() {
return embed("Test String").size();
return embed("Test String").length;
}
}

75
spring-ai-core/src/main/java/org/springframework/ai/model/EmbeddingUtils.java Normal file
View File

@@ -0,0 +1,75 @@
/*
* Copyright 2024 - 2024 the original author or authors.
*
* Licensed under the Apache License, Version 2.0 (the "License");
* you may not use this file except in compliance with the License.
* You may obtain a copy of the License at
*
* https://www.apache.org/licenses/LICENSE-2.0
*
* Unless required by applicable law or agreed to in writing, software
* distributed under the License is distributed on an "AS IS" BASIS,
* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
* See the License for the specific language governing permissions and
* limitations under the License.
*/
package org.springframework.ai.model;
import java.util.ArrayList;
import java.util.List;
/**
* @author Christian Tzolov
* @since 1.0.0
*/
public class EmbeddingUtils {
private static final float[] EMPTY_FLOAT_ARRAY = new float[0];
public static List<Float> doubleToFloat(final List<Double> doubles) {
return doubles.stream().map(f -> f.floatValue()).toList();
}
public static float[] toPrimitive(List<Float> 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<Float> toList(float[] floats) {
List<Float> output = new ArrayList<Float>();
for (float value : floats) {
output.add(value);
}
return output;
}
}

28
spring-ai-core/src/main/java/org/springframework/ai/vectorstore/SimpleVectorStore.java
View File

@@ -72,7 +72,7 @@ public class SimpleVectorStore implements VectorStore {
public void add(List<Document> documents) {
for (Document document : documents) {
logger.info("Calling EmbeddingModel for document id = {}", document.getId());
List<Double> 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<Double> 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<Double> 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<Double> vectorX, List<Double> 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<Double> vectorX, List<Double> 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<Double> vector) {
public static float norm(float[] vector) {
return dotProduct(vector, vector);
}

10
spring-ai-core/src/test/java/org/springframework/ai/embedding/AbstractEmbeddingModelTests.java
View File

@@ -48,17 +48,17 @@ public class AbstractEmbeddingModelTests {
EmbeddingModel dummy = new EmbeddingModel() {
@Override
public List<Double> 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<Double> embed(Document document) {
public float[] embed(Document document) {
throw new UnsupportedOperationException("Unimplemented method 'embed'");
}
@Override
public List<List<Double>> embed(List<String> texts) {
public List<float[]> embed(List<String> 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);
}

12
spring-ai-docs/src/main/antora/modules/ROOT/pages/api/embeddings.adoc
View File

@@ -50,14 +50,14 @@ public interface EmbeddingModel extends Model<EmbeddingRequest, EmbeddingRespons
* @param document the document to embed.
* @return the embedded vector.
*/
List<Double> 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<Double> 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<EmbeddingRequest, EmbeddingRespons
* @param texts list of texts to embed.
* @return list of list of embedded vectors.
*/
default List<List<Double>> embed(List<String> texts) {
default List<float[]> embed(List<String> 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<List<Double>> {
private List<Double> embedding;
public class Embedding implements ModelResult<float[]> {
private float[] embedding;
private Integer index;
private EmbeddingResultMetadata metadata;
// other methods omitted

2
spring-ai-docs/src/main/antora/modules/ROOT/pages/api/vectordbs.adoc
View File

@@ -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<Double>`, 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.

8
spring-ai-spring-boot-autoconfigure/src/test/java/org/springframework/ai/autoconfigure/transformers/TransformersEmbeddingModelAutoConfigurationIT.java
View File

@@ -50,10 +50,10 @@ public class TransformersEmbeddingModelAutoConfigurationIT {
EmbeddingModel embeddingModel = context.getBean(EmbeddingModel.class);
assertThat(embeddingModel).isInstanceOf(TransformersEmbeddingModel.class);
List<List<Double>> embeddings = embeddingModel.embed(List.of("Spring Framework", "Spring AI"));
List<float[]> 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<List<Double>> embeddings = embeddingModel.embed(List.of("Spring Framework", "Spring AI"));
List<float[]> 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
});
}

6
spring-ai-spring-boot-autoconfigure/src/test/java/org/springframework/ai/autoconfigure/vertexai/embedding/VertexAiTextEmbeddingModelAutoConfigurationIT.java
View File

@@ -60,10 +60,10 @@ public class VertexAiTextEmbeddingModelAutoConfigurationIT {
VertexAiTextEmbeddingModel embeddingModel = context.getBean(VertexAiTextEmbeddingModel.class);
assertThat(embeddingModel).isInstanceOf(VertexAiTextEmbeddingModel.class);
List<List<Double>> embeddings = embeddingModel.embed(List.of("Spring Framework", "Spring AI"));
List<float[]> 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);

14
vector-stores/spring-ai-azure-store/src/main/java/org/springframework/ai/vectorstore/azure/AzureVectorStore.java
View File

@@ -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<Float> toFloatList(List<Double> doubleList) {
return doubleList.stream().map(Double::floatValue).toList();
}
/**
* Internal data structure for retrieving and storing documents.
*/
private record AzureSearchDocument(String id, String content, List<Double> embedding, String metadata) {
private record AzureSearchDocument(String id, String content, List<Float> embedding, String metadata) {
}
@Override

18
vector-stores/spring-ai-cassandra-store/src/main/java/org/springframework/ai/vectorstore/CassandraVectorStore.java
View File

@@ -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<Object> 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<Double> 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;

15
vector-stores/spring-ai-chroma-store/src/main/java/org/springframework/ai/chroma/ChromaApi.java
View File

@@ -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<String> ids, List<List<Float>> embeddings, List<String> documents,
public record GetEmbeddingResponse(List<String> ids, List<float[]> embeddings, List<String> documents,
@JsonProperty("metadatas") List<Map<String, String>> 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<List<Float>> queryEmbeddings,
public record QueryRequest(@JsonProperty("query_embeddings") List<float[]> queryEmbeddings,
@JsonProperty("n_results") int nResults, Map<String, Object> where, List<Include> 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<Float> queryEmbedding, int nResults) {
public QueryRequest(float[] queryEmbedding, int nResults) {
this(List.of(queryEmbedding), nResults, Map.of(), Include.all);
}
public QueryRequest(List<Float> queryEmbedding, int nResults, Map<String, Object> where) {
public QueryRequest(float[] queryEmbedding, int nResults, Map<String, Object> 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<List<String>> ids, List<List<List<Float>>> embeddings,
List<List<String>> documents, @JsonProperty("metadatas") List<List<Map<String, Object>>> metadata,
List<List<Double>> distances) {
public record QueryResponse(List<List<String>> ids, List<List<float[]>> embeddings, List<List<String>> documents,
@JsonProperty("metadatas") List<List<Map<String, Object>>> metadata, List<List<Double>> distances) {
}
/**
* Single query embedding response.
*/
public record Embedding(String id, List<Float> embedding, String document, Map<String, Object> metadata,
public record Embedding(String id, float[] embedding, String document, Map<String, Object> metadata,
Double distances) {
}

8
vector-stores/spring-ai-chroma-store/src/main/java/org/springframework/ai/vectorstore/ChromaVectorStore.java
View File

@@ -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<Double> embedding = this.embeddingModel.embed(query);
float[] embedding = this.embeddingModel.embed(query);
Map<String, Object> 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);
}
}

24
vector-stores/spring-ai-chroma-store/src/main/java/org/springframework/ai/vectorstore/JsonUtils.java
View File

@@ -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<Float> toFloatList(List<Double> 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<Float> toFloatList(List<Double> 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<Double> embeddingDouble) {
float[] embeddingFloat = new float[embeddingDouble.size()];
public static float[] toFloatArray(List<Float> embedding) {
float[] embeddingFloat = new float[embedding.size()];
int i = 0;
for (Double d : embeddingDouble) {
for (Float d : embedding) {
embeddingFloat[i++] = d.floatValue();
}
return embeddingFloat;

42
vector-stores/spring-ai-chroma-store/src/test/java/org/springframework/ai/chroma/ChromaApiIT.java
View File

@@ -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");
}

10
vector-stores/spring-ai-elasticsearch-store/src/main/java/org/springframework/ai/vectorstore/ElasticsearchVectorStore.java
View File

@@ -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<Float> vectors = this.embeddingModel.embed(searchRequest.getQuery())
.stream()
.map(Double::floatValue)
.toList();
float[] vectors = this.embeddingModel.embed(searchRequest.getQuery());
SearchResponse<Document> 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")

17
vector-stores/spring-ai-gemfire-store/src/main/java/org/springframework/ai/vectorstore/GemFireVectorStore.java
View File

@@ -256,12 +256,12 @@ public class GemFireVectorStore implements VectorStore, InitializingBean {
private final String key;
private List<Float> vector;
private float[] vector;
@JsonInclude(JsonInclude.Include.NON_NULL)
private Map<String, Object> metadata;
public Embedding(@JsonProperty("key") String key, @JsonProperty("vector") List<Float> vector,
public Embedding(@JsonProperty("key") String key, @JsonProperty("vector") float[] vector,
String contentName, String content, @JsonProperty("metadata") Map<String, Object> metadata) {
this.key = key;
this.vector = vector;
@@ -273,7 +273,7 @@ public class GemFireVectorStore implements VectorStore, InitializingBean {
return key;
}
public List<Float> getVector() {
public float[] getVector() {
return vector;
}
@@ -289,7 +289,7 @@ public class GemFireVectorStore implements VectorStore, InitializingBean {
@JsonProperty("vector")
@NonNull
private final List<Float> 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<Float> 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<Float> 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<Float> 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<Double> vector = this.embeddingModel.embed(request.getQuery());
List<Float> floatVector = vector.stream().map(Double::floatValue).toList();
float[] floatVector = this.embeddingModel.embed(request.getQuery());
return client.post()
.uri("/" + indexName + QUERY)
.contentType(MediaType.APPLICATION_JSON)

9
vector-stores/spring-ai-hanadb-store/src/main/java/org/springframework/ai/vectorstore/HanaCloudVectorStore.java
View File

@@ -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(", ")) + "]";
}
}

14
vector-stores/spring-ai-milvus-store/src/main/java/org/springframework/ai/vectorstore/MilvusVectorStore.java
View File

@@ -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<List<Float>> embeddingArray = new ArrayList<>();
for (Document document : documents) {
List<Double> 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<InsertParam.Field> fields = new ArrayList<>();
@@ -327,7 +327,7 @@ public class MilvusVectorStore implements VectorStore, InitializingBean {
Assert.notNull(request.getQuery(), "Query string must not be null");
List<Double> 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<Float> toFloatList(List<Double> embeddingDouble) {
return embeddingDouble.stream().map(Number::floatValue).toList();
}
// ---------------------------------------------------------------------------------
// Initialization
// ---------------------------------------------------------------------------------

13
vector-stores/spring-ai-mongodb-atlas-store/src/main/java/org/springframework/ai/vectorstore/MongoDBAtlasVectorStore.java
View File

@@ -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<String, Object> metadata = mongoDocument.get(METADATA_FIELD_NAME, org.bson.Document.class);
List<Double> embedding = mongoDocument.getList(this.config.pathName, Double.class);
List<Float> 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<Document> documents) {
for (Document document : documents) {
List<Double> 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<Double> 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()

6
vector-stores/spring-ai-mongodb-atlas-store/src/main/java/org/springframework/ai/vectorstore/VectorSearchAggregation.java
View File

@@ -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<Double> embeddings, String path, int numCandidates, String index, int count,
record VectorSearchAggregation(List<Float> embeddings, String path, int numCandidates, String index, int count,
String filter) implements AggregationOperation {
@SuppressWarnings("null")

8
vector-stores/spring-ai-mongodb-atlas-store/src/test/java/org/springframework/ai/vectorstore/VectorSearchAggregationTest.java
View File

@@ -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")))

25
vector-stores/spring-ai-neo4j-store/src/main/java/org/springframework/ai/vectorstore/Neo4jVectorStore.java
View File

@@ -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<Double> 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();

6
vector-stores/spring-ai-opensearch-store/src/main/java/org/springframework/ai/vectorstore/OpenSearchVectorStore.java
View File

@@ -116,7 +116,7 @@ public class OpenSearchVectorStore implements VectorStore, InitializingBean {
public void add(List<Document> 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<Document> similaritySearch(List<Double> embedding, int topK, double similarityThreshold,
public List<Document> 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<Double> 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

18
vector-stores/spring-ai-oracle-store/src/main/java/org/springframework/ai/vectorstore/OracleVectorStore.java
View File

@@ -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<Double> 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<Double> toDoubleList(final double[] embeddings) {
final List<Double> result = new ArrayList<>(embeddings.length);
for (double v : embeddings) {
private List<Float> toFloatList(final float[] embeddings) {
final List<Float> result = new ArrayList<>(embeddings.length);
for (float v : embeddings) {
result.add(v);
}
return result;

20
vector-stores/spring-ai-pgvector-store/src/main/java/org/springframework/ai/vectorstore/PgVectorStore.java
View File

@@ -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<Double> embeddingDouble) {
float[] embeddingFloat = new float[embeddingDouble.size()];
private float[] toFloatArray(List<Float> 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<Double> 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<Double> 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<String, Object> toMap(PGobject pgObject) {

16
vector-stores/spring-ai-pinecone-store/src/main/java/org/springframework/ai/vectorstore/PineconeVectorStore.java
View File

@@ -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<Double> 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<Float> toFloatList(List<Double> doubleList) {
return doubleList.stream().map(d -> d.floatValue()).toList();
}
}

39
vector-stores/spring-ai-qdrant-store/src/main/java/org/springframework/ai/vectorstore/qdrant/QdrantVectorStore.java
View File

@@ -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<Double> 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<Float> toFloatList(List<Double> doubleList) {
return doubleList.stream().map(d -> d.floatValue()).toList();
}
@Override
public void afterPropertiesSet() throws Exception {

8
vector-stores/spring-ai-redis-store/src/main/java/org/springframework/ai/vectorstore/RedisVectorStore.java
View File

@@ -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<Double> embeddingDouble) {
float[] embeddingFloat = new float[embeddingDouble.size()];
private static float[] toFloatArray(List<Float> embedding) {
float[] embeddingFloat = new float[embedding.size()];
int i = 0;
for (Double d : embeddingDouble) {
for (Float d : embedding) {
embeddingFloat[i++] = d.floatValue();
}
return embeddingFloat;

14
vector-stores/spring-ai-typesense-store/src/main/java/org/springframework/ai/vectorstore/TypesenseVectorStore.java
View File

@@ -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<Double> 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<Double> 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<Float> 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);

22
vector-stores/spring-ai-weaviate-store/src/main/java/org/springframework/ai/vectorstore/WeaviateVectorStore.java
View File

@@ -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<Double> 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<Document> 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<Double> doubleList) {
return doubleList.stream().map(Number::floatValue).toList().toArray(new Float[0]);
}
}