Add Pinecone VectorStore

- Based on the official pinecone java library. Later expects that indices are created externally via Ops. - Map Document metadata to and from Pinecone's internal Struct. Later converts the metadata into pinecone json format. - Add integration tests and README.
2023-11-01 19:09:12 +01:00
parent 001ee990b3
commit 81a5eaf801
5 changed files with 810 additions and 0 deletions
--- a/pom.xml
+++ b/pom.xml
@@ -28,6 +28,8 @@
 		<module>document-readers/pdf-reader</module>
 		<module>document-readers/tika-reader</module>
 		<module>embedding-clients/transformers-embedding</module>
+		<module>vector-stores/spring-ai-pinecone</module>
+
 	</modules>

 	<organization>
@@ -82,6 +84,8 @@
 		<pgvector.version>0.1.3</pgvector.version>
 		<postgresql.version>42.6.0</postgresql.version>
 		<milvus.version>2.3.0</milvus.version>
+		<pinecone.version>0.6.0</pinecone.version>
+		<protobuf-java-util.version>3.24.4</protobuf-java-util.version>

 		<!-- testing dependecies -->
 		<testcontainers.version>1.19.0</testcontainers.version>
--- a/vector-stores/spring-ai-pinecone/README.md
+++ b/vector-stores/spring-ai-pinecone/README.md
@@ -0,0 +1,113 @@
+# Pinecone VectorStore
+
+This readme will walk you through setting up the Pinecone VectorStore to store document embeddings and perform similarity searches.
+
+## What is Pinecone?
+
+[Pinecone](https://www.pinecone.io/) is a popular cloud-based vector database, which allows you to store and search vectors efficiently.
+
+## Prerequisites
+
+1. Pinecone Account: Before you start, ensure you sign up for a [Pinecone account](https://app.pinecone.io/).
+2. Pinecone Project: Once registered, create a new project, an index, and generate an API key. You'll need these details for configuration.
+3. 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)
+
+## Configuraiton
+
+To set up PineconeVectorStore, gather the following details from your Pinecone account:
+
+* Pinecond API Key
+* Pinecone Environment
+* Pinecone Project ID
+* Pinecone Index Name
+* Pinecone Namespace
+
+> **Note**
+>  This information is available to you in the Pinecone UI portal.
+
+
+When setting up embeddings, select a vector dimension of 1526. This matches the dimensionality of OpenAI's model "text-embedding-ada-002", which we'll be using for this guide.
+
+Additionally, you'll need to provide your OpenAI API Key. Set it as an environment variable like so:
+
+```bash
+export SPRING_AI_OPENAI_API_KEY='Your_OpenAI_API_Key'
+```
+
+## Dependencies
+
+Add these dependencies to your project:
+
+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>
+    </dependency>
+```
+
+2. Pinecone 
+
+```xml
+    <dependency>
+        <groupId>org.springframework.experimental.ai</groupId>
+        <artifactId>spring-ai-pinecone</artifactId>
+        <version>0.7.0-SNAPSHOT</version>
+    </dependency>
+```
+
+## Sample Code
+
+To configure Pinecone in your application, you can use the following setup:
+
+```java
+@Bean
+public PineconeVectorStoreConfig pineconeVectorStoreConfig() {
+
+        return PineconeVectorStoreConfig.builder()
+        .withApiKey(System.getenv( <PINECONE_API_KEY> ))
+        .withEnvironment(gcp-starter)
+        .withProjectId(89309e6)
+        .withIndexName(spring-ai-test-index)
+        .withNamespace("") // Leave it empty as for free tier as later doesn't support namespaces.
+        .build();
+}
+```
+
+Integrate with OpenAI's embeddings by adding the Spring Boot OpenAI starter to your project.
+This provides you with an implementation of the Embeddings client:
+
+```java
+@Bean
+public VectorStore vectorStore(PineconeVectorStoreConfig config, EmbeddingClient embeddingClient) {
+    return new PineconeVectorStore(config, embeddingClient);
+}
+```
+
+In your main code, create some documents
+
+```java
+	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")));
+```
+
+Add the documents to your vector store:
+
+```java
+vectorStore.add(List.of(document));
+```
+
+And finally, retrieve documents similar to a query:
+
+```java
+List<Document> results = vectorStore.similaritySearch("Spring", 5);
+```
+
+If all goes well, you should retrieve the document containing the text "Spring AI rocks!!".
--- a/vector-stores/spring-ai-pinecone/pom.xml
+++ b/vector-stores/spring-ai-pinecone/pom.xml
@@ -0,0 +1,70 @@
+<?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 https://maven.apache.org/xsd/maven-4.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>
+        <relativePath>../../pom.xml</relativePath>
+    </parent>
+    <artifactId>spring-ai-pinecone</artifactId>
+    <packaging>jar</packaging>
+    <name>spring-ai-pinecone</name>
+    <description>spring-ai-pinecone</description>
+    <url>https://github.com/spring-projects-experimental/spring-ai</url>
+
+    <scm>
+        <url>https://github.com/spring-projects-experimental/spring-ai</url>
+        <connection>git://github.com/spring-projects-experimental/spring-ai.git</connection>
+        <developerConnection>git@github.com:spring-projects-experimental/spring-ai.git</developerConnection>
+    </scm>
+
+    <properties>
+        <maven.compiler.target>17</maven.compiler.target>
+        <maven.compiler.source>17</maven.compiler.source>
+    </properties>
+
+    <dependencies>
+        <dependency>
+            <groupId>org.springframework.experimental.ai</groupId>
+            <artifactId>spring-ai-core</artifactId>
+            <version>${project.parent.version}</version>
+        </dependency>
+
+        <dependency>
+            <groupId>io.pinecone</groupId>
+            <artifactId>pinecone-client</artifactId>
+            <version>${pinecone.version}</version>
+        </dependency>
+
+        <dependency>
+            <groupId>com.google.protobuf</groupId>
+            <artifactId>protobuf-java-util</artifactId>
+            <version>${protobuf-java-util.version}</version>
+        </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.boot</groupId>
+            <artifactId>spring-boot-starter-test</artifactId>
+            <scope>test</scope>
+        </dependency>
+
+        <dependency>
+            <groupId>org.awaitility</groupId>
+            <artifactId>awaitility</artifactId>
+            <version>3.0.0</version>
+            <scope>test</scope>
+        </dependency>
+
+    </dependencies>
+
+</project>
--- a/vector-stores/spring-ai-pinecone/src/main/java/org/springframework/ai/vectorstore/PineconeVectorStore.java
+++ b/vector-stores/spring-ai-pinecone/src/main/java/org/springframework/ai/vectorstore/PineconeVectorStore.java
@@ -0,0 +1,393 @@
+/*
+ * 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.time.Duration;
+import java.util.List;
+import java.util.Map;
+import java.util.Optional;
+
+import com.fasterxml.jackson.databind.ObjectMapper;
+import com.google.protobuf.Struct;
+import com.google.protobuf.Value;
+import com.google.protobuf.util.JsonFormat;
+import io.pinecone.PineconeClient;
+import io.pinecone.PineconeClientConfig;
+import io.pinecone.PineconeConnection;
+import io.pinecone.PineconeConnectionConfig;
+import io.pinecone.proto.DeleteRequest;
+import io.pinecone.proto.QueryRequest;
+import io.pinecone.proto.QueryResponse;
+import io.pinecone.proto.UpsertRequest;
+import io.pinecone.proto.Vector;
+
+import org.springframework.ai.document.Document;
+import org.springframework.ai.embedding.EmbeddingClient;
+import org.springframework.util.Assert;
+
+/**
+ * A VectorStore implementation backed by Pinecone, a cloud-based vector database. This
+ * store supports creating, updating, deleting, and similarity searching of documents in a
+ * Pinecone index.
+ *
+ * @author Christian Tzolov
+ */
+public class PineconeVectorStore implements VectorStore {
+
+	private static final String CONTENT_FIELD_NAME = "document_content";
+
+	private static final String DISTANCE_METADATA_FIELD_NAME = "distance";
+
+	private static final Double SIMILARITY_THRESHOLD_ALL = 0.0;
+
+	private final EmbeddingClient embeddingClient;
+
+	private final PineconeConnection pineconeConnection;
+
+	private final String pineconeNamespace;
+
+	private final int defaultSimilarityTopK;
+
+	private final ObjectMapper objectMapper;
+
+	/**
+	 * Configuration class for the PineconeVectorStore.
+	 */
+	public static final class PineconeVectorStoreConfig {
+
+		// The free tier (gcp-starter) doesn't support Namespaces.
+		// Leave the namespace empty (e.g. "") for the free tier.
+		private final String namespace;
+
+		private final PineconeConnectionConfig connectionConfig;
+
+		private final PineconeClientConfig clientConfig;
+
+		private final int defaultSimilarityTopK;
+
+		/**
+		 * Constructor using the builder.
+		 * @param builder The configuration builder.
+		 */
+		/**
+		 * Constructor using the builder.
+		 * @param builder The configuration builder.
+		 */
+		public PineconeVectorStoreConfig(Builder builder) {
+			this.namespace = builder.namespace;
+			this.defaultSimilarityTopK = builder.defaultSimilarityTopK;
+			this.connectionConfig = new PineconeConnectionConfig().withIndexName(builder.indexName);
+			this.clientConfig = new PineconeClientConfig().withApiKey(builder.apiKey)
+				.withEnvironment(builder.environment)
+				.withProjectName(builder.projectId)
+				.withApiKey(builder.apiKey)
+				.withServerSideTimeoutSec((int) builder.serverSideTimeout.toSeconds());
+		}
+
+		/**
+		 * Start building a new configuration.
+		 * @return The entry point for creating a new configuration.
+		 */
+		public static Builder builder() {
+			return new Builder();
+		}
+
+		/**
+		 * {@return the default config}
+		 */
+		public static PineconeVectorStoreConfig defaultConfig() {
+			return builder().build();
+		}
+
+		public static class Builder {
+
+			private String apiKey;
+
+			private String projectId;
+
+			private String environment;
+
+			private String indexName;
+
+			// 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.
+			 */
+			private Duration serverSideTimeout = Duration.ofSeconds(20);
+
+			private Builder() {
+			}
+
+			/**
+			 * Pinecone api key.
+			 * @param apiKey key to use.
+			 * @return this builder.
+			 */
+			public Builder withApiKey(String apiKey) {
+				this.apiKey = apiKey;
+				return this;
+			}
+
+			/**
+			 * Pinecone project id.
+			 * @param projectId Project id to use.
+			 * @return this builder.
+			 */
+			public Builder withProjectId(String projectId) {
+				this.projectId = projectId;
+				return this;
+			}
+
+			/**
+			 * Pinecone environment name.
+			 * @param environment Environment name (e.g. gcp-starter).
+			 * @return this builder.
+			 */
+			public Builder withEnvironment(String environment) {
+				this.environment = environment;
+				return this;
+			}
+
+			/**
+			 * Pinecone index name.
+			 * @param indexName Pinecone index name to use.
+			 * @return this builder.
+			 */
+			public Builder withIndexName(String indexName) {
+				this.indexName = indexName;
+				return this;
+			}
+
+			/**
+			 * Pinecone Namespace. The free-tier (gcp-starter) doesn't support Namespaces.
+			 * For free-tier leave the namespace empty.
+			 * @param namespace Pinecone namespace to use.
+			 * @return this builder.
+			 */
+			public Builder withNamespace(String namespace) {
+				this.namespace = namespace;
+				return this;
+			}
+
+			/**
+			 * Pinecone server side timeout.
+			 * @param serverSideTimeout server timeout to use.
+			 * @return this builder.
+			 */
+			public Builder withServerSideTimeout(Duration serverSideTimeout) {
+				this.serverSideTimeout = serverSideTimeout;
+				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}
+			 */
+			public PineconeVectorStoreConfig build() {
+				return new PineconeVectorStoreConfig(this);
+			}
+
+		}
+
+	}
+
+	/**
+	 * Constructs a new PineconeVectorStore.
+	 * @param config The configuration for the store.
+	 * @param embeddingClient The client for embedding operations.
+	 */
+	public PineconeVectorStore(PineconeVectorStoreConfig config, EmbeddingClient embeddingClient) {
+		Assert.notNull(config, "PineconeVectorStoreConfig must not be null");
+		Assert.notNull(embeddingClient, "EmbeddingClient must not be null");
+
+		this.embeddingClient = embeddingClient;
+		this.pineconeNamespace = config.namespace;
+		this.defaultSimilarityTopK = config.defaultSimilarityTopK;
+		this.pineconeConnection = new PineconeClient(config.clientConfig).connect(config.connectionConfig);
+		this.objectMapper = new ObjectMapper();
+	}
+
+	/**
+	 * Adds a list of documents to the vector store.
+	 * @param documents The list of documents to be added.
+	 */
+	@Override
+	public void add(List<Document> documents) {
+
+		List<Vector> upsertVectors = documents.stream().map(document -> {
+			// Compute and assign an embedding to the document.
+			document.setEmbedding(this.embeddingClient.embed(document));
+
+			return Vector.newBuilder()
+				.setId(document.getId())
+				.addAllValues(toFloatList(document.getEmbedding()))
+				.setMetadata(metadataToStruct(document))
+				.build();
+		}).toList();
+
+		UpsertRequest upsertRequest = UpsertRequest.newBuilder()
+			.addAllVectors(upsertVectors)
+			.setNamespace(this.pineconeNamespace)
+			.build();
+
+		this.pineconeConnection.getBlockingStub().upsert(upsertRequest);
+	}
+
+	/**
+	 * Converts the document metadata to a Protobuf Struct.
+	 * @param document The document containing metadata.
+	 * @return The metadata as a Protobuf Struct.
+	 */
+	private Struct metadataToStruct(Document document) {
+		try {
+			var structBuilder = Struct.newBuilder();
+			JsonFormat.parser()
+				.ignoringUnknownFields()
+				.merge(this.objectMapper.writeValueAsString(document.getMetadata()), structBuilder);
+			structBuilder.putFields(CONTENT_FIELD_NAME, contentValue(document));
+			return structBuilder.build();
+		}
+		catch (Exception e) {
+			throw new RuntimeException(e);
+		}
+	}
+
+	/**
+	 * Retrieves the content value of a document.
+	 * @param document The document.
+	 * @return The content value.
+	 */
+	private Value contentValue(Document document) {
+		return Value.newBuilder().setStringValue(document.getContent()).build();
+	}
+
+	/**
+	 * Deletes a list of documents by their IDs.
+	 * @param documentIds The list of document IDs to be deleted.
+	 * @return An optional boolean indicating the deletion status.
+	 */
+	@Override
+	public Optional<Boolean> delete(List<String> documentIds) {
+
+		DeleteRequest deleteRequest = DeleteRequest.newBuilder()
+			.setNamespace(this.pineconeNamespace) // ignored for free tier.
+			.addAllIds(documentIds)
+			.setDeleteAll(false)
+			.build();
+
+		this.pineconeConnection.getBlockingStub().delete(deleteRequest);
+
+		// The Pinecone delete API does not provide deletion status info.
+		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);
+	}
+
+	/**
+	 * 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);
+	}
+
+	/**
+	 * 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) {
+
+		List<Double> queryEmbedding = this.embeddingClient.embed(query);
+
+		QueryRequest queryRequest = QueryRequest.newBuilder()
+			.addAllVector(toFloatList(queryEmbedding))
+			.setTopK(topK)
+			.setIncludeMetadata(true)
+			.setNamespace(this.pineconeNamespace)
+			.build();
+
+		QueryResponse queryResponse = this.pineconeConnection.getBlockingStub().query(queryRequest);
+
+		return queryResponse.getMatchesList()
+			.stream()
+			.filter(scoredVector -> scoredVector.getScore() >= similarityThreshold)
+			.map(scoredVector -> {
+				var id = scoredVector.getId();
+				Struct metadataStruct = scoredVector.getMetadata();
+				var content = metadataStruct.getFieldsOrThrow(CONTENT_FIELD_NAME).getStringValue();
+				Map<String, Object> metadata = extractMetadata(metadataStruct);
+				metadata.put(DISTANCE_METADATA_FIELD_NAME, 1 - scoredVector.getScore());
+				return new Document(id, content, metadata);
+			})
+			.toList();
+	}
+
+	/**
+	 * Extracts metadata from a Protobuf Struct.
+	 * @param metadataStruct The Protobuf Struct containing metadata.
+	 * @return The metadata as a map.
+	 */
+	private Map<String, Object> extractMetadata(Struct metadataStruct) {
+		try {
+			String json = JsonFormat.printer().print(metadataStruct);
+			Map<String, Object> metadata = this.objectMapper.readValue(json, Map.class);
+			metadata.remove(CONTENT_FIELD_NAME);
+			return metadata;
+		}
+		catch (Exception e) {
+			throw new RuntimeException(e);
+		}
+	}
+
+	/**
+	 * Converts a list of doubles to a list of floats.
+	 * @param doubleList The list of doubles.
+	 * @return The converted list of floats.
+	 */
+	private List<Float> toFloatList(List<Double> doubleList) {
+		return doubleList.stream().map(d -> d.floatValue()).toList();
+	}
+
+}
--- a/vector-stores/spring-ai-pinecone/src/test/java/org/springframework/ai/vectorstore/PineconeVectorStoreIT.java
+++ b/vector-stores/spring-ai-pinecone/src/test/java/org/springframework/ai/vectorstore/PineconeVectorStoreIT.java
@@ -0,0 +1,230 @@
+/*
+ * 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.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.autoconfigure.openai.OpenAiAutoConfiguration;
+import org.springframework.ai.document.Document;
+import org.springframework.ai.embedding.EmbeddingClient;
+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 static org.assertj.core.api.Assertions.assertThat;
+import static org.hamcrest.Matchers.equalTo;
+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
+	// PINECONE_API_KEY with your pinecone credentials.
+	private static final String PINECONE_ENVIRONMENT = "gcp-starter";
+
+	private static final String PINECONE_PROJECT_ID = "89309e6";
+
+	private static final String PINECONE_INDEX_NAME = "spring-ai-test-index";
+
+	// NOTE: Leave it empty as for free tier as later doesn't support namespaces.
+	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")));
+
+	private final ApplicationContextRunner contextRunner = new ApplicationContextRunner()
+		.withUserConfiguration(TestApplication.class)
+		.withPropertyValues("spring.ai.openai.apiKey=" + System.getenv("OPENAI_API_KEY"));
+
+	@BeforeAll
+	public static void beforeAll() {
+		Awaitility.setDefaultPollInterval(10, TimeUnit.SECONDS);
+		Awaitility.setDefaultPollDelay(Duration.ZERO);
+		Awaitility.setDefaultTimeout(Duration.ONE_MINUTE);
+	}
+
+	@Test
+	public void addAndSearchTest() {
+
+		contextRunner.withConfiguration(AutoConfigurations.of(OpenAiAutoConfiguration.class)).run(context -> {
+
+			VectorStore vectorStore = context.getBean(VectorStore.class);
+
+			vectorStore.add(documents);
+
+			Awaitility.await().until(() -> {
+				return vectorStore.similaritySearch("Great", 1);
+			}, hasSize(1));
+
+			List<Document> results = vectorStore.similaritySearch("Great", 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");
+
+			// Remove all documents from the store
+			vectorStore.delete(documents.stream().map(doc -> doc.getId()).toList());
+
+			Awaitility.await().until(() -> {
+				return vectorStore.similaritySearch("Hello", 1);
+			}, hasSize(0));
+		});
+	}
+
+	@Test
+	public void documentUpdateTest() {
+
+		// Note ,using OpenAI to calculate embeddings
+		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));
+
+			Awaitility.await().until(() -> {
+				return vectorStore.similaritySearch("Spring", 5);
+			}, hasSize(1));
+
+			List<Document> results = vectorStore.similaritySearch("Spring", 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");
+
+			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));
+
+			Awaitility.await().until(() -> {
+				return vectorStore.similaritySearch("FooBar", 5).get(0).getContent();
+			}, equalTo("The World is Big and Salvation Lurks Around the Corner"));
+
+			results = vectorStore.similaritySearch("FooBar", 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");
+
+			// Remove all documents from the store
+			vectorStore.delete(List.of(document.getId()));
+			Awaitility.await().until(() -> {
+				return vectorStore.similaritySearch("FooBar", 1);
+			}, hasSize(0));
+
+		});
+	}
+
+	@Test
+	public void searchThresholdTest() {
+
+		contextRunner.withConfiguration(AutoConfigurations.of(OpenAiAutoConfiguration.class)).run(context -> {
+
+			VectorStore vectorStore = context.getBean(VectorStore.class);
+
+			vectorStore.add(documents);
+
+			Awaitility.await().until(() -> {
+				return vectorStore.similaritySearch("Great", 5);
+			}, hasSize(3));
+
+			List<Document> fullResult = vectorStore.similaritySearch("Great", 5, 0.0);
+
+			List<Float> distances = fullResult.stream().map(doc -> (Float) doc.getMetadata().get("distance")).toList();
+
+			assertThat(distances).hasSize(3);
+
+			float threshold = (distances.get(0) + distances.get(1)) / 2;
+
+			List<Document> results = vectorStore.similaritySearch("Great", 5, (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");
+
+			// Remove all documents from the store
+			vectorStore.delete(documents.stream().map(doc -> doc.getId()).toList());
+			Awaitility.await().until(() -> {
+				return vectorStore.similaritySearch("Hello", 1);
+			}, hasSize(0));
+		});
+	}
+
+	@SpringBootConfiguration
+	@EnableAutoConfiguration
+	public static class TestApplication {
+
+		@Bean
+		public PineconeVectorStoreConfig pineconeVectorStoreConfig() {
+
+			return PineconeVectorStoreConfig.builder()
+				.withApiKey(System.getenv("PINECONE_API_KEY"))
+				.withEnvironment(PINECONE_ENVIRONMENT)
+				.withProjectId(PINECONE_PROJECT_ID)
+				.withIndexName(PINECONE_INDEX_NAME)
+				.withNamespace(PINECONE_NAMESPACE)
+				.build();
+		}
+
+		@Bean
+		public VectorStore vectorStore(PineconeVectorStoreConfig config, EmbeddingClient embeddingClient) {
+			return new PineconeVectorStore(config, embeddingClient);
+		}
+
+	}
+
+}