Remove UUID Enforcement for ID Column in PGVectorStore

Signed-off-by: jitokim <[email protected]>
spring-projects · Dec 27, 2024 · 61409f0 · 61409f0
1 parent 34ac319
commit 61409f0
Show file tree

Hide file tree

Showing 3 changed files with 22 additions and 2 deletions.
diff --git a/spring-ai-core/src/main/java/org/springframework/ai/document/Document.java b/spring-ai-core/src/main/java/org/springframework/ai/document/Document.java
@@ -136,6 +136,10 @@ public Document(String text, Map<String, Object> metadata) {
 		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);
 	}

diff --git a/...vector-store/src/main/java/org/springframework/ai/vectorstore/pgvector/PgVectorStore.java b/...vector-store/src/main/java/org/springframework/ai/vectorstore/pgvector/PgVectorStore.java
@@ -155,6 +155,7 @@
  * @author Thomas Vitale
  * @author Soby Chacko
  * @author Sebastien Deleuze
+ * @author Jihoon Kim
  * @since 1.0.0
  */
 public class PgVectorStore extends AbstractObservationVectorStore implements InitializingBean {
@@ -313,13 +314,13 @@ private void insertOrUpdateBatch(List<Document> batch, List<Document> documents,
 			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);

diff --git a/...ctor-store/src/test/java/org/springframework/ai/vectorstore/pgvector/PgVectorStoreIT.java b/...ctor-store/src/test/java/org/springframework/ai/vectorstore/pgvector/PgVectorStoreIT.java
@@ -29,6 +29,7 @@
 
 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 @@
  * @author Muthukumaran Navaneethakrishnan
  * @author Christian Tzolov
  * @author Thomas Vitale
+ * @author Jihoon Kim
  */
 @Testcontainers
 @EnabledIfEnvironmentVariable(named = "OPENAI_API_KEY", matches = ".+")
@@ -166,6 +168,19 @@ public void addAndSearch(String distanceType) {
 			});
 	}
 
+	@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) {