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

Introduce checkstyle plugin

Browse Source 
- Based on https://github.com/spring-io/spring-javaformat
- In this iteration, checkstyles are only enabled for spring-ai-core
This commit is contained in:
Soby Chacko
2024-10-24 10:39:48 -04:00
committed by Mark Pollack
parent 33a72417e1
commit 8e758dbd00
1412 changed files with 26997 additions and 21963 deletions
Download Patch File Download Diff File Expand all files Collapse all files

16
vector-stores/spring-ai-azure-cosmos-db-store/pom.xml
View File

@@ -1,4 +1,20 @@
<?xml version="1.0" encoding="UTF-8"?>
<!--
~ Copyright 2023-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.
-->
<project xmlns="http://maven.apache.org/POM/4.0.0"
xmlns:xsi="http://www.w3.org/2001/XMLSchema-instance" xsi:schemaLocation="http://maven.apache.org/POM/4.0.0 http://maven.apache.org/maven-v4_0_0.xsd">
<modelVersion>4.0.0</modelVersion>

12
vector-stores/spring-ai-azure-cosmos-db-store/src/main/java/org/springframework/ai/vectorstore/CosmosDBFilterExpressionConverter.java
View File

@@ -1,11 +1,11 @@
/*
* Copyright 2024 the original author or authors.
* Copyright 2023-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
* 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,
@@ -16,9 +16,6 @@
package org.springframework.ai.vectorstore;
import static org.springframework.ai.vectorstore.filter.Filter.ExpressionType.AND;
import static org.springframework.ai.vectorstore.filter.Filter.ExpressionType.OR;
import java.util.Collection;
import java.util.Map;
import java.util.Optional;
@@ -29,6 +26,9 @@ import org.springframework.ai.vectorstore.filter.Filter;
import org.springframework.ai.vectorstore.filter.Filter.Key;
import org.springframework.ai.vectorstore.filter.converter.AbstractFilterExpressionConverter;
import static org.springframework.ai.vectorstore.filter.Filter.ExpressionType.AND;
import static org.springframework.ai.vectorstore.filter.Filter.ExpressionType.OR;
/**
* Converts {@link org.springframework.ai.vectorstore.filter.Filter.Expression} into
* Cosmos DB NoSQL API where clauses.
@@ -51,7 +51,7 @@ class CosmosDBFilterExpressionConverter extends AbstractFilterExpressionConverte
*/
private Optional<String> getMetadataField(String name) {
String metadataField = name;
return Optional.ofNullable(metadataFields.get(metadataField));
return Optional.ofNullable(this.metadataFields.get(metadataField));
}
@Override

49
vector-stores/spring-ai-azure-cosmos-db-store/src/main/java/org/springframework/ai/vectorstore/CosmosDBVectorStore.java
View File

@@ -1,11 +1,11 @@
/*
* Copyright 2024 the original author or authors.
* Copyright 2023-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
* 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,
@@ -24,20 +24,6 @@ import java.util.Optional;
import java.util.stream.Collectors;
import java.util.stream.IntStream;
import org.apache.commons.lang3.tuple.ImmutablePair;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import org.springframework.ai.document.Document;
import org.springframework.ai.embedding.BatchingStrategy;
import org.springframework.ai.embedding.EmbeddingModel;
import org.springframework.ai.embedding.EmbeddingOptionsBuilder;
import org.springframework.ai.embedding.TokenCountBatchingStrategy;
import org.springframework.ai.observation.conventions.VectorStoreProvider;
import org.springframework.ai.vectorstore.filter.Filter;
import org.springframework.ai.vectorstore.observation.AbstractObservationVectorStore;
import org.springframework.ai.vectorstore.observation.VectorStoreObservationContext;
import org.springframework.ai.vectorstore.observation.VectorStoreObservationConvention;
import com.azure.cosmos.CosmosAsyncClient;
import com.azure.cosmos.CosmosAsyncContainer;
import com.azure.cosmos.CosmosAsyncDatabase;
@@ -66,10 +52,23 @@ import com.azure.cosmos.util.CosmosPagedFlux;
import com.fasterxml.jackson.databind.JsonNode;
import com.fasterxml.jackson.databind.ObjectMapper;
import com.fasterxml.jackson.databind.node.ObjectNode;
import io.micrometer.observation.ObservationRegistry;
import org.apache.commons.lang3.tuple.ImmutablePair;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import reactor.core.publisher.Flux;
import org.springframework.ai.document.Document;
import org.springframework.ai.embedding.BatchingStrategy;
import org.springframework.ai.embedding.EmbeddingModel;
import org.springframework.ai.embedding.EmbeddingOptionsBuilder;
import org.springframework.ai.embedding.TokenCountBatchingStrategy;
import org.springframework.ai.observation.conventions.VectorStoreProvider;
import org.springframework.ai.vectorstore.filter.Filter;
import org.springframework.ai.vectorstore.observation.AbstractObservationVectorStore;
import org.springframework.ai.vectorstore.observation.VectorStoreObservationContext;
import org.springframework.ai.vectorstore.observation.VectorStoreObservationConvention;
/**
* @author Theo van Kraay
* @author Soby Chacko
@@ -81,14 +80,14 @@ public class CosmosDBVectorStore extends AbstractObservationVectorStore implemen
private final CosmosAsyncClient cosmosClient;
private CosmosAsyncContainer container;
private final EmbeddingModel embeddingModel;
private final CosmosDBVectorStoreConfig properties;
private final BatchingStrategy batchingStrategy;
private CosmosAsyncContainer container;
public CosmosDBVectorStore(ObservationRegistry observationRegistry,
VectorStoreObservationConvention customObservationConvention, CosmosAsyncClient cosmosClient,
CosmosDBVectorStoreConfig properties, EmbeddingModel embeddingModel) {
@@ -210,7 +209,7 @@ public class CosmosDBVectorStore extends AbstractObservationVectorStore implemen
CosmosItemOperation operation = CosmosBulkOperations
.getCreateItemOperation(mapCosmosDocument(doc, doc.getEmbedding()), new PartitionKey(doc.getId()));
return new ImmutablePair<>(doc.getId(), operation); // Pair the document ID
// with the operation
// with the operation
}).toList();
try {
@@ -233,7 +232,7 @@ public class CosmosDBVectorStore extends AbstractObservationVectorStore implemen
String errorMessage = String.format("Duplicate document id: %s", documentId);
logger.error(errorMessage);
throw new RuntimeException(errorMessage); // Throw an exception
// for status code 409
// for status code 409
}
else {
logger.info("Document added with status: {}", statusCode);
@@ -307,10 +306,10 @@ public class CosmosDBVectorStore extends AbstractObservationVectorStore implemen
if (filterExpression != null) {
CosmosDBFilterExpressionConverter filterExpressionConverter = new CosmosDBFilterExpressionConverter(
this.properties.getMetadataFieldsList()); // Use the expression
// directly as
// it handles the
// "metadata"
// fields internally
// directly as
// it handles the
// "metadata"
// fields internally
String filterQuery = filterExpressionConverter.convertExpression(filterExpression);
queryBuilder.append(" AND ").append(filterQuery);
}

12
vector-stores/spring-ai-azure-cosmos-db-store/src/main/java/org/springframework/ai/vectorstore/CosmosDBVectorStoreConfig.java
View File

@@ -1,11 +1,11 @@
/*
* Copyright 2024 the original author or authors.
* Copyright 2023-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
* 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,
@@ -51,15 +51,15 @@ public class CosmosDBVectorStoreConfig implements AutoCloseable {
this.vectorStoreThroughput = vectorStoreThroughput;
}
public String getMetadataFields() {
return this.metadataFields;
}
public void setMetadataFields(String metadataFields) {
this.metadataFields = metadataFields;
this.metadataFieldsList = List.of(metadataFields.split(","));
}
public String getMetadataFields() {
return this.metadataFields;
}
public List<String> getMetadataFieldsList() {
return this.metadataFieldsList;
}

41
vector-stores/spring-ai-azure-cosmos-db-store/src/test/java/org/springframework/ai/vectorstore/CosmosDBVectorStoreIT.java
View File

@@ -1,11 +1,11 @@
/*
* Copyright 2024 the original author or authors.
* Copyright 2023-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
* 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,
@@ -16,11 +16,17 @@
package org.springframework.ai.vectorstore;
import java.util.HashMap;
import java.util.List;
import java.util.Map;
import java.util.UUID;
import com.azure.cosmos.CosmosAsyncClient;
import com.azure.cosmos.CosmosClientBuilder;
import org.junit.jupiter.api.BeforeEach;
import org.junit.jupiter.api.Test;
import org.junit.jupiter.api.condition.EnabledIfEnvironmentVariable;
import org.springframework.ai.document.Document;
import org.springframework.ai.embedding.EmbeddingModel;
import org.springframework.ai.transformers.TransformersEmbeddingModel;
@@ -30,10 +36,6 @@ import org.springframework.boot.SpringBootConfiguration;
import org.springframework.boot.autoconfigure.EnableAutoConfiguration;
import org.springframework.boot.test.context.runner.ApplicationContextRunner;
import org.springframework.context.annotation.Bean;
import java.util.HashMap;
import java.util.List;
import java.util.Map;
import java.util.UUID;
import static org.assertj.core.api.Assertions.assertThat;
import static org.assertj.core.api.Assertions.assertThatThrownBy;
@@ -53,8 +55,8 @@ public class CosmosDBVectorStoreIT {
@BeforeEach
public void setup() {
contextRunner.run(context -> {
vectorStore = context.getBean(VectorStore.class);
this.contextRunner.run(context -> {
this.vectorStore = context.getBean(VectorStore.class);
});
}
@@ -66,25 +68,25 @@ public class CosmosDBVectorStoreIT {
Document document2 = new Document(UUID.randomUUID().toString(), "Sample content2", Map.of("key2", "value2"));
// Add the document to the vector store
vectorStore.add(List.of(document1, document2));
this.vectorStore.add(List.of(document1, document2));
// create duplicate docs and assert that second one throws exception
Document document3 = new Document(document1.getId(), "Sample content3", Map.of("key3", "value3"));
assertThatThrownBy(() -> vectorStore.add(List.of(document3))).isInstanceOf(Exception.class)
assertThatThrownBy(() -> this.vectorStore.add(List.of(document3))).isInstanceOf(Exception.class)
.hasMessageContaining("Duplicate document id: " + document1.getId());
// Perform a similarity search
List<Document> results = vectorStore.similaritySearch(SearchRequest.query("Sample content").withTopK(1));
List<Document> results = this.vectorStore.similaritySearch(SearchRequest.query("Sample content").withTopK(1));
// Verify the search results
assertThat(results).isNotEmpty();
assertThat(results.get(0).getId()).isEqualTo(document1.getId());
// Remove the documents from the vector store
vectorStore.delete(List.of(document1.getId(), document2.getId()));
this.vectorStore.delete(List.of(document1.getId(), document2.getId()));
// Perform a similarity search again
List<Document> results2 = vectorStore.similaritySearch(SearchRequest.query("Sample content").withTopK(1));
List<Document> results2 = this.vectorStore.similaritySearch(SearchRequest.query("Sample content").withTopK(1));
// Verify the search results
assertThat(results2).isEmpty();
@@ -124,16 +126,16 @@ public class CosmosDBVectorStoreIT {
Document document3 = new Document("3", "A document about the US", metadata3);
Document document4 = new Document("4", "A document about the US", metadata4);
vectorStore.add(List.of(document1, document2, document3, document4));
this.vectorStore.add(List.of(document1, document2, document3, document4));
FilterExpressionBuilder b = new FilterExpressionBuilder();
List<Document> results = vectorStore.similaritySearch(SearchRequest.query("The World")
List<Document> results = this.vectorStore.similaritySearch(SearchRequest.query("The World")
.withTopK(10)
.withFilterExpression((b.in("country", "UK", "NL")).build()));
assertThat(results).hasSize(2);
assertThat(results).extracting(Document::getId).containsExactlyInAnyOrder("1", "2");
List<Document> results2 = vectorStore.similaritySearch(SearchRequest.query("The World")
List<Document> results2 = this.vectorStore.similaritySearch(SearchRequest.query("The World")
.withTopK(10)
.withFilterExpression(
b.and(b.or(b.gte("year", 2021), b.eq("country", "NL")), b.ne("city", "Amsterdam")).build()));
@@ -141,17 +143,17 @@ public class CosmosDBVectorStoreIT {
assertThat(results2).hasSize(1);
assertThat(results2).extracting(Document::getId).containsExactlyInAnyOrder("1");
List<Document> results3 = vectorStore.similaritySearch(SearchRequest.query("The World")
List<Document> results3 = this.vectorStore.similaritySearch(SearchRequest.query("The World")
.withTopK(10)
.withFilterExpression(b.and(b.eq("country", "US"), b.eq("year", 2020)).build()));
assertThat(results3).hasSize(1);
assertThat(results3).extracting(Document::getId).containsExactlyInAnyOrder("4");
vectorStore.delete(List.of(document1.getId(), document2.getId(), document3.getId(), document4.getId()));
this.vectorStore.delete(List.of(document1.getId(), document2.getId(), document3.getId(), document4.getId()));
// Perform a similarity search again
List<Document> results4 = vectorStore.similaritySearch(SearchRequest.query("The World").withTopK(1));
List<Document> results4 = this.vectorStore.similaritySearch(SearchRequest.query("The World").withTopK(1));
// Verify the search results
assertThat(results4).isEmpty();
@@ -191,6 +193,7 @@ public class CosmosDBVectorStoreIT {
public VectorStoreObservationConvention observationConvention() {
// Replace with an actual observation convention or a mock if needed
return new VectorStoreObservationConvention() {
};
}

16
vector-stores/spring-ai-azure-cosmos-db-store/src/test/resources/application.properties
View File

@@ -1,3 +1,19 @@
#
# Copyright 2023-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.
#
spring.ai.vectorstore.cosmosdb.databaseName=db
spring.ai.vectorstore.cosmosdb.containerName=container
spring.ai.vectorstore.cosmosdb.key=${COSMOSDB_AI_ENDPOINT}

16
vector-stores/spring-ai-azure-store/pom.xml
View File

@@ -1,4 +1,20 @@
<?xml version="1.0" encoding="UTF-8"?>
<!--
~ Copyright 2023-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.
-->
<project xmlns="http://maven.apache.org/POM/4.0.0"
xmlns:xsi="http://www.w3.org/2001/XMLSchema-instance" xsi:schemaLocation="http://maven.apache.org/POM/4.0.0 http://maven.apache.org/maven-v4_0_0.xsd">
<modelVersion>4.0.0</modelVersion>

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

@@ -1,11 +1,11 @@
/*
* Copyright 2023 - 2024 the original author or authors.
* Copyright 2023-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
* 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,
@@ -13,6 +13,7 @@
* See the License for the specific language governing permissions and
* limitations under the License.
*/
package org.springframework.ai.vectorstore.azure;
import java.text.ParseException;

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

@@ -5,7 +5,7 @@
* 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
* 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,
@@ -23,27 +23,6 @@ import java.util.Map;
import java.util.Optional;
import java.util.stream.Collectors;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import org.springframework.ai.document.Document;
import org.springframework.ai.embedding.BatchingStrategy;
import org.springframework.ai.embedding.EmbeddingModel;
import org.springframework.ai.embedding.EmbeddingOptionsBuilder;
import org.springframework.ai.embedding.TokenCountBatchingStrategy;
import org.springframework.ai.model.EmbeddingUtils;
import org.springframework.ai.observation.conventions.VectorStoreProvider;
import org.springframework.ai.observation.conventions.VectorStoreSimilarityMetric;
import org.springframework.ai.vectorstore.SearchRequest;
import org.springframework.ai.vectorstore.filter.FilterExpressionConverter;
import org.springframework.ai.vectorstore.observation.AbstractObservationVectorStore;
import org.springframework.ai.vectorstore.observation.VectorStoreObservationContext;
import org.springframework.ai.vectorstore.observation.VectorStoreObservationContext.Builder;
import org.springframework.ai.vectorstore.observation.VectorStoreObservationConvention;
import org.springframework.beans.factory.InitializingBean;
import org.springframework.util.Assert;
import org.springframework.util.CollectionUtils;
import org.springframework.util.StringUtils;
import com.alibaba.fastjson2.JSONObject;
import com.alibaba.fastjson2.TypeReference;
import com.azure.core.util.Context;
@@ -63,8 +42,28 @@ import com.azure.search.documents.models.IndexingResult;
import com.azure.search.documents.models.SearchOptions;
import com.azure.search.documents.models.VectorSearchOptions;
import com.azure.search.documents.models.VectorizedQuery;
import io.micrometer.observation.ObservationRegistry;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import org.springframework.ai.document.Document;
import org.springframework.ai.embedding.BatchingStrategy;
import org.springframework.ai.embedding.EmbeddingModel;
import org.springframework.ai.embedding.EmbeddingOptionsBuilder;
import org.springframework.ai.embedding.TokenCountBatchingStrategy;
import org.springframework.ai.model.EmbeddingUtils;
import org.springframework.ai.observation.conventions.VectorStoreProvider;
import org.springframework.ai.observation.conventions.VectorStoreSimilarityMetric;
import org.springframework.ai.vectorstore.SearchRequest;
import org.springframework.ai.vectorstore.filter.FilterExpressionConverter;
import org.springframework.ai.vectorstore.observation.AbstractObservationVectorStore;
import org.springframework.ai.vectorstore.observation.VectorStoreObservationContext;
import org.springframework.ai.vectorstore.observation.VectorStoreObservationContext.Builder;
import org.springframework.ai.vectorstore.observation.VectorStoreObservationConvention;
import org.springframework.beans.factory.InitializingBean;
import org.springframework.util.Assert;
import org.springframework.util.CollectionUtils;
import org.springframework.util.StringUtils;
/**
* Uses Azure Cognitive Search as a backing vector store. Documents can be preloaded into
@@ -81,14 +80,14 @@ import io.micrometer.observation.ObservationRegistry;
*/
public class AzureVectorStore extends AbstractObservationVectorStore implements InitializingBean {
public static final String DEFAULT_INDEX_NAME = "spring_ai_azure_vector_store";
private static final Logger logger = LoggerFactory.getLogger(AzureVectorStore.class);
private static final String SPRING_AI_VECTOR_CONFIG = "spring-ai-vector-config";
private static final String SPRING_AI_VECTOR_PROFILE = "spring-ai-vector-profile";
public static final String DEFAULT_INDEX_NAME = "spring_ai_azure_vector_store";
private static final String ID_FIELD_NAME = "id";
private static final String CONTENT_FIELD_NAME = "content";
@@ -109,16 +108,8 @@ public class AzureVectorStore extends AbstractObservationVectorStore implements
private final EmbeddingModel embeddingModel;
private SearchClient searchClient;
private final FilterExpressionConverter filterExpressionConverter;
private int defaultTopK = DEFAULT_TOP_K;
private Double defaultSimilarityThreshold = DEFAULT_SIMILARITY_THRESHOLD;
private String indexName = DEFAULT_INDEX_NAME;
private final boolean initializeSchema;
private final BatchingStrategy batchingStrategy;
@@ -134,32 +125,13 @@ public class AzureVectorStore extends AbstractObservationVectorStore implements
*/
private final List<MetadataField> filterMetadataFields;
public record MetadataField(String name, SearchFieldDataType fieldType) {
private SearchClient searchClient;
public static MetadataField text(String name) {
return new MetadataField(name, SearchFieldDataType.STRING);
}
private int defaultTopK = DEFAULT_TOP_K;
public static MetadataField int32(String name) {
return new MetadataField(name, SearchFieldDataType.INT32);
}
private Double defaultSimilarityThreshold = DEFAULT_SIMILARITY_THRESHOLD;
public static MetadataField int64(String name) {
return new MetadataField(name, SearchFieldDataType.INT64);
}
public static MetadataField decimal(String name) {
return new MetadataField(name, SearchFieldDataType.DOUBLE);
}
public static MetadataField bool(String name) {
return new MetadataField(name, SearchFieldDataType.BOOLEAN);
}
public static MetadataField date(String name) {
return new MetadataField(name, SearchFieldDataType.DATE_TIME_OFFSET);
}
}
private String indexName = DEFAULT_INDEX_NAME;
/**
* Constructs a new AzureCognitiveSearchVectorStore.
@@ -320,7 +292,7 @@ public class AzureVectorStore extends AbstractObservationVectorStore implements
Assert.notNull(request, "The search request must not be null.");
var searchEmbedding = embeddingModel.embed(request.getQuery());
var searchEmbedding = this.embeddingModel.embed(request.getQuery());
final var vectorQuery = new VectorizedQuery(EmbeddingUtils.toList(searchEmbedding))
.setKNearestNeighborsCount(request.getTopK())
@@ -336,7 +308,7 @@ public class AzureVectorStore extends AbstractObservationVectorStore implements
searchOptions.setFilter(oDataFilter);
}
final var searchResults = searchClient.search(null, searchOptions, Context.NONE);
final var searchResults = this.searchClient.search(null, searchOptions, Context.NONE);
return searchResults.stream()
.filter(result -> result.getScore() >= request.getSimilarityThreshold())
@@ -346,6 +318,7 @@ public class AzureVectorStore extends AbstractObservationVectorStore implements
Map<String, Object> metadata = (StringUtils.hasText(entry.metadata()))
? JSONObject.parseObject(entry.metadata(), new TypeReference<Map<String, Object>>() {
}) : Map.of();
metadata.put(DISTANCE_METADATA_FIELD_NAME, 1 - (float) result.getScore());
@@ -359,12 +332,6 @@ public class AzureVectorStore extends AbstractObservationVectorStore implements
.collect(Collectors.toList());
}
/**
* Internal data structure for retrieving and storing documents.
*/
private record AzureSearchDocument(String id, String content, List<Float> embedding, String metadata) {
}
@Override
public void afterPropertiesSet() throws Exception {
@@ -426,4 +393,39 @@ public class AzureVectorStore extends AbstractObservationVectorStore implements
.withSimilarityMetric(this.initializeSchema ? VectorStoreSimilarityMetric.COSINE.value() : null);
}
public record MetadataField(String name, SearchFieldDataType fieldType) {
public static MetadataField text(String name) {
return new MetadataField(name, SearchFieldDataType.STRING);
}
public static MetadataField int32(String name) {
return new MetadataField(name, SearchFieldDataType.INT32);
}
public static MetadataField int64(String name) {
return new MetadataField(name, SearchFieldDataType.INT64);
}
public static MetadataField decimal(String name) {
return new MetadataField(name, SearchFieldDataType.DOUBLE);
}
public static MetadataField bool(String name) {
return new MetadataField(name, SearchFieldDataType.BOOLEAN);
}
public static MetadataField date(String name) {
return new MetadataField(name, SearchFieldDataType.DATE_TIME_OFFSET);
}
}
/**
* Internal data structure for retrieving and storing documents.
*/
private record AzureSearchDocument(String id, String content, List<Float> embedding, String metadata) {
}
}

7
vector-stores/spring-ai-azure-store/src/test/java/org/springframework/ai/vectorstore/azure/AzureAiSearchFilterExpressionConverterTests.java
View File

@@ -1,11 +1,11 @@
/*
* Copyright 2023 - 2024 the original author or authors.
* Copyright 2023-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
* 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,
@@ -13,6 +13,7 @@
* See the License for the specific language governing permissions and
* limitations under the License.
*/
package org.springframework.ai.vectorstore.azure;
import java.util.Date;
@@ -204,4 +205,4 @@ public class AzureAiSearchFilterExpressionConverterTests {
""");
}
}
}

51
vector-stores/spring-ai-azure-store/src/test/java/org/springframework/ai/vectorstore/azure/AzureVectorStoreIT.java
View File

@@ -5,7 +5,7 @@
* 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
* 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,
@@ -13,6 +13,7 @@
* See the License for the specific language governing permissions and
* limitations under the License.
*/
package org.springframework.ai.vectorstore.azure;
import java.io.IOException;
@@ -56,14 +57,14 @@ import static org.hamcrest.Matchers.hasSize;
@EnabledIfEnvironmentVariable(named = "AZURE_AI_SEARCH_ENDPOINT", matches = ".+")
public class AzureVectorStoreIT {
private final ApplicationContextRunner contextRunner = new ApplicationContextRunner()
.withUserConfiguration(Config.class);
List<Document> documents = List.of(
new Document("1", getText("classpath:/test/data/spring.ai.txt"), Map.of("meta1", "meta1")),
new Document("2", getText("classpath:/test/data/time.shelter.txt"), Map.of()),
new Document("3", getText("classpath:/test/data/great.depression.txt"), Map.of("meta2", "meta2")));
private final ApplicationContextRunner contextRunner = new ApplicationContextRunner()
.withUserConfiguration(Config.class);
@BeforeAll
public static void beforeAll() {
Awaitility.setDefaultPollInterval(2, TimeUnit.SECONDS);
@@ -71,14 +72,24 @@ public class AzureVectorStoreIT {
Awaitility.setDefaultTimeout(Duration.ofMinutes(1));
}
private static String getText(String uri) {
var resource = new DefaultResourceLoader().getResource(uri);
try {
return resource.getContentAsString(StandardCharsets.UTF_8);
}
catch (IOException e) {
throw new RuntimeException(e);
}
}
@Test
public void addAndSearchTest() {
contextRunner.run(context -> {
this.contextRunner.run(context -> {
VectorStore vectorStore = context.getBean(VectorStore.class);
vectorStore.add(documents);
vectorStore.add(this.documents);
Awaitility.await()
.until(() -> vectorStore.similaritySearch(SearchRequest.query("Great Depression").withTopK(1)),
@@ -88,14 +99,14 @@ public class AzureVectorStoreIT {
assertThat(results).hasSize(1);
Document resultDoc = results.get(0);
assertThat(resultDoc.getId()).isEqualTo(documents.get(2).getId());
assertThat(resultDoc.getId()).isEqualTo(this.documents.get(2).getId());
assertThat(resultDoc.getContent()).contains("The Great Depression (1929–1939) was an economic shock");
assertThat(resultDoc.getMetadata()).hasSize(2);
assertThat(resultDoc.getMetadata()).containsKey("meta2");
assertThat(resultDoc.getMetadata()).containsKey("distance");
// Remove all documents from the store
vectorStore.delete(documents.stream().map(doc -> doc.getId()).toList());
vectorStore.delete(this.documents.stream().map(doc -> doc.getId()).toList());
Awaitility.await()
.until(() -> vectorStore.similaritySearch(SearchRequest.query("Hello").withTopK(1)), hasSize(0));
@@ -105,7 +116,7 @@ public class AzureVectorStoreIT {
@Test
public void searchWithFilters() throws InterruptedException {
contextRunner.run(context -> {
this.contextRunner.run(context -> {
VectorStore vectorStore = context.getBean(VectorStore.class);
var bgDocument = new Document("1", "The World is Big and Salvation Lurks Around the Corner",
@@ -194,7 +205,7 @@ public class AzureVectorStoreIT {
@Test
public void documentUpdateTest() {
contextRunner.run(context -> {
this.contextRunner.run(context -> {
VectorStore vectorStore = context.getBean(VectorStore.class);
@@ -247,11 +258,11 @@ public class AzureVectorStoreIT {
@Test
public void searchThresholdTest() {
contextRunner.run(context -> {
this.contextRunner.run(context -> {
VectorStore vectorStore = context.getBean(VectorStore.class);
vectorStore.add(documents);
vectorStore.add(this.documents);
Awaitility.await()
.until(() -> vectorStore
@@ -272,13 +283,13 @@ public class AzureVectorStoreIT {
assertThat(results).hasSize(1);
Document resultDoc = results.get(0);
assertThat(resultDoc.getId()).isEqualTo(documents.get(2).getId());
assertThat(resultDoc.getId()).isEqualTo(this.documents.get(2).getId());
assertThat(resultDoc.getContent()).contains("The Great Depression (1929–1939) was an economic shock");
assertThat(resultDoc.getMetadata()).containsKey("meta2");
assertThat(resultDoc.getMetadata()).containsKey("distance");
// Remove all documents from the store
vectorStore.delete(documents.stream().map(doc -> doc.getId()).toList());
vectorStore.delete(this.documents.stream().map(doc -> doc.getId()).toList());
Awaitility.await()
.until(() -> vectorStore.similaritySearch(SearchRequest.query("Hello").withTopK(1)), hasSize(0));
});
@@ -309,14 +320,4 @@ public class AzureVectorStoreIT {
}
private static String getText(String uri) {
var resource = new DefaultResourceLoader().getResource(uri);
try {
return resource.getContentAsString(StandardCharsets.UTF_8);
}
catch (IOException e) {
throw new RuntimeException(e);
}
}
}
}

32
vector-stores/spring-ai-azure-store/src/test/java/org/springframework/ai/vectorstore/azure/AzureVectorStoreObservationIT.java
View File

@@ -1,11 +1,11 @@
/*
* Copyright 2024 the original author or authors.
* Copyright 2023-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
* 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,
@@ -13,9 +13,8 @@
* See the License for the specific language governing permissions and
* limitations under the License.
*/
package org.springframework.ai.vectorstore.azure;
import static org.assertj.core.api.Assertions.assertThat;
package org.springframework.ai.vectorstore.azure;
import java.io.IOException;
import java.nio.charset.StandardCharsets;
@@ -24,10 +23,17 @@ import java.util.List;
import java.util.Map;
import java.util.concurrent.TimeUnit;
import com.azure.core.credential.AzureKeyCredential;
import com.azure.search.documents.indexes.SearchIndexClient;
import com.azure.search.documents.indexes.SearchIndexClientBuilder;
import io.micrometer.observation.ObservationRegistry;
import io.micrometer.observation.tck.TestObservationRegistry;
import io.micrometer.observation.tck.TestObservationRegistryAssert;
import org.awaitility.Awaitility;
import org.junit.jupiter.api.BeforeAll;
import org.junit.jupiter.api.Test;
import org.junit.jupiter.api.condition.EnabledIfEnvironmentVariable;
import org.springframework.ai.document.Document;
import org.springframework.ai.embedding.EmbeddingModel;
import org.springframework.ai.embedding.TokenCountBatchingStrategy;
@@ -47,13 +53,7 @@ import org.springframework.boot.test.context.runner.ApplicationContextRunner;
import org.springframework.context.annotation.Bean;
import org.springframework.core.io.DefaultResourceLoader;
import com.azure.core.credential.AzureKeyCredential;
import com.azure.search.documents.indexes.SearchIndexClient;
import com.azure.search.documents.indexes.SearchIndexClientBuilder;
import io.micrometer.observation.ObservationRegistry;
import io.micrometer.observation.tck.TestObservationRegistry;
import io.micrometer.observation.tck.TestObservationRegistryAssert;
import static org.assertj.core.api.Assertions.assertThat;
/**
* Integration tests for observation instrumentation AbstractObservationVectorStore in
@@ -66,6 +66,9 @@ import io.micrometer.observation.tck.TestObservationRegistryAssert;
@EnabledIfEnvironmentVariable(named = "AZURE_AI_SEARCH_ENDPOINT", matches = ".+")
public class AzureVectorStoreObservationIT {
private final ApplicationContextRunner contextRunner = new ApplicationContextRunner()
.withUserConfiguration(Config.class);
List<Document> documents = List.of(
new Document(getText("classpath:/test/data/spring.ai.txt"), Map.of("meta1", "meta1")),
new Document(getText("classpath:/test/data/time.shelter.txt")),
@@ -81,9 +84,6 @@ public class AzureVectorStoreObservationIT {
}
}
private final ApplicationContextRunner contextRunner = new ApplicationContextRunner()
.withUserConfiguration(Config.class);
@BeforeAll
public static void beforeAll() {
Awaitility.setDefaultPollInterval(2, TimeUnit.SECONDS);
@@ -94,13 +94,13 @@ public class AzureVectorStoreObservationIT {
@Test
void observationVectorStoreAddAndQueryOperations() {
contextRunner.run(context -> {
this.contextRunner.run(context -> {
VectorStore vectorStore = context.getBean(VectorStore.class);
TestObservationRegistry observationRegistry = context.getBean(TestObservationRegistry.class);
vectorStore.add(documents);
vectorStore.add(this.documents);
TestObservationRegistryAssert.assertThat(observationRegistry)
.doesNotHaveAnyRemainingCurrentObservation()

16
vector-stores/spring-ai-cassandra-store/pom.xml
View File

@@ -1,4 +1,20 @@
<?xml version="1.0" encoding="UTF-8"?>
<!--
~ Copyright 2023-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.
-->
<project xmlns="http://maven.apache.org/POM/4.0.0"
xmlns:xsi="http://www.w3.org/2001/XMLSchema-instance" xsi:schemaLocation="http://maven.apache.org/POM/4.0.0 http://maven.apache.org/maven-v4_0_0.xsd">
<modelVersion>4.0.0</modelVersion>

13
vector-stores/spring-ai-cassandra-store/src/main/java/org/springframework/ai/cassandra/SchemaUtil.java
View File

@@ -1,30 +1,29 @@
/*
* Copyright 2023-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
* 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.
*
* See the NOTICE file distributed with this work for additional information
* regarding copyright ownership.
*/
package org.springframework.ai.cassandra;
import java.time.Duration;
import com.datastax.oss.driver.api.core.CqlSession;
import com.datastax.oss.driver.api.core.cql.SimpleStatement;
import com.datastax.oss.driver.api.querybuilder.SchemaBuilder;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import java.time.Duration;
/**
* @author Mick Semb Wever
* @since 1.0.0

37
vector-stores/spring-ai-cassandra-store/src/main/java/org/springframework/ai/chat/memory/CassandraChatMemory.java
View File

@@ -1,11 +1,11 @@
/*
* Copyright 2024 - 2024 the original author or authors.
* Copyright 2023-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
* 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,
@@ -13,8 +13,14 @@
* See the License for the specific language governing permissions and
* limitations under the License.
*/
package org.springframework.ai.chat.memory;
import java.time.Instant;
import java.util.ArrayList;
import java.util.List;
import java.util.concurrent.atomic.AtomicLong;
import com.datastax.oss.driver.api.core.cql.BoundStatementBuilder;
import com.datastax.oss.driver.api.core.cql.PreparedStatement;
import com.datastax.oss.driver.api.core.cql.Row;
@@ -31,18 +37,13 @@ import org.springframework.ai.chat.messages.AssistantMessage;
import org.springframework.ai.chat.messages.Message;
import org.springframework.ai.chat.messages.UserMessage;
import java.time.Instant;
import java.util.ArrayList;
import java.util.List;
import java.util.concurrent.atomic.AtomicLong;
/**
* Create a CassandraChatMemory like <code>
CassandraChatMemory.create(CassandraChatMemoryConfig.builder().withTimeToLive(Duration.ofDays(1)).build());
</code>
CassandraChatMemory.create(CassandraChatMemoryConfig.builder().withTimeToLive(Duration.ofDays(1)).build());
</code>
*
* For example @see org.springframework.ai.chat.memory.CassandraChatMemory
*
*
* @author Mick Semb Wever
* @since 1.0.0
*/
@@ -54,10 +55,6 @@ public final class CassandraChatMemory implements ChatMemory {
private final PreparedStatement addUserStmt, addAssistantStmt, getStmt, deleteStmt;
public static CassandraChatMemory create(CassandraChatMemoryConfig conf) {
return new CassandraChatMemory(conf);
}
public CassandraChatMemory(CassandraChatMemoryConfig config) {
this.conf = config;
this.conf.ensureSchemaExists();
@@ -67,6 +64,10 @@ public final class CassandraChatMemory implements ChatMemory {
this.deleteStmt = prepareDeleteStmt();
}
public static CassandraChatMemory create(CassandraChatMemoryConfig conf) {
return new CassandraChatMemory(conf);
}
@Override
public void add(String conversationId, List<Message> messages) {
final AtomicLong instantSeq = new AtomicLong(Instant.now().toEpochMilli());
@@ -90,8 +91,8 @@ public final class CassandraChatMemory implements ChatMemory {
PreparedStatement stmt;
switch (msg.getMessageType()) {
case USER -> stmt = addUserStmt;
case ASSISTANT -> stmt = addAssistantStmt;
case USER -> stmt = this.addUserStmt;
case ASSISTANT -> stmt = this.addAssistantStmt;
default -> throw new IllegalArgumentException("Cant add type " + msg);
}
@@ -115,7 +116,7 @@ public final class CassandraChatMemory implements ChatMemory {
public void clear(String sessionId) {
List<Object> primaryKeys = this.conf.primaryKeyTranslator.apply(sessionId);
BoundStatementBuilder builder = deleteStmt.boundStatementBuilder();
BoundStatementBuilder builder = this.deleteStmt.boundStatementBuilder();
for (int k = 0; k < primaryKeys.size(); ++k) {
SchemaColumn keyColumn = this.conf.getPrimaryKeyColumn(k);
@@ -129,7 +130,7 @@ public final class CassandraChatMemory implements ChatMemory {
public List<Message> get(String sessionId, int lastN) {
List<Object> primaryKeys = this.conf.primaryKeyTranslator.apply(sessionId);
BoundStatementBuilder builder = getStmt.boundStatementBuilder().setInt("lastN", lastN);
BoundStatementBuilder builder = this.getStmt.boundStatementBuilder().setInt("lastN", lastN);
for (int k = 0; k < primaryKeys.size(); ++k) {
SchemaColumn keyColumn = this.conf.getPrimaryKeyColumn(k);

250
vector-stores/spring-ai-cassandra-store/src/main/java/org/springframework/ai/chat/memory/CassandraChatMemoryConfig.java
View File

@@ -1,11 +1,11 @@
/*
* Copyright 2024 - 2024 the original author or authors.
* Copyright 2023-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
* 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,
@@ -13,8 +13,16 @@
* See the License for the specific language governing permissions and
* limitations under the License.
*/
package org.springframework.ai.chat.memory;
import java.net.InetSocketAddress;
import java.time.Duration;
import java.util.List;
import java.util.Map;
import java.util.UUID;
import java.util.function.Function;
import com.datastax.oss.driver.api.core.CqlSession;
import com.datastax.oss.driver.api.core.CqlSessionBuilder;
import com.datastax.oss.driver.api.core.cql.SimpleStatement;
@@ -32,42 +40,17 @@ import com.datastax.oss.driver.api.querybuilder.schema.CreateTableStart;
import com.datastax.oss.driver.api.querybuilder.schema.CreateTableWithOptions;
import com.datastax.oss.driver.shaded.guava.common.annotations.VisibleForTesting;
import com.datastax.oss.driver.shaded.guava.common.base.Preconditions;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import org.springframework.ai.cassandra.SchemaUtil;
import java.net.InetSocketAddress;
import java.time.Duration;
import java.util.List;
import java.util.Map;
import java.util.UUID;
import java.util.function.Function;
/**
* @author Mick Semb Wever
* @since 1.0.0
*/
public final class CassandraChatMemoryConfig {
private static final Logger logger = LoggerFactory.getLogger(CassandraChatMemoryConfig.class);
record Schema(String keyspace, String table, List<SchemaColumn> partitionKeys, List<SchemaColumn> clusteringKeys) {
}
public record SchemaColumn(String name, DataType type) {
public GenericType<Object> javaType() {
return CodecRegistry.DEFAULT.codecFor(type).getJavaType();
}
}
/** Given a string sessionId, return the value for each primary key column. */
public interface SessionIdToPrimaryKeysTranslator extends Function<String, List<Object>> {
}
public static final String DEFAULT_KEYSPACE_NAME = "springframework";
public static final String DEFAULT_TABLE_NAME = "ai_chat_memory";
@@ -82,6 +65,8 @@ public final class CassandraChatMemoryConfig {
public static final String DEFAULT_USER_COLUMN_NAME = "user";
private static final Logger logger = LoggerFactory.getLogger(CassandraChatMemoryConfig.class);
final CqlSession session;
final Schema schema;
@@ -90,16 +75,12 @@ public final class CassandraChatMemoryConfig {
final String userColumn;
final SessionIdToPrimaryKeysTranslator primaryKeyTranslator;
private final Integer timeToLiveSeconds;
private final boolean disallowSchemaChanges;
final SessionIdToPrimaryKeysTranslator primaryKeyTranslator;
public static Builder builder() {
return new Builder();
}
private CassandraChatMemoryConfig(Builder builder) {
this.session = builder.session;
this.schema = new Schema(builder.keyspace, builder.table, builder.partitionKeys, builder.clusteringKeys);
@@ -110,6 +91,10 @@ public final class CassandraChatMemoryConfig {
this.primaryKeyTranslator = builder.primaryKeyTranslator;
}
public static Builder builder() {
return new Builder();
}
SchemaColumn getPrimaryKeyColumn(int index) {
return index < this.schema.partitionKeys().size() ? this.schema.partitionKeys().get(index)
: this.schema.clusteringKeys().get(index - this.schema.partitionKeys().size());
@@ -121,6 +106,113 @@ public final class CassandraChatMemoryConfig {
this.session.execute(SchemaBuilder.dropKeyspace(this.schema.keyspace).ifExists().build());
}
void ensureSchemaExists() {
if (!this.disallowSchemaChanges) {
SchemaUtil.ensureKeyspaceExists(this.session, this.schema.keyspace);
ensureTableExists();
ensureTableColumnsExist();
SchemaUtil.checkSchemaAgreement(this.session);
}
else {
checkSchemaValid();
}
}
void checkSchemaValid() {
Preconditions.checkState(this.session.getMetadata().getKeyspace(this.schema.keyspace).isPresent(),
"keyspace %s does not exist", this.schema.keyspace);
Preconditions.checkState(this.session.getMetadata()
.getKeyspace(this.schema.keyspace)
.get()
.getTable(this.schema.table)
.isPresent(), "table %s does not exist");
TableMetadata tableMetadata = this.session.getMetadata()
.getKeyspace(this.schema.keyspace)
.get()
.getTable(this.schema.table)
.get();
Preconditions.checkState(tableMetadata.getColumn(this.assistantColumn).isPresent(), "column %s does not exist",
this.assistantColumn);
Preconditions.checkState(tableMetadata.getColumn(this.userColumn).isPresent(), "column %s does not exist",
this.userColumn);
}
private void ensureTableExists() {
if (this.session.getMetadata().getKeyspace(this.schema.keyspace).get().getTable(this.schema.table).isEmpty()) {
CreateTable createTable = null;
CreateTableStart createTableStart = SchemaBuilder.createTable(this.schema.keyspace, this.schema.table)
.ifNotExists();
for (SchemaColumn partitionKey : this.schema.partitionKeys) {
createTable = (null != createTable ? createTable : createTableStart).withPartitionKey(partitionKey.name,
partitionKey.type);
}
for (SchemaColumn clusteringKey : this.schema.clusteringKeys) {
createTable = createTable.withClusteringColumn(clusteringKey.name, clusteringKey.type);
}
String lastClusteringColumn = this.schema.clusteringKeys.get(this.schema.clusteringKeys.size() - 1).name();
CreateTableWithOptions createTableWithOptions = createTable.withColumn(this.userColumn, DataTypes.TEXT)
.withClusteringOrder(lastClusteringColumn, ClusteringOrder.DESC)
// TODO replace w/ SchemaBuilder.unifiedCompactionStrategy() is available
.withOption("compaction", Map.of("class", "UnifiedCompactionStrategy"));
if (null != this.timeToLiveSeconds) {
createTableWithOptions = createTableWithOptions.withDefaultTimeToLiveSeconds(this.timeToLiveSeconds);
}
this.session.execute(createTableWithOptions.build());
}
}
private void ensureTableColumnsExist() {
TableMetadata tableMetadata = this.session.getMetadata()
.getKeyspace(this.schema.keyspace())
.get()
.getTable(this.schema.table())
.get();
boolean addAssistantColumn = tableMetadata.getColumn(this.assistantColumn).isEmpty();
boolean addUserColumn = tableMetadata.getColumn(this.userColumn).isEmpty();
if (addAssistantColumn || addUserColumn) {
AlterTableAddColumn alterTable = SchemaBuilder.alterTable(this.schema.keyspace(), this.schema.table());
if (addAssistantColumn) {
alterTable = alterTable.addColumn(this.assistantColumn, DataTypes.TEXT);
}
if (addUserColumn) {
alterTable = alterTable.addColumn(this.userColumn, DataTypes.TEXT);
}
SimpleStatement stmt = ((AlterTableAddColumnEnd) alterTable).build();
logger.debug("Executing {}", stmt.getQuery());
this.session.execute(stmt);
}
}
/** Given a string sessionId, return the value for each primary key column. */
public interface SessionIdToPrimaryKeysTranslator extends Function<String, List<Object>> {
}
record Schema(String keyspace, String table, List<SchemaColumn> partitionKeys, List<SchemaColumn> clusteringKeys) {
}
public record SchemaColumn(String name, DataType type) {
public GenericType<Object> javaType() {
return CodecRegistry.DEFAULT.codecFor(this.type).getJavaType();
}
}
public static class Builder {
private CqlSession session = null;
@@ -226,14 +318,14 @@ public final class CassandraChatMemoryConfig {
public CassandraChatMemoryConfig build() {
int primaryKeyColumns = partitionKeys.size() + clusteringKeys.size();
int primaryKeyColumns = this.partitionKeys.size() + this.clusteringKeys.size();
int primaryKeysToBind = this.primaryKeyTranslator.apply(UUID.randomUUID().toString()).size();
Preconditions.checkArgument(primaryKeyColumns == primaryKeysToBind + 1,
"The primaryKeyTranslator must always return one less element than the number of primary keys in total. The last clustering key remains undefined, expecting to be the timestamp for messages within sessionId. The sessionId can map to any primary key column (though it should map to a partition key column).");
Preconditions.checkArgument(
clusteringKeys.get(clusteringKeys.size() - 1).name().equals(DEFAULT_EXCHANGE_ID_NAME),
this.clusteringKeys.get(this.clusteringKeys.size() - 1).name().equals(DEFAULT_EXCHANGE_ID_NAME),
"last clustering key must be the exchangeIdColumn");
return new CassandraChatMemoryConfig(this);
@@ -241,92 +333,4 @@ public final class CassandraChatMemoryConfig {
}
void ensureSchemaExists() {
if (!disallowSchemaChanges) {
SchemaUtil.ensureKeyspaceExists(this.session, this.schema.keyspace);
ensureTableExists();
ensureTableColumnsExist();
SchemaUtil.checkSchemaAgreement(this.session);
}
else {
checkSchemaValid();
}
}
void checkSchemaValid() {
Preconditions.checkState(session.getMetadata().getKeyspace(this.schema.keyspace).isPresent(),
"keyspace %s does not exist", this.schema.keyspace);
Preconditions.checkState(
session.getMetadata().getKeyspace(this.schema.keyspace).get().getTable(this.schema.table).isPresent(),
"table %s does not exist");
TableMetadata tableMetadata = session.getMetadata()
.getKeyspace(this.schema.keyspace)
.get()
.getTable(this.schema.table)
.get();
Preconditions.checkState(tableMetadata.getColumn(this.assistantColumn).isPresent(), "column %s does not exist",
this.assistantColumn);
Preconditions.checkState(tableMetadata.getColumn(this.userColumn).isPresent(), "column %s does not exist",
this.userColumn);
}
private void ensureTableExists() {
if (session.getMetadata().getKeyspace(schema.keyspace).get().getTable(this.schema.table).isEmpty()) {
CreateTable createTable = null;
CreateTableStart createTableStart = SchemaBuilder.createTable(this.schema.keyspace, this.schema.table)
.ifNotExists();
for (SchemaColumn partitionKey : this.schema.partitionKeys) {
createTable = (null != createTable ? createTable : createTableStart).withPartitionKey(partitionKey.name,
partitionKey.type);
}
for (SchemaColumn clusteringKey : this.schema.clusteringKeys) {
createTable = createTable.withClusteringColumn(clusteringKey.name, clusteringKey.type);
}
String lastClusteringColumn = this.schema.clusteringKeys.get(this.schema.clusteringKeys.size() - 1).name();
CreateTableWithOptions createTableWithOptions = createTable.withColumn(this.userColumn, DataTypes.TEXT)
.withClusteringOrder(lastClusteringColumn, ClusteringOrder.DESC)
// TODO replace w/ SchemaBuilder.unifiedCompactionStrategy() is available
.withOption("compaction", Map.of("class", "UnifiedCompactionStrategy"));
if (null != this.timeToLiveSeconds) {
createTableWithOptions = createTableWithOptions.withDefaultTimeToLiveSeconds(this.timeToLiveSeconds);
}
this.session.execute(createTableWithOptions.build());
}
}
private void ensureTableColumnsExist() {
TableMetadata tableMetadata = this.session.getMetadata()
.getKeyspace(this.schema.keyspace())
.get()
.getTable(this.schema.table())
.get();
boolean addAssistantColumn = tableMetadata.getColumn(this.assistantColumn).isEmpty();
boolean addUserColumn = tableMetadata.getColumn(this.userColumn).isEmpty();
if (addAssistantColumn || addUserColumn) {
AlterTableAddColumn alterTable = SchemaBuilder.alterTable(this.schema.keyspace(), this.schema.table());
if (addAssistantColumn) {
alterTable = alterTable.addColumn(this.assistantColumn, DataTypes.TEXT);
}
if (addUserColumn) {
alterTable = alterTable.addColumn(this.userColumn, DataTypes.TEXT);
}
SimpleStatement stmt = ((AlterTableAddColumnEnd) alterTable).build();
logger.debug("Executing {}", stmt.getQuery());
this.session.execute(stmt);
}
}
}

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

@@ -1,11 +1,11 @@
/*
* Copyright 2024 - 2024 the original author or authors.
* Copyright 2023-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
* 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,
@@ -13,8 +13,15 @@
* See the License for the specific language governing permissions and
* limitations under the License.
*/
package org.springframework.ai.vectorstore;
import java.util.Collection;
import java.util.Map;
import java.util.Optional;
import java.util.function.Function;
import java.util.stream.Collectors;
import com.datastax.oss.driver.api.core.metadata.schema.ColumnMetadata;
import com.datastax.oss.driver.api.core.type.DataTypes;
import com.datastax.oss.driver.api.core.type.codec.registry.CodecRegistry;
@@ -26,12 +33,6 @@ import org.springframework.ai.vectorstore.filter.Filter.Key;
import org.springframework.ai.vectorstore.filter.Filter.Value;
import org.springframework.ai.vectorstore.filter.converter.AbstractFilterExpressionConverter;
import java.util.Collection;
import java.util.Map;
import java.util.Optional;
import java.util.function.Function;
import java.util.stream.Collectors;
/**
* Converts {@link org.springframework.ai.vectorstore.filter.Filter.Expression} into CQL
* where clauses.
@@ -49,6 +50,24 @@ class CassandraFilterExpressionConverter extends AbstractFilterExpressionConvert
.collect(Collectors.toMap((c) -> c.getName().asInternal(), Function.identity()));
}
private static void doOperand(ExpressionType type, StringBuilder context) {
switch (type) {
case EQ -> context.append(" = ");
case NE -> context.append(" != ");
case GT -> context.append(" > ");
case GTE -> context.append(" >= ");
case IN -> context.append(" IN ");
case LT -> context.append(" < ");
case LTE -> context.append(" <= ");
// TODO SAI supports collections
// reach out to mck@apache.org if you'd like these implemented
// case CONTAINS -> context.append(" CONTAINS ");
// case CONTAINS_KEY -> context.append(" CONTAINS_KEY ");
default -> throw new UnsupportedOperationException(
String.format("Expression type %s not yet implemented. Patches welcome.", type));
}
}
@Override
protected void doKey(Key key, StringBuilder context) {
String keyName = key.key();
@@ -68,24 +87,6 @@ class CassandraFilterExpressionConverter extends AbstractFilterExpressionConvert
}
}
private static void doOperand(ExpressionType type, StringBuilder context) {
switch (type) {
case EQ -> context.append(" = ");
case NE -> context.append(" != ");
case GT -> context.append(" > ");
case GTE -> context.append(" >= ");
case IN -> context.append(" IN ");
case LT -> context.append(" < ");
case LTE -> context.append(" <= ");
// TODO SAI supports collections
// reach out to mck@apache.org if you'd like these implemented
// case CONTAINS -> context.append(" CONTAINS ");
// case CONTAINS_KEY -> context.append(" CONTAINS_KEY ");
default -> throw new UnsupportedOperationException(
String.format("Expression type %s not yet implemented. Patches welcome.", type));
}
}
private void doBinaryOperation(String operator, Filter.Expression expression, StringBuilder context) {
this.convertOperand(expression.left(), context);
context.append(operator);

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

@@ -1,11 +1,11 @@
/*
* Copyright 2024 the original author or authors.
* Copyright 2023-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
* 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,
@@ -16,6 +16,17 @@
package org.springframework.ai.vectorstore;
import java.util.ArrayList;
import java.util.HashMap;
import java.util.HashSet;
import java.util.List;
import java.util.Map;
import java.util.Optional;
import java.util.Set;
import java.util.concurrent.CompletableFuture;
import java.util.concurrent.ConcurrentHashMap;
import java.util.concurrent.ConcurrentMap;
import com.datastax.oss.driver.api.core.cql.BoundStatement;
import com.datastax.oss.driver.api.core.cql.BoundStatementBuilder;
import com.datastax.oss.driver.api.core.cql.PreparedStatement;
@@ -29,9 +40,7 @@ import com.datastax.oss.driver.api.querybuilder.delete.DeleteSelection;
import com.datastax.oss.driver.api.querybuilder.insert.InsertInto;
import com.datastax.oss.driver.api.querybuilder.insert.RegularInsert;
import com.datastax.oss.driver.shaded.guava.common.base.Preconditions;
import io.micrometer.observation.ObservationRegistry;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
@@ -50,17 +59,6 @@ import org.springframework.ai.vectorstore.observation.VectorStoreObservationCont
import org.springframework.ai.vectorstore.observation.VectorStoreObservationContext.Builder;
import org.springframework.ai.vectorstore.observation.VectorStoreObservationConvention;
import java.util.ArrayList;
import java.util.HashMap;
import java.util.HashSet;
import java.util.List;
import java.util.Map;
import java.util.Optional;
import java.util.Set;
import java.util.concurrent.CompletableFuture;
import java.util.concurrent.ConcurrentHashMap;
import java.util.concurrent.ConcurrentMap;
/**
* The CassandraVectorStore is for managing and querying vector data in an Apache
* Cassandra db. It offers functionalities like adding, deleting, and performing
@@ -108,16 +106,6 @@ import java.util.concurrent.ConcurrentMap;
*/
public class CassandraVectorStore extends AbstractObservationVectorStore implements AutoCloseable {
/**
* Indexes are automatically created with COSINE. This can be changed manually via
* cqlsh
*/
public enum Similarity {
COSINE, DOT_PRODUCT, EUCLIDEAN;
}
public static final String SIMILARITY_FIELD_NAME = "similarity_score";
public static final String DRIVER_PROFILE_UPDATES = "spring-ai-updates";
@@ -128,6 +116,10 @@ public class CassandraVectorStore extends AbstractObservationVectorStore impleme
private static final Logger logger = LoggerFactory.getLogger(CassandraVectorStore.class);
private static Map<Similarity, VectorStoreSimilarityMetric> SIMILARITY_TYPE_MAPPING = Map.of(Similarity.COSINE,
VectorStoreSimilarityMetric.COSINE, Similarity.EUCLIDEAN, VectorStoreSimilarityMetric.EUCLIDEAN,
Similarity.DOT_PRODUCT, VectorStoreSimilarityMetric.DOT);
private final CassandraVectorStoreConfig conf;
private final EmbeddingModel embeddingModel;
@@ -177,6 +169,15 @@ public class CassandraVectorStore extends AbstractObservationVectorStore impleme
this.batchingStrategy = batchingStrategy;
}
private static Float[] toFloatArray(float[] embedding) {
Float[] embeddingFloat = new Float[embedding.length];
int i = 0;
for (Float d : embedding) {
embeddingFloat[i++] = d.floatValue();
}
return embeddingFloat;
}
@Override
public void doAdd(List<Document> documents) {
var futures = new CompletableFuture[documents.size()];
@@ -275,7 +276,7 @@ public class CassandraVectorStore extends AbstractObservationVectorStore impleme
}
void checkSchemaValid() {
this.conf.checkSchemaValid(embeddingModel.dimensions());
this.conf.checkSchemaValid(this.embeddingModel.dimensions());
}
private Similarity getIndexSimilarity(TableMetadata metadata) {
@@ -289,7 +290,7 @@ public class CassandraVectorStore extends AbstractObservationVectorStore impleme
private PreparedStatement prepareDeleteStatement() {
Delete stmt = null;
DeleteSelection stmtStart = QueryBuilder.deleteFrom(conf.schema.keyspace(), conf.schema.table());
DeleteSelection stmtStart = QueryBuilder.deleteFrom(this.conf.schema.keyspace(), this.conf.schema.table());
for (var c : this.conf.schema.partitionKeys()) {
stmt = (null != stmt ? stmt : stmtStart).whereColumn(c.name()).isEqualTo(QueryBuilder.bindMarker(c.name()));
@@ -344,7 +345,7 @@ public class CassandraVectorStore extends AbstractObservationVectorStore impleme
String similarityFunction = new StringBuilder("similarity_").append(this.similarity.toString().toLowerCase())
.append('(')
.append(conf.schema.embedding())
.append(this.conf.schema.embedding())
.append(",?)")
.toString();
@@ -377,15 +378,6 @@ public class CassandraVectorStore extends AbstractObservationVectorStore impleme
return this.conf.primaryKeyTranslator.apply(primaryKeyValues);
}
private static Float[] toFloatArray(float[] embedding) {
Float[] embeddingFloat = new Float[embedding.length];
int i = 0;
for (Float d : embedding) {
embeddingFloat[i++] = d.floatValue();
}
return embeddingFloat;
}
@Override
public Builder createObservationContextBuilder(String operationName) {
return VectorStoreObservationContext.builder(VectorStoreProvider.CASSANDRA.value(), operationName)
@@ -395,10 +387,6 @@ public class CassandraVectorStore extends AbstractObservationVectorStore impleme
.withSimilarityMetric(getSimilarityMetric());
}
private static Map<Similarity, VectorStoreSimilarityMetric> SIMILARITY_TYPE_MAPPING = Map.of(Similarity.COSINE,
VectorStoreSimilarityMetric.COSINE, Similarity.EUCLIDEAN, VectorStoreSimilarityMetric.EUCLIDEAN,
Similarity.DOT_PRODUCT, VectorStoreSimilarityMetric.DOT);
private String getSimilarityMetric() {
if (!SIMILARITY_TYPE_MAPPING.containsKey(this.similarity)) {
return this.similarity.name();
@@ -406,4 +394,14 @@ public class CassandraVectorStore extends AbstractObservationVectorStore impleme
return SIMILARITY_TYPE_MAPPING.get(this.similarity).value();
}
/**
* Indexes are automatically created with COSINE. This can be changed manually via
* cqlsh
*/
public enum Similarity {
COSINE, DOT_PRODUCT, EUCLIDEAN;
}
}

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

@@ -1,11 +1,11 @@
/*
* Copyright 2024 - 2024 the original author or authors.
* Copyright 2023-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
* 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,
@@ -13,8 +13,20 @@
* See the License for the specific language governing permissions and
* limitations under the License.
*/
package org.springframework.ai.vectorstore;
import java.net.InetSocketAddress;
import java.util.Collections;
import java.util.HashSet;
import java.util.List;
import java.util.Optional;
import java.util.Set;
import java.util.concurrent.Executor;
import java.util.concurrent.Executors;
import java.util.function.Function;
import java.util.stream.Stream;
import com.datastax.oss.driver.api.core.CqlSession;
import com.datastax.oss.driver.api.core.CqlSessionBuilder;
import com.datastax.oss.driver.api.core.cql.SimpleStatement;
@@ -32,23 +44,11 @@ import com.datastax.oss.driver.api.querybuilder.schema.CreateTable;
import com.datastax.oss.driver.api.querybuilder.schema.CreateTableStart;
import com.datastax.oss.driver.shaded.guava.common.annotations.VisibleForTesting;
import com.datastax.oss.driver.shaded.guava.common.base.Preconditions;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import org.springframework.lang.Nullable;
import org.springframework.ai.cassandra.SchemaUtil;
import java.net.InetSocketAddress;
import java.util.concurrent.Executor;
import java.util.concurrent.Executors;
import java.util.function.Function;
import java.util.stream.Stream;
import java.util.Collections;
import java.util.HashSet;
import java.util.List;
import java.util.Optional;
import java.util.Set;
import org.springframework.lang.Nullable;
/**
* Configuration for the Cassandra vector store.
@@ -84,51 +84,6 @@ public class CassandraVectorStoreConfig implements AutoCloseable {
private static final Logger logger = LoggerFactory.getLogger(CassandraVectorStoreConfig.class);
record Schema(String keyspace, String table, List<SchemaColumn> partitionKeys, List<SchemaColumn> clusteringKeys,
String content, String embedding, String index, Set<SchemaColumn> metadataColumns) {
}
public record SchemaColumn(String name, DataType type, SchemaColumnTags... tags) {
public SchemaColumn(String name, DataType type) {
this(name, type, new SchemaColumnTags[0]);
}
public GenericType<Object> javaType() {
return CodecRegistry.DEFAULT.codecFor(type).getJavaType();
}
public boolean indexed() {
for (SchemaColumnTags t : tags) {
if (SchemaColumnTags.INDEXED == t) {
return true;
}
}
return false;
}
}
public enum SchemaColumnTags {
INDEXED
}
/**
* Given a string document id, return the value for each primary key column.
*
* It is a requirement that an empty {@code List<Object>} returns an example formatted
* id
*/
public interface DocumentIdTranslator extends Function<String, List<Object>> {
}
/** Given a list of primary key column values, return the document id. */
public interface PrimaryKeyTranslator extends Function<List<Object>, String> {
}
final CqlSession session;
final Schema schema;
@@ -181,6 +136,232 @@ public class CassandraVectorStoreConfig implements AutoCloseable {
this.session.execute(SchemaBuilder.dropKeyspace(this.schema.keyspace).ifExists().build());
}
void ensureSchemaExists(int vectorDimension) {
if (!this.disallowSchemaChanges) {
SchemaUtil.ensureKeyspaceExists(this.session, this.schema.keyspace);
ensureTableExists(vectorDimension);
ensureTableColumnsExist(vectorDimension);
ensureIndexesExists();
SchemaUtil.checkSchemaAgreement(this.session);
}
else {
checkSchemaValid(vectorDimension);
}
}
void checkSchemaValid(int vectorDimension) {
Preconditions.checkState(this.session.getMetadata().getKeyspace(this.schema.keyspace).isPresent(),
"keyspace %s does not exist", this.schema.keyspace);
Preconditions.checkState(this.session.getMetadata()
.getKeyspace(this.schema.keyspace)
.get()
.getTable(this.schema.table)
.isPresent(), "table %s does not exist");
TableMetadata tableMetadata = this.session.getMetadata()
.getKeyspace(this.schema.keyspace)
.get()
.getTable(this.schema.table)
.get();
Preconditions.checkState(tableMetadata.getColumn(this.schema.content).isPresent(), "column %s does not exist",
this.schema.content);
Preconditions.checkState(tableMetadata.getColumn(this.schema.embedding).isPresent(), "column %s does not exist",
this.schema.embedding);
for (SchemaColumn m : this.schema.metadataColumns) {
Optional<ColumnMetadata> column = tableMetadata.getColumn(m.name());
Preconditions.checkState(column.isPresent(), "column %s does not exist", m.name());
Preconditions.checkArgument(column.get().getType().equals(m.type()),
"Mismatching type on metadata column %s of %s vs %s", m.name(), column.get().getType(), m.type());
if (m.indexed()) {
Preconditions.checkState(
tableMetadata.getIndexes().values().stream().anyMatch((i) -> i.getTarget().equals(m.name())),
"index %s does not exist", m.name());
}
}
}
private void ensureIndexesExists() {
{
SimpleStatement indexStmt = SchemaBuilder.createIndex(this.schema.index)
.ifNotExists()
.custom("StorageAttachedIndex")
.onTable(this.schema.keyspace, this.schema.table)
.andColumn(this.schema.embedding)
.build();
logger.debug("Executing {}", indexStmt.getQuery());
this.session.execute(indexStmt);
}
Stream
.concat(this.schema.partitionKeys.stream(),
Stream.concat(this.schema.clusteringKeys.stream(), this.schema.metadataColumns.stream()))
.filter((cs) -> cs.indexed())
.forEach((metadata) -> {
SimpleStatement indexStmt = SchemaBuilder.createIndex(String.format("%s_idx", metadata.name()))
.ifNotExists()
.custom("StorageAttachedIndex")
.onTable(this.schema.keyspace, this.schema.table)
.andColumn(metadata.name())
.build();
logger.debug("Executing {}", indexStmt.getQuery());
this.session.execute(indexStmt);
});
}
private void ensureTableExists(int vectorDimension) {
if (this.session.getMetadata().getKeyspace(this.schema.keyspace).get().getTable(this.schema.table).isEmpty()) {
CreateTable createTable = null;
CreateTableStart createTableStart = SchemaBuilder.createTable(this.schema.keyspace, this.schema.table)
.ifNotExists();
for (SchemaColumn partitionKey : this.schema.partitionKeys) {
createTable = (null != createTable ? createTable : createTableStart).withPartitionKey(partitionKey.name,
partitionKey.type);
}
for (SchemaColumn clusteringKey : this.schema.clusteringKeys) {
createTable = createTable.withClusteringColumn(clusteringKey.name, clusteringKey.type);
}
createTable = createTable.withColumn(this.schema.content, DataTypes.TEXT);
for (SchemaColumn metadata : this.schema.metadataColumns) {
createTable = createTable.withColumn(metadata.name(), metadata.type());
}
// https://datastax-oss.atlassian.net/browse/JAVA-3118
// .withColumn(config.embedding, new DefaultVectorType(DataTypes.FLOAT,
// vectorDimension));
StringBuilder tableStmt = new StringBuilder(createTable.asCql());
tableStmt.setLength(tableStmt.length() - 1);
tableStmt.append(',')
.append(this.schema.embedding)
.append(" vector<float,")
.append(vectorDimension)
.append(">)");
logger.debug("Executing {}", tableStmt.toString());
this.session.execute(tableStmt.toString());
}
}
private void ensureTableColumnsExist(int vectorDimension) {
TableMetadata tableMetadata = this.session.getMetadata()
.getKeyspace(this.schema.keyspace)
.get()
.getTable(this.schema.table)
.get();
Set<SchemaColumn> newColumns = new HashSet<>();
boolean addContent = tableMetadata.getColumn(this.schema.content).isEmpty();
boolean addEmbedding = tableMetadata.getColumn(this.schema.embedding).isEmpty();
for (SchemaColumn metadata : this.schema.metadataColumns) {
Optional<ColumnMetadata> column = tableMetadata.getColumn(metadata.name());
if (column.isPresent()) {
Preconditions.checkArgument(column.get().getType().equals(metadata.type()),
"Cannot change type on metadata column %s from %s to %s", metadata.name(),
column.get().getType(), metadata.type());
}
else {
newColumns.add(metadata);
}
}
if (!newColumns.isEmpty() || addContent || addEmbedding) {
AlterTableAddColumn alterTable = SchemaBuilder.alterTable(this.schema.keyspace, this.schema.table);
for (SchemaColumn metadata : newColumns) {
alterTable = alterTable.addColumn(metadata.name(), metadata.type());
}
if (addContent) {
alterTable = alterTable.addColumn(this.schema.content, DataTypes.TEXT);
}
if (addEmbedding) {
// special case for embedding column, bc JAVA-3118, as above
StringBuilder alterTableStmt = new StringBuilder(((BuildableQuery) alterTable).asCql());
if (newColumns.isEmpty() && !addContent) {
alterTableStmt.append(" ADD (");
}
else {
alterTableStmt.setLength(alterTableStmt.length() - 1);
alterTableStmt.append(',');
}
alterTableStmt.append(this.schema.embedding)
.append(" vector<float,")
.append(vectorDimension)
.append(">)");
logger.debug("Executing {}", alterTableStmt.toString());
this.session.execute(alterTableStmt.toString());
}
else {
SimpleStatement stmt = ((AlterTableAddColumnEnd) alterTable).build();
logger.debug("Executing {}", stmt.getQuery());
this.session.execute(stmt);
}
}
}
public enum SchemaColumnTags {
INDEXED
}
/**
* Given a string document id, return the value for each primary key column.
*
* It is a requirement that an empty {@code List<Object>} returns an example formatted
* id
*/
public interface DocumentIdTranslator extends Function<String, List<Object>> {
}
/** Given a list of primary key column values, return the document id. */
public interface PrimaryKeyTranslator extends Function<List<Object>, String> {
}
record Schema(String keyspace, String table, List<SchemaColumn> partitionKeys, List<SchemaColumn> clusteringKeys,
String content, String embedding, String index, Set<SchemaColumn> metadataColumns) {
}
public record SchemaColumn(String name, DataType type, SchemaColumnTags... tags) {
public SchemaColumn(String name, DataType type) {
this(name, type, new SchemaColumnTags[0]);
}
public GenericType<Object> javaType() {
return CodecRegistry.DEFAULT.codecFor(this.type).getJavaType();
}
public boolean indexed() {
for (SchemaColumnTags t : this.tags) {
if (SchemaColumnTags.INDEXED == t) {
return true;
}
}
return false;
}
}
public static class Builder {
private CqlSession session = null;
@@ -383,183 +564,4 @@ public class CassandraVectorStoreConfig implements AutoCloseable {
}
void ensureSchemaExists(int vectorDimension) {
if (!this.disallowSchemaChanges) {
SchemaUtil.ensureKeyspaceExists(this.session, this.schema.keyspace);
ensureTableExists(vectorDimension);
ensureTableColumnsExist(vectorDimension);
ensureIndexesExists();
SchemaUtil.checkSchemaAgreement(session);
}
else {
checkSchemaValid(vectorDimension);
}
}
void checkSchemaValid(int vectorDimension) {
Preconditions.checkState(this.session.getMetadata().getKeyspace(this.schema.keyspace).isPresent(),
"keyspace %s does not exist", this.schema.keyspace);
Preconditions.checkState(this.session.getMetadata()
.getKeyspace(this.schema.keyspace)
.get()
.getTable(this.schema.table)
.isPresent(), "table %s does not exist");
TableMetadata tableMetadata = this.session.getMetadata()
.getKeyspace(this.schema.keyspace)
.get()
.getTable(this.schema.table)
.get();
Preconditions.checkState(tableMetadata.getColumn(this.schema.content).isPresent(), "column %s does not exist",
this.schema.content);
Preconditions.checkState(tableMetadata.getColumn(this.schema.embedding).isPresent(), "column %s does not exist",
this.schema.embedding);
for (SchemaColumn m : this.schema.metadataColumns) {
Optional<ColumnMetadata> column = tableMetadata.getColumn(m.name());
Preconditions.checkState(column.isPresent(), "column %s does not exist", m.name());
Preconditions.checkArgument(column.get().getType().equals(m.type()),
"Mismatching type on metadata column %s of %s vs %s", m.name(), column.get().getType(), m.type());
if (m.indexed()) {
Preconditions.checkState(
tableMetadata.getIndexes().values().stream().anyMatch((i) -> i.getTarget().equals(m.name())),
"index %s does not exist", m.name());
}
}
}
private void ensureIndexesExists() {
{
SimpleStatement indexStmt = SchemaBuilder.createIndex(this.schema.index)
.ifNotExists()
.custom("StorageAttachedIndex")
.onTable(this.schema.keyspace, this.schema.table)
.andColumn(this.schema.embedding)
.build();
logger.debug("Executing {}", indexStmt.getQuery());
this.session.execute(indexStmt);
}
Stream
.concat(this.schema.partitionKeys.stream(),
Stream.concat(this.schema.clusteringKeys.stream(), this.schema.metadataColumns.stream()))
.filter((cs) -> cs.indexed())
.forEach((metadata) -> {
SimpleStatement indexStmt = SchemaBuilder.createIndex(String.format("%s_idx", metadata.name()))
.ifNotExists()
.custom("StorageAttachedIndex")
.onTable(this.schema.keyspace, this.schema.table)
.andColumn(metadata.name())
.build();
logger.debug("Executing {}", indexStmt.getQuery());
this.session.execute(indexStmt);
});
}
private void ensureTableExists(int vectorDimension) {
if (this.session.getMetadata().getKeyspace(this.schema.keyspace).get().getTable(this.schema.table).isEmpty()) {
CreateTable createTable = null;
CreateTableStart createTableStart = SchemaBuilder.createTable(this.schema.keyspace, this.schema.table)
.ifNotExists();
for (SchemaColumn partitionKey : this.schema.partitionKeys) {
createTable = (null != createTable ? createTable : createTableStart).withPartitionKey(partitionKey.name,
partitionKey.type);
}
for (SchemaColumn clusteringKey : this.schema.clusteringKeys) {
createTable = createTable.withClusteringColumn(clusteringKey.name, clusteringKey.type);
}
createTable = createTable.withColumn(this.schema.content, DataTypes.TEXT);
for (SchemaColumn metadata : this.schema.metadataColumns) {
createTable = createTable.withColumn(metadata.name(), metadata.type());
}
// https://datastax-oss.atlassian.net/browse/JAVA-3118
// .withColumn(config.embedding, new DefaultVectorType(DataTypes.FLOAT,
// vectorDimension));
StringBuilder tableStmt = new StringBuilder(createTable.asCql());
tableStmt.setLength(tableStmt.length() - 1);
tableStmt.append(',')
.append(this.schema.embedding)
.append(" vector<float,")
.append(vectorDimension)
.append(">)");
logger.debug("Executing {}", tableStmt.toString());
this.session.execute(tableStmt.toString());
}
}
private void ensureTableColumnsExist(int vectorDimension) {
TableMetadata tableMetadata = this.session.getMetadata()
.getKeyspace(this.schema.keyspace)
.get()
.getTable(this.schema.table)
.get();
Set<SchemaColumn> newColumns = new HashSet<>();
boolean addContent = tableMetadata.getColumn(this.schema.content).isEmpty();
boolean addEmbedding = tableMetadata.getColumn(this.schema.embedding).isEmpty();
for (SchemaColumn metadata : this.schema.metadataColumns) {
Optional<ColumnMetadata> column = tableMetadata.getColumn(metadata.name());
if (column.isPresent()) {
Preconditions.checkArgument(column.get().getType().equals(metadata.type()),
"Cannot change type on metadata column %s from %s to %s", metadata.name(),
column.get().getType(), metadata.type());
}
else {
newColumns.add(metadata);
}
}
if (!newColumns.isEmpty() || addContent || addEmbedding) {
AlterTableAddColumn alterTable = SchemaBuilder.alterTable(this.schema.keyspace, this.schema.table);
for (SchemaColumn metadata : newColumns) {
alterTable = alterTable.addColumn(metadata.name(), metadata.type());
}
if (addContent) {
alterTable = alterTable.addColumn(this.schema.content, DataTypes.TEXT);
}
if (addEmbedding) {
// special case for embedding column, bc JAVA-3118, as above
StringBuilder alterTableStmt = new StringBuilder(((BuildableQuery) alterTable).asCql());
if (newColumns.isEmpty() && !addContent) {
alterTableStmt.append(" ADD (");
}
else {
alterTableStmt.setLength(alterTableStmt.length() - 1);
alterTableStmt.append(',');
}
alterTableStmt.append(this.schema.embedding)
.append(" vector<float,")
.append(vectorDimension)
.append(">)");
logger.debug("Executing {}", alterTableStmt.toString());
this.session.execute(alterTableStmt.toString());
}
else {
SimpleStatement stmt = ((AlterTableAddColumnEnd) alterTable).build();
logger.debug("Executing {}", stmt.getQuery());
this.session.execute(stmt);
}
}
}
}

5
vector-stores/spring-ai-cassandra-store/src/test/java/org/springframework/ai/CassandraImage.java
View File

@@ -1,11 +1,11 @@
/*
* Copyright 2024 the original author or authors.
* Copyright 2023-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
* 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,
@@ -13,6 +13,7 @@
* See the License for the specific language governing permissions and
* limitations under the License.
*/
package org.springframework.ai;
import org.testcontainers.utility.DockerImageName;

7
vector-stores/spring-ai-cassandra-store/src/test/java/org/springframework/ai/chat/memory/CassandraChatMemoryIT.java
View File

@@ -1,11 +1,11 @@
/*
* Copyright 2024 - 2024 the original author or authors.
* Copyright 2023-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
* 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,
@@ -13,6 +13,7 @@
* See the License for the specific language governing permissions and
* limitations under the License.
*/
package org.springframework.ai.chat.memory;
import java.time.Duration;
@@ -21,11 +22,11 @@ import com.datastax.oss.driver.api.core.CqlSession;
import com.datastax.oss.driver.api.core.CqlSessionBuilder;
import org.junit.jupiter.api.Assertions;
import org.junit.jupiter.api.Test;
import org.springframework.ai.CassandraImage;
import org.testcontainers.containers.CassandraContainer;
import org.testcontainers.junit.jupiter.Container;
import org.testcontainers.junit.jupiter.Testcontainers;
import org.springframework.ai.CassandraImage;
import org.springframework.boot.SpringBootConfiguration;
import org.springframework.boot.autoconfigure.EnableAutoConfiguration;
import org.springframework.boot.autoconfigure.jdbc.DataSourceAutoConfiguration;

27
vector-stores/spring-ai-cassandra-store/src/test/java/org/springframework/ai/vectorstore/CassandraFilterExpressionConverterTests.java
View File

@@ -1,11 +1,11 @@
/*
* Copyright 2024 - 2024 the original author or authors.
* Copyright 2023-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
* 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,
@@ -13,6 +13,7 @@
* See the License for the specific language governing permissions and
* limitations under the License.
*/
package org.springframework.ai.vectorstore;
import java.util.Collection;
@@ -20,8 +21,8 @@ import java.util.HashSet;
import java.util.List;
import java.util.Set;
import com.datastax.oss.driver.api.core.metadata.schema.ColumnMetadata;
import com.datastax.oss.driver.api.core.CqlIdentifier;
import com.datastax.oss.driver.api.core.metadata.schema.ColumnMetadata;
import com.datastax.oss.driver.api.core.type.DataTypes;
import com.datastax.oss.driver.internal.core.metadata.schema.DefaultColumnMetadata;
import org.junit.jupiter.api.Assertions;
@@ -47,6 +48,16 @@ import static org.springframework.ai.vectorstore.filter.Filter.ExpressionType.OR
*/
class CassandraFilterExpressionConverterTests {
private static final CqlIdentifier T = CqlIdentifier.fromInternal("test");
private static final Collection<ColumnMetadata> COLUMNS = Set.of(
new DefaultColumnMetadata(T, T, CqlIdentifier.fromInternal("id"), DataTypes.TEXT, false),
new DefaultColumnMetadata(T, T, CqlIdentifier.fromInternal("content"), DataTypes.TEXT, false),
new DefaultColumnMetadata(T, T, CqlIdentifier.fromInternal("country"), DataTypes.TEXT, false),
new DefaultColumnMetadata(T, T, CqlIdentifier.fromInternal("genre"), DataTypes.TEXT, false),
new DefaultColumnMetadata(T, T, CqlIdentifier.fromInternal("drama"), DataTypes.TEXT, false),
new DefaultColumnMetadata(T, T, CqlIdentifier.fromInternal("year"), DataTypes.SMALLINT, false));
@Test
void testEQOnPartition() {
@@ -199,14 +210,4 @@ class CassandraFilterExpressionConverterTests {
assertThat(vectorExpr).isEqualTo("\"'country 1 2 3'\" = 'BG'");
}
private static final CqlIdentifier T = CqlIdentifier.fromInternal("test");
private static final Collection<ColumnMetadata> COLUMNS = Set.of(
new DefaultColumnMetadata(T, T, CqlIdentifier.fromInternal("id"), DataTypes.TEXT, false),
new DefaultColumnMetadata(T, T, CqlIdentifier.fromInternal("content"), DataTypes.TEXT, false),
new DefaultColumnMetadata(T, T, CqlIdentifier.fromInternal("country"), DataTypes.TEXT, false),
new DefaultColumnMetadata(T, T, CqlIdentifier.fromInternal("genre"), DataTypes.TEXT, false),
new DefaultColumnMetadata(T, T, CqlIdentifier.fromInternal("drama"), DataTypes.TEXT, false),
new DefaultColumnMetadata(T, T, CqlIdentifier.fromInternal("year"), DataTypes.SMALLINT, false));
}

182
vector-stores/spring-ai-cassandra-store/src/test/java/org/springframework/ai/vectorstore/CassandraRichSchemaVectorStoreIT.java
View File

@@ -1,11 +1,11 @@
/*
* Copyright 2024 - 2024 the original author or authors.
* Copyright 2023-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
* 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,
@@ -13,6 +13,7 @@
* See the License for the specific language governing permissions and
* limitations under the License.
*/
package org.springframework.ai.vectorstore;
import java.io.IOException;
@@ -36,12 +37,12 @@ import org.junit.jupiter.api.Assertions;
import org.junit.jupiter.api.Test;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import org.springframework.ai.CassandraImage;
import org.testcontainers.containers.CassandraContainer;
import org.testcontainers.junit.jupiter.Container;
import org.testcontainers.junit.jupiter.Testcontainers;
import org.testcontainers.shaded.org.apache.commons.lang3.RandomStringUtils;
import org.springframework.ai.CassandraImage;
import org.springframework.ai.document.Document;
import org.springframework.ai.embedding.EmbeddingModel;
import org.springframework.ai.transformers.TransformersEmbeddingModel;
@@ -91,6 +92,64 @@ class CassandraRichSchemaVectorStoreIT {
private final ApplicationContextRunner contextRunner = new ApplicationContextRunner()
.withUserConfiguration(TestApplication.class);
static CassandraVectorStoreConfig.Builder storeBuilder(ApplicationContext context,
List<SchemaColumn> columnOverrides) throws IOException {
Optional<SchemaColumn> wikiOverride = columnOverrides.stream()
.filter((f) -> "wiki".equals(f.name()))
.findFirst();
Optional<SchemaColumn> langOverride = columnOverrides.stream()
.filter((f) -> "language".equals(f.name()))
.findFirst();
Optional<SchemaColumn> titleOverride = columnOverrides.stream()
.filter((f) -> "title".equals(f.name()))
.findFirst();
Optional<SchemaColumn> chunkNoOverride = columnOverrides.stream()
.filter((f) -> "chunk_no".equals(f.name()))
.findFirst();
SchemaColumn wikiSC = wikiOverride.orElse(new SchemaColumn("wiki", DataTypes.TEXT));
SchemaColumn langSC = langOverride.orElse(new SchemaColumn("language", DataTypes.TEXT));
SchemaColumn titleSC = titleOverride.orElse(new SchemaColumn("title", DataTypes.TEXT));
SchemaColumn chunkNoSC = chunkNoOverride.orElse(new SchemaColumn("chunk_no", DataTypes.INT));
List<SchemaColumn> partitionKeys = List.of(wikiSC, langSC, titleSC);
List<SchemaColumn> clusteringKeys = List.of(chunkNoSC);
CassandraVectorStoreConfig.Builder builder = CassandraVectorStoreConfig.builder()
.withCqlSession(context.getBean(CqlSession.class))
.withKeyspaceName("test_wikidata")
.withTableName("articles")
.withPartitionKeys(partitionKeys)
.withClusteringKeys(clusteringKeys)
.withContentColumnName("body")
.withEmbeddingColumnName("all_minilm_l6_v2_embedding")
.withIndexName("all_minilm_l6_v2_ann")
.addMetadataColumns(new SchemaColumn("revision", DataTypes.INT),
new SchemaColumn("id", DataTypes.INT, CassandraVectorStoreConfig.SchemaColumnTags.INDEXED))
// this store uses '§¶' as a deliminator in the document id between db columns
// 'title' and 'chunk_no'
.withPrimaryKeyTranslator((List<Object> primaryKeys) -> {
if (primaryKeys.isEmpty()) {
return "test§¶0";
}
return format("%s§¶%s", primaryKeys.get(2), primaryKeys.get(3));
})
.withDocumentIdTranslator((id) -> {
String[] parts = id.split("§¶");
String title = parts[0];
int chunk_no = 0 < parts.length ? Integer.parseInt(parts[1]) : 0;
return List.of("simplewiki", "en", title, chunk_no);
});
return builder;
}
@Test
void ensureSchemaCreation() {
this.contextRunner.run(context -> {
@@ -157,7 +216,7 @@ class CassandraRichSchemaVectorStoreIT {
@Test
void addAndSearch() {
contextRunner.run(context -> {
this.contextRunner.run(context -> {
try (CassandraVectorStore store = createStore(context, false).store()) {
store.add(documents);
@@ -192,7 +251,7 @@ class CassandraRichSchemaVectorStoreIT {
int docsPerAdd = 12; // 128;
int rounds = 3;
contextRunner.run(context -> {
this.contextRunner.run(context -> {
try (CassandraVectorStore store = new CassandraVectorStore(
storeBuilder(context, List.of()).withFixedThreadPoolExecutorSize(nThreads).build(),
@@ -231,7 +290,7 @@ class CassandraRichSchemaVectorStoreIT {
@Test
void searchWithPartitionFilter() throws InterruptedException {
contextRunner.run(context -> {
this.contextRunner.run(context -> {
try (CassandraVectorStore store = createStore(context, false).store()) {
store.add(documents);
@@ -282,7 +341,7 @@ class CassandraRichSchemaVectorStoreIT {
@Test
void unsearchableFilters() throws InterruptedException {
contextRunner.run(context -> {
this.contextRunner.run(context -> {
try (CassandraVectorStore store = createStore(context, false).store()) {
store.add(documents);
@@ -301,7 +360,7 @@ class CassandraRichSchemaVectorStoreIT {
@Test
void searchWithFilters() throws InterruptedException {
contextRunner.run(context -> {
this.contextRunner.run(context -> {
try (CassandraVectorStore store = createStore(context, false).store()) {
store.add(documents);
@@ -366,7 +425,7 @@ class CassandraRichSchemaVectorStoreIT {
@Test
void searchWithFilterOnPrimaryKeys() throws InterruptedException {
contextRunner.run(context -> {
this.contextRunner.run(context -> {
List<SchemaColumn> overrides = List.of(
new SchemaColumn("title", DataTypes.TEXT, CassandraVectorStoreConfig.SchemaColumnTags.INDEXED),
@@ -402,7 +461,7 @@ class CassandraRichSchemaVectorStoreIT {
@Test
void documentUpdate() {
contextRunner.run(context -> {
this.contextRunner.run(context -> {
try (CassandraVectorStore store = createStore(context, false).store()) {
store.add(documents);
@@ -453,7 +512,7 @@ class CassandraRichSchemaVectorStoreIT {
@Test
void searchWithThreshold() {
contextRunner.run(context -> {
this.contextRunner.run(context -> {
try (CassandraVectorStore store = createStore(context, false).store()) {
store.add(documents);
@@ -483,27 +542,6 @@ class CassandraRichSchemaVectorStoreIT {
});
}
@SpringBootConfiguration
@EnableAutoConfiguration(exclude = { DataSourceAutoConfiguration.class })
public static class TestApplication {
@Bean
public EmbeddingModel embeddingModel() {
// default is ONNX all-MiniLM-L6-v2
return new TransformersEmbeddingModel();
}
@Bean
public CqlSession cqlSession() {
return new CqlSessionBuilder()
// comment next two lines out to connect to a local C* cluster
.addContactPoint(cassandraContainer.getContactPoint())
.withLocalDatacenter(cassandraContainer.getLocalDatacenter())
.build();
}
}
private StoreWrapper<CassandraVectorStore, CassandraVectorStoreConfig> createStore(ApplicationContext context,
boolean disallowSchemaCreation) throws IOException {
@@ -526,64 +564,6 @@ class CassandraRichSchemaVectorStoreIT {
return new StoreWrapper(new CassandraVectorStore(conf, context.getBean(EmbeddingModel.class)), conf);
}
static CassandraVectorStoreConfig.Builder storeBuilder(ApplicationContext context,
List<SchemaColumn> columnOverrides) throws IOException {
Optional<SchemaColumn> wikiOverride = columnOverrides.stream()
.filter((f) -> "wiki".equals(f.name()))
.findFirst();
Optional<SchemaColumn> langOverride = columnOverrides.stream()
.filter((f) -> "language".equals(f.name()))
.findFirst();
Optional<SchemaColumn> titleOverride = columnOverrides.stream()
.filter((f) -> "title".equals(f.name()))
.findFirst();
Optional<SchemaColumn> chunkNoOverride = columnOverrides.stream()
.filter((f) -> "chunk_no".equals(f.name()))
.findFirst();
SchemaColumn wikiSC = wikiOverride.orElse(new SchemaColumn("wiki", DataTypes.TEXT));
SchemaColumn langSC = langOverride.orElse(new SchemaColumn("language", DataTypes.TEXT));
SchemaColumn titleSC = titleOverride.orElse(new SchemaColumn("title", DataTypes.TEXT));
SchemaColumn chunkNoSC = chunkNoOverride.orElse(new SchemaColumn("chunk_no", DataTypes.INT));
List<SchemaColumn> partitionKeys = List.of(wikiSC, langSC, titleSC);
List<SchemaColumn> clusteringKeys = List.of(chunkNoSC);
CassandraVectorStoreConfig.Builder builder = CassandraVectorStoreConfig.builder()
.withCqlSession(context.getBean(CqlSession.class))
.withKeyspaceName("test_wikidata")
.withTableName("articles")
.withPartitionKeys(partitionKeys)
.withClusteringKeys(clusteringKeys)
.withContentColumnName("body")
.withEmbeddingColumnName("all_minilm_l6_v2_embedding")
.withIndexName("all_minilm_l6_v2_ann")
.addMetadataColumns(new SchemaColumn("revision", DataTypes.INT),
new SchemaColumn("id", DataTypes.INT, CassandraVectorStoreConfig.SchemaColumnTags.INDEXED))
// this store uses '§¶' as a deliminator in the document id between db columns
// 'title' and 'chunk_no'
.withPrimaryKeyTranslator((List<Object> primaryKeys) -> {
if (primaryKeys.isEmpty()) {
return "test§¶0";
}
return format("%s§¶%s", primaryKeys.get(2), primaryKeys.get(3));
})
.withDocumentIdTranslator((id) -> {
String[] parts = id.split("§¶");
String title = parts[0];
int chunk_no = 0 < parts.length ? Integer.parseInt(parts[1]) : 0;
return List.of("simplewiki", "en", title, chunk_no);
});
return builder;
}
private void executeCqlFile(ApplicationContext context, String filename) throws IOException {
logger.info("executing {}", filename);
@@ -599,7 +579,29 @@ class CassandraRichSchemaVectorStoreIT {
}
}
@SpringBootConfiguration
@EnableAutoConfiguration(exclude = { DataSourceAutoConfiguration.class })
public static class TestApplication {
@Bean
public EmbeddingModel embeddingModel() {
// default is ONNX all-MiniLM-L6-v2
return new TransformersEmbeddingModel();
}
@Bean
public CqlSession cqlSession() {
return new CqlSessionBuilder()
// comment next two lines out to connect to a local C* cluster
.addContactPoint(cassandraContainer.getContactPoint())
.withLocalDatacenter(cassandraContainer.getLocalDatacenter())
.build();
}
}
public record StoreWrapper<K, V>(K store, V conf) {
}
}

61
vector-stores/spring-ai-cassandra-store/src/test/java/org/springframework/ai/vectorstore/CassandraVectorStoreIT.java
View File

@@ -1,11 +1,11 @@
/*
* Copyright 2024 - 2024 the original author or authors.
* Copyright 2023-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
* 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,
@@ -13,6 +13,7 @@
* See the License for the specific language governing permissions and
* limitations under the License.
*/
package org.springframework.ai.vectorstore;
import java.io.IOException;
@@ -29,11 +30,11 @@ import com.datastax.oss.driver.api.core.servererrors.SyntaxError;
import com.datastax.oss.driver.api.core.type.DataTypes;
import org.junit.jupiter.api.Assertions;
import org.junit.jupiter.api.Test;
import org.springframework.ai.CassandraImage;
import org.testcontainers.containers.CassandraContainer;
import org.testcontainers.junit.jupiter.Container;
import org.testcontainers.junit.jupiter.Testcontainers;
import org.springframework.ai.CassandraImage;
import org.springframework.ai.document.Document;
import org.springframework.ai.embedding.EmbeddingModel;
import org.springframework.ai.transformers.TransformersEmbeddingModel;
@@ -84,6 +85,26 @@ class CassandraVectorStoreIT {
}
}
private static CassandraVectorStoreConfig.Builder storeBuilder(CqlSession cqlSession) {
return CassandraVectorStoreConfig.builder()
.withCqlSession(cqlSession)
.withKeyspaceName("test_" + CassandraVectorStoreConfig.DEFAULT_KEYSPACE_NAME);
}
private static CassandraVectorStore createTestStore(ApplicationContext context, SchemaColumn... metadataFields) {
CassandraVectorStoreConfig.Builder builder = storeBuilder(context.getBean(CqlSession.class))
.addMetadataColumns(metadataFields);
return createTestStore(context, builder);
}
private static CassandraVectorStore createTestStore(ApplicationContext context,
CassandraVectorStoreConfig.Builder builder) {
CassandraVectorStoreConfig conf = builder.build();
conf.dropKeyspace();
return new CassandraVectorStore(conf, context.getBean(EmbeddingModel.class));
}
@Test
void ensureBeanGetsCreated() {
this.contextRunner.run(context -> {
@@ -96,7 +117,7 @@ class CassandraVectorStoreIT {
@Test
void addAndSearch() {
contextRunner.run(context -> {
this.contextRunner.run(context -> {
try (CassandraVectorStore store = createTestStore(context, new SchemaColumn("meta1", DataTypes.TEXT),
new SchemaColumn("meta2", DataTypes.TEXT))) {
@@ -132,7 +153,7 @@ class CassandraVectorStoreIT {
@Test
void addAndSearchReturnEmbeddings() {
contextRunner.run(context -> {
this.contextRunner.run(context -> {
CassandraVectorStoreConfig.Builder builder = storeBuilder(context.getBean(CqlSession.class))
.returnEmbeddings();
@@ -168,7 +189,7 @@ class CassandraVectorStoreIT {
@Test
void searchWithPartitionFilter() throws InterruptedException {
contextRunner.run(context -> {
this.contextRunner.run(context -> {
try (CassandraVectorStore store = createTestStore(context,
new SchemaColumn("year", DataTypes.SMALLINT, SchemaColumnTags.INDEXED))) {
@@ -224,7 +245,7 @@ class CassandraVectorStoreIT {
@Test
void unsearchableFilters() throws InterruptedException {
contextRunner.run(context -> {
this.contextRunner.run(context -> {
try (CassandraVectorStore store = context.getBean(CassandraVectorStore.class)) {
var bgDocument = new Document("The World is Big and Salvation Lurks Around the Corner",
@@ -251,7 +272,7 @@ class CassandraVectorStoreIT {
@Test
void searchWithFilters() throws InterruptedException {
contextRunner.run(context -> {
this.contextRunner.run(context -> {
try (CassandraVectorStore store = createTestStore(context,
new SchemaColumn("country", DataTypes.TEXT, SchemaColumnTags.INDEXED),
@@ -314,7 +335,7 @@ class CassandraVectorStoreIT {
@Test
void documentUpdate() {
contextRunner.run(context -> {
this.contextRunner.run(context -> {
try (CassandraVectorStore store = context.getBean(CassandraVectorStore.class)) {
Document document = new Document(UUID.randomUUID().toString(), "Spring AI rocks!!",
@@ -351,7 +372,7 @@ class CassandraVectorStoreIT {
@Test
void searchWithThreshold() {
contextRunner.run(context -> {
this.contextRunner.run(context -> {
try (CassandraVectorStore store = context.getBean(CassandraVectorStore.class)) {
store.add(documents());
@@ -414,24 +435,4 @@ class CassandraVectorStoreIT {
}
private static CassandraVectorStoreConfig.Builder storeBuilder(CqlSession cqlSession) {
return CassandraVectorStoreConfig.builder()
.withCqlSession(cqlSession)
.withKeyspaceName("test_" + CassandraVectorStoreConfig.DEFAULT_KEYSPACE_NAME);
}
private static CassandraVectorStore createTestStore(ApplicationContext context, SchemaColumn... metadataFields) {
CassandraVectorStoreConfig.Builder builder = storeBuilder(context.getBean(CqlSession.class))
.addMetadataColumns(metadataFields);
return createTestStore(context, builder);
}
private static CassandraVectorStore createTestStore(ApplicationContext context,
CassandraVectorStoreConfig.Builder builder) {
CassandraVectorStoreConfig conf = builder.build();
conf.dropKeyspace();
return new CassandraVectorStore(conf, context.getBean(EmbeddingModel.class));
}
}

44
vector-stores/spring-ai-cassandra-store/src/test/java/org/springframework/ai/vectorstore/CassandraVectorStoreObservationIT.java
View File

@@ -1,11 +1,11 @@
/*
* Copyright 2024 the original author or authors.
* Copyright 2023-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
* 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,
@@ -13,16 +13,25 @@
* See the License for the specific language governing permissions and
* limitations under the License.
*/
package org.springframework.ai.vectorstore;
import static org.assertj.core.api.Assertions.assertThat;
package org.springframework.ai.vectorstore;
import java.io.IOException;
import java.nio.charset.StandardCharsets;
import java.util.List;
import java.util.Map;
import com.datastax.oss.driver.api.core.CqlSession;
import com.datastax.oss.driver.api.core.CqlSessionBuilder;
import com.datastax.oss.driver.api.core.type.DataTypes;
import io.micrometer.observation.ObservationRegistry;
import io.micrometer.observation.tck.TestObservationRegistry;
import io.micrometer.observation.tck.TestObservationRegistryAssert;
import org.junit.jupiter.api.Test;
import org.testcontainers.containers.CassandraContainer;
import org.testcontainers.junit.jupiter.Container;
import org.testcontainers.junit.jupiter.Testcontainers;
import org.springframework.ai.CassandraImage;
import org.springframework.ai.document.Document;
import org.springframework.ai.embedding.EmbeddingModel;
@@ -40,17 +49,8 @@ import org.springframework.boot.autoconfigure.EnableAutoConfiguration;
import org.springframework.boot.test.context.runner.ApplicationContextRunner;
import org.springframework.context.annotation.Bean;
import org.springframework.core.io.DefaultResourceLoader;
import org.testcontainers.containers.CassandraContainer;
import org.testcontainers.junit.jupiter.Container;
import org.testcontainers.junit.jupiter.Testcontainers;
import com.datastax.oss.driver.api.core.CqlSession;
import com.datastax.oss.driver.api.core.CqlSessionBuilder;
import com.datastax.oss.driver.api.core.type.DataTypes;
import io.micrometer.observation.ObservationRegistry;
import io.micrometer.observation.tck.TestObservationRegistry;
import io.micrometer.observation.tck.TestObservationRegistryAssert;
import static org.assertj.core.api.Assertions.assertThat;
/**
* @author Christian Tzolov
@@ -80,16 +80,22 @@ public class CassandraVectorStoreObservationIT {
}
}
private static CassandraVectorStoreConfig.Builder storeBuilder(CqlSession cqlSession) {
return CassandraVectorStoreConfig.builder()
.withCqlSession(cqlSession)
.withKeyspaceName("test_" + CassandraVectorStoreConfig.DEFAULT_KEYSPACE_NAME);
}
@Test
void observationVectorStoreAddAndQueryOperations() {
contextRunner.run(context -> {
this.contextRunner.run(context -> {
VectorStore vectorStore = context.getBean(VectorStore.class);
TestObservationRegistry observationRegistry = context.getBean(TestObservationRegistry.class);
vectorStore.add(documents);
vectorStore.add(this.documents);
TestObservationRegistryAssert.assertThat(observationRegistry)
.doesNotHaveAnyRemainingCurrentObservation()
@@ -193,10 +199,4 @@ public class CassandraVectorStoreObservationIT {
}
private static CassandraVectorStoreConfig.Builder storeBuilder(CqlSession cqlSession) {
return CassandraVectorStoreConfig.builder()
.withCqlSession(cqlSession)
.withKeyspaceName("test_" + CassandraVectorStoreConfig.DEFAULT_KEYSPACE_NAME);
}
}

5
vector-stores/spring-ai-cassandra-store/src/test/java/org/springframework/ai/vectorstore/WikiVectorStoreExample.java
View File

@@ -1,11 +1,11 @@
/*
* Copyright 2024 - 2024 the original author or authors.
* Copyright 2023-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
* 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,
@@ -13,6 +13,7 @@
* See the License for the specific language governing permissions and
* limitations under the License.
*/
package org.springframework.ai.vectorstore;
import java.util.List;

16
vector-stores/spring-ai-cassandra-store/src/test/resources/test_wiki_full_schema.cql
View File

@@ -1,3 +1,19 @@
/*
* Copyright 2023-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.
*/
CREATE KEYSPACE IF NOT EXISTS test_wikidata WITH replication = {'class': 'SimpleStrategy', 'replication_factor': 1};
CREATE TABLE IF NOT EXISTS test_wikidata.articles (

16
vector-stores/spring-ai-cassandra-store/src/test/resources/test_wiki_partial_0_schema.cql
View File

@@ -1 +1,17 @@
/*
* Copyright 2023-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.
*/
CREATE KEYSPACE IF NOT EXISTS test_wikidata WITH replication = {'class': 'SimpleStrategy', 'replication_factor': 1};

16
vector-stores/spring-ai-cassandra-store/src/test/resources/test_wiki_partial_1_schema.cql
View File

@@ -1,3 +1,19 @@
/*
* Copyright 2023-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.
*/
CREATE KEYSPACE IF NOT EXISTS test_wikidata WITH replication = {'class': 'SimpleStrategy', 'replication_factor': 1};
CREATE TABLE IF NOT EXISTS test_wikidata.articles (

16
vector-stores/spring-ai-cassandra-store/src/test/resources/test_wiki_partial_2_schema.cql
View File

@@ -1,3 +1,19 @@
/*
* Copyright 2023-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.
*/
CREATE KEYSPACE IF NOT EXISTS test_wikidata WITH replication = {'class': 'SimpleStrategy', 'replication_factor': 1};
CREATE TABLE IF NOT EXISTS test_wikidata.articles (

16
vector-stores/spring-ai-cassandra-store/src/test/resources/test_wiki_partial_3_schema.cql
View File

@@ -1,3 +1,19 @@
/*
* Copyright 2023-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.
*/
CREATE KEYSPACE IF NOT EXISTS test_wikidata WITH replication = {'class': 'SimpleStrategy', 'replication_factor': 1};
CREATE TABLE IF NOT EXISTS test_wikidata.articles (

16
vector-stores/spring-ai-cassandra-store/src/test/resources/test_wiki_partial_4_schema.cql
View File

@@ -1,3 +1,19 @@
/*
* Copyright 2023-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.
*/
CREATE KEYSPACE IF NOT EXISTS test_wikidata WITH replication = {'class': 'SimpleStrategy', 'replication_factor': 1};
CREATE TABLE IF NOT EXISTS test_wikidata.articles (

16
vector-stores/spring-ai-chroma-store/pom.xml
View File

@@ -1,4 +1,20 @@
<?xml version="1.0" encoding="UTF-8"?>
<!--
~ Copyright 2023-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.
-->
<project xmlns="http://maven.apache.org/POM/4.0.0"
xmlns:xsi="http://www.w3.org/2001/XMLSchema-instance"
xsi:schemaLocation="http://maven.apache.org/POM/4.0.0 http://maven.apache.org/maven-v4_0_0.xsd">

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

@@ -1,11 +1,11 @@
/*
* Copyright 2023 - 2024 the original author or authors.
* Copyright 2023-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
* 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,
@@ -13,16 +13,20 @@
* See the License for the specific language governing permissions and
* limitations under the License.
*/
package org.springframework.ai.chroma;
import java.util.ArrayList;
import java.util.HashMap;
import java.util.List;
import java.util.Map;
import java.util.function.Consumer;
import java.util.regex.Matcher;
import java.util.regex.Pattern;
import com.fasterxml.jackson.annotation.JsonProperty;
import com.fasterxml.jackson.core.JsonProcessingException;
import com.fasterxml.jackson.databind.ObjectMapper;
import org.springframework.ai.chroma.ChromaApi.QueryRequest.Include;
import org.springframework.core.ParameterizedTypeReference;
import org.springframework.http.HttpHeaders;
@@ -36,10 +40,6 @@ import org.springframework.web.client.HttpServerErrorException;
import org.springframework.web.client.HttpStatusCodeException;
import org.springframework.web.client.RestClient;
import com.fasterxml.jackson.annotation.JsonProperty;
import com.fasterxml.jackson.core.JsonProcessingException;
import com.fasterxml.jackson.databind.ObjectMapper;
/**
* Single-class Chroma API implementation based on the (unofficial) Chroma REST API.
*
@@ -54,10 +54,10 @@ public class ChromaApi {
// Regular expression pattern that looks for a message.
private static Pattern MESSAGE_ERROR_PATTERN = Pattern.compile("\"message\":\"(.*?)\"");
private RestClient restClient;
private final ObjectMapper objectMapper;
private RestClient restClient;
private String keyToken;
public ChromaApi(String baseUrl) {
@@ -99,164 +99,6 @@ public class ChromaApi {
return this;
}
/**
* Chroma embedding collection.
*
* @param id Collection Id.
* @param name The name of the collection.
* @param metadata Metadata associated with the collection.
*/
public record Collection(String id, String name, Map<String, Object> metadata) {
}
/**
* Request to create a new collection with the given name and metadata.
*
* @param name The name of the collection to create.
* @param metadata Optional metadata to associate with the collection.
*/
public record CreateCollectionRequest(String name, Map<String, Object> metadata) {
public CreateCollectionRequest(String name) {
this(name, new HashMap<>(Map.of("hnsw:space", "cosine")));
}
}
/**
* Add embeddings to the chroma data store.
*
* @param ids The ids of the embeddings to add.
* @param embeddings The embeddings to add.
* @param metadata The metadata to associate with the embeddings. When querying, you
* can filter on this metadata.
* @param documents The documents contents to associate with the embeddings.
*/
public record AddEmbeddingsRequest(List<String> ids, List<float[]> embeddings,
@JsonProperty("metadatas") List<Map<String, Object>> metadata, List<String> documents) {
// Convenance for adding a single embedding.
public AddEmbeddingsRequest(String id, float[] embedding, Map<String, Object> metadata, String document) {
this(List.of(id), List.of(embedding), List.of(metadata), List.of(document));
}
}
/**
* Request to delete embedding from a collection.
*
* @param ids The ids of the embeddings to delete. (Optional)
* @param where Condition to filter items to delete based on metadata values.
* (Optional)
*/
public record DeleteEmbeddingsRequest(List<String> ids, Map<String, Object> where) {
public DeleteEmbeddingsRequest(List<String> ids) {
this(ids, Map.of());
}
}
/**
* Get embeddings from a collection.
*
* @param ids IDs of the embeddings to get.
* @param where Condition to filter results based on metadata values.
* @param limit Limit on the number of collection embeddings to get.
* @param offset Offset on the embeddings to get.
* @param include A list of what to include in the results. Can contain "embeddings",
* "metadatas", "documents", "distances". Ids are always included. Defaults to
* [metadatas, documents, distances].
*/
public record GetEmbeddingsRequest(List<String> ids, Map<String, Object> where, int limit, int offset,
List<Include> include) {
public GetEmbeddingsRequest(List<String> ids) {
this(ids, Map.of(), 10, 0, Include.all);
}
public GetEmbeddingsRequest(List<String> ids, Map<String, Object> where) {
this(ids, where, 10, 0, Include.all);
}
public GetEmbeddingsRequest(List<String> ids, Map<String, Object> where, int limit, int offset) {
this(ids, where, limit, offset, Include.all);
}
}
/**
* Object containing the get embedding results.
*
* @param ids List of document ids. One for each returned document.
* @param embeddings List of document embeddings. One for each returned document.
* @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<float[]> embeddings, List<String> documents,
@JsonProperty("metadatas") List<Map<String, String>> metadata) {
}
/**
* Request to get the nResults nearest neighbor embeddings for provided
* queryEmbeddings.
*
* @param queryEmbeddings The embeddings to get the closes neighbors of.
* @param nResults The number of neighbors to return for each query_embedding or
* query_texts.
* @param where Condition to filter results based on metadata values.
* @param include A list of what to include in the results. Can contain "embeddings",
* "metadatas", "documents", "distances". Ids are always included. Defaults to
* [metadatas, documents, distances].
*/
public record QueryRequest(@JsonProperty("query_embeddings") List<float[]> queryEmbeddings,
@JsonProperty("n_results") int nResults, Map<String, Object> where, List<Include> include) {
public enum Include {
@JsonProperty("metadatas")
METADATAS,
@JsonProperty("documents")
DOCUMENTS,
@JsonProperty("distances")
DISTANCES,
@JsonProperty("embeddings")
EMBEDDINGS;
public static final List<Include> all = List.of(METADATAS, DOCUMENTS, DISTANCES, EMBEDDINGS);
}
/**
* Convenience to query for a single embedding instead of a batch of embeddings.
*/
public QueryRequest(float[] queryEmbedding, int nResults) {
this(List.of(queryEmbedding), nResults, Map.of(), Include.all);
}
public QueryRequest(float[] queryEmbedding, int nResults, Map<String, Object> where) {
this(List.of(queryEmbedding), nResults, where, Include.all);
}
}
/**
* A QueryResponse object containing the query results.
*
* @param ids List of list of document ids. One for each returned document.
* @param embeddings List of list of document embeddings. One for each returned
* document.
* @param documents List of list of document contents. One for each returned document.
* @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<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, float[] embedding, String document, Map<String, Object> metadata,
Double distances) {
}
public List<Embedding> toEmbeddingResponseList(QueryResponse queryResponse) {
List<Embedding> result = new ArrayList<>();
@@ -271,10 +113,6 @@ public class ChromaApi {
return result;
}
//
// Chroma Client API (https://docs.trychroma.com/js_reference/Client)
//
public Collection createCollection(CreateCollectionRequest createCollectionRequest) {
return this.restClient.post()
@@ -330,10 +168,6 @@ public class ChromaApi {
}
}
private static class CollectionList extends ArrayList<Collection> {
}
public List<Collection> listCollections() {
return this.restClient.get()
@@ -344,10 +178,6 @@ public class ChromaApi {
.getBody();
}
//
// Chroma Collection API (https://docs.trychroma.com/js_reference/Collection)
//
public void upsertEmbeddings(String collectionId, AddEmbeddingsRequest embedding) {
this.restClient.post()
@@ -366,6 +196,7 @@ public class ChromaApi {
.body(deleteRequest)
.retrieve()
.toEntity(new ParameterizedTypeReference<List<String>>() {
})
.getBody();
}
@@ -391,6 +222,10 @@ public class ChromaApi {
.getBody();
}
//
// Chroma Client API (https://docs.trychroma.com/js_reference/Client)
//
public GetEmbeddingResponse getEmbeddings(String collectionId, GetEmbeddingsRequest getEmbeddingsRequest) {
return this.restClient.post()
@@ -442,4 +277,181 @@ public class ChromaApi {
return "";
}
/**
* Chroma embedding collection.
*
* @param id Collection Id.
* @param name The name of the collection.
* @param metadata Metadata associated with the collection.
*/
public record Collection(String id, String name, Map<String, Object> metadata) {
}
/**
* Request to create a new collection with the given name and metadata.
*
* @param name The name of the collection to create.
* @param metadata Optional metadata to associate with the collection.
*/
public record CreateCollectionRequest(String name, Map<String, Object> metadata) {
public CreateCollectionRequest(String name) {
this(name, new HashMap<>(Map.of("hnsw:space", "cosine")));
}
}
//
// Chroma Collection API (https://docs.trychroma.com/js_reference/Collection)
//
/**
* Add embeddings to the chroma data store.
*
* @param ids The ids of the embeddings to add.
* @param embeddings The embeddings to add.
* @param metadata The metadata to associate with the embeddings. When querying, you
* can filter on this metadata.
* @param documents The documents contents to associate with the embeddings.
*/
public record AddEmbeddingsRequest(List<String> ids, List<float[]> embeddings,
@JsonProperty("metadatas") List<Map<String, Object>> metadata, List<String> documents) {
// Convenance for adding a single embedding.
public AddEmbeddingsRequest(String id, float[] embedding, Map<String, Object> metadata, String document) {
this(List.of(id), List.of(embedding), List.of(metadata), List.of(document));
}
}
/**
* Request to delete embedding from a collection.
*
* @param ids The ids of the embeddings to delete. (Optional)
* @param where Condition to filter items to delete based on metadata values.
* (Optional)
*/
public record DeleteEmbeddingsRequest(List<String> ids, Map<String, Object> where) {
public DeleteEmbeddingsRequest(List<String> ids) {
this(ids, Map.of());
}
}
/**
* Get embeddings from a collection.
*
* @param ids IDs of the embeddings to get.
* @param where Condition to filter results based on metadata values.
* @param limit Limit on the number of collection embeddings to get.
* @param offset Offset on the embeddings to get.
* @param include A list of what to include in the results. Can contain "embeddings",
* "metadatas", "documents", "distances". Ids are always included. Defaults to
* [metadatas, documents, distances].
*/
public record GetEmbeddingsRequest(List<String> ids, Map<String, Object> where, int limit, int offset,
List<Include> include) {
public GetEmbeddingsRequest(List<String> ids) {
this(ids, Map.of(), 10, 0, Include.all);
}
public GetEmbeddingsRequest(List<String> ids, Map<String, Object> where) {
this(ids, where, 10, 0, Include.all);
}
public GetEmbeddingsRequest(List<String> ids, Map<String, Object> where, int limit, int offset) {
this(ids, where, limit, offset, Include.all);
}
}
/**
* Object containing the get embedding results.
*
* @param ids List of document ids. One for each returned document.
* @param embeddings List of document embeddings. One for each returned document.
* @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<float[]> embeddings, List<String> documents,
@JsonProperty("metadatas") List<Map<String, String>> metadata) {
}
/**
* Request to get the nResults nearest neighbor embeddings for provided
* queryEmbeddings.
*
* @param queryEmbeddings The embeddings to get the closes neighbors of.
* @param nResults The number of neighbors to return for each query_embedding or
* query_texts.
* @param where Condition to filter results based on metadata values.
* @param include A list of what to include in the results. Can contain "embeddings",
* "metadatas", "documents", "distances". Ids are always included. Defaults to
* [metadatas, documents, distances].
*/
public record QueryRequest(@JsonProperty("query_embeddings") List<float[]> queryEmbeddings,
@JsonProperty("n_results") int nResults, Map<String, Object> where, List<Include> include) {
/**
* Convenience to query for a single embedding instead of a batch of embeddings.
*/
public QueryRequest(float[] queryEmbedding, int nResults) {
this(List.of(queryEmbedding), nResults, Map.of(), Include.all);
}
public QueryRequest(float[] queryEmbedding, int nResults, Map<String, Object> where) {
this(List.of(queryEmbedding), nResults, where, Include.all);
}
public enum Include {
@JsonProperty("metadatas")
METADATAS,
@JsonProperty("documents")
DOCUMENTS,
@JsonProperty("distances")
DISTANCES,
@JsonProperty("embeddings")
EMBEDDINGS;
public static final List<Include> all = List.of(METADATAS, DOCUMENTS, DISTANCES, EMBEDDINGS);
}
}
/**
* A QueryResponse object containing the query results.
*
* @param ids List of list of document ids. One for each returned document.
* @param embeddings List of list of document embeddings. One for each returned
* document.
* @param documents List of list of document contents. One for each returned document.
* @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<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, float[] embedding, String document, Map<String, Object> metadata,
Double distances) {
}
private static class CollectionList extends ArrayList<Collection> {
}
}

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

@@ -1,11 +1,11 @@
/*
* Copyright 2023 - 2024 the original author or authors.
* Copyright 2023-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
* 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,
@@ -13,6 +13,7 @@
* See the License for the specific language governing permissions and
* limitations under the License.
*/
package org.springframework.ai.vectorstore;
import org.springframework.ai.vectorstore.filter.Filter;

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

@@ -22,6 +22,11 @@ import java.util.List;
import java.util.Map;
import java.util.Optional;
import com.fasterxml.jackson.core.JsonProcessingException;
import com.fasterxml.jackson.databind.ObjectMapper;
import com.fasterxml.jackson.databind.json.JsonMapper;
import io.micrometer.observation.ObservationRegistry;
import org.springframework.ai.chroma.ChromaApi;
import org.springframework.ai.chroma.ChromaApi.AddEmbeddingsRequest;
import org.springframework.ai.chroma.ChromaApi.DeleteEmbeddingsRequest;
@@ -43,11 +48,6 @@ import org.springframework.util.Assert;
import org.springframework.util.CollectionUtils;
import org.springframework.util.StringUtils;
import com.fasterxml.jackson.core.JsonProcessingException;
import com.fasterxml.jackson.databind.ObjectMapper;
import com.fasterxml.jackson.databind.json.JsonMapper;
import io.micrometer.observation.ObservationRegistry;
/**
* {@link ChromaVectorStore} is a concrete implementation of the {@link VectorStore}
* interface. It is responsible for adding, deleting, and searching documents based on
@@ -229,4 +229,4 @@ public class ChromaVectorStore extends AbstractObservationVectorStore implements
.withFieldName(this.initializeSchema ? DISTANCE_FIELD_NAME : null);
}
}
}

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

@@ -1,11 +1,11 @@
/*
* Copyright 2024 the original author or authors.
* Copyright 2023-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
* 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,
@@ -13,6 +13,7 @@
* See the License for the specific language governing permissions and
* limitations under the License.
*/
package org.springframework.ai;
import org.testcontainers.utility.DockerImageName;

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

@@ -1,11 +1,11 @@
/*
* Copyright 2023 - 2024 the original author or authors.
* Copyright 2023-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
* 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,
@@ -13,15 +13,18 @@
* See the License for the specific language governing permissions and
* limitations under the License.
*/
package org.springframework.ai.chroma;
import static org.assertj.core.api.Assertions.assertThat;
package org.springframework.ai.chroma;
import java.util.List;
import java.util.Map;
import org.junit.jupiter.api.BeforeEach;
import org.junit.jupiter.api.Test;
import org.testcontainers.chromadb.ChromaDBContainer;
import org.testcontainers.junit.jupiter.Container;
import org.testcontainers.junit.jupiter.Testcontainers;
import org.springframework.ai.ChromaImage;
import org.springframework.ai.chroma.ChromaApi.AddEmbeddingsRequest;
import org.springframework.ai.chroma.ChromaApi.Collection;
@@ -31,9 +34,8 @@ import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.boot.SpringBootConfiguration;
import org.springframework.boot.test.context.SpringBootTest;
import org.springframework.context.annotation.Bean;
import org.testcontainers.chromadb.ChromaDBContainer;
import org.testcontainers.junit.jupiter.Container;
import org.testcontainers.junit.jupiter.Testcontainers;
import static org.assertj.core.api.Assertions.assertThat;
/**
* @author Christian Tzolov
@@ -52,57 +54,58 @@ public class ChromaApiIT {
@BeforeEach
public void beforeEach() {
chroma.listCollections().stream().forEach(c -> chroma.deleteCollection(c.name()));
this.chroma.listCollections().stream().forEach(c -> this.chroma.deleteCollection(c.name()));
}
@Test
public void testClientWithMetadata() {
Map<String, Object> metadata = Map.of("hnsw:space", "cosine", "hnsw:M", 5);
var newCollection = chroma.createCollection(new ChromaApi.CreateCollectionRequest("TestCollection", metadata));
var newCollection = this.chroma
.createCollection(new ChromaApi.CreateCollectionRequest("TestCollection", metadata));
assertThat(newCollection).isNotNull();
assertThat(newCollection.name()).isEqualTo("TestCollection");
}
@Test
public void testClient() {
var newCollection = chroma.createCollection(new ChromaApi.CreateCollectionRequest("TestCollection"));
var newCollection = this.chroma.createCollection(new ChromaApi.CreateCollectionRequest("TestCollection"));
assertThat(newCollection).isNotNull();
assertThat(newCollection.name()).isEqualTo("TestCollection");
var getCollection = chroma.getCollection("TestCollection");
var getCollection = this.chroma.getCollection("TestCollection");
assertThat(getCollection).isNotNull();
assertThat(getCollection.name()).isEqualTo("TestCollection");
assertThat(getCollection.id()).isEqualTo(newCollection.id());
List<Collection> collections = chroma.listCollections();
List<Collection> collections = this.chroma.listCollections();
assertThat(collections).hasSize(1);
assertThat(collections.get(0).id()).isEqualTo(newCollection.id());
chroma.deleteCollection(newCollection.name());
assertThat(chroma.listCollections()).hasSize(0);
this.chroma.deleteCollection(newCollection.name());
assertThat(this.chroma.listCollections()).hasSize(0);
}
@Test
public void testCollection() {
var newCollection = chroma.createCollection(new ChromaApi.CreateCollectionRequest("TestCollection"));
assertThat(chroma.countEmbeddings(newCollection.id())).isEqualTo(0);
var newCollection = this.chroma.createCollection(new ChromaApi.CreateCollectionRequest("TestCollection"));
assertThat(this.chroma.countEmbeddings(newCollection.id())).isEqualTo(0);
var addEmbeddingRequest = new AddEmbeddingsRequest(List.of("id1", "id2"),
List.of(new float[] { 1f, 1f, 1f }, new float[] { 2f, 2f, 2f }),
List.of(Map.of(), Map.of("key1", "value1", "key2", true, "key3", 23.4)),
List.of("Hello World", "Big World"));
chroma.upsertEmbeddings(newCollection.id(), addEmbeddingRequest);
this.chroma.upsertEmbeddings(newCollection.id(), addEmbeddingRequest);
var addEmbeddingRequest2 = new AddEmbeddingsRequest("id3", new float[] { 3f, 3f, 3f },
Map.of("key1", "value1", "key2", true, "key3", 23.4), "Big World");
chroma.upsertEmbeddings(newCollection.id(), addEmbeddingRequest2);
this.chroma.upsertEmbeddings(newCollection.id(), addEmbeddingRequest2);
assertThat(chroma.countEmbeddings(newCollection.id())).isEqualTo(3);
assertThat(this.chroma.countEmbeddings(newCollection.id())).isEqualTo(3);
var queryResult = chroma.queryCollection(newCollection.id(),
new QueryRequest(new float[] { 1f, 1f, 1f }, 3, chroma.where("""
var queryResult = this.chroma.queryCollection(newCollection.id(),
new QueryRequest(new float[] { 1f, 1f, 1f }, 3, this.chroma.where("""
{
"key2" : { "$eq": true }
}
@@ -111,14 +114,14 @@ public class ChromaApiIT {
assertThat(queryResult.ids().get(0)).containsExactlyInAnyOrder("id2", "id3");
// Update existing embedding.
chroma.upsertEmbeddings(newCollection.id(), new AddEmbeddingsRequest("id3", new float[] { 6f, 6f, 6f },
this.chroma.upsertEmbeddings(newCollection.id(), new AddEmbeddingsRequest("id3", new float[] { 6f, 6f, 6f },
Map.of("key1", "value2", "key2", false, "key4", 23.4), "Small World"));
var result = chroma.getEmbeddings(newCollection.id(), new GetEmbeddingsRequest(List.of("id2")));
var result = this.chroma.getEmbeddings(newCollection.id(), new GetEmbeddingsRequest(List.of("id2")));
assertThat(result.ids().get(0)).isEqualTo("id2");
queryResult = chroma.queryCollection(newCollection.id(),
new QueryRequest(new float[] { 1f, 1f, 1f }, 3, chroma.where("""
queryResult = this.chroma.queryCollection(newCollection.id(),
new QueryRequest(new float[] { 1f, 1f, 1f }, 3, this.chroma.where("""
{
"key2" : { "$eq": true }
}
@@ -130,7 +133,7 @@ public class ChromaApiIT {
@Test
public void testQueryWhere() {
var collection = chroma.createCollection(new ChromaApi.CreateCollectionRequest("TestCollection"));
var collection = this.chroma.createCollection(new ChromaApi.CreateCollectionRequest("TestCollection"));
var add1 = new AddEmbeddingsRequest("id1", new float[] { 1f, 1f, 1f },
Map.of("country", "BG", "active", true, "price", 23.4, "year", 2020),
@@ -143,24 +146,24 @@ public class ChromaApiIT {
Map.of("country", "BG", "active", false, "price", 40.1, "year", 2023),
"The World is Big and Salvation Lurks Around the Corner");
chroma.upsertEmbeddings(collection.id(), add1);
chroma.upsertEmbeddings(collection.id(), add2);
chroma.upsertEmbeddings(collection.id(), add3);
this.chroma.upsertEmbeddings(collection.id(), add1);
this.chroma.upsertEmbeddings(collection.id(), add2);
this.chroma.upsertEmbeddings(collection.id(), add3);
assertThat(chroma.countEmbeddings(collection.id())).isEqualTo(3);
assertThat(this.chroma.countEmbeddings(collection.id())).isEqualTo(3);
var queryResult = chroma.queryCollection(collection.id(), new QueryRequest(new float[] { 1f, 1f, 1f }, 3));
var queryResult = this.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");
var chromaEmbeddings = chroma.toEmbeddingResponseList(queryResult);
var chromaEmbeddings = this.chroma.toEmbeddingResponseList(queryResult);
assertThat(chromaEmbeddings).hasSize(3);
assertThat(chromaEmbeddings).hasSize(3);
queryResult = chroma.queryCollection(collection.id(),
new QueryRequest(new float[] { 1f, 1f, 1f }, 3, chroma.where("""
queryResult = this.chroma.queryCollection(collection.id(),
new QueryRequest(new float[] { 1f, 1f, 1f }, 3, this.chroma.where("""
{
"$and" : [
{"country" : { "$eq": "BG"}},
@@ -171,8 +174,8 @@ public class ChromaApiIT {
assertThat(queryResult.ids().get(0)).hasSize(2);
assertThat(queryResult.ids().get(0)).containsExactlyInAnyOrder("id1", "id3");
queryResult = chroma.queryCollection(collection.id(),
new QueryRequest(new float[] { 1f, 1f, 1f }, 3, chroma.where("""
queryResult = this.chroma.queryCollection(collection.id(),
new QueryRequest(new float[] { 1f, 1f, 1f }, 3, this.chroma.where("""
{
"$and" : [
{"country" : { "$eq": "BG"}},

20
vector-stores/spring-ai-chroma-store/src/test/java/org/springframework/ai/vectorstore/BasicAuthChromaWhereIT.java
View File

@@ -1,11 +1,11 @@
/*
* Copyright 2023 - 2024 the original author or authors.
* Copyright 2023-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
* 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,
@@ -13,14 +13,18 @@
* See the License for the specific language governing permissions and
* limitations under the License.
*/
package org.springframework.ai.vectorstore;
import static org.assertj.core.api.Assertions.assertThat;
package org.springframework.ai.vectorstore;
import java.util.List;
import java.util.Map;
import org.junit.jupiter.api.Test;
import org.testcontainers.chromadb.ChromaDBContainer;
import org.testcontainers.junit.jupiter.Container;
import org.testcontainers.junit.jupiter.Testcontainers;
import org.testcontainers.utility.MountableFile;
import org.springframework.ai.ChromaImage;
import org.springframework.ai.chroma.ChromaApi;
import org.springframework.ai.document.Document;
@@ -32,10 +36,8 @@ import org.springframework.boot.test.context.runner.ApplicationContextRunner;
import org.springframework.context.annotation.Bean;
import org.springframework.http.client.SimpleClientHttpRequestFactory;
import org.springframework.web.client.RestClient;
import org.testcontainers.chromadb.ChromaDBContainer;
import org.testcontainers.junit.jupiter.Container;
import org.testcontainers.junit.jupiter.Testcontainers;
import org.testcontainers.utility.MountableFile;
import static org.assertj.core.api.Assertions.assertThat;
/**
* ChromaDB with Basic Authentication:
@@ -68,7 +70,7 @@ public class BasicAuthChromaWhereIT {
@Test
public void withInFiltersExpressions1() {
contextRunner.run(context -> {
this.contextRunner.run(context -> {
VectorStore vectorStore = context.getBean(VectorStore.class);

44
vector-stores/spring-ai-chroma-store/src/test/java/org/springframework/ai/vectorstore/ChromaVectorStoreIT.java
View File

@@ -1,11 +1,11 @@
/*
* Copyright 2023 - 2024 the original author or authors.
* Copyright 2023-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
* 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,
@@ -13,9 +13,8 @@
* See the License for the specific language governing permissions and
* limitations under the License.
*/
package org.springframework.ai.vectorstore;
import static org.assertj.core.api.Assertions.assertThat;
package org.springframework.ai.vectorstore;
import java.util.Collections;
import java.util.List;
@@ -23,6 +22,10 @@ import java.util.Map;
import java.util.UUID;
import org.junit.jupiter.api.Test;
import org.testcontainers.chromadb.ChromaDBContainer;
import org.testcontainers.junit.jupiter.Container;
import org.testcontainers.junit.jupiter.Testcontainers;
import org.springframework.ai.ChromaImage;
import org.springframework.ai.chroma.ChromaApi;
import org.springframework.ai.document.Document;
@@ -34,9 +37,8 @@ import org.springframework.boot.test.context.runner.ApplicationContextRunner;
import org.springframework.context.annotation.Bean;
import org.springframework.http.client.SimpleClientHttpRequestFactory;
import org.springframework.web.client.RestClient;
import org.testcontainers.chromadb.ChromaDBContainer;
import org.testcontainers.junit.jupiter.Container;
import org.testcontainers.junit.jupiter.Testcontainers;
import static org.assertj.core.api.Assertions.assertThat;
/**
* @author Christian Tzolov
@@ -49,6 +51,10 @@ public class ChromaVectorStoreIT {
@Container
static ChromaDBContainer chromaContainer = new ChromaDBContainer(ChromaImage.DEFAULT_IMAGE);
private final ApplicationContextRunner contextRunner = new ApplicationContextRunner()
.withUserConfiguration(TestApplication.class)
.withPropertyValues("spring.ai.openai.apiKey=" + System.getenv("OPENAI_API_KEY"));
List<Document> documents = List.of(
new Document("Spring AI rocks!! Spring AI rocks!! Spring AI rocks!! Spring AI rocks!! Spring AI rocks!!",
Collections.singletonMap("meta1", "meta1")),
@@ -57,29 +63,25 @@ public class ChromaVectorStoreIT {
"Great Depression Great Depression Great Depression Great Depression Great Depression Great Depression",
Collections.singletonMap("meta2", "meta2")));
private final ApplicationContextRunner contextRunner = new ApplicationContextRunner()
.withUserConfiguration(TestApplication.class)
.withPropertyValues("spring.ai.openai.apiKey=" + System.getenv("OPENAI_API_KEY"));
@Test
public void addAndSearch() {
contextRunner.run(context -> {
this.contextRunner.run(context -> {
VectorStore vectorStore = context.getBean(VectorStore.class);
vectorStore.add(documents);
vectorStore.add(this.documents);
List<Document> results = vectorStore.similaritySearch(SearchRequest.query("Great").withTopK(1));
assertThat(results).hasSize(1);
Document resultDoc = results.get(0);
assertThat(resultDoc.getId()).isEqualTo(documents.get(2).getId());
assertThat(resultDoc.getId()).isEqualTo(this.documents.get(2).getId());
assertThat(resultDoc.getContent()).isEqualTo(
"Great Depression Great Depression Great Depression Great Depression Great Depression Great Depression");
assertThat(resultDoc.getMetadata()).containsKeys("meta2", "distance");
// Remove all documents from the store
vectorStore.delete(documents.stream().map(doc -> doc.getId()).toList());
vectorStore.delete(this.documents.stream().map(doc -> doc.getId()).toList());
List<Document> results2 = vectorStore.similaritySearch(SearchRequest.query("Great").withTopK(1));
assertThat(results2).hasSize(0);
@@ -89,7 +91,7 @@ public class ChromaVectorStoreIT {
@Test
public void addAndSearchWithFilters() {
contextRunner.run(context -> {
this.contextRunner.run(context -> {
VectorStore vectorStore = context.getBean(VectorStore.class);
@@ -129,7 +131,7 @@ public class ChromaVectorStoreIT {
public void documentUpdateTest() {
// Note ,using OpenAI to calculate embeddings
contextRunner.run(context -> {
this.contextRunner.run(context -> {
VectorStore vectorStore = context.getBean(VectorStore.class);
@@ -170,11 +172,11 @@ public class ChromaVectorStoreIT {
@Test
public void searchThresholdTest() {
contextRunner.run(context -> {
this.contextRunner.run(context -> {
VectorStore vectorStore = context.getBean(VectorStore.class);
vectorStore.add(documents);
vectorStore.add(this.documents);
var request = SearchRequest.query("Great").withTopK(5);
List<Document> fullResult = vectorStore.similaritySearch(request.withSimilarityThresholdAll());
@@ -189,14 +191,14 @@ public class ChromaVectorStoreIT {
assertThat(results).hasSize(1);
Document resultDoc = results.get(0);
assertThat(resultDoc.getId()).isEqualTo(documents.get(2).getId());
assertThat(resultDoc.getId()).isEqualTo(this.documents.get(2).getId());
assertThat(resultDoc.getContent()).isEqualTo(
"Great Depression Great Depression Great Depression Great Depression Great Depression Great Depression");
assertThat(resultDoc.getMetadata()).containsKey("meta2");
assertThat(resultDoc.getMetadata()).containsKey("distance");
// Remove all documents from the store
vectorStore.delete(documents.stream().map(doc -> doc.getId()).toList());
vectorStore.delete(this.documents.stream().map(doc -> doc.getId()).toList());
});
}

25
vector-stores/spring-ai-chroma-store/src/test/java/org/springframework/ai/vectorstore/ChromaVectorStoreObservationIT.java
View File

@@ -1,11 +1,11 @@
/*
* Copyright 2024 the original author or authors.
* Copyright 2023-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
* 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,
@@ -13,16 +13,22 @@
* See the License for the specific language governing permissions and
* limitations under the License.
*/
package org.springframework.ai.vectorstore;
import static org.assertj.core.api.Assertions.assertThat;
package org.springframework.ai.vectorstore;
import java.io.IOException;
import java.nio.charset.StandardCharsets;
import java.util.List;
import java.util.Map;
import io.micrometer.observation.ObservationRegistry;
import io.micrometer.observation.tck.TestObservationRegistry;
import io.micrometer.observation.tck.TestObservationRegistryAssert;
import org.junit.jupiter.api.Test;
import org.testcontainers.chromadb.ChromaDBContainer;
import org.testcontainers.junit.jupiter.Container;
import org.testcontainers.junit.jupiter.Testcontainers;
import org.springframework.ai.ChromaImage;
import org.springframework.ai.chroma.ChromaApi;
import org.springframework.ai.document.Document;
@@ -42,13 +48,8 @@ import org.springframework.context.annotation.Bean;
import org.springframework.core.io.DefaultResourceLoader;
import org.springframework.http.client.SimpleClientHttpRequestFactory;
import org.springframework.web.client.RestClient;
import org.testcontainers.chromadb.ChromaDBContainer;
import org.testcontainers.junit.jupiter.Container;
import org.testcontainers.junit.jupiter.Testcontainers;
import io.micrometer.observation.ObservationRegistry;
import io.micrometer.observation.tck.TestObservationRegistry;
import io.micrometer.observation.tck.TestObservationRegistryAssert;
import static org.assertj.core.api.Assertions.assertThat;
/**
* @author Christian Tzolov
@@ -81,13 +82,13 @@ public class ChromaVectorStoreObservationIT {
@Test
void observationVectorStoreAddAndQueryOperations() {
contextRunner.run(context -> {
this.contextRunner.run(context -> {
ChromaVectorStore vectorStore = context.getBean(ChromaVectorStore.class);
TestObservationRegistry observationRegistry = context.getBean(TestObservationRegistry.class);
vectorStore.add(documents);
vectorStore.add(this.documents);
TestObservationRegistryAssert.assertThat(observationRegistry)
.doesNotHaveAnyRemainingCurrentObservation()

20
vector-stores/spring-ai-chroma-store/src/test/java/org/springframework/ai/vectorstore/TokenSecuredChromaWhereIT.java
View File

@@ -1,11 +1,11 @@
/*
* Copyright 2023 - 2024 the original author or authors.
* Copyright 2023-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
* 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,
@@ -13,14 +13,17 @@
* See the License for the specific language governing permissions and
* limitations under the License.
*/
package org.springframework.ai.vectorstore;
import static org.assertj.core.api.Assertions.assertThat;
package org.springframework.ai.vectorstore;
import java.util.List;
import java.util.Map;
import org.junit.jupiter.api.Test;
import org.testcontainers.chromadb.ChromaDBContainer;
import org.testcontainers.junit.jupiter.Container;
import org.testcontainers.junit.jupiter.Testcontainers;
import org.springframework.ai.ChromaImage;
import org.springframework.ai.chroma.ChromaApi;
import org.springframework.ai.document.Document;
@@ -32,9 +35,8 @@ import org.springframework.boot.test.context.runner.ApplicationContextRunner;
import org.springframework.context.annotation.Bean;
import org.springframework.http.client.SimpleClientHttpRequestFactory;
import org.springframework.web.client.RestClient;
import org.testcontainers.chromadb.ChromaDBContainer;
import org.testcontainers.junit.jupiter.Container;
import org.testcontainers.junit.jupiter.Testcontainers;
import static org.assertj.core.api.Assertions.assertThat;
/**
* ChromaDB with static API Token Authentication:
@@ -69,7 +71,7 @@ public class TokenSecuredChromaWhereIT {
@Test
public void withInFiltersExpressions1() {
contextRunner.run(context -> {
this.contextRunner.run(context -> {
VectorStore vectorStore = context.getBean(VectorStore.class);
@@ -93,7 +95,7 @@ public class TokenSecuredChromaWhereIT {
@Test
public void withInFiltersExpressions() {
contextRunner.run(context -> {
this.contextRunner.run(context -> {
VectorStore vectorStore = context.getBean(VectorStore.class);

16
vector-stores/spring-ai-elasticsearch-store/pom.xml
View File

@@ -1,4 +1,20 @@
<?xml version="1.0" encoding="UTF-8"?>
<!--
~ Copyright 2023-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.
-->
<project xmlns="http://maven.apache.org/POM/4.0.0"
xmlns:xsi="http://www.w3.org/2001/XMLSchema-instance" xsi:schemaLocation="http://maven.apache.org/POM/4.0.0 http://maven.apache.org/maven-v4_0_0.xsd">
<modelVersion>4.0.0</modelVersion>

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

@@ -1,11 +1,11 @@
/*
* Copyright 2023 - 2024 the original author or authors.
* Copyright 2023-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
* 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,
@@ -13,12 +13,8 @@
* See the License for the specific language governing permissions and
* limitations under the License.
*/
package org.springframework.ai.vectorstore;
import org.springframework.ai.vectorstore.filter.Filter;
import org.springframework.ai.vectorstore.filter.Filter.Expression;
import org.springframework.ai.vectorstore.filter.Filter.Key;
import org.springframework.ai.vectorstore.filter.converter.AbstractFilterExpressionConverter;
package org.springframework.ai.vectorstore;
import java.text.ParseException;
import java.text.SimpleDateFormat;
@@ -27,6 +23,11 @@ import java.util.List;
import java.util.TimeZone;
import java.util.regex.Pattern;
import org.springframework.ai.vectorstore.filter.Filter;
import org.springframework.ai.vectorstore.filter.Filter.Expression;
import org.springframework.ai.vectorstore.filter.Filter.Key;
import org.springframework.ai.vectorstore.filter.converter.AbstractFilterExpressionConverter;
/**
* ElasticsearchAiSearchFilterExpressionConverter is a class that converts
* Filter.Expression objects into Elasticsearch query string representation. It extends

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

@@ -5,7 +5,7 @@
* 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
* 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,
@@ -13,9 +13,8 @@
* See the License for the specific language governing permissions and
* limitations under the License.
*/
package org.springframework.ai.vectorstore;
import static java.lang.Math.sqrt;
package org.springframework.ai.vectorstore;
import java.io.IOException;
import java.util.List;
@@ -24,10 +23,22 @@ import java.util.Objects;
import java.util.Optional;
import java.util.stream.Collectors;
import co.elastic.clients.elasticsearch.ElasticsearchClient;
import co.elastic.clients.elasticsearch.core.BulkRequest;
import co.elastic.clients.elasticsearch.core.BulkResponse;
import co.elastic.clients.elasticsearch.core.SearchResponse;
import co.elastic.clients.elasticsearch.core.bulk.BulkResponseItem;
import co.elastic.clients.elasticsearch.core.search.Hit;
import co.elastic.clients.json.jackson.JacksonJsonpMapper;
import co.elastic.clients.transport.Version;
import co.elastic.clients.transport.rest_client.RestClientTransport;
import com.fasterxml.jackson.databind.DeserializationFeature;
import com.fasterxml.jackson.databind.ObjectMapper;
import io.micrometer.observation.ObservationRegistry;
import org.elasticsearch.client.RestClient;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import org.springframework.ai.document.Document;
import org.springframework.ai.embedding.BatchingStrategy;
import org.springframework.ai.embedding.EmbeddingModel;
@@ -45,18 +56,7 @@ import org.springframework.ai.vectorstore.observation.VectorStoreObservationConv
import org.springframework.beans.factory.InitializingBean;
import org.springframework.util.Assert;
import com.fasterxml.jackson.databind.DeserializationFeature;
import com.fasterxml.jackson.databind.ObjectMapper;
import co.elastic.clients.elasticsearch.ElasticsearchClient;
import co.elastic.clients.elasticsearch.core.BulkRequest;
import co.elastic.clients.elasticsearch.core.BulkResponse;
import co.elastic.clients.elasticsearch.core.SearchResponse;
import co.elastic.clients.elasticsearch.core.bulk.BulkResponseItem;
import co.elastic.clients.elasticsearch.core.search.Hit;
import co.elastic.clients.json.jackson.JacksonJsonpMapper;
import co.elastic.clients.transport.rest_client.RestClientTransport;
import io.micrometer.observation.ObservationRegistry;
import static java.lang.Math.sqrt;
/**
* The ElasticsearchVectorStore class implements the VectorStore interface and provides
@@ -79,6 +79,10 @@ public class ElasticsearchVectorStore extends AbstractObservationVectorStore imp
private static final Logger logger = LoggerFactory.getLogger(ElasticsearchVectorStore.class);
private static Map<SimilarityFunction, VectorStoreSimilarityMetric> SIMILARITY_TYPE_MAPPING = Map.of(
SimilarityFunction.cosine, VectorStoreSimilarityMetric.COSINE, SimilarityFunction.l2_norm,
VectorStoreSimilarityMetric.EUCLIDEAN, SimilarityFunction.dot_product, VectorStoreSimilarityMetric.DOT);
private final EmbeddingModel embeddingModel;
private final ElasticsearchClient elasticsearchClient;
@@ -176,14 +180,14 @@ public class ElasticsearchVectorStore extends AbstractObservationVectorStore imp
try {
float threshold = (float) searchRequest.getSimilarityThreshold();
// reverting l2_norm distance to its original value
if (options.getSimilarity().equals(SimilarityFunction.l2_norm)) {
if (this.options.getSimilarity().equals(SimilarityFunction.l2_norm)) {
threshold = 1 - threshold;
}
final float finalThreshold = threshold;
float[] vectors = this.embeddingModel.embed(searchRequest.getQuery());
SearchResponse<Document> res = elasticsearchClient.search(
sr -> sr.index(options.getIndexName())
SearchResponse<Document> res = this.elasticsearchClient.search(
sr -> sr.index(this.options.getIndexName())
.knn(knn -> knn.queryVector(EmbeddingUtils.toList(vectors))
.similarity(finalThreshold)
.k((long) searchRequest.getTopK())
@@ -215,7 +219,7 @@ public class ElasticsearchVectorStore extends AbstractObservationVectorStore imp
// more info on score/distance calculation
// https://www.elastic.co/guide/en/elasticsearch/reference/current/knn-search.html#knn-similarity-search
private float calculateDistance(Float score) {
switch (options.getSimilarity()) {
switch (this.options.getSimilarity()) {
case l2_norm:
// the returned value of l2_norm is the opposite of the other functions
// (closest to zero means more accurate), so to make it consistent
@@ -230,7 +234,7 @@ public class ElasticsearchVectorStore extends AbstractObservationVectorStore imp
public boolean indexExists() {
try {
return this.elasticsearchClient.indices().exists(ex -> ex.index(options.getIndexName())).value();
return this.elasticsearchClient.indices().exists(ex -> ex.index(this.options.getIndexName())).value();
}
catch (IOException e) {
throw new RuntimeException(e);
@@ -240,9 +244,10 @@ public class ElasticsearchVectorStore extends AbstractObservationVectorStore imp
private void createIndexMapping() {
try {
this.elasticsearchClient.indices()
.create(cr -> cr.index(options.getIndexName())
.mappings(map -> map.properties("embedding", p -> p.denseVector(
dv -> dv.similarity(options.getSimilarity().toString()).dims(options.getDimensions())))));
.create(cr -> cr.index(this.options.getIndexName())
.mappings(map -> map.properties("embedding",
p -> p.denseVector(dv -> dv.similarity(this.options.getSimilarity().toString())
.dims(this.options.getDimensions())))));
}
catch (IOException e) {
throw new RuntimeException(e);
@@ -267,10 +272,6 @@ public class ElasticsearchVectorStore extends AbstractObservationVectorStore imp
.withSimilarityMetric(getSimilarityMetric());
}
private static Map<SimilarityFunction, VectorStoreSimilarityMetric> SIMILARITY_TYPE_MAPPING = Map.of(
SimilarityFunction.cosine, VectorStoreSimilarityMetric.COSINE, SimilarityFunction.l2_norm,
VectorStoreSimilarityMetric.EUCLIDEAN, SimilarityFunction.dot_product, VectorStoreSimilarityMetric.DOT);
private String getSimilarityMetric() {
if (!SIMILARITY_TYPE_MAPPING.containsKey(this.options.getSimilarity())) {
return this.options.getSimilarity().name();

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

@@ -1,11 +1,11 @@
/*
* Copyright 2023 - 2024 the original author or authors.
* Copyright 2023-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
* 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,
@@ -13,6 +13,7 @@
* See the License for the specific language governing permissions and
* limitations under the License.
*/
package org.springframework.ai.vectorstore;
/**
@@ -40,7 +41,7 @@ public class ElasticsearchVectorStoreOptions {
private SimilarityFunction similarity = SimilarityFunction.cosine;
public String getIndexName() {
return indexName;
return this.indexName;
}
public void setIndexName(String indexName) {
@@ -48,7 +49,7 @@ public class ElasticsearchVectorStoreOptions {
}
public int getDimensions() {
return dimensions;
return this.dimensions;
}
public void setDimensions(int dims) {
@@ -56,7 +57,7 @@ public class ElasticsearchVectorStoreOptions {
}
public SimilarityFunction getSimilarity() {
return similarity;
return this.similarity;
}
public void setSimilarity(SimilarityFunction similarity) {

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

@@ -1,3 +1,19 @@
/*
* Copyright 2023-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.vectorstore;
/**

36
vector-stores/spring-ai-elasticsearch-store/src/test/java/org/springframework/ai/vectorstore/ElasticsearchAiSearchFilterExpressionConverterTest.java
View File

@@ -1,11 +1,11 @@
/*
* Copyright 2023 - 2024 the original author or authors.
* Copyright 2023-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
* 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,
@@ -13,15 +13,17 @@
* See the License for the specific language governing permissions and
* limitations under the License.
*/
package org.springframework.ai.vectorstore;
import org.junit.jupiter.api.Test;
import org.springframework.ai.vectorstore.filter.Filter;
import org.springframework.ai.vectorstore.filter.FilterExpressionConverter;
package org.springframework.ai.vectorstore;
import java.util.Date;
import java.util.List;
import org.junit.jupiter.api.Test;
import org.springframework.ai.vectorstore.filter.Filter;
import org.springframework.ai.vectorstore.filter.FilterExpressionConverter;
import static org.assertj.core.api.Assertions.assertThat;
import static org.springframework.ai.vectorstore.filter.Filter.ExpressionType.AND;
import static org.springframework.ai.vectorstore.filter.Filter.ExpressionType.EQ;
@@ -38,25 +40,25 @@ class ElasticsearchAiSearchFilterExpressionConverterTest {
@Test
public void testDate() {
String vectorExpr = converter.convertExpression(new Filter.Expression(EQ, new Filter.Key("activationDate"),
String vectorExpr = this.converter.convertExpression(new Filter.Expression(EQ, new Filter.Key("activationDate"),
new Filter.Value(new Date(1704637752148L))));
assertThat(vectorExpr).isEqualTo("metadata.activationDate:2024-01-07T14:29:12Z");
vectorExpr = converter.convertExpression(
vectorExpr = this.converter.convertExpression(
new Filter.Expression(EQ, new Filter.Key("activationDate"), new Filter.Value("1970-01-01T00:00:02Z")));
assertThat(vectorExpr).isEqualTo("metadata.activationDate:1970-01-01T00:00:02Z");
}
@Test
public void testEQ() {
String vectorExpr = converter
String vectorExpr = this.converter
.convertExpression(new Filter.Expression(EQ, new Filter.Key("country"), new Filter.Value("BG")));
assertThat(vectorExpr).isEqualTo("metadata.country:BG");
}
@Test
public void tesEqAndGte() {
String vectorExpr = converter.convertExpression(new Filter.Expression(AND,
String vectorExpr = this.converter.convertExpression(new Filter.Expression(AND,
new Filter.Expression(EQ, new Filter.Key("genre"), new Filter.Value("drama")),
new Filter.Expression(GTE, new Filter.Key("year"), new Filter.Value(2020))));
assertThat(vectorExpr).isEqualTo("metadata.genre:drama AND metadata.year:>=2020");
@@ -64,14 +66,14 @@ class ElasticsearchAiSearchFilterExpressionConverterTest {
@Test
public void tesIn() {
String vectorExpr = converter.convertExpression(new Filter.Expression(IN, new Filter.Key("genre"),
String vectorExpr = this.converter.convertExpression(new Filter.Expression(IN, new Filter.Key("genre"),
new Filter.Value(List.of("comedy", "documentary", "drama"))));
assertThat(vectorExpr).isEqualTo("(metadata.genre:comedy OR documentary OR drama)");
}
@Test
public void testNe() {
String vectorExpr = converter.convertExpression(
String vectorExpr = this.converter.convertExpression(
new Filter.Expression(OR, new Filter.Expression(GTE, new Filter.Key("year"), new Filter.Value(2020)),
new Filter.Expression(AND,
new Filter.Expression(EQ, new Filter.Key("country"), new Filter.Value("BG")),
@@ -81,7 +83,7 @@ class ElasticsearchAiSearchFilterExpressionConverterTest {
@Test
public void testGroup() {
String vectorExpr = converter.convertExpression(new Filter.Expression(AND,
String vectorExpr = this.converter.convertExpression(new Filter.Expression(AND,
new Filter.Group(new Filter.Expression(OR,
new Filter.Expression(GTE, new Filter.Key("year"), new Filter.Value(2020)),
new Filter.Expression(EQ, new Filter.Key("country"), new Filter.Value("BG")))),
@@ -92,7 +94,7 @@ class ElasticsearchAiSearchFilterExpressionConverterTest {
@Test
public void tesBoolean() {
String vectorExpr = converter.convertExpression(new Filter.Expression(AND,
String vectorExpr = this.converter.convertExpression(new Filter.Expression(AND,
new Filter.Expression(AND, new Filter.Expression(EQ, new Filter.Key("isOpen"), new Filter.Value(true)),
new Filter.Expression(GTE, new Filter.Key("year"), new Filter.Value(2020))),
new Filter.Expression(IN, new Filter.Key("country"), new Filter.Value(List.of("BG", "NL", "US")))));
@@ -103,7 +105,7 @@ class ElasticsearchAiSearchFilterExpressionConverterTest {
@Test
public void testDecimal() {
String vectorExpr = converter.convertExpression(new Filter.Expression(AND,
String vectorExpr = this.converter.convertExpression(new Filter.Expression(AND,
new Filter.Expression(GTE, new Filter.Key("temperature"), new Filter.Value(-15.6)),
new Filter.Expression(LTE, new Filter.Key("temperature"), new Filter.Value(20.13))));
@@ -112,11 +114,11 @@ class ElasticsearchAiSearchFilterExpressionConverterTest {
@Test
public void testComplexIdentifiers() {
String vectorExpr = converter
String vectorExpr = this.converter
.convertExpression(new Filter.Expression(EQ, new Filter.Key("\"country 1 2 3\""), new Filter.Value("BG")));
assertThat(vectorExpr).isEqualTo("metadata.country 1 2 3:BG");
vectorExpr = converter
vectorExpr = this.converter
.convertExpression(new Filter.Expression(EQ, new Filter.Key("'country 1 2 3'"), new Filter.Value("BG")));
assertThat(vectorExpr).isEqualTo("metadata.country 1 2 3:BG");
}

5
vector-stores/spring-ai-elasticsearch-store/src/test/java/org/springframework/ai/vectorstore/ElasticsearchImage.java
View File

@@ -1,11 +1,11 @@
/*
* Copyright 2024 the original author or authors.
* Copyright 2023-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
* 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,
@@ -13,6 +13,7 @@
* See the License for the specific language governing permissions and
* limitations under the License.
*/
package org.springframework.ai.vectorstore;
import org.testcontainers.utility.DockerImageName;

21
vector-stores/spring-ai-elasticsearch-store/src/test/java/org/springframework/ai/vectorstore/ElasticsearchVectorStoreIT.java
View File

@@ -1,11 +1,11 @@
/*
* Copyright 2023 - 2024 the original author or authors.
* Copyright 2023-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
* 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,
@@ -13,6 +13,7 @@
* See the License for the specific language governing permissions and
* limitations under the License.
*/
package org.springframework.ai.vectorstore;
import java.io.IOException;
@@ -121,7 +122,7 @@ class ElasticsearchVectorStoreIT {
assertThat(stats.total().docs().count()).isEqualTo(0L);
vectorStore.add(documents);
vectorStore.add(this.documents);
elasticsearchClient.indices().refresh();
stats = elasticsearchClient.indices()
.stats(s -> s.index("spring-ai-document-index"))
@@ -148,7 +149,7 @@ class ElasticsearchVectorStoreIT {
ElasticsearchVectorStore vectorStore = context.getBean("vectorStore_" + similarityFunction,
ElasticsearchVectorStore.class);
vectorStore.add(documents);
vectorStore.add(this.documents);
Awaitility.await()
.until(() -> vectorStore
@@ -160,14 +161,14 @@ class ElasticsearchVectorStoreIT {
assertThat(results).hasSize(1);
Document resultDoc = results.get(0);
assertThat(resultDoc.getId()).isEqualTo(documents.get(2).getId());
assertThat(resultDoc.getId()).isEqualTo(this.documents.get(2).getId());
assertThat(resultDoc.getContent()).contains("The Great Depression (1929–1939) was an economic shock");
assertThat(resultDoc.getMetadata()).hasSize(2);
assertThat(resultDoc.getMetadata()).containsKey("meta2");
assertThat(resultDoc.getMetadata()).containsKey("distance");
// Remove all documents from the store
vectorStore.delete(documents.stream().map(Document::getId).toList());
vectorStore.delete(this.documents.stream().map(Document::getId).toList());
Awaitility.await()
.until(() -> vectorStore
@@ -266,7 +267,7 @@ class ElasticsearchVectorStoreIT {
assertThat(results.get(0).getId()).isEqualTo(bgDocument2.getId());
// Remove all documents from the store
vectorStore.delete(documents.stream().map(Document::getId).toList());
vectorStore.delete(this.documents.stream().map(Document::getId).toList());
Awaitility.await()
.until(() -> vectorStore.similaritySearch(SearchRequest.query("The World").withTopK(1)), hasSize(0));
@@ -334,7 +335,7 @@ class ElasticsearchVectorStoreIT {
ElasticsearchVectorStore vectorStore = context.getBean("vectorStore_" + similarityFunction,
ElasticsearchVectorStore.class);
vectorStore.add(documents);
vectorStore.add(this.documents);
SearchRequest query = SearchRequest.query("Great Depression").withTopK(50).withSimilarityThresholdAll();
@@ -353,13 +354,13 @@ class ElasticsearchVectorStoreIT {
assertThat(results).hasSize(1);
Document resultDoc = results.get(0);
assertThat(resultDoc.getId()).isEqualTo(documents.get(2).getId());
assertThat(resultDoc.getId()).isEqualTo(this.documents.get(2).getId());
assertThat(resultDoc.getContent()).contains("The Great Depression (1929–1939) was an economic shock");
assertThat(resultDoc.getMetadata()).containsKey("meta2");
assertThat(resultDoc.getMetadata()).containsKey("distance");
// Remove all documents from the store
vectorStore.delete(documents.stream().map(Document::getId).toList());
vectorStore.delete(this.documents.stream().map(Document::getId).toList());
Awaitility.await()
.until(() -> vectorStore.similaritySearch(

46
vector-stores/spring-ai-elasticsearch-store/src/test/java/org/springframework/ai/vectorstore/ElasticsearchVectorStoreObservationIT.java
View File

@@ -1,11 +1,11 @@
/*
* Copyright 2024 the original author or authors.
* Copyright 2023-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
* 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,
@@ -13,9 +13,8 @@
* See the License for the specific language governing permissions and
* limitations under the License.
*/
package org.springframework.ai.vectorstore;
import static org.assertj.core.api.Assertions.assertThat;
package org.springframework.ai.vectorstore;
import java.io.IOException;
import java.nio.charset.StandardCharsets;
@@ -24,6 +23,15 @@ import java.util.List;
import java.util.Map;
import java.util.concurrent.TimeUnit;
import co.elastic.clients.elasticsearch.ElasticsearchClient;
import co.elastic.clients.elasticsearch.cat.indices.IndicesRecord;
import co.elastic.clients.json.jackson.JacksonJsonpMapper;
import co.elastic.clients.transport.rest_client.RestClientTransport;
import com.fasterxml.jackson.databind.DeserializationFeature;
import com.fasterxml.jackson.databind.ObjectMapper;
import io.micrometer.observation.ObservationRegistry;
import io.micrometer.observation.tck.TestObservationRegistry;
import io.micrometer.observation.tck.TestObservationRegistryAssert;
import org.apache.http.HttpHost;
import org.awaitility.Awaitility;
import org.elasticsearch.client.RestClient;
@@ -31,6 +39,10 @@ import org.junit.jupiter.api.BeforeAll;
import org.junit.jupiter.api.BeforeEach;
import org.junit.jupiter.api.Test;
import org.junit.jupiter.api.condition.EnabledIfEnvironmentVariable;
import org.testcontainers.elasticsearch.ElasticsearchContainer;
import org.testcontainers.junit.jupiter.Container;
import org.testcontainers.junit.jupiter.Testcontainers;
import org.springframework.ai.document.Document;
import org.springframework.ai.embedding.EmbeddingModel;
import org.springframework.ai.embedding.TokenCountBatchingStrategy;
@@ -48,21 +60,11 @@ import org.springframework.boot.autoconfigure.jdbc.DataSourceAutoConfiguration;
import org.springframework.boot.test.context.runner.ApplicationContextRunner;
import org.springframework.context.annotation.Bean;
import org.springframework.core.io.DefaultResourceLoader;
import org.testcontainers.elasticsearch.ElasticsearchContainer;
import org.testcontainers.junit.jupiter.Container;
import org.testcontainers.junit.jupiter.Testcontainers;
import com.fasterxml.jackson.databind.DeserializationFeature;
import com.fasterxml.jackson.databind.ObjectMapper;
import static org.assertj.core.api.Assertions.assertThat;
import static org.hamcrest.Matchers.greaterThan;
import co.elastic.clients.elasticsearch.ElasticsearchClient;
import co.elastic.clients.elasticsearch.cat.indices.IndicesRecord;
import co.elastic.clients.json.jackson.JacksonJsonpMapper;
import co.elastic.clients.transport.rest_client.RestClientTransport;
import io.micrometer.observation.ObservationRegistry;
import io.micrometer.observation.tck.TestObservationRegistry;
import io.micrometer.observation.tck.TestObservationRegistryAssert;
import static org.hamcrest.Matchers.greaterThan;;
;
/**
* @author Christian Tzolov
@@ -92,10 +94,6 @@ public class ElasticsearchVectorStoreObservationIT {
}
}
private ApplicationContextRunner getContextRunner() {
return new ApplicationContextRunner().withUserConfiguration(Config.class);
}
@BeforeAll
public static void beforeAll() {
Awaitility.setDefaultPollInterval(2, TimeUnit.SECONDS);
@@ -103,6 +101,10 @@ public class ElasticsearchVectorStoreObservationIT {
Awaitility.setDefaultTimeout(Duration.ofMinutes(1));
}
private ApplicationContextRunner getContextRunner() {
return new ApplicationContextRunner().withUserConfiguration(Config.class);
}
@BeforeEach
void cleanDatabase() {
getContextRunner().run(context -> {
@@ -124,7 +126,7 @@ public class ElasticsearchVectorStoreObservationIT {
TestObservationRegistry observationRegistry = context.getBean(TestObservationRegistry.class);
vectorStore.add(documents);
vectorStore.add(this.documents);
TestObservationRegistryAssert.assertThat(observationRegistry)
.doesNotHaveAnyRemainingCurrentObservation()

16
vector-stores/spring-ai-gemfire-store/pom.xml
View File

@@ -1,4 +1,20 @@
<?xml version="1.0" encoding="UTF-8"?>
<!--
~ Copyright 2023-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.
-->
<project xmlns="http://maven.apache.org/POM/4.0.0"
xmlns:xsi="http://www.w3.org/2001/XMLSchema-instance" xsi:schemaLocation="http://maven.apache.org/POM/4.0.0 https://maven.apache.org/xsd/maven-4.0.0.xsd">
<modelVersion>4.0.0</modelVersion>

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

@@ -16,17 +16,22 @@
package org.springframework.ai.vectorstore;
import static org.springframework.http.HttpStatus.BAD_REQUEST;
import static org.springframework.http.HttpStatus.NOT_FOUND;
import java.util.HashMap;
import java.util.List;
import java.util.Map;
import java.util.Optional;
import com.fasterxml.jackson.annotation.JsonCreator;
import com.fasterxml.jackson.annotation.JsonInclude;
import com.fasterxml.jackson.annotation.JsonProperty;
import com.fasterxml.jackson.core.JsonProcessingException;
import com.fasterxml.jackson.databind.ObjectMapper;
import com.fasterxml.jackson.databind.json.JsonMapper;
import io.micrometer.observation.ObservationRegistry;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import reactor.util.annotation.NonNull;
import org.springframework.ai.document.Document;
import org.springframework.ai.embedding.BatchingStrategy;
import org.springframework.ai.embedding.EmbeddingModel;
@@ -47,14 +52,8 @@ import org.springframework.web.reactive.function.client.WebClientException;
import org.springframework.web.reactive.function.client.WebClientResponseException;
import org.springframework.web.util.UriComponentsBuilder;
import com.fasterxml.jackson.annotation.JsonCreator;
import com.fasterxml.jackson.annotation.JsonInclude;
import com.fasterxml.jackson.annotation.JsonProperty;
import com.fasterxml.jackson.core.JsonProcessingException;
import com.fasterxml.jackson.databind.ObjectMapper;
import io.micrometer.observation.ObservationRegistry;
import reactor.util.annotation.NonNull;
import static org.springframework.http.HttpStatus.BAD_REQUEST;
import static org.springframework.http.HttpStatus.NOT_FOUND;
/**
* A VectorStore implementation backed by GemFire. This store supports creating, updating,

5
vector-stores/spring-ai-gemfire-store/src/test/java/org/springframework/ai/vectorstore/GemFireImage.java
View File

@@ -1,11 +1,11 @@
/*
* Copyright 2024 the original author or authors.
* Copyright 2023-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
* 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,
@@ -13,6 +13,7 @@
* See the License for the specific language governing permissions and
* limitations under the License.
*/
package org.springframework.ai.vectorstore;
import org.testcontainers.utility.DockerImageName;

52
vector-stores/spring-ai-gemfire-store/src/test/java/org/springframework/ai/vectorstore/GemFireVectorStoreIT.java
View File

@@ -5,7 +5,7 @@
* 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
* 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,
@@ -13,11 +13,8 @@
* See the License for the specific language governing permissions and
* limitations under the License.
*/
package org.springframework.ai.vectorstore;
import static java.util.concurrent.TimeUnit.MINUTES;
import static org.assertj.core.api.Assertions.assertThat;
import static org.hamcrest.Matchers.hasSize;
package org.springframework.ai.vectorstore;
import java.io.IOException;
import java.nio.charset.StandardCharsets;
@@ -34,6 +31,7 @@ import org.awaitility.Awaitility;
import org.junit.jupiter.api.AfterAll;
import org.junit.jupiter.api.BeforeAll;
import org.junit.jupiter.api.Test;
import org.springframework.ai.document.Document;
import org.springframework.ai.embedding.EmbeddingModel;
import org.springframework.ai.transformers.TransformersEmbeddingModel;
@@ -43,6 +41,10 @@ import org.springframework.boot.test.context.runner.ApplicationContextRunner;
import org.springframework.context.annotation.Bean;
import org.springframework.core.io.DefaultResourceLoader;
import static java.util.concurrent.TimeUnit.MINUTES;
import static org.assertj.core.api.Assertions.assertThat;
import static org.hamcrest.Matchers.hasSize;
/**
* @author Geet Rawat
* @author Soby Chacko
@@ -53,14 +55,22 @@ public class GemFireVectorStoreIT {
public static final String INDEX_NAME = "spring-ai-index1";
private static GemFireCluster gemFireCluster;
private static final int HTTP_SERVICE_PORT = 9090;
private static final int LOCATOR_COUNT = 1;
private static final int SERVER_COUNT = 1;
private static GemFireCluster gemFireCluster;
private final ApplicationContextRunner contextRunner = new ApplicationContextRunner()
.withUserConfiguration(TestApplication.class);
List<Document> documents = List.of(
new Document("1", getText("classpath:/test/data/spring.ai.txt"), Map.of("meta1", "meta1")),
new Document("2", getText("classpath:/test/data/time.shelter.txt"), Map.of()),
new Document("3", getText("classpath:/test/data/great.depression.txt"), Map.of("meta2", "meta2")));
@AfterAll
public static void stopGemFireCluster() {
gemFireCluster.close();
@@ -83,11 +93,6 @@ public class GemFireVectorStoreIT {
String.format("localhost[%d]", gemFireCluster.getLocatorPort()));
}
List<Document> documents = List.of(
new Document("1", getText("classpath:/test/data/spring.ai.txt"), Map.of("meta1", "meta1")),
new Document("2", getText("classpath:/test/data/time.shelter.txt"), Map.of()),
new Document("3", getText("classpath:/test/data/great.depression.txt"), Map.of("meta2", "meta2")));
public static String getText(String uri) {
var resource = new DefaultResourceLoader().getResource(uri);
try {
@@ -98,15 +103,12 @@ public class GemFireVectorStoreIT {
}
}
private final ApplicationContextRunner contextRunner = new ApplicationContextRunner()
.withUserConfiguration(TestApplication.class);
@Test
public void addAndDeleteEmbeddingTest() {
contextRunner.run(context -> {
this.contextRunner.run(context -> {
VectorStore vectorStore = context.getBean(VectorStore.class);
vectorStore.add(documents);
vectorStore.delete(documents.stream().map(doc -> doc.getId()).toList());
vectorStore.add(this.documents);
vectorStore.delete(this.documents.stream().map(doc -> doc.getId()).toList());
Awaitility.await()
.atMost(1, MINUTES)
.until(() -> vectorStore.similaritySearch(SearchRequest.query("Great Depression").withTopK(3)),
@@ -116,9 +118,9 @@ public class GemFireVectorStoreIT {
@Test
public void addAndSearchTest() {
contextRunner.run(context -> {
this.contextRunner.run(context -> {
VectorStore vectorStore = context.getBean(VectorStore.class);
vectorStore.add(documents);
vectorStore.add(this.documents);
Awaitility.await()
.atMost(1, MINUTES)
@@ -127,7 +129,7 @@ public class GemFireVectorStoreIT {
List<Document> results = vectorStore.similaritySearch(SearchRequest.query("Great Depression").withTopK(5));
Document resultDoc = results.get(0);
assertThat(resultDoc.getId()).isEqualTo(documents.get(2).getId());
assertThat(resultDoc.getId()).isEqualTo(this.documents.get(2).getId());
assertThat(resultDoc.getContent()).contains("The Great Depression (1929–1939)" + " was an economic shock");
assertThat(resultDoc.getMetadata()).hasSize(2);
assertThat(resultDoc.getMetadata()).containsKey("meta2");
@@ -137,7 +139,7 @@ public class GemFireVectorStoreIT {
@Test
public void documentUpdateTest() {
contextRunner.run(context -> {
this.contextRunner.run(context -> {
VectorStore vectorStore = context.getBean(VectorStore.class);
Document document = new Document(UUID.randomUUID().toString(), "Spring AI rocks!!",
@@ -175,9 +177,9 @@ public class GemFireVectorStoreIT {
@Test
public void searchThresholdTest() {
contextRunner.run(context -> {
this.contextRunner.run(context -> {
VectorStore vectorStore = context.getBean(VectorStore.class);
vectorStore.add(documents);
vectorStore.add(this.documents);
Awaitility.await()
.atMost(1, MINUTES)
@@ -198,7 +200,7 @@ public class GemFireVectorStoreIT {
assertThat(results).hasSize(1);
Document resultDoc = results.get(0);
assertThat(resultDoc.getId()).isEqualTo(documents.get(2).getId());
assertThat(resultDoc.getId()).isEqualTo(this.documents.get(2).getId());
assertThat(resultDoc.getContent()).contains("The Great Depression " + "(1929–1939) was an economic shock");
assertThat(resultDoc.getMetadata()).containsKey("meta2");
assertThat(resultDoc.getMetadata()).containsKey("distance");

49
vector-stores/spring-ai-gemfire-store/src/test/java/org/springframework/ai/vectorstore/GemFireVectorStoreObservationIT.java
View File

@@ -1,11 +1,11 @@
/*
* Copyright 2024 the original author or authors.
* Copyright 2023-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
* 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,
@@ -13,19 +13,26 @@
* See the License for the specific language governing permissions and
* limitations under the License.
*/
package org.springframework.ai.vectorstore;
import static org.assertj.core.api.Assertions.assertThat;
package org.springframework.ai.vectorstore;
import java.io.IOException;
import java.nio.charset.StandardCharsets;
import java.util.List;
import java.util.Map;
import com.github.dockerjava.api.model.ExposedPort;
import com.github.dockerjava.api.model.PortBinding;
import com.github.dockerjava.api.model.Ports;
import com.vmware.gemfire.testcontainers.GemFireCluster;
import io.micrometer.observation.ObservationRegistry;
import io.micrometer.observation.tck.TestObservationRegistry;
import io.micrometer.observation.tck.TestObservationRegistryAssert;
import org.awaitility.Awaitility;
import org.junit.jupiter.api.AfterAll;
import org.junit.jupiter.api.BeforeAll;
import org.junit.jupiter.api.Test;
import org.springframework.ai.document.Document;
import org.springframework.ai.embedding.EmbeddingModel;
import org.springframework.ai.embedding.TokenCountBatchingStrategy;
@@ -41,16 +48,8 @@ import org.springframework.boot.test.context.runner.ApplicationContextRunner;
import org.springframework.context.annotation.Bean;
import org.springframework.core.io.DefaultResourceLoader;
import com.github.dockerjava.api.model.ExposedPort;
import com.github.dockerjava.api.model.PortBinding;
import com.github.dockerjava.api.model.Ports;
import com.vmware.gemfire.testcontainers.GemFireCluster;
import io.micrometer.observation.ObservationRegistry;
import io.micrometer.observation.tck.TestObservationRegistry;
import io.micrometer.observation.tck.TestObservationRegistryAssert;
import static java.util.concurrent.TimeUnit.MINUTES;
import static org.assertj.core.api.Assertions.assertThat;
import static org.hamcrest.Matchers.hasSize;
/**
@@ -62,14 +61,22 @@ public class GemFireVectorStoreObservationIT {
public static final String TEST_INDEX_NAME = "spring-ai-index1";
private static GemFireCluster gemFireCluster;
private static final int HTTP_SERVICE_PORT = 9090;
private static final int LOCATOR_COUNT = 1;
private static final int SERVER_COUNT = 1;
private static GemFireCluster gemFireCluster;
private final ApplicationContextRunner contextRunner = new ApplicationContextRunner()
.withUserConfiguration(Config.class);
List<Document> documents = List.of(
new Document(getText("classpath:/test/data/spring.ai.txt"), Map.of("meta1", "meta1")),
new Document(getText("classpath:/test/data/time.shelter.txt")),
new Document(getText("classpath:/test/data/great.depression.txt"), Map.of("meta2", "meta2")));
@AfterAll
public static void stopGemFireCluster() {
gemFireCluster.close();
@@ -92,14 +99,6 @@ public class GemFireVectorStoreObservationIT {
String.format("localhost[%d]", gemFireCluster.getLocatorPort()));
}
private final ApplicationContextRunner contextRunner = new ApplicationContextRunner()
.withUserConfiguration(Config.class);
List<Document> documents = List.of(
new Document(getText("classpath:/test/data/spring.ai.txt"), Map.of("meta1", "meta1")),
new Document(getText("classpath:/test/data/time.shelter.txt")),
new Document(getText("classpath:/test/data/great.depression.txt"), Map.of("meta2", "meta2")));
public static String getText(String uri) {
var resource = new DefaultResourceLoader().getResource(uri);
try {
@@ -113,13 +112,13 @@ public class GemFireVectorStoreObservationIT {
@Test
void observationVectorStoreAddAndQueryOperations() {
contextRunner.run(context -> {
this.contextRunner.run(context -> {
VectorStore vectorStore = context.getBean(VectorStore.class);
TestObservationRegistry observationRegistry = context.getBean(TestObservationRegistry.class);
vectorStore.add(documents);
vectorStore.add(this.documents);
TestObservationRegistryAssert.assertThat(observationRegistry)
.doesNotHaveAnyRemainingCurrentObservation()

16
vector-stores/spring-ai-hanadb-store/pom.xml
View File

@@ -1,4 +1,20 @@
<?xml version="1.0" encoding="UTF-8"?>
<!--
~ Copyright 2023-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.
-->
<project xmlns="http://maven.apache.org/POM/4.0.0"
xmlns:xsi="http://www.w3.org/2001/XMLSchema-instance"
xsi:schemaLocation="http://maven.apache.org/POM/4.0.0 http://maven.apache.org/maven-v4_0_0.xsd">

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

@@ -15,14 +15,18 @@
*/
package org.springframework.ai.vectorstore;
import com.fasterxml.jackson.core.JsonProcessingException;
import java.util.Collections;
import java.util.List;
import java.util.Optional;
import java.util.stream.Collectors;
import com.fasterxml.jackson.core.JsonProcessingException;
import com.fasterxml.jackson.databind.ObjectMapper;
import com.fasterxml.jackson.databind.json.JsonMapper;
import io.micrometer.observation.ObservationRegistry;
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;
@@ -34,11 +38,6 @@ import org.springframework.ai.vectorstore.observation.VectorStoreObservationCont
import org.springframework.ai.vectorstore.observation.VectorStoreObservationContext.Builder;
import org.springframework.ai.vectorstore.observation.VectorStoreObservationConvention;
import java.util.Collections;
import java.util.List;
import java.util.Optional;
import java.util.stream.Collectors;
/**
* The <b>SAP HANA Cloud vector engine</b> offers multiple use cases in AI scenarios.
*

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

@@ -1,11 +1,11 @@
/*
* Copyright 2023 - 2024 the original author or authors.
* Copyright 2023-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
* 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,
@@ -13,6 +13,7 @@
* See the License for the specific language governing permissions and
* limitations under the License.
*/
package org.springframework.ai.vectorstore;
/**
@@ -37,11 +38,11 @@ public class HanaCloudVectorStoreConfig {
}
public String getTableName() {
return tableName;
return this.tableName;
}
public int getTopK() {
return topK;
return this.topK;
}
public static class HanaCloudVectorStoreConfigBuilder {
@@ -62,8 +63,8 @@ public class HanaCloudVectorStoreConfig {
public HanaCloudVectorStoreConfig build() {
HanaCloudVectorStoreConfig config = new HanaCloudVectorStoreConfig();
config.tableName = tableName;
config.topK = topK;
config.tableName = this.tableName;
config.topK = this.topK;
return config;
}

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

@@ -1,11 +1,11 @@
/*
* Copyright 2023 - 2024 the original author or authors.
* Copyright 2023-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
* 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,
@@ -13,6 +13,7 @@
* See the License for the specific language governing permissions and
* limitations under the License.
*/
package org.springframework.ai.vectorstore;
import jakarta.persistence.Column;
@@ -39,7 +40,7 @@ public abstract class HanaVectorEntity {
}
public String get_id() {
return _id;
return this._id;
}
}

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

@@ -1,11 +1,11 @@
/*
* Copyright 2023 - 2024 the original author or authors.
* Copyright 2023-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
* 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,
@@ -13,6 +13,7 @@
* See the License for the specific language governing permissions and
* limitations under the License.
*/
package org.springframework.ai.vectorstore;
import java.util.List;

7
vector-stores/spring-ai-hanadb-store/src/test/java/org/springframework/ai/vectorstore/CricketWorldCup.java
View File

@@ -1,11 +1,11 @@
/*
* Copyright 2023 - 2024 the original author or authors.
* Copyright 2023-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
* 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,
@@ -13,6 +13,7 @@
* See the License for the specific language governing permissions and
* limitations under the License.
*/
package org.springframework.ai.vectorstore;
import jakarta.persistence.Column;
@@ -31,7 +32,7 @@ public class CricketWorldCup extends HanaVectorEntity {
private String content;
public String getContent() {
return content;
return this.content;
}
}

26
vector-stores/spring-ai-hanadb-store/src/test/java/org/springframework/ai/vectorstore/CricketWorldCupHanaController.java
View File

@@ -1,11 +1,11 @@
/*
* Copyright 2023 - 2024 the original author or authors.
* Copyright 2023-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
* 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,
@@ -13,12 +13,21 @@
* See the License for the specific language governing permissions and
* limitations under the License.
*/
package org.springframework.ai.vectorstore;
import java.io.IOException;
import java.util.List;
import java.util.Map;
import java.util.function.Function;
import java.util.function.Supplier;
import java.util.stream.Collectors;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import org.springframework.ai.chat.model.ChatModel;
import org.springframework.ai.chat.messages.UserMessage;
import org.springframework.ai.chat.model.ChatModel;
import org.springframework.ai.chat.prompt.Prompt;
import org.springframework.ai.chat.prompt.SystemPromptTemplate;
import org.springframework.ai.document.Document;
@@ -33,13 +42,6 @@ import org.springframework.web.bind.annotation.RequestParam;
import org.springframework.web.bind.annotation.RestController;
import org.springframework.web.multipart.MultipartFile;
import java.io.IOException;
import java.util.List;
import java.util.Map;
import java.util.function.Function;
import java.util.function.Supplier;
import java.util.stream.Collectors;
/**
* @author Rahul Mittal
* @since 1.0.0
@@ -74,7 +76,7 @@ public class CricketWorldCupHanaController {
Function<List<Document>, List<Document>> splitter = new TokenTextSplitter();
List<Document> documents = splitter.apply(reader.get());
logger.info("{} documents created from pdf file: {}", documents.size(), pdf.getFilename());
hanaCloudVectorStore.accept(documents);
this.hanaCloudVectorStore.accept(documents);
return ResponseEntity.ok()
.body(String.format("%d documents created from pdf file: %s", documents.size(), pdf.getFilename()));
}
@@ -88,7 +90,7 @@ public class CricketWorldCupHanaController {
var userMessage = new UserMessage(message);
Prompt prompt = new Prompt(List.of(similarDocsMessage, userMessage));
String generation = chatModel.call(prompt).getResult().getOutput().getContent();
String generation = this.chatModel.call(prompt).getResult().getOutput().getContent();
logger.info("Generation: {}", generation);
return Map.of("generation", generation);
}

18
vector-stores/spring-ai-hanadb-store/src/test/java/org/springframework/ai/vectorstore/CricketWorldCupRepository.java
View File

@@ -1,11 +1,11 @@
/*
* Copyright 2023 - 2024 the original author or authors.
* Copyright 2023-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
* 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,
@@ -13,14 +13,16 @@
* See the License for the specific language governing permissions and
* limitations under the License.
*/
package org.springframework.ai.vectorstore;
import java.util.List;
import jakarta.persistence.EntityManager;
import jakarta.persistence.PersistenceContext;
import jakarta.transaction.Transactional;
import org.springframework.stereotype.Repository;
import java.util.List;
import org.springframework.stereotype.Repository;
/**
* @author Rahul Mittal
@@ -40,7 +42,7 @@ public class CricketWorldCupRepository implements HanaVectorRepository<CricketWo
VALUES(:_id, TO_REAL_VECTOR(:embedding), :content)
""", tableName);
entityManager.createNativeQuery(sql)
this.entityManager.createNativeQuery(sql)
.setParameter("_id", id)
.setParameter("embedding", embedding)
.setParameter("content", content)
@@ -54,7 +56,7 @@ public class CricketWorldCupRepository implements HanaVectorRepository<CricketWo
DELETE FROM %s WHERE _ID IN (:ids)
""", tableName);
return entityManager.createNativeQuery(sql).setParameter("ids", idList).executeUpdate();
return this.entityManager.createNativeQuery(sql).setParameter("ids", idList).executeUpdate();
}
@Override
@@ -64,7 +66,7 @@ public class CricketWorldCupRepository implements HanaVectorRepository<CricketWo
DELETE FROM %s
""", tableName);
return entityManager.createNativeQuery(sql).executeUpdate();
return this.entityManager.createNativeQuery(sql).executeUpdate();
}
@Override
@@ -74,7 +76,7 @@ public class CricketWorldCupRepository implements HanaVectorRepository<CricketWo
ORDER BY COSINE_SIMILARITY(EMBEDDING, TO_REAL_VECTOR(:queryEmbedding)) DESC
""", tableName);
return entityManager.createNativeQuery(sql, CricketWorldCup.class)
return this.entityManager.createNativeQuery(sql, CricketWorldCup.class)
.setParameter("topK", topK)
.setParameter("queryEmbedding", queryEmbedding)
.getResultList();

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

@@ -1,11 +1,11 @@
/*
* Copyright 2023 - 2024 the original author or authors.
* Copyright 2023-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
* 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,
@@ -13,6 +13,7 @@
* See the License for the specific language governing permissions and
* limitations under the License.
*/
package org.springframework.ai.vectorstore;
import java.util.List;
@@ -59,7 +60,7 @@ public class HanaCloudVectorStoreIT {
@Test
public void vectorStoreTest() {
contextRunner.run(context -> {
this.contextRunner.run(context -> {
VectorStore vectorStore = context.getBean(HanaCloudVectorStore.class);
int deleteCount = ((HanaCloudVectorStore) vectorStore).purgeEmbeddings();
@@ -128,4 +129,4 @@ public class HanaCloudVectorStoreIT {
}
}
}

25
vector-stores/spring-ai-hanadb-store/src/test/java/org/springframework/ai/vectorstore/HanaVectorStoreObservationIT.java
View File

@@ -1,11 +1,11 @@
/*
* Copyright 2024 the original author or authors.
* Copyright 2023-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
* 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,
@@ -13,9 +13,8 @@
* See the License for the specific language governing permissions and
* limitations under the License.
*/
package org.springframework.ai.vectorstore;
import static org.assertj.core.api.Assertions.assertThat;
package org.springframework.ai.vectorstore;
import java.io.IOException;
import java.nio.charset.StandardCharsets;
@@ -24,8 +23,12 @@ import java.util.Map;
import javax.sql.DataSource;
import io.micrometer.observation.ObservationRegistry;
import io.micrometer.observation.tck.TestObservationRegistry;
import io.micrometer.observation.tck.TestObservationRegistryAssert;
import org.junit.jupiter.api.Test;
import org.junit.jupiter.api.condition.EnabledIfEnvironmentVariable;
import org.springframework.ai.document.Document;
import org.springframework.ai.embedding.EmbeddingModel;
import org.springframework.ai.observation.conventions.SpringAiKind;
@@ -46,9 +49,7 @@ import org.springframework.orm.jpa.JpaVendorAdapter;
import org.springframework.orm.jpa.LocalContainerEntityManagerFactoryBean;
import org.springframework.orm.jpa.vendor.HibernateJpaVendorAdapter;
import io.micrometer.observation.ObservationRegistry;
import io.micrometer.observation.tck.TestObservationRegistry;
import io.micrometer.observation.tck.TestObservationRegistryAssert;
import static org.assertj.core.api.Assertions.assertThat;
/**
* @author Christian Tzolov
@@ -62,6 +63,9 @@ public class HanaVectorStoreObservationIT {
private static final String TEST_TABLE_NAME = "CRICKET_WORLD_CUP";
private final ApplicationContextRunner contextRunner = new ApplicationContextRunner()
.withUserConfiguration(Config.class);
List<Document> documents = List.of(
new Document(getText("classpath:/test/data/spring.ai.txt"), Map.of("meta1", "meta1")),
new Document(getText("classpath:/test/data/time.shelter.txt")),
@@ -77,19 +81,16 @@ public class HanaVectorStoreObservationIT {
}
}
private final ApplicationContextRunner contextRunner = new ApplicationContextRunner()
.withUserConfiguration(Config.class);
@Test
void observationVectorStoreAddAndQueryOperations() {
contextRunner.run(context -> {
this.contextRunner.run(context -> {
VectorStore vectorStore = context.getBean(VectorStore.class);
TestObservationRegistry observationRegistry = context.getBean(TestObservationRegistry.class);
vectorStore.add(documents);
vectorStore.add(this.documents);
TestObservationRegistryAssert.assertThat(observationRegistry)
.doesNotHaveAnyRemainingCurrentObservation()

16
vector-stores/spring-ai-hanadb-store/src/test/resources/application.properties
View File

@@ -1,3 +1,19 @@
#
# Copyright 2023-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.
#
spring.ai.openai.api-key=${OPENAI_API_KEY}
spring.ai.openai.embedding.options.model=text-embedding-ada-002

16
vector-stores/spring-ai-milvus-store/pom.xml
View File

@@ -1,4 +1,20 @@
<?xml version="1.0" encoding="UTF-8"?>
<!--
~ Copyright 2023-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.
-->
<project xmlns="http://maven.apache.org/POM/4.0.0"
xmlns:xsi="http://www.w3.org/2001/XMLSchema-instance" xsi:schemaLocation="http://maven.apache.org/POM/4.0.0 http://maven.apache.org/maven-v4_0_0.xsd">
<modelVersion>4.0.0</modelVersion>

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

@@ -1,11 +1,11 @@
/*
* Copyright 2023 - 2024 the original author or authors.
* Copyright 2023-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
* 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,
@@ -13,6 +13,7 @@
* See the License for the specific language governing permissions and
* limitations under the License.
*/
package org.springframework.ai.vectorstore;
import org.springframework.ai.vectorstore.filter.Filter.Expression;

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

@@ -5,7 +5,7 @@
* 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
* 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,
@@ -13,8 +13,15 @@
* See the License for the specific language governing permissions and
* limitations under the License.
*/
package org.springframework.ai.vectorstore;
import java.util.ArrayList;
import java.util.List;
import java.util.Map;
import java.util.Optional;
import java.util.stream.Collectors;
import com.alibaba.fastjson.JSONObject;
import io.micrometer.observation.ObservationRegistry;
import io.milvus.client.MilvusServiceClient;
@@ -44,6 +51,7 @@ import io.milvus.response.QueryResultsWrapper.RowRecord;
import io.milvus.response.SearchResultsWrapper;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import org.springframework.ai.document.Document;
import org.springframework.ai.embedding.BatchingStrategy;
import org.springframework.ai.embedding.EmbeddingModel;
@@ -60,12 +68,6 @@ import org.springframework.beans.factory.InitializingBean;
import org.springframework.util.Assert;
import org.springframework.util.StringUtils;
import java.util.ArrayList;
import java.util.List;
import java.util.Map;
import java.util.Optional;
import java.util.stream.Collectors;
/**
* @author Christian Tzolov
* @author Soby Chacko
@@ -73,8 +75,6 @@ import java.util.stream.Collectors;
*/
public class MilvusVectorStore extends AbstractObservationVectorStore implements InitializingBean {
private static final Logger logger = LoggerFactory.getLogger(MilvusVectorStore.class);
public static final int OPENAI_EMBEDDING_DIMENSION_SIZE = 1536;
public static final int INVALID_EMBEDDING_DIMENSION = -1;
@@ -97,6 +97,12 @@ public class MilvusVectorStore extends AbstractObservationVectorStore implements
public static final List<String> SEARCH_OUTPUT_FIELDS = List.of(DOC_ID_FIELD_NAME, CONTENT_FIELD_NAME,
METADATA_FIELD_NAME);
private static final Logger logger = LoggerFactory.getLogger(MilvusVectorStore.class);
private static Map<MetricType, VectorStoreSimilarityMetric> SIMILARITY_TYPE_MAPPING = Map.of(MetricType.COSINE,
VectorStoreSimilarityMetric.COSINE, MetricType.L2, VectorStoreSimilarityMetric.EUCLIDEAN, MetricType.IP,
VectorStoreSimilarityMetric.DOT);
public final FilterExpressionConverter filterExpressionConverter = new MilvusFilterExpressionConverter();
private final MilvusServiceClient milvusClient;
@@ -109,151 +115,6 @@ public class MilvusVectorStore extends AbstractObservationVectorStore implements
private final BatchingStrategy batchingStrategy;
/**
* Configuration for the Milvus vector store.
*/
public static class MilvusVectorStoreConfig {
private final String databaseName;
private final String collectionName;
private final int embeddingDimension;
private final IndexType indexType;
private final MetricType metricType;
private final String indexParameters;
/**
* Start building a new configuration.
* @return The entry point for creating a new configuration.
*/
public static Builder builder() {
return new Builder();
}
/**
* {@return the default config}
*/
public static MilvusVectorStoreConfig defaultConfig() {
return builder().build();
}
private MilvusVectorStoreConfig(Builder builder) {
this.databaseName = builder.databaseName;
this.collectionName = builder.collectionName;
this.embeddingDimension = builder.embeddingDimension;
this.indexType = builder.indexType;
this.metricType = builder.metricType;
this.indexParameters = builder.indexParameters;
}
public static class Builder {
private String databaseName = DEFAULT_DATABASE_NAME;
private String collectionName = DEFAULT_COLLECTION_NAME;
private int embeddingDimension = INVALID_EMBEDDING_DIMENSION;
private IndexType indexType = IndexType.IVF_FLAT;
private MetricType metricType = MetricType.COSINE;
private String indexParameters = "{\"nlist\":1024}";
private Builder() {
}
/**
* Configures the Milvus metric type to use. Leave {@literal null} or blank to
* use the metric metric: https://milvus.io/docs/metric.md#floating
* @param metricType the metric type to use
* @return this builder
*/
public Builder withMetricType(MetricType metricType) {
Assert.notNull(metricType, "Collection Name must not be empty");
Assert.isTrue(
metricType == MetricType.IP || metricType == MetricType.L2 || metricType == MetricType.COSINE,
"Only the text metric types IP and L2 are supported");
this.metricType = metricType;
return this;
}
/**
* Configures the Milvus index type to use. Leave {@literal null} or blank to
* use the default index.
* @param indexType the index type to use
* @return this builder
*/
public Builder withIndexType(IndexType indexType) {
this.indexType = indexType;
return this;
}
/**
* Configures the Milvus index parameters to use. Leave {@literal null} or
* blank to use the default index parameters.
* @param indexParameters the index parameters to use
* @return this builder
*/
public Builder withIndexParameters(String indexParameters) {
this.indexParameters = indexParameters;
return this;
}
/**
* Configures the Milvus database name to use. Leave {@literal null} or blank
* to use the default database.
* @param databaseName the database name to use
* @return this builder
*/
public Builder withDatabaseName(String databaseName) {
this.databaseName = databaseName;
return this;
}
/**
* Configures the Milvus collection name to use. Leave {@literal null} or
* blank to use the default collection name.
* @param collectionName the collection name to use
* @return this builder
*/
public Builder withCollectionName(String collectionName) {
this.collectionName = collectionName;
return this;
}
/**
* Configures the size of the embedding. Defaults to {@literal 1536}, inline
* with OpenAIs embeddings.
* @param newEmbeddingDimension The dimension of the embedding
* @return this builder
*/
public Builder withEmbeddingDimension(int newEmbeddingDimension) {
Assert.isTrue(newEmbeddingDimension >= 1 && newEmbeddingDimension <= 32768,
"Dimension has to be withing the boundaries 1 and 32768 (inclusively)");
this.embeddingDimension = newEmbeddingDimension;
return this;
}
/**
* {@return the immutable configuration}
*/
public MilvusVectorStoreConfig build() {
return new MilvusVectorStoreConfig(this);
}
}
}
public MilvusVectorStore(MilvusServiceClient milvusClient, EmbeddingModel embeddingModel,
boolean initializeSchema) {
this(milvusClient, embeddingModel, MilvusVectorStoreConfig.defaultConfig(), initializeSchema,
@@ -369,7 +230,7 @@ public class MilvusVectorStore extends AbstractObservationVectorStore implements
searchParamBuilder.withExpr(nativeFilterExpressions);
}
R<SearchResults> respSearch = milvusClient.search(searchParamBuilder.build());
R<SearchResults> respSearch = this.milvusClient.search(searchParamBuilder.build());
if (respSearch.getException() != null) {
throw new RuntimeException("Search failed!", respSearch.getException());
@@ -558,10 +419,6 @@ public class MilvusVectorStore extends AbstractObservationVectorStore implements
.withNamespace(this.config.databaseName);
}
private static Map<MetricType, VectorStoreSimilarityMetric> SIMILARITY_TYPE_MAPPING = Map.of(MetricType.COSINE,
VectorStoreSimilarityMetric.COSINE, MetricType.L2, VectorStoreSimilarityMetric.EUCLIDEAN, MetricType.IP,
VectorStoreSimilarityMetric.DOT);
private String getSimilarityMetric() {
if (!SIMILARITY_TYPE_MAPPING.containsKey(this.config.metricType)) {
return this.config.metricType.name();
@@ -569,4 +426,149 @@ public class MilvusVectorStore extends AbstractObservationVectorStore implements
return SIMILARITY_TYPE_MAPPING.get(this.config.metricType).value();
}
/**
* Configuration for the Milvus vector store.
*/
public static class MilvusVectorStoreConfig {
private final String databaseName;
private final String collectionName;
private final int embeddingDimension;
private final IndexType indexType;
private final MetricType metricType;
private final String indexParameters;
private MilvusVectorStoreConfig(Builder builder) {
this.databaseName = builder.databaseName;
this.collectionName = builder.collectionName;
this.embeddingDimension = builder.embeddingDimension;
this.indexType = builder.indexType;
this.metricType = builder.metricType;
this.indexParameters = builder.indexParameters;
}
/**
* Start building a new configuration.
* @return The entry point for creating a new configuration.
*/
public static Builder builder() {
return new Builder();
}
/**
* {@return the default config}
*/
public static MilvusVectorStoreConfig defaultConfig() {
return builder().build();
}
public static class Builder {
private String databaseName = DEFAULT_DATABASE_NAME;
private String collectionName = DEFAULT_COLLECTION_NAME;
private int embeddingDimension = INVALID_EMBEDDING_DIMENSION;
private IndexType indexType = IndexType.IVF_FLAT;
private MetricType metricType = MetricType.COSINE;
private String indexParameters = "{\"nlist\":1024}";
private Builder() {
}
/**
* Configures the Milvus metric type to use. Leave {@literal null} or blank to
* use the metric metric: https://milvus.io/docs/metric.md#floating
* @param metricType the metric type to use
* @return this builder
*/
public Builder withMetricType(MetricType metricType) {
Assert.notNull(metricType, "Collection Name must not be empty");
Assert.isTrue(
metricType == MetricType.IP || metricType == MetricType.L2 || metricType == MetricType.COSINE,
"Only the text metric types IP and L2 are supported");
this.metricType = metricType;
return this;
}
/**
* Configures the Milvus index type to use. Leave {@literal null} or blank to
* use the default index.
* @param indexType the index type to use
* @return this builder
*/
public Builder withIndexType(IndexType indexType) {
this.indexType = indexType;
return this;
}
/**
* Configures the Milvus index parameters to use. Leave {@literal null} or
* blank to use the default index parameters.
* @param indexParameters the index parameters to use
* @return this builder
*/
public Builder withIndexParameters(String indexParameters) {
this.indexParameters = indexParameters;
return this;
}
/**
* Configures the Milvus database name to use. Leave {@literal null} or blank
* to use the default database.
* @param databaseName the database name to use
* @return this builder
*/
public Builder withDatabaseName(String databaseName) {
this.databaseName = databaseName;
return this;
}
/**
* Configures the Milvus collection name to use. Leave {@literal null} or
* blank to use the default collection name.
* @param collectionName the collection name to use
* @return this builder
*/
public Builder withCollectionName(String collectionName) {
this.collectionName = collectionName;
return this;
}
/**
* Configures the size of the embedding. Defaults to {@literal 1536}, inline
* with OpenAIs embeddings.
* @param newEmbeddingDimension The dimension of the embedding
* @return this builder
*/
public Builder withEmbeddingDimension(int newEmbeddingDimension) {
Assert.isTrue(newEmbeddingDimension >= 1 && newEmbeddingDimension <= 32768,
"Dimension has to be withing the boundaries 1 and 32768 (inclusively)");
this.embeddingDimension = newEmbeddingDimension;
return this;
}
/**
* {@return the immutable configuration}
*/
public MilvusVectorStoreConfig build() {
return new MilvusVectorStoreConfig(this);
}
}
}
}

25
vector-stores/spring-ai-milvus-store/src/test/java/org/springframework/ai/vectorstore/MilvusEmbeddingDimensionsTests.java
View File

@@ -1,11 +1,11 @@
/*
* Copyright 2023 - 2024 the original author or authors.
* Copyright 2023-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
* 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,
@@ -13,6 +13,7 @@
* See the License for the specific language governing permissions and
* limitations under the License.
*/
package org.springframework.ai.vectorstore;
import io.milvus.client.MilvusServiceClient;
@@ -57,38 +58,40 @@ public class MilvusEmbeddingDimensionsTests {
.withEmbeddingDimension(explicitDimensions)
.build();
var dim = new MilvusVectorStore(milvusClient, embeddingModel, config, true, new TokenCountBatchingStrategy())
var dim = new MilvusVectorStore(this.milvusClient, this.embeddingModel, config, true,
new TokenCountBatchingStrategy())
.embeddingDimensions();
assertThat(dim).isEqualTo(explicitDimensions);
verify(embeddingModel, never()).dimensions();
verify(this.embeddingModel, never()).dimensions();
}
@Test
public void embeddingModelDimensions() {
when(embeddingModel.dimensions()).thenReturn(969);
when(this.embeddingModel.dimensions()).thenReturn(969);
MilvusVectorStoreConfig config = MilvusVectorStoreConfig.builder().build();
var dim = new MilvusVectorStore(milvusClient, embeddingModel, config, true, new TokenCountBatchingStrategy())
var dim = new MilvusVectorStore(this.milvusClient, this.embeddingModel, config, true,
new TokenCountBatchingStrategy())
.embeddingDimensions();
assertThat(dim).isEqualTo(969);
verify(embeddingModel, only()).dimensions();
verify(this.embeddingModel, only()).dimensions();
}
@Test
public void fallBackToDefaultDimensions() {
when(embeddingModel.dimensions()).thenThrow(new RuntimeException());
when(this.embeddingModel.dimensions()).thenThrow(new RuntimeException());
var dim = new MilvusVectorStore(milvusClient, embeddingModel, MilvusVectorStoreConfig.builder().build(), true,
new TokenCountBatchingStrategy())
var dim = new MilvusVectorStore(this.milvusClient, this.embeddingModel,
MilvusVectorStoreConfig.builder().build(), true, new TokenCountBatchingStrategy())
.embeddingDimensions();
assertThat(dim).isEqualTo(MilvusVectorStore.OPENAI_EMBEDDING_DIMENSION_SIZE);
verify(embeddingModel, only()).dimensions();
verify(this.embeddingModel, only()).dimensions();
}
@ParameterizedTest

23
vector-stores/spring-ai-milvus-store/src/test/java/org/springframework/ai/vectorstore/MilvusFilterExpressionConverterTests.java
View File

@@ -1,11 +1,11 @@
/*
* Copyright 2023 - 2024 the original author or authors.
* Copyright 2023-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
* 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,
@@ -13,6 +13,7 @@
* See the License for the specific language governing permissions and
* limitations under the License.
*/
package org.springframework.ai.vectorstore;
import java.util.List;
@@ -45,14 +46,14 @@ public class MilvusFilterExpressionConverterTests {
@Test
public void testEQ() {
// country == "BG"
String vectorExpr = converter.convertExpression(new Expression(EQ, new Key("country"), new Value("BG")));
String vectorExpr = this.converter.convertExpression(new Expression(EQ, new Key("country"), new Value("BG")));
assertThat(vectorExpr).isEqualTo("metadata[\"country\"] == \"BG\"");
}
@Test
public void tesEqAndGte() {
// genre == "drama" AND year >= 2020
String vectorExpr = converter
String vectorExpr = this.converter
.convertExpression(new Expression(AND, new Expression(EQ, new Key("genre"), new Value("drama")),
new Expression(GTE, new Key("year"), new Value(2020))));
assertThat(vectorExpr).isEqualTo("metadata[\"genre\"] == \"drama\" && metadata[\"year\"] >= 2020");
@@ -61,7 +62,7 @@ public class MilvusFilterExpressionConverterTests {
@Test
public void tesIn() {
// genre in ["comedy", "documentary", "drama"]
String vectorExpr = converter.convertExpression(
String vectorExpr = this.converter.convertExpression(
new Expression(IN, new Key("genre"), new Value(List.of("comedy", "documentary", "drama"))));
assertThat(vectorExpr).isEqualTo("metadata[\"genre\"] in [\"comedy\",\"documentary\",\"drama\"]");
}
@@ -69,7 +70,7 @@ public class MilvusFilterExpressionConverterTests {
@Test
public void testNe() {
// year >= 2020 OR country == "BG" AND city != "Sofia"
String vectorExpr = converter
String vectorExpr = this.converter
.convertExpression(new Expression(OR, new Expression(GTE, new Key("year"), new Value(2020)),
new Expression(AND, new Expression(EQ, new Key("country"), new Value("BG")),
new Expression(NE, new Key("city"), new Value("Sofia")))));
@@ -80,7 +81,7 @@ public class MilvusFilterExpressionConverterTests {
@Test
public void testGroup() {
// (year >= 2020 OR country == "BG") AND city NIN ["Sofia", "Plovdiv"]
String vectorExpr = converter.convertExpression(new Expression(AND,
String vectorExpr = this.converter.convertExpression(new Expression(AND,
new Group(new Expression(OR, new Expression(GTE, new Key("year"), new Value(2020)),
new Expression(EQ, new Key("country"), new Value("BG")))),
new Expression(NIN, new Key("city"), new Value(List.of("Sofia", "Plovdiv")))));
@@ -91,7 +92,7 @@ public class MilvusFilterExpressionConverterTests {
@Test
public void tesBoolean() {
// isOpen == true AND year >= 2020 AND country IN ["BG", "NL", "US"]
String vectorExpr = converter.convertExpression(new Expression(AND,
String vectorExpr = this.converter.convertExpression(new Expression(AND,
new Expression(AND, new Expression(EQ, new Key("isOpen"), new Value(true)),
new Expression(GTE, new Key("year"), new Value(2020))),
new Expression(IN, new Key("country"), new Value(List.of("BG", "NL", "US")))));
@@ -103,7 +104,7 @@ public class MilvusFilterExpressionConverterTests {
@Test
public void testDecimal() {
// temperature >= -15.6 && temperature <= +20.13
String vectorExpr = converter
String vectorExpr = this.converter
.convertExpression(new Expression(AND, new Expression(GTE, new Key("temperature"), new Value(-15.6)),
new Expression(LTE, new Key("temperature"), new Value(20.13))));
@@ -112,11 +113,11 @@ public class MilvusFilterExpressionConverterTests {
@Test
public void testComplexIdentifiers() {
String vectorExpr = converter
String vectorExpr = this.converter
.convertExpression(new Expression(EQ, new Key("\"country 1 2 3\""), new Value("BG")));
assertThat(vectorExpr).isEqualTo("metadata[\"country 1 2 3\"] == \"BG\"");
vectorExpr = converter.convertExpression(new Expression(EQ, new Key("'country 1 2 3'"), new Value("BG")));
vectorExpr = this.converter.convertExpression(new Expression(EQ, new Key("'country 1 2 3'"), new Value("BG")));
assertThat(vectorExpr).isEqualTo("metadata[\"country 1 2 3\"] == \"BG\"");
}

5
vector-stores/spring-ai-milvus-store/src/test/java/org/springframework/ai/vectorstore/MilvusImage.java
View File

@@ -1,11 +1,11 @@
/*
* Copyright 2024 the original author or authors.
* Copyright 2023-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
* 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,
@@ -13,6 +13,7 @@
* See the License for the specific language governing permissions and
* limitations under the License.
*/
package org.springframework.ai.vectorstore;
import org.testcontainers.utility.DockerImageName;

217
vector-stores/spring-ai-milvus-store/src/test/java/org/springframework/ai/vectorstore/MilvusVectorStoreIT.java
View File

@@ -1,11 +1,11 @@
/*
* Copyright 2023 - 2024 the original author or authors.
* Copyright 2023-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
* 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,
@@ -13,6 +13,7 @@
* See the License for the specific language governing permissions and
* limitations under the License.
*/
package org.springframework.ai.vectorstore;
import java.io.IOException;
@@ -31,6 +32,7 @@ import org.junit.jupiter.params.ParameterizedTest;
import org.junit.jupiter.params.provider.ValueSource;
import org.testcontainers.junit.jupiter.Container;
import org.testcontainers.junit.jupiter.Testcontainers;
import org.testcontainers.milvus.MilvusContainer;
import org.springframework.ai.document.Document;
import org.springframework.ai.embedding.EmbeddingModel;
@@ -45,7 +47,6 @@ import org.springframework.boot.autoconfigure.jdbc.DataSourceAutoConfiguration;
import org.springframework.boot.test.context.runner.ApplicationContextRunner;
import org.springframework.context.annotation.Bean;
import org.springframework.core.io.DefaultResourceLoader;
import org.testcontainers.milvus.MilvusContainer;
import static org.assertj.core.api.Assertions.assertThat;
@@ -88,30 +89,31 @@ public class MilvusVectorStoreIT {
@ValueSource(strings = { "COSINE", "L2", "IP" })
public void addAndSearch(String metricType) {
contextRunner.withPropertyValues("test.spring.ai.vectorstore.milvus.metricType=" + metricType).run(context -> {
this.contextRunner.withPropertyValues("test.spring.ai.vectorstore.milvus.metricType=" + metricType)
.run(context -> {
VectorStore vectorStore = context.getBean(VectorStore.class);
VectorStore vectorStore = context.getBean(VectorStore.class);
resetCollection(vectorStore);
resetCollection(vectorStore);
vectorStore.add(documents);
vectorStore.add(this.documents);
List<Document> results = vectorStore.similaritySearch(SearchRequest.query("Spring").withTopK(1));
List<Document> results = vectorStore.similaritySearch(SearchRequest.query("Spring").withTopK(1));
assertThat(results).hasSize(1);
Document resultDoc = results.get(0);
assertThat(resultDoc.getId()).isEqualTo(documents.get(0).getId());
assertThat(resultDoc.getContent()).contains(
"Spring AI provides abstractions that serve as the foundation for developing AI applications.");
assertThat(resultDoc.getMetadata()).hasSize(2);
assertThat(resultDoc.getMetadata()).containsKeys("meta1", "distance");
assertThat(results).hasSize(1);
Document resultDoc = results.get(0);
assertThat(resultDoc.getId()).isEqualTo(this.documents.get(0).getId());
assertThat(resultDoc.getContent()).contains(
"Spring AI provides abstractions that serve as the foundation for developing AI applications.");
assertThat(resultDoc.getMetadata()).hasSize(2);
assertThat(resultDoc.getMetadata()).containsKeys("meta1", "distance");
// Remove all documents from the store
vectorStore.delete(documents.stream().map(doc -> doc.getId()).toList());
// Remove all documents from the store
vectorStore.delete(this.documents.stream().map(doc -> doc.getId()).toList());
results = vectorStore.similaritySearch(SearchRequest.query("Spring").withTopK(1));
assertThat(results).hasSize(0);
});
results = vectorStore.similaritySearch(SearchRequest.query("Spring").withTopK(1));
assertThat(results).hasSize(0);
});
}
@ParameterizedTest(name = "{0} : {displayName} ")
@@ -121,135 +123,140 @@ public class MilvusVectorStoreIT {
// https://milvus.io/docs/json_data_type.md
contextRunner.withPropertyValues("test.spring.ai.vectorstore.milvus.metricType=" + metricType).run(context -> {
VectorStore vectorStore = context.getBean(VectorStore.class);
this.contextRunner.withPropertyValues("test.spring.ai.vectorstore.milvus.metricType=" + metricType)
.run(context -> {
VectorStore vectorStore = context.getBean(VectorStore.class);
resetCollection(vectorStore);
resetCollection(vectorStore);
var bgDocument = new Document("The World is Big and Salvation Lurks Around the Corner",
Map.of("country", "BG", "year", 2020));
var nlDocument = new Document("The World is Big and Salvation Lurks Around the Corner",
Map.of("country", "NL"));
var bgDocument2 = new Document("The World is Big and Salvation Lurks Around the Corner",
Map.of("country", "BG", "year", 2023));
var bgDocument = new Document("The World is Big and Salvation Lurks Around the Corner",
Map.of("country", "BG", "year", 2020));
var nlDocument = new Document("The World is Big and Salvation Lurks Around the Corner",
Map.of("country", "NL"));
var bgDocument2 = new Document("The World is Big and Salvation Lurks Around the Corner",
Map.of("country", "BG", "year", 2023));
vectorStore.add(List.of(bgDocument, nlDocument, bgDocument2));
vectorStore.add(List.of(bgDocument, nlDocument, bgDocument2));
List<Document> results = vectorStore.similaritySearch(SearchRequest.query("The World").withTopK(5));
assertThat(results).hasSize(3);
List<Document> results = vectorStore.similaritySearch(SearchRequest.query("The World").withTopK(5));
assertThat(results).hasSize(3);
results = vectorStore.similaritySearch(SearchRequest.query("The World")
.withTopK(5)
.withSimilarityThresholdAll()
.withFilterExpression("country == 'NL'"));
assertThat(results).hasSize(1);
assertThat(results.get(0).getId()).isEqualTo(nlDocument.getId());
results = vectorStore.similaritySearch(SearchRequest.query("The World")
.withTopK(5)
.withSimilarityThresholdAll()
.withFilterExpression("country == 'NL'"));
assertThat(results).hasSize(1);
assertThat(results.get(0).getId()).isEqualTo(nlDocument.getId());
results = vectorStore.similaritySearch(SearchRequest.query("The World")
.withTopK(5)
.withSimilarityThresholdAll()
.withFilterExpression("country == 'BG'"));
results = vectorStore.similaritySearch(SearchRequest.query("The World")
.withTopK(5)
.withSimilarityThresholdAll()
.withFilterExpression("country == 'BG'"));
assertThat(results).hasSize(2);
assertThat(results.get(0).getId()).isIn(bgDocument.getId(), bgDocument2.getId());
assertThat(results.get(1).getId()).isIn(bgDocument.getId(), bgDocument2.getId());
assertThat(results).hasSize(2);
assertThat(results.get(0).getId()).isIn(bgDocument.getId(), bgDocument2.getId());
assertThat(results.get(1).getId()).isIn(bgDocument.getId(), bgDocument2.getId());
results = vectorStore.similaritySearch(SearchRequest.query("The World")
.withTopK(5)
.withSimilarityThresholdAll()
.withFilterExpression("country == 'BG' && year == 2020"));
results = vectorStore.similaritySearch(SearchRequest.query("The World")
.withTopK(5)
.withSimilarityThresholdAll()
.withFilterExpression("country == 'BG' && year == 2020"));
assertThat(results).hasSize(1);
assertThat(results.get(0).getId()).isEqualTo(bgDocument.getId());
assertThat(results).hasSize(1);
assertThat(results.get(0).getId()).isEqualTo(bgDocument.getId());
results = vectorStore.similaritySearch(SearchRequest.query("The World")
.withTopK(5)
.withSimilarityThresholdAll()
.withFilterExpression("NOT(country == 'BG' && year == 2020)"));
results = vectorStore.similaritySearch(SearchRequest.query("The World")
.withTopK(5)
.withSimilarityThresholdAll()
.withFilterExpression("NOT(country == 'BG' && year == 2020)"));
assertThat(results).hasSize(2);
assertThat(results.get(0).getId()).isIn(nlDocument.getId(), bgDocument2.getId());
assertThat(results.get(1).getId()).isIn(nlDocument.getId(), bgDocument2.getId());
assertThat(results).hasSize(2);
assertThat(results.get(0).getId()).isIn(nlDocument.getId(), bgDocument2.getId());
assertThat(results.get(1).getId()).isIn(nlDocument.getId(), bgDocument2.getId());
});
});
}
@ParameterizedTest(name = "{0} : {displayName} ")
@ValueSource(strings = { "COSINE", "L2", "IP" })
public void documentUpdate(String metricType) {
contextRunner.withPropertyValues("test.spring.ai.vectorstore.milvus.metricType=" + metricType).run(context -> {
this.contextRunner.withPropertyValues("test.spring.ai.vectorstore.milvus.metricType=" + metricType)
.run(context -> {
VectorStore vectorStore = context.getBean(VectorStore.class);
VectorStore vectorStore = context.getBean(VectorStore.class);
resetCollection(vectorStore);
resetCollection(vectorStore);
Document document = new Document(UUID.randomUUID().toString(), "Spring AI rocks!!",
Collections.singletonMap("meta1", "meta1"));
Document document = new Document(UUID.randomUUID().toString(), "Spring AI rocks!!",
Collections.singletonMap("meta1", "meta1"));
vectorStore.add(List.of(document));
vectorStore.add(List.of(document));
List<Document> results = vectorStore.similaritySearch(SearchRequest.query("Spring").withTopK(5));
List<Document> results = vectorStore.similaritySearch(SearchRequest.query("Spring").withTopK(5));
assertThat(results).hasSize(1);
Document resultDoc = results.get(0);
assertThat(resultDoc.getId()).isEqualTo(document.getId());
assertThat(resultDoc.getContent()).isEqualTo("Spring AI rocks!!");
assertThat(resultDoc.getMetadata()).containsKey("meta1");
assertThat(resultDoc.getMetadata()).containsKey("distance");
assertThat(results).hasSize(1);
Document resultDoc = results.get(0);
assertThat(resultDoc.getId()).isEqualTo(document.getId());
assertThat(resultDoc.getContent()).isEqualTo("Spring AI rocks!!");
assertThat(resultDoc.getMetadata()).containsKey("meta1");
assertThat(resultDoc.getMetadata()).containsKey("distance");
Document sameIdDocument = new Document(document.getId(),
"The World is Big and Salvation Lurks Around the Corner",
Collections.singletonMap("meta2", "meta2"));
Document sameIdDocument = new Document(document.getId(),
"The World is Big and Salvation Lurks Around the Corner",
Collections.singletonMap("meta2", "meta2"));
vectorStore.add(List.of(sameIdDocument));
vectorStore.add(List.of(sameIdDocument));
results = vectorStore.similaritySearch(SearchRequest.query("FooBar").withTopK(5));
results = vectorStore.similaritySearch(SearchRequest.query("FooBar").withTopK(5));
assertThat(results).hasSize(1);
resultDoc = results.get(0);
assertThat(resultDoc.getId()).isEqualTo(document.getId());
assertThat(resultDoc.getContent()).isEqualTo("The World is Big and Salvation Lurks Around the Corner");
assertThat(resultDoc.getMetadata()).containsKey("meta2");
assertThat(resultDoc.getMetadata()).containsKey("distance");
assertThat(results).hasSize(1);
resultDoc = results.get(0);
assertThat(resultDoc.getId()).isEqualTo(document.getId());
assertThat(resultDoc.getContent()).isEqualTo("The World is Big and Salvation Lurks Around the Corner");
assertThat(resultDoc.getMetadata()).containsKey("meta2");
assertThat(resultDoc.getMetadata()).containsKey("distance");
vectorStore.delete(List.of(document.getId()));
vectorStore.delete(List.of(document.getId()));
});
});
}
@ParameterizedTest(name = "{0} : {displayName} ")
@ValueSource(strings = { "COSINE", "IP" })
public void searchWithThreshold(String metricType) {
contextRunner.withPropertyValues("test.spring.ai.vectorstore.milvus.metricType=" + metricType).run(context -> {
this.contextRunner.withPropertyValues("test.spring.ai.vectorstore.milvus.metricType=" + metricType)
.run(context -> {
VectorStore vectorStore = context.getBean(VectorStore.class);
VectorStore vectorStore = context.getBean(VectorStore.class);
resetCollection(vectorStore);
resetCollection(vectorStore);
vectorStore.add(documents);
vectorStore.add(this.documents);
List<Document> fullResult = vectorStore
.similaritySearch(SearchRequest.query("Spring").withTopK(5).withSimilarityThresholdAll());
List<Document> fullResult = vectorStore
.similaritySearch(SearchRequest.query("Spring").withTopK(5).withSimilarityThresholdAll());
List<Float> distances = fullResult.stream().map(doc -> (Float) doc.getMetadata().get("distance")).toList();
List<Float> distances = fullResult.stream()
.map(doc -> (Float) doc.getMetadata().get("distance"))
.toList();
assertThat(distances).hasSize(3);
assertThat(distances).hasSize(3);
float threshold = (distances.get(0) + distances.get(1)) / 2;
float threshold = (distances.get(0) + distances.get(1)) / 2;
List<Document> results = vectorStore
.similaritySearch(SearchRequest.query("Spring").withTopK(5).withSimilarityThreshold(1 - threshold));
List<Document> results = vectorStore
.similaritySearch(SearchRequest.query("Spring").withTopK(5).withSimilarityThreshold(1 - threshold));
assertThat(results).hasSize(1);
Document resultDoc = results.get(0);
assertThat(resultDoc.getId()).isEqualTo(documents.get(0).getId());
assertThat(resultDoc.getContent()).contains(
"Spring AI provides abstractions that serve as the foundation for developing AI applications.");
assertThat(resultDoc.getMetadata()).containsKeys("meta1", "distance");
assertThat(results).hasSize(1);
Document resultDoc = results.get(0);
assertThat(resultDoc.getId()).isEqualTo(this.documents.get(0).getId());
assertThat(resultDoc.getContent()).contains(
"Spring AI provides abstractions that serve as the foundation for developing AI applications.");
assertThat(resultDoc.getMetadata()).containsKeys("meta1", "distance");
});
});
}
@SpringBootConfiguration
@@ -265,7 +272,7 @@ public class MilvusVectorStoreIT {
.withCollectionName("test_vector_store")
.withDatabaseName("default")
.withIndexType(IndexType.IVF_FLAT)
.withMetricType(metricType)
.withMetricType(this.metricType)
.build();
return new MilvusVectorStore(milvusClient, embeddingModel, config, true, new TokenCountBatchingStrategy());
}
@@ -288,4 +295,4 @@ public class MilvusVectorStoreIT {
}
}
}

33
vector-stores/spring-ai-milvus-store/src/test/java/org/springframework/ai/vectorstore/MilvusVectorStoreObservationIT.java
View File

@@ -1,11 +1,11 @@
/*
* Copyright 2024 the original author or authors.
* Copyright 2023-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
* 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,
@@ -13,16 +13,26 @@
* See the License for the specific language governing permissions and
* limitations under the License.
*/
package org.springframework.ai.vectorstore;
import static org.assertj.core.api.Assertions.assertThat;
package org.springframework.ai.vectorstore;
import java.io.IOException;
import java.nio.charset.StandardCharsets;
import java.util.List;
import java.util.Map;
import io.micrometer.observation.ObservationRegistry;
import io.micrometer.observation.tck.TestObservationRegistry;
import io.micrometer.observation.tck.TestObservationRegistryAssert;
import io.milvus.client.MilvusServiceClient;
import io.milvus.param.ConnectParam;
import io.milvus.param.IndexType;
import io.milvus.param.MetricType;
import org.junit.jupiter.api.Test;
import org.testcontainers.junit.jupiter.Container;
import org.testcontainers.junit.jupiter.Testcontainers;
import org.testcontainers.milvus.MilvusContainer;
import org.springframework.ai.document.Document;
import org.springframework.ai.embedding.EmbeddingModel;
import org.springframework.ai.embedding.TokenCountBatchingStrategy;
@@ -40,17 +50,8 @@ import org.springframework.boot.autoconfigure.EnableAutoConfiguration;
import org.springframework.boot.test.context.runner.ApplicationContextRunner;
import org.springframework.context.annotation.Bean;
import org.springframework.core.io.DefaultResourceLoader;
import org.testcontainers.junit.jupiter.Container;
import org.testcontainers.junit.jupiter.Testcontainers;
import org.testcontainers.milvus.MilvusContainer;
import io.micrometer.observation.ObservationRegistry;
import io.micrometer.observation.tck.TestObservationRegistry;
import io.micrometer.observation.tck.TestObservationRegistryAssert;
import io.milvus.client.MilvusServiceClient;
import io.milvus.param.ConnectParam;
import io.milvus.param.IndexType;
import io.milvus.param.MetricType;
import static org.assertj.core.api.Assertions.assertThat;
/**
* @author Christian Tzolov
@@ -85,13 +86,13 @@ public class MilvusVectorStoreObservationIT {
@Test
void observationVectorStoreAddAndQueryOperations() {
contextRunner.run(context -> {
this.contextRunner.run(context -> {
VectorStore vectorStore = context.getBean(VectorStore.class);
TestObservationRegistry observationRegistry = context.getBean(TestObservationRegistry.class);
vectorStore.add(documents);
vectorStore.add(this.documents);
TestObservationRegistryAssert.assertThat(observationRegistry)
.doesNotHaveAnyRemainingCurrentObservation()

16
vector-stores/spring-ai-mongodb-atlas-store/pom.xml
View File

@@ -1,4 +1,20 @@
<?xml version="1.0" encoding="UTF-8"?>
<!--
~ Copyright 2023-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.
-->
<project xmlns="http://maven.apache.org/POM/4.0.0"
xmlns:xsi="http://www.w3.org/2001/XMLSchema-instance" xsi:schemaLocation="http://maven.apache.org/POM/4.0.0 http://maven.apache.org/maven-v4_0_0.xsd">
<modelVersion>4.0.0</modelVersion>

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

@@ -1,11 +1,11 @@
/*
* Copyright 2023 - 2024 the original author or authors.
* Copyright 2023-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
* 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,
@@ -13,6 +13,7 @@
* See the License for the specific language governing permissions and
* limitations under the License.
*/
package org.springframework.ai.vectorstore;
import org.springframework.ai.vectorstore.filter.Filter;

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

@@ -5,7 +5,7 @@
* 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
* 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,
@@ -16,14 +16,15 @@
package org.springframework.ai.vectorstore;
import static org.springframework.data.mongodb.core.query.Criteria.where;
import java.util.ArrayList;
import java.util.Collections;
import java.util.List;
import java.util.Map;
import java.util.Optional;
import com.mongodb.MongoCommandException;
import io.micrometer.observation.ObservationRegistry;
import org.springframework.ai.document.Document;
import org.springframework.ai.embedding.BatchingStrategy;
import org.springframework.ai.embedding.EmbeddingModel;
@@ -42,9 +43,7 @@ import org.springframework.data.mongodb.core.query.Criteria;
import org.springframework.data.mongodb.core.query.Query;
import org.springframework.util.Assert;
import com.mongodb.MongoCommandException;
import io.micrometer.observation.ObservationRegistry;
import static org.springframework.data.mongodb.core.query.Criteria.where;
/**
* @author Chris Smith
@@ -119,8 +118,8 @@ public class MongoDBAtlasVectorStore extends AbstractObservationVectorStore impl
}
// Create the collection if it does not exist
if (!mongoTemplate.collectionExists(this.config.collectionName)) {
mongoTemplate.createCollection(this.config.collectionName);
if (!this.mongoTemplate.collectionExists(this.config.collectionName)) {
this.mongoTemplate.createCollection(this.config.collectionName);
}
// Create search index
createSearchIndex();
@@ -128,7 +127,7 @@ public class MongoDBAtlasVectorStore extends AbstractObservationVectorStore impl
private void createSearchIndex() {
try {
mongoTemplate.executeCommand(createSearchIndexDefinition());
this.mongoTemplate.executeCommand(createSearchIndexDefinition());
}
catch (UncategorizedMongoDbException e) {
Throwable cause = e.getCause();
@@ -228,6 +227,15 @@ public class MongoDBAtlasVectorStore extends AbstractObservationVectorStore impl
.toList();
}
@Override
public VectorStoreObservationContext.Builder createObservationContextBuilder(String operationName) {
return VectorStoreObservationContext.builder(VectorStoreProvider.MONGODB.value(), operationName)
.withCollectionName(this.config.collectionName)
.withDimensions(this.embeddingModel.dimensions())
.withFieldName(this.config.pathName);
}
public static class MongoDBVectorStoreConfig {
private final String collectionName;
@@ -324,13 +332,4 @@ public class MongoDBAtlasVectorStore extends AbstractObservationVectorStore impl
}
@Override
public VectorStoreObservationContext.Builder createObservationContextBuilder(String operationName) {
return VectorStoreObservationContext.builder(VectorStoreProvider.MONGODB.value(), operationName)
.withCollectionName(this.config.collectionName)
.withDimensions(this.embeddingModel.dimensions())
.withFieldName(this.config.pathName);
}
}
}

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

@@ -1,11 +1,11 @@
/*
* Copyright 2023 - 2024 the original author or authors.
* Copyright 2023-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
* 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,
@@ -13,11 +13,13 @@
* See the License for the specific language governing permissions and
* limitations under the License.
*/
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;
@@ -28,15 +30,16 @@ record VectorSearchAggregation(List<Float> embeddings, String path, int numCandi
@SuppressWarnings("null")
@Override
public org.bson.Document toDocument(@NonNull AggregationOperationContext context) {
var vectorSearch = new Document("queryVector", embeddings).append("path", path)
.append("numCandidates", numCandidates)
.append("index", index)
.append("limit", count);
if (!filter.isEmpty()) {
vectorSearch.append("filter", Document.parse(filter));
var vectorSearch = new Document("queryVector", this.embeddings).append("path", this.path)
.append("numCandidates", this.numCandidates)
.append("index", this.index)
.append("limit", this.count);
if (!this.filter.isEmpty()) {
vectorSearch.append("filter", Document.parse(this.filter));
}
var doc = new Document("$vectorSearch", vectorSearch);
return context.getMappedObject(doc);
}
}
}

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

@@ -1,11 +1,11 @@
/*
* Copyright 2023 - 2024 the original author or authors.
* Copyright 2023-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
* 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,
@@ -13,6 +13,7 @@
* See the License for the specific language governing permissions and
* limitations under the License.
*/
package org.springframework.ai.vectorstore;
import java.util.List;
@@ -45,14 +46,14 @@ public class MongoDBAtlasFilterConverterTest {
@Test
public void testEQ() {
// country == "BG"
String vectorExpr = converter.convertExpression(new Expression(EQ, new Key("country"), new Value("BG")));
String vectorExpr = this.converter.convertExpression(new Expression(EQ, new Key("country"), new Value("BG")));
assertThat(vectorExpr).isEqualTo("{\"metadata.country\":{$eq:\"BG\"}}");
}
@Test
public void tesEqAndGte() {
// genre == "drama" AND year >= 2020
String vectorExpr = converter
String vectorExpr = this.converter
.convertExpression(new Expression(AND, new Expression(EQ, new Key("genre"), new Value("drama")),
new Expression(GTE, new Key("year"), new Value(2020))));
assertThat(vectorExpr)
@@ -62,7 +63,7 @@ public class MongoDBAtlasFilterConverterTest {
@Test
public void tesIn() {
// genre in ["comedy", "documentary", "drama"]
String vectorExpr = converter.convertExpression(
String vectorExpr = this.converter.convertExpression(
new Expression(IN, new Key("genre"), new Value(List.of("comedy", "documentary", "drama"))));
assertThat(vectorExpr).isEqualTo("{\"metadata.genre\":{$in:[\"comedy\",\"documentary\",\"drama\"]}}");
}
@@ -70,7 +71,7 @@ public class MongoDBAtlasFilterConverterTest {
@Test
public void testNe() {
// year >= 2020 OR country == "BG" AND city != "Sofia"
String vectorExpr = converter
String vectorExpr = this.converter
.convertExpression(new Expression(OR, new Expression(GTE, new Key("year"), new Value(2020)),
new Expression(AND, new Expression(EQ, new Key("country"), new Value("BG")),
new Expression(NE, new Key("city"), new Value("Sofia")))));
@@ -81,7 +82,7 @@ public class MongoDBAtlasFilterConverterTest {
@Test
public void testGroup() {
// (year >= 2020 OR country == "BG") AND city NIN ["Sofia", "Plovdiv"]
String vectorExpr = converter.convertExpression(new Expression(AND,
String vectorExpr = this.converter.convertExpression(new Expression(AND,
new Group(new Expression(OR, new Expression(GTE, new Key("year"), new Value(2020)),
new Expression(EQ, new Key("country"), new Value("BG")))),
new Expression(NIN, new Key("city"), new Value(List.of("Sofia", "Plovdiv")))));
@@ -92,7 +93,7 @@ public class MongoDBAtlasFilterConverterTest {
@Test
public void testBoolean() {
// isOpen == true AND year >= 2020 AND country IN ["BG", "NL", "US"]
String vectorExpr = converter.convertExpression(new Expression(AND,
String vectorExpr = this.converter.convertExpression(new Expression(AND,
new Expression(AND, new Expression(EQ, new Key("isOpen"), new Value(true)),
new Expression(GTE, new Key("year"), new Value(2020))),
new Expression(IN, new Key("country"), new Value(List.of("BG", "NL", "US")))));
@@ -104,7 +105,7 @@ public class MongoDBAtlasFilterConverterTest {
@Test
public void testDecimal() {
// temperature >= -15.6 && temperature <= +20.13
String vectorExpr = converter
String vectorExpr = this.converter
.convertExpression(new Expression(AND, new Expression(GTE, new Key("temperature"), new Value(-15.6)),
new Expression(LTE, new Key("temperature"), new Value(20.13))));
@@ -114,11 +115,11 @@ public class MongoDBAtlasFilterConverterTest {
@Test
public void testComplexIdentifiers() {
String vectorExpr = converter
String vectorExpr = this.converter
.convertExpression(new Expression(EQ, new Key("\"country 1 2 3\""), new Value("BG")));
assertThat(vectorExpr).isEqualTo("{\"metadata.country 1 2 3\":{$eq:\"BG\"}}");
vectorExpr = converter.convertExpression(new Expression(EQ, new Key("'country 1 2 3'"), new Value("BG")));
vectorExpr = this.converter.convertExpression(new Expression(EQ, new Key("'country 1 2 3'"), new Value("BG")));
assertThat(vectorExpr).isEqualTo("{\"metadata.country 1 2 3\":{$eq:\"BG\"}}");
}

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

@@ -5,7 +5,7 @@
* 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
* 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,
@@ -13,11 +13,22 @@
* See the License for the specific language governing permissions and
* limitations under the License.
*/
package org.springframework.ai.vectorstore;
import java.util.Arrays;
import java.util.Collections;
import java.util.List;
import java.util.Map;
import java.util.UUID;
import java.util.stream.Collectors;
import com.mongodb.client.MongoClient;
import org.junit.jupiter.api.BeforeEach;
import org.junit.jupiter.api.Test;
import org.testcontainers.junit.jupiter.Container;
import org.testcontainers.junit.jupiter.Testcontainers;
import org.testcontainers.mongodb.MongoDBAtlasLocalContainer;
import org.springframework.ai.document.Document;
import org.springframework.ai.embedding.EmbeddingModel;
@@ -34,17 +45,6 @@ import org.springframework.data.mongodb.core.convert.MongoCustomConversions;
import org.springframework.data.mongodb.core.mapping.MongoMappingContext;
import org.springframework.util.MimeType;
import org.testcontainers.junit.jupiter.Container;
import org.testcontainers.junit.jupiter.Testcontainers;
import org.testcontainers.mongodb.MongoDBAtlasLocalContainer;
import java.util.Arrays;
import java.util.Collections;
import java.util.List;
import java.util.Map;
import java.util.UUID;
import java.util.stream.Collectors;
import static org.assertj.core.api.Assertions.assertThat;
/**
@@ -66,7 +66,7 @@ class MongoDBAtlasVectorStoreIT {
@BeforeEach
public void beforeEach() {
contextRunner.run(context -> {
this.contextRunner.run(context -> {
MongoTemplate mongoTemplate = context.getBean(MongoTemplate.class);
mongoTemplate.getCollection("vector_store").deleteMany(new org.bson.Document());
});
@@ -74,7 +74,7 @@ class MongoDBAtlasVectorStoreIT {
@Test
void vectorStoreTest() {
contextRunner.run(context -> {
this.contextRunner.run(context -> {
VectorStore vectorStore = context.getBean(VectorStore.class);
List<Document> documents = List.of(
@@ -109,7 +109,7 @@ class MongoDBAtlasVectorStoreIT {
@Test
void documentUpdateTest() {
contextRunner.run(context -> {
this.contextRunner.run(context -> {
VectorStore vectorStore = context.getBean(VectorStore.class);
Document document = new Document(UUID.randomUUID().toString(), "Spring AI rocks!!",
@@ -144,7 +144,7 @@ class MongoDBAtlasVectorStoreIT {
@Test
void searchWithFilters() {
contextRunner.run(context -> {
this.contextRunner.run(context -> {
VectorStore vectorStore = context.getBean(VectorStore.class);
var bgDocument = new Document("The World is Big and Salvation Lurks Around the Corner",
@@ -228,6 +228,7 @@ class MongoDBAtlasVectorStoreIT {
@Bean
public Converter<MimeType, String> mimeTypeToStringConverter() {
return new Converter<MimeType, String>() {
@Override
public String convert(MimeType source) {
return source.toString();
@@ -238,6 +239,7 @@ class MongoDBAtlasVectorStoreIT {
@Bean
public Converter<String, MimeType> stringToMimeTypeConverter() {
return new Converter<String, MimeType>() {
@Override
public MimeType convert(String source) {
return MimeType.valueOf(source);
@@ -253,4 +255,4 @@ class MongoDBAtlasVectorStoreIT {
}
}
}

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

@@ -1,11 +1,11 @@
/*
* Copyright 2024 the original author or authors.
* Copyright 2023-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
* 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,
@@ -13,6 +13,7 @@
* See the License for the specific language governing permissions and
* limitations under the License.
*/
package org.springframework.ai.vectorstore;
import org.testcontainers.utility.DockerImageName;

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

@@ -1,11 +1,11 @@
/*
* Copyright 2024 the original author or authors.
* Copyright 2023-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
* 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,
@@ -13,9 +13,8 @@
* See the License for the specific language governing permissions and
* limitations under the License.
*/
package org.springframework.ai.vectorstore;
import static org.assertj.core.api.Assertions.assertThat;
package org.springframework.ai.vectorstore;
import java.io.IOException;
import java.nio.charset.StandardCharsets;
@@ -23,9 +22,17 @@ import java.util.Arrays;
import java.util.List;
import java.util.Map;
import com.mongodb.client.MongoClient;
import io.micrometer.observation.ObservationRegistry;
import io.micrometer.observation.tck.TestObservationRegistry;
import io.micrometer.observation.tck.TestObservationRegistryAssert;
import org.junit.jupiter.api.BeforeEach;
import org.junit.jupiter.api.Test;
import org.junit.jupiter.api.condition.EnabledIfEnvironmentVariable;
import org.testcontainers.junit.jupiter.Container;
import org.testcontainers.junit.jupiter.Testcontainers;
import org.testcontainers.mongodb.MongoDBAtlasLocalContainer;
import org.springframework.ai.document.Document;
import org.springframework.ai.embedding.EmbeddingModel;
import org.springframework.ai.embedding.TokenCountBatchingStrategy;
@@ -48,15 +55,7 @@ import org.springframework.data.mongodb.core.convert.MongoCustomConversions;
import org.springframework.data.mongodb.core.mapping.MongoMappingContext;
import org.springframework.util.MimeType;
import org.testcontainers.junit.jupiter.Container;
import org.testcontainers.junit.jupiter.Testcontainers;
import com.mongodb.client.MongoClient;
import io.micrometer.observation.ObservationRegistry;
import io.micrometer.observation.tck.TestObservationRegistry;
import io.micrometer.observation.tck.TestObservationRegistryAssert;
import org.testcontainers.mongodb.MongoDBAtlasLocalContainer;
import static org.assertj.core.api.Assertions.assertThat;
/**
* @author Christian Tzolov
@@ -93,7 +92,7 @@ public class MongoDbVectorStoreObservationIT {
@BeforeEach
public void beforeEach() {
contextRunner.run(context -> {
this.contextRunner.run(context -> {
MongoTemplate mongoTemplate = context.getBean(MongoTemplate.class);
mongoTemplate.getCollection("vector_store").deleteMany(new org.bson.Document());
});
@@ -102,13 +101,13 @@ public class MongoDbVectorStoreObservationIT {
@Test
void observationVectorStoreAddAndQueryOperations() {
contextRunner.run(context -> {
this.contextRunner.run(context -> {
VectorStore vectorStore = context.getBean(VectorStore.class);
TestObservationRegistry observationRegistry = context.getBean(TestObservationRegistry.class);
vectorStore.add(documents);
vectorStore.add(this.documents);
Thread.sleep(5000);
@@ -212,6 +211,7 @@ public class MongoDbVectorStoreObservationIT {
@Bean
public Converter<MimeType, String> mimeTypeToStringConverter() {
return new Converter<MimeType, String>() {
@Override
public String convert(MimeType source) {
return source.toString();
@@ -222,6 +222,7 @@ public class MongoDbVectorStoreObservationIT {
@Bean
public Converter<String, MimeType> stringToMimeTypeConverter() {
return new Converter<String, MimeType>() {
@Override
public MimeType convert(String source) {
return MimeType.valueOf(source);

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

@@ -1,11 +1,11 @@
/*
* Copyright 2023 - 2024 the original author or authors.
* Copyright 2023-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
* 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,
@@ -13,13 +13,15 @@
* See the License for the specific language governing permissions and
* limitations under the License.
*/
package org.springframework.ai.vectorstore;
import java.util.List;
import org.bson.Document;
import org.junit.jupiter.api.Test;
import org.springframework.data.mongodb.core.aggregation.Aggregation;
import java.util.List;
import org.springframework.data.mongodb.core.aggregation.Aggregation;
import static org.junit.jupiter.api.Assertions.assertEquals;
@@ -60,4 +62,4 @@ class VectorSearchAggregationTest {
assertEquals(expected, document);
}
}
}

16
vector-stores/spring-ai-neo4j-store/pom.xml
View File

@@ -1,4 +1,20 @@
<?xml version="1.0" encoding="UTF-8"?>
<!--
~ Copyright 2023-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.
-->
<project xmlns="http://maven.apache.org/POM/4.0.0"
xmlns:xsi="http://www.w3.org/2001/XMLSchema-instance" xsi:schemaLocation="http://maven.apache.org/POM/4.0.0 http://maven.apache.org/maven-v4_0_0.xsd">
<modelVersion>4.0.0</modelVersion>

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

@@ -5,7 +5,7 @@
* 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
* 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,
@@ -22,10 +22,12 @@ import java.util.Map;
import java.util.Optional;
import java.util.function.Predicate;
import io.micrometer.observation.ObservationRegistry;
import org.neo4j.cypherdsl.support.schema_name.SchemaNames;
import org.neo4j.driver.Driver;
import org.neo4j.driver.SessionConfig;
import org.neo4j.driver.Values;
import org.springframework.ai.document.Document;
import org.springframework.ai.embedding.BatchingStrategy;
import org.springframework.ai.embedding.EmbeddingModel;
@@ -40,8 +42,6 @@ import org.springframework.ai.vectorstore.observation.VectorStoreObservationConv
import org.springframework.beans.factory.InitializingBean;
import org.springframework.util.Assert;
import io.micrometer.observation.ObservationRegistry;
/**
* @author Gerrit Meier
* @author Michael Simons
@@ -51,222 +51,6 @@ import io.micrometer.observation.ObservationRegistry;
*/
public class Neo4jVectorStore extends AbstractObservationVectorStore implements InitializingBean {
/**
* An enum to configure the distance function used in the Neo4j vector index.
*/
public enum Neo4jDistanceType {
COSINE("cosine"), EUCLIDEAN("euclidean");
public final String name;
Neo4jDistanceType(String name) {
this.name = name;
}
}
/**
* Configuration for the Neo4j vector store.
*/
public static final class Neo4jVectorStoreConfig {
private final SessionConfig sessionConfig;
private final int embeddingDimension;
private final Neo4jDistanceType distanceType;
private final String embeddingProperty;
private final String label;
private final String indexName;
// needed for similarity search call
private final String indexNameNotSanitized;
private final String idProperty;
private final String constraintName;
/**
* Start building a new configuration.
* @return The entry point for creating a new configuration.
*/
public static Builder builder() {
return new Builder();
}
/**
* {@return the default config}
*/
public static Neo4jVectorStoreConfig defaultConfig() {
return builder().build();
}
private Neo4jVectorStoreConfig(Builder builder) {
this.sessionConfig = Optional.ofNullable(builder.databaseName)
.filter(Predicate.not(String::isBlank))
.map(SessionConfig::forDatabase)
.orElseGet(SessionConfig::defaultConfig);
this.embeddingDimension = builder.embeddingDimension;
this.distanceType = builder.distanceType;
this.embeddingProperty = SchemaNames.sanitize(builder.embeddingProperty).orElseThrow();
this.label = SchemaNames.sanitize(builder.label).orElseThrow();
this.indexNameNotSanitized = builder.indexName;
this.indexName = SchemaNames.sanitize(builder.indexName, true).orElseThrow();
this.constraintName = SchemaNames.sanitize(builder.constraintName).orElseThrow();
this.idProperty = SchemaNames.sanitize(builder.idProperty).orElseThrow();
}
public static class Builder {
private String databaseName;
private int embeddingDimension = DEFAULT_EMBEDDING_DIMENSION;
private Neo4jDistanceType distanceType = Neo4jDistanceType.COSINE;
private String label = DEFAULT_LABEL;
private String embeddingProperty = DEFAULT_EMBEDDING_PROPERTY;
private String indexName = DEFAULT_INDEX_NAME;
private String idProperty = DEFAULT_ID_PROPERTY;
private String constraintName = DEFAULT_CONSTRAINT_NAME;
private Builder() {
}
/**
* Configures the Neo4j database name to use. Leave {@literal null} or blank
* to use the default database.
* @param databaseName the database name to use
* @return this builder
*/
public Builder withDatabaseName(String databaseName) {
this.databaseName = databaseName;
return this;
}
/**
* Configures the size of the embedding. Defaults to {@literal 1536}, inline
* with OpenAIs embeddings.
* @param newEmbeddingDimension The dimension of the embedding
* @return this builder
*/
public Builder withEmbeddingDimension(int newEmbeddingDimension) {
Assert.isTrue(newEmbeddingDimension >= 1, "Dimension has to be positive.");
this.embeddingDimension = newEmbeddingDimension;
return this;
}
/**
* Configures the distance type to store in the index and to use in queries.
* @param newDistanceType The distance type, must not be {@literal null}
* @return this builder
*/
public Builder withDistanceType(Neo4jDistanceType newDistanceType) {
Assert.notNull(newDistanceType, "Distance type may not be null");
this.distanceType = newDistanceType;
return this;
}
/**
* Configures the node label to use for storing documents. Defaults to
* {@literal Document}.
* @param newLabel The label used on the nodes representing the document
* @return this builder
*/
public Builder withLabel(String newLabel) {
Assert.hasText(newLabel, "Content label may not be null or blank");
this.label = newLabel;
return this;
}
/**
* Configures the property of the node to use for storing embedding. Defaults
* to {@literal embedding}.
* @param newEmbeddingProperty The property of the nodes for storing the
* embedding
* @return this builder
*/
public Builder withEmbeddingProperty(String newEmbeddingProperty) {
Assert.hasText(newEmbeddingProperty, "Embedding property may not be null or blank");
this.embeddingProperty = newEmbeddingProperty;
return this;
}
/**
* Configures the vector index to be used. Defaults to
* {@literal spring-ai-document-index}.
* @param newIndexName The name of the index to be used for storing and
* searching data.
* @return this builder
*/
public Builder withIndexName(String newIndexName) {
Assert.hasText(newIndexName, "Index name may not be null or blank");
this.indexName = newIndexName;
return this;
}
/**
* Configures the id property to be used. Defaults to {@literal id}.
* @param newIdProperty The name of the id property of the {@link Document}
* entity
* @return this builder
*/
public Builder withIdProperty(String newIdProperty) {
Assert.hasText(newIdProperty, "Id property may not be null or blank");
this.idProperty = newIdProperty;
return this;
}
/**
* Configures the constraint name to be used. Defaults to
* {@literal Document_unique_idx}.
* @param newConstraintName The name of the unique constraint for the id
* property.
* @return this builder
*/
public Builder withConstraintName(String newConstraintName) {
Assert.hasText(newConstraintName, "Constraint name may not be null or blank");
this.constraintName = newConstraintName;
return this;
}
/**
* {@return the immutable configuration}
*/
public Neo4jVectorStoreConfig build() {
return new Neo4jVectorStoreConfig(this);
}
}
}
public static final int DEFAULT_EMBEDDING_DIMENSION = 1536;
public static final String DEFAULT_LABEL = "Document";
@@ -279,6 +63,10 @@ public class Neo4jVectorStore extends AbstractObservationVectorStore implements
public static final String DEFAULT_CONSTRAINT_NAME = DEFAULT_LABEL + "_unique_idx";
private static Map<Neo4jDistanceType, VectorStoreSimilarityMetric> SIMILARITY_TYPE_MAPPING = Map.of(
Neo4jDistanceType.COSINE, VectorStoreSimilarityMetric.COSINE, Neo4jDistanceType.EUCLIDEAN,
VectorStoreSimilarityMetric.EUCLIDEAN);
private final Neo4jVectorFilterExpressionConverter filterExpressionConverter = new Neo4jVectorFilterExpressionConverter();
private final Driver driver;
@@ -445,10 +233,6 @@ public class Neo4jVectorStore extends AbstractObservationVectorStore implements
.withSimilarityMetric(getSimilarityMetric());
}
private static Map<Neo4jDistanceType, VectorStoreSimilarityMetric> SIMILARITY_TYPE_MAPPING = Map.of(
Neo4jDistanceType.COSINE, VectorStoreSimilarityMetric.COSINE, Neo4jDistanceType.EUCLIDEAN,
VectorStoreSimilarityMetric.EUCLIDEAN);
private String getSimilarityMetric() {
if (!SIMILARITY_TYPE_MAPPING.containsKey(this.config.distanceType)) {
return this.config.distanceType.name();
@@ -456,4 +240,220 @@ public class Neo4jVectorStore extends AbstractObservationVectorStore implements
return SIMILARITY_TYPE_MAPPING.get(this.config.distanceType).value();
}
}
/**
* An enum to configure the distance function used in the Neo4j vector index.
*/
public enum Neo4jDistanceType {
COSINE("cosine"), EUCLIDEAN("euclidean");
public final String name;
Neo4jDistanceType(String name) {
this.name = name;
}
}
/**
* Configuration for the Neo4j vector store.
*/
public static final class Neo4jVectorStoreConfig {
private final SessionConfig sessionConfig;
private final int embeddingDimension;
private final Neo4jDistanceType distanceType;
private final String embeddingProperty;
private final String label;
private final String indexName;
// needed for similarity search call
private final String indexNameNotSanitized;
private final String idProperty;
private final String constraintName;
private Neo4jVectorStoreConfig(Builder builder) {
this.sessionConfig = Optional.ofNullable(builder.databaseName)
.filter(Predicate.not(String::isBlank))
.map(SessionConfig::forDatabase)
.orElseGet(SessionConfig::defaultConfig);
this.embeddingDimension = builder.embeddingDimension;
this.distanceType = builder.distanceType;
this.embeddingProperty = SchemaNames.sanitize(builder.embeddingProperty).orElseThrow();
this.label = SchemaNames.sanitize(builder.label).orElseThrow();
this.indexNameNotSanitized = builder.indexName;
this.indexName = SchemaNames.sanitize(builder.indexName, true).orElseThrow();
this.constraintName = SchemaNames.sanitize(builder.constraintName).orElseThrow();
this.idProperty = SchemaNames.sanitize(builder.idProperty).orElseThrow();
}
/**
* Start building a new configuration.
* @return The entry point for creating a new configuration.
*/
public static Builder builder() {
return new Builder();
}
/**
* {@return the default config}
*/
public static Neo4jVectorStoreConfig defaultConfig() {
return builder().build();
}
public static class Builder {
private String databaseName;
private int embeddingDimension = DEFAULT_EMBEDDING_DIMENSION;
private Neo4jDistanceType distanceType = Neo4jDistanceType.COSINE;
private String label = DEFAULT_LABEL;
private String embeddingProperty = DEFAULT_EMBEDDING_PROPERTY;
private String indexName = DEFAULT_INDEX_NAME;
private String idProperty = DEFAULT_ID_PROPERTY;
private String constraintName = DEFAULT_CONSTRAINT_NAME;
private Builder() {
}
/**
* Configures the Neo4j database name to use. Leave {@literal null} or blank
* to use the default database.
* @param databaseName the database name to use
* @return this builder
*/
public Builder withDatabaseName(String databaseName) {
this.databaseName = databaseName;
return this;
}
/**
* Configures the size of the embedding. Defaults to {@literal 1536}, inline
* with OpenAIs embeddings.
* @param newEmbeddingDimension The dimension of the embedding
* @return this builder
*/
public Builder withEmbeddingDimension(int newEmbeddingDimension) {
Assert.isTrue(newEmbeddingDimension >= 1, "Dimension has to be positive.");
this.embeddingDimension = newEmbeddingDimension;
return this;
}
/**
* Configures the distance type to store in the index and to use in queries.
* @param newDistanceType The distance type, must not be {@literal null}
* @return this builder
*/
public Builder withDistanceType(Neo4jDistanceType newDistanceType) {
Assert.notNull(newDistanceType, "Distance type may not be null");
this.distanceType = newDistanceType;
return this;
}
/**
* Configures the node label to use for storing documents. Defaults to
* {@literal Document}.
* @param newLabel The label used on the nodes representing the document
* @return this builder
*/
public Builder withLabel(String newLabel) {
Assert.hasText(newLabel, "Content label may not be null or blank");
this.label = newLabel;
return this;
}
/**
* Configures the property of the node to use for storing embedding. Defaults
* to {@literal embedding}.
* @param newEmbeddingProperty The property of the nodes for storing the
* embedding
* @return this builder
*/
public Builder withEmbeddingProperty(String newEmbeddingProperty) {
Assert.hasText(newEmbeddingProperty, "Embedding property may not be null or blank");
this.embeddingProperty = newEmbeddingProperty;
return this;
}
/**
* Configures the vector index to be used. Defaults to
* {@literal spring-ai-document-index}.
* @param newIndexName The name of the index to be used for storing and
* searching data.
* @return this builder
*/
public Builder withIndexName(String newIndexName) {
Assert.hasText(newIndexName, "Index name may not be null or blank");
this.indexName = newIndexName;
return this;
}
/**
* Configures the id property to be used. Defaults to {@literal id}.
* @param newIdProperty The name of the id property of the {@link Document}
* entity
* @return this builder
*/
public Builder withIdProperty(String newIdProperty) {
Assert.hasText(newIdProperty, "Id property may not be null or blank");
this.idProperty = newIdProperty;
return this;
}
/**
* Configures the constraint name to be used. Defaults to
* {@literal Document_unique_idx}.
* @param newConstraintName The name of the unique constraint for the id
* property.
* @return this builder
*/
public Builder withConstraintName(String newConstraintName) {
Assert.hasText(newConstraintName, "Constraint name may not be null or blank");
this.constraintName = newConstraintName;
return this;
}
/**
* {@return the immutable configuration}
*/
public Neo4jVectorStoreConfig build() {
return new Neo4jVectorStoreConfig(this);
}
}
}
}

5
vector-stores/spring-ai-neo4j-store/src/main/java/org/springframework/ai/vectorstore/filter/Neo4jVectorFilterExpressionConverter.java
View File

@@ -1,11 +1,11 @@
/*
* Copyright 2023 - 2024 the original author or authors.
* Copyright 2023-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
* 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,
@@ -13,6 +13,7 @@
* See the License for the specific language governing permissions and
* limitations under the License.
*/
package org.springframework.ai.vectorstore.filter;
import org.springframework.ai.vectorstore.filter.Filter.Expression;

5
vector-stores/spring-ai-neo4j-store/src/test/java/org/springframework/ai/vectorstore/Neo4jImage.java
View File

@@ -1,11 +1,11 @@
/*
* Copyright 2024 the original author or authors.
* Copyright 2023-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
* 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,
@@ -13,6 +13,7 @@
* See the License for the specific language governing permissions and
* limitations under the License.
*/
package org.springframework.ai.vectorstore;
import org.testcontainers.utility.DockerImageName;

15
vector-stores/spring-ai-neo4j-store/src/test/java/org/springframework/ai/vectorstore/Neo4jVectorStoreIT.java
View File

@@ -1,11 +1,11 @@
/*
* Copyright 2023 - 2024 the original author or authors.
* Copyright 2023-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
* 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,
@@ -13,6 +13,7 @@
* See the License for the specific language governing permissions and
* limitations under the License.
*/
package org.springframework.ai.vectorstore;
import java.util.Collections;
@@ -27,15 +28,15 @@ import org.junit.jupiter.api.condition.EnabledIfEnvironmentVariable;
import org.neo4j.driver.AuthTokens;
import org.neo4j.driver.Driver;
import org.neo4j.driver.GraphDatabase;
import org.springframework.ai.vectorstore.filter.FilterExpressionTextParser;
import org.testcontainers.containers.Neo4jContainer;
import org.testcontainers.junit.jupiter.Container;
import org.testcontainers.junit.jupiter.Testcontainers;
import org.springframework.ai.document.Document;
import org.springframework.ai.embedding.EmbeddingModel;
import org.springframework.ai.openai.api.OpenAiApi;
import org.springframework.ai.openai.OpenAiEmbeddingModel;
import org.springframework.ai.openai.api.OpenAiApi;
import org.springframework.ai.vectorstore.filter.FilterExpressionTextParser;
import org.springframework.boot.SpringBootConfiguration;
import org.springframework.boot.autoconfigure.EnableAutoConfiguration;
import org.springframework.boot.autoconfigure.jdbc.DataSourceAutoConfiguration;
@@ -57,6 +58,9 @@ class Neo4jVectorStoreIT {
@Container
static Neo4jContainer<?> neo4jContainer = new Neo4jContainer<>(Neo4jImage.DEFAULT_IMAGE).withRandomPassword();
private final ApplicationContextRunner contextRunner = new ApplicationContextRunner()
.withUserConfiguration(TestApplication.class);
List<Document> documents = List.of(
new Document("Spring AI rocks!! Spring AI rocks!! Spring AI rocks!! Spring AI rocks!! Spring AI rocks!!",
Collections.singletonMap("meta1", "meta1")),
@@ -65,9 +69,6 @@ class Neo4jVectorStoreIT {
"Great Depression Great Depression Great Depression Great Depression Great Depression Great Depression",
Collections.singletonMap("meta2", "meta2")));
private final ApplicationContextRunner contextRunner = new ApplicationContextRunner()
.withUserConfiguration(TestApplication.class);
@BeforeEach
void cleanDatabase() {
this.contextRunner

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

@@ -1,11 +1,11 @@
/*
* Copyright 2024 the original author or authors.
* Copyright 2023-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
* 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,
@@ -13,15 +13,17 @@
* See the License for the specific language governing permissions and
* limitations under the License.
*/
package org.springframework.ai.vectorstore;
import static org.assertj.core.api.Assertions.assertThat;
package org.springframework.ai.vectorstore;
import java.io.IOException;
import java.nio.charset.StandardCharsets;
import java.util.List;
import java.util.Map;
import io.micrometer.observation.ObservationRegistry;
import io.micrometer.observation.tck.TestObservationRegistry;
import io.micrometer.observation.tck.TestObservationRegistryAssert;
import org.junit.jupiter.api.BeforeEach;
import org.junit.jupiter.api.Test;
import org.junit.jupiter.api.condition.EnabledIfEnvironmentVariable;
@@ -29,6 +31,10 @@ import org.neo4j.cypherdsl.support.schema_name.SchemaNames;
import org.neo4j.driver.AuthTokens;
import org.neo4j.driver.Driver;
import org.neo4j.driver.GraphDatabase;
import org.testcontainers.containers.Neo4jContainer;
import org.testcontainers.junit.jupiter.Container;
import org.testcontainers.junit.jupiter.Testcontainers;
import org.springframework.ai.document.Document;
import org.springframework.ai.embedding.EmbeddingModel;
import org.springframework.ai.embedding.TokenCountBatchingStrategy;
@@ -46,13 +52,8 @@ import org.springframework.boot.autoconfigure.jdbc.DataSourceAutoConfiguration;
import org.springframework.boot.test.context.runner.ApplicationContextRunner;
import org.springframework.context.annotation.Bean;
import org.springframework.core.io.DefaultResourceLoader;
import org.testcontainers.containers.Neo4jContainer;
import org.testcontainers.junit.jupiter.Container;
import org.testcontainers.junit.jupiter.Testcontainers;
import io.micrometer.observation.ObservationRegistry;
import io.micrometer.observation.tck.TestObservationRegistry;
import io.micrometer.observation.tck.TestObservationRegistryAssert;
import static org.assertj.core.api.Assertions.assertThat;
/**
* @author Christian Tzolov
@@ -92,13 +93,13 @@ public class Neo4jVectorStoreObservationIT {
@Test
void observationVectorStoreAddAndQueryOperations() {
contextRunner.run(context -> {
this.contextRunner.run(context -> {
VectorStore vectorStore = context.getBean(VectorStore.class);
TestObservationRegistry observationRegistry = context.getBean(TestObservationRegistry.class);
vectorStore.add(documents);
vectorStore.add(this.documents);
TestObservationRegistryAssert.assertThat(observationRegistry)
.doesNotHaveAnyRemainingCurrentObservation()

25
vector-stores/spring-ai-neo4j-store/src/test/java/org/springframework/ai/vectorstore/filter/Neo4jVectorFilterExpressionConverterTests.java
View File

@@ -1,11 +1,11 @@
/*
* Copyright 2023 - 2024 the original author or authors.
* Copyright 2023-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
* 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,
@@ -13,6 +13,7 @@
* See the License for the specific language governing permissions and
* limitations under the License.
*/
package org.springframework.ai.vectorstore.filter;
import java.util.List;
@@ -46,14 +47,14 @@ public class Neo4jVectorFilterExpressionConverterTests {
@Test
public void testEQ() {
// country = "BG"
String vectorExpr = converter.convertExpression(new Expression(EQ, new Key("country"), new Value("BG")));
String vectorExpr = this.converter.convertExpression(new Expression(EQ, new Key("country"), new Value("BG")));
assertThat(vectorExpr).isEqualTo("node.`metadata.country` = \"BG\"");
}
@Test
public void tesEqAndGte() {
// genre = "drama" AND year >= 2020
String vectorExpr = converter
String vectorExpr = this.converter
.convertExpression(new Expression(AND, new Expression(EQ, new Key("genre"), new Value("drama")),
new Expression(GTE, new Key("year"), new Value(2020))));
assertThat(vectorExpr).isEqualTo("node.`metadata.genre` = \"drama\" AND node.`metadata.year` >= 2020");
@@ -62,7 +63,7 @@ public class Neo4jVectorFilterExpressionConverterTests {
@Test
public void tesIn() {
// genre in ["comedy", "documentary", "drama"]
String vectorExpr = converter.convertExpression(
String vectorExpr = this.converter.convertExpression(
new Expression(IN, new Key("genre"), new Value(List.of("comedy", "documentary", "drama"))));
assertThat(vectorExpr).isEqualTo("node.`metadata.genre` IN [\"comedy\",\"documentary\",\"drama\"]");
}
@@ -70,7 +71,7 @@ public class Neo4jVectorFilterExpressionConverterTests {
@Test
public void tesNIn() {
// genre in ["comedy", "documentary", "drama"]
String vectorExpr = converter.convertExpression(
String vectorExpr = this.converter.convertExpression(
new Expression(NIN, new Key("genre"), new Value(List.of("comedy", "documentary", "drama"))));
assertThat(vectorExpr).isEqualTo("NOT node.`metadata.genre` IN [\"comedy\",\"documentary\",\"drama\"]");
}
@@ -78,7 +79,7 @@ public class Neo4jVectorFilterExpressionConverterTests {
@Test
public void testNe() {
// year >= 2020 OR country = "BG" AND city <> "Sofia"
String vectorExpr = converter
String vectorExpr = this.converter
.convertExpression(new Expression(OR, new Expression(GTE, new Key("year"), new Value(2020)),
new Expression(AND, new Expression(EQ, new Key("country"), new Value("BG")),
new Expression(NE, new Key("city"), new Value("Sofia")))));
@@ -89,7 +90,7 @@ public class Neo4jVectorFilterExpressionConverterTests {
@Test
public void testGroup() {
// (year >= 2020 OR country = "BG") AND NOT city IN ["Sofia", "Plovdiv"]
String vectorExpr = converter.convertExpression(new Expression(AND,
String vectorExpr = this.converter.convertExpression(new Expression(AND,
new Group(new Expression(OR, new Expression(GTE, new Key("year"), new Value(2020)),
new Expression(EQ, new Key("country"), new Value("BG")))),
new Expression(NOT, new Expression(IN, new Key("city"), new Value(List.of("Sofia", "Plovdiv"))))));
@@ -100,7 +101,7 @@ public class Neo4jVectorFilterExpressionConverterTests {
@Test
public void testBoolean() {
// isOpen = true AND year >= 2020 AND country IN ["BG", "NL", "US"]
String vectorExpr = converter.convertExpression(new Expression(AND,
String vectorExpr = this.converter.convertExpression(new Expression(AND,
new Expression(AND, new Expression(EQ, new Key("isOpen"), new Value(true)),
new Expression(GTE, new Key("year"), new Value(2020))),
new Expression(IN, new Key("country"), new Value(List.of("BG", "NL", "US")))));
@@ -112,7 +113,7 @@ public class Neo4jVectorFilterExpressionConverterTests {
@Test
public void testDecimal() {
// temperature >= -15.6 AND temperature <= +20.13
String vectorExpr = converter
String vectorExpr = this.converter
.convertExpression(new Expression(AND, new Expression(GTE, new Key("temperature"), new Value(-15.6)),
new Expression(LTE, new Key("temperature"), new Value(20.13))));
@@ -122,7 +123,7 @@ public class Neo4jVectorFilterExpressionConverterTests {
@Test
public void testComplexIdentifiers() {
String vectorExpr = converter
String vectorExpr = this.converter
.convertExpression(new Expression(EQ, new Key("\"country 1 2 3\""), new Value("BG")));
assertThat(vectorExpr).isEqualTo("node.`metadata.country 1 2 3` = \"BG\"");
}
@@ -131,7 +132,7 @@ public class Neo4jVectorFilterExpressionConverterTests {
public void testComplexIdentifiers2() {
Filter.Expression expr = new FilterExpressionTextParser()
.parse("author in ['john', 'jill'] && 'article_type' == 'blog'");
String vectorExpr = converter.convertExpression(expr);
String vectorExpr = this.converter.convertExpression(expr);
assertThat(vectorExpr)
.isEqualTo("node.`metadata.author` IN [\"john\",\"jill\"] AND node.`metadata.'article_type'` = \"blog\"");
}

16
vector-stores/spring-ai-opensearch-store/pom.xml
View File

@@ -1,4 +1,20 @@
<?xml version="1.0" encoding="UTF-8"?>
<!--
~ Copyright 2023-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.
-->
<project xmlns="http://maven.apache.org/POM/4.0.0"
xmlns:xsi="http://www.w3.org/2001/XMLSchema-instance" xsi:schemaLocation="http://maven.apache.org/POM/4.0.0 http://maven.apache.org/maven-v4_0_0.xsd">
<modelVersion>4.0.0</modelVersion>

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

@@ -1,11 +1,11 @@
/*
* Copyright 2023 - 2024 the original author or authors.
* Copyright 2023-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
* 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,
@@ -13,12 +13,8 @@
* See the License for the specific language governing permissions and
* limitations under the License.
*/
package org.springframework.ai.vectorstore;
import org.springframework.ai.vectorstore.filter.Filter;
import org.springframework.ai.vectorstore.filter.Filter.Expression;
import org.springframework.ai.vectorstore.filter.Filter.Key;
import org.springframework.ai.vectorstore.filter.converter.AbstractFilterExpressionConverter;
package org.springframework.ai.vectorstore;
import java.text.ParseException;
import java.text.SimpleDateFormat;
@@ -27,6 +23,11 @@ import java.util.List;
import java.util.TimeZone;
import java.util.regex.Pattern;
import org.springframework.ai.vectorstore.filter.Filter;
import org.springframework.ai.vectorstore.filter.Filter.Expression;
import org.springframework.ai.vectorstore.filter.Filter.Key;
import org.springframework.ai.vectorstore.filter.converter.AbstractFilterExpressionConverter;
/**
* @author Jemin Huh
* @since 1.0.0

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

@@ -5,7 +5,7 @@
* 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
* 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,
@@ -16,6 +16,14 @@
package org.springframework.ai.vectorstore;
import java.io.IOException;
import java.io.StringReader;
import java.util.List;
import java.util.Objects;
import java.util.Optional;
import java.util.stream.Collectors;
import io.micrometer.observation.ObservationRegistry;
import org.opensearch.client.json.JsonData;
import org.opensearch.client.json.JsonpMapper;
import org.opensearch.client.opensearch.OpenSearchClient;
@@ -30,6 +38,7 @@ import org.opensearch.client.opensearch.indices.CreateIndexResponse;
import org.opensearch.client.transport.endpoints.BooleanResponse;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import org.springframework.ai.document.Document;
import org.springframework.ai.embedding.BatchingStrategy;
import org.springframework.ai.embedding.EmbeddingModel;
@@ -46,15 +55,6 @@ import org.springframework.ai.vectorstore.observation.VectorStoreObservationConv
import org.springframework.beans.factory.InitializingBean;
import org.springframework.util.Assert;
import io.micrometer.observation.ObservationRegistry;
import java.io.IOException;
import java.io.StringReader;
import java.util.List;
import java.util.Objects;
import java.util.Optional;
import java.util.stream.Collectors;
/**
* @author Jemin Huh
* @author Soby Chacko
@@ -67,8 +67,6 @@ public class OpenSearchVectorStore extends AbstractObservationVectorStore implem
public static final String COSINE_SIMILARITY_FUNCTION = "cosinesimil";
private static final Logger logger = LoggerFactory.getLogger(OpenSearchVectorStore.class);
public static final String DEFAULT_INDEX_NAME = "spring-ai-document-index";
public static final String DEFAULT_MAPPING_EMBEDDING_TYPE_KNN_VECTOR_DIMENSION_1536 = """
@@ -82,6 +80,8 @@ public class OpenSearchVectorStore extends AbstractObservationVectorStore implem
}
""";
private static final Logger logger = LoggerFactory.getLogger(OpenSearchVectorStore.class);
private final EmbeddingModel embeddingModel;
private final OpenSearchClient openSearchClient;
@@ -92,12 +92,12 @@ public class OpenSearchVectorStore extends AbstractObservationVectorStore implem
private final String mappingJson;
private String similarityFunction;
private final boolean initializeSchema;
private final BatchingStrategy batchingStrategy;
private String similarityFunction;
public OpenSearchVectorStore(OpenSearchClient openSearchClient, EmbeddingModel embeddingModel,
boolean initializeSchema) {
this(openSearchClient, embeddingModel, DEFAULT_MAPPING_EMBEDDING_TYPE_KNN_VECTOR_DIMENSION_1536,
@@ -245,7 +245,7 @@ public class OpenSearchVectorStore extends AbstractObservationVectorStore implem
}
private CreateIndexResponse createIndexMapping(String index, String mappingJson) {
JsonpMapper jsonpMapper = openSearchClient._transport().jsonpMapper();
JsonpMapper jsonpMapper = this.openSearchClient._transport().jsonpMapper();
try {
return this.openSearchClient.indices()
.create(new CreateIndexRequest.Builder().index(index)
@@ -285,4 +285,4 @@ public class OpenSearchVectorStore extends AbstractObservationVectorStore implem
return this.similarityFunction;
}
}
}

42
vector-stores/spring-ai-opensearch-store/src/test/java/org/springframework/ai/vectorstore/OpenSearchAiSearchFilterExpressionConverterTest.java
View File

@@ -1,11 +1,11 @@
/*
* Copyright 2023 - 2024 the original author or authors.
* Copyright 2023-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
* 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,
@@ -13,8 +13,17 @@
* See the License for the specific language governing permissions and
* limitations under the License.
*/
package org.springframework.ai.vectorstore;
import java.util.Date;
import java.util.List;
import org.junit.jupiter.api.Test;
import org.springframework.ai.vectorstore.filter.Filter;
import org.springframework.ai.vectorstore.filter.FilterExpressionConverter;
import static org.assertj.core.api.Assertions.assertThat;
import static org.springframework.ai.vectorstore.filter.Filter.ExpressionType.AND;
import static org.springframework.ai.vectorstore.filter.Filter.ExpressionType.EQ;
@@ -25,38 +34,31 @@ import static org.springframework.ai.vectorstore.filter.Filter.ExpressionType.NE
import static org.springframework.ai.vectorstore.filter.Filter.ExpressionType.NIN;
import static org.springframework.ai.vectorstore.filter.Filter.ExpressionType.OR;
import java.util.Date;
import java.util.List;
import org.junit.jupiter.api.Test;
import org.springframework.ai.vectorstore.filter.Filter;
import org.springframework.ai.vectorstore.filter.FilterExpressionConverter;
class OpenSearchAiSearchFilterExpressionConverterTest {
final FilterExpressionConverter converter = new OpenSearchAiSearchFilterExpressionConverter();
@Test
public void testDate() {
String vectorExpr = converter.convertExpression(new Filter.Expression(EQ, new Filter.Key("activationDate"),
String vectorExpr = this.converter.convertExpression(new Filter.Expression(EQ, new Filter.Key("activationDate"),
new Filter.Value(new Date(1704637752148L))));
assertThat(vectorExpr).isEqualTo("metadata.activationDate:2024-01-07T14:29:12Z");
vectorExpr = converter.convertExpression(
vectorExpr = this.converter.convertExpression(
new Filter.Expression(EQ, new Filter.Key("activationDate"), new Filter.Value("1970-01-01T00:00:02Z")));
assertThat(vectorExpr).isEqualTo("metadata.activationDate:1970-01-01T00:00:02Z");
}
@Test
public void testEQ() {
String vectorExpr = converter
String vectorExpr = this.converter
.convertExpression(new Filter.Expression(EQ, new Filter.Key("country"), new Filter.Value("BG")));
assertThat(vectorExpr).isEqualTo("metadata.country:BG");
}
@Test
public void tesEqAndGte() {
String vectorExpr = converter.convertExpression(new Filter.Expression(AND,
String vectorExpr = this.converter.convertExpression(new Filter.Expression(AND,
new Filter.Expression(EQ, new Filter.Key("genre"), new Filter.Value("drama")),
new Filter.Expression(GTE, new Filter.Key("year"), new Filter.Value(2020))));
assertThat(vectorExpr).isEqualTo("metadata.genre:drama AND metadata.year:>=2020");
@@ -64,14 +66,14 @@ class OpenSearchAiSearchFilterExpressionConverterTest {
@Test
public void tesIn() {
String vectorExpr = converter.convertExpression(new Filter.Expression(IN, new Filter.Key("genre"),
String vectorExpr = this.converter.convertExpression(new Filter.Expression(IN, new Filter.Key("genre"),
new Filter.Value(List.of("comedy", "documentary", "drama"))));
assertThat(vectorExpr).isEqualTo("(metadata.genre:comedy OR documentary OR drama)");
}
@Test
public void testNe() {
String vectorExpr = converter.convertExpression(
String vectorExpr = this.converter.convertExpression(
new Filter.Expression(OR, new Filter.Expression(GTE, new Filter.Key("year"), new Filter.Value(2020)),
new Filter.Expression(AND,
new Filter.Expression(EQ, new Filter.Key("country"), new Filter.Value("BG")),
@@ -81,7 +83,7 @@ class OpenSearchAiSearchFilterExpressionConverterTest {
@Test
public void testGroup() {
String vectorExpr = converter.convertExpression(new Filter.Expression(AND,
String vectorExpr = this.converter.convertExpression(new Filter.Expression(AND,
new Filter.Group(new Filter.Expression(OR,
new Filter.Expression(GTE, new Filter.Key("year"), new Filter.Value(2020)),
new Filter.Expression(EQ, new Filter.Key("country"), new Filter.Value("BG")))),
@@ -92,7 +94,7 @@ class OpenSearchAiSearchFilterExpressionConverterTest {
@Test
public void tesBoolean() {
String vectorExpr = converter.convertExpression(new Filter.Expression(AND,
String vectorExpr = this.converter.convertExpression(new Filter.Expression(AND,
new Filter.Expression(AND, new Filter.Expression(EQ, new Filter.Key("isOpen"), new Filter.Value(true)),
new Filter.Expression(GTE, new Filter.Key("year"), new Filter.Value(2020))),
new Filter.Expression(IN, new Filter.Key("country"), new Filter.Value(List.of("BG", "NL", "US")))));
@@ -103,7 +105,7 @@ class OpenSearchAiSearchFilterExpressionConverterTest {
@Test
public void testDecimal() {
String vectorExpr = converter.convertExpression(new Filter.Expression(AND,
String vectorExpr = this.converter.convertExpression(new Filter.Expression(AND,
new Filter.Expression(GTE, new Filter.Key("temperature"), new Filter.Value(-15.6)),
new Filter.Expression(LTE, new Filter.Key("temperature"), new Filter.Value(20.13))));
@@ -112,11 +114,11 @@ class OpenSearchAiSearchFilterExpressionConverterTest {
@Test
public void testComplexIdentifiers() {
String vectorExpr = converter
String vectorExpr = this.converter
.convertExpression(new Filter.Expression(EQ, new Filter.Key("\"country 1 2 3\""), new Filter.Value("BG")));
assertThat(vectorExpr).isEqualTo("metadata.country 1 2 3:BG");
vectorExpr = converter
vectorExpr = this.converter
.convertExpression(new Filter.Expression(EQ, new Filter.Key("'country 1 2 3'"), new Filter.Value("BG")));
assertThat(vectorExpr).isEqualTo("metadata.country 1 2 3:BG");
}

5
vector-stores/spring-ai-opensearch-store/src/test/java/org/springframework/ai/vectorstore/OpenSearchImage.java
View File

@@ -1,11 +1,11 @@
/*
* Copyright 2024 the original author or authors.
* Copyright 2023-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
* 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,
@@ -13,6 +13,7 @@
* See the License for the specific language governing permissions and
* limitations under the License.
*/
package org.springframework.ai.vectorstore;
import org.testcontainers.utility.DockerImageName;

43
vector-stores/spring-ai-opensearch-store/src/test/java/org/springframework/ai/vectorstore/OpenSearchVectorStoreIT.java
View File

@@ -5,7 +5,7 @@
* 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
* 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,
@@ -16,6 +16,17 @@
package org.springframework.ai.vectorstore;
import java.io.IOException;
import java.net.URISyntaxException;
import java.nio.charset.StandardCharsets;
import java.time.Duration;
import java.time.ZonedDateTime;
import java.util.Date;
import java.util.List;
import java.util.Map;
import java.util.UUID;
import java.util.concurrent.TimeUnit;
import org.apache.hc.core5.http.HttpHost;
import org.awaitility.Awaitility;
import org.junit.jupiter.api.BeforeAll;
@@ -27,6 +38,9 @@ import org.junit.jupiter.params.provider.ValueSource;
import org.opensearch.client.opensearch.OpenSearchClient;
import org.opensearch.client.transport.httpclient5.ApacheHttpClient5TransportBuilder;
import org.opensearch.testcontainers.OpensearchContainer;
import org.testcontainers.junit.jupiter.Container;
import org.testcontainers.junit.jupiter.Testcontainers;
import org.springframework.ai.document.Document;
import org.springframework.ai.embedding.EmbeddingModel;
import org.springframework.ai.openai.OpenAiEmbeddingModel;
@@ -38,19 +52,6 @@ import org.springframework.boot.autoconfigure.jdbc.DataSourceAutoConfiguration;
import org.springframework.boot.test.context.runner.ApplicationContextRunner;
import org.springframework.context.annotation.Bean;
import org.springframework.core.io.DefaultResourceLoader;
import org.testcontainers.junit.jupiter.Container;
import org.testcontainers.junit.jupiter.Testcontainers;
import java.io.IOException;
import java.net.URISyntaxException;
import java.nio.charset.StandardCharsets;
import java.time.Duration;
import java.time.ZonedDateTime;
import java.util.Date;
import java.util.List;
import java.util.Map;
import java.util.UUID;
import java.util.concurrent.TimeUnit;
import static org.assertj.core.api.Assertions.assertThat;
import static org.hamcrest.Matchers.equalTo;
@@ -126,7 +127,7 @@ class OpenSearchVectorStoreIT {
vectorStore.withSimilarityFunction(similarityFunction);
}
vectorStore.add(documents);
vectorStore.add(this.documents);
Awaitility.await()
.until(() -> vectorStore
@@ -138,14 +139,14 @@ class OpenSearchVectorStoreIT {
assertThat(results).hasSize(1);
Document resultDoc = results.get(0);
assertThat(resultDoc.getId()).isEqualTo(documents.get(2).getId());
assertThat(resultDoc.getId()).isEqualTo(this.documents.get(2).getId());
assertThat(resultDoc.getContent()).contains("The Great Depression (1929–1939) was an economic shock");
assertThat(resultDoc.getMetadata()).hasSize(2);
assertThat(resultDoc.getMetadata()).containsKey("meta2");
assertThat(resultDoc.getMetadata()).containsKey("distance");
// Remove all documents from the store
vectorStore.delete(documents.stream().map(Document::getId).toList());
vectorStore.delete(this.documents.stream().map(Document::getId).toList());
Awaitility.await()
.until(() -> vectorStore
@@ -245,7 +246,7 @@ class OpenSearchVectorStoreIT {
assertThat(results.get(0).getId()).isEqualTo(bgDocument2.getId());
// Remove all documents from the store
vectorStore.delete(documents.stream().map(Document::getId).toList());
vectorStore.delete(this.documents.stream().map(Document::getId).toList());
Awaitility.await()
.until(() -> vectorStore.similaritySearch(SearchRequest.query("The World").withTopK(1)), hasSize(0));
@@ -318,7 +319,7 @@ class OpenSearchVectorStoreIT {
vectorStore.withSimilarityFunction(similarityFunction);
}
vectorStore.add(documents);
vectorStore.add(this.documents);
SearchRequest query = SearchRequest.query("Great Depression")
.withTopK(50)
@@ -339,13 +340,13 @@ class OpenSearchVectorStoreIT {
assertThat(results).hasSize(1);
Document resultDoc = results.get(0);
assertThat(resultDoc.getId()).isEqualTo(documents.get(2).getId());
assertThat(resultDoc.getId()).isEqualTo(this.documents.get(2).getId());
assertThat(resultDoc.getContent()).contains("The Great Depression (1929–1939) was an economic shock");
assertThat(resultDoc.getMetadata()).containsKey("meta2");
assertThat(resultDoc.getMetadata()).containsKey("distance");
// Remove all documents from the store
vectorStore.delete(documents.stream().map(Document::getId).toList());
vectorStore.delete(this.documents.stream().map(Document::getId).toList());
Awaitility.await()
.until(() -> vectorStore

32
vector-stores/spring-ai-opensearch-store/src/test/java/org/springframework/ai/vectorstore/OpenSearchVectorStoreObservationIT.java
View File

@@ -1,11 +1,11 @@
/*
* Copyright 2024 the original author or authors.
* Copyright 2023-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
* 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,
@@ -13,9 +13,8 @@
* See the License for the specific language governing permissions and
* limitations under the License.
*/
package org.springframework.ai.vectorstore;
import static org.assertj.core.api.Assertions.assertThat;
package org.springframework.ai.vectorstore;
import java.io.IOException;
import java.net.URISyntaxException;
@@ -25,6 +24,9 @@ import java.util.List;
import java.util.Map;
import java.util.concurrent.TimeUnit;
import io.micrometer.observation.ObservationRegistry;
import io.micrometer.observation.tck.TestObservationRegistry;
import io.micrometer.observation.tck.TestObservationRegistryAssert;
import org.apache.hc.core5.http.HttpHost;
import org.awaitility.Awaitility;
import org.junit.jupiter.api.BeforeAll;
@@ -34,6 +36,9 @@ import org.junit.jupiter.api.condition.EnabledIfEnvironmentVariable;
import org.opensearch.client.opensearch.OpenSearchClient;
import org.opensearch.client.transport.httpclient5.ApacheHttpClient5TransportBuilder;
import org.opensearch.testcontainers.OpensearchContainer;
import org.testcontainers.junit.jupiter.Container;
import org.testcontainers.junit.jupiter.Testcontainers;
import org.springframework.ai.document.Document;
import org.springframework.ai.embedding.EmbeddingModel;
import org.springframework.ai.embedding.TokenCountBatchingStrategy;
@@ -51,13 +56,8 @@ import org.springframework.boot.autoconfigure.jdbc.DataSourceAutoConfiguration;
import org.springframework.boot.test.context.runner.ApplicationContextRunner;
import org.springframework.context.annotation.Bean;
import org.springframework.core.io.DefaultResourceLoader;
import org.testcontainers.junit.jupiter.Container;
import org.testcontainers.junit.jupiter.Testcontainers;
import io.micrometer.observation.ObservationRegistry;
import io.micrometer.observation.tck.TestObservationRegistry;
import io.micrometer.observation.tck.TestObservationRegistryAssert;
import static org.assertj.core.api.Assertions.assertThat;
import static org.hamcrest.Matchers.hasSize;
/**
@@ -87,10 +87,6 @@ public class OpenSearchVectorStoreObservationIT {
}
}
private ApplicationContextRunner getContextRunner() {
return new ApplicationContextRunner().withUserConfiguration(Config.class);
}
@BeforeAll
public static void beforeAll() {
Awaitility.setDefaultPollInterval(2, TimeUnit.SECONDS);
@@ -98,6 +94,10 @@ public class OpenSearchVectorStoreObservationIT {
Awaitility.setDefaultTimeout(Duration.ofMinutes(1));
}
private ApplicationContextRunner getContextRunner() {
return new ApplicationContextRunner().withUserConfiguration(Config.class);
}
@BeforeEach
void cleanDatabase() {
getContextRunner().run(context -> {
@@ -115,7 +115,7 @@ public class OpenSearchVectorStoreObservationIT {
TestObservationRegistry observationRegistry = context.getBean(TestObservationRegistry.class);
vectorStore.add(documents);
vectorStore.add(this.documents);
TestObservationRegistryAssert.assertThat(observationRegistry)
.doesNotHaveAnyRemainingCurrentObservation()
@@ -182,7 +182,7 @@ public class OpenSearchVectorStoreObservationIT {
observationRegistry.clear();
vectorStore.delete(documents.stream().map(Document::getId).toList());
vectorStore.delete(this.documents.stream().map(Document::getId).toList());
Awaitility.await()
.until(() -> vectorStore

16
vector-stores/spring-ai-oracle-store/pom.xml
View File

@@ -1,4 +1,20 @@
<?xml version="1.0" encoding="UTF-8"?>
<!--
~ Copyright 2023-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.
-->
<project xmlns="http://maven.apache.org/POM/4.0.0"
xmlns:xsi="http://www.w3.org/2001/XMLSchema-instance" xsi:schemaLocation="http://maven.apache.org/POM/4.0.0 http://maven.apache.org/maven-v4_0_0.xsd">
<modelVersion>4.0.0</modelVersion>

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

@@ -5,7 +5,7 @@
* 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
* 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,
@@ -16,9 +16,6 @@
package org.springframework.ai.vectorstore;
import static org.springframework.ai.vectorstore.OracleVectorStore.OracleVectorStoreDistanceType.DOT;
import static org.springframework.jdbc.core.StatementCreatorUtils.setParameterValue;
import java.io.ByteArrayOutputStream;
import java.sql.PreparedStatement;
import java.sql.ResultSet;
@@ -31,8 +28,16 @@ import java.util.List;
import java.util.Map;
import java.util.Optional;
import io.micrometer.observation.ObservationRegistry;
import oracle.jdbc.OracleType;
import oracle.sql.VECTOR;
import oracle.sql.json.OracleJsonFactory;
import oracle.sql.json.OracleJsonGenerator;
import oracle.sql.json.OracleJsonObject;
import oracle.sql.json.OracleJsonValue;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import org.springframework.ai.document.Document;
import org.springframework.ai.embedding.BatchingStrategy;
import org.springframework.ai.embedding.EmbeddingModel;
@@ -51,13 +56,8 @@ import org.springframework.jdbc.core.JdbcTemplate;
import org.springframework.jdbc.core.RowMapper;
import org.springframework.util.StringUtils;
import io.micrometer.observation.ObservationRegistry;
import oracle.jdbc.OracleType;
import oracle.sql.VECTOR;
import oracle.sql.json.OracleJsonFactory;
import oracle.sql.json.OracleJsonGenerator;
import oracle.sql.json.OracleJsonObject;
import oracle.sql.json.OracleJsonValue;
import static org.springframework.ai.vectorstore.OracleVectorStore.OracleVectorStoreDistanceType.DOT;
import static org.springframework.jdbc.core.StatementCreatorUtils.setParameterValue;
/**
* <p>
@@ -86,9 +86,461 @@ import oracle.sql.json.OracleJsonValue;
*/
public class OracleVectorStore extends AbstractObservationVectorStore implements InitializingBean {
public static final double SIMILARITY_THRESHOLD_EXACT_MATCH = 1.0d;
public static final String DEFAULT_TABLE_NAME = "SPRING_AI_VECTORS";
public static final OracleVectorStoreIndexType DEFAULT_INDEX_TYPE = OracleVectorStoreIndexType.IVF;
public static final OracleVectorStoreDistanceType DEFAULT_DISTANCE_TYPE = OracleVectorStoreDistanceType.COSINE;
public static final int DEFAULT_DIMENSIONS = -1;
public static final int DEFAULT_SEARCH_ACCURACY = -1;
private static final Logger logger = LoggerFactory.getLogger(OracleVectorStore.class);
public static final double SIMILARITY_THRESHOLD_EXACT_MATCH = 1.0d;
private static Map<OracleVectorStoreDistanceType, VectorStoreSimilarityMetric> SIMILARITY_TYPE_MAPPING = Map.of(
OracleVectorStoreDistanceType.COSINE, VectorStoreSimilarityMetric.COSINE,
OracleVectorStoreDistanceType.EUCLIDEAN, VectorStoreSimilarityMetric.EUCLIDEAN,
OracleVectorStoreDistanceType.DOT, VectorStoreSimilarityMetric.DOT);
public final FilterExpressionConverter filterExpressionConverter = new SqlJsonPathFilterExpressionConverter();
private final JdbcTemplate jdbcTemplate;
private final EmbeddingModel embeddingModel;
private final boolean initializeSchema;
private final boolean removeExistingVectorStoreTable;
/**
* Table name where vectors will be stored.
*/
private final String tableName;
/**
* Index type used to index the vectors. It can impact performance and database memory
* consumption.
*/
private final OracleVectorStoreIndexType indexType;
/**
* Distance type to use for computing vector distances.
*/
private final OracleVectorStoreDistanceType distanceType;
/**
* Expected number of dimensions for vectors. Enforcing vector dimensions is very
* useful to ensure future vector distance computations will be relevant.
*/
private final int dimensions;
private final boolean forcedNormalization;
private final int searchAccuracy;
private final BatchingStrategy batchingStrategy;
private final OracleJsonFactory osonFactory = new OracleJsonFactory();
private final ByteArrayOutputStream out = new ByteArrayOutputStream();
public OracleVectorStore(JdbcTemplate jdbcTemplate, EmbeddingModel embeddingModel) {
this(jdbcTemplate, embeddingModel, DEFAULT_TABLE_NAME, DEFAULT_INDEX_TYPE, DEFAULT_DISTANCE_TYPE,
DEFAULT_DIMENSIONS, DEFAULT_SEARCH_ACCURACY, false, false, false);
}
public OracleVectorStore(JdbcTemplate jdbcTemplate, EmbeddingModel embeddingModel, boolean initializeSchema) {
this(jdbcTemplate, embeddingModel, DEFAULT_TABLE_NAME, DEFAULT_INDEX_TYPE, DEFAULT_DISTANCE_TYPE,
DEFAULT_DIMENSIONS, DEFAULT_SEARCH_ACCURACY, initializeSchema, false, false);
}
public OracleVectorStore(JdbcTemplate jdbcTemplate, EmbeddingModel embeddingModel, String tableName,
OracleVectorStoreIndexType indexType, OracleVectorStoreDistanceType distanceType, int dimensions,
int searchAccuracy, boolean initializeSchema, boolean removeExistingVectorStoreTable,
boolean forcedNormalization) {
this(jdbcTemplate, embeddingModel, tableName, indexType, distanceType, dimensions, searchAccuracy,
initializeSchema, removeExistingVectorStoreTable, forcedNormalization, ObservationRegistry.NOOP, null,
new TokenCountBatchingStrategy());
}
public OracleVectorStore(JdbcTemplate jdbcTemplate, EmbeddingModel embeddingModel, String tableName,
OracleVectorStoreIndexType indexType, OracleVectorStoreDistanceType distanceType, int dimensions,
int searchAccuracy, boolean initializeSchema, boolean removeExistingVectorStoreTable,
boolean forcedNormalization, ObservationRegistry observationRegistry,
VectorStoreObservationConvention customObservationConvention, BatchingStrategy batchingStrategy) {
super(observationRegistry, customObservationConvention);
if (dimensions != DEFAULT_DIMENSIONS) {
if (dimensions <= 0) {
throw new RuntimeException("Number of dimensions must be strictly positive");
}
if (dimensions > 65535) {
throw new RuntimeException("Number of dimensions must be at most 65535");
}
}
if (searchAccuracy != DEFAULT_SEARCH_ACCURACY) {
if (searchAccuracy < 1) {
throw new RuntimeException("Search accuracy must be greater or equals to 1");
}
if (searchAccuracy > 100) {
throw new RuntimeException("Search accuracy must be lower or equals to 100");
}
}
this.jdbcTemplate = jdbcTemplate;
this.embeddingModel = embeddingModel;
this.tableName = tableName;
this.indexType = indexType;
this.distanceType = distanceType;
this.dimensions = dimensions;
this.searchAccuracy = searchAccuracy;
this.initializeSchema = initializeSchema;
this.removeExistingVectorStoreTable = removeExistingVectorStoreTable;
this.forcedNormalization = forcedNormalization;
this.batchingStrategy = batchingStrategy;
}
@Override
public void doAdd(final List<Document> documents) {
this.embeddingModel.embed(documents, EmbeddingOptionsBuilder.builder().build(), this.batchingStrategy);
this.jdbcTemplate.batchUpdate(getIngestStatement(), new BatchPreparedStatementSetter() {
@Override
public void setValues(PreparedStatement ps, int i) throws SQLException {
final Document document = documents.get(i);
final String content = document.getContent();
final byte[] json = toJson(document.getMetadata());
final VECTOR embeddingVector = toVECTOR(document.getEmbedding());
setParameterValue(ps, 1, Types.VARCHAR, document.getId());
setParameterValue(ps, 2, Types.VARCHAR, content);
setParameterValue(ps, 3, OracleType.JSON.getVendorTypeNumber(), json);
setParameterValue(ps, 4, OracleType.VECTOR.getVendorTypeNumber(), embeddingVector);
}
@Override
public int getBatchSize() {
return documents.size();
}
});
}
private String getIngestStatement() {
return String
.format("""
merge into %s target using (values(?, ?, ?, ?)) source (id, content, metadata, embedding) on (target.id = source.id)
when matched then update set target.content = source.content, target.metadata = source.metadata, target.embedding = source.embedding
when not matched then insert (target.id, target.content, target.metadata, target.embedding) values (source.id, source.content, source.metadata, source.embedding)""",
this.tableName);
}
/**
* Bind binary JSON from the client.
* @param m map of metadata
* @return the binary JSON ready to be inserted
*/
private byte[] toJson(final Map<String, Object> m) {
this.out.reset();
try (OracleJsonGenerator gen = this.osonFactory.createJsonBinaryGenerator(this.out)) {
gen.writeStartObject();
for (String key : m.keySet()) {
final Object o = m.get(key);
if (o instanceof String) {
gen.write(key, (String) o);
}
else if (o instanceof Integer) {
gen.write(key, (Integer) o);
}
else if (o instanceof Float) {
gen.write(key, (Float) o);
}
else if (o instanceof Double) {
gen.write(key, (Double) o);
}
else if (o instanceof Boolean) {
gen.write(key, (Boolean) o);
}
}
gen.writeEnd();
}
return this.out.toByteArray();
}
/**
* Converts a list of Double values into an Oracle VECTOR object ready to be inserted.
* Optionally normalize the vector beforehand (see forcedNormalization).
* @param floatList
* @return
* @throws SQLException
*/
private VECTOR toVECTOR(final float[] floatList) throws SQLException {
final double[] doubles = new double[floatList.length];
int i = 0;
for (double d : floatList) {
doubles[i++] = d;
}
if (this.forcedNormalization) {
return VECTOR.ofFloat64Values(normalize(doubles));
}
return VECTOR.ofFloat64Values(doubles);
}
/**
* Normalize a vector if requested.
* @param v vector to normalize
* @return the vector normalized
*/
private double[] normalize(final double[] v) {
double squaredSum = 0d;
for (double e : v) {
squaredSum += e * e;
}
final double magnitude = Math.sqrt(squaredSum);
if (magnitude > 0) {
final double multiplier = 1d / magnitude;
final int length = v.length;
for (int i = 0; i < length; i++) {
v[i] *= multiplier;
}
}
return v;
}
@Override
public Optional<Boolean> doDelete(final List<String> idList) {
final String sql = String.format("delete from %s where id=?", this.tableName);
final int[] argTypes = { Types.VARCHAR };
final List<Object[]> batchArgs = new ArrayList<>();
for (String id : idList) {
batchArgs.add(new Object[] { id });
}
final int[] deleteCounts = this.jdbcTemplate.batchUpdate(sql, batchArgs, argTypes);
int deleteCount = 0;
for (int detailedResult : deleteCounts) {
switch (detailedResult) {
case Statement.EXECUTE_FAILED:
break;
case 1:
case Statement.SUCCESS_NO_INFO:
deleteCount++;
break;
}
}
return Optional.of(deleteCount == idList.size());
}
@Override
public List<Document> doSimilaritySearch(SearchRequest request) {
try {
// From the provided query, generate a vector using the embedding model
final VECTOR embeddingVector = toVECTOR(this.embeddingModel.embed(request.getQuery()));
if (logger.isDebugEnabled()) {
this.jdbcTemplate.batchUpdate("insert into debug(embedding) values(?)",
new BatchPreparedStatementSetter() {
@Override
public void setValues(PreparedStatement ps, int i) throws SQLException {
setParameterValue(ps, 1, OracleType.VECTOR.getVendorTypeNumber(), embeddingVector);
}
@Override
public int getBatchSize() {
return 1;
}
});
}
final String nativeFilterExpression = (request.getFilterExpression() != null)
? this.filterExpressionConverter.convertExpression(request.getFilterExpression()) : "";
String jsonPathFilter = "";
if (request.getSimilarityThreshold() == SearchRequest.SIMILARITY_THRESHOLD_ACCEPT_ALL) {
if (StringUtils.hasText(nativeFilterExpression)) {
jsonPathFilter = String.format("where JSON_EXISTS( metadata, '%s' )\n", nativeFilterExpression);
}
final String sql = this.searchAccuracy == DEFAULT_SEARCH_ACCURACY ? String.format("""
select id, content, metadata, embedding, %sVECTOR_DISTANCE(embedding, ?, %s)%s as distance
from %s
%sorder by distance
fetch first %d rows only""", this.distanceType == DOT ? "(1+" : "", this.distanceType.name(),
this.distanceType == DOT ? ")/2" : "", this.tableName, jsonPathFilter, request.getTopK())
: String.format(
"""
select id, content, metadata, embedding, %sVECTOR_DISTANCE(embedding, ?, %s)%s as distance
from %s
%sorder by distance
fetch APPROXIMATE first %d rows only WITH TARGET ACCURACY %d""",
this.distanceType == DOT ? "(1+" : "", this.distanceType.name(),
this.distanceType == DOT ? ")/2" : "", this.tableName, jsonPathFilter,
request.getTopK(), this.searchAccuracy);
logger.debug("SQL query: " + sql);
return this.jdbcTemplate.query(sql, new DocumentRowMapper(), embeddingVector);
}
else if (request.getSimilarityThreshold() == SIMILARITY_THRESHOLD_EXACT_MATCH) {
if (StringUtils.hasText(nativeFilterExpression)) {
jsonPathFilter = String.format("where JSON_EXISTS( metadata, '%s' )\n", nativeFilterExpression);
}
final String sql = String.format("""
select id, content, metadata, embedding, %sVECTOR_DISTANCE(embedding, ?, %s)%s as distance
from %s
%sorder by distance
fetch EXACT first %d rows only""", this.distanceType == DOT ? "(1+" : "",
this.distanceType.name(), this.distanceType == DOT ? ")/2" : "", this.tableName, jsonPathFilter,
request.getTopK());
logger.debug("SQL query: " + sql);
return this.jdbcTemplate.query(sql, new DocumentRowMapper(), embeddingVector);
}
else {
if (!this.forcedNormalization
|| (this.distanceType != OracleVectorStoreDistanceType.COSINE && this.distanceType != DOT)) {
throw new RuntimeException(
"Similarity threshold filtering requires all vectors to be normalized, see the forcedNormalization parameter for this Vector store. Also only COSINE and DOT distance types are supported.");
}
final double distance = this.distanceType == DOT ? (1d - request.getSimilarityThreshold()) * 2d - 1d
: 1d - request.getSimilarityThreshold();
if (StringUtils.hasText(nativeFilterExpression)) {
jsonPathFilter = String.format(" and JSON_EXISTS( metadata, '%s' )", nativeFilterExpression);
}
final String sql = this.distanceType == DOT ? (this.searchAccuracy == DEFAULT_SEARCH_ACCURACY
? String.format(
"""
select id, content, metadata, embedding, (1+VECTOR_DISTANCE(embedding, ?, DOT))/2 as distance
from %s
where VECTOR_DISTANCE(embedding, ?, DOT) <= ?%s
order by distance
fetch first %d rows only""",
this.tableName, jsonPathFilter, request.getTopK())
: String.format(
"""
select id, content, metadata, embedding, (1+VECTOR_DISTANCE(embedding, ?, DOT))/2 as distance
from %s
where VECTOR_DISTANCE(embedding, ?, DOT) <= ?%s
order by distance
fetch APPROXIMATE first %d rows only WITH TARGET ACCURACY %d""",
this.tableName, jsonPathFilter, request.getTopK(), this.searchAccuracy)
) : (this.searchAccuracy == DEFAULT_SEARCH_ACCURACY ? String.format("""
select id, content, metadata, embedding, VECTOR_DISTANCE(embedding, ?, COSINE) as distance
from %s
where VECTOR_DISTANCE(embedding, ?, COSINE) <= ?%s
order by distance
fetch first %d rows only""", this.tableName, jsonPathFilter, request.getTopK())
: String.format(
"""
select id, content, metadata, embedding, VECTOR_DISTANCE(embedding, ?, COSINE) as distance
from %s
where VECTOR_DISTANCE(embedding, ?, COSINE) <= ?%s
order by distance
fetch APPROXIMATE first %d rows only WITH TARGET ACCURACY %d""",
this.tableName, jsonPathFilter, request.getTopK(), this.searchAccuracy));
logger.debug("SQL query: " + sql);
return this.jdbcTemplate.query(sql, new DocumentRowMapper(), embeddingVector, embeddingVector,
distance);
}
}
catch (SQLException sqle) {
throw new RuntimeException(sqle);
}
}
@Override
public void afterPropertiesSet() throws Exception {
if (this.initializeSchema) {
// Remove existing VectorStoreTable
if (this.removeExistingVectorStoreTable) {
this.jdbcTemplate.execute(String.format("drop table if exists %s purge", this.tableName));
}
this.jdbcTemplate.execute(String.format("""
create table if not exists %s (
id varchar2(36) default sys_guid() primary key,
content clob not null,
metadata json not null,
embedding vector(%s,FLOAT64) annotations(Distance '%s', IndexType '%s')
)""", this.tableName, this.dimensions == DEFAULT_DIMENSIONS ? "*" : String.valueOf(this.dimensions),
this.distanceType.name(), this.indexType.name()));
if (logger.isDebugEnabled()) {
this.jdbcTemplate.execute(String.format("""
create table if not exists debug (
id varchar2(36) default sys_guid() primary key,
embedding vector(%s,FLOAT64) annotations(Distance '%s')
)""", this.dimensions == DEFAULT_DIMENSIONS ? "*" : String.valueOf(this.dimensions),
this.distanceType.name()));
}
switch (this.indexType) {
case IVF:
this.jdbcTemplate.execute(String.format("""
create vector index if not exists vector_index_%s on %s (embedding)
organization neighbor partitions
distance %s
with target accuracy %d
parameters (type IVF, neighbor partitions 10)""", this.tableName,
this.tableName, this.distanceType.name(),
this.searchAccuracy == DEFAULT_SEARCH_ACCURACY ? 95 : this.searchAccuracy));
break;
/*
* TODO: Enable for 23.5 case HNSW:
* this.jdbcTemplate.execute(String.format(""" create vector index if not
* exists vector_index_%s on %s (embedding) organization inmemory neighbor
* graph distance %s with target accuracy %d parameters (type HNSW,
* neighbors 40, efconstruction 500)""", tableName, tableName,
* distanceType.name(), searchAccuracy == DEFAULT_SEARCH_ACCURACY ? 95 :
* searchAccuracy)); break;
*/
}
}
}
public String getTableName() {
return this.tableName;
}
@Override
public Builder createObservationContextBuilder(String operationName) {
return VectorStoreObservationContext.builder(VectorStoreProvider.ORACLE.value(), operationName)
.withDimensions(this.embeddingModel.dimensions())
.withCollectionName(this.getTableName())
.withSimilarityMetric(getSimilarityMetric());
}
private String getSimilarityMetric() {
if (!SIMILARITY_TYPE_MAPPING.containsKey(this.distanceType)) {
return this.distanceType.name();
}
return SIMILARITY_TYPE_MAPPING.get(this.distanceType).value();
}
public enum OracleVectorStoreIndexType {
@@ -174,255 +626,6 @@ public class OracleVectorStore extends AbstractObservationVectorStore implements
}
public static final String DEFAULT_TABLE_NAME = "SPRING_AI_VECTORS";
public static final OracleVectorStoreIndexType DEFAULT_INDEX_TYPE = OracleVectorStoreIndexType.IVF;
public static final OracleVectorStoreDistanceType DEFAULT_DISTANCE_TYPE = OracleVectorStoreDistanceType.COSINE;
public static final int DEFAULT_DIMENSIONS = -1;
public static final int DEFAULT_SEARCH_ACCURACY = -1;
private final JdbcTemplate jdbcTemplate;
private final EmbeddingModel embeddingModel;
private final boolean initializeSchema;
private final boolean removeExistingVectorStoreTable;
public final FilterExpressionConverter filterExpressionConverter = new SqlJsonPathFilterExpressionConverter();
/**
* Table name where vectors will be stored.
*/
private final String tableName;
/**
* Index type used to index the vectors. It can impact performance and database memory
* consumption.
*/
private final OracleVectorStoreIndexType indexType;
/**
* Distance type to use for computing vector distances.
*/
private final OracleVectorStoreDistanceType distanceType;
/**
* Expected number of dimensions for vectors. Enforcing vector dimensions is very
* useful to ensure future vector distance computations will be relevant.
*/
private final int dimensions;
private final boolean forcedNormalization;
private final int searchAccuracy;
private final BatchingStrategy batchingStrategy;
public OracleVectorStore(JdbcTemplate jdbcTemplate, EmbeddingModel embeddingModel) {
this(jdbcTemplate, embeddingModel, DEFAULT_TABLE_NAME, DEFAULT_INDEX_TYPE, DEFAULT_DISTANCE_TYPE,
DEFAULT_DIMENSIONS, DEFAULT_SEARCH_ACCURACY, false, false, false);
}
public OracleVectorStore(JdbcTemplate jdbcTemplate, EmbeddingModel embeddingModel, boolean initializeSchema) {
this(jdbcTemplate, embeddingModel, DEFAULT_TABLE_NAME, DEFAULT_INDEX_TYPE, DEFAULT_DISTANCE_TYPE,
DEFAULT_DIMENSIONS, DEFAULT_SEARCH_ACCURACY, initializeSchema, false, false);
}
public OracleVectorStore(JdbcTemplate jdbcTemplate, EmbeddingModel embeddingModel, String tableName,
OracleVectorStoreIndexType indexType, OracleVectorStoreDistanceType distanceType, int dimensions,
int searchAccuracy, boolean initializeSchema, boolean removeExistingVectorStoreTable,
boolean forcedNormalization) {
this(jdbcTemplate, embeddingModel, tableName, indexType, distanceType, dimensions, searchAccuracy,
initializeSchema, removeExistingVectorStoreTable, forcedNormalization, ObservationRegistry.NOOP, null,
new TokenCountBatchingStrategy());
}
public OracleVectorStore(JdbcTemplate jdbcTemplate, EmbeddingModel embeddingModel, String tableName,
OracleVectorStoreIndexType indexType, OracleVectorStoreDistanceType distanceType, int dimensions,
int searchAccuracy, boolean initializeSchema, boolean removeExistingVectorStoreTable,
boolean forcedNormalization, ObservationRegistry observationRegistry,
VectorStoreObservationConvention customObservationConvention, BatchingStrategy batchingStrategy) {
super(observationRegistry, customObservationConvention);
if (dimensions != DEFAULT_DIMENSIONS) {
if (dimensions <= 0) {
throw new RuntimeException("Number of dimensions must be strictly positive");
}
if (dimensions > 65535) {
throw new RuntimeException("Number of dimensions must be at most 65535");
}
}
if (searchAccuracy != DEFAULT_SEARCH_ACCURACY) {
if (searchAccuracy < 1) {
throw new RuntimeException("Search accuracy must be greater or equals to 1");
}
if (searchAccuracy > 100) {
throw new RuntimeException("Search accuracy must be lower or equals to 100");
}
}
this.jdbcTemplate = jdbcTemplate;
this.embeddingModel = embeddingModel;
this.tableName = tableName;
this.indexType = indexType;
this.distanceType = distanceType;
this.dimensions = dimensions;
this.searchAccuracy = searchAccuracy;
this.initializeSchema = initializeSchema;
this.removeExistingVectorStoreTable = removeExistingVectorStoreTable;
this.forcedNormalization = forcedNormalization;
this.batchingStrategy = batchingStrategy;
}
@Override
public void doAdd(final List<Document> documents) {
this.embeddingModel.embed(documents, EmbeddingOptionsBuilder.builder().build(), this.batchingStrategy);
this.jdbcTemplate.batchUpdate(getIngestStatement(), new BatchPreparedStatementSetter() {
@Override
public void setValues(PreparedStatement ps, int i) throws SQLException {
final Document document = documents.get(i);
final String content = document.getContent();
final byte[] json = toJson(document.getMetadata());
final VECTOR embeddingVector = toVECTOR(document.getEmbedding());
setParameterValue(ps, 1, Types.VARCHAR, document.getId());
setParameterValue(ps, 2, Types.VARCHAR, content);
setParameterValue(ps, 3, OracleType.JSON.getVendorTypeNumber(), json);
setParameterValue(ps, 4, OracleType.VECTOR.getVendorTypeNumber(), embeddingVector);
}
@Override
public int getBatchSize() {
return documents.size();
}
});
}
private String getIngestStatement() {
return String
.format("""
merge into %s target using (values(?, ?, ?, ?)) source (id, content, metadata, embedding) on (target.id = source.id)
when matched then update set target.content = source.content, target.metadata = source.metadata, target.embedding = source.embedding
when not matched then insert (target.id, target.content, target.metadata, target.embedding) values (source.id, source.content, source.metadata, source.embedding)""",
tableName);
}
private final OracleJsonFactory osonFactory = new OracleJsonFactory();
private final ByteArrayOutputStream out = new ByteArrayOutputStream();
/**
* Bind binary JSON from the client.
* @param m map of metadata
* @return the binary JSON ready to be inserted
*/
private byte[] toJson(final Map<String, Object> m) {
out.reset();
try (OracleJsonGenerator gen = osonFactory.createJsonBinaryGenerator(out)) {
gen.writeStartObject();
for (String key : m.keySet()) {
final Object o = m.get(key);
if (o instanceof String) {
gen.write(key, (String) o);
}
else if (o instanceof Integer) {
gen.write(key, (Integer) o);
}
else if (o instanceof Float) {
gen.write(key, (Float) o);
}
else if (o instanceof Double) {
gen.write(key, (Double) o);
}
else if (o instanceof Boolean) {
gen.write(key, (Boolean) o);
}
}
gen.writeEnd();
}
return out.toByteArray();
}
/**
* Converts a list of Double values into an Oracle VECTOR object ready to be inserted.
* Optionally normalize the vector beforehand (see forcedNormalization).
* @param floatList
* @return
* @throws SQLException
*/
private VECTOR toVECTOR(final float[] floatList) throws SQLException {
final double[] doubles = new double[floatList.length];
int i = 0;
for (double d : floatList) {
doubles[i++] = d;
}
if (forcedNormalization) {
return VECTOR.ofFloat64Values(normalize(doubles));
}
return VECTOR.ofFloat64Values(doubles);
}
/**
* Normalize a vector if requested.
* @param v vector to normalize
* @return the vector normalized
*/
private double[] normalize(final double[] v) {
double squaredSum = 0d;
for (double e : v) {
squaredSum += e * e;
}
final double magnitude = Math.sqrt(squaredSum);
if (magnitude > 0) {
final double multiplier = 1d / magnitude;
final int length = v.length;
for (int i = 0; i < length; i++) {
v[i] *= multiplier;
}
}
return v;
}
@Override
public Optional<Boolean> doDelete(final List<String> idList) {
final String sql = String.format("delete from %s where id=?", tableName);
final int[] argTypes = { Types.VARCHAR };
final List<Object[]> batchArgs = new ArrayList<>();
for (String id : idList) {
batchArgs.add(new Object[] { id });
}
final int[] deleteCounts = jdbcTemplate.batchUpdate(sql, batchArgs, argTypes);
int deleteCount = 0;
for (int detailedResult : deleteCounts) {
switch (detailedResult) {
case Statement.EXECUTE_FAILED:
break;
case 1:
case Statement.SUCCESS_NO_INFO:
deleteCount++;
break;
}
}
return Optional.of(deleteCount == idList.size());
}
private static class DocumentRowMapper implements RowMapper<Document> {
@Override
@@ -459,195 +662,4 @@ public class OracleVectorStore extends AbstractObservationVectorStore implements
}
@Override
public List<Document> doSimilaritySearch(SearchRequest request) {
try {
// From the provided query, generate a vector using the embedding model
final VECTOR embeddingVector = toVECTOR(embeddingModel.embed(request.getQuery()));
if (logger.isDebugEnabled()) {
this.jdbcTemplate.batchUpdate("insert into debug(embedding) values(?)",
new BatchPreparedStatementSetter() {
@Override
public void setValues(PreparedStatement ps, int i) throws SQLException {
setParameterValue(ps, 1, OracleType.VECTOR.getVendorTypeNumber(), embeddingVector);
}
@Override
public int getBatchSize() {
return 1;
}
});
}
final String nativeFilterExpression = (request.getFilterExpression() != null)
? this.filterExpressionConverter.convertExpression(request.getFilterExpression()) : "";
String jsonPathFilter = "";
if (request.getSimilarityThreshold() == SearchRequest.SIMILARITY_THRESHOLD_ACCEPT_ALL) {
if (StringUtils.hasText(nativeFilterExpression)) {
jsonPathFilter = String.format("where JSON_EXISTS( metadata, '%s' )\n", nativeFilterExpression);
}
final String sql = searchAccuracy == DEFAULT_SEARCH_ACCURACY ? String.format("""
select id, content, metadata, embedding, %sVECTOR_DISTANCE(embedding, ?, %s)%s as distance
from %s
%sorder by distance
fetch first %d rows only""", distanceType == DOT ? "(1+" : "", distanceType.name(),
distanceType == DOT ? ")/2" : "", tableName, jsonPathFilter, request.getTopK())
: String.format(
"""
select id, content, metadata, embedding, %sVECTOR_DISTANCE(embedding, ?, %s)%s as distance
from %s
%sorder by distance
fetch APPROXIMATE first %d rows only WITH TARGET ACCURACY %d""",
distanceType == DOT ? "(1+" : "", distanceType.name(), distanceType == DOT ? ")/2" : "",
tableName, jsonPathFilter, request.getTopK(), searchAccuracy);
logger.debug("SQL query: " + sql);
return this.jdbcTemplate.query(sql, new DocumentRowMapper(), embeddingVector);
}
else if (request.getSimilarityThreshold() == SIMILARITY_THRESHOLD_EXACT_MATCH) {
if (StringUtils.hasText(nativeFilterExpression)) {
jsonPathFilter = String.format("where JSON_EXISTS( metadata, '%s' )\n", nativeFilterExpression);
}
final String sql = String.format("""
select id, content, metadata, embedding, %sVECTOR_DISTANCE(embedding, ?, %s)%s as distance
from %s
%sorder by distance
fetch EXACT first %d rows only""", distanceType == DOT ? "(1+" : "", distanceType.name(),
distanceType == DOT ? ")/2" : "", tableName, jsonPathFilter, request.getTopK());
logger.debug("SQL query: " + sql);
return this.jdbcTemplate.query(sql, new DocumentRowMapper(), embeddingVector);
}
else {
if (!forcedNormalization
|| (distanceType != OracleVectorStoreDistanceType.COSINE && distanceType != DOT)) {
throw new RuntimeException(
"Similarity threshold filtering requires all vectors to be normalized, see the forcedNormalization parameter for this Vector store. Also only COSINE and DOT distance types are supported.");
}
final double distance = distanceType == DOT ? (1d - request.getSimilarityThreshold()) * 2d - 1d
: 1d - request.getSimilarityThreshold();
if (StringUtils.hasText(nativeFilterExpression)) {
jsonPathFilter = String.format(" and JSON_EXISTS( metadata, '%s' )", nativeFilterExpression);
}
final String sql = distanceType == DOT ? (searchAccuracy == DEFAULT_SEARCH_ACCURACY ? String.format("""
select id, content, metadata, embedding, (1+VECTOR_DISTANCE(embedding, ?, DOT))/2 as distance
from %s
where VECTOR_DISTANCE(embedding, ?, DOT) <= ?%s
order by distance
fetch first %d rows only""", tableName, jsonPathFilter, request.getTopK()) : String.format("""
select id, content, metadata, embedding, (1+VECTOR_DISTANCE(embedding, ?, DOT))/2 as distance
from %s
where VECTOR_DISTANCE(embedding, ?, DOT) <= ?%s
order by distance
fetch APPROXIMATE first %d rows only WITH TARGET ACCURACY %d""", tableName, jsonPathFilter,
request.getTopK(), searchAccuracy)
) : (searchAccuracy == DEFAULT_SEARCH_ACCURACY ? String.format("""
select id, content, metadata, embedding, VECTOR_DISTANCE(embedding, ?, COSINE) as distance
from %s
where VECTOR_DISTANCE(embedding, ?, COSINE) <= ?%s
order by distance
fetch first %d rows only""", tableName, jsonPathFilter, request.getTopK()) : String.format("""
select id, content, metadata, embedding, VECTOR_DISTANCE(embedding, ?, COSINE) as distance
from %s
where VECTOR_DISTANCE(embedding, ?, COSINE) <= ?%s
order by distance
fetch APPROXIMATE first %d rows only WITH TARGET ACCURACY %d""", tableName, jsonPathFilter,
request.getTopK(), searchAccuracy));
logger.debug("SQL query: " + sql);
return this.jdbcTemplate.query(sql, new DocumentRowMapper(), embeddingVector, embeddingVector,
distance);
}
}
catch (SQLException sqle) {
throw new RuntimeException(sqle);
}
}
@Override
public void afterPropertiesSet() throws Exception {
if (this.initializeSchema) {
// Remove existing VectorStoreTable
if (this.removeExistingVectorStoreTable) {
this.jdbcTemplate.execute(String.format("drop table if exists %s purge", tableName));
}
this.jdbcTemplate.execute(String.format("""
create table if not exists %s (
id varchar2(36) default sys_guid() primary key,
content clob not null,
metadata json not null,
embedding vector(%s,FLOAT64) annotations(Distance '%s', IndexType '%s')
)""", tableName, dimensions == DEFAULT_DIMENSIONS ? "*" : String.valueOf(dimensions),
distanceType.name(), indexType.name()));
if (logger.isDebugEnabled()) {
this.jdbcTemplate.execute(String.format("""
create table if not exists debug (
id varchar2(36) default sys_guid() primary key,
embedding vector(%s,FLOAT64) annotations(Distance '%s')
)""", dimensions == DEFAULT_DIMENSIONS ? "*" : String.valueOf(dimensions),
distanceType.name()));
}
switch (indexType) {
case IVF:
this.jdbcTemplate.execute(String.format("""
create vector index if not exists vector_index_%s on %s (embedding)
organization neighbor partitions
distance %s
with target accuracy %d
parameters (type IVF, neighbor partitions 10)""", tableName, tableName,
distanceType.name(), searchAccuracy == DEFAULT_SEARCH_ACCURACY ? 95 : searchAccuracy));
break;
/*
* TODO: Enable for 23.5 case HNSW:
* this.jdbcTemplate.execute(String.format(""" create vector index if not
* exists vector_index_%s on %s (embedding) organization inmemory neighbor
* graph distance %s with target accuracy %d parameters (type HNSW,
* neighbors 40, efconstruction 500)""", tableName, tableName,
* distanceType.name(), searchAccuracy == DEFAULT_SEARCH_ACCURACY ? 95 :
* searchAccuracy)); break;
*/
}
}
}
public String getTableName() {
return tableName;
}
@Override
public Builder createObservationContextBuilder(String operationName) {
return VectorStoreObservationContext.builder(VectorStoreProvider.ORACLE.value(), operationName)
.withDimensions(this.embeddingModel.dimensions())
.withCollectionName(this.getTableName())
.withSimilarityMetric(getSimilarityMetric());
}
private static Map<OracleVectorStoreDistanceType, VectorStoreSimilarityMetric> SIMILARITY_TYPE_MAPPING = Map.of(
OracleVectorStoreDistanceType.COSINE, VectorStoreSimilarityMetric.COSINE,
OracleVectorStoreDistanceType.EUCLIDEAN, VectorStoreSimilarityMetric.EUCLIDEAN,
OracleVectorStoreDistanceType.DOT, VectorStoreSimilarityMetric.DOT);
private String getSimilarityMetric() {
if (!SIMILARITY_TYPE_MAPPING.containsKey(this.distanceType)) {
return this.distanceType.name();
}
return SIMILARITY_TYPE_MAPPING.get(this.distanceType).value();
}
}

Some files were not shown because too many files have changed in this diff Show More