From 5e0c909aebc907345ef5a2346d854cf54139baea Mon Sep 17 00:00:00 2001
From: Ilayaperumal Gopinathan <151690+ilayaperumalg@users.noreply.github.com>
Date: Fri, 24 Jan 2025 12:37:59 +0000
Subject: [PATCH] Revert "Remove UUID Enforcement for ID Column in
 PGVectorStore"

This reverts commit bac507c62a9b2feeb21c9c0c263cd43c5a037cb7.
---
 .../org/springframework/ai/document/Document.java |  4 ----
 .../ai/vectorstore/pgvector/PgVectorStore.java    |  5 ++---
 .../ai/vectorstore/pgvector/PgVectorStoreIT.java  | 15 ---------------
 3 files changed, 2 insertions(+), 22 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 506bbdc3478..7d02ae50e6a 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,10 +136,6 @@ 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-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 4e1fa76669e..19b70507db8 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,7 +152,6 @@
  * @author Thomas Vitale
  * @author Soby Chacko
  * @author Sebastien Deleuze
- * @author Jihoon Kim
  * @since 1.0.0
  */
 public class PgVectorStore extends AbstractObservationVectorStore implements InitializingBean {
@@ -273,13 +272,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, id);
+				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);
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 ea5766371be..82a70237a2f 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,7 +29,6 @@
 
 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;
@@ -68,7 +67,6 @@
  * @author Muthukumaran Navaneethakrishnan
  * @author Christian Tzolov
  * @author Thomas Vitale
- * @author Jihoon Kim
  */
 @Testcontainers
 @EnabledIfEnvironmentVariable(named = "OPENAI_API_KEY", matches = ".+")
@@ -168,19 +166,6 @@ 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) {