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

Merge SimplePersistentVectorStore and InMemoryVector store to be SimpleVectorStore

Browse Source 
- Doc updates

Fixes  #146
This commit is contained in:
Omkar Shetkar
2023-12-16 13:29:12 +05:30
committed by Mark Pollack
parent 9ab857a8ab
commit dd64a50d2f
8 changed files with 267 additions and 282 deletions
Download Patch File Download Diff File Expand all files Collapse all files

4
models/spring-ai-openai/src/test/java/org/springframework/ai/openai/acme/AcmeIT.java
View File

@@ -22,7 +22,7 @@ import org.springframework.ai.prompt.messages.UserMessage;
import org.springframework.ai.reader.JsonReader;
import org.springframework.ai.retriever.VectorStoreRetriever;
import org.springframework.ai.transformer.splitter.TokenTextSplitter;
import org.springframework.ai.vectorstore.InMemoryVectorStore;
import org.springframework.ai.vectorstore.SimpleVectorStore;
import org.springframework.ai.vectorstore.VectorStore;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.beans.factory.annotation.Value;
@@ -67,7 +67,7 @@ public class AcmeIT extends AbstractIT {
// Step 2 - Create embeddings and save to vector store
logger.info("Creating Embeddings...");
VectorStore vectorStore = new InMemoryVectorStore(embeddingClient);
VectorStore vectorStore = new SimpleVectorStore(embeddingClient);
vectorStore.accept(textSplitter.apply(jsonReader.get()));

6
models/spring-ai-openai/src/test/java/org/springframework/ai/openai/vectorstore/SimplePersistentVectorStoreIT.java
View File

@@ -6,7 +6,7 @@ import org.junit.jupiter.api.io.TempDir;
import org.springframework.ai.document.Document;
import org.springframework.ai.embedding.EmbeddingClient;
import org.springframework.ai.reader.JsonReader;
import org.springframework.ai.vectorstore.SimplePersistentVectorStore;
import org.springframework.ai.vectorstore.SimpleVectorStore;
import org.springframework.ai.reader.JsonMetadataGenerator;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.beans.factory.annotation.Value;
@@ -34,14 +34,14 @@ public class SimplePersistentVectorStoreIT {
JsonReader jsonReader = new JsonReader(bikesJsonResource, new ProductMetadataGenerator(), "price", "name",
"shortDescription", "description", "tags");
List<Document> documents = jsonReader.get();
SimplePersistentVectorStore vectorStore = new SimplePersistentVectorStore(this.embeddingClient);
SimpleVectorStore vectorStore = new SimpleVectorStore(this.embeddingClient);
vectorStore.add(documents);
File tempFile = new File(workingDir.toFile(), "temp.txt");
vectorStore.save(tempFile);
assertThat(tempFile).isNotEmpty();
assertThat(tempFile).content().contains("Velo 99 XR1 AXS");
SimplePersistentVectorStore vectorStore2 = new SimplePersistentVectorStore(this.embeddingClient);
SimpleVectorStore vectorStore2 = new SimpleVectorStore(this.embeddingClient);
vectorStore2.load(tempFile);
List<Document> similaritySearch = vectorStore2.similaritySearch("Velo 99 XR1 AXS");

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

@@ -1,124 +0,0 @@
package org.springframework.ai.vectorstore;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import org.springframework.ai.document.Document;
import org.springframework.ai.embedding.EmbeddingClient;
import java.util.*;
import java.util.concurrent.ConcurrentHashMap;
/***
* @author Raphael Yu
* @author Dingmeng Xue
* @author Mark Pollack
* @author Christian Tzolov
*/
public class InMemoryVectorStore implements VectorStore {
private static final Logger logger = LoggerFactory.getLogger(InMemoryVectorStore.class);
protected Map<String, Document> store = new ConcurrentHashMap<>();
protected EmbeddingClient embeddingClient;
public InMemoryVectorStore(EmbeddingClient embeddingClient) {
Objects.requireNonNull(embeddingClient, "EmbeddingClient must not be null");
this.embeddingClient = embeddingClient;
}
@Override
public void add(List<Document> documents) {
for (Document document : documents) {
logger.info("Calling EmbeddingClient for document id = " + document.getId());
List<Double> embedding = this.embeddingClient.embed(document);
document.setEmbedding(embedding);
this.store.put(document.getId(), document);
}
}
@Override
public Optional<Boolean> delete(List<String> idList) {
for (String id : idList) {
this.store.remove(id);
}
return Optional.of(true);
}
@Override
public List<Document> similaritySearch(SearchRequest request) {
if (request.getFilterExpression() != null) {
throw new UnsupportedOperationException(
"The [" + this.getClass() + "] doesn't support metadata filtering!");
}
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(request.getTopK())
.map(s -> this.store.get(s.key))
.toList();
return similarities;
}
private List<Double> getUserQueryEmbedding(String query) {
List<Double> userQueryEmbedding = this.embeddingClient.embed(query);
return userQueryEmbedding;
}
public static class Similarity {
private String key;
private double similarity;
public Similarity(String key, double similarity) {
this.key = key;
this.similarity = similarity;
}
}
public class EmbeddingMath {
public static double cosineSimilarity(List<Double> vectorX, List<Double> vectorY) {
if (vectorX.size() != vectorY.size()) {
throw new IllegalArgumentException("Vectors lengths must be equal");
}
double dotProduct = dotProduct(vectorX, vectorY);
double normX = norm(vectorX);
double normY = norm(vectorY);
if (normX == 0 || normY == 0) {
throw new IllegalArgumentException("Vectors cannot have zero norm");
}
return dotProduct / (Math.sqrt(normX) * Math.sqrt(normY));
}
public static double dotProduct(List<Double> vectorX, List<Double> vectorY) {
if (vectorX.size() != vectorY.size()) {
throw new IllegalArgumentException("Vectors lengths must be equal");
}
double result = 0;
for (int i = 0; i < vectorX.size(); ++i) {
result += vectorX.get(i) * vectorY.get(i);
}
return result;
}
public static double norm(List<Double> vector) {
return dotProduct(vector, vector);
}
}
}

96
spring-ai-core/src/main/java/org/springframework/ai/vectorstore/SimplePersistentVectorStore.java
View File

@@ -1,96 +0,0 @@
package org.springframework.ai.vectorstore;
import com.fasterxml.jackson.core.JsonProcessingException;
import com.fasterxml.jackson.core.type.TypeReference;
import com.fasterxml.jackson.databind.ObjectMapper;
import com.fasterxml.jackson.databind.ObjectWriter;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import org.springframework.ai.document.Document;
import org.springframework.ai.embedding.EmbeddingClient;
import org.springframework.core.io.Resource;
import org.springframework.util.StreamUtils;
import java.io.File;
import java.io.FileOutputStream;
import java.io.IOException;
import java.io.OutputStream;
import java.nio.charset.Charset;
import java.util.HashMap;
import java.util.Map;
/**
* Adds simple serialization/deserialization to the data stored in the InMemoryVectorStore
*/
public class SimplePersistentVectorStore extends InMemoryVectorStore {
private static final Logger logger = LoggerFactory.getLogger(SimplePersistentVectorStore.class);
public SimplePersistentVectorStore(EmbeddingClient embeddingClient) {
super(embeddingClient);
}
public void save(File file) {
String json = getVectorDbAsJson();
try {
if (!file.exists()) {
logger.info("Creating new vector store file: " + file);
file.createNewFile();
}
else {
logger.info("Replacing existing vector store file: " + file);
file.delete();
file.createNewFile();
}
}
catch (IOException ex) {
throw new RuntimeException(ex);
}
try (OutputStream stream = new FileOutputStream(file)) {
StreamUtils.copy(json, Charset.forName("UTF-8"), stream);
}
catch (IOException e) {
throw new RuntimeException(e);
}
}
public void load(File file) {
TypeReference<HashMap<String, Document>> typeRef = new TypeReference<>() {
};
ObjectMapper objectMapper = new ObjectMapper();
try {
Map<String, Document> deserializedMap = objectMapper.readValue(file, typeRef);
this.store = deserializedMap;
}
catch (IOException ex) {
throw new RuntimeException(ex);
}
}
public void load(Resource resource) {
TypeReference<HashMap<String, Document>> typeRef = new TypeReference<>() {
};
ObjectMapper objectMapper = new ObjectMapper();
try {
Map<String, Document> deserializedMap = objectMapper.readValue(resource.getInputStream(), typeRef);
this.store = deserializedMap;
}
catch (IOException ex) {
throw new RuntimeException(ex);
}
}
private String getVectorDbAsJson() {
ObjectMapper objectMapper = new ObjectMapper();
ObjectWriter objectWriter = objectMapper.writerWithDefaultPrettyPrinter();
String json;
try {
json = objectWriter.writeValueAsString(this.store);
}
catch (JsonProcessingException e) {
throw new RuntimeException("Error serializing documentMap to JSON.", e);
}
return json;
}
}

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

@@ -0,0 +1,225 @@
package org.springframework.ai.vectorstore;
import com.fasterxml.jackson.core.JsonProcessingException;
import com.fasterxml.jackson.core.type.TypeReference;
import com.fasterxml.jackson.databind.ObjectMapper;
import com.fasterxml.jackson.databind.ObjectWriter;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import org.springframework.ai.document.Document;
import org.springframework.ai.embedding.EmbeddingClient;
import org.springframework.core.io.Resource;
import java.io.*;
import java.nio.charset.StandardCharsets;
import java.util.*;
import java.util.concurrent.ConcurrentHashMap;
/**
* SimpleVectorStore is a simple implementation of the VectorStore interface.
*
* It also provides methods to save the current state of the vectors to a file, and to
* load vectors from a file.
*
* For a deeper understanding of the mathematical concepts and computations involved in
* calculating similarity scores among vectors, refer to this
* [resource](https://docs.spring.io/spring-ai/reference/api/vectordbs.html#_understanding_vectors).
*
* @author Raphael Yu
* @author Dingmeng Xue
* @author Mark Pollack
* @author Christian Tzolov
*/
public class SimpleVectorStore implements VectorStore {
private static final Logger logger = LoggerFactory.getLogger(SimpleVectorStore.class);
protected Map<String, Document> store = new ConcurrentHashMap<>();
protected EmbeddingClient embeddingClient;
public SimpleVectorStore(EmbeddingClient embeddingClient) {
Objects.requireNonNull(embeddingClient, "EmbeddingClient must not be null");
this.embeddingClient = embeddingClient;
}
@Override
public void add(List<Document> documents) {
for (Document document : documents) {
logger.info("Calling EmbeddingClient for document id = {}", document.getId());
List<Double> embedding = this.embeddingClient.embed(document);
document.setEmbedding(embedding);
this.store.put(document.getId(), document);
}
}
@Override
public Optional<Boolean> delete(List<String> idList) {
for (String id : idList) {
this.store.remove(id);
}
return Optional.of(true);
}
@Override
public List<Document> similaritySearch(SearchRequest request) {
if (request.getFilterExpression() != null) {
throw new UnsupportedOperationException(
"The [" + this.getClass() + "] doesn't support metadata filtering!");
}
List<Double> userQueryEmbedding = getUserQueryEmbedding(request.getQuery());
return this.store.values()
.stream()
.map(entry -> new Similarity(entry.getId(),
EmbeddingMath.cosineSimilarity(userQueryEmbedding, entry.getEmbedding())))
.filter(s -> s.score >= request.getSimilarityThreshold())
.sorted(Comparator.<Similarity>comparingDouble(s -> s.score).reversed())
.limit(request.getTopK())
.map(s -> this.store.get(s.key))
.toList();
}
/**
* Serialize the vector store content into a file in JSON format.
* @param file the file to save the vector store content
*/
public void save(File file) {
String json = getVectorDbAsJson();
try {
if (!file.exists()) {
logger.info("Creating new vector store file: {}", file);
file.createNewFile();
}
else {
logger.info("Overwriting existing vector store file: {}", file);
}
try (OutputStream stream = new FileOutputStream(file);
Writer writer = new OutputStreamWriter(stream, StandardCharsets.UTF_8)) {
writer.write(json);
writer.flush();
}
}
catch (IOException ex) {
logger.error("IOException occurred while saving vector store file.", ex);
throw new RuntimeException(ex);
}
catch (SecurityException ex) {
logger.error("SecurityException occurred while saving vector store file.", ex);
throw new RuntimeException(ex);
}
catch (NullPointerException ex) {
logger.error("NullPointerException occurred while saving vector store file.", ex);
throw new RuntimeException(ex);
}
}
/**
* Deserialize the vector store content from a file in JSON format into memory.
* @param file the file to load the vector store content
*/
public void load(File file) {
TypeReference<HashMap<String, Document>> typeRef = new TypeReference<>() {
};
ObjectMapper objectMapper = new ObjectMapper();
try {
Map<String, Document> deserializedMap = objectMapper.readValue(file, typeRef);
this.store = deserializedMap;
}
catch (IOException ex) {
throw new RuntimeException(ex);
}
}
/**
* Deserialize the vector store content from a resource in JSON format into memory.
* @param resource the resource to load the vector store content
*/
public void load(Resource resource) {
TypeReference<HashMap<String, Document>> typeRef = new TypeReference<>() {
};
ObjectMapper objectMapper = new ObjectMapper();
try {
Map<String, Document> deserializedMap = objectMapper.readValue(resource.getInputStream(), typeRef);
this.store = deserializedMap;
}
catch (IOException ex) {
throw new RuntimeException(ex);
}
}
private String getVectorDbAsJson() {
ObjectMapper objectMapper = new ObjectMapper();
ObjectWriter objectWriter = objectMapper.writerWithDefaultPrettyPrinter();
String json;
try {
json = objectWriter.writeValueAsString(this.store);
}
catch (JsonProcessingException e) {
throw new RuntimeException("Error serializing documentMap to JSON.", e);
}
return json;
}
private List<Double> getUserQueryEmbedding(String query) {
return this.embeddingClient.embed(query);
}
public static class Similarity {
private String key;
private double score;
public Similarity(String key, double score) {
this.key = key;
this.score = score;
}
}
public class EmbeddingMath {
private EmbeddingMath() {
throw new UnsupportedOperationException("This is a utility class and cannot be instantiated");
}
public static double cosineSimilarity(List<Double> vectorX, List<Double> vectorY) {
if (vectorX == null || vectorY == null) {
throw new RuntimeException("Vectors must not be null");
}
if (vectorX.size() != vectorY.size()) {
throw new IllegalArgumentException("Vectors lengths must be equal");
}
double dotProduct = dotProduct(vectorX, vectorY);
double normX = norm(vectorX);
double normY = norm(vectorY);
if (normX == 0 || normY == 0) {
throw new IllegalArgumentException("Vectors cannot have zero norm");
}
return dotProduct / (Math.sqrt(normX) * Math.sqrt(normY));
}
public static double dotProduct(List<Double> vectorX, List<Double> vectorY) {
if (vectorX.size() != vectorY.size()) {
throw new IllegalArgumentException("Vectors lengths must be equal");
}
double result = 0;
for (int i = 0; i < vectorX.size(); ++i) {
result += vectorX.get(i) * vectorY.get(i);
}
return result;
}
public static double norm(List<Double> vector) {
return dotProduct(vector, vector);
}
}
}

77
spring-ai-docs/src/main/antora/modules/ROOT/pages/api/etl-pipeline.adoc
View File

@@ -11,49 +11,37 @@ The RAG use-case is designed to augment the capabilities of generative models by
=== DocumentReader
Provides a source of documents from diverse origins.
```java
public interface DocumentReader extends Supplier<List<Document>> {
}
```
==== Available Implementations
*JsonReader*: Parses documents in JSON format.
*TextReader*: Processes plain text documents.
*PagePdfDocumentReader*: Uses Apache PdfBox library to parse PDF documents
*ParagraphPdfDocumentReader*: Uses the PDF catalog (e.g. TOC) information to split the input PDF into text paragraphs and output a single `Document` per paragraph.
*TikaDocumentReader*: Uses Apache Tika to extract text from a variety of document
* formats, such as PDF, DOC/DOCX, PPT/PPTX, and HTML. For a comprehensive list of supported formats, refer to the https://tika.apache.org/2.9.0/formats.html[Tika documentation].
=== DocumentTransformer
Transforms a batch of documents as part of the processing workflow.
```java
public interface DocumentTransformer extends Function<List<Document>, List<Document>> {
}
```
=== DocumentWriter
```java
public interface DocumentWriter extends Consumer<List<Document>> {
}
```
=== Available Implementations
==== DocumentReader Interface
*Supplier<List<Document>>*::
+ Provides a source of documents from diverse origins.
*JsonReader*::
+ Parses documents in JSON format.
*TextReader*::
+ Processes plain text documents.
*Document*::
+ Represents the core data structure manipulated throughout the pipeline.
=== DocumentTransformer Interface
*Function<List<Document>, List<Document>>*::
+ Transforms a batch of documents as part of the processing workflow.
==== Available Implementations
*TextSplitter*::
+ Divides documents to fit the AI model's context window.
@@ -70,34 +58,21 @@ public interface DocumentWriter extends Consumer<List<Document>> {
*SummaryMetadataEnricher*::
+ Enriches documents with summarization metadata for enhanced retrieval.
=== DocumentWriter Interface
=== DocumentWriter
*Consumer<List<Document>>*::
+ Manages the final stage of the ETL process, preparing documents for storage.
Manages the final stage of the ETL process, preparing documents for storage.
*VectorStore*::
+ The abstracted interface for vector database interactions.
```java
public interface DocumentWriter extends Consumer<List<Document>> {
*MilvusVectorStore*::
+ An implementation for the Milvus vector database.
}
```
*PgVectorStore*::
+ Provides vector storage capabilities using PostgreSQL.
== Available Implementations
*SimplePersistentVectorStore*::
+ A straightforward approach to persistent vector storage.
There is an implementation for each of the Vector Stores that Spring AI supports, e.g. `PineconeVectorStore`.
*InMemoryVectorStore*::
+ Enables rapid access with in-memory storage solutions.
*Neo4jVectorStore*::
+ Leverages the Neo4j graph database for vector storage.
*RedisVectorStore*::
+ Provides vector storage capabilities using Redis.
== Using PDF Reader
See xref:api/vectordbs.adoc[Vector DB Documentation] for a full listing.
== Using PagePdfDocumentReader

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

@@ -87,14 +87,15 @@ country == 'UK' && year >= 2020 && isActive == true.
These are the available implementations of the `VectorStore` interface:
* `InMemoryVectorStore` and `SimplePersistentVectorStore`.
* Pinecone: https://www.pinecone.io/[PineCone] vector store.
* PgVector [`PgVectorStore`]: The https://github.com/pgvector/pgvector[PostgreSQL/PGVector] vector store.
* Azure Vector Search [`AzureVectorStore`] the https://learn.microsoft.com/en-us/azure/search/vector-search-overview[Azure] vector store
* Chroma [`ChromaVectorStore`]: https://www.trychroma.com/[Chroma] vector store.
* Milvus [`MilvusVectorStore`]: The https://milvus.io/[Milvus] vector store
* Neo4j [`Neo4jVectorStore`]: The https://neo4j.com/[Neo4j] vector store
* PgVector [`PgVectorStore`]: The https://github.com/pgvector/pgvector[PostgreSQL/PGVector] vector store.
* Pinecone: https://www.pinecone.io/[PineCone] vector store.
* Redis [`RedisVectorStore`]: The https://redis.io/[Redis] vector store
* Simple Vector Store [`SimpleVectorStore`]: A simple implementation of persistent vector storage, good for educational purposes.
* Weaviate [`WeaviateVectorStore`] The https://weaviate.io/[Weaviate] vector store
* Azure Vector Search [`AzureVectorStore`] the https://learn.microsoft.com/en-us/azure/search/vector-search-overview[Azure] vector store
* Redisj [`RedisVectorStore`]: The https://redis.io/[Redis] vector store
More implementations may be supported in future releases.

6
vector-stores/spring-ai-chroma/src/main/java/org/springframework/ai/vectorsore/ChromaVectorStore.java
View File

@@ -38,7 +38,11 @@ import org.springframework.util.CollectionUtils;
import org.springframework.util.StringUtils;
/**
* @author Christian Tzolov
* {@link ChromaVectorStore} is a concrete implementation of the {@link VectorStore}
* interface. It is responsible for adding, deleting, and searching documents based on
* their similarity to a query, using the {@link ChromaApi} and {@link EmbeddingClient}
* for embedding calculations. For more information about how it does this, see the
* official <a href="https://www.trychroma.com/">Chroma website</a>.
*/
public class ChromaVectorStore implements VectorStore, InitializingBean {