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

Simplify VectorStore interface

Browse Source 
- Collapses all VectorStore similiaritySearch methdos into one with SearchRequest builder.
 - Fix all affected code and tests.
 - Bump the project version to 0.7.1.
 - Add tests
 - Add autoconfigurations for milvus, pinecone and pgvecor stores.
 - Improve and unify the VectorStore ITs.
 - Make use of TrasformersEmbeddingClient for auto-configurations ITs.
This commit is contained in:
Christian Tzolov
2023-11-10 21:33:03 +01:00
committed by Mark Pollack
parent e0a9d4a8fe
commit 952ff6c319
62 changed files with 2404 additions and 563 deletions
Download Patch File Download Diff File Expand all files Collapse all files

28
README.md
View File

@@ -11,15 +11,15 @@ Let's make your `@Beans` intelligent!
* [Documentation](https://docs.spring.io/spring-ai/reference/)
* [Issues](https://github.com/spring-projects-experimental/spring-ai/issues)
* [Discussions](https://github.com/spring-projects-experimental/spring-ai/discussions) - Go here if you have a question, suggestion, or feedback!
* [Discussions](https://github.com/spring-projects-experimental/spring-ai/discussions) - Go here if you have a question, suggestion, or feedback!
* [JavaDocs](https://docs.spring.io/spring-ai/docs/current-SNAPSHOT/)
## Educational Resources
* Follow the [Workshop](#workshop)
* Overview of Spring AI @ Devoxx 2023
* Follow the [Workshop](#workshop)
* Overview of Spring AI @ Devoxx 2023
<br>[![Watch the Devoxx 2023 video](https://img.youtube.com/vi/7OY9fKVxAFQ/default.jpg)](https://www.youtube.com/watch?v=7OY9fKVxAFQ)
* Introducing Spring AI - Add Generative AI to your Spring Applications
* Introducing Spring AI - Add Generative AI to your Spring Applications
<br>[![Watch the video](https://img.youtube.com/vi/1g_wuincUdU/default.jpg)](https://www.youtube.com/watch?v=1g_wuincUdU)
## Dependencies
@@ -49,7 +49,7 @@ And the Spring Boot Starter depending on if you are using Azure Open AI or Open
<dependency>
<groupId>org.springframework.experimental.ai</groupId>
<artifactId>spring-ai-azure-openai-spring-boot-starter</artifactId>
<version>0.7.0-SNAPSHOT</version>
<version>0.7.1-SNAPSHOT</version>
</dependency>
```
@@ -59,7 +59,7 @@ And the Spring Boot Starter depending on if you are using Azure Open AI or Open
<dependency>
<groupId>org.springframework.experimental.ai</groupId>
<artifactId>spring-ai-openai-spring-boot-starter</artifactId>
<version>0.7.0-SNAPSHOT</version>
<version>0.7.1-SNAPSHOT</version>
</dependency>
```
@@ -71,7 +71,7 @@ And the Spring Boot Starter depending on if you are using Azure Open AI or Open
<dependency>
<groupId>org.springframework.experimental.ai</groupId>
<artifactId>spring-ai-openai-spring-boot-starter</artifactId>
<version>0.7.0-SNAPSHOT</version>
<version>0.7.1-SNAPSHOT</version>
</dependency>
```
@@ -96,11 +96,11 @@ These Python libraries share foundational themes with Spring projects, such as:
Taking inspiration from these libraries, the Spring AI project aims to provide a similar experience for Spring developers in the AI domain.
Note, that the Spring AI API is not a direct port of either LangChain or LlamaIndex. You will see significant differences in the API if you are familiar with those two projects, though concepts and ideas are fairly portable.
Note, that the Spring AI API is not a direct port of either LangChain or LlamaIndex. You will see significant differences in the API if you are familiar with those two projects, though concepts and ideas are fairly portable.
## Feature Overview
This is a high level feature overview.
This is a high level feature overview.
The features that are implemented lay the foundation, with subsequent more complex features building upon them.
You can find more details in the [Reference Documentation](https://docs.spring.io/spring-ai/reference/)
@@ -134,7 +134,7 @@ For implementation details, visit the [OutputParser API guide](https://docs.spri
Incorporating proprietary data into Generative AI without retraining the model has been a breakthrough.
Retraining models, especially those with billions of parameters, is challenging due to the specialized hardware required.
The 'In-context' learning technique provides a simpler method to infuse your pre-trained model with data, whether from text files, HTML, or database results.
The 'In-context' learning technique provides a simpler method to infuse your pre-trained model with data, whether from text files, HTML, or database results.
The right techniques are critical for developing successful solutions.
@@ -150,9 +150,9 @@ The subsequent classes and interfaces support RAG's data preparation.
**Documents:**
The `Document` class encapsulates your data, including text and metadata, for the AI model.
The `Document` class encapsulates your data, including text and metadata, for the AI model.
While a Document can represent extensive content, such as an entire file, the RAG approach
segments content into smaller pieces for inclusion in the prompt.
segments content into smaller pieces for inclusion in the prompt.
The ETL process uses the interfaces `DocumentReader`, `DocumentTransformer`, and `DocumentWriter`, ending with data storage in a Vector Database.
This database later discerns the pieces of data that are pertinent to a user's query.
@@ -160,7 +160,7 @@ This database later discerns the pieces of data that are pertinent to a user's q
**Document Readers:**
Document Readers produce a `List<Document>` from diverse sources like PDFs, Markdown files, and Word documents.
Given that many sources are unstructured, Document Readers often segment based on content semantics, avoiding splits within tables or code sections.
Given that many sources are unstructured, Document Readers often segment based on content semantics, avoiding splits within tables or code sections.
After the initial creation of the `List<Document>`, the data flows through transformers for further refinement.
**Document Transformers:**
@@ -171,7 +171,7 @@ Each model has a context-window indicating its input and output data limits. Typ
**Document Writers:**
The final ETL step within RAG involves committing the data segments to a Vector Database.
The final ETL step within RAG involves committing the data segments to a Vector Database.
Though the `DocumentWriter` interface isn't exclusively for Vector Database writing, it the main type of implementation.
**Vector Stores:** Vector Databases are instrumental in incorporating your data with AI models.

4
document-readers/pdf-reader/pom.xml
View File

@@ -5,7 +5,7 @@
<parent>
<groupId>org.springframework.experimental.ai</groupId>
<artifactId>spring-ai</artifactId>
<version>0.7.0-SNAPSHOT</version>
<version>0.7.1-SNAPSHOT</version>
<relativePath>../../pom.xml</relativePath>
</parent>
<artifactId>spring-ai-pdf-document-reader</artifactId>
@@ -36,7 +36,7 @@
<groupId>commons-logging</groupId>
<artifactId>commons-logging</artifactId>
</exclusion>
</exclusions>
</exclusions>
</dependency>
<!-- TESTING -->

2
document-readers/tika-reader/pom.xml
View File

@@ -5,7 +5,7 @@
<parent>
<groupId>org.springframework.experimental.ai</groupId>
<artifactId>spring-ai</artifactId>
<version>0.7.0-SNAPSHOT</version>
<version>0.7.1-SNAPSHOT</version>
<relativePath>../../pom.xml</relativePath>
</parent>
<artifactId>spring-ai-tika-document-reader</artifactId>

2
embedding-clients/spring-ai-postgresml-embedding-client/pom.xml
View File

@@ -5,7 +5,7 @@
<parent>
<groupId>org.springframework.experimental.ai</groupId>
<artifactId>spring-ai</artifactId>
<version>0.7.0-SNAPSHOT</version>
<version>0.7.1-SNAPSHOT</version>
<relativePath>../../pom.xml</relativePath>
</parent>
<artifactId>spring-ai-postgresml-embedding-client</artifactId>

2
embedding-clients/transformers-embedding/README.md
View File

@@ -36,7 +36,7 @@ Add the `transformers-embedding` project to your maven dependencies:
<dependency>
<groupId>org.springframework.experimental.ai</groupId>
<artifactId>transformers-embedding</artifactId>
<version>0.7.0-SNAPSHOT</version>
<version>0.7.1-SNAPSHOT</version>
</dependency>
```

2
embedding-clients/transformers-embedding/pom.xml
View File

@@ -5,7 +5,7 @@
<parent>
<groupId>org.springframework.experimental.ai</groupId>
<artifactId>spring-ai</artifactId>
<version>0.7.0-SNAPSHOT</version>
<version>0.7.1-SNAPSHOT</version>
<relativePath>../../pom.xml</relativePath>
</parent>
<artifactId>transformers-embedding</artifactId>

4
pom.xml
View File

@@ -4,7 +4,7 @@
<groupId>org.springframework.experimental.ai</groupId>
<artifactId>spring-ai</artifactId>
<version>0.7.0-SNAPSHOT</version>
<version>0.7.1-SNAPSHOT</version>
<packaging>pom</packaging>
<url>https://github.com/spring-projects-experimental/spring-ai</url>
@@ -87,7 +87,7 @@
<pdfbox.version>3.0.0</pdfbox.version>
<pgvector.version>0.1.3</pgvector.version>
<postgresql.version>42.6.0</postgresql.version>
<milvus.version>2.3.0</milvus.version>
<milvus.version>2.3.3</milvus.version>
<pinecone.version>0.6.0</pinecone.version>
<protobuf-java-util.version>3.24.4</protobuf-java-util.version>

2
spring-ai-azure-openai/pom.xml
View File

@@ -4,7 +4,7 @@
<parent>
<groupId>org.springframework.experimental.ai</groupId>
<artifactId>spring-ai</artifactId>
<version>0.7.0-SNAPSHOT</version>
<version>0.7.1-SNAPSHOT</version>
</parent>
<artifactId>spring-ai-azure-openai</artifactId>
<packaging>jar</packaging>

2
spring-ai-core/pom.xml
View File

@@ -5,7 +5,7 @@
<parent>
<groupId>org.springframework.experimental.ai</groupId>
<artifactId>spring-ai</artifactId>
<version>0.7.0-SNAPSHOT</version>
<version>0.7.1-SNAPSHOT</version>
</parent>
<artifactId>spring-ai-core</artifactId>
<packaging>jar</packaging>

39
spring-ai-core/src/main/java/org/springframework/ai/ResourceUtils.java Normal file
View File

@@ -0,0 +1,39 @@
/*
* 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;
import java.io.IOException;
import java.nio.charset.StandardCharsets;
import org.springframework.core.io.DefaultResourceLoader;
/**
* @author Christian Tzolov
*/
public class ResourceUtils {
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);
}
}
}

18
spring-ai-core/src/main/java/org/springframework/ai/retriever/VectorStoreRetriever.java
View File

@@ -1,13 +1,14 @@
package org.springframework.ai.retriever;
import org.springframework.ai.document.Document;
import org.springframework.ai.document.DocumentRetriever;
import org.springframework.ai.vectorstore.VectorStore;
import java.util.List;
import java.util.Objects;
import java.util.Optional;
import org.springframework.ai.document.Document;
import org.springframework.ai.document.DocumentRetriever;
import org.springframework.ai.vectorstore.SearchRequest;
import org.springframework.ai.vectorstore.VectorStore;
public class VectorStoreRetriever implements DocumentRetriever {
private VectorStore vectorStore;
@@ -47,12 +48,13 @@ public class VectorStoreRetriever implements DocumentRetriever {
@Override
public List<Document> retrieve(String query) {
SearchRequest request = SearchRequest.query(query).withTopK(this.k);
if (threshold.isPresent()) {
return this.vectorStore.similaritySearch(query, this.k, this.threshold.get());
}
else {
return this.vectorStore.similaritySearch(query, this.k);
request.withSimilarityThreshold(this.threshold.get());
}
return this.vectorStore.similaritySearch(request);
}
}

32
spring-ai-core/src/main/java/org/springframework/ai/vectorstore/InMemoryVectorStore.java
View File

@@ -12,6 +12,7 @@ import java.util.concurrent.ConcurrentHashMap;
* @author Raphael Yu
* @author Dingmeng Xue
* @author Mark Pollack
* @author Christian Tzolov
*/
public class InMemoryVectorStore implements VectorStore {
@@ -45,36 +46,23 @@ public class InMemoryVectorStore implements VectorStore {
}
@Override
public List<Document> similaritySearch(String query) {
return similaritySearch(query, 4);
}
public List<Document> similaritySearch(SearchRequest request) {
if (request.getFilterExpression() != null) {
throw new UnsupportedOperationException(
"The [" + this.getClass() + "] doesn't support metadata filtering!");
}
@Override
public List<Document> similaritySearch(String query, int k) {
List<Double> userQueryEmbedding = getUserQueryEmbedding(query);
List<Double> userQueryEmbedding = getUserQueryEmbedding(request.getQuery());
var similarities = this.store.values()
.stream()
.map(entry -> new Similarity(entry.getId(),
EmbeddingMath.cosineSimilarity(userQueryEmbedding, entry.getEmbedding())))
.filter(s -> s.similarity >= request.getSimilarityThreshold())
.sorted(Comparator.<Similarity>comparingDouble(s -> s.similarity).reversed())
.limit(k)
.map(s -> store.get(s.key))
.limit(request.getTopK())
.map(s -> this.store.get(s.key))
.toList();
return similarities;
}
@Override
public List<Document> similaritySearch(String query, int k, double threshold) {
List<Double> userQueryEmbedding = getUserQueryEmbedding(query);
var similarities = this.store.values()
.stream()
.map(entry -> new Similarity(entry.getId(),
EmbeddingMath.cosineSimilarity(userQueryEmbedding, entry.getEmbedding())))
.filter(s -> s.similarity >= threshold)
.sorted(Comparator.<Similarity>comparingDouble(s -> s.similarity).reversed())
.limit(k)
.map(s -> store.get(s.key))
.toList();
return similarities;
}

251
spring-ai-core/src/main/java/org/springframework/ai/vectorstore/SearchRequest.java Normal file
View File

@@ -0,0 +1,251 @@
/*
* 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 org.springframework.ai.document.Document;
import org.springframework.ai.vectorstore.filter.Filter;
import org.springframework.ai.vectorstore.filter.FilterExpressionBuilder;
import org.springframework.ai.vectorstore.filter.FilterExpressionTextParser;
import org.springframework.util.Assert;
/**
* Similarity search request builder. Use the {@link #query(String)}, {@link #defaults()}
* or {@link #from(SearchRequest)} factory methods to create a new {@link SearchRequest}
* instance and then apply the 'with' methods to alter the default values.
*
* @author Christian Tzolov
*/
public class SearchRequest {
/**
* Similarity threshold that accepts all search scores. A threshold value of 0.0 means
* any similarity is accepted or disable the similarity threshold filtering. A
* threshold value of 1.0 means an exact match is required.
*/
public static final double SIMILARITY_THRESHOLD_ACCEPT_ALL = 0.0;
/**
* Default value for the top 'k' similar results to return.
*/
public static final int DEFAULT_TOP_K = 4;
public String query;
private int topK = DEFAULT_TOP_K;
private double similarityThreshold = SIMILARITY_THRESHOLD_ACCEPT_ALL;
private Filter.Expression filterExpression;
private SearchRequest(String query) {
this.query = query;
}
/**
* Create a new {@link SearchRequest} builder instance with specified embedding query
* string.
* @param query Text to use for embedding similarity comparison.
* @return Returns new {@link SearchRequest} builder instance.
*/
public static SearchRequest query(String query) {
Assert.notNull(query, "Query can not be null.");
return new SearchRequest(query);
}
/**
* Create a new {@link SearchRequest} builder instance with an empty embedding query
* string. Use the {@link #withQuery(String query)} to set/update the embedding query
* text.
* @return Returns new {@link SearchRequest} builder instance.
*/
public static SearchRequest defaults() {
return new SearchRequest("");
}
/**
* Copy an existing {@link SearchRequest} instance.
* @param originalSearchRequest {@link SearchRequest} instance to copy.
* @return Returns new {@link SearchRequest} builder instance.
*/
public static SearchRequest from(SearchRequest originalSearchRequest) {
return new SearchRequest(originalSearchRequest.getQuery()).withTopK(originalSearchRequest.getTopK())
.withSimilarityThreshold(originalSearchRequest.getSimilarityThreshold())
.withFilterExpression(originalSearchRequest.getFilterExpression());
}
/**
* @param query Text to use for embedding similarity comparison.
* @return this builder.
*/
public SearchRequest withQuery(String query) {
Assert.notNull(query, "Query can not be null.");
this.query = query;
return this;
}
/**
* @param topK the top 'k' similar results to return.
* @return this builder.
*/
public SearchRequest withTopK(int topK) {
Assert.isTrue(topK >= 0, "TopK should be positive.");
this.topK = topK;
return this;
}
/**
* Similarity threshold score to filter the search response by. Only documents with
* similarity score equal or greater than the 'threshold' will be returned. Note that
* this is a post processing step performed on the client not the server side. A
* threshold value of 0.0 means any similarity is accepted or disable the similarity
* threshold filtering. A threshold value of 1.0 means an exact match is required.
* @param threshold The lower bound of the similarity score.
* @return this builder.
*/
public SearchRequest withSimilarityThreshold(double threshold) {
Assert.isTrue(threshold >= 0 && threshold <= 1, "Similarity threshold must be in [0,1] range.");
this.similarityThreshold = threshold;
return this;
}
/**
* Sets disables the similarity threshold by setting it to 0.0 - all results are
* accepted.
* @return this builder.
*/
public SearchRequest withSimilarityThresholdAll() {
return withSimilarityThreshold(SIMILARITY_THRESHOLD_ACCEPT_ALL);
}
/**
* Retrieves documents by query embedding similarity and matching the filters. Value
* of 'null' means that no metadata filters will be applied to the search.
*
* For example if the {@link Document#getMetadata()} schema is:
*
* <pre>{@code
* &#123;
* "country": <Text>,
* "city": <Text>,
* "year": <Number>,
* "price": <Decimal>,
* "isActive": <Boolean>
* &#125;
* }</pre>
*
* you can constrain the search result to only UK countries with isActive=true and
* year equal or greater 2020. You can build this such metadata filter
* programmatically like this:
*
* <pre>{@code
* var exp = new Filter.Expression(AND,
* new Expression(EQ, new Key("country"), new Value("UK")),
* new Expression(AND,
* new Expression(GTE, new Key("year"), new Value(2020)),
* new Expression(EQ, new Key("isActive"), new Value(true))));
* }</pre>
*
* The {@link Filter.Expression} is portable across all vector stores.<br/>
*
*
* The {@link FilterExpressionBuilder} is a DSL creating expressions programmatically:
*
* <pre>{@code
* var b = new FilterExpressionBuilder();
* var exp = b.and(
* b.eq("country", "UK"),
* b.and(
* b.gte("year", 2020),
* b.eq("isActive", true)));
* }</pre>
*
* The {@link FilterExpressionTextParser} converts textual, SQL like filter expression
* language into {@link Filter.Expression}:
*
* <pre>{@code
* var parser = new FilterExpressionTextParser();
* var exp = parser.parse("country == 'UK' && isActive == true && year >=2020");
* }</pre>
* @param expression {@link Filter.Expression} instance used to define the metadata
* filter criteria. The 'null' value stands for no expression filters.
* @return this builder.
*/
public SearchRequest withFilterExpression(Filter.Expression expression) {
this.filterExpression = expression;
return this;
}
/**
* Document metadata filter expression. For example if your
* {@link Document#getMetadata()} has a schema like:
*
* <pre>{@code
* &#123;
* "country": <Text>,
* "city": <Text>,
* "year": <Number>,
* "price": <Decimal>,
* "isActive": <Boolean>
* &#125;
* }</pre>
*
* then you can constrain the search result with metadata filter expressions like:
*
* <pre>{@code
* country == 'UK' && year >= 2020 && isActive == true
* Or
* country == 'BG' && (city NOT IN ['Sofia', 'Plovdiv'] || price < 134.34)
* }</pre>
*
* This ensures that the response contains only embeddings that match the specified
* filer criteria. <br/>
*
* The declarative, SQL like, filter syntax is portable across all vector stores
* supporting the filter search feature.<br/>
*
* The {@link FilterExpressionTextParser} is used to convert the text filter
* expression into {@link Filter.Expression}.
* @param textExpression declarative, portable, SQL like, metadata filter syntax. The
* 'null' value stands for no expression filters.
* @return this.builder
*/
public SearchRequest withFilterExpression(String textExpression) {
this.filterExpression = (textExpression != null) ? Filter.parser().parse(textExpression) : null;
return this;
}
public String getQuery() {
return query;
}
public int getTopK() {
return topK;
}
public double getSimilarityThreshold() {
return similarityThreshold;
}
public Filter.Expression getFilterExpression() {
return filterExpression;
}
public boolean hasFilterExpression() {
return this.filterExpression != null;
}
}

118
spring-ai-core/src/main/java/org/springframework/ai/vectorstore/VectorStore.java
View File

@@ -5,10 +5,15 @@ import java.util.Optional;
import org.springframework.ai.document.Document;
import org.springframework.ai.document.DocumentWriter;
import org.springframework.ai.vectorstore.filter.Filter;
import org.springframework.ai.vectorstore.filter.FilterExpressionBuilder;
import org.springframework.ai.vectorstore.filter.FilterExpressionTextParser;
/**
* The {@code VectorStore} interface defines the operations for managing and querying
* documents in a vector database. It extends {@link DocumentWriter} to support document
* writing operations. Vector databases are specialized for AI applications, performing
* similarity searches based on vector representations of data rather than exact matches.
* This interface allows for adding, deleting, and searching documents based on their
* similarity to a given query.
*/
public interface VectorStore extends DocumentWriter {
/**
@@ -30,107 +35,24 @@ public interface VectorStore extends DocumentWriter {
*/
Optional<Boolean> delete(List<String> idList);
List<Document> similaritySearch(String query);
List<Document> similaritySearch(String query, int k);
/**
* @param query The query to send, it will be converted to an embedding based on the
* configuration of the vector store.
* @param k the top 'k' similar results
* @param threshold the lower bound of the similarity score
* @return similar documents
*/
List<Document> similaritySearch(String query, int k, double threshold);
/**
* Retrieves documents by query embedding similarity and metadata filters to retrieve
* exactly the number of nearest-neighbor results that match the filters.
*
* For example if your {@link Document#getMetadata()} has a schema like:
*
* <pre>{@code
* &#123;
* "country": <Text>,
* "city": <Text>,
* "year": <Number>,
* "price": <Decimal>,
* "isActive": <Boolean>
* &#125;
* }</pre>
*
* then you can constrain the search result with metadata filter expressions
* equivalent to (country == 'UK' AND year >= 2020 AND isActive == true). You can
* build this filter programmatically like this:
*
* <pre>{@code
*
* new Filter.Expression(AND,
* new Expression(EQ, new Key("country"), new Value("UK")),
* new Expression(AND,
* new Expression(GTE, new Key("year"), new Value(2020)),
* new Expression(EQ, new Key("isActive"), new Value(true))));
*
* }</pre>
*
* and it will ensure that the response contains only embeddings that match the
* specified filer criteria. <br/>
*
* The {@link Filter.Expression} is portable across all vector stores that offer
* metadata filtering. The {@link FilterExpressionBuilder} is expression DSL and
* {@link FilterExpressionTextParser} is text expression parser that build
* {@link Filter.Expression}.
* @param topK the top 'k' similar results to return.
* @param similarityThreshold the lower bound of the similarity score
* @param filterExpression portable metadata filter expression.
* @return similar documents that match the requested similarity threshold and filter.
* exactly the number of nearest-neighbor results that match the request criteria.
* @param request Search request for set search parameters, such as the query text,
* topK, similarity threshold and metadata filter expressions.
* @return Returns documents th match the query request conditions.
*/
default List<Document> similaritySearch(String query, int topK, double similarityThreshold,
Filter.Expression filterExpression) {
throw new UnsupportedOperationException("This vector store doesn't support search filtering");
}
List<Document> similaritySearch(SearchRequest request);
/**
* Retrieves documents by query embedding similarity and metadata filters to retrieve
* exactly the number of nearest-neighbor results that match the filters.
*
* For example if your {@link Document#getMetadata()} has a schema like:
*
* <pre>{@code
* &#123;
* "country": <Text>,
* "city": <Text>,
* "year": <Number>,
* "price": <Decimal>,
* "isActive": <Boolean>
* &#125;
* }</pre>
*
* then you can constrain the search result with metadata filter expressions like:
*
* <pre>{@code
*country == 'UK' && year >= 2020 && isActive == true
* Or
*country == 'BG' && (city NOT IN ['Sofia', 'Plovdiv'] || price < 134.34)
* }</pre>
*
* This ensures that the response contains only embeddings that match the specified
* filer criteria. <br/>
*
* The declarative, SQL like, filter syntax is portable across all vector stores
* supporting the filter search feature.<br/>
*
* The {@link FilterExpressionTextParser} is used to convert the text filter
* expression into {@link Filter.Expression}.
* @param topK the top 'k' similar results to return.
* @param similarityThreshold the lower bound of the similarity score
* @param filterExpression portable metadata filter expression.
* @return similar documents that match the requested similarity threshold and filter.
* Retrieves documents by query embedding similarity using the default
* {@link SearchRequest}'s' search criteria.
* @param query Text to use for embedding similarity comparison.
* @return Returns a list of documents that have embeddings similar to the query text
* embedding.
*/
default List<Document> similaritySearch(String query, int topK, double similarityThreshold,
String filterExpression) {
var filterExpressionObject = Filter.parser().parse(filterExpression);
return similaritySearch(query, topK, similarityThreshold, filterExpressionObject);
default List<Document> similaritySearch(String query) {
return this.similaritySearch(SearchRequest.query(query));
}
}

2
spring-ai-core/src/main/java/org/springframework/ai/vectorstore/filter/FilterExpressionTextParser.java
View File

@@ -40,7 +40,7 @@ import org.springframework.util.Assert;
/**
*
* Parse a text, vector-store agnostic, filter expression language into
* Parse a textual, vector-store agnostic, filter expression language into
* {@link Filter.Expression}.
*
* The vector-store agnostic, filter expression language is defined by a formal ANTLR4

2
spring-ai-core/src/main/java/org/springframework/ai/vectorstore/filter/converter/AbstractFilterExpressionConverter.java
View File

@@ -22,6 +22,7 @@ import org.springframework.ai.vectorstore.filter.Filter;
import org.springframework.ai.vectorstore.filter.Filter.ExpressionType;
import org.springframework.ai.vectorstore.filter.Filter.Group;
import org.springframework.ai.vectorstore.filter.Filter.Operand;
import org.springframework.util.Assert;
/**
* @author Christian Tzolov
@@ -29,6 +30,7 @@ import org.springframework.ai.vectorstore.filter.Filter.Operand;
public abstract class AbstractFilterExpressionConverter {
public String convert(Operand operand) {
Assert.notNull(operand, "Operand can't be null");
var context = new StringBuilder();
this.convert(operand, context);
return context.toString();

1
spring-ai-core/src/main/resources/embedding/embedding-model-dimensions.properties
View File

@@ -17,3 +17,4 @@ code-search-ada-code-001=1024
code-search-ada-text-001=1024
code-search-babbage-code-001=2048
code-search-babbage-text-001=2048
sentence-transformers/all-MiniLM-L6-v2=384

145
spring-ai-core/src/test/java/org/springframework/ai/vectorstore/filter/SearchRequestTests.java Normal file
View File

@@ -0,0 +1,145 @@
/*
* 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.filter;
import org.junit.jupiter.api.Test;
import org.springframework.ai.vectorstore.SearchRequest;
import org.springframework.ai.vectorstore.filter.FilterExpressionTextParser.FilterExpressionParseException;
import static org.assertj.core.api.Assertions.assertThat;
import static org.assertj.core.api.Assertions.assertThatThrownBy;
/**
* @author Christian Tzolov
*/
public class SearchRequestTests {
@Test
public void createDefaults() {
var emptyRequest = SearchRequest.defaults();
assertThat(emptyRequest.getQuery()).isEqualTo("");
checkDefaults(emptyRequest);
}
@Test
public void createQuery() {
var emptyRequest = SearchRequest.query("New Query");
assertThat(emptyRequest.getQuery()).isEqualTo("New Query");
checkDefaults(emptyRequest);
}
@Test
public void createFrom() {
var originalRequest = SearchRequest.query("New Query")
.withTopK(696)
.withSimilarityThreshold(0.678)
.withFilterExpression("country == 'NL'");
var newRequest = SearchRequest.from(originalRequest);
assertThat(newRequest).isNotSameAs(originalRequest);
assertThat(newRequest.getQuery()).isEqualTo(originalRequest.getQuery());
assertThat(newRequest.getTopK()).isEqualTo(originalRequest.getTopK());
assertThat(newRequest.getFilterExpression()).isEqualTo(originalRequest.getFilterExpression());
assertThat(newRequest.getSimilarityThreshold()).isEqualTo(originalRequest.getSimilarityThreshold());
}
@Test
public void withQuery() {
var emptyRequest = SearchRequest.defaults();
assertThat(emptyRequest.getQuery()).isEqualTo("");
emptyRequest.withQuery("New Query");
assertThat(emptyRequest.getQuery()).isEqualTo("New Query");
}
@Test()
public void withSimilarityThreshold() {
var request = SearchRequest.query("Test").withSimilarityThreshold(0.678);
assertThat(request.getSimilarityThreshold()).isEqualTo(0.678);
request.withSimilarityThreshold(0.9);
assertThat(request.getSimilarityThreshold()).isEqualTo(0.9);
assertThatThrownBy(() -> {
request.withSimilarityThreshold(-1);
}).isInstanceOf(IllegalArgumentException.class)
.hasMessageContaining("Similarity threshold must be in [0,1] range.");
assertThatThrownBy(() -> {
request.withSimilarityThreshold(1.1);
}).isInstanceOf(IllegalArgumentException.class)
.hasMessageContaining("Similarity threshold must be in [0,1] range.");
}
@Test()
public void withTopK() {
var request = SearchRequest.query("Test").withTopK(66);
assertThat(request.getTopK()).isEqualTo(66);
request.withTopK(89);
assertThat(request.getTopK()).isEqualTo(89);
assertThatThrownBy(() -> {
request.withTopK(-1);
}).isInstanceOf(IllegalArgumentException.class).hasMessageContaining("TopK should be positive.");
}
@Test()
public void withFilterExpression() {
var request = SearchRequest.query("Test").withFilterExpression("country == 'BG' && year >= 2022");
assertThat(request.getFilterExpression()).isEqualTo(new Filter.Expression(Filter.ExpressionType.AND,
new Filter.Expression(Filter.ExpressionType.EQ, new Filter.Key("country"), new Filter.Value("BG")),
new Filter.Expression(Filter.ExpressionType.GTE, new Filter.Key("year"), new Filter.Value(2022))));
assertThat(request.hasFilterExpression()).isTrue();
request.withFilterExpression("active == true");
assertThat(request.getFilterExpression()).isEqualTo(
new Filter.Expression(Filter.ExpressionType.EQ, new Filter.Key("active"), new Filter.Value(true)));
assertThat(request.hasFilterExpression()).isTrue();
request.withFilterExpression(Filter.builder().eq("country", "NL").build());
assertThat(request.getFilterExpression()).isEqualTo(
new Filter.Expression(Filter.ExpressionType.EQ, new Filter.Key("country"), new Filter.Value("NL")));
assertThat(request.hasFilterExpression()).isTrue();
request.withFilterExpression((String) null);
assertThat(request.getFilterExpression()).isNull();
assertThat(request.hasFilterExpression()).isFalse();
request.withFilterExpression((Filter.Expression) null);
assertThat(request.getFilterExpression()).isNull();
assertThat(request.hasFilterExpression()).isFalse();
assertThatThrownBy(() -> {
request.withFilterExpression("FooBar");
}).isInstanceOf(FilterExpressionParseException.class)
.hasMessageContaining("Error: no viable alternative at input 'FooBar'");
}
private void checkDefaults(SearchRequest request) {
assertThat(request.getFilterExpression()).isNull();
assertThat(request.getSimilarityThreshold()).isEqualTo(SearchRequest.SIMILARITY_THRESHOLD_ACCEPT_ALL);
assertThat(request.getTopK()).isEqualTo(SearchRequest.DEFAULT_TOP_K);
}
}

2
spring-ai-docs/pom.xml
View File

@@ -4,7 +4,7 @@
<parent>
<groupId>org.springframework.experimental.ai</groupId>
<artifactId>spring-ai</artifactId>
<version>0.7.0-SNAPSHOT</version>
<version>0.7.1-SNAPSHOT</version>
</parent>
<artifactId>spring-ai-docs</artifactId>
<name>Spring AI Docs</name>

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

@@ -146,7 +146,7 @@ Add the Spring Boot starter to you project's dependencies:
<dependency>
<groupId>org.springframework.experimental.ai</groupId>
<artifactId>spring-ai-azure-openai-spring-boot-starter</artifactId>
<version>0.7.0-SNAPSHOT</version>
<version>0.7.1-SNAPSHOT</version>
</dependency>
----
@@ -187,7 +187,7 @@ There is not yet a Spring Boot Starter for this client implementation, so you sh
<dependency>
<groupId>org.springframework.experimental.ai</groupId>
<artifactId>spring-ai-huggingface</artifactId>
<version>0.7.0-SNAPSHOT</version>
<version>0.7.1-SNAPSHOT</version>
</dependency>
----
@@ -208,7 +208,7 @@ There is not yet a Spring Boot Starter for this client implementation, so you sh
<dependency>
<groupId>org.springframework.experimental.ai</groupId>
<artifactId>spring-ai-ollama</artifactId>
<version>0.7.0-SNAPSHOT</version>
<version>0.7.1-SNAPSHOT</version>
</dependency>
----

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

@@ -28,24 +28,43 @@ Here is the `VectorStore` interface definition:
```java
public interface VectorStore {
void add(List<Document> documents);
void add(List<Document> documents);
Optional<Boolean> delete(List<String> idList);
Optional<Boolean> delete(List<String> idList);
List<Document> similaritySearch(String query);
List<Document> similaritySearch(String query);
List<Document> similaritySearch(String query, int k);
List<Document> similaritySearch(String query, int k, double threshold);
List<Document> similaritySearch(String query, int topK, double similarityThreshold,
Filter.Expression filterExpression);
List<Document> similaritySearch(String query, int topK, double similarityThreshold,
String filterExpression);
List<Document> similaritySearch(SearchRequest request);
}
```
and the related `SearchRequest` builder:
```java
public class SearchRequest {
public final String query;
private int topK = 4;
private double similarityThreshold = SIMILARITY_THRESHOLD_ALL;
private Filter.Expression filterExpression;
public static SearchRequest query(String query) { return new SearchRequest(query); }
private SearchRequest(String query) { this.query = query; }
public SearchRequest withTopK(int topK) {...}
public SearchRequest withSimilarityThreshold(double threshold) {...}
public SearchRequest withSimilarityThresholdAll() {...}
public SearchRequest withFilterExpression(Filter.Expression expression) {...}
public SearchRequest withFilterExpression(String textExpression) {...}
public String getQuery() {...}
public int getTopK() {...}
public double getSimilarityThreshold() {...}
public Filter.Expression getFilterExpression() {...}
}
```
To insert data into the vector database, encapsulate it within a `Document` object.
The `Document` class encapsulates content from a data source, such as a PDF or Word document, and includes text represented as a string.
It also contains metadata in the form of key-value pairs, including details such as the filename.
@@ -59,7 +78,7 @@ The `similaritySearch` methods in the interface allow for retrieving documents s
* `k`: An integer that specifies the maximum number of similar documents to return. This is often referred to as a 'top K' search, or 'K nearest neighbors' (KNN).
* `threshold`: A double value ranging from 0 to 1, where values closer to 1 indicate higher similarity. By default, if you set a threshold of 0.75, for instance, only documents with a similarity above this value are returned.
* `Filter.Expression`: A class used for passing a fluent DSL (Domain-Specific Language) expression that functions similarly to a 'where' clause in SQL, but it applies exclusively to the metadata key-value pairs of a `Document`.
* `filterExpression`: An external DSL based on ANTLR4 that accepts filter expressions as strings. For example, with metadata keys like country, year, and `isActive`, you could use an expression such as
* `filterExpression`: An external DSL based on ANTLR4 that accepts filter expressions as strings. For example, with metadata keys like country, year, and `isActive`, you could use an expression such as
``` java
country == 'UK' && year >= 2020 && isActive == true.
```

4
spring-ai-docs/src/main/antora/modules/ROOT/pages/getting-started.adoc
View File

@@ -62,7 +62,7 @@ Add the Spring Boot Starter, depending on whether you use Azure Open AI or Open
<dependency>
<groupId>org.springframework.experimental.ai</groupId>
<artifactId>spring-ai-azure-openai-spring-boot-starter</artifactId>
<version>0.7.0-SNAPSHOT</version>
<version>0.7.1-SNAPSHOT</version>
</dependency>
----
@@ -72,7 +72,7 @@ Add the Spring Boot Starter, depending on whether you use Azure Open AI or Open
<dependency>
<groupId>org.springframework.experimental.ai</groupId>
<artifactId>spring-ai-openai-spring-boot-starter</artifactId>
<version>0.7.0-SNAPSHOT</version>
<version>0.7.1-SNAPSHOT</version>
</dependency>
----

2
spring-ai-huggingface/pom.xml
View File

@@ -4,7 +4,7 @@
<parent>
<groupId>org.springframework.experimental.ai</groupId>
<artifactId>spring-ai</artifactId>
<version>0.7.0-SNAPSHOT</version>
<version>0.7.1-SNAPSHOT</version>
</parent>
<artifactId>spring-ai-huggingface</artifactId>
<packaging>jar</packaging>

2
spring-ai-ollama/pom.xml
View File

@@ -6,7 +6,7 @@
<parent>
<groupId>org.springframework.experimental.ai</groupId>
<artifactId>spring-ai</artifactId>
<version>0.7.0-SNAPSHOT</version>
<version>0.7.1-SNAPSHOT</version>
</parent>
<artifactId>spring-ai-ollama</artifactId>

2
spring-ai-openai/pom.xml
View File

@@ -4,7 +4,7 @@
<parent>
<groupId>org.springframework.experimental.ai</groupId>
<artifactId>spring-ai</artifactId>
<version>0.7.0-SNAPSHOT</version>
<version>0.7.1-SNAPSHOT</version>
</parent>
<artifactId>spring-ai-openai</artifactId>
<packaging>jar</packaging>

91
spring-ai-spring-boot-autoconfigure/pom.xml
View File

@@ -1,10 +1,11 @@
<?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">
<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.experimental.ai</groupId>
<artifactId>spring-ai</artifactId>
<version>0.7.0-SNAPSHOT</version>
<version>0.7.1-SNAPSHOT</version>
</parent>
<artifactId>spring-ai-spring-boot-autoconfigure</artifactId>
<packaging>jar</packaging>
@@ -48,6 +49,45 @@
<optional>true</optional>
</dependency>
<!-- Pinecone Vector Store-->
<dependency>
<groupId>org.springframework.experimental.ai</groupId>
<artifactId>spring-ai-pinecone</artifactId>
<version>${project.parent.version}</version>
<optional>true</optional>
</dependency>
<!-- Milvus Vector Store -->
<dependency>
<groupId>org.springframework.experimental.ai</groupId>
<artifactId>spring-ai-milvus-store</artifactId>
<version>${project.parent.version}</version>
<optional>true</optional>
</dependency>
<!-- PG Vector Store-->
<dependency>
<groupId>org.springframework.experimental.ai</groupId>
<artifactId>spring-ai-pgvector-store</artifactId>
<version>${project.parent.version}</version>
<optional>true</optional>
</dependency>
<dependency>
<groupId>com.pgvector</groupId>
<artifactId>pgvector</artifactId>
<version>${pgvector.version}</version>
<optional>true</optional>
</dependency>
<dependency>
<groupId>org.postgresql</groupId>
<artifactId>postgresql</artifactId>
<version>${postgresql.version}</version>
<optional>true</optional>
</dependency>
<dependency>
<groupId>org.springframework.boot</groupId>
<artifactId>spring-boot-configuration-processor</artifactId>
@@ -61,11 +101,58 @@
</dependency>
<!-- test dependencies -->
<dependency>
<groupId>org.springframework.experimental.ai</groupId>
<artifactId>spring-ai-test</artifactId>
<version>${project.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>org.springframework.boot</groupId>
<artifactId>spring-boot-starter-jdbc</artifactId>
<scope>test</scope>
</dependency>
<dependency>
<groupId>org.postgresql</groupId>
<artifactId>postgresql</artifactId>
<scope>test</scope>
</dependency>
<dependency>
<groupId>org.springframework.experimental.ai</groupId>
<artifactId>transformers-embedding</artifactId>
<version>${parent.version}</version>
<scope>test</scope>
</dependency>
<dependency>
<groupId>org.testcontainers</groupId>
<artifactId>testcontainers</artifactId>
<version>${testcontainers.version}</version>
<scope>test</scope>
</dependency>
<dependency>
<groupId>org.testcontainers</groupId>
<artifactId>junit-jupiter</artifactId>
<version>${testcontainers.version}</version>
<scope>test</scope>
</dependency>
<dependency>
<groupId>org.awaitility</groupId>
<artifactId>awaitility</artifactId>
<version>3.0.0</version>
<scope>test</scope>
</dependency>
</dependencies>
</project>

280
spring-ai-spring-boot-autoconfigure/src/main/java/org/springframework/ai/autoconfigure/vectorstore/milvus/MilvusServiceClientProperties.java Normal file
View File

@@ -0,0 +1,280 @@
/*
* 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.milvus;
import java.util.concurrent.TimeUnit;
import org.springframework.boot.context.properties.ConfigurationProperties;
import static org.springframework.ai.autoconfigure.vectorstore.milvus.MilvusServiceClientProperties.CONFIG_PREFIX;
/**
* Parameters for Milvus client connection.
*
* @author Christian Tzolov
*/
@ConfigurationProperties(CONFIG_PREFIX)
public class MilvusServiceClientProperties {
public static final String CONFIG_PREFIX = "spring.ai.vectorstore.milvus.client";
/**
* Milvus host name/address.
*/
private String host = "localhost";
/**
* Milvus the connection port. Value must be greater than zero and less than 65536.
*/
private int port = 19530;
/**
* The uri of Milvus instance
*/
private String uri;
/**
* Token serving as the key for identification and authentication purposes.
*/
private String token;
/**
* Connection timeout value of client channel. The timeout value must be greater than
* zero.
*/
private long connectTimeoutMs = 10000;
/**
* Keep-alive time value of client channel. The keep-alive value must be greater than
* zero. Default is 55000 ms.
*/
private long keepAliveTimeMs = 55000;
/**
* The keep-alive timeout value of client channel. The timeout value must be greater
* than zero. Default value is 20000 ms.
*/
private long keepAliveTimeoutMs = 20000;
/**
* Enables the keep-alive function for client channel.
*/
// private boolean keepAliveWithoutCalls = false;
/**
* Deadline for how long you are willing to wait for a reply from the server. With a
* deadline setting, the client will wait when encounter fast RPC fail caused by
* network fluctuations. The deadline value must be larger than or equal to zero.
* Default value is 0, deadline is disabled.
*/
private long rpcDeadlineMs = 0; // Disabling deadline
/**
* The client.key path for tls two-way authentication, only takes effect when "secure"
* is True.
*/
private String clientKeyPath;
/**
* The client.pem path for tls two-way authentication, only takes effect when "secure"
* is True.
*/
private String clientPemPath;
/**
* The ca.pem path for tls two-way authentication, only takes effect when "secure" is
* True.
*/
private String caPemPath;
/**
* server.pem path for tls one-way authentication, only takes effect when "secure" is
* True.
*/
private String serverPemPath;
/**
* Sets the target name override for SSL host name checking, only takes effect when
* "secure" is True. Note: this value is passed to grpc.ssl_target_name_override
*/
private String serverName;
/**
* Secure the authorization for this connection, set to True to enable TLS.
*/
protected boolean secure = false;
/**
* Idle timeout value of client channel. The timeout value must be larger than zero.
*/
private long idleTimeoutMs = TimeUnit.MILLISECONDS.convert(24, TimeUnit.HOURS);
/**
* The username and password for this connection
*/
private String username = "root";
/**
* The password for this connection
*/
private String password = "milvus";
public String getHost() {
return host;
}
public void setHost(String host) {
this.host = host;
}
public int getPort() {
return port;
}
public void setPort(int port) {
this.port = port;
}
public String getUri() {
return uri;
}
public void setUri(String uri) {
this.uri = uri;
}
public String getToken() {
return token;
}
public void setToken(String token) {
this.token = token;
}
public long getConnectTimeoutMs() {
return connectTimeoutMs;
}
public void setConnectTimeoutMs(long connectTimeoutMs) {
this.connectTimeoutMs = connectTimeoutMs;
}
public long getKeepAliveTimeMs() {
return keepAliveTimeMs;
}
public void setKeepAliveTimeMs(long keepAliveTimeMs) {
this.keepAliveTimeMs = keepAliveTimeMs;
}
public long getKeepAliveTimeoutMs() {
return keepAliveTimeoutMs;
}
public void setKeepAliveTimeoutMs(long keepAliveTimeoutMs) {
this.keepAliveTimeoutMs = keepAliveTimeoutMs;
}
// public boolean isKeepAliveWithoutCalls() {
// return keepAliveWithoutCalls;
// }
// public void setKeepAliveWithoutCalls(boolean keepAliveWithoutCalls) {
// this.keepAliveWithoutCalls = keepAliveWithoutCalls;
// }
public long getRpcDeadlineMs() {
return rpcDeadlineMs;
}
public void setRpcDeadlineMs(long rpcDeadlineMs) {
this.rpcDeadlineMs = rpcDeadlineMs;
}
public String getClientKeyPath() {
return clientKeyPath;
}
public void setClientKeyPath(String clientKeyPath) {
this.clientKeyPath = clientKeyPath;
}
public String getClientPemPath() {
return clientPemPath;
}
public void setClientPemPath(String clientPemPath) {
this.clientPemPath = clientPemPath;
}
public String getCaPemPath() {
return caPemPath;
}
public void setCaPemPath(String caPemPath) {
this.caPemPath = caPemPath;
}
public String getServerPemPath() {
return serverPemPath;
}
public void setServerPemPath(String serverPemPath) {
this.serverPemPath = serverPemPath;
}
public String getServerName() {
return serverName;
}
public void setServerName(String serverName) {
this.serverName = serverName;
}
public boolean isSecure() {
return secure;
}
public void setSecure(boolean secure) {
this.secure = secure;
}
public long getIdleTimeoutMs() {
return idleTimeoutMs;
}
public void setIdleTimeoutMs(long idleTimeoutMs) {
this.idleTimeoutMs = idleTimeoutMs;
}
public String getUsername() {
return username;
}
public void setUsername(String username) {
this.username = username;
}
public String getPassword() {
return password;
}
public void setPassword(String password) {
this.password = password;
}
}

107
spring-ai-spring-boot-autoconfigure/src/main/java/org/springframework/ai/autoconfigure/vectorstore/milvus/MilvusVectorStoreAutoConfiguration.java Normal file
View File

@@ -0,0 +1,107 @@
/*
* 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.milvus;
import java.util.concurrent.TimeUnit;
import io.milvus.client.MilvusServiceClient;
import io.milvus.param.ConnectParam;
import org.springframework.ai.embedding.EmbeddingClient;
import org.springframework.ai.vectorstore.MilvusVectorStore;
import org.springframework.ai.vectorstore.MilvusVectorStore.MilvusVectorStoreConfig;
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;
import org.springframework.util.StringUtils;
/**
* @author Christian Tzolov
*/
@AutoConfiguration
@ConditionalOnClass({ MilvusVectorStore.class, EmbeddingClient.class })
@EnableConfigurationProperties({ MilvusServiceClientProperties.class, MilvusVectorStoreProperties.class })
public class MilvusVectorStoreAutoConfiguration {
@Bean
@ConditionalOnMissingBean
public VectorStore vectorStore(MilvusServiceClient milvusClient, EmbeddingClient embeddingClient,
MilvusVectorStoreProperties properties) {
MilvusVectorStoreConfig config = MilvusVectorStoreConfig.builder()
.withCollectionName(properties.getCollectionName())
.withDatabaseName(properties.getDatabaseName())
.withIndexType(properties.getIndexType())
.withMetricType(properties.getMetricType())
.withIndexParameters(properties.getIndexParameters())
.build();
return new MilvusVectorStore(milvusClient, embeddingClient, config);
}
@Bean
@ConditionalOnMissingBean
public MilvusServiceClient milvusClient(MilvusVectorStoreProperties serverProperties,
MilvusServiceClientProperties clientProperties) {
var builder = ConnectParam.newBuilder()
.withHost(clientProperties.getHost())
.withPort(clientProperties.getPort())
.withDatabaseName(serverProperties.getDatabaseName())
.withConnectTimeout(clientProperties.getConnectTimeoutMs(), TimeUnit.MILLISECONDS)
.withKeepAliveTime(clientProperties.getKeepAliveTimeMs(), TimeUnit.MILLISECONDS)
.withKeepAliveTimeout(clientProperties.getKeepAliveTimeoutMs(), TimeUnit.MILLISECONDS)
.withRpcDeadline(clientProperties.getRpcDeadlineMs(), TimeUnit.MILLISECONDS)
.withSecure(clientProperties.isSecure())
.withIdleTimeout(clientProperties.getIdleTimeoutMs(), TimeUnit.MILLISECONDS)
.withAuthorization(clientProperties.getUsername(), clientProperties.getPassword());
if (clientProperties.isSecure() && StringUtils.hasText(clientProperties.getUri())) {
builder.withUri(clientProperties.getUri());
}
if (clientProperties.isSecure() && StringUtils.hasText(clientProperties.getToken())) {
builder.withUri(clientProperties.getToken());
}
if (clientProperties.isSecure() && StringUtils.hasText(clientProperties.getClientKeyPath())) {
builder.withUri(clientProperties.getClientKeyPath());
}
if (clientProperties.isSecure() && StringUtils.hasText(clientProperties.getClientPemPath())) {
builder.withUri(clientProperties.getClientPemPath());
}
if (clientProperties.isSecure() && StringUtils.hasText(clientProperties.getCaPemPath())) {
builder.withUri(clientProperties.getCaPemPath());
}
if (clientProperties.isSecure() && StringUtils.hasText(clientProperties.getServerPemPath())) {
builder.withUri(clientProperties.getServerPemPath());
}
if (clientProperties.isSecure() && StringUtils.hasText(clientProperties.getServerName())) {
builder.withUri(clientProperties.getServerName());
}
return new MilvusServiceClient(builder.build());
}
}

109
spring-ai-spring-boot-autoconfigure/src/main/java/org/springframework/ai/autoconfigure/vectorstore/milvus/MilvusVectorStoreProperties.java Normal file
View File

@@ -0,0 +1,109 @@
/*
* 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.milvus;
import io.milvus.param.IndexType;
import io.milvus.param.MetricType;
import org.springframework.ai.vectorstore.MilvusVectorStore;
import org.springframework.boot.context.properties.ConfigurationProperties;
import org.springframework.util.Assert;
import static org.springframework.ai.autoconfigure.vectorstore.milvus.MilvusVectorStoreProperties.CONFIG_PREFIX;
/**
* @author Christian Tzolov
*/
@ConfigurationProperties(CONFIG_PREFIX)
public class MilvusVectorStoreProperties {
public static final String CONFIG_PREFIX = "spring.ai.vectorstore.milvus";
/**
* The database name
*/
private String databaseName = MilvusVectorStore.DEFAULT_DATABASE_NAME;
private String collectionName = MilvusVectorStore.DEFAULT_COLLECTION_NAME;
private int embeddingDimension = MilvusVectorStore.OPENAI_EMBEDDING_DIMENSION_SIZE;
private IndexType indexType = IndexType.IVF_FLAT;
private MetricType metricType = MetricType.COSINE;
private String indexParameters = "{\"nlist\":1024}";
public static String getConfigPrefix() {
return CONFIG_PREFIX;
}
public String getDatabaseName() {
return databaseName;
}
public void setDatabaseName(String databaseName) {
Assert.hasText(databaseName, "Database name should not be empty.");
this.databaseName = databaseName;
}
public String getCollectionName() {
return collectionName;
}
public void setCollectionName(String collectionName) {
Assert.hasText(collectionName, "Collection name should not be empty.");
this.collectionName = collectionName;
}
public int getEmbeddingDimension() {
return embeddingDimension;
}
public void setEmbeddingDimension(int embeddingDimension) {
Assert.isTrue(embeddingDimension > 0, "Embedding dimension should be a positive value.");
this.embeddingDimension = embeddingDimension;
}
public IndexType getIndexType() {
return indexType;
}
public void setIndexType(IndexType indexType) {
Assert.notNull(indexType, "Index type can not be null");
this.indexType = indexType;
}
public MetricType getMetricType() {
return metricType;
}
public void setMetricType(MetricType metricType) {
Assert.notNull(metricType, "MetricType can not be null");
this.metricType = metricType;
}
public String getIndexParameters() {
return indexParameters;
}
public void setIndexParameters(String indexParameters) {
Assert.notNull(indexParameters, "indexParameters can not be null");
this.indexParameters = indexParameters;
}
}

49
spring-ai-spring-boot-autoconfigure/src/main/java/org/springframework/ai/autoconfigure/vectorstore/pgvector/PgVectorStoreAutoConfiguration.java Normal file
View File

@@ -0,0 +1,49 @@
/*
* 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.pgvector;
import javax.sql.DataSource;
import org.springframework.ai.embedding.EmbeddingClient;
import org.springframework.ai.vectorstore.PgVectorStore;
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.autoconfigure.jdbc.JdbcTemplateAutoConfiguration;
import org.springframework.boot.context.properties.EnableConfigurationProperties;
import org.springframework.context.annotation.Bean;
import org.springframework.jdbc.core.JdbcTemplate;
/**
* @author Christian Tzolov
*/
@AutoConfiguration(after = JdbcTemplateAutoConfiguration.class)
@ConditionalOnClass({ PgVectorStore.class, DataSource.class, JdbcTemplate.class })
@EnableConfigurationProperties(PgVectorStoreProperties.class)
public class PgVectorStoreAutoConfiguration {
@Bean
@ConditionalOnMissingBean
public VectorStore vectorStore(JdbcTemplate jdbcTemplate, EmbeddingClient embeddingClient,
PgVectorStoreProperties properties) {
return new PgVectorStore(jdbcTemplate, embeddingClient, properties.getDimensions(),
properties.getDistanceType(), properties.isRemoveExistingVectorStoreTable(), properties.getIndexType());
}
}

77
spring-ai-spring-boot-autoconfigure/src/main/java/org/springframework/ai/autoconfigure/vectorstore/pgvector/PgVectorStoreProperties.java Normal file
View File

@@ -0,0 +1,77 @@
/*
* 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.pgvector;
import org.springframework.ai.vectorstore.PgVectorStore;
import org.springframework.ai.vectorstore.PgVectorStore.PgDistanceType;
import org.springframework.ai.vectorstore.PgVectorStore.PgIndexType;
import org.springframework.boot.context.properties.ConfigurationProperties;
import static org.springframework.ai.autoconfigure.vectorstore.pgvector.PgVectorStoreProperties.CONFIG_PREFIX;
/**
* @author Christian Tzolov
*/
@ConfigurationProperties(CONFIG_PREFIX)
public class PgVectorStoreProperties {
public static final String CONFIG_PREFIX = "spring.ai.vectorstore.pgvector";
private int dimensions = PgVectorStore.INVALID_EMBEDDING_DIMENSION;
private PgIndexType indexType = PgIndexType.HNSW;
private PgDistanceType distanceType = PgDistanceType.CosineDistance;
private boolean removeExistingVectorStoreTable = false;
public static String getConfigPrefix() {
return CONFIG_PREFIX;
}
public int getDimensions() {
return dimensions;
}
public void setDimensions(int dimensions) {
this.dimensions = dimensions;
}
public PgIndexType getIndexType() {
return indexType;
}
public void setIndexType(PgIndexType createIndexMethod) {
this.indexType = createIndexMethod;
}
public PgDistanceType getDistanceType() {
return distanceType;
}
public void setDistanceType(PgDistanceType distanceType) {
this.distanceType = distanceType;
}
public boolean isRemoveExistingVectorStoreTable() {
return removeExistingVectorStoreTable;
}
public void setRemoveExistingVectorStoreTable(boolean removeExistingVectorStoreTable) {
this.removeExistingVectorStoreTable = removeExistingVectorStoreTable;
}
}

53
spring-ai-spring-boot-autoconfigure/src/main/java/org/springframework/ai/autoconfigure/vectorstore/pinecone/PineconeVectorStoreAutoConfiguration.java Normal file
View File

@@ -0,0 +1,53 @@
/*
* 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.pinecone;
import org.springframework.ai.embedding.EmbeddingClient;
import org.springframework.ai.vectorstore.PineconeVectorStore;
import org.springframework.ai.vectorstore.PineconeVectorStore.PineconeVectorStoreConfig;
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({ PineconeVectorStore.class, EmbeddingClient.class })
@EnableConfigurationProperties(PineconeVectorStoreProperties.class)
public class PineconeVectorStoreAutoConfiguration {
@Bean
@ConditionalOnMissingBean
public VectorStore vectorStore(EmbeddingClient embeddingClient, PineconeVectorStoreProperties properties) {
var config = PineconeVectorStoreConfig.builder()
.withApiKey(properties.getApiKey())
.withEnvironment(properties.getEnvironment())
.withProjectId(properties.getProjectId())
.withIndexName(properties.getIndexName())
.withNamespace(properties.getNamespace())
.withServerSideTimeout(properties.getServerSideTimeout())
.build();
return new PineconeVectorStore(config, embeddingClient);
}
}

92
spring-ai-spring-boot-autoconfigure/src/main/java/org/springframework/ai/autoconfigure/vectorstore/pinecone/PineconeVectorStoreProperties.java Normal file
View File

@@ -0,0 +1,92 @@
/*
* 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.pinecone;
import java.time.Duration;
import org.springframework.boot.context.properties.ConfigurationProperties;
import static org.springframework.ai.autoconfigure.vectorstore.pinecone.PineconeVectorStoreProperties.CONFIG_PREFIX;
/**
* @author Christian Tzolov
*/
@ConfigurationProperties(CONFIG_PREFIX)
public class PineconeVectorStoreProperties {
public static final String CONFIG_PREFIX = "spring.ai.vectorstore.pinecone";
private String apiKey;
private String environment = "gcp-starter";
private String projectId;
private String indexName;
private String namespace = "";
private Duration serverSideTimeout = Duration.ofSeconds(20);
public String getApiKey() {
return this.apiKey;
}
public void setApiKey(String apiKey) {
this.apiKey = apiKey;
}
public String getEnvironment() {
return this.environment;
}
public void setEnvironment(String environment) {
this.environment = environment;
}
public String getProjectId() {
return this.projectId;
}
public void setProjectId(String projectId) {
this.projectId = projectId;
}
public String getNamespace() {
return this.namespace;
}
public void setNamespace(String namespace) {
this.namespace = namespace;
}
public String getIndexName() {
return this.indexName;
}
public void setIndexName(String indexName) {
this.indexName = indexName;
}
public Duration getServerSideTimeout() {
return this.serverSideTimeout;
}
public void setServerSideTimeout(Duration serverSideTimeout) {
this.serverSideTimeout = serverSideTimeout;
}
}

5
spring-ai-spring-boot-autoconfigure/src/main/resources/META-INF/spring/org.springframework.boot.autoconfigure.AutoConfiguration.imports
View File

@@ -1,2 +1,5 @@
org.springframework.ai.autoconfigure.openai.OpenAiAutoConfiguration
org.springframework.ai.autoconfigure.azure.openai.AzureOpenAiAutoConfiguration
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

139
spring-ai-spring-boot-autoconfigure/src/test/java/org/springframework/ai/autoconfigure/vectorstore/milvus/MilvusVectorStoreAutoConfigurationIT.java Normal file
View File

@@ -0,0 +1,139 @@
/*
* 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.milvus;
import java.io.File;
import java.io.IOException;
import java.nio.charset.StandardCharsets;
import java.time.Duration;
import java.util.List;
import java.util.Map;
import java.util.UUID;
import org.junit.jupiter.api.AfterAll;
import org.junit.jupiter.api.BeforeAll;
import org.junit.jupiter.api.Test;
import org.springframework.core.io.DefaultResourceLoader;
import org.testcontainers.containers.DockerComposeContainer;
import org.testcontainers.containers.wait.strategy.Wait;
import org.testcontainers.junit.jupiter.Testcontainers;
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.springframework.util.FileSystemUtils;
import static org.assertj.core.api.Assertions.assertThat;
/**
* @author Christian Tzolov
*/
@Testcontainers
public class MilvusVectorStoreAutoConfigurationIT {
private static DockerComposeContainer milvusContainer;
private static final File TEMP_FOLDER = new File("target/test-" + UUID.randomUUID().toString());
List<Document> documents = List.of(
new Document(getText("classpath:/test/data/spring.ai.txt"), Map.of("spring", "great")),
new Document(getText("classpath:/test/data/time.shelter.txt")),
new Document(getText("classpath:/test/data/great.depression.txt"), Map.of("depression", "bad")));
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);
}
}
@BeforeAll
public static void beforeAll() {
FileSystemUtils.deleteRecursively(TEMP_FOLDER);
TEMP_FOLDER.mkdirs();
milvusContainer = new DockerComposeContainer(new File("src/test/resources/milvus/docker-compose.yml"))
.withEnv("DOCKER_VOLUME_DIRECTORY", TEMP_FOLDER.getAbsolutePath())
.withExposedService("standalone", 19530)
.withExposedService("standalone", 9091,
Wait.forHttp("/healthz").forPort(9091).forStatusCode(200).forStatusCode(401))
.waitingFor("standalone", Wait.forLogMessage(".*Proxy successfully started.*\\s", 1)
.withStartupTimeout(Duration.ofSeconds(100)));
milvusContainer.start();
}
@AfterAll
public static void afterAll() {
milvusContainer.stop();
FileSystemUtils.deleteRecursively(TEMP_FOLDER);
}
private final ApplicationContextRunner contextRunner = new ApplicationContextRunner()
.withConfiguration(AutoConfigurations.of(MilvusVectorStoreAutoConfiguration.class))
.withUserConfiguration(Config.class);
@Test
public void addAndSearch() {
contextRunner.withPropertyValues("spring.ai.vectorstore.milvus.metricType=COSINE",
"spring.ai.vectorstore.milvus.indexType=IVF_FLAT",
"spring.ai.vectorstore.milvus.embeddingDimension=384",
"spring.ai.vectorstore.milvus.collectionName=myTestCollection",
"spring.ai.vectorstore.milvus.client.host=localhost", "spring.ai.vectorstore.milvus.client.port=19530")
.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("spring", "distance");
// 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).hasSize(0);
});
}
@Configuration(proxyBeanMethods = false)
static class Config {
@Bean
public EmbeddingClient embeddingClient() {
return new TransformersEmbeddingClient();
}
}
}

115
spring-ai-spring-boot-autoconfigure/src/test/java/org/springframework/ai/autoconfigure/vectorstore/pgvector/PgVectorStoreAutoConfigurationIT.java Normal file
View File

@@ -0,0 +1,115 @@
/*
* 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.pgvector;
import java.io.IOException;
import java.nio.charset.StandardCharsets;
import java.util.List;
import java.util.Map;
import org.junit.jupiter.api.Test;
import org.springframework.core.io.DefaultResourceLoader;
import org.testcontainers.containers.GenericContainer;
import org.testcontainers.junit.jupiter.Container;
import org.testcontainers.junit.jupiter.Testcontainers;
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.autoconfigure.jdbc.DataSourceAutoConfiguration;
import org.springframework.boot.autoconfigure.jdbc.JdbcTemplateAutoConfiguration;
import org.springframework.boot.test.context.runner.ApplicationContextRunner;
import org.springframework.context.annotation.Bean;
import org.springframework.context.annotation.Configuration;
import static org.assertj.core.api.Assertions.assertThat;
/**
* @author Christian Tzolov
*/
@Testcontainers
public class PgVectorStoreAutoConfigurationIT {
@Container
static GenericContainer<?> postgresContainer = new GenericContainer<>("ankane/pgvector")
.withEnv("POSTGRES_USER", "postgres")
.withEnv("POSTGRES_PASSWORD", "postgres")
.withExposedPorts(5432);
List<Document> documents = List.of(
new Document(getText("classpath:/test/data/spring.ai.txt"), Map.of("spring", "great")),
new Document(getText("classpath:/test/data/time.shelter.txt")),
new Document(getText("classpath:/test/data/great.depression.txt"), Map.of("depression", "bad")));
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);
}
}
private final ApplicationContextRunner contextRunner = new ApplicationContextRunner()
.withConfiguration(AutoConfigurations.of(PgVectorStoreAutoConfiguration.class,
JdbcTemplateAutoConfiguration.class, DataSourceAutoConfiguration.class))
.withUserConfiguration(Config.class)
.withPropertyValues("spring.ai.vectorstore.pgvector.distanceType=CosineDistance",
// JdbcTemplate configuration
String.format("spring.datasource.url=jdbc:postgresql://localhost:%d/%s",
postgresContainer.getMappedPort(5432), "postgres"),
"spring.datasource.username=postgres", "spring.datasource.password=postgres");
@Test
public void addAndSearch() {
contextRunner.run(context -> {
VectorStore vectorStore = context.getBean(VectorStore.class);
vectorStore.add(documents);
List<Document> results = vectorStore
.similaritySearch(SearchRequest.query("What is Great Depression?").withTopK(1));
assertThat(results).hasSize(1);
Document resultDoc = results.get(0);
assertThat(resultDoc.getId()).isEqualTo(documents.get(2).getId());
assertThat(resultDoc.getMetadata()).containsKeys("depression", "distance");
// Remove all documents from the store
vectorStore.delete(documents.stream().map(doc -> doc.getId()).toList());
results = vectorStore.similaritySearch(SearchRequest.query("Great Depression").withTopK(1));
assertThat(results).hasSize(0);
});
}
@Configuration(proxyBeanMethods = false)
static class Config {
@Bean
public EmbeddingClient embeddingClient() {
return new TransformersEmbeddingClient();
}
}
}

56
spring-ai-spring-boot-autoconfigure/src/test/java/org/springframework/ai/autoconfigure/vectorstore/pgvector/PgVectorStorePropertiesTests.java Normal file
View File

@@ -0,0 +1,56 @@
/*
* 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.pgvector;
import static org.assertj.core.api.Assertions.assertThat;
import org.junit.jupiter.api.Test;
import org.springframework.ai.vectorstore.PgVectorStore;
import org.springframework.ai.vectorstore.PgVectorStore.PgDistanceType;
import org.springframework.ai.vectorstore.PgVectorStore.PgIndexType;
/**
* @author Christian Tzolov
*/
public class PgVectorStorePropertiesTests {
@Test
public void defaultValues() {
var props = new PgVectorStoreProperties();
assertThat(props.getDimensions()).isEqualTo(PgVectorStore.INVALID_EMBEDDING_DIMENSION);
assertThat(props.getDistanceType()).isEqualTo(PgDistanceType.CosineDistance);
assertThat(props.getIndexType()).isEqualTo(PgIndexType.HNSW);
assertThat(props.isRemoveExistingVectorStoreTable()).isFalse();
}
@Test
public void customValues() {
var props = new PgVectorStoreProperties();
props.setDimensions(1536);
props.setDistanceType(PgDistanceType.EuclideanDistance);
props.setIndexType(PgIndexType.IVFFLAT);
props.setRemoveExistingVectorStoreTable(true);
assertThat(props.getDimensions()).isEqualTo(1536);
assertThat(props.getDistanceType()).isEqualTo(PgDistanceType.EuclideanDistance);
assertThat(props.getIndexType()).isEqualTo(PgIndexType.IVFFLAT);
assertThat(props.isRemoveExistingVectorStoreTable()).isTrue();
}
}

124
spring-ai-spring-boot-autoconfigure/src/test/java/org/springframework/ai/autoconfigure/vectorstore/pinecone/PineconeVectorStoreAutoConfigurationIT.java Normal file
View File

@@ -0,0 +1,124 @@
/*
* 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.pinecone;
import java.io.IOException;
import java.nio.charset.StandardCharsets;
import java.util.List;
import java.util.Map;
import java.util.concurrent.TimeUnit;
import org.awaitility.Awaitility;
import org.awaitility.Duration;
import org.junit.jupiter.api.BeforeAll;
import org.junit.jupiter.api.Test;
import org.junit.jupiter.api.condition.EnabledIfEnvironmentVariable;
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.springframework.core.io.DefaultResourceLoader;
import static org.assertj.core.api.Assertions.assertThat;
import static org.hamcrest.Matchers.hasSize;
/**
* @author Christian Tzolov
*/
@EnabledIfEnvironmentVariable(named = "PINECONE_API_KEY", matches = ".+")
public class PineconeVectorStoreAutoConfigurationIT {
List<Document> documents = List.of(
new Document("1", getText("classpath:/test/data/spring.ai.txt"), Map.of("spring", "great")),
new Document("2", getText("classpath:/test/data/time.shelter.txt"), Map.of()),
new Document("3", getText("classpath:/test/data/great.depression.txt"), Map.of("depression", "bad")));
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);
}
}
private final ApplicationContextRunner contextRunner = new ApplicationContextRunner()
.withConfiguration(AutoConfigurations.of(PineconeVectorStoreAutoConfiguration.class))
.withUserConfiguration(Config.class)
.withPropertyValues("spring.ai.vectorstore.pinecone.apiKey=" + System.getenv("PINECONE_API_KEY"),
"spring.ai.vectorstore.pinecone.environment=gcp-starter",
"spring.ai.vectorstore.pinecone.projectId=814621f",
"spring.ai.vectorstore.pinecone.indexName=spring-ai-test-index");
@BeforeAll
public static void beforeAll() {
Awaitility.setDefaultPollInterval(2, TimeUnit.SECONDS);
Awaitility.setDefaultPollDelay(Duration.ZERO);
Awaitility.setDefaultTimeout(Duration.ONE_MINUTE);
}
@Test
public void addAndSearchTest() {
contextRunner.run(context -> {
VectorStore vectorStore = context.getBean(VectorStore.class);
vectorStore.add(documents);
Awaitility.await().until(() -> {
return vectorStore.similaritySearch(SearchRequest.query("Spring").withTopK(1));
}, hasSize(1));
List<Document> results = vectorStore.similaritySearch(SearchRequest.query("Spring").withTopK(1));
assertThat(results).hasSize(1);
Document resultDoc = results.get(0);
assertThat(resultDoc.getId()).isEqualTo(documents.get(0).getId());
assertThat(resultDoc.getContent()).contains(
"Spring AI provides abstractions that serve as the foundation for developing AI applications.");
assertThat(resultDoc.getMetadata()).hasSize(2);
assertThat(resultDoc.getMetadata()).containsKeys("spring", "distance");
// Remove all documents from the store
vectorStore.delete(documents.stream().map(doc -> doc.getId()).toList());
Awaitility.await().until(() -> {
return vectorStore.similaritySearch(SearchRequest.query("Spring").withTopK(1));
}, hasSize(0));
});
}
@Configuration(proxyBeanMethods = false)
static class Config {
@Bean
public EmbeddingClient embeddingClient() {
return new TransformersEmbeddingClient();
}
}
}

59
spring-ai-spring-boot-autoconfigure/src/test/java/org/springframework/ai/autoconfigure/vectorstore/pinecone/PineconeVectorStorePropertiesTests.java Normal file
View File

@@ -0,0 +1,59 @@
/*
* 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.pinecone;
import java.time.Duration;
import org.junit.jupiter.api.Test;
import static org.assertj.core.api.Assertions.assertThat;
/**
* @author Christian Tzolov
*/
public class PineconeVectorStorePropertiesTests {
@Test
public void defaultValues() {
var props = new PineconeVectorStoreProperties();
assertThat(props.getEnvironment()).isEqualTo("gcp-starter");
assertThat(props.getNamespace()).isEqualTo("");
assertThat(props.getApiKey()).isNull();
assertThat(props.getProjectId()).isNull();
assertThat(props.getIndexName()).isNull();
assertThat(props.getServerSideTimeout()).isEqualTo(Duration.ofSeconds(20));
}
@Test
public void customValues() {
var props = new PineconeVectorStoreProperties();
props.setApiKey("key");
props.setEnvironment("env");
props.setIndexName("index");
props.setNamespace("namespace");
props.setProjectId("project");
props.setServerSideTimeout(Duration.ofSeconds(60));
assertThat(props.getEnvironment()).isEqualTo("env");
assertThat(props.getNamespace()).isEqualTo("namespace");
assertThat(props.getApiKey()).isEqualTo("key");
assertThat(props.getProjectId()).isEqualTo("project");
assertThat(props.getIndexName()).isEqualTo("index");
assertThat(props.getServerSideTimeout()).isEqualTo(Duration.ofSeconds(60));
}
}

62
spring-ai-spring-boot-autoconfigure/src/test/resources/milvus/docker-compose.yml Normal file
View File

@@ -0,0 +1,62 @@
version: '3.5'
services:
etcd:
image: quay.io/coreos/etcd:v3.5.5
ports:
- "2379:2379"
environment:
- ETCD_AUTO_COMPACTION_MODE=revision
- ETCD_AUTO_COMPACTION_RETENTION=1000
- ETCD_QUOTA_BACKEND_BYTES=4294967296
- ETCD_SNAPSHOT_COUNT=50000
volumes:
- ${DOCKER_VOLUME_DIRECTORY:-.}/volumes/etcd:/etcd
command: etcd -advertise-client-urls=http://127.0.0.1:2379 -listen-client-urls http://0.0.0.0:2379 --data-dir /etcd
healthcheck:
test: ["CMD", "etcdctl", "endpoint", "health"]
interval: 30s
timeout: 20s
retries: 3
minio:
image: minio/minio:RELEASE.2023-11-11T08-14-41Z
environment:
MINIO_ACCESS_KEY: minioadmin
MINIO_SECRET_KEY: minioadmin
ports:
- "9001:9001"
- "9000:9000"
volumes:
- ${DOCKER_VOLUME_DIRECTORY:-.}/volumes/minio:/minio_data
command: minio server /minio_data --console-address ":9001"
healthcheck:
test: ["CMD", "curl", "-f", "http://localhost:9000/minio/health/live"]
interval: 30s
timeout: 20s
retries: 3
standalone:
image: milvusdb/milvus:v2.3.1
command: ["milvus", "run", "standalone"]
environment:
ETCD_ENDPOINTS: etcd:2379
MINIO_ADDRESS: minio:9000
volumes:
- ${DOCKER_VOLUME_DIRECTORY:-.}/volumes/milvus:/var/lib/milvus
healthcheck:
test: ["CMD", "curl", "-f", "http://localhost:9091/healthz"]
interval: 30s
start_period: 90s
timeout: 20s
retries: 3
ports:
- "19530:19530"
- "9091:9091"
depends_on:
- "etcd"
- "minio"
networks:
default:
name: milvus

2
spring-ai-spring-boot-starters/spring-ai-starter-azure-openai/pom.xml
View File

@@ -4,7 +4,7 @@
<parent>
<groupId>org.springframework.experimental.ai</groupId>
<artifactId>spring-ai</artifactId>
<version>0.7.0-SNAPSHOT</version>
<version>0.7.1-SNAPSHOT</version>
<relativePath>../../pom.xml</relativePath>
</parent>
<artifactId>spring-ai-azure-openai-spring-boot-starter</artifactId>

2
spring-ai-spring-boot-starters/spring-ai-starter-openai/pom.xml
View File

@@ -4,7 +4,7 @@
<parent>
<groupId>org.springframework.experimental.ai</groupId>
<artifactId>spring-ai</artifactId>
<version>0.7.0-SNAPSHOT</version>
<version>0.7.1-SNAPSHOT</version>
<relativePath>../../pom.xml</relativePath>
</parent>
<artifactId>spring-ai-openai-spring-boot-starter</artifactId>

2
spring-ai-test/pom.xml
View File

@@ -4,7 +4,7 @@
<parent>
<groupId>org.springframework.experimental.ai</groupId>
<artifactId>spring-ai</artifactId>
<version>0.7.0-SNAPSHOT</version>
<version>0.7.1-SNAPSHOT</version>
</parent>
<artifactId>spring-ai-test</artifactId>
<packaging>jar</packaging>

4
spring-ai-test/src/main/resources/test/data/great.depression.txt Normal file
View File

@@ -0,0 +1,4 @@
The Great Depression (1929–1939) was an economic shock that affected most countries across the world. It was a period of economic depression that became evident after a major fall in stock prices in the United States.[1] The economic contagion began around September 1929 and led to the Wall Street stock market crash of October 24 (Black Thursday). It was the longest, deepest, and most widespread depression of the 20th century.[2]
Between 1929 and 1932, worldwide gross domestic product (GDP) fell by an estimated 15%. By comparison, worldwide GDP fell by less than 1% from 2008 to 2009 during the Great Recession.[3] Some economies started to recover by the mid-1930s. However, in many countries,[specify] the negative effects of the Great Depression lasted until the beginning of World War II. Devastating effects were seen in both rich and poor countries with falling personal income, prices, tax revenues, and profits. International trade fell by more than 50%, unemployment in the U.S. rose to 23% and in some countries rose as high as 33%.[4]
Cities around the world were hit hard, especially those dependent on heavy industry. Construction was virtually halted in many countries. Farming communities and rural areas suffered as crop prices fell by about 60%.[5][6][7] Faced with plummeting demand and few job alternatives, areas dependent on primary sector industries suffered the most.[8]
Economic historians usually consider the catalyst of the Great Depression to be the sudden devastating collapse of U.S. stock market prices, starting on October 24, 1929. However, some dispute this conclusion, seeing the stock crash less as a cause of the Depression and more as a symptom of the rising nervousness of investors partly due to gradual price declines caused by falling sales of consumer goods (as a result of overproduction because of new production techniques, falling exports and income inequality, among other factors) that had already been underway as part of a gradual Depression

6
spring-ai-test/src/main/resources/test/data/spring.ai.txt Normal file
View File

@@ -0,0 +1,6 @@
The Spring AI project aims to streamline the development of applications that incorporate artificial intelligence functionality without unnecessary complexity.
The project draws inspiration from notable Python projects, such as LangChain and LlamaIndex, but Spring AI is not a direct port of those projects. The project was founded with the belief that the next wave of Generative AI applications will not be only for Python developers but will be ubiquitous across many programming languages.
At its core, Spring AI provides abstractions that serve as the foundation for developing AI applications. These abstractions have multiple implementations, enabling easy component swapping with minimal code changes. For example, Spring AI introduces the AiClient interface with implementations for OpenAI and Azure OpenAI.
In addition to these core abstractions, Spring AI aims to provide higher-level functionalities to address common use cases such as “Q&A over your documentation” or “Chat with your documentation.” As the complexity of the use cases increases, the Spring AI project will integrate with other projects in the Spring Ecosystem, such as Spring Integration, Spring Batch, and Spring Data.
To simplify setup, Spring Boot starters are available to help set up essential dependencies and classes. There is also a collection of sample applications to help you explore the project’s features. Lastly, the new Spring CLI project also enables you to get started quickly by using the spring boot new ai command for new projects or spring boot add ai for adding AI capabilities to your existing application.

2
spring-ai-test/src/main/resources/test/data/time.shelter.txt Normal file
View File

@@ -0,0 +1,2 @@
Somewhere in the Andes, they believe to this very day that the future is behind you. It comes up from behind your back, surprising and unforeseeable, while the past is always before your eyes, that which has already happened. When they talk about the past, the people of the Aymara tribe point in front of them. You walk forward facing the past and you turn back toward the future.
― Georgi Gospodinov, Time Shelter

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

@@ -5,7 +5,7 @@
<parent>
<groupId>org.springframework.experimental.ai</groupId>
<artifactId>spring-ai</artifactId>
<version>0.7.0-SNAPSHOT</version>
<version>0.7.1-SNAPSHOT</version>
<relativePath>../../pom.xml</relativePath>
</parent>
<artifactId>spring-ai-milvus-store</artifactId>
@@ -36,11 +36,20 @@
<!-- TESTING -->
<dependency>
<groupId>org.springframework.experimental.ai</groupId>
<artifactId>spring-ai-openai-spring-boot-starter</artifactId>
<artifactId>spring-ai-openai</artifactId>
<!-- <artifactId>transformers-embedding</artifactId> -->
<version>${parent.version}</version>
<scope>test</scope>
</dependency>
<dependency>
<groupId>org.springframework.experimental.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>

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

@@ -54,7 +54,6 @@ import org.slf4j.LoggerFactory;
import org.springframework.ai.document.Document;
import org.springframework.ai.embedding.EmbeddingClient;
import org.springframework.ai.vectorstore.filter.Filter;
import org.springframework.ai.vectorstore.filter.converter.MilvusFilterExpressionConverter;
import org.springframework.beans.factory.InitializingBean;
import org.springframework.util.Assert;
@@ -150,7 +149,7 @@ public class MilvusVectorStore implements VectorStore, InitializingBean {
private IndexType indexType = IndexType.IVF_FLAT;
private MetricType metricType = MetricType.L2;
private MetricType metricType = MetricType.COSINE;
private String indexParameters = "{\"nlist\":1024}";
@@ -159,13 +158,14 @@ public class MilvusVectorStore implements VectorStore, InitializingBean {
/**
* Configures the Milvus metric type to use. Leave {@literal null} or blank to
* use the metric metric.
* use the metric metric: https://milvus.io/docs/metric.md#floating
* @param metricType the metric type to use
* @return this builder
*/
public Builder withMetricType(MetricType metricType) {
Assert.notNull(metricType, "Collection Name must not be empty");
Assert.isTrue(metricType == MetricType.IP || metricType == MetricType.L2,
Assert.isTrue(
metricType == MetricType.IP || metricType == MetricType.L2 || metricType == MetricType.COSINE,
"Only the text metric types IP and L2 are supported");
this.metricType = metricType;
@@ -321,43 +321,26 @@ public class MilvusVectorStore implements VectorStore, InitializingBean {
}
@Override
public List<Document> similaritySearch(String query) {
return this.similaritySearch(query, 4);
}
public List<Document> similaritySearch(SearchRequest request) {
@Override
public List<Document> similaritySearch(String query, int topK) {
return similaritySearch(query, topK, 0.0D);
}
String nativeFilterExpressions = (request.getFilterExpression() != null)
? this.filterExpressionConverter.convert(request.getFilterExpression()) : "";
@Override
public List<Document> similaritySearch(String query, int topK, double similarityThreshold) {
return internalSimilaritySearch(query, topK, similarityThreshold, "");
}
Assert.notNull(request.getQuery(), "Query string must not be null");
@Override
public List<Document> similaritySearch(String query, int k, double threshold, Filter.Expression filterExpression) {
String pgVectorFilterExpression = this.filterExpressionConverter.convert(filterExpression);
return this.internalSimilaritySearch(query, k, threshold, pgVectorFilterExpression);
}
List<Document> internalSimilaritySearch(String query, int topK, double similarityThreshold,
String filterExpressions) {
Assert.notNull(query, "Query string must not be null");
List<Double> embedding = this.embeddingClient.embed(query);
List<Double> embedding = this.embeddingClient.embed(request.getQuery());
var searchParamBuilder = SearchParam.newBuilder()
.withCollectionName(this.config.collectionName)
.withConsistencyLevel(ConsistencyLevelEnum.STRONG)
.withMetricType(this.config.metricType)
.withOutFields(SEARCH_OUTPUT_FIELDS)
.withTopK(topK)
.withTopK(request.getTopK())
.withVectors(List.of(toFloatList(embedding)))
.withVectorFieldName(EMBEDDING_FIELD_NAME);
if (StringUtils.hasText(filterExpressions)) {
searchParamBuilder.withExpr(filterExpressions);
if (StringUtils.hasText(nativeFilterExpressions)) {
searchParamBuilder.withExpr(nativeFilterExpressions);
}
R<SearchResults> respSearch = milvusClient.search(searchParamBuilder.build());
@@ -368,9 +351,9 @@ public class MilvusVectorStore implements VectorStore, InitializingBean {
SearchResultsWrapper wrapperSearch = new SearchResultsWrapper(respSearch.getData().getResults());
return wrapperSearch.getRowRecords()
return wrapperSearch.getRowRecords(0)
.stream()
.filter(rowRecord -> getResultSimilarity(rowRecord) >= similarityThreshold)
.filter(rowRecord -> getResultSimilarity(rowRecord) >= request.getSimilarityThreshold())
.map(rowRecord -> {
String docId = (String) rowRecord.get(DOC_ID_FIELD_NAME);
String content = (String) rowRecord.get(CONTENT_FIELD_NAME);
@@ -384,7 +367,8 @@ public class MilvusVectorStore implements VectorStore, InitializingBean {
private float getResultSimilarity(RowRecord rowRecord) {
Float distance = (Float) rowRecord.get(DISTANCE_FIELD_NAME);
return (this.config.metricType == MetricType.IP) ? distance : (1 - distance);
return (this.config.metricType == MetricType.IP || this.config.metricType == MetricType.COSINE) ? distance
: (1 - distance);
}
private List<Float> toFloatList(List<Double> embeddingDouble) {

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

@@ -17,31 +17,40 @@
package org.springframework.ai.vectorstore;
import java.io.File;
import java.io.IOException;
import java.nio.charset.StandardCharsets;
import java.time.Duration;
import java.util.Collections;
import java.util.List;
import java.util.Map;
import java.util.UUID;
import com.theokanning.openai.client.OpenAiApi;
import com.theokanning.openai.service.OpenAiService;
import io.milvus.client.MilvusServiceClient;
import io.milvus.param.ConnectParam;
import io.milvus.param.IndexType;
import io.milvus.param.MetricType;
import org.junit.jupiter.api.AfterAll;
import org.junit.jupiter.api.BeforeAll;
import org.junit.jupiter.api.condition.EnabledIfEnvironmentVariable;
import org.junit.jupiter.params.ParameterizedTest;
import org.junit.jupiter.params.provider.ValueSource;
import org.springframework.core.io.DefaultResourceLoader;
import org.testcontainers.containers.DockerComposeContainer;
import org.testcontainers.containers.wait.strategy.Wait;
import org.testcontainers.junit.jupiter.Testcontainers;
import retrofit2.Retrofit;
import retrofit2.adapter.rxjava2.RxJava2CallAdapterFactory;
import retrofit2.converter.jackson.JacksonConverterFactory;
import org.springframework.ai.autoconfigure.openai.OpenAiAutoConfiguration;
import org.springframework.ai.ResourceUtils;
import org.springframework.ai.document.Document;
import org.springframework.ai.embedding.EmbeddingClient;
import org.springframework.ai.openai.embedding.OpenAiEmbeddingClient;
import org.springframework.ai.vectorstore.MilvusVectorStore.MilvusVectorStoreConfig;
import org.springframework.beans.factory.annotation.Value;
import org.springframework.boot.SpringBootConfiguration;
import org.springframework.boot.autoconfigure.AutoConfigurations;
import org.springframework.boot.autoconfigure.EnableAutoConfiguration;
import org.springframework.boot.autoconfigure.jdbc.DataSourceAutoConfiguration;
import org.springframework.boot.test.context.runner.ApplicationContextRunner;
@@ -54,6 +63,7 @@ import static org.assertj.core.api.Assertions.assertThat;
* @author Christian Tzolov
*/
@Testcontainers
@EnabledIfEnvironmentVariable(named = "OPENAI_API_KEY", matches = ".+")
public class MilvusVectorStoreIT {
private static DockerComposeContainer milvusContainer;
@@ -61,16 +71,22 @@ public class MilvusVectorStoreIT {
private static final File TEMP_FOLDER = new File("target/test-" + UUID.randomUUID().toString());
private final ApplicationContextRunner contextRunner = new ApplicationContextRunner()
.withUserConfiguration(TestApplication.class)
.withPropertyValues("spring.ai.openai.apiKey=" + System.getenv("OPENAI_API_KEY"));
.withUserConfiguration(TestApplication.class);
List<Document> documents = List.of(
new Document("Spring AI rocks!! Spring AI rocks!! Spring AI rocks!! Spring AI rocks!! Spring AI rocks!!",
Collections.singletonMap("meta1", "meta1")),
new Document("Hello World Hello World Hello World Hello World Hello World Hello World Hello World"),
new Document(
"Great Depression Great Depression Great Depression Great Depression Great Depression Great Depression",
Collections.singletonMap("meta2", "meta2")));
new Document(getText("classpath:/test/data/spring.ai.txt"), Map.of("meta1", "meta1")),
new Document(getText("classpath:/test/data/time.shelter.txt")),
new Document(getText("classpath:/test/data/great.depression.txt"), Map.of("meta2", "meta2")));
public static String getText(String uri) {
var resource = new DefaultResourceLoader().getResource(uri);
try {
return resource.getContentAsString(StandardCharsets.UTF_8);
}
catch (IOException e) {
throw new RuntimeException(e);
}
}
@BeforeAll
public static void beforeAll() {
@@ -99,161 +115,159 @@ public class MilvusVectorStoreIT {
}
@ParameterizedTest(name = "{0} : {displayName} ")
@ValueSource(strings = { "L2", "IP" })
@ValueSource(strings = { "COSINE", "L2", "IP" })
public void addAndSearch(String metricType) {
contextRunner.withConfiguration(AutoConfigurations.of(OpenAiAutoConfiguration.class))
.withPropertyValues("test.spring.ai.vectorstore.milvus.metricType=" + metricType)
.run(context -> {
contextRunner.withPropertyValues("test.spring.ai.vectorstore.milvus.metricType=" + metricType).run(context -> {
VectorStore vectorStore = context.getBean(VectorStore.class);
VectorStore vectorStore = context.getBean(VectorStore.class);
resetCollection(vectorStore);
resetCollection(vectorStore);
vectorStore.add(documents);
vectorStore.add(documents);
List<Document> results = vectorStore.similaritySearch("Great", 1);
List<Document> results = vectorStore.similaritySearch(SearchRequest.query("Spring").withTopK(1));
assertThat(results).hasSize(1);
Document resultDoc = results.get(0);
assertThat(resultDoc.getId()).isEqualTo(documents.get(2).getId());
assertThat(resultDoc.getContent()).isEqualTo(
"Great Depression Great Depression Great Depression Great Depression Great Depression Great Depression");
assertThat(resultDoc.getMetadata()).hasSize(2);
assertThat(resultDoc.getMetadata()).containsKey("meta2");
assertThat(resultDoc.getMetadata()).containsKey("distance");
assertThat(results).hasSize(1);
Document resultDoc = results.get(0);
assertThat(resultDoc.getId()).isEqualTo(documents.get(0).getId());
assertThat(resultDoc.getContent()).contains(
"Spring AI provides abstractions that serve as the foundation for developing AI applications.");
assertThat(resultDoc.getMetadata()).hasSize(2);
assertThat(resultDoc.getMetadata()).containsKeys("meta1", "distance");
// Remove all documents from the store
vectorStore.delete(documents.stream().map(doc -> doc.getId()).toList());
// Remove all documents from the store
vectorStore.delete(documents.stream().map(doc -> doc.getId()).toList());
List<Document> results2 = vectorStore.similaritySearch("Hello", 1);
assertThat(results2).hasSize(0);
});
results = vectorStore.similaritySearch(SearchRequest.query("Spring").withTopK(1));
assertThat(results).hasSize(0);
});
}
@ParameterizedTest(name = "{0} : {displayName} ")
@ValueSource(strings = { "L2" })
// @ValueSource(strings = { "IP", "L2" })
@ValueSource(strings = { "COSINE" })
// @ValueSource(strings = { "COSINE", "IP", "L2" })
public void searchWithFilters(String metricType) throws InterruptedException {
// https://milvus.io/docs/json_data_type.md
final double THRESHOLD_ALL = 0.0;
contextRunner.withPropertyValues("test.spring.ai.vectorstore.milvus.metricType=" + metricType).run(context -> {
VectorStore vectorStore = context.getBean(VectorStore.class);
contextRunner.withConfiguration(AutoConfigurations.of(OpenAiAutoConfiguration.class))
.withPropertyValues("test.spring.ai.vectorstore.milvus.metricType=" + metricType)
.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));
var bgDocument = new Document("The World is Big and Salvation Lurks Around the Corner",
Map.of("country", "BG", "year", 2020));
var nlDocument = new Document("The World is Big and Salvation Lurks Around the Corner",
Map.of("country", "NL"));
var bgDocument2 = new Document("The World is Big and Salvation Lurks Around the Corner",
Map.of("country", "BG", "year", 2023));
vectorStore.add(List.of(bgDocument, nlDocument, bgDocument2));
vectorStore.add(List.of(bgDocument, nlDocument, bgDocument2));
List<Document> results = vectorStore.similaritySearch(SearchRequest.query("The World").withTopK(5));
assertThat(results).hasSize(3);
List<Document> results = vectorStore.similaritySearch("The World", 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("The World", 5, THRESHOLD_ALL, "country == 'NL'");
assertThat(results).hasSize(1);
assertThat(results.get(0).getId()).isEqualTo(nlDocument.getId());
results = vectorStore.similaritySearch(SearchRequest.query("The World")
.withTopK(5)
.withSimilarityThresholdAll()
.withFilterExpression("country == 'BG'"));
results = vectorStore.similaritySearch("The World", 5, THRESHOLD_ALL, "country == 'BG'");
assertThat(results).hasSize(2);
assertThat(results.get(0).getId()).isIn(bgDocument.getId(), bgDocument2.getId());
assertThat(results.get(1).getId()).isIn(bgDocument.getId(), bgDocument2.getId());
assertThat(results).hasSize(2);
assertThat(results.get(0).getId()).isIn(bgDocument.getId(), bgDocument2.getId());
assertThat(results.get(1).getId()).isIn(bgDocument.getId(), bgDocument2.getId());
results = vectorStore.similaritySearch("The World", 5, THRESHOLD_ALL,
"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("country == 'BG' && year == 2020"));
assertThat(results).hasSize(1);
assertThat(results.get(0).getId()).isEqualTo(bgDocument.getId());
});
}
@ParameterizedTest(name = "{0} : {displayName} ")
@ValueSource(strings = { "L2", "IP" })
@ValueSource(strings = { "COSINE", "L2", "IP" })
public void documentUpdate(String metricType) {
contextRunner.withConfiguration(AutoConfigurations.of(OpenAiAutoConfiguration.class))
.withPropertyValues("test.spring.ai.vectorstore.milvus.metricType=" + metricType)
.run(context -> {
contextRunner.withPropertyValues("test.spring.ai.vectorstore.milvus.metricType=" + metricType).run(context -> {
VectorStore vectorStore = context.getBean(VectorStore.class);
VectorStore vectorStore = context.getBean(VectorStore.class);
resetCollection(vectorStore);
resetCollection(vectorStore);
Document document = new Document(UUID.randomUUID().toString(), "Spring AI rocks!!",
Collections.singletonMap("meta1", "meta1"));
Document document = new Document(UUID.randomUUID().toString(), "Spring AI rocks!!",
Collections.singletonMap("meta1", "meta1"));
vectorStore.add(List.of(document));
vectorStore.add(List.of(document));
List<Document> results = vectorStore.similaritySearch("Spring", 5);
List<Document> results = vectorStore.similaritySearch(SearchRequest.query("Spring").withTopK(5));
assertThat(results).hasSize(1);
Document resultDoc = results.get(0);
assertThat(resultDoc.getId()).isEqualTo(document.getId());
assertThat(resultDoc.getContent()).isEqualTo("Spring AI rocks!!");
assertThat(resultDoc.getMetadata()).containsKey("meta1");
assertThat(resultDoc.getMetadata()).containsKey("distance");
assertThat(results).hasSize(1);
Document resultDoc = results.get(0);
assertThat(resultDoc.getId()).isEqualTo(document.getId());
assertThat(resultDoc.getContent()).isEqualTo("Spring AI rocks!!");
assertThat(resultDoc.getMetadata()).containsKey("meta1");
assertThat(resultDoc.getMetadata()).containsKey("distance");
Document sameIdDocument = new Document(document.getId(),
"The World is Big and Salvation Lurks Around the Corner",
Collections.singletonMap("meta2", "meta2"));
Document sameIdDocument = new Document(document.getId(),
"The World is Big and Salvation Lurks Around the Corner",
Collections.singletonMap("meta2", "meta2"));
vectorStore.add(List.of(sameIdDocument));
vectorStore.add(List.of(sameIdDocument));
results = vectorStore.similaritySearch("FooBar", 5);
results = vectorStore.similaritySearch(SearchRequest.query("FooBar").withTopK(5));
assertThat(results).hasSize(1);
resultDoc = results.get(0);
assertThat(resultDoc.getId()).isEqualTo(document.getId());
assertThat(resultDoc.getContent()).isEqualTo("The World is Big and Salvation Lurks Around the Corner");
assertThat(resultDoc.getMetadata()).containsKey("meta2");
assertThat(resultDoc.getMetadata()).containsKey("distance");
assertThat(results).hasSize(1);
resultDoc = results.get(0);
assertThat(resultDoc.getId()).isEqualTo(document.getId());
assertThat(resultDoc.getContent()).isEqualTo("The World is Big and Salvation Lurks Around the Corner");
assertThat(resultDoc.getMetadata()).containsKey("meta2");
assertThat(resultDoc.getMetadata()).containsKey("distance");
vectorStore.delete(List.of(document.getId()));
vectorStore.delete(List.of(document.getId()));
});
});
}
@ParameterizedTest(name = "{0} : {displayName} ")
@ValueSource(strings = { "L2", "IP" })
@ValueSource(strings = { "COSINE", "L2", "IP" })
public void searchWithThreshold(String metricType) {
contextRunner.withConfiguration(AutoConfigurations.of(OpenAiAutoConfiguration.class))
.withPropertyValues("test.spring.ai.vectorstore.milvus.metricType=" + metricType)
.run(context -> {
contextRunner.withPropertyValues("test.spring.ai.vectorstore.milvus.metricType=" + metricType).run(context -> {
VectorStore vectorStore = context.getBean(VectorStore.class);
VectorStore vectorStore = context.getBean(VectorStore.class);
resetCollection(vectorStore);
resetCollection(vectorStore);
vectorStore.add(documents);
vectorStore.add(documents);
List<Document> fullResult = vectorStore.similaritySearch("Great", 5, 0.0);
List<Document> fullResult = vectorStore
.similaritySearch(SearchRequest.query("Spring").withTopK(5).withSimilarityThresholdAll());
List<Float> distances = fullResult.stream()
.map(doc -> (Float) doc.getMetadata().get("distance"))
.toList();
List<Float> distances = fullResult.stream().map(doc -> (Float) doc.getMetadata().get("distance")).toList();
assertThat(distances).hasSize(3);
assertThat(distances).hasSize(3);
float threshold = (distances.get(0) + distances.get(1)) / 2;
float threshold = (distances.get(0) + distances.get(1)) / 2;
List<Document> results = vectorStore.similaritySearch("Great", 5, (1 - threshold));
List<Document> results = vectorStore
.similaritySearch(SearchRequest.query("Spring").withTopK(5).withSimilarityThreshold(1 - threshold));
assertThat(results).hasSize(1);
Document resultDoc = results.get(0);
assertThat(resultDoc.getId()).isEqualTo(documents.get(2).getId());
assertThat(resultDoc.getContent()).isEqualTo(
"Great Depression Great Depression Great Depression Great Depression Great Depression Great Depression");
assertThat(resultDoc.getMetadata()).containsKey("meta2");
assertThat(resultDoc.getMetadata()).containsKey("distance");
assertThat(results).hasSize(1);
Document resultDoc = results.get(0);
assertThat(resultDoc.getId()).isEqualTo(documents.get(0).getId());
assertThat(resultDoc.getContent()).contains(
"Spring AI provides abstractions that serve as the foundation for developing AI applications.");
assertThat(resultDoc.getMetadata()).containsKeys("meta1", "distance");
});
});
}
@SpringBootConfiguration
@@ -282,6 +296,25 @@ public class MilvusVectorStoreIT {
.build());
}
@Bean
public EmbeddingClient embeddingClient() {
Retrofit retrofit = new Retrofit.Builder().baseUrl("https://api.openai.com")
.client(OpenAiService.defaultClient(System.getenv("OPENAI_API_KEY"), Duration.ofSeconds(60)))
.addConverterFactory(JacksonConverterFactory.create(OpenAiService.defaultObjectMapper()))
.addCallAdapterFactory(RxJava2CallAdapterFactory.create())
.build();
OpenAiApi api = retrofit.create(OpenAiApi.class);
return new OpenAiEmbeddingClient(new OpenAiService(api), "text-embedding-ada-002");
}
// @Bean
// public EmbeddingClient embeddingClient() {
// return new TransformersEmbeddingClient();
// }
}
}

2
vector-stores/spring-ai-milvus-store/src/test/resources/docker-compose.yml
View File

@@ -20,7 +20,7 @@ services:
retries: 3
minio:
image: minio/minio:RELEASE.2023-09-23T03-47-50Z
image: minio/minio:RELEASE.2023-11-11T08-14-41Z
environment:
MINIO_ACCESS_KEY: minioadmin
MINIO_SECRET_KEY: minioadmin

16
vector-stores/spring-ai-neo4j-store/README.md
View File

@@ -6,7 +6,7 @@ This readme walks you through setting up `Neo4jVectorStore` to store document em
[Neo4j](https://neo4j.com) is an open source NoSQL graph database.
It is a fully transactional database (ACID) that stores data structured as graphs consisting of nodes, connected by relationships.
Inspired by the structure of the real world, it allows for high query performance on complex data, while remaining intuitive and simple for the developer.
Inspired by the structure of the real world, it allows for high query performance on complex data, while remaining intuitive and simple for the developer.
## What is Neo4j Vector Search?
@@ -18,7 +18,7 @@ Those indexes are powered by Lucene using a Hierarchical Navigable Small World G
1. OpenAI Account: Create an account at [OpenAI Signup](https://platform.openai.com/signup) and generate the token at [API Keys](https://platform.openai.com/account/api-keys).
2. A running Neo4j (5.13+) instance
2. A running Neo4j (5.13+) instance
1. [Docker](https://hub.docker.com/_/neo4j) image _neo4j:5.13_
2. [Neo4j Desktop](https://neo4j.com/download/)
3. [Neo4j Aura](https://neo4j.com/cloud/aura-free/)
@@ -53,22 +53,22 @@ To acquire Spring AI artifacts, declare the Spring Snapshot repository:
Add these dependencies to your project:
1. OpenAI: Required for calculating embeddings.
1. OpenAI: Required for calculating embeddings.
```xml
<dependency>
<groupId>org.springframework.experimental.ai</groupId>
<artifactId>spring-ai-openai-spring-boot-starter</artifactId>
<version>0.7.0-SNAPSHOT</version>
<version>0.7.1-SNAPSHOT</version>
</dependency>
```
2. Neo4j Vector Store
```xml
<dependency>
<groupId>org.springframework.experimental.ai</groupId>
<artifactId>spring-ai-neo4j-store</artifactId>
<version>0.7.0-SNAPSHOT</version>
<version>0.7.1-SNAPSHOT</version>
</dependency>
```
@@ -112,7 +112,7 @@ vectorStore.add(List.of(document));
And finally, retrieve documents similar to a query:
```java
List<Document> results = vectorStore.similaritySearch("Spring", 5);
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!!" as the first result.
@@ -124,7 +124,7 @@ The default configuration should fit for most of the basic use-cases, but if you
The default params
* embedding dimension = 1536
* distance type = cosine
* distance type = cosine
* document node label = "Document"
* node property for embedding = "embedding"
* database name = "neo4j"

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

@@ -5,7 +5,7 @@
<parent>
<groupId>org.springframework.experimental.ai</groupId>
<artifactId>spring-ai</artifactId>
<version>0.7.0-SNAPSHOT</version>
<version>0.7.1-SNAPSHOT</version>
<relativePath>../../pom.xml</relativePath>
</parent>
<artifactId>spring-ai-neo4j-store</artifactId>
@@ -41,7 +41,7 @@
<!-- TESTING -->
<dependency>
<groupId>org.springframework.experimental.ai</groupId>
<artifactId>spring-ai-openai-spring-boot-starter</artifactId>
<artifactId>spring-ai-openai</artifactId>
<version>${parent.version}</version>
<scope>test</scope>
</dependency>

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

@@ -258,23 +258,17 @@ public class Neo4jVectorStore implements VectorStore, InitializingBean {
}
@Override
public List<Document> similaritySearch(String query) {
return this.similaritySearch(query, 5);
}
public List<Document> similaritySearch(SearchRequest request) {
if (request.getFilterExpression() != null) {
throw new UnsupportedOperationException(
"The [" + this.getClass() + "] doesn't support metadata filtering!");
}
@Override
public List<Document> similaritySearch(String query, int k) {
return this.similaritySearch(query, k, 0);
}
@Override
public List<Document> similaritySearch(String query, int k, double threshold) {
Assert.isTrue(k > 0, "The number of documents to returned must be greater than zero");
Assert.isTrue(threshold >= 0 && threshold <= 1,
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.");
var embedding = Values.value(toFloatArray(this.embeddingClient.embed(query)));
var embedding = Values.value(toFloatArray(this.embeddingClient.embed(request.getQuery())));
try (var session = this.driver.session(this.config.sessionConfig)) {
return session
.run("""
@@ -282,8 +276,9 @@ public class Neo4jVectorStore implements VectorStore, InitializingBean {
YIELD node, score
WHERE score >= $threshold
RETURN node, score
""", Map.of("indexName", INDEX_NAME, "numberOfNearestNeighbours", k, "embeddingValue",
embedding, "threshold", threshold))
""",
Map.of("indexName", INDEX_NAME, "numberOfNearestNeighbours", request.getTopK(),
"embeddingValue", embedding, "threshold", request.getSimilarityThreshold()))
.list(Neo4jVectorStore::recordToDocument);
}
}

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

@@ -1,11 +1,15 @@
package org.springframework.ai.vectorstore;
import java.time.Duration;
import java.util.Collections;
import java.util.List;
import java.util.UUID;
import com.theokanning.openai.client.OpenAiApi;
import com.theokanning.openai.service.OpenAiService;
import org.junit.jupiter.api.BeforeEach;
import org.junit.jupiter.api.Test;
import org.junit.jupiter.api.condition.EnabledIfEnvironmentVariable;
import org.neo4j.driver.AuthTokens;
import org.neo4j.driver.Driver;
import org.neo4j.driver.GraphDatabase;
@@ -13,12 +17,14 @@ import org.testcontainers.containers.Neo4jContainer;
import org.testcontainers.junit.jupiter.Container;
import org.testcontainers.junit.jupiter.Testcontainers;
import org.testcontainers.utility.DockerImageName;
import retrofit2.Retrofit;
import retrofit2.adapter.rxjava2.RxJava2CallAdapterFactory;
import retrofit2.converter.jackson.JacksonConverterFactory;
import org.springframework.ai.autoconfigure.openai.OpenAiAutoConfiguration;
import org.springframework.ai.document.Document;
import org.springframework.ai.embedding.EmbeddingClient;
import org.springframework.ai.openai.embedding.OpenAiEmbeddingClient;
import org.springframework.boot.SpringBootConfiguration;
import org.springframework.boot.autoconfigure.AutoConfigurations;
import org.springframework.boot.autoconfigure.EnableAutoConfiguration;
import org.springframework.boot.autoconfigure.jdbc.DataSourceAutoConfiguration;
import org.springframework.boot.test.context.runner.ApplicationContextRunner;
@@ -29,8 +35,10 @@ import static org.assertj.core.api.Assertions.assertThat;
/**
* @author Gerrit Meier
* @author Michael Simons
* @author Christian Tzolov
*/
@Testcontainers
@EnabledIfEnvironmentVariable(named = "OPENAI_API_KEY", matches = ".+")
class Neo4jVectorStoreIT {
// Neo4j 5.12 has a bug wrt checking limits, so either 5.11 or anything higher than
@@ -48,8 +56,7 @@ class Neo4jVectorStoreIT {
Collections.singletonMap("meta2", "meta2")));
private final ApplicationContextRunner contextRunner = new ApplicationContextRunner()
.withUserConfiguration(TestApplication.class)
.withPropertyValues("spring.ai.openai.apiKey=" + System.getenv("OPENAI_API_KEY"));
.withUserConfiguration(TestApplication.class);
@BeforeEach
void cleanDatabase() {
@@ -59,13 +66,13 @@ class Neo4jVectorStoreIT {
@Test
void addAndSearchTest() {
this.contextRunner.withConfiguration(AutoConfigurations.of(OpenAiAutoConfiguration.class)).run(context -> {
this.contextRunner.run(context -> {
VectorStore vectorStore = context.getBean(VectorStore.class);
vectorStore.add(this.documents);
List<Document> results = vectorStore.similaritySearch("Great", 1);
List<Document> results = vectorStore.similaritySearch(SearchRequest.query("Great").withTopK(1));
assertThat(results).hasSize(1);
Document resultDoc = results.get(0);
@@ -78,7 +85,7 @@ class Neo4jVectorStoreIT {
// Remove all documents from the store
vectorStore.delete(this.documents.stream().map(Document::getId).toList());
List<Document> results2 = vectorStore.similaritySearch("Great", 1);
List<Document> results2 = vectorStore.similaritySearch(SearchRequest.query("Great").withTopK(1));
assertThat(results2).isEmpty();
});
}
@@ -86,7 +93,7 @@ class Neo4jVectorStoreIT {
@Test
void documentUpdateTest() {
this.contextRunner.withConfiguration(AutoConfigurations.of(OpenAiAutoConfiguration.class)).run(context -> {
this.contextRunner.run(context -> {
VectorStore vectorStore = context.getBean(VectorStore.class);
@@ -95,7 +102,7 @@ class Neo4jVectorStoreIT {
vectorStore.add(List.of(document));
List<Document> results = vectorStore.similaritySearch("Spring", 5);
List<Document> results = vectorStore.similaritySearch(SearchRequest.query("Spring").withTopK(5));
assertThat(results).hasSize(1);
Document resultDoc = results.get(0);
@@ -110,7 +117,7 @@ class Neo4jVectorStoreIT {
vectorStore.add(List.of(sameIdDocument));
results = vectorStore.similaritySearch("FooBar", 5);
results = vectorStore.similaritySearch(SearchRequest.query("FooBar").withTopK(5));
assertThat(results).hasSize(1);
resultDoc = results.get(0);
@@ -125,13 +132,14 @@ class Neo4jVectorStoreIT {
@Test
void searchThresholdTest() {
this.contextRunner.withConfiguration(AutoConfigurations.of(OpenAiAutoConfiguration.class)).run(context -> {
this.contextRunner.run(context -> {
VectorStore vectorStore = context.getBean(VectorStore.class);
vectorStore.add(this.documents);
List<Document> fullResult = vectorStore.similaritySearch("Great", 5, 0);
List<Document> fullResult = vectorStore
.similaritySearch(SearchRequest.query("Great").withTopK(5).withSimilarityThresholdAll());
List<Float> distances = fullResult.stream().map(doc -> (Float) doc.getMetadata().get("distance")).toList();
@@ -139,7 +147,8 @@ class Neo4jVectorStoreIT {
float threshold = (distances.get(0) + distances.get(1)) / 2;
List<Document> results = vectorStore.similaritySearch("Great", 5, 1 - threshold);
List<Document> results = vectorStore
.similaritySearch(SearchRequest.query("Great").withTopK(5).withSimilarityThreshold(1 - threshold));
assertThat(results).hasSize(1);
Document resultDoc = results.get(0);
@@ -199,6 +208,20 @@ class Neo4jVectorStoreIT {
AuthTokens.basic("neo4j", neo4jContainer.getAdminPassword()));
}
@Bean
public EmbeddingClient embeddingClient() {
Retrofit retrofit = new Retrofit.Builder().baseUrl("https://api.openai.com")
.client(OpenAiService.defaultClient(System.getenv("OPENAI_API_KEY"), Duration.ofSeconds(60)))
.addConverterFactory(JacksonConverterFactory.create(OpenAiService.defaultObjectMapper()))
.addCallAdapterFactory(RxJava2CallAdapterFactory.create())
.build();
OpenAiApi api = retrofit.create(OpenAiApi.class);
return new OpenAiEmbeddingClient(new OpenAiService(api), "text-embedding-ada-002");
}
}
}

6
vector-stores/spring-ai-pgvector-store/README.md
View File

@@ -83,7 +83,7 @@ Add these dependencies to your project:
<dependency>
<groupId>org.springframework.experimental.ai</groupId>
<artifactId>spring-ai-openai-spring-boot-starter</artifactId>
<version>0.7.0-SNAPSHOT</version>
<version>0.7.1-SNAPSHOT</version>
</dependency>
```
@@ -93,7 +93,7 @@ Add these dependencies to your project:
<dependency>
<groupId>org.springframework.experimental.ai</groupId>
<artifactId>spring-ai-pgvector-store</artifactId>
<version>0.7.0-SNAPSHOT</version>
<version>0.7.1-SNAPSHOT</version>
</dependency>
```
@@ -139,7 +139,7 @@ vectorStore.add(List.of(document));
And finally, retrieve documents similar to a query:
```java
List<Document> results = vectorStore.similaritySearch("Spring", 5);
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!!".

12
vector-stores/spring-ai-pgvector-store/pom.xml
View File

@@ -5,7 +5,7 @@
<parent>
<groupId>org.springframework.experimental.ai</groupId>
<artifactId>spring-ai</artifactId>
<version>0.7.0-SNAPSHOT</version>
<version>0.7.1-SNAPSHOT</version>
<relativePath>../../pom.xml</relativePath>
</parent>
<artifactId>spring-ai-pgvector-store</artifactId>
@@ -53,7 +53,15 @@
<!-- TESTING -->
<dependency>
<groupId>org.springframework.experimental.ai</groupId>
<artifactId>spring-ai-openai-spring-boot-starter</artifactId>
<artifactId>spring-ai-openai</artifactId>
<version>${parent.version}</version>
<scope>test</scope>
</dependency>
<dependency>
<groupId>org.springframework.experimental.ai</groupId>
<artifactId>spring-ai-test</artifactId>
<version>${parent.version}</version>
<scope>test</scope>
</dependency>

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

@@ -33,7 +33,6 @@ import org.slf4j.LoggerFactory;
import org.springframework.ai.document.Document;
import org.springframework.ai.embedding.EmbeddingClient;
import org.springframework.ai.vectorstore.filter.Filter;
import org.springframework.ai.vectorstore.filter.converter.PgVectorFilterExpressionConverter;
import org.springframework.beans.factory.InitializingBean;
import org.springframework.jdbc.core.JdbcTemplate;
@@ -102,11 +101,23 @@ public class PgVectorStore implements VectorStore, InitializingBean {
}
/**
* Defaults to CosineDistance. But if vectors are normalized to length 1 (like OpenAI
* embeddings), use inner product (NegativeInnerProduct) for best performance.
*/
public enum PgDistanceType {
// NOTE: works only if If vectors are normalized to length 1 (like OpenAI
// embeddings), use inner product for best performance.
// The Sentence transformers are NOT normalized:
// https://github.com/UKPLab/sentence-transformers/issues/233
EuclideanDistance("<->", "vector_l2_ops",
"SELECT *, embedding <-> ? AS distance FROM %s WHERE embedding <-> ? < ? %s ORDER BY distance LIMIT ? "),
// NOTE: works only if If vectors are normalized to length 1 (like OpenAI
// embeddings), use inner product for best performance.
// The Sentence transformers are NOT normalized:
// https://github.com/UKPLab/sentence-transformers/issues/233
NegativeInnerProduct("<#>", "vector_ip_ops",
"SELECT *, (1 + (embedding <#> ?)) AS distance FROM %s WHERE (1 + (embedding <#> ?)) < ? %s ORDER BY distance LIMIT ? "),
@@ -259,44 +270,24 @@ public class PgVectorStore implements VectorStore, InitializingBean {
}
@Override
public List<Document> similaritySearch(String query) {
return this.similaritySearch(query, 4);
}
public List<Document> similaritySearch(SearchRequest request) {
@Override
public List<Document> similaritySearch(String query, int topK) {
return this.similaritySearch(query, topK, 0.0 /** ALL */
);
}
@Override
public List<Document> similaritySearch(String query, int topK, double similarityThreshold) {
return this.internalSimilaritySearch(query, topK, similarityThreshold, "");
}
@Override
public List<Document> similaritySearch(String query, int k, double threshold, Filter.Expression filterExpression) {
String pgVectorFilterExpression = this.filterExpressionConverter.convert(filterExpression);
return this.internalSimilaritySearch(query, k, threshold, pgVectorFilterExpression);
}
List<Document> internalSimilaritySearch(String query, int topK, double similarityThreshold,
String filterExpression) {
String nativeFilterExpression = (request.getFilterExpression() != null)
? this.filterExpressionConverter.convert(request.getFilterExpression()) : "";
String jsonPathFilter = "";
if (StringUtils.hasText(filterExpression)) {
jsonPathFilter = " AND metadata::jsonb @@ '" + filterExpression + "'::jsonpath ";
if (StringUtils.hasText(nativeFilterExpression)) {
jsonPathFilter = " AND metadata::jsonb @@ '" + nativeFilterExpression + "'::jsonpath ";
}
double distance = 1 - similarityThreshold;
double distance = 1 - request.getSimilarityThreshold();
PGvector queryEmbedding = getQueryEmbedding(query);
PGvector queryEmbedding = getQueryEmbedding(request.getQuery());
return this.jdbcTemplate.query(
String.format(this.getDistanceType().similaritySearchSqlTemplate, VECTOR_TABLE_NAME, jsonPathFilter),
new DocumentRowMapper(this.objectMapper), queryEmbedding, queryEmbedding, distance, topK);
new DocumentRowMapper(this.objectMapper), queryEmbedding, queryEmbedding, distance, request.getTopK());
}
public List<Double> embeddingDistance(String query) {

128
vector-stores/spring-ai-pgvector-store/src/test/java/org/springframework/ai/vectorstore/PgVectorStoreIT.java
View File

@@ -16,6 +16,9 @@
package org.springframework.ai.vectorstore;
import java.io.IOException;
import java.nio.charset.StandardCharsets;
import java.time.Duration;
import java.util.Collections;
import java.util.Iterator;
import java.util.List;
@@ -24,22 +27,29 @@ import java.util.UUID;
import javax.sql.DataSource;
import com.theokanning.openai.client.OpenAiApi;
import com.theokanning.openai.service.OpenAiService;
import com.zaxxer.hikari.HikariDataSource;
import org.junit.Assert;
import org.junit.jupiter.api.condition.EnabledIfEnvironmentVariable;
import org.junit.jupiter.params.ParameterizedTest;
import org.junit.jupiter.params.provider.ValueSource;
import org.springframework.core.io.DefaultResourceLoader;
import org.testcontainers.containers.GenericContainer;
import org.testcontainers.junit.jupiter.Container;
import org.testcontainers.junit.jupiter.Testcontainers;
import retrofit2.Retrofit;
import retrofit2.adapter.rxjava2.RxJava2CallAdapterFactory;
import retrofit2.converter.jackson.JacksonConverterFactory;
import org.springframework.ai.autoconfigure.openai.OpenAiAutoConfiguration;
import org.springframework.ai.ResourceUtils;
import org.springframework.ai.document.Document;
import org.springframework.ai.embedding.EmbeddingClient;
import org.springframework.ai.openai.embedding.OpenAiEmbeddingClient;
import org.springframework.ai.vectorstore.PgVectorStore.PgIndexType;
import org.springframework.ai.vectorstore.filter.FilterExpressionTextParser.FilterExpressionParseException;
import org.springframework.beans.factory.annotation.Value;
import org.springframework.boot.SpringBootConfiguration;
import org.springframework.boot.autoconfigure.AutoConfigurations;
import org.springframework.boot.autoconfigure.EnableAutoConfiguration;
import org.springframework.boot.autoconfigure.jdbc.DataSourceAutoConfiguration;
import org.springframework.boot.autoconfigure.jdbc.DataSourceProperties;
@@ -57,6 +67,7 @@ import static org.assertj.core.api.Assertions.assertThat;
* @author Christian Tzolov
*/
@Testcontainers
@EnabledIfEnvironmentVariable(named = "OPENAI_API_KEY", matches = ".+")
public class PgVectorStoreIT {
@Container
@@ -66,17 +77,23 @@ public class PgVectorStoreIT {
.withExposedPorts(5432);
List<Document> documents = List.of(
new Document("Spring AI rocks!! Spring AI rocks!! Spring AI rocks!! Spring AI rocks!! Spring AI rocks!!",
Collections.singletonMap("meta1", "meta1")),
new Document("Hello World Hello World Hello World Hello World Hello World Hello World Hello World"),
new Document(
"Great Depression Great Depression Great Depression Great Depression Great Depression Great Depression",
Collections.singletonMap("meta2", "meta2")));
new Document(getText("classpath:/test/data/spring.ai.txt"), Map.of("meta1", "meta1")),
new Document(getText("classpath:/test/data/time.shelter.txt")),
new Document(getText("classpath:/test/data/great.depression.txt"), Map.of("meta2", "meta2")));
public static String getText(String uri) {
var resource = new DefaultResourceLoader().getResource(uri);
try {
return resource.getContentAsString(StandardCharsets.UTF_8);
}
catch (IOException e) {
throw new RuntimeException(e);
}
}
private final ApplicationContextRunner contextRunner = new ApplicationContextRunner()
.withUserConfiguration(TestApplication.class)
.withPropertyValues("spring.ai.openai.apiKey=" + System.getenv("OPENAI_API_KEY"),
"test.spring.ai.vectorstore.pgvector.distanceType=CosineDistance",
.withPropertyValues("test.spring.ai.vectorstore.pgvector.distanceType=CosineDistance",
// JdbcTemplate configuration
String.format("app.datasource.url=jdbc:postgresql://localhost:%d/%s",
@@ -92,27 +109,26 @@ public class PgVectorStoreIT {
@ParameterizedTest(name = "{0} : {displayName} ")
@ValueSource(strings = { "CosineDistance", "EuclideanDistance", "NegativeInnerProduct" })
public void addAndSearch(String distanceType) {
contextRunner.withConfiguration(AutoConfigurations.of(OpenAiAutoConfiguration.class))
.withPropertyValues("test.spring.ai.vectorstore.pgvector.distanceType=" + distanceType)
contextRunner.withPropertyValues("test.spring.ai.vectorstore.pgvector.distanceType=" + distanceType)
.run(context -> {
VectorStore vectorStore = context.getBean(VectorStore.class);
vectorStore.add(documents);
List<Document> results = vectorStore.similaritySearch("Great", 1);
List<Document> results = vectorStore
.similaritySearch(SearchRequest.query("What is Great Depression").withTopK(1));
assertThat(results).hasSize(1);
Document resultDoc = results.get(0);
assertThat(resultDoc.getId()).isEqualTo(documents.get(2).getId());
assertThat(resultDoc.getContent()).isEqualTo(
"Great Depression Great Depression Great Depression Great Depression Great Depression Great Depression");
assertThat(resultDoc.getMetadata()).containsKeys("meta2", "distance");
// Remove all documents from the store
vectorStore.delete(documents.stream().map(doc -> doc.getId()).toList());
List<Document> results2 = vectorStore.similaritySearch("Great", 1);
List<Document> results2 = vectorStore
.similaritySearch(SearchRequest.query("Great Depression").withTopK(1));
assertThat(results2).hasSize(0);
dropTable(context);
@@ -123,10 +139,7 @@ public class PgVectorStoreIT {
@ValueSource(strings = { "CosineDistance", "EuclideanDistance", "NegativeInnerProduct" })
public void searchWithFilters(String distanceType) {
final double THRESHOLD_ALL = 0.0;
contextRunner.withConfiguration(AutoConfigurations.of(OpenAiAutoConfiguration.class))
.withPropertyValues("test.spring.ai.vectorstore.pgvector.distanceType=" + distanceType)
contextRunner.withPropertyValues("test.spring.ai.vectorstore.pgvector.distanceType=" + distanceType)
.run(context -> {
VectorStore vectorStore = context.getBean(VectorStore.class);
@@ -140,45 +153,44 @@ public class PgVectorStoreIT {
vectorStore.add(List.of(bgDocument, nlDocument, bgDocument2));
List<Document> results = vectorStore.similaritySearch("The World", 5);
SearchRequest searchRequest = SearchRequest.query("The World").withTopK(5).withSimilarityThresholdAll();
List<Document> results = vectorStore.similaritySearch(searchRequest);
assertThat(results).hasSize(3);
results = vectorStore.similaritySearch("The World", 5, THRESHOLD_ALL, "country == 'NL'");
results = vectorStore.similaritySearch(searchRequest.withFilterExpression("country == 'NL'"));
assertThat(results).hasSize(1);
assertThat(results.get(0).getId()).isEqualTo(nlDocument.getId());
results = vectorStore.similaritySearch("The World", 5, THRESHOLD_ALL, "country == 'BG'");
results = vectorStore.similaritySearch(searchRequest.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("The World", 5, THRESHOLD_ALL,
"country == 'BG' && year == 2020");
results = vectorStore
.similaritySearch(searchRequest.withFilterExpression("country == 'BG' && year == 2020"));
assertThat(results).hasSize(1);
assertThat(results.get(0).getId()).isEqualTo(bgDocument.getId());
results = vectorStore.similaritySearch("The World", 5, THRESHOLD_ALL,
"(country == 'BG' && year == 2020) || (country == 'NL')");
results = vectorStore.similaritySearch(
searchRequest.withFilterExpression("(country == 'BG' && year == 2020) || (country == 'NL')"));
assertThat(results).hasSize(2);
assertThat(results.get(0).getId()).isIn(bgDocument.getId(), nlDocument.getId());
assertThat(results.get(1).getId()).isIn(bgDocument.getId(), nlDocument.getId());
try {
vectorStore.similaritySearch("The World", 5, THRESHOLD_ALL, "country == NL");
vectorStore.similaritySearch(searchRequest.withFilterExpression("country == NL"));
Assert.fail("Invalid filter expression should have been cached!");
}
catch (FilterExpressionParseException e) {
assertThat(e.getMessage()).contains("Line: 1:17, Error: no viable alternative at input 'NL'");
}
try {
results = ((PgVectorStore) vectorStore).internalSimilaritySearch("The World", 5, THRESHOLD_ALL,
"Invalid Expression");
Assert.fail("Malicious jsonpath expressions should be detected!");
}
catch (Exception e) {
}
// Remove all documents from the store
dropTable(context);
});
@@ -188,8 +200,7 @@ public class PgVectorStoreIT {
@ValueSource(strings = { "CosineDistance", "EuclideanDistance", "NegativeInnerProduct" })
public void documentUpdate(String distanceType) {
contextRunner.withConfiguration(AutoConfigurations.of(OpenAiAutoConfiguration.class))
.withPropertyValues("test.spring.ai.vectorstore.pgvector.distanceType=" + distanceType)
contextRunner.withPropertyValues("test.spring.ai.vectorstore.pgvector.distanceType=" + distanceType)
.run(context -> {
VectorStore vectorStore = context.getBean(VectorStore.class);
@@ -199,7 +210,7 @@ public class PgVectorStoreIT {
vectorStore.add(List.of(document));
List<Document> results = vectorStore.similaritySearch("Spring", 5);
List<Document> results = vectorStore.similaritySearch(SearchRequest.query("Spring").withTopK(5));
assertThat(results).hasSize(1);
Document resultDoc = results.get(0);
@@ -213,7 +224,7 @@ public class PgVectorStoreIT {
vectorStore.add(List.of(sameIdDocument));
results = vectorStore.similaritySearch("FooBar", 5);
results = vectorStore.similaritySearch(SearchRequest.query("FooBar").withTopK(5));
assertThat(results).hasSize(1);
resultDoc = results.get(0);
@@ -227,38 +238,35 @@ public class PgVectorStoreIT {
@ParameterizedTest(name = "{0} : {displayName} ")
@ValueSource(strings = { "CosineDistance", "EuclideanDistance", "NegativeInnerProduct" })
// @ValueSource(strings = { "CosineDistance" })
public void searchWithThreshold(String distanceType) {
contextRunner.withConfiguration(AutoConfigurations.of(OpenAiAutoConfiguration.class))
.withPropertyValues("test.spring.ai.vectorstore.pgvector.distanceType=" + distanceType)
contextRunner.withPropertyValues("test.spring.ai.vectorstore.pgvector.distanceType=" + distanceType)
.run(context -> {
VectorStore vectorStore = context.getBean(VectorStore.class);
vectorStore.add(documents);
List<Document> fullResult = vectorStore.similaritySearch("Great", 5, 0.0);
List<Float> distances = fullResult.stream()
.map(doc -> (Float) doc.getMetadata().get("distance"))
.toList();
List<Document> fullResult = vectorStore
.similaritySearch(SearchRequest.query("Time Shelter").withTopK(5).withSimilarityThresholdAll());
assertThat(fullResult).hasSize(3);
assertThat(isSortedByDistance(fullResult)).isTrue();
fullResult.stream().forEach(doc -> System.out.println(doc.getMetadata().get("distance")));
List<Float> distances = fullResult.stream()
.map(doc -> (Float) doc.getMetadata().get("distance"))
.toList();
float threshold = (distances.get(0) + distances.get(1)) / 2;
List<Document> results = vectorStore.similaritySearch("Great", 5, (1 - threshold));
List<Document> results = vectorStore.similaritySearch(
SearchRequest.query("Time Shelter").withTopK(5).withSimilarityThreshold(1 - threshold));
assertThat(results).hasSize(1);
Document resultDoc = results.get(0);
assertThat(resultDoc.getId()).isEqualTo(documents.get(2).getId());
assertThat(resultDoc.getContent()).isEqualTo(
"Great Depression Great Depression Great Depression Great Depression Great Depression Great Depression");
assertThat(resultDoc.getMetadata()).containsKeys("meta2", "distance");
assertThat(resultDoc.getId()).isEqualTo(documents.get(1).getId());
dropTable(context);
});
@@ -314,6 +322,20 @@ public class PgVectorStoreIT {
return dataSourceProperties.initializeDataSourceBuilder().type(HikariDataSource.class).build();
}
@Bean
public EmbeddingClient embeddingClient() {
Retrofit retrofit = new Retrofit.Builder().baseUrl("https://api.openai.com")
.client(OpenAiService.defaultClient(System.getenv("OPENAI_API_KEY"), Duration.ofSeconds(60)))
.addConverterFactory(JacksonConverterFactory.create(OpenAiService.defaultObjectMapper()))
.addCallAdapterFactory(RxJava2CallAdapterFactory.create())
.build();
OpenAiApi api = retrofit.create(OpenAiApi.class);
return new OpenAiEmbeddingClient(new OpenAiService(api), "text-embedding-ada-002");
}
}
}

6
vector-stores/spring-ai-pinecone/README.md
View File

@@ -59,7 +59,7 @@ Add these dependencies to your project:
<dependency>
<groupId>org.springframework.experimental.ai</groupId>
<artifactId>spring-ai-openai-spring-boot-starter</artifactId>
<version>0.7.0-SNAPSHOT</version>
<version>0.7.1-SNAPSHOT</version>
</dependency>
```
@@ -69,7 +69,7 @@ Add these dependencies to your project:
<dependency>
<groupId>org.springframework.experimental.ai</groupId>
<artifactId>spring-ai-pinecone</artifactId>
<version>0.7.0-SNAPSHOT</version>
<version>0.7.1-SNAPSHOT</version>
</dependency>
```
@@ -119,7 +119,7 @@ vectorStore.add(List.of(document));
And finally, retrieve documents similar to a query:
```java
List<Document> results = vectorStore.similaritySearch("Spring", 5);
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!!".

21
vector-stores/spring-ai-pinecone/pom.xml
View File

@@ -5,7 +5,7 @@
<parent>
<groupId>org.springframework.experimental.ai</groupId>
<artifactId>spring-ai</artifactId>
<version>0.7.0-SNAPSHOT</version>
<version>0.7.1-SNAPSHOT</version>
<relativePath>../../pom.xml</relativePath>
</parent>
<artifactId>spring-ai-pinecone</artifactId>
@@ -45,12 +45,19 @@
</dependency>
<!-- TESTING -->
<dependency>
<groupId>org.springframework.experimental.ai</groupId>
<artifactId>spring-ai-openai-spring-boot-starter</artifactId>
<version>${parent.version}</version>
<scope>test</scope>
</dependency>
<dependency>
<groupId>org.springframework.experimental.ai</groupId>
<artifactId>transformers-embedding</artifactId>
<version>${parent.version}</version>
<scope>test</scope>
</dependency>
<dependency>
<groupId>org.springframework.experimental.ai</groupId>
<artifactId>spring-ai-test</artifactId>
<version>${parent.version}</version>
<scope>test</scope>
</dependency>
<dependency>
<groupId>org.springframework.boot</groupId>

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

@@ -36,7 +36,6 @@ import io.pinecone.proto.Vector;
import org.springframework.ai.document.Document;
import org.springframework.ai.embedding.EmbeddingClient;
import org.springframework.ai.vectorstore.filter.Filter;
import org.springframework.ai.vectorstore.filter.converter.PineconeFilterExpressionConverter;
import org.springframework.util.Assert;
import org.springframework.util.StringUtils;
@@ -54,8 +53,6 @@ public class PineconeVectorStore implements VectorStore {
private static final String DISTANCE_METADATA_FIELD_NAME = "distance";
private static final Double SIMILARITY_THRESHOLD_ALL = 0.0;
public final PineconeFilterExpressionConverter filterExpressionConverter = new PineconeFilterExpressionConverter();
private final EmbeddingClient embeddingClient;
@@ -64,8 +61,6 @@ public class PineconeVectorStore implements VectorStore {
private final String pineconeNamespace;
private final int defaultSimilarityTopK;
private final ObjectMapper objectMapper;
/**
@@ -81,7 +76,7 @@ public class PineconeVectorStore implements VectorStore {
private final PineconeClientConfig clientConfig;
private final int defaultSimilarityTopK;
// private final int defaultSimilarityTopK;
/**
* Constructor using the builder.
@@ -93,7 +88,7 @@ public class PineconeVectorStore implements VectorStore {
*/
public PineconeVectorStoreConfig(Builder builder) {
this.namespace = builder.namespace;
this.defaultSimilarityTopK = builder.defaultSimilarityTopK;
// this.defaultSimilarityTopK = builder.defaultSimilarityTopK;
this.connectionConfig = new PineconeConnectionConfig().withIndexName(builder.indexName);
this.clientConfig = new PineconeClientConfig().withApiKey(builder.apiKey)
.withEnvironment(builder.environment)
@@ -130,8 +125,6 @@ public class PineconeVectorStore implements VectorStore {
// The free-tier (gcp-starter) doesn't support Namespaces!
private String namespace = "";
private int defaultSimilarityTopK = 5;
/**
* Optional server-side timeout in seconds for all operations. Default: 20
* seconds.
@@ -202,16 +195,6 @@ public class PineconeVectorStore implements VectorStore {
return this;
}
/**
* Pinecone default top K similarity search response size.
* @param defaultSimilarityTopK default top K to use.
* @return this builder.
*/
public Builder withDefaultTopK(int defaultSimilarityTopK) {
this.defaultSimilarityTopK = defaultSimilarityTopK;
return this;
}
/**
* {@return the immutable configuration}
*/
@@ -234,7 +217,6 @@ public class PineconeVectorStore implements VectorStore {
this.embeddingClient = embeddingClient;
this.pineconeNamespace = config.namespace;
this.defaultSimilarityTopK = config.defaultSimilarityTopK;
this.pineconeConnection = new PineconeClient(config.clientConfig).connect(config.connectionConfig);
this.objectMapper = new ObjectMapper();
}
@@ -313,65 +295,37 @@ public class PineconeVectorStore implements VectorStore {
return Optional.of(true);
}
/**
* Searches for documents similar to the given query. Uses the default topK value.
* @param query The query string.
* @return A list of similar documents.
*/
@Override
public List<Document> similaritySearch(String query) {
return similaritySearch(query, this.defaultSimilarityTopK);
}
public List<Document> similaritySearch(SearchRequest request) {
// return this.internalSimilaritySearch(request.getQuery(), request.getTopK(),
// request.getSimilarityThreshold(),
// request.getFilterExpression());
// }
/**
* Searches for documents similar to the given query.
* @param query The query string.
* @param topK The maximum number of results to return.
* @return A list of similar documents.
*/
@Override
public List<Document> similaritySearch(String query, int topK) {
return similaritySearch(query, topK, SIMILARITY_THRESHOLD_ALL);
}
// List<Document> internalSimilaritySearch(String query, int topK, double
// similarityThreshold,
// Filter.Expression filterExpression) {
/**
* Searches for documents similar to the given query.
* @param query The query string.
* @param topK The maximum number of results to return.
* @param similarityThreshold The similarity threshold for results.
* @return A list of similar documents.
*/
@Override
public List<Document> similaritySearch(String query, int topK, double similarityThreshold) {
String nativeExpressionFilters = (request.getFilterExpression() != null)
? this.filterExpressionConverter.convert(request.getFilterExpression()) : "";
return internalSimilaritySearch(query, topK, similarityThreshold, "");
}
@Override
public List<Document> similaritySearch(String query, int k, double threshold, Filter.Expression filterExpression) {
String pgVectorFilterExpression = this.filterExpressionConverter.convert(filterExpression);
return this.internalSimilaritySearch(query, k, threshold, pgVectorFilterExpression);
}
List<Document> internalSimilaritySearch(String query, int topK, double similarityThreshold, String filters) {
List<Double> queryEmbedding = this.embeddingClient.embed(query);
List<Double> queryEmbedding = this.embeddingClient.embed(request.getQuery());
var queryRequestBuilder = QueryRequest.newBuilder()
.addAllVector(toFloatList(queryEmbedding))
.setTopK(topK)
.setTopK(request.getTopK())
.setIncludeMetadata(true)
.setNamespace(this.pineconeNamespace);
if (StringUtils.hasText(filters)) {
queryRequestBuilder.setFilter(metadataFiltersToStruct(filters));
if (StringUtils.hasText(nativeExpressionFilters)) {
queryRequestBuilder.setFilter(metadataFiltersToStruct(nativeExpressionFilters));
}
QueryResponse queryResponse = this.pineconeConnection.getBlockingStub().query(queryRequestBuilder.build());
return queryResponse.getMatchesList()
.stream()
.filter(scoredVector -> scoredVector.getScore() >= similarityThreshold)
.filter(scoredVector -> scoredVector.getScore() >= request.getSimilarityThreshold())
.map(scoredVector -> {
var id = scoredVector.getId();
Struct metadataStruct = scoredVector.getMetadata();

98
vector-stores/spring-ai-pinecone/src/test/java/org/springframework/ai/vectorstore/PineconeVectorStoreIT.java
View File

@@ -16,6 +16,8 @@
package org.springframework.ai.vectorstore;
import java.io.IOException;
import java.nio.charset.StandardCharsets;
import java.util.Collections;
import java.util.List;
import java.util.Map;
@@ -28,15 +30,16 @@ import org.junit.jupiter.api.BeforeAll;
import org.junit.jupiter.api.Test;
import org.junit.jupiter.api.condition.EnabledIfEnvironmentVariable;
import org.springframework.ai.autoconfigure.openai.OpenAiAutoConfiguration;
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.PineconeVectorStore.PineconeVectorStoreConfig;
import org.springframework.boot.SpringBootConfiguration;
import org.springframework.boot.autoconfigure.AutoConfigurations;
import org.springframework.boot.autoconfigure.EnableAutoConfiguration;
import org.springframework.boot.test.context.runner.ApplicationContextRunner;
import org.springframework.context.annotation.Bean;
import org.springframework.core.io.DefaultResourceLoader;
import static org.assertj.core.api.Assertions.assertThat;
import static org.hamcrest.Matchers.equalTo;
@@ -46,7 +49,6 @@ import static org.hamcrest.Matchers.hasSize;
* @author Christian Tzolov
*/
@EnabledIfEnvironmentVariable(named = "PINECONE_API_KEY", matches = ".+")
@EnabledIfEnvironmentVariable(named = "OPENAI_API_KEY", matches = ".+")
public class PineconeVectorStoreIT {
// Replace the PINECONE_ENVIRONMENT, PINECONE_PROJECT_ID, PINECONE_INDEX_NAME and
@@ -61,16 +63,22 @@ public class PineconeVectorStoreIT {
private static final String PINECONE_NAMESPACE = "";
List<Document> documents = List.of(
new Document("Spring AI rocks!! Spring AI rocks!! Spring AI rocks!! Spring AI rocks!! Spring AI rocks!!",
Collections.singletonMap("meta1", "meta1")),
new Document("Hello World Hello World Hello World Hello World Hello World Hello World Hello World"),
new Document(
"Great Depression Great Depression Great Depression Great Depression Great Depression Great Depression",
Collections.singletonMap("meta2", "meta2")));
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);
}
}
private final ApplicationContextRunner contextRunner = new ApplicationContextRunner()
.withUserConfiguration(TestApplication.class)
.withPropertyValues("spring.ai.openai.apiKey=" + System.getenv("OPENAI_API_KEY"));
.withUserConfiguration(TestApplication.class);
@BeforeAll
public static void beforeAll() {
@@ -82,23 +90,22 @@ public class PineconeVectorStoreIT {
@Test
public void addAndSearchTest() {
contextRunner.withConfiguration(AutoConfigurations.of(OpenAiAutoConfiguration.class)).run(context -> {
contextRunner.run(context -> {
VectorStore vectorStore = context.getBean(VectorStore.class);
vectorStore.add(documents);
Awaitility.await().until(() -> {
return vectorStore.similaritySearch("Great", 1);
return vectorStore.similaritySearch(SearchRequest.query("Great Depression").withTopK(1));
}, hasSize(1));
List<Document> results = vectorStore.similaritySearch("Great", 1);
List<Document> results = vectorStore.similaritySearch(SearchRequest.query("Great Depression").withTopK(1));
assertThat(results).hasSize(1);
Document resultDoc = results.get(0);
assertThat(resultDoc.getId()).isEqualTo(documents.get(2).getId());
assertThat(resultDoc.getContent()).isEqualTo(
"Great Depression Great Depression Great Depression Great Depression Great Depression Great Depression");
assertThat(resultDoc.getContent()).contains("The Great Depression (1929–1939) was an economic shock");
assertThat(resultDoc.getMetadata()).hasSize(2);
assertThat(resultDoc.getMetadata()).containsKey("meta2");
assertThat(resultDoc.getMetadata()).containsKey("distance");
@@ -107,7 +114,7 @@ public class PineconeVectorStoreIT {
vectorStore.delete(documents.stream().map(doc -> doc.getId()).toList());
Awaitility.await().until(() -> {
return vectorStore.similaritySearch("Hello", 1);
return vectorStore.similaritySearch(SearchRequest.query("Hello").withTopK(1));
}, hasSize(0));
});
}
@@ -118,9 +125,7 @@ public class PineconeVectorStoreIT {
// Pinecone metadata filtering syntax:
// https://docs.pinecone.io/docs/metadata-filtering
final double THRESHOLD_ALL = 0.0;
contextRunner.withConfiguration(AutoConfigurations.of(OpenAiAutoConfiguration.class)).run(context -> {
contextRunner.run(context -> {
VectorStore vectorStore = context.getBean(VectorStore.class);
@@ -131,18 +136,24 @@ public class PineconeVectorStoreIT {
vectorStore.add(List.of(bgDocument, nlDocument));
SearchRequest searchRequest = SearchRequest.query("The World");
Awaitility.await().until(() -> {
return vectorStore.similaritySearch("The World", 1);
return vectorStore.similaritySearch(searchRequest.withTopK(1));
}, hasSize(1));
List<Document> results = vectorStore.similaritySearch("The World", 5);
List<Document> results = vectorStore.similaritySearch(searchRequest.withTopK(5));
assertThat(results).hasSize(2);
results = vectorStore.similaritySearch("The World", 5, THRESHOLD_ALL, "country == 'Bulgaria'");
results = vectorStore.similaritySearch(searchRequest.withTopK(5)
.withSimilarityThresholdAll()
.withFilterExpression("country == 'Bulgaria'"));
assertThat(results).hasSize(1);
assertThat(results.get(0).getId()).isEqualTo(bgDocument.getId());
results = vectorStore.similaritySearch("The World", 5, THRESHOLD_ALL, "country == 'Netherland'");
results = vectorStore.similaritySearch(searchRequest.withTopK(5)
.withSimilarityThresholdAll()
.withFilterExpression("country == 'Netherland'"));
assertThat(results).hasSize(1);
assertThat(results.get(0).getId()).isEqualTo(nlDocument.getId());
@@ -150,7 +161,7 @@ public class PineconeVectorStoreIT {
vectorStore.delete(List.of(bgDocument, nlDocument).stream().map(doc -> doc.getId()).toList());
Awaitility.await().until(() -> {
return vectorStore.similaritySearch("The World", 1);
return vectorStore.similaritySearch(searchRequest.withTopK(1));
}, hasSize(0));
});
}
@@ -159,7 +170,7 @@ public class PineconeVectorStoreIT {
public void documentUpdateTest() {
// Note ,using OpenAI to calculate embeddings
contextRunner.withConfiguration(AutoConfigurations.of(OpenAiAutoConfiguration.class)).run(context -> {
contextRunner.run(context -> {
VectorStore vectorStore = context.getBean(VectorStore.class);
@@ -168,11 +179,13 @@ public class PineconeVectorStoreIT {
vectorStore.add(List.of(document));
SearchRequest springSearchRequest = SearchRequest.query("Spring").withTopK(5);
Awaitility.await().until(() -> {
return vectorStore.similaritySearch("Spring", 5);
return vectorStore.similaritySearch(springSearchRequest);
}, hasSize(1));
List<Document> results = vectorStore.similaritySearch("Spring", 5);
List<Document> results = vectorStore.similaritySearch(springSearchRequest);
assertThat(results).hasSize(1);
Document resultDoc = results.get(0);
@@ -187,11 +200,13 @@ public class PineconeVectorStoreIT {
vectorStore.add(List.of(sameIdDocument));
SearchRequest fooBarSearchRequest = SearchRequest.query("FooBar").withTopK(5);
Awaitility.await().until(() -> {
return vectorStore.similaritySearch("FooBar", 5).get(0).getContent();
return vectorStore.similaritySearch(fooBarSearchRequest).get(0).getContent();
}, equalTo("The World is Big and Salvation Lurks Around the Corner"));
results = vectorStore.similaritySearch("FooBar", 5);
results = vectorStore.similaritySearch(fooBarSearchRequest);
assertThat(results).hasSize(1);
resultDoc = results.get(0);
@@ -203,7 +218,7 @@ public class PineconeVectorStoreIT {
// Remove all documents from the store
vectorStore.delete(List.of(document.getId()));
Awaitility.await().until(() -> {
return vectorStore.similaritySearch("FooBar", 1);
return vectorStore.similaritySearch(fooBarSearchRequest);
}, hasSize(0));
});
@@ -212,17 +227,19 @@ public class PineconeVectorStoreIT {
@Test
public void searchThresholdTest() {
contextRunner.withConfiguration(AutoConfigurations.of(OpenAiAutoConfiguration.class)).run(context -> {
contextRunner.run(context -> {
VectorStore vectorStore = context.getBean(VectorStore.class);
vectorStore.add(documents);
Awaitility.await().until(() -> {
return vectorStore.similaritySearch("Great", 5);
return vectorStore
.similaritySearch(SearchRequest.query("Depression").withTopK(50).withSimilarityThresholdAll());
}, hasSize(3));
List<Document> fullResult = vectorStore.similaritySearch("Great", 5, 0.0);
List<Document> fullResult = vectorStore
.similaritySearch(SearchRequest.query("Depression").withTopK(5).withSimilarityThresholdAll());
List<Float> distances = fullResult.stream().map(doc -> (Float) doc.getMetadata().get("distance")).toList();
@@ -230,20 +247,20 @@ public class PineconeVectorStoreIT {
float threshold = (distances.get(0) + distances.get(1)) / 2;
List<Document> results = vectorStore.similaritySearch("Great", 5, (1 - threshold));
List<Document> results = vectorStore
.similaritySearch(SearchRequest.query("Depression").withTopK(5).withSimilarityThreshold(1 - threshold));
assertThat(results).hasSize(1);
Document resultDoc = results.get(0);
assertThat(resultDoc.getId()).isEqualTo(documents.get(2).getId());
assertThat(resultDoc.getContent()).isEqualTo(
"Great Depression Great Depression Great Depression Great Depression Great Depression Great Depression");
assertThat(resultDoc.getContent()).contains("The Great Depression (1929–1939) was an economic shock");
assertThat(resultDoc.getMetadata()).containsKey("meta2");
assertThat(resultDoc.getMetadata()).containsKey("distance");
// Remove all documents from the store
vectorStore.delete(documents.stream().map(doc -> doc.getId()).toList());
Awaitility.await().until(() -> {
return vectorStore.similaritySearch("Hello", 1);
return vectorStore.similaritySearch(SearchRequest.query("Hello").withTopK(1));
}, hasSize(0));
});
}
@@ -269,6 +286,11 @@ public class PineconeVectorStoreIT {
return new PineconeVectorStore(config, embeddingClient);
}
@Bean
public EmbeddingClient embeddingClient() {
return new TransformersEmbeddingClient();
}
}
}