Skip to content

Commit 94bf663

Browse files
sobychackosobychacko
committed
PgVectorStore builder refactoring
1 parent 3c8817f commit 94bf663

File tree

7 files changed

+64
-60
lines changed

7 files changed

+64
-60
lines changed

‎spring-ai-spring-boot-autoconfigure/src/main/java/org/springframework/ai/autoconfigure/vectorstore/pgvector/PgVectorStoreAutoConfiguration.java‎

Lines changed: 2 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -62,7 +62,8 @@ public PgVectorStore vectorStore(JdbcTemplate jdbcTemplate, EmbeddingModel embed
6262

6363
var initializeSchema = properties.isInitializeSchema();
6464

65-
return PgVectorStore.builder(jdbcTemplate)
65+
return PgVectorStore.builder()
66+
.jdbcTemplate(jdbcTemplate)
6667
.embeddingModel(embeddingModel)
6768
.schemaName(properties.getSchemaName())
6869
.vectorTableName(properties.getTableName())

‎vector-stores/spring-ai-pgvector-store/src/main/java/org/springframework/ai/pg/vectorstore/PgVectorStore.java‎

Lines changed: 48 additions & 52 deletions
Original file line numberDiff line numberDiff line change
@@ -125,60 +125,54 @@ public class PgVectorStore extends AbstractObservationVectorStore implements Ini
125125

126126
private final int maxDocumentBatchSize;
127127

128-
// @Deprecated(forRemoval = true, since = "1.0.0-M5")
129-
// public PgVectorStore(JdbcTemplate jdbcTemplate, EmbeddingModel embeddingModel) {
130-
// this(jdbcTemplate, embeddingModel, INVALID_EMBEDDING_DIMENSION,
131-
// PgDistanceType.COSINE_DISTANCE, false,
132-
// PgIndexType.NONE, false);
133-
// }
134-
//
135-
// @Deprecated(forRemoval = true, since = "1.0.0-M5")
136-
// public PgVectorStore(JdbcTemplate jdbcTemplate, EmbeddingModel embeddingModel, int
137-
// dimensions) {
138-
// this(jdbcTemplate, embeddingModel, dimensions, PgDistanceType.COSINE_DISTANCE,
139-
// false, PgIndexType.NONE, false);
140-
// }
141-
//
142-
// @Deprecated(forRemoval = true, since = "1.0.0-M5")
143-
// public PgVectorStore(JdbcTemplate jdbcTemplate, EmbeddingModel embeddingModel, int
144-
// dimensions,
145-
// PgDistanceType distanceType, boolean removeExistingVectorStoreTable, PgIndexType
146-
// createIndexMethod,
147-
// boolean initializeSchema) {
148-
//
149-
// this(DEFAULT_TABLE_NAME, jdbcTemplate, embeddingModel, dimensions, distanceType,
150-
// removeExistingVectorStoreTable,
151-
// createIndexMethod, initializeSchema);
152-
// }
153-
//
154-
// @Deprecated(forRemoval = true, since = "1.0.0-M5")
155-
// public PgVectorStore(String vectorTableName, JdbcTemplate jdbcTemplate,
156-
// EmbeddingModel embeddingModel,
157-
// int dimensions, PgDistanceType distanceType, boolean
158-
// removeExistingVectorStoreTable,
159-
// PgIndexType createIndexMethod, boolean initializeSchema) {
160-
//
161-
//
162-
// this(builder(jdbcTemplate).schemaName(DEFAULT_SCHEMA_NAME)
163-
// .vectorTableName(vectorTableName)
164-
// .vectorTableValidationsEnabled(DEFAULT_SCHEMA_VALIDATION)
165-
// .dimensions(dimensions)
166-
// .distanceType(distanceType)
167-
// .removeExistingVectorStoreTable(removeExistingVectorStoreTable)
168-
// .indexType(createIndexMethod)
169-
// .initializeSchema(initializeSchema));
170-
// }
128+
@Deprecated(forRemoval = true, since = "1.0.0-M5")
129+
public PgVectorStore(JdbcTemplate jdbcTemplate, EmbeddingModel embeddingModel) {
130+
this(jdbcTemplate, embeddingModel, INVALID_EMBEDDING_DIMENSION, PgDistanceType.COSINE_DISTANCE, false,
131+
PgIndexType.NONE, false);
132+
}
133+
134+
@Deprecated(forRemoval = true, since = "1.0.0-M5")
135+
public PgVectorStore(JdbcTemplate jdbcTemplate, EmbeddingModel embeddingModel, int dimensions) {
136+
this(jdbcTemplate, embeddingModel, dimensions, PgDistanceType.COSINE_DISTANCE, false, PgIndexType.NONE, false);
137+
}
138+
139+
@Deprecated(forRemoval = true, since = "1.0.0-M5")
140+
public PgVectorStore(JdbcTemplate jdbcTemplate, EmbeddingModel embeddingModel, int dimensions,
141+
PgDistanceType distanceType, boolean removeExistingVectorStoreTable, PgIndexType createIndexMethod,
142+
boolean initializeSchema) {
143+
144+
this(DEFAULT_TABLE_NAME, jdbcTemplate, embeddingModel, dimensions, distanceType, removeExistingVectorStoreTable,
145+
createIndexMethod, initializeSchema);
146+
}
147+
148+
@Deprecated(forRemoval = true, since = "1.0.0-M5")
149+
public PgVectorStore(String vectorTableName, JdbcTemplate jdbcTemplate, EmbeddingModel embeddingModel,
150+
int dimensions, PgDistanceType distanceType, boolean removeExistingVectorStoreTable,
151+
PgIndexType createIndexMethod, boolean initializeSchema) {
152+
153+
this(builder().jdbcTemplate(jdbcTemplate)
154+
.schemaName(DEFAULT_SCHEMA_NAME)
155+
.vectorTableName(vectorTableName)
156+
.vectorTableValidationsEnabled(DEFAULT_SCHEMA_VALIDATION)
157+
.dimensions(dimensions)
158+
.distanceType(distanceType)
159+
.removeExistingVectorStoreTable(removeExistingVectorStoreTable)
160+
.indexType(createIndexMethod)
161+
.initializeSchema(initializeSchema));
162+
}
171163

172164
/**
173165
* @param builder {@link VectorStore.Builder} for pg vector store
174166
*/
175-
private PgVectorStore(PgVectorStoreBuilder builder) {
167+
protected PgVectorStore(PgVectorStoreBuilder builder) {
176168
super(builder);
169+
170+
Assert.notNull(builder.jdbcTemplate, "JdbcTemplate must not be null");
171+
177172
this.objectMapper = JsonMapper.builder().addModules(JacksonUtils.instantiateAvailableModules()).build();
178173

179-
String vectorTableName1 = builder.vectorTableName;
180-
this.vectorTableName = (null == vectorTableName1 || vectorTableName1.isEmpty()) ? DEFAULT_TABLE_NAME
181-
: vectorTableName1.trim();
174+
String vectorTable = builder.vectorTableName;
175+
this.vectorTableName = (null == vectorTable || vectorTable.isEmpty()) ? DEFAULT_TABLE_NAME : vectorTable.trim();
182176
logger.info("Using the vector table name: {}. Is empty: {}", this.vectorTableName,
183177
(this.vectorTableName == null || this.vectorTableName.isEmpty()));
184178

@@ -203,8 +197,8 @@ public PgDistanceType getDistanceType() {
203197
return this.distanceType;
204198
}
205199

206-
public static PgVectorStoreBuilder builder(JdbcTemplate jdbcTemplate) {
207-
return new PgVectorStoreBuilder(jdbcTemplate);
200+
public static PgVectorStoreBuilder builder() {
201+
return new PgVectorStoreBuilder();
208202
}
209203

210204
@Override
@@ -542,7 +536,7 @@ private Map<String, Object> toMap(PGobject pgObject) {
542536

543537
public static class PgVectorStoreBuilder extends AbstractVectorStoreBuilder<PgVectorStoreBuilder> {
544538

545-
private final JdbcTemplate jdbcTemplate;
539+
private JdbcTemplate jdbcTemplate;
546540

547541
private String schemaName = PgVectorStore.DEFAULT_SCHEMA_NAME;
548542

@@ -564,9 +558,10 @@ public static class PgVectorStoreBuilder extends AbstractVectorStoreBuilder<PgVe
564558

565559
private int maxDocumentBatchSize = MAX_DOCUMENT_BATCH_SIZE;
566560

567-
public PgVectorStoreBuilder(JdbcTemplate jdbcTemplate) {
561+
public PgVectorStoreBuilder jdbcTemplate(JdbcTemplate jdbcTemplate) {
568562
Assert.notNull(jdbcTemplate, "JdbcTemplate must not be null");
569563
this.jdbcTemplate = jdbcTemplate;
564+
return this;
570565
}
571566

572567
public PgVectorStoreBuilder schemaName(String schemaName) {
@@ -727,7 +722,8 @@ public Builder withMaxDocumentBatchSize(int maxDocumentBatchSize) {
727722
}
728723

729724
public PgVectorStore build() {
730-
return PgVectorStore.builder(this.jdbcTemplate)
725+
return PgVectorStore.builder()
726+
.jdbcTemplate(this.jdbcTemplate)
731727
.embeddingModel(this.embeddingModel)
732728
.schemaName(this.schemaName)
733729
.vectorTableName(this.vectorTableName)

‎vector-stores/spring-ai-pgvector-store/src/test/java/org/springframework/ai/pg/vectorstore/PgVectorEmbeddingDimensionsTests.java‎

Lines changed: 6 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -47,7 +47,8 @@ public void explicitlySetDimensions() {
4747

4848
final int explicitDimensions = 696;
4949

50-
PgVectorStore pgVectorStore = PgVectorStore.builder(this.jdbcTemplate)
50+
PgVectorStore pgVectorStore = PgVectorStore.builder()
51+
.jdbcTemplate(this.jdbcTemplate)
5152
.embeddingModel(this.embeddingModel)
5253
.dimensions(explicitDimensions)
5354
.build();
@@ -61,7 +62,8 @@ public void explicitlySetDimensions() {
6162
public void embeddingModelDimensions() {
6263
given(this.embeddingModel.dimensions()).willReturn(969);
6364

64-
PgVectorStore pgVectorStore = PgVectorStore.builder(this.jdbcTemplate)
65+
PgVectorStore pgVectorStore = PgVectorStore.builder()
66+
.jdbcTemplate(this.jdbcTemplate)
6567
.embeddingModel(this.embeddingModel)
6668
.build();
6769
var dim = pgVectorStore.embeddingDimensions();
@@ -76,7 +78,8 @@ public void fallBackToDefaultDimensions() {
7678

7779
given(this.embeddingModel.dimensions()).willThrow(new RuntimeException());
7880

79-
PgVectorStore pgVectorStore = PgVectorStore.builder(this.jdbcTemplate)
81+
PgVectorStore pgVectorStore = PgVectorStore.builder()
82+
.jdbcTemplate(this.jdbcTemplate)
8083
.embeddingModel(this.embeddingModel)
8184
.build();
8285
var dim = pgVectorStore.embeddingDimensions();

‎vector-stores/spring-ai-pgvector-store/src/test/java/org/springframework/ai/pg/vectorstore/PgVectorStoreCustomNamesIT.java‎

Lines changed: 2 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -196,7 +196,8 @@ public static class TestApplication {
196196
@Bean
197197
public VectorStore vectorStore(JdbcTemplate jdbcTemplate, EmbeddingModel embeddingModel) {
198198

199-
return PgVectorStore.builder(jdbcTemplate)
199+
return PgVectorStore.builder()
200+
.jdbcTemplate(jdbcTemplate)
200201
.embeddingModel(embeddingModel)
201202
.schemaName(this.schemaName)
202203
.vectorTableName(this.vectorTableName)

‎vector-stores/spring-ai-pgvector-store/src/test/java/org/springframework/ai/pg/vectorstore/PgVectorStoreIT.java‎

Lines changed: 2 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -356,7 +356,8 @@ public static class TestApplication {
356356

357357
@Bean
358358
public VectorStore vectorStore(JdbcTemplate jdbcTemplate, EmbeddingModel embeddingModel) {
359-
return PgVectorStore.builder(jdbcTemplate)
359+
return PgVectorStore.builder()
360+
.jdbcTemplate(jdbcTemplate)
360361
.embeddingModel(embeddingModel)
361362
.dimensions(PgVectorStore.INVALID_EMBEDDING_DIMENSION)
362363
.distanceType(this.distanceType)

‎vector-stores/spring-ai-pgvector-store/src/test/java/org/springframework/ai/pg/vectorstore/PgVectorStoreObservationIT.java‎

Lines changed: 2 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -187,7 +187,8 @@ public TestObservationRegistry observationRegistry() {
187187
@Bean
188188
public VectorStore vectorStore(JdbcTemplate jdbcTemplate, EmbeddingModel embeddingModel,
189189
ObservationRegistry observationRegistry) {
190-
return PgVectorStore.builder(jdbcTemplate)
190+
return PgVectorStore.builder()
191+
.jdbcTemplate(jdbcTemplate)
191192
.embeddingModel(embeddingModel)
192193
.distanceType(PgVectorStore.PgDistanceType.COSINE_DISTANCE)
193194
.indexType(PgIndexType.HNSW)

‎vector-stores/spring-ai-pgvector-store/src/test/java/org/springframework/ai/pg/vectorstore/PgVectorStoreTests.java‎

Lines changed: 2 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -79,7 +79,8 @@ void shouldAddDocumentsInBatchesAndEmbedOnce() {
7979
// Given
8080
var jdbcTemplate = mock(JdbcTemplate.class);
8181
var embeddingModel = mock(EmbeddingModel.class);
82-
var pgVectorStore = PgVectorStore.builder(jdbcTemplate)
82+
var pgVectorStore = PgVectorStore.builder()
83+
.jdbcTemplate(jdbcTemplate)
8384
.embeddingModel(embeddingModel)
8485
.maxDocumentBatchSize(1000)
8586
.build();

0 commit comments

Comments
 (0)