Support for SimpleVectorStore with metdata filter expressions
Signed-off-by: Jemin Huh <hjm1980@gmail.com>
This commit is contained in:
@@ -32,6 +32,7 @@ import java.util.Map;
|
||||
import java.util.Objects;
|
||||
import java.util.Optional;
|
||||
import java.util.concurrent.ConcurrentHashMap;
|
||||
import java.util.function.Predicate;
|
||||
|
||||
import com.fasterxml.jackson.core.JsonProcessingException;
|
||||
import com.fasterxml.jackson.core.type.TypeReference;
|
||||
@@ -46,9 +47,14 @@ import org.springframework.ai.embedding.EmbeddingModel;
|
||||
import org.springframework.ai.observation.conventions.VectorStoreProvider;
|
||||
import org.springframework.ai.observation.conventions.VectorStoreSimilarityMetric;
|
||||
import org.springframework.ai.util.JacksonUtils;
|
||||
import org.springframework.ai.vectorstore.filter.FilterExpressionConverter;
|
||||
import org.springframework.ai.vectorstore.filter.converter.SimpleVectorStoreFilterExpressionConverter;
|
||||
import org.springframework.ai.vectorstore.observation.AbstractObservationVectorStore;
|
||||
import org.springframework.ai.vectorstore.observation.VectorStoreObservationContext;
|
||||
import org.springframework.core.io.Resource;
|
||||
import org.springframework.expression.ExpressionParser;
|
||||
import org.springframework.expression.spel.standard.SpelExpressionParser;
|
||||
import org.springframework.expression.spel.support.StandardEvaluationContext;
|
||||
|
||||
/**
|
||||
* SimpleVectorStore is a simple implementation of the VectorStore interface.
|
||||
@@ -67,6 +73,7 @@ import org.springframework.core.io.Resource;
|
||||
* @author Sebastien Deleuze
|
||||
* @author Ilayaperumal Gopinathan
|
||||
* @author Thomas Vitale
|
||||
* @author Jemin Huh
|
||||
*/
|
||||
public class SimpleVectorStore extends AbstractObservationVectorStore {
|
||||
|
||||
@@ -74,11 +81,17 @@ public class SimpleVectorStore extends AbstractObservationVectorStore {
|
||||
|
||||
private final ObjectMapper objectMapper;
|
||||
|
||||
private final ExpressionParser expressionParser;
|
||||
|
||||
private final FilterExpressionConverter filterExpressionConverter;
|
||||
|
||||
protected Map<String, SimpleVectorStoreContent> store = new ConcurrentHashMap<>();
|
||||
|
||||
protected SimpleVectorStore(SimpleVectorStoreBuilder builder) {
|
||||
super(builder);
|
||||
this.objectMapper = JsonMapper.builder().addModules(JacksonUtils.instantiateAvailableModules()).build();
|
||||
this.expressionParser = new SpelExpressionParser();
|
||||
this.filterExpressionConverter = new SimpleVectorStoreFilterExpressionConverter();
|
||||
}
|
||||
|
||||
/**
|
||||
@@ -114,14 +127,11 @@ public class SimpleVectorStore extends AbstractObservationVectorStore {
|
||||
|
||||
@Override
|
||||
public List<Document> doSimilaritySearch(SearchRequest request) {
|
||||
if (request.getFilterExpression() != null) {
|
||||
throw new UnsupportedOperationException(
|
||||
"The [" + this.getClass() + "] doesn't support metadata filtering!");
|
||||
}
|
||||
|
||||
Predicate<SimpleVectorStoreContent> documentFilterPredicate = doFilterPredicate(request);
|
||||
float[] userQueryEmbedding = getUserQueryEmbedding(request.getQuery());
|
||||
return this.store.values()
|
||||
.stream()
|
||||
.filter(documentFilterPredicate)
|
||||
.map(content -> content
|
||||
.toDocument(EmbeddingMath.cosineSimilarity(userQueryEmbedding, content.getEmbedding())))
|
||||
.filter(document -> document.getScore() >= request.getSimilarityThreshold())
|
||||
@@ -130,6 +140,16 @@ public class SimpleVectorStore extends AbstractObservationVectorStore {
|
||||
.toList();
|
||||
}
|
||||
|
||||
private Predicate<SimpleVectorStoreContent> doFilterPredicate(SearchRequest request) {
|
||||
return request.hasFilterExpression() ? document -> {
|
||||
StandardEvaluationContext context = new StandardEvaluationContext();
|
||||
context.setVariable("metadata", document.getMetadata());
|
||||
return this.expressionParser
|
||||
.parseExpression(this.filterExpressionConverter.convertExpression(request.getFilterExpression()))
|
||||
.getValue(context, Boolean.class);
|
||||
} : document -> true;
|
||||
}
|
||||
|
||||
/**
|
||||
* Serialize the vector store content into a file in JSON format.
|
||||
* @param file the file to save the vector store content
|
||||
|
||||
@@ -0,0 +1,147 @@
|
||||
/*
|
||||
* 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.filter.converter;
|
||||
|
||||
import org.springframework.ai.vectorstore.filter.Filter;
|
||||
import org.springframework.ai.vectorstore.filter.Filter.Expression;
|
||||
|
||||
import java.text.ParseException;
|
||||
import java.text.SimpleDateFormat;
|
||||
import java.util.Date;
|
||||
import java.util.List;
|
||||
import java.util.TimeZone;
|
||||
import java.util.regex.Pattern;
|
||||
|
||||
/**
|
||||
* Converts {@link Expression} into SpEL metadata filter expression format.
|
||||
* (https://docs.spring.io/spring-framework/reference/core/expressions.html)
|
||||
*
|
||||
* @author Jemin Huh
|
||||
*/
|
||||
public class SimpleVectorStoreFilterExpressionConverter extends AbstractFilterExpressionConverter {
|
||||
|
||||
private static final Pattern DATE_FORMAT_PATTERN = Pattern.compile("\\d{4}-\\d{2}-\\d{2}T\\d{2}:\\d{2}:\\d{2}Z");
|
||||
|
||||
private final SimpleDateFormat dateFormat;
|
||||
|
||||
public SimpleVectorStoreFilterExpressionConverter() {
|
||||
this.dateFormat = new SimpleDateFormat("yyyy-MM-dd'T'HH:mm:ss'Z'");
|
||||
this.dateFormat.setTimeZone(TimeZone.getTimeZone("UTC"));
|
||||
}
|
||||
|
||||
@Override
|
||||
protected void doExpression(Filter.Expression expression, StringBuilder context) {
|
||||
this.convertOperand(expression.left(), context);
|
||||
context.append(getOperationSymbol(expression));
|
||||
this.convertOperand(expression.right(), context);
|
||||
}
|
||||
|
||||
private String getOperationSymbol(Filter.Expression exp) {
|
||||
return switch (exp.type()) {
|
||||
case AND -> " and ";
|
||||
case OR -> " or ";
|
||||
case EQ -> " == ";
|
||||
case LT -> " < ";
|
||||
case LTE -> " <= ";
|
||||
case GT -> " > ";
|
||||
case GTE -> " >= ";
|
||||
case NE -> " != ";
|
||||
case IN -> " in ";
|
||||
case NIN -> " not in ";
|
||||
default -> throw new RuntimeException("Not supported expression type: " + exp.type());
|
||||
};
|
||||
}
|
||||
|
||||
@Override
|
||||
protected void doKey(Filter.Key key, StringBuilder context) {
|
||||
var identifier = hasOuterQuotes(key.key()) ? removeOuterQuotes(key.key()) : key.key();
|
||||
context.append("#metadata['").append(identifier).append("']");
|
||||
}
|
||||
|
||||
@Override
|
||||
protected void doValue(Filter.Value filterValue, StringBuilder context) {
|
||||
if (filterValue.value() instanceof List<?> list) {
|
||||
var formattedList = new StringBuilder("{");
|
||||
int c = 0;
|
||||
for (Object v : list) {
|
||||
this.doSingleValue(v, formattedList);
|
||||
if (c++ < list.size() - 1) {
|
||||
this.doAddValueRangeSpitter(filterValue, formattedList);
|
||||
}
|
||||
}
|
||||
formattedList.append("}");
|
||||
|
||||
if (context.lastIndexOf("in ") == -1) {
|
||||
context.append(formattedList);
|
||||
}
|
||||
else {
|
||||
appendSpELContains(formattedList, context);
|
||||
}
|
||||
}
|
||||
else {
|
||||
this.doSingleValue(filterValue.value(), context);
|
||||
}
|
||||
}
|
||||
|
||||
private void appendSpELContains(StringBuilder formattedList, StringBuilder context) {
|
||||
int metadataStart = context.lastIndexOf("#metadata");
|
||||
if (metadataStart == -1)
|
||||
throw new RuntimeException("Wrong SpEL expression: " + context);
|
||||
|
||||
int metadataEnd = context.indexOf(" ", metadataStart);
|
||||
String metadata = context.substring(metadataStart, metadataEnd);
|
||||
context.setLength(context.lastIndexOf("in "));
|
||||
context.delete(metadataStart, metadataEnd + 1);
|
||||
context.append(formattedList).append(".contains(").append(metadata).append(")");
|
||||
}
|
||||
|
||||
@Override
|
||||
protected void doSingleValue(Object value, StringBuilder context) {
|
||||
if (value instanceof Date date) {
|
||||
context.append("'");
|
||||
context.append(this.dateFormat.format(date));
|
||||
context.append("'");
|
||||
}
|
||||
else if (value instanceof String text) {
|
||||
context.append("'");
|
||||
if (DATE_FORMAT_PATTERN.matcher(text).matches()) {
|
||||
try {
|
||||
Date date = this.dateFormat.parse(text);
|
||||
context.append(this.dateFormat.format(date));
|
||||
}
|
||||
catch (ParseException e) {
|
||||
throw new IllegalArgumentException("Invalid date type:" + text, e);
|
||||
}
|
||||
}
|
||||
else {
|
||||
context.append(text);
|
||||
}
|
||||
context.append("'");
|
||||
}
|
||||
else {
|
||||
context.append(value);
|
||||
}
|
||||
}
|
||||
|
||||
@Override
|
||||
protected void doGroup(Filter.Group group, StringBuilder context) {
|
||||
context.append("(");
|
||||
super.doGroup(group, context);
|
||||
context.append(")");
|
||||
}
|
||||
|
||||
}
|
||||
@@ -0,0 +1,239 @@
|
||||
/*
|
||||
* 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;
|
||||
|
||||
import org.junit.jupiter.api.BeforeEach;
|
||||
import org.junit.jupiter.api.Test;
|
||||
import org.junit.jupiter.api.io.CleanupMode;
|
||||
import org.junit.jupiter.api.io.TempDir;
|
||||
import org.springframework.ai.document.Document;
|
||||
import org.springframework.ai.embedding.EmbeddingModel;
|
||||
import org.springframework.ai.vectorstore.filter.Filter;
|
||||
|
||||
import java.nio.file.Path;
|
||||
import java.util.Arrays;
|
||||
import java.util.Date;
|
||||
import java.util.List;
|
||||
import java.util.Map;
|
||||
|
||||
import static org.assertj.core.api.Assertions.assertThat;
|
||||
import static org.mockito.ArgumentMatchers.any;
|
||||
import static org.mockito.Mockito.mock;
|
||||
import static org.mockito.Mockito.when;
|
||||
import static org.springframework.ai.vectorstore.filter.Filter.ExpressionType.AND;
|
||||
import static org.springframework.ai.vectorstore.filter.Filter.ExpressionType.EQ;
|
||||
import static org.springframework.ai.vectorstore.filter.Filter.ExpressionType.GTE;
|
||||
import static org.springframework.ai.vectorstore.filter.Filter.ExpressionType.LTE;
|
||||
|
||||
/**
|
||||
* @author Jemin Huh
|
||||
*/
|
||||
class SimpleVectorStoreWithFilterTests {
|
||||
|
||||
@TempDir(cleanup = CleanupMode.ON_SUCCESS)
|
||||
Path tempDir;
|
||||
|
||||
private SimpleVectorStore vectorStore;
|
||||
|
||||
private EmbeddingModel mockEmbeddingModel;
|
||||
|
||||
@BeforeEach
|
||||
void setUp() {
|
||||
this.mockEmbeddingModel = mock(EmbeddingModel.class);
|
||||
when(this.mockEmbeddingModel.dimensions()).thenReturn(3);
|
||||
when(this.mockEmbeddingModel.embed(any(String.class))).thenReturn(new float[] { 0.1f, 0.2f, 0.3f });
|
||||
when(this.mockEmbeddingModel.embed(any(Document.class))).thenReturn(new float[] { 0.1f, 0.2f, 0.3f });
|
||||
this.vectorStore = SimpleVectorStore.builder(this.mockEmbeddingModel).build();
|
||||
}
|
||||
|
||||
@Test
|
||||
void shouldAddAndRetrieveDocumentWithFilter() {
|
||||
Document doc = Document.builder()
|
||||
.id("1")
|
||||
.text("test content")
|
||||
.metadata(Map.of("country", "BG", "year", 2020, "activationDate", "1970-01-01T00:00:02Z"))
|
||||
.build();
|
||||
|
||||
this.vectorStore.add(List.of(doc));
|
||||
|
||||
List<Document> results = this.vectorStore.similaritySearch(
|
||||
SearchRequest.builder().query("test content").filterExpression("country == 'BG'").build());
|
||||
assertThat(results).hasSize(1).first().satisfies(result -> {
|
||||
assertThat(result.getId()).isEqualTo("1");
|
||||
assertThat(result.getText()).isEqualTo("test content");
|
||||
assertThat(result.getMetadata()).hasSize(4);
|
||||
assertThat(result.getMetadata()).containsEntry("distance", 2.220446049250313E-16);
|
||||
});
|
||||
|
||||
results = this.vectorStore.similaritySearch(
|
||||
SearchRequest.builder().query("test content").filterExpression("country == 'KR'").build());
|
||||
assertThat(results).hasSize(0);
|
||||
|
||||
results = this.vectorStore.similaritySearch(SearchRequest.builder()
|
||||
.query("test content")
|
||||
.filterExpression("country == 'BG' && year == 2020")
|
||||
.build());
|
||||
assertThat(results).hasSize(1).first().satisfies(result -> {
|
||||
assertThat(result.getId()).isEqualTo("1");
|
||||
assertThat(result.getText()).isEqualTo("test content");
|
||||
assertThat(result.getMetadata()).hasSize(4);
|
||||
assertThat(result.getMetadata()).containsEntry("distance", 2.220446049250313E-16);
|
||||
});
|
||||
|
||||
results = this.vectorStore.similaritySearch(SearchRequest.builder()
|
||||
.query("test content")
|
||||
.filterExpression("country == 'BG' && year == 2024")
|
||||
.build());
|
||||
assertThat(results).hasSize(0);
|
||||
|
||||
results = this.vectorStore.similaritySearch(
|
||||
SearchRequest.builder().query("test content").filterExpression("country in ['BG', 'NL']").build());
|
||||
assertThat(results).hasSize(1).first().satisfies(result -> {
|
||||
assertThat(result.getId()).isEqualTo("1");
|
||||
assertThat(result.getText()).isEqualTo("test content");
|
||||
assertThat(result.getMetadata()).hasSize(4);
|
||||
assertThat(result.getMetadata()).containsEntry("distance", 2.220446049250313E-16);
|
||||
});
|
||||
|
||||
results = this.vectorStore.similaritySearch(
|
||||
SearchRequest.builder().query("test content").filterExpression("country in ['KR', 'NL']").build());
|
||||
assertThat(results).hasSize(0);
|
||||
|
||||
results = this.vectorStore.similaritySearch(SearchRequest.builder()
|
||||
.query("test content")
|
||||
.filterExpression(
|
||||
new Filter.Expression(EQ, new Filter.Key("activationDate"), new Filter.Value(new Date(2000))))
|
||||
.build());
|
||||
assertThat(results).hasSize(1).first().satisfies(result -> {
|
||||
assertThat(result.getId()).isEqualTo("1");
|
||||
assertThat(result.getText()).isEqualTo("test content");
|
||||
assertThat(result.getMetadata()).hasSize(4);
|
||||
assertThat(result.getMetadata()).containsEntry("distance", 2.220446049250313E-16);
|
||||
});
|
||||
|
||||
results = this.vectorStore.similaritySearch(SearchRequest.builder()
|
||||
.query("test content")
|
||||
.filterExpression(
|
||||
new Filter.Expression(EQ, new Filter.Key("activationDate"), new Filter.Value(new Date(3000))))
|
||||
.build());
|
||||
assertThat(results).hasSize(0);
|
||||
|
||||
}
|
||||
|
||||
@Test
|
||||
void shouldAddMultipleDocumentsWithFilter() {
|
||||
List<Document> docs = Arrays.asList(
|
||||
Document.builder()
|
||||
.id("1")
|
||||
.text("first")
|
||||
.metadata(Map.of("country", "BG", "year", 2020, "activationDate", "1970-01-01T00:00:02Z"))
|
||||
.build(),
|
||||
Document.builder()
|
||||
.id("2")
|
||||
.text("second")
|
||||
.metadata(Map.of("country", "KR", "year", 2022, "activationDate", "1970-01-01T00:00:03Z"))
|
||||
.build());
|
||||
|
||||
this.vectorStore.add(docs);
|
||||
|
||||
List<Document> results = this.vectorStore.similaritySearch("first");
|
||||
assertThat(results).hasSize(2).extracting(Document::getId).containsExactlyInAnyOrder("1", "2");
|
||||
|
||||
results = this.vectorStore
|
||||
.similaritySearch(SearchRequest.builder().query("first").filterExpression("country == 'BG'").build());
|
||||
assertThat(results).hasSize(1).first().satisfies(result -> {
|
||||
assertThat(result.getId()).isEqualTo("1");
|
||||
assertThat(result.getText()).isEqualTo("first");
|
||||
assertThat(result.getMetadata()).hasSize(4);
|
||||
assertThat(result.getMetadata()).containsEntry("distance", 2.220446049250313E-16);
|
||||
});
|
||||
|
||||
results = this.vectorStore
|
||||
.similaritySearch(SearchRequest.builder().query("first").filterExpression("country == 'NL'").build());
|
||||
assertThat(results).hasSize(0);
|
||||
|
||||
results = this.vectorStore.similaritySearch(
|
||||
SearchRequest.builder().query("first").filterExpression("country == 'BG' && year == 2020").build());
|
||||
assertThat(results).hasSize(1).first().satisfies(result -> {
|
||||
assertThat(result.getId()).isEqualTo("1");
|
||||
assertThat(result.getText()).isEqualTo("first");
|
||||
assertThat(result.getMetadata()).hasSize(4);
|
||||
assertThat(result.getMetadata()).containsEntry("distance", 2.220446049250313E-16);
|
||||
});
|
||||
|
||||
results = this.vectorStore.similaritySearch(
|
||||
SearchRequest.builder().query("first").filterExpression("country == 'KR' && year == 2022").build());
|
||||
assertThat(results).hasSize(1).first().satisfies(result -> {
|
||||
assertThat(result.getId()).isEqualTo("2");
|
||||
assertThat(result.getText()).isEqualTo("second");
|
||||
assertThat(result.getMetadata()).hasSize(4);
|
||||
assertThat(result.getMetadata()).containsEntry("distance", 2.220446049250313E-16);
|
||||
});
|
||||
|
||||
results = this.vectorStore.similaritySearch(SearchRequest.builder()
|
||||
.query("test content")
|
||||
.filterExpression("country == 'KR' && year == 2024")
|
||||
.build());
|
||||
assertThat(results).hasSize(0);
|
||||
|
||||
results = this.vectorStore.similaritySearch(
|
||||
SearchRequest.builder().query("first").filterExpression("country in ['BG', 'NL']").build());
|
||||
assertThat(results).hasSize(1).first().satisfies(result -> {
|
||||
assertThat(result.getId()).isEqualTo("1");
|
||||
assertThat(result.getText()).isEqualTo("first");
|
||||
assertThat(result.getMetadata()).hasSize(4);
|
||||
assertThat(result.getMetadata()).containsEntry("distance", 2.220446049250313E-16);
|
||||
});
|
||||
|
||||
results = this.vectorStore.similaritySearch(
|
||||
SearchRequest.builder().query("first").filterExpression("country in ['KR', 'NL']").build());
|
||||
assertThat(results).hasSize(1);
|
||||
|
||||
results = this.vectorStore.similaritySearch(SearchRequest.builder()
|
||||
.query("first")
|
||||
.filterExpression(
|
||||
new Filter.Expression(EQ, new Filter.Key("activationDate"), new Filter.Value(new Date(2000))))
|
||||
.build());
|
||||
assertThat(results).hasSize(1).first().satisfies(result -> {
|
||||
assertThat(result.getId()).isEqualTo("1");
|
||||
assertThat(result.getText()).isEqualTo("first");
|
||||
assertThat(result.getMetadata()).hasSize(4);
|
||||
assertThat(result.getMetadata()).containsEntry("distance", 2.220446049250313E-16);
|
||||
});
|
||||
|
||||
results = this.vectorStore.similaritySearch(SearchRequest.builder()
|
||||
.query("first")
|
||||
.filterExpression(new Filter.Expression(AND,
|
||||
new Filter.Expression(GTE, new Filter.Key("activationDate"), new Filter.Value(new Date(2000))),
|
||||
new Filter.Expression(LTE, new Filter.Key("activationDate"), new Filter.Value(new Date(3000)))))
|
||||
.build());
|
||||
assertThat(results).hasSize(2).first().satisfies(result -> {
|
||||
assertThat(result.getId()).isEqualTo("1");
|
||||
assertThat(result.getText()).isEqualTo("first");
|
||||
assertThat(result.getMetadata()).hasSize(4);
|
||||
assertThat(result.getMetadata()).containsEntry("distance", 2.220446049250313E-16);
|
||||
});
|
||||
|
||||
results = this.vectorStore.similaritySearch(SearchRequest.builder()
|
||||
.query("test content")
|
||||
.filterExpression(
|
||||
new Filter.Expression(EQ, new Filter.Key("activationDate"), new Filter.Value(new Date(3000))))
|
||||
.build());
|
||||
assertThat(results).hasSize(1);
|
||||
}
|
||||
|
||||
}
|
||||
@@ -0,0 +1,199 @@
|
||||
/*
|
||||
* 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.filter.converter;
|
||||
|
||||
import org.junit.jupiter.api.Assertions;
|
||||
import org.junit.jupiter.api.Test;
|
||||
import org.springframework.ai.vectorstore.filter.Filter;
|
||||
import org.springframework.ai.vectorstore.filter.FilterExpressionConverter;
|
||||
import org.springframework.expression.ExpressionParser;
|
||||
import org.springframework.expression.spel.standard.SpelExpressionParser;
|
||||
import org.springframework.expression.spel.support.StandardEvaluationContext;
|
||||
|
||||
import java.util.Date;
|
||||
import java.util.List;
|
||||
import java.util.Map;
|
||||
|
||||
import static org.assertj.core.api.Assertions.assertThat;
|
||||
import static org.springframework.ai.vectorstore.filter.Filter.ExpressionType.*;
|
||||
|
||||
/**
|
||||
* @author Jemin Huh
|
||||
*/
|
||||
public class SimpleVectorStoreFilterExpressionConverterTests {
|
||||
|
||||
final FilterExpressionConverter converter = new SimpleVectorStoreFilterExpressionConverter();
|
||||
|
||||
@Test
|
||||
public void testDate() {
|
||||
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'");
|
||||
|
||||
StandardEvaluationContext context = new StandardEvaluationContext();
|
||||
ExpressionParser parser = new SpelExpressionParser();
|
||||
context.setVariable("metadata",
|
||||
Map.of("activationDate", "2024-01-07T14:29:12Z", "year", 2020, "country", "BG"));
|
||||
Assertions.assertTrue(parser.parseExpression(vectorExpr).getValue(context, Boolean.class));
|
||||
|
||||
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'");
|
||||
|
||||
context.setVariable("metadata",
|
||||
Map.of("activationDate", "1970-01-01T00:00:02Z", "year", 2020, "country", "BG"));
|
||||
Assertions.assertTrue(parser.parseExpression(vectorExpr).getValue(context, Boolean.class));
|
||||
|
||||
}
|
||||
|
||||
@Test
|
||||
public void testEQ() {
|
||||
String vectorExpr = this.converter
|
||||
.convertExpression(new Filter.Expression(EQ, new Filter.Key("country"), new Filter.Value("BG")));
|
||||
assertThat(vectorExpr).isEqualTo("#metadata['country'] == 'BG'");
|
||||
|
||||
StandardEvaluationContext context = new StandardEvaluationContext();
|
||||
ExpressionParser parser = new SpelExpressionParser();
|
||||
context.setVariable("metadata", Map.of("city", "Seoul", "year", 2020, "country", "BG"));
|
||||
Assertions.assertTrue(parser.parseExpression(vectorExpr).getValue(context, Boolean.class));
|
||||
|
||||
}
|
||||
|
||||
@Test
|
||||
public void tesEqAndGte() {
|
||||
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");
|
||||
|
||||
StandardEvaluationContext context = new StandardEvaluationContext();
|
||||
ExpressionParser parser = new SpelExpressionParser();
|
||||
context.setVariable("metadata", Map.of("genre", "drama", "year", 2020, "country", "BG"));
|
||||
Assertions.assertTrue(parser.parseExpression(vectorExpr).getValue(context, Boolean.class));
|
||||
|
||||
}
|
||||
|
||||
@Test
|
||||
public void tesIn() {
|
||||
String vectorExpr = this.converter.convertExpression(new Filter.Expression(IN, new Filter.Key("genre"),
|
||||
new Filter.Value(List.of("comedy", "documentary", "drama"))));
|
||||
assertThat(vectorExpr).isEqualTo("{'comedy','documentary','drama'}.contains(#metadata['genre'])");
|
||||
|
||||
StandardEvaluationContext context = new StandardEvaluationContext();
|
||||
ExpressionParser parser = new SpelExpressionParser();
|
||||
context.setVariable("metadata", Map.of("genre", "drama", "year", 2020, "country", "BG"));
|
||||
Assertions.assertTrue(parser.parseExpression(vectorExpr).getValue(context, Boolean.class));
|
||||
|
||||
}
|
||||
|
||||
@Test
|
||||
public void testNe() {
|
||||
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")),
|
||||
new Filter.Expression(NE, new Filter.Key("city"), new Filter.Value("Sofia")))));
|
||||
assertThat(vectorExpr)
|
||||
.isEqualTo("#metadata['year'] >= 2020 or #metadata['country'] == 'BG' and #metadata['city'] != 'Sofia'");
|
||||
|
||||
StandardEvaluationContext context = new StandardEvaluationContext();
|
||||
ExpressionParser parser = new SpelExpressionParser();
|
||||
context.setVariable("metadata", Map.of("city", "Seoul", "year", 2020, "country", "BG"));
|
||||
Assertions.assertTrue(parser.parseExpression(vectorExpr).getValue(context, Boolean.class));
|
||||
|
||||
}
|
||||
|
||||
@Test
|
||||
public void testGroup() {
|
||||
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")))),
|
||||
new Filter.Expression(NIN, new Filter.Key("city"), new Filter.Value(List.of("Sofia", "Plovdiv")))));
|
||||
assertThat(vectorExpr).isEqualTo(
|
||||
"(#metadata['year'] >= 2020 or #metadata['country'] == 'BG') and not {'Sofia','Plovdiv'}.contains(#metadata['city'])");
|
||||
|
||||
StandardEvaluationContext context = new StandardEvaluationContext();
|
||||
ExpressionParser parser = new SpelExpressionParser();
|
||||
context.setVariable("metadata", Map.of("city", "Seoul", "year", 2020, "country", "BG"));
|
||||
Assertions.assertTrue(parser.parseExpression(vectorExpr).getValue(context, Boolean.class));
|
||||
|
||||
}
|
||||
|
||||
@Test
|
||||
public void tesBoolean() {
|
||||
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")))));
|
||||
|
||||
assertThat(vectorExpr).isEqualTo(
|
||||
"#metadata['isOpen'] == true and #metadata['year'] >= 2020 and {'BG','NL','US'}.contains(#metadata['country'])");
|
||||
|
||||
StandardEvaluationContext context = new StandardEvaluationContext();
|
||||
ExpressionParser parser = new SpelExpressionParser();
|
||||
context.setVariable("metadata", Map.of("isOpen", true, "year", 2020, "country", "NL"));
|
||||
Assertions.assertTrue(parser.parseExpression(vectorExpr).getValue(context, Boolean.class));
|
||||
|
||||
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(NIN, new Filter.Key("country"), new Filter.Value(List.of("BG", "NL", "US")))));
|
||||
|
||||
assertThat(vectorExpr).isEqualTo(
|
||||
"#metadata['isOpen'] == true and #metadata['year'] >= 2020 and not {'BG','NL','US'}.contains(#metadata['country'])");
|
||||
|
||||
context.setVariable("metadata", Map.of("isOpen", true, "year", 2020, "country", "KR"));
|
||||
Assertions.assertTrue(parser.parseExpression(vectorExpr).getValue(context, Boolean.class));
|
||||
}
|
||||
|
||||
@Test
|
||||
public void testDecimal() {
|
||||
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))));
|
||||
|
||||
assertThat(vectorExpr).isEqualTo("#metadata['temperature'] >= -15.6 and #metadata['temperature'] <= 20.13");
|
||||
|
||||
StandardEvaluationContext context = new StandardEvaluationContext();
|
||||
ExpressionParser parser = new SpelExpressionParser();
|
||||
context.setVariable("metadata", Map.of("temperature", -15.6));
|
||||
Assertions.assertTrue(parser.parseExpression(vectorExpr).getValue(context, Boolean.class));
|
||||
context.setVariable("metadata", Map.of("temperature", 20.13));
|
||||
Assertions.assertTrue(parser.parseExpression(vectorExpr).getValue(context, Boolean.class));
|
||||
context.setVariable("metadata", Map.of("temperature", -1.6));
|
||||
Assertions.assertTrue(parser.parseExpression(vectorExpr).getValue(context, Boolean.class));
|
||||
|
||||
}
|
||||
|
||||
@Test
|
||||
public void testComplexIdentifiers() {
|
||||
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 = 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'");
|
||||
|
||||
StandardEvaluationContext context = new StandardEvaluationContext();
|
||||
ExpressionParser parser = new SpelExpressionParser();
|
||||
context.setVariable("metadata", Map.of("country 1 2 3", "BG"));
|
||||
Assertions.assertTrue(parser.parseExpression(vectorExpr).getValue(context, Boolean.class));
|
||||
}
|
||||
|
||||
}
|
||||
Reference in New Issue
Block a user