diff --git a/spring-ai-vector-store/src/main/java/org/springframework/ai/vectorstore/RerankingAdvisor.java b/spring-ai-vector-store/src/main/java/org/springframework/ai/vectorstore/RerankingAdvisor.java new file mode 100644 index 00000000000..847d5b11740 --- /dev/null +++ b/spring-ai-vector-store/src/main/java/org/springframework/ai/vectorstore/RerankingAdvisor.java @@ -0,0 +1,22 @@ +package org.springframework.ai.vectorstore; + +import java.util.List; + +import org.springframework.ai.document.Document; + +/** + * Defines a pluggable advisor for reranking search results in + * {@link SearchMode#HYBRID_RERANKED} mode. Implementations refine the order and selection + * of documents based on the search request. + */ +public interface RerankingAdvisor { + + /** + * Reranks the provided search results according to the search request. + * @param results The initial list of documents to rerank. + * @param request The search request containing query and mode information. + * @return A reranked list of documents. + */ + List rerank(List results, SearchRequest request); + +} diff --git a/spring-ai-vector-store/src/main/java/org/springframework/ai/vectorstore/SearchMode.java b/spring-ai-vector-store/src/main/java/org/springframework/ai/vectorstore/SearchMode.java new file mode 100644 index 00000000000..a581f61db64 --- /dev/null +++ b/spring-ai-vector-store/src/main/java/org/springframework/ai/vectorstore/SearchMode.java @@ -0,0 +1,26 @@ +package org.springframework.ai.vectorstore; + +/** + * Enum defining the search modes supported by the {@link VectorStore} interface. Each + * mode specifies a different type of search operation for querying documents. + */ +public enum SearchMode { + + /** + * Vector-based similarity search using embeddings (e.g., cosine similarity). + */ + VECTOR, + /** + * Keyword-based full-text search (e.g., TF-IDF or BM25). + */ + FULL_TEXT, + /** + * Hybrid search combining vector and full-text search (e.g., using rank fusion). + */ + HYBRID, + /** + * Hybrid search with additional reranking for enhanced relevance. + */ + HYBRID_RERANKED + +} diff --git a/spring-ai-vector-store/src/main/java/org/springframework/ai/vectorstore/SearchRequest.java b/spring-ai-vector-store/src/main/java/org/springframework/ai/vectorstore/SearchRequest.java index 4eae298a86d..a99440ffaa0 100644 --- a/spring-ai-vector-store/src/main/java/org/springframework/ai/vectorstore/SearchRequest.java +++ b/spring-ai-vector-store/src/main/java/org/springframework/ai/vectorstore/SearchRequest.java @@ -59,6 +59,15 @@ public class SearchRequest { @Nullable private Filter.Expression filterExpression; + private SearchMode searchMode = SearchMode.VECTOR; + + @Nullable + private RerankingAdvisor rerankingAdvisor; + + private double scoreThreshold = SIMILARITY_THRESHOLD_ACCEPT_ALL; + + private int resultLimit = topK; + /** * Copy an existing {@link SearchRequest.Builder} instance. * @param originalSearchRequest {@link SearchRequest} instance to copy. @@ -68,7 +77,11 @@ public static Builder from(SearchRequest originalSearchRequest) { return builder().query(originalSearchRequest.getQuery()) .topK(originalSearchRequest.getTopK()) .similarityThreshold(originalSearchRequest.getSimilarityThreshold()) - .filterExpression(originalSearchRequest.getFilterExpression()); + .filterExpression(originalSearchRequest.getFilterExpression()) + .searchMode(originalSearchRequest.getSearchMode()) + .rerankingAdvisor(originalSearchRequest.getRerankingAdvisor()) + .scoreThreshold(originalSearchRequest.getScoreThreshold()) + .resultLimit(originalSearchRequest.getResultLimit()); } public SearchRequest() { @@ -79,6 +92,10 @@ protected SearchRequest(SearchRequest original) { this.topK = original.topK; this.similarityThreshold = original.similarityThreshold; this.filterExpression = original.filterExpression; + this.searchMode = original.searchMode; + this.rerankingAdvisor = original.rerankingAdvisor; + this.scoreThreshold = original.scoreThreshold; + this.resultLimit = original.resultLimit; } public String getQuery() { @@ -102,10 +119,45 @@ public boolean hasFilterExpression() { return this.filterExpression != null; } + /** + * Returns the search mode. + * @return The search mode. + */ + public SearchMode getSearchMode() { + return this.searchMode; + } + + /** + * Returns the reranking advisor. + * @return The reranking advisor, or null if none set. + */ + @Nullable + public RerankingAdvisor getRerankingAdvisor() { + return this.rerankingAdvisor; + } + + /** + * Returns the score threshold for filtering results. + * @return The score threshold. + */ + public double getScoreThreshold() { + return this.scoreThreshold; + } + + /** + * Returns the maximum number of results to return. + * @return The result limit. + */ + public int getResultLimit() { + return this.resultLimit; + } + @Override public String toString() { return "SearchRequest{" + "query='" + this.query + '\'' + ", topK=" + this.topK + ", similarityThreshold=" - + this.similarityThreshold + ", filterExpression=" + this.filterExpression + '}'; + + this.similarityThreshold + ", filterExpression=" + this.filterExpression + ", searchMode=" + + this.searchMode + ", rerankingAdvisor=" + this.rerankingAdvisor + ", scoreThreshold=" + + this.scoreThreshold + ", resultLimit=" + this.resultLimit + '}'; } @Override @@ -118,13 +170,16 @@ public boolean equals(Object o) { } SearchRequest that = (SearchRequest) o; return this.topK == that.topK && Double.compare(that.similarityThreshold, this.similarityThreshold) == 0 + && Double.compare(that.scoreThreshold, this.scoreThreshold) == 0 && this.resultLimit == that.resultLimit && Objects.equals(this.query, that.query) - && Objects.equals(this.filterExpression, that.filterExpression); + && Objects.equals(this.filterExpression, that.filterExpression) && this.searchMode == that.searchMode + && Objects.equals(this.rerankingAdvisor, that.rerankingAdvisor); } @Override public int hashCode() { - return Objects.hash(this.query, this.topK, this.similarityThreshold, this.filterExpression); + return Objects.hash(this.query, this.topK, this.similarityThreshold, this.filterExpression, this.searchMode, + this.rerankingAdvisor, this.scoreThreshold, this.resultLimit); } /** @@ -287,6 +342,50 @@ public Builder filterExpression(@Nullable String textExpression) { return this; } + /** + * Sets the search mode (e.g., vector, full-text, hybrid, or reranked hybrid). + * @param searchMode The search mode. + * @return This builder for chaining. + */ + public Builder searchMode(SearchMode searchMode) { + this.searchRequest.searchMode = searchMode != null ? searchMode : SearchMode.VECTOR; + return this; + } + + /** + * Sets the reranking advisor. + * @param rerankingAdvisor The reranking advisor, or null for no reranking. + * @return This builder. + */ + public Builder rerankingAdvisor(@Nullable RerankingAdvisor rerankingAdvisor) { + this.searchRequest.rerankingAdvisor = rerankingAdvisor; + return this; + } + + /** + * Sets the score threshold for filtering results. + * @param scoreThreshold The lower bound of the score. + * @return This builder. + * @throws IllegalArgumentException if threshold is negative. + */ + public Builder scoreThreshold(double scoreThreshold) { + Assert.isTrue(scoreThreshold >= 0, "Score threshold must be non-negative."); + this.searchRequest.scoreThreshold = scoreThreshold; + return this; + } + + /** + * Sets the maximum number of results to return. + * @param resultLimit The maximum number of results. + * @return This builder. + * @throws IllegalArgumentException if resultLimit is negative. + */ + public Builder resultLimit(int resultLimit) { + Assert.isTrue(resultLimit >= 0, "Result limit must be positive."); + this.searchRequest.resultLimit = resultLimit; + return this; + } + public SearchRequest build() { return this.searchRequest; } diff --git a/spring-ai-vector-store/src/main/java/org/springframework/ai/vectorstore/VectorStore.java b/spring-ai-vector-store/src/main/java/org/springframework/ai/vectorstore/VectorStore.java index df1a11f614d..a4f1299b842 100644 --- a/spring-ai-vector-store/src/main/java/org/springframework/ai/vectorstore/VectorStore.java +++ b/spring-ai-vector-store/src/main/java/org/springframework/ai/vectorstore/VectorStore.java @@ -84,26 +84,52 @@ default void delete(String filterExpression) { this.delete(textExpression); } + /** + * Retrieves documents based on the specified search mode, query, and criteria. + * Supports vector similarity, full-text, hybrid, and reranked hybrid searches as + * defined by {@link SearchRequest#getSearchMode()}. + * @param request Search request specifying query text, search mode, topK, similarity + * threshold, and optional metadata filter expressions. + * @return List of documents matching the request criteria, or empty list if no + * matches. + */ + List search(SearchRequest request); + /** * Retrieves documents by query embedding similarity and metadata filters to retrieve * exactly the number of nearest-neighbor results that match the request criteria. + * Delegates to {@link #search(SearchRequest)} with {@link SearchMode#VECTOR}. * @param request Search request for set search parameters, such as the query text, * topK, similarity threshold and metadata filter expressions. - * @return Returns documents th match the query request conditions. + * @return Returns documents that match the query request conditions, or null if no + * matches. + * @deprecated Use {@link #search(SearchRequest)} with {@link SearchMode#VECTOR} + * instead. */ - @Nullable - List similaritySearch(SearchRequest request); + @Deprecated + default List similaritySearch(SearchRequest request) { + return search(SearchRequest.builder() + .query(request.getQuery()) + .topK(request.getTopK()) + .similarityThreshold(request.getSimilarityThreshold()) + .filterExpression(request.getFilterExpression()) + .searchMode(SearchMode.VECTOR) + .build()); + } /** * Retrieves documents by query embedding similarity using the default - * {@link SearchRequest}'s' search criteria. + * {@link SearchRequest}'s search criteria. Delegates to + * {@link #search(SearchRequest)} with {@link SearchMode#VECTOR}. * @param query Text to use for embedding similarity comparison. * @return Returns a list of documents that have embeddings similar to the query text - * embedding. + * embedding, or null if no matches. + * @deprecated Use {@link #search(SearchRequest)} with {@link SearchMode#VECTOR} + * instead. */ - @Nullable + @Deprecated default List similaritySearch(String query) { - return this.similaritySearch(SearchRequest.builder().query(query).build()); + return search(SearchRequest.builder().query(query).searchMode(SearchMode.VECTOR).build()); } /** diff --git a/spring-ai-vector-store/src/main/java/org/springframework/ai/vectorstore/observation/AbstractObservationVectorStore.java b/spring-ai-vector-store/src/main/java/org/springframework/ai/vectorstore/observation/AbstractObservationVectorStore.java index ca3c3ae9185..81eb1940722 100644 --- a/spring-ai-vector-store/src/main/java/org/springframework/ai/vectorstore/observation/AbstractObservationVectorStore.java +++ b/spring-ai-vector-store/src/main/java/org/springframework/ai/vectorstore/observation/AbstractObservationVectorStore.java @@ -24,6 +24,8 @@ import org.springframework.ai.embedding.BatchingStrategy; import org.springframework.ai.embedding.EmbeddingModel; import org.springframework.ai.vectorstore.AbstractVectorStoreBuilder; +import org.springframework.ai.vectorstore.RerankingAdvisor; +import org.springframework.ai.vectorstore.SearchMode; import org.springframework.ai.vectorstore.SearchRequest; import org.springframework.ai.vectorstore.VectorStore; import org.springframework.ai.vectorstore.filter.Filter; @@ -111,9 +113,9 @@ public void delete(Filter.Expression filterExpression) { } @Override + @Deprecated @Nullable public List similaritySearch(SearchRequest request) { - VectorStoreObservationContext searchObservationContext = this .createObservationContextBuilder(VectorStoreObservationContext.Operation.QUERY.value()) .queryRequest(request) @@ -123,12 +125,50 @@ public List similaritySearch(SearchRequest request) { .observation(this.customObservationConvention, DEFAULT_OBSERVATION_CONVENTION, () -> searchObservationContext, this.observationRegistry) .observe(() -> { - var documents = this.doSimilaritySearch(request); + SearchRequest vectorRequest = SearchRequest.builder() + .query(request.getQuery()) + .topK(request.getTopK()) + .similarityThreshold(request.getSimilarityThreshold()) + .filterExpression(request.getFilterExpression()) + .searchMode(SearchMode.VECTOR) + .build(); + var documents = search(vectorRequest); searchObservationContext.setQueryResponse(documents); return documents; }); } + /** + * Retrieves documents based on the specified search mode, query, and criteria. + * Default implementation supports {@link SearchMode#VECTOR} by delegating to + * {@link #doSimilaritySearch(SearchRequest)}. Other modes throw an exception. + * @param request Search request specifying query text, search mode, topK, similarity + * threshold, and optional metadata filter expressions. + * @return List of documents matching the request criteria, or empty list if no + * matches. + * @throws UnsupportedOperationException if the search mode is not supported. + */ + @Override + public List search(SearchRequest request) { + VectorStoreObservationContext searchObservationContext = this + .createObservationContextBuilder(VectorStoreObservationContext.Operation.QUERY.value()) + .queryRequest(request) + .build(); + + return VectorStoreObservationDocumentation.AI_VECTOR_STORE + .observation(this.customObservationConvention, DEFAULT_OBSERVATION_CONVENTION, + () -> searchObservationContext, this.observationRegistry) + .observe(() -> { + if (request.getSearchMode() != SearchMode.VECTOR) { + throw new UnsupportedOperationException( + "Search mode " + request.getSearchMode() + " not supported"); + } + var documents = doSimilaritySearch(request); + searchObservationContext.setQueryResponse(documents); + return documents != null ? documents : List.of(); + }); + } + /** * Perform the actual add operation. * @param documents the documents to add diff --git a/vector-stores/spring-ai-azure-store/src/main/java/org/springframework/ai/vectorstore/azure/AzureSemanticRerankingAdvisor.java b/vector-stores/spring-ai-azure-store/src/main/java/org/springframework/ai/vectorstore/azure/AzureSemanticRerankingAdvisor.java new file mode 100644 index 00000000000..b3b8138d93a --- /dev/null +++ b/vector-stores/spring-ai-azure-store/src/main/java/org/springframework/ai/vectorstore/azure/AzureSemanticRerankingAdvisor.java @@ -0,0 +1,33 @@ +package org.springframework.ai.vectorstore.azure; + +import java.util.Comparator; +import java.util.List; +import java.util.stream.Collectors; + +import org.springframework.ai.document.Document; +import org.springframework.ai.vectorstore.RerankingAdvisor; +import org.springframework.ai.vectorstore.SearchRequest; + +/** + * Reranking advisor for Azure AI Search's + * {@link org.springframework.ai.vectorstore.SearchMode#HYBRID_RERANKED} mode, filtering + * and sorting results based on reranker scores. + */ +public class AzureSemanticRerankingAdvisor implements RerankingAdvisor { + + /** + * Reranks search results by filtering documents based on the similarity threshold and + * sorting by score in descending order. + * @param results The initial list of documents. + * @param request The search request. + * @return The reranked list of documents. + */ + @Override + public List rerank(List results, SearchRequest request) { + return results.stream() + .filter(doc -> doc.getScore() >= request.getScoreThreshold()) + .sorted(Comparator.comparingDouble(Document::getScore).reversed()) + .collect(Collectors.toList()); + } + +} diff --git a/vector-stores/spring-ai-azure-store/src/main/java/org/springframework/ai/vectorstore/azure/AzureVectorStore.java b/vector-stores/spring-ai-azure-store/src/main/java/org/springframework/ai/vectorstore/azure/AzureVectorStore.java index 0f86bd10c9f..df229dd4bdf 100644 --- a/vector-stores/spring-ai-azure-store/src/main/java/org/springframework/ai/vectorstore/azure/AzureVectorStore.java +++ b/vector-stores/spring-ai-azure-store/src/main/java/org/springframework/ai/vectorstore/azure/AzureVectorStore.java @@ -18,6 +18,7 @@ import java.util.ArrayList; import java.util.Collections; +import java.util.HashMap; import java.util.List; import java.util.Map; import java.util.Optional; @@ -39,9 +40,12 @@ import com.azure.search.documents.indexes.models.VectorSearchProfile; import com.azure.search.documents.models.IndexDocumentsResult; import com.azure.search.documents.models.IndexingResult; +import com.azure.search.documents.models.QueryType; import com.azure.search.documents.models.SearchOptions; +import com.azure.search.documents.models.SemanticSearchOptions; import com.azure.search.documents.models.VectorSearchOptions; import com.azure.search.documents.models.VectorizedQuery; +import com.azure.search.documents.util.SearchPagedIterable; import org.slf4j.Logger; import org.slf4j.LoggerFactory; @@ -53,10 +57,13 @@ import org.springframework.ai.observation.conventions.VectorStoreProvider; import org.springframework.ai.observation.conventions.VectorStoreSimilarityMetric; import org.springframework.ai.vectorstore.AbstractVectorStoreBuilder; +import org.springframework.ai.vectorstore.RerankingAdvisor; +import org.springframework.ai.vectorstore.SearchMode; import org.springframework.ai.vectorstore.SearchRequest; import org.springframework.ai.vectorstore.filter.FilterExpressionConverter; import org.springframework.ai.vectorstore.observation.AbstractObservationVectorStore; import org.springframework.ai.vectorstore.observation.VectorStoreObservationContext; +import org.springframework.ai.vectorstore.observation.VectorStoreObservationDocumentation; import org.springframework.beans.factory.InitializingBean; import org.springframework.lang.Nullable; import org.springframework.util.Assert; @@ -126,6 +133,8 @@ public class AzureVectorStore extends AbstractObservationVectorStore implements private String indexName; + private RerankingAdvisor rerankingAdvisor; + /** * Protected constructor that accepts a builder instance. This is the preferred way to * create new AzureVectorStore instances. @@ -201,59 +210,16 @@ public void doDelete(List documentIds) { this.searchClient.deleteDocuments(searchDocumentIds); } - @Override - public List similaritySearch(String query) { - return this.similaritySearch(SearchRequest.builder() - .query(query) - .topK(this.defaultTopK) - .similarityThreshold(this.defaultSimilarityThreshold) - .build()); - } - @Override public List doSimilaritySearch(SearchRequest request) { - Assert.notNull(request, "The search request must not be null."); - - var searchEmbedding = this.embeddingModel.embed(request.getQuery()); - - final var vectorQuery = new VectorizedQuery(EmbeddingUtils.toList(searchEmbedding)) - .setKNearestNeighborsCount(request.getTopK()) - // Set the fields to compare the vector against. This is a comma-delimited - // list of field names. - .setFields(EMBEDDING_FIELD_NAME); - - var searchOptions = new SearchOptions() - .setVectorSearchOptions(new VectorSearchOptions().setQueries(vectorQuery)); - - if (request.hasFilterExpression()) { - String oDataFilter = this.filterExpressionConverter.convertExpression(request.getFilterExpression()); - searchOptions.setFilter(oDataFilter); - } - - final var searchResults = this.searchClient.search(null, searchOptions, Context.NONE); - - return searchResults.stream() - .filter(result -> result.getScore() >= request.getSimilarityThreshold()) - .map(result -> { - - final AzureSearchDocument entry = result.getDocument(AzureSearchDocument.class); - - Map metadata = (StringUtils.hasText(entry.metadata())) - ? JSONObject.parseObject(entry.metadata(), new TypeReference>() { - - }) : Map.of(); - - metadata.put(DocumentMetadata.DISTANCE.value(), 1.0 - result.getScore()); - - return Document.builder() - .id(entry.id()) - .text(entry.content) - .metadata(metadata) - .score(result.getScore()) - .build(); - }) - .collect(Collectors.toList()); + return search(SearchRequest.builder() + .query(request.getQuery()) + .topK(request.getTopK()) + .similarityThreshold(request.getSimilarityThreshold()) + .filterExpression(request.getFilterExpression()) + .searchMode(SearchMode.VECTOR) + .build()); } @Override @@ -309,6 +275,83 @@ public void afterPropertiesSet() throws Exception { this.searchClient = this.searchIndexClient.getSearchClient(this.indexName); } + @Override + public List search(SearchRequest request) { + Assert.notNull(request, "The search request must not be null."); + + SearchOptions options = new SearchOptions(); + int limit = request.getResultLimit(); + + switch (request.getSearchMode()) { + case VECTOR: + var vectorQuery = new VectorizedQuery(EmbeddingUtils.toList(embeddingModel.embed(request.getQuery()))) + .setKNearestNeighborsCount(request.getTopK()) + .setFields(EMBEDDING_FIELD_NAME); + options.setVectorSearchOptions(new VectorSearchOptions().setQueries(vectorQuery)); + options.setTop(request.getTopK()); + break; + case FULL_TEXT: + options.setQueryType(QueryType.FULL); + options.setTop(request.getTopK()); + break; + case HYBRID: + var hybridVectorQuery = new VectorizedQuery( + EmbeddingUtils.toList(embeddingModel.embed(request.getQuery()))) + .setKNearestNeighborsCount(request.getTopK()) + .setFields(EMBEDDING_FIELD_NAME); + options.setVectorSearchOptions(new VectorSearchOptions().setQueries(hybridVectorQuery)); + options.setQueryType(QueryType.FULL); + options.setTop(request.getTopK()); + break; + case HYBRID_RERANKED: + var rerankedVectorQuery = new VectorizedQuery( + EmbeddingUtils.toList(embeddingModel.embed(request.getQuery()))) + .setKNearestNeighborsCount(request.getTopK()) + .setFields(EMBEDDING_FIELD_NAME); + options.setVectorSearchOptions(new VectorSearchOptions().setQueries(rerankedVectorQuery)); + options.setQueryType(QueryType.SEMANTIC); + SemanticSearchOptions semanticSearchOptions = new SemanticSearchOptions(); + semanticSearchOptions.setSemanticConfigurationName("semanticConfiguration"); // TODO: + // make + // configurable + options.setSemanticSearchOptions(semanticSearchOptions); + options.setTop(limit); + break; + default: + throw new IllegalArgumentException("Unsupported search mode: " + request.getSearchMode()); + } + + if (request.hasFilterExpression()) { + options.setFilter(filterExpressionConverter.convertExpression(request.getFilterExpression())); + } + + SearchPagedIterable results = searchClient.search(request.getQuery(), options, + com.azure.core.util.Context.NONE); + List documents = results.stream().filter(result -> { + double score = request.getSearchMode() == SearchMode.HYBRID_RERANKED && result.getSemanticSearch() != null + ? result.getSemanticSearch().getRerankerScore() : result.getScore(); + return score >= request.getScoreThreshold(); + }).map(result -> { + final AzureSearchDocument entry = result.getDocument(AzureSearchDocument.class); + Map metadata = StringUtils.hasText(entry.metadata()) + ? JSONObject.parseObject(entry.metadata(), new TypeReference>() { + }) : new HashMap<>(); + metadata.put(DocumentMetadata.DISTANCE.value(), 1.0 - result.getScore()); + double score = request.getSearchMode() == SearchMode.HYBRID_RERANKED && result.getSemanticSearch() != null + ? result.getSemanticSearch().getRerankerScore() : result.getScore(); + if (request.getSearchMode() == SearchMode.HYBRID_RERANKED && result.getSemanticSearch() != null) { + metadata.put("re-rank_score", result.getSemanticSearch().getRerankerScore()); + } + return Document.builder().id(entry.id()).text(entry.content()).metadata(metadata).score(score).build(); + }).collect(Collectors.toList()); + + if (request.getRerankingAdvisor() != null && request.getSearchMode() == SearchMode.HYBRID_RERANKED) { + documents = request.getRerankingAdvisor().rerank(documents, request); + } + + return documents; + } + @Override public VectorStoreObservationContext.Builder createObservationContextBuilder(String operationName) { diff --git a/vector-stores/spring-ai-azure-store/src/test/java/org/springframework/ai/vectorstore/azure/AzureVectorStoreIT.java b/vector-stores/spring-ai-azure-store/src/test/java/org/springframework/ai/vectorstore/azure/AzureVectorStoreIT.java index 1e4bd4aea04..48b6b727aee 100644 --- a/vector-stores/spring-ai-azure-store/src/test/java/org/springframework/ai/vectorstore/azure/AzureVectorStoreIT.java +++ b/vector-stores/spring-ai-azure-store/src/test/java/org/springframework/ai/vectorstore/azure/AzureVectorStoreIT.java @@ -40,6 +40,7 @@ import org.springframework.ai.document.DocumentMetadata; import org.springframework.ai.embedding.EmbeddingModel; import org.springframework.ai.transformers.TransformersEmbeddingModel; +import org.springframework.ai.vectorstore.SearchMode; import org.springframework.ai.vectorstore.SearchRequest; import org.springframework.ai.vectorstore.VectorStore; import org.springframework.ai.vectorstore.azure.AzureVectorStore.MetadataField; @@ -119,6 +120,131 @@ public void addAndSearchTest() { }); } + @Test + public void addAndFullTextSearchTest() { + this.contextRunner.run(context -> { + VectorStore vectorStore = context.getBean(VectorStore.class); + + vectorStore.add(this.documents); + + Awaitility.await() + .until(() -> vectorStore.search(SearchRequest.builder() + .query("Great Depression") + .topK(1) + .searchMode(SearchMode.FULL_TEXT) + .build()), hasSize(1)); + + List results = vectorStore.search(SearchRequest.builder() + .query("Great Depression") + .topK(1) + .searchMode(SearchMode.FULL_TEXT) + .scoreThreshold(0.0) + .build()); + + assertThat(results).hasSize(1); + Document resultDoc = results.get(0); + assertThat(resultDoc.getId()).isEqualTo(this.documents.get(2).getId()); + assertThat(resultDoc.getText()).contains("The Great Depression (1929–1939) was an economic shock"); + assertThat(resultDoc.getMetadata()).hasSize(2); + assertThat(resultDoc.getMetadata()).containsKey("meta2"); + assertThat(resultDoc.getMetadata()).containsKey(DocumentMetadata.DISTANCE.value()); + + // Remove all documents from the store + vectorStore.delete(this.documents.stream().map(doc -> doc.getId()).toList()); + + Awaitility.await() + .until(() -> vectorStore + .search(SearchRequest.builder().query("Hello").topK(1).searchMode(SearchMode.FULL_TEXT).build()), + hasSize(0)); + }); + } + + @Test + public void addAndHybridSearchTest() { + this.contextRunner.run(context -> { + VectorStore vectorStore = context.getBean(VectorStore.class); + + vectorStore.add(this.documents); + + Awaitility.await() + .until(() -> vectorStore.search(SearchRequest.builder() + .query("Great Depression") + .topK(1) + .searchMode(SearchMode.HYBRID) + .build()), hasSize(1)); + + List results = vectorStore.search(SearchRequest.builder() + .query("Great Depression") + .topK(1) + .searchMode(SearchMode.HYBRID) + .scoreThreshold(0.0) + .build()); + + assertThat(results).hasSize(1); + Document resultDoc = results.get(0); + assertThat(resultDoc.getId()).isEqualTo(this.documents.get(2).getId()); + assertThat(resultDoc.getText()).contains("The Great Depression (1929–1939) was an economic shock"); + assertThat(resultDoc.getMetadata()).hasSize(2); + assertThat(resultDoc.getMetadata()).containsKey("meta2"); + assertThat(resultDoc.getMetadata()).containsKey(DocumentMetadata.DISTANCE.value()); + + // Remove all documents from the store + vectorStore.delete(this.documents.stream().map(doc -> doc.getId()).toList()); + + Awaitility.await() + .until(() -> vectorStore + .search(SearchRequest.builder().query("Hello").topK(1).searchMode(SearchMode.HYBRID).build()), + hasSize(0)); + }); + } + + @Test + public void addAndHybridRerankedSearchTest() { + this.contextRunner.run(context -> { + VectorStore vectorStore = context.getBean(VectorStore.class); + + vectorStore.add(this.documents); + + Awaitility.await() + .until(() -> vectorStore.search(SearchRequest.builder() + .query("Great Depression") + .topK(1) + .searchMode(SearchMode.HYBRID_RERANKED) + .rerankingAdvisor(new AzureSemanticRerankingAdvisor()) + .build()), hasSize(1)); + + List results = vectorStore.search(SearchRequest.builder() + .query("Great Depression") + .topK(1) + .searchMode(SearchMode.HYBRID_RERANKED) + .rerankingAdvisor(new AzureSemanticRerankingAdvisor()) + .scoreThreshold(2.0) + .resultLimit(10) + .build()); + + assertThat(results).hasSize(1); + Document resultDoc = results.get(0); + assertThat(resultDoc.getId()).isEqualTo(this.documents.get(2).getId()); + assertThat(resultDoc.getText()).contains("The Great Depression (1929–1939) was an economic shock"); + assertThat(resultDoc.getMetadata()).hasSize(3); + assertThat(resultDoc.getMetadata()).containsKey("meta2"); + assertThat(resultDoc.getMetadata()).containsKey(DocumentMetadata.DISTANCE.value()); + assertThat(resultDoc.getMetadata()).containsKey("re-rank_score"); + assertThat(resultDoc.getScore()).isGreaterThanOrEqualTo(2.0); + + // Remove all documents from the store + vectorStore.delete(this.documents.stream().map(doc -> doc.getId()).toList()); + + Awaitility.await() + .until(() -> vectorStore.search(SearchRequest.builder() + .query("Hello") + .topK(1) + .searchMode(SearchMode.HYBRID_RERANKED) + .rerankingAdvisor(new AzureSemanticRerankingAdvisor()) + .build()), hasSize(0)); + }); + } + @Test public void searchWithFilters() throws InterruptedException {