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 ffccfb353e..abcd200d2d 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 @@ -635,7 +635,13 @@ public Document mapRow(ResultSet rs, int rowNum) throws SQLException { Float distance = rs.getFloat(COLUMN_DISTANCE); Map metadata = toMap(pgMetadata); - metadata.put(DocumentMetadata.DISTANCE.value(), distance); + if (metadata.containsKey(DocumentMetadata.DISTANCE.value())) { + logger.warn("Skipping computed similarity distance for document " + id + + " because its metadata already contains a '" + DocumentMetadata.DISTANCE.value() + "' key"); + } + else { + metadata.put(DocumentMetadata.DISTANCE.value(), distance); + } // @formatter:off return Document.builder() diff --git a/vector-stores/spring-ai-pgvector-store/src/test/java/org/springframework/ai/vectorstore/pgvector/PgVectorStoreTests.java b/vector-stores/spring-ai-pgvector-store/src/test/java/org/springframework/ai/vectorstore/pgvector/PgVectorStoreTests.java index 8cf4c3990e..43a54a6a2f 100644 --- a/vector-stores/spring-ai-pgvector-store/src/test/java/org/springframework/ai/vectorstore/pgvector/PgVectorStoreTests.java +++ b/vector-stores/spring-ai-pgvector-store/src/test/java/org/springframework/ai/vectorstore/pgvector/PgVectorStoreTests.java @@ -16,6 +16,7 @@ package org.springframework.ai.vectorstore.pgvector; +import java.sql.ResultSet; import java.util.Collections; import java.util.List; @@ -24,8 +25,10 @@ import org.junit.jupiter.params.provider.CsvSource; import org.mockito.ArgumentCaptor; import org.mockito.ArgumentMatchers; +import org.postgresql.util.PGobject; import org.springframework.ai.document.Document; +import org.springframework.ai.document.DocumentMetadata; import org.springframework.ai.embedding.EmbeddingModel; import org.springframework.ai.vectorstore.SearchRequest; import org.springframework.ai.vectorstore.filter.Filter; @@ -206,4 +209,59 @@ void similaritySearchDoublesSingleQuotesInsideJsonPathSqlLiteral() { assertThat(sqlCaptor.getValue()).contains("O''Brien"); } + @Test + @SuppressWarnings("unchecked") + void rowMapperSkipsDistanceMetadataWhenUserAlreadyUsesThatKey() throws Exception { + var jdbcTemplate = mock(JdbcTemplate.class); + var embeddingModel = mock(EmbeddingModel.class); + ArgumentCaptor> rowMapperCaptor = ArgumentCaptor.forClass(RowMapper.class); + when(jdbcTemplate.query(anyString(), rowMapperCaptor.capture(), any(), any(), any(), any())) + .thenReturn(List.of()); + + var store = PgVectorStore.builder(jdbcTemplate, embeddingModel).build(); + store.doSimilaritySearch(SearchRequest.builder().query("hello").topK(5).similarityThresholdAll().build()); + + var pgMetadata = new PGobject(); + pgMetadata.setType("json"); + pgMetadata.setValue("{\"distance\": \"12.5 miles from depot\"}"); + + var resultSet = mock(ResultSet.class); + when(resultSet.getString("id")).thenReturn("doc-1"); + when(resultSet.getString("content")).thenReturn("hello world"); + when(resultSet.getObject("metadata", PGobject.class)).thenReturn(pgMetadata); + when(resultSet.getFloat("distance")).thenReturn(0.42f); + + Document document = rowMapperCaptor.getValue().mapRow(resultSet, 0); + + assertThat(document.getMetadata()).containsEntry(DocumentMetadata.DISTANCE.value(), "12.5 miles from depot"); + assertThat(document.getScore()).isEqualTo(1.0 - 0.42f); + } + + @Test + @SuppressWarnings("unchecked") + void rowMapperAddsDistanceMetadataWhenAbsent() throws Exception { + var jdbcTemplate = mock(JdbcTemplate.class); + var embeddingModel = mock(EmbeddingModel.class); + ArgumentCaptor> rowMapperCaptor = ArgumentCaptor.forClass(RowMapper.class); + when(jdbcTemplate.query(anyString(), rowMapperCaptor.capture(), any(), any(), any(), any())) + .thenReturn(List.of()); + + var store = PgVectorStore.builder(jdbcTemplate, embeddingModel).build(); + store.doSimilaritySearch(SearchRequest.builder().query("hello").topK(5).similarityThresholdAll().build()); + + var pgMetadata = new PGobject(); + pgMetadata.setType("json"); + pgMetadata.setValue("{\"author\": \"jane\"}"); + + var resultSet = mock(ResultSet.class); + when(resultSet.getString("id")).thenReturn("doc-1"); + when(resultSet.getString("content")).thenReturn("hello world"); + when(resultSet.getObject("metadata", PGobject.class)).thenReturn(pgMetadata); + when(resultSet.getFloat("distance")).thenReturn(0.42f); + + Document document = rowMapperCaptor.getValue().mapRow(resultSet, 0); + + assertThat(document.getMetadata()).containsEntry(DocumentMetadata.DISTANCE.value(), 0.42f); + } + }