From bac507c62a9b2feeb21c9c0c263cd43c5a037cb7 Mon Sep 17 00:00:00 2001
From: jitokim <pigberger70@gmail.com>
Date: Sat, 28 Dec 2024 00:25:56 +0900
Subject: [PATCH] Remove UUID Enforcement for ID Column in PGVectorStore

Signed-off-by: jitokim <pigberger70@gmail.com>
---
 .../org/springframework/ai/document/Document.java |  4 ++++
 .../ai/vectorstore/pgvector/PgVectorStore.java    |  5 +++--
 .../ai/vectorstore/pgvector/PgVectorStoreIT.java  | 15 +++++++++++++++
 3 files changed, 22 insertions(+), 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
index 7d02ae50e..506bbdc34 100644
--- 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 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);
 	}
diff --git a/vector-stores/spring-ai-pgvector-store/src/main/java/org/springframework/ai/vectorstore/pgvector/PgVectorStore.java b/vector-stores/spring-ai-pgvector-store/src/main/java/org/springframework/ai/vectorstore/pgvector/PgVectorStore.java
index 19b70507d..4e1fa7666 100644
--- a/vector-stores/spring-ai-pgvector-store/src/main/java/org/springframework/ai/vectorstore/pgvector/PgVectorStore.java
+++ b/vector-stores/spring-ai-pgvector-store/src/main/java/org/springframework/ai/vectorstore/pgvector/PgVectorStore.java
@@ -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);
diff --git a/vector-stores/spring-ai-pgvector-store/src/test/java/org/springframework/ai/vectorstore/pgvector/PgVectorStoreIT.java b/vector-stores/spring-ai-pgvector-store/src/test/java/org/springframework/ai/vectorstore/pgvector/PgVectorStoreIT.java
index 82a70237a..ea5766371 100644
--- a/vector-stores/spring-ai-pgvector-store/src/test/java/org/springframework/ai/vectorstore/pgvector/PgVectorStoreIT.java
+++ b/vector-stores/spring-ai-pgvector-store/src/test/java/org/springframework/ai/vectorstore/pgvector/PgVectorStoreIT.java
@@ -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) {