From f2dfe2b7dfb0d7e37582cd6b6536d95ad3547046 Mon Sep 17 00:00:00 2001 From: Christian Tzolov Date: Tue, 10 Oct 2023 08:41:41 +0200 Subject: [PATCH] PgVectorStore - reduce the SmartLifeCycle to InitializingBean As the LifeCycle#start() occures later in time than the InitializingBean#afterPropertySet() it could cose some initialization issues with the vector clients. --- .../ai/vectorstore/PgVectorStore.java | 33 +++---------------- .../ai/vectorstore/PgVectorStoreIT.java | 10 ++++++ 2 files changed, 14 insertions(+), 29 deletions(-) diff --git a/vector-stores/spring-ai-pgvector-store/src/main/java/org/springframework/ai/vectorstore/PgVectorStore.java b/vector-stores/spring-ai-pgvector-store/src/main/java/org/springframework/ai/vectorstore/PgVectorStore.java index 271de91d9..cddb6565f 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/vectorstore/PgVectorStore.java @@ -22,7 +22,6 @@ import java.util.List; import java.util.Map; import java.util.Optional; import java.util.UUID; -import java.util.concurrent.atomic.AtomicBoolean; import java.util.stream.IntStream; import com.fasterxml.jackson.core.JsonProcessingException; @@ -34,7 +33,7 @@ import org.slf4j.LoggerFactory; import org.springframework.ai.document.Document; import org.springframework.ai.embedding.EmbeddingClient; -import org.springframework.context.SmartLifecycle; +import org.springframework.beans.factory.InitializingBean; import org.springframework.jdbc.core.JdbcTemplate; import org.springframework.jdbc.core.RowMapper; import org.springframework.lang.Nullable; @@ -45,7 +44,7 @@ import org.springframework.lang.Nullable; * * @author Christian Tzolov */ -public class PgVectorStore implements VectorStore, SmartLifecycle { +public class PgVectorStore implements VectorStore, InitializingBean { private static final Logger logger = LoggerFactory.getLogger(PgVectorStore.class); @@ -291,13 +290,10 @@ public class PgVectorStore implements VectorStore, SmartLifecycle { } // --------------------------------------------------------------------------------- - // SmartLifecycle + // Initialize // --------------------------------------------------------------------------------- - private AtomicBoolean isRunning = new AtomicBoolean(false); - @Override - public void start() { - + public void afterPropertiesSet() throws Exception { // Enable the PGVector, JSONB and UUID support. this.jdbcTemplate.execute("CREATE EXTENSION IF NOT EXISTS vector"); this.jdbcTemplate.execute("CREATE EXTENSION IF NOT EXISTS hstore"); @@ -317,7 +313,6 @@ public class PgVectorStore implements VectorStore, SmartLifecycle { + " (embedding " + this.getDistanceType().index + ")"); } - this.isRunning.set(true); } int embeddingDimensions() { @@ -339,24 +334,4 @@ public class PgVectorStore implements VectorStore, SmartLifecycle { return OPENAI_EMBEDDING_DIMENSION_SIZE; } - @Override - public void stop() { - // Remove existing VectorStoreTable - if (this.removeExistingVectorStoreTable) { - this.jdbcTemplate.execute("DROP TABLE IF EXISTS " + VECTOR_TABLE_NAME); - } - - this.isRunning.set(false); - } - - @Override - public boolean isRunning() { - return this.isRunning.get(); - } - - @Override - public boolean isAutoStartup() { - return true; - } - } \ No newline at end of file diff --git a/vector-stores/spring-ai-pgvector-store/src/test/java/org/springframework/ai/vectorstore/PgVectorStoreIT.java b/vector-stores/spring-ai-pgvector-store/src/test/java/org/springframework/ai/vectorstore/PgVectorStoreIT.java index d88506c2f..45342a741 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/vectorstore/PgVectorStoreIT.java @@ -42,6 +42,7 @@ import org.springframework.boot.autoconfigure.jdbc.DataSourceAutoConfiguration; import org.springframework.boot.autoconfigure.jdbc.DataSourceProperties; import org.springframework.boot.context.properties.ConfigurationProperties; import org.springframework.boot.test.context.runner.ApplicationContextRunner; +import org.springframework.context.ApplicationContext; import org.springframework.context.annotation.Bean; import org.springframework.context.annotation.Primary; import org.springframework.jdbc.core.JdbcTemplate; @@ -80,6 +81,11 @@ public class PgVectorStoreIT { "app.datasource.username=postgres", "app.datasource.password=postgres", "app.datasource.type=com.zaxxer.hikari.HikariDataSource"); + private static void dropTable(ApplicationContext context) { + JdbcTemplate jdbcTemplate = context.getBean(JdbcTemplate.class); + jdbcTemplate.execute("DROP TABLE IF EXISTS vector_store"); + } + @ParameterizedTest @ValueSource(strings = { "CosineDistance", "EuclideanDistance", "NegativeInnerProduct" }) public void addAndSearchTest(String distanceType) { @@ -106,6 +112,7 @@ public class PgVectorStoreIT { List results2 = vectorStore.similaritySearch("Great", 1); assertThat(results2).hasSize(0); + dropTable(context); }); } @@ -145,6 +152,8 @@ public class PgVectorStoreIT { assertThat(resultDoc.getId()).isEqualTo(document.getId()); assertThat(resultDoc.getContent()).isEqualTo("The World is Big and Salvation Lurks Around the Corner"); assertThat(resultDoc.getMetadata()).containsKeys("meta2", "distance"); + + dropTable(context); }); } @@ -183,6 +192,7 @@ public class PgVectorStoreIT { "Great Depression Great Depression Great Depression Great Depression Great Depression Great Depression"); assertThat(resultDoc.getMetadata()).containsKeys("meta2", "distance"); + dropTable(context); }); }