Skip to content

Commit bac507c

Browse files
jitokimilayaperumalg
authored andcommitted
Remove UUID Enforcement for ID Column in PGVectorStore
Signed-off-by: jitokim <[email protected]>
1 parent bb6b8cc commit bac507c

File tree

3 files changed

+22
-2
lines changed
  • spring-ai-core/src/main/java/org/springframework/ai/document
  • vector-stores/spring-ai-pgvector-store/src

3 files changed

+22
-2
lines changed

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

Lines changed: 4 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -136,6 +136,10 @@ public Document(String text, Map<String, Object> metadata) {
136136
this(new RandomIdGenerator().generateId(), text, null, metadata, null);
137137
}
138138

139+
public Document(String id, String text) {
140+
this(id, text, new HashMap<>());
141+
}
142+
139143
public Document(String id, String text, Map<String, Object> metadata) {
140144
this(id, text, null, metadata, null);
141145
}

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

Lines changed: 3 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -152,6 +152,7 @@
152152
* @author Thomas Vitale
153153
* @author Soby Chacko
154154
* @author Sebastien Deleuze
155+
* @author Jihoon Kim
155156
* @since 1.0.0
156157
*/
157158
public class PgVectorStore extends AbstractObservationVectorStore implements InitializingBean {
@@ -272,13 +273,13 @@ private void insertOrUpdateBatch(List<Document> batch, List<Document> documents,
272273
public void setValues(PreparedStatement ps, int i) throws SQLException {
273274

274275
var document = batch.get(i);
276+
var id = document.getId();
275277
var content = document.getText();
276278
var json = toJson(document.getMetadata());
277279
var embedding = embeddings.get(documents.indexOf(document));
278280
var pGvector = new PGvector(embedding);
279281

280-
StatementCreatorUtils.setParameterValue(ps, 1, SqlTypeValue.TYPE_UNKNOWN,
281-
UUID.fromString(document.getId()));
282+
StatementCreatorUtils.setParameterValue(ps, 1, SqlTypeValue.TYPE_UNKNOWN, id);
282283
StatementCreatorUtils.setParameterValue(ps, 2, SqlTypeValue.TYPE_UNKNOWN, content);
283284
StatementCreatorUtils.setParameterValue(ps, 3, SqlTypeValue.TYPE_UNKNOWN, json);
284285
StatementCreatorUtils.setParameterValue(ps, 4, SqlTypeValue.TYPE_UNKNOWN, pGvector);

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

Lines changed: 15 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -29,6 +29,7 @@
2929

3030
import com.zaxxer.hikari.HikariDataSource;
3131
import org.junit.Assert;
32+
import org.junit.jupiter.api.Test;
3233
import org.junit.jupiter.api.condition.EnabledIfEnvironmentVariable;
3334
import org.junit.jupiter.params.ParameterizedTest;
3435
import org.junit.jupiter.params.provider.Arguments;
@@ -67,6 +68,7 @@
6768
* @author Muthukumaran Navaneethakrishnan
6869
* @author Christian Tzolov
6970
* @author Thomas Vitale
71+
* @author Jihoon Kim
7072
*/
7173
@Testcontainers
7274
@EnabledIfEnvironmentVariable(named = "OPENAI_API_KEY", matches = ".+")
@@ -166,6 +168,19 @@ public void addAndSearch(String distanceType) {
166168
});
167169
}
168170

171+
@Test
172+
public void shouldAllowNonUuidFormat() {
173+
this.contextRunner.withPropertyValues("test.spring.ai.vectorstore.pgvector.distanceType=" + "COSINE_DISTANCE")
174+
.run(context -> {
175+
176+
VectorStore vectorStore = context.getBean(VectorStore.class);
177+
178+
vectorStore.add(List.of(new Document("NOT_UUID", "TEXT")));
179+
180+
dropTable(context);
181+
});
182+
}
183+
169184
@ParameterizedTest(name = "Filter expression {0} should return {1} records ")
170185
@MethodSource("provideFilters")
171186
public void searchWithInFilter(String expression, Integer expectedRecords) {

0 commit comments

Comments
 (0)