From f2b33a77b078a6ea704aabb7421e01340863bbd0 Mon Sep 17 00:00:00 2001 From: sinsy <550569627@qq.com> Date: Tue, 23 Jul 2024 18:45:46 +0800 Subject: [PATCH 1/4] feat: add pg hybridSearch --- .../ai/vectorstore/VectorStore.java | 4 +++ .../ai/vectorstore/PgVectorStore.java | 29 ++++++++++++++++--- 2 files changed, 29 insertions(+), 4 deletions(-) diff --git a/spring-ai-core/src/main/java/org/springframework/ai/vectorstore/VectorStore.java b/spring-ai-core/src/main/java/org/springframework/ai/vectorstore/VectorStore.java index 8c4ff0c2a9c..c0f47fdf9e2 100644 --- a/spring-ai-core/src/main/java/org/springframework/ai/vectorstore/VectorStore.java +++ b/spring-ai-core/src/main/java/org/springframework/ai/vectorstore/VectorStore.java @@ -70,4 +70,8 @@ default List similaritySearch(String query) { return this.similaritySearch(SearchRequest.query(query)); } + default List hybridSearch(String query) { + return this.similaritySearch(SearchRequest.query(query)); + } + } diff --git a/vector-stores/spring-ai-pgvector-store/src/main/java/org/springframework/ai/vectorstore/PgVectorStore.java b/vector-stores/spring-ai-pgvector-store/src/main/java/org/springframework/ai/vectorstore/PgVectorStore.java index bf2c662f716..9f026afc892 100644 --- a/vector-stores/spring-ai-pgvector-store/src/main/java/org/springframework/ai/vectorstore/PgVectorStore.java +++ b/vector-stores/spring-ai-pgvector-store/src/main/java/org/springframework/ai/vectorstore/PgVectorStore.java @@ -239,6 +239,20 @@ public List similaritySearch(SearchRequest request) { new DocumentRowMapper(this.objectMapper), queryEmbedding, queryEmbedding, distance, request.getTopK()); } + @Override + public List hybridSearch(String query) { + SearchRequest searchRequest = SearchRequest.query(query); + + List similaritySearch = similaritySearch(searchRequest); + + String sql = "SELECT *, ts_rank_cd(to_tsvector(content), plainto_tsquery(?)) AS rank FROM " + getFullyQualifiedTableName() + " ORDER BY rank DESC LIMIT ?"; + + List keyLikeSearch = this.jdbcTemplate.query(sql, new DocumentRowMapper(this.objectMapper, true), query, searchRequest.getTopK()); + + similaritySearch.addAll(keyLikeSearch); + return similaritySearch; + } + public List embeddingDistance(String query) { return this.jdbcTemplate.query( "SELECT embedding " + this.comparisonOperator() + " ? AS distance FROM " + getFullyQualifiedTableName(), @@ -425,21 +439,28 @@ private static class DocumentRowMapper implements RowMapper { private ObjectMapper objectMapper; + private boolean isHybrid = false; + public DocumentRowMapper(ObjectMapper objectMapper) { this.objectMapper = objectMapper; } + public DocumentRowMapper(ObjectMapper objectMapper, boolean isHybrid) { + this.objectMapper = objectMapper; + this.isHybrid = isHybrid; + } + @Override public Document mapRow(ResultSet rs, int rowNum) throws SQLException { String id = rs.getString(COLUMN_ID); String content = rs.getString(COLUMN_CONTENT); PGobject pgMetadata = rs.getObject(COLUMN_METADATA, PGobject.class); PGobject embedding = rs.getObject(COLUMN_EMBEDDING, PGobject.class); - Float distance = rs.getFloat(COLUMN_DISTANCE); - Map metadata = toMap(pgMetadata); - metadata.put(COLUMN_DISTANCE, distance); - + if (!isHybrid) { + Float distance = rs.getFloat(COLUMN_DISTANCE); + metadata.put(COLUMN_DISTANCE, distance); + } Document document = new Document(id, content, metadata); document.setEmbedding(toDoubleList(embedding)); From 5623c9c899e41a8c8ea441abc7d6ead757b81c16 Mon Sep 17 00:00:00 2001 From: sinsy <550569627@qq.com> Date: Wed, 24 Jul 2024 10:01:06 +0800 Subject: [PATCH 2/4] add docs & fix ci --- .../org/springframework/ai/vectorstore/VectorStore.java | 7 +++++++ .../org/springframework/ai/vectorstore/PgVectorStore.java | 6 ++++-- 2 files changed, 11 insertions(+), 2 deletions(-) diff --git a/spring-ai-core/src/main/java/org/springframework/ai/vectorstore/VectorStore.java b/spring-ai-core/src/main/java/org/springframework/ai/vectorstore/VectorStore.java index c0f47fdf9e2..9a3479e731b 100644 --- a/spring-ai-core/src/main/java/org/springframework/ai/vectorstore/VectorStore.java +++ b/spring-ai-core/src/main/java/org/springframework/ai/vectorstore/VectorStore.java @@ -70,6 +70,13 @@ default List similaritySearch(String query) { return this.similaritySearch(SearchRequest.query(query)); } + /** + * Retrieves documents by query embedding similarity using the default, + * If the subclass implements this method, use hybrid search, similar search + full text search. + * {@link SearchRequest}'s' search criteria. + * @param query Text to use for embedding similarity comparison. + * @return Returns a list of documents that have embeddings similar to the query text + */ default List hybridSearch(String query) { return this.similaritySearch(SearchRequest.query(query)); } diff --git a/vector-stores/spring-ai-pgvector-store/src/main/java/org/springframework/ai/vectorstore/PgVectorStore.java b/vector-stores/spring-ai-pgvector-store/src/main/java/org/springframework/ai/vectorstore/PgVectorStore.java index 9f026afc892..769153dfbb0 100644 --- a/vector-stores/spring-ai-pgvector-store/src/main/java/org/springframework/ai/vectorstore/PgVectorStore.java +++ b/vector-stores/spring-ai-pgvector-store/src/main/java/org/springframework/ai/vectorstore/PgVectorStore.java @@ -245,9 +245,11 @@ public List hybridSearch(String query) { List similaritySearch = similaritySearch(searchRequest); - String sql = "SELECT *, ts_rank_cd(to_tsvector(content), plainto_tsquery(?)) AS rank FROM " + getFullyQualifiedTableName() + " ORDER BY rank DESC LIMIT ?"; + String sql = "SELECT *, ts_rank_cd(to_tsvector(content), plainto_tsquery(?)) AS rank FROM " + + getFullyQualifiedTableName() + " ORDER BY rank DESC LIMIT ?"; - List keyLikeSearch = this.jdbcTemplate.query(sql, new DocumentRowMapper(this.objectMapper, true), query, searchRequest.getTopK()); + List keyLikeSearch = this.jdbcTemplate.query(sql, new DocumentRowMapper(this.objectMapper, true), + query, searchRequest.getTopK()); similaritySearch.addAll(keyLikeSearch); return similaritySearch; From 1652b1f4c07d848113123f04cbabc6ac363f8b0f Mon Sep 17 00:00:00 2001 From: sinsy <550569627@qq.com> Date: Wed, 24 Jul 2024 10:03:25 +0800 Subject: [PATCH 3/4] fix ci --- .../org/springframework/ai/vectorstore/VectorStore.java | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/spring-ai-core/src/main/java/org/springframework/ai/vectorstore/VectorStore.java b/spring-ai-core/src/main/java/org/springframework/ai/vectorstore/VectorStore.java index 9a3479e731b..1846dce68be 100644 --- a/spring-ai-core/src/main/java/org/springframework/ai/vectorstore/VectorStore.java +++ b/spring-ai-core/src/main/java/org/springframework/ai/vectorstore/VectorStore.java @@ -71,9 +71,9 @@ default List similaritySearch(String query) { } /** - * Retrieves documents by query embedding similarity using the default, - * If the subclass implements this method, use hybrid search, similar search + full text search. - * {@link SearchRequest}'s' search criteria. + * Retrieves documents by query embedding similarity using the default, If the + * subclass implements this method, use hybrid search, similar search + full text + * search. {@link SearchRequest}'s' search criteria. * @param query Text to use for embedding similarity comparison. * @return Returns a list of documents that have embeddings similar to the query text */ From da3697ca2aeed1d6d43992ba37c4f46b4dcd55a6 Mon Sep 17 00:00:00 2001 From: sinsy <550569627@qq.com> Date: Wed, 31 Jul 2024 10:18:05 +0800 Subject: [PATCH 4/4] add ci --- .../ai/vectorstore/VectorStore.java | 7 ++++--- .../PgVectorStoreAutoConfigurationIT.java | 19 ++++++++++++------- .../ai/vectorstore/PgVectorStore.java | 5 ++--- 3 files changed, 18 insertions(+), 13 deletions(-) diff --git a/spring-ai-core/src/main/java/org/springframework/ai/vectorstore/VectorStore.java b/spring-ai-core/src/main/java/org/springframework/ai/vectorstore/VectorStore.java index 1846dce68be..fb4a7796fec 100644 --- a/spring-ai-core/src/main/java/org/springframework/ai/vectorstore/VectorStore.java +++ b/spring-ai-core/src/main/java/org/springframework/ai/vectorstore/VectorStore.java @@ -74,11 +74,12 @@ default List similaritySearch(String query) { * Retrieves documents by query embedding similarity using the default, If the * subclass implements this method, use hybrid search, similar search + full text * search. {@link SearchRequest}'s' search criteria. - * @param query Text to use for embedding similarity comparison. + * @param request request for set search parameters, such as the query text, topK, + * similarity threshold and metadata filter expressions. * @return Returns a list of documents that have embeddings similar to the query text */ - default List hybridSearch(String query) { - return this.similaritySearch(SearchRequest.query(query)); + default List hybridSearch(SearchRequest request) { + return this.similaritySearch(request); } } diff --git a/spring-ai-spring-boot-autoconfigure/src/test/java/org/springframework/ai/autoconfigure/vectorstore/pgvector/PgVectorStoreAutoConfigurationIT.java b/spring-ai-spring-boot-autoconfigure/src/test/java/org/springframework/ai/autoconfigure/vectorstore/pgvector/PgVectorStoreAutoConfigurationIT.java index df5be884cf6..d0dfdea0e0a 100644 --- a/spring-ai-spring-boot-autoconfigure/src/test/java/org/springframework/ai/autoconfigure/vectorstore/pgvector/PgVectorStoreAutoConfigurationIT.java +++ b/spring-ai-spring-boot-autoconfigure/src/test/java/org/springframework/ai/autoconfigure/vectorstore/pgvector/PgVectorStoreAutoConfigurationIT.java @@ -95,18 +95,23 @@ public void addAndSearch() { vectorStore.add(documents); - List results = vectorStore + List similarityResults = vectorStore .similaritySearch(SearchRequest.query("What is Great Depression?").withTopK(1)); + assertThat(similarityResults).hasSize(1); + Document similarityResultDoc = similarityResults.get(0); + assertThat(similarityResultDoc.getId()).isEqualTo(documents.get(2).getId()); + assertThat(similarityResultDoc.getMetadata()).containsKeys("depression", "distance"); - assertThat(results).hasSize(1); - Document resultDoc = results.get(0); - assertThat(resultDoc.getId()).isEqualTo(documents.get(2).getId()); - assertThat(resultDoc.getMetadata()).containsKeys("depression", "distance"); + List hybridResults = vectorStore + .hybridSearch(SearchRequest.query("What is Great Depression?").withTopK(1)); + assertThat(hybridResults).hasSize(2); + Document hybridResultDoc = hybridResults.get(1); + assertThat(hybridResultDoc.getMetadata()).containsKeys("depression"); // Remove all documents from the store vectorStore.delete(documents.stream().map(doc -> doc.getId()).toList()); - results = vectorStore.similaritySearch(SearchRequest.query("Great Depression").withTopK(1)); - assertThat(results).hasSize(0); + similarityResults = vectorStore.similaritySearch(SearchRequest.query("Great Depression").withTopK(1)); + assertThat(similarityResults).hasSize(0); }); } diff --git a/vector-stores/spring-ai-pgvector-store/src/main/java/org/springframework/ai/vectorstore/PgVectorStore.java b/vector-stores/spring-ai-pgvector-store/src/main/java/org/springframework/ai/vectorstore/PgVectorStore.java index 769153dfbb0..b4e5b93f9f9 100644 --- a/vector-stores/spring-ai-pgvector-store/src/main/java/org/springframework/ai/vectorstore/PgVectorStore.java +++ b/vector-stores/spring-ai-pgvector-store/src/main/java/org/springframework/ai/vectorstore/PgVectorStore.java @@ -240,8 +240,7 @@ public List similaritySearch(SearchRequest request) { } @Override - public List hybridSearch(String query) { - SearchRequest searchRequest = SearchRequest.query(query); + public List hybridSearch(SearchRequest searchRequest) { List similaritySearch = similaritySearch(searchRequest); @@ -249,7 +248,7 @@ public List hybridSearch(String query) { + getFullyQualifiedTableName() + " ORDER BY rank DESC LIMIT ?"; List keyLikeSearch = this.jdbcTemplate.query(sql, new DocumentRowMapper(this.objectMapper, true), - query, searchRequest.getTopK()); + searchRequest.getQuery(), searchRequest.getTopK()); similaritySearch.addAll(keyLikeSearch); return similaritySearch;