Add support for customizable PgVectorStore schema, table, and index names (#747)
This change allows users to specify custom names, facilitating management of multiple vector databases within a single database instance. Key changes: - Implement configurable schema, table, and index names for PgVectorStore - Add properties to set custom schema and table names - Introduce optional schema/table & field validation for custom configurations - Include additional tests for new configurations and existing deployments Additional improvements: - Rename properties to schemaName, tableName, and schemaValidation - Update pgvector documentation with new properties - Add schema/table name tests to PgVectorStoreAutoConfigurationIT and PgVectorStorePropertiesTests - Create standalone PgVectorSchemaValidator class for schema/table validation - Add missing 'CREATE SCHEMA IF NOT EXISTS' when initializeSchema=true - Remove redundant code and classes Resolves #747 Co-authored-by: Christian Tzolov <ctzolov@vmware.com>
This commit is contained in:
Muthukumaran Navaneethakrishnan
committed by
Christian Tzolov
Christian Tzolov
parent
067a33dbe2
commit
f0ca61252b
@@ -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.
|
||||
|
||||
@@ -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();
|
||||
}
|
||||
|
||||
}
|
||||
|
||||
@@ -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;
|
||||
}
|
||||
|
||||
}
|
||||
|
||||
@@ -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<Document> 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<Document> 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();
|
||||
}
|
||||
|
||||
}
|
||||
|
||||
}
|
||||
@@ -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<Document> 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);
|
||||
}
|
||||
|
||||
}
|
||||
|
||||
@@ -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");
|
||||
}
|
||||
|
||||
}
|
||||
|
||||
@@ -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<String> 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<Map<String, Object>> 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<String> availableColumns = new ArrayList<>();
|
||||
for (Map<String, Object> 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);
|
||||
|
||||
}
|
||||
}
|
||||
|
||||
}
|
||||
@@ -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<Document> 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<String, Object> map) {
|
||||
try {
|
||||
return objectMapper.writeValueAsString(map);
|
||||
}
|
||||
catch (JsonProcessingException e) {
|
||||
throw new RuntimeException(e);
|
||||
}
|
||||
}
|
||||
|
||||
private float[] toFloatArray(List<Double> embeddingDouble) {
|
||||
float[] embeddingFloat = new float[embeddingDouble.size()];
|
||||
int i = 0;
|
||||
for (Double d : embeddingDouble) {
|
||||
embeddingFloat[i++] = d.floatValue();
|
||||
}
|
||||
return embeddingFloat;
|
||||
}
|
||||
|
||||
@Override
|
||||
public Optional<Boolean> delete(List<String> 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<Document> 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<Double> embeddingDistance(String query) {
|
||||
return this.jdbcTemplate.query(
|
||||
"SELECT embedding " + this.comparisonOperator() + " ? AS distance FROM " + getFullyQualifiedTableName(),
|
||||
new RowMapper<Double>() {
|
||||
@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<Double> 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<Document> 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<String, Object> map) {
|
||||
try {
|
||||
return objectMapper.writeValueAsString(map);
|
||||
}
|
||||
catch (JsonProcessingException e) {
|
||||
throw new RuntimeException(e);
|
||||
}
|
||||
}
|
||||
|
||||
private float[] toFloatArray(List<Double> embeddingDouble) {
|
||||
float[] embeddingFloat = new float[embeddingDouble.size()];
|
||||
int i = 0;
|
||||
for (Double d : embeddingDouble) {
|
||||
embeddingFloat[i++] = d.floatValue();
|
||||
}
|
||||
return embeddingFloat;
|
||||
}
|
||||
|
||||
@Override
|
||||
public Optional<Boolean> delete(List<String> 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<Document> 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<Double> embeddingDistance(String query) {
|
||||
return this.jdbcTemplate.query(
|
||||
"SELECT embedding " + this.comparisonOperator() + " ? AS distance FROM " + VECTOR_TABLE_NAME,
|
||||
new RowMapper<Double>() {
|
||||
@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<Double> 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);
|
||||
}
|
||||
|
||||
}
|
||||
|
||||
}
|
||||
@@ -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")));
|
||||
}
|
||||
|
||||
}
|
||||
|
||||
}
|
||||
@@ -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);
|
||||
}
|
||||
|
||||
}
|
||||
Reference in New Issue
Block a user