From 3c8817f2cb13f75d52dbbd2dea229548d4b88a6b Mon Sep 17 00:00:00 2001 From: Soby Chacko Date: Sat, 7 Dec 2024 18:52:53 -0500 Subject: [PATCH 1/2] refactor: move PgVectorStore to dedicated package and update builder pattern - Move PgVectorStore and related classes to org.springframework.ai.pg.vectorstore package - Update builder pattern to use more idiomatic method names (e.g. withSchemaName -> schemaName) - Deprecate existing constructors and old Builder class in favor of new static builder() method - Update tests to reflect the new builder style usage --- .../RetrievalAugmentationAdvisorIT.java | 2 +- .../VectorStoreDocumentRetrieverIT.java | 2 +- .../PgVectorStoreAutoConfiguration.java | 28 +- .../pgvector/PgVectorStoreProperties.java | 6 +- .../PgVectorStoreAutoConfigurationIT.java | 2 +- .../PgVectorStorePropertiesTests.java | 6 +- .../PgVectorFilterExpressionConverter.java | 2 +- .../vectorstore/PgVectorSchemaValidator.java | 2 +- .../{ => pg}/vectorstore/PgVectorStore.java | 239 +++++++++++++----- .../ai/pg/vectorstore/package-info.java | 25 ++ .../PgVectorEmbeddingDimensionsTests.java | 18 +- ...gVectorFilterExpressionConverterTests.java | 2 +- .../{ => pg}/vectorstore/PgVectorImage.java | 2 +- .../PgVectorStoreCustomNamesIT.java | 23 +- .../{ => pg}/vectorstore/PgVectorStoreIT.java | 16 +- .../PgVectorStoreObservationIT.java | 17 +- .../vectorstore/PgVectorStoreTests.java | 6 +- .../PgVectorStoreWithChatMemoryAdvisorIT.java | 2 +- 18 files changed, 280 insertions(+), 120 deletions(-) rename vector-stores/spring-ai-pgvector-store/src/main/java/org/springframework/ai/{ => pg}/vectorstore/PgVectorFilterExpressionConverter.java (98%) rename vector-stores/spring-ai-pgvector-store/src/main/java/org/springframework/ai/{ => pg}/vectorstore/PgVectorSchemaValidator.java (99%) rename vector-stores/spring-ai-pgvector-store/src/main/java/org/springframework/ai/{ => pg}/vectorstore/PgVectorStore.java (75%) create mode 100644 vector-stores/spring-ai-pgvector-store/src/main/java/org/springframework/ai/pg/vectorstore/package-info.java rename vector-stores/spring-ai-pgvector-store/src/test/java/org/springframework/ai/{ => pg}/vectorstore/PgVectorEmbeddingDimensionsTests.java (76%) rename vector-stores/spring-ai-pgvector-store/src/test/java/org/springframework/ai/{ => pg}/vectorstore/PgVectorFilterExpressionConverterTests.java (99%) rename vector-stores/spring-ai-pgvector-store/src/test/java/org/springframework/ai/{ => pg}/vectorstore/PgVectorImage.java (94%) rename vector-stores/spring-ai-pgvector-store/src/test/java/org/springframework/ai/{ => pg}/vectorstore/PgVectorStoreCustomNamesIT.java (93%) rename vector-stores/spring-ai-pgvector-store/src/test/java/org/springframework/ai/{ => pg}/vectorstore/PgVectorStoreIT.java (96%) rename vector-stores/spring-ai-pgvector-store/src/test/java/org/springframework/ai/{ => pg}/vectorstore/PgVectorStoreObservationIT.java (94%) rename vector-stores/spring-ai-pgvector-store/src/test/java/org/springframework/ai/{ => pg}/vectorstore/PgVectorStoreTests.java (95%) rename vector-stores/spring-ai-pgvector-store/src/test/java/org/springframework/ai/{ => pg}/vectorstore/PgVectorStoreWithChatMemoryAdvisorIT.java (99%) diff --git a/spring-ai-integration-tests/src/test/java/org/springframework/ai/integration/tests/client/advisor/RetrievalAugmentationAdvisorIT.java b/spring-ai-integration-tests/src/test/java/org/springframework/ai/integration/tests/client/advisor/RetrievalAugmentationAdvisorIT.java index a971ddba6f7..11ea5f4d332 100644 --- a/spring-ai-integration-tests/src/test/java/org/springframework/ai/integration/tests/client/advisor/RetrievalAugmentationAdvisorIT.java +++ b/spring-ai-integration-tests/src/test/java/org/springframework/ai/integration/tests/client/advisor/RetrievalAugmentationAdvisorIT.java @@ -38,7 +38,7 @@ import org.springframework.ai.rag.retrieval.search.VectorStoreDocumentRetriever; import org.springframework.ai.reader.markdown.MarkdownDocumentReader; import org.springframework.ai.reader.markdown.config.MarkdownDocumentReaderConfig; -import org.springframework.ai.vectorstore.PgVectorStore; +import org.springframework.ai.pg.vectorstore.PgVectorStore; import org.springframework.beans.factory.annotation.Autowired; import org.springframework.beans.factory.annotation.Value; import org.springframework.boot.test.context.SpringBootTest; diff --git a/spring-ai-integration-tests/src/test/java/org/springframework/ai/integration/tests/rag/retrieval/search/VectorStoreDocumentRetrieverIT.java b/spring-ai-integration-tests/src/test/java/org/springframework/ai/integration/tests/rag/retrieval/search/VectorStoreDocumentRetrieverIT.java index df7f9627a91..7e5e10aa921 100644 --- a/spring-ai-integration-tests/src/test/java/org/springframework/ai/integration/tests/rag/retrieval/search/VectorStoreDocumentRetrieverIT.java +++ b/spring-ai-integration-tests/src/test/java/org/springframework/ai/integration/tests/rag/retrieval/search/VectorStoreDocumentRetrieverIT.java @@ -29,7 +29,7 @@ import org.springframework.ai.rag.Query; import org.springframework.ai.rag.retrieval.search.DocumentRetriever; import org.springframework.ai.rag.retrieval.search.VectorStoreDocumentRetriever; -import org.springframework.ai.vectorstore.PgVectorStore; +import org.springframework.ai.pg.vectorstore.PgVectorStore; import org.springframework.ai.vectorstore.filter.Filter; import org.springframework.beans.factory.annotation.Autowired; import org.springframework.boot.test.context.SpringBootTest; diff --git a/spring-ai-spring-boot-autoconfigure/src/main/java/org/springframework/ai/autoconfigure/vectorstore/pgvector/PgVectorStoreAutoConfiguration.java b/spring-ai-spring-boot-autoconfigure/src/main/java/org/springframework/ai/autoconfigure/vectorstore/pgvector/PgVectorStoreAutoConfiguration.java index a32f88acd12..2b6a8974c28 100644 --- a/spring-ai-spring-boot-autoconfigure/src/main/java/org/springframework/ai/autoconfigure/vectorstore/pgvector/PgVectorStoreAutoConfiguration.java +++ b/spring-ai-spring-boot-autoconfigure/src/main/java/org/springframework/ai/autoconfigure/vectorstore/pgvector/PgVectorStoreAutoConfiguration.java @@ -23,7 +23,7 @@ import org.springframework.ai.embedding.BatchingStrategy; import org.springframework.ai.embedding.EmbeddingModel; import org.springframework.ai.embedding.TokenCountBatchingStrategy; -import org.springframework.ai.vectorstore.PgVectorStore; +import org.springframework.ai.pg.vectorstore.PgVectorStore; import org.springframework.ai.vectorstore.observation.VectorStoreObservationConvention; import org.springframework.beans.factory.ObjectProvider; import org.springframework.boot.autoconfigure.AutoConfiguration; @@ -62,18 +62,20 @@ public PgVectorStore vectorStore(JdbcTemplate jdbcTemplate, EmbeddingModel embed var initializeSchema = properties.isInitializeSchema(); - return new PgVectorStore.Builder(jdbcTemplate, embeddingModel).withSchemaName(properties.getSchemaName()) - .withVectorTableName(properties.getTableName()) - .withVectorTableValidationsEnabled(properties.isSchemaValidation()) - .withDimensions(properties.getDimensions()) - .withDistanceType(properties.getDistanceType()) - .withRemoveExistingVectorStoreTable(properties.isRemoveExistingVectorStoreTable()) - .withIndexType(properties.getIndexType()) - .withInitializeSchema(initializeSchema) - .withObservationRegistry(observationRegistry.getIfUnique(() -> ObservationRegistry.NOOP)) - .withSearchObservationConvention(customObservationConvention.getIfAvailable(() -> null)) - .withBatchingStrategy(batchingStrategy) - .withMaxDocumentBatchSize(properties.getMaxDocumentBatchSize()) + return PgVectorStore.builder(jdbcTemplate) + .embeddingModel(embeddingModel) + .schemaName(properties.getSchemaName()) + .vectorTableName(properties.getTableName()) + .vectorTableValidationsEnabled(properties.isSchemaValidation()) + .dimensions(properties.getDimensions()) + .distanceType(properties.getDistanceType()) + .removeExistingVectorStoreTable(properties.isRemoveExistingVectorStoreTable()) + .indexType(properties.getIndexType()) + .initializeSchema(initializeSchema) + .observationRegistry(observationRegistry.getIfUnique(() -> ObservationRegistry.NOOP)) + .customObservationConvention(customObservationConvention.getIfAvailable(() -> null)) + .batchingStrategy(batchingStrategy) + .maxDocumentBatchSize(properties.getMaxDocumentBatchSize()) .build(); } diff --git a/spring-ai-spring-boot-autoconfigure/src/main/java/org/springframework/ai/autoconfigure/vectorstore/pgvector/PgVectorStoreProperties.java b/spring-ai-spring-boot-autoconfigure/src/main/java/org/springframework/ai/autoconfigure/vectorstore/pgvector/PgVectorStoreProperties.java index 08d91cfaf21..766a4b88885 100644 --- a/spring-ai-spring-boot-autoconfigure/src/main/java/org/springframework/ai/autoconfigure/vectorstore/pgvector/PgVectorStoreProperties.java +++ b/spring-ai-spring-boot-autoconfigure/src/main/java/org/springframework/ai/autoconfigure/vectorstore/pgvector/PgVectorStoreProperties.java @@ -17,9 +17,9 @@ package org.springframework.ai.autoconfigure.vectorstore.pgvector; import org.springframework.ai.autoconfigure.vectorstore.CommonVectorStoreProperties; -import org.springframework.ai.vectorstore.PgVectorStore; -import org.springframework.ai.vectorstore.PgVectorStore.PgDistanceType; -import org.springframework.ai.vectorstore.PgVectorStore.PgIndexType; +import org.springframework.ai.pg.vectorstore.PgVectorStore; +import org.springframework.ai.pg.vectorstore.PgVectorStore.PgDistanceType; +import org.springframework.ai.pg.vectorstore.PgVectorStore.PgIndexType; import org.springframework.boot.context.properties.ConfigurationProperties; /** diff --git a/spring-ai-spring-boot-autoconfigure/src/test/java/org/springframework/ai/autoconfigure/vectorstore/pgvector/PgVectorStoreAutoConfigurationIT.java b/spring-ai-spring-boot-autoconfigure/src/test/java/org/springframework/ai/autoconfigure/vectorstore/pgvector/PgVectorStoreAutoConfigurationIT.java index e6f12372ba7..fa7b72c3afe 100644 --- a/spring-ai-spring-boot-autoconfigure/src/test/java/org/springframework/ai/autoconfigure/vectorstore/pgvector/PgVectorStoreAutoConfigurationIT.java +++ b/spring-ai-spring-boot-autoconfigure/src/test/java/org/springframework/ai/autoconfigure/vectorstore/pgvector/PgVectorStoreAutoConfigurationIT.java @@ -33,7 +33,7 @@ import org.springframework.ai.embedding.EmbeddingModel; import org.springframework.ai.observation.conventions.VectorStoreProvider; import org.springframework.ai.transformers.TransformersEmbeddingModel; -import org.springframework.ai.vectorstore.PgVectorStore; +import org.springframework.ai.pg.vectorstore.PgVectorStore; import org.springframework.ai.vectorstore.SearchRequest; import org.springframework.ai.vectorstore.observation.VectorStoreObservationContext; import org.springframework.boot.autoconfigure.AutoConfigurations; diff --git a/spring-ai-spring-boot-autoconfigure/src/test/java/org/springframework/ai/autoconfigure/vectorstore/pgvector/PgVectorStorePropertiesTests.java b/spring-ai-spring-boot-autoconfigure/src/test/java/org/springframework/ai/autoconfigure/vectorstore/pgvector/PgVectorStorePropertiesTests.java index f52f4ce782c..cd8b1f8062e 100644 --- a/spring-ai-spring-boot-autoconfigure/src/test/java/org/springframework/ai/autoconfigure/vectorstore/pgvector/PgVectorStorePropertiesTests.java +++ b/spring-ai-spring-boot-autoconfigure/src/test/java/org/springframework/ai/autoconfigure/vectorstore/pgvector/PgVectorStorePropertiesTests.java @@ -18,9 +18,9 @@ import org.junit.jupiter.api.Test; -import org.springframework.ai.vectorstore.PgVectorStore; -import org.springframework.ai.vectorstore.PgVectorStore.PgDistanceType; -import org.springframework.ai.vectorstore.PgVectorStore.PgIndexType; +import org.springframework.ai.pg.vectorstore.PgVectorStore; +import org.springframework.ai.pg.vectorstore.PgVectorStore.PgDistanceType; +import org.springframework.ai.pg.vectorstore.PgVectorStore.PgIndexType; import static org.assertj.core.api.Assertions.assertThat; diff --git a/vector-stores/spring-ai-pgvector-store/src/main/java/org/springframework/ai/vectorstore/PgVectorFilterExpressionConverter.java b/vector-stores/spring-ai-pgvector-store/src/main/java/org/springframework/ai/pg/vectorstore/PgVectorFilterExpressionConverter.java similarity index 98% rename from vector-stores/spring-ai-pgvector-store/src/main/java/org/springframework/ai/vectorstore/PgVectorFilterExpressionConverter.java rename to vector-stores/spring-ai-pgvector-store/src/main/java/org/springframework/ai/pg/vectorstore/PgVectorFilterExpressionConverter.java index 814de5b4c49..d21622f352e 100644 --- a/vector-stores/spring-ai-pgvector-store/src/main/java/org/springframework/ai/vectorstore/PgVectorFilterExpressionConverter.java +++ b/vector-stores/spring-ai-pgvector-store/src/main/java/org/springframework/ai/pg/vectorstore/PgVectorFilterExpressionConverter.java @@ -14,7 +14,7 @@ * limitations under the License. */ -package org.springframework.ai.vectorstore; +package org.springframework.ai.pg.vectorstore; import java.util.List; diff --git a/vector-stores/spring-ai-pgvector-store/src/main/java/org/springframework/ai/vectorstore/PgVectorSchemaValidator.java b/vector-stores/spring-ai-pgvector-store/src/main/java/org/springframework/ai/pg/vectorstore/PgVectorSchemaValidator.java similarity index 99% rename from vector-stores/spring-ai-pgvector-store/src/main/java/org/springframework/ai/vectorstore/PgVectorSchemaValidator.java rename to vector-stores/spring-ai-pgvector-store/src/main/java/org/springframework/ai/pg/vectorstore/PgVectorSchemaValidator.java index b854217f473..de809fcbb35 100644 --- a/vector-stores/spring-ai-pgvector-store/src/main/java/org/springframework/ai/vectorstore/PgVectorSchemaValidator.java +++ b/vector-stores/spring-ai-pgvector-store/src/main/java/org/springframework/ai/pg/vectorstore/PgVectorSchemaValidator.java @@ -14,7 +14,7 @@ * limitations under the License. */ -package org.springframework.ai.vectorstore; +package org.springframework.ai.pg.vectorstore; import java.util.ArrayList; import java.util.List; 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/pg/vectorstore/PgVectorStore.java similarity index 75% rename from vector-stores/spring-ai-pgvector-store/src/main/java/org/springframework/ai/vectorstore/PgVectorStore.java rename to vector-stores/spring-ai-pgvector-store/src/main/java/org/springframework/ai/pg/vectorstore/PgVectorStore.java index f721a5cff4c..9db988ca0e2 100644 --- 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/pg/vectorstore/PgVectorStore.java @@ -14,7 +14,7 @@ * limitations under the License. */ -package org.springframework.ai.vectorstore; +package org.springframework.ai.pg.vectorstore; import java.sql.PreparedStatement; import java.sql.ResultSet; @@ -43,6 +43,9 @@ import org.springframework.ai.observation.conventions.VectorStoreProvider; import org.springframework.ai.observation.conventions.VectorStoreSimilarityMetric; import org.springframework.ai.util.JacksonUtils; +import org.springframework.ai.vectorstore.AbstractVectorStoreBuilder; +import org.springframework.ai.vectorstore.SearchRequest; +import org.springframework.ai.vectorstore.VectorStore; import org.springframework.ai.vectorstore.filter.FilterExpressionConverter; import org.springframework.ai.vectorstore.observation.AbstractObservationVectorStore; import org.springframework.ai.vectorstore.observation.VectorStoreObservationContext; @@ -54,6 +57,7 @@ import org.springframework.jdbc.core.SqlTypeValue; import org.springframework.jdbc.core.StatementCreatorUtils; import org.springframework.lang.Nullable; +import org.springframework.util.Assert; import org.springframework.util.StringUtils; /** @@ -99,8 +103,6 @@ public class PgVectorStore extends AbstractObservationVectorStore implements Ini private final JdbcTemplate jdbcTemplate; - private final EmbeddingModel embeddingModel; - private final String schemaName; private final boolean schemaValidation; @@ -123,77 +125,88 @@ public class PgVectorStore extends AbstractObservationVectorStore implements Ini private final int maxDocumentBatchSize; - public PgVectorStore(JdbcTemplate jdbcTemplate, EmbeddingModel embeddingModel) { - this(jdbcTemplate, embeddingModel, INVALID_EMBEDDING_DIMENSION, PgDistanceType.COSINE_DISTANCE, false, - PgIndexType.NONE, false); - } - - public PgVectorStore(JdbcTemplate jdbcTemplate, EmbeddingModel embeddingModel, int dimensions) { - this(jdbcTemplate, embeddingModel, dimensions, PgDistanceType.COSINE_DISTANCE, false, PgIndexType.NONE, false); - } - - public PgVectorStore(JdbcTemplate jdbcTemplate, EmbeddingModel embeddingModel, int dimensions, - PgDistanceType distanceType, boolean removeExistingVectorStoreTable, PgIndexType createIndexMethod, - boolean initializeSchema) { - - this(DEFAULT_TABLE_NAME, jdbcTemplate, embeddingModel, dimensions, distanceType, removeExistingVectorStoreTable, - createIndexMethod, initializeSchema); - } - - public PgVectorStore(String vectorTableName, JdbcTemplate jdbcTemplate, EmbeddingModel embeddingModel, - int dimensions, PgDistanceType distanceType, boolean removeExistingVectorStoreTable, - PgIndexType createIndexMethod, boolean initializeSchema) { - - this(DEFAULT_SCHEMA_NAME, vectorTableName, DEFAULT_SCHEMA_VALIDATION, jdbcTemplate, embeddingModel, dimensions, - distanceType, removeExistingVectorStoreTable, createIndexMethod, initializeSchema); - } - - private PgVectorStore(String schemaName, String vectorTableName, boolean vectorTableValidationsEnabled, - JdbcTemplate jdbcTemplate, EmbeddingModel embeddingModel, int dimensions, PgDistanceType distanceType, - boolean removeExistingVectorStoreTable, PgIndexType createIndexMethod, boolean initializeSchema) { - - this(schemaName, vectorTableName, vectorTableValidationsEnabled, jdbcTemplate, embeddingModel, dimensions, - distanceType, removeExistingVectorStoreTable, createIndexMethod, initializeSchema, - ObservationRegistry.NOOP, null, new TokenCountBatchingStrategy(), MAX_DOCUMENT_BATCH_SIZE); - } - - private PgVectorStore(String schemaName, String vectorTableName, boolean vectorTableValidationsEnabled, - JdbcTemplate jdbcTemplate, EmbeddingModel embeddingModel, int dimensions, PgDistanceType distanceType, - boolean removeExistingVectorStoreTable, PgIndexType createIndexMethod, boolean initializeSchema, - ObservationRegistry observationRegistry, VectorStoreObservationConvention customObservationConvention, - BatchingStrategy batchingStrategy, int maxDocumentBatchSize) { - - super(observationRegistry, customObservationConvention); + // @Deprecated(forRemoval = true, since = "1.0.0-M5") + // public PgVectorStore(JdbcTemplate jdbcTemplate, EmbeddingModel embeddingModel) { + // this(jdbcTemplate, embeddingModel, INVALID_EMBEDDING_DIMENSION, + // PgDistanceType.COSINE_DISTANCE, false, + // PgIndexType.NONE, false); + // } + // + // @Deprecated(forRemoval = true, since = "1.0.0-M5") + // public PgVectorStore(JdbcTemplate jdbcTemplate, EmbeddingModel embeddingModel, int + // dimensions) { + // this(jdbcTemplate, embeddingModel, dimensions, PgDistanceType.COSINE_DISTANCE, + // false, PgIndexType.NONE, false); + // } + // + // @Deprecated(forRemoval = true, since = "1.0.0-M5") + // public PgVectorStore(JdbcTemplate jdbcTemplate, EmbeddingModel embeddingModel, int + // dimensions, + // PgDistanceType distanceType, boolean removeExistingVectorStoreTable, PgIndexType + // createIndexMethod, + // boolean initializeSchema) { + // + // this(DEFAULT_TABLE_NAME, jdbcTemplate, embeddingModel, dimensions, distanceType, + // removeExistingVectorStoreTable, + // createIndexMethod, initializeSchema); + // } + // + // @Deprecated(forRemoval = true, since = "1.0.0-M5") + // public PgVectorStore(String vectorTableName, JdbcTemplate jdbcTemplate, + // EmbeddingModel embeddingModel, + // int dimensions, PgDistanceType distanceType, boolean + // removeExistingVectorStoreTable, + // PgIndexType createIndexMethod, boolean initializeSchema) { + // + // + // this(builder(jdbcTemplate).schemaName(DEFAULT_SCHEMA_NAME) + // .vectorTableName(vectorTableName) + // .vectorTableValidationsEnabled(DEFAULT_SCHEMA_VALIDATION) + // .dimensions(dimensions) + // .distanceType(distanceType) + // .removeExistingVectorStoreTable(removeExistingVectorStoreTable) + // .indexType(createIndexMethod) + // .initializeSchema(initializeSchema)); + // } + /** + * @param builder {@link VectorStore.Builder} for pg vector store + */ + private PgVectorStore(PgVectorStoreBuilder builder) { + super(builder); this.objectMapper = JsonMapper.builder().addModules(JacksonUtils.instantiateAvailableModules()).build(); - this.vectorTableName = (null == vectorTableName || vectorTableName.isEmpty()) ? DEFAULT_TABLE_NAME - : vectorTableName.trim(); + String vectorTableName1 = builder.vectorTableName; + this.vectorTableName = (null == vectorTableName1 || vectorTableName1.isEmpty()) ? DEFAULT_TABLE_NAME + : vectorTableName1.trim(); logger.info("Using the vector table name: {}. Is empty: {}", this.vectorTableName, - (vectorTableName == null || vectorTableName.isEmpty())); + (this.vectorTableName == null || this.vectorTableName.isEmpty())); this.vectorIndexName = this.vectorTableName.equals(DEFAULT_TABLE_NAME) ? DEFAULT_VECTOR_INDEX_NAME : this.vectorTableName + "_index"; - this.schemaName = schemaName; - this.schemaValidation = vectorTableValidationsEnabled; - - this.jdbcTemplate = jdbcTemplate; - this.embeddingModel = embeddingModel; - this.dimensions = dimensions; - this.distanceType = distanceType; - this.removeExistingVectorStoreTable = removeExistingVectorStoreTable; - this.createIndexMethod = createIndexMethod; - this.initializeSchema = initializeSchema; - this.schemaValidator = new PgVectorSchemaValidator(jdbcTemplate); - this.batchingStrategy = batchingStrategy; - this.maxDocumentBatchSize = maxDocumentBatchSize; + this.schemaName = builder.schemaName; + this.schemaValidation = builder.vectorTableValidationsEnabled; + + this.jdbcTemplate = builder.jdbcTemplate; + this.dimensions = builder.dimensions; + this.distanceType = builder.distanceType; + this.removeExistingVectorStoreTable = builder.removeExistingVectorStoreTable; + this.createIndexMethod = builder.indexType; + this.initializeSchema = builder.initializeSchema; + this.schemaValidator = new PgVectorSchemaValidator(this.jdbcTemplate); + this.batchingStrategy = builder.batchingStrategy; + this.maxDocumentBatchSize = builder.maxDocumentBatchSize; } public PgDistanceType getDistanceType() { return this.distanceType; } + public static PgVectorStoreBuilder builder(JdbcTemplate jdbcTemplate) { + return new PgVectorStoreBuilder(jdbcTemplate); + } + @Override public void doAdd(List documents) { List embeddings = this.embeddingModel.embed(documents, EmbeddingOptionsBuilder.builder().build(), @@ -527,6 +540,93 @@ private Map toMap(PGobject pgObject) { } + public static class PgVectorStoreBuilder extends AbstractVectorStoreBuilder { + + private final JdbcTemplate jdbcTemplate; + + private String schemaName = PgVectorStore.DEFAULT_SCHEMA_NAME; + + private String vectorTableName = PgVectorStore.DEFAULT_TABLE_NAME; + + private boolean vectorTableValidationsEnabled = PgVectorStore.DEFAULT_SCHEMA_VALIDATION; + + private int dimensions = PgVectorStore.INVALID_EMBEDDING_DIMENSION; + + private PgDistanceType distanceType = PgDistanceType.COSINE_DISTANCE; + + private boolean removeExistingVectorStoreTable = false; + + private PgIndexType indexType = PgIndexType.HNSW; + + private boolean initializeSchema; + + private BatchingStrategy batchingStrategy = new TokenCountBatchingStrategy(); + + private int maxDocumentBatchSize = MAX_DOCUMENT_BATCH_SIZE; + + public PgVectorStoreBuilder(JdbcTemplate jdbcTemplate) { + Assert.notNull(jdbcTemplate, "JdbcTemplate must not be null"); + this.jdbcTemplate = jdbcTemplate; + } + + public PgVectorStoreBuilder schemaName(String schemaName) { + this.schemaName = schemaName; + return this; + } + + public PgVectorStoreBuilder vectorTableName(String vectorTableName) { + this.vectorTableName = vectorTableName; + return this; + } + + public PgVectorStoreBuilder vectorTableValidationsEnabled(boolean vectorTableValidationsEnabled) { + this.vectorTableValidationsEnabled = vectorTableValidationsEnabled; + return this; + } + + public PgVectorStoreBuilder dimensions(int dimensions) { + this.dimensions = dimensions; + return this; + } + + public PgVectorStoreBuilder distanceType(PgDistanceType distanceType) { + this.distanceType = distanceType; + return this; + } + + public PgVectorStoreBuilder removeExistingVectorStoreTable(boolean removeExistingVectorStoreTable) { + this.removeExistingVectorStoreTable = removeExistingVectorStoreTable; + return this; + } + + public PgVectorStoreBuilder indexType(PgIndexType indexType) { + this.indexType = indexType; + return this; + } + + public PgVectorStoreBuilder initializeSchema(boolean initializeSchema) { + this.initializeSchema = initializeSchema; + return this; + } + + public PgVectorStoreBuilder batchingStrategy(BatchingStrategy batchingStrategy) { + this.batchingStrategy = batchingStrategy; + return this; + } + + public PgVectorStoreBuilder maxDocumentBatchSize(int maxDocumentBatchSize) { + this.maxDocumentBatchSize = maxDocumentBatchSize; + return this; + } + + public PgVectorStore build() { + validate(); + return new PgVectorStore(this); + } + + } + + @Deprecated(forRemoval = true, since = "1.0.0-M5") public static class Builder { private final JdbcTemplate jdbcTemplate; @@ -558,7 +658,6 @@ public static class Builder { @Nullable private VectorStoreObservationConvention searchObservationConvention; - // Builder constructor with mandatory parameters public Builder(JdbcTemplate jdbcTemplate, EmbeddingModel embeddingModel) { if (jdbcTemplate == null || embeddingModel == null) { throw new IllegalArgumentException("JdbcTemplate and EmbeddingModel must not be null"); @@ -628,11 +727,19 @@ public Builder withMaxDocumentBatchSize(int maxDocumentBatchSize) { } public PgVectorStore build() { - return new PgVectorStore(this.schemaName, this.vectorTableName, this.vectorTableValidationsEnabled, - this.jdbcTemplate, this.embeddingModel, this.dimensions, this.distanceType, - this.removeExistingVectorStoreTable, this.indexType, this.initializeSchema, - this.observationRegistry, this.searchObservationConvention, this.batchingStrategy, - this.maxDocumentBatchSize); + return PgVectorStore.builder(this.jdbcTemplate) + .embeddingModel(this.embeddingModel) + .schemaName(this.schemaName) + .vectorTableName(this.vectorTableName) + .vectorTableValidationsEnabled(this.vectorTableValidationsEnabled) + .dimensions(this.dimensions) + .distanceType(this.distanceType) + .removeExistingVectorStoreTable(this.removeExistingVectorStoreTable) + .indexType(this.indexType) + .initializeSchema(this.initializeSchema) + .batchingStrategy(this.batchingStrategy) + .maxDocumentBatchSize(this.maxDocumentBatchSize) + .build(); } } diff --git a/vector-stores/spring-ai-pgvector-store/src/main/java/org/springframework/ai/pg/vectorstore/package-info.java b/vector-stores/spring-ai-pgvector-store/src/main/java/org/springframework/ai/pg/vectorstore/package-info.java new file mode 100644 index 00000000000..04b28f34280 --- /dev/null +++ b/vector-stores/spring-ai-pgvector-store/src/main/java/org/springframework/ai/pg/vectorstore/package-info.java @@ -0,0 +1,25 @@ +/* + * Copyright 2023-2024 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. + */ + +/** + * Provides the API for embedding observations. + */ +@NonNullApi +@NonNullFields +package org.springframework.ai.pg.vectorstore; + +import org.springframework.lang.NonNullApi; +import org.springframework.lang.NonNullFields; diff --git a/vector-stores/spring-ai-pgvector-store/src/test/java/org/springframework/ai/vectorstore/PgVectorEmbeddingDimensionsTests.java b/vector-stores/spring-ai-pgvector-store/src/test/java/org/springframework/ai/pg/vectorstore/PgVectorEmbeddingDimensionsTests.java similarity index 76% rename from vector-stores/spring-ai-pgvector-store/src/test/java/org/springframework/ai/vectorstore/PgVectorEmbeddingDimensionsTests.java rename to vector-stores/spring-ai-pgvector-store/src/test/java/org/springframework/ai/pg/vectorstore/PgVectorEmbeddingDimensionsTests.java index 587b2cb8e02..601b2b70df0 100644 --- a/vector-stores/spring-ai-pgvector-store/src/test/java/org/springframework/ai/vectorstore/PgVectorEmbeddingDimensionsTests.java +++ b/vector-stores/spring-ai-pgvector-store/src/test/java/org/springframework/ai/pg/vectorstore/PgVectorEmbeddingDimensionsTests.java @@ -14,7 +14,7 @@ * limitations under the License. */ -package org.springframework.ai.vectorstore; +package org.springframework.ai.pg.vectorstore; import org.junit.jupiter.api.Test; import org.junit.jupiter.api.extension.ExtendWith; @@ -47,7 +47,11 @@ public void explicitlySetDimensions() { final int explicitDimensions = 696; - var dim = new PgVectorStore(this.jdbcTemplate, this.embeddingModel, explicitDimensions).embeddingDimensions(); + PgVectorStore pgVectorStore = PgVectorStore.builder(this.jdbcTemplate) + .embeddingModel(this.embeddingModel) + .dimensions(explicitDimensions) + .build(); + var dim = pgVectorStore.embeddingDimensions(); assertThat(dim).isEqualTo(explicitDimensions); verify(this.embeddingModel, never()).dimensions(); @@ -57,7 +61,10 @@ public void explicitlySetDimensions() { public void embeddingModelDimensions() { given(this.embeddingModel.dimensions()).willReturn(969); - var dim = new PgVectorStore(this.jdbcTemplate, this.embeddingModel).embeddingDimensions(); + PgVectorStore pgVectorStore = PgVectorStore.builder(this.jdbcTemplate) + .embeddingModel(this.embeddingModel) + .build(); + var dim = pgVectorStore.embeddingDimensions(); assertThat(dim).isEqualTo(969); @@ -69,7 +76,10 @@ public void fallBackToDefaultDimensions() { given(this.embeddingModel.dimensions()).willThrow(new RuntimeException()); - var dim = new PgVectorStore(this.jdbcTemplate, this.embeddingModel).embeddingDimensions(); + PgVectorStore pgVectorStore = PgVectorStore.builder(this.jdbcTemplate) + .embeddingModel(this.embeddingModel) + .build(); + var dim = pgVectorStore.embeddingDimensions(); assertThat(dim).isEqualTo(PgVectorStore.OPENAI_EMBEDDING_DIMENSION_SIZE); verify(this.embeddingModel, only()).dimensions(); diff --git a/vector-stores/spring-ai-pgvector-store/src/test/java/org/springframework/ai/vectorstore/PgVectorFilterExpressionConverterTests.java b/vector-stores/spring-ai-pgvector-store/src/test/java/org/springframework/ai/pg/vectorstore/PgVectorFilterExpressionConverterTests.java similarity index 99% rename from vector-stores/spring-ai-pgvector-store/src/test/java/org/springframework/ai/vectorstore/PgVectorFilterExpressionConverterTests.java rename to vector-stores/spring-ai-pgvector-store/src/test/java/org/springframework/ai/pg/vectorstore/PgVectorFilterExpressionConverterTests.java index b2ef770cca8..c3344756f6a 100644 --- a/vector-stores/spring-ai-pgvector-store/src/test/java/org/springframework/ai/vectorstore/PgVectorFilterExpressionConverterTests.java +++ b/vector-stores/spring-ai-pgvector-store/src/test/java/org/springframework/ai/pg/vectorstore/PgVectorFilterExpressionConverterTests.java @@ -14,7 +14,7 @@ * limitations under the License. */ -package org.springframework.ai.vectorstore; +package org.springframework.ai.pg.vectorstore; import java.util.List; diff --git a/vector-stores/spring-ai-pgvector-store/src/test/java/org/springframework/ai/vectorstore/PgVectorImage.java b/vector-stores/spring-ai-pgvector-store/src/test/java/org/springframework/ai/pg/vectorstore/PgVectorImage.java similarity index 94% rename from vector-stores/spring-ai-pgvector-store/src/test/java/org/springframework/ai/vectorstore/PgVectorImage.java rename to vector-stores/spring-ai-pgvector-store/src/test/java/org/springframework/ai/pg/vectorstore/PgVectorImage.java index 2dead75a532..3f6e57ef7ac 100644 --- a/vector-stores/spring-ai-pgvector-store/src/test/java/org/springframework/ai/vectorstore/PgVectorImage.java +++ b/vector-stores/spring-ai-pgvector-store/src/test/java/org/springframework/ai/pg/vectorstore/PgVectorImage.java @@ -14,7 +14,7 @@ * limitations under the License. */ -package org.springframework.ai.vectorstore; +package org.springframework.ai.pg.vectorstore; import org.testcontainers.utility.DockerImageName; diff --git a/vector-stores/spring-ai-pgvector-store/src/test/java/org/springframework/ai/vectorstore/PgVectorStoreCustomNamesIT.java b/vector-stores/spring-ai-pgvector-store/src/test/java/org/springframework/ai/pg/vectorstore/PgVectorStoreCustomNamesIT.java similarity index 93% rename from vector-stores/spring-ai-pgvector-store/src/test/java/org/springframework/ai/vectorstore/PgVectorStoreCustomNamesIT.java rename to vector-stores/spring-ai-pgvector-store/src/test/java/org/springframework/ai/pg/vectorstore/PgVectorStoreCustomNamesIT.java index 3741cf5a453..b0b7e609ba6 100644 --- a/vector-stores/spring-ai-pgvector-store/src/test/java/org/springframework/ai/vectorstore/PgVectorStoreCustomNamesIT.java +++ b/vector-stores/spring-ai-pgvector-store/src/test/java/org/springframework/ai/pg/vectorstore/PgVectorStoreCustomNamesIT.java @@ -14,7 +14,7 @@ * limitations under the License. */ -package org.springframework.ai.vectorstore; +package org.springframework.ai.pg.vectorstore; import java.util.Random; @@ -30,7 +30,8 @@ import org.springframework.ai.embedding.EmbeddingModel; import org.springframework.ai.openai.OpenAiEmbeddingModel; import org.springframework.ai.openai.api.OpenAiApi; -import org.springframework.ai.vectorstore.PgVectorStore.PgIndexType; +import org.springframework.ai.pg.vectorstore.PgVectorStore.PgIndexType; +import org.springframework.ai.vectorstore.VectorStore; import org.springframework.beans.factory.annotation.Value; import org.springframework.boot.SpringBootConfiguration; import org.springframework.boot.autoconfigure.EnableAutoConfiguration; @@ -195,14 +196,16 @@ public static class TestApplication { @Bean public VectorStore vectorStore(JdbcTemplate jdbcTemplate, EmbeddingModel embeddingModel) { - return new PgVectorStore.Builder(jdbcTemplate, embeddingModel).withSchemaName(this.schemaName) - .withVectorTableName(this.vectorTableName) - .withVectorTableValidationsEnabled(this.schemaValidation) - .withDimensions(this.dimensions) - .withDistanceType(PgVectorStore.PgDistanceType.COSINE_DISTANCE) - .withRemoveExistingVectorStoreTable(true) - .withIndexType(PgIndexType.HNSW) - .withInitializeSchema(true) + return PgVectorStore.builder(jdbcTemplate) + .embeddingModel(embeddingModel) + .schemaName(this.schemaName) + .vectorTableName(this.vectorTableName) + .vectorTableValidationsEnabled(this.schemaValidation) + .dimensions(this.dimensions) + .distanceType(PgVectorStore.PgDistanceType.COSINE_DISTANCE) + .removeExistingVectorStoreTable(true) + .indexType(PgIndexType.HNSW) + .initializeSchema(true) .build(); } 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/pg/vectorstore/PgVectorStoreIT.java similarity index 96% rename from vector-stores/spring-ai-pgvector-store/src/test/java/org/springframework/ai/vectorstore/PgVectorStoreIT.java rename to vector-stores/spring-ai-pgvector-store/src/test/java/org/springframework/ai/pg/vectorstore/PgVectorStoreIT.java index 83c2e49e782..547d3ce7e48 100644 --- 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/pg/vectorstore/PgVectorStoreIT.java @@ -14,7 +14,7 @@ * limitations under the License. */ -package org.springframework.ai.vectorstore; +package org.springframework.ai.pg.vectorstore; import java.io.IOException; import java.nio.charset.StandardCharsets; @@ -43,7 +43,9 @@ import org.springframework.ai.embedding.EmbeddingModel; import org.springframework.ai.openai.OpenAiEmbeddingModel; import org.springframework.ai.openai.api.OpenAiApi; -import org.springframework.ai.vectorstore.PgVectorStore.PgIndexType; +import org.springframework.ai.pg.vectorstore.PgVectorStore.PgIndexType; +import org.springframework.ai.vectorstore.SearchRequest; +import org.springframework.ai.vectorstore.VectorStore; import org.springframework.ai.vectorstore.filter.FilterExpressionTextParser.FilterExpressionParseException; import org.springframework.beans.factory.annotation.Value; import org.springframework.boot.SpringBootConfiguration; @@ -354,8 +356,14 @@ public static class TestApplication { @Bean public VectorStore vectorStore(JdbcTemplate jdbcTemplate, EmbeddingModel embeddingModel) { - return new PgVectorStore(jdbcTemplate, embeddingModel, PgVectorStore.INVALID_EMBEDDING_DIMENSION, - this.distanceType, true, PgIndexType.HNSW, true); + return PgVectorStore.builder(jdbcTemplate) + .embeddingModel(embeddingModel) + .dimensions(PgVectorStore.INVALID_EMBEDDING_DIMENSION) + .distanceType(this.distanceType) + .initializeSchema(true) + .indexType(PgIndexType.HNSW) + .removeExistingVectorStoreTable(true) + .build(); } @Bean diff --git a/vector-stores/spring-ai-pgvector-store/src/test/java/org/springframework/ai/vectorstore/PgVectorStoreObservationIT.java b/vector-stores/spring-ai-pgvector-store/src/test/java/org/springframework/ai/pg/vectorstore/PgVectorStoreObservationIT.java similarity index 94% rename from vector-stores/spring-ai-pgvector-store/src/test/java/org/springframework/ai/vectorstore/PgVectorStoreObservationIT.java rename to vector-stores/spring-ai-pgvector-store/src/test/java/org/springframework/ai/pg/vectorstore/PgVectorStoreObservationIT.java index 963b256e915..d5efd762a32 100644 --- a/vector-stores/spring-ai-pgvector-store/src/test/java/org/springframework/ai/vectorstore/PgVectorStoreObservationIT.java +++ b/vector-stores/spring-ai-pgvector-store/src/test/java/org/springframework/ai/pg/vectorstore/PgVectorStoreObservationIT.java @@ -14,7 +14,7 @@ * limitations under the License. */ -package org.springframework.ai.vectorstore; +package org.springframework.ai.pg.vectorstore; import java.io.IOException; import java.nio.charset.StandardCharsets; @@ -41,7 +41,9 @@ import org.springframework.ai.openai.OpenAiChatModel; import org.springframework.ai.openai.OpenAiEmbeddingModel; import org.springframework.ai.openai.api.OpenAiApi; -import org.springframework.ai.vectorstore.PgVectorStore.PgIndexType; +import org.springframework.ai.pg.vectorstore.PgVectorStore.PgIndexType; +import org.springframework.ai.vectorstore.SearchRequest; +import org.springframework.ai.vectorstore.VectorStore; import org.springframework.ai.vectorstore.observation.DefaultVectorStoreObservationConvention; import org.springframework.ai.vectorstore.observation.VectorStoreObservationDocumentation.HighCardinalityKeyNames; import org.springframework.ai.vectorstore.observation.VectorStoreObservationDocumentation.LowCardinalityKeyNames; @@ -185,11 +187,12 @@ public TestObservationRegistry observationRegistry() { @Bean public VectorStore vectorStore(JdbcTemplate jdbcTemplate, EmbeddingModel embeddingModel, ObservationRegistry observationRegistry) { - return new PgVectorStore.Builder(jdbcTemplate, embeddingModel) - .withDistanceType(PgVectorStore.PgDistanceType.COSINE_DISTANCE) - .withIndexType(PgIndexType.HNSW) - .withObservationRegistry(observationRegistry) - .withInitializeSchema(true) + return PgVectorStore.builder(jdbcTemplate) + .embeddingModel(embeddingModel) + .distanceType(PgVectorStore.PgDistanceType.COSINE_DISTANCE) + .indexType(PgIndexType.HNSW) + .observationRegistry(observationRegistry) + .initializeSchema(true) .build(); } diff --git a/vector-stores/spring-ai-pgvector-store/src/test/java/org/springframework/ai/vectorstore/PgVectorStoreTests.java b/vector-stores/spring-ai-pgvector-store/src/test/java/org/springframework/ai/pg/vectorstore/PgVectorStoreTests.java similarity index 95% rename from vector-stores/spring-ai-pgvector-store/src/test/java/org/springframework/ai/vectorstore/PgVectorStoreTests.java rename to vector-stores/spring-ai-pgvector-store/src/test/java/org/springframework/ai/pg/vectorstore/PgVectorStoreTests.java index 993f8856ead..c16e25a741f 100644 --- a/vector-stores/spring-ai-pgvector-store/src/test/java/org/springframework/ai/vectorstore/PgVectorStoreTests.java +++ b/vector-stores/spring-ai-pgvector-store/src/test/java/org/springframework/ai/pg/vectorstore/PgVectorStoreTests.java @@ -14,7 +14,7 @@ * limitations under the License. */ -package org.springframework.ai.vectorstore; +package org.springframework.ai.pg.vectorstore; import java.util.Collections; @@ -79,7 +79,9 @@ void shouldAddDocumentsInBatchesAndEmbedOnce() { // Given var jdbcTemplate = mock(JdbcTemplate.class); var embeddingModel = mock(EmbeddingModel.class); - var pgVectorStore = new PgVectorStore.Builder(jdbcTemplate, embeddingModel).withMaxDocumentBatchSize(1000) + var pgVectorStore = PgVectorStore.builder(jdbcTemplate) + .embeddingModel(embeddingModel) + .maxDocumentBatchSize(1000) .build(); // Testing with 9989 documents diff --git a/vector-stores/spring-ai-pgvector-store/src/test/java/org/springframework/ai/vectorstore/PgVectorStoreWithChatMemoryAdvisorIT.java b/vector-stores/spring-ai-pgvector-store/src/test/java/org/springframework/ai/pg/vectorstore/PgVectorStoreWithChatMemoryAdvisorIT.java similarity index 99% rename from vector-stores/spring-ai-pgvector-store/src/test/java/org/springframework/ai/vectorstore/PgVectorStoreWithChatMemoryAdvisorIT.java rename to vector-stores/spring-ai-pgvector-store/src/test/java/org/springframework/ai/pg/vectorstore/PgVectorStoreWithChatMemoryAdvisorIT.java index 0765e9cab79..c7adc74d5c9 100644 --- a/vector-stores/spring-ai-pgvector-store/src/test/java/org/springframework/ai/vectorstore/PgVectorStoreWithChatMemoryAdvisorIT.java +++ b/vector-stores/spring-ai-pgvector-store/src/test/java/org/springframework/ai/pg/vectorstore/PgVectorStoreWithChatMemoryAdvisorIT.java @@ -14,7 +14,7 @@ * limitations under the License. */ -package org.springframework.ai.vectorstore; +package org.springframework.ai.pg.vectorstore; import java.util.List; import java.util.Map; From 94bf6638deee954ed1ed89e79dc33fd440cc4a6d Mon Sep 17 00:00:00 2001 From: Soby Chacko Date: Tue, 10 Dec 2024 20:49:31 -0500 Subject: [PATCH 2/2] PgVectorStore builder refactoring --- .../PgVectorStoreAutoConfiguration.java | 3 +- .../ai/pg/vectorstore/PgVectorStore.java | 100 +++++++++--------- .../PgVectorEmbeddingDimensionsTests.java | 9 +- .../PgVectorStoreCustomNamesIT.java | 3 +- .../ai/pg/vectorstore/PgVectorStoreIT.java | 3 +- .../PgVectorStoreObservationIT.java | 3 +- .../ai/pg/vectorstore/PgVectorStoreTests.java | 3 +- 7 files changed, 64 insertions(+), 60 deletions(-) diff --git a/spring-ai-spring-boot-autoconfigure/src/main/java/org/springframework/ai/autoconfigure/vectorstore/pgvector/PgVectorStoreAutoConfiguration.java b/spring-ai-spring-boot-autoconfigure/src/main/java/org/springframework/ai/autoconfigure/vectorstore/pgvector/PgVectorStoreAutoConfiguration.java index 2b6a8974c28..5211ca08feb 100644 --- a/spring-ai-spring-boot-autoconfigure/src/main/java/org/springframework/ai/autoconfigure/vectorstore/pgvector/PgVectorStoreAutoConfiguration.java +++ b/spring-ai-spring-boot-autoconfigure/src/main/java/org/springframework/ai/autoconfigure/vectorstore/pgvector/PgVectorStoreAutoConfiguration.java @@ -62,7 +62,8 @@ public PgVectorStore vectorStore(JdbcTemplate jdbcTemplate, EmbeddingModel embed var initializeSchema = properties.isInitializeSchema(); - return PgVectorStore.builder(jdbcTemplate) + return PgVectorStore.builder() + .jdbcTemplate(jdbcTemplate) .embeddingModel(embeddingModel) .schemaName(properties.getSchemaName()) .vectorTableName(properties.getTableName()) diff --git a/vector-stores/spring-ai-pgvector-store/src/main/java/org/springframework/ai/pg/vectorstore/PgVectorStore.java b/vector-stores/spring-ai-pgvector-store/src/main/java/org/springframework/ai/pg/vectorstore/PgVectorStore.java index 9db988ca0e2..26403585eed 100644 --- a/vector-stores/spring-ai-pgvector-store/src/main/java/org/springframework/ai/pg/vectorstore/PgVectorStore.java +++ b/vector-stores/spring-ai-pgvector-store/src/main/java/org/springframework/ai/pg/vectorstore/PgVectorStore.java @@ -125,60 +125,54 @@ public class PgVectorStore extends AbstractObservationVectorStore implements Ini private final int maxDocumentBatchSize; - // @Deprecated(forRemoval = true, since = "1.0.0-M5") - // public PgVectorStore(JdbcTemplate jdbcTemplate, EmbeddingModel embeddingModel) { - // this(jdbcTemplate, embeddingModel, INVALID_EMBEDDING_DIMENSION, - // PgDistanceType.COSINE_DISTANCE, false, - // PgIndexType.NONE, false); - // } - // - // @Deprecated(forRemoval = true, since = "1.0.0-M5") - // public PgVectorStore(JdbcTemplate jdbcTemplate, EmbeddingModel embeddingModel, int - // dimensions) { - // this(jdbcTemplate, embeddingModel, dimensions, PgDistanceType.COSINE_DISTANCE, - // false, PgIndexType.NONE, false); - // } - // - // @Deprecated(forRemoval = true, since = "1.0.0-M5") - // public PgVectorStore(JdbcTemplate jdbcTemplate, EmbeddingModel embeddingModel, int - // dimensions, - // PgDistanceType distanceType, boolean removeExistingVectorStoreTable, PgIndexType - // createIndexMethod, - // boolean initializeSchema) { - // - // this(DEFAULT_TABLE_NAME, jdbcTemplate, embeddingModel, dimensions, distanceType, - // removeExistingVectorStoreTable, - // createIndexMethod, initializeSchema); - // } - // - // @Deprecated(forRemoval = true, since = "1.0.0-M5") - // public PgVectorStore(String vectorTableName, JdbcTemplate jdbcTemplate, - // EmbeddingModel embeddingModel, - // int dimensions, PgDistanceType distanceType, boolean - // removeExistingVectorStoreTable, - // PgIndexType createIndexMethod, boolean initializeSchema) { - // - // - // this(builder(jdbcTemplate).schemaName(DEFAULT_SCHEMA_NAME) - // .vectorTableName(vectorTableName) - // .vectorTableValidationsEnabled(DEFAULT_SCHEMA_VALIDATION) - // .dimensions(dimensions) - // .distanceType(distanceType) - // .removeExistingVectorStoreTable(removeExistingVectorStoreTable) - // .indexType(createIndexMethod) - // .initializeSchema(initializeSchema)); - // } + @Deprecated(forRemoval = true, since = "1.0.0-M5") + public PgVectorStore(JdbcTemplate jdbcTemplate, EmbeddingModel embeddingModel) { + this(jdbcTemplate, embeddingModel, INVALID_EMBEDDING_DIMENSION, PgDistanceType.COSINE_DISTANCE, false, + PgIndexType.NONE, false); + } + + @Deprecated(forRemoval = true, since = "1.0.0-M5") + public PgVectorStore(JdbcTemplate jdbcTemplate, EmbeddingModel embeddingModel, int dimensions) { + this(jdbcTemplate, embeddingModel, dimensions, PgDistanceType.COSINE_DISTANCE, false, PgIndexType.NONE, false); + } + + @Deprecated(forRemoval = true, since = "1.0.0-M5") + public PgVectorStore(JdbcTemplate jdbcTemplate, EmbeddingModel embeddingModel, int dimensions, + PgDistanceType distanceType, boolean removeExistingVectorStoreTable, PgIndexType createIndexMethod, + boolean initializeSchema) { + + this(DEFAULT_TABLE_NAME, jdbcTemplate, embeddingModel, dimensions, distanceType, removeExistingVectorStoreTable, + createIndexMethod, initializeSchema); + } + + @Deprecated(forRemoval = true, since = "1.0.0-M5") + public PgVectorStore(String vectorTableName, JdbcTemplate jdbcTemplate, EmbeddingModel embeddingModel, + int dimensions, PgDistanceType distanceType, boolean removeExistingVectorStoreTable, + PgIndexType createIndexMethod, boolean initializeSchema) { + + this(builder().jdbcTemplate(jdbcTemplate) + .schemaName(DEFAULT_SCHEMA_NAME) + .vectorTableName(vectorTableName) + .vectorTableValidationsEnabled(DEFAULT_SCHEMA_VALIDATION) + .dimensions(dimensions) + .distanceType(distanceType) + .removeExistingVectorStoreTable(removeExistingVectorStoreTable) + .indexType(createIndexMethod) + .initializeSchema(initializeSchema)); + } /** * @param builder {@link VectorStore.Builder} for pg vector store */ - private PgVectorStore(PgVectorStoreBuilder builder) { + protected PgVectorStore(PgVectorStoreBuilder builder) { super(builder); + + Assert.notNull(builder.jdbcTemplate, "JdbcTemplate must not be null"); + this.objectMapper = JsonMapper.builder().addModules(JacksonUtils.instantiateAvailableModules()).build(); - String vectorTableName1 = builder.vectorTableName; - this.vectorTableName = (null == vectorTableName1 || vectorTableName1.isEmpty()) ? DEFAULT_TABLE_NAME - : vectorTableName1.trim(); + String vectorTable = builder.vectorTableName; + this.vectorTableName = (null == vectorTable || vectorTable.isEmpty()) ? DEFAULT_TABLE_NAME : vectorTable.trim(); logger.info("Using the vector table name: {}. Is empty: {}", this.vectorTableName, (this.vectorTableName == null || this.vectorTableName.isEmpty())); @@ -203,8 +197,8 @@ public PgDistanceType getDistanceType() { return this.distanceType; } - public static PgVectorStoreBuilder builder(JdbcTemplate jdbcTemplate) { - return new PgVectorStoreBuilder(jdbcTemplate); + public static PgVectorStoreBuilder builder() { + return new PgVectorStoreBuilder(); } @Override @@ -542,7 +536,7 @@ private Map toMap(PGobject pgObject) { public static class PgVectorStoreBuilder extends AbstractVectorStoreBuilder { - private final JdbcTemplate jdbcTemplate; + private JdbcTemplate jdbcTemplate; private String schemaName = PgVectorStore.DEFAULT_SCHEMA_NAME; @@ -564,9 +558,10 @@ public static class PgVectorStoreBuilder extends AbstractVectorStoreBuilder