Spring/spring-ai
1
0
Fork 0
Code Issues Packages Projects Releases Wiki Activity

Refactor ID handling for different IdType formats

Browse Source 
- Add handling for UUID, TEXT, INTEGER, SERIAL, BIGSERIAL formats in `convertIdToPgType` function.
- Implemented type conversion logic based on the IdType value (UUID, TEXT, INTEGER, SERIAL, BIGSERIAL).
- Add unit tests to validate correct conversion for UUID and non-UUID IdType formats.
  - `testToPgTypeWithUuidIdType`: Validates UUID handling.
  - `testToPgTypeWithNonUuidIdType`: Validates handling for non-UUID IdTypes.

Signed-off-by: jitokim <pigberger70@gmail.com>
This commit is contained in:
jitokim
2025-01-24 23:42:34 +09:00
committed by Ilayaperumal Gopinathan
parent f40945bd63
commit 4dbe73433a
2 changed files with 101 additions and 6 deletions
Download Patch File Download Diff File Expand all files Collapse all files

42
vector-stores/spring-ai-pgvector-store/src/main/java/org/springframework/ai/vectorstore/pgvector/PgVectorStore.java
View File

@@ -35,10 +35,8 @@ import org.slf4j.LoggerFactory;
import org.springframework.ai.document.Document;
import org.springframework.ai.document.DocumentMetadata;
import org.springframework.ai.embedding.BatchingStrategy;
import org.springframework.ai.embedding.EmbeddingModel;
import org.springframework.ai.embedding.EmbeddingOptionsBuilder;
import org.springframework.ai.embedding.TokenCountBatchingStrategy;
import org.springframework.ai.observation.conventions.VectorStoreProvider;
import org.springframework.ai.observation.conventions.VectorStoreSimilarityMetric;
import org.springframework.ai.util.JacksonUtils;
@@ -153,6 +151,7 @@ import org.springframework.util.StringUtils;
* @author Thomas Vitale
* @author Soby Chacko
* @author Sebastien Deleuze
* @author Jihoon Kim
* @since 1.0.0
*/
public class PgVectorStore extends AbstractObservationVectorStore implements InitializingBean {
@@ -163,6 +162,8 @@ public class PgVectorStore extends AbstractObservationVectorStore implements Ini
public static final String DEFAULT_TABLE_NAME = "vector_store";
public static final PgIdType DEFAULT_ID_TYPE = PgIdType.UUID;
public static final String DEFAULT_VECTOR_INDEX_NAME = "spring_ai_vector_index";
public static final String DEFAULT_SCHEMA_NAME = "public";
@@ -188,6 +189,8 @@ public class PgVectorStore extends AbstractObservationVectorStore implements Ini
private final String schemaName;
private final PgIdType idType;
private final boolean schemaValidation;
private final boolean initializeSchema;
@@ -225,6 +228,7 @@ public class PgVectorStore extends AbstractObservationVectorStore implements Ini
: this.vectorTableName + "_index";
this.schemaName = builder.schemaName;
this.idType = builder.idType;
this.schemaValidation = builder.vectorTableValidationsEnabled;
this.jdbcTemplate = builder.jdbcTemplate;
@@ -273,13 +277,13 @@ public class PgVectorStore extends AbstractObservationVectorStore implements Ini
public void setValues(PreparedStatement ps, int i) throws SQLException {
var document = batch.get(i);
var id = convertIdToPgType(document.getId());
var content = document.getText();
var json = toJson(document.getMetadata());
var embedding = embeddings.get(documents.indexOf(document));
var pGvector = new PGvector(embedding);
StatementCreatorUtils.setParameterValue(ps, 1, SqlTypeValue.TYPE_UNKNOWN,
UUID.fromString(document.getId()));
StatementCreatorUtils.setParameterValue(ps, 1, SqlTypeValue.TYPE_UNKNOWN, id);
StatementCreatorUtils.setParameterValue(ps, 2, SqlTypeValue.TYPE_UNKNOWN, content);
StatementCreatorUtils.setParameterValue(ps, 3, SqlTypeValue.TYPE_UNKNOWN, json);
StatementCreatorUtils.setParameterValue(ps, 4, SqlTypeValue.TYPE_UNKNOWN, pGvector);
@@ -304,6 +308,19 @@ public class PgVectorStore extends AbstractObservationVectorStore implements Ini
}
}
private Object convertIdToPgType(String id) {
if (this.initializeSchema) {
return UUID.fromString(id);
}
return switch (getIdType()) {
case UUID -> UUID.fromString(id);
case TEXT -> id;
case INTEGER, SERIAL -> Integer.valueOf(id);
case BIGSERIAL -> Long.valueOf(id);
};
}
@Override
public Optional<Boolean> doDelete(List<String> idList) {
int updateCount = 0;
@@ -429,6 +446,10 @@ public class PgVectorStore extends AbstractObservationVectorStore implements Ini
return this.schemaName + "." + this.vectorTableName;
}
private PgIdType getIdType() {
return this.idType;
}
private String getVectorTableName() {
return this.vectorTableName;
}
@@ -513,6 +534,12 @@ public class PgVectorStore extends AbstractObservationVectorStore implements Ini
}
public enum PgIdType {
UUID, TEXT, INTEGER, SERIAL, BIGSERIAL
}
/**
* Defaults to CosineDistance. But if vectors are normalized to length 1 (like OpenAI
* embeddings), use inner product (NegativeInnerProduct) for best performance.
@@ -608,6 +635,8 @@ public class PgVectorStore extends AbstractObservationVectorStore implements Ini
private String vectorTableName = PgVectorStore.DEFAULT_TABLE_NAME;
private PgIdType idType = PgVectorStore.DEFAULT_ID_TYPE;
private boolean vectorTableValidationsEnabled = PgVectorStore.DEFAULT_SCHEMA_VALIDATION;
private int dimensions = PgVectorStore.INVALID_EMBEDDING_DIMENSION;
@@ -638,6 +667,11 @@ public class PgVectorStore extends AbstractObservationVectorStore implements Ini
return this;
}
public PgVectorStoreBuilder idType(PgIdType idType) {
this.idType = idType;
return this;
}
public PgVectorStoreBuilder vectorTableValidationsEnabled(boolean vectorTableValidationsEnabled) {
this.vectorTableValidationsEnabled = vectorTableValidationsEnabled;
return this;

65
vector-stores/spring-ai-pgvector-store/src/test/java/org/springframework/ai/vectorstore/pgvector/PgVectorStoreIT.java
View File

@@ -19,6 +19,7 @@ package org.springframework.ai.vectorstore.pgvector;
import java.io.IOException;
import java.nio.charset.StandardCharsets;
import java.util.Collections;
import java.util.HashMap;
import java.util.Iterator;
import java.util.List;
import java.util.Map;
@@ -42,14 +43,16 @@ import org.testcontainers.junit.jupiter.Testcontainers;
import org.springframework.ai.document.Document;
import org.springframework.ai.document.DocumentMetadata;
import org.springframework.ai.document.id.RandomIdGenerator;
import org.springframework.ai.embedding.EmbeddingModel;
import org.springframework.ai.openai.OpenAiEmbeddingModel;
import org.springframework.ai.openai.api.OpenAiApi;
import org.springframework.ai.vectorstore.pgvector.PgVectorStore.PgIdType;
import org.springframework.ai.vectorstore.pgvector.PgVectorStore.PgIndexType;
import org.springframework.ai.vectorstore.SearchRequest;
import org.springframework.ai.vectorstore.VectorStore;
import org.springframework.ai.vectorstore.filter.Filter;
import org.springframework.ai.vectorstore.filter.FilterExpressionTextParser.FilterExpressionParseException;
import org.springframework.ai.vectorstore.pgvector.PgVectorStore.PgIndexType;
import org.springframework.beans.factory.annotation.Value;
import org.springframework.boot.SpringBootConfiguration;
import org.springframework.boot.autoconfigure.EnableAutoConfiguration;
@@ -70,6 +73,7 @@ import static org.assertj.core.api.Assertions.assertThat;
* @author Muthukumaran Navaneethakrishnan
* @author Christian Tzolov
* @author Thomas Vitale
* @author Jihoon Kim
*/
@Testcontainers
@EnabledIfEnvironmentVariable(named = "OPENAI_API_KEY", matches = ".+")
@@ -106,6 +110,27 @@ public class PgVectorStoreIT {
}
}
private static void initSchema(ApplicationContext context) {
PgVectorStore vectorStore = context.getBean(PgVectorStore.class);
JdbcTemplate jdbcTemplate = context.getBean(JdbcTemplate.class);
// Enable the PGVector, JSONB and UUID support.
jdbcTemplate.execute("CREATE EXTENSION IF NOT EXISTS vector");
jdbcTemplate.execute("CREATE EXTENSION IF NOT EXISTS hstore");
jdbcTemplate.execute("CREATE EXTENSION IF NOT EXISTS \"uuid-ossp\"");
jdbcTemplate.execute(String.format("CREATE SCHEMA IF NOT EXISTS %s", PgVectorStore.DEFAULT_SCHEMA_NAME));
jdbcTemplate.execute(String.format("""
CREATE TABLE IF NOT EXISTS %s.%s (
id text PRIMARY KEY,
content text,
metadata json,
embedding vector(%d)
)
""", PgVectorStore.DEFAULT_SCHEMA_NAME, PgVectorStore.DEFAULT_TABLE_NAME,
vectorStore.embeddingDimensions()));
}
private static void dropTable(ApplicationContext context) {
JdbcTemplate jdbcTemplate = context.getBean(JdbcTemplate.class);
jdbcTemplate.execute("DROP TABLE IF EXISTS vector_store");
@@ -169,6 +194,35 @@ public class PgVectorStoreIT {
});
}
@Test
public void testToPgTypeWithUuidIdType() {
this.contextRunner.withPropertyValues("test.spring.ai.vectorstore.pgvector.distanceType=" + "COSINE_DISTANCE")
.run(context -> {
VectorStore vectorStore = context.getBean(VectorStore.class);
vectorStore.add(List.of(new Document(new RandomIdGenerator().generateId(), "TEXT", new HashMap<>())));
dropTable(context);
});
}
@Test
public void testToPgTypeWithNonUuidIdType() {
this.contextRunner.withPropertyValues("test.spring.ai.vectorstore.pgvector.distanceType=" + "COSINE_DISTANCE")
.withPropertyValues("test.spring.ai.vectorstore.pgvector.initializeSchema=" + false)
.withPropertyValues("test.spring.ai.vectorstore.pgvector.idType=" + "TEXT")
.run(context -> {
VectorStore vectorStore = context.getBean(VectorStore.class);
initSchema(context);
vectorStore.add(List.of(new Document("NOT_UUID", "TEXT", new HashMap<>())));
dropTable(context);
});
}
@ParameterizedTest(name = "Filter expression {0} should return {1} records ")
@MethodSource("provideFilters")
public void searchWithInFilter(String expression, Integer expectedRecords) {
@@ -498,12 +552,19 @@ public class PgVectorStoreIT {
@Value("${test.spring.ai.vectorstore.pgvector.distanceType}")
PgVectorStore.PgDistanceType distanceType;
@Value("${test.spring.ai.vectorstore.pgvector.initializeSchema:true}")
boolean initializeSchema;
@Value("${test.spring.ai.vectorstore.pgvector.idType:UUID}")
PgIdType idType;
@Bean
public VectorStore vectorStore(JdbcTemplate jdbcTemplate, EmbeddingModel embeddingModel) {
return PgVectorStore.builder(jdbcTemplate, embeddingModel)
.dimensions(PgVectorStore.INVALID_EMBEDDING_DIMENSION)
.idType(idType)
.distanceType(this.distanceType)
.initializeSchema(true)
.initializeSchema(initializeSchema)
.indexType(PgIndexType.HNSW)
.removeExistingVectorStoreTable(true)
.build();