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:
2
pom.xml
2
pom.xml
@@ -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>
|
||||
|
||||
|
||||
@@ -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[]
|
||||
|
||||
@@ -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
|
||||
|
||||
|
||||
|
||||
@@ -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.
|
||||
|
||||
|
||||
@@ -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]
|
||||
----
|
||||
@@ -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>
|
||||
|
||||
@@ -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);
|
||||
}
|
||||
|
||||
}
|
||||
@@ -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;
|
||||
}
|
||||
|
||||
}
|
||||
@@ -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
|
||||
|
||||
@@ -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();
|
||||
}
|
||||
|
||||
}
|
||||
|
||||
}
|
||||
@@ -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:");
|
||||
}
|
||||
|
||||
}
|
||||
@@ -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>
|
||||
80
vector-stores/spring-ai-redis/pom.xml
Normal file
80
vector-stores/spring-ai-redis/pom.xml
Normal 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>
|
||||
@@ -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;
|
||||
|
||||
}
|
||||
|
||||
}
|
||||
|
||||
}
|
||||
@@ -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;
|
||||
}
|
||||
|
||||
}
|
||||
@@ -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}");
|
||||
}
|
||||
|
||||
}
|
||||
@@ -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();
|
||||
}
|
||||
|
||||
}
|
||||
|
||||
}
|
||||
Reference in New Issue
Block a user