diff --git a/spring-ai-docs/src/main/antora/modules/ROOT/pages/api/vectordbs/pgvector.adoc b/spring-ai-docs/src/main/antora/modules/ROOT/pages/api/vectordbs/pgvector.adoc index 2374f69b5..44d610a61 100644 --- a/spring-ai-docs/src/main/antora/modules/ROOT/pages/api/vectordbs/pgvector.adoc +++ b/spring-ai-docs/src/main/antora/modules/ROOT/pages/api/vectordbs/pgvector.adoc @@ -142,9 +142,15 @@ You can use the following properties in your Spring Boot configuration to custom |`spring.ai.vectorstore.pgvector.dimensions`| Embeddings dimension. If not specified explicitly the PgVectorStore will retrieve the dimensions form the provided `EmbeddingModel`. Dimensions are set to the embedding column the on table creation. If you change the dimensions your would have to re-create the vector_store table as well. | - |`spring.ai.vectorstore.pgvector.remove-existing-vector-store-table` | Deletes the existing `vector_store` table on start up. | false |`spring.ai.vectorstore.pgvector.initialize-schema` | Whether to initialize the required schema | false +|`spring.ai.vectorstore.pgvector.schema-name` | Vector store schema name | `public` +|`spring.ai.vectorstore.pgvector.table-name` | Vector store table name | `vector_store` +|`spring.ai.vectorstore.pgvector.schema-validation` | Enables schema and table name validation to ensure they are valid and existing objects. | false |=== +TIP: If you configure a custom schema and/or table name, consider enabling schema validation by setting `spring.ai.vectorstore.pgvector.schema-validation=true`. +This ensures the correctness of the names and reduces the risk of SQL injection attacks. + == Metadata filtering You can leverage the generic, portable link:https://docs.spring.io/spring-ai/reference/api/vectordbs.html#_metadata_filters[metadata filters] with the PgVector store. 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 8ccde1a76..7a81d9a54 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 @@ -19,7 +19,6 @@ import javax.sql.DataSource; import org.springframework.ai.embedding.EmbeddingModel; import org.springframework.ai.vectorstore.PgVectorStore; -import org.springframework.beans.factory.InitializingBean; import org.springframework.boot.autoconfigure.AutoConfiguration; import org.springframework.boot.autoconfigure.condition.ConditionalOnClass; import org.springframework.boot.autoconfigure.condition.ConditionalOnMissingBean; @@ -42,8 +41,16 @@ public class PgVectorStoreAutoConfiguration { public PgVectorStore vectorStore(JdbcTemplate jdbcTemplate, EmbeddingModel embeddingModel, PgVectorStoreProperties properties) { var initializeSchema = properties.isInitializeSchema(); - return new PgVectorStore(jdbcTemplate, embeddingModel, properties.getDimensions(), properties.getDistanceType(), - properties.isRemoveExistingVectorStoreTable(), properties.getIndexType(), initializeSchema); + + 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) + .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 927a82193..783041c76 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 @@ -23,6 +23,7 @@ import org.springframework.boot.context.properties.ConfigurationProperties; /** * @author Christian Tzolov + * @author Muthukumaran Navaneethakrishnan */ @ConfigurationProperties(PgVectorStoreProperties.CONFIG_PREFIX) public class PgVectorStoreProperties extends CommonVectorStoreProperties { @@ -37,6 +38,13 @@ public class PgVectorStoreProperties extends CommonVectorStoreProperties { private boolean removeExistingVectorStoreTable = false; + // Dynamically generate table name in PgVectorStore to allow backward compatibility + private String tableName = PgVectorStore.DEFAULT_TABLE_NAME; + + private String schemaName = PgVectorStore.DEFAULT_SCHEMA_NAME; + + private boolean schemaValidation = PgVectorStore.DEFAULT_SCHEMA_VALIDATION; + public int getDimensions() { return dimensions; } @@ -62,11 +70,35 @@ public class PgVectorStoreProperties extends CommonVectorStoreProperties { } public boolean isRemoveExistingVectorStoreTable() { - return removeExistingVectorStoreTable; + return this.removeExistingVectorStoreTable; } public void setRemoveExistingVectorStoreTable(boolean removeExistingVectorStoreTable) { this.removeExistingVectorStoreTable = removeExistingVectorStoreTable; } + public String getTableName() { + return this.tableName; + } + + public void setTableName(String vectorTableName) { + this.tableName = vectorTableName; + } + + public String getSchemaName() { + return this.schemaName; + } + + public void setSchemaName(String schemaName) { + this.schemaName = schemaName; + } + + public boolean isSchemaValidation() { + return this.schemaValidation; + } + + public void setSchemaValidation(boolean schemaValidation) { + this.schemaValidation = schemaValidation; + } + } diff --git a/spring-ai-spring-boot-autoconfigure/src/test/java/org/springframework/ai/autoconfigure/vectorstore/pgvector/PgVectorStoreAutoConfiguration2IT.java b/spring-ai-spring-boot-autoconfigure/src/test/java/org/springframework/ai/autoconfigure/vectorstore/pgvector/PgVectorStoreAutoConfiguration2IT.java deleted file mode 100644 index c31a59368..000000000 --- a/spring-ai-spring-boot-autoconfigure/src/test/java/org/springframework/ai/autoconfigure/vectorstore/pgvector/PgVectorStoreAutoConfiguration2IT.java +++ /dev/null @@ -1,114 +0,0 @@ -/* - * 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. - */ -package org.springframework.ai.autoconfigure.vectorstore.pgvector; - -import java.io.IOException; -import java.nio.charset.StandardCharsets; -import java.util.List; -import java.util.Map; - -import org.junit.jupiter.api.Test; -import org.testcontainers.containers.GenericContainer; -import org.testcontainers.containers.PostgreSQLContainer; -import org.testcontainers.junit.jupiter.Container; -import org.testcontainers.junit.jupiter.Testcontainers; -import org.testcontainers.utility.DockerImageName; -import org.springframework.ai.document.Document; -import org.springframework.ai.embedding.EmbeddingModel; -import org.springframework.ai.transformers.TransformersEmbeddingModel; -import org.springframework.ai.vectorstore.SearchRequest; -import org.springframework.ai.vectorstore.VectorStore; -import org.springframework.boot.autoconfigure.AutoConfigurations; -import org.springframework.boot.autoconfigure.jdbc.DataSourceAutoConfiguration; -import org.springframework.boot.autoconfigure.jdbc.JdbcTemplateAutoConfiguration; -import org.springframework.boot.test.context.runner.ApplicationContextRunner; -import org.springframework.context.annotation.Bean; -import org.springframework.context.annotation.Configuration; -import org.springframework.core.io.DefaultResourceLoader; - -import static org.assertj.core.api.Assertions.assertThat; - -/** - * @author Christian Tzolov - */ -@Testcontainers -public class PgVectorStoreAutoConfiguration2IT { - - @Container - @SuppressWarnings("resource") - static PostgreSQLContainer postgresContainer = new PostgreSQLContainer<>("pgvector/pgvector:pg16") - .withUsername("postgres") - .withPassword("postgres"); - - List documents = List.of( - new Document(getText("classpath:/test/data/spring.ai.txt"), Map.of("spring", "great")), - new Document(getText("classpath:/test/data/time.shelter.txt")), - new Document(getText("classpath:/test/data/great.depression.txt"), Map.of("depression", "bad"))); - - public static String getText(String uri) { - var resource = new DefaultResourceLoader().getResource(uri); - try { - return resource.getContentAsString(StandardCharsets.UTF_8); - } - catch (IOException e) { - throw new RuntimeException(e); - } - } - - private final ApplicationContextRunner contextRunner = new ApplicationContextRunner() - .withConfiguration(AutoConfigurations.of(PgVectorStoreAutoConfiguration.class, - JdbcTemplateAutoConfiguration.class, DataSourceAutoConfiguration.class)) - .withUserConfiguration(Config.class) - .withPropertyValues("spring.ai.vectorstore.pgvector.distanceType=COSINE_DISTANCE", - // JdbcTemplate configuration - String.format("spring.datasource.url=jdbc:postgresql://%s:%d/%s", postgresContainer.getHost(), - postgresContainer.getMappedPort(5432), "postgres"), - "spring.datasource.username=postgres", "spring.datasource.password=postgres"); - - @Test - public void addAndSearch() { - contextRunner.run(context -> { - - VectorStore vectorStore = context.getBean(VectorStore.class); - - vectorStore.add(documents); - - List results = vectorStore - .similaritySearch(SearchRequest.query("What is Great Depression?").withTopK(1)); - - assertThat(results).hasSize(1); - Document resultDoc = results.get(0); - assertThat(resultDoc.getId()).isEqualTo(documents.get(2).getId()); - assertThat(resultDoc.getMetadata()).containsKeys("depression", "distance"); - - // Remove all documents from the store - vectorStore.delete(documents.stream().map(doc -> doc.getId()).toList()); - results = vectorStore.similaritySearch(SearchRequest.query("Great Depression").withTopK(1)); - assertThat(results).hasSize(0); - }); - } - - @Configuration(proxyBeanMethods = false) - static class Config { - - @Bean - public EmbeddingModel embeddingModel() { - return new TransformersEmbeddingModel(); - } - - } - -} 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 346145760..78b514c92 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 @@ -15,42 +15,44 @@ */ package org.springframework.ai.autoconfigure.vectorstore.pgvector; +import static org.assertj.core.api.Assertions.assertThat; + import java.io.IOException; import java.nio.charset.StandardCharsets; import java.util.List; import java.util.Map; import org.junit.jupiter.api.Test; -import org.testcontainers.containers.GenericContainer; -import org.testcontainers.junit.jupiter.Container; -import org.testcontainers.junit.jupiter.Testcontainers; - +import org.junit.jupiter.params.ParameterizedTest; +import org.junit.jupiter.params.provider.ValueSource; import org.springframework.ai.document.Document; import org.springframework.ai.embedding.EmbeddingModel; import org.springframework.ai.transformers.TransformersEmbeddingModel; +import org.springframework.ai.vectorstore.PgVectorStore; import org.springframework.ai.vectorstore.SearchRequest; -import org.springframework.ai.vectorstore.VectorStore; import org.springframework.boot.autoconfigure.AutoConfigurations; import org.springframework.boot.autoconfigure.jdbc.DataSourceAutoConfiguration; import org.springframework.boot.autoconfigure.jdbc.JdbcTemplateAutoConfiguration; import org.springframework.boot.test.context.runner.ApplicationContextRunner; +import org.springframework.context.ApplicationContext; import org.springframework.context.annotation.Bean; import org.springframework.context.annotation.Configuration; import org.springframework.core.io.DefaultResourceLoader; - -import static org.assertj.core.api.Assertions.assertThat; +import org.springframework.jdbc.core.JdbcTemplate; +import org.testcontainers.containers.PostgreSQLContainer; +import org.testcontainers.junit.jupiter.Container; +import org.testcontainers.junit.jupiter.Testcontainers; /** * @author Christian Tzolov + * @author Muthukumaran Navaneethakrishnan */ @Testcontainers public class PgVectorStoreAutoConfigurationIT { @Container - static GenericContainer postgresContainer = new GenericContainer<>("pgvector/pgvector:0.7.2-pg16") - .withEnv("POSTGRES_USER", "postgres") - .withEnv("POSTGRES_PASSWORD", "postgres") - .withExposedPorts(5432); + @SuppressWarnings("resource") + static PostgreSQLContainer postgresContainer = new PostgreSQLContainer<>("pgvector/pgvector:pg16"); List documents = List.of( new Document(getText("classpath:/test/data/spring.ai.txt"), Map.of("spring", "great")), @@ -72,16 +74,22 @@ public class PgVectorStoreAutoConfigurationIT { JdbcTemplateAutoConfiguration.class, DataSourceAutoConfiguration.class)) .withUserConfiguration(Config.class) .withPropertyValues("spring.ai.vectorstore.pgvector.distanceType=COSINE_DISTANCE", - // JdbcTemplate configuration + "spring.ai.vectorstore.pgvector.initialize-schema=true", String.format("spring.datasource.url=jdbc:postgresql://%s:%d/%s", postgresContainer.getHost(), - postgresContainer.getMappedPort(5432), "postgres"), - "spring.datasource.username=postgres", "spring.datasource.password=postgres"); + postgresContainer.getMappedPort(5432), postgresContainer.getDatabaseName()), + "spring.datasource.username=" + postgresContainer.getUsername(), + "spring.datasource.password=" + postgresContainer.getPassword()); @Test public void addAndSearch() { + contextRunner.run(context -> { - VectorStore vectorStore = context.getBean(VectorStore.class); + PgVectorStore vectorStore = context.getBean(PgVectorStore.class); + + assertThat(isFullyQualifiedTableExists(context, PgVectorStore.DEFAULT_SCHEMA_NAME, + PgVectorStore.DEFAULT_TABLE_NAME)) + .isTrue(); vectorStore.add(documents); @@ -100,6 +108,35 @@ public class PgVectorStoreAutoConfigurationIT { }); } + @ParameterizedTest(name = "{0} : {displayName} ") + @ValueSource(strings = { "public:vector_store", "my_schema:my_table" }) + public void customSchemaNames(String schemaTableName) { + String schemaName = schemaTableName.split(":")[0]; + String tableName = schemaTableName.split(":")[1]; + + contextRunner + .withPropertyValues("spring.ai.vectorstore.pgvector.schema-name=" + schemaName, + "spring.ai.vectorstore.pgvector.table-name=" + tableName) + .run(context -> { + assertThat(isFullyQualifiedTableExists(context, schemaName, tableName)).isTrue(); + }); + } + + @ParameterizedTest(name = "{0} : {displayName} ") + @ValueSource(strings = { "public:vector_store", "my_schema:my_table" }) + public void disableSchemaInitialization(String schemaTableName) { + String schemaName = schemaTableName.split(":")[0]; + String tableName = schemaTableName.split(":")[1]; + + contextRunner + .withPropertyValues("spring.ai.vectorstore.pgvector.schema-name=" + schemaName, + "spring.ai.vectorstore.pgvector.table-name=" + tableName, + "spring.ai.vectorstore.pgvector.initialize-schema=false") + .run(context -> { + assertThat(isFullyQualifiedTableExists(context, schemaName, tableName)).isFalse(); + }); + } + @Configuration(proxyBeanMethods = false) static class Config { @@ -110,4 +147,11 @@ public class PgVectorStoreAutoConfigurationIT { } + private static boolean isFullyQualifiedTableExists(ApplicationContext context, String schemaName, + String tableName) { + JdbcTemplate jdbcTemplate = context.getBean(JdbcTemplate.class); + String sql = "SELECT EXISTS (SELECT FROM information_schema.tables WHERE table_schema = ? AND table_name = ?)"; + return jdbcTemplate.queryForObject(sql, Boolean.class, schemaName, tableName); + } + } 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 db36d5e2e..a4e4ddef3 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 @@ -35,6 +35,11 @@ public class PgVectorStorePropertiesTests { assertThat(props.getDistanceType()).isEqualTo(PgDistanceType.COSINE_DISTANCE); assertThat(props.getIndexType()).isEqualTo(PgIndexType.HNSW); assertThat(props.isRemoveExistingVectorStoreTable()).isFalse(); + + assertThat(props.isSchemaValidation()).isFalse(); + assertThat(props.getSchemaName()).isEqualTo(PgVectorStore.DEFAULT_SCHEMA_NAME); + assertThat(props.getTableName()).isEqualTo(PgVectorStore.DEFAULT_TABLE_NAME); + } @Test @@ -46,10 +51,18 @@ public class PgVectorStorePropertiesTests { props.setIndexType(PgIndexType.IVFFLAT); props.setRemoveExistingVectorStoreTable(true); + props.setSchemaValidation(true); + props.setSchemaName("my_vector_schema"); + props.setTableName("my_vector_table"); + assertThat(props.getDimensions()).isEqualTo(1536); assertThat(props.getDistanceType()).isEqualTo(PgDistanceType.EUCLIDEAN_DISTANCE); assertThat(props.getIndexType()).isEqualTo(PgIndexType.IVFFLAT); assertThat(props.isRemoveExistingVectorStoreTable()).isTrue(); + + assertThat(props.isSchemaValidation()).isTrue(); + assertThat(props.getSchemaName()).isEqualTo("my_vector_schema"); + assertThat(props.getTableName()).isEqualTo("my_vector_table"); } } 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/vectorstore/PgVectorSchemaValidator.java new file mode 100644 index 000000000..f8017d97b --- /dev/null +++ b/vector-stores/spring-ai-pgvector-store/src/main/java/org/springframework/ai/vectorstore/PgVectorSchemaValidator.java @@ -0,0 +1,150 @@ +/* + * Copyright 2024 - 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. + */ +package org.springframework.ai.vectorstore; + +import java.util.ArrayList; +import java.util.List; +import java.util.Map; + +import org.slf4j.Logger; +import org.slf4j.LoggerFactory; +import org.springframework.dao.DataAccessException; +import org.springframework.jdbc.core.JdbcTemplate; + +/** + * @author Muthukumaran Navaneethakrishnan + * @author Christian Tzolov + * @since 1.0.0 + */ +public class PgVectorSchemaValidator { + + private static final Logger logger = LoggerFactory.getLogger(PgVectorSchemaValidator.class); + + private final JdbcTemplate jdbcTemplate; + + public PgVectorSchemaValidator(JdbcTemplate jdbcTemplate) { + this.jdbcTemplate = jdbcTemplate; + } + + static boolean isValidNameForDatabaseObject(String name) { + + if (name == null) { + return false; + } + + // Check if the table or schema has Only alphanumeric characters and underscores + // and should be less than 64 characters + if (!name.matches("^[a-zA-Z0-9_]{1,64}$")) { + return false; + } + + // Check to ensure the table or schema name is not purely numeric + if (name.matches("^[0-9]+$")) { + return false; + } + + return true; + + } + + public boolean isTableExists(String schemaName, String tableName) { + String sql = "SELECT 1 FROM information_schema.tables WHERE table_schema = ? AND table_name = ?"; + try { + // Query for a single integer value, if it exists, table exists + jdbcTemplate.queryForObject(sql, Integer.class, schemaName, tableName); + return true; + } + catch (DataAccessException e) { + return false; + } + } + + void validateTableSchema(String schemaName, String tableName) { + + if (!isValidNameForDatabaseObject(schemaName)) { + throw new IllegalArgumentException( + "Schema name should only contain alphanumeric characters and underscores"); + } + if (!isValidNameForDatabaseObject(tableName)) { + throw new IllegalArgumentException( + "Table name should only contain alphanumeric characters and underscores"); + } + + if (!isTableExists(schemaName, tableName)) { + throw new IllegalStateException("Table " + tableName + " does not exist in schema " + schemaName); + } + + try { + logger.info("Validating PGVectorStore schema for table: {} in schema: {}", tableName, schemaName); + + List expectedColumns = new ArrayList<>(); + expectedColumns.add("id"); + expectedColumns.add("content"); + expectedColumns.add("metadata"); + expectedColumns.add("embedding"); + + // Query to check if the table exists with the required fields and types + // Include the schema name in the query to target the correct table + String query = "SELECT column_name, data_type FROM information_schema.columns " + + "WHERE table_schema = ? AND table_name = ?"; + List> columns = jdbcTemplate.queryForList(query, + new Object[] { schemaName, tableName }); + + if (columns.isEmpty()) { + throw new IllegalStateException("Error while validating table schema, Table " + tableName + + " does not exist in schema " + schemaName); + } + + // Check each column against expected fields + List availableColumns = new ArrayList<>(); + for (Map column : columns) { + String columnName = (String) column.get("column_name"); + availableColumns.add(columnName); + + } + + expectedColumns.removeAll(availableColumns); + + if (expectedColumns.isEmpty()) { + logger.info("PG VectorStore schema validation successful"); + } + else { + throw new IllegalStateException("Missing fields " + expectedColumns); + } + + } + catch (DataAccessException | IllegalStateException e) { + logger.error("Error while validating table schema" + e.getMessage()); + logger + .error("Failed to operate with the specified table in the database. To resolve this issue, please ensure the following steps are completed:\n" + + "1. Ensure the necessary PostgreSQL extensions are enabled. Run the following SQL commands:\n" + + " CREATE EXTENSION IF NOT EXISTS vector;\n" + " CREATE EXTENSION IF NOT EXISTS hstore;\n" + + " CREATE EXTENSION IF NOT EXISTS \"uuid-ossp\";\n" + + "2. Verify that the table exists with the appropriate structure. If it does not exist, create it using a SQL command similar to the following, replacing 'embedding_dimensions' with the appropriate size based on your vector embeddings:\n" + + String.format(" CREATE TABLE IF NOT EXISTS %s (\n" + + " id uuid DEFAULT uuid_generate_v4() PRIMARY KEY,\n" + " content text,\n" + + " metadata json,\n" + + " embedding vector(embedding_dimensions) // Replace 'embedding_dimensions' with your specific value\n" + + " );\n", schemaName + "." + tableName) + + "3. Create an appropriate index for the vector embedding to optimize performance. Adjust the index type and options based on your usage. Example SQL for creating an index:\n" + + String.format(" CREATE INDEX ON %s USING HNSW (embedding vector_cosine_ops);\n", tableName) + + "\nPlease adjust these commands based on your specific configuration and the capabilities of your vector database system."); + throw new IllegalStateException(e); + + } + } + +} 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 515a1b424..dacde2fc9 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 @@ -24,13 +24,9 @@ import java.util.Optional; import java.util.UUID; import java.util.stream.IntStream; -import com.fasterxml.jackson.core.JsonProcessingException; -import com.fasterxml.jackson.databind.ObjectMapper; -import com.pgvector.PGvector; import org.postgresql.util.PGobject; import org.slf4j.Logger; import org.slf4j.LoggerFactory; - import org.springframework.ai.document.Document; import org.springframework.ai.embedding.EmbeddingModel; import org.springframework.ai.vectorstore.filter.FilterExpressionConverter; @@ -43,12 +39,17 @@ import org.springframework.jdbc.core.StatementCreatorUtils; import org.springframework.lang.Nullable; import org.springframework.util.StringUtils; +import com.fasterxml.jackson.core.JsonProcessingException; +import com.fasterxml.jackson.databind.ObjectMapper; +import com.pgvector.PGvector; + /** * Uses the "vector_store" table to store the Spring AI vector data. The table and the * vector index will be auto-created if not available. * * @author Christian Tzolov * @author Josh Long + * @author Muthukumaran Navaneethakrishnan */ public class PgVectorStore implements VectorStore, InitializingBean { @@ -58,16 +59,30 @@ public class PgVectorStore implements VectorStore, InitializingBean { public static final int INVALID_EMBEDDING_DIMENSION = -1; - public static final String VECTOR_TABLE_NAME = "vector_store"; + public final static String DEFAULT_TABLE_NAME = "vector_store"; - public static final String VECTOR_INDEX_NAME = "spring_ai_vector_index"; + public final static String DEFAULT_VECTOR_INDEX_NAME = "spring_ai_vector_index"; + + public final static String DEFAULT_SCHEMA_NAME = "public"; + + public final static boolean DEFAULT_SCHEMA_VALIDATION = false; public final FilterExpressionConverter filterExpressionConverter = new PgVectorFilterExpressionConverter(); + private final String vectorTableName; + + private final String vectorIndexName; + private final JdbcTemplate jdbcTemplate; private final EmbeddingModel embeddingModel; + private final String schemaName; + + private final boolean schemaValidation; + + private final boolean initializeSchema; + private int dimensions; private PgDistanceType distanceType; @@ -78,7 +93,255 @@ public class PgVectorStore implements VectorStore, InitializingBean { private PgIndexType createIndexMethod; - private final boolean initializeSchema; + private PgVectorSchemaValidator schemaValidator; + + 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.vectorTableName = (null == vectorTableName || vectorTableName.isEmpty()) ? DEFAULT_TABLE_NAME + : vectorTableName.trim(); + logger.info("Using the vector table name: {}", + this.vectorTableName + " is empty" + (null == vectorTableName || 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); + } + + public PgDistanceType getDistanceType() { + return distanceType; + } + + @Override + public void add(List documents) { + + int size = documents.size(); + + this.jdbcTemplate.batchUpdate( + "INSERT INTO " + getFullyQualifiedTableName() + + " (id, content, metadata, embedding) VALUES (?, ?, ?::jsonb, ?) " + "ON CONFLICT (id) DO " + + "UPDATE SET content = ? , metadata = ?::jsonb , embedding = ? ", + new BatchPreparedStatementSetter() { + @Override + public void setValues(PreparedStatement ps, int i) throws SQLException { + + var document = documents.get(i); + var content = document.getContent(); + var json = toJson(document.getMetadata()); + var pGvector = new PGvector(toFloatArray(embeddingModel.embed(document))); + + StatementCreatorUtils.setParameterValue(ps, 1, SqlTypeValue.TYPE_UNKNOWN, + UUID.fromString(document.getId())); + StatementCreatorUtils.setParameterValue(ps, 2, SqlTypeValue.TYPE_UNKNOWN, content); + StatementCreatorUtils.setParameterValue(ps, 3, SqlTypeValue.TYPE_UNKNOWN, json); + StatementCreatorUtils.setParameterValue(ps, 4, SqlTypeValue.TYPE_UNKNOWN, pGvector); + StatementCreatorUtils.setParameterValue(ps, 5, SqlTypeValue.TYPE_UNKNOWN, content); + StatementCreatorUtils.setParameterValue(ps, 6, SqlTypeValue.TYPE_UNKNOWN, json); + StatementCreatorUtils.setParameterValue(ps, 7, SqlTypeValue.TYPE_UNKNOWN, pGvector); + } + + @Override + public int getBatchSize() { + return size; + } + }); + } + + private String toJson(Map map) { + try { + return objectMapper.writeValueAsString(map); + } + catch (JsonProcessingException e) { + throw new RuntimeException(e); + } + } + + private float[] toFloatArray(List embeddingDouble) { + float[] embeddingFloat = new float[embeddingDouble.size()]; + int i = 0; + for (Double d : embeddingDouble) { + embeddingFloat[i++] = d.floatValue(); + } + return embeddingFloat; + } + + @Override + public Optional delete(List idList) { + int updateCount = 0; + for (String id : idList) { + int count = jdbcTemplate.update("DELETE FROM " + getFullyQualifiedTableName() + " WHERE id = ?", + UUID.fromString(id)); + updateCount = updateCount + count; + } + + return Optional.of(updateCount == idList.size()); + } + + @Override + public List similaritySearch(SearchRequest request) { + + String nativeFilterExpression = (request.getFilterExpression() != null) + ? this.filterExpressionConverter.convertExpression(request.getFilterExpression()) : ""; + + String jsonPathFilter = ""; + + if (StringUtils.hasText(nativeFilterExpression)) { + jsonPathFilter = " AND metadata::jsonb @@ '" + nativeFilterExpression + "'::jsonpath "; + } + + double distance = 1 - request.getSimilarityThreshold(); + + PGvector queryEmbedding = getQueryEmbedding(request.getQuery()); + + return this.jdbcTemplate.query( + String.format(this.getDistanceType().similaritySearchSqlTemplate, getFullyQualifiedTableName(), + jsonPathFilter), + new DocumentRowMapper(this.objectMapper), queryEmbedding, queryEmbedding, distance, request.getTopK()); + } + + public List embeddingDistance(String query) { + return this.jdbcTemplate.query( + "SELECT embedding " + this.comparisonOperator() + " ? AS distance FROM " + getFullyQualifiedTableName(), + new RowMapper() { + @Override + @Nullable + public Double mapRow(ResultSet rs, int rowNum) throws SQLException { + return rs.getDouble(DocumentRowMapper.COLUMN_DISTANCE); + } + + }, getQueryEmbedding(query)); + } + + private PGvector getQueryEmbedding(String query) { + List embedding = this.embeddingModel.embed(query); + return new PGvector(toFloatArray(embedding)); + } + + private String comparisonOperator() { + return this.getDistanceType().operator; + } + + // --------------------------------------------------------------------------------- + // Initialize + // --------------------------------------------------------------------------------- + @Override + public void afterPropertiesSet() throws Exception { + + logger.info("Initializing PGVectorStore schema for table: {} in schema: {}", this.getVectorTableName(), + this.getSchemaName()); + + logger.info("vectorTableValidationsEnabled {}", this.schemaValidation); + + if (this.schemaValidation) { + this.schemaValidator.validateTableSchema(this.getSchemaName(), this.getVectorTableName()); + } + + if (!this.initializeSchema) { + logger.debug("Skipping the schema initialization for the table: " + this.getFullyQualifiedTableName()); + return; + } + + // 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"); + this.jdbcTemplate.execute("CREATE EXTENSION IF NOT EXISTS \"uuid-ossp\""); + + this.jdbcTemplate.execute(String.format("CREATE SCHEMA IF NOT EXISTS %s", this.getSchemaName())); + + // Remove existing VectorStoreTable + if (this.removeExistingVectorStoreTable) { + this.jdbcTemplate.execute(String.format("DROP TABLE IF EXISTS %s", this.getFullyQualifiedTableName())); + } + + this.jdbcTemplate.execute(String.format(""" + CREATE TABLE IF NOT EXISTS %s ( + id uuid DEFAULT uuid_generate_v4() PRIMARY KEY, + content text, + metadata json, + embedding vector(%d) + ) + """, this.getFullyQualifiedTableName(), this.embeddingDimensions())); + + if (this.createIndexMethod != PgIndexType.NONE) { + this.jdbcTemplate.execute(String.format(""" + CREATE INDEX IF NOT EXISTS %s ON %s USING %s (embedding %s) + """, this.getVectorIndexName(), this.getFullyQualifiedTableName(), this.createIndexMethod, + this.getDistanceType().index)); + } + } + + private String getFullyQualifiedTableName() { + return this.schemaName + "." + this.vectorTableName; + } + + private String getVectorTableName() { + return this.vectorTableName; + } + + private String getSchemaName() { + return this.schemaName; + } + + private String getVectorIndexName() { + return this.vectorIndexName; + } + + int embeddingDimensions() { + // The manually set dimensions have precedence over the computed one. + if (this.dimensions > 0) { + return this.dimensions; + } + + try { + int embeddingDimensions = this.embeddingModel.dimensions(); + if (embeddingDimensions > 0) { + return embeddingDimensions; + } + } + catch (Exception e) { + logger.warn("Failed to obtain the embedding dimensions from the embedding model and fall backs to default:" + + OPENAI_EMBEDDING_DIMENSION_SIZE, e); + } + return OPENAI_EMBEDDING_DIMENSION_SIZE; + } /** * By default, pgvector performs exact nearest neighbor search, which provides perfect @@ -199,192 +462,83 @@ public class PgVectorStore implements VectorStore, InitializingBean { } - public PgVectorStore(JdbcTemplate jdbcTemplate, EmbeddingModel embeddingModel) { - this(jdbcTemplate, embeddingModel, INVALID_EMBEDDING_DIMENSION, PgVectorStore.PgDistanceType.COSINE_DISTANCE, - false, PgIndexType.NONE, false); - } + public static class Builder { - public PgVectorStore(JdbcTemplate jdbcTemplate, EmbeddingModel embeddingModel, int dimensions) { - this(jdbcTemplate, embeddingModel, dimensions, PgVectorStore.PgDistanceType.COSINE_DISTANCE, false, - PgIndexType.NONE, false); - } + private final JdbcTemplate jdbcTemplate; - public PgVectorStore(JdbcTemplate jdbcTemplate, EmbeddingModel embeddingModel, int dimensions, - PgDistanceType distanceType, boolean removeExistingVectorStoreTable, PgIndexType createIndexMethod, - boolean initializeSchema) { + private final EmbeddingModel embeddingModel; - this.jdbcTemplate = jdbcTemplate; - this.embeddingModel = embeddingModel; - this.dimensions = dimensions; - this.distanceType = distanceType; - this.removeExistingVectorStoreTable = removeExistingVectorStoreTable; - this.createIndexMethod = createIndexMethod; - this.initializeSchema = initializeSchema; - } + private String schemaName = PgVectorStore.DEFAULT_SCHEMA_NAME; - public PgDistanceType getDistanceType() { - return distanceType; - } + private String vectorTableName; - @Override - public void add(List documents) { + private boolean vectorTableValidationsEnabled = PgVectorStore.DEFAULT_SCHEMA_VALIDATION; - int size = documents.size(); + private int dimensions = PgVectorStore.INVALID_EMBEDDING_DIMENSION; - this.jdbcTemplate.batchUpdate( - "INSERT INTO " + VECTOR_TABLE_NAME + " (id, content, metadata, embedding) VALUES (?, ?, ?::jsonb, ?) " - + "ON CONFLICT (id) DO " + "UPDATE SET content = ? , metadata = ?::jsonb , embedding = ? ", - new BatchPreparedStatementSetter() { - @Override - public void setValues(PreparedStatement ps, int i) throws SQLException { + private PgDistanceType distanceType = PgDistanceType.COSINE_DISTANCE; - var document = documents.get(i); - var content = document.getContent(); - var json = toJson(document.getMetadata()); - var pGvector = new PGvector(toFloatArray(embeddingModel.embed(document))); + private boolean removeExistingVectorStoreTable = false; - StatementCreatorUtils.setParameterValue(ps, 1, SqlTypeValue.TYPE_UNKNOWN, - UUID.fromString(document.getId())); - StatementCreatorUtils.setParameterValue(ps, 2, SqlTypeValue.TYPE_UNKNOWN, content); - StatementCreatorUtils.setParameterValue(ps, 3, SqlTypeValue.TYPE_UNKNOWN, json); - StatementCreatorUtils.setParameterValue(ps, 4, SqlTypeValue.TYPE_UNKNOWN, pGvector); - StatementCreatorUtils.setParameterValue(ps, 5, SqlTypeValue.TYPE_UNKNOWN, content); - StatementCreatorUtils.setParameterValue(ps, 6, SqlTypeValue.TYPE_UNKNOWN, json); - StatementCreatorUtils.setParameterValue(ps, 7, SqlTypeValue.TYPE_UNKNOWN, pGvector); - } + private PgIndexType indexType = PgIndexType.HNSW; - @Override - public int getBatchSize() { - return size; - } - }); - } + private boolean initializeSchema; - private String toJson(Map map) { - try { - return objectMapper.writeValueAsString(map); - } - catch (JsonProcessingException e) { - throw new RuntimeException(e); - } - } - - private float[] toFloatArray(List embeddingDouble) { - float[] embeddingFloat = new float[embeddingDouble.size()]; - int i = 0; - for (Double d : embeddingDouble) { - embeddingFloat[i++] = d.floatValue(); - } - return embeddingFloat; - } - - @Override - public Optional delete(List idList) { - int updateCount = 0; - for (String id : idList) { - int count = jdbcTemplate.update("DELETE FROM " + VECTOR_TABLE_NAME + " WHERE id = ?", UUID.fromString(id)); - updateCount = updateCount + count; - } - - return Optional.of(updateCount == idList.size()); - } - - @Override - public List similaritySearch(SearchRequest request) { - - String nativeFilterExpression = (request.getFilterExpression() != null) - ? this.filterExpressionConverter.convertExpression(request.getFilterExpression()) : ""; - - String jsonPathFilter = ""; - - if (StringUtils.hasText(nativeFilterExpression)) { - jsonPathFilter = " AND metadata::jsonb @@ '" + nativeFilterExpression + "'::jsonpath "; - } - - double distance = 1 - request.getSimilarityThreshold(); - - PGvector queryEmbedding = getQueryEmbedding(request.getQuery()); - - return this.jdbcTemplate.query( - String.format(this.getDistanceType().similaritySearchSqlTemplate, VECTOR_TABLE_NAME, jsonPathFilter), - new DocumentRowMapper(this.objectMapper), queryEmbedding, queryEmbedding, distance, request.getTopK()); - } - - public List embeddingDistance(String query) { - return this.jdbcTemplate.query( - "SELECT embedding " + this.comparisonOperator() + " ? AS distance FROM " + VECTOR_TABLE_NAME, - new RowMapper() { - @Override - @Nullable - public Double mapRow(ResultSet rs, int rowNum) throws SQLException { - return rs.getDouble(DocumentRowMapper.COLUMN_DISTANCE); - } - - }, getQueryEmbedding(query)); - } - - private PGvector getQueryEmbedding(String query) { - List embedding = this.embeddingModel.embed(query); - return new PGvector(toFloatArray(embedding)); - } - - private String comparisonOperator() { - return this.getDistanceType().operator; - } - - // --------------------------------------------------------------------------------- - // Initialize - // --------------------------------------------------------------------------------- - @Override - public void afterPropertiesSet() throws Exception { - - if (!this.initializeSchema) { - return; - } - - // 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"); - this.jdbcTemplate.execute("CREATE EXTENSION IF NOT EXISTS \"uuid-ossp\""); - - // Remove existing VectorStoreTable - if (this.removeExistingVectorStoreTable) { - this.jdbcTemplate.execute("DROP TABLE IF EXISTS " + VECTOR_TABLE_NAME); - } - - this.jdbcTemplate.execute(String.format(""" - CREATE TABLE IF NOT EXISTS %s ( - id uuid DEFAULT uuid_generate_v4() PRIMARY KEY, - content text, - metadata json, - embedding vector(%d) - ) - """, VECTOR_TABLE_NAME, this.embeddingDimensions())); - - if (this.createIndexMethod != PgIndexType.NONE) { - this.jdbcTemplate.execute(String.format(""" - CREATE INDEX IF NOT EXISTS %s ON %s USING %s (embedding %s) - """, VECTOR_INDEX_NAME, VECTOR_TABLE_NAME, this.createIndexMethod, this.getDistanceType().index)); - } - } - - int embeddingDimensions() { - // The manually set dimensions have precedence over the computed one. - if (this.dimensions > 0) { - return this.dimensions; - } - - try { - int embeddingDimensions = this.embeddingModel.dimensions(); - if (embeddingDimensions > 0) { - return embeddingDimensions; + // 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"); } + this.jdbcTemplate = jdbcTemplate; + this.embeddingModel = embeddingModel; } - catch (Exception e) { - logger.warn("Failed to obtain the embedding dimensions from the embedding model and fall backs to default:" - + OPENAI_EMBEDDING_DIMENSION_SIZE, e); + + public Builder withSchemaName(String schemaName) { + this.schemaName = schemaName; + return this; } - return OPENAI_EMBEDDING_DIMENSION_SIZE; + + public Builder withVectorTableName(String vectorTableName) { + this.vectorTableName = vectorTableName; + return this; + } + + public Builder withVectorTableValidationsEnabled(boolean vectorTableValidationsEnabled) { + this.vectorTableValidationsEnabled = vectorTableValidationsEnabled; + return this; + } + + public Builder withDimensions(int dimensions) { + this.dimensions = dimensions; + return this; + } + + public Builder withDistanceType(PgDistanceType distanceType) { + this.distanceType = distanceType; + return this; + } + + public Builder withRemoveExistingVectorStoreTable(boolean removeExistingVectorStoreTable) { + this.removeExistingVectorStoreTable = removeExistingVectorStoreTable; + return this; + } + + public Builder withIndexType(PgIndexType indexType) { + this.indexType = indexType; + return this; + } + + public Builder withInitializeSchema(boolean initializeSchema) { + this.initializeSchema = initializeSchema; + return this; + } + + public PgVectorStore build() { + return new PgVectorStore(schemaName, vectorTableName, vectorTableValidationsEnabled, jdbcTemplate, + embeddingModel, dimensions, distanceType, removeExistingVectorStoreTable, indexType, + initializeSchema); + } + } } \ No newline at end of file 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/vectorstore/PgVectorStoreCustomNamesIT.java new file mode 100644 index 000000000..a255b0116 --- /dev/null +++ b/vector-stores/spring-ai-pgvector-store/src/test/java/org/springframework/ai/vectorstore/PgVectorStoreCustomNamesIT.java @@ -0,0 +1,244 @@ +/* + * 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. + */ +package org.springframework.ai.vectorstore; + +import com.zaxxer.hikari.HikariDataSource; +import org.junit.jupiter.api.Test; +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.beans.factory.annotation.Value; +import org.springframework.boot.SpringBootConfiguration; +import org.springframework.boot.autoconfigure.EnableAutoConfiguration; +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; +import org.testcontainers.containers.PostgreSQLContainer; +import org.testcontainers.junit.jupiter.Container; +import org.testcontainers.junit.jupiter.Testcontainers; + +import javax.sql.DataSource; +import java.util.Random; + +import static org.assertj.core.api.Assertions.assertThat; + +/** + * @author Muthukumaran Navaneethakrishnan + */ +@Testcontainers +public class PgVectorStoreCustomNamesIT { + + @Container + @SuppressWarnings("resource") + static PostgreSQLContainer postgresContainer = new PostgreSQLContainer<>("pgvector/pgvector:pg16") + .withUsername("postgres") + .withPassword("postgres"); + + private final ApplicationContextRunner contextRunner = new ApplicationContextRunner() + .withUserConfiguration(TestApplication.class) + .withPropertyValues("test.spring.ai.vectorstore.pgvector.distanceType=COSINE_DISTANCE", + + // JdbcTemplate configuration + String.format("app.datasource.url=jdbc:postgresql://%s:%d/%s", postgresContainer.getHost(), + postgresContainer.getMappedPort(5432), "postgres"), + "app.datasource.username=postgres", "app.datasource.password=postgres", + "app.datasource.type=com.zaxxer.hikari.HikariDataSource"); + + private static void dropTableByName(ApplicationContext context, String name) { + JdbcTemplate jdbcTemplate = context.getBean(JdbcTemplate.class); + jdbcTemplate.execute("DROP TABLE IF EXISTS " + name); + } + + private static boolean isIndexExists(ApplicationContext context, String schemaName, String tableName, + String indexName) { + JdbcTemplate jdbcTemplate = context.getBean(JdbcTemplate.class); + String sql = "SELECT EXISTS (SELECT 1 FROM pg_indexes WHERE schemaname = ? AND tablename = ? AND indexname = ?)"; + return jdbcTemplate.queryForObject(sql, Boolean.class, schemaName, tableName, indexName); + } + + @SuppressWarnings("null") + private static boolean isTableExists(ApplicationContext context, String tableName) { + JdbcTemplate jdbcTemplate = context.getBean(JdbcTemplate.class); + return jdbcTemplate.queryForObject( + "SELECT EXISTS (SELECT FROM information_schema.tables WHERE table_name = '" + tableName + "')", + Boolean.class); + } + + private static boolean isSchemaExists(ApplicationContext context, String schemaName) { + JdbcTemplate jdbcTemplate = context.getBean(JdbcTemplate.class); + String sql = "SELECT EXISTS (SELECT FROM information_schema.schemata WHERE schema_name = ?)"; + return jdbcTemplate.queryForObject(sql, Boolean.class, schemaName); + } + + @Test + public void shouldCreateDefaultTableAndIndexIfNotPresentInConfig() { + contextRunner.withPropertyValues("test.spring.ai.vectorstore.pgvector.schemaValidation=false").run(context -> { + assertThat(context).hasNotFailed(); + assertThat(isTableExists(context, "vector_store")).isTrue(); + assertThat(isSchemaExists(context, "public")).isTrue(); + dropTableByName(context, "vector_store"); + + }); + } + + @Test + public void shouldCreateTableAndIndexIfNotPresentInDatabase() { + String tableName = "new_vector_table"; + contextRunner.withPropertyValues("test.spring.ai.vectorstore.pgvector.vectorTableName=" + tableName) + .run(context -> { + assertThat(isTableExists(context, tableName)).isTrue(); + assertThat(isIndexExists(context, "public", tableName, tableName + "_index")).isTrue(); + assertThat(isTableExists(context, "vector_store")).isFalse(); + dropTableByName(context, tableName); + }); + } + + @Test + public void shouldFailWhenCustomTableIsAbsentAndValidationEnabled() { + + String tableName = "customvectortable"; + + contextRunner + .withPropertyValues("test.spring.ai.vectorstore.pgvector.vectorTableName=" + tableName, + "test.spring.ai.vectorstore.pgvector.schemaValidation=true") + + .run(context -> { + + assertThat(context).hasFailed(); + assertThat(context.getStartupFailure()).hasCauseInstanceOf(IllegalStateException.class) + .hasMessageContaining(tableName + " does not exist"); + + }); + } + + @Test + public void shouldFailOnSQLInjectionAttemptInTableName() { + + String tableName = "users; DROP TABLE users;"; + + contextRunner + .withPropertyValues("test.spring.ai.vectorstore.pgvector.vectorTableName=" + tableName, + "test.spring.ai.vectorstore.pgvector.schemaValidation=true") + + .run(context -> { + + assertThat(context).hasFailed(); + assertThat(context.getStartupFailure()).hasCauseInstanceOf(IllegalArgumentException.class) + .hasMessageContaining("Table name should only contain alphanumeric characters and underscores"); + + }); + + } + + @Test + public void shouldFailOnSQLInjectionAttemptInSchemaName() { + + String schemaName = "public; DROP TABLE users;"; + String tableName = "customvectortable"; + + contextRunner + .withPropertyValues("test.spring.ai.vectorstore.pgvector.vectorTableName=" + tableName, + "test.spring.ai.vectorstore.pgvector.schemaName=" + schemaName, + "test.spring.ai.vectorstore.pgvector.schemaValidation=true") + + .run(context -> { + + assertThat(context).hasFailed(); + assertThat(context.getStartupFailure()).hasCauseInstanceOf(IllegalArgumentException.class) + .hasMessageContaining("Schema name should only contain alphanumeric characters and underscores"); + + }); + + } + + @SpringBootConfiguration + @EnableAutoConfiguration(exclude = { DataSourceAutoConfiguration.class }) + public static class TestApplication { + + @Value("${test.spring.ai.vectorstore.pgvector.vectorTableName:}") + String vectorTableName; + + @Value("${test.spring.ai.vectorstore.pgvector.schemaName:public}") + String schemaName; + + @Value("${test.spring.ai.vectorstore.pgvector.schemaValidation:false}") + boolean schemaValidation; + + int dimensions = 768; + + @Bean + public VectorStore vectorStore(JdbcTemplate jdbcTemplate, EmbeddingModel embeddingModel) { + + return new PgVectorStore.Builder(jdbcTemplate, embeddingModel).withSchemaName(schemaName) + .withVectorTableName(vectorTableName) + .withVectorTableValidationsEnabled(schemaValidation) + .withDimensions(dimensions) + .withDistanceType(PgVectorStore.PgDistanceType.COSINE_DISTANCE) + .withRemoveExistingVectorStoreTable(true) + .withIndexType(PgIndexType.HNSW) + .withInitializeSchema(true) + .build(); + } + + public Float[] generateFloatArray(int size, float min, float max) { + float[] result = new float[size]; + Random random = new Random(); + + for (int i = 0; i < size; i++) { + result[i] = min + random.nextFloat() * (max - min); + } + Float[] embeddingObjects = new Float[result.length]; + for (int i = 0; i < result.length; i++) { + embeddingObjects[i] = result[i]; // Auto-boxing float to Float + } + + return embeddingObjects; + } + + // + @Bean + public JdbcTemplate myJdbcTemplate(DataSource dataSource) { + JdbcTemplate jdbcTemplate = new JdbcTemplate(dataSource); + + return jdbcTemplate; + } + + @Bean + @Primary + @ConfigurationProperties("app.datasource") + public DataSourceProperties dataSourceProperties() { + return new DataSourceProperties(); + } + + @Bean + public HikariDataSource dataSource(DataSourceProperties dataSourceProperties) { + return dataSourceProperties.initializeDataSourceBuilder().type(HikariDataSource.class).build(); + } + + @Bean + public EmbeddingModel embeddingModel() { + return new OpenAiEmbeddingModel(new OpenAiApi(System.getenv("OPENAI_API_KEY"))); + } + + } + +} 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/vectorstore/PgVectorStoreTests.java new file mode 100644 index 000000000..f5b69a922 --- /dev/null +++ b/vector-stores/spring-ai-pgvector-store/src/test/java/org/springframework/ai/vectorstore/PgVectorStoreTests.java @@ -0,0 +1,60 @@ +/* + * 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. + */ +package org.springframework.ai.vectorstore; + +import org.junit.jupiter.params.ParameterizedTest; +import org.junit.jupiter.params.provider.CsvSource; + +import static org.assertj.core.api.Assertions.assertThat; + +/** + * @author Muthukumaran Navaneethakrishnan + */ + +public class PgVectorStoreTests { + + @ParameterizedTest(name = "{0} - Verifies valid Table name") + @CsvSource({ + // Standard valid cases + "customvectorstore, true", "user_data, true", "test123, true", "valid_table_name, true", + + // Edge cases + "'', false", // Empty string + " , false", // Spaces only + "custom vector store, false", // Spaces in name + "customvectorstore;, false", // Semicolon appended + "customvectorstore--, false", // SQL comment appended + "drop table users;, false", // SQL command as a name + "customvectorstore;drop table users;, false", // Valid name followed by + // command + "customvectorstore#, false", // Hash character included + "customvectorstore$, false", // Dollar sign included + "1, false", // Numeric only + "customvectorstore or 1=1, false", // SQL Injection attempt + "customvectorstore;--, false", // Ending with comment + "custom_vector_store; DROP TABLE users;, false", // Injection with valid part + "'customvectorstore\u0000', false", // Null byte included + "'customvectorstore\n', false", // Newline character + "12345678901234567890123456789012345678901234567890123456789012345, false" // More + // than + // 64 + // characters + }) + public void isValidTable(String tableName, Boolean expected) { + assertThat(PgVectorSchemaValidator.isValidNameForDatabaseObject(tableName)).isEqualTo(expected); + } + +}