Vector store impls: code cleanup and checkstyle fixes

sobychacko · sobychacko · commit 3f0557b8b3e4 · 2025-02-11T16:16:57.000-05:00
Signed-off-by: Soby Chacko &lt;soby.chacko@broadcom.com&gt;
diff --git a/vector-stores/spring-ai-azure-cosmos-db-store/src/main/java/org/springframework/ai/vectorstore/cosmosdb/CosmosDBVectorStore.java b/vector-stores/spring-ai-azure-cosmos-db-store/src/main/java/org/springframework/ai/vectorstore/cosmosdb/CosmosDBVectorStore.java
@@ -57,7 +57,6 @@
 import reactor.core.publisher.Flux;
 
 import org.springframework.ai.document.Document;
-import org.springframework.ai.embedding.BatchingStrategy;
 import org.springframework.ai.embedding.EmbeddingModel;
 import org.springframework.ai.embedding.EmbeddingOptionsBuilder;
 import org.springframework.ai.observation.conventions.VectorStoreProvider;
diff --git a/vector-stores/spring-ai-chroma-store/src/main/java/org/springframework/ai/chroma/vectorstore/ChromaVectorStore.java b/vector-stores/spring-ai-chroma-store/src/main/java/org/springframework/ai/chroma/vectorstore/ChromaVectorStore.java
@@ -20,7 +20,6 @@
 import java.util.HashMap;
 import java.util.List;
 import java.util.Map;
-import java.util.Optional;
 
 import com.fasterxml.jackson.core.JsonProcessingException;
 import com.fasterxml.jackson.databind.ObjectMapper;
diff --git a/vector-stores/spring-ai-chroma-store/src/test/java/org/springframework/ai/chroma/vectorstore/ChromaVectorStoreIT.java b/vector-stores/spring-ai-chroma-store/src/test/java/org/springframework/ai/chroma/vectorstore/ChromaVectorStoreIT.java
@@ -19,7 +19,6 @@
 import java.util.Collections;
 import java.util.List;
 import java.util.Map;
-import java.util.Optional;
 import java.util.UUID;
 
 import org.junit.jupiter.api.Test;
diff --git a/vector-stores/spring-ai-elasticsearch-store/src/main/java/org/springframework/ai/vectorstore/elasticsearch/ElasticsearchVectorStore.java b/vector-stores/spring-ai-elasticsearch-store/src/main/java/org/springframework/ai/vectorstore/elasticsearch/ElasticsearchVectorStore.java
@@ -26,7 +26,6 @@
 import co.elastic.clients.elasticsearch.ElasticsearchClient;
 import co.elastic.clients.elasticsearch.core.BulkRequest;
 import co.elastic.clients.elasticsearch.core.BulkResponse;
-import co.elastic.clients.elasticsearch.core.DeleteByQueryResponse;
 import co.elastic.clients.elasticsearch.core.SearchResponse;
 import co.elastic.clients.elasticsearch.core.bulk.BulkResponseItem;
 import co.elastic.clients.elasticsearch.core.search.Hit;
@@ -41,10 +40,8 @@
 
 import org.springframework.ai.document.Document;
 import org.springframework.ai.document.DocumentMetadata;
-import org.springframework.ai.embedding.BatchingStrategy;
 import org.springframework.ai.embedding.EmbeddingModel;
 import org.springframework.ai.embedding.EmbeddingOptionsBuilder;
-import org.springframework.ai.embedding.TokenCountBatchingStrategy;
 import org.springframework.ai.model.EmbeddingUtils;
 import org.springframework.ai.observation.conventions.VectorStoreProvider;
 import org.springframework.ai.observation.conventions.VectorStoreSimilarityMetric;
diff --git a/vector-stores/spring-ai-elasticsearch-store/src/test/java/org/springframework/ai/vectorstore/elasticsearch/ElasticsearchVectorStoreIT.java b/vector-stores/spring-ai-elasticsearch-store/src/test/java/org/springframework/ai/vectorstore/elasticsearch/ElasticsearchVectorStoreIT.java
@@ -54,7 +54,6 @@
 import org.springframework.ai.openai.OpenAiEmbeddingModel;
 import org.springframework.ai.openai.api.OpenAiApi;
 import org.springframework.ai.vectorstore.SearchRequest;
-import org.springframework.ai.vectorstore.filter.Filter;
 import org.springframework.ai.vectorstore.filter.Filter.Expression;
 import org.springframework.ai.vectorstore.filter.Filter.ExpressionType;
 import org.springframework.ai.vectorstore.filter.Filter.Key;
diff --git a/vector-stores/spring-ai-gemfire-store/src/main/java/org/springframework/ai/vectorstore/gemfire/GemFireVectorStore.java b/vector-stores/spring-ai-gemfire-store/src/main/java/org/springframework/ai/vectorstore/gemfire/GemFireVectorStore.java
@@ -19,7 +19,6 @@
 import java.util.HashMap;
 import java.util.List;
 import java.util.Map;
-import java.util.Optional;
 
 import com.fasterxml.jackson.annotation.JsonCreator;
 import com.fasterxml.jackson.annotation.JsonInclude;
@@ -32,10 +31,8 @@
 
 import org.springframework.ai.document.Document;
 import org.springframework.ai.document.DocumentMetadata;
-import org.springframework.ai.embedding.BatchingStrategy;
 import org.springframework.ai.embedding.EmbeddingModel;
 import org.springframework.ai.embedding.EmbeddingOptionsBuilder;
-import org.springframework.ai.embedding.TokenCountBatchingStrategy;
 import org.springframework.ai.observation.conventions.VectorStoreProvider;
 import org.springframework.ai.util.JacksonUtils;
 import org.springframework.ai.vectorstore.AbstractVectorStoreBuilder;
diff --git a/vector-stores/spring-ai-hanadb-store/src/main/java/org/springframework/ai/vectorstore/hanadb/HanaCloudVectorStore.java b/vector-stores/spring-ai-hanadb-store/src/main/java/org/springframework/ai/vectorstore/hanadb/HanaCloudVectorStore.java
@@ -18,7 +18,6 @@
 
 import java.util.Collections;
 import java.util.List;
-import java.util.Optional;
 import java.util.stream.Collectors;
 
 import com.fasterxml.jackson.core.JsonProcessingException;
diff --git a/vector-stores/spring-ai-mariadb-store/src/main/java/org/springframework/ai/vectorstore/mariadb/MariaDBVectorStore.java b/vector-stores/spring-ai-mariadb-store/src/main/java/org/springframework/ai/vectorstore/mariadb/MariaDBVectorStore.java
@@ -32,10 +32,8 @@
 import org.slf4j.LoggerFactory;
 
 import org.springframework.ai.document.Document;
-import org.springframework.ai.embedding.BatchingStrategy;
 import org.springframework.ai.embedding.EmbeddingModel;
 import org.springframework.ai.embedding.EmbeddingOptionsBuilder;
-import org.springframework.ai.embedding.TokenCountBatchingStrategy;
 import org.springframework.ai.observation.conventions.VectorStoreProvider;
 import org.springframework.ai.observation.conventions.VectorStoreSimilarityMetric;
 import org.springframework.ai.util.JacksonUtils;
diff --git a/vector-stores/spring-ai-milvus-store/src/main/java/org/springframework/ai/vectorstore/milvus/MilvusVectorStore.java b/vector-stores/spring-ai-milvus-store/src/main/java/org/springframework/ai/vectorstore/milvus/MilvusVectorStore.java
@@ -54,10 +54,8 @@
 
 import org.springframework.ai.document.Document;
 import org.springframework.ai.document.DocumentMetadata;
-import org.springframework.ai.embedding.BatchingStrategy;
 import org.springframework.ai.embedding.EmbeddingModel;
 import org.springframework.ai.embedding.EmbeddingOptionsBuilder;
-import org.springframework.ai.embedding.TokenCountBatchingStrategy;
 import org.springframework.ai.model.EmbeddingUtils;
 import org.springframework.ai.observation.conventions.VectorStoreProvider;
 import org.springframework.ai.observation.conventions.VectorStoreSimilarityMetric;
diff --git a/vector-stores/spring-ai-opensearch-store/src/test/java/org/springframework/ai/vectorstore/opensearch/OpenSearchVectorStoreIT.java b/vector-stores/spring-ai-opensearch-store/src/test/java/org/springframework/ai/vectorstore/opensearch/OpenSearchVectorStoreIT.java
@@ -19,14 +19,12 @@
 import java.io.IOException;
 import java.net.URISyntaxException;
 import java.nio.charset.StandardCharsets;
-import java.text.SimpleDateFormat;
 import java.time.Duration;
 import java.time.ZonedDateTime;
 import java.util.Date;
 import java.util.List;
 import java.util.Map;
 import java.util.Optional;
-import java.util.TimeZone;
 import java.util.UUID;
 import java.util.concurrent.TimeUnit;
 import java.util.stream.Collectors;
diff --git a/vector-stores/spring-ai-oracle-store/src/main/java/org/springframework/ai/vectorstore/oracle/OracleVectorStore.java b/vector-stores/spring-ai-oracle-store/src/main/java/org/springframework/ai/vectorstore/oracle/OracleVectorStore.java
@@ -39,10 +39,8 @@
 
 import org.springframework.ai.document.Document;
 import org.springframework.ai.document.DocumentMetadata;
-import org.springframework.ai.embedding.BatchingStrategy;
 import org.springframework.ai.embedding.EmbeddingModel;
 import org.springframework.ai.embedding.EmbeddingOptionsBuilder;
-import org.springframework.ai.embedding.TokenCountBatchingStrategy;
 import org.springframework.ai.observation.conventions.VectorStoreProvider;
 import org.springframework.ai.observation.conventions.VectorStoreSimilarityMetric;
 import org.springframework.ai.vectorstore.AbstractVectorStoreBuilder;
diff --git a/vector-stores/spring-ai-pgvector-store/src/test/java/org/springframework/ai/vectorstore/pgvector/PgVectorStoreIT.java b/vector-stores/spring-ai-pgvector-store/src/test/java/org/springframework/ai/vectorstore/pgvector/PgVectorStoreIT.java
@@ -47,12 +47,12 @@
 import org.springframework.ai.embedding.EmbeddingModel;
 import org.springframework.ai.openai.OpenAiEmbeddingModel;
 import org.springframework.ai.openai.api.OpenAiApi;
-import org.springframework.ai.vectorstore.pgvector.PgVectorStore.PgIdType;
-import org.springframework.ai.vectorstore.pgvector.PgVectorStore.PgIndexType;
 import org.springframework.ai.vectorstore.SearchRequest;
 import org.springframework.ai.vectorstore.VectorStore;
 import org.springframework.ai.vectorstore.filter.Filter;
 import org.springframework.ai.vectorstore.filter.FilterExpressionTextParser.FilterExpressionParseException;
+import org.springframework.ai.vectorstore.pgvector.PgVectorStore.PgIdType;
+import org.springframework.ai.vectorstore.pgvector.PgVectorStore.PgIndexType;
 import org.springframework.beans.factory.annotation.Value;
 import org.springframework.boot.SpringBootConfiguration;
 import org.springframework.boot.autoconfigure.EnableAutoConfiguration;
@@ -560,9 +560,9 @@ public static class TestApplication {
 		public VectorStore vectorStore(JdbcTemplate jdbcTemplate, EmbeddingModel embeddingModel) {
 			return PgVectorStore.builder(jdbcTemplate, embeddingModel)
 				.dimensions(PgVectorStore.INVALID_EMBEDDING_DIMENSION)
-				.idType(idType)
+				.idType(this.idType)
 				.distanceType(this.distanceType)
-				.initializeSchema(initializeSchema)
+				.initializeSchema(this.initializeSchema)
 				.indexType(PgIndexType.HNSW)
 				.removeExistingVectorStoreTable(true)
 				.build();
diff --git a/vector-stores/spring-ai-pinecone-store/src/main/java/org/springframework/ai/vectorstore/pinecone/PineconeVectorStore.java b/vector-stores/spring-ai-pinecone-store/src/main/java/org/springframework/ai/vectorstore/pinecone/PineconeVectorStore.java
@@ -419,6 +419,58 @@ private Builder(EmbeddingModel embeddingModel, String apiKey, String projectId,
 			this.indexName = indexName;
 		}
 
+		/**
+		 * Sets the Pinecone namespace. Note: The free-tier (gcp-starter) doesn't support
+		 * Namespaces.
+		 * @param namespace The namespace to use (leave empty for free tier)
+		 * @return The builder instance
+		 */
+		public Builder namespace(@Nullable String namespace) {
+			this.namespace = namespace != null ? namespace : "";
+			return this;
+		}
+
+		/**
+		 * Sets the content field name.
+		 * @param contentFieldName The content field name to use
+		 * @return The builder instance
+		 */
+		public Builder contentFieldName(@Nullable String contentFieldName) {
+			this.contentFieldName = contentFieldName != null ? contentFieldName : CONTENT_FIELD_NAME;
+			return this;
+		}
+
+		/**
+		 * Sets the distance metadata field name.
+		 * @param distanceMetadataFieldName The distance metadata field name to use
+		 * @return The builder instance
+		 */
+		public Builder distanceMetadataFieldName(@Nullable String distanceMetadataFieldName) {
+			this.distanceMetadataFieldName = distanceMetadataFieldName != null ? distanceMetadataFieldName
+					: DocumentMetadata.DISTANCE.value();
+			return this;
+		}
+
+		/**
+		 * Sets the server-side timeout.
+		 * @param serverSideTimeout The timeout duration to use
+		 * @return The builder instance
+		 */
+		public Builder serverSideTimeout(@Nullable Duration serverSideTimeout) {
+			this.serverSideTimeout = serverSideTimeout != null ? serverSideTimeout : Duration.ofSeconds(20);
+			return this;
+		}
+
+		/**
+		 * Builds a new PineconeVectorStore instance with the configured properties.
+		 * @return A new PineconeVectorStore instance
+		 * @throws IllegalStateException if the builder is in an invalid state
+		 */
+		@Override
+		public PineconeVectorStore build() {
+			return new PineconeVectorStore(this);
+		}
+
 		/**
 		 * First step interface requiring API key configuration.
 		 */
@@ -487,19 +539,30 @@ public interface BuilderWithIndexName {
 		 */
 		public static class StepBuilder {
 
+			/**
+			 * Initiates the step builder sequence with the embedding model.
+			 * @param embeddingModel The embedding model to use
+			 * @return The first step for API key configuration
+			 * @throws IllegalArgumentException if embeddingModel is null
+			 */
+			static BuilderWithApiKey start(EmbeddingModel embeddingModel) {
+				Assert.notNull(embeddingModel, "EmbeddingModel must not be null");
+				return new ApiKeyStep(embeddingModel);
+			}
+
 			private record ApiKeyStep(EmbeddingModel embeddingModel) implements BuilderWithApiKey {
 				@Override
 				public BuilderWithProjectId apiKey(String apiKey) {
 					Assert.hasText(apiKey, "ApiKey must not be null or empty");
-					return new ProjectIdStep(embeddingModel, apiKey);
+					return new ProjectIdStep(this.embeddingModel, apiKey);
 				}
 			}
 
 			private record ProjectIdStep(EmbeddingModel embeddingModel, String apiKey) implements BuilderWithProjectId {
 				@Override
 				public BuilderWithEnvironment projectId(String projectId) {
 					Assert.hasText(projectId, "ProjectId must not be null or empty");
-					return new EnvironmentStep(embeddingModel, apiKey, projectId);
+					return new EnvironmentStep(this.embeddingModel, this.apiKey, projectId);
 				}
 			}
 
@@ -508,7 +571,7 @@ private record EnvironmentStep(EmbeddingModel embeddingModel, String apiKey,
 				@Override
 				public BuilderWithIndexName environment(String environment) {
 					Assert.hasText(environment, "Environment must not be null or empty");
-					return new IndexNameStep(embeddingModel, apiKey, projectId, environment);
+					return new IndexNameStep(this.embeddingModel, this.apiKey, this.projectId, environment);
 				}
 			}
 
@@ -517,73 +580,10 @@ private record IndexNameStep(EmbeddingModel embeddingModel, String apiKey, Strin
 				@Override
 				public Builder indexName(String indexName) {
 					Assert.hasText(indexName, "IndexName must not be null or empty");
-					return new Builder(embeddingModel, apiKey, projectId, environment, indexName);
+					return new Builder(this.embeddingModel, this.apiKey, this.projectId, this.environment, indexName);
 				}
 			}
 
-			/**
-			 * Initiates the step builder sequence with the embedding model.
-			 * @param embeddingModel The embedding model to use
-			 * @return The first step for API key configuration
-			 * @throws IllegalArgumentException if embeddingModel is null
-			 */
-			static BuilderWithApiKey start(EmbeddingModel embeddingModel) {
-				Assert.notNull(embeddingModel, "EmbeddingModel must not be null");
-				return new ApiKeyStep(embeddingModel);
-			}
-
-		}
-
-		/**
-		 * Sets the Pinecone namespace. Note: The free-tier (gcp-starter) doesn't support
-		 * Namespaces.
-		 * @param namespace The namespace to use (leave empty for free tier)
-		 * @return The builder instance
-		 */
-		public Builder namespace(@Nullable String namespace) {
-			this.namespace = namespace != null ? namespace : "";
-			return this;
-		}
-
-		/**
-		 * Sets the content field name.
-		 * @param contentFieldName The content field name to use
-		 * @return The builder instance
-		 */
-		public Builder contentFieldName(@Nullable String contentFieldName) {
-			this.contentFieldName = contentFieldName != null ? contentFieldName : CONTENT_FIELD_NAME;
-			return this;
-		}
-
-		/**
-		 * Sets the distance metadata field name.
-		 * @param distanceMetadataFieldName The distance metadata field name to use
-		 * @return The builder instance
-		 */
-		public Builder distanceMetadataFieldName(@Nullable String distanceMetadataFieldName) {
-			this.distanceMetadataFieldName = distanceMetadataFieldName != null ? distanceMetadataFieldName
-					: DocumentMetadata.DISTANCE.value();
-			return this;
-		}
-
-		/**
-		 * Sets the server-side timeout.
-		 * @param serverSideTimeout The timeout duration to use
-		 * @return The builder instance
-		 */
-		public Builder serverSideTimeout(@Nullable Duration serverSideTimeout) {
-			this.serverSideTimeout = serverSideTimeout != null ? serverSideTimeout : Duration.ofSeconds(20);
-			return this;
-		}
-
-		/**
-		 * Builds a new PineconeVectorStore instance with the configured properties.
-		 * @return A new PineconeVectorStore instance
-		 * @throws IllegalStateException if the builder is in an invalid state
-		 */
-		@Override
-		public PineconeVectorStore build() {
-			return new PineconeVectorStore(this);
 		}
 
 	}
diff --git a/vector-stores/spring-ai-qdrant-store/src/main/java/org/springframework/ai/vectorstore/qdrant/QdrantVectorStore.java b/vector-stores/spring-ai-qdrant-store/src/main/java/org/springframework/ai/vectorstore/qdrant/QdrantVectorStore.java
@@ -31,7 +31,6 @@
 import io.qdrant.client.grpc.Points.PointStruct;
 import io.qdrant.client.grpc.Points.ScoredPoint;
 import io.qdrant.client.grpc.Points.SearchPoints;
-import io.qdrant.client.grpc.Points.UpdateStatus;
 import org.slf4j.Logger;
 import org.slf4j.LoggerFactory;
 
diff --git a/vector-stores/spring-ai-typesense-store/src/test/java/org/springframework/ai/vectorstore/typesense/TypesenseVectorStoreIT.java b/vector-stores/spring-ai-typesense-store/src/test/java/org/springframework/ai/vectorstore/typesense/TypesenseVectorStoreIT.java
@@ -28,7 +28,6 @@
 import java.util.stream.Collectors;
 
 import org.junit.jupiter.api.Test;
-
 import org.testcontainers.junit.jupiter.Container;
 import org.testcontainers.junit.jupiter.Testcontainers;
 import org.testcontainers.typesense.TypesenseContainer;
