diff --git a/.gitignore b/.gitignore index 476f773c8..ca5d03678 100644 --- a/.gitignore +++ b/.gitignore @@ -26,10 +26,11 @@ out /.gradletasknamecache **/*.flattened-pom.xml -vscode +vscode settings.json node node_modules package-lock.json package.json +.vscode \ No newline at end of file diff --git a/README.md b/README.md index dce550567..6913f0b0c 100644 --- a/README.md +++ b/README.md @@ -149,10 +149,10 @@ To build with only unit tests ``` To build including integration tests. -Set API key environment variables for OpenAI and Azure OpenAI before running. +Set API key environment variables for OpenAI and Azure OpenAI before running. ```shell -./mvnw clean package -Pintegration-tests +./mvnw clean verify -Pintegration-tests ``` To build the docs diff --git a/pom.xml b/pom.xml index 8f05ace92..9675749f9 100644 --- a/pom.xml +++ b/pom.xml @@ -19,6 +19,7 @@ spring-ai-spring-boot-starters/spring-ai-starter-openai spring-ai-spring-boot-starters/spring-ai-starter-azure-openai spring-ai-docs + vector-stores/spring-ai-pgvector-store @@ -61,7 +62,7 @@ 17 - + 3.1.2 4.0.2 0.12.0 @@ -127,9 +128,6 @@ ${maven-surefire-plugin.version} ${surefireArgLine} - - **/*IntegrationTests.java - @@ -247,11 +245,6 @@ org.apache.maven.plugins maven-failsafe-plugin ${maven-failsafe-plugin.version} - - - **/*IntegrationTests.java - - diff --git a/spring-ai-core/pom.xml b/spring-ai-core/pom.xml index cb15c5f3e..b2c74e939 100644 --- a/spring-ai-core/pom.xml +++ b/spring-ai-core/pom.xml @@ -1,5 +1,6 @@ - + 4.0.0 org.springframework.experimental.ai @@ -18,6 +19,10 @@ git@github.com:spring-projects-experimental/spring-ai.git + + 4.31.1 + + @@ -46,21 +51,15 @@ com.github.victools jsonschema-generator - 4.31.1 + ${jsonschema.version} com.github.victools jsonschema-module-jackson - 4.31.1 + ${jsonschema.version} - - - - - - org.springframework.boot diff --git a/spring-ai-core/src/main/java/org/springframework/ai/document/Document.java b/spring-ai-core/src/main/java/org/springframework/ai/document/Document.java index ae6b6745d..ce1e89f3c 100644 --- a/spring-ai-core/src/main/java/org/springframework/ai/document/Document.java +++ b/spring-ai-core/src/main/java/org/springframework/ai/document/Document.java @@ -6,14 +6,10 @@ import java.util.*; public class Document { - private static String DEFAULT_TEXT_TEMPLATE = "{metadata_string}\n\n{text}"; - - private static String DEFAULT_METADATA_TEMPLATE = "{key}: {value}"; - /** * Unique ID */ - private String id = UUID.randomUUID().toString(); + private final String id; private List embedding = new ArrayList<>(); @@ -25,28 +21,24 @@ public class Document { // Type; introduce when support images, now only text. + // TODO: Rename to `content` instead. private String text; - private MetadataMode metadataMode = MetadataMode.NONE; - - private List excludedMetadataKeysForEmbedding; - - private List excludedMetadataKeysForLlm; - - private List relatedIds; - - private String textTemplate = DEFAULT_TEXT_TEMPLATE; - - private String metadataTemplate = DEFAULT_METADATA_TEMPLATE; - - private String metadataSeparator = "\n"; - public Document(String text) { + this(UUID.randomUUID().toString(), text); + } + + public Document(String id, String text) { + this.id = id; this.text = text; - this.metadata = metadata; } public Document(String text, Map metadata) { + this(UUID.randomUUID().toString(), text, metadata); + } + + public Document(String id, String text, Map metadata) { + this.id = id; this.text = text; this.metadata = metadata; } @@ -59,6 +51,43 @@ public class Document { return this.text; } + public Map getMetadata() { + return metadata; + } + + public List getEmbedding() { + return embedding; + } + + public void setEmbedding(List embedding) { + this.embedding = embedding; + } + + @Override + public String toString() { + return "Document{" + "id='" + id + '\'' + ", metadata=" + metadata + ", text='" + text + '\'' + '}'; + } + + // TODO: Consider moving the following methods & fields in a seprarate + // dedicated class. (e.g. DocumentService, DocumentUtil or alike)¬ + + // private List excludedMetadataKeysForEmbedding; + // private List relatedIds; + + private static String DEFAULT_TEXT_TEMPLATE = "{metadata_string}\n\n{text}"; + + private static String DEFAULT_METADATA_TEMPLATE = "{key}: {value}"; + + private final String textTemplate = DEFAULT_TEXT_TEMPLATE; + + private final String metadataTemplate = DEFAULT_METADATA_TEMPLATE; + + private final String metadataSeparator = "\n"; + + private MetadataMode metadataMode = MetadataMode.NONE; + + private List excludedMetadataKeysForLlm; + public String getContent() { return getContent(MetadataMode.ALL); } @@ -103,45 +132,28 @@ public class Document { return String.join(getMetadataSeparator(), metadataStringList); } - public String getTextTemplate() { + private String getTextTemplate() { return textTemplate; } - public String getMetadataTemplate() { + private String getMetadataTemplate() { return metadataTemplate; } - public String getMetadataSeparator() { + private String getMetadataSeparator() { return metadataSeparator; } - public Map getMetadata() { - return metadata; - } + // public void setTextTemplate(String textTemplate) { + // this.textTemplate = textTemplate; + // } - public List getEmbedding() { - return embedding; - } + // public void setMetadataTemplate(String metadataTemplate) { + // this.metadataTemplate = metadataTemplate; + // } - public void setEmbedding(List embedding) { - this.embedding = embedding; - } - - @Override - public String toString() { - return "Document{" + "id='" + id + '\'' + ", metadata=" + metadata + ", text='" + text + '\'' + '}'; - } - - public void setTextTemplate(String textTemplate) { - this.textTemplate = textTemplate; - } - - public void setMetadataTemplate(String metadataTemplate) { - this.metadataTemplate = metadataTemplate; - } - - public void setMetadataSeparator(String metadataSeparator) { - this.metadataSeparator = metadataSeparator; - } + // public void setMetadataSeparator(String metadataSeparator) { + // this.metadataSeparator = metadataSeparator; + // } } diff --git a/spring-ai-docs/src/main/antora/modules/ROOT/images/lawofcosines.png b/spring-ai-docs/src/main/antora/modules/ROOT/images/lawofcosines.png new file mode 100644 index 000000000..46ceadd92 Binary files /dev/null and b/spring-ai-docs/src/main/antora/modules/ROOT/images/lawofcosines.png differ diff --git a/spring-ai-docs/src/main/antora/modules/ROOT/images/pythagorean-triangle.png b/spring-ai-docs/src/main/antora/modules/ROOT/images/pythagorean-triangle.png new file mode 100644 index 000000000..4013db1e9 Binary files /dev/null and b/spring-ai-docs/src/main/antora/modules/ROOT/images/pythagorean-triangle.png differ diff --git a/spring-ai-docs/src/main/antora/modules/ROOT/images/vector_2d_coordinates.png b/spring-ai-docs/src/main/antora/modules/ROOT/images/vector_2d_coordinates.png new file mode 100644 index 000000000..9bc501867 Binary files /dev/null and b/spring-ai-docs/src/main/antora/modules/ROOT/images/vector_2d_coordinates.png differ diff --git a/spring-ai-docs/src/main/antora/modules/ROOT/images/vector_similarity.png b/spring-ai-docs/src/main/antora/modules/ROOT/images/vector_similarity.png new file mode 100644 index 000000000..9b8a01bb8 Binary files /dev/null and b/spring-ai-docs/src/main/antora/modules/ROOT/images/vector_similarity.png differ diff --git a/spring-ai-docs/src/main/antora/modules/ROOT/pages/api/vectordbs.adoc b/spring-ai-docs/src/main/antora/modules/ROOT/pages/api/vectordbs.adoc index b29015307..3c2a48959 100644 --- a/spring-ai-docs/src/main/antora/modules/ROOT/pages/api/vectordbs.adoc +++ b/spring-ai-docs/src/main/antora/modules/ROOT/pages/api/vectordbs.adoc @@ -1 +1,202 @@ -= Vector Databases \ No newline at end of file += Vector Databases + +== Introduction +Vector Databases are a specialized type of database that plays an essential role in AI applications. + +In Vector Databases, queries differ from traditional relational databases. +Instead of exact matches, they perform similarity searches. +When given a vector as a query, a Vector Database returns vectors that are "similar" to the query vector. +Further details on how this similarity is calculated at a high level in provided in a later section. + +Vector Databases are used to integrate your data with AI Models. +The first step in their usage is to load your data into a Vector Database. +Then, when a user query is to be sent to the AI model, a set of similar documents is retrieved first. +These documents then serve as the context for the user's question and are sent to the AI model along with the user's query. +This technique is known as Retrieval Augmented Generation. + +In the following sections, we will describe the Spring AI interface for using multiple Vector Database implementations and some high-level sample usage. + +The last section attempts to demystify the underlying approach of similarity search of Vector Databases. + +== Spring AI Vector Database Support + +Spring AI provides an abstract API over Vector Databases using the interface `VectorStore`. + +The API of the Vector store is shown below + +```java +public interface VectorStore { + + void add(List documents); + + Optional delete(List idList); + + List similaritySearch(String query); + + List similaritySearch(String query, int k); + + List similaritySearch(String query, int k, double threshold); +``` + +The Spring AI library uses the `Document` class to model the `String` you want to store in the Vector Database along with the vector representing that `String`. +The vector that representing the `String` is of the type `List` and is often called the string's embedding. + +The `List is computed from the `String` using an implementation of the class `EmbeddingClient`. +An Embedding Model is an example of an AI model, that transforms from `String` to a `List`. +The embeddings are used internally by other AI Model implementations, such as models that convert `text` to `text` like ChatGPT. + +`EmbeddingClient` that computes the embedding is passed as a constructor argument to `VectorStore` implementation. + + +The `VectorStore` implementations supported by Spring AI are: + +* InMemoryVectorStore +* SimplePersistentVectorStore +* PgVector - A Vector Store build on PostgreSQL. + +More are implementations are coming, with Pinecone being the next implementation. + +If you have a Vector Database that needs to be supported by Spring AI, please open an issue on GitHub or, even better, submit a Pull Request with an implementation. + +== Usage + +To compute the embeddings for a Vector Database, you need to pick an Embedding Model that matches the higher-level AI model being used. + +For example, with OpenAI's ChatGPT, we use the `OpenAiEmbeddingClient` and the model name `text-embedding-ada-002`. + +The Spring Boot Starter's autoconfiguation for OpenAI makes an implementation of `EmbeddingClient` available in the Spring Application Context for Dependency Injection. + +The general usage of loading data into a vector store is something you do as a batch-type job, first loading data into Spring AI's `Document` class and then calling the `save` method. + +Given a String `sourceFile` that represents a JSON file with data we want to load into the Vector Database, we use Spring AI's `JsonLoader` to load specific fields in the JSON file, which splits them up into small pieces and then passes those small pieces to the vector store implementation. +The `VectorStore` implementation computes the embeddings and stores the JSON and the embedding in the Vector Database. + +```java + @Autowired + VectorStore vectorStore; + + void load(String sourceFile) {} + JsonLoader jsonLoader = new JsonLoader(new FileSystemResource(sourceFile), + "price", "name", "shortDescription", "description", "tags"); + List documents = jsonLoader.load(); + this.vectorStore.add(documents); + } +``` + +Later, when a user question is to be passed into the AI Model, a similarity search is done to retrieve similar documents, which are then 'stuffed' into the prompt as context for the user's question. + +```java + + String question = + List similarDocuments = store.similaritySearch(question); +``` + +There are additional options to be passed into the `similaritySearch` method that defines how many documents to retrieve and a threshold of the similarity search. + +== Understanding Vectors + +Vectors have dimensionality and a direction. +For example, a picture of a two-dimensional vector stem:[\vec{a}] in the cartesian coordinate system pictured as an arrow. + +image::vector_2d_coordinates.png[] + +The head of the vector stem:[\vec{a}] is at the point stem:[(a_1, a_2)] +The *x* coordinate value is stem:[a_1] and the *y* coordinate value is stem:[a_2] and are also referred to as the components of the vector. + +== Similarity + +Several mathematical formulas can be used to determine if two vectors are similar. + +One of the most intuitive to visualize and understand is cosine similarity. + +Look at the following pictures that show three sets of graphs. + +image::vector_similarity.png[] + +The vectors stem:[\vec{A}] and stem:[\vec{B}] are considered similar, when they are pointing close to each other, as in the first diagram. +The vectors are considered unrelated when pointing perpendicular to each other and opposite when they point away from each other. + +The angle between them, stem:[\theta], is a good measure of their similarity. +How can the angle stem:[\theta] be computed? + +We are all familiar with the https://en.wikipedia.org/wiki/Pythagorean_theorem#History[Pythagorean Theorem] + +image:pythagorean-triangle.png[] + +What about when the angle between *a* and *b* is not 90 degrees? + +Enter the https://en.wikipedia.org/wiki/Law_of_cosines[Law of cosines] + + +.Law of Cosines +**** +stem:[a^2 + b^2 - 2ab\cos\theta = c^2] +**** + +Showing this as a vector diagram + +image:lawofcosines.png[] + + +The magnitude of this vector is defined in terms of its components as: + +.Magnitude +**** +stem:[\vec{A} * \vec{A} = ||\vec{A}||^2 = A_1^2 + A_2^2 ] +**** + +and the dot product between two vectors stem:[\vec{A}] and stem:[\vec{B}] is defined in terms of its components as: + + +.Dot Product +**** +stem:[\vec{A} * \vec{B} = A_1B_1 + A_2B_2] +**** + +Rewriting the Law of Cosines with vector magnitudes and dot products gives: + +.Law of Cosines in Vector form +**** +stem:[||\vec{A}||^2 + ||\vec{B}||^2 - 2||\vec{A}||||\vec{B}||\cos\theta = ||\vec{C}||^2] +**** + + +Replacing stem:[||\vec{C}||^2] with stem:[||\vec{B} - \vec{A}||^2] gives: + +.Law of Cosines in Vector form only in terms of stem:[\vec{A}] and stem:[\vec{B}] + +**** +stem:[||\vec{A}||^2 + ||\vec{B}||^2 - 2||\vec{A}||||\vec{B}||\cos\theta = ||\vec{B} - \vec{A}||^2] +**** + + +https://towardsdatascience.com/cosine-similarity-how-does-it-measure-the-similarity-maths-behind-and-usage-in-python-50ad30aad7db[Expanding this out] gives us the formula for https://en.wikipedia.org/wiki/Cosine_similarity[Cosine Similarity]. + +.Cosine Similarity +**** +stem:[similarity(vec{A},vec{B}) = \cos(\theta) = \frac{\vec{A}\cdot\vec{B}}{||\vec{A}\||\cdot||\vec{B}||] +**** + +This formula works for dimensions higher than 2 or 3, though it is hard to visualize, https://projector.tensorflow.org/[but can be done to some extent]. +It is common for vectors in AI/ML applications to have hundreds or a thousand dimensions. + +The similarity function in higher dimensions using the components of the vector is shown below. +It expands the two-dimensional definitions of Magnitude and Dot Product given previously to *N* dimensions using the https://en.wikipedia.org/wiki/Summation[Summation mathematical syntax]. + +.Cosine Similarity with vector components +**** +stem:[similarity(vec{A},vec{B}) = \cos(\theta) = \frac{ \sum_{i=1}^{n} {A_i B_i} }{ \sqrt{\sum_{i=1}^{n}{A_i^2} \cdot \sum_{i=1}^{n}{B_i^2}}] +**** + +This is the key formula used in the simple implementation of a Vector Store and can be found in the `InMemoryVectorStore` implementation. + + + + + + + + + + + diff --git a/spring-ai-openai/src/test/java/org/springframework/ai/openai/acme/AcmeIntegrationTest.java b/spring-ai-openai/src/test/java/org/springframework/ai/openai/acme/AcmeIT.java similarity index 95% rename from spring-ai-openai/src/test/java/org/springframework/ai/openai/acme/AcmeIntegrationTest.java rename to spring-ai-openai/src/test/java/org/springframework/ai/openai/acme/AcmeIT.java index 6f3e1257a..c762d9aa4 100644 --- a/spring-ai-openai/src/test/java/org/springframework/ai/openai/acme/AcmeIntegrationTest.java +++ b/spring-ai-openai/src/test/java/org/springframework/ai/openai/acme/AcmeIT.java @@ -8,7 +8,7 @@ import org.springframework.ai.client.AiResponse; import org.springframework.ai.document.Document; import org.springframework.ai.loader.impl.JsonLoader; import org.springframework.ai.openai.embedding.OpenAiEmbeddingClient; -import org.springframework.ai.openai.testutils.AbstractIntegrationTest; +import org.springframework.ai.openai.testutils.AbstractIT; import org.springframework.ai.prompt.Prompt; import org.springframework.ai.prompt.SystemPromptTemplate; import org.springframework.ai.prompt.messages.Message; @@ -28,9 +28,9 @@ import java.util.stream.Collectors; import static org.assertj.core.api.Assertions.assertThat; @SpringBootTest -public class AcmeIntegrationTest extends AbstractIntegrationTest { +public class AcmeIT extends AbstractIT { - private static final Logger logger = LoggerFactory.getLogger(AcmeIntegrationTest.class); + private static final Logger logger = LoggerFactory.getLogger(AcmeIT.class); @Value("classpath:/data/acme/bikes.json") private Resource bikesResource; diff --git a/spring-ai-openai/src/test/java/org/springframework/ai/openai/client/ClientIntegrationTests.java b/spring-ai-openai/src/test/java/org/springframework/ai/openai/client/ClientIT.java similarity index 96% rename from spring-ai-openai/src/test/java/org/springframework/ai/openai/client/ClientIntegrationTests.java rename to spring-ai-openai/src/test/java/org/springframework/ai/openai/client/ClientIT.java index 5a07d2ce9..05a15cdbd 100644 --- a/spring-ai-openai/src/test/java/org/springframework/ai/openai/client/ClientIntegrationTests.java +++ b/spring-ai-openai/src/test/java/org/springframework/ai/openai/client/ClientIT.java @@ -3,7 +3,7 @@ package org.springframework.ai.openai.client; import org.junit.jupiter.api.Test; import org.springframework.ai.client.AiResponse; import org.springframework.ai.client.Generation; -import org.springframework.ai.openai.testutils.AbstractIntegrationTest; +import org.springframework.ai.openai.testutils.AbstractIT; import org.springframework.ai.parser.BeanOutputParser; import org.springframework.ai.parser.ListOutputParser; import org.springframework.ai.parser.MapOutputParser; @@ -24,7 +24,7 @@ import java.util.Map; import static org.assertj.core.api.Assertions.assertThat; @SpringBootTest -class ClientIntegrationTests extends AbstractIntegrationTest { +class ClientIT extends AbstractIT { @Value("classpath:/prompts/system-message.st") private Resource systemResource; diff --git a/spring-ai-openai/src/test/java/org/springframework/ai/openai/embedding/EmbeddingIntegrationTest.java b/spring-ai-openai/src/test/java/org/springframework/ai/openai/embedding/EmbeddingIT.java similarity index 97% rename from spring-ai-openai/src/test/java/org/springframework/ai/openai/embedding/EmbeddingIntegrationTest.java rename to spring-ai-openai/src/test/java/org/springframework/ai/openai/embedding/EmbeddingIT.java index afb4165ea..5011134b3 100644 --- a/spring-ai-openai/src/test/java/org/springframework/ai/openai/embedding/EmbeddingIntegrationTest.java +++ b/spring-ai-openai/src/test/java/org/springframework/ai/openai/embedding/EmbeddingIT.java @@ -10,7 +10,7 @@ import java.util.List; import static org.assertj.core.api.Assertions.assertThat; @SpringBootTest -class EmbeddingIntegrationTest { +class EmbeddingIT { @Autowired private OpenAiEmbeddingClient embeddingClient; diff --git a/spring-ai-openai/src/test/java/org/springframework/ai/openai/testutils/AbstractIntegrationTest.java b/spring-ai-openai/src/test/java/org/springframework/ai/openai/testutils/AbstractIT.java similarity index 97% rename from spring-ai-openai/src/test/java/org/springframework/ai/openai/testutils/AbstractIntegrationTest.java rename to spring-ai-openai/src/test/java/org/springframework/ai/openai/testutils/AbstractIT.java index 0d3bd07ff..1c06eb36e 100644 --- a/spring-ai-openai/src/test/java/org/springframework/ai/openai/testutils/AbstractIntegrationTest.java +++ b/spring-ai-openai/src/test/java/org/springframework/ai/openai/testutils/AbstractIT.java @@ -18,9 +18,9 @@ import java.util.Map; import static org.assertj.core.api.Assertions.assertThat; import static org.assertj.core.api.Assertions.fail; -public class AbstractIntegrationTest { +public abstract class AbstractIT { - private static final Logger logger = LoggerFactory.getLogger(AbstractIntegrationTest.class); + private static final Logger logger = LoggerFactory.getLogger(AbstractIT.class); @Autowired protected OpenAiClient openAiClient; diff --git a/vector-stores/spring-ai-pgvector-store/README.md b/vector-stores/spring-ai-pgvector-store/README.md new file mode 100644 index 000000000..05bd0edfc --- /dev/null +++ b/vector-stores/spring-ai-pgvector-store/README.md @@ -0,0 +1,16 @@ + +# PGvector VectorStore + +Pgvector is an open-source extension for PostgreSQL that enables storing and searching over machine learning-generated embeddings. It provides different capabilities that let users identify both exact and approximate nearest neighbors. It is designed to work seamlessly with other PostgreSQL features, including indexing and querying. + +## Start Postgres+PGVecgor DB: + +``` +docker run -it --rm --name postgres -p 5432:5432 -e POSTGRES_USER=postgres -e POSTGRES_PASSWORD=postgres ankane/pgvector +``` + +You can connect to this server like this: + +``` +psql -U postgres -h localhost -p 5432 +``` \ No newline at end of file diff --git a/vector-stores/spring-ai-pgvector-store/pom.xml b/vector-stores/spring-ai-pgvector-store/pom.xml new file mode 100644 index 000000000..b640e5ae4 --- /dev/null +++ b/vector-stores/spring-ai-pgvector-store/pom.xml @@ -0,0 +1,94 @@ + + + 4.0.0 + + org.springframework.experimental.ai + spring-ai + 0.2.0-SNAPSHOT + ../../pom.xml + + spring-ai-pgvector-store + jar + Spring AI Vector Store - pgvector + Spring AI PGVector Vector Store + https://github.com/spring-projects-experimental/spring-ai + + + https://github.com/spring-projects-experimental/spring-ai + git://github.com/spring-projects-experimental/spring-ai.git + git@github.com:spring-projects-experimental/spring-ai.git + + + + 0.1.3 + 42.6.0 + 0.2.0-SNAPSHOT + + 1.19.0 + 4.0.3 + + + + + org.springframework.experimental.ai + spring-ai-core + ${spring-ai.version} + + + + + com.pgvector + pgvector + ${pgvector.version} + + + + org.postgresql + postgresql + ${postgresql.version} + + + + org.springframework + spring-jdbc + + + + + org.springframework.experimental.ai + spring-ai-openai-spring-boot-starter + ${spring-ai.version} + test + + + + org.springframework.boot + spring-boot-starter-test + test + + + + org.testcontainers + testcontainers + ${testcontainers.version} + test + + + + org.testcontainers + junit-jupiter + ${testcontainers.version} + test + + + + com.zaxxer + HikariCP + ${hikari-cp.version} + test + + + + + diff --git a/vector-stores/spring-ai-pgvector-store/src/main/java/org/springframework/ai/vectorstore/PgVectorStore.java b/vector-stores/spring-ai-pgvector-store/src/main/java/org/springframework/ai/vectorstore/PgVectorStore.java new file mode 100644 index 000000000..e3b3a50b0 --- /dev/null +++ b/vector-stores/spring-ai-pgvector-store/src/main/java/org/springframework/ai/vectorstore/PgVectorStore.java @@ -0,0 +1,261 @@ +/* + * Copyright 2023-2023 the original author or authors. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * https://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package org.springframework.ai.vectorstore; + +import java.sql.ResultSet; +import java.sql.SQLException; +import java.util.List; +import java.util.Map; +import java.util.Optional; +import java.util.UUID; +import java.util.stream.Collectors; +import java.util.stream.IntStream; + +import com.fasterxml.jackson.core.JsonProcessingException; +import com.fasterxml.jackson.databind.ObjectMapper; +import com.pgvector.PGvector; +import org.postgresql.util.PGobject; + +import org.springframework.ai.document.Document; +import org.springframework.ai.embedding.EmbeddingClient; +import org.springframework.jdbc.core.JdbcTemplate; +import org.springframework.jdbc.core.RowMapper; + +/** + * @author Christian Tzolov + */ +public class PgVectorStore implements VectorStore { + + public static final int OPENAI_EMBEDDING_DIMENSION_SIZE = 1536; + + private final JdbcTemplate jdbcTemplate; + + private final EmbeddingClient embeddingClient; + + private int dimensions; + + private PgDistanceType distanceType; + + private ObjectMapper objectMapper = new ObjectMapper(); + + /** + * By default, pgvector performs exact nearest neighbor search, which provides perfect + * recall. You can add an index to use approximate nearest neighbor search, which + * trades some recall for speed. Unlike typical indexes, you will see different + * results for queries after adding an approximate index. + */ + public enum PgIndexType { + + /** + * Performs exact nearest neighbor search, which provides perfect recall. + */ + NONE, + /** + * An IVFFlat index divides vectors into lists, and then searches a subset of + * those lists that are closest to the query vector. It has faster build times and + * uses less memory than HNSW, but has lower query performance (in terms of + * speed-recall tradeoff). + */ + IVFFLAT, + /** + * An HNSW index creates a multilayer graph. It has slower build times and uses + * more memory than IVFFlat, but has better query performance (in terms of + * speed-recall tradeoff). There’s no training step like IVFFlat, so the index can + * be created without any data in the table. + */ + HNSW; + + } + + public enum PgDistanceType { + + EuclideanDistance("<->", "vector_l2_ops"), + + NegativeInnerProduct("<#>", "vector_ip_ops"), + + CosineDistance("<=>", "vector_cosine_ops"); + + public final String operator; + + public final String index; + + PgDistanceType(String operator, String index) { + this.operator = operator; + this.index = index; + } + + } + + private static class DocumentRowMapper implements RowMapper { + + private static final String COLUMN_EMBEDDING = "embedding"; + + private static final String COLUMN_METADATA = "metadata"; + + private static final String COLUMN_ID = "id"; + + private static final String COLUMN_CONTENT = "content"; + + private ObjectMapper objectMapper; + + public DocumentRowMapper(ObjectMapper objectMapper) { + this.objectMapper = objectMapper; + } + + @Override + public Document mapRow(ResultSet rs, int rowNum) throws SQLException { + String id = rs.getString(COLUMN_ID); + String content = rs.getString(COLUMN_CONTENT); + PGobject metadata = rs.getObject(COLUMN_METADATA, PGobject.class); + PGobject embedding = rs.getObject(COLUMN_EMBEDDING, PGobject.class); + + Document document = new Document(id, content, toMap(metadata)); + document.setEmbedding(toDoubleList(embedding)); + + return document; + } + + private List toDoubleList(PGobject embedding) throws SQLException { + float[] floatArray = new PGvector(embedding.getValue()).toArray(); + List doubleEmbedding = IntStream.range(0, floatArray.length) + .mapToDouble(i -> floatArray[i]) + .boxed() + .collect(Collectors.toList()); + return doubleEmbedding; + + } + + private Map toMap(PGobject pgObject) { + + String source = pgObject.getValue(); + try { + return (Map) objectMapper.readValue(source, Map.class); + } + catch (JsonProcessingException e) { + throw new RuntimeException(e); + } + } + + } + + public PgVectorStore(JdbcTemplate jdbcTemplate, EmbeddingClient embeddingClient) { + this(jdbcTemplate, embeddingClient, OPENAI_EMBEDDING_DIMENSION_SIZE, + PgVectorStore.PgDistanceType.CosineDistance, false, PgIndexType.NONE); + } + + public PgVectorStore(JdbcTemplate jdbcTemplate, EmbeddingClient embeddingClient, int dimensions, + PgDistanceType distanceType, boolean removeExistingVectorStoreTable, PgIndexType createIndexMethod) { + this.jdbcTemplate = jdbcTemplate; + this.embeddingClient = embeddingClient; + this.dimensions = dimensions; + this.distanceType = distanceType; + + // Add PGVector support. + this.jdbcTemplate.execute("CREATE EXTENSION IF NOT EXISTS vector"); + // Add JSONB support. + this.jdbcTemplate.execute("CREATE EXTENSION IF NOT EXISTS hstore"); + // Add UUID support. + this.jdbcTemplate.execute("CREATE EXTENSION IF NOT EXISTS \"uuid-ossp\""); + + // Remove existing VectorStoreTable + if (removeExistingVectorStoreTable) { + this.jdbcTemplate.execute("DROP TABLE IF EXISTS vector_store"); + } + + // TODO: we create id of type UUID, while the Document's id is String!!! + // TODO: remove injection! + this.jdbcTemplate + .execute("CREATE TABLE IF NOT EXISTS vector_store ( " + "id uuid DEFAULT uuid_generate_v4 () PRIMARY KEY, " + + "content text, " + "metadata json, " + "embedding vector(" + this.dimensions + "))"); + + if (createIndexMethod != PgIndexType.NONE) { + this.jdbcTemplate.execute("CREATE INDEX ON vector_store USING " + createIndexMethod.name() + " (embedding " + + distanceType.index + ")"); + } + + this.jdbcTemplate + .execute("CREATE TABLE IF NOT EXISTS vector_store ( " + "id uuid DEFAULT uuid_generate_v4 () PRIMARY KEY, " + + "content text, " + "metadata json, " + "embedding vector(" + this.dimensions + "))"); + } + + @Override + public void add(List documents) { + for (Document document : documents) { + List embedding = this.embeddingClient.embed(document); + document.setEmbedding(embedding); + + UUID id = UUID.fromString(document.getId()); + String content = document.getText(); // TODO: shall we use the text of text + + // metadata? + Map metadata = document.getMetadata(); + PGvector pgEmbedding = new PGvector(toFloatArray(embedding)); + + jdbcTemplate.update( + "INSERT INTO vector_store (id, content, metadata, embedding) VALUES (?, ?, ?::jsonb, ?) " + + "ON CONFLICT (id) DO " + "UPDATE SET content = ? , metadata = ?::jsonb , embedding = ? ", + id, content, metadata, pgEmbedding, content, metadata, pgEmbedding); + } + } + + private float[] toFloatArray(List embeddingDouble) { + float[] embeddingFloat = new float[embeddingDouble.size()]; + int i = 0; + for (Double d : embeddingDouble) { + embeddingFloat[i++] = d.floatValue(); + } + return embeddingFloat; + } + + @Override + public Optional delete(List idList) { + int updateCount = 0; + for (String id : idList) { + int count = jdbcTemplate.update("DELETE FROM vector_store WHERE id = ?", UUID.fromString(id)); + + updateCount = updateCount + count; + } + + return Optional.of(updateCount == idList.size()); + } + + @Override + public List similaritySearch(String query) { + return this.similaritySearch(query, 4); + } + + @Override + public List similaritySearch(String query, int k) { + PGvector queryEmbedding = getQueryEmbedding(query); + return this.jdbcTemplate.query( + "SELECT * FROM vector_store ORDER BY embedding " + this.distanceType.operator + " ? LIMIT ?", + new DocumentRowMapper(this.objectMapper), queryEmbedding, k); + } + + @Override + public List similaritySearch(String query, int k, double threshold) { + PGvector queryEmbedding = getQueryEmbedding(query); + return this.jdbcTemplate.query( + "SELECT * FROM vector_store ORDER BY embedding " + this.distanceType.operator + " ? < ? LIMIT ? ", + new DocumentRowMapper(this.objectMapper), queryEmbedding, threshold, k); + } + + private PGvector getQueryEmbedding(String query) { + List embedding = this.embeddingClient.embed(query); + return new PGvector(toFloatArray(embedding)); + } + +} \ No newline at end of file diff --git a/vector-stores/spring-ai-pgvector-store/src/test/java/org/springframework/ai/vectorstore/PgVectorStoreIT.java b/vector-stores/spring-ai-pgvector-store/src/test/java/org/springframework/ai/vectorstore/PgVectorStoreIT.java new file mode 100644 index 000000000..03453cfc3 --- /dev/null +++ b/vector-stores/spring-ai-pgvector-store/src/test/java/org/springframework/ai/vectorstore/PgVectorStoreIT.java @@ -0,0 +1,173 @@ +/* + * Copyright 2023-2023 the original author or authors. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * https://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package org.springframework.ai.vectorstore; + +import java.util.Collections; +import java.util.List; +import java.util.UUID; +import java.util.stream.Collectors; + +import javax.sql.DataSource; + +import com.zaxxer.hikari.HikariDataSource; +import org.junit.jupiter.api.Test; +import org.testcontainers.containers.GenericContainer; +import org.testcontainers.junit.jupiter.Container; +import org.testcontainers.junit.jupiter.Testcontainers; + +import org.springframework.ai.autoconfigure.openai.OpenAiAutoConfiguration; +import org.springframework.ai.document.Document; +import org.springframework.ai.embedding.EmbeddingClient; +import org.springframework.ai.vectorstore.PgVectorStore.PgIndexType; +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; +import org.springframework.boot.context.properties.ConfigurationProperties; +import org.springframework.boot.test.context.runner.ApplicationContextRunner; +import org.springframework.context.annotation.Bean; +import org.springframework.context.annotation.Primary; +import org.springframework.jdbc.core.JdbcTemplate; + +import static org.assertj.core.api.Assertions.assertThat; + +/** + * @author Christian Tzolov + */ +@Testcontainers +public class PgVectorStoreIT { + + @Container + static GenericContainer postgresContainer = new GenericContainer<>("ankane/pgvector") + .withEnv("POSTGRES_USER", "postgres") + .withEnv("POSTGRES_PASSWORD", "postgres") + .withExposedPorts(5432); + + private final ApplicationContextRunner contextRunner = new ApplicationContextRunner() + .withUserConfiguration(TestApplication.class) + .withPropertyValues("spring.datasource.type=com.zaxxer.hikari.HikariDataSource", + "spring.ai.openai.apiKey=" + System.getenv("SPRING_AI_OPENAI_API_KEY"), + + // JdbcTemplate configuration + String.format("app.datasource.url=jdbc:postgresql://localhost:%d/%s", + postgresContainer.getMappedPort(5432), "postgres"), + "app.datasource.username=postgres", "app.datasource.password=postgres", + "app.datasource.type=com.zaxxer.hikari.HikariDataSource"); + + @Test + public void vectorStoreTest() { + contextRunner.withConfiguration(AutoConfigurations.of(OpenAiAutoConfiguration.class)).run(context -> { + + VectorStore vectorStore = context.getBean(VectorStore.class); + + List 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"))); + + vectorStore.add(documents); + + List results = vectorStore.similaritySearch("Great", 1); + + assertThat(results).hasSize(1); + Document resultDoc = results.get(0); + assertThat(resultDoc.getId()).isEqualTo(documents.get(2).getId()); + assertThat(resultDoc.getText()).isEqualTo( + "Great Depression Great Depression Great Depression Great Depression Great Depression Great Depression"); + assertThat(resultDoc.getMetadata()).isEqualTo(Collections.singletonMap("meta2", "meta2")); + + // Remove all documents from the store + vectorStore.delete(documents.stream().map(doc -> doc.getId()).collect(Collectors.toList())); + + List results2 = vectorStore.similaritySearch("Great", 1); + assertThat(results2).hasSize(0); + + }); + } + + @Test + public void documentUpdateTest() { + + contextRunner.withConfiguration(AutoConfigurations.of(OpenAiAutoConfiguration.class)).run(context -> { + + VectorStore vectorStore = context.getBean(VectorStore.class); + + Document document = new Document(UUID.randomUUID().toString(), "Spring AI rocks!!", + Collections.singletonMap("meta1", "meta1")); + + vectorStore.add(List.of(document)); + + List results = vectorStore.similaritySearch("Spring", 5); + + assertThat(results).hasSize(1); + Document resultDoc = results.get(0); + assertThat(resultDoc.getId()).isEqualTo(document.getId()); + assertThat(resultDoc.getText()).isEqualTo("Spring AI rocks!!"); + assertThat(resultDoc.getMetadata()).isEqualTo(Collections.singletonMap("meta1", "meta1")); + + Document sameIdDocument = new Document(document.getId(), + "The World is Big and Salvation Lurks Around the Corner", + Collections.singletonMap("meta2", "meta2")); + + vectorStore.add(List.of(sameIdDocument)); + + results = vectorStore.similaritySearch("FooBar", 5); + + assertThat(results).hasSize(1); + resultDoc = results.get(0); + assertThat(resultDoc.getId()).isEqualTo(document.getId()); + assertThat(resultDoc.getText()).isEqualTo("The World is Big and Salvation Lurks Around the Corner"); + assertThat(resultDoc.getMetadata()).isEqualTo(Collections.singletonMap("meta2", "meta2")); + + }); + } + + @SpringBootConfiguration + @EnableAutoConfiguration(exclude = { DataSourceAutoConfiguration.class }) + public static class TestApplication { + + @Bean + public VectorStore vectorStore(JdbcTemplate jdbcTemplate, EmbeddingClient embeddingClient) { + return new PgVectorStore(jdbcTemplate, embeddingClient, 1536, PgVectorStore.PgDistanceType.CosineDistance, + true, PgIndexType.HNSW); + } + + @Bean + public JdbcTemplate myJdbcTemplate(DataSource dataSource) { + return new JdbcTemplate(dataSource); + } + + @Bean + @Primary + @ConfigurationProperties("app.datasource") + public DataSourceProperties dataSourceProperties() { + return new DataSourceProperties(); + } + + @Bean + public HikariDataSource dataSource(DataSourceProperties dataSourceProperties) { + return dataSourceProperties.initializeDataSourceBuilder().type(HikariDataSource.class).build(); + } + + } + +}