Add support for Redis vector store

- Added autoconfiguration for Redis vector store
- Added spring boot starter for Redis vector store
- Supports portable metadata filter expressions

Fixes #11
This commit is contained in:
jruaux
2023-12-17 16:07:05 -08:00
committed by Mark Pollack
parent 55d748b86e
commit a232166cf3
17 changed files with 1653 additions and 0 deletions

View File

@@ -30,6 +30,7 @@
<module>spring-ai-spring-boot-starters/spring-ai-starter-pinecone-store</module>
<module>spring-ai-spring-boot-starters/spring-ai-starter-azure-store</module>
<module>spring-ai-spring-boot-starters/spring-ai-starter-weaviate-store</module>
<module>spring-ai-spring-boot-starters/spring-ai-starter-redis</module>
<module>spring-ai-docs</module>
<module>vector-stores/spring-ai-pgvector-store</module>
<module>vector-stores/spring-ai-milvus-store</module>
@@ -42,6 +43,7 @@
<module>vector-stores/spring-ai-chroma</module>
<module>vector-stores/spring-ai-azure</module>
<module>vector-stores/spring-ai-weaviate</module>
<module>vector-stores/spring-ai-redis</module>
<module>spring-ai-vertex-ai</module>
<module>spring-ai-spring-boot-starters/spring-ai-starter-vertex-ai</module>

View File

@@ -16,6 +16,7 @@
*** xref:api/vectordbs/neo4j.adoc[]
*** xref:api/vectordbs/pgvector.adoc[]
*** xref:api/vectordbs/weaviate.adoc[]
*** xref:api/vectordbs/redis.adoc[]
** xref:api/testing.adoc[]
* Appendices
** xref:glossary.adoc[]

View File

@@ -93,6 +93,10 @@ public interface DocumentWriter extends Consumer<List<Document>> {
*Neo4jVectorStore*::
+ Leverages the Neo4j graph database for vector storage.
*RedisVectorStore*::
+ Provides vector storage capabilities using Redis.
== Using PDF Reader

View File

@@ -94,6 +94,7 @@ These are the available implementations of the `VectorStore` interface:
* Neo4j [`Neo4jVectorStore`]: The https://neo4j.com/[Neo4j] vector store
* Weaviate [`WeaviateVectorStore`] The https://weaviate.io/[Weaviate] vector store
* Azure Vector Search [`AzureVectorStore`] the https://learn.microsoft.com/en-us/azure/search/vector-search-overview[Azure] vector store
* Redisj [`RedisVectorStore`]: The https://redis.io/[Redis] vector store
More implementations may be supported in future releases.

View File

@@ -0,0 +1,183 @@
= Redis
This section walks you through setting up `RedisVectorStore` to store document embeddings and perform similarity searches.
== What is Redis?
link:https://redis.io[Redis] is an open source (BSD licensed), in-memory data structure store used as a database, cache, message broker, and streaming engine. Redis provides data structures such as strings, hashes, lists, sets, sorted sets with range queries, bitmaps, hyperloglogs, geospatial indexes, and streams.
== What is Redis Vector Search?
link:https://redis.io/docs/interact/search-and-query/[Redis Search and Query] extends the core features of Redis OSS and allows you to use Redis as a vector database:
* Store vectors and the associated metadata within hashes or JSON documents
* Retrieve vectors
* Perform vector searches
== Prerequisites
1. `EmbeddingClient` instance to compute the document embeddings. Several options are available:
- `Transformers Embedding` - computes the embedding in your local environment. Follow the ONNX Transformers Embedding instructions.
- `OpenAI Embedding` - uses the OpenAI embedding endpoint. You need to create an account at link:https://platform.openai.com/signup[OpenAI Signup] and generate the api-key token at link:https://platform.openai.com/account/api-keys[API Keys].
- You can also use the `Azure OpenAI Embedding`.
2. A Redis Stack instance
a. https://app.redislabs.com/#/[Redis Cloud] (recommended)
b. link:https://hub.docker.com/r/redis/redis-stack[Docker] image _redis/redis-stack:latest_
== Dependencies
Add these dependencies to your project:
* Embedding Client boot starter, required for calculating embeddings.
* Transformers Embedding (Local) and follow the ONNX Transformers Embedding instructions.
[source,xml]
----
<dependency>
<groupId>org.springframework.ai</groupId>
<artifactId>spring-ai-transformers-embedding-spring-boot-starter</artifactId>
<version>0.8.0-SNAPSHOT</version>
</dependency>
----
or use OpenAI (Cloud)
[source,xml]
----
<dependency>
<groupId>org.springframework.ai</groupId>
<artifactId>spring-ai-openai-spring-boot-starter</artifactId>
<version>0.8.0-SNAPSHOT</version>
</dependency>
----
You'll need to provide your OpenAI API Key. Set it as an environment variable like so:
[source,bash]
----
export SPRING_AI_OPENAI_API_KEY='Your_OpenAI_API_Key'
----
* Add the Redis Vector Store and Jedis dependencies
[source,xml]
----
<dependency>
<groupId>org.springframework.ai</groupId>
<artifactId>spring-ai-redis-store</artifactId>
<version>0.8.0-SNAPSHOT</version>
</dependency>
<dependency>
<groupId>redis.clients</groupId>
<artifactId>jedis</artifactId>
<version>5.1.0</version>
</dependency>
----
== Usage
Create a RedisVectorStore instance connected to your Redis database:
[source,java]
----
@Bean
public VectorStore vectorStore(EmbeddingClient embeddingClient) {
RedisVectorStoreConfig config = RedisVectorStoreConfig.builder()
.withURI("redis://localhost:6379")
// Define the metadata fields to be used
// in the similarity search filters.
.withMetadataFields(
MetadataField.tag("country"),
MetadataField.numeric("year"))
.build();
return new RedisVectorStore(config, embeddingClient);
}
----
> [NOTE]
> You must list explicitly all metadata field names and types (`TAG`, `TEXT`, or `NUMERIC`) for any metadata field used in filter expression.
> The `withMetadataFields` above registers filterable metadata fields: `country` of type `TAG`, `year` of type `NUMERIC`.
>
Then in your main code, create some documents:
[source,java]
----
List<Document> documents = List.of(
new Document("Spring AI rocks!! Spring AI rocks!! Spring AI rocks!! Spring AI rocks!! Spring AI rocks!!", Map.of("country", "UK", "year", 2020)),
new Document("The World is Big and Salvation Lurks Around the Corner", Map.of()),
new Document("You walk forward facing the past and you turn back toward the future.", Map.of("country", "NL", "year", 2023)));
----
Now add the documents to your vector store:
[source,java]
----
vectorStore.add(List.of(document));
----
And finally, retrieve documents similar to a query:
[source,java]
----
List<Document> results = vectorStore.similaritySearch(
SearchRequest
.query("Spring")
.withTopK(5));
----
If all goes well, you should retrieve the document containing the text "Spring AI rocks!!".
=== Metadata filtering
You can leverage the generic, portable link:https://docs.spring.io/spring-ai/reference/api/vectordbs.html#_metadata_filters[metadata filters] with RedisVectorStore as well.
For example, you can use either the text expression language:
[source,java]
----
vectorStore.similaritySearch(
SearchRequest
.query("The World")
.withTopK(TOP_K)
.withSimilarityThreshold(SIMILARITY_THRESHOLD)
.withFilterExpression("country in ['UK', 'NL'] && year >= 2020"));
----
or programmatically using the expression DSL:
[source,java]
----
FilterExpressionBuilder b = Filter.builder();
vectorStore.similaritySearch(
SearchRequest
.query("The World")
.withTopK(TOP_K)
.withSimilarityThreshold(SIMILARITY_THRESHOLD)
.withFilterExpression(b.and(
b.in("country", "UK", "NL"),
b.gte("year", 2020)).build()));
----
The portable filter expressions get automatically converted into link:https://redis.io/docs/interact/search-and-query/query/[Redis search queries].
For example, the following portable filter expression:
[source,sql]
----
country in ['UK', 'NL'] && year >= 2020
----
is converted into Redis query:
[source]
----
@country:{UK | NL} @year:[2020 inf]
----

View File

@@ -132,6 +132,21 @@
<version>${project.parent.version}</version>
<optional>true</optional>
</dependency>
<!-- Redis Vector Store-->
<dependency>
<groupId>org.springframework.ai</groupId>
<artifactId>spring-ai-redis</artifactId>
<version>${project.parent.version}</version>
<optional>true</optional>
</dependency>
<!-- Override Jedis version -->
<dependency>
<groupId>redis.clients</groupId>
<artifactId>jedis</artifactId>
<version>5.1.0</version>
</dependency>
<!-- Vertex AI LLM -->
<dependency>
@@ -186,6 +201,13 @@
<version>${testcontainers.version}</version>
<scope>test</scope>
</dependency>
<dependency>
<groupId>com.redis</groupId>
<artifactId>testcontainers-redis</artifactId>
<version>2.0.1</version>
<scope>test</scope>
</dependency>
<dependency>
<groupId>org.awaitility</groupId>

View File

@@ -0,0 +1,50 @@
/*
* Copyright 2023-2023 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.autoconfigure.vectorstore.redis;
import org.springframework.ai.embedding.EmbeddingClient;
import org.springframework.ai.vectorstore.RedisVectorStore;
import org.springframework.ai.vectorstore.RedisVectorStore.RedisVectorStoreConfig;
import org.springframework.ai.vectorstore.VectorStore;
import org.springframework.boot.autoconfigure.AutoConfiguration;
import org.springframework.boot.autoconfigure.condition.ConditionalOnClass;
import org.springframework.boot.autoconfigure.condition.ConditionalOnMissingBean;
import org.springframework.boot.context.properties.EnableConfigurationProperties;
import org.springframework.context.annotation.Bean;
/**
* @author Christian Tzolov
*/
@AutoConfiguration
@ConditionalOnClass({ RedisVectorStore.class, EmbeddingClient.class })
@EnableConfigurationProperties(RedisVectorStoreProperties.class)
public class RedisVectorStoreAutoConfiguration {
@Bean
@ConditionalOnMissingBean
public VectorStore vectorStore(EmbeddingClient embeddingClient, RedisVectorStoreProperties properties) {
var config = RedisVectorStoreConfig.builder()
.withURI(properties.getUri())
.withIndexName(properties.getIndex())
.withPrefix(properties.getPrefix())
.build();
return new RedisVectorStore(config, embeddingClient);
}
}

View File

@@ -0,0 +1,61 @@
/*
* Copyright 2023-2023 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.autoconfigure.vectorstore.redis;
import org.springframework.boot.context.properties.ConfigurationProperties;
import static org.springframework.ai.autoconfigure.vectorstore.redis.RedisVectorStoreProperties.CONFIG_PREFIX;
/**
* @author Julien Ruaux
*/
@ConfigurationProperties(CONFIG_PREFIX)
public class RedisVectorStoreProperties {
public static final String CONFIG_PREFIX = "spring.ai.vectorstore.redis";
private String uri = "redis://localhost:6379";
private String index;
private String prefix;
public String getUri() {
return uri;
}
public void setUri(String uri) {
this.uri = uri;
}
public String getIndex() {
return index;
}
public void setIndex(String name) {
this.index = name;
}
public String getPrefix() {
return prefix;
}
public void setPrefix(String prefix) {
this.prefix = prefix;
}
}

View File

@@ -4,6 +4,7 @@ org.springframework.ai.autoconfigure.azure.openai.AzureOpenAiAutoConfiguration
org.springframework.ai.autoconfigure.vectorstore.pgvector.PgVectorStoreAutoConfiguration
org.springframework.ai.autoconfigure.vectorstore.pinecone.PineconeVectorStoreAutoConfiguration
org.springframework.ai.autoconfigure.vectorstore.milvus.MilvusVectorStoreAutoConfiguration
org.springframework.ai.autoconfigure.vectorstore.redis.RedisVectorStoreAutoConfiguration
org.springframework.ai.autoconfigure.embedding.transformer.TransformersEmbeddingClientAutoConfiguration
org.springframework.ai.autoconfigure.huggingface.HuggingfaceAutoConfiguration
org.springframework.ai.autoconfigure.vectorstore.chroma.ChromaVectorStoreAutoConfiguration

View File

@@ -0,0 +1,94 @@
/*
* Copyright 2023-2023 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.autoconfigure.vectorstore.redis;
import static org.assertj.core.api.Assertions.assertThat;
import java.util.List;
import java.util.Map;
import org.junit.jupiter.api.Test;
import org.springframework.ai.ResourceUtils;
import org.springframework.ai.document.Document;
import org.springframework.ai.embedding.EmbeddingClient;
import org.springframework.ai.embedding.TransformersEmbeddingClient;
import org.springframework.ai.vectorstore.SearchRequest;
import org.springframework.ai.vectorstore.VectorStore;
import org.springframework.boot.autoconfigure.AutoConfigurations;
import org.springframework.boot.test.context.runner.ApplicationContextRunner;
import org.springframework.context.annotation.Bean;
import org.springframework.context.annotation.Configuration;
import org.testcontainers.junit.jupiter.Container;
import org.testcontainers.junit.jupiter.Testcontainers;
import com.redis.testcontainers.RedisStackContainer;
/**
* @author Julien Ruaux
*/
@Testcontainers
class RedisVectorStoreAutoConfigurationIT {
@Container
static RedisStackContainer redisContainer = new RedisStackContainer(
RedisStackContainer.DEFAULT_IMAGE_NAME.withTag(RedisStackContainer.DEFAULT_TAG));
List<Document> documents = List.of(
new Document(ResourceUtils.getText("classpath:/test/data/spring.ai.txt"), Map.of("spring", "great")),
new Document(ResourceUtils.getText("classpath:/test/data/time.shelter.txt")), new Document(
ResourceUtils.getText("classpath:/test/data/great.depression.txt"), Map.of("depression", "bad")));
private final ApplicationContextRunner contextRunner = new ApplicationContextRunner()
.withConfiguration(AutoConfigurations.of(RedisVectorStoreAutoConfiguration.class))
.withUserConfiguration(Config.class)
.withPropertyValues("spring.ai.vectorstore.redis.index=myIdx")
.withPropertyValues("spring.ai.vectorstore.redis.prefix=doc:");
@Test
void addAndSearch() {
contextRunner.withPropertyValues("spring.ai.vectorstore.redis.uri=" + redisContainer.getRedisURI())
.run(context -> {
VectorStore vectorStore = context.getBean(VectorStore.class);
vectorStore.add(documents);
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.");
// Remove all documents from the store
vectorStore.delete(documents.stream().map(doc -> doc.getId()).toList());
results = vectorStore.similaritySearch(SearchRequest.query("Spring").withTopK(1));
assertThat(results).isEmpty();
});
}
@Configuration(proxyBeanMethods = false)
static class Config {
@Bean
public EmbeddingClient embeddingClient() {
return new TransformersEmbeddingClient();
}
}
}

View File

@@ -0,0 +1,48 @@
/*
* Copyright 2023-2023 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.autoconfigure.vectorstore.redis;
import static org.assertj.core.api.Assertions.assertThat;
import org.junit.jupiter.api.Test;
/**
* @author Julien Ruaux
*/
class RedisVectorStorePropertiesTests {
@Test
void defaultValues() {
var props = new RedisVectorStoreProperties();
assertThat(props.getUri()).isEqualTo("redis://localhost:6379");
assertThat(props.getIndex()).isNull();
assertThat(props.getPrefix()).isNull();
}
@Test
void customValues() {
var props = new RedisVectorStoreProperties();
props.setUri("redis://redis.com:12345");
props.setIndex("myIdx");
props.setPrefix("doc:");
assertThat(props.getUri()).isEqualTo("redis://redis.com:12345");
assertThat(props.getIndex()).isEqualTo("myIdx");
assertThat(props.getPrefix()).isEqualTo("doc:");
}
}

View File

@@ -0,0 +1,42 @@
<?xml version="1.0" encoding="UTF-8"?>
<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>
<parent>
<groupId>org.springframework.ai</groupId>
<artifactId>spring-ai</artifactId>
<version>0.8.0-SNAPSHOT</version>
<relativePath>../../pom.xml</relativePath>
</parent>
<artifactId>spring-ai-redis-spring-boot-starter</artifactId>
<packaging>jar</packaging>
<name>Spring AI Starter - Redis Vector Store</name>
<description>Spring AI Redis Vector Store Auto Configuration</description>
<url>https://github.com/spring-projects/spring-ai</url>
<scm>
<url>https://github.com/spring-projects/spring-ai</url>
<connection>git://github.com/spring-projects/spring-ai.git</connection>
<developerConnection>git@github.com:spring-projects/spring-ai.git</developerConnection>
</scm>
<dependencies>
<dependency>
<groupId>org.springframework.boot</groupId>
<artifactId>spring-boot-starter</artifactId>
</dependency>
<dependency>
<groupId>org.springframework.ai</groupId>
<artifactId>spring-ai-spring-boot-autoconfigure</artifactId>
<version>${project.parent.version}</version>
</dependency>
<dependency>
<groupId>org.springframework.ai</groupId>
<artifactId>spring-ai-redis</artifactId>
<version>${project.parent.version}</version>
</dependency>
</dependencies>
</project>

View File

@@ -0,0 +1,80 @@
<?xml version="1.0" encoding="UTF-8"?>
<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>
<parent>
<groupId>org.springframework.ai</groupId>
<artifactId>spring-ai</artifactId>
<version>0.8.0-SNAPSHOT</version>
<relativePath>../../pom.xml</relativePath>
</parent>
<artifactId>spring-ai-redis</artifactId>
<packaging>jar</packaging>
<name>spring-ai-redis</name>
<description>Spring AI Redis Vector Store</description>
<url>https://github.com/spring-projects/spring-ai</url>
<scm>
<url>https://github.com/spring-projects-experimental/spring-ai</url>
<connection>git://github.com/spring-projects-experimental/spring-ai.git</connection>
<developerConnection>git@github.com:spring-projects-experimental/spring-ai.git</developerConnection>
</scm>
<properties>
<testcontainers-redis.version>2.0.1</testcontainers-redis.version>
<jedis.version>5.1.0</jedis.version>
<maven.compiler.target>17</maven.compiler.target>
<maven.compiler.source>17</maven.compiler.source>
</properties>
<dependencies>
<dependency>
<groupId>org.springframework.ai</groupId>
<artifactId>spring-ai-core</artifactId>
<version>${parent.version}</version>
</dependency>
<dependency>
<groupId>redis.clients</groupId>
<artifactId>jedis</artifactId>
<version>${jedis.version}</version>
</dependency>
<!-- TESTING -->
<dependency>
<groupId>org.springframework.ai</groupId>
<artifactId>transformers-embedding</artifactId>
<version>${parent.version}</version>
<scope>test</scope>
</dependency>
<dependency>
<groupId>org.springframework.ai</groupId>
<artifactId>spring-ai-test</artifactId>
<version>${parent.version}</version>
<scope>test</scope>
</dependency>
<dependency>
<groupId>org.springframework.boot</groupId>
<artifactId>spring-boot-starter-test</artifactId>
<scope>test</scope>
</dependency>
<dependency>
<groupId>com.redis</groupId>
<artifactId>testcontainers-redis</artifactId>
<version>${testcontainers-redis.version}</version>
<scope>test</scope>
</dependency>
<dependency>
<groupId>org.testcontainers</groupId>
<artifactId>junit-jupiter</artifactId>
<scope>test</scope>
</dependency>
</dependencies>
</project>

View File

@@ -0,0 +1,214 @@
/*
* Copyright 2023-2023 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 java.text.MessageFormat;
import java.util.List;
import java.util.Map;
import java.util.function.Function;
import java.util.stream.Collectors;
import org.springframework.ai.vectorstore.RedisVectorStore.MetadataField;
import org.springframework.ai.vectorstore.filter.Filter.Expression;
import org.springframework.ai.vectorstore.filter.Filter.ExpressionType;
import org.springframework.ai.vectorstore.filter.Filter.Group;
import org.springframework.ai.vectorstore.filter.Filter.Key;
import org.springframework.ai.vectorstore.filter.Filter.Value;
import org.springframework.ai.vectorstore.filter.converter.AbstractFilterExpressionConverter;
/**
* Converts {@link Expression} into Redis search filter expression format.
* (https://redis.io/docs/interact/search-and-query/query/)
*
* @author Julien Ruaux
*/
public class RedisFilterExpressionConverter extends AbstractFilterExpressionConverter {
public static final NumericBoundary POSITIVE_INFINITY = new NumericBoundary(Double.POSITIVE_INFINITY, true);
public static final NumericBoundary NEGATIVE_INFINITY = new NumericBoundary(Double.NEGATIVE_INFINITY, true);
private Map<String, MetadataField> metadataFields;
public RedisFilterExpressionConverter(List<MetadataField> metadataFields) {
this.metadataFields = metadataFields.stream()
.collect(Collectors.toMap(MetadataField::name, Function.identity()));
}
@Override
protected void doStartGroup(Group group, StringBuilder context) {
context.append("(");
}
@Override
protected void doEndGroup(Group group, StringBuilder context) {
context.append(")");
}
@Override
protected void doKey(Key key, StringBuilder context) {
context.append("@").append(key.key()).append(":");
}
@Override
protected void doExpression(Expression expression, StringBuilder context) {
switch (expression.type()) {
case NIN:
doExpression(negate(ExpressionType.IN, expression), context);
break;
case NE:
doExpression(negate(ExpressionType.EQ, expression), context);
break;
case AND:
doBinaryOperation(" ", expression, context);
break;
case OR:
doBinaryOperation(" | ", expression, context);
break;
case NOT:
context.append("-");
convertOperand(expression.left(), context);
break;
default:
doField(expression, context);
break;
}
}
private Expression negate(ExpressionType expressionType, Expression expression) {
return new Expression(ExpressionType.NOT, new Expression(expressionType, expression.left(), expression.right()),
null);
}
private void doBinaryOperation(String delimiter, Expression expression, StringBuilder context) {
this.convertOperand(expression.left(), context);
context.append(delimiter);
this.convertOperand(expression.right(), context);
}
private void doField(Expression expression, StringBuilder context) {
Key key = (Key) expression.left();
doKey(key, context);
MetadataField field = metadataFields.getOrDefault(key.key(), MetadataField.tag(key.key()));
Value value = (Value) expression.right();
switch (field.fieldType()) {
case NUMERIC:
Numeric numeric = numeric(expression, value);
context.append("[");
context.append(numeric.lower());
context.append(" ");
context.append(numeric.upper());
context.append("]");
break;
case TAG:
context.append("{");
context.append(stringValue(expression, value));
context.append("}");
break;
case TEXT:
context.append("(");
context.append(stringValue(expression, value));
context.append(")");
break;
default:
throw new UnsupportedOperationException(
MessageFormat.format("Field type {0} not supported", field.fieldType()));
}
}
private Object stringValue(Expression expression, Value value) {
String delimiter = tagValueDelimiter(expression);
if (value.value() instanceof List<?> list) {
return String.join(delimiter, list.stream().map(String::valueOf).toList());
}
return value.value();
}
private String tagValueDelimiter(Expression expression) {
switch (expression.type()) {
case IN:
return " | ";
case EQ:
return " ";
default:
throw new UnsupportedOperationException(
MessageFormat.format("Tag operand {0} not supported", expression.type()));
}
}
private Numeric numeric(Expression expression, Value value) {
switch (expression.type()) {
case EQ:
return new Numeric(inclusive(value), inclusive(value));
case GT:
return new Numeric(exclusive(value), POSITIVE_INFINITY);
case GTE:
return new Numeric(inclusive(value), POSITIVE_INFINITY);
case LT:
return new Numeric(NEGATIVE_INFINITY, exclusive(value));
case LTE:
return new Numeric(NEGATIVE_INFINITY, inclusive(value));
default:
throw new UnsupportedOperationException(MessageFormat
.format("Expression type {0} not supported for numeric fields", expression.type()));
}
}
private NumericBoundary inclusive(Value value) {
return new NumericBoundary(value.value(), false);
}
private NumericBoundary exclusive(Value value) {
return new NumericBoundary(value.value(), true);
}
static record Numeric(NumericBoundary lower, NumericBoundary upper) {
}
static record NumericBoundary(Object value, boolean exclusive) {
private static final String INFINITY = "inf";
private static final String MINUS_INFINITY = "-inf";
private static final String INCLUSIVE_FORMAT = "%s";
private static final String EXCLUSIVE_FORMAT = "(%s";
@Override
public String toString() {
if (this == NEGATIVE_INFINITY) {
return MINUS_INFINITY;
}
if (this == POSITIVE_INFINITY) {
return INFINITY;
}
return String.format(formatString(), value);
}
private String formatString() {
if (exclusive) {
return EXCLUSIVE_FORMAT;
}
return INCLUSIVE_FORMAT;
}
}
}

View File

@@ -0,0 +1,473 @@
/*
* Copyright 2023-2023 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 java.text.MessageFormat;
import java.util.ArrayList;
import java.util.Arrays;
import java.util.HashMap;
import java.util.List;
import java.util.Map;
import java.util.Optional;
import java.util.function.Function;
import java.util.function.Predicate;
import java.util.stream.Collectors;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import org.springframework.ai.document.Document;
import org.springframework.ai.embedding.EmbeddingClient;
import org.springframework.ai.vectorstore.filter.converter.FilterExpressionConverter;
import org.springframework.beans.factory.InitializingBean;
import org.springframework.util.Assert;
import org.springframework.util.CollectionUtils;
import redis.clients.jedis.JedisPooled;
import redis.clients.jedis.Pipeline;
import redis.clients.jedis.json.Path2;
import redis.clients.jedis.search.FTCreateParams;
import redis.clients.jedis.search.IndexDataType;
import redis.clients.jedis.search.Query;
import redis.clients.jedis.search.RediSearchUtil;
import redis.clients.jedis.search.Schema.FieldType;
import redis.clients.jedis.search.SearchResult;
import redis.clients.jedis.search.schemafields.NumericField;
import redis.clients.jedis.search.schemafields.SchemaField;
import redis.clients.jedis.search.schemafields.TagField;
import redis.clients.jedis.search.schemafields.TextField;
import redis.clients.jedis.search.schemafields.VectorField;
import redis.clients.jedis.search.schemafields.VectorField.VectorAlgorithm;
/**
* The RedisVectorStore is for managing and querying vector data in a Redis database. It
* offers functionalities like adding, deleting, and performing similarity searches on
* documents.
*
* The store utilizes RedisJSON and RediSearch to handle JSON documents and to index and
* search vector data. It supports various vector algorithms (e.g., FLAT, HSNW) for
* efficient similarity searches. Additionally, it allows for custom metadata fields in
* the documents to be stored alongside the vector and content data.
*
* This class requires a RedisVectorStoreConfig configuration object for initialization,
* which includes settings like Redis URI, index name, field names, and vector algorithms.
* It also requires an EmbeddingClient to convert documents into embeddings before storing
* them.
*
* @author Julien Ruaux
* @see VectorStore
* @see RedisVectorStoreConfig
* @see EmbeddingClient
*/
public class RedisVectorStore implements VectorStore, InitializingBean {
public enum Algorithm {
FLAT, HSNW
}
public record MetadataField(String name, FieldType fieldType) {
public static MetadataField text(String name) {
return new MetadataField(name, FieldType.TEXT);
}
public static MetadataField numeric(String name) {
return new MetadataField(name, FieldType.NUMERIC);
}
public static MetadataField tag(String name) {
return new MetadataField(name, FieldType.TAG);
}
}
/**
* Configuration for the Redis vector store.
*/
public static final class RedisVectorStoreConfig {
private final String uri;
private final String indexName;
private final String prefix;
private final String contentFieldName;
private final String embeddingFieldName;
private final Algorithm vectorAlgorithm;
private final List<MetadataField> metadataFields;
/**
* 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 RedisVectorStoreConfig defaultConfig() {
return builder().build();
}
private RedisVectorStoreConfig(Builder builder) {
this.uri = builder.uri;
this.indexName = builder.indexName;
this.prefix = builder.prefix;
this.contentFieldName = builder.contentFieldName;
this.embeddingFieldName = builder.embeddingFieldName;
this.vectorAlgorithm = builder.vectorAlgorithm;
this.metadataFields = builder.metadataFields;
}
public static class Builder {
private String uri = DEFAULT_URI;
private String indexName = DEFAULT_INDEX_NAME;
private String prefix = DEFAULT_PREFIX;
private String contentFieldName = DEFAULT_CONTENT_FIELD_NAME;
private String embeddingFieldName = DEFAULT_EMBEDDING_FIELD_NAME;
private Algorithm vectorAlgorithm = DEFAULT_VECTOR_ALGORITHM;
private List<MetadataField> metadataFields = new ArrayList<>();
private Builder() {
}
/**
* Configures the Redis URI to use.
* @param uri the Redis URI to use
* @return this builder
*/
public Builder withURI(String uri) {
this.uri = uri;
return this;
}
/**
* Configures the Redis index name to use.
* @param name the index name to use
* @return this builder
*/
public Builder withIndexName(String name) {
this.indexName = name;
return this;
}
/**
* Configures the Redis key prefix to use (default: "embedding:").
* @param prefix the prefix to use
* @return this builder
*/
public Builder withPrefix(String prefix) {
this.prefix = prefix;
return this;
}
/**
* Configures the Redis content field name to use.
* @param name the content field name to use
* @return this builder
*/
public Builder withContentFieldName(String name) {
this.contentFieldName = name;
return this;
}
/**
* Configures the Redis embedding field name to use.
* @param name the embedding field name to use
* @return this builder
*/
public Builder withEmbeddingFieldName(String name) {
this.embeddingFieldName = name;
return this;
}
/**
* Configures the Redis vector algorithmto use.
* @param algorithm the vector algorithm to use
* @return this builder
*/
public Builder withVectorAlgorithm(Algorithm algorithm) {
this.vectorAlgorithm = algorithm;
return this;
}
public Builder withMetadataFields(MetadataField... fields) {
return withMetadataFields(Arrays.asList(fields));
}
public Builder withMetadataFields(List<MetadataField> fields) {
this.metadataFields = fields;
return this;
}
/**
* {@return the immutable configuration}
*/
public RedisVectorStoreConfig build() {
return new RedisVectorStoreConfig(this);
}
}
}
public static final String DEFAULT_URI = "redis://localhost:6379";
public static final String DEFAULT_INDEX_NAME = "spring-ai-index";
public static final String DEFAULT_CONTENT_FIELD_NAME = "content";
public static final String DEFAULT_EMBEDDING_FIELD_NAME = "embedding";
public static final String DEFAULT_PREFIX = "embedding:";
public static final Algorithm DEFAULT_VECTOR_ALGORITHM = Algorithm.HSNW;
private static final String QUERY_FORMAT = "%s=>[KNN %s @%s $%s AS %s]";
private static final Path2 JSON_SET_PATH = Path2.of("$");
private static final String JSON_PATH_PREFIX = "$.";
private static final Logger logger = LoggerFactory.getLogger(RedisVectorStore.class);
private static final Predicate<Object> RESPONSE_OK = Predicate.isEqual("OK");
private static final Predicate<Object> RESPONSE_DEL_OK = Predicate.isEqual(1l);
private static final String VECTOR_TYPE_FLOAT32 = "FLOAT32";
private static final String EMBEDDING_PARAM_NAME = "BLOB";
public static final String DISTANCE_FIELD_NAME = "vector_score";
private static final String DEFAULT_DISTANCE_METRIC = "COSINE";
private final JedisPooled jedis;
private final EmbeddingClient embeddingClient;
private final RedisVectorStoreConfig config;
private FilterExpressionConverter filterExpressionConverter;
public RedisVectorStore(RedisVectorStoreConfig config, EmbeddingClient embeddingClient) {
Assert.notNull(config, "Config must not be null");
Assert.notNull(embeddingClient, "Embedding client must not be null");
this.jedis = new JedisPooled(config.uri);
this.embeddingClient = embeddingClient;
this.config = config;
}
public JedisPooled getJedis() {
return jedis;
}
@Override
public void add(List<Document> documents) {
Pipeline pipeline = jedis.pipelined();
for (Document document : documents) {
var embedding = this.embeddingClient.embed(document);
document.setEmbedding(embedding);
var fields = new HashMap<String, Object>();
fields.put(config.embeddingFieldName, embedding);
fields.put(config.contentFieldName, document.getContent());
fields.putAll(document.getMetadata());
pipeline.jsonSetWithEscape(key(document.getId()), JSON_SET_PATH, fields);
}
List<Object> responses = pipeline.syncAndReturnAll();
Optional<Object> errResponse = responses.stream().filter(Predicate.not(RESPONSE_OK)).findAny();
if (errResponse.isPresent()) {
String message = MessageFormat.format("Could not add document: {0}", errResponse.get());
if (logger.isErrorEnabled()) {
logger.error(message);
}
throw new RuntimeException(message);
}
}
private String key(String id) {
return config.prefix + id;
}
@Override
public Optional<Boolean> delete(List<String> idList) {
Pipeline pipeline = jedis.pipelined();
for (String id : idList) {
pipeline.jsonDel(key(id));
}
List<Object> responses = pipeline.syncAndReturnAll();
Optional<Object> errResponse = responses.stream().filter(Predicate.not(RESPONSE_DEL_OK)).findAny();
if (errResponse.isPresent()) {
if (logger.isErrorEnabled()) {
logger.error("Could not delete document: {}", errResponse.get());
}
return Optional.of(false);
}
return Optional.of(true);
}
@Override
public List<Document> similaritySearch(SearchRequest request) {
Assert.isTrue(request.getTopK() > 0, "The number of documents to returned must be greater than zero");
Assert.isTrue(request.getSimilarityThreshold() >= 0 && request.getSimilarityThreshold() <= 1,
"The similarity score is bounded between 0 and 1; least to most similar respectively.");
String filter = nativeExpressionFilter(request);
String queryString = String.format(QUERY_FORMAT, filter, request.getTopK(), config.embeddingFieldName,
EMBEDDING_PARAM_NAME, DISTANCE_FIELD_NAME);
List<String> returnFields = new ArrayList<>();
config.metadataFields.stream().map(MetadataField::name).forEach(returnFields::add);
returnFields.add(config.embeddingFieldName);
returnFields.add(config.contentFieldName);
returnFields.add(DISTANCE_FIELD_NAME);
var embedding = toFloatArray(this.embeddingClient.embed(request.getQuery()));
Query query = new Query(queryString).addParam(EMBEDDING_PARAM_NAME, RediSearchUtil.toByteArray(embedding))
.returnFields(returnFields.toArray(new String[0]))
.setSortBy(DISTANCE_FIELD_NAME, true)
.dialect(2);
SearchResult result = jedis.ftSearch(config.indexName, query);
return result.getDocuments()
.stream()
.filter(d -> similarityScore(d) >= request.getSimilarityThreshold())
.map(this::toDocument)
.toList();
}
private Document toDocument(redis.clients.jedis.search.Document doc) {
var id = doc.getId().substring(config.prefix.length());
var content = doc.hasProperty(config.contentFieldName) ? doc.getString(config.contentFieldName) : null;
Map<String, Object> metadata = config.metadataFields.stream()
.map(MetadataField::name)
.filter(doc::hasProperty)
.collect(Collectors.toMap(Function.identity(), doc::getString));
metadata.put(DISTANCE_FIELD_NAME, 1 - similarityScore(doc));
return new Document(id, content, metadata);
}
private float similarityScore(redis.clients.jedis.search.Document doc) {
return (2 - Float.parseFloat(doc.getString(DISTANCE_FIELD_NAME))) / 2;
}
private String nativeExpressionFilter(SearchRequest request) {
if (request.getFilterExpression() == null) {
return "*";
}
return "(" + filterExpressionConverter.convertExpression(request.getFilterExpression()) + ")";
}
@Override
public void afterPropertiesSet() {
// If index already exists don't do anything
if (jedis.ftList().contains(config.indexName)) {
return;
}
String response = jedis.ftCreate(config.indexName,
FTCreateParams.createParams().on(IndexDataType.JSON).addPrefix(config.prefix), schemaFields());
if (!RESPONSE_OK.test(response)) {
String message = MessageFormat.format("Could not create index: {0}", response);
throw new RuntimeException(message);
}
filterExpressionConverter = new RedisFilterExpressionConverter(config.metadataFields);
}
private Iterable<SchemaField> schemaFields() {
Map<String, Object> vectorAttrs = new HashMap<>();
vectorAttrs.put("DIM", embeddingClient.dimensions());
vectorAttrs.put("DISTANCE_METRIC", DEFAULT_DISTANCE_METRIC);
vectorAttrs.put("TYPE", VECTOR_TYPE_FLOAT32);
List<SchemaField> fields = new ArrayList<>();
fields.add(TextField.of(jsonPath(config.contentFieldName)).as(config.contentFieldName).weight(1.0));
fields.add(VectorField.builder()
.fieldName(jsonPath(config.embeddingFieldName))
.algorithm(vectorAlgorithm())
.attributes(vectorAttrs)
.as(config.embeddingFieldName)
.build());
if (!CollectionUtils.isEmpty(config.metadataFields)) {
for (MetadataField field : config.metadataFields) {
fields.add(schemaField(field));
}
}
return fields;
}
private SchemaField schemaField(MetadataField field) {
String fieldName = jsonPath(field.name);
switch (field.fieldType) {
case NUMERIC:
return NumericField.of(fieldName).as(field.name);
case TAG:
return TagField.of(fieldName).as(field.name);
case TEXT:
return TextField.of(fieldName).as(field.name);
default:
throw new IllegalArgumentException(
MessageFormat.format("Field {0} has unsupported type {1}", field.name, field.fieldType));
}
}
private VectorAlgorithm vectorAlgorithm() {
if (config.vectorAlgorithm == Algorithm.HSNW) {
return VectorAlgorithm.HNSW;
}
return VectorAlgorithm.FLAT;
}
private String jsonPath(String field) {
return JSON_PATH_PREFIX + field;
}
private static float[] toFloatArray(List<Double> embeddingDouble) {
float[] embeddingFloat = new float[embeddingDouble.size()];
int i = 0;
for (Double d : embeddingDouble) {
embeddingFloat[i++] = d.floatValue();
}
return embeddingFloat;
}
}

View File

@@ -0,0 +1,129 @@
/*
* Copyright 2023-2023 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 static org.assertj.core.api.Assertions.assertThat;
import static org.springframework.ai.vectorstore.RedisVectorStore.MetadataField.tag;
import static org.springframework.ai.vectorstore.RedisVectorStore.MetadataField.numeric;
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.IN;
import static org.springframework.ai.vectorstore.filter.Filter.ExpressionType.LTE;
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.Arrays;
import java.util.List;
import org.junit.jupiter.api.Test;
import org.springframework.ai.vectorstore.RedisVectorStore.MetadataField;
import org.springframework.ai.vectorstore.filter.Filter.Expression;
import org.springframework.ai.vectorstore.filter.Filter.Group;
import org.springframework.ai.vectorstore.filter.Filter.Key;
import org.springframework.ai.vectorstore.filter.Filter.Value;
/**
* @author Julien Ruaux
*/
class RedisFilterExpressionConverterTests {
private static RedisFilterExpressionConverter converter(MetadataField... fields) {
return new RedisFilterExpressionConverter(Arrays.asList(fields));
}
@Test
void testEQ() {
// country == "BG"
String vectorExpr = converter(tag("country"))
.convertExpression(new Expression(EQ, new Key("country"), new Value("BG")));
assertThat(vectorExpr).isEqualTo("@country:{BG}");
}
@Test
void tesEqAndGte() {
// genre == "drama" AND year >= 2020
String vectorExpr = converter(tag("genre"), numeric("year"))
.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("@genre:{drama} @year:[2020 inf]");
}
@Test
void tesIn() {
// genre in ["comedy", "documentary", "drama"]
String vectorExpr = converter(tag("genre")).convertExpression(
new Expression(IN, new Key("genre"), new Value(List.of("comedy", "documentary", "drama"))));
assertThat(vectorExpr).isEqualTo("@genre:{comedy | documentary | drama}");
}
@Test
void testNe() {
// year >= 2020 OR country == "BG" AND city != "Sofia"
String vectorExpr = converter(numeric("year"), tag("country"), tag("city"))
.convertExpression(new Expression(OR, new Expression(GTE, new Key("year"), new Value(2020)),
new Group(new Expression(AND, new Expression(EQ, new Key("country"), new Value("BG")),
new Expression(NE, new Key("city"), new Value("Sofia"))))));
assertThat(vectorExpr).isEqualTo("@year:[2020 inf] | (@country:{BG} -@city:{Sofia})");
}
@Test
void testGroup() {
// (year >= 2020 OR country == "BG") AND city NIN ["Sofia", "Plovdiv"]
String vectorExpr = converter(numeric("year"), tag("country"), tag("city"))
.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")))));
assertThat(vectorExpr).isEqualTo("(@year:[2020 inf] | @country:{BG}) -@city:{Sofia | Plovdiv}");
}
@Test
void tesBoolean() {
// isOpen == true AND year >= 2020 AND country IN ["BG", "NL", "US"]
String vectorExpr = converter(numeric("year"), tag("country"), tag("isOpen"))
.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")))));
assertThat(vectorExpr).isEqualTo("@isOpen:{true} @year:[2020 inf] @country:{BG | NL | US}");
}
@Test
void testDecimal() {
// temperature >= -15.6 && temperature <= +20.13
String vectorExpr = converter(numeric("temperature"))
.convertExpression(new Expression(AND, new Expression(GTE, new Key("temperature"), new Value(-15.6)),
new Expression(LTE, new Key("temperature"), new Value(20.13))));
assertThat(vectorExpr).isEqualTo("@temperature:[-15.6 inf] @temperature:[-inf 20.13]");
}
@Test
void testComplexIdentifiers() {
String vectorExpr = converter(tag("country 1 2 3"))
.convertExpression(new Expression(EQ, new Key("\"country 1 2 3\""), new Value("BG")));
assertThat(vectorExpr).isEqualTo("@\"country 1 2 3\":{BG}");
vectorExpr = converter(tag("country 1 2 3"))
.convertExpression(new Expression(EQ, new Key("'country 1 2 3'"), new Value("BG")));
assertThat(vectorExpr).isEqualTo("@'country 1 2 3':{BG}");
}
}

View File

@@ -0,0 +1,248 @@
package org.springframework.ai.vectorstore;
import static org.assertj.core.api.Assertions.assertThat;
import java.io.IOException;
import java.nio.charset.StandardCharsets;
import java.util.Collections;
import java.util.List;
import java.util.Map;
import java.util.UUID;
import org.junit.jupiter.api.BeforeEach;
import org.junit.jupiter.api.Test;
import org.springframework.ai.document.Document;
import org.springframework.ai.embedding.EmbeddingClient;
import org.springframework.ai.embedding.TransformersEmbeddingClient;
import org.springframework.ai.vectorstore.RedisVectorStore.MetadataField;
import org.springframework.ai.vectorstore.RedisVectorStore.RedisVectorStoreConfig;
import org.springframework.boot.SpringBootConfiguration;
import org.springframework.boot.autoconfigure.EnableAutoConfiguration;
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 com.redis.testcontainers.RedisStackContainer;
/**
* @author Julien Ruaux
*/
@Testcontainers
class RedisVectorStoreIT {
@Container
static RedisStackContainer redisContainer = new RedisStackContainer(
RedisStackContainer.DEFAULT_IMAGE_NAME.withTag(RedisStackContainer.DEFAULT_TAG));
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")));
public 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);
}
}
@BeforeEach
void cleanDatabase() {
this.contextRunner.run(context -> context.getBean(RedisVectorStore.class).getJedis().flushAll());
}
@Test
void ensureIndexGetsCreated() {
this.contextRunner.run(context -> {
assertThat(context.getBean(RedisVectorStore.class)
.getJedis()
.ftList()
.contains(RedisVectorStore.DEFAULT_INDEX_NAME));
});
}
@Test
void addAndSearch() {
contextRunner.run(context -> {
VectorStore vectorStore = context.getBean(VectorStore.class);
vectorStore.add(documents);
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", RedisVectorStore.DISTANCE_FIELD_NAME);
// Remove all documents from the store
vectorStore.delete(documents.stream().map(doc -> doc.getId()).toList());
results = vectorStore.similaritySearch(SearchRequest.query("Spring").withTopK(1));
assertThat(results).isEmpty();
});
}
@Test
void searchWithFilters() throws InterruptedException {
contextRunner.run(context -> {
VectorStore vectorStore = context.getBean(VectorStore.class);
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));
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 == '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());
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());
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());
});
}
@Test
void documentUpdate() {
contextRunner.run(context -> {
VectorStore vectorStore = context.getBean(VectorStore.class);
Document document = new Document(UUID.randomUUID().toString(), "Spring AI rocks!!",
Collections.singletonMap("meta1", "meta1"));
vectorStore.add(List.of(document));
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(RedisVectorStore.DISTANCE_FIELD_NAME);
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));
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(RedisVectorStore.DISTANCE_FIELD_NAME);
vectorStore.delete(List.of(document.getId()));
});
}
@Test
void searchWithThreshold() {
contextRunner.run(context -> {
VectorStore vectorStore = context.getBean(VectorStore.class);
vectorStore.add(documents);
List<Document> fullResult = vectorStore
.similaritySearch(SearchRequest.query("Spring").withTopK(5).withSimilarityThresholdAll());
List<Float> distances = fullResult.stream()
.map(doc -> (Float) doc.getMetadata().get(RedisVectorStore.DISTANCE_FIELD_NAME))
.toList();
assertThat(distances).hasSize(3);
float threshold = (distances.get(0) + distances.get(1)) / 2;
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", RedisVectorStore.DISTANCE_FIELD_NAME);
});
}
@SpringBootConfiguration
@EnableAutoConfiguration(exclude = { DataSourceAutoConfiguration.class })
public static class TestApplication {
@Bean
public RedisVectorStore vectorStore(EmbeddingClient embeddingClient) {
return new RedisVectorStore(RedisVectorStoreConfig.builder()
.withURI(redisContainer.getRedisURI())
.withMetadataFields(MetadataField.tag("meta1"), MetadataField.tag("meta2"),
MetadataField.tag("country"), MetadataField.numeric("year"))
.build(), embeddingClient);
}
@Bean
public EmbeddingClient embeddingClient() {
return new TransformersEmbeddingClient();
}
}
}