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

Remove UUID Enforcement for ID Column in PGVectorStore

Browse Source 
Signed-off-by: jitokim <pigberger70@gmail.com>
This commit is contained in:
jitokim
2024-12-28 00:25:56 +09:00
committed by Ilayaperumal Gopinathan
parent bb6b8cc20f
commit bac507c62a
3 changed files with 22 additions and 2 deletions
Download Patch File Download Diff File Expand all files Collapse all files

4
spring-ai-core/src/main/java/org/springframework/ai/document/Document.java
View File

@@ -136,6 +136,10 @@ public class Document {
this(new RandomIdGenerator().generateId(), text, null, metadata, null);
}
public Document(String id, String text) {
this(id, text, new HashMap<>());
}
public Document(String id, String text, Map<String, Object> metadata) {
this(id, text, null, metadata, null);
}

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

@@ -152,6 +152,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 {
@@ -272,13 +273,13 @@ public class PgVectorStore extends AbstractObservationVectorStore implements Ini
public void setValues(PreparedStatement ps, int i) throws SQLException {
var document = batch.get(i);
var id = 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);

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

@@ -29,6 +29,7 @@ import javax.sql.DataSource;
import com.zaxxer.hikari.HikariDataSource;
import org.junit.Assert;
import org.junit.jupiter.api.Test;
import org.junit.jupiter.api.condition.EnabledIfEnvironmentVariable;
import org.junit.jupiter.params.ParameterizedTest;
import org.junit.jupiter.params.provider.Arguments;
@@ -67,6 +68,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 = ".+")
@@ -166,6 +168,19 @@ public class PgVectorStoreIT {
});
}
@Test
public void shouldAllowNonUuidFormat() {
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("NOT_UUID", "TEXT")));
dropTable(context);
});
}
@ParameterizedTest(name = "Filter expression {0} should return {1} records ")
@MethodSource("provideFilters")
public void searchWithInFilter(String expression, Integer expectedRecords) {